ps2_slm_rl/src/model/loader.py

"""
src/model/loader.py

Loads Phi-3-Mini with optional LoRA wrapping and prints module names.

HALLUCINATION NOTE:
  - The target_modules list for Phi-3-mini is based on the published architecture.
    Phi-3-mini uses a fused 'gate_up_proj' (not separate 'gate_proj' and 'up_proj').
    Run --print-modules to verify against the actual loaded model before training.
  - 'flash_attention_2' requires CUDA >= 11.8 and an Ampere+ GPU (A100, RTX 3090+).
    Kaggle P100 does NOT support it. Use 'eager' or 'sdpa' on Kaggle.
  - 'bfloat16' is not supported on Pascal GPUs (P100). Use 'float16' instead.
"""

from __future__ import annotations

import argparse
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, TaskType, get_peft_model, PeftModel


PHI3_MINI_ID = "microsoft/Phi-3-mini-4k-instruct"

# Verified Phi-3-mini module names via architecture inspection.
# 'gate_up_proj' is the fused projection (Phi-3 specific).
PHI3_LORA_TARGET_MODULES = [
    "q_proj",
    "k_proj",
    "v_proj",
    "o_proj",
    "gate_up_proj",
    "down_proj",
]


def get_dtype(dtype_str: str) -> torch.dtype:
    mapping = {
        "bfloat16": torch.bfloat16,
        "float16": torch.float16,
        "float32": torch.float32,
    }
    if dtype_str not in mapping:
        raise ValueError(f"Unknown dtype: {dtype_str}. Choose from {list(mapping)}")
    return mapping[dtype_str]


def load_base_model(
    model_id: str = PHI3_MINI_ID,
    torch_dtype: str = "bfloat16",
    attn_implementation: str = "eager",
    device_map: str = "auto",
) -> AutoModelForCausalLM:
    """
    Load the base causal LM.

    Args:
        model_id:            HuggingFace model ID or local path.
        torch_dtype:         "bfloat16", "float16", or "float32".
        attn_implementation: "eager", "sdpa", or "flash_attention_2".
                             Use "eager" on Kaggle P100.
        device_map:          "auto" for multi-GPU, "cuda:0" for single GPU.
    """
    dtype = get_dtype(torch_dtype)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        attn_implementation=attn_implementation,
        device_map=device_map,
    )
    model.config.use_cache = False  # required for gradient checkpointing
    return model


def load_tokenizer(model_id: str = PHI3_MINI_ID) -> AutoTokenizer:
    tok = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
    # Phi-3 tokenizer has pad_token; if missing, set to eos_token
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"  # required for decoder-only models in batch generation
    return tok


def make_lora_config(
    r: int = 16,
    lora_alpha: int = 32,
    lora_dropout: float = 0.05,
    bias: str = "none",
    target_modules: list[str] | None = None,
) -> LoraConfig:
    if target_modules is None:
        target_modules = PHI3_LORA_TARGET_MODULES
    return LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=r,
        lora_alpha=lora_alpha,
        lora_dropout=lora_dropout,
        bias=bias,
        target_modules=target_modules,
    )


def load_model_with_lora(
    model_id: str = PHI3_MINI_ID,
    lora_cfg: LoraConfig | None = None,
    torch_dtype: str = "bfloat16",
    attn_implementation: str = "eager",
) -> tuple[AutoModelForCausalLM, AutoTokenizer]:
    """
    Load base model and wrap with LoRA for training.
    Returns (peft_model, tokenizer).
    """
    model = load_base_model(model_id, torch_dtype, attn_implementation)
    tok = load_tokenizer(model_id)

    if lora_cfg is None:
        lora_cfg = make_lora_config()

    model = get_peft_model(model, lora_cfg)
    model.print_trainable_parameters()
    return model, tok


def load_sft_checkpoint_for_grpo(
    base_model_id: str = PHI3_MINI_ID,
    sft_checkpoint_path: str = "checkpoints/sft",
    torch_dtype: str = "bfloat16",
    attn_implementation: str = "eager",
) -> tuple[AutoModelForCausalLM, AutoModelForCausalLM, AutoTokenizer]:
    """
    Load the SFT checkpoint as the trainable policy and a frozen copy as the
    reference model for KL divergence computation.

    Returns (policy_model, ref_model, tokenizer).

    HALLUCINATION NOTE: GRPOTrainer in trl 1.4.0 accepts a 'model' and infers
    the reference model internally when using PEFT (it keeps the frozen base).
    Passing a separate ref_model is supported but not required with PEFT.
    See scripts/train_grpo.py for how this is used.
    """
    tok = load_tokenizer(base_model_id)
    dtype = get_dtype(torch_dtype)

    # Policy: SFT checkpoint with LoRA weights (will continue training)
    base = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=dtype,
        attn_implementation=attn_implementation,
        device_map="auto",
    )
    base.config.use_cache = False
    policy = PeftModel.from_pretrained(base, sft_checkpoint_path, is_trainable=True)

    # Reference: frozen copy of the SFT checkpoint
    ref_base = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=dtype,
        attn_implementation=attn_implementation,
        device_map="auto",
    )
    ref_base.config.use_cache = False
    ref_model = PeftModel.from_pretrained(ref_base, sft_checkpoint_path, is_trainable=False)
    for param in ref_model.parameters():
        param.requires_grad = False

    return policy, ref_model, tok


def merge_and_save(
    base_model_id: str,
    lora_checkpoint_path: str,
    output_path: str,
    torch_dtype: str = "bfloat16",
) -> None:
    """
    Merge LoRA weights into the base model and save for inference.
    The merged model has no adapter overhead.
    """
    dtype = get_dtype(torch_dtype)
    base = AutoModelForCausalLM.from_pretrained(
        base_model_id, torch_dtype=dtype, device_map="cpu"
    )
    model = PeftModel.from_pretrained(base, lora_checkpoint_path)
    model = model.merge_and_unload()
    model.save_pretrained(output_path)

    tok = load_tokenizer(base_model_id)
    tok.save_pretrained(output_path)
    print(f"Merged model saved to {output_path}")


def print_model_modules(model_id: str = PHI3_MINI_ID) -> None:
    """Print all named modules — use this to verify target_modules before training."""
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
    for name, module in model.named_modules():
        if any(x in name for x in ["proj", "fc", "linear", "embed", "head"]):
            print(f"  {name}: {module.__class__.__name__}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--print-modules", action="store_true",
                        help="Print module names for the model")
    parser.add_argument("--model-id", default=PHI3_MINI_ID)
    args = parser.parse_args()

    if args.print_modules:
        print(f"\nModules in {args.model_id}:")
        print_model_modules(args.model_id)