A fictional system turns delivery messages into short JSON events. A locally adapted 1.5B teacher passes its next-token distributions to a 0.5B student, evaluated against a student trained only on the same labels.
For one message, show the teacher, SFT student, and distilled student, their field errors and inference costs. Show soft probabilities for a few tokens at the same prefix, with temperature.
Download the complete project, revision r001
A bounded task and held-out families
ParcelEvents-v1 contains 400 fictional parcel families, five messages per family, or 2,000 examples. A seed-50 permutation splits families into 280 training, 60 validation, and 60 test families: 1,400/300/300 messages. Five situations cover carrier receipt, transfer between hubs, confirmed delivery to a recipient, arrival with missing identifier and time, and correction of an unconfirmed draft. Arrival at a hub is not delivery to a recipient.
Responses have exactly five fields: parcel, event, location, time, and delivered. parcel/location/time may be null; delivered is Boolean. Four events are allowed: received, transferred, arrived_at_hub, and delivered. Data share eight hubs and sentence patterns across partitions. Identities and families remain disjoint; this regularity does not represent the diversity of real logistics documents. Of 300 test cases, 60 state a delivery and 240 do not; each of the three nullable items is missing in 60 cases.
Align coordinates before comparing distributions
The Qwen2.5-1.5B-Instruct teacher and Qwen2.5-0.5B-Instruct student each produce 151936 logits. Equal dimension does not prove coordinate 42 means the same token. The check compares the complete token-to-ID dictionary, added and special tokens, their IDs, the chat template, serialized BPE configuration, and generation configuration.
The acquired tokenizers must describe exactly the same 151665 effective IDs and the same 271 unused output coordinates. Those coordinates remain in the full KL. They are not removed to make two distributions match artificially. Any vocabulary, template, or special-token mismatch stops the course without silent remapping.
One prompt-and-response construction provides input_ids and labels. The check actually encodes all 2,000 examples with both tokenizers and requires the same sequence and hash. Prompts have at most 320 tokens, responses with EOS at most 96, and context 416. Oversized targets are rejected; no gold truncation is hidden. EOS is 151645 and PAD 151643. Padding is on the right and its mask applies to positions.
Learn the teacher locally
The teacher has 28 blocks, width 1536, 12 query heads, and two KV heads. Its BF16 base stays frozen. We code LoRA: W(x)+(alpha/r)B(Ax), rank 8, alpha 16, zero dropout, on the seven q/k/v/o/gate/up/down projections in each of 28 blocks. The 196 projections add 9,232,384 FP32 parameters. The addition is computed in FP32 then returns to the BF16 base dtype. Tests compare values and gradients with PEFT; no PEFT trainer is called.
The teacher learns 1,400 training responses for two epochs. With microbatch 1 and accumulation of eight examples, this makes 350 AdamW updates at 1e−4. Validation follows cross-entropy each epoch; the protocol retains the final state fixed in advance, without selecting a checkpoint from test results. The teacher then becomes eval and fully frozen for distillation. Weight fingerprints and absence of gradients check this immutability.
Teacher forcing and the coded loss
At response token t, both models receive the same prompt and preceding gold tokens. This is off-policy teacher forcing: the teacher evaluates that supervised prefix, not a free student response. No giant cache of corpus logits or top-k approximation is used.
With z as teacher logits and u as student logits, p_T=softmax(z/T) and q_T=softmax(u/T). The directed divergence is KL(p_T || q_T)=sum_v p_T(v)[log p_T(v)−log q_T(v)]. Reversing arguments gives another loss. We use stable log_softmax and all 151936 coordinates. The teacher is computed under no_grad per microbatch; divergence sends a gradient only to the student.
The loss is (1−alpha) CE(q_1,target) + alpha T² KL(p_T || q_T). CE uses student logits at temperature 1. The T² factor compensates the scale of the distillation signal; it guarantees neither better prediction nor identical gradients at every temperature. Tests compare the formula and gradients with torch.nn.functional.kl_div on asymmetric FP64 distributions, check direction, zero for identical distributions, alpha=0/1 limits, and exclusion of the teacher graph.
Logits and labels are shifted to predict the next token, then only response tokens, including EOS, are selected. Prompt and PAD are excluded. Each accumulation divides by the total active tokens in eight examples. Averaging eight already-averaged losses would weight short and long responses equally. At alpha=0, the path is exactly CE and does not load the teacher.
Two students and an exact resume
The student has 24 blocks, width 896, 14 query heads, two KV heads, and 494,032,768 unique parameters. All are trained in FP32; embeddings and output projection share the same Parameter. The alpha=0 SFT control and alpha=0.5/T=2 distillation start from the same public weights, seed, and training order. Each runs two epochs, 350 updates at 2e−5.
AdamW uses weight_decay=0.01, gradient clipping at 1, 35 warm-up updates, then linear decay to 10% of the initial rate. Eval mode retains gradients and disables dropout sources; it is not no_grad. The SFT control calls no teacher. Teacher forwards in KD, optimized-model calculation/backward, optimization, and validation costs are retained separately.
Each branch compares interruption after one epoch, two continuous epochs, and continuation of the remaining epoch. Weights, Adam states, scheduler, permutation, RNG, and numerical histories must be exactly equal. Durations remain outside state comparison. Local states are losslessly compressed at gzip level 1; final weights are also saved and reloaded in Safetensors. The total with exercise retains 15 resume states and five final weights. These Torch pickles are produced locally; do not replace them with third-party files.
Generate freely and measure the actual scope
The adapted teacher, initial student, SFT, and KD each generate 300 test responses without a gold prefix. The two exercise variants do likewise: 1,800 test responses, plus separately identified warm-ups. Generation is greedy over the full vocabulary, without KV cache or forced schema, capped at 96 new tokens. IDs, raw texts, EOS or truncation, log probabilities, durations, and two full-logit probes are retained.
The student service physically mounts neither teacher nor its adapters. Reloading the student uses its configuration, tokenizer, and learned weights only. Inference mounts no labels. Evaluation then combines outputs and targets without modifying weights.
The parser distinguishes parseable JSON, strict schema, and complete response: a truncated output fails even when it contains a well-formed object. Duplicate keys, NaN, extra fields, Markdown, and unknown IDs are not repaired. Errors remain in the denominator of 300 cases for exact match and each field. False-delivery and wrongly-filled-field rates apply to valid complete responses and also report eligible total populations and coverage. An invalid output is not certified safe.
Each model receives the same warm-up training message outside the test. Reports give latency per request and token, lengths with EOS, median, and empirical p95 (rank 285 of 300). The first two requests include full-logit probe cost, stated in the report. Loading and integrity checks remain distinct. Unique parameters, dtypes, logical bytes, and actual served-file size are measured. The BF16 teacher with FP32 adapters and fully FP32 student also differ in precision: a parameter ratio is no guarantee of memory or speed.
Exercise: shift alpha and T
The main SFT is the real alpha=0 comparator. Two additional students start from the same base: alpha=1/T=2 and alpha=1/T=4, each with 350 updates and its own exact continuation. No checkpoint is simply renamed. Data, teacher, and main results are protected by fingerprints; the teacher is not retrained for the exercise.
Alpha=1 exclusively imitates the teacher on gold prefixes, including its errors. Raising T redistributes mass and gradients; it does not reveal the truth of the delivered field. Compare the same 300 events, errors, and costs. No improvement or optimal temperature is imposed by the course.
Limits
One seed, a regular synthetic corpus, and a small set of types do not constitute a business evaluation. Lower KL on gold prefixes does not prove better free generation. The teacher can transfer its errors. Local measurements without controlled concurrent load are not a service-performance commitment. The practical exercise sends no parcel, message, or external action.
Build events and verify assets
Data and tokenizer alignment
Common error: comparing only vocab_size ignores an ID permutation; compare dictionaries and encoded examples.
Invariant: disjoint families and complete alignment report; clear stop on 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
Form truly identical prefixes
Teacher forcing and response mask
Common error: adding a template twice on the teacher side shifts positions; hash tokens from the shared batch.
Invariant: IDs, masks, and shifts identical for teacher and 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
Adapt the teacher and measure its quality
A useful teacher, not an oracle
Common error: using test labels to choose the best teacher invalidates transfer; validation selection only.
Invariant: real trained adapters, frozen teacher, and initial results retained.
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)
Code KL and train both students
Soft targets and gradient
Common error: torch.kl_div in the wrong order reverses divergence; test asymmetric distributions.
Invariant: correct analytic KL, immutable teacher, two genuinely updated students, and resumability.
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
Serve the student without depending on the teacher
Free quality and inference cost
Common error: a shorter generation appears faster; publish lengths and cost per request and token.
Invariant: no teacher request during student inference; 300 IDs per branch.
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()}
Explore transfer weighting
The role of alpha and T
Common error: reusing the SFT checkpoint while renaming it alpha=1 is not distillation; require update traceability.
Invariant: alpha=0 matches SFT and every variant has its own identity.
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']}
Sources and reference scope
Hinton et al., Distilling the Knowledge in a Neural Network motivates temperature and the T² factor. PyTorch 2.8 kl_div provides the formula reference used by tests. TRL GKD 0.23.1 notably describes student-generated prefixes: our exercise uses off-policy gold prefixes and does not claim equivalence with that trainer. Qwen models and dependencies have revisions and hashes in public locks.
