From b75ebc7e0dfbbde7b8f863a0769a131172ffd1c8 Mon Sep 17 00:00:00 2001 From: Kbediako Date: Thu, 23 Jul 2026 11:51:15 +1000 Subject: [PATCH] Preserve sampled actions in float32 --- src/kernels.cu | 5 +- src/pufferlib.cu | 35 ++-- tests/profile_kernels.cu | 416 ++++++++++++++++++++++++++++++++++++++- 3 files changed, 435 insertions(+), 21 deletions(-) diff --git a/src/kernels.cu b/src/kernels.cu index b3baae5089..c33610ef62 100644 --- a/src/kernels.cu +++ b/src/kernels.cu @@ -176,8 +176,9 @@ __device__ __forceinline__ void copy_bytes(const char* __restrict__ src, */ // Transpose dims 0,1: [A, B, C] -> [B, A, C]. For 2D, pass C=1. -__global__ void transpose_102(precision_t* __restrict__ dst, - const precision_t* __restrict__ src, int A, int B, int C) { +template +__global__ void transpose_102(T* __restrict__ dst, + const T* __restrict__ src, int A, int B, int C) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = A * B * C; if (idx >= total) { diff --git a/src/pufferlib.cu b/src/pufferlib.cu index 4af2e2b640..f067a3a71c 100644 --- a/src/pufferlib.cu +++ b/src/pufferlib.cu @@ -50,7 +50,7 @@ typedef struct { // a constant subset of agents struct RolloutBuf { PrecisionTensor observations; // (horizon, agents, input_size) - PrecisionTensor actions; // (horizon, agents, num_atns) + FloatTensor actions; // (horizon, agents, num_atns) PrecisionTensor values; // (horizon, agents) PrecisionTensor logprobs; // ... PrecisionTensor rewards; @@ -95,7 +95,7 @@ void register_rollout_buffers(RolloutBuf& bufs, Allocator* alloc, int T, int B, struct TrainGraph { PrecisionTensor mb_state; // (layers, B, hidden) PrecisionTensor mb_obs; // (B, T, input_size) - PrecisionTensor mb_actions; // (B, T, num_atns) + FloatTensor mb_actions; // (B, T, num_atns) PrecisionTensor mb_logprobs; // (B, T) PrecisionTensor mb_advantages; // ... PrecisionTensor mb_values; @@ -142,7 +142,7 @@ void register_train_buffers(TrainGraph& bufs, Allocator* alloc, int B, int T, in struct PPOGraphArgs { precision_t* out_ratio; precision_t* out_newvalue; - const precision_t* actions; + const float* actions; const precision_t* old_logprobs; const precision_t* advantages; const precision_t* prio; @@ -224,7 +224,8 @@ void register_prio_buffers(PrioBuffers& bufs, Allocator* alloc, int B, int minib // Slice: select dim0 index t, then narrow dim0 from start for count. // 3D (T, B, F) -> (count, F); 2D (T, B) -> (count,) -inline PrecisionTensor puf_slice(PrecisionTensor& p, int t, int start, int count) { +template +inline Tensor puf_slice(Tensor& p, int t, int start, int count) { if (ndim(p.shape) == 3) { long B = p.shape[1], F = p.shape[2]; return {.data = p.data + (t*B + start)*F, .shape = {count, F}}; @@ -466,7 +467,8 @@ __global__ void sample_logits( PrecisionTensor dec_out, // (B, logits_dim + 1 for values) PrecisionTensor logstd_puf, // (1, od) - continuous actions only IntTensor act_sizes_puf, // (num_atns,) action head sizes - precision_t* __restrict__ actions, // (B, num_atns) + float* __restrict__ actions, // (B, num_atns) + float* __restrict__ env_actions, // (B, num_atns) precision_t* __restrict__ logprobs, // (B,) precision_t* __restrict__ value_out, // (B,) curandStatePhilox4_32_10_t* __restrict__ rng_states, @@ -509,13 +511,16 @@ __global__ void sample_logits( float noise = curand_normal(&state); float action = finite_or_clamp(mean + std * noise, -1.0e6f, 1.0e6f); + // Preserve the previous reduced-precision continuous action semantics. precision_t stored_action_p = from_float(action); float stored_action = to_float(stored_action_p); // Log probability: -0.5 * ((action - mean) / std)^2 - 0.5 * log(2*pi) - log(std) float normalized = (stored_action - mean) / std; float log_prob = -0.5f * normalized * normalized - 0.5f * LOG_2PI - log_std; - actions[idx * num_atns + h] = stored_action_p; + int action_idx = idx * num_atns + h; + actions[action_idx] = stored_action; + env_actions[action_idx] = stored_action; total_log_prob += log_prob; } } else { @@ -573,8 +578,11 @@ __global__ void sample_logits( float sampled_logit = masked_logit(logits, logits_base, logits_offset, sampled_action, action_mask, mask_base); float log_prob = sampled_logit - logsumexp; - // Write action for this head - actions[idx * num_atns + h] = from_float(sampled_action); + // Float32 preserves large categorical indices that BF16 cannot represent exactly. + int action_idx = idx * num_atns + h; + float action = static_cast(sampled_action); + actions[action_idx] = action; + env_actions[action_idx] = action; total_log_prob += log_prob; // Advance to next action head @@ -680,7 +688,7 @@ extern "C" void net_callback_wrapper(void* ctx, int buf, int t) { int sub_start = start + bank_off; PrecisionTensor obs_b = puf_slice(rollouts.observations, t, sub_start, bank_size); - PrecisionTensor act_b = puf_slice(rollouts.actions, t, sub_start, bank_size); + FloatTensor act_b = puf_slice(rollouts.actions, t, sub_start, bank_size); PrecisionTensor lp_b = puf_slice(rollouts.logprobs, t, sub_start, bank_size); PrecisionTensor val_b = puf_slice(rollouts.values, t, sub_start, bank_size); PrecisionTensor mask_b = {}; @@ -701,13 +709,10 @@ extern "C" void net_callback_wrapper(void* ctx, int buf, int t) { // Offset RNG by bank_off so banks don't collide on per-buffer rng slots. sample_logits<<>>( dec_puf, p_logstd, pufferl->act_sizes_puf, - act_b.data, lp_b.data, val_b.data, + act_b.data, env.actions.data + (long)sub_start * act_cols, + lp_b.data, val_b.data, pufferl->rng_states[buf] + bank_off, mask_b.data, mask_stride_b); - - cast<<>>( - env.actions.data + (long)sub_start * act_cols, - act_b.data, numel(act_b.shape)); } if (capturing) { @@ -1453,7 +1458,7 @@ __global__ void select_copy(RolloutBuf rollouts, TrainGraph graph, // Compute row byte counts from tensor shapes int obs_row_bytes = (numel(rollouts.observations.shape) / rollouts.observations.shape[0]) * sizeof(precision_t); - int act_row_bytes = (numel(rollouts.actions.shape) / rollouts.actions.shape[0]) * sizeof(precision_t); + int act_row_bytes = (numel(rollouts.actions.shape) / rollouts.actions.shape[0]) * sizeof(float); int lp_row_bytes = (numel(rollouts.logprobs.shape) / rollouts.logprobs.shape[0]) * sizeof(precision_t); int horizon = rollouts.values.shape[1]; diff --git a/tests/profile_kernels.cu b/tests/profile_kernels.cu index 269102be41..f66ed1f3c6 100644 --- a/tests/profile_kernels.cu +++ b/tests/profile_kernels.cu @@ -4,10 +4,19 @@ #include #include #include +#include #include "pufferlib.cu" #include "ini.h" +static_assert(std::is_same_v, FloatTensor>); +static_assert(std::is_same_v, FloatTensor>); +static_assert(std::is_same_v); +static_assert(std::is_same_v, PrecisionTensor>); +static_assert(std::is_same_v, PrecisionTensor>); +static_assert(std::is_same_v, PrecisionTensor>); +static_assert(std::is_same_v, PrecisionTensor>); + const int WARMUP_ITERS = 100; const int TIMING_ITERS = 1000; @@ -35,6 +44,7 @@ void print_usage(const char* prog) { printf(" logcoeffsvals - log_coeffs_and_values fwd+bwd\n"); printf(" fusedscan - Fused scan (checkpointed) kernel only\n"); printf(" samplelogits - Sample logits kernel only\n"); + printf(" actiontransport - Verify exact discrete and rounded continuous action transport\n"); printf(" ppoloss - PPO loss fused fwd+bwd kernel\n"); printf(" im2col - im2col + col2im (nmmo3 conv sizes, B=1024)\n"); printf(" envspeed - Environment step throughput\n"); @@ -322,7 +332,8 @@ struct PPOProfile { PPOGraphArgs ga; FloatTensor loss, losses_acc, ppo_partials; FloatTensor grad_logits_t, grad_values_t, adv_mean_t, adv_var_t, ent_coef_t; - PrecisionTensor logits_t, actions_t, old_logprobs_t, advantages_t, prio_t, values_t, returns_t; + FloatTensor actions_t; + PrecisionTensor logits_t, old_logprobs_t, advantages_t, prio_t, values_t, returns_t; PrecisionTensor ratio_t, newvalue_t; IntTensor act_sizes_t; Allocator alloc; @@ -405,7 +416,7 @@ PPOProfile* create_ppoloss(int N, int T, int A) { float_to_device(p->logits_t.data, buf, NT * fused_cols); // actions for (int i = 0; i < NT; ++i) buf[i] = (float)(rand() % A); - float_to_device(p->actions_t.data, buf, NT); + cudaMemcpy(p->actions_t.data, buf, NT * sizeof(float), cudaMemcpyHostToDevice); // old_logprobs for (int i = 0; i < NT; ++i) buf[i] = rand1() * 2.0f; float_to_device(p->old_logprobs_t.data, buf, NT); @@ -473,7 +484,8 @@ void profile_ppoloss(int N, int T, int A) { struct SampleLogitsProfile { PrecisionTensor dec_out, logstd; IntTensor act_sizes; - PrecisionTensor actions_t, logprobs_t, value_out_t; + FloatTensor actions_t, env_actions_t; + PrecisionTensor logprobs_t, value_out_t; curandStatePhilox4_32_10_t* rng_states; Allocator alloc; int B, A; @@ -488,6 +500,7 @@ SampleLogitsProfile* create_samplelogits(int B, int A) { p->logstd = {.shape = {0}}; // empty for discrete p->act_sizes = {.shape = {1}}; p->actions_t = {.shape = {B}}; + p->env_actions_t = {.shape = {B}}; p->logprobs_t = {.shape = {B}}; p->value_out_t = {.shape = {B}}; @@ -495,6 +508,7 @@ SampleLogitsProfile* create_samplelogits(int B, int A) { alloc_register(&p->alloc, &p->dec_out); alloc_register(&p->alloc, &p->act_sizes); alloc_register(&p->alloc, &p->actions_t); + alloc_register(&p->alloc, &p->env_actions_t); alloc_register(&p->alloc, &p->logprobs_t); alloc_register(&p->alloc, &p->value_out_t); alloc_create(&p->alloc); @@ -515,7 +529,8 @@ SampleLogitsProfile* create_samplelogits(int B, int A) { void run_samplelogits(SampleLogitsProfile* p) { sample_logits<<B), BLOCK_SIZE>>>( p->dec_out, p->logstd, p->act_sizes, - p->actions_t.data, p->logprobs_t.data, p->value_out_t.data, + p->actions_t.data, p->env_actions_t.data, + p->logprobs_t.data, p->value_out_t.data, p->rng_states, nullptr, 0); } @@ -530,6 +545,396 @@ void profile_samplelogits(int B, int A) { free(p); } +static float round_to_bf16(float value) { + uint32_t bits; + memcpy(&bits, &value, sizeof(bits)); + bits += 0x7fffu + ((bits >> 16) & 1u); + bits &= 0xffff0000u; + memcpy(&value, &bits, sizeof(value)); + return value; +} + +static void require_equal(const char* path, const std::vector& actual, + const std::vector& expected) { + for (size_t i = 0; i < expected.size(); ++i) { + if (actual[i] != expected[i]) { + throw std::runtime_error(std::string(path) + " changed action " + + std::to_string(expected[i]) + " to " + std::to_string(actual[i])); + } + } +} + +static std::vector precision_to_host(const precision_t* device, int count) { + std::vector stored(count); + std::vector result(count); + cudaMemcpy(stored.data(), device, count * sizeof(precision_t), cudaMemcpyDeviceToHost); + for (int i = 0; i < count; ++i) result[i] = to_float(stored[i]); + return result; +} + +void test_action_transport() { + constexpr int T = 2; + constexpr int B = 4; + constexpr int H = 2; + constexpr int MB = 2; + constexpr int A0 = 3572; + constexpr int A1 = 7; + constexpr int A = A0 + A1; + constexpr int S = T * B; + const std::vector expected = { + 257, 1, 511, 2, 2049, 3, 3571, 4, + 511, 5, 2049, 6, 3571, 0, 257, 1, + }; + + PrecisionTensor dec_out = {.shape = {S, A + 1}}; + PrecisionTensor action_mask = {.shape = {S, A}}; + PrecisionTensor logprobs = {.shape = {S}}; + PrecisionTensor values = {.shape = {S}}; + FloatTensor rollout_actions = {.shape = {S, H}}; + FloatTensor env_actions = {.shape = {S, H}}; + IntTensor act_sizes = {.shape = {H}}; + + PrecisionTensor advantages = {.shape = {B, T}}; + FloatTensor priorities = {.shape = {MB}}; + IntTensor indices = {.shape = {MB}}; + + Allocator alloc = {}; + RolloutBuf transposed; + TrainGraph minibatch; + register_rollout_buffers(transposed, &alloc, B, T, 1, H, 0); + register_train_buffers(minibatch, &alloc, MB, T, 1, 1, H, 1, 0); + alloc_register(&alloc, &dec_out); + alloc_register(&alloc, &action_mask); + alloc_register(&alloc, &logprobs); + alloc_register(&alloc, &values); + alloc_register(&alloc, &rollout_actions); + alloc_register(&alloc, &env_actions); + alloc_register(&alloc, &act_sizes); + alloc_register(&alloc, &advantages); + alloc_register(&alloc, &priorities); + alloc_register(&alloc, &indices); + if (alloc_create(&alloc) != cudaSuccess) { + throw std::runtime_error("action transport test allocation failed"); + } + + std::vector logits(S * (A + 1), 0.0f); + std::vector mask(S * A, 0.0f); + for (int i = 0; i < S; ++i) { + mask[i * A + (int)expected[i * H]] = 1.0f; + mask[i * A + A0 + (int)expected[i * H + 1]] = 1.0f; + } + float_to_device(dec_out.data, logits.data(), logits.size()); + float_to_device(action_mask.data, mask.data(), mask.size()); + cudaMemset(transposed.observations.data, 0, B * T * sizeof(precision_t)); + cudaMemset(transposed.logprobs.data, 0, B * T * sizeof(precision_t)); + cudaMemset(transposed.values.data, 0, B * T * sizeof(precision_t)); + cudaMemset(advantages.data, 0, B * T * sizeof(precision_t)); + const float host_priorities[MB] = {1.0f, 1.0f}; + const int action_sizes[H] = {A0, A1}; + const int host_indices[MB] = {2, 0}; + cudaMemcpy(priorities.data, host_priorities, sizeof(host_priorities), cudaMemcpyHostToDevice); + cudaMemcpy(act_sizes.data, action_sizes, sizeof(action_sizes), cudaMemcpyHostToDevice); + cudaMemcpy(indices.data, host_indices, sizeof(host_indices), cudaMemcpyHostToDevice); + + curandStatePhilox4_32_10_t* rng_states = nullptr; + cudaMalloc(&rng_states, S * sizeof(*rng_states)); + rng_init<<>>(rng_states, 42, S); + PrecisionTensor no_logstd = {}; + sample_logits<<>>( + dec_out, no_logstd, act_sizes, rollout_actions.data, env_actions.data, + logprobs.data, values.data, rng_states, action_mask.data, A); + transpose_102<<>>( + transposed.actions.data, rollout_actions.data, T, B, H); + select_copy<<>>( + transposed, minibatch, indices.data, advantages.data, priorities.data); + + std::vector expected_transposed(S * H); + for (int b = 0; b < B; ++b) { + for (int t = 0; t < T; ++t) { + for (int h = 0; h < H; ++h) { + expected_transposed[(b * T + t) * H + h] = expected[(t * B + b) * H + h]; + } + } + } + std::vector expected_minibatch(MB * T * H); + for (int mb = 0; mb < MB; ++mb) { + std::copy_n(expected_transposed.begin() + host_indices[mb] * T * H, T * H, + expected_minibatch.begin() + mb * T * H); + } + + std::vector rollout(S * H), env(S * H), train(S * H), ppo(MB * T * H); + cudaMemcpy(rollout.data(), rollout_actions.data, rollout.size() * sizeof(float), + cudaMemcpyDeviceToHost); + cudaMemcpy(env.data(), env_actions.data, env.size() * sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(train.data(), transposed.actions.data, train.size() * sizeof(float), + cudaMemcpyDeviceToHost); + cudaMemcpy(ppo.data(), minibatch.mb_actions.data, ppo.size() * sizeof(float), + cudaMemcpyDeviceToHost); + require_equal("rollout", rollout, expected); + require_equal("environment", env, expected); + require_equal("transpose", train, expected_transposed); + require_equal("PPO minibatch", ppo, expected_minibatch); + + // Synthetic primary + two frozen-bank slices. This isolates the captured + // sampler's pointer offsets; it is not a full WeightBank integration test. + constexpr int BANKS = 3; + const int bank_starts[BANKS + 1] = {0, 2, 5, S}; + cudaStream_t graph_stream; + cudaGraph_t action_graph; + cudaGraphExec_t action_graph_exec; + if (cudaStreamCreateWithFlags(&graph_stream, cudaStreamNonBlocking) != cudaSuccess || + cudaStreamBeginCapture(graph_stream, cudaStreamCaptureModeGlobal) != cudaSuccess) { + throw std::runtime_error("synthetic bank graph capture setup failed"); + } + for (int bank = 0; bank < BANKS; ++bank) { + int start = bank_starts[bank]; + int count = bank_starts[bank + 1] - start; + PrecisionTensor bank_out = { + .data = dec_out.data + (long)start * (A + 1), .shape = {count, A + 1}}; + sample_logits<<>>( + bank_out, no_logstd, act_sizes, + rollout_actions.data + (long)start * H, + env_actions.data + (long)start * H, + logprobs.data + start, values.data + start, rng_states + start, + action_mask.data + (long)start * A, A); + } + if (cudaStreamEndCapture(graph_stream, &action_graph) != cudaSuccess || + cudaGraphInstantiate(&action_graph_exec, action_graph, 0) != cudaSuccess) { + throw std::runtime_error("synthetic bank graph creation failed"); + } + for (int replay = 0; replay < 2; ++replay) { + cudaMemsetAsync(rollout_actions.data, 0xff, S * H * sizeof(float), graph_stream); + cudaMemsetAsync(env_actions.data, 0xff, S * H * sizeof(float), graph_stream); + if (cudaGraphLaunch(action_graph_exec, graph_stream) != cudaSuccess || + cudaStreamSynchronize(graph_stream) != cudaSuccess) { + throw std::runtime_error("synthetic bank graph replay failed"); + } + cudaMemcpy(rollout.data(), rollout_actions.data, rollout.size() * sizeof(float), + cudaMemcpyDeviceToHost); + cudaMemcpy(env.data(), env_actions.data, env.size() * sizeof(float), + cudaMemcpyDeviceToHost); + require_equal("synthetic bank graph rollout", rollout, expected); + require_equal("synthetic bank graph environment", env, expected); + } + cudaGraphExecDestroy(action_graph_exec); + cudaGraphDestroy(action_graph); + cudaStreamDestroy(graph_stream); + + PPOProfile* discrete_ppo = create_ppoloss(1, 1, A0); + std::vector ppo_logits(A0 + 1, 0.0f); + float_to_device(discrete_ppo->logits_t.data, ppo_logits.data(), ppo_logits.size()); + const float one = 1.0f; + const float zero = 0.0f; + const float old_logprob = -logf((float)A0); + float_to_device(discrete_ppo->old_logprobs_t.data, &old_logprob, 1); + float_to_device(discrete_ppo->advantages_t.data, &one, 1); + float_to_device(discrete_ppo->prio_t.data, &one, 1); + float_to_device(discrete_ppo->values_t.data, &zero, 1); + float_to_device(discrete_ppo->returns_t.data, &zero, 1); + cudaMemcpy(discrete_ppo->adv_mean_t.data, &zero, sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(discrete_ppo->adv_var_t.data, &one, sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(discrete_ppo->ent_coef_t.data, &zero, sizeof(float), cudaMemcpyHostToDevice); + discrete_ppo->ga.actions = minibatch.mb_actions.data; + run_ppoloss(discrete_ppo); + std::vector discrete_grad(A0); + cudaMemcpy(discrete_grad.data(), discrete_ppo->grad_logits_t.data, + A0 * sizeof(float), cudaMemcpyDeviceToHost); + int selected_action = (int)expected_minibatch[0]; + if (!(discrete_grad[selected_action] < 0.0f && + discrete_grad[selected_action - 1] > 0.0f)) { + throw std::runtime_error("PPO did not consume the selected exact large action"); + } + alloc_free(&discrete_ppo->alloc); + free(discrete_ppo); + + constexpr int CB = 32; + constexpr int CH = 2; + PrecisionTensor continuous_out = {.shape = {CB, CH + 1}}; + PrecisionTensor logstd = {.shape = {CH}}; + PrecisionTensor continuous_logprobs = {.shape = {CB}}; + PrecisionTensor continuous_values = {.shape = {CB}}; + FloatTensor continuous_rollout = {.shape = {CB, CH}}; + FloatTensor continuous_env = {.shape = {CB, CH}}; + IntTensor continuous_size = {.shape = {CH}}; + Allocator continuous_alloc = {}; + alloc_register(&continuous_alloc, &continuous_out); + alloc_register(&continuous_alloc, &logstd); + alloc_register(&continuous_alloc, &continuous_logprobs); + alloc_register(&continuous_alloc, &continuous_values); + alloc_register(&continuous_alloc, &continuous_rollout); + alloc_register(&continuous_alloc, &continuous_env); + alloc_register(&continuous_alloc, &continuous_size); + if (alloc_create(&continuous_alloc) != cudaSuccess) { + throw std::runtime_error("continuous action transport test allocation failed"); + } + + std::vector continuous_logits(CB * (CH + 1)); + for (int i = 0; i < CB; ++i) { + continuous_logits[i * (CH + 1)] = 0.1234567f; + continuous_logits[i * (CH + 1) + 1] = -0.2345678f; + continuous_logits[i * (CH + 1) + 2] = 0.0f; + } + const float continuous_logstd[CH] = {-1.25f, -0.7f}; + const int continuous_sizes[CH] = {1, 1}; + float_to_device(continuous_out.data, continuous_logits.data(), continuous_logits.size()); + float_to_device(logstd.data, continuous_logstd, CH); + cudaMemcpy(continuous_size.data, continuous_sizes, sizeof(continuous_sizes), + cudaMemcpyHostToDevice); + cudaFree(rng_states); + cudaMalloc(&rng_states, CB * sizeof(*rng_states)); + rng_init<<>>(rng_states, 73, CB); + sample_logits<<>>( + continuous_out, logstd, continuous_size, + continuous_rollout.data, continuous_env.data, + continuous_logprobs.data, continuous_values.data, + rng_states, nullptr, 0); + + std::vector continuous_rollout_host(CB * CH), continuous_env_host(CB * CH); + cudaMemcpy(continuous_rollout_host.data(), continuous_rollout.data, + CB * CH * sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(continuous_env_host.data(), continuous_env.data, + CB * CH * sizeof(float), cudaMemcpyDeviceToHost); + require_equal("continuous environment", continuous_env_host, continuous_rollout_host); + bool retained_non_bf16 = false; + if (USE_BF16) { + for (float action : continuous_rollout_host) { + if (action != round_to_bf16(action)) { + throw std::runtime_error("continuous action bypassed BF16 rounding"); + } + } + } else { + for (float action : continuous_rollout_host) { + retained_non_bf16 |= action != round_to_bf16(action); + } + if (!retained_non_bf16) { + throw std::runtime_error("FP32 continuous probe did not retain extra precision"); + } + } + + std::vector continuous_logprobs_host = + precision_to_host(continuous_logprobs.data, CB); + std::vector recomputed_logprobs(CB); + const float means[CH] = { + USE_BF16 ? round_to_bf16(0.1234567f) : 0.1234567f, + USE_BF16 ? round_to_bf16(-0.2345678f) : -0.2345678f, + }; + const float logstds[CH] = { + USE_BF16 ? round_to_bf16(continuous_logstd[0]) : continuous_logstd[0], + USE_BF16 ? round_to_bf16(continuous_logstd[1]) : continuous_logstd[1], + }; + for (int i = 0; i < CB; ++i) { + float expected_logprob = 0.0f; + for (int h = 0; h < CH; ++h) { + float normalized = (continuous_rollout_host[i * CH + h] - means[h]) / expf(logstds[h]); + expected_logprob += + -0.5f * normalized * normalized - 0.9189385332046727f - logstds[h]; + } + recomputed_logprobs[i] = expected_logprob; + float tolerance = USE_BF16 ? 0.02f : 2e-5f; + if (fabsf(continuous_logprobs_host[i] - expected_logprob) > tolerance) { + throw std::runtime_error("continuous transported action/logprob mismatch"); + } + } + + PrecisionTensor ppo_advantages = {.shape = {1}}; + PrecisionTensor ppo_priorities = {.shape = {1}}; + PrecisionTensor ppo_values = {.shape = {1}}; + PrecisionTensor ppo_returns = {.shape = {1}}; + PrecisionTensor ppo_ratio = {.shape = {1}}; + PrecisionTensor ppo_newvalue = {.shape = {1}}; + FloatTensor ppo_grad_logits = {.shape = {CH}}; + FloatTensor ppo_grad_logstd = {.shape = {CH}}; + FloatTensor ppo_grad_values = {.shape = {1}}; + FloatTensor ppo_adv_mean = {.shape = {1}}; + FloatTensor ppo_adv_var = {.shape = {1}}; + FloatTensor ppo_ent_coef = {.shape = {1}}; + FloatTensor ppo_partials = {.shape = {LOSS_N + 1}}; + Allocator ppo_alloc = {}; + alloc_register(&ppo_alloc, &ppo_advantages); + alloc_register(&ppo_alloc, &ppo_priorities); + alloc_register(&ppo_alloc, &ppo_values); + alloc_register(&ppo_alloc, &ppo_returns); + alloc_register(&ppo_alloc, &ppo_ratio); + alloc_register(&ppo_alloc, &ppo_newvalue); + alloc_register(&ppo_alloc, &ppo_grad_logits); + alloc_register(&ppo_alloc, &ppo_grad_logstd); + alloc_register(&ppo_alloc, &ppo_grad_values); + alloc_register(&ppo_alloc, &ppo_adv_mean); + alloc_register(&ppo_alloc, &ppo_adv_var); + alloc_register(&ppo_alloc, &ppo_ent_coef); + alloc_register(&ppo_alloc, &ppo_partials); + if (alloc_create(&ppo_alloc) != cudaSuccess) { + throw std::runtime_error("continuous PPO probe allocation failed"); + } + float_to_device(ppo_advantages.data, &one, 1); + float_to_device(ppo_priorities.data, &one, 1); + float_to_device(ppo_values.data, &zero, 1); + float_to_device(ppo_returns.data, &zero, 1); + cudaMemcpy(ppo_adv_mean.data, &zero, sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(ppo_adv_var.data, &one, sizeof(float), cudaMemcpyHostToDevice); + cudaMemcpy(ppo_ent_coef.data, &zero, sizeof(float), cudaMemcpyHostToDevice); + PPOKernelArgs ka = { + .grad_logits = ppo_grad_logits.data, + .grad_logstd = ppo_grad_logstd.data, + .grad_values_pred = ppo_grad_values.data, + .logits = continuous_out.data, + .logstd = logstd.data, + .values_pred = continuous_out.data + CH, + .adv_mean = ppo_adv_mean.data, + .adv_var = ppo_adv_var.data, + .act_sizes = continuous_size.data, + .num_atns = CH, + .clip_coef = 0.1f, + .vf_clip_coef = 0.1f, + .vf_coef = 0.5f, + .ent_coef = ppo_ent_coef.data, + .T_seq = 1, + .A_total = CH, + .N = 1, + .logits_stride_n = CH + 1, + .logits_stride_t = CH + 1, + .logits_stride_a = 1, + .values_stride_n = CH + 1, + .values_stride_t = CH + 1, + .is_continuous = true, + }; + PPOGraphArgs ga = { + .out_ratio = ppo_ratio.data, + .out_newvalue = ppo_newvalue.data, + .actions = continuous_rollout.data, + .old_logprobs = continuous_logprobs.data, + .advantages = ppo_advantages.data, + .prio = ppo_priorities.data, + .values = ppo_values.data, + .returns = ppo_returns.data, + }; + ppo_loss_compute<<<1, PPO_THREADS>>>(ppo_partials.data, ka, ga); + std::vector continuous_grad_logits(CH), continuous_grad_logstd(CH); + cudaMemcpy(continuous_grad_logits.data(), ppo_grad_logits.data, CH * sizeof(float), + cudaMemcpyDeviceToHost); + cudaMemcpy(continuous_grad_logstd.data(), ppo_grad_logstd.data, CH * sizeof(float), + cudaMemcpyDeviceToHost); + float expected_ratio = expf(recomputed_logprobs[0] - continuous_logprobs_host[0]); + for (int h = 0; h < CH; ++h) { + float action = continuous_rollout_host[h]; + float variance = expf(2.0f * logstds[h]); + float diff = action - means[h]; + float expected_mean_grad = -expected_ratio * diff / variance; + float expected_std_grad = -expected_ratio * (diff * diff / variance - 1.0f); + if (!isfinite(continuous_grad_logits[h]) || !isfinite(continuous_grad_logstd[h]) || + fabsf(continuous_grad_logits[h] - expected_mean_grad) > 2e-4f || + fabsf(continuous_grad_logstd[h] - expected_std_grad) > 2e-4f) { + throw std::runtime_error("continuous PPO did not consume transported float actions"); + } + } + + alloc_free(&ppo_alloc); + cudaFree(rng_states); + alloc_free(&continuous_alloc); + alloc_free(&alloc); + printf("action transport: passed\n"); +} + struct Im2ColProfile { PrecisionTensor input, col, grad_input; Allocator alloc; @@ -722,6 +1127,8 @@ int main(int argc, char** argv) { profile_fusedscan(BT, T_, H_); if (strcmp(profile, "kernels") == 0 || strcmp(profile, "samplelogits") == 0 || run_all) profile_samplelogits(BR, A_); + if (strcmp(profile, "kernels") == 0 || strcmp(profile, "actiontransport") == 0 || run_all) + test_action_transport(); if (strcmp(profile, "kernels") == 0 || strcmp(profile, "ppoloss") == 0 || run_all) profile_ppoloss(BT, T_, A_); if (strcmp(profile, "kernels") == 0 || strcmp(profile, "im2col") == 0 || run_all) { @@ -738,6 +1145,7 @@ int main(int argc, char** argv) { && strcmp(profile, "logcoeffsvals") != 0 && strcmp(profile, "fusedscan") != 0 && strcmp(profile, "samplelogits") != 0 + && strcmp(profile, "actiontransport") != 0 && strcmp(profile, "ppoloss") != 0 && strcmp(profile, "im2col") != 0 && strcmp(profile, "envspeed") != 0