| """Full-model GPT-1-style task fine-tuning of original WiggleGPT.
|
|
|
| Native model_bio.py is imported unchanged. A final-state classifier is added;
|
| its language-model objective uses the native tied projection. No CPU forwards.
|
| """
|
| from __future__ import annotations
|
|
|
| import argparse
|
| from collections import Counter
|
| from datetime import datetime
|
| from zoneinfo import ZoneInfo
|
| import gc
|
| import hashlib
|
| import json
|
| import math
|
| import os
|
| from pathlib import Path
|
| import random
|
| import shutil
|
| import signal
|
| import sys
|
| import time
|
|
|
| ROOT = Path(__file__).resolve().parent
|
| NATIVE = ROOT / 'native'
|
| sys.path.insert(0, str(NATIVE))
|
| from task_runtime import load_model
|
| import model_bio
|
| import torch
|
| from torch import nn
|
| from torch.nn import functional as F
|
|
|
| STOP = False
|
| MARKERS = {'start': 50257, 'delimiter': 50258, 'classify': 50259, 'pad': 50256}
|
|
|
| def stop_requested(signum, frame):
|
| global STOP
|
| STOP = True
|
| print('Stop requested; finishing the update and saving latest.pt.', flush=True)
|
|
|
| def digest(path):
|
| with Path(path).open('rb') as stream:
|
| return hashlib.file_digest(stream, 'sha256').hexdigest()
|
|
|
| class TaskModel(nn.Module):
|
| def __init__(self, base, task, dropout):
|
| super().__init__()
|
| self.base = base
|
| self.task = task
|
| self.dropout = dropout
|
|
|
| with torch.device('cuda:0'):
|
| self.classifier = nn.Linear(base.config.n_embd, 2 if task == 'sst2' else 1)
|
| nn.init.normal_(self.classifier.weight, mean=0.0, std=0.02)
|
| nn.init.zeros_(self.classifier.bias)
|
| self.hidden = None
|
| self.base.transformer.ln_f.register_forward_hook(self.capture_hidden)
|
|
|
| def capture_hidden(self, module, args, output):
|
| self.hidden = output
|
|
|
| def forward(self, tokens, lengths, auxiliary=False):
|
| batch, choices, width = tokens.shape
|
| flat = tokens.reshape(batch * choices, width)
|
| self.base(flat)
|
| hidden = self.hidden
|
| self.hidden = None
|
| positions = lengths.reshape(-1) - 1
|
| features = hidden[torch.arange(batch * choices, device=flat.device), positions]
|
| features = features.reshape(batch, choices, -1)
|
|
|
| if self.training and self.dropout:
|
| mask = F.dropout(torch.ones((batch, 1, features.shape[-1]), device=features.device,
|
| dtype=features.dtype), p=self.dropout, training=True)
|
| features = features * mask
|
| scores = self.classifier(features).squeeze(1) if self.task == 'sst2' else self.classifier(features).squeeze(-1)
|
| lm_loss = None
|
| if auxiliary:
|
| logits = self.base.lm_head(hidden[:, :-1])
|
| targets = flat[:, 1:].long()
|
| per_token = F.cross_entropy(logits.float().reshape(-1, logits.shape[-1]),
|
| targets.reshape(-1), reduction='none').view_as(targets)
|
| mask = torch.arange(width - 1, device=flat.device)[None, :] < (lengths.reshape(-1, 1) - 1)
|
|
|
| lm_loss = ((per_token * mask).sum(-1) / mask.sum(-1).clamp_min(1)).mean()
|
| return scores.float(), lm_loss
|
|
|
| def load_data(task, device):
|
| prefix = 'sst2' if task == 'sst2' else 'race_middle'
|
| splits = ['train', 'validation'] + (['test'] if task != 'sst2' else [])
|
| data = {}
|
| for split in splits:
|
| path = ROOT / 'data' / f'{prefix}_{split}_encoded.pt'
|
| if not path.exists():
|
|
|
| path = ROOT / 'data' / f'{prefix}_{split}.pt'
|
| cache = torch.load(path, map_location='cpu', weights_only=False, mmap=True)
|
|
|
| ids = cache.get('input_ids', cache.get('tokens'))
|
| data[split] = {'input_ids': ids.to(device), 'lengths': cache['lengths'].to(device),
|
| 'labels': cache['labels'].to(device), 'cache_path': str(path),
|
| 'metadata': {k:v for k,v in cache.items() if not torch.is_tensor(v)}}
|
| return data
|
|
|
| def select_batch(data, index):
|
| return data['input_ids'][index], data['lengths'][index], data['labels'][index]
|
|
|
| def evaluate(model, data, batch_size, fixed_indices=None):
|
| model.eval()
|
| predictions, losses = [], []
|
| if fixed_indices is None:
|
| fixed_indices = torch.arange(data['labels'].shape[0], device='cuda:0')
|
|
|
| with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16):
|
| for start in range(0, len(fixed_indices), batch_size):
|
| indices = fixed_indices[start:start + batch_size]
|
| x, lengths, labels = select_batch(data, indices)
|
| logits, _ = model(x, lengths)
|
| predictions.append(logits.argmax(-1))
|
| losses.append(F.cross_entropy(logits, labels, reduction='sum'))
|
| predicted = torch.cat(predictions)
|
| labels = data['labels'][fixed_indices]
|
| correct = int((predicted == labels).sum())
|
| counts = torch.bincount(labels, minlength=2 if model.task == 'sst2' else 4).cpu().tolist()
|
| return {'correct': correct, 'total': len(labels), 'accuracy': correct / len(labels),
|
| 'loss': float(torch.stack(losses).sum() / len(labels)),
|
| 'predictions': predicted.cpu().tolist(), 'gold_label_counts': counts,
|
| 'majority_accuracy': max(counts) / len(labels)}
|
|
|
| def optimizer_for(model, config):
|
| decay, no_decay = [], []
|
| protected = ('omega', 'phi', 'baseline', 'norm', 'bias', 'scale', 'phase')
|
| for name, parameter in model.named_parameters():
|
| if parameter.ndim >= 2 and not parameter.is_complex() and not any(word in name for word in protected):
|
| decay.append(parameter)
|
| else:
|
| no_decay.append(parameter)
|
| print(f'Decay: {sum(p.numel() for p in decay):,}; protected/no decay: {sum(p.numel() for p in no_decay):,}', flush=True)
|
| return torch.optim.AdamW([{'params':decay, 'weight_decay':config['weight_decay']},
|
| {'params':no_decay, 'weight_decay':0.0}],
|
| lr=config['learning_rate'], betas=tuple(config['betas']), eps=config['epsilon'], foreach=True)
|
|
|
| def schedule(step, total, config):
|
| warmup = max(1, round(total * config['warmup_fraction']))
|
| if step < warmup:
|
| return config['learning_rate'] * (step + 1) / warmup
|
| return config['learning_rate'] * max(0.0, (total - step) / max(1, total - warmup))
|
|
|
| def atomic_save(path, payload):
|
| temporary = path.with_suffix('.pending.pt')
|
| torch.save(payload, temporary)
|
| temporary.replace(path)
|
|
|
| def rng_state(generator):
|
| return {'python': random.getstate(), 'torch': torch.get_rng_state(),
|
| 'cuda': torch.cuda.get_rng_state(0), 'batch_generator': generator.get_state()}
|
|
|
| def restore_rng(state, generator):
|
| random.setstate(state['python'])
|
| torch.set_rng_state(state['torch'].cpu())
|
| torch.cuda.set_rng_state(state['cuda'].cpu(), 0)
|
| generator.set_state(state['batch_generator'].cpu())
|
|
|
| def full_summary(result):
|
| return {k:v for k,v in result.items() if k != 'predictions'}
|
|
|
| def raw_rows(task, split):
|
| path = ROOT / 'data' / f'{task}_{split}.jsonl'
|
| return [json.loads(line) for line in path.read_text(encoding='utf-8').splitlines()]
|
|
|
| def run_task(task, config, resume=None):
|
| task_config = config['tasks'][task]
|
| random.seed(config['seed'])
|
| torch.manual_seed(config['seed'])
|
| torch.cuda.manual_seed_all(config['seed'])
|
| base, native_config, encoder, metadata = load_model((ROOT / config['starting_checkpoint']).resolve())
|
| free, total = torch.cuda.mem_get_info()
|
| cap = min(int(config['max_process_gib'] * 1024**3),
|
| total - int(config['desktop_reserve_gib'] * 1024**3),
|
| free + torch.cuda.memory_reserved() - int(0.5 * 1024**3))
|
| torch.cuda.set_per_process_memory_fraction(cap / total, 0)
|
| base.config.gradient_checkpointing = task_config['gradient_checkpointing']
|
| metadata['checkpoint_sha256'] = digest((ROOT / config['starting_checkpoint']).resolve())
|
| model = TaskModel(base, task, config['classifier_dropout'])
|
|
|
| with torch.no_grad():
|
| marker_ids = torch.tensor([MARKERS['start'],MARKERS['delimiter'],MARKERS['classify']], device='cuda:0')
|
| base.transformer.wte.weight.index_copy_(0, marker_ids,
|
| torch.randn((3, base.config.n_embd), device='cuda:0') * 0.02)
|
| optimizer = optimizer_for(model, config)
|
| data = load_data(task, 'cuda:0')
|
| n_train = data['train']['labels'].shape[0]
|
| per_epoch = math.ceil(n_train / config['effective_batch_size'])
|
| total_steps = per_epoch * config['epochs']
|
| generator = torch.Generator(device='cuda:0').manual_seed(config['seed'] + 123)
|
| if resume:
|
| checkpoint = torch.load(resume, map_location='cpu', weights_only=False, mmap=True)
|
| model.load_state_dict(checkpoint['model'], strict=True)
|
| optimizer.load_state_dict(checkpoint['optimizer'])
|
| state = checkpoint['state']
|
| order = checkpoint['order'].to('cuda:0') if checkpoint['order'] is not None else None
|
| restore_rng(checkpoint['rng'], generator)
|
| run_dir = Path(resume).resolve().parent
|
| del checkpoint
|
| else:
|
| run_dir = ROOT / 'runs' / f'train_tasks_{datetime.now(ZoneInfo("Europe/London")).strftime("%Y-%m-%d_%H-%M-%S")}_{task}'
|
| run_dir.mkdir(parents=True, exist_ok=False)
|
| state = {'step':0, 'epoch':0, 'offset':0, 'best_accuracy':-1.0, 'best_loss':float('inf'),
|
| 'best_step':0, 'nonfinite_updates':0, 'consecutive_nonfinite':0, 'last_eval_step':-1,
|
| 'elapsed_seconds':0.0, 'lr_scale':1.0}
|
| order = None
|
| (run_dir / 'config.json').write_text(json.dumps(config, indent=2), encoding='utf-8')
|
| shutil.copy2(__file__, run_dir / 'train_tasks.py')
|
| shutil.copy2(model_bio.__file__, run_dir / 'model_bio.py')
|
| provenance = {'model_source_sha256':digest(model_bio.__file__), 'trainer_sha256':digest(__file__),
|
| 'starting_checkpoint':metadata, 'torch':torch.__version__, 'markers':MARKERS,
|
| 'native_backbone_dropout':base.config.dropout,
|
| 'data':{split:d['metadata'] for split,d in data.items()},
|
| 'all_backbone_parameters_train':True}
|
| (run_dir/'provenance.json').write_text(json.dumps(provenance, indent=2), encoding='utf-8')
|
| log = (run_dir / (run_dir.name + '.log')).open('a', encoding='utf-8', buffering=1)
|
| def emit(event):
|
| event.update(task=task)
|
| line = json.dumps(event, ensure_ascii=False)
|
| print(line, flush=True)
|
| log.write(line + '\n')
|
| with (run_dir/'metrics.jsonl').open('a',encoding='utf-8') as stream:
|
| stream.write(line + '\n')
|
| def save(name):
|
| atomic_save(run_dir/name, {'model':model.state_dict(), 'optimizer':optimizer.state_dict(),
|
| 'state':dict(state), 'order':order, 'rng':rng_state(generator), 'config':config,
|
| 'task':task, 'model_args':vars(base.config), 'markers':MARKERS,
|
| 'source_hashes':{'model':digest(model_bio.__file__),'trainer':digest(__file__)}})
|
| evaluation_batch = 16 if task == 'sst2' else 4
|
| heldout = 'validation' if task == 'sst2' else 'test'
|
| before_path = run_dir/'before.json'
|
| start_time = time.perf_counter()
|
| elapsed_prior = state['elapsed_seconds']
|
| emit({'event':'start','step':state['step'],'train_examples':n_train,'epochs':config['epochs'],
|
| 'total_steps':total_steps,'microbatch':task_config['microbatch_size'],
|
| 'effective_batch':config['effective_batch_size'],'memory_cap_gib':round(cap/1024**3,2),
|
| 'trainable_parameters':sum(p.numel() for p in model.parameters() if p.requires_grad),
|
| 'device':torch.cuda.get_device_name(0)})
|
| if not before_path.exists():
|
| before = evaluate(model, data[heldout], evaluation_batch)
|
| before_path.write_text(json.dumps(before,indent=2),encoding='utf-8')
|
| emit({'event':'before','heldout':heldout,**full_summary(before)})
|
| else:
|
| before = json.loads(before_path.read_text(encoding='utf-8'))
|
| train_fixed = torch.arange(min(128,n_train), device='cuda:0')
|
| def evaluate_and_save():
|
| validation = evaluate(model, data['validation'], evaluation_batch)
|
| train_eval = evaluate(model, data['train'], evaluation_batch, train_fixed)
|
| changed = (validation['accuracy'] > state['best_accuracy'] or
|
| (validation['accuracy'] == state['best_accuracy'] and validation['loss'] < state['best_loss']))
|
| if changed:
|
| state['best_accuracy'] = validation['accuracy']
|
| state['best_loss'] = validation['loss']
|
| state['best_step'] = state['step']
|
| state['last_eval_step'] = state['step']
|
| state['elapsed_seconds'] = elapsed_prior + time.perf_counter() - start_time
|
|
|
| omegas = torch.cat([p.detach().reshape(-1) for name,p in model.named_parameters() if name.endswith('.omega')])
|
| emit({'event':'eval','step':state['step'],'epoch':state['epoch'],'offset':state['offset'],
|
| 'train':full_summary(train_eval),'validation':full_summary(validation),
|
| 'best_step':state['best_step'],'omega_std':float(omegas.std(correction=0)),
|
| 'peak_allocated_gib':round(torch.cuda.max_memory_allocated()/1024**3,3),
|
| 'sample':{'gold':int(data['validation']['labels'][0]),'predicted':validation['predictions'][0]},
|
| 'elapsed_seconds':round(state['elapsed_seconds'],1)})
|
| if changed:
|
| save('best.pt')
|
| save('latest.pt')
|
| model.train()
|
| if state['last_eval_step'] < 0:
|
| evaluate_and_save()
|
| model.train()
|
| accumulated = torch.zeros(3,device='cuda:0')
|
| logged_steps = 0
|
| log_start = time.perf_counter()
|
| stopping_cause = 'epochs_complete'
|
| while state['epoch'] < config['epochs']:
|
| if order is None:
|
| order = torch.randperm(n_train,device='cuda:0',generator=generator)
|
| while state['offset'] < n_train:
|
| end = min(n_train,state['offset']+config['effective_batch_size'])
|
| indices = order[state['offset']:end]
|
| effective = len(indices)
|
| lr = schedule(state['step'], total_steps, config) * state['lr_scale']
|
| for group in optimizer.param_groups:
|
| group['lr'] = lr
|
| optimizer.zero_grad(set_to_none=True)
|
| batch_values = torch.zeros(3,device='cuda:0')
|
| for position in range(0,effective,task_config['microbatch_size']):
|
| micro = indices[position:position+task_config['microbatch_size']]
|
| x,lengths,labels = select_batch(data['train'],micro)
|
|
|
| with torch.autocast('cuda',dtype=torch.bfloat16):
|
| logits,lm_loss = model(x,lengths,auxiliary=True)
|
| classification_loss = F.cross_entropy(logits,labels)
|
| loss = classification_loss + config['lm_coefficient'] * lm_loss
|
| weighted = loss * (len(micro)/effective)
|
| weighted.backward()
|
| batch_values += torch.stack((loss.detach(),classification_loss.detach(),lm_loss.detach())) * (len(micro)/effective)
|
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(),config['grad_clip'])
|
| if torch.isfinite(norm):
|
| optimizer.step()
|
| state['consecutive_nonfinite'] = 0
|
| else:
|
| state['nonfinite_updates'] += 1
|
| state['consecutive_nonfinite'] += 1
|
| if state['consecutive_nonfinite'] >= 2:
|
| state['lr_scale'] *= 0.5
|
| emit({'event':'nonfinite_update_skipped','step':state['step']+1,
|
| 'count':state['nonfinite_updates'],'lr_scale':state['lr_scale']})
|
| state['step'] += 1
|
| state['offset'] = end
|
| accumulated += batch_values
|
| logged_steps += 1
|
| if state['step'] == 1 or state['step'] % config['log_interval'] == 0:
|
|
|
| values = (accumulated/logged_steps).cpu().tolist()
|
| state['elapsed_seconds'] = elapsed_prior + time.perf_counter()-start_time
|
| emit({'event':'train','step':state['step'],'total_steps':total_steps,
|
| 'epoch_fraction':round(state['epoch']+end/n_train,4),'loss':values[0],
|
| 'classification_loss':values[1],'lm_loss':values[2],'lr':lr,
|
| 'updates_per_second':round(logged_steps/(time.perf_counter()-log_start),3),
|
| 'peak_allocated_gib':round(torch.cuda.max_memory_allocated()/1024**3,3),
|
| 'nonfinite_updates':state['nonfinite_updates'],'elapsed_seconds':round(state['elapsed_seconds'],1)})
|
| accumulated.zero_(); logged_steps=0;log_start=time.perf_counter()
|
| if state['step'] % task_config['eval_interval'] == 0:
|
| evaluate_and_save()
|
| log_start=time.perf_counter()
|
| elif state['step'] % config['checkpoint_interval'] == 0:
|
| state['elapsed_seconds'] = elapsed_prior + time.perf_counter()-start_time
|
| save('latest.pt')
|
| if STOP or state['consecutive_nonfinite'] >= config['max_consecutive_nonfinite']:
|
| stopping_cause = 'user_interrupt' if STOP else 'nonfinite_limit'
|
| state['elapsed_seconds'] = elapsed_prior + time.perf_counter()-start_time
|
| save('latest.pt')
|
| emit({'event':'stopped','stopping_cause':stopping_cause,'step':state['step']})
|
| log.close()
|
| return False
|
| state['epoch'] += 1
|
| state['offset'] = 0
|
| order = None
|
| if state['last_eval_step'] != state['step']:
|
| evaluate_and_save()
|
| else:
|
| save('latest.pt')
|
|
|
| best = torch.load(run_dir/'best.pt',map_location='cpu',weights_only=False,mmap=True)
|
| model.load_state_dict(best['model'],strict=True)
|
| best_step = best['state']['step']
|
| del best
|
| after = evaluate(model,data[heldout],evaluation_batch)
|
| (run_dir/'after.json').write_text(json.dumps(after,indent=2),encoding='utf-8')
|
| rows = raw_rows(task,heldout)
|
| examples = []
|
| for row_index,(row,pb,pa) in enumerate(zip(rows,before['predictions'],after['predictions'])):
|
| example = {'id':f'{task}/{heldout}/{row_index}','input':row.get('text',row.get('article','')),
|
| 'prompt':row.get('question',''),'options':row.get('options',['negative','positive']),
|
| 'gold':row['label'],'predicted_before':pb,'predicted_after':pa,
|
| 'correct_before':pb==row['label'],'correct_after':pa==row['label']}
|
| examples.append(example)
|
| (run_dir/'heldout_predictions.json').write_text(json.dumps(examples,ensure_ascii=False,indent=2),encoding='utf-8')
|
| probe_count = 100 if task=='sst2' else 40
|
| probe = {'before_correct':sum(e['correct_before'] for e in examples[:probe_count]),
|
| 'after_correct':sum(e['correct_after'] for e in examples[:probe_count]),'total':probe_count}
|
| state['elapsed_seconds'] = elapsed_prior+time.perf_counter()-start_time
|
| summary = {'id':task,'name':'SST-2 sentiment' if task=='sst2' else 'RACE middle-school reading',
|
| 'train_examples':n_train,'validation_examples':len(data['validation']['labels']),
|
| 'test_examples':len(data['test']['labels']) if task!='sst2' else 1821,
|
| 'heldout_split':heldout,'baseline':{'label':'Before task training (same classifier)',**full_summary(before)},
|
| 'finetuned':full_summary(after),'chance_accuracy':0.5 if task=='sst2' else 0.25,
|
| 'majority_accuracy':after['majority_accuracy'],'paired_earlier_probe':probe,
|
| 'training':{'epochs':config['epochs'],'max_steps':total_steps,'steps':state['step'],
|
| 'batch_size':config['effective_batch_size'],'microbatch_size':task_config['microbatch_size'],
|
| 'gradient_accumulation':config['effective_batch_size']//task_config['microbatch_size'],
|
| 'learning_rate':config['learning_rate'],'max_length':config['max_length'],
|
| 'lm_coefficient':config['lm_coefficient'],'stopping_cause':stopping_cause,
|
| 'classifier_dropout':config['classifier_dropout'],'native_backbone_dropout':base.config.dropout,
|
| 'nonfinite_updates':state['nonfinite_updates'],'best_step':best_step,
|
| 'peak_allocated_gib':round(torch.cuda.max_memory_allocated()/1024**3,3),
|
| 'trainable_parameters':sum(p.numel() for p in model.parameters()),
|
| 'head':'768-to-2 linear' if task=='sst2' else 'shared 768-to-1 linear over four answers',
|
| 'checkpoint_selection':'highest validation accuracy, lower validation classification loss breaks ties',
|
| 'details':'Full native backbone trained. Zero decay on oscillator and 1-D parameters. Linear GPT-1 warmup/decay, AdamW, gradient clip1.0. Original backbone dropout retained; classifier dropout0.1. Three unused padding-vocabulary slots initialized for task markers.'},
|
| 'examples':examples,'checkpoint':{'name':str((run_dir/'best.pt').relative_to(ROOT)),
|
| 'sha256':digest(run_dir/'best.pt')},
|
| 'elapsed_seconds':state['elapsed_seconds']}
|
| (run_dir/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2),encoding='utf-8')
|
| emit({'event':'complete','step':state['step'],'best_step':best_step,'stopping_cause':stopping_cause,
|
| 'heldout':heldout,'before':full_summary(before),'after':full_summary(after),
|
| 'paired_probe':probe,'elapsed_seconds':round(state['elapsed_seconds'],1)})
|
| log.close()
|
| del model,base,optimizer,data
|
| gc.collect();torch.cuda.empty_cache()
|
| return True
|
|
|
| def write_comparison():
|
| tasks=[]
|
| for task in ['sst2','race_middle']:
|
| paths=sorted((ROOT/'runs').glob(f'*_{task}/summary.json'),key=lambda path:path.stat().st_mtime)
|
| if not paths:
|
| return
|
| tasks.append(json.loads(paths[-1].read_text(encoding='utf-8')))
|
| data={'date':datetime.now(ZoneInfo('Europe/London')).isoformat(),
|
| 'scope':'Full SST-2 training and validation; full RACE-middle training, validation selection and held-out test. Historical GPT-1 figures use different training coverage and test sets.',
|
| 'method':'Two independent full-model adaptations from original pretrained WiggleGPT. Classifier cross-entropy plus0.5 times token-average auxiliary language-model cross-entropy, with GPT-1-style input markers and linear task heads.',
|
| 'control_method':'Identical task inputs and initialized classifier evaluated before task training; after scores use validation-selected trained checkpoint. Earlier language-model likelihood probes are separately identified.',
|
| 'historical':{'gpt1_race_accuracy':59.0,'gpt1_race_middle_accuracy':62.9,'gpt1_sst2_accuracy':91.3,
|
| 'source_url':'https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf',
|
| 'note':'GPT-1 SST-2 91.3% is GLUE test, while this run reports public validation. GPT-1 RACE-M62.9% used combined RACE-middle/high training, while Wiggle used middle-only. Historical rows give context and are not a matched-model replication.'},
|
| 'tasks':tasks,'elapsed_seconds':sum(t['elapsed_seconds'] for t in tasks)}
|
| (ROOT/'comparison.json').write_text(json.dumps(data,ensure_ascii=False,indent=2),encoding='utf-8')
|
|
|
| def main():
|
| parser=argparse.ArgumentParser(description=__doc__)
|
| parser.add_argument('--task',choices=['all','sst2','race_middle'],default='all')
|
| parser.add_argument('--resume',type=Path)
|
| parser.add_argument('--config',type=Path,default=ROOT/'train_config.json')
|
| parser.add_argument('--starting-checkpoint',type=Path,help='Original pretrained checkpoint; overrides config for a new run.')
|
| args=parser.parse_args()
|
| config=json.loads(args.config.read_text(encoding='utf-8'))
|
| if args.starting_checkpoint:
|
| config['starting_checkpoint']=str(args.starting_checkpoint.expanduser().resolve())
|
| torch.set_num_threads(config['cpu_threads'])
|
| try:
|
| import psutil
|
| psutil.Process().nice(psutil.BELOW_NORMAL_PRIORITY_CLASS if os.name == 'nt' else 10)
|
| except (ImportError, OSError):
|
| print('Optional process priority setting unavailable; continuing with configured CPU threads.',flush=True)
|
| if not torch.cuda.is_available():
|
| raise RuntimeError('CUDA required; CPU model execution is disabled.')
|
| torch.cuda.set_device(0)
|
| torch.set_float32_matmul_precision('high')
|
| signal.signal(signal.SIGINT,stop_requested)
|
| if args.resume:
|
| resume_meta=torch.load(args.resume,map_location='cpu',weights_only=False,mmap=True)
|
| selected=[resume_meta['task']]
|
| del resume_meta
|
| else:
|
| selected=['sst2','race_middle'] if args.task=='all' else [args.task]
|
| for task in selected:
|
| if not run_task(task,config,args.resume):
|
| return
|
| write_comparison()
|
| print('All requested task adaptations complete.',flush=True)
|
|
|
| if __name__=='__main__':
|
| main()
|
|
|
