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

#!/usr/bin/env python3
"""Fine-tune Qwen3-VL-Embedding with ms-swift — equivalent to train_contrastors.py.

Uses ms-swift's embedding training pipeline (InfoNCE loss, cross-GPU negative sharing)
instead of our custom GradCache training loop.

Equivalence notes vs train_contrastors.py:
  - Same model (Qwen3-VL-Embedding-2B), same LoRA targets (q/k/v/o_proj)
  - Same loss (InfoNCE with in-batch + hard negatives)
  - Same instructions (QUERY_INSTRUCTION / DOC_INSTRUCTION in data JSONL)
  - Temperature: FIXED at 0.07 (swift has no learnable LogitScale)
  - No GradCache: memory bounded by batch size (use DeepSpeed ZeRO-2 to compensate)
  - No custom retrieval eval (R@1/5/10): swift only does loss-based eval
    → Run retrieval eval separately after training

Data format: use convert_data_for_swift.py to convert from contrastors format.

Single GPU:
    CUDA_VISIBLE_DEVICES=3 uv run python train_swift.py

Multi-GPU:
    CUDA_VISIBLE_DEVICES=1,2 uv run python train_swift.py --nproc-per-node 2

Resume:
    uv run python train_swift.py --resume training/output_swift/vX-XXX/checkpoint-50

Best config (matching train_contrastors.py defaults):
    CUDA_VISIBLE_DEVICES=1,2 uv run python train_swift.py \\
        --train-jsonl data/train_hn_swift.jsonl \\
        --eval-jsonl data/eval_swift.jsonl \\
        --num-hard-negatives 5 \\
        --batch-size 16 \\
        --lr 1e-5 \\
        --max-steps 50 \\
        --warmup-steps 20 \\
        --eval-steps 25 \\
        --save-steps 50 \\
        --nproc-per-node 2
"""

import argparse
import os


def main():
    parser = argparse.ArgumentParser(
        description="Fine-tune Qwen3-VL-Embedding with ms-swift (InfoNCE)"
    )

    # Model
    parser.add_argument("--model", default="Qwen/Qwen3-VL-Embedding-2B")

    # Data (swift format — use convert_data_for_swift.py first)
    parser.add_argument("--train-jsonl", default="data/train_hn_swift.jsonl")
    parser.add_argument("--eval-jsonl", default="data/eval_swift.jsonl")

    # Training
    parser.add_argument("--batch-size", type=int, default=4, help="Per-GPU batch size")
    parser.add_argument("--lr", type=float, default=2e-5)
    parser.add_argument("--warmup-steps", type=int, default=50)
    parser.add_argument("--max-steps", type=int, default=500)
    parser.add_argument("--scheduler", choices=["cosine", "constant"], default="cosine")
    parser.add_argument(
        "--temperature",
        type=float,
        default=0.07,
        help="Fixed InfoNCE temperature (not learnable in swift)",
    )
    parser.add_argument("--max-grad-norm", type=float, default=1.0)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument(
        "--num-hard-negatives",
        type=int,
        default=0,
        help="Hard negatives per query (requires swift-format data with negative_messages)",
    )

    # LoRA (match train_contrastors.py defaults)
    parser.add_argument("--lora-r", type=int, default=32)
    parser.add_argument("--lora-alpha", type=int, default=32)
    parser.add_argument("--lora-dropout", type=float, default=0.05)

    # Resolution
    parser.add_argument(
        "--max-num-visual-tokens",
        type=int,
        default=4096,
        help="Max visual tokens → converted to max_pixels for processor",
    )

    # Eval / Save
    parser.add_argument("--eval-steps", type=int, default=100)
    parser.add_argument("--save-steps", type=int, default=100)
    parser.add_argument("--save-total-limit", type=int, default=5)
    parser.add_argument("--logging-steps", type=int, default=5)

    # Output
    parser.add_argument("--output-dir", default="training/output_swift")

    # Resume
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Path to checkpoint directory to resume from",
    )

    # Distributed
    parser.add_argument(
        "--nproc-per-node",
        type=int,
        default=1,
        help="Number of GPUs (sets NPROC_PER_NODE for swift)",
    )
    parser.add_argument(
        "--deepspeed",
        default=None,
        help="DeepSpeed config: 'zero2', 'zero3', or path to JSON",
    )

    # Wandb
    parser.add_argument("--wandb-project", default="wiki-screenshot-training")
    parser.add_argument("--no-wandb", action="store_true")

    # Freeze
    parser.add_argument(
        "--freeze-vit",
        action="store_true",
        default=True,
        help="Freeze vision encoder (default: True)",
    )
    parser.add_argument("--no-freeze-vit", dest="freeze_vit", action="store_false")

    args = parser.parse_args()

    # --- Environment variables for swift InfoNCE ---
    os.environ["INFONCE_TEMPERATURE"] = str(args.temperature)
    os.environ["INFONCE_USE_BATCH"] = (
        "True"  # in-batch negatives (like train_contrastors.py)
    )
    if args.num_hard_negatives > 0:
        os.environ["INFONCE_HARD_NEGATIVES"] = str(args.num_hard_negatives)

    if args.nproc_per_node > 1:
        os.environ["NPROC_PER_NODE"] = str(args.nproc_per_node)

    if args.no_wandb:
        os.environ["WANDB_DISABLED"] = "true"
    else:
        os.environ["WANDB_PROJECT"] = args.wandb_project

    # Convert max_num_visual_tokens → max_pixels
    # Qwen3-VL: each visual token covers (patch_size * spatial_merge_size)^2 = 28^2 = 784 pixels
    pixels_per_token = 28 * 28  # patch_size=14, spatial_merge_size=2
    max_pixels = args.max_num_visual_tokens * pixels_per_token

    # --- Build SftArguments ---
    # Import after env vars are set so swift picks them up
    from swift import SftArguments, sft_main

    # Map scheduler name
    lr_scheduler_type = (
        "cosine" if args.scheduler == "cosine" else "constant_with_warmup"
    )

    sft_args = SftArguments(
        # Model
        model=args.model,
        task_type="embedding",
        # LoRA — match train_contrastors.py: q_proj, k_proj, v_proj, o_proj
        tuner_type="lora",
        lora_rank=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=args.lora_dropout,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        freeze_vit=args.freeze_vit,
        # Loss
        loss_type="infonce",
        # Data
        dataset=[args.train_jsonl],
        val_dataset=[args.eval_jsonl],
        split_dataset_ratio=0.0,  # We provide val_dataset explicitly
        # Training hyperparams
        per_device_train_batch_size=args.batch_size,
        per_device_eval_batch_size=args.batch_size,
        learning_rate=args.lr,
        lr_scheduler_type=lr_scheduler_type,
        warmup_steps=args.warmup_steps,
        max_steps=args.max_steps,
        max_grad_norm=args.max_grad_norm,
        weight_decay=args.weight_decay,
        # Precision
        torch_dtype="bfloat16",
        # Resolution
        max_pixels=max_pixels,
        # Eval / Save / Log
        eval_strategy="steps",
        eval_steps=args.eval_steps,
        save_steps=args.save_steps,
        save_total_limit=args.save_total_limit,
        logging_steps=args.logging_steps,
        dataloader_drop_last=True,
        # Output
        output_dir=args.output_dir,
        # Resume
        resume_from_checkpoint=args.resume,
        # DeepSpeed
        deepspeed=args.deepspeed,
        # Misc
        dataloader_num_workers=4,
    )

    # --- Run training ---
    result = sft_main(sft_args)

    # Print results
    if result:
        print(f"\n{'=' * 60}")
        print("Training complete!")
        if hasattr(result, "last_model_checkpoint") and result.last_model_checkpoint:
            print(f"Last checkpoint: {result.last_model_checkpoint}")
        if hasattr(result, "best_model_checkpoint") and result.best_model_checkpoint:
            print(f"Best checkpoint: {result.best_model_checkpoint}")
        print(f"{'=' * 60}")


if __name__ == "__main__":
    main()