gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(rl): normalize aux loss by n_valid in step_batched_from_device

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 22:36:06 +02:00
parent 0e70bf96fe
commit 6ce61deef0
1 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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))