Skip to content

feat(diffusion): add SFT loss hub and pre-encoded data manager - #90

Draft
zhihengy wants to merge 2 commits into
mainfrom
feat/sft
Draft

feat(diffusion): add SFT loss hub and pre-encoded data manager#90
zhihengy wants to merge 2 commits into
mainfrom
feat/sft

Conversation

@zhihengy

@zhihengy zhihengy commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

What

Diffusion SFT on (video, prompt) datasets, reusing the PR #63 hook architecture end to end. No sglang engines are involved: --loss-type sft_loss flips on the existing debug_train_only wiring, so placement groups, weight sync, and the driver loop need no changes.

How

  • loss_hub/sft.py: prepare_sft_batch samples grid sigmas (phase-pure per micro-batch for dual-expert Wan2.2 via component_for_timestep), corrupts cached clean latents with x_t = (1-σ)x₀ + σε, and collates cached cond through the family TrainPipelineConfig; sft_loss_formula is velocity MSE against ε - x₀. Both plug into the existing --custom-prepare-train-batch-path / --custom-loss-function-path slots (auto-filled by --loss-type sft_loss).
  • miles/ray/sft_data_manager.py: replaces RolloutManager behind the same driver interface. generate(rollout_id) is stateless — per-epoch seeded shuffle, contiguous slice, TrainDataDPSplitter, and a flow-shifted training sigma grid emitted as scheduler meta — so resume needs no cursor state.
  • scripts/sft_encode_wan.py: one-time Ray encode job (UMT5 + Wan VAE from diffusers) producing one .pt per sample with the pair schema the prepare hook consumes.
  • scripts/run-diffusion-sft-wan22.sh: 4-GPU Wan2.2 dual-expert LoRA SFT recipe.
  • arguments: --sft-data-path; sft_loss validation rejects RL-only combos (--eval-interval, --diffusion-kl-beta, --diffusion-recompute-old-log-prob, --ref-mode, --ema-shadow, --n-samples-per-prompt != 1) and requires --diffusion-flow-shift for the training grid.

The rollout-side data_conversion_hub (expand hook) is intentionally not used: its (samples, rewards) signature is rollout-specific and SFT has neither.

Validation

  • tests/fast/backends/fsdp_utils/test_loss_hub_sft.py (5 tests, pass): corruption/target identity, timestep scaling, dual-expert phase purity, loss values.
  • SftDataManager exercised directly: shift math (σ(0.5) → 0.7507 at shift 3.0), terminal sigma, DP shard shapes, exact one-epoch coverage, different shuffle across epochs.
  • Full parse_args round-trip for the recipe arg set (hooks auto-filled, debug_train_only=True, gbs derivation) plus rejection of each invalid combo.
  • Not yet run on GPU — draft for design review first.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant