diff --git a/unirl/algorithms/diffusionopd.py b/unirl/algorithms/diffusionopd.py index 9d7e6f11c..10a6d50d4 100644 --- a/unirl/algorithms/diffusionopd.py +++ b/unirl/algorithms/diffusionopd.py @@ -190,8 +190,7 @@ def compute_loss_and_backward( ) teacher_means = gather_sde_field(segment.sde_means, segment.sde_indices, target_steps, field_name="sde_means") - # fp32 before squaring: bf16 loses the shrinking teacher-student delta. - student_f32 = student_means.float() + # Replay means are fp32 by contract; align the stored teacher anchor before squaring. teacher_f32 = teacher_means.to(device=student_means.device, dtype=torch.float32) sigma_t = _transition_sigma( @@ -202,7 +201,7 @@ def compute_loss_and_backward( device=student_means.device, add_coefficient=self.add_kl_coefficient, ) - kl_per_elem = _gaussian_kl_div(student_f32, teacher_f32, sigma_t) + kl_per_elem = _gaussian_kl_div(student_means, teacher_f32, sigma_t) kl_per_sample_step = kl_per_elem.mean(dim=tuple(range(2, kl_per_elem.ndim))) # [B, S'] loss = kl_per_sample_step.mean() diff --git a/unirl/models/bagel/diffusion.py b/unirl/models/bagel/diffusion.py index 466723904..b55d181e9 100644 --- a/unirl/models/bagel/diffusion.py +++ b/unirl/models/bagel/diffusion.py @@ -387,7 +387,7 @@ def diffuse( if log_prob is not None: sde_logp_list.append(log_prob.to(dtype=self.logprob_dtype)) if prev_mean is not None: - sde_means_list.append(prev_mean.detach().to(dtype=self.trajectory_dtype)) + sde_means_list.append(prev_mean.detach()) positions_collected = [p for p, _ in stored_pairs] latents_stacked = torch.stack([t for _, t in stored_pairs], dim=0).unsqueeze(0) @@ -469,7 +469,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=0).unsqueeze(0).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=0).unsqueeze(0).to(dtype=self.trajectory_dtype) + means_t = torch.stack(prev_sample_means, dim=0).unsqueeze(0) return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def build_forward_kwargs( diff --git a/unirl/models/boogu_image/diffusion.py b/unirl/models/boogu_image/diffusion.py index 42875d164..78b3893ce 100644 --- a/unirl/models/boogu_image/diffusion.py +++ b/unirl/models/boogu_image/diffusion.py @@ -436,7 +436,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/flux2_klein/diffusion.py b/unirl/models/flux2_klein/diffusion.py index 438fe89c5..b80622595 100644 --- a/unirl/models/flux2_klein/diffusion.py +++ b/unirl/models/flux2_klein/diffusion.py @@ -463,7 +463,7 @@ def replay( self.model.transformer.train() log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) @staticmethod diff --git a/unirl/models/hunyuan_image3/diffusion.py b/unirl/models/hunyuan_image3/diffusion.py index 4d134ebae..21610de4c 100644 --- a/unirl/models/hunyuan_image3/diffusion.py +++ b/unirl/models/hunyuan_image3/diffusion.py @@ -713,7 +713,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/hunyuan_video10/diffusion.py b/unirl/models/hunyuan_video10/diffusion.py index 659fe8cbf..c261dc1af 100644 --- a/unirl/models/hunyuan_video10/diffusion.py +++ b/unirl/models/hunyuan_video10/diffusion.py @@ -436,7 +436,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/hunyuan_video15/diffusion.py b/unirl/models/hunyuan_video15/diffusion.py index 139a7134f..1a5adda68 100644 --- a/unirl/models/hunyuan_video15/diffusion.py +++ b/unirl/models/hunyuan_video15/diffusion.py @@ -522,7 +522,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/ltx2/diffusion.py b/unirl/models/ltx2/diffusion.py index f80f313d4..5a5e9e64b 100644 --- a/unirl/models/ltx2/diffusion.py +++ b/unirl/models/ltx2/diffusion.py @@ -450,7 +450,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) diff --git a/unirl/models/qwen_image/diffusion.py b/unirl/models/qwen_image/diffusion.py index 9c7d89b8a..6fdf0042a 100644 --- a/unirl/models/qwen_image/diffusion.py +++ b/unirl/models/qwen_image/diffusion.py @@ -483,7 +483,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def _batched_step_kwargs(self, segment: LatentSegment, params: DiffusionSamplingParams) -> Dict[str, Any]: diff --git a/unirl/models/sd3/diffusion.py b/unirl/models/sd3/diffusion.py index c269f6ddf..ef7de65fc 100644 --- a/unirl/models/sd3/diffusion.py +++ b/unirl/models/sd3/diffusion.py @@ -397,7 +397,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) @staticmethod diff --git a/unirl/models/sensenova_u1/diffusion.py b/unirl/models/sensenova_u1/diffusion.py index eeb1d2e66..dbbced1a9 100644 --- a/unirl/models/sensenova_u1/diffusion.py +++ b/unirl/models/sensenova_u1/diffusion.py @@ -316,7 +316,7 @@ def _diffuse_one( if log_prob is not None: log_probs.append(log_prob.to(dtype=self.logprob_dtype)) if mean is not None: - means.append(mean.detach().to(dtype=self.trajectory_dtype)) + means.append(mean.detach()) return LatentSegment( latents=torch.stack([value for _, value in stored], dim=1), @@ -444,7 +444,7 @@ def replay( return ReplayResult( log_probs=torch.cat(batch_log_probs, dim=0).to(dtype=self.logprob_dtype), - prev_sample_means=torch.cat(batch_means, dim=0).to(dtype=self.trajectory_dtype), + prev_sample_means=torch.cat(batch_means, dim=0), ) def predict_noise_at_step( diff --git a/unirl/models/types/batched_replay.py b/unirl/models/types/batched_replay.py index be40d60c7..7fa6ba728 100644 --- a/unirl/models/types/batched_replay.py +++ b/unirl/models/types/batched_replay.py @@ -67,7 +67,7 @@ def _replay_batched_steps( means_t = None if prev_mean_all is not None: tail = prev_mean_all.shape[1:] - means_t = prev_mean_all.view(S, B, *tail).transpose(0, 1).contiguous().to(dtype=self.trajectory_dtype) + means_t = prev_mean_all.view(S, B, *tail).transpose(0, 1).contiguous() return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) diff --git a/unirl/models/types/replay_result.py b/unirl/models/types/replay_result.py index 5e4a02310..b5f69ca64 100644 --- a/unirl/models/types/replay_result.py +++ b/unirl/models/types/replay_result.py @@ -19,7 +19,8 @@ class ReplayResult: prev_sample_means: Optional[torch.Tensor] = None """The SDE transition's mean μ_θ at each replayed step. Shape ``[B, S', *latent_shape]`` for diffusion. Used by GRPO's KL penalty. - ``None`` when the stage doesn't produce it.""" + Kept at the SDE kernel's native fp32; see Gotchas in + ``unirl/sde/README.md``. ``None`` when the stage doesn't produce it.""" logits: Optional[torch.Tensor] = None """Per-step token logits at each replayed position. Shape diff --git a/unirl/models/wan21/diffusion.py b/unirl/models/wan21/diffusion.py index 6a1591f54..c7ec5ba28 100644 --- a/unirl/models/wan21/diffusion.py +++ b/unirl/models/wan21/diffusion.py @@ -403,7 +403,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/wan22/diffusion.py b/unirl/models/wan22/diffusion.py index 50cd08ff0..a4dd6c9c1 100644 --- a/unirl/models/wan22/diffusion.py +++ b/unirl/models/wan22/diffusion.py @@ -447,7 +447,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) def predict_noise_at_step( diff --git a/unirl/models/z_image/diffusion.py b/unirl/models/z_image/diffusion.py index 1b8c774ff..bf6c605ba 100644 --- a/unirl/models/z_image/diffusion.py +++ b/unirl/models/z_image/diffusion.py @@ -397,7 +397,7 @@ def replay( prev_sample_means.append(prev_mean) log_probs_t = torch.stack(log_probs, dim=1).to(dtype=self.logprob_dtype) - means_t = torch.stack(prev_sample_means, dim=1).to(dtype=self.trajectory_dtype) if prev_sample_means else None + means_t = torch.stack(prev_sample_means, dim=1) if prev_sample_means else None return ReplayResult(log_probs=log_probs_t, prev_sample_means=means_t) @staticmethod diff --git a/unirl/sde/README.md b/unirl/sde/README.md index 3ac0365d8..0f1278cf8 100644 --- a/unirl/sde/README.md +++ b/unirl/sde/README.md @@ -109,6 +109,13 @@ MixGRPO keeps `FlowSDEStrategy` and adds a `WindowScheduler` under - **The dtype round-trip in `_finalize_logp` is not a no-op cast.** It simulates trajectory *storage* precision so replay-time log-prob matches sampling-time log-prob. Delete it as dead code and the ratio drifts. Skipped for `eta<1e-7`. +- **`prev_sample_means` leave replay at the kernel's native fp32 — don't narrow + them to `trajectory_precision`.** Log-prob scores the fp32 mean; FlowDPPO's + Gaussian-KL mask (`(Δmean)²/(2σ²)` vs the 1e-5 threshold) scores whatever + replay returns, so a bf16 round-trip re-scores a *different* distribution: + ULP-scale rounding dwarfs the threshold, and sub-ULP policy deltas collapse to + KL=0. `log_probs` keeps its explicit `logprob_dtype` cast; only the mean must + not follow the trajectory dtype. - **`ensure_sample_sigmas` takes height/width/steps with no defaults, on purpose** — a silent `1024×1024` mis-derives μ for dynamic-shift models rendering at anything else (e.g. WAN T2V at 480×832).