[blacklist-twx]

What does this PR do?

MiniMax H3 is CFG-distilled. Ordinary flow-matching SFT fits its conditional prediction directly to a raw velocity target and can erode its distilled guidance behavior during prolonged fine-tuning. This PR adds opt-in CFG-calibrated flow matching to the existing offline FL2VA and prepared visual Ref2VA training paths, plus a Ref2VA LoRA recipe.

Given conditional velocity p, detached unconditional velocity u from the same current model, and the raw FM target y = noise - clean:

p_calibrated = (p + (s - 1) * stopgrad(u)) / s
loss = s ** (2 - k) * MSE(p_calibrated, y)

Both predictions use the same noisy latents, timesteps, and DiT reference latent anchors. Only the conditional prediction receives gradients. Curvature power k controls the 1/s**k curvature factor: 2 retains the inverse-CFG attenuation and 0 removes it. These are settings of one objective, not separate methods. Scale 1 retains the existing single-forward FM path.

The approach is informed by DiffSynth-Studio's H3 training and AI Toolkit's guidance-loss implementation. Calibration mitigates objective mismatch; it does not guarantee generation quality or establish the original model's distillation scale or conditioning-drop policy.

Enabling calibration adds one no-gradient DiT forward. Our prior training measurements showed about 30% additional step time; this is not a throughput benchmark of this public branch, and overhead depends on hardware and branch sequence lengths.

Checklist Before Starting

• Related: [model, config] feat: add MiniMax H3 model support #1075 introduced H3; [model, trainer] feat: pack MiniMax H3 offline training #1204 added prepared visual Ref2VA and packed offline training.
• Supersedes [model] feat: add H3 CFG-calibrated Ref2VA training #1256.
• Title follows the required [model] feat: ... format.

Test

• Extended existing CI-enumerated tests in tests/models/test_minimax_h3_packing.py and tests/models/test_model_forward_no_implicit_sync.py; both are listed in the GPU/NPU unit workflows.
• Local CPU H3 packing/CFG suite: 93 passed. Covers paired-cache creation, metadata-only runtime validation, FL2VA/Ref2VA, schedules, curvature, detached negative predictions, packing, and the unchanged default path.
• Four H20 GPUs: reduced-size, randomly initialized H3 with 56 heads, FP32/SDPA and gradient checkpointing. Tested FSDP2 SP1/DP4 and SP2/DP2, FL2VA/Ref2VA, CFG disabled and enabled with keep-visual/drop-visual negatives. 12 cases, three optimizer updates each (36 total); finite losses/gradients and nonzero LoRA updates in all cases.
• Forward loss matched the same-weight, same-input unsharded reference within tolerance; maximum absolute error 1.35899e-5. SP1 gradients matched the reference. SP2/reference gradient-norm ratios were approximately 0.5 with CFG both off and on: no additional CFG-specific norm-scaling discrepancy was observed. This does not establish full SP gradient equivalence; the scaling observation remains a separate unresolved diagnostic item, without attribution to a particular upstream change.
• GPU implicit-device-sync suite: 8 passed, 18 deselected, using separate small fixtures.
• make quality, git diff --check, and documentation task-path checks passed.
• Detailed smoke-test report.
• The distributed smoke used a manual model/condition/collator/FSDP2/optimizer harness, not a full DiTTrainer job. It does not validate pretrained generation quality, real encoder execution, BF16, multi-node behavior, checkpoint resume, or long-run convergence. The full repository test suite was not validated; shared Trainer/resume tests on this CPU environment encounter a missing CUDA Triton dependency.

API and Usage Example

model:
  condition_model_cfg:
    training_cfg_scale: 4.0
    training_cfg_schedule: constant  # constant | sigma
    training_cfg_curvature_power: 2.0  # 0..2

Paired negative caches are prepared through add_unconditional_cache():

• Default: negative prompt " " (one space), retaining the sample's Qwen visual context.
• Experimental opt-in: drop_visual=True drops Qwen visual context while retaining DiT reference latent anchors. This changes the guidance direction; it is not an equivalent or merely cheaper form of the default.
• Both policies use the same paired fields, unconditional_prompt_embeds and unconditional_packed, and the same training loss. No shared-embedding loader, cfg_unconditional_mode, or unconditional file-path setting is required.
• Value checks run on CPU during cache creation; training performs metadata checks only. External cache producers must validate values before saving them.

The optional sigma schedule uses s = 1 + (training_cfg_scale - 1) * sigma separately for video/audio, where sigma = 1 - t is the effective noise level from the existing schedulers. At scale 4, sigma=0.1 gives s=1.3, while sigma=0.9 gives s=3.7: stronger calibration at high noise, with more room for ordinary FM learning at low noise. Existing shifts affect how often each region is sampled. This is a protection/learning tradeoff, not a quality guarantee or a new noise sampler. The recipe defaults to constant.

Design & Code Changes

• Model-local objective, paired-conditioning helpers, and focused tests; both FL2VA and visual Ref2VA can enable calibration.
• Reuse native prompt-encoding helpers and preserve positive-branch noise, timesteps and reference anchors.
• Remove the separate shared-empty path, validation tokens, and per-step value-validation hook. The common DiTTrainer is unchanged from the PR base.
• Document the negative-conditioning rationale and limits of preliminary quality observations; further quality experiments can follow separately.
• Reuse existing LoRA, checkpoint, FSDP2 and SP infrastructure. No added noise sampler or sigma-bin loss weighting.
• Add the model guide, offline Ref2VA recipe, and README Supported Models entry.

Checklist Before Submitting

• Read the contribution guide.
• Applied quality, formatting and diff checks.
• Updated model documentation and recipe.
• No tasks/ scripts moved or renamed; documentation task-path check passed.
• Extended existing CI-enumerated tests and documented the separate accelerator smoke coverage and limitations.

Summary by CodeRabbit

• New Features

  • Added CFG-aware LoRA training support for MiniMax H3 Ref2VA, with configurable guidance scale, schedule, and curvature.
  • Added support for preparing negative-prompt training data while retaining visual reference inputs by default.
  • Added shared prompt encoding for text, keyframe, and reference-based inputs.

• Documentation

  • Added MiniMax H3 training guidance, an offline training configuration, and a recorded smoke-test report.
  • Listed MiniMax H3 in the supported models table.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: ByteDance-Seed/VeOmni/.coderabbit.yaml
• Review profile: CHILL
• Plan: Advanced
• Run ID: df928077-3b85-4872-9a44-d4667b539b76

---

---

📥 Commits

Reviewing files that changed from the base of the PR and between b2fd747 and a362f57.

---

---

📒 Files selected for processing (9)

• configs/dit/minimax_h3_ref2va_cfg_offline.yaml
• docs/examples/minimax_h3.md
• docs/examples/minimax_h3_smoke_test.md
• docs/index.md
• tests/models/test_minimax_h3_packing.py
• tests/models/test_model_forward_no_implicit_sync.py
• veomni/models/diffusers/minimax_h3/minimax_h3_condition/configuration_minimax_h3_condition.py
• veomni/models/diffusers/minimax_h3/minimax_h3_condition/modeling_minimax_h3_condition.py
• veomni/models/diffusers/minimax_h3/minimax_h3_core/cfg_conditioning.py

---

---

💤 Files with no reviewable changes (2)

• configs/dit/minimax_h3_ref2va_cfg_offline.yaml
• veomni/models/diffusers/minimax_h3/minimax_h3_condition/configuration_minimax_h3_condition.py

---

---

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.

---

---

---

📝 Walkthrough

📝 Walkthrough

Walkthrough

MiniMax H3 now supports CFG-conditioned training with configurable scales and calibrated flow-matching losses. The change adds unconditional-cache preparation and shared prompt-encoding helpers, plus an offline Ref2VA configuration, documentation, and tests.

Changes

MiniMax H3 CFG training

Layer / File(s)	Summary
CFG settings and loss objectives veomni/models/diffusers/minimax_h3/minimax_h3_condition/configuration_minimax_h3_condition.py, veomni/models/diffusers/minimax_h3/minimax_h3_core/training_objectives.py, tests/models/test_minimax_h3_packing.py	The condition config validates CFG scale, schedule, and curvature. Shared helpers resolve video and audio scales and compute calibrated MSE. Tests cover objective behavior, invalid settings, and noise sampling.
Prompt encoding and unconditional-cache preparation veomni/models/diffusers/minimax_h3/minimax_h3_core/minimax_h3_text_encoder.py, veomni/models/diffusers/minimax_h3/minimax_h3_core/cfg_conditioning.py, veomni/models/diffusers/minimax_h3/minimax_h3_condition/modeling_minimax_h3_condition.py, veomni/models/diffusers/minimax_h3/inference.py, tests/models/test_minimax_h3_packing.py	Shared helpers encode text and reference prompts, validate packed layouts, and prepare negative-branch inputs. The condition model creates unconditional caches. Tests cover cache construction, prompt layouts, and metadata validation.
CFG forward passes and calibrated losses veomni/models/diffusers/minimax_h3/minimax_h3_condition/modeling_minimax_h3_condition.py, veomni/models/diffusers/minimax_h3/minimax_h3_transformer/modeling_minimax_h3_transformer.py, tests/models/test_minimax_h3_packing.py, tests/models/test_model_forward_no_implicit_sync.py	The condition model attaches CFG inputs and scales to prepared samples. The transformer computes unconditional predictions without gradients and uses them for calibrated video and audio losses. Tests cover packed execution, checkpointing, and synchronization modes.
Offline recipe and usage validation configs/dit/minimax_h3_ref2va_cfg_offline.yaml, docs/examples/minimax_h3.md, docs/examples/minimax_h3_smoke_test.md, docs/index.md, README.md, tests/models/test_minimax_h3_packing.py	The offline Ref2VA configuration sets CFG-aware LoRA training parameters. Documentation describes cache requirements, CFG settings, and validation boundaries. The smoke-test report records SP1 and SP2 results and limitations. Tests cover cache serialization and the recipe.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant ConditionModel as MiniMaxH3ConditionModel
  participant Transformer as MiniMaxH3DiTModel
  participant Objective as cfg_calibrated_mse
  ConditionModel->>Transformer: provide conditional and unconditional inputs
  Transformer->>Transformer: compute unconditional predictions without gradients
  Transformer->>Transformer: compute conditional predictions
  Transformer->>Objective: pass predictions, targets, scales, and curvature
  Objective-->>Transformer: return calibrated modality losses

Loading

Suggested reviewers: nancyfyong, coach257

---

---

Merge Risk: ⚪ Minimal · up to a362f

No actionable PR-introduced issue is established. The reported SP2 scaling discrepancy remains a separate diagnostic, not a demonstrated CFG regression.

Security Architecture Review

Security architecture risk: 🔵 Low · up to a362f

The change is opt-in and preserves the ordinary training path by default. No introduced security vulnerability was established. Remaining uncertainty concerns external cache producers, cache integrity, and production recovery behavior.

Retained concerns
No architecture-level concerns identified.

Security review details

Security Blast Radius

• inferred — The newly consumed cache values can affect losses and learned parameters in training jobs that explicitly enable calibration. Inspected preparation uses the current sample rather than a global cache or cross-job lookup. Wider tenant, service, credential, and environment exposure depends on deployment facts not supplied here.

Trust Boundaries and Controls

• observed — The cache producer is responsible for value validation, while the training consumer enforces metadata compatibility. These controls enforce numerical and layout contracts, not producer authentication or artifact provenance. The public cache helper rejects unsupported training references and validates its output; external producer compliance remains unknown.

Resilience and Maintainability Implications

• observed — The negative pass is scoped to a gradient-disabled context and the loss detaches its prediction, preserving the intended conditional-only gradient boundary. Preparation and cache construction use new result containers rather than installing shared cross-job state. These local safeguards do not establish interruption or checkpoint-resume correctness for a complete production job.

---

Pre-merge checks | 4 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 63 functions across 10 files. (4 skipped:…	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely identifies the main change: adding CFG-calibrated Ref2VA training for MiniMax H3.

---

Full details: Docstring Coverage

Explanation

Docstring coverage is 28.57% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 63 functions across 10 files. (4 skipped: 4 unsupported.)

---

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

---

---

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

---

_{Comment @coderabbitai help to get the list of available commands.}

[Coach257]

Thanks for the PR. I checked out b877e04 and reviewed it locally. Overall the objective is implemented correctly and the default path is unchanged; the inline comments are mostly about a sync-contract regression, the sigma schedule's effective behavior, and data/doc gaps.

Verified locally (CPU)

• cfg_calibrated_mse: recomputed independently that k=0 equals MSE(p, s*y - (s-1)*stopgrad(u)) and that u receives no gradient. Signs are consistent (p, u are both -logits, target is noise - clean).
• prepare_unconditional reuses the already-noised video/audio rows and reference anchors and remaps per-token timesteps; no extra RNG draws.
• replace_text_layout / the geometry check match build_packed_fl2va / build_packed_ref2va (text on axis 0 as 0..text_len, visual/audio coordinates offset by text_len).
• training_cfg_scale=1 keeps the old loss and RNG sequence; pack_samples ignores the new keys; disabling checkpointing for the no_grad branch is correct since the DiT only honors the explicit flag.
• tests/models/test_minimax_h3_packing.py: 82 passed. All H3 tests + tests/trainer/test_dit_microbatch.py: 114 passed, 21 skipped (GPU-only). make quality passes.

Not verified: anything on GPU (including test_no_implicit_sync_in_minimax_h3_forward_backward), FSDP2/SP, or end-to-end training.

[Coach257]

Thanks for the quick update. I re-reviewed b2fd747; the previous round is addressed:

• Implicit syncs: fixed. On an A100, test_no_implicit_sync_in_minimax_h3_forward_backward passes all 12 cases (off / per_sample / shared_empty x single / packed x SDPA / FA2). As a negative control I injected one bool(tensor.all()) into prepare_unconditional; both CFG cases fail as expected, so the gate really covers the new path.
• sigma vs constant (now the recipe default), the documented scale distribution, the " " negative prompt, the cache helpers, and removing video_sigma_bucket_weights / training_cfg_loss_mode all look good. CPU tests (H3 + DiT trainer + resume): 122 passed; make quality passes.

Two follow-up suggestions below. Together they make the PR smaller and remove the shared DiTTrainer change entirely:

1. Remove shared_empty. It is not a cheaper version of per_sample; it is a different CFG objective (details inline).
2. Drop per-step value validation. Validate once when the cache is written, and keep only metadata checks at train time. Metadata checks run fine on device tensors without syncs, so prepare_condition_batch, _cfg_batch_token and the DiTTrainer.preforward hook are no longer needed. No legacy-layout compatibility is required.

With the trainer change gone, the [model] title would be accurate again (right now the commit says [model, trainer], the PR title says [model]). A short multi-GPU FSDP2 smoke run is still the remaining open item from the last round.

[blacklist-twx]

@Coach257 Thank you for the prompt and detailed review, and for taking the time to run the checks yourself! I’m addressing your latest comments locally and will update the PR after validation.
Negative-conditioning policy
I agree that the default should retain the Qwen visual context and replace only the text prompt, consistent with VeOmni’s existing negative-conditioning path.
I would also like to retain an explicit, opt-in experimental option to drop the Qwen visual context. Importantly, this would not remove the DiT reference latent anchors: it changes only the Qwen conditioning used by the negative pass. I agree that this changes the guidance direction and should not be presented as an equivalent or merely cheaper version of the default.
For context, the related open-source implementations take somewhat different approaches:

• DiffSynth: its positive and negative prompt-encoding paths share the keyframes/reference blocks while using different text prompts. This supports retaining Qwen visual context as the default. Source
• AI Toolkit training: its guidance-loss path uses cached unconditional prompt embeddings, constructed through the static prompt encoder without per-sample reference images in the initial encoding call. This provides a related precedent for a text-only negative-conditioning experiment, rather than proof that the objectives are equivalent. Cache construction, guidance-loss usage
• AI Toolkit and vLLM-Omni inference: both implement H3’s CFG-distilled, single-pass inference without a separate negative prediction. However, this does not establish how the original distillation teacher dropped conditioning. AI Toolkit, vLLM-Omni
  In our limited internal experiments, removing Qwen visual context from the negative pass appeared to preserve generation quality better during fine-tuning. The datasets and training horizons were limited, so I would treat this as a preliminary observation, not a general recommendation.
  My proposed scope is therefore: retain visual context by default, document dropping it as an experimental alternative, and use the same paired-cache schema and loss implementation for both. This would avoid reintroducing a separate shared-embedding loader or changes to the common Trainer. I’d appreciate your feedback on whether this is an acceptable scope for the PR.
  Validation
  I’m also preparing a small multi-GPU smoke test covering FSDP2 and sequence parallelism. Before running it, I’d appreciate your guidance on the expected coverage: whether a reduced-size H3 model with synthetic cached inputs is sufficient to validate distributed forward/backward and parameter updates, or whether you would prefer a short run with pretrained weights and real cached samples.
  I can provide the exact code revision, configuration, hardware/software environment, and per-step loss and gradient norms. Please let me know if there are additional cases or checks you would find useful.

[Coach257]

@blacklist-twx Thanks for the detailed reply. The proposed scope sounds good to me:

• Default: keep the Qwen visual context and replace only the text prompt.
• Dropping the Qwen visual context (DiT anchors kept) as a documented, opt-in experimental option, clearly described as a different guidance direction.
• Both options share the paired-cache schema and the same loss; no separate shared-embedding loader and no change to the common DiTTrainer. Please also move the per-step value validation to cache creation as suggested, keeping only metadata checks at train time.

On effectiveness: feel free to keep running your experiments on the visual-context-free negative. Sharing the results here (or in the docs) once you have them would be great, but it doesn't need to block this PR.

On the multi-GPU smoke test: a reduced-size H3 model with synthetic cached inputs is sufficient for correctness. The useful checks are loss / grad-norm parity on identical inputs for single GPU vs. FSDP2, and SP off vs. on, with CFG both off and on, plus confirming that the LoRA parameters actually update. Pretrained weights and real samples aren't needed for this PR.

[blacklist-twx]

@Coach257 Thanks for the review. I’ve pushed the updates addressing your comments and added the four-GPU FSDP2/SP test report, including its limitations. The 93 focused tests and documentation build pass.

Please take another look when convenient. Any further suggestions are welcome.