Add HF-FSDP VLM training backend + SGLang rollout bridge

TL;DR

Add a new HF-FSDP training backend in slime for VLM SFT / pre-alignment, while reusing SGLang for VLM inference in rollout. Keep RL wiring compatible with veRL-style actor/learner, but implemented natively in slime’s DataBuffer → Rollout → Training pipeline.

Motivation

• Reuse the mature HuggingFace ecosystem (models, processors, datasets) for vision-language training.

• Use PyTorch FSDP (v2 / fully_shard) for simple, fast, and stable multi-GPU scaling.

• Leverage SGLang for high-throughput VLM inference during rollout and evaluation.

• Keep slime’s current decoupled architecture; no Megatron/DeepSpeed dependency for this path.

Scope (MVP)

• New trainer backend: --trainer fsdp_vlm

• Target models (HF native): Qwen2-VL / Qwen2.5-VL, MiniCPM-V-4_5 / 4.0, Llama-3.2-Vision

• Training stage: SFT / pre-alignment only (RL later)

• Data: LLAVA/Qwen-style multimodal chat JSON or HF datasets with images column

• Rollout: call SGLang VLM endpoint with image(s) + prompt

• Checkpoints: torch.distributed.checkpoint (sharded), easy to merge for release

Proposed Directory Layout

slime_plugins/
  trainers/
    fsdp_vlm/
      __init__.py          # entry & registry
      builder.py           # HF AutoModel/AutoProcessor, freezing knobs
      dataset.py           # VLM dataset + collator (images + text)
      engine.py            # FSDP training loop, grad accum, ckpt
      checkpoint.py        # sharded save/load wrappers
      eval.py              # quick eval & sampling hooks (HF generate or SGLang)

Key Design Decisions

• FSDP settings (PyTorch 2.x): ShardingStrategy.FULL_SHARD, bf16 mixed precision, transformer_auto_wrap_policy, activation checkpointing, backward_prefetch=PRE, limit_all_gathers=True, use_orig_params=True, sync_module_states=True.

• Freezing: flags --freeze-vision, --freeze-mm-proj (default freeze vision tower for stable SFT).

• LoRA (optional): prefer FSDP + LoRA (non-quantized) via PEFT; avoid FSDP + QLoRA for MVP.

• Label masking: image placeholder tokens get label -100.

• Data IO: lazy image decode (PIL/accimage), worker prefetch; later upgrade to WebDataset/DALI if needed.

• Rollout: add VLMRolloutEngine that forwards images (URLs/base64) + prompts to SGLang; store responses in DataBuffer.

New CLI Flags (initial set)

• VLM:

--vlm_model_name_or_path

--processor_name_or_path

--freeze_vision, --freeze_mm_proj

--image_size, --image_strategy {longest,pad_square}

--pack_multi_image

--chat_template {qwen,llava,minicpm,llama32v}

• FSDP / Dist:

--fsdp_enable

--fsdp_sharding {full,hybrid} (default: full)

--fsdp_mp_dtype {bf16}

--fsdp_auto_wrap_cls "Qwen2DecoderLayer,LlamaDecoderLayer,MiniCPMDecoderLayer"

--grad_ckpt, --grad_accum_steps, --clip_grad_norm

--fsdp_checkpoint {sharded,full} (train: sharded; publish: full)

• LoRA (optional):

--lora_r, --lora_alpha, --lora_dropout

--lora_target_modules "q_proj,k_proj,v_proj,o_proj,..."

--lora_ignore_modules "vision_tower,mm_proj"

Acceptance Criteria

• torchrun … --trainer fsdp_vlm trains Qwen2-VL-7B-Instruct or MiniCPM-V-4_5 on a small JSON dataset end-to-end (single node, ≥4 GPUs).

• FSDP fully-sharded training runs with bf16, activation checkpointing, auto-wrap, grad accumulation.

• Sharded checkpoints saved/restored via torch.distributed.checkpoint; full consolidated weights can be exported for inference.

• Collator correctly handles multiple images, applies image token label masking, and uses the model’s Processor.

• Minimal rollout path exists: VLMRolloutEngine can send image(s)+prompt to SGLang and buffer the response.

• A tiny eval command runs HF generate or calls SGLang to sample outputs.

Non-Goals (for this issue)

• Full RL (GRPO/RLAIF-V) pipeline

• Quantization/QLoRA

• Video training (multi-frame) and DALI/WebDataset optimizations

Starter Command (PoC)

CUDA_VISIBLE_DEVICES=0,1,2,3 \
torchrun --nnodes=1 --nproc_per_node=4 --rdzv_backend=c10d --rdzv_endpoint=localhost:29500 \
  -m slime.train \
    --trainer fsdp_vlm \
    --vlm_model_name_or_path Qwen/Qwen2-VL-7B-Instruct \
    --processor_name_or_path Qwen/Qwen2-VL-7B-Instruct \
    --train_file /path/to/train.json \
    --eval_file  /path/to/val.json \
    --image_strategy longest \
    --freeze_vision true \
    --fsdp_enable true \
    --fsdp_sharding full \
    --fsdp_mp_dtype bf16 \
    --fsdp_auto_wrap_cls "Qwen2DecoderLayer" \
    --grad_ckpt true \
    --per_device_train_batch_size 1 \
    --grad_accum_steps 16 \
    --lr 2e-5 \
    --save_dir /ckpts/slime_vlm_fsdp

Task Checklist

• Registry: add --trainer fsdp_vlm & dynamic import in train.py

• builder.py: HF AutoModel/Processor loading; freezing knobs; chat template handling

• dataset.py: JSON/HF-dataset loader + collator (multi-image, label masking)

• engine.py: FSDP init (policy, mp, ckpting), train loop (autocast bf16, grad accum, clip)

• checkpoint.py: sharded save/load helpers; optional full merge util

• eval.py: small eval hooks (HF generate + optional SGLang call)

• rollout engine: minimal VLMRolloutEngine using SGLang for inference

• Docs: README snippet with PoC command & data format example

• Smoke tests: single-node 4xGPU run on tiny dataset; verify loss decreases & ckpt reload works

Pitfalls to Watch

• Avoid auto-wrapping vision tower; only wrap decoder layers.

• Ensure image placeholders are masked with -100 in labels.

• FSDP + LoRA: register LoRA via PEFT before wrapping; exclude from unnecessary flatten; save adapters separately.

• Memory: tune grad-ckpt + accum to fit per-GPU budget.

References (high-level)

• PyTorch FSDP2 best practices (auto-wrap, bf16, sharded checkpoints)

• Qwen-VL / MiniCPM-V / Llama-3.2-Vision HF training examples

• SGLang VLM inference interface

• veRL actor/learner orchestration concepts

Start by reading the existing slime.train entry point and comparing it with the proposed slime_plugins/trainers/fsdp_vlm files. Use the provided four-GPU starter command and tiny dataset as the initial validation path. Done means the listed backend, sharded checkpointing, multimodal collator, minimal SGLang rollout, evaluation hook, documentation, and smoke tests all meet the acceptance criteria.