import argparse import importlib.util import math import re import sys from pathlib import Path import mlx.core as mx import mlx.nn as nn import mlx.optimizers as optim import numpy as np import yaml from mlx_lm.tuner.callbacks import WandBCallback from mlx_lm.tuner.utils import ( build_schedule, load_adapters, print_trainable_parameters, ) from mlx_lm.utils import load, load_tokenizer from .chat_template import ensure_chat_template from .trainer.cpo_trainer import CPOTrainingArgs, evaluate_cpo, train_cpo from .trainer.datasets import CacheDataset, load_dataset from .trainer.dpo_trainer import DPOTrainingArgs, evaluate_dpo, train_dpo from .trainer.dsla_trainer import DSLATrainingArgs, evaluate_dsla, train_dsla from .trainer.ftpo_trainer import FTPOTrainingArgs, evaluate_ftpo, train_ftpo from .trainer.grpo_reward_functions import ( get_default_reward_functions, get_reward_function, list_available_reward_functions, ) from .trainer.grpo_trainer import GRPOTrainingArgs, evaluate_grpo, train_grpo from .trainer.klpo_trainer import KLPOTrainingArgs, evaluate_klpo, train_klpo from .trainer.online_dpo_trainer import ( OnlineDPOTrainingArgs, evaluate_online_dpo, train_online_dpo, ) from .trainer.orpo_trainer import ORPOTrainingArgs, evaluate_orpo, train_orpo from .trainer.ppo_trainer import PPOTrainingArgs, evaluate_ppo, train_ppo from .trainer.rlhf_reinforce_trainer import ( RLHFReinforceTrainingArgs, evaluate_rlhf_reinforce, train_rlhf_reinforce, ) from .trainer.sft_trainer import ( SFTTrainingArgs, TrainingCallback, evaluate_sft, get_sft_loss, train_sft, ) from .trainer.xpo_trainer import XPOTrainingArgs, evaluate_xpo, train_xpo from .utils import from_pretrained, save_pretrained_merged, save_to_lmstudio_merged from .visuals import ( Colors, print_banner, print_error, print_info, print_section, print_success, print_warning, ) yaml_loader = yaml.SafeLoader yaml_loader.add_implicit_resolver( "tag:yaml.org,2002:float", re.compile( """^(?: [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)? |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+) |\\.[0-9_]+(?:[eE][-+][0-9]+)? |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]* |[-+]?\\.(?:inf|Inf|INF) |\\.(?:nan|NaN|NAN))$""", re.X, ), list("-+0123456789."), ) CONFIG_DEFAULTS = { "model": "mlx_model", "train": False, "load_in_4bits": False, "load_in_6bits": False, "load_in_8bits": False, "load_in_mxfp4": False, "train_type": "lora", "train_mode": "sft", "sft_loss_type": "nll", "optimizer": "adam", "optimizer_config": {"adam": {}, "adamw": {}, "muon": {}}, "data": "data/", "seed": 0, "num_layers": -1, "batch_size": 1, "iters": None, "epochs": None, "gradient_accumulation_steps": 1, "micro_batch_size": None, "val_batches": 25, "learning_rate": 1e-5, "steps_per_report": 10, "steps_per_eval": 200, "resume_adapter_file": None, "adapter_path": "adapters", "save_every": 100, "test": False, "test_batches": 500, "max_seq_length": 2048, "config": None, "grad_checkpoint": False, "efficient_long_context": False, "recurrence_chunk_size": 64, "lr_schedule": None, "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0}, "mask_prompt": False, "fuse": True, "beta": 0.1, "reward_scaling": 1.0, "dpo_cpo_loss_type": "sigmoid", "dsla_loss": "dpo", "latent_weight": 0.1, "latent_margin": 0.05, "latent_gamma": 10.0, "latent_variant": "both", "latent_pooling": "answer_mean", "latent_layer": "final", "delta": 50.0, "reference_model_path": None, "lambda_mse_target": 0.05, "tau_mse_target": 1.0, "lambda_mse": 0.4, "clip_epsilon_logits": 2.0, "judge": None, "judge_config": {}, "alpha": 1e-5, "group_size": 4, "epsilon": 1e-4, "epsilon_high": None, "max_completion_length": 512, "temperature": 0.8, "reward_weights": None, "reward_functions": None, "reward_functions_file": None, "grpo_loss_type": "grpo", "importance_sampling_level": "token", "klpo_route": "token", "klpo_kl_estimator": "mc", "klpo_mc_samples": 128, "klpo_top_k": 128, "klpo_tail_floor": 1e-6, "lm_studio_name": None, "qat_enable": False, "qat_bits": 8, "qat_group_size": 64, "qat_mode": "affine", "qat_start_step": 1, "qat_interval": 1, } def load_reward_functions_from_file(file_path): """Load reward functions from a Python file""" if not file_path or not Path(file_path).exists(): return None try: print(f"Loading custom reward functions from {file_path}") spec = importlib.util.spec_from_file_location("custom_rewards", file_path) custom_rewards = importlib.util.module_from_spec(spec) sys.modules["custom_rewards"] = custom_rewards spec.loader.exec_module(custom_rewards) print("Successfully loaded custom reward functions") return True except Exception as e: print(f"Error loading custom reward functions: {e}") return None def calculate_iters(train_set, batch_size, epochs) -> int: num_samples = len(train_set) batches_per_epoch = math.ceil(num_samples / batch_size) iters = epochs * batches_per_epoch print( f"[INFO] Calculated {iters} iterations from {epochs} epochs (dataset size: {num_samples}, batch size: {batch_size})" ) return iters def load_reference_model(args): """Load reference model, reusing main model if no separate path specified""" if args.reference_model_path: print(f"Loading pretrained reference model from {args.reference_model_path}") model, _ = load(args.reference_model_path) else: print("Loading pretrained reference model (using main model)") model, _ = load(args.model) return model.freeze() def load_judge_model(args, reference_model=None): """Load judge model, reusing reference model if paths match""" if args.judge is None: raise ValueError( f"Training mode '{args.train_mode}' requires a judge model. " "Set 'judge' in the YAML config or pass --judge." ) if args.judge == args.reference_model_path and reference_model is not None: print("Loading judge model (reusing reference model)") tokenizer = load_tokenizer(args.judge) ensure_chat_template(tokenizer) return reference_model, tokenizer print(f"Loading pretrained judge model from {args.judge}") model, tokenizer = load(args.judge) ensure_chat_template(tokenizer) return model.freeze(), tokenizer def build_parser(): parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.") parser.add_argument( "--model", type=str, help="The path to the local model directory or Hugging Face repo.", ) parser.add_argument( "--lm-studio-name", type=str, help="The name to use when sending the trained model to LM Studio.", ) parser.add_argument( "--load-in-4bits", action="store_true", help="Load the model in 4-bit quantization.", default=None, ) parser.add_argument( "--load-in-6bits", action="store_true", help="Load the model in 6-bit quantization.", default=None, ) parser.add_argument( "--load-in-8bits", action="store_true", help="Load the model in 8-bit quantization.", default=None, ) parser.add_argument( "--load-in-mxfp4", action="store_true", help="Load the model in mixed FP4 quantization.", default=None, ) parser.add_argument( "--train", action="store_true", help="Do training", default=None ) parser.add_argument( "--data", type=str, help="Directory with {train, valid, test}.jsonl files or the name of a Hugging Face dataset", ) parser.add_argument( "--train-type", type=str, choices=["lora", "dora", "full"], help="Type of fine-tuning to perform: lora, dora, or full.", ) parser.add_argument( "--train-mode", type=str, default=None, choices=[ "sft", "dpo", "dsla", "ftpo", "cpo", "orpo", "grpo", "klpo", "online_dpo", "xpo", "rlhf_reinforce", "ppo", ], help="Training mode", ) parser.add_argument( "--optimizer", type=str, choices=["adam", "adamw", "muon"], default=None, help="Optimizer to use for training", ) parser.add_argument( "--sft-loss-type", type=str, choices=["nll", "chunked_nll", "dft"], default="nll", help="SFT loss type: NLL, chunked NLL, or dynamic fine-tuning loss.", ) parser.add_argument( "--mask-prompt", action="store_true", help="Mask the prompt in the loss when training", default=None, ) parser.add_argument( "--num-layers", type=int, help="Number of layers to fine-tune. Default is 16, use -1 for all.", ) parser.add_argument("--batch-size", type=int, help="Minibatch size.") parser.add_argument("--iters", type=int, help="Iterations to train for.") parser.add_argument( "--epochs", type=int, help="Epochs to train for. Ignored if --iters is provided.", ) parser.add_argument( "--gradient-accumulation-steps", type=int, help="Number of gradient accumulation steps.", default=1, ) parser.add_argument( "--val-batches", type=int, help="Number of validation batches, -1 uses the entire validation set.", ) parser.add_argument("--learning-rate", type=float, help="Optimizer learning rate.") parser.add_argument( "--steps-per-report", type=int, help="Number of training steps between loss reporting.", ) parser.add_argument( "--steps-per-eval", type=int, help="Number of training steps between validations.", ) parser.add_argument( "--resume-adapter-file", type=str, help="Load path to resume training from the given fine-tuned weights.", ) parser.add_argument( "--adapter-path", type=str, help="Save/load path for the fine-tuned weights." ) parser.add_argument( "--save-every", type=int, help="Save the model every N iterations." ) parser.add_argument( "--test", action="store_true", help="Evaluate on the test set after training", default=None, ) parser.add_argument( "--test-batches", type=int, help="Number of test set batches, -1 uses the entire test set.", ) parser.add_argument("--max-seq-length", type=int, help="Maximum sequence length.") parser.add_argument( "-c", "--config", type=str, help="A YAML configuration file with the training options", ) parser.add_argument( "--grad-checkpoint", action="store_true", help="Use gradient checkpointing to reduce memory use.", default=None, ) parser.add_argument( "--efficient-long-context", action="store_true", help="Use efficient long context processing (Experimental, only supported in SFT, DPO, CPO, ORPO).", default=None, ) parser.add_argument( "--recurrence-chunk-size", type=int, help="Chunk size for memory-safe recurrent training fallbacks.", default=None, ) parser.add_argument( "--wandb", type=str, default=None, help="WandB project name to report training metrics. Disabled if None.", ) parser.add_argument("--seed", type=int, help="The PRNG seed") parser.add_argument( "--fuse", action="store_true", help="Fuse and save the trained model.", default=None, ) parser.add_argument( "--beta", type=float, help="Temperature parameter for ORPO training.", default=None, ) parser.add_argument( "--reward-scaling", type=float, help="Reward scaling factor for ORPO training, not implemented.", default=1.0, ) parser.add_argument( "--dpo-cpo-loss-type", type=str, help="DPO loss type: 'sigmoid', 'hinge', 'ipo', or 'dpop'.", choices=["sigmoid", "hinge", "ipo", "dpop"], default="sigmoid", ) parser.add_argument( "--delta", type=float, help="Delta parameter for DPOP loss type.", default=50.0 ) parser.add_argument( "--dsla-loss", choices=["dpo", "orpo", "cpo"], default="dpo", help="Underlying preference objective for --train-mode dsla.", ) parser.add_argument("--latent-weight", type=float, default=0.1) parser.add_argument("--latent-margin", type=float, default=0.05) parser.add_argument("--latent-gamma", type=float, default=10.0) parser.add_argument( "--latent-variant", choices=["similarity", "direction", "both"], default="both" ) parser.add_argument( "--latent-pooling", choices=["answer_mean", "last_token", "last_k_mean", "prompt_answer_mean"], default="answer_mean", ) parser.add_argument("--latent-layer", default="final") parser.add_argument( "--reference-model-path", type=str, help="Path to reference model weights. If None, uses the same model.", default=None, ) parser.add_argument("--lambda-mse-target", type=float, default=None) parser.add_argument("--tau-mse-target", type=float, default=None) parser.add_argument("--lambda-mse", type=float, default=None) parser.add_argument("--clip-epsilon-logits", type=float, default=None) parser.add_argument( "--judge", type=str, help="Judge to use can be a model ID or 'human'.", default=None, ) parser.add_argument( "--alpha", type=list[float], help="Judge to use can be a model ID or 'human'.", default=[1e-5], ) parser.add_argument( "--group-size", type=int, help="Number of generations.", default=4 ) parser.add_argument( "--max-completion-length", type=int, help="Maximum length of the prompt.", default=512, ) parser.add_argument( "--epsilon", type=float, help="The Epsilon for numerical stability.", default=1e-4, ) parser.add_argument( "--temperature", type=float, help="Temperature for sampling.", default=1.0 ) parser.add_argument( "--micro-batch-size", type=int, default=None, help="Bound online DPO/RLHF/XPO scoring memory with this microbatch size.", ) parser.add_argument( "--reward-weights", type=str, help="Weights for each reward function.", default=None, ) parser.add_argument( "--reward-functions", type=str, help="Comma-separated list of reward function names to use.", default=None, ) parser.add_argument( "--reward-functions-file", type=str, help="Path to a Python file containing custom reward functions.", default=None, ) parser.add_argument( "--list-reward-functions", action="store_true", help="List all available reward functions and exit", ) parser.add_argument( "--grpo-loss-type", type=str, help="GRPO loss type: 'grpo', 'bnpo', or 'dr_grpo'.", choices=["grpo", "bnpo", "dr_grpo"], default="grpo", ) parser.add_argument( "--epsilon-high", type=float, help="Upper-bound epsilon value for clipping.", default=None, ) parser.add_argument( "--importance-sampling-level", type=str, choices=["token", "sequence"], default="token", help="Level of importance sampling to use (default: token).", ) parser.add_argument( "--klpo-route", type=str, choices=["token", "sequence"], default="token", help="KLPO regression route.", ) parser.add_argument( "--klpo-kl-estimator", type=str, choices=["mc", "topk", "binary", "full"], default="mc", help="KLPO conditional KL estimator.", ) parser.add_argument( "--klpo-mc-samples", type=int, default=128, help="KLPO auxiliary MC-KL draws per visited prefix.", ) parser.add_argument( "--klpo-top-k", type=int, default=128, help="KLPO TopK-KL sampler head size.", ) parser.add_argument( "--klpo-tail-floor", type=float, default=1e-6, help="KLPO TopK-KL tail probability floor.", ) parser.add_argument( "--qat-enable", action="store_true", default=None, help="Enable minimal QAT-style projection in SFT training.", ) parser.add_argument( "--qat-bits", type=int, default=None, help="Bit-width used by QAT projection (SFT only).", ) parser.add_argument( "--qat-group-size", type=int, default=None, help="Group size used by QAT projection (SFT only).", ) parser.add_argument( "--qat-mode", type=str, choices=["affine"], default=None, help="QAT projection mode (SFT only).", ) parser.add_argument( "--qat-start-step", type=int, default=None, help="Apply QAT projection starting from this optimizer step (SFT only).", ) parser.add_argument( "--qat-interval", type=int, default=None, help="Apply QAT projection every N optimizer steps (SFT only).", ) return parser def _dsla_training_args(args, **overrides): """Map CLI/config settings to the standalone DSLA trainer.""" values = { name: getattr(args, name) for name in DSLATrainingArgs.__dataclass_fields__ if hasattr(args, name) } values.update( loss_type=args.dsla_loss, steps_per_save=args.save_every, seq_step_size=512 if args.efficient_long_context else None, ) values.update(overrides) return DSLATrainingArgs(**values) def train_model( args, model: nn.Module, tokenizer, adapter_file: Path, reference_model: nn.Module = None, judge_model: nn.Module = None, judge_tokenizer=None, train_set: CacheDataset = None, valid_set: CacheDataset = None, training_callback: TrainingCallback = None, ): mx.random.seed(args.seed) if args.iters is None: if args.epochs is not None: args.iters = calculate_iters( train_set=train_set, batch_size=args.batch_size, epochs=args.epochs ) else: args.iters = SFTTrainingArgs().iters print_info( f"Neither iters nor epochs was provided; defaulting to {args.iters} iterations." ) if args.resume_adapter_file is not None: print_warning( f"Loading fine-tuned weights from {Colors.CYAN}{args.resume_adapter_file}{Colors.RESET}" ) model.load_weights(args.resume_adapter_file, strict=False) print_trainable_parameters(model) lr = build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate optimizer_config = args.optimizer_config.get(args.optimizer.lower(), {}) opt_class = {"adam": optim.Adam, "adamw": optim.AdamW, "muon": optim.Muon}[ args.optimizer.lower() ] opt = opt_class(learning_rate=lr, **optimizer_config) print_info(f"Training mode: {Colors.YELLOW}{args.train_mode.upper()}{Colors.RESET}") # Training mode dispatch if args.train_mode == "dsla": train_dsla( model=model, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=_dsla_training_args(args, adapter_file=adapter_file), ref_model=reference_model, training_callback=training_callback, ) elif args.train_mode == "orpo": train_orpo( model=model, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=ORPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, beta=args.beta, seq_step_size=512 if args.efficient_long_context else None, reward_scaling=args.reward_scaling, gradient_accumulation_steps=args.gradient_accumulation_steps, qat_enable=args.qat_enable, qat_bits=args.qat_bits, qat_group_size=args.qat_group_size, qat_mode=args.qat_mode, qat_start_step=args.qat_start_step, qat_interval=args.qat_interval, ), training_callback=training_callback, ) elif args.train_mode == "ftpo": train_ftpo( model=model, ref_model=reference_model, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=FTPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, gradient_accumulation_steps=args.gradient_accumulation_steps, lambda_mse_target=args.lambda_mse_target, tau_mse_target=args.tau_mse_target, lambda_mse=args.lambda_mse, clip_epsilon_logits=args.clip_epsilon_logits, ), training_callback=training_callback, ) elif args.train_mode == "dpo": train_dpo( model=model, ref_model=reference_model, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=DPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, beta=args.beta, loss_type=args.dpo_cpo_loss_type, delta=args.delta, reference_model_path=args.reference_model_path, seq_step_size=512 if args.efficient_long_context else None, gradient_accumulation_steps=args.gradient_accumulation_steps, qat_enable=args.qat_enable, qat_bits=args.qat_bits, qat_group_size=args.qat_group_size, qat_mode=args.qat_mode, qat_start_step=args.qat_start_step, qat_interval=args.qat_interval, ), training_callback=training_callback, ) elif args.train_mode in ["online_dpo", "ppo", "rlhf_reinforce", "xpo"]: train_func = { "online_dpo": (train_online_dpo, OnlineDPOTrainingArgs), "ppo": (train_ppo, PPOTrainingArgs), "rlhf_reinforce": (train_rlhf_reinforce, RLHFReinforceTrainingArgs), "xpo": (train_xpo, XPOTrainingArgs), }[args.train_mode] train_args_kwargs = { "batch_size": args.batch_size, "iters": args.iters, "val_batches": args.val_batches, "steps_per_report": args.steps_per_report, "steps_per_eval": args.steps_per_eval, "steps_per_save": args.save_every, "adapter_file": adapter_file, "max_seq_length": args.max_seq_length, "grad_checkpoint": args.grad_checkpoint, "recurrence_chunk_size": args.recurrence_chunk_size, "beta": args.beta, "reference_model_path": args.reference_model_path, "gradient_accumulation_steps": args.gradient_accumulation_steps, "micro_batch_size": args.micro_batch_size, "judge": args.judge, "max_completion_length": args.max_completion_length, } if args.train_mode in ["online_dpo", "xpo"]: train_args_kwargs.update( {"loss_type": args.dpo_cpo_loss_type, "delta": args.delta} ) if args.train_mode == "ppo": train_args_kwargs.update( { "loss_type": args.dpo_cpo_loss_type, "delta": args.delta, "epsilon": args.epsilon, "temperature": args.temperature, } ) if args.train_mode == "xpo": train_args_kwargs["alpha"] = args.alpha if args.train_mode == "online_dpo": train_args_kwargs["temperature"] = args.temperature train_func[0]( model=model, tokenizer=tokenizer, ref_model=reference_model, judge_model=judge_model, judge_tokenizer=judge_tokenizer, judge_config=args.judge_config, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=train_func[1](**train_args_kwargs), training_callback=training_callback, ) elif args.train_mode == "cpo": train_cpo( model=model, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, args=CPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, beta=args.beta, loss_type=args.dpo_cpo_loss_type, delta=args.delta, seq_step_size=512 if args.efficient_long_context else None, reference_model_path=args.reference_model_path, gradient_accumulation_steps=args.gradient_accumulation_steps, ), training_callback=training_callback, ) elif args.train_mode == "klpo": if args.reward_functions_file: load_reward_functions_from_file(args.reward_functions_file) reward_funcs = get_default_reward_functions() if args.reward_functions: func_names = [name.strip() for name in args.reward_functions.split(",")] try: reward_funcs = [get_reward_function(name) for name in func_names] print_success(f"Using custom reward functions: {', '.join(func_names)}") except KeyError as e: print_error(f"Error: {e!s}") print_info( f"Available reward functions: {list_available_reward_functions()}" ) return train_klpo( model=model, tokenizer=tokenizer, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, reward_funcs=reward_funcs, args=KLPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, max_completion_length=args.max_completion_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, beta=args.beta, route=args.klpo_route, kl_estimator=args.klpo_kl_estimator, mc_samples=args.klpo_mc_samples, top_k=args.klpo_top_k, tail_floor=args.klpo_tail_floor, temperature=args.temperature, gradient_accumulation_steps=args.gradient_accumulation_steps, reward_weights=( [float(x) for x in args.reward_weights.strip("[]").split(",")] if args.reward_weights else None ), ), training_callback=training_callback, ) elif args.train_mode == "grpo": if args.reward_functions_file: load_reward_functions_from_file(args.reward_functions_file) reward_funcs = get_default_reward_functions() if args.reward_functions: func_names = [name.strip() for name in args.reward_functions.split(",")] try: reward_funcs = [get_reward_function(name) for name in func_names] print_success(f"Using custom reward functions: {', '.join(func_names)}") except KeyError as e: print_error(f"Error: {e!s}") print_info( f"Available reward functions: {list_available_reward_functions()}" ) return train_grpo( model=model, ref_model=reference_model, tokenizer=tokenizer, optimizer=opt, train_dataset=train_set, val_dataset=valid_set, reward_funcs=reward_funcs, args=GRPOTrainingArgs( batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, max_completion_length=args.max_completion_length, grad_checkpoint=args.grad_checkpoint, recurrence_chunk_size=args.recurrence_chunk_size, beta=args.beta, group_size=args.group_size, epsilon=args.epsilon, epsilon_high=args.epsilon_high, reference_model_path=args.reference_model_path, temperature=args.temperature, gradient_accumulation_steps=args.gradient_accumulation_steps, reward_weights=( [float(x) for x in args.reward_weights.strip("[]").split(",")] if args.reward_weights else None ), importance_sampling_level=args.importance_sampling_level, grpo_loss_type=args.grpo_loss_type, ), training_callback=training_callback, ) elif args.train_mode == "sft": train_sft( model=model, args=SFTTrainingArgs( loss_type=args.sft_loss_type, batch_size=args.batch_size, iters=args.iters, val_batches=args.val_batches, steps_per_report=args.steps_per_report, steps_per_eval=args.steps_per_eval, steps_per_save=args.save_every, adapter_file=adapter_file, max_seq_length=args.max_seq_length, grad_checkpoint=args.grad_checkpoint, gradient_accumulation_steps=args.gradient_accumulation_steps, seq_step_size=512 if args.efficient_long_context else None, recurrence_chunk_size=args.recurrence_chunk_size, qat_enable=args.qat_enable, qat_bits=args.qat_bits, qat_group_size=args.qat_group_size, qat_mode=args.qat_mode, qat_start_step=args.qat_start_step, qat_interval=args.qat_interval, ), optimizer=opt, train_dataset=train_set, val_dataset=valid_set, training_callback=training_callback, ) else: raise ValueError(f"The train mode {args.train_mode} does not exist.") def evaluate_model( args, model: nn.Module, tokenizer, reference_model: nn.Module = None, judge_model: nn.Module = None, judge_tokenizer=None, test_set: CacheDataset = None, ): """Evaluate model on test set based on training mode""" print_section(f"Evaluating {args.train_mode.upper()} Model") if args.train_mode == "dsla": test_loss, test_rewards, _, test_metrics = evaluate_dsla( model=model, dataset=test_set, args=_dsla_training_args(args), ref_model=reference_model, num_batches=args.test_batches, ) print(f"{Colors.BOLD}DSLA Test loss:{Colors.RESET} {test_loss:.3f}") print( f"Chosen reward: {test_rewards[0]:.3f}, rejected reward: {test_rewards[1]:.3f}" ) for metric_name, metric_value in test_metrics.items(): print(f" {metric_name}: {float(metric_value):.3f}") elif args.train_mode == "orpo": test_loss, test_rewards, _, test_metrics = evaluate_orpo( model=model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, beta=args.beta, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}\n" f" {Colors.YELLOW}Rewards:{Colors.RESET} {test_rewards[0]:.3f}, {test_rewards[1]:.3f}" ) print(f"\n{Colors.CYAN}ORPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "ftpo": test_loss, test_metrics = evaluate_ftpo( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, args=FTPOTrainingArgs( lambda_mse_target=args.lambda_mse_target, tau_mse_target=args.tau_mse_target, lambda_mse=args.lambda_mse, clip_epsilon_logits=args.clip_epsilon_logits, ), ) print(f"{Colors.BOLD}FTPO Test loss:{Colors.RESET} {test_loss:.3f}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "dpo": test_loss, _, _, test_metrics = evaluate_dpo( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, beta=args.beta, delta=args.delta, loss_type=args.dpo_cpo_loss_type, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}DPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "cpo": test_loss, _, _, test_metrics = evaluate_cpo( model=model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, beta=args.beta, delta=args.delta, loss_type=args.dpo_cpo_loss_type, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}CPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "online_dpo": test_loss, _, _, test_metrics = evaluate_online_dpo( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, beta=args.beta, delta=args.delta, max_seq_length=args.max_seq_length, loss_type=args.dpo_cpo_loss_type, judge_config=args.judge_config, judge_model=judge_model, judge_tokenizer=judge_tokenizer, tokenizer=tokenizer, max_tokens=args.max_completion_length, temperature=args.temperature, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}Online DPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "ppo": test_loss, _, _, test_metrics = evaluate_ppo( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, beta=args.beta, epsilon=args.epsilon, max_seq_length=args.max_seq_length, loss_type=args.dpo_cpo_loss_type, judge_config=args.judge_config, judge_model=judge_model, judge_tokenizer=judge_tokenizer, tokenizer=tokenizer, max_tokens=args.max_completion_length, temperature=args.temperature, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}PPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "rlhf_reinforce": test_loss, _, test_metrics = evaluate_rlhf_reinforce( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, beta=args.beta, max_seq_length=args.max_seq_length, judge_config=args.judge_config, judge_model=judge_model, judge_tokenizer=judge_tokenizer, tokenizer=tokenizer, max_tokens=args.max_completion_length, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}RLHF Reinforce Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "xpo": test_loss, _, _, test_metrics = evaluate_xpo( model=model, ref_model=reference_model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, beta=args.beta, delta=args.delta, max_seq_length=args.max_seq_length, loss_type=args.dpo_cpo_loss_type, judge_config=args.judge_config, alpha=args.alpha, judge_model=judge_model, judge_tokenizer=judge_tokenizer, tokenizer=tokenizer, max_tokens=args.max_completion_length, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}" ) print(f"\n{Colors.CYAN}XPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "klpo": if args.reward_functions_file: load_reward_functions_from_file(args.reward_functions_file) reward_funcs = get_default_reward_functions() if args.reward_functions: func_names = [name.strip() for name in args.reward_functions.split(",")] try: reward_funcs = [get_reward_function(name) for name in func_names] except KeyError as e: print_error(f"Error: {e!s}") print_info( f"Available reward functions: {list_available_reward_functions()}" ) return test_loss, test_ntokens, test_metrics = evaluate_klpo( model=model, dataset=test_set, tokenizer=tokenizer, batch_size=args.batch_size, num_batches=args.test_batches, beta=args.beta, route=args.klpo_route, kl_estimator=args.klpo_kl_estimator, mc_samples=args.klpo_mc_samples, top_k=args.klpo_top_k, tail_floor=args.klpo_tail_floor, max_seq_length=args.max_seq_length, max_tokens=args.max_completion_length, temperature=args.temperature, reward_funcs=reward_funcs, reward_weights=( [float(x) for x in args.reward_weights.strip("[]").split(",")] if args.reward_weights else None ), ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}\n" f" {Colors.YELLOW}Tokens:{Colors.RESET} {test_ntokens}" ) print(f"\n{Colors.CYAN}KLPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "grpo": if args.reward_functions_file: load_reward_functions_from_file(args.reward_functions_file) reward_funcs = get_default_reward_functions() if args.reward_functions: func_names = [name.strip() for name in args.reward_functions.split(",")] try: reward_funcs = [get_reward_function(name) for name in func_names] except KeyError as e: print_error(f"Error: {e!s}") print_info( f"Available reward functions: {list_available_reward_functions()}" ) return from .trainer.grpo_trainer import iterate_batches, loss_fn test_loss, test_ntokens, test_metrics = evaluate_grpo( model=model, dataset=test_set, loss_fn=loss_fn, ref_model=reference_model, reward_funcs=reward_funcs, tokenizer=tokenizer, group_size=args.group_size, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, max_tokens=args.max_completion_length, beta=args.beta, epsilon=args.epsilon, epsilon_high=args.epsilon_high, iterate_batches=iterate_batches, grpo_loss_type=args.grpo_loss_type, end_answer_token=getattr(args, "end_answer_token", None), temperature=args.temperature, top_p=getattr(args, "top_p", 1.0), top_k=getattr(args, "top_k", -1), min_p=getattr(args, "min_p", 0.0), ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}\n" f" {Colors.YELLOW}Tokens:{Colors.RESET} {test_ntokens}" ) print(f"\n{Colors.CYAN}GRPO Test Metrics:{Colors.RESET}") for metric_name, metric_value in test_metrics.items(): print( f" {Colors.WHITE}{metric_name}:{Colors.RESET} {float(metric_value):.3f}" ) elif args.train_mode == "sft": test_loss, test_ntokens = evaluate_sft( model=model, dataset=test_set, batch_size=args.batch_size, num_batches=args.test_batches, max_seq_length=args.max_seq_length, loss=get_sft_loss(args.sft_loss_type), recurrence_chunk_size=args.recurrence_chunk_size, ) test_ppl = math.exp(test_loss) print( f"{Colors.BOLD}Test Results:{Colors.RESET}\n" f" {Colors.YELLOW}Loss:{Colors.RESET} {test_loss:.3f}\n" f" {Colors.YELLOW}Perplexity:{Colors.RESET} {test_ppl:.3f}\n" f" {Colors.YELLOW}Tokens:{Colors.RESET} {test_ntokens}" ) def build_lora_config(args): if args.train_type not in ["lora", "dora"]: return None lora_parameters = dict(getattr(args, "lora_parameters", {}) or {}) return { "rank": lora_parameters.get("rank", 8), "dropout": lora_parameters.get("dropout", 0.0), "scale": lora_parameters.get("scale", 10.0), "use_dora": args.train_type == "dora", "num_layers": getattr(args, "num_layers", None), } def run(args, training_callback: TrainingCallback = None): np.random.seed(args.seed) if args.wandb is not None: training_callback = WandBCallback( project_name=args.wandb, log_dir=args.adapter_path, config=vars(args), wrapped_callback=training_callback, ) quantization_config = None if args.load_in_4bits: quantization_config = {"bits": 4, "group_size": 128} elif args.load_in_6bits: quantization_config = {"bits": 6, "group_size": 128} elif args.load_in_8bits: quantization_config = {"bits": 8, "group_size": 128} elif args.load_in_mxfp4: quantization_config = {"bits": 4, "group_size": 32, "mode": "mxfp4"} print_info(f"Loading model: {Colors.CYAN}{args.model}{Colors.RESET}") model, tokenizer, adapter_file = from_pretrained( model=args.model, new_adapter_path=args.adapter_path, lora_config=build_lora_config(args), quantized_load=quantization_config, ) reference_model = ( load_reference_model(args) if args.train_mode in ["dpo", "ftpo", "grpo", "online_dpo", "ppo", "rlhf_reinforce", "xpo"] or (args.train_mode == "dsla" and args.dsla_loss == "dpo") else None ) judge_model, judge_tokenizer = ( load_judge_model(args, reference_model) if args.train_mode in ["online_dpo", "ppo", "rlhf_reinforce", "xpo"] else (None, None) ) print_info("Loading datasets") train_set, valid_set, test_set = load_dataset(args, tokenizer) if args.test and not args.train: if args.adapter_path != "": load_adapters(model, args.adapter_path) elif args.train: print_section("Training") train_model( args=args, model=model, tokenizer=tokenizer, adapter_file=adapter_file, reference_model=reference_model, judge_model=judge_model, judge_tokenizer=judge_tokenizer, train_set=CacheDataset(train_set), valid_set=CacheDataset(valid_set), training_callback=training_callback, ) else: raise ValueError("Must provide at least one of --train or --test") if args.test: print_section("Testing") evaluate_model( args=args, model=model, tokenizer=tokenizer, reference_model=reference_model, judge_model=judge_model, judge_tokenizer=judge_tokenizer, test_set=CacheDataset(test_set), ) mx.clear_cache() del reference_model, judge_model, judge_tokenizer if args.fuse and args.train: print_section("Fusing Model") if args.lm_studio_name is not None: save_to_lmstudio_merged( model=model, tokenizer=tokenizer, new_model_name=args.lm_studio_name, de_quantize=True, ) else: save_pretrained_merged( model=model, tokenizer=tokenizer, save_path=args.adapter_path, de_quantize=( False if ( args.load_in_4bits or args.load_in_6bits or args.load_in_8bits or args.load_in_mxfp4 ) else True ), ) print_success( f"Model fused and saved to {Colors.CYAN}{args.adapter_path}{Colors.RESET}" ) def main(args=None): import os import types os.environ["TOKENIZERS_PARALLELISM"] = "true" print_banner() if args is None: parser = build_parser() args = parser.parse_args() elif isinstance(args, dict): default_args = vars(build_parser().parse_args([])) default_args.update(args) args = types.SimpleNamespace(**default_args) if args.config: with open(args.config, "r") as f: config_args = yaml.load(f, Loader=yaml_loader) for k, v in config_args.items(): if getattr(args, k, None) is None: setattr(args, k, v) for k, v in CONFIG_DEFAULTS.items(): if getattr(args, k, None) is None: setattr(args, k, v) print_section("Configuration Summary") print(f"{Colors.WHITE}Model:{Colors.RESET} {args.model}") print(f"{Colors.WHITE}Training Mode:{Colors.RESET} {args.train_mode.upper()}") print(f"{Colors.WHITE}Training Type:{Colors.RESET} {args.train_type}") print(f"{Colors.WHITE}Batch Size:{Colors.RESET} {args.batch_size}") print(f"{Colors.WHITE}Learning Rate:{Colors.RESET} {args.learning_rate}") print(f"{Colors.WHITE}Optimizer:{Colors.RESET} {args.optimizer}") if args.train_mode == "sft" and args.qat_enable: print( f"{Colors.WHITE}QAT:{Colors.RESET} enabled " f"(bits={args.qat_bits}, group_size={args.qat_group_size}, " f"mode={args.qat_mode}, start={args.qat_start_step}, interval={args.qat_interval})" ) if args.load_in_4bits: print(f"{Colors.WHITE}Quantization:{Colors.RESET} 4-bit") elif args.load_in_6bits: print(f"{Colors.WHITE}Quantization:{Colors.RESET} 6-bit") elif args.load_in_8bits: print(f"{Colors.WHITE}Quantization:{Colors.RESET} 8-bit") run(args) if __name__ == "__main__": main()