scripts/train.py

"""Single-GPU ROCm/CUDA LoRA/full fine-tuning on variable-sized decision groups."""
import argparse
import json
import os
import shutil
import time
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import numpy as np
import torch
from transformers import Trainer, TrainerCallback, TrainingArguments, ProgressCallback, set_seed
from transformers.trainer_pt_utils import LengthGroupedSampler
from mugi_decision.progress_callback import StableProgressCallback
from mugi_decision.optimizations import enable_fast_trainable_tokens
from mugi_decision.runtime import gpu_device, load_model, dtype_for, save_json, gpu_metadata
from mugi_decision.frozen import FrozenDecisions, FrozenCollator, model_encoding_signature
from mugi_decision.resume import validate_resume

LORA_TARGETS = ['q_proj','k_proj','v_proj','o_proj','in_proj_qkv','in_proj_z','in_proj_b','in_proj_a','out_proj','gate_proj','up_proj','down_proj']

def main():
    p = argparse.ArgumentParser()
    p.add_argument('--backend', choices=['auto','rocm','cuda'], default=None)
    p.add_argument('--progress', action=argparse.BooleanOptionalAction, default=True)
    p.add_argument('--eta-window', type=int, default=250, help='Sliding window of optimizer updates for stable speed/ETA')
    p.add_argument('--model', default=str(Path(__file__).resolve().parents[1]))
    p.add_argument('--data', default='data/ready', help='Completed offline dataset directory with manifest.json')
    p.add_argument('--no-eval', action='store_true')
    p.add_argument('--initial-eval', action='store_true')
    p.add_argument('--output', default='runs/mugi-lora')
    p.add_argument('--mode', choices=['lora','full'], default='lora')
    p.add_argument('--dtype', choices=['bf16','fp32'], default='bf16')
    p.add_argument('--max-length', type=int, help='Must equal prepared manifest; change lengths with build_training_set.py')
    p.add_argument('--batch-groups', type=int, default=1)
    p.add_argument('--group-by-length', action='store_true', help='Group the same frozen rows by precomputed length to reduce padding')
    p.add_argument('--gradient-accumulation', type=int, default=16)
    p.add_argument('--epoch-snapshots', action='store_true', help='Evaluate and preserve each completed epoch adapter separately')
    p.add_argument('--epochs', type=float, default=1)
    p.add_argument('--max-steps', type=int, default=-1)
    p.add_argument('--lr', type=float, default=1e-4)
    p.add_argument('--optimizer', choices=['adamw_torch','adamw_torch_fused'], default='adamw_torch_fused')
    p.add_argument('--native-token-lookup', action='store_true', help='Use the original PEFT full-vocabulary merge path')
    p.add_argument('--keep-padding-mask', action='store_true', help='Disable the equivalent causal right-padding fast path')
    p.add_argument('--rank', type=int, default=16)
    p.add_argument('--save-steps', type=int, default=500)
    p.add_argument('--logging-steps', type=int, default=10)
    p.add_argument('--seed', type=int, default=42)
    p.add_argument('--best-model-metric', choices=['accuracy_macro_language','accuracy_train_language_weighted'], default='accuracy_macro_language')
    p.add_argument('--init-adapter', help='Load trained LoRA/head/marker weights with fresh optimizer and schedule')
    p.add_argument('--resume', help='A checkpoint directory created by this trainer')
    p.add_argument('--no-gradient-checkpointing', action='store_true')
    p.add_argument('--checkpoint-linear-only', action='store_true', help='Checkpoint 18 DeltaNet decoder layers; leave 6 full-attention layers without recomputation')
    p.add_argument('--checkpoint-linear-stride', type=int, default=1, choices=[1, 2, 3],
                   help='Checkpoint every Nth DeltaNet layer; 2 selects nine layers, 3 selects six')
    p.add_argument('--adaptive-checkpoint-tokens', type=int, default=0,
                   help='Skip recomputation below this padded branch-token count; requires a validated GPU memory budget')
    p.add_argument('--autocast-lora', action='store_true', help='Avoid redundant FP32 activation copies inside BF16 LoRA autocast')
    p.add_argument('--runtime-cumsum', action='store_true', help='Avoid FLA scalar cumsum recompilation for each candidate count')
    p.add_argument('--fused-norms', action='store_true', help='Use CUDA fusion for frozen Qwen RMSNorms without per-length compilation')
    p.add_argument('--allow-performance-resume', action='store_true',
                   help='Allow explicit performance changes while keeping effective batch/data/LR/sampler unchanged')
    p.add_argument('--cpu-threads', type=int, help='Explicit CPU intra-op limit for this process')
    p.add_argument('--cpu-node', type=int, help='Linux CPU affinity to the GPU-local NUMA node')
    p.add_argument('--benchmark-steps', type=int,
                   help='Stop after N updates, only with a separate benchmark-* output directory')
    args = p.parse_args()
    if args.epoch_snapshots and (args.no_eval or args.mode != 'lora'):
        p.error('epoch-snapshots requires LoRA mode and evaluation')
    if args.init_adapter and (args.resume or args.mode != 'lora'):
        p.error('init-adapter requires LoRA mode and cannot be combined with resume')
    if args.batch_groups < 1 or args.gradient_accumulation < 1:
        p.error('batch-groups and gradient-accumulation must be positive')
    if args.checkpoint_linear_stride != 1 and not args.checkpoint_linear_only:
        p.error('checkpoint-linear-stride requires checkpoint-linear-only')
    if args.adaptive_checkpoint_tokens < 0 or (args.adaptive_checkpoint_tokens and
            (not args.checkpoint_linear_only or args.no_gradient_checkpointing)):
        p.error('adaptive checkpointing requires checkpoint-linear-only and a nonnegative threshold')
    if args.benchmark_steps is not None:
        if args.benchmark_steps < 1 or not Path(args.output).name.startswith('benchmark-'):
            p.error('benchmark-steps requires a positive count and a separate benchmark-* output directory')
    if args.cpu_threads is not None:
        if args.cpu_threads < 1:
            p.error('cpu-threads must be positive')
    if args.cpu_node is not None:
        if not hasattr(os, 'sched_setaffinity'):
            p.error('cpu-node requires Linux CPU affinity')
        spans = Path(f'/sys/devices/system/node/node{args.cpu_node}/cpulist').read_text().strip().split(',')
        cpus = set()
        for span in spans:
            bounds = list(map(int, span.split('-')))
            cpus.update(range(bounds[0], bounds[-1]+1))
        cpus &= os.sched_getaffinity(0)
        if not cpus:
            p.error('Requested NUMA node has no allowed CPUs')
        os.sched_setaffinity(0, cpus)
    if args.cpu_threads is not None:
        torch.set_num_threads(args.cpu_threads)
    if args.eta_window < 1:
        p.error('eta-window must be positive')
    if args.checkpoint_linear_only and args.no_gradient_checkpointing:
        p.error('checkpoint-linear-only requires gradient checkpointing')
    output = Path(args.output)
    if output.exists() and any(output.iterdir()) and not args.resume:
        raise SystemExit('Output directory is not empty. Use --resume or a new --output.')
    print('Checking frozen dataset hashes; no runtime tokenization or filtering...', flush=True)
    signature = model_encoding_signature(args.model)
    train_data = FrozenDecisions(args.data, 'train', signature)
    eval_data = None if args.no_eval else FrozenDecisions(args.data, 'validation', signature)
    prepared_length = train_data.manifest['max_length']
    if args.max_length is not None and args.max_length != prepared_length:
        raise ValueError('max-length differs from prepared data. Build a new dataset offline.')
    args.max_length = prepared_length
    performance_changes = {}
    initialization = None
    if args.resume:
        source_config = Path(args.resume).resolve().parent/'run_config.json'
        if not source_config.is_file():
            raise ValueError('Resume checkpoint needs its original run_config.json in the parent directory')
        previous = json.loads(source_config.read_text(encoding='utf-8'))
        initialization = previous.get('initialization')
        performance_changes = validate_resume(previous, vars(args), train_data.manifest_hash,
                                              args.allow_performance_resume)
        if performance_changes:
            print('Explicit performance resume changes: '+json.dumps(performance_changes), flush=True)
    device = gpu_device(backend=args.backend)
    if args.runtime_cumsum:
        if torch.version.hip:
            raise ValueError("Runtime cumsum optimization requires CUDA")
        from mugi_decision.runtime_cumsum import enable_runtime_cumsum
        enable_runtime_cumsum()
    set_seed(args.seed)
    model, tokenizer = load_model(args.model, device, dtype_for(args.dtype))
    if args.init_adapter:
        from peft import PeftModel
        from mugi_decision.initialization import adapter_initialization
        initialization = adapter_initialization(args.init_adapter, args.model, args.rank,
                                                model.config.decision_token_ids, LORA_TARGETS)
        model = PeftModel.from_pretrained(model, args.init_adapter, is_trainable=True,
                                         local_files_only=True)
        print('Weight-only initialization: '+json.dumps(initialization), flush=True)
    elif args.mode == 'lora':
        from peft import LoraConfig, get_peft_model
        model = get_peft_model(model, LoraConfig(r=args.rank, lora_alpha=2*args.rank,
            lora_dropout=0.05, target_modules=LORA_TARGETS, modules_to_save=['score'],
            trainable_token_indices={'backbone.embed_tokens': model.config.decision_token_ids},
            bias='none'))
    fast_token_layers = 0 if args.native_token_lookup else enable_fast_trainable_tokens(model)
    fused_norm_layers = 0
    if args.fused_norms:
        from mugi_decision.cuda_optimizations import enable_fused_norms
        fused_norm_layers = enable_fused_norms(model)
        if not fused_norm_layers:
            raise ValueError('Fused norms require frozen Qwen3.5 norm weights (LoRA mode)')
    autocast_lora_layers = 0
    if args.autocast_lora:
        from mugi_decision.cuda_optimizations import enable_autocast_lora
        autocast_lora_layers = enable_autocast_lora(model)
        if not autocast_lora_layers or args.dtype != 'bf16':
            raise ValueError('Autocast LoRA requires LoRA mode and BF16 training')
    checkpoint_policy = None
    if args.adaptive_checkpoint_tokens:
        from mugi_decision.checkpointing import AdaptiveLinearCheckpointing
        checkpoint_policy = AdaptiveLinearCheckpointing(model, args.checkpoint_linear_stride, args.adaptive_checkpoint_tokens)
    # Trainer must normalize loss per decision group, not per candidate branch/token.
    model.accepts_loss_kwargs = False
    trainable = sum(x.numel() for x in model.parameters() if x.requires_grad)
    print(json.dumps({'train':train_data.summary,'eval':eval_data.summary if eval_data else None,
                      'trainable_parameters':trainable, 'hip':torch.version.hip, 'fast_token_layers':fast_token_layers,
                      'optimizer':args.optimizer,'causal_padding_fast_path':not args.keep_padding_mask}, ensure_ascii=False), flush=True)
    settings = TrainingArguments(output_dir=str(output), per_device_train_batch_size=args.batch_groups,
        per_device_eval_batch_size=args.batch_groups, gradient_accumulation_steps=args.gradient_accumulation,
        learning_rate=args.lr, num_train_epochs=args.epochs, max_steps=args.max_steps,
        warmup_steps=0.03, weight_decay=0.01, max_grad_norm=1.0,
        bf16=args.dtype=='bf16', fp16=False, optim=args.optimizer,
        gradient_checkpointing=not args.no_gradient_checkpointing,
        gradient_checkpointing_kwargs={'use_reentrant':False},
        logging_steps=args.logging_steps, save_steps=args.save_steps, save_total_limit=2,
        eval_strategy='steps' if eval_data else 'no', eval_steps=args.save_steps,
        remove_unused_columns=False, label_names=['labels'],
        dataloader_num_workers=0, dataloader_pin_memory=False,
        report_to='none', seed=args.seed, data_seed=args.seed,
        disable_tqdm=not args.progress,
        eval_accumulation_steps=8, logging_nan_inf_filter=False,
        load_best_model_at_end=bool(eval_data), metric_for_best_model=args.best_model_metric,
        greater_is_better=True)
    # Keep evaluation order fixed; shuffle candidates only in the training collator.
    class DecisionTrainer(Trainer):
        def _get_train_sampler(self, train_dataset=None):
            dataset = self.train_dataset if train_dataset is None else train_dataset
            if args.group_by_length:
                class RecordedSampler(LengthGroupedSampler):
                    def __iter__(sampler_self):
                        indices = list(super().__iter__())
                        if args.benchmark_steps is not None:
                            save_json(output/'sampler_indices.json', indices)
                        return iter(indices)
                return RecordedSampler(
                    batch_size=self.args.train_batch_size * self.args.gradient_accumulation_steps,
                    lengths=dataset.lengths)
            return super()._get_train_sampler(train_dataset)

        def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
            if checkpoint_policy is not None and model.training:
                checkpoint_policy(inputs['input_ids'])
            loss, outputs = super().compute_loss(model, inputs, return_outputs=True,
                                                num_items_in_batch=num_items_in_batch)
            if not torch.isfinite(loss).all():
                raise FloatingPointError('Non-finite decision loss; stopping without skipping the sample')
            return (loss, outputs) if return_outputs else loss
        def get_eval_dataloader(self, eval_dataset=None):
            old = self.data_collator
            try:
                self.data_collator = FrozenCollator(tokenizer, augment=False, elide_padding_mask=not args.keep_padding_mask)
                return super().get_eval_dataloader(eval_dataset)
            finally:
                self.data_collator = old
    def preprocess(logits, labels):
        # ModelOutput also contains mask and flat_scores; only group logits matter.
        return (logits[0] if isinstance(logits, tuple) else logits).argmax(-1)
    def metrics(result):
        correct = np.asarray(result.predictions) == np.asarray(result.label_ids)
        if len(correct) != len(eval_data):
            raise ValueError('Evaluation predictions do not align with frozen records')
        report = {'accuracy':float(correct.mean())}
        group_names = ['language','source','options']
        language_scores=[]
        for dimension, name in enumerate(group_names):
            values = [group[dimension] for group in eval_data.groups]
            for value in sorted(set(values), key=str):
                selected = np.asarray([x == value for x in values])
                accuracy = float(correct[selected].mean())
                report[f'accuracy_{name}_{value}'] = accuracy
                if name == 'language': language_scores.append(accuracy)
        report['accuracy_macro_language'] = float(np.mean(language_scores))
        if args.best_model_metric == 'accuracy_train_language_weighted':
            from mugi_decision.language_metrics import train_weighted_language_accuracy
            report['accuracy_train_language_weighted'] = train_weighted_language_accuracy(
                {language: report['accuracy_language_' + language] for language in train_data.summary['languages'] if 'accuracy_language_' + language in report},
                train_data.summary['languages'])
        from mugi_decision.eval_cohorts import cohort_metrics, named_cohort_metrics
        report.update(cohort_metrics(correct, eval_data.groups,
                                     eval_data.manifest.get('legacy_validation_rows', 0)))
        report.update(named_cohort_metrics(correct, eval_data.groups,
                                          eval_data.manifest.get('validation_cohorts', [])))
        report['random_choice_baseline'] = float(np.mean([1/x[2] for x in eval_data.groups]))
        return report
    stable_progress = StableProgressCallback(window=args.eta_window, enabled=args.progress)
    class Progress(TrainerCallback):
        def __init__(self):
            self.started = time.monotonic()
            self.logs = {}
        def on_train_begin(self, training_args, state, control, **kwargs):
            self.start_step = state.global_step
            state.train_batch_size = training_args.train_batch_size
            if args.checkpoint_linear_only:
                # Trainer enables all layers before on_train_begin, including on resume.
                base = model.get_base_model() if args.mode == 'lora' else model
                layers = base.backbone.layers
                linear_index = 0
                for layer in layers:
                    linear = hasattr(layer, 'linear_attn')
                    layer.gradient_checkpointing = linear and linear_index % args.checkpoint_linear_stride == 0
                    linear_index += int(linear)
                enabled = sum(bool(layer.gradient_checkpointing) for layer in layers)
                print(f'Gradient checkpointing: {enabled}/{len(layers)} decoder layers (DeltaNet stride {args.checkpoint_linear_stride})', flush=True)
        def on_log(self, args, state, control, logs=None, **kwargs):
            self.logs.update(logs or {})
            for key in ('loss','grad_norm'):
                if key in self.logs and not np.isfinite(float(self.logs[key])):
                    raise FloatingPointError(f'Non-finite {key}; stopped')
        def on_step_end(self, training_args, state, control, **kwargs):
            if args.benchmark_steps is not None and state.global_step-self.start_step >= args.benchmark_steps:
                control.should_training_stop = True
                control.should_save = True
            save_json(output/'progress.json', {'status':'training','step':state.global_step,
                'max_steps':state.max_steps,'epoch':state.epoch,
                'elapsed_seconds':time.monotonic()-self.started,
                'peak_allocated_gib':torch.cuda.max_memory_allocated(device)/2**30,
                'peak_reserved_gib':torch.cuda.max_memory_reserved(device)/2**30,
                'latest_logs':self.logs, 'eta':stable_progress.snapshot})
    progress_callback = Progress()
    callbacks = [stable_progress, progress_callback]
    if args.epoch_snapshots:
        from mugi_decision.epoch_snapshots import EpochSnapshots
        callbacks.append(EpochSnapshots(output, args.model, train_data.manifest_hash))
    trainer = DecisionTrainer(model=model, args=settings, train_dataset=train_data, eval_dataset=eval_data,
        data_collator=FrozenCollator(tokenizer, augment=True, elide_padding_mask=not args.keep_padding_mask), processing_class=tokenizer,
        compute_metrics=metrics, preprocess_logits_for_metrics=preprocess, callbacks=callbacks)
    for callback in list(trainer.callback_handler.callbacks):
        if type(callback) is ProgressCallback:
            trainer.remove_callback(callback)
    if trainer.model_accepts_loss_kwargs:
        raise RuntimeError('Trainer must normalize the per-question mean loss during gradient accumulation')
    output.mkdir(parents=True, exist_ok=True)
    save_json(output/'run_config.json', dict(vars(args), base_model=str(Path(args.model).resolve()),
              dataset_manifest_sha256=train_data.manifest_hash, initialization=initialization, runtime=gpu_metadata(device),
              train_summary=train_data.summary, eval_summary=eval_data.summary if eval_data else None,
              trainable_parameters=trainable, fast_token_layers=fast_token_layers,
              fused_norm_layers=fused_norm_layers, autocast_lora_layers=autocast_lora_layers, performance_resume_changes=performance_changes,
              cpu_threads=torch.get_num_threads(), cpu_affinity=sorted(os.sched_getaffinity(0)) if hasattr(os, 'sched_getaffinity') else None,
              special_tokens_trainable=args.mode=='full' or bool(model.peft_config)))
    if args.initial_eval and eval_data and not args.resume:
        save_json(output/'progress.json', {'status':'initial_evaluation','step':0,'eval_rows':len(eval_data)})
        trainer.save_metrics('initial_eval', trainer.evaluate())
    torch.cuda.reset_peak_memory_stats(device)
    try:
        result = trainer.train(resume_from_checkpoint=args.resume)
    except BaseException as exc:
        save_json(output/'progress.json', {'status':'failed','step':trainer.state.global_step,'error':repr(exc)})
        raise
    if args.benchmark_steps is not None:
        completed = trainer.state.global_step-progress_callback.start_step
        runtime = result.metrics['train_runtime']
        save_json(output/'benchmark_results.json', {
            'start_step':progress_callback.start_step, 'end_step':trainer.state.global_step,
            'total_planned_steps':trainer.state.max_steps,
            'updates':completed, 'questions':completed*args.batch_groups*args.gradient_accumulation,
            'runtime_seconds':runtime,
            'runtime_questions_per_second':completed*args.batch_groups*args.gradient_accumulation/runtime,
            'steady_window':stable_progress.snapshot,
            'checkpoint_batches':checkpoint_policy.batches if checkpoint_policy else None,
            'peak_allocated_gib':torch.cuda.max_memory_allocated(device)/2**30,
            'peak_reserved_gib':torch.cuda.max_memory_reserved(device)/2**30})
    result.metrics['peak_allocated_gib'] = torch.cuda.max_memory_allocated(device) / 2**30
    result.metrics['peak_reserved_gib'] = torch.cuda.max_memory_reserved(device) / 2**30
    trainer.save_model(str(output/'final'))
    tokenizer.save_pretrained(output/'final')
    if (Path(args.model)/'LICENSE').exists():
        shutil.copyfile(Path(args.model)/'LICENSE', output/'final'/'LICENSE')
    save_json(output/'final'/'decision_training.json', {'base_model':str(Path(args.model).resolve()),'mode':args.mode,'steps':trainer.state.global_step,'state':'finetuned',
        'dataset_manifest_sha256':train_data.manifest_hash, 'best_checkpoint':trainer.state.best_model_checkpoint,
        'initialization':initialization})
    trainer.save_metrics('train', result.metrics)
    if eval_data:
        trainer.save_metrics('eval', trainer.evaluate())
    save_json(output/'progress.json', {'status':'completed','step':trainer.state.global_step,
        'final':str((output/'final').resolve()),'best_checkpoint':trainer.state.best_model_checkpoint})
    print('Saved:', output/'final')

if __name__ == '__main__': main()