[DMD2 7/9] feat(dmd2): train the student with the distribution matching gradient from the critic group's scores

docs/user-guide/cli-reference.md

@@ -284,7 +284,7 @@ See [Dtype Control](../advanced/dtype-control.md).
 | `--diffusion-nft-timestep-strategy-path` | str | `…nft.drop_random_steps` | `(args, num_steps, generator) -> list[int]` returning the denoising steps each sample trains on, in training order; `generator` is seeded per sample. Built in: `drop_random_steps` (random order), `drop_final_steps` (schedule order, noisiest first). The final clean output (σ = 0) is never a training step. |
 | `--diffusion-nft-num-dropped-timesteps` | int | `0` | Steps the built-in strategies drop: random ones per sample for `drop_random_steps`, the final (lowest-noise) ones for `drop_final_steps`. |
 | `--no-diffusion-nft-adaptive-weight` | flag | off | |
-| `--dmd2-critic-steps-per-rollout` | int | `5` | Critic optimizer steps per rollout under `dmd2`, before the actor's single step. Each DP shard must split into this many steps. The last step's samples are scored by the critic before that step's update. |
+| `--dmd2-critic-steps-per-rollout` | int | `5` | Critic optimizer steps per rollout under `dmd2`, before the actor's single step. Each DP shard must split into this many steps. The last step's samples are scored by the critic before that step's update, and the actor trains on them only. |
 | `--dmd2-teacher-guidance-scale` | float | – | Required by `dmd2`: the teacher's CFG scale when scoring the student's samples; `1` runs it unguided. |
 | `--dmd2-score-sigma-range` | float float | `0.02 0.98` | `MIN MAX` of the unshifted sigma at which `dmd2` scores the student's samples: the last critic step draws grid points in this range and reads them through the visual `--fsdp-flow-shift`. The default is the reference DMD2 range. |
 | `--dmd2-critic-loss-space` | str | `velocity` | `velocity` or `x0`: the space of the critic's flow-matching loss. `x0` regresses x0 = x_σ − σ·v, which weights the velocity MSE by σ². |

docs/user-guide/customization.md

@@ -234,7 +234,7 @@ def buffer_filter(
 | `nft` | `data_conversion_hub.nft.expand_samples_to_train_pairs` | `loss_hub.nft.prepare_nft_batch` | `nft_loss_formula` |
 | `sft_loss` | none; convert `sft_rollout.convert_samples_to_train_data` | `loss_hub.sft.prepare_sft_batch` | `sft_loss_formula` |
 | `custom_loss` | your expand (or convert) | your prepare | your formula |
-| `dmd2` | `data_conversion_hub.dmd2.expand_samples_to_train_pairs` | your prepare | your formula |
+| `dmd2` | `data_conversion_hub.dmd2.expand_samples_to_train_pairs` | `loss_hub.dmd2.prepare_generator_batch` | `dmd2_loss_formula` |
 
 `sft_loss` auto-fills nothing: it requires the convert, rollout-log, prepare and loss
 paths, plus `--train-only` and a rollout function that runs no engines (e.g.
@@ -244,6 +244,9 @@ before the actor on each rollout: it starts from `--critic-load` (default `--tea
 learns the student's samples with the SFT loss at `--critic-lr` (default `--lr`). On its last step
 it keeps its own x0 predictions from that step's forward, before the update, and after training hands
 them with the teacher's x0 at the same noisy inputs to the actor as `ctx.external_data["critic"]`.
+The actor trains on those samples only (its train-input hook defaults to
+`loss_hub.dmd2.select_scored_pairs`): it recomputes the student's x0 at the rollout's step and moves
+it along their difference.
 
 DiT forward always stays in the FSDP actor; the loss hook only computes the
 objective.

miles/backends/fsdp_utils/loss_hub/dmd2.py

@@ -1,4 +1,4 @@
-"""DMD2 critic-group hooks: the critic (fake score) learns the student's samples and scores those of its last step."""
+"""DMD2 hooks: the critic group learns and scores the student's samples; the actor follows the scores."""
 
 from __future__ import annotations
 
@@ -10,6 +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.utils.metric_buffer import MetricBuffer
 
 
@@ -153,3 +154,83 @@ def _predict_velocity(
             true_cfg_scale=None,
         )
     return pred["visual"].float()
+
+
+def select_scored_pairs(ctx: DiffusionLossContext, train_pairs: list[dict]) -> list[dict]:
+    """The pairs of the critic's last step, which it scored; the actor trains on these only."""
+    scores = ctx.external_data["critic"]
+    scored = [pair for pair in train_pairs if pair["sample_index"] in scores]
+    if len(scored) != len(scores):
+        raise ValueError(
+            f"the critic scored {len(scores)} samples, but only {len(scored)} of them are in this rank's train pairs"
+        )
+    return scored
+
+
+def prepare_generator_batch(
+    ctx: DiffusionLossContext,
+    batch: list[dict],
+    *,
+    pad_to_len: int | None = None,
+) -> PreparedBatch:
+    """The student's step ``i`` on its own rollout input ``x_i``, unguided."""
+    if len(ctx.models) != 1:
+        raise ValueError("DMD2 currently supports a single DiT component")
+    device = ctx.device
+    config = ctx.train_pipeline_config
+    component_name, model = next(iter(ctx.models.items()))
+    latents = torch.stack([pair["latent"] for pair in batch]).to(device=device, dtype=torch.float32)
+    timesteps = torch.tensor([pair["timestep"] for pair in batch], device=device, dtype=torch.float32)
+    sigmas = torch.tensor([pair["sigma"] for pair in batch], device=device, dtype=torch.float32)
+    pos_list = [config.prepare_cond_kwargs(pair["denoising_env"].pos_cond_kwargs, device) for pair in batch]
+    # Rollout trajectories carry only the visual stream.
+    return PreparedBatch(
+        latents={"visual": latents},
+        timesteps={"visual": timesteps},
+        timesteps_for_model={"visual": config.process_timestep_as_input(timesteps)},
+        model=model,
+        component_name=component_name,
+        guidance_scale=1.0,
+        use_cfg=False,
+        cfg_batching=False,
+        true_cfg_scale=None,
+        pos_cond=config.collate_cond_for_sample_batch(pos_list, device, pad_to_len=pad_to_len),
+        neg_cond=None,
+        joint_cond=None,
+        advantage=torch.ones(len(batch), device=device, dtype=torch.float32),
+        extras={"sigmas": sigmas},
+    )
+
+
+def dmd2_loss_formula(
+    ctx: DiffusionLossContext,
+    batch: list[dict],
+    prepared: PreparedBatch,
+    *,
+    new_pred: dict[str, torch.Tensor],
+    old_pred: dict[str, torch.Tensor] | None,
+    ref_pred: dict[str, torch.Tensor] | None,
+    metrics: MetricBuffer,
+    write_old_log_prob: bool = False,
+    old_log_prob_from_new: bool = False,
+) -> torch.Tensor:
+    """Distribution matching: x0's gradient is ``g / numel``, ``g = (x0_fake - x0_real) / mean|x0 - x0_real|``."""
+    latents = prepared.latents["visual"]
+    sigmas = prepared.extras["sigmas"].view(len(batch), *([1] * (latents.ndim - 1)))
+    x0 = latents - sigmas * new_pred["visual"].float()
+    scores = ctx.external_data["critic"]
+    x0_real = torch.stack([scores[pair["sample_index"]]["x0_real"] for pair in batch]).to(x0.device)
+    x0_fake = torch.stack([scores[pair["sample_index"]]["x0_fake"] for pair in batch]).to(x0.device)
+    reduce_dims = tuple(range(1, x0.ndim))
+    with torch.no_grad():
+        grad = torch.nan_to_num((x0_fake - x0_real) / (x0 - x0_real).abs().mean(dim=reduce_dims, keepdim=True))
+    per_pair = 0.5 * (x0 - (x0 - grad).detach()).square().mean(dim=reduce_dims)
+    loss_sum = per_pair.sum()
+
+    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)
+    return loss_sum

miles/backends/fsdp_utils/metrics.py

@@ -33,6 +33,10 @@
     "nft_adv_mean": MetricReduce.MEAN,
     "nft_t_mean": MetricReduce.MEAN,
     "nft_num_timesteps": MetricReduce.MEAN,
+    "dmd2_grad_abs_mean": MetricReduce.MEAN,
+    "x0_mean_abs_diff": MetricReduce.MEAN,
+    "x0_max_abs_diff": MetricReduce.MAX,
+    "x0_rel_max": MetricReduce.MAX,
 }
 
 

miles/utils/arguments.py

@@ -900,7 +900,8 @@ def add_algo_arguments(parser):
                 help=(
                     "Train objective shortcut. Diffusion default policy_loss = Flow-GRPO prepare+formula. "
                     "nft auto-fills prepare / loss-formula paths and selects the NFT train-data converter. "
-                    "dmd2 selects the DMD2 train-data converter and trains a critic group next to the actor. "
+                    "dmd2 auto-fills the DMD2 converter / prepare / loss-formula / train-input paths and trains a "
+                    "critic group before the actor. "
                     "DiT forward always stays in the FSDP actor."
                 ),
             )
@@ -976,7 +977,10 @@ def add_algo_arguments(parser):
                 "--dmd2-critic-steps-per-rollout",
                 type=int,
                 default=5,
-                help="DMD2 critic optimizer steps per rollout, taken before the actor's single step (TTUR).",
+                help=(
+                    "DMD2 critic optimizer steps per rollout, taken before the actor's single step (TTUR); "
+                    "the critic scores its last step's samples and the actor trains on those only."
+                ),
             )
             parser.add_argument(
                 "--dmd2-teacher-guidance-scale",
@@ -1737,6 +1741,12 @@ def set_default_diffusion_args(args) -> None:
             args.critic_load = args.teacher_load
         if args.critic_lr is None:
             args.critic_lr = args.lr
+        if args.custom_prepare_train_batch_path is None:
+            args.custom_prepare_train_batch_path = "miles.backends.fsdp_utils.loss_hub.dmd2.prepare_generator_batch"
+        if args.custom_loss_function_path is None:
+            args.custom_loss_function_path = "miles.backends.fsdp_utils.loss_hub.dmd2.dmd2_loss_formula"
+        if args.custom_train_input_function_path is None:
+            args.custom_train_input_function_path = "miles.backends.fsdp_utils.loss_hub.dmd2.select_scored_pairs"
 
     if args.ref_mode is None:
         args.ref_mode = "lora_base" if args.diffusion_kl_beta > 0 else "none"
@@ -1884,11 +1894,6 @@ def validate_dmd2_args(args) -> None:
         )
     if args.ref_mode != "none":
         raise ValueError("--loss-type dmd2 uses no reference model; drop --ref-mode and --diffusion-kl-beta")
-    if args.custom_prepare_train_batch_path is None or args.custom_loss_function_path is None:
-        raise ValueError(
-            "--loss-type dmd2 has no built-in generator loss yet; pass --custom-prepare-train-batch-path "
-            "and --custom-loss-function-path"
-        )
     if args.num_steps_per_rollout != 1:
         raise ValueError(
             f"--loss-type dmd2 takes one actor step per rollout, got --num-steps-per-rollout "

miles/utils/trainer_config.py

@@ -24,6 +24,7 @@ def compute_trainer_configs(args: Namespace) -> list[TrainerConfig]:
         "global_batch_size": args.global_batch_size // critic_steps,
         "custom_prepare_train_batch_path": "miles.backends.fsdp_utils.loss_hub.dmd2.prepare_critic_batch",
         "custom_loss_function_path": "miles.backends.fsdp_utils.loss_hub.dmd2.critic_loss_formula",
+        "custom_train_input_function_path": None,
         "custom_train_output_function_path": "miles.backends.fsdp_utils.loss_hub.dmd2.score_last_step_samples",
     }
     return [