Blame for task_finetuning/train_tasks.py

"""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
        # Added linear classifier is the requested GPT-1 task head, not a backbone replacement.
        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)
        # GPT-1's multiple-choice head shares a dropout feature mask across options.
        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)
            # Average within each candidate sequence, then across all candidates.
            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():
            # Preparation script may use the compact .pt suffix.
            path = ROOT / 'data' / f'{prefix}_{split}.pt'
        cache = torch.load(path, map_location='cpu', weights_only=False, mmap=True)
        # Both conventional names are accepted; contents stay unchanged.
        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')
    # AMP-COMPLEX: native weights/activations are real; native RoPE explicitly casts its complex rotation to float32.
    with torch.inference_mode(), torch.autocast('cuda', dtype=torch.bfloat16):  # lint: ok AMP-COMPLEX real native weights; RoPE explicitly uses float32 complex rotation
        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'])
    # Reuse three existing padding slots; real GPT-2 IDs and tokenizer are unchanged.
    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
        # Once-per-evaluation diagnostic over twelve oscillator tensors, not Python-built training batches.
        omegas = torch.cat([p.detach().reshape(-1) for name,p in model.named_parameters() if name.endswith('.omega')])  # lint: ok STACK-COMPREHENSION evaluation statistics over twelve oscillator tensors
        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)
                # AMP-COMPLEX: real-valued model; native rotary complex operations remain explicitly float32.
                with torch.autocast('cuda',dtype=torch.bfloat16):  # lint: ok AMP-COMPLEX real native weights; RoPE explicitly uses float32 complex rotation
                    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:
                # One host transfer per logging interval; hot microbatches remain on CUDA.
                values = (accumulated/logged_steps).cpu().tolist()  # lint: ok SYNC-IN-LOOP once per logging interval
                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')
    # Final held-out scores use the best validation-selected checkpoint.
    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()