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