[model] feat: add H3 CFG-calibrated Ref2VA training

.agents/knowledge/architecture.md

@@ -211,7 +211,12 @@ YAML Config -> VeOmniArguments -> Trainer
 ## DiT Fixed Microbatches
 
 `DiTTrainer` honors `train.micro_batch_size` and keeps `dyn_bsz=false`.
-`DiTDataCollator` produces dict-of-lists microbatches; the existing
+`DiTDataCollator` produces dict-of-lists microbatches.
+For offline training, `preforward` calls the optional condition-model
+`prepare_condition_batch(batch)` hook on CPU before recursively transferring
+tensors to the accelerator. H3 uses it to validate cached CFG embeddings and
+layouts; models without the hook keep the same transfer behavior. The hook
+returns the batch and must not sample training noise. The subsequent
 `get_condition` / `process_condition` / model-forward path is unchanged.
 Models return sample-mean scalar losses, and the trainer divides by the number
 of accumulation microbatches. Packing and SP/CP handling remain model-owned.

README.md

@@ -79,6 +79,7 @@ Our guiding principles when building VeOmni are:
 | [GPT-OSS](https://huggingface.co/openai/gpt-oss-120b)    | 120B                          | [gpt_oss_120b_lora_ep4.yaml](configs/text/gpt_oss_120b_lora_ep4.yaml) |
 | [Qwen2-3 Omni](https://huggingface.co/Qwen)              | 7B/30BA3B                     | [qwen25_omni.yaml](configs/multimodal/qwen25_omni/qwen25_omni.yaml)   |
 | [Wan](https://huggingface.co/Wan-AI)                     | Wan2.1-I2V-14B-480P           | [wan_sft.yaml](configs/dit/wan_sft.yaml)                              |
+| [MiniMax H3](docs/examples/minimax_h3.md)                | Checkpoint-dependent          | [FL2VA](configs/dit/minimax_h3_fl2va_offline.yaml); [Ref2VA](configs/dit/minimax_h3_ref2va_cfg_offline.yaml) |
 | [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3)     | Checkpoint-dependent          | [ltx2_av_lora.yaml](configs/dit/ltx2_av_lora.yaml); [training guide](docs/examples/ltx-2.3.md) |
 
 Support new models to VeOmni see [Support New Models](https://veomni.readthedocs.io/en/latest/usage/support_new_models/guide_and_checklist.html)

configs/dit/minimax_h3_ref2va_cfg_offline.yaml

@@ -0,0 +1,88 @@
+# Prepared visual Ref2VA + CFG-calibrated FM LoRA training.
+# See docs/examples/minimax_h3.md for the decoded cache contract.
+# Replace the model/cache paths before running:
+# bash train.sh tasks/train_dit.py configs/dit/minimax_h3_ref2va_cfg_offline.yaml
+model:
+  model_path: /path/to/MiniMax-H3/Ref2VA/transformer
+  config_path: /path/to/MiniMax-H3/Ref2VA/transformer/config.json
+  condition_model_path: /path/to/MiniMax-H3/Ref2VA
+  condition_model_cfg:
+    skip_encoder_load: true
+    sigma_shift_video: 12.0
+    sigma_shift_audio: 3.0
+    training_cfg_scale: 4.0
+    # Same scale for video/audio; sigma is an opt-in per-modality heuristic.
+    training_cfg_schedule: constant
+    # Relative conditional curvature: 1 / scale**power; 0 disables attenuation.
+    training_cfg_curvature_power: 2.0
+    cfg_unconditional_mode: per_sample
+    # Alternatively: shared_empty plus a local cfg_unconditional_path.
+
+  lora_config:
+    rank: 64
+    alpha: 64
+    lora_dropout: 0.0
+    # Restrict to main DiT blocks, not the Qwen token refiner.
+    lora_modules: 'dit\.blocks\.\d+\.(attn\.(qkv_proj|out_proj)|mlp\.(fc1|fc2))'
+
+  ops_implementation:
+    attn_implementation: eager
+    rotary_pos_emb_implementation: eager
+    rms_norm_implementation: eager
+    swiglu_mlp_implementation: eager
+    cross_entropy_loss_implementation: eager
+    moe_implementation: eager
+    load_balancing_loss_implementation: eager
+
+  optimizer:
+    type: adamw
+    lr: 2.0e-5
+    lr_decay_style: linear
+    lr_min: 2.0e-5
+    lr_warmup_ratio: 0.0
+    weight_decay: 0.01
+    max_grad_norm: 1.0
+
+  accelerator:
+    init_device: meta
+    ulysses_size: 1
+    gradient_checkpointing:
+      enable: true
+    fsdp_config:
+      fsdp_mode: fsdp2
+      mixed_precision:
+        enable: true
+        param_dtype: bfloat16
+        reduce_dtype: float32
+        cast_forward_inputs: false
+
+data:
+  train_path: /path/to/prepared_ref2va_cache
+  data_type: diffusion
+  source_name: minimax_h3
+  data_transform: dit_offline
+  datasets_type: iterable
+  dataset_repeat: true
+  log_sample: false
+  text_keys: text
+  dataloader:
+    type: native
+    num_workers: 8
+    drop_last: true
+
+train:
+  training_task: offline_training
+  global_batch_size: 8
+  micro_batch_size: 1
+  dyn_bsz: false
+  seed: 42
+  num_train_epochs: 1
+  max_steps: 100
+  checkpoint:
+    output_dir: output/minimax_h3_ref2va_cfg
+    save_steps: 25
+    save_epochs: 0
+    save_hf_weights: false
+  wandb:
+    enable: false
+    project: minimax_h3_ref2va_cfg

docs/examples/minimax_h3.md

@@ -379,3 +379,212 @@ video, audio = pipe(
 **Important**:
 
 - `num_frames` must satisfy `(N-5) % 17 == 0`, otherwise the Video VAE raises an error
+
+## Prepared Ref2VA with CFG-aware LoRA
+
+Use [minimax_h3_ref2va_cfg_offline.yaml](../../configs/dit/minimax_h3_ref2va_cfg_offline.yaml)
+with a compatible Ref2VA checkpoint and prepared visual-reference caches.
+It uses the existing `DiTTrainer`, native LoRA, FSDP2 and model-owned packing.
+It does **not** add raw-reference encoding, audio-reference support or a
+Training Adapter downloader. A training/de-distillation adapter is separate from
+the trainable LoRA; this recipe does not load one implicitly.
+
+```text
+Prepared video/audio latents + paired conditions
+                         |
+            Shared timestep and noise
+                /                 \
+Negative layout                   Positive layout
+(same targets and visual anchors)       |
+        |                               |
+Current DiT (no gradient)       Current DiT (gradient enabled)
+                \                 /
+                CFG-aware FM objective
+                         |
+          Existing trainer and optimizer
+```
+
+### Cache contract
+
+The existing `dit_offline` transform deserializes **every column** using Python
+pickle. Only load trusted caches. The following table describes the decoded
+objects, not a new byte-storage schema or a direct Magnus-table reader.
+
+| Field | Decoded value | Purpose |
+| --- | --- | --- |
+| `input_latents` | Float tensor `[1,24,T,H,W]` | Clean video target |
+| `audio_input_latents` | Float tensor `[C,32,Ta]` | Clean audio target |
+| `prompt_embeds` | Float tensor `[L,D]` | Conditional Qwen representation |
+| `packed` | Dictionary from `build_packed_ref2va` | Conditional positions and segment bounds |
+| `ref_visual_anchor` | Float tensor `[cond_rows,96]` | Prepared reference image/video rows |
+| `has_audio` | Boolean | Set false for placeholder audio; its loss is zero |
+| `unconditional_prompt_embeds` | Float tensor `[Lu,D]` | Required in `per_sample` mode |
+| `unconditional_packed` | Dictionary from the same layout builder | Required in `per_sample` mode; uses `Lu` and its own text token tags |
+
+`per_sample` supports an independently encoded no-text condition retaining
+reference visual tokens in the Qwen representation. It cannot derive those
+embeddings by truncating or zeroing the conditional embedding. Both packed
+layouts must have identical target/reference geometry, reference order and
+compact `[text | cond | audio | video]` structure. Text lengths may differ.
+Use the same source references and compatible encoder checkpoint for both caches;
+structural checks cannot prove that two embeddings have the same semantic origin.
+
+The negative prompt is exactly `" "` (**one space**, not `""`), matching the
+native inference pipeline's default `negative_prompt`. For `per_sample`, use
+the same FL2VA/Ref2VA presentation, encoder checkpoint, resized keyframes or
+prepared reference blocks as the conditional cache, including reference order
+and video timestamps. For `shared_empty`, encode that space with the text-only
+T2VA presentation and no references. This aligns with native inference's
+default string; it does not establish the negative condition used in the
+original model's distillation.
+
+### Generate unconditional caches
+
+The condition model exposes two offline helpers that reuse native inference's
+Qwen presentation and encoding code. Run them while preparing the dataset,
+using the already-loaded condition model in evaluation mode with its Qwen
+encoder on the desired device. No DiT or VAE forward is needed for these calls.
+The helpers return CPU tensors; the original sample and latent anchors are
+preserved.
+
+```python
+from safetensors.torch import save_file
+
+# `condition_model` is the encoder-loaded MiniMaxH3ConditionModel used to
+# produce the conditional caches (skip_encoder_load=False).
+condition_model.eval()
+
+# Shared text-only cache. Native T2VA presentation; exactly one space.
+empty = condition_model.encode_unconditional(negative_prompt=" ")
+save_file(empty, "unconditional.safetensors")
+
+# Per-sample Ref2VA cache. `sample` is the decoded CPU sample from the table
+# above; refs are its original prepared reference blocks, in the same order.
+# An image block contains kind="image", prepared_image=<resized PIL image>.
+# A visual-video block contains kind="video", prepared_frames=<24 FPS frames>,
+# ref_audio_t=0. Use the same preparation as for the positive cache.
+sample_with_negative = condition_model.add_unconditional_cache(
+    sample, negative_prompt=" ", ref_blocks=refs
+)
+
+# For FL2VA, instead supply the original resized keyframe images:
+fl2va_with_negative = condition_model.add_unconditional_cache(
+    fl2va_sample, negative_prompt=" ", keyframe_images=keyframe_images
+)
+# Save each augmented sample using the existing OfflineEmbeddingSaver.save().
+```
+
+`add_unconditional_cache` writes `unconditional_prompt_embeds` and
+`unconditional_packed`. It rebuilds the Qwen prefix length/tags and shifts the
+target/reference positions consistently. It cannot recover original references
+from a conditional embedding; keep the prepared source images/frames available
+during dataset creation. Reference audio is not supported by this recipe.
+
+For a shared pure-empty Qwen embedding, configure:
+
+```yaml
+model:
+  condition_model_cfg:
+    training_cfg_scale: 4.0
+    training_cfg_schedule: constant
+    training_cfg_curvature_power: 2.0
+    cfg_unconditional_mode: shared_empty
+    cfg_unconditional_path: /path/to/unconditional.safetensors
+```
+
+The local file must contain finite `prompt_embeds [Lu,D]` and integer
+`text_token_tags [Lu]` containing only text tags (`1`). It is loaded once per
+condition-model instance. The negative branch removes the conditional Qwen
+prefix, including its visual tokens, but **retains independent visual latent
+anchors**. It reuses the exact positive branch's noisy video/audio and reference
+rows; it never calls the noise sampler a second time. This also works with
+prepared FL2VA inputs and their existing keyframe anchors.
+
+### CFG-calibrated FM and loss settings
+
+Let `p` and `u` be conditional and detached unconditional FM velocities from the
+**same current model**, including its trainable LoRA, and `y = noise - clean`.
+The negative pass runs first under `no_grad`; it is not a frozen-base teacher.
+
+There is one objective, with a configurable curvature exponent `k`:
+
+```text
+p_calibrated = (p + (s - 1) * stopgrad(u)) / s
+loss = s ** (2 - k) * MSE(p_calibrated, y)
+```
+
+| `training_cfg_curvature_power` | Relative conditional curvature | At `s=4` |
+| --- | --- | --- |
+| `2.0` (default) | `1/s²` | `1/16` |
+| `1.0` | `1/s` | `1/4` |
+| `0.0` | `1` | `1` |
+
+Any finite value in `[0, 2]` is supported. CFG-calibrated FM is a descriptive
+name for the existing CFG-aware transformation, not a novel algorithm: `k=2`
+is the inverse-CFG formula and `k=0` is algebraically equivalent to
+`MSE(p, s*y - (s-1)*stopgrad(u))` (the target-guidance form).
+The optimum is unchanged for fixed `u` and `s`, but the gradient scale changes.
+With sigma-dependent scales this also reweights different samples/modalities;
+it cannot in general be replaced by a single learning-rate change. Loss values
+are not directly comparable across curvature settings, and this objective does
+not guarantee preservation of generation quality.
+
+This objective distills guidance into the conditional prediction. At inference,
+the default `cfg_scale=1` uses that prediction directly; applying additional
+inference-time CFG can compound the guidance and should be validated separately.
+
+`sigma = 1 - t` is the effective noise level from the existing scheduler. With
+`training_cfg_schedule: sigma`, video and audio use their respective sigmas:
+`s = 1 + (training_cfg_scale - 1) * sigma`. For a configured scale of `4`,
+`sigma=0.1` gives `s=1.3`, while `sigma=0.9` gives `s=3.7`. This concentrates
+stronger CFG calibration at high noise and leaves low-noise examples closer to
+ordinary FM, giving them more room to learn from new data. `constant` instead
+applies the configured scale at every noise level, equally to video and audio,
+and is the default in the example recipe.
+
+The existing shifts strongly favor high sigma. Enumerating all 1,000 scheduler
+indices at scale `4`, before model-dtype rounding, gives:
+
+| Modality / shift | Mean scale | Median scale | Scale > 3 | Scale < 2 |
+| --- | ---: | ---: | ---: | ---: |
+| Video / 12 | 3.53 | 3.77 | 85.8% | 3.9% |
+| Audio / 3 | 3.03 | 3.25 | 60.0% | 14.2% |
+
+Therefore sigma scheduling gives relatively little low-noise relief for video
+with the shipped shifts. Separate modality scales are an **opt-in training
+heuristic** based on their effective corruption levels, not a reconstruction
+of a known original distillation scale. Native inference applies a single CFG
+scale to both modalities. A quality ablation is needed to establish whether
+per-modality scheduling helps; the formula alone does not show this. Neither
+schedule changes the noise sampling distribution or guarantees quality.
+
+| Condition config key | Default | Meaning |
+| --- | --- | --- |
+| `training_cfg_scale` | `1.0` | Scale one disables the extra branch and preserves upstream RNG/loss behavior |
+| `training_cfg_schedule` | `constant` | `constant` uses the configured scale everywhere; `sigma` varies it with each modality's effective noise level |
+| `training_cfg_curvature_power` | `2.0` | Curvature exponent `k` in `[0,2]`; relative curvature `1/s**k` |
+| `cfg_unconditional_mode` | `per_sample` | `per_sample` or `shared_empty` |
+| `cfg_unconditional_path` | unset | Local shared-empty safetensors, required only when enabled in that mode |
+
+Noise sampling is unchanged from upstream: a uniform scheduler index is shared
+by video and audio, with their respective existing shifts. Existing scheduler
+weights are retained; no additional sigma-bin weights are applied.
+No mixed raw-FM objective, stochastic CFG gate, step warmup or drift anchor is
+silently enabled by this configuration.
+
+### Validation boundary
+
+For offline training, `DiTTrainer.preforward` calls the condition model's optional
+`prepare_condition_batch` hook before device transfer. H3 checks embedding
+finiteness and packed geometry on CPU there. Device-side processing only remaps
+the validated rows and reuses the positive branch's unique timesteps. A custom
+training loop must call this hook **before** moving the batch to the accelerator;
+direct CPU calls to `process_condition` also perform the validation. Serialized
+validation markers are not accepted as proof of a checked cache.
+
+CPU tiny-model tests cover the objective gradients, paired noise, reference
+geometry, cached embeddings, checkpoint recomputation and packed/serial loss
+equivalence. The accelerator sync test also covers both enabled CFG data modes;
+running that test requires a GPU. Existing H3 packing restrictions remain, including no multi-sample
+checkpoint offload. Accelerator SP/FSDP2, pretrained quality and throughput still
+require hardware validation before treating this recipe as production-tested.