diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index c6f69491..a883b82d 100644 --- a/docs/user-guide/cli-reference.md +++ b/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 σ². | diff --git a/docs/user-guide/customization.md b/docs/user-guide/customization.md index e872452e..1285acf7 100644 --- a/docs/user-guide/customization.md +++ b/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. diff --git a/miles/backends/fsdp_utils/loss_hub/dmd2.py b/miles/backends/fsdp_utils/loss_hub/dmd2.py index 001690a7..2746b5f4 100644 --- a/miles/backends/fsdp_utils/loss_hub/dmd2.py +++ b/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 diff --git a/miles/backends/fsdp_utils/metrics.py b/miles/backends/fsdp_utils/metrics.py index 32b41939..b51f17d8 100644 --- a/miles/backends/fsdp_utils/metrics.py +++ b/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, } diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2c787003..4ec0333b 100644 --- a/miles/utils/arguments.py +++ b/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 " diff --git a/miles/utils/trainer_config.py b/miles/utils/trainer_config.py index 69966bf8..0c14de64 100644 --- a/miles/utils/trainer_config.py +++ b/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 [ diff --git a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py index 32228c3a..6aed8b94 100644 --- a/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py +++ b/tests/fast/backends/fsdp_utils/test_loss_hub_dmd2.py @@ -1,4 +1,4 @@ -"""DMD2 critic-group hooks: the critic (fake score) and the frozen teacher share this group; the actor holds the student. +"""DMD2 hooks: the critic (fake score) and the frozen teacher share the critic group; the actor holds the student. rollout pair: x0 = the student's x0 prediction at its stop step i, denoising_env = the rollout cond │ @@ -15,6 +15,13 @@ score_last_step_samples, without grad after the critic's steps, per kept micro-batch: teacher, CFG --dmd2-teacher-guidance-scale over the rollout's neg cond ──► x0_real = x_s - sigma v_real ──► {sample_index: {"x0_real", "x0_fake"}}; the teacher is resharded so sleep moves its shards + │ ctx.external_data["critic"] on the paired actor rank + select_scored_pairs ──► the actor's train-input hook keeps only the scored pairs, i.e. the critic's last step + prepare_generator_batch ──► the student's unguided step i on the rollout input x_i + dmd2_loss_formula: + x0 = x_i - sigma_i v_student (with grad; the rollout's x0_i only feeds a metric) + g = nan_to_num((x0_fake - x0_real) / mean|x0 - x0_real|) per pair + loss = 0.5 mean((x0 - sg(x0 - g))^2) ──► d/dx0 = g / numel, the reference implementation's gradient What each test pins: test_critic_batch_is_the_sft_batch_on_the_student_x0 earlier steps: inputs and targets == SFT on (x0, pos cond) @@ -27,6 +34,13 @@ test_teacher_runs_cfg_and_the_critic_runs_unguided CFG combine for the teacher only test_teacher_cfg_needs_the_rollout_negative_cond rejected instead of a forward without text test_scoring_reshards_the_teacher teacher resharded after its no-grad forwards + test_actor_keeps_the_critic_last_step_pairs same shard ──► the actor trains on the scored 1/K only + test_scores_outside_the_shard_are_rejected a broken rank pairing fails instead of shrinking the batch + test_generator_batch_is_the_student_step_on_its_input x_i, t_i, sigma_i and the positive cond, no CFG + test_generator_needs_a_single_component multi-DiT families are rejected + test_surrogate_gradient_matches_the_reference same gradient as the reference's batch-mean MSE + test_critic_matching_the_teacher_gives_no_gradient x0_fake == x0_real ──► zero loss and gradient + test_gaussian_student_moves_toward_the_teacher analytic scores: sign and normalization move mu """ from tests.ci.ci_register import register_cpu_ci @@ -43,6 +57,8 @@ from miles.backends.fsdp_utils.configs.sd3 import SD3TrainPipelineConfig 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, FlowNoiseSchedule +from miles.backends.fsdp_utils.metrics import new_metric_buffer +from miles.utils.train_data_utils import build_microbatch_schedule from miles.utils.types import CondKwargs, DenoisingEnv FLOW_SHIFT = 3.0 @@ -308,3 +324,164 @@ def test_scoring_reshards_the_teacher(monkeypatch): dmd2.score_last_step_samples(ctx, _pairs(2)) assert resharded == [ctx.teacher_models["transformer"]] + + +def test_actor_keeps_the_critic_last_step_pairs(): + ctx = _ctx() + pairs = _pairs(9) + schedule = build_microbatch_schedule( + num_pairs_per_optim_step=3, num_optim_steps_per_rollout=CRITIC_STEPS, micro_batch_size=3 + ) + for optim_step_idx, [(pair_lo, pair_hi)] in enumerate(schedule): + _critic_step(ctx, pairs[pair_lo:pair_hi], optim_step_idx) + + actor_ctx = _ctx() + actor_ctx.external_data["critic"] = dmd2.score_last_step_samples(ctx, pairs) + + assert dmd2.select_scored_pairs(actor_ctx, pairs) == pairs[6:] + + +def test_scores_outside_the_shard_are_rejected(): + pairs = _pairs(4) + ctx = _ctx() + ctx.external_data["critic"] = {pair["sample_index"]: {} for pair in pairs[2:]} + with pytest.raises(ValueError, match="only 1 of them"): + dmd2.select_scored_pairs(ctx, pairs[:3]) + + +def _generator_pairs(num_pairs): + return [ + { + "latent": torch.randn(LATENT_SHAPE, generator=torch.Generator().manual_seed(n)).to(torch.bfloat16), + "x0": torch.zeros(LATENT_SHAPE), + "timestep": 750.0 - 250.0 * n, + "sigma": 0.75 - 0.25 * n, + "denoising_env": DenoisingEnv(pos_cond_kwargs=_cond(1.0 + n), neg_cond_kwargs=None), + "sample_index": 100 + n, + } + for n in range(num_pairs) + ] + + +def _metrics(): + return new_metric_buffer(None, torch.device("cpu"), ["transformer"]) + + +def _loss(ctx, pairs, x0): + """dmd2_loss_formula with the student predicting ``x0`` from each pair's rollout input.""" + prepared = dmd2.prepare_generator_batch(ctx, pairs) + sigmas = prepared.extras["sigmas"].view(-1, 1, 1, 1) + new_pred = {"visual": (prepared.latents["visual"] - x0) / sigmas} + metrics = _metrics() + return ( + dmd2.dmd2_loss_formula(ctx, pairs, prepared, new_pred=new_pred, old_pred=None, ref_pred=None, metrics=metrics), + metrics, + ) + + +def _scores(pairs, x0_real, x0_fake): + return { + pair["sample_index"]: {"x0_real": real, "x0_fake": fake} + for pair, real, fake in zip(pairs, x0_real, x0_fake, strict=True) + } + + +def test_generator_batch_is_the_student_step_on_its_input(): + ctx = _ctx() + pairs = _generator_pairs(2) + + prepared = dmd2.prepare_generator_batch(ctx, pairs) + + assert prepared.model is ctx.models["transformer"] + assert torch.equal(prepared.latents["visual"], torch.stack([pair["latent"] for pair in pairs]).float()) + assert prepared.timesteps["visual"].tolist() == [750.0, 500.0] + assert prepared.extras["sigmas"].tolist() == [0.75, 0.5] + assert prepared.pos_cond["encoder_hidden_states"][:, 0, 0].tolist() == [1.0, 2.0] + assert not prepared.use_cfg and prepared.neg_cond is None + + +def test_generator_needs_a_single_component(): + ctx = _ctx() + ctx.models["transformer_2"] = ctx.models["transformer"] + with pytest.raises(ValueError, match="single DiT component"): + dmd2.prepare_generator_batch(ctx, _generator_pairs(1)) + + +def test_surrogate_gradient_matches_the_reference(): + ctx = _ctx() + pairs = _generator_pairs(2) + x0 = torch.randn(2, *LATENT_SHAPE, requires_grad=True) + x0_real, x0_fake = torch.randn(2, *LATENT_SHAPE), torch.randn(2, *LATENT_SHAPE) + ctx.external_data = {"critic": _scores(pairs, x0_real, x0_fake)} + + loss_sum, _ = _loss(ctx, pairs, x0) + (loss_sum / len(pairs)).backward() + + # tianweiy/DMD2 compute_distribution_matching_loss, on the same x0 and predictions. + reference_x0 = x0.detach().clone().requires_grad_() + with torch.no_grad(): + p_real, p_fake = reference_x0 - x0_real, reference_x0 - x0_fake + grad = torch.nan_to_num((p_real - p_fake) / p_real.abs().mean(dim=[1, 2, 3], keepdim=True)) + torch.nn.functional.mse_loss(reference_x0, (reference_x0 - grad).detach(), reduction="mean").mul(0.5).backward() + + torch.testing.assert_close(x0.grad, reference_x0.grad) + + +def test_critic_matching_the_teacher_gives_no_gradient(): + ctx = _ctx() + pairs = _generator_pairs(2) + x0 = torch.randn(2, *LATENT_SHAPE, requires_grad=True) + x0_pred = torch.randn(2, *LATENT_SHAPE) + ctx.external_data = {"critic": _scores(pairs, x0_pred, x0_pred)} + + loss_sum, metrics = _loss(ctx, pairs, x0) + loss_sum.backward() + + assert float(loss_sum) == 0.0 + assert torch.count_nonzero(x0.grad) == 0 + assert metrics.reduce()["dmd2_grad_abs_mean"] == 0.0 + + +def test_gaussian_student_moves_toward_the_teacher(): + # Data N(mu, 1) per element, noised as x_s = a x + b eps: E[x | x_s] = mu + a (x_s - a mu) / (a^2 + b^2). + def posterior_mean(mu, x_s, a, b): + return mu + a * (x_s - a * mu) / (a**2 + b**2) + + ctx = _ctx() + pairs = _generator_pairs(2) + teacher_mu, student_mu = 1.0, torch.tensor(-1.0, requires_grad=True) + optimizer = torch.optim.SGD([student_mu], lr=0.5) + generator = torch.Generator().manual_seed(0) + gaps = [] + for _ in range(20): + optimizer.zero_grad() + x0 = student_mu + torch.randn(2, *LATENT_SHAPE, generator=generator) + a, b = 0.5, 0.5 + x_s = a * x0.detach() + b * torch.randn(x0.shape, generator=generator) + ctx.external_data = { + "critic": _scores( + pairs, + posterior_mean(teacher_mu, x_s, a, b), + posterior_mean(float(student_mu), x_s, a, b), + ) + } + loss_sum, _ = _loss(ctx, pairs, x0) + (loss_sum / len(pairs)).backward() + optimizer.step() + gaps.append(abs(teacher_mu - float(student_mu))) + + assert gaps == sorted(gaps, reverse=True) + assert gaps[-1] < 0.5 * gaps[0] + + +def test_trainer_and_rollout_x0_diff_is_recorded(): + ctx = _ctx() + pairs = _generator_pairs(2) + x0 = torch.stack([pair["x0"] for pair in pairs]) + 0.25 + ctx.external_data = {"critic": _scores(pairs, x0, x0)} + + _, metrics = _loss(ctx, pairs, x0) + + reduced = metrics.reduce() + assert reduced["x0_mean_abs_diff"] == pytest.approx(0.25) + assert reduced["x0_max_abs_diff"] == pytest.approx(0.25) diff --git a/tests/fast/utils/test_dmd2_arguments.py b/tests/fast/utils/test_dmd2_arguments.py index d493e3e9..f878644a 100644 --- a/tests/fast/utils/test_dmd2_arguments.py +++ b/tests/fast/utils/test_dmd2_arguments.py @@ -1,12 +1,13 @@ -"""--loss-type dmd2 validation: reject what the critic group and the one-step actor cannot run, never fill args. +"""--loss-type dmd2 defaults and validation; validation rejects what DMD2 cannot run and never fills args. - defaults ──► --critic-load = --teacher-load, --critic-lr = --lr, unless the user set them + defaults ──► converter / generator prepare / generator loss / train-input paths, + --critic-load = --teacher-load, --critic-lr = --lr, unless the user set them critic group ──► scores with --teacher-load (no adapters), learns on a visual --fsdp-flow-shift grid, scores in a 0 <= MIN < MAX <= 1 --dmd2-score-sigma-range, takes --dmd2-critic-steps-per-rollout steps over each DP shard other losses ──► --critic-load / --critic-lr have no consumer and are rejected teacher score ──► CFG --dmd2-teacher-guidance-scale >= 1 - actor ──► unguided rollouts, one step per rollout, generator hooks given until a built-in lands + actor ──► unguided rollouts, one step per rollout """ from tests.ci.ci_register import register_cpu_ci @@ -35,8 +36,6 @@ def _args(**overrides): diffusion_true_cfg_scale=None, use_lora=False, ref_mode="none", - custom_prepare_train_batch_path="my.prepare_generator", - custom_loss_function_path="my.generator_loss", num_steps_per_rollout=1, global_batch_size=40, actor_num_gpus_per_node=4, @@ -61,6 +60,7 @@ def test_dmd2_critic_defaults_to_the_teacher_and_the_actor_lr(user, expected): custom_expand_samples_to_train_pairs_path=None, custom_prepare_train_batch_path=None, custom_loss_function_path=None, + custom_train_input_function_path=None, ref_mode=None, diffusion_kl_beta=0.0, teacher_load="teacher", @@ -78,6 +78,37 @@ def test_critic_flags_without_dmd2_are_rejected(flag): validate_dmd2_args(Namespace(loss_type="nft", critic_load=None, critic_lr=None)) +@pytest.mark.parametrize("custom_loss", [None, "my.generator_loss"]) +def test_dmd2_fills_only_unset_hook_paths(custom_loss): + args = Namespace( + sglang_tp_size=1, + sglang_sp_degree=None, + sglang_enable_cfg_parallel=False, + rollout_num_gpus_per_engine=1, + loss_type="dmd2", + custom_expand_samples_to_train_pairs_path=None, + custom_prepare_train_batch_path=None, + custom_loss_function_path=custom_loss, + custom_train_input_function_path=None, + ref_mode=None, + diffusion_kl_beta=0.0, + teacher_load="teacher", + lr=2e-6, + critic_load=None, + critic_lr=None, + ) + set_default_diffusion_args(args) + assert args.custom_expand_samples_to_train_pairs_path.endswith( + "data_conversion_hub.dmd2.expand_samples_to_train_pairs" + ) + assert args.custom_prepare_train_batch_path.endswith("loss_hub.dmd2.prepare_generator_batch") + assert args.custom_loss_function_path == ( + custom_loss or "miles.backends.fsdp_utils.loss_hub.dmd2.dmd2_loss_formula" + ) + assert args.custom_train_input_function_path.endswith("loss_hub.dmd2.select_scored_pairs") + assert args.ref_mode == "none" + + def test_valid_dmd2_args_pass_unchanged(): args = _args() original = vars(args).copy() @@ -102,7 +133,6 @@ def test_valid_dmd2_args_pass_unchanged(): ({"use_lora": True}, "drop --use-lora"), ({"teacher_lora_adapter_paths": ["teacher-adapter"]}, "drop --use-lora and --teacher-lora-adapter-path"), ({"ref_mode": "ref"}, "uses no reference model"), - ({"custom_loss_function_path": None}, "no built-in generator loss yet"), ({"num_steps_per_rollout": 2}, "one actor step per rollout"), ({"dmd2_critic_steps_per_rollout": 0}, "must be >= 1"), # 40 samples over 2 DP shards = 20 per shard, which 3 critic steps cannot split. diff --git a/tests/fast/utils/test_trainer_config.py b/tests/fast/utils/test_trainer_config.py index 7b28a3d4..e8b8de87 100644 --- a/tests/fast/utils/test_trainer_config.py +++ b/tests/fast/utils/test_trainer_config.py @@ -47,6 +47,7 @@ def test_dmd2_critic_trains_first_from_its_own_checkpoint(load, save, critic_dir dmd2_critic_steps_per_rollout=5, custom_prepare_train_batch_path="my.prepare_generator", custom_loss_function_path="my.generator_loss", + custom_train_input_function_path="my.select_pairs", custom_train_output_function_path=None, ) original = vars(args).copy() @@ -60,6 +61,7 @@ def test_dmd2_critic_trains_first_from_its_own_checkpoint(load, save, critic_dir assert (critic.num_steps_per_rollout, critic.global_batch_size) == (5, 8) assert critic.custom_prepare_train_batch_path.endswith("loss_hub.dmd2.prepare_critic_batch") assert critic.custom_loss_function_path.endswith("loss_hub.dmd2.critic_loss_formula") + assert critic.custom_train_input_function_path is None assert critic.custom_train_output_function_path.endswith("loss_hub.dmd2.score_last_step_samples") assert vars(actor) == {**original, "teacher_load": None} assert vars(args) == original