Un sistema fittizio trasforma messaggi di consegna in brevi eventi JSON. Un insegnante da 1,5B viene adattato localmente, poi trasmette le distribuzioni del token successivo a uno studente da 0,5B, valutato rispetto a uno studente addestrato soltanto sulle stesse label.
Per lo stesso messaggio, mostrate insegnante, studente SFT e studente distillato, i loro errori di campo e i costi di inferenza. Mostrate le probabilità soft di alcuni token sullo stesso prefisso, con temperatura.
Scaricare il progetto completo, revisione r001
Un compito delimitato e famiglie riservate
ParcelEvents-v1 contiene 400 famiglie fittizie di pacchi, cinque messaggi per famiglia, cioè 2000 esempi. Una permutazione con seed 50 ripartisce le famiglie in 280 train, 60 validation e 60 test: 1400/300/300 messaggi. Le cinque situazioni coprono ricezione da parte del corriere, trasferimento tra hub, consegna confermata al destinatario, arrivo con identificativo e ora assenti e correzione di una bozza non confermata. Un arrivo all'hub non è una consegna al destinatario.
Le risposte possiedono esattamente cinque campi: parcel, event, location, time e delivered. I campi parcel/location/time possono valere null; delivered è un booleano. Sono ammessi quattro eventi: received, transferred, arrived_at_hub e delivered. I dati condividono otto hub e schemi di frase tra le partizioni. Identità e famiglie restano disgiunte; questa regolarità non rappresenta la diversità dei documenti logistici reali. Nei 300 casi test, 60 indicano una consegna e 240 non la indicano; ciascuna delle tre informazioni nullable manca in 60 casi.
Allineare le coordinate prima di confrontare le distribuzioni
Il professore Qwen2.5-1.5B-Instruct e lo studente Qwen2.5-0.5B-Instruct producono ciascuno 151936 logits. Una dimensione uguale non prova che la coordinata 42 significhi lo stesso token. Il controllo confronta l'intero dizionario token→ID, i token aggiunti e speciali, i loro ID, il template di conversazione, la configurazione BPE serializzata e la configurazione di generazione.
I tokenizzatori acquisiti devono descrivere esattamente gli stessi 151665 ID effettivi e le stesse 271 coordinate di output inutilizzate. Queste coordinate restano incluse nella KL completa. Non le eliminiamo per far coincidere artificialmente due distribuzioni. Un disaccordo su vocabolario, template o token speciale interrompe il percorso senza remappatura silenziosa.
Un'unica costruzione di prompt e risposta fornisce input_ids e labels. Il controllo codifica realmente tutti i 2000 esempi con i due tokenizzatori ed esige la stessa sequenza e lo stesso hash. Prompt al massimo 320 token, risposta con EOS al massimo 96, contesto 416. Un superamento dei target viene rifiutato; nessun troncamento del gold è nascosto. EOS vale 151645 e PAD 151643. Il padding è a destra e la sua maschera riguarda le posizioni.
Apprendere il professore localmente
Il professore ha 28 blocchi, larghezza 1536, 12 teste di query e due teste KV. La sua base BF16 resta congelata. Codifichiamo LoRA: W(x)+(alpha/r)B(Ax), rango 8, alpha 16, dropout nullo, sulle sette proiezioni q/k/v/o/gate/up/down di ciascuno dei 28 blocchi. Le 196 proiezioni aggiungono 9 232 384 parametri FP32. L'addizione viene calcolata in FP32, poi ritorna al dtype BF16 della base. I test confrontano valori e gradienti con PEFT; nessun trainer PEFT viene chiamato.
Il professore apprende le 1400 risposte train per due epoche. Con microbatch 1 e accumulo di otto esempi, ciò produce 350 aggiornamenti AdamW a 1e−4. La validazione segue la CE a ogni epoca; il protocollo conserva lo stato finale fissato in anticipo, senza selezionare un checkpoint in base ai risultati test. Il professore passa poi a eval e resta interamente congelato per la distillazione. Le impronte dei pesi e l'assenza di gradienti controllano questa immutabilità.
Teacher forcing e loss codata
Al token di risposta t, i due modelli ricevono lo stesso prompt e gli stessi token gold precedenti. È teacher forcing fuori politica: il professore valuta questo prefisso supervisionato, non una risposta libera dello studente. Non vengono usate né una gigantesca cache di logits del corpus né un'approssimazione top-k.
Con z i logits del professore e u quelli dello studente, p_T=softmax(z/T), q_T=softmax(u/T). La divergenza orientata è KL(p_T || q_T)=somma_v p_T(v)[log p_T(v)−log q_T(v)]. Invertire gli argomenti produce un'altra loss. Usiamo log_softmax stabile e tutte le 151936 coordinate. Il professore viene calcolato sotto no_grad per microbatch; la divergenza trasmette un gradiente soltanto allo studente.
La loss vale (1−alpha) CE(q_1, target) + alpha T² KL(p_T || q_T). La CE usa i logits dello studente a temperatura 1. Il fattore T² compensa la scala del segnale di distillazione; non garantisce né una previsione migliore né un'identità perfetta dei gradienti per ogni temperatura. I test confrontano la formula e i suoi gradienti con torch.nn.functional.kl_div su distribuzioni FP64 asimmetriche, verificano il verso, la nullità per due distribuzioni identiche, i limiti alpha=0/1 e l'esclusione del grafo del professore.
Spostiamo logits e labels per predire il token successivo, poi selezioniamo soltanto i token di risposta, EOS incluso. Prompt e PAD sono esclusi. Ogni accumulo divide per il numero totale di token attivi degli otto esempi. Fare la media di otto loss già mediate darebbe lo stesso peso a una risposta corta e a una lunga. Con alpha=0, il percorso è esattamente la CE e non carica il professore.
Due studenti e una ripresa esatta
Lo studente ha 24 blocchi, larghezza 896, 14 teste di query, due teste KV e 494 032 768 parametri unici. Tutti vengono addestrati in FP32; embeddings e proiezione di output condividono lo stesso Parameter. Il testimone SFT alpha=0 e la distillazione alpha=0,5/T=2 ripartono dagli stessi pesi pubblici, dallo stesso seed e dallo stesso ordine train. Ciascuno esegue due epoche, 350 aggiornamenti a 2e−5.
AdamW usa weight_decay=0,01, clipping del gradiente a 1, warm-up di 35 update poi decrescita lineare fino al 10% del tasso iniziale. La modalità eval conserva i gradienti e disattiva le fonti di dropout; non è sinonimo di no_grad. Il testimone SFT non chiama alcun professore. I costi di forward del professore in KD, calcolo/backward del modello ottimizzato, ottimizzazione e validazione vengono conservati separatamente.
Ogni ramo confronta un'interruzione dopo un'epoca, due epoche continue e la continuazione dell'epoca rimanente. Pesi, stati Adam, scheduler, permutazione, RNG e storie numeriche devono essere esattamente uguali. Le durate restano fuori dal confronto di stato. Gli stati locali sono compressi senza perdita, livello gzip 1; i pesi finali vengono inoltre salvati e ricaricati in Safetensors. Il totale con esercizio conserva 15 stati di ripresa e cinque pesi finali. Questi pickle Torch sono prodotti localmente; non sostituiteli con file di terzi.
Generare liberamente e misurare il perimetro reale
Professore adattato, studente iniziale, SFT e KD generano ciascuno 300 risposte test, senza prefisso gold. Le due varianti dell'esercizio fanno lo stesso: 1800 risposte test, oltre ai warm-up identificati separatamente. La generazione è greedy sull'intero vocabolario, senza cache KV, senza schema forzato, con un limite di 96 nuovi token. ID, testi grezzi, EOS o taglio, log-probabilità, durate e due sonde di logits completi vengono conservati.
Il servizio studente non ha fisicamente alcun montaggio del professore o dei suoi adattatori. Il ricaricamento dello studente usa la sua configurazione, il suo tokenizer e i suoi soli pesi appresi. L'inferenza non monta alcuna label. La valutazione riunisce poi uscite e target, senza modificare i pesi.
Il parser distingue JSON parsabile, schema rigoroso e risposta completa: un'uscita tagliata fallisce anche se contiene un oggetto ben formato. Chiavi duplicate, NaN, campi aggiuntivi, Markdown e ID sconosciuti non vengono riparati. Gli errori restano al denominatore dei 300 casi per exact match e per ogni campo. I tassi di falsa consegna o di campo compilato a torto riguardano le risposte complete valide e indicano anche le popolazioni totali idonee e la copertura. Un'uscita non valida non è certificata sicura.
Ogni modello riceve lo stesso messaggio train di warm-up, fuori dal test. Riportiamo latenza per richiesta e per token, lunghezze con EOS, mediana e p95 empirico (rango 285 su 300). Le prime due richieste includono il costo delle sonde di logits, indicato nel rapporto. Il caricamento e le verifiche di integrità restano distinti. Vengono misurati parametri unici, dtype, byte logici e dimensione reale dei file serviti. Il professore BF16 con adattatori FP32 e lo studente interamente FP32 differiscono anche per precisione: un rapporto di parametri non è un rapporto garantito di memoria o velocità.
Esercizio: spostare alpha e T
L'SFT principale è il vero comparatore alpha=0. Due studenti aggiuntivi ripartono dalla stessa base: alpha=1/T=2 e alpha=1/T=4, ciascuno con 350 update e la propria continuazione esatta. Nessun checkpoint viene semplicemente rinominato. Dati, professore e risultati principali sono protetti da impronte; il professore non viene riaddestrato per l'esercizio.
Alpha=1 imita esclusivamente il professore sui prefissi gold, compresi i suoi errori. Aumentare T ridistribuisce la massa e i gradienti; non rivela la verità del campo delivered. Confrontate gli stessi 300 eventi, errori e costi. Il programma non impone alcun miglioramento né una temperatura ottimale.
Limiti
Un seed, un corpus sintetico regolare e un piccolo insieme di tipi non costituiscono una valutazione aziendale. La diminuzione della KL sui prefissi gold non prova una migliore generazione libera. Il professore può trasferire i suoi errori. Le misure locali senza carico concorrente controllato non sono un impegno sulle prestazioni di servizio. Il tutorial non invia alcun pacco, messaggio né azione esterna.
Costruire gli eventi e verificare gli asset
Allineamento di dati e tokenizzatori
Errore frequente: confrontare soltanto vocab_size ignora una permutazione degli ID: confrontare i dizionari e gli esempi codificati.
Invariante: famiglie disgiunte e rapporto di allineamento completo; arresto chiaro in caso di mismatch.
File 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
Formare prefissi davvero identici
Teacher forcing e maschera di risposta
Errore frequente: un template aggiunto due volte sul lato insegnante sposta le posizioni: calcolare l'hash dei token del batch comune.
Invariante: ID, maschere e spostamenti identici per teacher/student.
File 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
Adattare l'insegnante e misurare la sua qualità
Insegnante utile, non oracolo
Errore frequente: usare label test per scegliere il teacher migliore invalida il trasferimento: solo selezione validation.
Invariante: adattatori reali addestrati, teacher congelato e risultati iniziali conservati.
File 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)
Codificare la KL e addestrare i due studenti
Target soft e gradiente
Errore frequente: torch.kl_div preso nel verso sbagliato inverte la divergenza: test su distribuzioni asimmetriche.
Invariante: KL analitica corretta, teacher immutabile, due studenti davvero aggiornati e ripresa possibile.
File 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
Servire lo studente senza dipendere dall'insegnante
Qualità libera e costo di inferenza
Errore frequente: una generazione più corta sembra più veloce: pubblicare lunghezze e costo per richiesta e token.
Invariante: nessuna richiesta al teacher durante l'inferenza dello studente; 300 ID per ramo.
File 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()}
Esplorare la ponderazione del trasferimento
Ruolo di alpha e T
Errore frequente: riutilizzare il checkpoint SFT rinominandolo alpha=1 non è una distillazione: esigere la traccia degli aggiornamenti.
Invariante: alpha=0 coincide con SFT e ogni variante possiede la propria identità.
File 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']}
Fonti e perimetro di riferimento
Hinton et al., Distilling the Knowledge in a Neural Network motiva temperatura e fattore T². PyTorch 2.8 kl_div fornisce il riferimento della formula usata dai test. TRL GKD 0.23.1 descrive in particolare prefissi generati dallo studente: il nostro esercizio usa i prefissi gold fuori politica e non rivendica l'equivalenza di questo trainer. I modelli Qwen e le dipendenze hanno revisioni e hash nei lock pubblici.
