gemma_code_llm/train_lora.py

from __future__ import annotations

import argparse
import inspect
from pathlib import Path

from .data import TokenizedCodeDataset, load_jsonl_records, split_records
from .modeling import load_model_and_tokenizer


def _csv(value: str) -> list[str]:
    return [part.strip() for part in value.split(",") if part.strip()]


def _import_training_stack():
    try:
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
        from transformers import DataCollatorForSeq2Seq, Trainer, TrainingArguments
    except Exception as exc:
        raise RuntimeError(
            "Missing training dependencies. Install them with: pip install -r requirements.txt"
        ) from exc
    return LoraConfig, get_peft_model, prepare_model_for_kbit_training, DataCollatorForSeq2Seq, Trainer, TrainingArguments


def _build_training_args(TrainingArguments, args, torch):
    signature = inspect.signature(TrainingArguments.__init__).parameters
    cuda = torch.cuda.is_available()
    use_bf16 = cuda and torch.cuda.is_bf16_supported()
    use_fp16 = cuda and not use_bf16

    kwargs = {
        "output_dir": args.output_dir,
        "per_device_train_batch_size": args.batch_size,
        "gradient_accumulation_steps": args.gradient_accumulation_steps,
        "num_train_epochs": args.epochs,
        "learning_rate": args.learning_rate,
        "logging_steps": args.logging_steps,
        "save_steps": args.save_steps,
        "save_total_limit": args.save_total_limit,
        "warmup_ratio": args.warmup_ratio,
        "weight_decay": args.weight_decay,
        "optim": "adamw_torch",
        "report_to": "none",
        "remove_unused_columns": False,
        "gradient_checkpointing": args.gradient_checkpointing,
    }

    if "bf16" in signature:
        kwargs["bf16"] = use_bf16
    if "fp16" in signature:
        kwargs["fp16"] = use_fp16

    if args.validation_size > 0:
        strategy_name = "eval_strategy" if "eval_strategy" in signature else "evaluation_strategy"
        kwargs[strategy_name] = "steps"
        kwargs["eval_steps"] = args.eval_steps

    return TrainingArguments(**kwargs)


def main() -> int:
    parser = argparse.ArgumentParser(description="Fine-tune Gemma for code generation with LoRA.")
    parser.add_argument("--data", required=True, help="Path to JSONL training data.")
    parser.add_argument("--base-model", default="google/gemma-3-1b-it", help="Base Hugging Face model id.")
    parser.add_argument("--output-dir", default="outputs/gemma-code-lora", help="Where to save the LoRA adapter.")
    parser.add_argument("--max-records", type=int, default=None, help="Optional limit for quick tests.")
    parser.add_argument("--validation-size", type=float, default=0.0, help="Fraction of data for validation.")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-length", type=int, default=2048)
    parser.add_argument("--epochs", type=float, default=2.0)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    parser.add_argument("--warmup-ratio", type=float, default=0.03)
    parser.add_argument("--weight-decay", type=float, default=0.0)
    parser.add_argument("--logging-steps", type=int, default=10)
    parser.add_argument("--save-steps", type=int, default=100)
    parser.add_argument("--eval-steps", type=int, default=100)
    parser.add_argument("--save-total-limit", type=int, default=2)
    parser.add_argument("--lora-r", type=int, default=16)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)
    parser.add_argument(
        "--target-modules",
        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
        help="Comma-separated LoRA target module names.",
    )
    parser.add_argument("--quantization", choices=["none", "4bit", "8bit"], default="none")
    parser.add_argument("--dtype", choices=["auto", "float32", "float16", "bfloat16"], default="auto")
    parser.add_argument("--gradient-checkpointing", action="store_true")
    parser.add_argument("--trust-remote-code", action="store_true")
    parser.add_argument("--resume-from-checkpoint", default=None)
    args = parser.parse_args()

    LoraConfig, get_peft_model, prepare_model_for_kbit_training, DataCollatorForSeq2Seq, Trainer, TrainingArguments = (
        _import_training_stack()
    )

    records = load_jsonl_records(args.data, max_records=args.max_records)
    train_records, eval_records = split_records(records, args.validation_size, args.seed)

    model, tokenizer, torch = load_model_and_tokenizer(
        args.base_model,
        quantization=args.quantization,
        dtype=args.dtype,
        trust_remote_code=args.trust_remote_code,
        for_training=True,
    )

    model.config.use_cache = False
    if args.quantization != "none":
        model = prepare_model_for_kbit_training(model)

    lora_config = LoraConfig(
        task_type="CAUSAL_LM",
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        target_modules=_csv(args.target_modules),
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    train_dataset = TokenizedCodeDataset(train_records, tokenizer, args.max_length)
    eval_dataset = TokenizedCodeDataset(eval_records, tokenizer, args.max_length) if eval_records else None
    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        model=model,
        label_pad_token_id=-100,
        pad_to_multiple_of=8,
    )

    training_args = _build_training_args(TrainingArguments, args, torch)
    trainer_kwargs = {
        "model": model,
        "args": training_args,
        "train_dataset": train_dataset,
        "eval_dataset": eval_dataset,
        "data_collator": data_collator,
    }

    trainer_signature = inspect.signature(Trainer.__init__).parameters
    if "processing_class" in trainer_signature:
        trainer_kwargs["processing_class"] = tokenizer
    else:
        trainer_kwargs["tokenizer"] = tokenizer

    trainer = Trainer(**trainer_kwargs)
    trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))
    print(f"Saved LoRA adapter to {output_dir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())