feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic - #96
Open
zhihengy wants to merge 3 commits into
Open
feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic#96zhihengy wants to merge 3 commits into
zhihengy wants to merge 3 commits into
Conversation
…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>
zhihengy
added a commit
that referenced
this pull request
Aug 6, 2026
…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
reviewed
Aug 6, 2026
| embeds[:, int(inputs.attention_mask[0].sum()) :] = 0 | ||
|
|
||
| return { | ||
| "latent": latent[0].to(torch.float16).cpu(), |
Collaborator
There was a problem hiding this comment.
The dtypes here need to be double-checked.
Collaborator
Author
There was a problem hiding this comment.
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.
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
New
miles/rollout/encoder_hub/package holding frozen-encoder logic per model family, dispatched by the resolvedargs.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-checkpointHF name/path),encode_sample(encoder, pixels, prompt, generator)(one media/prompt pair → clean latent + cond kwargs), andvalidate_args(args)(family-specific encode constraints — for Wan2.2, the(num_frames-1) % 4 == 0VAE 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 two-PR stack:
main: additive module, no callers, no behavior change;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.Merge this first, then #90.