Download source/src/vimeml/training/joint_v41.py from Voltline/vimeml-tiny-ja-v4.2: direct link, hf CLI and curl.
- Browser
- Download file 24.2 kB
-
https://huggingface.co/Voltline/vimeml-tiny-ja-v4.2/resolve/main/source/src/vimeml/training/joint_v41.py
- Command line
-
hf download hf://Voltline/vimeml-tiny-ja-v4.2/source/src/vimeml/training/joint_v41.py
-
curl -L -o joint_v41.py https://huggingface.co/Voltline/vimeml-tiny-ja-v4.2/resolve/main/source/src/vimeml/training/joint_v41.py
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}] | |
| 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() | |