src/loratune.py

#!/usr/bin/env python3
"""Centralized Alpaca LoRA finetuning for post-pruned models."""

import argparse
import itertools
import json
import os
from types import SimpleNamespace
from pathlib import Path

import torch
from contextlib import nullcontext

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig
from transformers.models.auto.configuration_auto import CONFIG_MAPPING
import ppl_eval

from fuse_layers_data import FixedSeqDataset, load_instruction_records
from fuse_layers_distill import LoRALinear, apply_lora_adapters, merge_lora_adapters

try:
    from tqdm import tqdm
except Exception:  # pragma: no cover
    tqdm = None


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run centralized Alpaca LoRA finetuning.")
    parser.add_argument("--base_model", required=True, help="Path or HF model id to finetune")
    parser.add_argument("--output_dir", required=True, help="Directory to save merged model")
    parser.add_argument("--device", default="cuda", help="Training device")
    parser.add_argument(
        "--dtype",
        default="bfloat16",
        choices=["float32", "float16", "bfloat16"],
        help="Model load/training dtype",
    )
    parser.add_argument("--trust_remote_code", action="store_true", help="Enable trust_remote_code")
    parser.add_argument("--seed", type=int, default=42, help="Random seed")

    parser.add_argument(
        "--instruction_dataset",
        default="yahma/alpaca-cleaned",
        help="HF dataset name for Alpaca-style instruction data",
    )
    parser.add_argument("--instruction_config", default=None, help="Optional dataset config")
    parser.add_argument("--instruction_split", default="train", help="Dataset split")
    parser.add_argument("--instruction_field_instruction", default="instruction")
    parser.add_argument("--instruction_field_input", default="input")
    parser.add_argument("--instruction_field_output", default="output")
    parser.add_argument("--max_samples", type=int, default=0, help="Limit instruction samples (0 = all)")
    parser.add_argument("--seq_len", type=int, default=1024, help="Training sequence length")
    parser.add_argument("--batch_size", type=int, default=64, help="Global batch size")
    parser.add_argument("--micro_batch_size", type=int, default=4, help="Per-step micro-batch size")
    parser.add_argument("--epochs", type=float, default=1.0, help="Training epochs")
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Learning rate")
    parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay")
    parser.add_argument("--max_grad_norm", type=float, default=1.0, help="Gradient clipping norm")
    parser.add_argument("--log_steps", type=int, default=100, help="Log every N optimizer steps")
    parser.add_argument(
        "--save_steps",
        type=int,
        default=200,
        help="Save LoRA adapter checkpoints every N optimizer steps (0 = disable)",
    )
    parser.add_argument(
        "--no_wikitext2_ppl_on_log",
        dest="wikitext2_ppl_on_log",
        action="store_false",
        help="Disable Wikitext-2 perplexity evaluation at loss log steps",
    )
    parser.set_defaults(wikitext2_ppl_on_log=True)
    parser.add_argument("--wikitext2_ppl_seq_len", type=int, default=128)
    parser.add_argument("--wikitext2_ppl_batch_size", type=int, default=8)
    parser.add_argument("--wikitext2_ppl_max_batches", type=int, default=None)

    parser.add_argument("--lora_rank", type=int, default=8, help="LoRA rank")
    parser.add_argument("--lora_alpha", type=float, default=16.0, help="LoRA alpha")
    parser.add_argument("--lora_dropout", type=float, default=0.0, help="LoRA dropout")
    parser.add_argument(
        "--lora_target_modules",
        nargs="*",
        default=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "down_proj", "up_proj"],
        help="Linear module suffixes to LoRA-wrap",
    )

    return parser.parse_args()


def get_dtype(name: str) -> torch.dtype:
    return {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
    }[name]


def seed_all(seed: int) -> None:
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def normalize_config(config):
    layer_types = getattr(config, "layer_types", None)
    num_hidden_layers = getattr(config, "num_hidden_layers", None)
    if layer_types is not None and num_hidden_layers is not None and len(layer_types) != num_hidden_layers:
        config.layer_types = list(layer_types[:num_hidden_layers])
    if getattr(config, "_attn_implementation", None) is None:
        config._attn_implementation = "eager"
    return config


def load_normalized_config(base_model: str, trust_remote_code: bool):
    config_dict, unused_kwargs = PretrainedConfig.get_config_dict(base_model, trust_remote_code=trust_remote_code)
    layer_types = config_dict.get("layer_types")
    num_hidden_layers = config_dict.get("num_hidden_layers")
    if layer_types is not None and num_hidden_layers is not None and len(layer_types) != num_hidden_layers:
        config_dict["layer_types"] = list(layer_types[:num_hidden_layers])
    if config_dict.get("_attn_implementation") is None:
        config_dict["_attn_implementation"] = "eager"
    model_type = config_dict["model_type"]
    config_class = CONFIG_MAPPING[model_type]
    config = config_class.from_dict(config_dict, **unused_kwargs)
    return normalize_config(config)


def validate_local_model_dir(base_path: Path) -> None:
    if not base_path.exists() or not base_path.is_dir():
        return

    has_config = (base_path / "config.json").is_file()
    has_weights = any(
        (base_path / name).is_file()
        for name in (
            "model.safetensors",
            "model.safetensors.index.json",
            "pytorch_model.bin",
            "pytorch_model.bin.index.json",
        )
    )
    if has_config and has_weights:
        return

    raise SystemExit(
        "Local --base_model points to an incomplete HF model directory: "
        f"{base_path}. Expected at least config.json and model weights. "
        "Set --base_model/BASE_MODEL to a saved HF model directory."
    )


def load_base_artifacts(args: argparse.Namespace):
    base_path = Path(args.base_model)
    if base_path.is_file() and base_path.suffix == ".bin":
        checkpoint = torch.load(str(base_path), map_location="cpu", weights_only=False)
        if not isinstance(checkpoint, dict) or "model" not in checkpoint or "tokenizer" not in checkpoint:
            raise SystemExit("Expected a .bin checkpoint dict with `model` and `tokenizer` entries.")
        model = checkpoint["model"]
        tokenizer = checkpoint["tokenizer"]
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
        return model, tokenizer

    validate_local_model_dir(base_path)
    tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=args.trust_remote_code)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token or tokenizer.unk_token
    config = load_normalized_config(args.base_model, trust_remote_code=args.trust_remote_code)
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model,
        config=config,
        torch_dtype=get_dtype(args.dtype),
        trust_remote_code=args.trust_remote_code,
    )
    return model, tokenizer


def build_training_loader(args: argparse.Namespace, tokenizer) -> torch.utils.data.DataLoader:
    num_samples = args.max_samples if args.max_samples > 0 else 0
    records = load_instruction_records(args, num_samples)
    if not records:
        raise SystemExit("No instruction records were loaded.")
    dataset = FixedSeqDataset(records, tokenizer, args.seq_len)
    return torch.utils.data.DataLoader(dataset, batch_size=args.micro_batch_size, shuffle=True)


def save_lora_adapters(
    model: torch.nn.Module, args: argparse.Namespace, subdir: str = "lora_adapter"
) -> str:
    adapter_dir = os.path.join(args.output_dir, subdir)
    os.makedirs(adapter_dir, exist_ok=True)

    adapter_state = {}
    adapter_modules = {}
    for module_name, module in model.named_modules():
        if not isinstance(module, LoRALinear):
            continue
        adapter_modules[module_name] = {
            "rank": module.rank,
            "alpha": module.alpha,
            "scaling": module.scaling,
            "dropout": getattr(module.dropout, "p", 0.0),
            "base_layer_class": type(module.base).__name__,
            "in_features": module.base.in_features,
            "out_features": module.base.out_features,
        }
        adapter_state[f"{module_name}.lora_A.weight"] = module.lora_A.weight.detach().cpu()
        adapter_state[f"{module_name}.lora_B.weight"] = module.lora_B.weight.detach().cpu()

    torch.save(adapter_state, os.path.join(adapter_dir, "adapter_model.bin"))
    with open(os.path.join(adapter_dir, "adapter_config.json"), "w", encoding="utf-8") as handle:
        json.dump(
            {
                "base_model": args.base_model,
                "lora_rank": args.lora_rank,
                "lora_alpha": args.lora_alpha,
                "lora_dropout": args.lora_dropout,
                "lora_target_modules": list(args.lora_target_modules),
                "batch_size": args.batch_size,
                "micro_batch_size": args.micro_batch_size,
                "grad_accum_steps": args.grad_accum_steps,
                "modules": adapter_modules,
            },
            handle,
            indent=2,
        )
    return adapter_dir


def prepare_wikitext2_eval(args: argparse.Namespace, model, tokenizer):
    if not args.wikitext2_ppl_on_log:
        return None
    return ppl_eval.prepare_ppl_dataloaders(
        tokenizer=tokenizer,
        datasets=["wikitext"],
        configs=["wikitext-2-raw-v1"],
        split="test",
        text_field=None,
        num_samples=0,
        seq_len=args.wikitext2_ppl_seq_len,
        batch_size=args.wikitext2_ppl_batch_size,
        seed=args.seed,
        shuffle=False,
        model_family="auto",
        add_bos="auto",
        cache_dir=None,
        num_workers=0,
        model=model,
    )


def train(model: torch.nn.Module, dataloader, args: argparse.Namespace, wikitext2_eval_dataloaders=None) -> dict:
    lora_args = SimpleNamespace(
        lora_rank=args.lora_rank,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        lora_target_modules=args.lora_target_modules,
        lora_respect_exclude_pairs=False,
        layer_path=None,
        exclude_pairs=None,
    )
    lora_modules = apply_lora_adapters(model, lora_args)
    lora_params = [param for module in lora_modules for param in module.lora_parameters()]

    optimizer = torch.optim.AdamW(
        lora_params,
        lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    model.train()

    device = torch.device(args.device)
    device_type = device.type
    amp_dtype = None
    if args.dtype == "float16":
        amp_dtype = torch.float16
    elif args.dtype == "bfloat16":
        amp_dtype = torch.bfloat16
    use_amp = amp_dtype is not None and device_type == "cuda"
    use_scaler = use_amp and amp_dtype == torch.float16
    scaler = torch.cuda.amp.GradScaler() if use_scaler else None

    full_epochs = int(args.epochs)
    fractional = args.epochs - full_epochs
    epoch_plan = [None] * full_epochs
    if fractional > 1e-8:
        frac_batches = max(1, int(round(fractional * len(dataloader))))
        epoch_plan.append(frac_batches)

    optimizer.zero_grad(set_to_none=True)
    optimizer_step = 0
    seen_batches = 0
    last_loss = None
    ppl_history = []

    for epoch_idx, max_batches in enumerate(epoch_plan, start=1):
        iterator = dataloader if max_batches is None else itertools.islice(dataloader, max_batches)
        if tqdm is not None:
            iterator = tqdm(iterator, desc=f"LoRA epoch {epoch_idx}", unit="batch", total=max_batches)
        for batch in iterator:
            input_ids = batch[0].to(args.device)
            attention_mask = batch[1].to(args.device)

            autocast_ctx = (
                torch.autocast(device_type=device_type, dtype=amp_dtype)
                if use_amp
                else nullcontext()
            )
            with autocast_ctx:
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
                logits = outputs.logits[:, :-1, :].contiguous()
                labels = input_ids[:, 1:].contiguous()
                mask = attention_mask[:, 1:].contiguous()
                ce_flat = torch.nn.functional.cross_entropy(
                    logits.view(-1, logits.size(-1)),
                    labels.view(-1),
                    reduction="none",
                )
                denom = mask.sum()
                if denom.item() == 0:
                    continue
                loss = (ce_flat * mask.reshape(-1).to(ce_flat.dtype)).sum() / denom

            last_loss = float(loss.detach().item())
            scaled_loss = loss / max(args.grad_accum_steps, 1)
            if use_scaler:
                scaler.scale(scaled_loss).backward()
            else:
                scaled_loss.backward()

            seen_batches += 1
            if seen_batches % max(args.grad_accum_steps, 1) != 0:
                continue

            if args.max_grad_norm is not None:
                if use_scaler:
                    scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(lora_params, args.max_grad_norm)
            if use_scaler:
                scaler.step(optimizer)
                scaler.update()
            else:
                optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            optimizer_step += 1

            if args.log_steps and optimizer_step % args.log_steps == 0:
                print(f"[loratune] step={optimizer_step} loss={last_loss:.6f}")
                if wikitext2_eval_dataloaders is not None:
                    prev_mode = model.training
                    model.eval()
                    ppl_results = ppl_eval.evaluate_ppl_dataloaders(
                        model,
                        wikitext2_eval_dataloaders,
                        args.device,
                        max_batches=args.wikitext2_ppl_max_batches,
                    )
                    ppl_history.append({"step": optimizer_step, "ppl": ppl_results})
                    print(f"[loratune] ppl step={optimizer_step} {ppl_results}")
                    if prev_mode:
                        model.train()

            if args.save_steps and optimizer_step % args.save_steps == 0:
                checkpoint_dir = save_lora_adapters(
                    model,
                    args,
                    subdir=os.path.join("checkpoints", f"step_{optimizer_step}"),
                )
                print(f"[loratune] saved adapter checkpoint to {checkpoint_dir}")

    adapter_dir = save_lora_adapters(model, args)
    merge_lora_adapters(model)
    return {
        "adapter_dir": adapter_dir,
        "optimizer_steps": optimizer_step,
        "seen_batches": seen_batches,
        "last_loss": last_loss,
        "wikitext2_ppl_history": ppl_history,
    }


def main() -> None:
    args = parse_args()
    if args.batch_size < 1:
        raise SystemExit("--batch_size must be >= 1")
    if args.micro_batch_size < 1:
        raise SystemExit("--micro_batch_size must be >= 1")
    args.grad_accum_steps = args.batch_size // args.micro_batch_size
    if args.grad_accum_steps < 1:
        raise SystemExit("--batch_size must be >= --micro_batch_size")

    seed_all(args.seed)
    os.makedirs(args.output_dir, exist_ok=True)

    model, tokenizer = load_base_artifacts(args)
    if args.dtype != "float32":
        model = model.to(get_dtype(args.dtype))
    model.to(args.device)

    dataloader = build_training_loader(args, tokenizer)
    wikitext2_eval_dataloaders = prepare_wikitext2_eval(args, model, tokenizer)
    metrics = train(model, dataloader, args, wikitext2_eval_dataloaders=wikitext2_eval_dataloaders)

    model.save_pretrained(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    with open(os.path.join(args.output_dir, "loratune_metrics.json"), "w", encoding="utf-8") as handle:
        json.dump(
            {
                "base_model": args.base_model,
                "instruction_dataset": args.instruction_dataset,
                "seq_len": args.seq_len,
                "batch_size": args.batch_size,
                "micro_batch_size": args.micro_batch_size,
                "grad_accum_steps": args.grad_accum_steps,
                "epochs": args.epochs,
                "learning_rate": args.learning_rate,
                "save_steps": args.save_steps,
                "lora_rank": args.lora_rank,
                "lora_alpha": args.lora_alpha,
                "lora_dropout": args.lora_dropout,
                **metrics,
            },
            handle,
            indent=2,
        )


if __name__ == "__main__":
    main()