diff --git a/src/models.cu b/src/models.cu index 2ba14caa43..921c322dc0 100644 --- a/src/models.cu +++ b/src/models.cu @@ -83,6 +83,14 @@ struct EncoderActivations { PrecisionTensor out, saved_input, wgrad_scratch; }; +#ifndef ALIAS_TRAIN_LINEAR_INPUTS +#define ALIAS_TRAIN_LINEAR_INPUTS 1 +#endif + +#ifndef MINGRU_FUSE_HIGHWAY_GRAD_ADD +#define MINGRU_FUSE_HIGHWAY_GRAD_ADD 1 +#endif + // The core of 4.0 is the MinGRU fused scan operation. This allows us to parallelize // training across the sequence dimension and scale to longer sequences __device__ __forceinline__ void log_coeffs_and_values_fwd(float gate, float hidden, @@ -143,6 +151,7 @@ struct PrefixScan { precision_t* input_ptr = nullptr; // (B, T, H) original input before projection (for highway gate) int B = 0, T = 0, H = 0; FloatTensor a_star, s_vals, log_values_buf; + FloatTensor backward_chunks; PrecisionTensor out, next_state; PrecisionTensor grad_combined, grad_state; PrecisionTensor grad_input; // (B, T, H) highway gate gradient w.r.t. input @@ -150,7 +159,7 @@ struct PrefixScan { // Checkpointing trades off partial recomputation for memory bandwidth. #define CHECKPOINT_INTERVAL 4 -__global__ void mingru_scan_forward(PrefixScan scan) { +__device__ __forceinline__ void mingru_scan_forward_impl(PrefixScan scan) { int T_seq = scan.T, H = scan.H, B = scan.B; precision_t* __restrict__ out = scan.out.data; precision_t* __restrict__ next_state = scan.next_state.data; @@ -235,8 +244,16 @@ __global__ void mingru_scan_forward(PrefixScan scan) { next_state[bH + h] = from_float(scan_result); } +__global__ void mingru_scan_forward_ref(PrefixScan scan) { + mingru_scan_forward_impl(scan); +} + +__global__ void mingru_scan_forward(PrefixScan scan) { + mingru_scan_forward_impl(scan); +} + // Reads sparse checkpoints from forward pass, recomputes intermediate values in chunks -__global__ void mingru_scan_backward(PrefixScan scan, +__global__ void mingru_scan_backward_ref(PrefixScan scan, const precision_t* __restrict__ grad_out, const precision_t* __restrict__ grad_next_state) { int T_seq = scan.T, H = scan.H, B = scan.B; @@ -385,6 +402,164 @@ __global__ void mingru_scan_backward(PrefixScan scan, grad_state[state_idx] = from_float(grad_z_0 / to_float(state[state_idx])); } +// Split the reverse scan into independent chunks. In exact arithmetic, +// +// acc_t = alpha_t * acc_{t+1} + u_t +// carry_t = carry_{t+1} + u_t - k_t * acc_t +// +// composes across a chunk as the triangular affine transfer +// +// acc_start = A * acc_end + B +// carry_start = carry_end + D * acc_end + E. +// +// Floating-point addition is not affine, however, and composing those maps +// changes near-cancelled gate gradients. The tiled kernel therefore builds +// chunk-local values in parallel, replays the cheap carries in reference order, +// then applies all chunks in parallel before the block exits. This keeps exact +// fp32 rounding without global synchronization between phases. +#define MINGRU_BACKWARD_CHUNK 4 + +#define MINGRU_BACKWARD_H_TILE 16 +__global__ void mingru_scan_backward_tiled(PrefixScan scan, + const precision_t* __restrict__ grad_out, + const precision_t* __restrict__ grad_next_state) { + int T_seq = scan.T, H = scan.H; + int h_tiles = (H + MINGRU_BACKWARD_H_TILE - 1) / MINGRU_BACKWARD_H_TILE; + int b = blockIdx.x / h_tiles; + int h_tile = blockIdx.x - b * h_tiles; + int chunk = threadIdx.x / MINGRU_BACKWARD_H_TILE; + int lane = threadIdx.x - chunk * MINGRU_BACKWARD_H_TILE; + int h = h_tile * MINGRU_BACKWARD_H_TILE + lane; + bool valid = h < H; + + extern __shared__ float shared[]; + int shared_steps = T_seq * MINGRU_BACKWARD_H_TILE; + float* shared_u = shared; + float* shared_k = shared_u + shared_steps; + float* shared_s = shared_k + shared_steps; + + int bHT = b * H * T_seq; + int cbase = 3 * bHT; + int H3 = 3 * H; + int H2 = 2 * H; + int state_idx = b * H + h; + int out_base = bHT + h; + int buf_base = b * (T_seq + 1) * H + h; + int chunk_start = chunk * MINGRU_BACKWARD_CHUNK + 1; + int chunk_end = min(T_seq, chunk_start + MINGRU_BACKWARD_CHUNK - 1); + + if (valid) { + int checkpoint_idx = buf_base + (chunk_start - 1) * H; + float a_star = scan.a_star.data[checkpoint_idx]; + float s = scan.s_vals.data[checkpoint_idx]; + float log_value = scan.log_values_buf.data[checkpoint_idx]; + for (int t = chunk_start; t <= chunk_end; ++t) { + int t_offset = (t - 1) * H3; + int input_idx = out_base + (t - 1) * H; + int shared_idx = (t - 1) * MINGRU_BACKWARD_H_TILE + lane; + float hidden_val = to_float(scan.combined_ptr[cbase + h + t_offset]); + float gate_val = to_float(scan.combined_ptr[cbase + H + h + t_offset]); + float proj_val = to_float(scan.combined_ptr[cbase + H2 + h + t_offset]); + float log_coeff; + log_coeffs_and_values_fwd( + gate_val, hidden_val, &log_coeff, &log_value); + a_star += log_coeff; + float z = log_value - a_star; + s = logaddexp(s, z); + float scan_result = __expf(a_star + s); + float grad_out_val = to_float(grad_out[input_idx]); + float proj_sigmoid = sigmoid(proj_val); + float grad_scan_from_next = (t == T_seq) + ? to_float(grad_next_state[state_idx]) : 0.0f; + shared_u[shared_idx] = (grad_scan_from_next + + grad_out_val * proj_sigmoid) * scan_result; + shared_k[shared_idx] = __expf(z - s); + shared_s[shared_idx] = s; + + float x_val = to_float(scan.input_ptr[input_idx]); + scan.grad_combined.data[cbase + H2 + h + t_offset] = from_float( + grad_out_val * (scan_result - x_val) + * proj_sigmoid * (1.0f - proj_sigmoid)); + scan.grad_input.data[input_idx] = from_float( + grad_out_val * (1.0f - proj_sigmoid)); + } + } + __syncthreads(); + + if (chunk == 0 && valid) { + float acc = 0.0f; + float carry_grad_a = 0.0f; + float s_next = shared_s[(T_seq - 1) * MINGRU_BACKWARD_H_TILE + lane]; + for (int t = T_seq; t > 0; --t) { + int boundary_chunk = (t - 1) / MINGRU_BACKWARD_CHUNK; + if (t == T_seq || t % MINGRU_BACKWARD_CHUNK == 0) { + float* boundary = &scan.backward_chunks.data[ + (boundary_chunk * scan.B * H + state_idx) * 2]; + boundary[0] = acc; + boundary[1] = carry_grad_a; + } + int shared_idx = (t - 1) * MINGRU_BACKWARD_H_TILE + lane; + float grad_log_h = shared_u[shared_idx]; + float s_t = shared_s[shared_idx]; + float alpha_t = (t == T_seq) ? 0.0f : __expf(s_t - s_next); + if (t == T_seq) { + acc = grad_log_h; + } else { + acc = __fmaf_rn(acc, alpha_t, grad_log_h); + } + float grad_z = acc * shared_k[shared_idx]; + float grad_a = __fsub_rn( + __fadd_rn(grad_log_h, carry_grad_a), grad_z); + carry_grad_a = grad_a; + shared_s[shared_idx] = alpha_t; + s_next = s_t; + } + acc *= __expf(scan.s_vals.data[buf_base] - s_next); + float grad_z_0 = acc; + scan.grad_state.data[state_idx] = from_float( + grad_z_0 / to_float(scan.state_ptr[state_idx])); + } + __syncthreads(); + + if (valid) { + float* boundary = &scan.backward_chunks.data[(chunk * scan.B * H + state_idx) * 2]; + float acc = boundary[0]; + float carry_grad_a = boundary[1]; + for (int t = chunk_end; t >= chunk_start; --t) { + int t_offset = (t - 1) * H3; + int shared_idx = (t - 1) * MINGRU_BACKWARD_H_TILE + lane; + float grad_log_h = shared_u[shared_idx]; + if (t == T_seq) { + acc = grad_log_h; + } else { + acc = __fmaf_rn(acc, shared_s[shared_idx], grad_log_h); + } + float grad_z = acc * shared_k[shared_idx]; + float grad_a = __fsub_rn( + __fadd_rn(grad_log_h, carry_grad_a), grad_z); + carry_grad_a = grad_a; + + float hidden_val = to_float(scan.combined_ptr[cbase + h + t_offset]); + float gate_val = to_float(scan.combined_ptr[cbase + H + h + t_offset]); + float grad_g, grad_h; + log_coeffs_and_values_bwd( + grad_a, grad_z, gate_val, hidden_val, &grad_g, &grad_h); + scan.grad_combined.data[cbase + h + t_offset] = from_float(grad_h); + scan.grad_combined.data[cbase + H + h + t_offset] = from_float(grad_g); + } + } +} + +static void mingru_scan_backward(PrefixScan scan, + const precision_t* grad_out, const precision_t* grad_next_state, cudaStream_t stream) { + int chunks = (scan.T + MINGRU_BACKWARD_CHUNK - 1) / MINGRU_BACKWARD_CHUNK; + int h_tiles = (scan.H + MINGRU_BACKWARD_H_TILE - 1) / MINGRU_BACKWARD_H_TILE; + int threads = chunks * MINGRU_BACKWARD_H_TILE; + int shared_bytes = 3 * scan.T * MINGRU_BACKWARD_H_TILE * sizeof(float); + mingru_scan_backward_tiled<<>>( + scan, grad_out, grad_next_state); +} + __global__ void sum_rows_to_precision_kernel(precision_t* __restrict__ dst, const float* __restrict__ src, int R, int C) { int col = blockIdx.x * blockDim.x + threadIdx.x; @@ -412,7 +587,13 @@ __global__ void assemble_decoder_grad( static PrecisionTensor encoder_forward(void* w, void* activations, PrecisionTensor input, cudaStream_t stream) { EncoderWeights* ew = (EncoderWeights*)w; EncoderActivations* a = (EncoderActivations*)activations; - if (a->saved_input.data) puf_copy(&a->saved_input, &input, stream); + if (numel(a->saved_input.shape) > 1) { +#if ALIAS_TRAIN_LINEAR_INPUTS + a->saved_input.data = input.data; +#else + puf_copy(&a->saved_input, &input, stream); +#endif + } puf_mm(&input, &ew->weight, &a->out, stream); return a->out; } @@ -446,7 +627,9 @@ static void encoder_reg_train(void* w, void* activations, Allocator* acts, Alloc .wgrad_scratch = {.shape = {ew->out_dim, ew->in_dim}}, }; alloc_register(acts,&a->out); +#if !ALIAS_TRAIN_LINEAR_INPUTS alloc_register(acts,&a->saved_input); +#endif alloc_register(grads,&a->wgrad_scratch); } @@ -485,8 +668,12 @@ struct DecoderActivations { static PrecisionTensor decoder_forward(void* w, void* activations, PrecisionTensor input, cudaStream_t stream) { DecoderWeights* dw = (DecoderWeights*)w; DecoderActivations* a = (DecoderActivations*)activations; - if (a->saved_input.data) { + if (numel(a->saved_input.shape) > 1) { +#if ALIAS_TRAIN_LINEAR_INPUTS + a->saved_input.data = input.data; +#else puf_copy(&a->saved_input, &input, stream); +#endif } puf_mm(&input, &dw->weight, &a->out, stream); return a->out; @@ -524,7 +711,9 @@ static void decoder_reg_train(void* w, void* activations, Allocator* acts, Alloc .logstd_scratch = {.shape = {1, dw->output_dim}}, }; alloc_register(acts,&a->out); +#if !ALIAS_TRAIN_LINEAR_INPUTS alloc_register(acts,&a->saved_input); +#endif alloc_register(acts,&a->grad_out); alloc_register(acts,&a->grad_input); alloc_register(grads,&a->wgrad_scratch); @@ -626,6 +815,7 @@ static void mingru_reg_train(void* w, void* activations, Allocator* acts, Alloca MinGRUWeights* m = (MinGRUWeights*)w; MinGRUActivations* a = (MinGRUActivations*)activations; int H = m->hidden, TT = m->horizon, B = B_TT / TT; + int backward_chunks = (TT + MINGRU_BACKWARD_CHUNK - 1) / MINGRU_BACKWARD_CHUNK; a->num_layers = m->num_layers; a->saved_inputs = (PrecisionTensor*)calloc(m->num_layers, sizeof(PrecisionTensor)); a->scan_bufs = (PrefixScan*)calloc(m->num_layers, sizeof(PrefixScan)); @@ -633,7 +823,9 @@ static void mingru_reg_train(void* w, void* activations, Allocator* acts, Alloca a->wgrad_scratch = (PrecisionTensor*)calloc(m->num_layers, sizeof(PrecisionTensor)); a->grad_input_buf = {.shape = {B_TT, H}}; a->grad_next_state = {.shape = {B, 1, H}}; +#if !MINGRU_FUSE_HIGHWAY_GRAD_ADD alloc_register(acts,&a->grad_input_buf); +#endif alloc_register(acts,&a->grad_next_state); for (int i = 0; i < m->num_layers; i++) { a->scan_bufs[i] = { @@ -641,6 +833,7 @@ static void mingru_reg_train(void* w, void* activations, Allocator* acts, Alloca .a_star = {.shape = {B, TT + 1, H}}, .s_vals = {.shape = {B, TT + 1, H}}, .log_values_buf = {.shape = {B, TT + 1, H}}, + .backward_chunks = {.shape = {backward_chunks, B, H, 2}}, .out = {.shape = {B, TT, H}}, .next_state = {.shape = {B, 1, H}}, .grad_combined = {.shape = {B, TT, 3 * H}}, @@ -650,13 +843,16 @@ static void mingru_reg_train(void* w, void* activations, Allocator* acts, Alloca a->saved_inputs[i] = {.shape = {B, TT, H}}; a->combined_bufs[i] = {.shape = {B_TT, 3 * H}}; a->wgrad_scratch[i] = {.shape = {3 * H, H}}; +#if !ALIAS_TRAIN_LINEAR_INPUTS alloc_register(acts,&a->saved_inputs[i]); +#endif alloc_register(acts,&a->combined_bufs[i]); alloc_register(acts,&a->scan_bufs[i].out); alloc_register(acts,&a->scan_bufs[i].next_state); alloc_register(acts,&a->scan_bufs[i].a_star); alloc_register(acts,&a->scan_bufs[i].s_vals); alloc_register(acts,&a->scan_bufs[i].log_values_buf); + alloc_register(acts,&a->scan_bufs[i].backward_chunks); alloc_register(acts,&a->scan_bufs[i].grad_combined); alloc_register(acts,&a->scan_bufs[i].grad_state); alloc_register(acts,&a->scan_bufs[i].grad_input); @@ -724,7 +920,11 @@ static PrecisionTensor mingru_forward_train(void* w, PrecisionTensor x, Precisio MinGRUActivations* a = (MinGRUActivations*)activations; int B = x.shape[0]; for (int i = 0; i < m->num_layers; i++) { +#if ALIAS_TRAIN_LINEAR_INPUTS + a->saved_inputs[i].data = x.data; +#else puf_copy(&a->saved_inputs[i], &x, stream); +#endif PrecisionTensor state_i = mingru_state_layer(m, state, i); puf_mm(&x, &m->weights[i], &a->combined_bufs[i], stream); a->scan_bufs[i].combined_ptr = a->combined_bufs[i].data; @@ -741,14 +941,19 @@ static PrecisionTensor mingru_backward(void* w, PrecisionTensor grad, void* acti MinGRUActivations* a = (MinGRUActivations*)activations; for (int i = m->num_layers - 1; i >= 0; i--) { PrefixScan& scan = a->scan_bufs[i]; - mingru_scan_backward<<>>( - scan, grad.data, a->grad_next_state.data); + mingru_scan_backward(scan, grad.data, a->grad_next_state.data, stream); puf_mm_tn(&scan.grad_combined, &a->saved_inputs[i], &a->wgrad_scratch[i], stream); +#if MINGRU_FUSE_HIGHWAY_GRAD_ADD + puf_addmm_nn(&scan.grad_combined, &m->weights[i], &scan.grad_input, + 1.0f, 1.0f, stream); + grad = scan.grad_input; +#else puf_mm_nn(&scan.grad_combined, &m->weights[i], &a->grad_input_buf, stream); int n = numel(scan.grad_input.shape); add_kernel<<>>( a->grad_input_buf.data, scan.grad_input.data, n); grad = a->grad_input_buf; +#endif } return grad; } diff --git a/src/pufferlib.cu b/src/pufferlib.cu index 4af2e2b640..193447b48d 100644 --- a/src/pufferlib.cu +++ b/src/pufferlib.cu @@ -204,8 +204,9 @@ void register_ppo_buffers(PPOBuffersPuf& bufs, Allocator* alloc, int N, int T, i // Prioritized replay over single-epoch data. These kernels are // the least cleaned because we will likely have a better method in 5.0 +#define PRIO_CDF_BLOCK_SIZE 1024 struct PrioBuffers { - FloatTensor prio_probs, cdf, mb_prio; + FloatTensor prio_probs, cdf, cdf_block_sums, mb_prio; IntTensor idx; }; @@ -213,11 +214,13 @@ void register_prio_buffers(PrioBuffers& bufs, Allocator* alloc, int B, int minib bufs = (PrioBuffers){ .prio_probs = {.shape = {B}}, .cdf = {.shape = {B}}, + .cdf_block_sums = {.shape = {(B + PRIO_CDF_BLOCK_SIZE - 1) / PRIO_CDF_BLOCK_SIZE}}, .mb_prio = {.shape = {minibatch_segments}}, .idx = {.shape = {minibatch_segments}}, }; alloc_register(alloc, &bufs.prio_probs); alloc_register(alloc, &bufs.cdf); + alloc_register(alloc, &bufs.cdf_block_sums); alloc_register(alloc, &bufs.idx); alloc_register(alloc, &bufs.mb_prio); } @@ -1196,7 +1199,7 @@ __global__ void compute_prio_imp_weights( } } -__global__ void build_cdf( +__global__ void build_cdf_serial( float* __restrict__ cdf, const float* __restrict__ probs, int B) { if (blockIdx.x == 0 && threadIdx.x == 0) { float cum = 0.0f; @@ -1207,6 +1210,75 @@ __global__ void build_cdf( } } +template +__device__ __forceinline__ float prio_inclusive_scan(float value, float* warp_sums) { + int lane = threadIdx.x % PRIO_WARP_SIZE; + int warp_id = threadIdx.x / PRIO_WARP_SIZE; + + for (int offset = 1; offset < PRIO_WARP_SIZE; offset *= 2) { + float other = __shfl_up_sync(PRIO_FULL_MASK, value, offset); + if (lane >= offset) { + value += other; + } + } + if (lane == PRIO_WARP_SIZE - 1) { + warp_sums[warp_id] = value; + } + __syncthreads(); + + if (warp_id == 0) { + float warp_sum = lane < num_warps ? warp_sums[lane] : 0.0f; + for (int offset = 1; offset < PRIO_WARP_SIZE; offset *= 2) { + float other = __shfl_up_sync(PRIO_FULL_MASK, warp_sum, offset); + if (lane >= offset) { + warp_sum += other; + } + } + if (lane < num_warps) { + warp_sums[lane] = warp_sum; + } + } + __syncthreads(); + + if (warp_id > 0) { + value += warp_sums[warp_id - 1]; + } + return value; +} + +__global__ void build_cdf_scan_blocks( + float* __restrict__ cdf, float* __restrict__ block_sums, + const float* __restrict__ probs, int B) { + __shared__ float warp_sums[PRIO_CDF_BLOCK_SIZE / PRIO_WARP_SIZE]; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + float value = idx < B ? probs[idx] : 0.0f; + value = prio_inclusive_scan(value, warp_sums); + if (idx < B) { + cdf[idx] = value; + } + if (threadIdx.x == PRIO_CDF_BLOCK_SIZE - 1) { + block_sums[blockIdx.x] = value; + } +} + +__global__ void scan_cdf_block_sums(float* block_sums, int num_blocks) { + __shared__ float warp_sums[PRIO_NUM_WARPS]; + int tx = threadIdx.x; + float value = tx < num_blocks ? block_sums[tx] : 0.0f; + value = prio_inclusive_scan(value, warp_sums); + if (tx < num_blocks) { + block_sums[tx] = value; + } +} + +__global__ void add_cdf_block_offsets( + float* cdf, const float* __restrict__ block_sums, int B) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < B && blockIdx.x > 0) { + cdf[idx] += block_sums[blockIdx.x - 1]; + } +} + __global__ void advance_rng_offset(int64_t* __restrict__ offset_ptr, int64_t delta) { if (blockIdx.x == 0 && threadIdx.x == 0) { *offset_ptr += delta; @@ -1247,7 +1319,17 @@ void prio_replay_cuda(PrecisionTensor& advantages, float prio_alpha, compute_prio_normalize<<<1, PRIO_BLOCK_SIZE, 0, stream>>>( bufs.prio_probs.data, B); //int block = fmaxf(((minibatch_segments + 31) / 32) * 32, 32); - build_cdf<<<1, 1, 0, stream>>>(bufs.cdf.data, bufs.prio_probs.data, B); + int cdf_blocks = (B + PRIO_CDF_BLOCK_SIZE - 1) / PRIO_CDF_BLOCK_SIZE; + // scan_cdf_block_sums is a single block: B caps at PRIO_CDF_BLOCK_SIZE^2 + assert(cdf_blocks <= PRIO_CDF_BLOCK_SIZE && "build_cdf: B too large for two-level scan"); + build_cdf_scan_blocks<<>>( + bufs.cdf.data, bufs.cdf_block_sums.data, bufs.prio_probs.data, B); + if (cdf_blocks > 1) { + scan_cdf_block_sums<<<1, PRIO_BLOCK_SIZE, 0, stream>>>( + bufs.cdf_block_sums.data, cdf_blocks); + add_cdf_block_offsets<<>>( + bufs.cdf.data, bufs.cdf_block_sums.data, B); + } int threads = 256; int blocks = (minibatch_segments + threads - 1) / threads; multinomial_sample<<>>(