feat(diffusion): add SFT loss hub and pre-encoded data manager by zhihengy · Pull Request #90 · radixark/miles_diffusion

zhihengy · 2026-08-03T06:37:40Z

Stack: this PR is based on #96 (feat/encoder-hub), which introduces the encoder_hub module this PR wires into. Merge #96 into main first, then this PR.

What

Diffusion SFT on (video, prompt) datasets, delivered entirely through existing plugin seams: a rollout function + convert/log hooks on the RolloutManager side, and the PR #63 prepare/loss hooks on the trainer side. No sglang engines run: --loss-type sft_loss implies a formalized train_only mode. RolloutManager, the driver loop, and placement orchestration are structurally untouched, so an upstream merge sees additive plugin files plus a few flag renames.

How

Data plane — miles/rollout/sft_rollout.py, all plugged via existing --*-path slots:

generate_rollout (via --rollout-function-path): pulls (prompt, metadata.video or metadata.image) rows from the standard RolloutDataSource (--prompt-data/--input-key; epoch iteration, per-epoch shuffle, and checkpoint resume come from the framework), checks a per-sample content-addressed cache (filename = hash(checkpoint, encode geometry, video path+size+mtime, prompt)), and lazily encodes misses through a persistent encoder actor pool. First epoch pays the encode once; epoch 2+ is all hits; a fully-warm run never loads the encoders at all.
The pool occupies the architectural slot sglang engines hold in RL — literally: it seats itself via the rollout placement view (same bundle-index contract engines use), at 0.3 GPU/bundle beside the 0.7 training actors. Under SFT the rollout view keeps its seats (only debug_train_only empties them), and the pool is topology-agnostic: a future disaggregated encoder pool only needs the topology gate relaxed. Writes are temp-then-rename (a crash never leaves a loadable-looking entry) and VAE posterior sampling is seeded from the cache key, so rebuilds are bit-deterministic.
convert_samples_to_train_data (via the full-override --custom-convert-samples-to-train-data-path): emits train pairs + the flow-shifted sigma grid; the reward/advantage code path never executes.
log_rollout_data (via --custom-rollout-log-function-path): logs rollout/sft_cache_miss, rollout/sft_encode_seconds, rollout/sft_epoch on the rollout/step axis and short-circuits the RL reward stats.

Image datasets: Wan trains images as single-frame videos. Rows with metadata.image (or any image-extension path) read via PIL as [C,1,H,W] and require --sft-num-frames 1; the VAE, cache addressing, and loss path are shape-agnostic. Validated e2e (16-image dataset, 100 steps, 23s encode).

Trainer side — loss_hub/sft.py (unchanged across data-plane iterations):

prepare_sft_batch samples grid sigmas phase-pure per micro-batch (dual-expert Wan2.2 via component_for_timestep, mass-weighted so the marginal stays uniform), corrupts cached latents with x_t = (1-σ)x₀ + σε, and collates cond through the family TrainPipelineConfig; sft_loss_formula is velocity MSE vs ε - x₀.

Family seam: miles/rollout/encoder_hub (base PR #96) — frozen-encoder loading/encoding per model family, dispatched by args.diffusion_model_family, loading from the explicit --sft-encoder-checkpoint. TrainPipelineConfig stays training-backend-only; families without an encoder_hub entry are rejected at argument validation.

Small upstream-friendly cleanups:

--train-only formalizes the "no rollout engines" mode as an explicit flag (--debug-train-only kept as an argparse alias): it gates engine startup, router launch (also fixing the stray router the debug mode used to spawn), weight sync, and eval; sft_loss validates it is set rather than inferring it.
rm_hub: set/get_reward_placement_group renamed to set/get_manager_placement_group (old names kept as aliases) — it publishes the manager's placement for colocated actor pools, and the encoder pool is now its second consumer.

Validation

CPU: 9 tests — loss-hub math (corruption identity, timestep scaling, grid-sampling uniformity, dual-expert phase purity, loss values) + per-sample cache-key invalidation axes (geometry/checkpoint/prompt/video size/mtime, per-sample independence).
Deterministic re-encode verified bit-identical on GPU (same seed → same latent and cond; different seed differs).
E2E on 16-video X-Fun demo (Wan2.2-TI2V-5B LoRA, 4x H200, cold cache): round 0 lazily creates the pool and encodes 16 samples in 57s, rounds 1-99 are all cache hits (~3s/round), 400 steps complete, final train loss 0.1998, clean exit, no temp-file leftovers.
Full parse_args round-trip: the sample script passes the five plugin paths explicitly and sft_loss validates them (missing path / RL-default rollout function rejected with the exact flag to set); invalid combos (eval, KL, recompute, ref-mode, EMA, n_samples>1) rejected at validation.

GPU validation: head-to-head vs VideoX-Fun + before/after

Trained Wan2.2-TI2V-5B LoRA on PAI/X-Fun-Videos-Demo (16 open-source videos, 480x832@24fps, detailed captions), settings matched to the scripts/wan2.2/README_TRAIN_LORA.md quick-start of VideoX-Fun as the reference implementation. Both runs in wandb project SFT: miles-d vs videox-fun.

setting	value (both frameworks)
model	Wan2.2-TI2V-5B (single DiT; diffusers layout for miles, original for VideoX-Fun)
data	16 videos, 480x480 center-crop, 77 frames @ stride 2
sigma sampling	uniform over the shift-5.0 grid
loss	velocity MSE vs `eps - x0`
LoRA	r=64, alpha=32 on to_q/to_k/to_v + ffn (no out-proj), kaiming init
optimizer	AdamW lr 1e-4 constant, betas (0.9, 0.999), wd 3e-2, eps 1e-10, grad clip 0.05
batch	global 4 (4 GPUs x 1), 100 epochs = 400 steps, seed 42, bf16 forward / fp32 master

Caveat: 16 videos x 100 epochs is the best case for a warm cache (first build took ~60s here and is excluded); the wall-clock gap does not generalize to huge-dataset few-epoch runs, and the before/after results demonstrate fit to the training set, not held-out quality (held-out fixed-(t, eps) eval is planned as the next step).

Smoothed train loss: miles 0.290 -> 0.232, VideoX-Fun 0.267 -> 0.251 (VideoX keeps a hardcoded 10% text dropout and per-epoch temporal jitter, which raises its floor; miles trains fixed offline clips). Wall clock 7 min vs 24 min — the offline-encode payoff (no per-step UMT5 forward).

Fixed-(t, eps) loss on the training set (same frozen noise + timesteps for both models — isolates model improvement from sampling variance):

sigma	base	after SFT	delta
0.3	0.3277	0.3152	-3.8%
0.5	0.2190	0.2050	-6.4%
0.7	0.1820	0.1624	-10.8%
0.9	0.2266	0.1802	-20.5%

Before/after generations (training prompts, same seed 42, 480x480x33f, 40 UniPC steps, cfg 5.0; top = base, bottom = after 400-step LoRA SFT — outputs shift toward the training clips' composition and lighting):

Side-by-side animations: p0 | p3 (base left, SFT right; assets on branch sft-pr90-assets).

…AE sampling

…log hooks with lazy encoder pool

Rockdu · 2026-08-05T20:53:52Z

+            # lives in each row's metadata dict under "video" or "image" (images train as single
+            # frames and require --sft-num-frames 1). Encoded pairs are cached next to the jsonl
+            # under .sft_cache/, one content-addressed file per sample.
+            parser.add_argument("--sft-height", type=int, default=None, help="SFT encode height (center crop)")


May reuse --diffusion-height and --diffusion-width here

Same with num-frames; All of these specification parameters should be understood as describing the dimensions of the generated media, whether the media is obtained by preprocessing the data source or generated directly by the rollout engine.

done (ddc5648) — dropped --sft-height/--sft-width/--sft-num-frames; SFT now reads --diffusion-height/--diffusion-width/--diffusion-output-num-frames. --sft-frame-stride stays (preprocessing-only, no rollout analogue).

Rockdu · 2026-08-05T20:59:15Z



 def set_default_diffusion_args(args) -> None:
+    if args.loss_type == "sft_loss":


We should follow Miles LLM- give args through our sample instead of making a default custom function. These default args can be redundant in the codebase

done (9dc9b8e) — removed the auto-wiring block; the sample script (scripts/run-diffusion-sft-wan22.sh) now passes the five plugin paths explicitly, and validation rejects sft_loss when a path is missing / still the RL default (with the exact flag to set), so misconfigs fail at parse time instead of deep in the RL data path. Validation-not-defaulting keeps the command line describing what actually runs.

Rockdu · 2026-08-05T21:10:26Z

+        from miles.utils.misc import load_function
+
+        sft_cfg_cls = load_function(args.train_pipeline_config_path)
+        if sft_cfg_cls.encode_sft_sample is TrainPipelineConfig.encode_sft_sample:


The encoders are a relatively independent part of the miles, and TrainPipelineConfig should only organize everything that happens in the training engine. We shouldn't make training engine and other parts over-coupled

done — encoder logic moved out of TrainPipelineConfig into miles/rollout/encoder_hub, split as a stacked base PR #96 (this PR's base is now feat/encoder-hub). Argument validation no longer probes the TPC class; it dispatches encoder_hub by args.diffusion_model_family (6486fe6).

Rockdu · 2026-08-05T21:11:19Z

+                f"--loss-type sft_loss is not supported for {sft_cfg_cls.__name__}: "
+                "it does not implement load_sft_encoder/encode_sft_sample"
+            )
+        if args.diffusion_flow_shift is None:


Let's rename this to fsdp_flow_shift since it's regenerated at training side

done (ddc5648) — added --fsdp-flow-shift for the training-side sigma grid and SFT now requires it; kept --diffusion-flow-shift untouched as the rollout-engine launch parameter for RL.

Rockdu · 2026-08-05T21:30:32Z


+    # Formal "no rollout engines" mode: skips engine/router startup, weight sync, and
+    # the rollout placement view. Implied by debug_train_only and by SFT.
+    args.train_only = args.debug_train_only or args.loss_type == "sft_loss"


train_only rename here makes more sense, but we should do as few silent inferences on args as we can since Miles wants to expose every option explicitly to users. Here, let's make debug_train_only an alias of train_only and only do verification for args.train_only==True when loss_type="sft_loss"

done (bcef656) — --train-only is now a real flag with --debug-train-only as an argparse alias (same dest, old scripts unchanged), the codebase reads only args.train_only, and sft_loss validates train_only==True instead of inferring it; the sample script passes --train-only explicitly.

Rockdu · 2026-08-05T21:54:16Z

        """Postprocess the model after FSDP wrap + weight materialization (default: no-op)."""
        return None
+
+    def load_sft_encoder(self, args, device: torch.device):


SFT encoders loading: according to Miles' philosophy, encoders should be directly passed as an HF checkpoint name argument option

done — encoders now load from the explicit --sft-encoder-checkpoint (HF name or path with tokenizer/text_encoder/vae subfolders), passed in the sample script; it also replaces hf_checkpoint in the per-sample cache key since the encoder weights are what determine cache content (#96 + 6486fe6).

Rockdu · 2026-08-05T21:56:28Z

+        """Load this family's frozen encode components (tokenizer/text encoder/VAE) for SFT caching."""
+        raise NotImplementedError(f"{type(self).__name__} does not implement SFT encoding")
+
+    def encode_sft_sample(self, encoder, pixels: torch.Tensor, prompt: str, generator: torch.Generator) -> dict:


Also, this process function should be moved to rollout and not coupled with the training side; maybe let's create an encoder_hub for all encoder logic

done — created miles/rollout/encoder_hub as the stacked base PR #96: load_encoder / encode_sample / validate_args per family, rm_hub-style lazy dispatch. The SFT encode actor pool consumes it from the rollout side; both encoder methods are removed from TrainPipelineConfig (6486fe6).

Rockdu · 2026-08-05T22:08:35Z

        return noise_pred_neg + scale * (noise_pred_pos - noise_pred_neg)
+
+    @classmethod
+    def validate_args(cls, args) -> None:


We should put this validation on the encoder side as well, disentangled from the trainer TPC. (sft_num_frames - 1) % 4 == 0 is a Wan encoder/VAE-specific constraint, not a global SFT rule. As we already resolved diffusion_model_family in arguments.py (this arg is designed to make model-specific logic decentralized for easier maintenance), we should directly do model-specific operations according to the args.diffusion_model_family

We want to keep the train pipeline config only as the driver for the training backend

done — the 4k+1 constraint moved to encoder_hub/wan2_2.validate_args, dispatched via args.diffusion_model_family at argument validation; Wan's TPC validate_args override is deleted and TPC is back to training-backend-only (#96 + 6486fe6). One accepted limitation: a custom --train-pipeline-config-path resolves family=None, which has no encoder_hub entry, so SFT rejects it until an override is actually needed.

Rockdu · 2026-08-05T22:32:11Z

+        timesteps_for_model = timesteps
+
+    cond_list = [{key: value.to(device) for key, value in pair["cond_kwargs"].items()} for pair in batch]
+    pos_cond = cast_cond_to_dtype(


Rockdu TODO: integrate SFT&NFT into input dtype precision control

Rockdu · 2026-08-05T22:36:28Z

+from miles.utils.metric_buffer import MetricBuffer
+
+
+def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]:


TODO: centralize forward noising logic (in NFT/SFT) into diffusers/local-maintained schedulers and add new args for train-side scheduler designation

…--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>

…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>

…mes 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>

…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>

zhihengy · 2026-08-06T00:52:56Z

Review resolution note — encoder comments are resolved in a stacked PR.

The encoder-related review comments (TPC coupling, encoder as explicit HF checkpoint arg, encoder_hub on the rollout side, Wan 4k+1 validation dispatched by model family) are addressed by splitting the encoder logic into a new base PR:

feat(diffusion): encoder_hub — family-dispatched frozen-encoder logic #96 (feat/encoder-hub → main): introduces miles/rollout/encoder_hub/ — get_encoder(diffusion_model_family) dispatch, with wan2_2 providing load_encoder (from the explicit --sft-encoder-checkpoint), encode_sample, and validate_args (the 4k+1 VAE constraint). Purely additive, no callers.
this PR (base switched main → feat/encoder-hub): commit 6486fe6 removes load_sft_encoder/encode_sft_sample and the Wan validate_args override from TrainPipelineConfig, wires argument validation and the SFT encode actor pool to encoder_hub, and adds --sft-encoder-checkpoint (which also replaces hf_checkpoint in the per-sample cache key, since encoder weights are what determine cache content).

Merge order: #96 first, then this PR. Per-thread replies are under the corresponding inline comments (some may show as outdated after the refactor).

Rockdu · 2026-08-06T03:26:09Z

    return {
-        "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids),
-        "rollout": (pg, rollout_pg_reordered_bundle_indices, rollout_pg_reordered_gpu_ids),
+        "actor": (pg, all_reordered_bundle_indices, all_reordered_gpu_ids),


From my understanding, this assumes all settings are collocated training, which is not good for later framework evolution.

Rockdu · 2026-08-06T04:24:41Z

+IMAGE_EXTENSIONS = {".bmp", ".jpeg", ".jpg", ".png", ".webp"}
+
+
+def read_media_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor:


Seems like this can be replaced by any general data preprocessing. Do we want to make a custom function here? Just a discussion, we can keep this here for now

Rockdu · 2026-08-06T07:49:45Z

+    num_grid = len(ctx.scheduler.timesteps)
+    if len(ctx.models) == 1:
+        component_name, model = next(iter(ctx.models.items()))
+        return component_name, model, torch.randint(num_grid, (bsz,))


This may break determinism for load/save

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

zhihengy added 2 commits August 3, 2026 17:50

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

6867abd

feat(diffusion): add frame stride to SFT encode script

0dbb5e2

zhihengy force-pushed the feat/sft branch from a79b4d3 to 0dbb5e2 Compare August 3, 2026 17:50

zhihengy added 6 commits August 4, 2026 00:02

feat(diffusion): auto-build SFT cache via family encode hooks

cd743e1

fix(sft): per-sample content-addressed cache, atomic writes, seeded V…

17fdec6

…AE sampling

refactor(sft): plug into RolloutManager via rollout-function/convert/…

a6dcec3

…log hooks with lazy encoder pool

feat(sft): image datasets for Wan (single-frame media branch)

48e6d37

refactor(sft): pick dual-expert via uniform anchor index

0220401

refactor(sft): seat encoder pool via the rollout placement view

03b36e1

zhihengy marked this pull request as ready for review August 5, 2026 02:59

zhihengy requested review from Rockdu and guapisolo as code owners August 5, 2026 02:59

chore(sft): encoder pool inherits the engine's 0.3 GPU share

a2db7a7

Rockdu reviewed Aug 5, 2026

View reviewed changes

zhihengy and others added 3 commits August 5, 2026 23:32

zhihengy mentioned this pull request Aug 6, 2026

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

Open

zhihengy and others added 2 commits August 6, 2026 00:43

Merge feat/encoder-hub: encoder_hub base for SFT rewiring

55f79fb

zhihengy changed the base branch from main to feat/encoder-hub August 6, 2026 00:49

Rockdu reviewed Aug 6, 2026

View reviewed changes

zhihengy added 2 commits August 6, 2026 06:03

Merge feat/encoder-hub: dtype rationale comment

a23e884

Merge feat/encoder-hub: align encode dtypes with sglang-d

a6d546c

Rockdu reviewed Aug 6, 2026

View reviewed changes

style: black-format the train_only exclusion assert

ca5982b

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



		def set_default_diffusion_args(args) -> None:
		if args.loss_type == "sft_loss":

		from miles.utils.metric_buffer import MetricBuffer


		def sample_grid_indices(ctx: DiffusionLossContext, bsz: int) -> tuple[str, nn.Module, torch.Tensor]:

		IMAGE_EXTENSIONS = {".bmp", ".jpeg", ".jpg", ".png", ".webp"}


		def read_media_clip(path: str, *, height: int, width: int, num_frames: int, frame_stride: int) -> torch.Tensor:

Uh oh!

Conversation

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

Uh oh!

What

How

Validation

GPU validation: head-to-head vs VideoX-Fun + before/after

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhihengy commented Aug 6, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

zhihengy commented Aug 3, 2026 •

edited

Loading

zhihengy Aug 6, 2026 •

edited

Loading