From 6ce61deef03e0ab10920a3a6c8a00aa771cd1060 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 22:36:06 +0200 Subject: [PATCH] fix(rl): normalize aux loss by n_valid in step_batched_from_device MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Raw aux loss sum was 11M (unnormalized across B×K positions). Now divides by n_valid per direction (prof/size), matching the original step_batched normalization. aux: 11M → 116. Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/src/trainer/perception.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index b6dca160b..f6914812c 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -3941,7 +3941,13 @@ impl PerceptionTrainer { let sl = std::ptr::read_volatile(self.aux_loss_size_long_host_d.host_ptr); let ps = std::ptr::read_volatile(self.aux_loss_prof_short_host_d.host_ptr); let ss = std::ptr::read_volatile(self.aux_loss_size_short_host_d.host_ptr); - pl + sl + ps + ss + let n_pl = std::ptr::read_volatile(self.aux_valid_prof_long_host_d.host_ptr) as f32; + let n_sl = std::ptr::read_volatile(self.aux_valid_size_long_host_d.host_ptr) as f32; + let n_ps = std::ptr::read_volatile(self.aux_valid_prof_short_host_d.host_ptr) as f32; + let n_ss = std::ptr::read_volatile(self.aux_valid_size_short_host_d.host_ptr) as f32; + let prof = if n_pl + n_ps > 0.0 { (pl + ps) / (n_pl + n_ps) } else { 0.0 }; + let size = if n_sl + n_ss > 0.0 { (sl + ss) / (n_sl + n_ss) } else { 0.0 }; + prof + size }; Ok((loss, aux_loss))