feat(minimax-h3): DiffusionNFT on the vLLM-Omni rollout backend at 32-GPU scale - #5
Draft
nussejzz wants to merge 39 commits into
Draft
feat(minimax-h3): DiffusionNFT on the vLLM-Omni rollout backend at 32-GPU scale#5nussejzz wants to merge 39 commits into
nussejzz wants to merge 39 commits into
Conversation
Adds unirl/models/minimax_h3/ covering MiniMax-H3's t2va task (text -> video + stereo audio) on the FL2VA checkpoint, plus a trainside recipe. MiniMax-H3 is a 33B dense omni-modal transformer that denoises video and audio jointly in one packed sequence. Three checkpoint contracts drive most of the design: - Guidance is DISTILLED: no CFG, no negative prompt, one forward per step. Conditions carry a single `text` slot with no negative twin. - Velocity is DATA-ward (x0 = x_t + sigma*v), the opposite of diffusers. Negating the transformer output once in predict_noise makes the whole existing FlowSDEStrategy / log-prob / replay stack apply unchanged. - Two schedules: video shift=12, audio shift=3. Only the video grid is pinned onto the Part; audio is derived in-stage from the same static formula, so LatentSegment/DiffusionSamplingParams are untouched. H3 is also the repo's first MIXED-DTYPE checkpoint: _keep_in_fp32_modules stays fp32 while the block stack is bf16. finalize_meta_init grows an opt-in `keep_in_fp32` argument (default None reproduces the historical uniform cast exactly, so no existing bundle changes behaviour), and the recipe uses root_wrap: false so those modules sit outside every FSDP group and MixedPrecisionPolicy never all-gathers a bf16 compute copy. Batch-1 throughout: the packed sequence carries unbatched per-row metadata, so callers chunk via forward_batch_size/micro_batch_size, the bagel navit discipline. Audio decodes to length-first [L, 2] stereo rather than LTX-2's mono downmix. Vendors 5 files (~3.1k LOC) from the unmerged diffusers PR #14355 at abc5e9bf, since no released diffusers ships H3 and the PR is a draft. See vendor/VENDOR_COMMIT.txt for the file map and re-vendor procedure. Verified (no checkpoint required): - sigma grids bit-identical to the reference MiniMaxH3Scheduler for both shifts at n=10/16/50 (max|d| = 0.000e+00) - 10-step Euler rollout matches the reference under the negation to 9.5e-07; the un-negated control diverges by 7.5e-02 - packed layout for 768x768/124f: 21982 rows = 256 text + 414 audio + 21312 video, float64 positions, text rows share the video timestep - pre-commit run --files ... all hooks pass NOT yet verified: anything requiring the 250GB checkpoint on a GPU pod -- full ODE parity against a reference clip, dtype survival through FSDP wrap, rope.inv_freq survival through to_empty, rollout<->replay logprob agreement.
…bundle.device Two e2e bring-up failures on an 8xH20 pod, both in the same seam. MiniMaxH3Bundle was a plain @DataClass -- the only bundle in the repo not subclassing models.types.bundle.Bundle. The recipe declares it as a top-level role, so Worker.add_remote() calls role.setup(...) on it and every rank died with AttributeError: 'MiniMaxH3Bundle' object has no attribute 'setup' after fully loading 135 GB of weights. Remote.setup() then rebinds self.device to the worker's device STRING ('cuda:0'), so the stage's device-gated autocast raised on str.type. Gate on the dtype instead, which is what every other diffusion stage here does. Found by running the recipe, not by reading it.
…om params Second and third e2e bring-up failures. MiniMaxH3Pipeline was a bare class -- the only pipeline in the repo not subclassing models.types.pipeline.Pipeline -- so the recipe's top-level pipeline role died in Worker.add_remote() with the same missing-setup AttributeError the bundle just hit. Behind it sat a harder one: replay() required an explicit geometry kwarg, but StageAlgorithm's contract is replay(conditions, segment=, params=, step_indices=) and FlowGRPO has no geometry to pass. Every training step would have raised. Geometry is a pure function of the shared (height, width, num_frames) already on params, so derive it when absent and keep the argument as an override. Also gives the step/stage their DiffusionStep/DiffusionStage protocol bases, matching ltx2 and every other diffusion package here.
… one FSDPBackend.__init__ has no root_wrap parameter -- it reads fsdp_cfg.root_wrap and forwards it to fsdp_wrap. With the key one level too high the recipe died with TypeError before any wrap happened, which would have left the fp32-pinned modules inside the mp_policy had it silently defaulted instead.
… outside FSDP fsdp_wrap(root_wrap=false) rejected the recipe with 24 trainable params outside every fully_shard group, all under token_refiner.refiner_blocks.*. The target suffixes (attn.to_q, ff.net.0.proj, ...) are not unique to MiniMaxH3TransformerBlock: MiniMaxH3TokenRefinerBlock, a plain pre-norm block over the text stream, carries the same ones and is not in block_class_names. The recipe's comment claimed the targets were 'all strictly INSIDE MiniMaxH3TransformerBlock' -- that was wrong. module_prefix: transformer_blocks scopes them to the 50-layer denoising stack, which is what the recipe means to adapt. Wrapping the refiner instead would also have satisfied the guard but silently widens what is trained.
…n its own device
Two bugs in one stage, both found on the first real rollout.
transformers 5.x nests Qwen3-VL as model.{visual,language_model}, so the
decoder stack is at model.language_model.layers, not model.layers. The bound
check died with AttributeError: 'Qwen3VLModel' object has no attribute
'layers' before a single denoising step ran. Probed against the real
checkpoint: 64 decoder layers, hidden_states is a 65-tuple, hidden_states[50]
is (1, T, 5120) and differs from last_hidden_state -- i.e. layer 50 really is
the intermediate H3 conditions on, and 5120 matches the DiT's text_dim.
The guard was also off by one (> vs >=): hidden_states has num_layers+1
entries, so index 50 is valid at exactly 50 layers.
Second, the recipe sets aux_components_on_cpu, which parks the 32B conditioner
on CPU while bundle.device stays the train device -- so input_ids and
mm_token_type_ids were being built on CUDA for a CPU module. Build them on the
encoder's own device and move only the resulting hidden state back.
to_list() returns list[Text] dataclass wrappers, so the tokenizer got a Text instance and raised 'text input must be of type str'. ltx2 and wan21 both read texts.texts; do the same.
…ls, trajectory positions Four bugs in the rollout path, all found by running one generate() on a single GPU rather than through the trainer. - compute_trajectory_positions(sde_indices, num_steps) takes two arguments; the stage passed one, so generate() died before the first forward. - The audio x_T was (rows, 2) instead of (rows, 32). MINIMAX_H3_AUDIO_CHANNELS is the STEREO count -- it is why audio occupies two channel-major row blocks and is already folded into num_audio_rows -- not the per-row feature width. The transformer's audio_proj_in is Linear(32, hidden) and the audio VAE declares latent_channels=32 with 32-long latents_mean/std. Added MINIMAX_H3_AUDIO_LATENT_CHANNELS next to the video constant so the two cannot be confused again. - patchify_video_latents returns 2-D (batch*rows, C) -- the reference pipeline is strictly batch-1 so it folds the batch away -- but the transformer indexes rows on dim 1. Restore the batch axis explicitly. - The video decode never passed the required "channels" argument to unpatchify_video_tokens. Verified against the checkpoint: geometry resolves to 768x768/124f -> latent (24, 37, 48, 48), 21312 video rows + 414 audio rows, and the Qwen3-VL layer-50 conditioning is (1, T, 5120), matching the DiT's text_dim.
…ature make_denoise_step_generators takes (base_seed, step_index, sample_ids); the stage called it once before the loop with (keys, base_seed, device), so generate() raised TypeError on the unexpected 'keys'. Placement was wrong too, not just the kwargs. The seed tuple includes step_index, so the generators must be rebuilt inside the loop. Hoisting them out would have each step draw from a running stream rather than its own seed -- reproducible within a process, but no longer byte-identical across engines, which is exactly the rollout/replay agreement flow-GRPO's ratio depends on. Now built per step and only when eta > 0, mirroring ltx2.
…nditioner aux_components_on_cpu moved the 32B conditioner AND both VAEs to CPU. The conditioner belongs there -- it is ~64 GB bf16, larger than the per-rank DiT shard, and measured 0.8s to embed a prompt from CPU, which is noise against a 22k-row denoising loop. The VAEs do not. Decoding one 124-frame 768x768 sample through the fp32 video VAE on CPU ran past 7 minutes without finishing (154 cores saturated), against ~2 minutes for the entire 10-step denoise. It would have been the dominant cost of every rollout. Together the two VAEs are only ~10 GB fp32 next to a 7.75 GB/rank shard, so there is room. Split into vae_components_on_cpu, defaulting False, and corrected the recipe comment that asserted the VAEs had to be parked to fit.
pipeline.generate read sample.sample_ids for the denoise seed keys, but Sample exposes root_group_ids()/split() and has no sample_ids -- the field is on Part (sample.py:90). Rollout died with AttributeError on the first generate(). Use gen.sample_ids, which the function already holds as sample.parts[-1]. NOTE: unirl/models/ltx2/pipeline.py:319 has the identical unguarded line and will raise the same way on any LTX-2 generate. Left alone deliberately -- it is a separate model, out of scope here, and a drive-by fix would ship unverified.
Part.fill returns a Part, but TrainsideRolloutEngine._generate_core does `chunk.parts[-1]` on generate()'s return value (engine.py:153), so rollout died with AttributeError: 'Part' object has no attribute 'parts' -- after running the full denoise and both VAE decodes, i.e. the entire generate path is sound and only the hand-back was wrong. Wrap it the way sd3, wan21 and ltx2 all do: Sample(parts=[*sample.parts[:-1], filled], reward_compute_s=...)
transformers 5.x rejects the deprecated `audios=` keyword outright. Because CLAPRewardScorer runs inside the composite's try/except, it does not crash -- every sample is flagged as a scoring failure instead, and the trainer stops with "Reward computation flagged 2 of 2 sample(s) as failure. First few: [(0, 'You passed keyword argument `audios` which is deprecated...')]". Shared reward code rather than an H3 file, but it is on the t2av path and blocks any recipe blending `clap` -- the ltx2 audioreward recipe included. The substitution is the one the library's own message prescribes.
…on 95 GB compute_loss_and_backward OOM'd on an H20 after a ~10-minute rollout, trying to allocate 1.16 GiB against a 95 GB card. replay() builds one autograd graph spanning every replayed step. At ~22k packed rows each step retains roughly 12 GB of activation-checkpoint boundaries (50 layers x 21982 rows x 5376 hidden, bf16), so all 10 steps want >100 GB per rank on top of the ~26 GB baseline (7.75 GB shard + 10 GB VAEs + reward models). sde_indices selects which steps record SDE log-probs, and therefore which get replayed and trained on -- the standard flow-GRPO timestep-subset lever, no code change. Four steps fits. LTX-2 never needed this because its packed sequence is about an order of magnitude shorter.
The flat reward curve was not an optimisation problem. At eta=0.7 the rollout produced blue/white noise, so every reward was scored on garbage. Isolated it by decoding real rollouts. The model and the port are fine: at eta=0 the same code generates a coherent, prompt-faithful 124-frame clip (verified at both 10 and 24 steps). Step count is not the variable -- eta is. Swept at 768x768/124f, 10 steps, judged on decoded frames: 0.00 clean · 0.05 clean · 0.10 clean · 0.20 marginal · 0.30 damaging 0.70 unrecognisable The cause is that H3 is guidance-distilled: its trajectory is sharp and it was never trained to denoise from states the SDE noise knocks it onto. The usual flow-GRPO eta is simply not transferable to a distilled checkpoint. Also confirms the velocity negation against the REAL model, which the earlier synthetic-velocity check could not: negated gives a final latent std of 1.04 (correct for these unit-variance VAE latents), un-negated diverges to 2.96.
…t eyeballing 0.1 was picked by looking at frames, which can only see quality. Scoring the sweep with the actual reward (videopickscore, 4 samples per eta sharing one x_T, so the spread is exactly the diversity the policy's log-prob covers) shows 0.1 is the better sampler and the worse policy: eta quality explore 0.10 0.9186 0.0034 quality ceiling, almost no signal for GRPO to rank 0.20 0.8799 0.0121 3.6x the signal for -4% quality <- chosen 0.30 0.8577 0.0083 dominated by 0.2 on BOTH axes 0.40 0.8020 0.0209 most signal, -13% quality 0.2 is the knee. The 0.3 dip is not structure -- at K=4 the std estimate is itself noisy -- but 0.3 being no better than 0.2 on quality either makes 0.2 the safe pick.
…ame 0
videopickscore represents a whole clip by ONE frame, so which frame is a
modelling choice rather than a detail. On a MiniMax-H3 rollout that opens with
a reveal, frame 0 had std 0.015 -- effectively blank -- and scoring it cost
-0.0815 against mid-clip:
frame 0 20 60 123
score .6526 .7075 .7340 .7347
std .015 .156 .263 .271
For scale, that penalty is 6.7x the entire policy-attributable reward spread at
eta 0.2 and 24x at eta 0.1. It is also sample-dependent -- a clip that opens on
content pays nothing -- so it enters the advantages as variance, swamping what
the policy actually controls.
Adds VideoPickScoreSpec.frame_selection ("first" | "middle"), defaulting to
"first" so every existing recipe is byte-identical. VideoPickScoreScorer needs
its own __init__ because PickScoreRewardScorer.__init__ consumes only
device/batch_size/processor_id/model_id. T2AVCompositeSpec forwards the field
to inner specs that declare it, alongside device/batch_size.
Measured on one clip; the effect size is large and the mechanism is plain, but
it is one clip.
…efault to 8 GRPO normalises rewards within a group. For a group of two, (r - mean)/std is algebraically always +-1, so the update carries no magnitude at all -- only which of the two sibling samples scored higher. With the measured policy- attributable spread (~0.012 at eta 0.2) sitting below the within-rollout reward std (~0.020), that sign is close to a coin flip. Three arms x ~80 steps at group 2, with the eta, dead-clap and blank-frame defects all already fixed, gave slopes indistinguishable from zero (t = +0.45, -2.22, -0.03). No reward-side fix can rescue an estimator that throws away magnitude before the gradient sees it. Costs 4x the rollout of 2 (~55 min/step vs ~15 on 1x8 H20).
… better quality FlowSDE adds exploration noise ON TOP of the level the sigma schedule prescribes, handing the denoiser states noisier than anything it trained on. H3 is guidance-distilled and has the least capacity of any model class to spend steps cleaning that up, so the residue lands in the output as speckle. CPS reallocates WITHIN the budget (det^2 + std^2 == sigma_next^2 exactly), keeping every state on-schedule. Measured, K=4 per eta sharing one x_T, real videopickscore on the mid frame: kernel eta quality explore latent std video CPS 0.7 0.8776 0.0162 1.001 clean FlowSDE 0.2 0.8748 0.0039 1.028 clean <- previous default FlowSDE 0.7 0.7282 0.0195 1.228 speckle CPS at 0.7 wins on BOTH axes against the previous default. Note the trap in the last row: FlowSDE at 0.7 has the highest reward spread in the sweep and the worst video, because that spread is speckle variance rather than scene diversity. Tuning eta on in-group reward std alone selects the setting that destroys the sample. Config-only: denoise() dispatches generically through step() + compute_log_prob.
…t -- drop to 1e-12 Adam's step is lr*m/(sqrt(v)+eps), scale-invariant only while sqrt(v) >> eps. This adapter has 166,297,600 trainable params against a grad_norm of ~5e-5, so the per-coordinate gradient is ~3.9e-9 -- below the 1e-8 default. Epsilon then dominates the denominator and the update shrinks in proportion to the gradient, degrading Adam toward SGD exactly when scale-invariance was being relied on. Measured: ||lora_B|| after 2 rollouts is 0.194 under CPS vs 0.707 under FlowSDE at matched group size. That 3.6x gap is an artefact of epsilon, not of the kernel -- CPS's log-prob omits the 1/(2*sigma^2) factor, so its gradient is smaller and falls further below eps. At 1e-8 even the FlowSDE runs were getting only ~60% of a full step. 1e-12 puts eps ~4 orders below sqrt(v), negligible again, and stays far above fp32 underflow (v ~ 1.5e-17).
…ormers>=5 The scorer failed at import on transformers 5.x with ImportError: cannot import name 'AutoModelForTextToVideo' so videoalign was unusable on any environment tracking current transformers (this repo pins 5.11.0 for the Qwen3-VL conditioners). The Qwen2-VL backbone it wraps now resolves through AutoModelForImageTextToText; verified on 5.11.0 that .from_config(Qwen2VLConfig) yields the same Qwen2VLForConditionalGeneration the old class produced. The import-time fallback keeps transformers 4.x working.
…nfig in transformers>=5
…t 0.0 rewards The published weights (KlingTeam/VideoReward) are a single checkpoint-*/model.pth -- a .pth in a SUBDIRECTORY. _load_checkpoint globbed only '*.safetensors' and 'adapter_*.safetensors' non-recursively, so it matched nothing, built an empty state dict, and load_state_dict(..., strict=False) accepted that without error. The reward heads stayed randomly initialised and every score came back exactly 0.0 -- a plausible-looking reward that is pure noise, which would have trained against a constant. Search recursively, accept .pth/.bin (unwrapping the usual state_dict/model/ module containers), and fail loudly: raise when no weights file is found, and raise when zero model keys or zero reward-head keys matched.
Add grouped H3 rollout, staged LoRA publication, row-level reward dispatch, and FSDP/HSDP recipes so collocated training stays within memory while using all reward and training ranks. Co-authored-by: Cursor <cursoragent@cursor.com>
Move the CUDA stack to cu130, pin the rebased upstream integration, and replace legacy stage YAML boot with current direct diffusion-stage arguments while retaining grouped LoRA and row-level reward behavior. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
Allow the shared package range to cover both stacks while pinning SGLang to 5.6 and current vLLM-Omni to 5.10+, preserving independent solvability. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
Update the temporary source pin after replaying the upstream dependency PR onto the latest main branch. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
Apply current import formatting and collapse legacy multiline docstrings after rebasing the trainside model package onto the latest lint rules. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
Move the checkpoint, packing, replay, conditioner, and mixed-dtype invariants out of multiline docstrings into the model README. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
Target vLLM 0.27 and Torch 2.13 cu130, consume the public unique-reply RPC API, make the qualified four-GPU topology explicit, and restore VAE tile-parallel setup for the custom pipeline path. Signed-off-by: DingZuhao <e1583181@u.nus.edu> Co-authored-by: Cursor <cursoragent@cursor.com>
…t layout add_lora() only registers CPU-side logical weights, so a name or shape mismatch between the Diffusers adapter and the fused vllm DiT is silent: the engine keeps serving base weights and the only symptom is a reward curve that never moves. Translate the H3 adapter into the serving layout in one place (transfer/minimax_h3_lora.py: blocks/out_proj/mlp renames plus the GEGLU half-swap that Diffusers stores as [up, gate] and serving consumes as [gate, up]), require every expected block/slice to carry both A and B, and compare per-layer checksums against the awake engine rather than trusting the load call. Checksum collection now replies from a single rank instead of every TP rank. Co-authored-by: Cursor <cursoragent@cursor.com>
Four defects that only surface at 32-GPU scale with heterogeneous prompts: - Global CUDA autocast ran the fp32-pinned projections, timestep MLP and rope in bf16 despite their storage dtype. H3 conditions on t = 1 - sigma, where bf16 spacing near 1.0 is ~0.008, so this was a real precision loss on the low-noise end. Disable autocast around exactly those modules. - The output adapter concatenated per-request prompt embeddings directly, which raises on any batch whose prompts tokenize to different lengths (hit at 32 prompts spanning 9-58 tokens). Pad to the batch maximum, carry an attention mask, and trim back to the true length before the packed forward, so the DiT still sees varlen input and never attends to padding. - SDE exploration keys were derived from the engine request id, which contains a fresh UUID per attempt, so a retried or resumed sample explored a different trajectory than the one being replayed. Derive them from the logical sample id instead. - Sparse GRPO replay stores only SDE boundaries, so decode could be handed a segment without the terminal latent. Always store it. Geometry now requires an explicit sampler_kwargs.allow_nonstandard_canvas opt-in for anything off the released 768-short-edge distribution, and the FlowGRPO parity gate reports per-step log-prob, transition-mean and velocity-RMS drift so a cross-engine mismatch is attributable rather than a single scalar. Co-authored-by: Cursor <cursoragent@cursor.com>
…ewards VideoPickScore scored a single frame, which cannot separate a clip that is good throughout from one that is good only where it happened to sample. Add uniform multi-frame scoring with either a plain mean or a top-k/mean blend, so a recipe can choose how much of the score comes from the best moments. The T2AV composite previously let an inner scorer fail and still contributed its zero to the weighted sum without renormalizing, which silently trains against a fraction of the intended reward. Propagate inner failures and non-finite scores instead, and forward the shared frame-selection, CLAP id and ImageBind checkpoint settings so a reward can be pinned to a local path rather than resolved relative to the worker cwd. Co-authored-by: Cursor <cursoragent@cursor.com>
…rantees Recipes for the topology that passed the capacity, LoRA-sync, parity and throughput gates: trainer HSDP shard8 x replicate4 over 32 GPUs, eight four-GPU rollout replicas at DiT TP2 x Ulysses2 with text-encoder TP4, and a resident reward service at DP32, all time-sharing the same GPUs. The _tp4 and _hsdp4 variants exist so the comparison stays reproducible; TP2 x UP2 is the fastest verified arm (91.1s vs 112.7s median generate at K8), not a proven optimum. README states what is verified and, explicitly, that reward convergence is not. Supporting fixes for long runs: - Resuming advanced the data source twice when both a checkpoint and data_source.start_batch were set, so the run silently skipped prompts. Reconcile the two cursors and reject a start_batch ahead of the checkpoint. - A checkpoint save that wrote no model state returned successfully, leaving a directory that only fails on load. Assert an artifact exists, and wait for the writer on every save rather than only the last one. - Colocated store setup timed out at 30s during 32-rank NCCL bring-up. - Trainer master weights are fp32 for this mixed-dtype checkpoint. Co-authored-by: Cursor <cursoragent@cursor.com>
…lout An EMA/DiffusionNFT run had only ever been exercised with a trainside rollout, where the trainer serves its own shadow weights. Pointing one at a separate serving engine on the same GPUs surfaces two failures, both fatal on the second rollout. Adapter scoping. EmaLoraConfig had no module_prefix, so its target suffixes could only be given unscoped. On MiniMax-H3 the denoising stack and the token refiner share `attn.to_q` and `ff.net.*`, so the trainer adapted 364 modules where the serving engine has slots for the denoising stack's 350, and weight sync rejected the payload for the 14 extras. Route inject_nft through the same resolve_target_modules_pattern that inject_lora already uses. Scoping is the honest fix rather than dropping the extras at sync time: an off-policy objective whose trainable set is wider than the set that generated its samples is not the objective the recipe asked for. Cache release. offload() ends in an empty_cache, and it was the only thing that did. A run that keeps the train state resident -- which an EMA algorithm must, since _validate_residency_config forbids it from offloading here -- never returned the reserve, so the rollout's physical allocation failed once the first optimizer step had grown it: the engine's own allocator reported out of memory. Release the caching allocator before waking a colocated external rollout, independently of whether weights move. Co-authored-by: Cursor <cursoragent@cursor.com>
DiffusionNFT does not walk a trajectory: it noises the rollout's clean final latent to a chosen level and takes one MSE there. That needs two things from a stage, neither of which the H3 stage had. predict_noise_at_step runs one packed forward at an arbitrary (xt, sigma). Both streams take the same sigma: the video/audio shift split describes the reverse schedule that generate walks, and a forward-process objective jumps to a noise level instead of walking it, so there is no second grid to be on. nft_clean_latents packs both streams into the x0 the loss trains on. The caller's default is segment.latents[:, -1], which on this model is the video half only -- audio would then receive no gradient at all while the reward scores how well the two agree. The algorithm picks the override up off the stage, so models with a single stream keep the default. Co-authored-by: Cursor <cursoragent@cursor.com>
…ackend The rollout decided what to record from the SDE schedule: a non-empty sde_indices meant "this feeds an algorithm", which held while every training objective walked an SDE. A forward-process objective does not, and neither does evaluation, so the two became indistinguishable by schedule alone -- the engine returned decoded media only and the algorithm received a segment of None. Three consequences of an empty schedule, made explicit: - The request now states whether it wants the trajectory. SDE steps still imply it, so Flow-GRPO recipes are unchanged and eval still infers the opposite from the same empty schedule. - The terminal position is always captured. It is the entire capture set for a forward-process rollout, and decode needs it even when sparse GRPO replay asks only for earlier boundaries. - Old-policy log-probs are optional rather than required, and a schedule with no SDE steps that carries them anyway is now an error rather than data the segment would silently keep. Co-authored-by: Cursor <cursoragent@cursor.com>
Same objective as the trainside NFT recipe, on the collocated serving stack the Flow-GRPO recipes use: eight four-GPU rollout replicas at DiT TP2 x Ulysses2, resident reward, FSDP32 trainer, 8 prompts x 16 samples. Standalone rather than layered on the Flow-GRPO chain: those recipes set old_logp_source / clip_range / max_rollout_replay_logp_absdiff, and Hydra would merge them into this algorithm block and hand DiffusionNFT constructor arguments it does not accept. Four settings are load-bearing and commented as such in the file. eta 0 with num_sde_steps 0 is the forward-process switch. init_same_noise must be false or a deterministic rollout gives all 16 siblings one x_T, one sample, and an advantage of exactly zero. Both offload policies are off because the trainer rejects either alongside an EMA algorithm; affordable because a LoRA run keeps only ~2.4 GB/rank resident. record_trajectory is what tells the engine this empty schedule belongs to training rather than evaluation. Tests cover the contracts that fail quietly rather than loudly: dual-stream packing round-trips through the resolved geometry, nft_clean_latents carries both streams and rejects a segment missing audio, the terminal position is captured whether or not there are SDE boundaries, and EMA LoRA scoping leaves the token refiner alone. Co-authored-by: Cursor <cursoragent@cursor.com>
Owner
Author
|
Result of the 200-rollout run referenced in the Test Plan: stopped at rollout 112. The objective trains, but this recipe's reward composition is reward-hackable and the video collapses. Reward rose monotonically, 0.13 -> 0.41 (2.6x). Video collapsed into saturated Two supporting signals from the same run:
The second contributing factor is that the reference model never anchors. This does not change what the PR claims -- the objective runs end-to-end on the |
Owner
Author
|
Correction to my previous comment: the second half of that diagnosis was wrong. I read The two terms are equal in value but their gradients w.r.t. and Magnitudes are consistent with this reading: What still holds, and is measured rather than inferred:
Also worth stating for anyone reading the metrics: DiffusionNFT has no PPO |
CjhHa1
force-pushed
the
feat/minimax-h3-vllmomni-32gpu
branch
from
August 28, 2026 08:13
12e9c1a to
2a67fff
Compare
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.
Summary
Runs the DiffusionNFT objective for MiniMax-H3 t2va on the collocated vLLM-Omni
serving stack, where it had only ever run trainside.
DiffusionNFT is a forward-process objective: it never walks an SDE, it noises
the rollout's clean final latent to a chosen level and takes one MSE there. That
breaks an assumption the rollout path had baked in — that a non-empty
sde_indicesis what distinguishes "this rollout feeds an algorithm" from "thisis an evaluation pass". With no SDE steps the two look identical, so the engine
returned decoded media only and the algorithm got a segment of
None. Therequest now says which it is; SDE steps still imply it, so Flow-GRPO recipes are
untouched and eval still infers the opposite from the same empty schedule.
Stacked on Tencent-Hunyuan#378 (Tencent-Hunyuan/UniRL), which contributes the vLLM-Omni H3
rollout backend this builds on. Only the four commits here are new.
Four changes, in the order a reviewer probably wants them:
EmaLoraConfiggainedmodule_prefix, because unscopedsuffixes adapted the token refiner too — 364 modules where the serving engine
has slots for the denoising stack's 350, and weight sync rejected the payload.
Scoping beats dropping the extras at sync time: an off-policy objective whose
trainable set is wider than the set that generated its samples is not the
objective the recipe asked for. Separately,
offload()was the only thingthat ever called
empty_cache(), so a run keeping the train state resident —which an EMA algorithm must, since
_validate_residency_configforbids itfrom offloading here — starved the rollout's physical allocation on the
second wake.
predict_noise_at_stepfor one packed forward at an arbitrary(xt, sigma), andnft_clean_latentssox0carries audio as well asvideo. The caller's default is the video half only, which would leave the
audio stream with no gradient while the reward scores how well the two agree.
log-probs became optional, and a no-SDE schedule that carries them anyway is
now an error.
trainer, 8 prompts x 16 samples. Standalone rather than layered on the
Flow-GRPO chain, whose
old_logp_source/clip_rangekeys Hydra would mergeinto this
algorithmblock and hand DiffusionNFT arguments it rejects.Related Issue
N/A
Test Plan
Environment: 4 nodes x 8 H20 (32 GPUs), MiniMax-H3 33B t2va, 256x448 / 107
frames / 10 steps, CLAP + ImageBind reward, prompts from the wan_dancegrpo
train split.
pytest— 25 passed, 9 of them new intests/test_h3_nft_forward_process.py(dual-stream packing round-trip and geometry guard,
nft_clean_latentscarrying both streams and rejecting a segment missing audio, the
forward-process switch resolving to an empty schedule, terminal-position
capture with and without SDE boundaries, EMA LoRA scoping leaving the token
refiner alone).
python -m unirl.train_diffusion --config-name=diffusion/minimax_h3/minimax_h3_t2va_vllmomni_32c_nft --cfg job --resolve— resolves, exit 0.
rollouts generated, scored, and took an optimizer step
(
reward=0.1331 loss=346.4989 gn=0.0575, thenreward=0.3007 loss=383.6635 gn=0.1497). The second rollout is the one thatexercises weight sync and the rollout re-wake, i.e. the two failures fixed in
commit 1.
progress:
rollout 1/200 reward=0.1582 loss=345.7430 gn=0.0394,rollout 2/200 reward=0.2912 loss=382.8397 gn=0.0713, ~7.4 min/step, noerrors. Reward convergence over the full run is NOT claimed here — see
Reviewer Notes.
Compatibility / Risk
existing Flow-GRPO recipe:
record_trajectoryonly adds a second way to askfor the trajectory, and the terminal position is one extra stored latent.
EmaLoraConfig.module_prefixdefaults to"", so existing EMA recipes keeptheir current unscoped behaviour.
release_cached_memory()is a new backend method on the BROADCAST path; itfires only for a colocated external rollout that is not offloading, so no
existing configuration reaches it.
([Core][Diffusion] Propagate RPC policy and support packed LoRA targets vllm-project/vllm-omni#6351), same as feat(minimax-h3): vLLM-Omni rollout backend for H3 FlowGRPO at 32-GPU scale Tencent-Hunyuan/UniRL#378.
weight is a recipe choice rather than a default.
Reviewer Notes
to filter the refiner adapters out at sync time, which would have let the
recipe keep an unscoped target list. I did not, because it silently trains
parameters the rollout never sees, and this objective is off-policy — the
samples come from the EMA shadow, so a trainable set wider than the served set
changes what is being optimised rather than just wasting work.
shadow that is still a hard copy (
ema_flat_steps: 75) say nothing aboutconvergence, and the earlier Flow-GRPO work on this checkpoint showed
ImageBind falling while CLAP stayed flat. This PR claims the objective runs
correctly end-to-end on this backend, not that it improves the reward. Happy
to hold it as a draft until the 200-rollout run finishes.
Checklist
Made with Cursor