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

fix(bf16): staging pipeline + upload_batch_gpu BF16 boundaries

Browse Source 
- upload_batch: exclude actions from bf16 staging, upload i32 separately
- upload_batch_gpu: replace broken f32 DtoD with bf16 DtoD + proper
  actions bf16→i32 conversion via host roundtrip (64 elements, negligible)
- Remove dead mixed-precision conversion paths (bf16_states_buf etc.)

Smoke test now reaches CUDA graph capture + forward + loss + backward.
Next blocker: dqn_grad_norm_kernel launch args (shared memory or
buffer size from bf16 conversion).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 15:28:43 +01:00
parent d35189f15e
commit 937aa45ba9
1 changed files with 74 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

151
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -1746,7 +1746,7 @@ impl GpuDqnTrainer {
// ── Allocate consolidated transfer buffers ─────────────────
// Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights
let upload_staging_len = b * config.state_dim * 2 + b * 4; // 2*B*SD + 4*B
let upload_staging_len = b * config.state_dim * 2 + b * 3; // 2*B*SD + 3*B (no actions — uploaded separately as i32)
let upload_staging_buf = alloc_f32(&stream, upload_staging_len, "upload_staging")?;
// Readback: total_loss(1) + grad_norm(1) + td_errors(B)
@@ -2969,7 +2969,7 @@ impl GpuDqnTrainer {
/// buffers via `memcpy_dtod_async`. This turns 6 PCIe round-trips into 1,
/// saving ~12-30us per batch on H100 PCIe Gen5.
///
/// Layout: [states(B*SD) | next_states(B*SD) | actions_as_f32(B) | rewards(B) | dones(B) | is_weights(B)]
/// Layout: BF16 staging [states | next_states | rewards | dones | is_weights], actions uploaded separately as i32.
fn upload_batch(
&mut self,
states: &[f32],
@@ -2981,61 +2981,50 @@ impl GpuDqnTrainer {
) -> Result<(), MLError> {
let b = self.config.batch_size;
let sd = self.config.state_dim;
let bf16_size = std::mem::size_of::<half::bf16>();
// ── Pack all data into pre-allocated host buffer (zero malloc) ──
// ── Pack f32 data (no actions) into staging buffer → convert to BF16 ──
self.upload_staging_host.clear();
self.upload_staging_host.extend_from_slice(states); // B * SD f32s
self.upload_staging_host.extend_from_slice(next_states); // B * SD f32s
// Reinterpret i32 actions as f32 bits (same 4 bytes, no conversion)
for &a in actions {
self.upload_staging_host.push(f32::from_bits(a as u32));
}
self.upload_staging_host.extend_from_slice(rewards); // B f32s
self.upload_staging_host.extend_from_slice(dones); // B f32s
self.upload_staging_host.extend_from_slice(is_weights); // B f32s
self.upload_staging_host.extend_from_slice(states); // B * SD
self.upload_staging_host.extend_from_slice(next_states); // B * SD
self.upload_staging_host.extend_from_slice(rewards); // B
self.upload_staging_host.extend_from_slice(dones); // B
self.upload_staging_host.extend_from_slice(is_weights); // B
// ── Single HtoD transfer ──────────────────────────────────
// Single HtoD: f32 host → bf16 GPU
super::htod_f32_to_bf16(&self.stream, &self.upload_staging_host, &mut self.upload_staging_buf)?;
// ── Scatter from staging to individual buffers via DtoD ───
// ── Scatter BF16 staging to individual buffers via DtoD ───
let staging_base = raw_device_ptr(&self.upload_staging_buf, &self.stream);
let f32_size = std::mem::size_of::<half::bf16>();
let mut byte_offset: u64 = 0;
// states: B * SD elements
let states_bytes = b * sd * f32_size;
let states_dst = raw_device_ptr(&self.states_buf, &self.stream);
dtod_copy(states_dst, staging_base + byte_offset, states_bytes, &self.stream, 0, "upload_scatter")?;
// states: B * SD bf16 elements
let states_bytes = b * sd * bf16_size;
dtod_copy(raw_device_ptr(&self.states_buf, &self.stream), staging_base + byte_offset, states_bytes, &self.stream, 0, "upload_scatter")?;
byte_offset += states_bytes as u64;
// next_states: B * SD elements
let next_states_bytes = b * sd * f32_size;
let next_states_dst = raw_device_ptr(&self.next_states_buf, &self.stream);
dtod_copy(next_states_dst, staging_base + byte_offset, next_states_bytes, &self.stream, 1, "upload_scatter")?;
// next_states: B * SD bf16 elements
let next_states_bytes = b * sd * bf16_size;
dtod_copy(raw_device_ptr(&self.next_states_buf, &self.stream), staging_base + byte_offset, next_states_bytes, &self.stream, 1, "upload_scatter")?;
byte_offset += next_states_bytes as u64;
// actions: B elements as i32 (4 bytes each, not bf16)
let actions_bytes = b * std::mem::size_of::<i32>();
let actions_dst = raw_device_ptr_i32(&self.actions_buf, &self.stream);
dtod_copy(actions_dst, staging_base + byte_offset, actions_bytes, &self.stream, 2, "upload_scatter")?;
byte_offset += actions_bytes as u64;
// rewards: B elements
let rewards_bytes = b * f32_size;
let rewards_dst = raw_device_ptr(&self.rewards_buf, &self.stream);
dtod_copy(rewards_dst, staging_base + byte_offset, rewards_bytes, &self.stream, 3, "upload_scatter")?;
// rewards: B bf16 elements
let rewards_bytes = b * bf16_size;
dtod_copy(raw_device_ptr(&self.rewards_buf, &self.stream), staging_base + byte_offset, rewards_bytes, &self.stream, 2, "upload_scatter")?;
byte_offset += rewards_bytes as u64;
// dones: B elements
let dones_bytes = b * f32_size;
let dones_dst = raw_device_ptr(&self.dones_buf, &self.stream);
dtod_copy(dones_dst, staging_base + byte_offset, dones_bytes, &self.stream, 4, "upload_scatter")?;
// dones: B bf16 elements
let dones_bytes = b * bf16_size;
dtod_copy(raw_device_ptr(&self.dones_buf, &self.stream), staging_base + byte_offset, dones_bytes, &self.stream, 3, "upload_scatter")?;
byte_offset += dones_bytes as u64;
// is_weights: B elements
let is_weights_bytes = b * f32_size;
let is_weights_dst = raw_device_ptr(&self.is_weights_buf, &self.stream);
dtod_copy(is_weights_dst, staging_base + byte_offset, is_weights_bytes, &self.stream, 5, "upload_scatter")?;
// is_weights: B bf16 elements
let is_weights_bytes = b * bf16_size;
dtod_copy(raw_device_ptr(&self.is_weights_buf, &self.stream), staging_base + byte_offset, is_weights_bytes, &self.stream, 4, "upload_scatter")?;
// ── Actions: separate i32 upload (not bf16) ──
self.stream.memcpy_htod(actions, &mut self.actions_buf)
.map_err(|e| MLError::ModelError(format!("actions HtoD: {e}")))?;
Ok(())
}
@@ -3059,45 +3048,53 @@ impl GpuDqnTrainer {
let b = self.config.batch_size;
let sd = self.config.state_dim;
// GpuBatch stores GpuTensor (CudaSlice<half::bf16>) for all fields.
// States may be BF16-typed GpuTensors whose underlying data is u16-reinterpreted f32.
// Check shape to determine if we have BF16 data that needs conversion.
let states_data = gpu_batch.states.data();
let next_states_data = gpu_batch.next_states.data();
// GpuBatch stores GpuTensor (CudaSlice<half::bf16>) for states/next_states/rewards/dones/weights.
// Actions are CudaSlice<u32> — separate from GpuTensor.
// All data is already BF16 on GPU — direct DtoD copy, no conversion.
let bf16_size = std::mem::size_of::<half::bf16>();
// If states are stored as BF16 (u16 packed in f32 storage with half the elements),
// use the bf16 path. Otherwise copy f32 directly.
let states_numel = gpu_batch.states.numel();
if states_numel == b * sd {
// F32 states -- direct DtoD copy to states_buf
dtod_from_slice_f32(states_data, &self.states_buf, b * sd, &self.stream, "states")?;
dtod_from_slice_f32(next_states_data, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
} else {
// BF16 states -- reinterpret f32 storage as u16 and convert
let states_u16: &CudaSlice<u16> = unsafe {
&*(states_data as *const CudaSlice<half::bf16> as *const CudaSlice<u16>)
};
let next_states_u16: &CudaSlice<u16> = unsafe {
&*(next_states_data as *const CudaSlice<half::bf16> as *const CudaSlice<u16>)
};
dtod_from_bf16(states_u16, &self.bf16_states_buf, b * sd, &self.stream, "states")?;
dtod_from_bf16(next_states_u16, &self.bf16_next_states_buf, b * sd, &self.stream, "next_states")?;
launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_states_buf, &self.states_buf, b * sd, &self.stream, "states")?;
launch_bf16_to_f32(&self.bf16_to_f32_kernel, &self.bf16_next_states_buf, &self.next_states_buf, b * sd, &self.stream, "next_states")?;
// States + next_states: bf16 DtoD
dtod_copy(
raw_device_ptr(&self.states_buf, &self.stream),
raw_device_ptr(gpu_batch.states.data(), &self.stream),
b * sd * bf16_size, &self.stream, 0, "states",
)?;
dtod_copy(
raw_device_ptr(&self.next_states_buf, &self.stream),
raw_device_ptr(gpu_batch.next_states.data(), &self.stream),
b * sd * bf16_size, &self.stream, 1, "next_states",
)?;
// Rewards, dones, IS-weights: bf16 DtoD
dtod_copy(
raw_device_ptr(&self.rewards_buf, &self.stream),
raw_device_ptr(gpu_batch.rewards.data(), &self.stream),
b * bf16_size, &self.stream, 2, "rewards",
)?;
dtod_copy(
raw_device_ptr(&self.dones_buf, &self.stream),
raw_device_ptr(gpu_batch.dones.data(), &self.stream),
b * bf16_size, &self.stream, 3, "dones",
)?;
dtod_copy(
raw_device_ptr(&self.is_weights_buf, &self.stream),
raw_device_ptr(gpu_batch.weights.data(), &self.stream),
b * bf16_size, &self.stream, 4, "is_weights",
)?;
// Actions: GpuTensor<bf16> → i32 buf.
// Actions are integers (0-80) stored as bf16 in GpuTensor.
// Download bf16 → convert to i32 → upload. Batch_size elements only (64), negligible.
{
let actions_bf16 = gpu_batch.actions.data();
let mut host_bf16 = vec![half::bf16::ZERO; b];
self.stream.memcpy_dtoh(actions_bf16, &mut host_bf16)
.map_err(|e| MLError::ModelError(format!("actions DtoH: {e}")))?;
let host_i32: Vec<i32> = host_bf16.iter().map(|v| v.to_f32() as i32).collect();
self.stream.memcpy_htod(&host_i32, &mut self.actions_buf)
.map_err(|e| MLError::ModelError(format!("actions HtoD: {e}")))?;
}
// Rewards, dones, IS-weights: already F32 -- direct DtoD
dtod_from_slice_f32(gpu_batch.rewards.data(), &self.rewards_buf, b, &self.stream, "rewards")?;
dtod_from_slice_f32(gpu_batch.dones.data(), &self.dones_buf, b, &self.stream, "dones")?;
dtod_from_slice_f32(gpu_batch.weights.data(), &self.is_weights_buf, b, &self.stream, "is_weights")?;
// Actions: GpuTensor f32 -> I32 buf. Reinterpret f32 as u32 then DtoD as i32.
let actions_f32 = gpu_batch.actions.data();
let actions_u32: &CudaSlice<u32> = unsafe {
&*(actions_f32 as *const CudaSlice<half::bf16> as *const CudaSlice<u32>)
};
dtod_from_u32_to_i32(actions_u32, &self.actions_buf, b, &self.stream, "actions")?;
Ok(())
}