"""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()