Một hệ thống giả định biến các tin nhắn giao hàng thành những sự kiện JSON ngắn. Một mô hình giáo viên 1,5B được tinh chỉnh cục bộ rồi truyền các phân phối token kế tiếp cho một mô hình học viên 0,5B; học viên này được đánh giá bên cạnh một học viên chỉ huấn luyện trên cùng các nhãn.
Với cùng một tin nhắn, hãy hiển thị giáo viên, học viên SFT và học viên được chưng cất, các lỗi trường của chúng và chi phí suy luận. Hãy cho thấy xác suất mềm của vài token trên cùng một tiền tố, với nhiệt độ.
Tải dự án hoàn chỉnh, bản sửa đổi r001
Một nhiệm vụ có giới hạn và các họ dữ liệu được giữ riêng
ParcelEvents-v1 có 400 họ bưu kiện giả định, năm tin nhắn cho mỗi họ, tức 2000 ví dụ. Một hoán vị với seed 50 phân chia các họ thành 280 train, 60 validation và 60 test: 1400/300/300 tin nhắn. Năm tình huống gồm nhận bởi hãng vận chuyển, chuyển giữa các hub, giao xác nhận cho người nhận, đến nơi khi thiếu mã định danh và giờ, cùng sửa một bản nháp chưa xác nhận. Đến một hub không phải là giao cho người nhận.
Câu trả lời có đúng năm trường: parcel, event, location, time và delivered. Các trường parcel/location/time có thể là null; delivered là boolean. Bốn sự kiện được phép là received, transferred, arrived_at_hub và delivered. Dữ liệu dùng chung tám hub và các mẫu câu giữa các phần. Danh tính và họ vẫn tách biệt; tính đều đặn này không đại diện cho sự đa dạng của tài liệu logistics thực tế. Trong 300 trường hợp test, 60 trường hợp chỉ báo đã giao và 240 trường hợp không; mỗi trong ba thông tin nullable thiếu trong 60 trường hợp.
Căn chỉnh các tọa độ trước khi so sánh phân phối
Giáo viên Qwen2.5-1.5B-Instruct và học viên Qwen2.5-0.5B-Instruct đều tạo ra 151936 logits. Cùng kích thước không chứng minh tọa độ 42 mang cùng token. Việc kiểm tra so sánh toàn bộ từ điển token→ID, token được thêm và token đặc biệt, ID của chúng, mẫu hội thoại, cấu hình BPE tuần tự hóa và cấu hình sinh.
Các tokenizer đã nhận phải mô tả chính xác cùng 151665 ID hiệu lực và cùng 271 tọa độ đầu ra không dùng. Các tọa độ này vẫn được đưa vào KL đầy đủ. Chúng không bị bỏ đi để làm hai phân phối khớp nhau một cách giả tạo. Bất đồng về từ vựng, mẫu hoặc token đặc biệt dừng quy trình mà không ánh xạ lại âm thầm.
Một cách xây dựng prompt và câu trả lời duy nhất cung cấp input_ids và labels. Kiểm tra thực sự mã hóa 2000 ví dụ bằng hai tokenizer và yêu cầu cùng chuỗi cùng hash. Prompt tối đa 320 token, câu trả lời có EOS tối đa 96, ngữ cảnh 416. Vượt quá target bị từ chối; không có việc cắt gold âm thầm. EOS là 151645 và PAD là 151643. Padding ở bên phải và mặt nạ của nó áp dụng theo vị trí.
Huấn luyện giáo viên cục bộ
Giáo viên có 28 khối, độ rộng 1536, 12 đầu truy vấn và hai đầu KV. Nền BF16 của nó vẫn đóng băng. Chúng tôi mã hóa LoRA: W(x)+(alpha/r)B(Ax), hạng 8, alpha 16, dropout bằng không, trên bảy phép chiếu q/k/v/o/gate/up/down của từng khối trong 28 khối. 196 phép chiếu thêm 9 232 384 tham số FP32. Phép cộng được tính ở FP32 rồi quay về dtype BF16 của nền. Các kiểm thử so sánh giá trị và gradient với PEFT; không trainer PEFT nào được gọi.
Giáo viên học 1400 câu trả lời train trong hai epoch. Với microbatch 1 và tích lũy tám ví dụ, điều đó tạo 350 cập nhật AdamW ở 1e−4. Validation theo dõi CE sau mỗi epoch; giao thức giữ trạng thái cuối đã cố định trước, không chọn checkpoint theo kết quả test. Sau đó giáo viên chuyển sang eval và bị đóng băng hoàn toàn cho chưng cất. Hash trọng số và việc không có gradient kiểm tra tính bất biến này.
Teacher forcing và loss được mã hóa
Tại token câu trả lời t, hai mô hình nhận cùng prompt và cùng các token gold trước đó. Đây là teacher forcing ngoài chính sách: giáo viên đánh giá tiền tố có giám sát này, không phải câu trả lời tự do của học viên. Không dùng cache logits khổng lồ của corpus hay xấp xỉ top-k.
Với z là logits của giáo viên và u là của học viên, p_T=softmax(z/T), q_T=softmax(u/T). Phân kỳ có hướng là KL(p_T || q_T)=sum_v p_T(v)[log p_T(v)−log q_T(v)]. Đảo các đối số tạo ra một loss khác. Chúng tôi dùng log_softmax ổn định và toàn bộ 151936 tọa độ. Giáo viên được tính dưới no_grad theo microbatch; phân kỳ chỉ truyền gradient đến học viên.
Loss là (1−alpha) CE(q_1, target) + alpha T² KL(p_T || q_T). CE dùng logits của học viên ở nhiệt độ 1. Hệ số T² bù cho thang của tín hiệu chưng cất; nó không bảo đảm dự đoán tốt hơn hay gradient giống hệt ở mọi nhiệt độ. Các kiểm thử so sánh công thức và gradient với torch.nn.functional.kl_div trên các phân phối FP64 bất đối xứng, kiểm tra chiều, giá trị không cho hai phân phối giống nhau, các giới hạn alpha=0/1 và việc loại đồ thị của giáo viên.
Chúng tôi dịch logits và labels để dự đoán token kế tiếp, rồi chỉ chọn token câu trả lời, gồm EOS. Prompt và PAD bị loại. Mỗi lần tích lũy chia cho tổng số token hoạt động của tám ví dụ. Lấy trung bình tám loss đã được trung bình sẽ cho câu trả lời ngắn và dài cùng trọng số. Với alpha=0, đường đi chính xác là CE và không tải giáo viên.
Hai học viên và một lần tiếp tục chính xác
Học viên có 24 khối, độ rộng 896, 14 đầu truy vấn, hai đầu KV và 494 032 768 tham số duy nhất. Tất cả được huấn luyện ở FP32; embeddings và phép chiếu đầu ra dùng chung cùng Parameter. Đối chứng SFT alpha=0 và chưng cất alpha=0,5/T=2 bắt đầu từ cùng trọng số công khai, cùng seed và cùng thứ tự train. Mỗi nhánh thực hiện hai epoch, 350 cập nhật ở 2e−5.
AdamW dùng weight_decay=0,01, clipping gradient ở 1, warm-up 35 update rồi giảm tuyến tính đến 10% tỷ lệ ban đầu. Chế độ eval giữ gradient và tắt các nguồn dropout; nó không đồng nghĩa với no_grad. Đối chứng SFT không gọi giáo viên nào. Chi phí forward của giáo viên trong KD, tính/backward của mô hình được tối ưu, tối ưu hóa và validation được lưu riêng.
Mỗi nhánh so sánh một lần dừng sau một epoch, hai epoch liên tục và phần tiếp tục của epoch còn lại. Trọng số, trạng thái Adam, scheduler, hoán vị, RNG và lịch sử số phải hoàn toàn bằng nhau. Thời lượng nằm ngoài so sánh trạng thái. Các trạng thái cục bộ được nén không mất dữ liệu, gzip mức 1; trọng số cuối cũng được lưu và nạp lại bằng Safetensors. Tổng với bài tập giữ 15 trạng thái tiếp tục và năm trọng số cuối. Các pickle Torch này được tạo cục bộ; đừng thay bằng tệp của bên thứ ba.
Sinh tự do và đo phạm vi thực tế
Giáo viên đã thích nghi, học viên ban đầu, SFT và KD mỗi bên sinh 300 câu trả lời test, không có tiền tố gold. Hai biến thể của bài tập cũng làm vậy: 1800 câu trả lời test, cộng các warm-up được nhận diện riêng. Việc sinh là greedy trên toàn bộ từ vựng, không cache KV, không schema bắt buộc, giới hạn 96 token mới. ID, văn bản thô, EOS hay cắt, log-xác suất, thời lượng và hai probe logits đầy đủ được giữ lại.
Dịch vụ học viên về mặt vật lý không gắn giáo viên hay adapter của nó. Việc nạp lại học viên dùng cấu hình, tokenizer và chỉ các trọng số đã học của chính nó. Suy luận không gắn label nào. Sau đó đánh giá ghép đầu ra và target mà không sửa trọng số.
Bộ phân tích phân biệt JSON có thể phân tích, schema nghiêm ngặt và câu trả lời hoàn chỉnh: một đầu ra bị cắt thất bại dù chứa đối tượng đúng dạng. Khóa trùng, NaN, trường bổ sung, Markdown và ID chưa biết không được sửa. Lỗi vẫn ở mẫu số 300 trường hợp cho exact match và từng trường. Tỷ lệ giao sai hoặc điền trường sai áp dụng cho câu trả lời đầy đủ hợp lệ và cũng chỉ rõ tổng quần thể đủ điều kiện cùng độ phủ. Một đầu ra không hợp lệ không được chứng nhận an toàn.
Mỗi mô hình nhận cùng tin nhắn train để warm-up, ngoài tập test. Chúng tôi báo cáo độ trễ theo yêu cầu và theo token, độ dài có EOS, trung vị và p95 thực nghiệm (hạng 285 trên 300). Hai yêu cầu đầu bao gồm chi phí probe logits, được nêu trong báo cáo. Việc nạp và kiểm tra toàn vẹn vẫn tách biệt. Tham số duy nhất, dtype, byte logic và kích thước thực của tệp được phục vụ đều được đo. Giáo viên BF16 với adapter FP32 và học viên hoàn toàn FP32 cũng khác về độ chính xác: tỷ lệ tham số không phải là tỷ lệ bộ nhớ hoặc tốc độ được bảo đảm.
Bài tập: thay đổi alpha và T
SFT chính là đối chiếu thực sự alpha=0. Hai học viên bổ sung bắt đầu từ cùng nền: alpha=1/T=2 và alpha=1/T=4, mỗi bên có 350 update và phần tiếp tục chính xác của riêng mình. Không checkpoint nào chỉ được đổi tên. Dữ liệu, giáo viên và kết quả chính được bảo vệ bằng hash; giáo viên không được huấn luyện lại cho bài tập.
Alpha=1 chỉ bắt chước giáo viên trên các tiền tố gold, gồm cả lỗi của nó. Tăng T phân phối lại khối lượng và gradient; điều đó không tiết lộ sự thật của trường delivered. Hãy so sánh cùng 300 sự kiện, lỗi và chi phí. Chương trình không áp đặt cải thiện nào hay nhiệt độ tối ưu.
Giới hạn
Một seed, corpus tổng hợp đều đặn và tập nhỏ các loại không tạo thành đánh giá nghiệp vụ. KL giảm trên tiền tố gold không chứng minh sinh tự do tốt hơn. Giáo viên có thể truyền lỗi của mình. Các phép đo cục bộ không có tải đồng thời được kiểm soát không phải cam kết hiệu năng dịch vụ. Hướng dẫn không gửi bưu kiện, tin nhắn hay hành động bên ngoài nào.
Xây dựng sự kiện và kiểm tra tài sản
Căn chỉnh dữ liệu và tokenizer
Lỗi thường gặp: chỉ so sánh vocab_size bỏ qua hoán vị ID: hãy so sánh các từ điển và ví dụ đã mã hóa.
Bất biến: các họ tách biệt và báo cáo căn chỉnh đầy đủ; dừng rõ ràng khi mismatch.
Tệp src/distillation_generative/data.py
"""Original synthetic parcel events, grouped by full parcel/scenario family."""
import hashlib,json,random
from pathlib import Path
HUBS=['East','West','North','South','Central','Coastal','Valley','Airport']
FIELDS=['parcel','event','location','time','delivered']
EVENTS=['received','transferred','arrived_at_hub','delivered']
def generate_parcel_messages():
order=list(range(400));random.Random(50).shuffle(order);split={family:('train' if i<280 else 'validation' if i<340 else 'test') for i,family in enumerate(order)};pages=[];labels=[]
for family in range(400):
parcel=f'P{family:04d}';location=HUBS[family%8];old_location=HUBS[(family+3)%8];minutes=(family*37+65)%(24*60);clock=f'{minutes//60:02d}:{minutes%60:02d}'
for variant in range(5):
identifier=f'PAR{family:03d}-M{variant}';target={'parcel':parcel,'event':None,'location':location,'time':clock,'delivered':False}
if variant==0:
message=f'Parcel {parcel} was received by the carrier at {location} at {clock}. It has not been delivered to the recipient.';target['event']='received'
elif variant==1:
message=f'Parcel {parcel} left {old_location} earlier. Its latest confirmed event is a transfer to {location} at {clock}. The recipient has not received it.';target['event']='transferred'
elif variant==2:
message=f'The recipient confirmed delivery of parcel {parcel} at {location} at {clock}. This is a completed delivery, not a hub arrival.';target.update(event='delivered',delivered=True)
elif variant==3:
message=f'A parcel reached hub {location}. Its parcel identifier and event time were not reported. The recipient has not received it. Internal case note {family:03d} is not a parcel identifier.';target.update(parcel=None,event='arrived_at_hub',time=None)
else:
message=f'An unverified draft said parcel {parcel} was delivered at {old_location}. Correction: parcel {parcel} only arrived at a hub at {clock}; the hub name was not reported. Delivery has not occurred.';target.update(event='arrived_at_hub',location=None)
pages.append({'document_id':identifier,'family':f'PAR{family:03d}','split':split[family],'variant':variant,'statement':message,'content_sha256':hashlib.sha256(message.encode()).hexdigest()});labels.append({'document_id':identifier,'target':target,'response':json.dumps(target,ensure_ascii=False,separators=(',',':'))})
return pages,labels
def audit_corpus(pages,labels):
if len(pages)!=2000 or len(labels)!=2000 or len({r['document_id'] for r in pages})!=2000 or len({r['statement'] for r in pages})!=2000:raise ValueError('Expected2000unique original messages')
families={};gold={r['document_id']:r for r in labels}
for p in pages:
families.setdefault(p['family'],[]).append(p)
if p['content_sha256']!=hashlib.sha256(p['statement'].encode()).hexdigest() or set(gold[p['document_id']]['target'])!=set(FIELDS) or json.loads(gold[p['document_id']]['response'])!=gold[p['document_id']]['target']:raise ValueError('Corpus integrity changed')
if len(families)!=400 or any(len(rows)!=5 or len({r['split'] for r in rows})!=1 for rows in families.values()):raise ValueError('Whole parcel family must stay in one split')
counts={s:sum(p['split']==s for p in pages) for s in ['train','validation','test']}
if counts!={'train':1400,'validation':300,'test':300}:raise ValueError('Frozen split sizes changed')
return {'dataset':'ParcelEvents-v1','messages':2000,'families':400,'messages_by_split':counts,'families_by_split':{s:sum(rows[0]['split']==s for rows in families.values()) for s in counts},'shared_event_templates':True,'shared_hubs':HUBS,'real_parcels':False,'missing_field_counts':{field:sum(g['target'][field] is None for g in labels) for field in FIELDS},'delivery_labels':sum(g['target']['delivered'] for g in labels)}
def prepare(config):
from .config import save_json,sha
root=Path('data/generated')
if config['seed']!=50:raise ValueError('Frozen synthetic seed50 required')
if (root/'manifest.json').exists():raise ValueError('Preserve existing data; use a new extraction')
pages,labels=generate_parcel_messages();audit=audit_corpus(pages,labels);gold={g['document_id']:g for g in labels}
save_json(root/'input/pages.json',pages);save_json(root/'labels/gold.json',labels)
for split in ['train','validation']:
selected=[p for p in pages if p['split']==split];save_json(root/f'{split}/pages.json',selected);save_json(root/f'labels/{split}.json',[gold[p['document_id']] for p in selected])
save_json('outputs/data-audit.json',audit);save_json(root/'manifest.json',{**audit,'files':{str(f.relative_to(root)):sha(f) for f in root.rglob('*') if f.is_file()}});return audit
Tạo các tiền tố thực sự giống hệt
Teacher forcing và mặt nạ câu trả lời
Lỗi thường gặp: một template được thêm hai lần phía giáo viên làm lệch vị trí: hãy băm các token của batch chung.
Bất biến: ID, mặt nạ và độ dịch giống nhau cho teacher/student.
Tệp src/distillation_generative/collator.py
import hashlib,json
import torch
SYSTEM='Extract the latest confirmed parcel event from this fictional message. Return only one JSON object with keys parcel, event, location, time, delivered. Use null for missing parcel, location or time. Event is received, transferred, arrived_at_hub or delivered. Delivered is a boolean. Ignore unverified drafts and do not invent a delivery.'
def prompt_ids(tokenizer,statement):
messages=[{'role':'system','content':SYSTEM},{'role':'user','content':statement}]
return messages,tokenizer.apply_chat_template(messages,tokenize=True,add_generation_prompt=True)
def completion_row(tokenizer,statement,response,config,identifier):
messages,prompt=prompt_ids(tokenizer,statement);target=tokenizer.encode(response,add_special_tokens=False)+[151645]
if not prompt or len(prompt)>config['max_prompt_tokens'] or not 1<len(target)<=config['max_response_tokens_including_eos'] or len(prompt)+len(target)>config['context']:raise ValueError(identifier+': context budget exceeded; do not truncate gold')
ids=prompt+target;labels=[-100]*len(prompt)+target
return {'document_id':identifier,'messages':messages,'input_ids':ids,'labels':labels,'prompt_tokens':len(prompt),'response_tokens':len(target),'template_version':'parcel-events-v1','shared_tokens_sha256':hashlib.sha256(json.dumps({'input_ids':ids,'labels':labels},separators=(',',':')).encode()).hexdigest()}
def collate(rows,device='cuda',pad_to=None):
if not rows:raise ValueError('Empty batch')
width=max(len(r['input_ids']) for r in rows) if pad_to is None else pad_to;ids=torch.full((len(rows),width),151643,dtype=torch.long,device=device);mask=torch.zeros_like(ids,dtype=torch.bool);labels=torch.full_like(ids,-100)
for i,row in enumerate(rows):
n=len(row['input_ids'])
if n>width:raise ValueError('Padding would truncate an example')
ids[i,:n]=torch.tensor(row['input_ids'],device=device);mask[i,:n]=True;labels[i,:n]=torch.tensor(row['labels'],device=device)
return ids,mask,labels
Thích nghi giáo viên và đo chất lượng của nó
Giáo viên hữu ích, không phải nhà tiên tri
Lỗi thường gặp: dùng nhãn test để chọn teacher tốt nhất làm vô hiệu chuyển giao: chỉ chọn bằng validation.
Bất biến: adapter thực sự được huấn luyện, teacher đóng băng và kết quả ban đầu được giữ lại.
Tệp src/distillation_generative/teacher.py
"""Local teacher training shares the explicit accumulation loop with students."""
from .train import train,verify_resume
def train_teacher_adapter(config,mode):
return train(config,"teacher",mode)
Mã hóa KL và huấn luyện hai học viên
Target mềm và gradient
Lỗi thường gặp: lấy torch.kl_div theo chiều sai làm đảo phân kỳ: kiểm thử trên các phân phối bất đối xứng.
Bất biến: KL giải tích đúng, teacher bất biến, hai học viên thực sự được cập nhật và có thể tiếp tục.
Tệp src/distillation_generative/train.py
"""Explicit microbatch/accumulation training; no external trainer."""
import math,time
from pathlib import Path
import torch
from .config import read_json,save_json,sha,seed_all,rng_state,restore_rng
from .models import load_base,load_tokenizer,reload_teacher,adapter_state,copy_teacher_adapter,adapter_identity,frozen_teacher_identity,model_identity,parameter_manifest,save_teacher,save_student
from .collator import completion_row,collate
from .loss import masked_distillation_terms
from .state import save_resume_state,load_resume_state,equal
VARIANTS={'teacher':(0.,2.),'sft':(0.,2.),'kd':(.5,2.),'alpha1-t2':(1.,2.),'alpha1-t4':(1.,4.)}
def learning_rate(step,total,peak):
warm=max(1,int(total*.1))
if step<warm:return peak*(step+1)/warm
return peak*(1-min(1.,(step-warm)/max(1,total-warm-1))*.9)
def assert_learning_scope():
for name in ['data/generated/input','data/generated/labels/gold.json','outputs/inference']:
if Path(name).exists():raise ValueError('Learning must not expose heldout inputs/gold/inference: '+name)
alignment=read_json('outputs/learning/alignment.json')
if alignment['status']!='PASS' or alignment['encoded_examples_verified']!=2000 or not alignment['all_token_ids_identical']:raise ValueError('Actual complete tokenizer alignment required before training')
def encoded_rows(tokenizer,config,split):
pages=read_json(f'data/generated/{split}/pages.json');gold={g['document_id']:g for g in read_json(f'data/generated/labels/{split}.json')}
return [completion_row(tokenizer,p['statement'],gold[p['document_id']]['response'],config,p['document_id']) for p in pages]
def terms_for_row(model,teacher,row,alpha,temperature):
ids,mask,labels=collate([row]);teacher_logits=None;teacher_seconds=0.
if alpha>0:
if teacher is None or teacher.training or any(p.requires_grad or p.grad is not None for p in teacher.parameters()):raise ValueError('Distillation teacher must be frozen and in eval mode')
torch.cuda.synchronize();start=time.perf_counter()
with torch.no_grad():teacher_logits=teacher(input_ids=ids,attention_mask=mask,use_cache=False).logits.float()
torch.cuda.synchronize();teacher_seconds=time.perf_counter()-start
student_logits=model(input_ids=ids,attention_mask=mask,use_cache=False).logits
return masked_distillation_terms(student_logits,labels,teacher_logits,alpha,temperature),teacher_seconds
@torch.no_grad()
def validation_scores(model,teacher,rows,alpha,temperature):
records=[]
for row in rows:
terms,_=terms_for_row(model,teacher,row,alpha,temperature);records.append({'document_id':row['document_id'],'tokens':terms['tokens'],'ce_sum':float(terms['ce_sum']),'kl_sum':float(terms['kl_sum']),'mixed_sum':float(terms['mixed_sum']),'token_logprobs':terms['student_token_logprobs'].cpu().tolist(),'token_kl':terms['teacher_token_kl'].cpu().tolist(),'target_ids':terms['targets'].cpu().tolist()})
den=sum(r['tokens'] for r in records)
return {'examples':len(rows),'tokens':den,'ce':sum(r['ce_sum'] for r in records)/den,'kl':sum(r['kl_sum'] for r in records)/den,'mixed':sum(r['mixed_sum'] for r in records)/den,'used_for_selection':False,'rows':records}
def train(config,branch,mode):
if branch not in VARIANTS or mode not in ['interrupt','continuous','resume']:raise ValueError('Unknown training branch or continuation mode')
assert_learning_scope();seed_all(50);alpha,temperature=VARIANTS[branch];role='teacher' if branch=='teacher' else 'student';tokenizer=load_tokenizer(role);rows=encoded_rows(tokenizer,config,'train');validation=encoded_rows(tokenizer,config,'validation')
if len(rows)!=1400 or len(validation)!=300:raise ValueError('Frozen training populations changed')
model=load_base(role);teacher=None;teacher_meta=None
if alpha>0:teacher,teacher_meta=reload_teacher('assets/teacher-checkpoints/teacher-resume.safetensors',frozen=True)
parameters=[p for p in model.parameters() if p.requires_grad];expected=9232384 if role=='teacher' else 494032768
if sum(p.numel() for p in parameters)!=expected or len({id(p) for p in parameters})!=len(parameters):raise ValueError('Unexpected trainable population or duplicate tied parameter')
initial_parameters=parameter_manifest(model);initial=model_identity(model);frozen_base=frozen_teacher_identity(model) if role=='teacher' else None;teacher_before=model_identity(teacher) if teacher is not None else None;peak=1e-4 if role=='teacher' else 2e-5;total=350
identity={'branch':branch,'config':config,'alpha':alpha,'temperature':temperature,'initial_model_digest':initial,'teacher_checkpoint':None if teacher_meta is None else teacher_meta['sha256'],'teacher_digest':teacher_before,'alignment_sha256':sha('outputs/learning/alignment.json'),'data_sha256':{n:sha('data/generated/'+n) for n in ['train/pages.json','labels/train.json','validation/pages.json','labels/validation.json']},'model_sha256':sha('assets/'+role+'/model.safetensors'),'tokenizer_sha256':sha('assets/'+role+'/tokenizer.json'),'code_sha256':{n:sha('src/distillation_generative/'+n) for n in ['models.py','collator.py','loss.py','train.py','state.py','alignment.py']}}
optimizer=torch.optim.AdamW(parameters,lr=peak,weight_decay=.01);sampler=torch.Generator().manual_seed(150);root=Path('assets')/('teacher-checkpoints' if role=='teacher' else 'student-checkpoints');destination=root/f'{branch}-{mode}.pt.gz';updates=first_epoch=0;history=[];validation_history=[]
if destination.exists():raise ValueError('Preserve existing training checkpoint')
if mode=='resume':
state=load_resume_state(root/f'{branch}-interrupt.pt.gz')
if state['identity']!=identity:raise ValueError('Resume teacher/student/data/objective identity changed')
if role=='teacher':copy_teacher_adapter(model,state['model'])
else:model.load_state_dict(state['model'],strict=True)
optimizer.load_state_dict(state['optimizer']);sampler.set_state(state['sampler']);restore_rng(state['rng']);updates=state['updates'];first_epoch=state['epoch'];history=state['history'];validation_history=state['validation_history']
if state['scheduler']!={'next_update':updates,'next_lr':learning_rate(updates,total,peak)}:raise ValueError('Resume scheduler changed')
del state
stop=1 if mode=='interrupt' else 2;phase={'teacher_forward_seconds':0.,'optimized_model_forward_loss_backward_seconds':0.,'optimizer_seconds':0.,'validation_seconds':0.};torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();start_run=time.perf_counter()
for epoch in range(first_epoch,stop):
order=torch.randperm(1400,generator=sampler).tolist()
for offset in range(0,1400,8):
batch=[rows[i] for i in order[offset:offset+8]];den=sum(r['response_tokens'] for r in batch);optimizer.zero_grad(set_to_none=True);lr=learning_rate(updates,total,peak)
for group in optimizer.param_groups:group['lr']=lr
trace=[];loss_total=0.
for row in batch:
torch.cuda.synchronize();started=time.perf_counter();terms,teacher_seconds=terms_for_row(model,teacher,row,alpha,temperature);loss=terms['mixed_sum']/den;loss.backward();torch.cuda.synchronize();elapsed=time.perf_counter()-started;phase['teacher_forward_seconds']+=teacher_seconds;phase['optimized_model_forward_loss_backward_seconds']+=elapsed-teacher_seconds;probe=None
if updates==0 and len(trace)<2:
probe={'target':int(terms['targets'][0]),'student_logits':terms['student_selected_logits'][0].detach().cpu().tolist(),'teacher_logits':None if terms['teacher_selected_logits'] is None else terms['teacher_selected_logits'][0].detach().cpu().tolist()}
trace.append({'document_id':row['document_id'],'shared_tokens_sha256':row['shared_tokens_sha256'],'tokens':terms['tokens'],'target_ids':terms['targets'].detach().cpu().tolist(),'token_logprobs':terms['student_token_logprobs'].detach().cpu().tolist(),'token_kl':terms['teacher_token_kl'].detach().cpu().tolist(),'ce_sum':float(terms['ce_sum'].detach()),'kl_sum':float(terms['kl_sum'].detach()),'mixed_sum':float(terms['mixed_sum'].detach()),'first_token_probe':probe});loss_total+=float(loss.detach());del terms,loss
if any(p.grad is None or not torch.isfinite(p.grad).all() for p in parameters):raise ValueError('Missing/nonfinite trainable gradients')
if teacher is not None and any(p.grad is not None or p.requires_grad for p in teacher.parameters()):raise ValueError('Teacher received a student gradient')
started=time.perf_counter();norm=float(torch.nn.utils.clip_grad_norm_(parameters,1.,error_if_nonfinite=True));optimizer.step();torch.cuda.synchronize();phase['optimizer_seconds']+=time.perf_counter()-started;updates+=1;history.append({'epoch':epoch+1,'update':updates,'lr':lr,'response_token_denominator':den,'rows':trace,'loss':loss_total,'unclipped_gradient_norm':norm})
started=time.perf_counter();validation_history.append({'epoch':epoch+1,**validation_scores(model,teacher,validation,alpha,temperature)});torch.cuda.synchronize();phase['validation_seconds']+=time.perf_counter()-started
torch.cuda.synchronize();seconds=time.perf_counter()-start_run;final_parameters=parameter_manifest(model);final=model_identity(model)
if final==initial or not any(h['unclipped_gradient_norm']>0 for h in history):raise ValueError('No actual training update')
if role=='teacher' and frozen_teacher_identity(model)!=frozen_base:raise ValueError('Teacher base changed during LoRA')
if teacher is not None and model_identity(teacher)!=teacher_before:raise ValueError('Distillation teacher changed')
saved=adapter_state(model) if role=='teacher' else model.state_dict();state={'identity':identity,'model':{n:t.detach().cpu().clone() for n,t in saved.items()},'optimizer':optimizer.state_dict(),'sampler':sampler.get_state(),'rng':rng_state(),'epoch':stop,'updates':updates,'history':history,'validation_history':validation_history,'scheduler':{'next_update':updates,'next_lr':learning_rate(updates,total,peak)}};storage=save_resume_state(destination,state)
result={'branch':branch,'mode':mode,'identity':identity,'epochs':stop,'updates':updates,'initial_parameters':initial_parameters,'final_parameters':final_parameters,'initial_model_digest':initial,'final_model_digest':final,'trainable_parameters':expected,'history':history,'validation_history':validation_history,'teacher_unchanged':True,'teacher_gradients':0,'teacher_forward_calls':(stop-first_epoch)*1400 if alpha>0 else 0,'validation_teacher_forward_calls':(stop-first_epoch)*300 if alpha>0 else 0,'frozen_base_unchanged':True if role=='teacher' else None,'seconds':seconds,'phase_seconds':phase,'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'checkpoint_storage':storage,'checkpoint_sha256':sha(destination),'training_test_visible':False};save_json(f'outputs/learning/{branch}/{mode}.json',result)
if mode=='resume':
saver=save_teacher if role=='teacher' else save_student;saver(model,root/f'{branch}-resume.safetensors',identity)
return {'branch':branch,'mode':mode,'epochs':stop,'updates':updates,'seconds':seconds}
def verify_resume(branch):
if branch not in VARIANTS:raise ValueError('Unknown variant')
root=Path('assets')/('teacher-checkpoints' if branch=='teacher' else 'student-checkpoints');a=root/f'{branch}-continuous.pt.gz';b=root/f'{branch}-resume.pt.gz';continuous=load_resume_state(a);resumed=load_resume_state(b)
if not equal(continuous,resumed):raise ValueError('Exact model/teacher/optimizer/sampler/RNG continuation failed')
result={'status':'PASS','branch':branch,'all_state_exact':True,'updates':resumed['updates'],'continuous_sha256':sha(a),'resume_sha256':sha(b)};save_json(f'outputs/learning/{branch}/resume-check.json',result);return result
Phục vụ học viên mà không phụ thuộc giáo viên
Chất lượng tự do và chi phí suy luận
Lỗi thường gặp: một lần sinh ngắn hơn có vẻ nhanh hơn: công bố độ dài và chi phí theo yêu cầu và token.
Bất biến: không yêu cầu teacher nào trong suy luận học viên; 300 ID mỗi nhánh.
Tệp src/distillation_generative/evaluate.py
import json,time,statistics
from pathlib import Path
import torch
from .config import read_json,save_json,sha
from .models import load_base,load_tokenizer,reload_teacher,reload_student,model_identity,parameter_manifest
from .collator import prompt_ids
from .event_schema import FIELDS,parse_event,score_event
@torch.no_grad()
def generate_one(model,tokenizer,page,config,probe=False):
started=time.perf_counter();messages,prompt=prompt_ids(tokenizer,page['statement']);tokenization_seconds=time.perf_counter()-started
if not prompt or len(prompt)>config['max_prompt_tokens'] or len(prompt)+config['max_new_tokens']>config['context']:raise ValueError('Prompt exceeds frozen generation budget')
ids=torch.tensor([prompt],device='cuda',dtype=torch.long);generated=[];selected=[];first=None;torch.cuda.synchronize();start=time.perf_counter();first_token_seconds=None
for step in range(config['max_new_tokens']):
logits=model(input_ids=ids,attention_mask=torch.ones_like(ids),use_cache=False).logits[0,-1].float();token=int(logits.argmax());logprob=float(torch.log_softmax(logits,-1)[token])
if not torch.isfinite(logits).all() or not torch.isfinite(torch.tensor(logprob)):raise ValueError('Nonfinite generated distribution')
if step==0:
torch.cuda.synchronize();first_token_seconds=time.perf_counter()-start
if probe:first={'prefix_ids':prompt,'target':token,'logits':logits.cpu().tolist(),'selected_logprob':logprob}
generated.append(token);selected.append(logprob);ids=torch.cat([ids,torch.tensor([[token]],device='cuda')],dim=1)
if token==151645:break
torch.cuda.synchronize();generation_seconds=time.perf_counter()-start;finish='eos' if generated[-1]==151645 else 'length';decode_start=time.perf_counter();text=tokenizer.decode(generated[:-1] if finish=='eos' else generated,skip_special_tokens=False,clean_up_tokenization_spaces=False);decode_seconds=time.perf_counter()-decode_start
return {'document_id':page['document_id'],'messages':messages,'prompt_ids':prompt,'generated_ids':generated,'generated_tokens':len(generated),'token_logprobs':selected,'text':text,'tokenization_valid':all(tokenizer.convert_ids_to_tokens(t) is not None for t in generated),'finish_reason':finish,'mode':'greedy','max_new_tokens':96,'first_token_probe':first,'tokenization_seconds':tokenization_seconds,'generation_seconds':generation_seconds,'first_token_seconds':first_token_seconds,'decode_seconds':decode_seconds,'seconds_per_generated_token':generation_seconds/len(generated),'kv_cache':False}
@torch.no_grad()
def generate_events(config,branch):
allowed=['initial','teacher','sft','kd','alpha1-t2','alpha1-t4']
if branch not in allowed:raise ValueError('Unknown serving branch')
if Path('data/generated/labels').exists() or Path('outputs/learning').exists():raise ValueError('Serving must not mount labels or learning outputs')
role='teacher' if branch=='teacher' else 'student'
if role=='student' and (Path('assets/teacher').exists() or Path('assets/teacher-checkpoints').exists()):raise ValueError('Student serving must be physically independent of the teacher')
pages=read_json('data/generated/input/pages.json');test=[p for p in pages if p['split']=='test'];warmup_page=next(p for p in pages if p['split']=='train')
if len(test)!=300 or len({p['document_id'] for p in test})!=300:raise ValueError('Exactly300reserved messages required')
started=time.perf_counter();tokenizer=load_tokenizer(role);tokenizer_load_seconds=time.perf_counter()-started;started=time.perf_counter()
if branch=='teacher':model,metadata=reload_teacher('assets/teacher-checkpoints/teacher-resume.safetensors');served=['assets/teacher/model.safetensors','assets/teacher-checkpoints/teacher-resume.safetensors']
elif branch=='initial':model=load_base('student');metadata={'sha256':sha('assets/student/model.safetensors')};served=['assets/student/model.safetensors']
else:model,metadata=reload_student(f'assets/student-checkpoints/{branch}-resume.safetensors');served=[f'assets/student-checkpoints/{branch}-resume.safetensors']
model.eval().requires_grad_(False);torch.cuda.synchronize();load_seconds=time.perf_counter()-started;identity=model_identity(model);manifest=parameter_manifest(model);parameters=sum(p.numel() for p in model.parameters());logical_bytes=sum(p.numel()*p.element_size() for p in model.parameters());files={name:{'sha256':sha(name),'size_bytes':Path(name).stat().st_size} for name in served}
# Same fixed train statement warms each model; it is outside all300test denominators.
warmup=generate_one(model,tokenizer,warmup_page,config);torch.cuda.reset_peak_memory_stats();torch.cuda.synchronize();started=time.perf_counter();rows=[]
for index,page in enumerate(test):rows.append({'page':page,'generation':generate_one(model,tokenizer,page,config,probe=index<2)})
torch.cuda.synchronize();seconds=time.perf_counter()-started
if model_identity(model)!=identity:raise ValueError('Inference changed served weights')
loops=[r['generation']['generation_seconds'] for r in rows];lengths=[r['generation']['generated_tokens'] for r in rows];result={'branch':branch,'rows':rows,'cases':300,'labels_visible':False,'teacher_visible':role=='teacher','teacher_calls_during_student_serving':0,'checkpoint_sha256':metadata['sha256'],'model_digest':identity,'parameters':parameters,'logical_parameter_bytes':logical_bytes,'parameter_dtypes':sorted({r['dtype'] for r in manifest}),'served_weight_files':files,'served_weight_file_bytes':sum(v['size_bytes'] for v in files.values()),'tokenizer_load_seconds':tokenizer_load_seconds,'model_load_seconds_including_integrity_checks':load_seconds,'warmup':warmup,'test_seconds':seconds,'latency':{'mean_request_generation_seconds':sum(loops)/300,'median_request_generation_seconds':statistics.median(loops),'p95_request_generation_seconds':sorted(loops)[284],'total_generated_tokens_including_eos':sum(lengths),'mean_generated_tokens_including_eos':sum(lengths)/300,'weighted_seconds_per_generated_token':sum(loops)/sum(lengths),'tokenization_seconds':sum(r['generation']['tokenization_seconds'] for r in rows),'decode_seconds':sum(r['generation']['decode_seconds'] for r in rows),'first_two_requests_include_full_logit_probe_copy':True},'peak_allocated_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'same_no_kv_cache_greedy_algorithm':True};save_json(f'outputs/inference/{branch}.json',result);return {k:v for k,v in result.items() if k not in ['rows','warmup','served_weight_files']}
def summarize(rows):
n=len(rows);valid=[r for r in rows if r['parsed']['complete_valid']];not_delivered=[r for r in rows if r['score']['false_delivery_eligible']];valid_not_delivered=[r for r in not_delivered if r['parsed']['complete_valid']];missing=sum(r['score']['missing_fields'] for r in rows);valid_missing=sum(r['score']['missing_fields'] for r in valid);filled=sum(r['score']['filled_when_absent'] for r in rows);false=sum(r['score']['false_delivery'] for r in rows)
return {'cases':n,'json_parsable':sum(r['parsed']['json_parsable'] for r in rows),'schema_valid_before_eos_check':sum(r['parsed']['schema_valid'] for r in rows),'complete_valid':len(valid),'complete_valid_rate':len(valid)/n,'exact_match':sum(r['score']['exact_match'] for r in rows),'exact_match_rate':sum(r['score']['exact_match'] for r in rows)/n,'per_field_correct':{f:sum(r['score']['fields_correct'][f] for r in rows) for f in FIELDS},'per_field_denominator':n,'false_delivery':false,'not_delivered_reference_cases':len(not_delivered),'valid_not_delivered_reference_cases':len(valid_not_delivered),'false_delivery_rate_all_eligible':false/len(not_delivered) if not_delivered else None,'false_delivery_rate_valid_eligible':false/len(valid_not_delivered) if valid_not_delivered else None,'filled_when_absent':filled,'missing_reference_fields':missing,'missing_reference_fields_in_valid_outputs':valid_missing,'filled_when_absent_rate_valid_fields':filled/valid_missing if valid_missing else None,'without_eos':sum(r['finish_reason']!='eos' for r in rows),'undecodable_token':sum(not r['tokenization_valid'] for r in rows),'invalid_outputs_not_certified':n-len(valid)}
def evaluate(config,exercise=False):
gold={g['document_id']:g for g in read_json('data/generated/labels/gold.json')};branches=['alpha1-t2','alpha1-t4'] if exercise else ['initial','teacher','sft','kd'];result={}
for branch in branches:
raw=read_json(f'outputs/inference/{branch}.json');rows=[]
for item in raw['rows']:
page=item['page'];g=item['generation'];target=gold[page['document_id']]['target'];parsed=parse_event(g['text'],g['finish_reason'],g['tokenization_valid']);score=score_event(parsed,target);rows.append({'document_id':page['document_id'],'target':target,'text':g['text'],'finish_reason':g['finish_reason'],'tokenization_valid':g['tokenization_valid'],'parsed':parsed,'score':score})
result[branch]={'metrics':summarize(rows),'rows':rows,'raw_sha256':sha(f'outputs/inference/{branch}.json'),'costs':{k:raw[k] for k in ['parameters','logical_parameter_bytes','served_weight_file_bytes','parameter_dtypes','tokenizer_load_seconds','model_load_seconds_including_integrity_checks','test_seconds','latency','peak_allocated_bytes','peak_reserved_bytes']},'warmup_document_id':raw['warmup']['document_id'],'warmup_generation_seconds':raw['warmup']['generation_seconds']}
output={'policies':result,'scope':'Original fictitious messages, fixed300testcases; malformed/truncated/undecodable outputs remain in all-case denominators. False-delivery/absent-field rates on valid outputs are conditional and report coverage. FP32student versusBF16teacher+FP32LoRA is not a pure model-size latency experiment. No teacher calls during student serving.'};save_json('outputs/exercise/results.json' if exercise else 'outputs/results.json',output);return {b:v['metrics'] for b,v in result.items()}
Khảo sát trọng số của chuyển giao
Vai trò của alpha và T
Lỗi thường gặp: dùng lại checkpoint SFT rồi đổi tên alpha=1 không phải là chưng cất: phải yêu cầu dấu vết cập nhật.
Bất biến: alpha=0 khớp với SFT và mỗi biến thể có danh tính riêng.
Tệp src/distillation_generative/exercise.py
from pathlib import Path
from .config import read_json,save_json,sha
def begin(config):
path=Path('outputs/exercise/protected-main.json')
if path.exists():raise ValueError('Preserve an existing exercise')
files=[Path('configs/minimum.yaml'),Path('outputs/results.json'),Path('outputs/tokenizer-alignment.json')]+list(Path('outputs/inference').glob('*.json'))+list(Path('assets/teacher-checkpoints').glob('*.safetensors'))+list(Path('assets/student-checkpoints').glob('*.safetensors'))
identities={str(p):sha(p) for p in files};save_json(path,identities);return {'protected_files':len(identities),'exercise':'Freshalpha1T2andalpha1T4; actual mainSFT isalpha0 comparator, no renamed weights'}
def finish(config):
protected=read_json('outputs/exercise/protected-main.json')
if any(sha(path)!=value for path,value in protected.items()):raise ValueError('Exercise changed main artifacts')
main=read_json('outputs/results.json')['policies'];extra=read_json('outputs/exercise/results.json')['policies'];variants={n:{'metrics':v['metrics'],'costs':v['costs']} for n,v in {**main,**extra}.items() if n in ['sft','kd','alpha1-t2','alpha1-t4']};budgets={}
for branch in variants:
run=read_json(f'outputs/learning/{branch}/continuous.json');budgets[branch]={'updates':run['updates'],'alpha':run['identity']['alpha'],'temperature':run['identity']['temperature'],'seconds':run['seconds'],'teacher_forward_seconds':run['phase_seconds']['teacher_forward_seconds'],'final_model_digest':run['final_model_digest']}
if run['updates']!=350 or run['identity']['initial_model_digest']!=read_json('outputs/learning/sft/continuous.json')['identity']['initial_model_digest']:raise ValueError('Exercise must start from same student and fixed350updates')
if set(variants)!={'sft','kd','alpha1-t2','alpha1-t4'} or len({b['final_model_digest'] for b in budgets.values()})!=4:raise ValueError('Distinct actually learned branch states required')
result={'status':'PASS','variants':variants,'budgets':budgets,'main_unchanged':True,'protected_sha256':protected,'alpha0_is_actual_main_sft':True,'fresh_exercise_training':True,'teacher_refit':False,'test_selection':False,'quality_gain_required':False};save_json('outputs/exercise/distillation-mixture.json',result);return {k:v for k,v in result.items() if k not in ['protected_sha256','variants']}
Nguồn và phạm vi tham chiếu
Hinton et al., Distilling the Knowledge in a Neural Network giải thích nhiệt độ và hệ số T². PyTorch 2.8 kl_div cung cấp công thức tham chiếu dùng trong kiểm thử. TRL GKD 0.23.1 mô tả đặc biệt các tiền tố do học viên sinh: bài tập của chúng tôi dùng tiền tố gold ngoài chính sách và không tuyên bố tương đương trainer này. Các mô hình Qwen và phụ thuộc có revision và hash trong các lock công khai.
