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

fix(per-horizon-cfc): sync trunk bucket metadata from transition output

Browse Source 
URGENT correctness fix surfaced by Task 13 implementer.

Bug: training-time Phase 2 dispatch sites (cfc_step_per_branch_fwd at
~3139, h_mag_per_bucket for Controller D at ~3211-3212) read
self.trunk.bucket_*_d which are zero-initialized at trunk construction
and only populated via load_checkpoint. The Phase 1→2 transition
populates a SEPARATE BucketRoutingMetadata owned by the trainer
(self.bucket_routing_metadata) but never syncs it into the trunk fields.
With zero-init bucket_dim_k, the per-branch kernel's uniform predicate
`tid >= bucket_dim_k[branch]` early-returns ALL threads -> Phase 2
silently no-ops during training, Smoke 1 fails before producing signal.

Fix: at transition, DtoD-copy metadata buffers into trunk fields BEFORE
the `Some(metadata)` move (so `metadata.*` borrows remain valid):
- metadata.bucket_channel_offset_d -> trunk.bucket_channel_offset_d
- metadata.bucket_dim_k_d          -> trunk.bucket_dim_k_d
- metadata.bucket_id_per_channel_d -> trunk.bucket_id_per_channel_d
- metadata.heads_w_skip_offset_d   -> trunk.heads_w_skip_offset_d

Trunk fields remain the single source of truth for both training-time
dispatch AND checkpoint serialization. Total sync cost: 4 DtoD memcpys,
~256 bytes total at the one-shot transition (off the hot path).

Verification:
- `cargo check -p ml-alpha --all-targets`: clean
- `cargo test -p ml-alpha --lib`: 33 passed, 0 failed
- 3 read sites (perception.rs:3139, 3211-3212, 4725-4726) verified
  unchanged; all 4 new memcpy_dtod_async calls bracketed in the
  transition block (perception.rs:2224, 2234, 2244, 2254).
- Block-diagonal heads grad-mask init (step 7) continues to read
  `metadata.bucket_id_per_channel_d` via `as_ref()` after the move;
  ordering preserved per pearl_canary_input_freshness_launch_order.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-21 19:16:27 +02:00
parent ed8b53b474
commit bf4d0c899f
1 changed files with 67 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2196,12 +2196,76 @@ impl PerceptionTrainer {
}
}
// 5. Persist routing metadata in trainer state. Tasks
// 5. Sync trunk's bucket metadata fields from the trainer-owned
// BucketRoutingMetadata. Without this, training-time Phase 2
// dispatch sites that read `self.trunk.bucket_*_d`
// (cfc_step_per_branch_fwd at the per-branch dispatch and
// h_mag_per_bucket_kernel for Controller D) would see the
// zero-initialized trunk buffers — the per-branch kernel's
// uniform predicate `tid >= bucket_dim_k[branch]` would
// early-return ALL threads → silent Phase 2 no-op.
//
// Trunk fields are the single source of truth for BOTH
// training-time dispatch AND checkpoint serialization. Cost:
// 4 DtoD memcpys, ~256 bytes total at the one-shot
// transition (off the hot path). Sync MUST happen BEFORE
// the `Some(metadata)` move below (borrow of `metadata.*_d`
// requires `metadata` to still be live), and BEFORE the
// heads_w_skip_mask_init kernel at step 5.5 (which still
// reads `metadata.bucket_id_per_channel_d` directly through
// `self.bucket_routing_metadata.as_ref()` after the move).
// Per `pearl_canary_input_freshness_launch_order`: producer
// (sync) precedes every consumer.
unsafe {
let s = self.stream.cu_stream();
// bucket_channel_offset_d: [N_HORIZONS + 1] × u32
let (src, _gs) = metadata.bucket_channel_offset_d.device_ptr(&self.stream);
let (dst, _gd) = self.trunk.bucket_channel_offset_d.device_ptr_mut(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
dst,
src,
(N_HORIZONS + 1) * std::mem::size_of::<u32>(),
s,
)
.context("sync bucket_channel_offset to trunk")?;
// bucket_dim_k_d: [N_HORIZONS] × u32
let (src, _gs) = metadata.bucket_dim_k_d.device_ptr(&self.stream);
let (dst, _gd) = self.trunk.bucket_dim_k_d.device_ptr_mut(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
dst,
src,
N_HORIZONS * std::mem::size_of::<u32>(),
s,
)
.context("sync bucket_dim_k to trunk")?;
// bucket_id_per_channel_d: [HIDDEN_DIM] × u8
let (src, _gs) = metadata.bucket_id_per_channel_d.device_ptr(&self.stream);
let (dst, _gd) = self.trunk.bucket_id_per_channel_d.device_ptr_mut(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
dst,
src,
HIDDEN_DIM * std::mem::size_of::<u8>(),
s,
)
.context("sync bucket_id_per_channel to trunk")?;
// heads_w_skip_offset_d: [N_HORIZONS + 1] × u32
let (src, _gs) = metadata.heads_w_skip_offset_d.device_ptr(&self.stream);
let (dst, _gd) = self.trunk.heads_w_skip_offset_d.device_ptr_mut(&self.stream);
cudarc::driver::result::memcpy_dtod_async(
dst,
src,
(N_HORIZONS + 1) * std::mem::size_of::<u32>(),
s,
)
.context("sync heads_w_skip_offset to trunk")?;
}
// 6. Persist routing metadata in trainer state. Tasks
// 10–12 read its slices to drive Phase 2 dispatch +
// Controllers B / C / D.
self.bucket_routing_metadata = Some(metadata);
// 5.5. Block-diagonal heads grad-mask init (follow-up to
// 7. Block-diagonal heads grad-mask init (follow-up to
// Task 10 Option B). One-shot: build the routing
// mask AND zero-multiply off-bucket entries of
// `heads_w_skip_d` in place. From here on, the per-
@@ -2236,7 +2300,7 @@ impl PerceptionTrainer {
}
}
// 6. Latch trainer into Phase 2.
// 8. Latch trainer into Phase 2.
//
// NOTE: Task 9 leaves the per-step dispatch path
// unchanged (Phase 1 single-CfC, 640-float heads_w_skip).