train_rl.py

"""
RL training script using GRPO (Group Relative Policy Optimization) for LoReFT on reasoning tasks.

This implements the actual training from the paper using policy gradients with importance sampling.

Usage:
    # Full fine-tuning with GRPO
    python train_rl.py --mode full_finetune --dataset gsm8k --lr 1e-5

    # LoRA rank 1 with GRPO
    python train_rl.py --mode lora --lora_rank 1 --dataset gsm8k --lr 1e-4

    # LoReFT with GRPO
    python train_rl.py --mode loreft --reft_rank 4 --dataset gsm8k --lr 5e-5

    # Learning rate sweep
    python train_rl.py --mode lora --lora_rank 1 --dataset gsm8k \
        --lr_sweep 1e-7,1e-6,1e-5,1e-4,1e-3 --use_wandb
"""

import argparse
import json
import os
import yaml

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    set_seed,
)
from trl import GRPOConfig, GRPOTrainer
import wandb

from preft import (
    ADAPTER_REGISTRY,
    PreftConfig,
    get_preft_model,
)
from preft.trainer.callbacks import LoRALoggingCallback, ReFTLoggingCallback
from rl_training.config import ExperimentConfig, build_experiment_config, build_run_name
from rl_training.data import prepare_dataset
from rl_training.rewards import check_answer_correct, compute_reward, extract_answer, make_reward_function
from rl_training.trainer_trl_hooks import (
    ReFTGRPOTrainerTRLHooks,
    vllm_reft_fork_available,
)
# from rl_training.tokenizer_compat import normalize_checkpoint_tokenizer_config


# Try importing PEFT for LoRA
try:
    from peft import LoraConfig, get_peft_model
    PEFT_AVAILABLE = True
except ImportError:
    PEFT_AVAILABLE = False

def setup_model(config: ExperimentConfig):
    """Setup model based on experiment configuration."""
    print(f"\n{'='*60}")
    print(f"Setting up model: {config.mode}")
    print(f"{'='*60}")

    # Load base model
    print(f"Loading model: {config.model_name}")
    attn_impl = "flash_attention_2" if config.use_flash_attn else "sdpa"
    model = AutoModelForCausalLM.from_pretrained(
        config.model_name,
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        attn_implementation=attn_impl,
    )

    if config.mode == "full_finetune":
        # Full fine-tuning: enable all parameters
        print("Enabling all parameters for full fine-tuning...")
        for param in model.parameters():
            param.requires_grad = True

    elif config.mode == "lora":
        # LoRA fine-tuning
        # Default alpha=32 to keep cross-rank comparisons apples-to-apples
        # (the rank*2 scaling makes small ranks under-effective). Sweeps that
        # explicitly want rank-scaled alpha — e.g., the prefill_lora_alpha
        # sweep — pass --lora_alpha and override.
        effective_alpha = config.lora_alpha if config.lora_alpha is not None else 32
        target_modules = (
            config.lora_modules or ["q_proj", "v_proj", "k_proj", "o_proj"])
        print(f"Applying LoRA: rank={config.lora_rank}, alpha={effective_alpha}, "
              f"dropout=0, backend={config.lora_backend}, "
              f"position={config.lora_position}, target_modules={target_modules}")

        for param in model.parameters():
            param.requires_grad = False

        if config.lora_backend == "preft":
            # Preft LoRA path: gating lives in LoRALinear.forward; saved
            # adapter dir is already peft-compatible (vLLM-loadable).
            from preft import PreftConfig as _PreftConfig
            preft_cfg = _PreftConfig(
                layer_indices=list(range(model.config.num_hidden_layers)),
                use_residual=False,
                use_lora=True,
                lora_target_modules=target_modules,
                lora_alpha=float(effective_alpha),
                lora_dropout=0.0,
                lora_rank=config.lora_rank,
                position=config.lora_position,
                hidden_size=model.config.hidden_size,
                base_model_name_or_path=config.model_name,
            )
            model = get_preft_model(model, preft_cfg)
            model.print_trainable_parameters()
        else:
            if not PEFT_AVAILABLE:
                raise RuntimeError("peft required for LoRA backend='peft'. "
                                   "Install with: pip install peft, or use "
                                   "--lora_backend preft.")
            lora_config = LoraConfig(
                r=config.lora_rank,
                lora_alpha=effective_alpha,
                target_modules=target_modules,
                lora_dropout=0.0,
                bias="none",
                task_type="CAUSAL_LM",
            )
            model = get_peft_model(model, lora_config)
            model.print_trainable_parameters()

    else:
        # ReFT intervention
        print(f"Applying ReFT: rank={config.reft_rank}, layers={config.reft_layers}")

        preft_config = PreftConfig(
            layer_indices=config.reft_layers,
            adapter_type=config.mode,
            low_rank_dim=config.reft_rank,
            position=config.reft_position,
            use_residual=config.use_residual,
            frozen=False,
            logging=config.use_wandb,
        )
        model = get_preft_model(model, preft_config)
        model.print_trainable_parameters()

    return model


def _load_lora_init_weights(model, init_dir: str, lora_backend: str):
    """Copy LoRA weights from a peft-format adapter dir into *model* in-place.

    Used by ``--lora_init_from`` to pin the initial LoRA state across paired
    peft/preft runs so the only remaining source of trajectory divergence
    is vLLM rollout RNG (not init-RNG order, which differs between backends).

    Accepts either a directory containing ``adapter_model.safetensors`` or
    a path to that file directly.
    """
    import os as _os
    from safetensors.torch import load_file as _load_file

    weights_path = init_dir
    if _os.path.isdir(init_dir):
        weights_path = _os.path.join(init_dir, "adapter_model.safetensors")
    if not _os.path.exists(weights_path):
        raise FileNotFoundError(
            f"--lora_init_from: missing weights file at {weights_path}")

    state_dict = _load_file(weights_path)
    print(f"  Loading initial LoRA from {weights_path} "
          f"({len(state_dict)} tensors)")

    if lora_backend == "preft":
        # PreftModel._load_adapter_state_dict accepts peft on-disk keys directly.
        actual = model.module if hasattr(model, "module") else model
        actual._load_adapter_state_dict(state_dict)
    else:
        # peft path: set_peft_model_state_dict handles the .default. insertion.
        from peft import set_peft_model_state_dict
        actual = model.module if hasattr(model, "module") else model
        set_peft_model_state_dict(actual, state_dict)
    print(f"  Loaded initial LoRA weights ({lora_backend} backend)")


def train_with_grpo(config: ExperimentConfig):
    """Train model using GRPO (Group Relative Policy Optimization)."""
    # Suppress the TF32 warning — vLLM triggers it, we want high throughput not high precision here
    import torch as _torch
    _torch.set_float32_matmul_precision("high")

    set_seed(config.seed)

    # Setup run name
    if config.run_name is None:
        config.run_name = build_run_name(config)

    if config.use_vllm and config.mode in ADAPTER_REGISTRY:
        if not vllm_reft_fork_available():
            raise RuntimeError(
                "ReFT with --use_vllm now requires the forked vLLM integration API. "
                "Install the local forked vLLM package before launching training."
            )

    print(f"\n{'='*60}")
    print(f"Starting GRPO training: {config.run_name}")
    print(f"{'='*60}")
    print(f"Mode: {config.mode}")
    print(f"Dataset: {config.dataset}")
    print(f"Learning rate: {config.learning_rate:.0e}")

    # Initialize wandb if requested
    if config.use_wandb:
        wandb.init(
            entity="reft-aryaman",
            project=f"grpo-{config.dataset}-{config.mode}",
            name=config.run_name,
            config=config.__dict__,
        )

    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(config.model_name)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Base Llama models ship with chat_template=None. Use a SIMPLE template
    # that adds no special tokens — just `User:`/`Assistant:` labels around
    # bos/eos. Matches train_sft.py fallback.
    if tokenizer.chat_template is None:
        tokenizer.chat_template = (
            "{{ bos_token }}"
            "{% for m in messages %}"
            "{% if m['role'] == 'user' %}User: {{ m['content'] }}\n"
            "{% elif m['role'] == 'assistant' %}Assistant: {{ m['content'] }}{{ eos_token }}"
            "{% elif m['role'] == 'system' %}{{ m['content'] }}\n"
            "{% endif %}"
            "{% endfor %}"
            "{% if add_generation_prompt %}Assistant: {% endif %}"
        )
        print(f"  chat_template set to simple User:/Assistant: format")

    # Load dataset
    print(f"\nLoading {config.dataset} dataset...")
    train_dataset = prepare_dataset(config.dataset, tokenizer, split="train",
                                     max_samples=config.max_train_samples)
    print(f"Loaded {len(train_dataset)} training examples")

    # Load eval dataset
    eval_dataset = None
    if config.eval_steps and config.eval_steps > 0:
        if config.dataset == "tulu":
            eval_split = "train"  # tulu has no test split; use held-out from train
            eval_dataset_name = config.dataset
        elif config.dataset == "mbpp":
            # MBPP is the train task; HumanEval is the standard held-out eval.
            eval_split = "test"
            eval_dataset_name = "humaneval"
        else:
            eval_split = "test"
            eval_dataset_name = config.dataset
        eval_dataset = prepare_dataset(eval_dataset_name, tokenizer, split=eval_split,
                                       max_samples=config.max_eval_samples)
        print(f"Loaded {len(eval_dataset)} eval examples "
              f"(dataset={eval_dataset_name}, split={eval_split})")

    # Setup model
    model = setup_model(config)

    # Optional: override LoRA init weights from a peft-format dir so paired
    # peft/preft parity runs start from byte-identical weights.
    if config.lora_init_from and config.mode == "lora":
        _load_lora_init_weights(model, config.lora_init_from, config.lora_backend)

    # Create output directory
    output_dir = os.path.join(config.output_dir, config.run_name)
    os.makedirs(output_dir, exist_ok=True)

    # Debug: Check dataset format
    print(f"\n=== Dataset Debug ===")
    print(f"Dataset columns: {train_dataset.column_names}")
    print(f"First example: {train_dataset[0]}")
    print("=" * 40)

    # GRPO training config
    training_args = GRPOConfig(
        output_dir=output_dir,
        num_train_epochs=config.num_train_epochs,
        per_device_train_batch_size=config.batch_size,
        gradient_accumulation_steps=config.gradient_accumulation_steps,
        learning_rate=config.learning_rate,
        warmup_ratio=config.warmup_ratio,
        weight_decay=config.weight_decay,
        logging_steps=10,
        save_strategy="steps",
        save_steps=500,
        save_total_limit=1,
        report_to="wandb" if config.use_wandb else "none",
        seed=config.seed,
        remove_unused_columns=False,
        gradient_checkpointing=config.gradient_checkpointing,
        bf16=True,
        max_completion_length=config.max_completion_length,
        num_generations=config.num_completions,
        loss_type=config.loss_type,
        mask_truncated_completions=config.mask_truncated_completions,
        use_vllm=config.use_vllm,
        vllm_mode=config.vllm_mode,
        vllm_gpu_memory_utilization=config.vllm_gpu_memory_utilization,
        vllm_max_model_length=config.max_prompt_length + config.max_completion_length,
        vllm_server_port=config.vllm_server_port,
        vllm_group_port=config.vllm_group_port,
        vllm_server_timeout=config.vllm_server_timeout,
        vllm_importance_sampling_correction=config.vllm_mode == "server",
        eval_strategy="steps" if eval_dataset is not None else "no",
        eval_steps=config.eval_steps if eval_dataset is not None else None,
        num_generations_eval=1,
        per_device_eval_batch_size=config.batch_size,
        use_liger_kernel=config.use_liger_kernel,
        beta=config.beta_value if config.teacher_adapter_path else 0.0,
    )

    # Attach lora_position to training_args so hooks can read it
    training_args.lora_position = config.lora_position
    training_args.teacher_adapter_path = config.teacher_adapter_path

    reward_function = make_reward_function(config.dataset)

    # Use the preft-aware trainer whenever the model is a PreftModel.  Its
    # `_save` routes through model.save_pretrained() (adapter-only); HF Trainer's
    # default `_save` tries to safetensors-serialize the full base model and
    # crashes on tied embed_tokens / lm_head weights.
    from preft import PreftModel
    is_preft_model = isinstance(model, PreftModel) or hasattr(model, "preft_config")
    if (
        is_preft_model
        or config.mode in ADAPTER_REGISTRY
        or config.lora_position == "prefill"
        or config.teacher_adapter_path
    ):
        trainer_cls = ReFTGRPOTrainerTRLHooks
    else:
        trainer_cls = GRPOTrainer
    callbacks = []
    if config.mode == "lora" and config.use_wandb:
        callbacks.append(
            LoRALoggingCallback(
                model=model,
                lora_rank=config.lora_rank,
                lora_alpha=config.lora_alpha if config.lora_alpha is not None else 32,
            )
        )
    if config.mode in ADAPTER_REGISTRY and config.use_wandb:
        callbacks.append(ReFTLoggingCallback(model=model))

    trainer = trainer_cls(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        reward_funcs=reward_function,
        callbacks=callbacks or None,
    )

    # Verify vLLM ReFT setup
    if config.use_vllm and config.mode in ADAPTER_REGISTRY:
        from rl_training.trainer_trl_hooks import vllm_reft_fork_available as _vllm_check, trl_grpo_supports_reft_hooks as _trl_check
        print(f"\n=== vLLM ReFT Diagnostics ===")
        print(f"  vllm.reft fork available: {_vllm_check()}")
        print(f"  TRL supports ReFT hooks: {_trl_check()}")
        print(f"  _is_reft_model: {trainer._is_reft_model()}")
        llm = trainer._get_llm()
        if llm is not None:
            try:
                vllm_model = llm.llm_engine.model_executor.driver_worker.model_runner.model
                reft_layers = sum(1 for layer in vllm_model.model.layers if hasattr(layer, 'reft_adapter'))
                print(f"  vLLM ReFT adapters: {reft_layers}/{len(vllm_model.model.layers)} layers")
            except Exception as e:
                print(f"  vLLM model inspection failed: {e}")
        else:
            print(f"  vLLM LLM not yet initialized (will init on first generate)")
        print(f"{'='*28}\n")

    # Train
    print("\nStarting GRPO training...")
    resume_ckpt = None
    if config.resume_from_checkpoint:
        # Look for any checkpoint-* subdir written by the HF Trainer
        existing = sorted(
            (d for d in os.listdir(output_dir) if d.startswith("checkpoint-")),
            key=lambda d: int(d.split("-")[1]),
        ) if os.path.isdir(output_dir) else []
        if existing:
            resume_ckpt = os.path.join(output_dir, existing[-1])
            print(f"Resuming from checkpoint: {resume_ckpt}")
        else:
            print("No checkpoint found in output_dir — starting from scratch")
    trainer.train(resume_from_checkpoint=resume_ckpt)

    # Final eval — on the FULL test set (no max_samples cap), so the reported
    # number is comparable across runs and subjects.
    if eval_dataset is not None:
        eval_split = "test"
        full_eval_name = "humaneval" if config.dataset == "mbpp" else config.dataset
        full_eval_dataset = prepare_dataset(
            full_eval_name, tokenizer, split=eval_split, max_samples=None,
        )
        print(f"\nRunning final evaluation on full {eval_split} split "
              f"({len(full_eval_dataset)} examples)...")
        metrics = trainer.evaluate(eval_dataset=full_eval_dataset)
        print(f"Final eval metrics: {metrics}")
        # Persist for offline plotting / overrides in plots/_util.py.
        full_eval_path = os.path.join(output_dir, "full_eval.json")
        with open(full_eval_path, "w") as f:
            json.dump({
                "run_name": config.run_name,
                "dataset": config.dataset,
                "eval_split": eval_split,
                "num_examples": len(full_eval_dataset),
                "eval_reward": metrics.get("eval_reward")
                                  or metrics.get("eval/reward")
                                  or metrics.get("eval_rewards/reward"),
                "metrics": metrics,
            }, f, indent=2)
        print(f"Wrote full-set eval to {full_eval_path}")

    # Save final adapter to the bare output_dir so downstream eval scripts
    # don't have to crawl into checkpoint-N/. (HF Trainer also wrote the
    # latest checkpoint there; this just mirrors the final state to the
    # canonical location. Run-name collisions across positions can no
    # longer happen now that build_run_name encodes reft_position.)
    print(f"\nSaving model to {output_dir}")
    if config.mode in ("full_finetune", "lora") or config.mode in ADAPTER_REGISTRY:
        model.save_pretrained(output_dir)

    tokenizer.save_pretrained(output_dir)

    # Save config
    with open(os.path.join(output_dir, "training_config.json"), "w") as f:
        json.dump(config.__dict__, f, indent=2)

    if config.use_wandb:
        wandb.finish()

    print(f"Training complete! Model saved to {output_dir}")


def run_lr_sweep(args):
    """Run learning rate sweep."""
    learning_rates = [float(lr) for lr in args.lr_sweep.split(",")]
    print(f"\nRunning LR sweep with rates: {learning_rates}")

    for lr in learning_rates:
        config = build_experiment_config(args, learning_rate=lr)

        try:
            train_with_grpo(config)
        except Exception as e:
            print(f"Error training with lr={lr}: {e}")
            import traceback

            traceback.print_exc()


def main():
    parser = argparse.ArgumentParser(description="Train with GRPO on reasoning tasks")

    # Mode
    parser.add_argument(
        "--mode", type=str, choices=["full_finetune", "lora"] + list(ADAPTER_REGISTRY), help="Training mode"
    )

    # Model and dataset
    parser.add_argument("--model", type=str, default="meta-llama/Llama-3.1-8B", help="Model name or path")
    parser.add_argument("--dataset", type=str, default="gsm8k",
                        choices=["gsm8k", "math", "tulu", "mbpp"],
                        help="Training dataset. (HumanEval is eval-only; "
                             "selected automatically when --dataset=mbpp.)")

    # Learning rate
    parser.add_argument("--lr", type=float, default=5e-5, help="Learning rate")
    parser.add_argument("--lr_sweep", type=str, default=None, help="Comma-separated learning rates for sweep")

    # LoRA config
    parser.add_argument("--lora_rank", type=int, default=8)
    parser.add_argument("--lora_alpha", type=int, default=None)
    parser.add_argument("--lora_modules", type=str, default="q_proj,v_proj,k_proj,o_proj,gate_proj,up_proj,down_proj")
    parser.add_argument("--lora_position", type=str, default="all", choices=["all", "prefill"])
    parser.add_argument("--lora_backend", type=str, default="peft",
                        choices=["peft", "preft"],
                        help="Library hosting the LoRA modules. 'peft' uses "
                             "peft.LoraConfig + prefill_only_wrapper monkey-patch. "
                             "'preft' uses PreftConfig(use_lora=True) — gating in "
                             "LoRALinear.forward, weight sync via "
                             "PreftModel.sync_lora_to_vllm. Saved checkpoints are "
                             "peft-format either way (vLLM-loadable).")
    parser.add_argument("--lora_init_from", type=str, default=None,
                        help="Path to a peft-format adapter directory whose "
                             "lora_A/lora_B weights are copied into the freshly-"
                             "built model before training. Eliminates init-RNG "
                             "drift between paired peft/preft parity runs.")

    # LoReFT config
    parser.add_argument("--reft_rank", type=int, default=4)
    parser.add_argument("--reft_layers", type=str, default="all")
    parser.add_argument("--reft_position", type=str, default="last")
    parser.add_argument("--use_residual", action="store_true", default=True, help="Intervene on block output default")
    # Training config
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_completions", type=int, default=4, help="K samples per prompt for GRPO")
    parser.add_argument("--max_prompt_length", type=int, default=2048)
    parser.add_argument("--max_completion_length", type=int, default=2048)
    parser.add_argument("--warmup_ratio", type=float, default=0.1)
    parser.add_argument("--weight_decay", type=float, default=0.0)
    parser.add_argument("--beta_value", type=float, default=0.5, help="The coefficient of KL divergence")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
    parser.add_argument(
        "--loss_type", type=str, default="dapo", choices=["grpo", "dr_grpo", "dapo", "bnpo", "cispo", "sapo"]
    )
    parser.add_argument(
        "--mask_truncated_completions", action="store_true", help="Mask truncated completions in loss computation"
    )

    # Memory
    parser.add_argument("--gradient_checkpointing", action="store_true", help="Enable gradient checkpointing to reduce activation memory (slower but uses ~60-70%% less activation memory)")
    parser.add_argument("--use_flash_attn", action="store_true", help="Use flash_attention_2 instead of sdpa")
    parser.add_argument("--use_liger_kernel", action="store_true", help="Use Liger fused kernels for memory efficiency")

    # vLLM
    parser.add_argument("--use_vllm", action="store_true", help="Use vLLM for faster generation. For ReFT modes, applies adapters during prefill only.")
    parser.add_argument("--vllm_mode", type=str, default="colocate", choices=["colocate", "server"],
                        help="vLLM mode: 'colocate' shares GPU, 'server' uses a separate GPU for generation")
    parser.add_argument("--vllm_gpu_memory_utilization", type=float, default=0.3,
                        help="Fraction of GPU memory for vLLM KV cache (colocate mode)")
    parser.add_argument("--vllm_server_port", type=int, default=8000,
                        help="Port of the vLLM server (server mode)")
    parser.add_argument("--vllm_group_port", type=int, default=51216,
                        help="Port for the weight update group (server mode)")
    parser.add_argument("--vllm_server_timeout", type=float, default=240.0,
                        help="Timeout in seconds waiting for vLLM server (server mode)")

    # Eval
    parser.add_argument("--eval_steps", type=int, default=None, help="Evaluate every N steps (disabled if not set)")
    parser.add_argument("--max_eval_samples", type=int, default=256, help="Max eval samples from test split")
    parser.add_argument("--max_train_samples", type=int, default=None, help="Max training samples (default: all)")

    # Output
    parser.add_argument("--output_dir", type=str, default="./outputs")
    parser.add_argument("--run_name", type=str, default=None,
                        help="Override the auto-generated wandb display name / output dir.")
    parser.add_argument("--resume_from_checkpoint", action="store_true",
                        help="Resume from latest checkpoint in output_dir/run_name if one exists")
    parser.add_argument("--use_wandb", action="store_true")
    parser.add_argument("--seed", type=int, default=42)

    # Config file
    parser.add_argument(
        "--config", type=str, default=None, help="Path to YAML config file. CLI args override config values."
    )

    # Parse: load YAML defaults first, then let CLI args override
    args, remaining = parser.parse_known_args()
    if args.config:
        with open(args.config) as f:
            config_dict = yaml.safe_load(f)
        # Set YAML values as defaults, then re-parse so CLI args take priority
        parser.set_defaults(**config_dict)
        # use_wandb in YAML is a bool, but argparse expects store_true
        if config_dict.get("use_wandb"):
            parser.set_defaults(use_wandb=True)
    args = parser.parse_args()

    if args.mode is None:
        parser.error("--mode is required (either via CLI or config file)")

    if args.lr_sweep:
        run_lr_sweep(args)
    else:
        config = build_experiment_config(args)
        train_with_grpo(config)


if __name__ == "__main__":
    main()