Voltline's picture
Release VimeML V4.2 epoch20 FP32 and Core ML INT8 (GPL-2.0)
c83c617 verified
Raw History Blame Contribute Delete
24.2 kB
"""V4.1 random-initialized LM/chat/conversion pretraining, with exact epoch quotas."""
import copy
import hashlib
import json
import math
from pathlib import Path
import random
import signal
import time
import torch
from vimeml.training.data import SentenceWindowDataset,file_sha,write_json
from vimeml.training.data_mixture import mixture_dataset
from vimeml.training.v41_data import IndexedConversionWindows,conversion_quota
from vimeml.training.model_v4 import TinyGPTV4,GPTV4Config
from vimeml.training.train import optimizer_for,learning_rate,atomic_checkpoint,rng_state,restore_rng,resolve_precision
from vimeml.posttraining.data import ConversionDataset,ConversionCodec,TOKENIZER_SHA
from vimeml.posttraining.streams import ConversionWindows,CommittedStream,loader_for
from vimeml.posttraining.runtime import update,restore_checkpoint,TASK_CONTRACT
from vimeml.posttraining.train import validate,checkpoint_payload
from vimeml.posttraining.infer import read_snapshot,fp32_evaluation,evaluate_snapshot
ROOT=Path(__file__).resolve().parents[3]
def read(path):
return json.loads(Path(path).read_text('utf-8'))
def build_plan(config,root=ROOT,verify_files=True):
root=Path(root)
if config.get('stage')!='v4.1' or config.get('architecture')!='tiny_gpt_v4' or config.get('initialization')!={'mode':'random'}:
raise ValueError('V4.1 requires explicit random initialization; no source checkpoint')
GPTV4Config(**config['model'])
s=dict(config['training'])
if s['epochs']!=16 or s['lm_weight']!=.8 or s['conversion_weight']!=.2:
raise ValueError('Use the frozen 16-epoch V4.1 joint objective')
for key in ('lm_batch','conversion_batch','lm_microbatch','conversion_microbatch','epochs','num_workers',
'cpu_threads','log_every','eval_every','checkpoint_every','validation_batch','bucket_multiplier','warmup_steps'):
if type(s[key]) is not int or s[key]<1:
raise ValueError(f'Invalid training integer: {key}')
if type(s['seed']) is not int or not 0<=s['seed']<2**63:
raise ValueError('Invalid random seed')
for key in ('learning_rate','min_learning_rate','lm_weight','conversion_weight','beta1','beta2','weight_decay','grad_clip'):
if not math.isfinite(s[key]) or s[key]<=0:
raise ValueError(f'Invalid training scalar: {key}')
if (s['device']!='cuda' or s['precision']!='bf16' or not config['runtime']['compile_backbone']
or not 0<s['min_learning_rate']<=s['learning_rate'] or s['beta1']>=1 or s['beta2']>=1
or not 0<=s['prefix_crop_probability']<1 or s['conversion_microbatch']<s['conversion_batch']
or s['lm_microbatch']>s['lm_batch']):
raise ValueError('V4.1 CUDA/precision/optimizer/microbatch contract mismatch')
mix_dir=root/config['data_mixture']['directory']
mix=read(mix_dir/'manifest.json')
request=mix['request']
if (mix['status']!='complete' or request['epochs']!=16 or len(mix['epochs'])!=16
or request['seed']!=s['seed'] or request['crop_probability']!=s['prefix_crop_probability']
or request['crop_minimum']!=8 or request['chat_fraction']!=.05):
raise ValueError('Frozen mixture differs from training configuration')
if file_sha(mix_dir/'manifest.json')!=config['dataset_sha256']:
raise ValueError('V4.1 dataset manifest changed')
required={}
signatures={'dataset':config['dataset_sha256']}
for bank in request['joined_stores']:
folder,index=root/bank['token_dir'],root/bank['index_dir']
manifest=read(folder/'manifest.json')
if file_sha(folder/'manifest.json')!=bank['token_manifest_sha256'] or file_sha(index/'manifest.json')!=bank['window_manifest_sha256']:
raise ValueError('Changed LM/chat bank manifest')
if manifest['tokenizer_model_sha256']!=TOKENIZER_SHA or manifest['vocab_size']!=8192:
raise ValueError('LM/chat vocabulary mismatch')
windows=read(index/'manifest.json')
for split in ('train','validation'):
for suffix in ('tokens.bin','offsets.bin','sources.bin'):
required[folder/f'{split}.{suffix}']=manifest['output_sha256'][f'{split}.{suffix}']
required[index/f'{split}.windows.npz']=windows['splits'][split]['index_sha256']
required[folder/'manifest.json']=bank['token_manifest_sha256']
required[index/'manifest.json']=bank['window_manifest_sha256']
primary=request['joined_stores'][0]
if config['token_dir']!=primary['token_dir'] or config['index_dir']!=primary['index_dir']:
raise ValueError('Primary model tokenizer/LM bank mismatch')
signatures.update(tokens=primary['token_manifest_sha256'],windows=primary['window_manifest_sha256'])
conversion=root/config['conversion_dir']
cv=read(conversion/'manifest.json')
acceptance=read(conversion/'acceptance.json')
signatures['conversion']=file_sha(conversion/'manifest.json')
if signatures['conversion']!=request['conversion_manifest_sha256'] or cv['contract']['tokenizer_sha256']!=TOKENIZER_SHA or acceptance['manifest_sha256']!=signatures['conversion'] or acceptance['status']!='pass':
raise ValueError('Conversion bank identity mismatch')
for name,digest in cv['output_sha256'].items():
if name.endswith('.bin'):
required[conversion/name]=digest
required[conversion/'manifest.json']=signatures['conversion']
required[conversion/'acceptance.json']=file_sha(conversion/'acceptance.json')
for name in ('old_keep','conversion_keep'):
required[root/request[name]]=request[name+'_sha256']
for record in mix['epochs']:
required[mix_dir/record['file']]=record['sha256']
required[mix_dir/'manifest.json']=config['dataset_sha256']
required[mix_dir/'request.json']=file_sha(mix_dir/'request.json')
tokenizer=root/config['tokenizer_dir']/'tokenizer.model'
if file_sha(tokenizer)!=TOKENIZER_SHA:
raise ValueError('Tokenizer changed')
required[tokenizer]=TOKENIZER_SHA
standard=read_snapshot('standard-development',root)
required[root/'outputs/evaluation-v4-20261015/inputs/standard-development.json']=file_sha(root/'outputs/evaluation-v4-20261015/inputs/standard-development.json')
code={p.relative_to(root).as_posix():file_sha(p) for p in sorted((root/'src/vimeml').rglob('*.py'))}
code['scripts/training/train_v41.py']=file_sha(root/'scripts/training/train_v41.py')
signatures['code']=hashlib.sha256(json.dumps(code,sort_keys=True).encode()).hexdigest()
if verify_files:
for path,digest in required.items():
if file_sha(path)!=digest:
raise ValueError('Changed training dependency: '+str(path))
s['updates_per_epoch']=[math.ceil(r['windows']/s['lm_batch']) for r in mix['epochs']]
s['max_steps']=sum(s['updates_per_epoch'])
if request['conversion_examples']<max(s['updates_per_epoch']):
raise ValueError('Not enough conversion examples for the proportional update schedule')
return {'stage':'v4.1','parameters':15_070_848,'settings':s,'signatures':signatures,'code_sha256':code,
'required_files':{p.relative_to(root).as_posix():h for p,h in required.items()},
'conversion_examples':request['conversion_examples'],'request':request,
'standard_cases':len(standard),'task_contract':TASK_CONTRACT}
def initial_state():
return {'epoch':0,'epoch_batch':0,'step':0,'lm_cursor':{'epoch':0,'batch':0,'offset':0},
'conversion_cursor':{'epoch':0,'batch':0,'offset':0},'epoch_results':[],
'best_conversion':None,'best_lm':None,'best_rerank_key':None,'best_ordinary_key':None,
'training_seconds':0.,'totals':{n:{k:0 for k in ('examples','input_tokens','supervised_tokens','padded_positions')}
for n in ('lm','conversion')}}
def validate_state(state,plan):
e,b=state['epoch'],state['epoch_batch']
updates=plan['settings']['updates_per_epoch']
if not 0<=e<=len(updates) or b<0 or (e==len(updates) and b) or (e<len(updates) and b>updates[e]):
raise ValueError('Invalid committed LM epoch cursor')
if state['step']!=sum(updates[:e])+b:
raise ValueError('Step does not match committed epoch updates')
expected=e*plan['conversion_examples']+(b*plan['conversion_examples']//updates[e] if e<len(updates) else 0)
if state['totals']['conversion']['examples']!=expected:
raise ValueError('Conversion exposure is not exactly proportional to committed LM updates')
if state['lm_cursor']!={'epoch':e,'batch':b,'offset':0}:
raise ValueError('LM stream cursor mismatch')
epochs=[r['epoch'] for r in state['epoch_results']]
allowed=list(range(1,e+1))
if e<len(updates) and b==updates[e] and len(epochs)==e+1:
allowed.append(e+1)
if epochs!=allowed:
raise ValueError('Invalid completed/pending validation history')
def progress(state,plan,status='running',latest=None):
updates=plan['settings']['updates_per_epoch']
e=state['epoch']
fraction=state['epoch_batch']/updates[e] if e<len(updates) else 0
return {'status':status,'step':state['step'],'max_steps':plan['settings']['max_steps'],
'completed_lm_epochs':e+fraction,'completed_conversion_epochs':state['totals']['conversion']['examples']/plan['conversion_examples'],
'lm_cursor':state['lm_cursor'],'conversion_cursor':state['conversion_cursor'],
'totals':state['totals'],'training_seconds':state['training_seconds'],'latest_update':latest,
'full_validation':state['epoch_results'][-1] if state['epoch_results'] else None,
'best_conversion_loss':state['best_conversion'],'best_lm_bpc':state['best_lm'],
'best_rerank_key':state['best_rerank_key'],'best_ordinary_key':state['best_ordinary_key']}
def combine_conversion(parts):
"""Rejoin a cursor boundary into one small conversion forward/backward."""
if len(parts)==1:
return parts
lengths=torch.cat([p['lengths'] for p in parts])
width=math.ceil(int(lengths.max())/8)*8
pinned=parts[0]['input_ids'].is_pinned()
ids=torch.zeros((len(lengths),width),dtype=torch.long,pin_memory=pinned)
labels=torch.full_like(ids,-100)
start=0
for part in parts:
rows=len(part['lengths'])
columns=min(width,part['input_ids'].shape[1])
ids[start:start+rows,:columns]=part['input_ids'][:,:columns]
labels[start:start+rows,:columns]=part['labels'][:,:columns]
start+=rows
return [{'input_ids':ids,'labels':labels,'lengths':lengths}]
@torch.inference_mode()
def evaluate_ordinary(model,codec,rows,destination):
from vimeml.training.infer_v4 import JapaneseLMV4
from vimeml.training.evaluate_ime import score_candidates
from vimeml.benchmarks.evaluate_ajimee import eligibility
lm=JapaneseLMV4.__new__(JapaneseLMV4)
lm.model,lm.processor,lm.device=model,codec.sp,next(model.parameters()).device
lm.special={'pad':0,'unk':1,'bos':2,'eos':3}
records=[]
for row in rows:
order=row['orders']['azookey']
reason=eligibility(lm,row['left_context'],order)
scores=None
if reason:
ranked=order
else:
scored=score_candidates(lm,row['left_context'],order)
scores=[r['log_probability_sum'] for r in scored['candidates']]
ranked=[order[i] for i in sorted(range(len(order)),key=lambda i:(-scores[i],i))]
records.append({'id':row['id'],'correct':ranked[0] in row['answers'],'ranked':ranked,'scores':scores,'fallback':reason})
destination.mkdir(parents=True,exist_ok=True)
with (destination/'scores.jsonl').open('w',encoding='utf-8') as f:
for row in records:
f.write(json.dumps(row,ensure_ascii=False)+'\n')
result={'correct':sum(r['correct'] for r in records),'total':len(rows),'fallbacks':sum(r['fallback'] is not None for r in records),'precision':'fp32','score':'unchanged ordinary LM contextual suffix sum'}
write_json(destination/'metrics.json',result)
return result
def run(config,resume=False,root=ROOT):
from torch.utils.tensorboard import SummaryWriter
root=Path(root)
output,log_dir=root/config['output_dir'],root/config['log_dir']
if not resume and ((output.exists() and any(output.iterdir())) or (log_dir.exists() and any(log_dir.iterdir()))):
raise ValueError('Fresh V4.1 requires empty output and log directories')
print('Verifying frozen V4.1 union, tokenizer and per-epoch maps...',flush=True)
plan=build_plan(config,root,verify_files=True)
s=plan['settings']
torch.set_num_threads(s['cpu_threads'])
random.seed(s['seed'])
torch.manual_seed(s['seed'])
device=torch.device(s['device'])
if not torch.cuda.is_available():
raise ValueError('A GPU instance is required for formal V4.1 training')
precision=resolve_precision(device,s['precision'])
torch.set_float32_matmul_precision('high')
model=TinyGPTV4(GPTV4Config(**config['model'])).to(device)
optimizer=optimizer_for(model,s,device)
state=initial_state()
initialization={'mode':'random','seed':s['seed'],'source_checkpoint':None,'optimizer':'fresh_AdamW'}
restored=None
if resume:
saved=torch.load(output/'last.pt',map_location='cpu',weights_only=True)
state,restored,initialization=restore_checkpoint(saved,model,optimizer,config,plan['signatures'],precision)
validate_state(state,plan)
del saved
if state['epoch']==s['epochs']:
print('V4.1 is already complete',flush=True)
return progress(state,plan,'complete')
output.mkdir(parents=True,exist_ok=True)
write_json(output/'config.json',config)
write_json(output/'plan.json',plan)
from vimeml.training.runtime_v2 import compile_backbone
compilation=compile_backbone(model)
write_json(output/'environment.json',{'torch':str(torch.__version__),'cuda':torch.version.cuda,
'gpu':torch.cuda.get_device_name(),'precision':precision,'parameters':model.parameter_count(),
'initialization':initialization,'compilation':compilation,'validation_precision':'primitive_FP32'})
lm_train=mixture_dataset(config,root)
conv_train=IndexedConversionWindows(ConversionDataset(root/config['conversion_dir'],'train'),root/plan['request']['conversion_keep'])
conv_val=ConversionWindows(ConversionDataset(root/config['conversion_dir'],'validation'))
lm_val=SentenceWindowDataset(root/config['token_dir'],root/config['index_dir'],'validation')
chat_bank=plan['request']['joined_stores'][1]
chat_val=SentenceWindowDataset(root/chat_bank['token_dir'],root/chat_bank['index_dir'],'validation')
options={'workers':s['num_workers'],'multiplier':s['bucket_multiplier'],'pin_memory':True}
lm_loader=loader_for(lm_train,s['lm_batch'],seed=s['seed'],**options)
conv_loader=loader_for(conv_train,s['conversion_batch'],seed=s['seed']+1,**options)
conv_stream=CommittedStream(conv_loader,state['conversion_cursor'])
def val_loader(data):
return loader_for(data,s['validation_batch'],workers=0,seed=s['seed'],shuffle=False,pin_memory=True)
validation=val_loader(conv_val)
old_validation,chat_validation=val_loader(lm_val),val_loader(chat_val)
subset=sorted(random.Random(s['seed']+1000019).sample(range(len(conv_val)),min(4096,len(conv_val))))
subset_loader=loader_for(conv_val,s['validation_batch'],workers=0,seed=s['seed'],indices=subset,shuffle=False,pin_memory=True)
codec=ConversionCodec(root/config['tokenizer_dir']/'tokenizer.model')
standard=read_snapshot('standard-development',root)
writer=SummaryWriter(str(log_dir),purge_step=state['step']+1 if resume else None,flush_secs=5)
metrics=(output/'metrics.jsonl').open('a',encoding='utf-8',buffering=1)
requested=False
handlers={}
def request_stop(signum,frame):
nonlocal requested
if requested:
raise KeyboardInterrupt('Second interrupt; keep last committed checkpoint')
requested=True
print('Saving after current committed update/validation',flush=True)
for sig in (signal.SIGINT,signal.SIGTERM):
handlers[sig]=signal.getsignal(sig)
signal.signal(sig,request_stop)
def save(name='last'):
payload=checkpoint_payload(model,optimizer,state,config,plan,precision,initialization)
payload['posttraining']['stage']='v4.1'
payload['training_regime']='random_initialized_joint_pretraining'
atomic_checkpoint(output/f'{name}.pt',payload)
if restored is not None:
restore_rng(restored)
write_json(output/'progress.json',progress(state,plan,'running' if resume else 'initializing'))
try:
if not resume:
save()
while state['epoch']<s['epochs'] and not requested:
e=state['epoch']
lm_stream=CommittedStream(lm_loader,state['lm_cursor'])
updates=s['updates_per_epoch'][e]
model.train()
while state['epoch_batch']<updates and not requested:
began=time.perf_counter()
cursor=state['epoch_batch']
lm_batch=lm_stream.next_batch()
nconv=conversion_quota(cursor,updates,plan['conversion_examples'])
conversion=combine_conversion(conv_stream.next_group(nconv))
rate=learning_rate(state['step']+1,s)
for group in optimizer.param_groups:
group['lr']=rate
result=update(model,optimizer,conversion,[lm_batch],s,device,precision)
torch.cuda.synchronize(device)
state['step']+=1
state['epoch_batch']+=1
state['lm_cursor']=lm_stream.state_dict()
state['conversion_cursor']=conv_stream.state_dict()
seconds=time.perf_counter()-began
state['training_seconds']+=seconds
for task in ('lm','conversion'):
for field in state['totals'][task]:
state['totals'][task][field]+=result[field][task]
if state['step']%s['log_every']==0:
latest={'step':state['step'],'learning_rate':rate,'seconds':seconds,**result}
metrics.write(json.dumps(latest)+'\n')
for task in ('lm','conversion'):
writer.add_scalar(f'{task}/train_loss',result['losses'][task],state['step'])
writer.add_scalar(f'{task}/supervised_tokens',state['totals'][task]['supervised_tokens'],state['step'])
writer.add_scalar(f'{task}/input_tokens',state['totals'][task]['input_tokens'],state['step'])
writer.add_scalar('optimizer/learning_rate',rate,state['step'])
writer.add_scalar('optimizer/grad_norm',result['grad_norm'],state['step'])
write_json(output/'progress.json',progress(state,plan,latest=latest))
print(f"step={state['step']}/{s['max_steps']} epoch={e+1} lm={result['losses']['lm']:.4f} conv={result['losses']['conversion']:.4f} seconds={seconds:.3f}",flush=True)
if state['step']%s['checkpoint_every']==0:
save()
if state['step']%s['eval_every']==0 and state['epoch_batch']<updates:
with fp32_evaluation(model):
subset_result=validate(model,subset_loader,device)
writer.add_scalar('conversion/fixed_subset_loss',subset_result['loss'],state['step'])
write_json(output/'subset-validation.json',{'step':state['step'],**subset_result})
if requested:
break
if not any(r['epoch']==e+1 for r in state['epoch_results']):
save()
write_json(output/'progress.json',progress(state,plan,'validating'))
destination=output/'epoch-evaluation'/f'epoch-{e+1}'
destination.mkdir(parents=True,exist_ok=True)
print(f'[epoch {e+1}] full old LM, chat, conversion and frozen Standard',flush=True)
before=rng_state()
with fp32_evaluation(model):
cv=validate(model,validation,device,conv_val.data.manifest['splits']['validation']['supervised_tokens_full_reading'])
old_stats=read(root/config['token_dir']/'manifest.json')['splits']['validation']
chat_stats=read(root/chat_bank['token_dir']/'manifest.json')['splits']['validation']
lv=validate(model,old_validation,device,old_stats['prediction_pairs'],old_stats['characters'])
chat=validate(model,chat_validation,device,chat_stats['prediction_pairs'],chat_stats['characters'])
ime=evaluate_snapshot(model,codec,standard,destination/'standard-development')
ordinary=evaluate_ordinary(model,codec,standard,destination/'standard-ordinary')
restore_rng(before)
report={'epoch':e+1,'step':state['step'],'conversion':cv,'lm':lv,'chat':chat,'ime':ime,'ordinary_ime':ordinary}
roles=[]
if state['best_conversion'] is None or cv['loss']<state['best_conversion']:
state['best_conversion']=cv['loss'];roles.append('best-conversion')
if state['best_lm'] is None or lv['bpc']<state['best_lm']:
state['best_lm']=lv['bpc'];roles.append('best-lm')
key=[ime['correct'],-cv['loss'],-state['step']]
if state['best_rerank_key'] is None or key>state['best_rerank_key']:
state['best_rerank_key']=key;roles+=['best-rerank','best']
ordinary_key=[ordinary['correct'],-lv['bpc'],-state['step']]
if state['best_ordinary_key'] is None or ordinary_key>state['best_ordinary_key']:
state['best_ordinary_key']=ordinary_key;roles.append('best-ordinary')
state['epoch_results'].append(report)
write_json(destination/'summary.json',report)
write_json(output/'full-validation.json',report)
for tag,value in [('conversion/full_validation_loss',cv['loss']),('lm/full_validation_bpc',lv['bpc']),
('chat/full_validation_bpc',chat['bpc']),('ime/conditional_correct',ime['correct']),('ime/ordinary_correct',ordinary['correct'])]:
writer.add_scalar(tag,value,state['step'])
for role in roles:
save(role)
save(f'epoch-{e+1}')
print(f"[epoch {e+1}] LM BPC={lv['bpc']:.6f} chat BPC={chat['bpc']:.6f} conv={cv['loss']:.6f} conditional={ime['correct']}/728 ordinary={ordinary['correct']}/728",flush=True)
state['epoch']+=1
state['epoch_batch']=0
state['lm_cursor']={'epoch':state['epoch'],'batch':0,'offset':0}
validate_state(state,plan)
save()
terminal='complete' if state['epoch']==s['epochs'] else 'interrupted'
save()
result=progress(state,plan,terminal)
result['checkpoints']={n:str(output/f'{n}.pt') for n in ('best','best-rerank','best-ordinary','best-lm','best-conversion','last') if (output/f'{n}.pt').exists()}
result['peak_allocated_MiB']=torch.cuda.max_memory_allocated(device)/2**20
result['peak_reserved_MiB']=torch.cuda.max_memory_reserved(device)/2**20
write_json(output/'summary.json',result)
write_json(output/'progress.json',result)
return result
except BaseException as error:
write_json(output/'failure.json',{'status':'failed','active':True,'error_type':type(error).__name__,
'error':str(error),'last_committed_in_memory_step':state['step'],'safe_resume_checkpoint':str(output/'last.pt')})
raise
finally:
writer.close();metrics.close()
for sig,handler in handlers.items():
signal.signal(sig,handler)
lm_train.close()
for data in (conv_train.data,conv_val.data,lm_val,chat_val):
data.close()