startrail-org--pixelrag/train/train.py

#!/usr/bin/env python3
"""InfoNCE training loop with LoRA fine-tuning, checkpointing, and wandb logging.

Usage:
    uv run python -m training.train --gpu-id 2

Resume from checkpoint:
    uv run python -m training.train --gpu-id 2 --resume training/checkpoints/run_001/step_200.pt
"""

import argparse
import logging
import signal
import threading
import time
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.nn.utils import clip_grad_norm_
from torch.utils.data import DataLoader
from transformers import (
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
)

from training.dataset import QueryChunkDataset, make_collate_fn
from training.evaluate import run_eval
from training.model import load_model_for_training, load_processor, pool_and_normalize

import sys

LOG_PATH = Path("training/train.log")

# Direct file writer that bypasses wandb's stdout/stderr capture
_log_fd = None


def _log(msg: str):
    """Write log line directly to file descriptor (wandb-proof)."""
    global _log_fd
    if _log_fd is None:
        _log_fd = open(LOG_PATH, "w", buffering=1)  # line-buffered
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    _log_fd.write(f"{ts} {msg}\n")
    _log_fd.flush()


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    stream=sys.stdout,
)
logger = logging.getLogger(__name__)


def info_nce_loss(
    q_emb: torch.Tensor, i_emb: torch.Tensor, temperature: float = 0.07
) -> torch.Tensor:
    """In-batch negatives InfoNCE loss.

    q_emb: (B, D) L2-normalized query embeddings
    i_emb: (B, D) L2-normalized image embeddings
    Diagonal entries are positive pairs.
    """
    logits = q_emb @ i_emb.T / temperature  # (B, B)
    labels = torch.arange(logits.size(0), device=logits.device)
    return F.cross_entropy(logits, labels)


def save_checkpoint(
    model, optimizer, scheduler, step, config, best_recall_10=0.0, loss_history=None
):
    """Save LoRA weights + optimizer/scheduler state."""
    ckpt_dir = Path(config.checkpoint_dir) / config.run_name
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    path = ckpt_dir / f"step_{step}.pt"
    torch.save(
        {
            "step": step,
            "model_state_dict": {k: v.cpu() for k, v in model.state_dict().items()},
            "optimizer_state_dict": optimizer.state_dict(),
            "scheduler_state_dict": scheduler.state_dict(),
            "best_recall_10": best_recall_10,
            "loss_history": loss_history or [],
            "config": vars(config),
        },
        path,
    )
    _log(f"Checkpoint saved: {path}")


def train(config):
    import wandb

    device = f"cuda:{config.gpu_id}"

    # wandb — disable console capture so logger output reaches stdout
    wandb.init(
        project="wiki-embedding",
        name=config.run_name,
        config=vars(config),
        settings=wandb.Settings(console="off"),
    )

    # Model
    _log(f"Loading model {config.model} on GPU {config.gpu_id}...")
    model = load_model_for_training(
        config.model,
        config.gpu_id,
        lora_r=config.lora_r,
        lora_alpha=config.lora_alpha,
    )
    processor = load_processor(
        config.model, min_pixels=config.min_pixels, max_pixels=config.max_pixels
    )

    # Data
    dataset = QueryChunkDataset(config.train_jsonl)
    collate_fn = make_collate_fn(processor, device=device)
    loader = DataLoader(
        dataset,
        batch_size=config.batch_size,
        shuffle=True,
        num_workers=0,
        collate_fn=collate_fn,
    )

    # Optimizer + scheduler
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr, weight_decay=0.01)
    if config.scheduler == "cosine":
        scheduler = get_cosine_schedule_with_warmup(
            optimizer,
            num_warmup_steps=config.warmup_steps,
            num_training_steps=config.max_steps,
        )
    else:
        scheduler = get_constant_schedule_with_warmup(
            optimizer,
            num_warmup_steps=config.warmup_steps,
        )

    # Resume
    start_step = 0
    best_recall_10 = 0.0
    loss_history = []
    if config.resume:
        _log(f"Resuming from {config.resume}")
        ckpt = torch.load(config.resume, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model_state_dict"], strict=False)
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        scheduler.load_state_dict(ckpt["scheduler_state_dict"])
        start_step = ckpt["step"]
        best_recall_10 = ckpt.get("best_recall_10", 0.0)
        loss_history = ckpt.get("loss_history", [])
        _log(f"Resumed from step {start_step}")

    # Graceful shutdown
    shutdown = threading.Event()
    original_sigterm = signal.getsignal(signal.SIGTERM)
    original_sigint = signal.getsignal(signal.SIGINT)

    def _shutdown_handler(signum, frame):
        _log(f"Received signal {signum}, shutting down gracefully...")
        shutdown.set()

    signal.signal(signal.SIGTERM, _shutdown_handler)
    signal.signal(signal.SIGINT, _shutdown_handler)

    # Train loop
    step = start_step
    model.train()
    skipped_batches = 0

    _log(f"Starting training from step {step}, max_steps={config.max_steps}")

    try:
        while step < config.max_steps:
            for batch in loader:
                if shutdown.is_set():
                    _log("Shutdown requested, saving checkpoint...")
                    save_checkpoint(
                        model,
                        optimizer,
                        scheduler,
                        step,
                        config,
                        best_recall_10,
                        loss_history,
                    )
                    wandb.finish()
                    return

                if batch is None:
                    skipped_batches += 1
                    wandb.log({"skipped_batches": skipped_batches}, step=step)
                    continue

                t0 = time.time()
                try:
                    q_inputs, i_inputs = batch

                    # Forward: query embeddings
                    q_out = model(**q_inputs, output_hidden_states=True)
                    q_emb = pool_and_normalize(
                        q_out.hidden_states[-1], q_inputs["attention_mask"]
                    )

                    # Forward: image embeddings
                    i_out = model(**i_inputs, output_hidden_states=True)
                    i_emb = pool_and_normalize(
                        i_out.hidden_states[-1], i_inputs["attention_mask"]
                    )

                    loss = info_nce_loss(q_emb, i_emb, temperature=config.temperature)
                    loss.backward()
                    clip_grad_norm_(model.parameters(), 1.0)
                    optimizer.step()
                    scheduler.step()
                    optimizer.zero_grad()

                except torch.cuda.OutOfMemoryError:
                    _log(f"WARNING OOM at step {step}, skipping batch")
                    torch.cuda.empty_cache()
                    optimizer.zero_grad()
                    skipped_batches += 1
                    wandb.log({"skipped_batches": skipped_batches}, step=step)
                    continue

                dt = time.time() - t0
                step += 1
                loss_val = loss.item()
                loss_history.append(loss_val)
                lr = scheduler.get_last_lr()[0]
                gpu_mem = torch.cuda.memory_allocated(config.gpu_id) / 1e9

                wandb.log(
                    {
                        "loss": loss_val,
                        "lr": lr,
                        "batch_time": dt,
                        "pairs_per_sec": config.batch_size / dt,
                        "gpu_mem_gb": gpu_mem,
                    },
                    step=step,
                )
                _log(
                    f"step={step} loss={loss_val:.4f} lr={lr:.2e} "
                    f"batch_time={dt:.1f}s pairs/s={config.batch_size / dt:.1f} "
                    f"gpu_mem={gpu_mem:.1f}GB"
                )

                # Eval
                if step % config.eval_every == 0:
                    r1, r10, mrr = run_eval(
                        model,
                        processor,
                        config.eval_jsonl,
                        device,
                        batch_size=config.batch_size,
                    )
                    wandb.log(
                        {
                            "recall@1": r1,
                            "recall@10": r10,
                            "mrr": mrr,
                        },
                        step=step,
                    )
                    _log(
                        f"eval step={step} recall@1={r1:.3f} recall@10={r10:.3f} mrr={mrr:.3f}"
                    )
                    if r10 > best_recall_10:
                        best_recall_10 = r10

                # Checkpoint
                if step % config.save_every == 0:
                    save_checkpoint(
                        model,
                        optimizer,
                        scheduler,
                        step,
                        config,
                        best_recall_10,
                        loss_history,
                    )

                if step >= config.max_steps:
                    break

    finally:
        # Restore signal handlers
        signal.signal(signal.SIGTERM, original_sigterm)
        signal.signal(signal.SIGINT, original_sigint)

    # Final checkpoint
    save_checkpoint(
        model, optimizer, scheduler, step, config, best_recall_10, loss_history
    )
    wandb.finish()
    _log("Training complete!")


def main():
    parser = argparse.ArgumentParser(
        description="LoRA fine-tuning for Qwen3-VL embeddings"
    )
    parser.add_argument("--model", default="Qwen/Qwen3-VL-Embedding-2B")
    parser.add_argument("--gpu-id", type=int, default=2)
    parser.add_argument("--train-jsonl", default="training/data/train.jsonl")
    parser.add_argument("--eval-jsonl", default="training/data/eval.jsonl")
    parser.add_argument("--checkpoint-dir", default="training/checkpoints")
    parser.add_argument("--run-name", default="run_001")
    parser.add_argument("--resume", type=str, default=None)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument(
        "--warmup-steps",
        type=int,
        default=None,
        help="Warmup steps (default: 5% of max_steps)",
    )
    parser.add_argument("--max-steps", type=int, default=500)
    parser.add_argument(
        "--scheduler", choices=["cosine", "constant"], default="constant"
    )
    parser.add_argument("--temperature", type=float, default=0.07)
    parser.add_argument("--eval-every", type=int, default=100)
    parser.add_argument("--save-every", type=int, default=100)
    parser.add_argument("--lora-r", type=int, default=32)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument(
        "--min-pixels",
        type=int,
        default=128 * 28 * 28,
        help="Min image pixels for processor (default: 128*28*28)",
    )
    parser.add_argument(
        "--max-pixels",
        type=int,
        default=256 * 28 * 28,
        help="Max image pixels for processor (default: 256*28*28)",
    )
    config = parser.parse_args()
    if config.warmup_steps is None:
        config.warmup_steps = max(1, (config.max_steps + 19) // 20)
    train(config)


if __name__ == "__main__":
    main()