feat(ml-alpha): Phase 1B-B — rollout collection loop + GAE behind FOXHUNT_USE_ROLLOUT flag
Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B Builds on commit5cd2f8703(1B-A foundation). Wires the RolloutBuffer into the per-step training loop via a new RolloutCollection helper, gated behind FOXHUNT_USE_ROLLOUT=1 env flag. Mode 0 (default unset) is bit-equal to HEAD; Mode 1 (flag=1) snapshots per-step data into the rollout buffer and invokes GAE every T_rollout steps. This phase is a STRUCTURAL validation, not an empirical pnl test. Phase 1B-C will replace the per-step training in Mode 1 with multi- epoch PPO over the rollout buffer. New components: * crates/ml-alpha/cuda/rollout_pack.cu (42 LOC): Single-thread-per-batch deterministic f32 -> u8 packer for dones. Used because RolloutBuffer.dones_d is u8 (per spec for memory) but the trainer's dones_d is f32 (per existing pipeline). * crates/ml-alpha/src/trainer/rollout_collection.rs (375 LOC): RolloutCollection helper with three methods: - new(ctx): loads rollout_pack cubin - snapshot_step(trainer, buf, t, b_size, t_max): DtoD scatter via cuMemcpy2DAsync_v2 for rewards/v_t/actions/log_pi/h_t at offset b*t_max+t in the [B × T] buffer; rollout_pack kernel for dones - copy_bootstrap_v(trainer, buf, b_size): DtoD V(s_T) from trainer.v_pred_tp1_d into buf.v_t_bootstrap_d at rollout end All transfers use mapped-pinned / DtoD only — no raw memcpy_dtoh on regular slices (`feedback_no_htod_htoh_only_mapped_pinned`). * crates/ml-alpha/examples/alpha_rl_train.rs (+137 LOC at lines 59-66, 410-466, 591-666): - FOXHUNT_USE_ROLLOUT env-flag parse at startup - Conditional RolloutBuffer + RolloutCollection allocation - Per-step snapshot_step() after step_with_lobsim_gpu - Every T_rollout steps: copy_bootstrap_v + compute_gae + stderr log `[rollout] step=N GAE complete (T=…, γ=…, λ=0.95): adv_mean=… adv_abs_mean=… returns_mean=…` Adaptations from the dispatch (documented inline): 1. T_rollout fallback to 32 when slot 758 is unbootstrapped. Slot 758 (RL_PPO_ROLLOUT_HORIZON_INDEX) was allocated in 1B-A but no kernel bootstraps it — `read_isv_host(758)` returns 0.0. Binary uses `.clamp(32, 1024)`, flooring to T=32 for the smoke. Phase 1B-C will add the trainer bootstrap entry (or a controller kernel) to make T_rollout=256 the production default. 2. Mode 1 still runs per-step training as a side effect; the rollout buffer is populated as a parallel snapshot. The "stop training in Mode 1" separation would require a ~1500 LOC refactor of step_with_lobsim_gpu (single-step/training-coupled per pearl_foxhunt_trainer_is_genuinely_single_step). The dispatch explicitly authorized this "skeleton with correct semantics" adaptation — the load-bearing 1B-B gate is "collection + GAE pipeline works structurally", which is verified. Phase 1B-C will replace the per-step training with multi-epoch PPO over the rollout buffer. Validation (all gates PASS): * cargo build --release --example alpha_rl_train -p ml-alpha: exit 0 (~60s) * cargo test rollout_buffer_invariants --release: 5/5 PASS (no regression) * Mode 0 (flag unset, default): ./scripts/determinism-check.sh --quick exit 0 - DETERMINISTIC: all checksums.* leaves match across all 200 rows - pipeline output bit-equal to HEAD5cd2f8703baseline * Mode 1 (FOXHUNT_USE_ROLLOUT=1): smoke completes in 2m 53s - 500 train + 100 eval steps, exit 0, completed_clean=true - 16 GAE windows logged (T=32 each) - adv_abs_mean range [4.4e-3, 1.92e-1] — always > 1e-3 sanity gate - eval_summary written: total_pnl_usd=-$101,487.70, n_trades=301, wr=0.346 * Cross-source consistency BOTH modes within $50 (Mode 0: $0.00 diff; Mode 1: $0.01 diff). Phase 0 contract preserved. * Pre-commit hook PASS (0 memcpy_dtoh raw violations, 0 atomicAdd). Next: Phase 1B-C will (a) bootstrap slot 758 to production T_rollout=256, (b) replace Mode 1's per-step training with multi-epoch PPO over the rollout buffer using the existing PPO surrogate + V regression kernels operating on minibatches of [B × T] flattened to [B*T]. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -52,6 +52,7 @@ const KERNELS: &[&str] = &[
|
||||
"ema_update_per_step", // RL Phase R3: generic per-step EMA producer (slot-parameterised) for continuous EMAs (kl_pi, entropy_observed, advantage_var_ratio, trade_duration)
|
||||
"compute_advantage_return", // RL Phase R3: element-wise A_t = r + γ(1-done)·V(s_{t+1}) − V(s_t), R_t = r + γ(1-done)·V(s_{t+1}); reads γ from ISV[400]
|
||||
"gae_backward_sweep", // Phase 1B-A (2026-06-02): GAE backward sweep over [B × T] rollout trajectories — A_t = δ_t + γλ·A_{t+1}·non_terminal, returns_t = A_t + V_t; deterministic single-thread-per-batch sequential sweep; replaces compute_advantage_return when wired in Phase 1B-B+
|
||||
"rollout_pack", // Phase 1B-B (2026-06-02): per-step f32→u8 dones packer; writes `dones_f32 [B]` into rollout buffer's `dones_u8_bt [B × T]` at offset `b * T + t_cursor`; single-thread-per-batch, deterministic
|
||||
"rl_action_kernel", // RL Phase R4: Thompson sampler over C51 atoms; one block per batch, N_ACTIONS threads; per-batch xorshift32 PRNG state; replaces host Thompson loop per feedback_cpu_is_read_only
|
||||
"argmax_expected_q", // RL Phase R4: argmax over expected Q per action; Bellman-target argmax (Double-DQN); pairs with rl_action_kernel per pearl_thompson_for_distributional_action_selection
|
||||
"log_pi_at_action", // RL Phase R4: per-batch log π(action_b) via log-softmax + lookup; PPO importance-ratio path
|
||||
|
||||
42
crates/ml-alpha/cuda/rollout_pack.cu
Normal file
42
crates/ml-alpha/cuda/rollout_pack.cu
Normal file
@@ -0,0 +1,42 @@
|
||||
// rollout_pack.cu — Per-step pack helpers for Phase 1B-B rollout collection.
|
||||
//
|
||||
// Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2
|
||||
// Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B
|
||||
//
|
||||
// The trainer's per-step `dones_d` is `f32` in {0.0, 1.0} (legacy contract
|
||||
// driven by `compute_advantage_return` and the K-loop sampled buffers).
|
||||
// The rollout buffer's `dones_d` is `u8` (compact storage — 4× smaller,
|
||||
// matches the `gae_backward_sweep` kernel's `uint8_t*` input). This kernel
|
||||
// packs a per-step `[B]` f32 done vector into the rollout buffer's
|
||||
// `[B × T]` u8 slice at offset `bt = b * T + t` for the current write
|
||||
// cursor `t`.
|
||||
//
|
||||
// Single-thread-per-batch kernel (one thread per env). Deterministic by
|
||||
// construction — no atomic ops, no reductions, no shared memory. The
|
||||
// fmaxf clamp + ≥0.5 thresholding makes the conversion robust to fp32
|
||||
// noise (the trainer never produces non-{0, 1} dones, but the pack is
|
||||
// defensive).
|
||||
//
|
||||
// Determinism contract: each thread writes one distinct element, so the
|
||||
// kernel is bit-equal across runs for identical inputs. Compatible with
|
||||
// `pearl_determinism_achieved`.
|
||||
//
|
||||
// Per `feedback_no_atomicadd`: no atomic ops.
|
||||
// Per `feedback_cpu_is_read_only`: pure device-side, no host roundtrip.
|
||||
// Per `feedback_no_nvrtc`: pre-compiled cubin via build.rs.
|
||||
|
||||
#include <stdint.h>
|
||||
|
||||
extern "C" __global__ void rollout_pack_dones_f32_to_u8(
|
||||
const float* __restrict__ dones_f32, // [B] — per-step trainer dones (f32 in {0, 1})
|
||||
uint8_t* __restrict__ dones_u8_bt, // [B × T] — rollout buffer dones (u8)
|
||||
const int B,
|
||||
const int T,
|
||||
const int t_cursor // current write cursor in [0, T)
|
||||
) {
|
||||
const int b = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (b >= B) return;
|
||||
const int idx = b * T + t_cursor;
|
||||
const float d = dones_f32[b];
|
||||
dones_u8_bt[idx] = (d >= 0.5f) ? (uint8_t)1 : (uint8_t)0;
|
||||
}
|
||||
@@ -58,14 +58,17 @@ use ml_alpha::heads::HORIZONS;
|
||||
// the per-N-step stderr ticker.
|
||||
use ml_alpha::rl::isv_slots::{
|
||||
RL_CONF_GATE_FIRED_COUNT_INDEX, RL_FRD_GATE_FIRED_COUNT_INDEX, RL_GAMMA_INDEX,
|
||||
RL_HEAT_CAP_FIRED_COUNT_INDEX, RL_PER_ALPHA_INDEX, RL_PPO_CLIP_INDEX, RL_PYRAMID_ADD_COUNT_INDEX,
|
||||
RL_REWARD_SCALE_INDEX,
|
||||
RL_HEAT_CAP_FIRED_COUNT_INDEX, RL_PER_ALPHA_INDEX, RL_PPO_CLIP_INDEX,
|
||||
RL_PPO_ROLLOUT_HORIZON_INDEX, RL_PYRAMID_ADD_COUNT_INDEX, RL_REWARD_SCALE_INDEX,
|
||||
};
|
||||
use ml_alpha::heads::HIDDEN_DIM as ROLLOUT_HIDDEN_DIM;
|
||||
use ml_alpha::trainer::diag_staging::DiagStaging;
|
||||
use ml_alpha::trainer::integrated::{
|
||||
read_slice_d_into, DiagInputs, IntegratedStepStats, IntegratedTrainer, IntegratedTrainerConfig,
|
||||
};
|
||||
use ml_alpha::trainer::perception::PerceptionTrainerConfig;
|
||||
use ml_alpha::trainer::rollout_buffer::RolloutBuffer;
|
||||
use ml_alpha::trainer::rollout_collection::RolloutCollection;
|
||||
use ml_backtesting::sim::LobSimCuda;
|
||||
use ml_core::device::MlDevice;
|
||||
use serde::Serialize;
|
||||
@@ -420,6 +423,70 @@ fn main() -> Result<()> {
|
||||
cli.seq_len, cli.n_backtests, cli.per_capacity
|
||||
);
|
||||
|
||||
// ── Phase 1B-B rollout collection gate (FOXHUNT_USE_ROLLOUT). ────
|
||||
//
|
||||
// Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2
|
||||
// Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B
|
||||
//
|
||||
// TEMPORARY VALIDATION FLAG — per `feedback_no_feature_flags`,
|
||||
// permanent feature flags are forbidden. This env toggle is the
|
||||
// same pattern as `FOXHUNT_DETERMINISTIC` (Phase 2.3 determinism
|
||||
// foundation). REMOVED in Phase 1B-E once rollout-buffer + GAE is
|
||||
// the default path.
|
||||
//
|
||||
// Gate semantics:
|
||||
// * `FOXHUNT_USE_ROLLOUT=0` (default, or unset): legacy path
|
||||
// unchanged. RolloutBuffer NOT allocated. Bit-equal to HEAD.
|
||||
// * `FOXHUNT_USE_ROLLOUT=1`: per-step snapshot into RolloutBuffer
|
||||
// after each `step_with_lobsim_gpu`. Every T_rollout steps run
|
||||
// GAE backward sweep + log advantage / returns stats to stderr.
|
||||
// Training pipeline runs UNCHANGED (1B-B is structural validation
|
||||
// of collection + GAE wiring, not training replacement — that's
|
||||
// 1B-C).
|
||||
let use_rollout: bool = std::env::var("FOXHUNT_USE_ROLLOUT")
|
||||
.ok()
|
||||
.and_then(|v| v.parse::<u8>().ok())
|
||||
.map(|n| n > 0)
|
||||
.unwrap_or(false);
|
||||
|
||||
let t_rollout_default: usize = {
|
||||
let v = trainer.read_isv_host(RL_PPO_ROLLOUT_HORIZON_INDEX) as usize;
|
||||
v.clamp(32, 1024)
|
||||
};
|
||||
|
||||
let (mut rollout_buf, rollout_collection): (Option<RolloutBuffer>, Option<RolloutCollection>) =
|
||||
if use_rollout {
|
||||
let ctx = dev
|
||||
.cuda_context()
|
||||
.context("FOXHUNT_USE_ROLLOUT=1: dev.cuda_context()")?;
|
||||
let stream = trainer.stream.clone();
|
||||
let buf = RolloutBuffer::new(
|
||||
ctx,
|
||||
&stream,
|
||||
cli.n_backtests,
|
||||
t_rollout_default,
|
||||
ROLLOUT_HIDDEN_DIM,
|
||||
)
|
||||
.context("FOXHUNT_USE_ROLLOUT=1: RolloutBuffer::new")?;
|
||||
let coll = RolloutCollection::new(ctx)
|
||||
.context("FOXHUNT_USE_ROLLOUT=1: RolloutCollection::new")?;
|
||||
eprintln!(
|
||||
"[rollout] FOXHUNT_USE_ROLLOUT=1 → RolloutBuffer + GAE active \
|
||||
(B={} T={} HIDDEN_DIM={}, γ from ISV[{}], λ=0.95)",
|
||||
cli.n_backtests,
|
||||
t_rollout_default,
|
||||
ROLLOUT_HIDDEN_DIM,
|
||||
RL_GAMMA_INDEX,
|
||||
);
|
||||
(Some(buf), Some(coll))
|
||||
} else {
|
||||
eprintln!(
|
||||
"[rollout] FOXHUNT_USE_ROLLOUT unset/=0 → legacy per-step path \
|
||||
(RolloutBuffer not allocated; bit-equal to HEAD)"
|
||||
);
|
||||
(None, None)
|
||||
};
|
||||
|
||||
let mut summary = AlphaRlTrainSummary {
|
||||
n_steps_planned: cli.n_steps,
|
||||
seq_len: cli.seq_len,
|
||||
@@ -505,6 +572,10 @@ fn main() -> Result<()> {
|
||||
let mut heat_cap_total: u64 = 0;
|
||||
|
||||
let t_start = std::time::Instant::now();
|
||||
// Phase 1B-B: rollout write-cursor (counts steps within the CURRENT
|
||||
// T_rollout window; resets to 0 after each GAE invocation). Only
|
||||
// advances when `use_rollout` is true.
|
||||
let mut rollout_t: usize = 0;
|
||||
for step in 0..cli.n_steps {
|
||||
// GPU-resident data loading: sample_and_gather + gather_next +
|
||||
// gather_current + gather_frd_labels all run as GPU kernels
|
||||
@@ -513,6 +584,68 @@ fn main() -> Result<()> {
|
||||
.step_with_lobsim_gpu(&mut gpu_data_loader, &gpu_dataset, &mut sim)
|
||||
.with_context(|| format!("step_with_lobsim_gpu at step {step}"))?;
|
||||
|
||||
// ── Phase 1B-B rollout snapshot + GAE (gated). ───────────────
|
||||
//
|
||||
// Stream-ordered after step_with_lobsim_gpu's writes (same
|
||||
// trainer.stream). The snapshot is DtoD-only (cuMemcpy2DAsync +
|
||||
// pack kernel) — zero host roundtrip, zero atomic ops.
|
||||
//
|
||||
// Every T_rollout steps: invoke GAE backward sweep, log the
|
||||
// advantage / returns summary stats to stderr, then reset the
|
||||
// write cursor and `current_t`. This is STRUCTURAL VALIDATION
|
||||
// per the 1B-B dispatch — the resulting advantages_d / returns_d
|
||||
// are not yet consumed by training (1B-C wires that in).
|
||||
if let (Some(buf), Some(coll)) = (rollout_buf.as_mut(), rollout_collection.as_ref()) {
|
||||
coll.snapshot_step(
|
||||
&mut trainer,
|
||||
buf,
|
||||
rollout_t,
|
||||
cli.n_backtests,
|
||||
t_rollout_default,
|
||||
)
|
||||
.with_context(|| {
|
||||
format!("rollout snapshot_step at step={step} rollout_t={rollout_t}")
|
||||
})?;
|
||||
buf.current_t = rollout_t + 1;
|
||||
rollout_t += 1;
|
||||
|
||||
if rollout_t >= t_rollout_default {
|
||||
// Bootstrap V(s_T) from the last step's v_pred_tp1_d
|
||||
// (V on h_{t+1}), then run GAE backward sweep.
|
||||
coll.copy_bootstrap_v(&mut trainer, buf, cli.n_backtests)
|
||||
.context("rollout copy_bootstrap_v")?;
|
||||
let gamma_isv = trainer.read_isv_host(RL_GAMMA_INDEX);
|
||||
// λ = 0.95 hardcoded for 1B-B (1B-C allocates a slot).
|
||||
buf.compute_gae(&trainer.stream, gamma_isv, 0.95_f32)
|
||||
.context("rollout compute_gae")?;
|
||||
|
||||
// Diagnostic: read advantages / returns means to verify
|
||||
// the buffer was populated non-trivially. Uses
|
||||
// `read_slice_d_into` (mapped-pinned, stream-ordered
|
||||
// sync per `feedback_no_htod_htoh_only_mapped_pinned`).
|
||||
let bt = cli.n_backtests * t_rollout_default;
|
||||
let mut adv_host: Vec<f32> = vec![0.0; bt];
|
||||
let mut ret_host: Vec<f32> = vec![0.0; bt];
|
||||
read_slice_d_into(&trainer.stream, &buf.advantages_d, &mut adv_host)
|
||||
.context("rollout read advantages_d")?;
|
||||
read_slice_d_into(&trainer.stream, &buf.returns_d, &mut ret_host)
|
||||
.context("rollout read returns_d")?;
|
||||
let adv_mean: f64 =
|
||||
adv_host.iter().map(|&x| x as f64).sum::<f64>() / bt as f64;
|
||||
let adv_abs_mean: f64 =
|
||||
adv_host.iter().map(|&x| (x as f64).abs()).sum::<f64>() / bt as f64;
|
||||
let ret_mean: f64 =
|
||||
ret_host.iter().map(|&x| x as f64).sum::<f64>() / bt as f64;
|
||||
eprintln!(
|
||||
"[rollout] step={step} GAE complete (T={t_rollout_default}, γ={gamma_isv:.4}, λ=0.95): \
|
||||
adv_mean={adv_mean:.4e} adv_abs_mean={adv_abs_mean:.4e} returns_mean={ret_mean:.4e}"
|
||||
);
|
||||
|
||||
rollout_t = 0;
|
||||
buf.reset();
|
||||
}
|
||||
}
|
||||
|
||||
// Gate G8: NaN abort. Per `feedback_stop_on_anomaly` +
|
||||
// `feedback_kill_runs_on_anomaly_quickly` the cluster smoke
|
||||
// tail-watcher kills the workflow on non-zero exit; here we
|
||||
|
||||
@@ -8,3 +8,4 @@ pub mod optim;
|
||||
pub mod perception;
|
||||
pub mod raw_launch;
|
||||
pub mod rollout_buffer;
|
||||
pub mod rollout_collection;
|
||||
|
||||
375
crates/ml-alpha/src/trainer/rollout_collection.rs
Normal file
375
crates/ml-alpha/src/trainer/rollout_collection.rs
Normal file
@@ -0,0 +1,375 @@
|
||||
//! Phase 1B-B: T-step env collection populating the RolloutBuffer.
|
||||
//!
|
||||
//! Spec: docs/superpowers/specs/2026-06-02-trainer-rollout-buffer-gae.md §1.2
|
||||
//! Plan: docs/superpowers/plans/2026-06-02-trainer-rollout-buffer-gae-implementation.md §Phase 1B-B
|
||||
//!
|
||||
//! Runs `T_rollout` env steps via the existing `step_with_lobsim_gpu`
|
||||
//! pipeline and snapshots per-step `(h_t, a_t, log_pi_old_t, V(s_t),
|
||||
//! r_t, done_t)` into the `RolloutBuffer` at offset `b * T + t`. After
|
||||
//! the loop completes, copies `V(s_{t+1})` from the last step's
|
||||
//! `v_pred_tp1_d` into `v_t_bootstrap_d` (the GAE boundary condition).
|
||||
//! Caller then invokes `RolloutBuffer::compute_gae(γ, λ)`.
|
||||
//!
|
||||
//! ## STRUCTURAL note — Phase 1B-B vs 1B-C split
|
||||
//!
|
||||
//! `step_with_lobsim_gpu` is single-step / training-coupled per
|
||||
//! `pearl_foxhunt_trainer_is_genuinely_single_step` — it does
|
||||
//! forward + env step + PER replay training + Adam in a single call.
|
||||
//! There is no clean "forward-only" path without a ~1500 LOC refactor.
|
||||
//! Per `feedback_investigation_first_falsification_methodology` STOP-on-
|
||||
//! unexpected-finding rule, we do not refactor in 1B-B.
|
||||
//!
|
||||
//! Consequently:
|
||||
//! * 1B-B (this module): training STILL runs per env-step (as it does
|
||||
//! today). The rollout buffer is populated as a SIDE EFFECT after
|
||||
//! each `step_with_lobsim_gpu` call. GAE runs at the end.
|
||||
//! This validates the COLLECTION + GAE pipeline structurally.
|
||||
//! * 1B-C (future): adds multi-epoch PPO update over the rollout
|
||||
//! buffer. THAT is where the per-step training is split — multi-
|
||||
//! epoch PPO replaces the per-step PPO update path.
|
||||
//!
|
||||
//! Phase 1B-B's falsification therefore validates:
|
||||
//! 1. Rollout buffer fills correctly (advantages/returns non-trivial)
|
||||
//! 2. GAE backward sweep produces sensible distributions
|
||||
//! 3. Default `FOXHUNT_USE_ROLLOUT=0` is BIT-EQUAL to HEAD
|
||||
//! 4. `FOXHUNT_USE_ROLLOUT=1` smoke completes without NaN
|
||||
//!
|
||||
//! Per `feedback_no_atomicadd`: no atomic ops.
|
||||
//! Per `feedback_cpu_is_read_only`: pure device-side population.
|
||||
//! Per `feedback_no_htod_htoh_only_mapped_pinned`: all CPU↔GPU paths via
|
||||
//! mapped-pinned helpers; this module uses DtoD-only via the existing
|
||||
//! `raw_memcpy_dtod_async` (same pattern as the trainer hot path).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::driver::{
|
||||
CudaContext, CudaFunction, CudaModule, DevicePtr, DevicePtrMut, LaunchConfig,
|
||||
PushKernelArg,
|
||||
};
|
||||
|
||||
use crate::heads::HIDDEN_DIM;
|
||||
use crate::rl::reward::RlLobBackend;
|
||||
use crate::trainer::integrated::{IntegratedStepStats, IntegratedTrainer};
|
||||
use crate::trainer::rollout_buffer::RolloutBuffer;
|
||||
|
||||
const ROLLOUT_PACK_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/rollout_pack.cubin"));
|
||||
|
||||
/// Helpers for the rollout-collection path. Holds the per-step
|
||||
/// `rollout_pack_dones_f32_to_u8` kernel handle (loaded once at
|
||||
/// construction) and exposes `collect()` as the public entry point.
|
||||
pub struct RolloutCollection {
|
||||
/// Per-step `dones_f32 [B] → dones_u8_bt [B × T]` packer.
|
||||
pack_dones_kernel: CudaFunction,
|
||||
/// Module owning the pack kernel — kept alive for the kernel's lifetime.
|
||||
_pack_module: Arc<CudaModule>,
|
||||
}
|
||||
|
||||
impl RolloutCollection {
|
||||
/// Load the rollout-pack kernel cubin. Allocates no device memory of
|
||||
/// its own; the `RolloutBuffer` is the data owner.
|
||||
pub fn new(ctx: &Arc<CudaContext>) -> Result<Self> {
|
||||
let pack_module = ctx
|
||||
.load_cubin(ROLLOUT_PACK_CUBIN.to_vec())
|
||||
.context("load rollout_pack cubin")?;
|
||||
let pack_dones_kernel = pack_module
|
||||
.load_function("rollout_pack_dones_f32_to_u8")
|
||||
.context("load rollout_pack_dones_f32_to_u8")?;
|
||||
Ok(Self {
|
||||
pack_dones_kernel,
|
||||
_pack_module: pack_module,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run `t_rollout` env steps via `step_with_lobsim_gpu`, snapshotting
|
||||
/// per-step state into `buf`. After this returns, `buf.v_t_bootstrap_d`
|
||||
/// holds `V(s_T)` and the caller should invoke
|
||||
/// `buf.compute_gae(stream, γ, λ)`.
|
||||
///
|
||||
/// Returns the final step's `IntegratedStepStats` so the caller can
|
||||
/// drive its existing diag / NaN-abort logic without restructuring.
|
||||
///
|
||||
/// The per-step training pipeline runs unchanged (single-step PPO +
|
||||
/// PER replay) — see module docs for the 1B-B/1B-C split rationale.
|
||||
pub fn collect(
|
||||
&self,
|
||||
trainer: &mut IntegratedTrainer,
|
||||
gpu_loader: &mut crate::data::gpu_dataset::GpuDataLoader,
|
||||
gpu_dataset: &crate::data::gpu_dataset::GpuDataset,
|
||||
sim: &mut dyn RlLobBackend,
|
||||
buf: &mut RolloutBuffer,
|
||||
t_rollout: usize,
|
||||
) -> Result<IntegratedStepStats> {
|
||||
anyhow::ensure!(
|
||||
t_rollout > 0 && t_rollout <= buf.t_rollout,
|
||||
"RolloutCollection::collect: t_rollout={t_rollout} out of range \
|
||||
(buf.t_rollout={})",
|
||||
buf.t_rollout
|
||||
);
|
||||
anyhow::ensure!(
|
||||
buf.hidden_dim == HIDDEN_DIM,
|
||||
"RolloutCollection::collect: buf.hidden_dim={} != HIDDEN_DIM={}",
|
||||
buf.hidden_dim,
|
||||
HIDDEN_DIM
|
||||
);
|
||||
anyhow::ensure!(
|
||||
buf.b_size == trainer.perception.h_t_view().len() / HIDDEN_DIM,
|
||||
"RolloutCollection::collect: buf.b_size={} != trainer h_t batch \
|
||||
(h_t.len={} / HIDDEN_DIM={})",
|
||||
buf.b_size,
|
||||
trainer.perception.h_t_view().len(),
|
||||
HIDDEN_DIM
|
||||
);
|
||||
|
||||
buf.reset();
|
||||
let b_size = buf.b_size;
|
||||
let t_max = buf.t_rollout;
|
||||
|
||||
let mut last_stats: Option<IntegratedStepStats> = None;
|
||||
for t in 0..t_rollout {
|
||||
// Drive ONE env step via the existing single-step pipeline.
|
||||
// This populates the trainer's per-step `[B]` buffers
|
||||
// (rewards_d / dones_d / actions_d / log_pi_old_d /
|
||||
// v_pred_d / v_pred_tp1_d) and `perception.h_t_view()` with
|
||||
// values from this step's forward + env transition.
|
||||
let stats = trainer
|
||||
.step_with_lobsim_gpu(gpu_loader, gpu_dataset, sim)
|
||||
.with_context(|| {
|
||||
format!("RolloutCollection::collect: step_with_lobsim_gpu at t={t}")
|
||||
})?;
|
||||
|
||||
// Snapshot the per-step `[B]` buffers into the rollout
|
||||
// buffer at offset `b * T + t` for each batch element.
|
||||
// All copies are DtoD-async on the trainer's main stream,
|
||||
// so they're stream-ordered after the step's writes — no
|
||||
// cross-stream sync needed.
|
||||
self.snapshot_step(trainer, buf, t, b_size, t_max)?;
|
||||
|
||||
buf.current_t = t + 1;
|
||||
last_stats = Some(stats);
|
||||
}
|
||||
|
||||
// Bootstrap V(s_T): the last step's `v_pred_tp1_d` IS V(s_{t+1})
|
||||
// for the final action — exactly the GAE bootstrap target. Copy
|
||||
// `[B]` f32 DtoD into `buf.v_t_bootstrap_d`.
|
||||
self.copy_bootstrap_v(trainer, buf, b_size)?;
|
||||
|
||||
Ok(last_stats.expect("t_rollout > 0 guarantees ≥1 iteration"))
|
||||
}
|
||||
|
||||
/// DtoD-snapshot the trainer's per-step `[B]` buffers into the
|
||||
/// rollout buffer's `[B × T]` slots at write cursor `t`.
|
||||
///
|
||||
/// Each copy is `B × sizeof(T)` bytes on the trainer's main stream.
|
||||
/// Stream-ordered after the producing `step_with_lobsim_gpu` writes.
|
||||
///
|
||||
/// PUBLIC for integration with the existing main training loop —
|
||||
/// the loop drives `step_with_lobsim_gpu` itself (so the diag /
|
||||
/// NaN-abort / counter machinery stays intact) and calls
|
||||
/// `snapshot_step` after each step.
|
||||
pub fn snapshot_step(
|
||||
&self,
|
||||
trainer: &mut IntegratedTrainer,
|
||||
buf: &mut RolloutBuffer,
|
||||
t: usize,
|
||||
b_size: usize,
|
||||
t_max: usize,
|
||||
) -> Result<()> {
|
||||
let stream = &trainer.stream;
|
||||
let raw_stream = stream.cu_stream();
|
||||
|
||||
// `gae_backward_sweep.cu` indexes `rewards[b * T + t]` — i.e.
|
||||
// `[B, T]` row-major (batch-major). Per-step writes therefore
|
||||
// scatter a contiguous `[B]` buffer into column `t` of `[B, T]`:
|
||||
// one element per batch row, with destination stride
|
||||
// `T * sizeof(elem)`. cuMemcpy2D is the canonical DtoD-only idiom
|
||||
// for this layout:
|
||||
// * src pitch = elem_sz (contiguous [B] source)
|
||||
// * dst pitch = T * elem_sz (per-batch row pitch of [B, T])
|
||||
// * width = elem_sz (one element per row)
|
||||
// * height = B (one row per batch element)
|
||||
// No host roundtrip, no atomic ops; deterministic per
|
||||
// `pearl_determinism_achieved` (single-direction DtoD).
|
||||
// The u8 dones use a dedicated pack kernel (`rollout_pack.cu`)
|
||||
// that does the f32→u8 conversion + `b * T + t` scatter in one
|
||||
// launch.
|
||||
|
||||
let sz_f32 = std::mem::size_of::<f32>();
|
||||
let sz_i32 = std::mem::size_of::<i32>();
|
||||
|
||||
// Helper: scatter `[B]` of `elem_sz` bytes into `[B, T]` column t.
|
||||
// Implements cuMemcpy2D with src stride = elem_sz (contiguous),
|
||||
// dst stride = T * elem_sz (row pitch of the [B, T] buffer).
|
||||
let scatter_to_column = |dst_base: u64, src: u64, elem_sz: usize| -> Result<()> {
|
||||
// src offset 0, dst offset = t * elem_sz (column t in row b=0)
|
||||
let dst_off = (t as u64) * (elem_sz as u64);
|
||||
unsafe {
|
||||
let mut copy = std::mem::MaybeUninit::<
|
||||
cudarc::driver::sys::CUDA_MEMCPY2D_v2,
|
||||
>::zeroed();
|
||||
let cref = &mut *copy.as_mut_ptr();
|
||||
cref.srcMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE;
|
||||
cref.srcDevice = src;
|
||||
cref.srcPitch = elem_sz; // contiguous source [B] → row pitch = elem_sz
|
||||
cref.dstMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE;
|
||||
cref.dstDevice = dst_base + dst_off;
|
||||
cref.dstPitch = (t_max as usize) * elem_sz; // row pitch of [B, T] = T * elem_sz
|
||||
cref.WidthInBytes = elem_sz; // one element per row
|
||||
cref.Height = b_size;
|
||||
let r = cudarc::driver::sys::cuMemcpy2DAsync_v2(
|
||||
copy.as_ptr(),
|
||||
raw_stream,
|
||||
);
|
||||
if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
|
||||
anyhow::bail!(
|
||||
"RolloutCollection::scatter_to_column cuMemcpy2DAsync \
|
||||
(t={t}, elem_sz={elem_sz}): {:?}",
|
||||
r
|
||||
);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
|
||||
// rewards_d: f32 [B] → rollout.rewards_d[b, t]
|
||||
let rewards_dst = buf.rewards_d.device_ptr_mut(stream).0;
|
||||
scatter_to_column(rewards_dst, trainer.rewards_d.device_ptr(stream).0, sz_f32)?;
|
||||
|
||||
// v_t_d: f32 [B] (= trainer.v_pred_d) → rollout.v_t_d[b, t]
|
||||
let v_t_dst = buf.v_t_d.device_ptr_mut(stream).0;
|
||||
scatter_to_column(v_t_dst, trainer.v_pred_d.device_ptr(stream).0, sz_f32)?;
|
||||
|
||||
// actions_d: i32 [B] → rollout.actions_d[b, t]
|
||||
let actions_dst = buf.actions_d.device_ptr_mut(stream).0;
|
||||
scatter_to_column(actions_dst, trainer.actions_d.device_ptr(stream).0, sz_i32)?;
|
||||
|
||||
// log_pi_old_d: f32 [B] → rollout.log_pi_old_d[b, t]
|
||||
let log_pi_dst = buf.log_pi_old_d.device_ptr_mut(stream).0;
|
||||
scatter_to_column(log_pi_dst, trainer.log_pi_old_d.device_ptr(stream).0, sz_f32)?;
|
||||
|
||||
// h_t_d: f32 [B × HIDDEN_DIM] → rollout.h_t_d[b, t, :]
|
||||
//
|
||||
// Layout: rollout.h_t_d is `[B × T × HIDDEN_DIM]` row-major in
|
||||
// (b, t, h). One batch row at offset (b, t, *) starts at byte
|
||||
// `(b * T + t) * HIDDEN_DIM * sizeof(f32)`. Per-batch rows of
|
||||
// length HIDDEN_DIM are not contiguous across `b` (stride =
|
||||
// T * HIDDEN_DIM * 4 bytes). Source `perception.h_t_view()` is
|
||||
// contiguous `[B × HIDDEN_DIM]` (stride = HIDDEN_DIM * 4 bytes).
|
||||
//
|
||||
// cuMemcpy2D with:
|
||||
// * src stride = HIDDEN_DIM * sz_f32 (per-batch row pitch)
|
||||
// * dst stride = T * HIDDEN_DIM * sz_f32 (per-batch row pitch in [B, T, H])
|
||||
// * WidthInBytes = HIDDEN_DIM * sz_f32 (full hidden row)
|
||||
// * Height = B
|
||||
// copies B rows of HIDDEN_DIM floats from contiguous src to
|
||||
// strided dst at column t.
|
||||
{
|
||||
let h_src = trainer.perception.h_t_view().device_ptr(stream).0;
|
||||
let h_dst_base = buf.h_t_d.device_ptr_mut(stream).0;
|
||||
let row_bytes = HIDDEN_DIM * sz_f32;
|
||||
let dst_off = (t as u64) * (row_bytes as u64);
|
||||
unsafe {
|
||||
let mut copy = std::mem::MaybeUninit::<
|
||||
cudarc::driver::sys::CUDA_MEMCPY2D_v2,
|
||||
>::zeroed();
|
||||
let cref = &mut *copy.as_mut_ptr();
|
||||
cref.srcMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE;
|
||||
cref.srcDevice = h_src;
|
||||
cref.srcPitch = row_bytes;
|
||||
cref.dstMemoryType = cudarc::driver::sys::CUmemorytype::CU_MEMORYTYPE_DEVICE;
|
||||
cref.dstDevice = h_dst_base + dst_off;
|
||||
cref.dstPitch = (t_max as usize) * row_bytes;
|
||||
cref.WidthInBytes = row_bytes;
|
||||
cref.Height = b_size;
|
||||
let r = cudarc::driver::sys::cuMemcpy2DAsync_v2(
|
||||
copy.as_ptr(),
|
||||
raw_stream,
|
||||
);
|
||||
if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS {
|
||||
anyhow::bail!(
|
||||
"RolloutCollection::snapshot_step h_t cuMemcpy2DAsync \
|
||||
(t={t}): {:?}",
|
||||
r
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dones_d: f32 [B] {0.0, 1.0} → rollout.dones_d[b, t] (u8 {0, 1})
|
||||
// Pack kernel writes `b * T + t` directly per the GAE indexing.
|
||||
self.launch_pack_dones(trainer, buf, t, b_size, t_max)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `rollout_pack_dones_f32_to_u8` to convert + scatter the
|
||||
/// per-step `[B]` f32 dones into the `[B × T]` u8 rollout buffer at
|
||||
/// column `t`. Single-thread-per-batch kernel; one launch per step.
|
||||
fn launch_pack_dones(
|
||||
&self,
|
||||
trainer: &mut IntegratedTrainer,
|
||||
buf: &mut RolloutBuffer,
|
||||
t: usize,
|
||||
b_size: usize,
|
||||
t_max: usize,
|
||||
) -> Result<()> {
|
||||
let stream = &trainer.stream;
|
||||
let b = b_size as i32;
|
||||
let t_i = t_max as i32;
|
||||
let t_cursor = t as i32;
|
||||
|
||||
let threads: u32 = 256;
|
||||
let blocks: u32 = ((b_size as u32) + threads - 1) / threads;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (blocks.max(1), 1, 1),
|
||||
block_dim: (threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
|
||||
let src = trainer.dones_d.device_ptr(stream).0;
|
||||
let dst = buf.dones_d.device_ptr_mut(stream).0;
|
||||
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.pack_dones_kernel)
|
||||
.arg(&src)
|
||||
.arg(&dst)
|
||||
.arg(&b)
|
||||
.arg(&t_i)
|
||||
.arg(&t_cursor)
|
||||
.launch(cfg)
|
||||
.with_context(|| {
|
||||
format!("rollout_pack_dones_f32_to_u8 launch (t={t})")
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Copy the last step's `v_pred_tp1_d` ([B] f32 — `V(s_{t+1})` for
|
||||
/// the final action) into `buf.v_t_bootstrap_d`. This is the
|
||||
/// boundary condition for the GAE backward sweep at `t = T - 1`.
|
||||
///
|
||||
/// Contiguous DtoD; `B × sizeof(f32)` bytes. PUBLIC for the same
|
||||
/// reason as `snapshot_step`.
|
||||
pub fn copy_bootstrap_v(
|
||||
&self,
|
||||
trainer: &mut IntegratedTrainer,
|
||||
buf: &mut RolloutBuffer,
|
||||
b_size: usize,
|
||||
) -> Result<()> {
|
||||
use crate::trainer::raw_launch::raw_memcpy_dtod_async;
|
||||
|
||||
let stream = &trainer.stream;
|
||||
let raw_stream = stream.cu_stream();
|
||||
let src = trainer.v_pred_tp1_d.device_ptr(stream).0;
|
||||
let dst = buf.v_t_bootstrap_d.device_ptr_mut(stream).0;
|
||||
let nbytes = b_size * std::mem::size_of::<f32>();
|
||||
unsafe {
|
||||
raw_memcpy_dtod_async(dst, src, nbytes, raw_stream).map_err(|e| {
|
||||
anyhow::anyhow!("copy_bootstrap_v DtoD: {:?}", e)
|
||||
})?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user