[DMD2 9/10] feat(dmd2): add the MiniMax H3 4-step DMD2 recipe with eval media and DMD2 metrics

docs/models/h3/h3.md

@@ -61,6 +61,42 @@ python3 scripts/run_diffusion_grpo_h3_t2va_2gpu.py \
   --rollout-batch-size 4 --n-samples-per-prompt 16 --num-rollout 200
 ```
 
+### 4.3 DMD2 4-step distillation (2 nodes x 8 GPUs)
+
+Recipe: `scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py` — FastVideo's FastH3 Preview v1
+Dense / Data-Free recipe on `--loss-type dmd2`: full-parameter student and critic, frozen
+teacher, lr 2e-6, 4 student steps with an Euler rollout stopped at a random step, an x0-space
+critic, and 5 critic steps of 64 samples per student step. Prompts are not in the repo: the
+script downloads `vidprom_filtered_extended.txt` from the Hub (`gdhe17/Self-Forcing` at a pinned
+revision and sha256) into `--data-dir` and splits it into fixed files with pinned hashes.
+
+**Status:** [○ NV — Not verified](../../user-guide/recipe-verification.md#nv)
+
+```bash
+# after the multi-node bring-up, on the head node
+MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py
+
+# references on the same eval prompts and seeds: the teacher at 50 steps, FastH3 at 4 steps
+MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference teacher
+MILES_SCRIPT_EXTERNAL_RAY=1 python3 scripts/run_diffusion_dmd2_h3_t2va_16gpu_multinode.py --reference fasth3
+```
+
+What to watch:
+
+| Signal | Where | Reading |
+|---|---|---|
+| Student media | W&B `eval_media/`, `<output_dir>/<run>/eval_media/eval_<rollout>/` | Rollout 0 is the base model at 4 steps; later evals render the same 26 prompts and seeds. Compare with the two reference runs' `eval_0/` (same file names, `index.jsonl` lists prompt and seed) |
+| Critic | `train/critic-loss` | Rising to a plateau is the normal DMD2 shape; judge by the media |
+| Distribution matching | `train/dmd2_grad_abs_mean`, `dmd2_score_gap_abs_mean` (critic vs teacher x0), `dmd2_teacher_dist_abs_mean` (student vs teacher x0, g's normalizer), each also `_step{i}` | A step whose curves drift away from the others is where the student breaks |
+| Student output | `train/dmd2_x0_mean`, `train/dmd2_x0_std` | Latent statistics drifting from rollout 0 flag saturation |
+| Rollout vs train | `train/x0_mean_abs_diff` (and `_step{i}`), `x0_rel_max` | The trainer's recomputed x0 against the rollout's; should stay small |
+| Memory | `train/peak_memory_allocated_gb` (this process), `train/device_memory_used_gb` (whole GPU), both with a `critic-` twin | Headroom on each group's steps |
+
+Training rollouts carry no media: with `--rm-type none` the engine skips their VAE decode, so the
+student is judged on the eval renders. The engine returns audio only with an sglang-miles change that
+adds it to the rollout response; until then the saved videos are silent. Only the video stream is
+distilled, so listen for drift.
+
 ## 5. Recipe notes
 
 SFT: the encode pool (`miles/rollout/encoder_hub/h3/`) imports the engine's own recipes

docs/user-guide/cli-reference.md

@@ -182,9 +182,9 @@ See [Dtype Control](../advanced/dtype-control.md).
 | `--diffusion-num-steps` | int | `10` | |
 | `--diffusion-flow-shift` | float | – | Generation-schedule shift, sent to the engine. |
 | `--rollout-microgroup-size` | int | `1` | |
-| `--diffusion-fps` | float | – | Video only. |
+| `--diffusion-fps` | float | – | Video only. Also the fps logged and saved media play at (8 when unset); H3 drops it from the request, so set it to 24 there. |
 | `--diffusion-output-num-frames` | int | `1` | |
-| `--rollout-video-dtype` | enum | `keep` | `keep` / `uint8`: dtype of the decoded video in rollout responses. `uint8` quantizes engine-side by truncation (×255, then clamp), cutting the response ~4x. That matches the PickScore reward path; HPS, OCR and the OpenAI-API reward round instead, so keep `keep` for them and for any consumer that needs the float tensor. |
+| `--rollout-video-dtype` | enum | `keep` | `keep` / `uint8`: dtype of the decoded video in rollout and eval responses. `uint8` quantizes engine-side by truncation (×255, then clamp), cutting the response ~4x. That matches the PickScore reward path; HPS, OCR and the OpenAI-API reward round instead, so keep `keep` for them and for any consumer that needs the float tensor. |
 | `--diffusion-guidance-scale` | float | `4.0` | |
 | `--diffusion-guidance-scale-2` | float | – | Wan2.2 low-noise expert; **required** when training it. |
 | `--diffusion-true-cfg-scale` | float | – | |
@@ -296,7 +296,7 @@ See [Dtype Control](../advanced/dtype-control.md).
 
 | Flag | Type | Default | Notes |
 |---|---|---|---|
-| `--rm-type` | str | – | `pickscore` / `hps` / `hpsv3` / `ocr` / `openai_api` / `custom_api`. Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. |
+| `--rm-type` | str | – | `pickscore` / `hps` / `hpsv3` / `ocr` / `openai_api` / `custom_api`, or `none` (constant 0, for a loss that reads no reward such as `dmd2`; training rollouts then skip the engine's VAE decode, which needs an engine with `rollout_return_generated_output`, while eval still decodes). Overridable per sample via `metadata.rm_type`. Ignored when `--custom-rm-path` is set. |
 | `--custom-api-rm-config` | str | – | Required for `--rm-type custom_api`: a YAML file or inline `base64:<payload>` with `actor_class`, `actor_kwargs`, `max_concurrency`. A flat mapping configures the default OpenAI-compatible image actor. |
 | `--openai-api-rm-config` | str | – | Required for `--rm-type openai_api`; same format as `--custom-api-rm-config`. |
 | `--group-rm` | flag | off | Score a whole prompt group at once. |
@@ -408,6 +408,7 @@ Every one takes a dotted path.
 | `--train-only` | flag | off | No engines, no weight sync, no eval. (`--debug-train-only` is a legacy alias.) |
 | `--debug-skip-optimizer-step` | flag | off | No backward/step — weights never drift. Use it to measure pure forward divergence from the engine. |
 | `--save-debug-rollout-data` | path template | – | Fills `{rollout_id}`. |
+| `--save-eval-media-dir` | path | – | Keeps every eval sample's video (with its audio when the engine returns it) or image under `<dir>/eval_<rollout_id>/`, with an `index.jsonl` of prompts, seeds and rewards. Independent of `--wandb-log-num-images`. |
 | `--load-debug-rollout-data` | path template | – | Replays a rollout with **no engines started**. |
 | `--load-debug-rollout-data-subsample` | float | – | |
 | `--save-debug-train-data` | path template | – | Fills `{rollout_id}` and `{rank}`. |

docs/user-guide/recipe-verification.md

@@ -62,3 +62,4 @@ count as verification.
     Flow-GRPO + PickScore.
   - `run_diffusion_sft_wan22.py` — Wan2.2 4-GPU LoRA SFT.
   - `run_diffusion_sft_h3_ref2va.py` — MiniMax H3 8-GPU Ref2VA LoRA SFT.
+  - `run_diffusion_dmd2_h3_t2va_16gpu_multinode.py` — MiniMax H3 16-GPU t2va 4-step DMD2.

docs/user-guide/rewards.md

@@ -13,7 +13,7 @@ For `--custom-rm-path`, `--custom-reward-post-process-path`, and other
 
 | Stage | Flag | Role |
 |---|---|---|
-| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `hps`, `hpsv3`, `ocr`, `openai_api`, `custom_api`); ignored when `--custom-rm-path` is set |
+| Reward type | `--rm-type` | Selects built-in scorer (`pickscore`, `hps`, `hpsv3`, `ocr`, `openai_api`, `custom_api`), or `none` for a constant 0 when the loss reads no reward (DMD2), which also skips decoding training rollouts; ignored when `--custom-rm-path` is set |
 | API reward config | `--openai-api-rm-config`, `--custom-api-rm-config` | YAML file or inline `base64:<payload>` for the two API rewards |
 | Per-sample override | `metadata.rm_type` in JSONL | Overrides global `--rm-type` |
 | Custom reward / norm | see [Customization](customization.md) | `--custom-rm-path`, `--custom-reward-post-process-path` |

miles/backends/fsdp_utils/actor.py

@@ -461,6 +461,7 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
                     device,
                     self.models,
                     sigma_buckets=self.args.log_loss_sigma_bucket,
+                    num_rollout_steps=self.args.diffusion_num_steps,
                 )
 
                 for pair_lo, pair_hi in microbatch_ranges:
@@ -490,6 +491,12 @@ def _train_core(self, rollout_id: int, rollout_data, external_data: dict[str, An
                     self.scaler.step(self.optimizer)
                     self.scaler.update()
                     self.lr_scheduler.step()
+                    # Peak since the last reset, so a rollout's first step also covers scoring and wake-up before it.
+                    metrics.emit_max("peak_memory_allocated_gb", torch.cuda.max_memory_allocated() / 2**30)
+                    # Device-wide, so it also counts the other processes on this GPU (sleeping groups, the engine).
+                    free_bytes, total_bytes = torch.cuda.mem_get_info()
+                    metrics.emit_max("device_memory_used_gb", (total_bytes - free_bytes) / 2**30)
+                    torch.cuda.reset_peak_memory_stats()
                 else:
                     self.optimizer.zero_grad(set_to_none=True)
                 self.global_step += 1

miles/backends/fsdp_utils/loss_hub/dmd2.py

@@ -10,7 +10,7 @@
 from miles.backends.fsdp_utils.input_dtype_policy import apply_input_dtype_policy
 from miles.backends.fsdp_utils.loss_hub.sft import prepare_sft_batch, sft_loss_formula
 from miles.backends.fsdp_utils.loss_hub.types import DiffusionLossContext, PreparedBatch
-from miles.backends.fsdp_utils.metrics import record_rollout_train_abs_diff
+from miles.backends.fsdp_utils.metrics import record_rollout_train_abs_diff, step_metric_key
 from miles.utils.metric_buffer import MetricBuffer
 
 
@@ -230,7 +230,19 @@ def dmd2_loss_formula(
     with torch.no_grad():
         bsz = len(batch)
         metrics.emit_mean("loss", total=loss_sum, count=bsz)
-        metrics.emit_mean("dmd2_grad_abs_mean", total=grad.abs().mean(dim=reduce_dims).sum(), count=bsz)
         rollout_x0 = torch.stack([pair["x0"] for pair in batch]).to(device=x0.device, dtype=torch.float32)
         record_rollout_train_abs_diff(metrics, "x0", x0, rollout_x0)
+        per_pair = {
+            "dmd2_grad_abs_mean": grad.abs().mean(dim=reduce_dims),
+            "dmd2_score_gap_abs_mean": (x0_fake - x0_real).abs().mean(dim=reduce_dims),
+            "dmd2_teacher_dist_abs_mean": (x0 - x0_real).abs().mean(dim=reduce_dims),
+            "x0_mean_abs_diff": (x0 - rollout_x0).abs().mean(dim=reduce_dims),
+        }
+        for key in ("dmd2_grad_abs_mean", "dmd2_score_gap_abs_mean", "dmd2_teacher_dist_abs_mean"):
+            metrics.emit_mean(key, total=per_pair[key].sum(), count=bsz)
+        metrics.emit_mean("dmd2_x0_mean", total=x0.flatten(1).mean(dim=1).sum(), count=bsz)
+        metrics.emit_mean("dmd2_x0_std", total=x0.flatten(1).std(dim=1).sum(), count=bsz)
+        for index, pair in enumerate(batch):
+            for key, values in per_pair.items():
+                metrics.emit_mean(step_metric_key(key, pair["step_index"]), total=values[index], count=1)
     return loss_sum

miles/backends/fsdp_utils/metrics.py

@@ -37,11 +37,29 @@
     "nft_t_mean": MetricReduce.MEAN,
     "nft_num_timesteps": MetricReduce.MEAN,
     "dmd2_grad_abs_mean": MetricReduce.MEAN,
+    "dmd2_score_gap_abs_mean": MetricReduce.MEAN,
+    "dmd2_teacher_dist_abs_mean": MetricReduce.MEAN,
+    "dmd2_x0_mean": MetricReduce.MEAN,
+    "dmd2_x0_std": MetricReduce.MEAN,
     "x0_mean_abs_diff": MetricReduce.MEAN,
     "x0_max_abs_diff": MetricReduce.MAX,
     "x0_rel_max": MetricReduce.MAX,
+    "peak_memory_allocated_gb": MetricReduce.MAX,
+    "device_memory_used_gb": MetricReduce.MAX,
 }
 
+# Also logged per rollout denoising step, under step_metric_key; a step no pair trained on is not logged.
+STEP_SCHEMA = {
+    "dmd2_grad_abs_mean": MetricReduce.MEAN,
+    "dmd2_score_gap_abs_mean": MetricReduce.MEAN,
+    "dmd2_teacher_dist_abs_mean": MetricReduce.MEAN,
+    "x0_mean_abs_diff": MetricReduce.MEAN,
+}
+
+
+def step_metric_key(key: str, step_index: int) -> str:
+    return f"{key}_step{step_index}"
+
 
 def sigma_bucket_key(stream_name: str, bucket: int, num_buckets: int) -> str:
     # Keep integer edges as floats (0.0, 1.0) so key names stay stable across bucket counts.
@@ -53,10 +71,13 @@ def edge(value: float) -> str:
 
 
 def new_metric_buffer(
-    group, device: torch.device, components: Collection[str], sigma_buckets: int = 0
+    group, device: torch.device, components: Collection[str], sigma_buckets: int = 0, num_rollout_steps: int = 0
 ) -> MetricBuffer:
     """Buffer for one optimizer step, reducing over the DP group."""
     schema = dict(SCHEMA)
+    for key, reduce in STEP_SCHEMA.items():
+        for step_index in range(num_rollout_steps):
+            schema[step_metric_key(key, step_index)] = reduce
     # SFT loss by sigma bucket: velocity-MSE magnitude varies strongly with the
     # corruption level, so only bucketed curves are comparable across steps.
     for stream_name in ("visual", "audio"):