Un sistema ficticio transforma mensajes de entrega en eventos JSON breves. Un profesor de 1,5B se adapta localmente y transmite sus distribuciones del siguiente token a un estudiante de 0,5B, evaluado frente a otro estudiante entrenado solo con las mismas etiquetas.
Para un mismo mensaje, muestre al profesor, al estudiante SFT y al estudiante destilado, sus errores de campo y sus costes de inferencia. Muestre las probabilidades suaves de algunos tokens en el mismo prefijo, con temperatura.
Descargar el proyecto completo, revisión r001
Una tarea acotada y familias reservadas
ParcelEvents-v1 contiene 400 familias ficticias de paquetes, cinco mensajes por familia: 2000 ejemplos. Una permutación con semilla 50 reparte las familias en 280 de entrenamiento, 60 de validación y 60 de prueba: 1400/300/300 mensajes. Las cinco situaciones cubren recepción por el transportista, transferencia entre hubs, entrega confirmada al destinatario, llegada sin identificador ni hora y corrección de un borrador no confirmado. Una llegada al hub no es una entrega al destinatario.
Las respuestas tienen exactamente cinco campos: parcel, event, location, time y delivered. Los campos parcel/location/time pueden valer null; delivered es un booleano. Se admiten cuatro eventos: received, transferred, arrived_at_hub y delivered. Los datos comparten ocho hubs y patrones de frase entre particiones. Las identidades y familias siguen siendo disjuntas; esta regularidad no representa la diversidad de documentos logísticos reales. En los 300 casos de prueba, 60 indican una entrega y 240 no; cada una de las tres informaciones que puede ser nula falta en 60 casos.
Alinear las coordenadas antes de comparar las distribuciones
El profesor Qwen2.5-1.5B-Instruct y el estudiante Qwen2.5-0.5B-Instruct producen cada uno 151936 logits. Una dimensión igual no demuestra que la coordenada 42 signifique el mismo token. El control compara el diccionario completo token→ID, los tokens añadidos y especiales, sus ID, la plantilla de conversación, la configuración BPE serializada y la configuración de generación.
Los tokenizers adquiridos deben describir exactamente los mismos 151665 ID efectivos y las mismas 271 coordenadas de salida no utilizadas. Estas coordenadas siguen incluidas en la KL completa. No se eliminan para hacer coincidir artificialmente dos distribuciones. Un desacuerdo de vocabulario, plantilla o token especial detiene el recorrido sin remapeo silencioso.
Una única construcción del prompt y de la respuesta proporciona input_ids y labels. El control codifica realmente los 2000 ejemplos con ambos tokenizers y exige la misma secuencia y el mismo hash. El prompt tiene como máximo 320 tokens, la respuesta con EOS 96 y el contexto 416. Se rechaza un exceso de los objetivos; no hay truncamiento oculto del gold. EOS vale 151645 y PAD 151643. El padding está a la derecha y su máscara se aplica a las posiciones.
Aprender el profesor localmente
El profesor tiene 28 bloques, anchura 1536, 12 cabezas de consulta y dos cabezas KV. Su base BF16 permanece congelada. Codificamos LoRA: W(x)+(alpha/r)B(Ax), rango 8, alpha 16, dropout nulo, en las siete proyecciones q/k/v/o/gate/up/down de cada uno de los 28 bloques. Las 196 proyecciones añaden 9 232 384 parámetros FP32. La suma se calcula en FP32 y luego vuelve al dtype BF16 de la base. Las pruebas comparan valores y gradientes con PEFT; no se llama a ningún trainer PEFT.
El profesor aprende las 1400 respuestas de entrenamiento durante dos épocas. Con microbatch 1 y acumulación de ocho ejemplos, esto da 350 actualizaciones AdamW a 1e−4. La validación sigue la CE en cada época; el protocolo conserva el estado final fijado de antemano, sin seleccionar un checkpoint según los resultados de prueba. El profesor pasa después a eval y queda totalmente congelado para la destilación. Las huellas de los pesos y la ausencia de gradientes controlan esta inmutabilidad.
Teacher forcing y pérdida codificada
En el token de respuesta t, ambos modelos reciben el mismo prompt y los mismos tokens gold anteriores. Es teacher forcing fuera de política: el profesor evalúa este prefijo supervisado, no una respuesta libre del estudiante. No se usan una caché gigantesca de logits del corpus ni aproximación top-k.
Con z como logits del profesor y u los del estudiante, p_T=softmax(z/T), q_T=softmax(u/T). La divergencia orientada es KL(p_T || q_T)=suma_v p_T(v)[log p_T(v)−log q_T(v)]. Invertir los argumentos da otra pérdida. Usamos log_softmax estable y las 151936 coordenadas. El profesor se calcula bajo no_grad por microbatch; la divergencia transmite un gradiente solo al estudiante.
La pérdida vale (1−alpha) CE(q_1, objetivo) + alpha T² KL(p_T || q_T). La CE utiliza los logits del estudiante a temperatura 1. El factor T² compensa la escala de la señal de destilación; no garantiza ni una predicción mejor ni una identidad perfecta de los gradientes para cualquier temperatura. Las pruebas comparan la fórmula y sus gradientes con torch.nn.functional.kl_div en distribuciones FP64 asimétricas, verifican el sentido, el valor nulo para dos distribuciones idénticas, los límites alpha=0/1 y la exclusión del grafo del profesor.
Se desplazan logits y etiquetas para predecir el siguiente token y solo se seleccionan los tokens de respuesta, incluido EOS. Se excluyen prompt y PAD. Cada acumulación divide por el número total de tokens activos de los ocho ejemplos. Promediar ocho pérdidas ya promediadas daría el mismo peso a una respuesta corta y a una larga. Con alpha=0, el recorrido es exactamente la CE y no carga al profesor.
Dos estudiantes y una reanudación exacta
El estudiante tiene 24 bloques, anchura 896, 14 cabezas de consulta, dos cabezas KV y 494 032 768 parámetros únicos. Todos se entrenan en FP32; los embeddings y la proyección de salida comparten el mismo Parameter. El control SFT alpha=0 y la destilación alpha=0,5/T=2 parten de los mismos pesos públicos, la misma semilla y el mismo orden de entrenamiento. Cada uno realiza dos épocas, 350 actualizaciones a 2e−5.
AdamW usa weight_decay=0,01, recorte de gradiente a 1, calentamiento de 35 actualizaciones y una disminución lineal hasta el 10 % de la tasa inicial. El modo eval conserva los gradientes y desactiva las fuentes de dropout; no es sinónimo de no_grad. El control SFT no llama a ningún profesor. Los costes del forward del profesor en KD, cálculo/backward del modelo optimizado, optimización y validación se conservan por separado.
Cada rama compara una interrupción después de una época, dos épocas continuas y la continuación de la época restante. Los pesos, estados Adam, scheduler, permutación, RNG e historiales numéricos deben ser exactamente iguales. Las duraciones quedan fuera de la comparación de estado. Los estados locales se comprimen sin pérdida, nivel gzip 1; los pesos finales también se guardan y recargan en Safetensors. El total con el ejercicio conserva 15 estados de reanudación y cinco pesos finales. Estos pickles Torch se producen localmente; no los sustituya por archivos de terceros.
Generar libremente y medir el alcance real
El profesor adaptado, el estudiante inicial, SFT y KD generan cada uno 300 respuestas de prueba, sin prefijo gold. Las dos variantes del ejercicio hacen lo mismo: 1800 respuestas de prueba, más los warm-ups identificados por separado. La generación es codiciosa sobre todo el vocabulario, sin caché KV, sin esquema forzado y con un límite de 96 tokens nuevos. Se conservan los ID, textos brutos, EOS o corte, log-probabilidades, duraciones y dos sondas de logits completos.
El servicio del estudiante no tiene físicamente ningún montaje del profesor ni de sus adaptadores. La recarga del estudiante usa su configuración, su tokenizer y únicamente sus pesos aprendidos. La inferencia no monta ninguna etiqueta. La evaluación reúne después las salidas y los objetivos, sin modificar los pesos.
El analizador distingue JSON analizable, esquema estricto y respuesta completa: una salida cortada falla incluso si contiene un objeto bien formado. Las claves duplicadas, NaN, campos adicionales, Markdown e ID desconocidos no se reparan. Los errores permanecen en el denominador de los 300 casos para exact match y cada campo. Las tasas de entrega falsa o de campo rellenado erróneamente se calculan sobre las respuestas completas válidas e indican también las poblaciones elegibles totales y la cobertura. Una salida no válida no se certifica como segura.
Cada modelo recibe el mismo mensaje de entrenamiento de warm-up, fuera de prueba. Se informa de la latencia por solicitud y por token, longitudes con EOS, mediana y p95 empírico (rango 285 de 300). Las dos primeras solicitudes incluyen el coste de las sondas de logits, señalado en el informe. La carga y sus comprobaciones de integridad permanecen separadas. Se miden parámetros únicos, dtypes, bytes lógicos y tamaño real de los archivos servidos. El profesor BF16 con adaptadores FP32 y el estudiante enteramente FP32 también difieren en precisión: una relación de parámetros no es una relación garantizada de memoria o velocidad.
Ejercicio: desplazar alpha y T
El SFT principal es el verdadero comparador alpha=0. Dos estudiantes adicionales parten de la misma base: alpha=1/T=2 y alpha=1/T=4, cada uno con 350 actualizaciones y su propia continuación exacta. No se renombra simplemente ningún checkpoint. Los datos, el profesor y los resultados principales están protegidos por huellas; el profesor no se reentrena para el ejercicio.
Alpha=1 imita exclusivamente al profesor en los prefijos gold, incluidos sus errores. Aumentar T redistribuye la masa y los gradientes; no revela la verdad del campo delivered. Compare los mismos 300 eventos, errores y costes. El programa no impone ninguna mejora ni temperatura óptima.
Límites
Una semilla, un corpus sintético regular y un pequeño conjunto de tipos no constituyen una evaluación de negocio. La disminución de KL en prefijos gold no demuestra una generación libre mejor. El profesor puede transferir sus errores. Las mediciones locales sin carga concurrente controlada no son un compromiso de rendimiento del servicio. El TP no envía ningún paquete, mensaje ni acción externa.
Construir los eventos y verificar los activos
Alineación de datos y tokenizers
Error frecuente: comparar solo vocab_size ignora una permutación de ID; comparar los diccionarios y los ejemplos codificados.
Invariante: familias disjuntas e informe de alineación completo; parada clara en caso de mismatch.
Archivo 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
Formar prefijos realmente idénticos
Teacher forcing y máscara de respuesta
Error frecuente: una plantilla añadida dos veces en el profesor desplaza las posiciones; aplicar hash a los tokens del batch común.
Invariante: ID, máscaras y desplazamientos idénticos para teacher/student.
Archivo 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
Adaptar el profesor y medir su calidad
Profesor útil, no oráculo
Error frecuente: usar etiquetas de prueba para elegir el mejor teacher invalida la transferencia; selección solo en validación.
Invariante: adaptadores reales entrenados, profesor congelado y resultados iniciales conservados.
Archivo 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)
Codificar la KL y entrenar los dos estudiantes
Objetivos suaves y gradiente
Error frecuente: torch.kl_div tomado en el orden incorrecto invierte la divergencia; probar distribuciones asimétricas.
Invariante: KL analítica correcta, profesor inmutable, dos estudiantes realmente actualizados y reanudación posible.
Archivo 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
Servir al estudiante sin depender del profesor
Calidad libre y coste de inferencia
Error frecuente: una generación más corta parece más rápida; publicar longitudes y coste por solicitud y token.
Invariante: ninguna solicitud al profesor durante la inferencia del estudiante; 300 ID por rama.
Archivo 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()}
Explorar la ponderación de la transferencia
Papel de alpha y T
Error frecuente: reutilizar el checkpoint SFT renombrándolo alpha=1 no es una destilación; exigir el rastro de las actualizaciones.
Invariante: alpha=0 coincide con SFT y cada variante tiene su propia identidad.
Archivo 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']}
Fuentes y alcance de referencia
Hinton et al., Distilling the Knowledge in a Neural Network motiva la temperatura y el factor T². PyTorch 2.8 kl_div proporciona la referencia de fórmula usada por las pruebas. TRL GKD 0.23.1 describe, entre otros aspectos, prefijos generados por el estudiante: nuestro ejercicio emplea prefijos gold fuera de política y no afirma equivalencia con este trainer. Los modelos Qwen y las dependencias tienen sus revisiones y hashes en los locks públicos.
