feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic by zhihengy · Pull Request #96 · radixark/miles_diffusion

zhihengy · 2026-08-06T00:43:24Z

What

New miles/rollout/encoder_hub/ package holding frozen-encoder logic per model family, dispatched by the resolved args.diffusion_model_family (same decentralization mechanism as the TPC registry). Motivated by review on #90: TrainPipelineConfig should only drive the training backend; encoders are an independent, rollout-side concern.

Each family module provides three functions: load_encoder(args, device) (frozen tokenizer/text-encoder/VAE from the explicit --sft-encoder-checkpoint HF name/path), encode_sample(encoder, pixels, prompt, generator) (one media/prompt pair → clean latent + cond kwargs), and validate_args(args) (family-specific encode constraints — for Wan2.2, the (num_frames-1) % 4 == 0 VAE temporal-stride rule, which is an encoder property, not a global SFT rule).

Wan2.2 is the only entry; unsupported families are rejected at dispatch with a clear error.

Stack

This PR is the base of a three-PR stack (bottom → top):

this PR (feat/encoder-hub) → main: additive module, no callers, no behavior change;
feat(diffusion): add SFT loss hub and pre-encoded data manager #90 (feat/sft) → this PR: SFT wires argument validation and its encode actor pool to this hub, defines --sft-encoder-checkpoint, and removes the encoder methods from TrainPipelineConfig;
fix(sft): scope deterministic training randomness #102 (fix/sft-rng-scoping) → feat(diffusion): add SFT loss hub and pre-encoded data manager #90: scopes deterministic training randomness in the SFT loss hub.

Merge order: this PR → #90 → #102.

…training-only Review feedback (PR #90): TrainPipelineConfig drives the training backend only. Encoder loading/encoding now lives in miles/rollout/encoder_hub (stacked base PR #96), dispatched by args.diffusion_model_family; the Wan-specific 4k+1 frame constraint validates there too. Encoders load from the explicit --sft-encoder-checkpoint, which also replaces hf_checkpoint in the per-sample cache key since it is what determines cache content. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Rockdu · 2026-08-06T04:48:39Z

+    embeds[:, int(inputs.attention_mask[0].sum()) :] = 0
+
+    return {
+        "latent": latent[0].to(torch.float16).cpu(),


The dtypes here need to be double-checked.

done (173abc6) — aligned with sglang-d's Wan pipeline config (vae_precision fp32, text_encoder_precisions fp32, DiT bf16): UMT5 now computes in fp32 like the rollout engine, VAE stays fp32, and both cached tensors store bf16 — the precision the DiT boundary sees on the rollout path. Note for anyone with an existing .sft_cache: the cache key does not encode dtypes, so clear it manually to re-encode.

…ogic Review feedback on PR #90: TrainPipelineConfig should only drive the training backend, so frozen-encoder loading/encoding lives in its own rollout-side hub, dispatched by the resolved diffusion model family. Wan2.2 provides UMT5+VAE loading (from the explicit --sft-encoder-checkpoint introduced by the consumer PR), sample encoding, and its 4k+1 frame-count constraint. No callers in this PR; the SFT PR (#90) is stacked on top and wires argument validation and the encode actor pool to this hub. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Match the rollout engine's Wan pipeline config (vae_precision fp32, text_encoder_precisions fp32, DiT bf16): UMT5 now computes in fp32, and both cached tensors store bf16 — the precision the DiT boundary sees on the rollout path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…training-only Review feedback (PR #90): TrainPipelineConfig drives the training backend only. Encoder loading/encoding now lives in miles/rollout/encoder_hub (stacked base PR #96), dispatched by args.diffusion_model_family; the Wan-specific 4k+1 frame constraint validates there too. Encoders load from the explicit --sft-encoder-checkpoint, which also replaces hf_checkpoint in the per-sample cache key since it is what determines cache content. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* feat(diffusion): add SFT loss hub and pre-encoded data manager * feat(diffusion): add frame stride to SFT encode script * feat(diffusion): auto-build SFT cache via family encode hooks * fix(sft): per-sample content-addressed cache, atomic writes, seeded VAE sampling * refactor(sft): plug into RolloutManager via rollout-function/convert/log hooks with lazy encoder pool * feat(sft): image datasets for Wan (single-frame media branch) * refactor(sft): pick dual-expert via uniform anchor index * refactor(sft): seat encoder pool via the rollout placement view * chore(sft): encoder pool inherits the engine's 0.3 GPU share * refactor(sft): reuse --diffusion-height/width/output-num-frames; add --fsdp-flow-shift Review feedback (PR #90): the media geometry args describe the generated media regardless of whether it comes from preprocessing or the rollout engine, so SFT reuses --diffusion-height/--diffusion-width/ --diffusion-output-num-frames instead of its own --sft-* trio. The SFT training sigma grid is regenerated on the training side, so its shift is now --fsdp-flow-shift (fsdp_* namespace), leaving --diffusion-flow-shift as the rollout-engine launch parameter. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(sft): validate explicit plugin paths instead of auto-wiring defaults Review feedback (PR #90): follow Miles LLM convention — the sample script passes the five SFT plugin paths explicitly and the framework validates the combination instead of silently rewriting args. Misconfigured runs now fail at argument validation with the exact flag to set, rather than deep in the RL data path. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor: make --train-only an explicit flag; --debug-train-only becomes its alias Review feedback (PR #90): no silent inference of train_only from loss_type. --train-only is now a real user-facing flag (argparse alias keeps --debug-train-only working) and sft_loss validates it is set instead of setting it. The rollout placement view keeps its seats unconditionally: engine startup is gated by args.train_only, and rollout-side actor pools (the SFT encoder pool) seat there. The debug_rollout_only/train_only exclusion assert now runs before the debug_rollout_only reconfiguration so the combo fails with the intended message. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * refactor(sft): move frozen-encoder logic to encoder_hub; TPC back to training-only Review feedback (PR #90): TrainPipelineConfig drives the training backend only. Encoder loading/encoding now lives in miles/rollout/encoder_hub (stacked base PR #96), dispatched by args.diffusion_model_family; the Wan-specific 4k+1 frame constraint validates there too. Encoders load from the explicit --sft-encoder-checkpoint, which also replaces hf_checkpoint in the per-sample cache key since it is what determines cache content. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * style: black-format the train_only exclusion assert Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * feat(sft): support disaggregated encoder GPUs Treat explicit --rollout-num-gpus in train-only jobs as a dedicated rollout-side producer pool while preserving the colocated default. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(ci): lazy-load rollout manager from placement helper Keep CPU placement tests from importing optional SGLang runtime modules during collection. Co-authored-by: Cursor <cursoragent@cursor.com> * Revert "fix(ci): lazy-load rollout manager from placement helper" This reverts commit 4d2450c. * fix(ci): expose SGLang source package to CPU tests Put the checked-out Python package ahead of the workspace root so multimodal_gen imports resolve under SGLang's current layout. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(sglang): defer diffusion-only imports until engine startup Use SGLang's common process helper at module load so CPU-side orchestration can import without GPU-only diffusion dependencies. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Cursor <cursoragent@cursor.com>

zhihengy requested review from Rockdu and guapisolo as code owners August 6, 2026 00:43

zhihengy mentioned this pull request Aug 6, 2026

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

Merged

Rockdu reviewed Aug 6, 2026

View reviewed changes

zhihengy and others added 3 commits August 7, 2026 11:49

docs(encoder_hub): record measured rationale for cache storage dtypes

2a3c9ce

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

zhihengy force-pushed the feat/encoder-hub branch from 173abc6 to 28a93be Compare August 7, 2026 11:50

zhihengy mentioned this pull request Aug 8, 2026

fix(sft): scope deterministic training randomness #102

Merged

zhihengy merged commit adc285e into main Aug 8, 2026
15 checks passed

Navigation Menu

Uh oh!

feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic - #96

feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic#96
zhihengy merged 3 commits into
mainfrom
feat/encoder-hub

zhihengy commented Aug 6, 2026 •

edited

Loading

Uh oh!

Rockdu Aug 6, 2026

Uh oh!

zhihengy Aug 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

zhihengy commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Stack

Uh oh!

Rockdu Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

zhihengy Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zhihengy commented Aug 6, 2026 •

edited

Loading