From 729f6a29e151a4a4a340e824a74ac8627f8eefcb Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 22 Jul 2026 22:55:22 +0000 Subject: [PATCH 1/2] feat(diffusion): --diffusion-recompute-old-log-prob for impl-consistent PPO ratio Optional (default off). Before each optimizer window, recompute every pair's old log-prob with the trainer's own forward at pre-update weights, overwriting the rollout-engine values. The PPO ratio then compares log-probs from one implementation, so the trainer/rollout numerical gap no longer enters the ratio as per-pair multiplicative noise; at each window's first optim step the log-ratio is exactly 0 (flow_grpo convention). The pass mirrors the training microbatch schedule and legacy pad so batch shape cannot re-introduce the noise it removes. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 33 ++++++++++++++++++++++++++++-- miles/utils/arguments.py | 9 ++++++++ 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index e3c3ba46..f8b917de 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -381,6 +381,26 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: parallel_state=self.parallel_state, ) + # ------------- Recompute old log-probs (impl-consistent PPO ratio) ------------- + if self.args.diffusion_recompute_old_log_prob: + with timer("recompute_old_log_prob"), torch.no_grad(): + for microbatch_ranges in microbatch_schedule: + legacy_pad_to_len = self._maybe_legacy_window_pad_len(train_pairs, microbatch_ranges) + for pair_lo, pair_hi in microbatch_ranges: + self._forward_train_pair_batch( + train_pairs[pair_lo:pair_hi], + use_cfg=use_cfg, + guidance_scale=guidance_scale, + true_cfg_scale=true_cfg_scale, + clip_range=clip_range, + noise_level=noise_level, + num_train_timesteps=num_train_timesteps, + log_stats=defaultdict(list), + device=device, + pad_to_len=legacy_pad_to_len, + write_old_log_prob=True, + ) + # ------------- Forward / Backward ------------- with timer("actor_train"): for microbatch_ranges in microbatch_schedule: @@ -461,8 +481,12 @@ def _forward_train_pair_batch( device: torch.device, kl_beta: float = 0.0, pad_to_len: int | None = None, - ) -> torch.Tensor: - """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses.""" + write_old_log_prob: bool = False, + ) -> torch.Tensor | None: + """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses. + + With ``write_old_log_prob``, the computed log-prob overwrites each pair's + ``log_prob_old`` and no loss is returned (caller wraps in ``no_grad``).""" forward_dtype = self._forward_dtype train_pipeline_config = self.train_pipeline_config bsz = len(batch) @@ -589,6 +613,11 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: noise_level=noise_level, ) + if write_old_log_prob: + for pair, log_prob in zip(batch, log_prob_new_microbatch): + pair["log_prob_old"] = log_prob.cpu() + return None + log_prob_new = log_prob_new_microbatch # (bsz,) -- sde_step_with_logprob means over non-batch dims log_prob_old = log_prob_old_microbatch # (bsz,) ratio = torch.exp(log_prob_new - log_prob_old) # (bsz,) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c123cf51..4db94d5f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -130,6 +130,15 @@ def add_train_arguments(parser): default=5.0, help="Max absolute value for advantage clipping in diffusion training.", ) + parser.add_argument( + "--diffusion-recompute-old-log-prob", + action="store_true", + help=( + "Recompute old log-probs with the trainer forward (pre-update weights) " + "instead of using rollout-stored values, making the PPO ratio " + "implementation-consistent." + ), + ) parser.add_argument( "--diffusion-kl-beta", type=float, From 0e6dba4d5f78899c449cbfe2300c25e6b1b05e92 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 23 Jul 2026 00:02:52 +0000 Subject: [PATCH 2/2] perf(diffusion): skip redundant window-0 old-log-prob recompute --- miles/backends/fsdp_utils/actor.py | 18 +++++++++++++----- miles/utils/arguments.py | 3 ++- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f8b917de..e00837e5 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -384,7 +384,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: # ------------- Recompute old log-probs (impl-consistent PPO ratio) ------------- if self.args.diffusion_recompute_old_log_prob: with timer("recompute_old_log_prob"), torch.no_grad(): - for microbatch_ranges in microbatch_schedule: + # Skip window 0: its training forward runs on the same pre-update weights and doubles as the recompute. + for microbatch_ranges in microbatch_schedule[1:]: legacy_pad_to_len = self._maybe_legacy_window_pad_len(train_pairs, microbatch_ranges) for pair_lo, pair_hi in microbatch_ranges: self._forward_train_pair_batch( @@ -403,9 +404,11 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: # ------------- Forward / Backward ------------- with timer("actor_train"): - for microbatch_ranges in microbatch_schedule: + for optim_step_idx, microbatch_ranges in enumerate(microbatch_schedule): self.optimizer.zero_grad(set_to_none=True) + old_log_prob_from_new = self.args.diffusion_recompute_old_log_prob and optim_step_idx == 0 + num_local_pairs = sum(pair_hi - pair_lo for pair_lo, pair_hi in microbatch_ranges) # LEGACY 2D parity: pad cond to the whole-window width. TODO: remove with legacy 2D path. @@ -427,6 +430,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: device=device, kl_beta=kl_beta, pad_to_len=legacy_pad_to_len, + old_log_prob_from_new=old_log_prob_from_new, ) if not self.args.debug_skip_optimizer_step: # ShardedGradScaler keeps fp16 policy grads from underflowing @@ -482,11 +486,15 @@ def _forward_train_pair_batch( kl_beta: float = 0.0, pad_to_len: int | None = None, write_old_log_prob: bool = False, + old_log_prob_from_new: bool = False, ) -> torch.Tensor | None: """One DiT forward + PPO loss over ``len(batch)`` train pairs. Returns sum of per-pair losses. With ``write_old_log_prob``, the computed log-prob overwrites each pair's - ``log_prob_old`` and no loss is returned (caller wraps in ``no_grad``).""" + ``log_prob_old`` and no loss is returned (caller wraps in ``no_grad``). + + With ``old_log_prob_from_new``, the PPO ratio uses this forward's detached + log-prob as old (valid only under pre-update weights).""" forward_dtype = self._forward_dtype train_pipeline_config = self.train_pipeline_config bsz = len(batch) @@ -614,12 +622,12 @@ def _compute_noise_pred(disable_adapter: bool = False) -> torch.Tensor: ) if write_old_log_prob: - for pair, log_prob in zip(batch, log_prob_new_microbatch): + for pair, log_prob in zip(batch, log_prob_new_microbatch, strict=True): pair["log_prob_old"] = log_prob.cpu() return None log_prob_new = log_prob_new_microbatch # (bsz,) -- sde_step_with_logprob means over non-batch dims - log_prob_old = log_prob_old_microbatch # (bsz,) + log_prob_old = log_prob_new.detach() if old_log_prob_from_new else log_prob_old_microbatch # (bsz,) ratio = torch.exp(log_prob_new - log_prob_old) # (bsz,) unclipped = -advantage * ratio clipped = -advantage * torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4db94d5f..e38fbd8f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -136,7 +136,8 @@ def add_train_arguments(parser): help=( "Recompute old log-probs with the trainer forward (pre-update weights) " "instead of using rollout-stored values, making the PPO ratio " - "implementation-consistent." + "implementation-consistent. The first optimizer window skips the extra " + "pass and reuses its training forward's log-prob (ratio == 1)." ), ) parser.add_argument(