diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 2d0e52ee1..0046fce21 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -540,6 +540,25 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: torch::Tensor tile_begin, torch::Tensor tile_end, torch::Tensor seg_first_tile); +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end); +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile); +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end); +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype); +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index); +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index); #endif // PyBind11 Module Registration @@ -848,5 +867,28 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate", py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile")); + m.def("h3_rmsnorm_backward_partials", &h3_rmsnorm_backward_partials, + "WS1 RMSNorm dweight / table-gradient tile partials over explicit row lists (SP)", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), py::arg("shift"), + py::arg("scale"), py::arg("index"), py::arg("dw_rows"), py::arg("dw_begin"), + py::arg("dw_end"), py::arg("seg_rows") = py::none(), py::arg("seg_begin") = py::none(), + py::arg("seg_end") = py::none()); + m.def("h3_rmsnorm_fold_partials", &h3_rmsnorm_fold_partials, + "WS1 ascending folds of RMSNorm dweight and per-segment table-gradient partials", + py::arg("dw_partial"), py::arg("weight"), py::arg("seg_partial") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_grad_partials", &h3_gate_grad_partials, + "WS1 d_gate tile partials over explicit row lists (SP)", py::arg("grad"), py::arg("y"), + py::arg("rows"), py::arg("tile_begin"), py::arg("tile_end")); + m.def("h3_gate_grad_fold", &h3_gate_grad_fold, + "WS1 per-segment ascending fold of d_gate partials, cast once", py::arg("partial"), + py::arg("seg_first_tile"), py::arg("dtype")); + m.def("h3_rmsnorm_backward_dx", &h3_rmsnorm_backward_dx, + "Row-local dx of the WS1 RMSNorm(+modulation) backward", py::arg("grad"), py::arg("x"), + py::arg("weight"), py::arg("rstd"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_gate_residual_backward_dy", &h3_gate_residual_backward_dy, + "Row-local d_y = grad * gate[index] of the WS1 gated-residual backward", + py::arg("grad"), py::arg("gate"), py::arg("index")); #endif } diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 1f635796d..3290eaf5d 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -212,9 +212,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(), "grad must be a contiguous (C, S, H) CUDA tensor"); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->dim() == 1 && t->is_contiguous() && + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && + t->dim() == 1 && t->is_contiguous() && t->scalar_type() == at::kLong, - "tile metadata must be contiguous int64 tensors on ", grad.device()); + "tile metadata must be contiguous 1-D int64 tensors on grad device"); } TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); @@ -225,9 +226,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); - if (num_rows == 0) return out; // a zero-sized grid is not a valid launch auto partial = torch::empty({std::max(tiles, 1), width}, grad.options().dtype(at::kFloat)); auto stream = at::cuda::getCurrentCUDAStream(); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 81051fcc6..79c5e64e4 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -16,6 +16,8 @@ // 16-bit values is exact in FP32, so this equals the eager VJP bitwise); // d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented // sum over positions sorted stably by row, in fixed tiles folded in order. +// The tile partials and the fold are also bound separately for sequence +// parallel callers (each tile computed where its rows live, folded in order). #include #include @@ -262,3 +264,96 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dy, dgate}; } + +// Sequence parallel: the WS1 d_gate tile partials over explicit row lists of +// (grad, y) rows -> (tiles, N) FP32. +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + for (const auto* t : {&rows, &tile_begin, &tile_end}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on grad's device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin and tile_end differ in length"); + const int64_t n = grad.size(1); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); + if (tiles == 0) return partial; + const int threads = 256; + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + // The row and tile lists are caller-built: every row must be a row of grad and every tile must + // lie inside the list before the kernel reads through them. + h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), rows.data_ptr(), rows.numel(), + grad.size(0), tile_begin.data_ptr(), tile_end.data_ptr(), + tiles, nullptr, 0, 0); + H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(y.data_ptr()), + rows.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), + partial.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// The WS1 d_gate fold: each segment's tiles in ascending order, cast once to dtype. +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat, + "partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(seg_first_tile.is_cuda() && seg_first_tile.device() == partial.device() && + seg_first_tile.scalar_type() == at::kLong && seg_first_tile.dim() == 1 && + seg_first_tile.is_contiguous(), + "seg_first_tile must be a contiguous int64 CUDA tensor on partial's device"); + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); + const int64_t n = partial.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + const c10::cuda::CUDAGuard guard(partial.device()); + auto out = torch::empty({segments, n}, partial.options().dtype(dtype)); + if (segments == 0) return out; // a zero-sized grid is not a valid launch + const int threads = 256; + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile.data_ptr(), segments, partial.size(0)); + H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Row-local part of the WS1 backward only: d_y = cast(grad * gate[index]). +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + const Gate g = make_gate(gate, index, grad); + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const c10::cuda::CUDAGuard guard(grad.device()); + auto dy = torch::empty_like(grad); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward_dy", [&] { + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + reinterpret_cast(grad.data_ptr()), nullptr, reinterpret_cast(dy.data_ptr()), + rows, n, g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dy; +} diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index dc4230fdb..3aec2ac47 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -26,6 +26,10 @@ // autograd's RMSNorm backward, which forms rstd^3 differently. // dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order // dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) +// +// The tile partials and the folds are also bound separately (`*_partials`, +// `*_fold_*`) so a sequence-parallel caller can compute each WS1 tile on the +// rank that holds its rows and fold the gathered partials in WS1 order. #include #include @@ -113,9 +117,9 @@ struct Modulation { const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr const void* scale; int64_t row_stride; + int64_t num_rows; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] - int64_t num_rows; }; template @@ -195,18 +199,25 @@ __global__ void __launch_bounds__(32 * kWarps) } } -// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd. +// Tiles are the fixed 256-row blocks, or (sequence parallel) explicit row lists +// rows[tile_begin[tile] .. tile_end[tile]) that hold the same rows in the same order. template __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* __restrict__ x, const float* __restrict__ rstd, float* __restrict__ partial, int64_t rows, int n, - Modulation mod) { + Modulation mod, + const int64_t* __restrict__ row_list = nullptr, + const int64_t* __restrict__ tile_begin = nullptr, + const int64_t* __restrict__ tile_end = nullptr) { const int j = blockIdx.y * blockDim.x + threadIdx.x; if (j >= n) return; - const int64_t r0 = static_cast(blockIdx.x) * kRowTile; - const int64_t r1 = min(r0 + kRowTile, rows); + const int64_t tile = blockIdx.x; + const int64_t p0 = row_list != nullptr ? tile_begin[tile] : tile * kRowTile; + const int64_t p1 = row_list != nullptr ? tile_end[tile] : min(p0 + kRowTile, rows); float acc = 0.0f; - for (int64_t r = r0; r < r1; ++r) { + for (int64_t p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; // Read unchecked: this loop is the kernel's whole cost, and an assert in it slows the backward by // about 11%. The entry points run h3_tiles::check_index_range over the S index values first. const T* scale = nullptr; @@ -279,7 +290,7 @@ Modulation make_modulation(const c10::optional& shift, int64_t n) { // No modulation: kernels test `index == nullptr`. seq = 1 (not 0) keeps `r % seq` defined for // any code that computes it before looking at the index. - Modulation mod{nullptr, nullptr, 0, nullptr, 1, 0}; + Modulation mod{nullptr, nullptr, 0, 0, nullptr, 1}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -308,9 +319,9 @@ Modulation make_modulation(const c10::optional& shift, mod.shift = sh.data_ptr(); mod.scale = sc.data_ptr(); mod.row_stride = sh.stride(0); + mod.num_rows = sh.size(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); - mod.num_rows = sh.size(0); return mod; } @@ -472,3 +483,168 @@ std::vector h3_rmsnorm_backward( } return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } + +namespace { + +void check_tiles(const torch::Tensor& rows, const torch::Tensor& begin, const torch::Tensor& end, + const torch::Tensor& like) { + for (const auto* t : {&rows, &begin, &end}) { + TORCH_CHECK(t->is_cuda() && t->device() == like.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on x's device"); + } + TORCH_CHECK(begin.numel() == end.numel(), "tile_begin and tile_end differ in length"); +} + +} // namespace + +// Sequence parallel: the dweight and (with modulation) table-gradient tile +// partials of the WS1 backward, for explicit tiles over the rows of x. `index` +// gives each row's table row (one entry per row of x). Returns +// {dweight_partial (tiles, N)} or {dweight_partial, table_partial (seg_tiles, 2N)}. +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.dim() == 1 && rstd.is_contiguous() && rstd.size(0) == x.size(0), + "rstd must be contiguous float32 (M,) statistics on x's device"); + check_tiles(dw_rows, dw_begin, dw_end, x); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + TORCH_CHECK(mod.index == nullptr || mod.seq == rows, "index needs one entry per row of x"); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + const int64_t tiles = dw_begin.numel(); + TORCH_CHECK(2 * n <= std::numeric_limits::max() && (2 * n + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + // The row and tile lists are caller-built: every row must be a row of x and every tile must lie + // inside its list before the kernels read through them. + h3_tiles::check_lists(stream, dw_rows.data_ptr(), dw_rows.numel(), rows, + dw_begin.data_ptr(), dw_end.data_ptr(), tiles, nullptr, 0, + 0); + // The dweight partials read the modulation index inside their row loop: check its values first + // (this entry point can be called without the dx entry point having run). + if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.num_rows); + if (tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, + 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + rstd.data_ptr(), dw_partial.data_ptr(), rows, static_cast(n), mod, + dw_rows.data_ptr(), dw_begin.data_ptr(), dw_end.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + if (mod.index == nullptr) return {dw_partial}; + + TORCH_CHECK(seg_rows.has_value() && seg_begin.has_value() && seg_end.has_value(), + "modulated partials need the segment tiles"); + check_tiles(*seg_rows, *seg_begin, *seg_end, x); + const int64_t seg_tiles = seg_begin->numel(); + const int64_t width = 2 * n; + h3_tiles::check_lists(stream, seg_rows->data_ptr(), seg_rows->numel(), rows, + seg_begin->data_ptr(), seg_end->data_ptr(), seg_tiles, + nullptr, 0, 0); + auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_partials", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + seg_rows->data_ptr(), seg_begin->data_ptr(), + seg_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dw_partial, seg_partial}; +} + +// The WS1 folds: dweight = ascending fold of all tiles, cast to weight's dtype; +// table rows = ascending fold of each segment's tiles, FP32 (R, 2N). +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile) { + TORCH_CHECK(dw_partial.is_cuda() && dw_partial.dim() == 2 && dw_partial.is_contiguous() && + dw_partial.scalar_type() == at::kFloat, + "dw_partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(weight.is_cuda() && weight.device() == dw_partial.device() && weight.dim() == 1 && + weight.size(0) == dw_partial.size(1), + "weight must be (N,) on dw_partial's device"); + const int64_t n = dw_partial.size(1); + const c10::cuda::CUDAGuard guard(dw_partial.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + auto dweight = torch::empty_like(weight); + H3_DISPATCH(weight.scalar_type(), "h3_rmsnorm_fold_partials", [&] { + fold_tiles_kernel<<>>( + dw_partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), + dw_partial.size(0), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!seg_partial.has_value()) return {dweight}; + + TORCH_CHECK(seg_first_tile.has_value() && seg_first_tile->is_cuda() && + seg_first_tile->device() == dw_partial.device() && + seg_first_tile->scalar_type() == at::kLong && seg_first_tile->is_contiguous(), + "segment fold needs seg_first_tile (int64, on dw_partial's device)"); + TORCH_CHECK(seg_partial->is_cuda() && seg_partial->device() == dw_partial.device() && + seg_partial->dim() == 2 && seg_partial->is_contiguous() && + seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, + "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + TORCH_CHECK(seg_first_tile->numel() >= 1, "seg_first_tile must hold R + 1 entries"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + TORCH_CHECK(width <= std::numeric_limits::max() && (width + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (segments > 0) { // a zero-sized grid is not a valid launch + h3_tiles::check_lists(stream, nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile->data_ptr(), segments, seg_partial->size(0)); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, + stream>>>(seg_partial->data_ptr(), + seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +// Row-local part of the WS1 backward only: dx (sequence-parallel callers +// compute the cross-row reductions with h3_rmsnorm_backward_partials). +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.dim() == 1 && rstd.is_contiguous() && rstd.size(0) == x.size(0), + "rstd must be contiguous float32 (M,) statistics on x's device"); + const Modulation mod = make_modulation(shift, scale, index, x, x.size(1)); + const c10::cuda::CUDAGuard guard(x.device()); + auto dx = torch::empty_like(x); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_dx", [&] { + rmsnorm_modulate_dx_kernel<<(x.size(0)), dim3(32, kWarps), 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + reinterpret_cast(dx.data_ptr()), static_cast(x.size(1)), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dx; +} diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 24c473e5f..17901f91e 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -79,6 +79,48 @@ inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +// Sequence-parallel entry points take explicit lists. `rows` (n_rows entries, each a row of an +// M-row activation, so in [0, row_limit)) and the tiles [tile_begin, tile_end) over it, and/or the +// segment boundaries seg_first_tile (segments + 1 entries, in [0, seg_limit], non-decreasing). A +// null pointer skips its part. Same device-side asserts, same single small launch. +static __global__ void check_lists_kernel(const int64_t* __restrict__ rows, int64_t n_rows, + int64_t row_limit, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, int64_t tiles, + const int64_t* __restrict__ seg_first_tile, + int64_t segments, int64_t seg_limit) { + const int64_t first_thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + if (rows != nullptr) { + for (int64_t i = first_thread; i < n_rows; i += stride) { + CUDA_KERNEL_ASSERT(rows[i] >= 0 && rows[i] < row_limit); + } + } + if (tile_begin != nullptr) { + for (int64_t i = first_thread; i < tiles; i += stride) { + CUDA_KERNEL_ASSERT(tile_begin[i] >= 0 && tile_begin[i] <= tile_end[i] && + tile_end[i] <= n_rows); + } + } + if (seg_first_tile != nullptr) { + for (int64_t i = first_thread; i <= segments; i += stride) { + CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= seg_limit); + CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + } + } +} + +inline void check_lists(cudaStream_t stream, const int64_t* rows, int64_t n_rows, int64_t row_limit, + const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, + const int64_t* seg_first_tile, int64_t segments, int64_t seg_limit) { + // At least one block, so an empty list is a no-op and not an invalid launch. + const int64_t items = std::max({n_rows, tiles, segments + 1, int64_t{1}}); + const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); + check_lists_kernel<<>>(rows, n_rows, row_limit, tile_begin, tile_end, + tiles, seg_first_tile, segments, seg_limit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + // Table-row indices read inside a kernel's loop: `n` entries, each a row of a table with `limit` // rows. The assert would cost the loop its scheduling (the same 6-12% as above, measured on the // rmsnorm dweight partials), so the entry point checks the S index values once, here, before it diff --git a/docs/.nav.yml b/docs/.nav.yml index 9481a07d1..fcf820f03 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -41,6 +41,7 @@ nav: - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md - operators/h3-tp-adaln-3mod.md + - operators/h3-sp-norm-adaln.md - Developer Guide: - contributing/README.md - Contributor Guide: contributing/contributor-guide.md diff --git a/docs/operators/README.md b/docs/operators/README.md index f0ea0a86d..89dcdee6d 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -40,5 +40,6 @@ Every operator page should include: - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) - [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) +- [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) - [H3 Flash Attention (optional fast backend)](h3-flash-attention.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 1854f32d3..029211e2c 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -62,7 +62,7 @@ python benchmarks/models/benchmark_h3_conditioning.py --op final_adaln_out B200, B = 1, T = 3, pinned `norm_out` weights. Backward timings exclude input and forward-graph setup: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.19 ms | 0.21 ms | 1.37 ms | 0.83 ms | | 32768 | 0.42 ms | 0.82 ms | 2.53 ms | 4.45 ms | @@ -71,6 +71,11 @@ and forward-graph setup: At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + ## Evidence ![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/figure.png) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..72755e37f --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,117 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `ecde5bb`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. + +## Tests + +```bash +python -m pytest tests/distributed/sp/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-sp-norm-adaln-b200/report.json +``` + +The rank-backward tests run SP 2–8 on one GPU, with ranks as threads around the plain backward +functions. They cover odd `S`, `B = 2`, shards smaller than a tile, block and interleaved packing, +and bf16 and fp32, and compare every rank against WS1. The NCCL tests run the autograd region end +to end and need as many GPUs as ranks. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- With interleaved packing, up to a full shard of rows can move in the backward. The result is + still exact. +- Attention and FFN are other rows. Until they land, the region uses a seeded stand-in for the + sublayer output. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index c910add64..d2a740225 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -516,3 +516,81 @@ def h3_gate_residual_backward( checked on the device before use. """ ... +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Sequence parallel: FP32 tile partials of the rmsnorm weight (and table) gradients. + + Same operands as ``h3_rmsnorm_backward`` plus explicit int64 row lists: ``dw_rows`` with tiles + ``[dw_begin, dw_end)`` over it, and with modulation ``seg_rows`` with ``[seg_begin, seg_end)`` + tiles; ``index`` then has one entry per row of ``x``. Returns ``[dw_partial (tiles, N)]`` or + ``[dw_partial, seg_partial (seg_tiles, 2N)]``. Lists and index values are checked on the device + before use; ``h3_rmsnorm_fold_partials`` folds the gathered partials. + """ + ... + +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Fold gathered ``h3_rmsnorm_backward_partials`` in tile order. + + Return ``[dweight]`` in ``weight``'s dtype, or ``[dweight, d_shift, d_scale]`` (FP32 ``(R, N)``) + when ``seg_partial`` and ``seg_first_tile`` (``R + 1`` entries) are given. + """ + ... + +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: + """Sequence parallel: FP32 ``(tiles, N)`` partials of ``sum(grad * y)`` over tiles of ``rows``. + + ``rows`` is an int64 list of rows of ``grad``/``y`` and ``[tile_begin, tile_end)`` cut it into + tiles; the lists are checked on the device before use. + """ + ... + +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + """Fold gathered ``h3_gate_grad_partials`` per table row; return ``(R, N)`` in ``dtype``. + + ``seg_first_tile`` has ``R + 1`` entries into the tile axis of ``partial``, checked on the device. + """ + ... + +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: + """The row-local ``dx`` of ``h3_rmsnorm_backward`` alone, in ``x``'s dtype (sequence parallel).""" + ... + +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: + """``d_y = grad * gate[index[r % S]]``, the exact part of ``h3_gate_residual_backward``.""" + ... diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py new file mode 100644 index 000000000..80cbad474 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sequence-parallel H3 RMSNorm / AdaLN modulation / gated residual (RFC #420 ``sp_norm_adaln``). + +Ownership: rank ``r`` of ``sp`` holds packed positions ``[r * S // sp, (r + 1) * S // sp)`` +of every batch item. The full ``(S,)`` row index is replicated; it is small, and every +rank derives the same reduction plan from it. + +Forward and the row-local gradients (``dx``, ``d_sublayer``, ``d_residual``) are the WS1 +kernels on the local rows. Rows are independent, so they are byte-equal by construction. + +The cross-row reductions (``d_norm_weight``, ``d_shift``/``d_scale``, ``d_gate``) are WS1 +two-level folds: FP32 sums over fixed 256-element tiles, then an ascending fold of the +tiles. Every rank builds the global WS1 tile list. A tile is computed by the rank that +holds its first row; the tile's rows from other ranks are all-gathered beforehand (only +rows of tiles that straddle a shard boundary move). The tile partials are all-gathered, +put back in WS1 tile order, and every rank runs the WS1 fold. The partial and fold kernels +are the WS1 kernels, so every gradient equals WS1 on every rank. The only collectives are +rank-ordered all-gathers (copies). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + BACKWARD_TILE, + _segment_tiles, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +_SYMBOLS = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward_dx", + "h3_rmsnorm_backward_partials", + "h3_rmsnorm_fold_partials", + "h3_gate_residual_forward", + "h3_gate_residual_backward_dy", + "h3_gate_grad_partials", + "h3_gate_grad_fold", +) +BACKWARD_IMPL = "row_local_ws1+gather_straddling_rows+ws1_tile_partials+gather+ws1_fold" + + +def sp_norm_adaln_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +@dataclass(frozen=True) +class SPRowLayout: + """Which packed positions each sequence-parallel rank holds.""" + + seq_len: int + batch: int + sp: int + rank: int + + def bounds(self, rank: int) -> tuple[int, int]: + return rank * self.seq_len // self.sp, (rank + 1) * self.seq_len // self.sp + + @property + def lo(self) -> int: + return self.bounds(self.rank)[0] + + @property + def hi(self) -> int: + return self.bounds(self.rank)[1] + + @property + def local_len(self) -> int: + return self.hi - self.lo + + +def sp_row_layout(seq_len: int, batch: int, sp: int, rank: int) -> SPRowLayout: + if sp < 1 or not 0 <= rank < sp: + raise ValueError(f"need 0 <= rank < sp, got rank={rank}, sp={sp}") + if batch < 1 or seq_len < sp: + raise ValueError(f"every rank needs a row: seq_len={seq_len}, sp={sp}, batch={batch}") + return SPRowLayout(seq_len, batch, sp, rank) + + +def _ordinal_by_group(group: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per-group counts, and each element's position within its group (in input order).""" + + counts = torch.bincount(group, minlength=groups) + order = torch.argsort(group, stable=True) + ordinal = torch.empty_like(group) + first = torch.cumsum(counts, 0) - counts + ordinal[order] = torch.arange(group.numel(), device=group.device) - first[group[order]] + return counts, ordinal + + +class _Plan: + """Global WS1 tiles of one backward, and this rank's part in computing them. + + ``families`` are ``(rows, tile_begin, tile_end)`` over global flattened rows + ``b * S + s``; consecutive tiles cover ``rows`` in order. Every rank builds the same + plan from replicated metadata, so collective calls and shapes agree. + """ + + def __init__(self, layout: SPRowLayout, families, device) -> None: + self.layout, sp = layout, layout.sp + self._his = torch.tensor([layout.bounds(r)[1] for r in range(sp)], device=device) + + # Rows of a tile computed by another rank: the only activation rows that move. + tile_owner, shared = [], [] + for rows, begin, end in families: + owner = self.owner(rows[begin]) + tile_owner.append(owner) + shared.append(rows[torch.repeat_interleave(owner, end - begin) != self.owner(rows)]) + shared = torch.unique(torch.cat(shared)) + shared_owner = self.owner(shared) + counts, slot = _ordinal_by_group(shared_owner, sp) + self.max_send = int(counts.max()) if shared.numel() else 0 + + # Buffer row of every global row this rank reads: its own rows first, then + # every rank's send block (padded to max_send) in rank order. + own = self.global_rows(layout.rank) + buffer_index = torch.full((layout.batch * layout.seq_len,), -1, device=device) + buffer_index[own] = torch.arange(own.numel(), device=device) + foreign = shared_owner != layout.rank + buffer_index[shared[foreign]] = ( + own.numel() + shared_owner[foreign] * self.max_send + slot[foreign] + ) + mine = ~foreign + self.send_local = self.local_index(shared[mine][torch.argsort(slot[mine])]) + self.recv_global = torch.zeros(sp * self.max_send, dtype=torch.long, device=device) + self.recv_global[shared_owner * self.max_send + slot] = shared # padding: never read + + self.families = [ + self._my_tiles(rows, begin, end, owner, buffer_index) + for (rows, begin, end), owner in zip(families, tile_owner) + ] + + def owner(self, rows: torch.Tensor) -> torch.Tensor: + return torch.searchsorted(self._his, rows % self.layout.seq_len, right=True) + + def global_rows(self, rank: int) -> torch.Tensor: + lo, hi = self.layout.bounds(rank) + s = torch.arange(lo, hi, device=self._his.device) + b = torch.arange(self.layout.batch, device=self._his.device) + return (b[:, None] * self.layout.seq_len + s[None, :]).reshape(-1) + + def local_index(self, rows: torch.Tensor) -> torch.Tensor: + b, s = rows // self.layout.seq_len, rows % self.layout.seq_len + return b * self.layout.local_len + s - self.layout.lo + + def _my_tiles(self, rows, begin, end, owner, buffer_index): + mine = torch.nonzero(owner == self.layout.rank).flatten() + lengths = (end - begin)[mine] + my_end = torch.cumsum(lengths, 0) + my_begin = my_end - lengths + elems = torch.repeat_interleave(begin[mine] - my_begin, lengths) + elems = elems + torch.arange(elems.numel(), device=rows.device) + counts, ordinal = _ordinal_by_group(owner, self.layout.sp) + width = int(counts.max()) if owner.numel() else 0 + return { + "tiles": (buffer_index[rows[elems]].contiguous(), my_begin, my_end), + "width": width, # tiles per rank in the padded partial gather + "gather_index": owner * width + ordinal, # WS1 tile order within that gather + } + + def exchange(self, collective, *tensors: torch.Tensor) -> list[torch.Tensor]: + """For each ``(M_local, ...)`` tensor: local rows, then the gathered send blocks.""" + + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_zeros((self.max_send, *t.shape[1:])) + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family: int, partial: torch.Tensor) -> torch.Tensor: + """This rank's tile partials -> every tile's partial, in WS1 tile order.""" + + fam = self.families[family] + padded = partial.new_zeros((fam["width"], *partial.shape[1:])) + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + +def _dweight_tiles(total: int, device): + begin = torch.arange(0, total, BACKWARD_TILE, device=device) + return torch.arange(total, device=device), begin, torch.clamp(begin + BACKWARD_TILE, max=total) + + +@dataclass(frozen=True) +class SPPlan: + """Everything an SP backward needs that depends only on the layout and the row index. + + Built once per ``(row index, table rows)`` and shared by every norm and gated + residual that uses that index (every block of a forward pass). + """ + + layout: SPRowLayout + plan: _Plan + local_index: torch.Tensor | None # (S_local,) table row of each local position + index_buf: torch.Tensor | None # table row of each exchange-buffer row + seg_first_tile: torch.Tensor | None + num_rows: int | None + + +def sp_plan( + layout: SPRowLayout, index_full=None, num_rows: int | None = None, *, device=None +) -> SPPlan: + """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" + + device = index_full.device if index_full is not None else torch.device(device or "cuda") + families = [_dweight_tiles(layout.batch * layout.seq_len, device)] + if index_full is None: + return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) + plan = _Plan(layout, [*families, tuple(seg_tiles)], device) + local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) + received = index_full[plan.recv_global % layout.seq_len] + index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() + return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) + + +def sp_rmsnorm_backward(grad, x, weight, rstd, shift, scale, sp: SPPlan, collective): + """One rank's backward on its ``(B * S_local, H)`` rows: ``(dx, dweight, d_shift, d_scale)``. + + ``dx`` is row-local; the other gradients are the full WS1 reductions, identical on + every rank. ``shift``/``scale`` are ``None`` for the plain norm. + """ + + weight, plan, modulated = weight.contiguous(), sp.plan, sp.local_index is not None + dx = _C.h3_rmsnorm_backward_dx(grad, x, weight, rstd, shift, scale, sp.local_index) + g_buf, x_buf, rstd_buf = plan.exchange(collective, grad, x, rstd) + seg = plan.families[1]["tiles"] if modulated else (None, None, None) + partials = _C.h3_rmsnorm_backward_partials( + g_buf, x_buf, weight, rstd_buf, shift, scale, sp.index_buf, + *plan.families[0]["tiles"], *seg, + ) # fmt: skip + gathered = [plan.gather_partials(collective, i, p) for i, p in enumerate(partials)] + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_rmsnorm_backward_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + if not modulated: + (dweight,) = _C.h3_rmsnorm_fold_partials(gathered[0], weight) + return dx, dweight, None, None + dweight, d_shift, d_scale = _C.h3_rmsnorm_fold_partials( + gathered[0], weight, gathered[1], sp.seg_first_tile + ) + return dx, dweight, d_shift.to(shift.dtype), d_scale.to(scale.dtype) + + +def sp_gate_residual_backward(grad, y, gate, sp: SPPlan, collective): + """One rank's backward on its rows: ``(d_sublayer, d_gate)``; ``d_residual`` is ``grad``.""" + + plan = sp.plan + dy = _C.h3_gate_residual_backward_dy(grad, gate, sp.local_index) + g_buf, y_buf = plan.exchange(collective, grad, y) + partial = _C.h3_gate_grad_partials(g_buf, y_buf, *plan.families[1]["tiles"]) + dgate = _C.h3_gate_grad_fold( + plan.gather_partials(collective, 1, partial), sp.seg_first_tile, gate.dtype + ) + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_gate_grad_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + return dy, dgate + + +class _SPRMSNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, sp, collective, eps): + x2 = x.contiguous().view(-1, x.shape[-1]) + y, rstd = _C.h3_rmsnorm_forward( + x2, weight.contiguous(), float(eps), shift, scale, sp.local_index + ) + ctx.save_for_backward(x2, weight, rstd, shift, scale) + ctx.sp, ctx.collective, ctx.x_shape = sp, collective, x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + dx, dweight, d_shift, d_scale = sp_rmsnorm_backward( + g2, x2, weight, rstd, shift, scale, ctx.sp, ctx.collective + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None, None + + +class _SPGateResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, sp, collective): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, sp.local_index) + ctx.save_for_backward(y2, gate) + ctx.sp, ctx.collective, ctx.shape = sp, collective, residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + dy, dgate = sp_gate_residual_backward(g2, y2, gate, ctx.sp, ctx.collective) + return grad, dy.view(ctx.shape), dgate, None, None + + +class H3SPNormAdaLNCudaOp: + """Sequence-parallel ``h3_rmsnorm`` (+ modulation) and ``adaln_gate_residual``. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` and + ``backend_id`` (e.g. ``DeterministicCollective``). Activations are this rank's + ``(B, S_local, H)`` rows; the norm weight, the table views and the ``(S,)`` row + index are the full, replicated tensors. Gradients of the replicated tensors come + back identical on every rank, equal to WS1. + """ + + op_class = "reduction" + backward_impl = BACKWARD_IMPL + + def __init__(self, collective, seq_len: int, batch: int) -> None: + if not sp_norm_adaln_available(): + raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") + self.collective = collective + self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None, *, device=None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + device = index.device if index is not None else torch.device(device or "cuda") + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + if index is not None and index.is_inference(): + # An inference tensor has no version counter (reading it raises), so an in-place edit + # could not be seen: build the plan every time instead of reusing a cached one. + return sp_plan(self.layout, index, num_rows, device=device) + version = -1 if index is None else index._version + for obj, ver, rows, cached_device, plan in self._plans: + if obj is index and ver == version and rows == num_rows and cached_device == device: + return plan + plan = sp_plan(self.layout, index, num_rows, device=device) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, device, plan), *self._plans[:3]] + return plan + + def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: + expected = (self.layout.batch, self.layout.local_len) + if not isinstance(x, torch.Tensor) or x.dim() != 3 or tuple(x.shape[:2]) != expected: + shape = tuple(x.shape) if isinstance(x, torch.Tensor) else type(x).__name__ + raise ValueError(f"activations must be this rank's {expected} + (H,) rows, got {shape}") + if not x.is_cuda: + raise ValueError("H3SPNormAdaLNCudaOp needs CUDA tensors") + if index is None: + return + if index.dim() != 1 or index.shape[0] != self.layout.seq_len: + raise ValueError( + f"index must be the full ({self.layout.seq_len},) row index, " + f"got {tuple(index.shape)}" + ) + if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): + raise IndexError(f"index must be in [0, {num_rows})") + + def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + self._check(x, None, None) + validate_h3_rmsnorm(x, weight, eps) + return _SPRMSNorm.apply( + x, weight, None, None, self.plan(None, None, device=x.device), self.collective, eps + ) + + def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + self._check(x, index, shift.shape[0]) + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index[self.layout.lo : self.layout.hi]) + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + sp = self.plan(index, shift.shape[0]) + return _SPRMSNorm.apply(x, weight, shift, scale, sp, self.collective, eps) + + def gate_residual(self, residual, y, gate, index) -> torch.Tensor: + self._check(residual, index, gate.shape[0]) + validate_h3_gate_residual(residual, y, gate, index[self.layout.lo : self.layout.hi]) + if gate.stride(1) != 1: + gate = gate.contiguous() + sp = self.plan(index, gate.shape[0]) + return _SPGateResidual.apply(residual, y, gate, sp, self.collective) + + def readback(self) -> dict: + return { + "op": "sp_norm_adaln", + "backward_impl": BACKWARD_IMPL, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_256_element_tiles_then_ascending_fold", + "sp": self.layout.sp, + "rank": self.layout.rank, + "positions": [self.layout.lo, self.layout.hi], + "batch": self.layout.batch, + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py index 28768e881..22ed0a114 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -85,6 +85,8 @@ def slots(self) -> list[tuple[int, int, int, int]]: def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: if tp < 1 or not 0 <= rank < tp: raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total <= 0: + raise ValueError(f"N must be positive, got {n_total}") if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") if n_total % (tp * DINPUT_CHUNK): diff --git a/rl_engine/validation/distributed/h3_ws2.py b/rl_engine/validation/distributed/h3_ws2.py index 1cc831986..7342774a0 100644 --- a/rl_engine/validation/distributed/h3_ws2.py +++ b/rl_engine/validation/distributed/h3_ws2.py @@ -124,6 +124,191 @@ def step(op, args): return {"equality": equality, "perf": perf, "readback": tp_op.readback()} +def _region(gate_fn, norm_fn, residual, y, weight, table): + """One block's MLP-side SP region: ``norm2(residual + gate_msa * y)`` with modulation.""" + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.chunk(6, dim=-1) + hidden = gate_fn(residual, y, gate_msa) + return norm_fn(hidden, weight, shift_mlp, scale_mlp) + + +def ws1_sp_region(residual, y, weight, table, index, grad) -> dict[str, torch.Tensor]: + """The SP region on one GPU with the WS1 ops (the reference for ``sp_norm_adaln``).""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import H3GateResidualCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, weight, table)] + out = _region( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad) + names = ("d_residual", "d_y", "d_weight", "d_table") + return {"out": out.detach(), **{n: leaf.grad for n, leaf in zip(names, leaves)}} + + +def sp_region_rank(collective, residual, y, weight, table, index, grad) -> dict[str, Any]: + """One SP rank: its rows of the region and its (replicated) parameter gradients.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + op = H3SPNormAdaLNCudaOp(collective, seq_len=residual.shape[1], batch=residual.shape[0]) + rows = slice(op.layout.lo, op.layout.hi) + local = [t[:, rows].contiguous() for t in (residual, y)] + leaves = [t.detach().clone().requires_grad_(True) for t in (*local, weight, table)] + out = _region( + lambda r, yy, g: op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: op.norm_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad[:, rows].contiguous()) + names = ("d_residual", "d_y", "d_weight", "d_table") + return { + "out": out.detach(), + **{n: leaf.grad for n, leaf in zip(names, leaves)}, + "readback": op.readback(), + } + + +def sp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Rows: each rank's slice equals WS1's; parameters: every rank's copy equals WS1's.""" + + def rows(key): + return all( + torch.equal(r[key], ws1[key][:, slice(*r["readback"]["positions"])]) for r in ranks + ) + + return { + "out": rows("out"), + "d_residual": rows("d_residual"), + "d_y": rows("d_y"), + "d_weight": all(torch.equal(r["d_weight"], ws1["d_weight"]) for r in ranks), + "d_table": all(torch.equal(r["d_table"], ws1["d_table"]) for r in ranks), + } + + +def sp_case(batch, seq, hidden=5376, num_t=3, *, layout="block", dtype=torch.bfloat16, seed=0): + """Inputs of the SP region; identical on every rank (CPU generator).""" + + from rl_engine.validation.models.h3_cases import h3_block_layout, h3_packed_layout + + g = torch.Generator(device="cpu").manual_seed(seed) + make = h3_block_layout if layout == "block" else h3_packed_layout + ti, tags = make(seq, num_t, seed=seed) + return { + "residual": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + "y": (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda(), + "weight": (torch.rand(hidden, generator=g) + 0.5).to(dtype).cuda(), + "table": (torch.randn(3 * num_t, 6 * hidden, generator=g) * 0.5).to(dtype).cuda(), + "index": (ti * 3 + tags).cuda(), + "grad": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + } + + +def naive_sp_mismatch(case: dict[str, torch.Tensor], sp: int) -> dict[str, float]: + """Fraction of parameter-gradient elements where a naive SP backward differs from WS1. + + Naive: each rank runs the WS1 backward on its own rows, then the per-rank results + are summed in rank order (what a plain all-reduce of local gradients does). + """ + + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_row_layout + + ws1 = ws1_sp_region(**case) + parts = [] + for rank in range(sp): + lay = sp_row_layout(case["residual"].shape[1], case["residual"].shape[0], sp, rank) + part = slice(lay.lo, lay.hi) + local = {k: (v[:, part] if k in ("residual", "y", "grad") else v) for k, v in case.items()} + local["index"] = case["index"][part] + parts.append(ws1_sp_region(**local)) + out = {} + for key in ("d_weight", "d_table"): + total = parts[0][key].float() + for p in parts[1:]: + total = total + p[key].float() + out[key] = (total.to(ws1[key].dtype) != ws1[key]).float().mean().item() + return out + + +def sp_region_sweep(collective, cases: list[dict[str, Any]], iters: int = 20) -> dict[str, Any]: + """Evidence for one SP rank: byte equality, rows exchanged and time, per case.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + _dweight_tiles, + _Plan, + sp_row_layout, + ) + + results = [] + for spec in cases: + case = sp_case(**spec) + batch, seq = case["residual"].shape[:2] + ws1 = ws1_sp_region(**case) + mine = sp_region_rank(collective, **case) + equal = sp_matches_ws1(ws1, [mine]) + lay = sp_row_layout(seq, batch, collective.world_size, collective.rank) + tiles = _segment_tiles(case["index"].repeat(batch), case["table"].shape[0]) + plan = _Plan(lay, [_dweight_tiles(batch * seq, "cuda"), tuple(tiles[:3])], "cuda") + del ws1, mine + torch.cuda.empty_cache() + results.append( + { + **spec, + "equal": equal, + "local_rows": batch * lay.local_len, + "rows_sent": int(plan.send_local.numel()), + "rows_gathered_per_rank": plan.max_send, + **_region_times(collective, case, lay, iters), + } + ) + return {"cases": results} + + +def _region_times(collective, case, lay, iters) -> dict[str, float]: + """Forward + backward time of the region (leaves prepared once, grads reset per call).""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import H3GateResidualCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + from rl_engine.validation.models.h3_report import time_us + + index, rows = case["index"], slice(lay.lo, lay.hi) + sp_op = H3SPNormAdaLNCudaOp(collective, seq_len=lay.seq_len, batch=lay.batch) + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + + def leaves(local): + acts = [case[k][:, rows] if local else case[k] for k in ("residual", "y")] + return [ + t.detach().contiguous().requires_grad_(True) + for t in (*acts, case["weight"], case["table"]) + ] + + def step(fns, args, grad): + for leaf in args: + leaf.grad = None + _region(*fns, *args).backward(grad) + + sp_fns = ( + lambda r, yy, g: sp_op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: sp_op.norm_modulated(h, w, sh, sc, index), + ) + ws1_fns = ( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + ) + sp_args, ws1_args = leaves(True), leaves(False) + sp_grad = case["grad"][:, rows].contiguous() + return { + "sp_forward_backward_us": time_us(lambda: step(sp_fns, sp_args, sp_grad), 3, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_fns, ws1_args, case["grad"]), 3, iters), + } + + # --- real multi-process groups ------------------------------------------------- @@ -183,33 +368,40 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa for r in range(world) ] started = [] + deadline = time.monotonic() + timeout + status = {} try: for proc in procs: proc.start() started.append(proc) - deadline = time.monotonic() + timeout - pending = set(range(world)) - while pending: + while len(status) < world or any(proc.is_alive() for proc in procs): remaining = deadline - time.monotonic() if remaining <= 0: - raise TimeoutError(f"timed out waiting for ranks {sorted(pending)}") + raise TimeoutError(f"rank results timed out after {timeout}s") try: result = queue.get(timeout=min(0.1, remaining)) except Empty: - for rank in pending: - if procs[rank].exitcode is not None: - raise RuntimeError( - f"rank {rank} exited with code {procs[rank].exitcode} " - "without reporting a result" - ) from None - continue + exited = [ + (rank, proc.exitcode) + for rank, proc in enumerate(procs) + if proc.exitcode is not None and (proc.exitcode != 0 or rank not in status) + ] + if not exited: + continue + # A worker may have flushed its final message between the + # timed read and the exit-code check. + try: + result = queue.get_nowait() + except Empty: + rank, code = exited[0] + raise RuntimeError( + f"rank {rank} exited with code {code} without completing" + ) from None if "error" in result: raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - pending.remove(result["rank"]) + status[result["rank"]] = result for rank, proc in enumerate(procs): - proc.join(timeout=max(0, deadline - time.monotonic())) - if proc.is_alive(): - raise TimeoutError(f"rank {rank} did not exit after reporting a result") + proc.join() if proc.exitcode != 0: raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] diff --git a/rl_engine/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py index 8eb509d33..62184cc75 100644 --- a/rl_engine/validation/models/h3_cases.py +++ b/rl_engine/validation/models/h3_cases.py @@ -38,3 +38,27 @@ def h3_packed_layout( timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) return timestep_indices.to(device), token_tags.to(device) + + +def h3_block_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for block-structured packing, as H3 packs requests. + + Each timestep owns one contiguous run of text, video and audio blocks (about + 5% / 80% / 15% of its tokens), so table rows form contiguous position runs. + ``h3_packed_layout`` is the interleaved stress case. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + cuts = torch.sort(torch.randperm(seq_len - 1, generator=generator)[: num_timesteps - 1] + 1) + edges = [0, *cuts.values.tolist(), seq_len] + timestep_indices = torch.empty(seq_len, dtype=torch.long) + token_tags = torch.empty(seq_len, dtype=torch.long) + for t, (lo, hi) in enumerate(zip(edges, edges[1:])): + text, audio = (hi - lo) // 20, (hi - lo) * 3 // 20 + timestep_indices[lo:hi] = t + token_tags[lo:hi] = 0 # video + token_tags[lo : lo + text] = 1 + token_tags[hi - audio : hi] = 2 + return timestep_indices.to(device), token_tags.to(device) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 46fe86c32..db2291225 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -607,11 +607,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def prepare_backward(fn, residual=residual, y=y, table=table): + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad - def backward(output, grad=grad): + def backward(state): + output, grad = state output.backward(grad) cases.append( @@ -717,11 +718,12 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def prepare_backward(fn, tensors=(x, nw, temb, w, b)): + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): leaves = [t.detach().requires_grad_(True) for t in tensors] - return fn(*leaves) + return fn(*leaves), grad - def backward(output, grad=grad): + def backward(state): + output, grad = state output.backward(grad) cases.append( diff --git a/tests/distributed/sp/__init__.py b/tests/distributed/sp/__init__.py new file mode 100644 index 000000000..b1bbeb05c --- /dev/null +++ b/tests/distributed/sp/__init__.py @@ -0,0 +1 @@ +"""tests.distributed.sp package.""" diff --git a/tests/distributed/sp/test_h3_sp_norm_adaln.py b/tests/distributed/sp/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..72be21164 --- /dev/null +++ b/tests/distributed/sp/test_h3_sp_norm_adaln.py @@ -0,0 +1,445 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``sp_norm_adaln``: sequence-parallel norm / modulation / gated residual. + +* ownership: contiguous position shards per batch item; bad layouts fail closed; +* backward (one GPU, ranks as threads around the plain backward functions): every + rank's row gradients equal WS1's rows, and its ``dweight``/``d_shift``/``d_scale``/ + ``d_gate`` equal WS1's, for SP 2..8, odd S, B > 1, interleaved modality rows and + bf16/fp32; +* end to end (real NCCL, autograd): the MLP-side block region + ``norm2(residual + gate_msa * y)`` with modulation is byte-equal to WS1 on every rank. +""" + +from __future__ import annotations + +import os +import threading +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + SPRowLayout, + sp_row_layout, +) +from rl_engine.validation.distributed.h3_ws2 import ( + run_world, + sp_case, + sp_matches_ws1, + sp_region_rank, + ws1_sp_region, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +ALL_TRUE = dict.fromkeys(("out", "d_residual", "d_y", "d_weight", "d_table"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_norm_adaln_available, + ) + + if not sp_norm_adaln_available(): + pytest.skip("rl_engine._C lacks the H3 SP symbols") + + +class _ThreadGroup: + """Ranks as threads on one GPU; ``all_gather`` is the rank-ordered concatenation. + + Only for the plain backward functions: autograd runs every backward of a device + on one worker thread, so autograd-driven ranks need real processes. + """ + + def __init__(self, world: int) -> None: + self.barrier, self.slots = threading.Barrier(world), [None] * world + + def handle(self, rank: int): + def all_gather(t): + self.slots[rank] = t + self.barrier.wait() + out = torch.cat(self.slots) + self.barrier.wait() + return out + + return SimpleNamespace(rank=rank, world_size=len(self.slots), all_gather=all_gather) + + def run(self, fn): + results, errors = [None] * len(self.slots), [] + + def body(rank): + try: + results[rank] = fn(self.handle(rank)) + except BaseException as exc: # re-raised below + errors.append(exc) + self.barrier.abort() + + threads = [threading.Thread(target=body, args=(r,)) for r in range(len(self.slots))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise next( + (e for e in errors if not isinstance(e, threading.BrokenBarrierError)), errors[0] + ) + return results + + +class TestLayout: + @pytest.mark.parametrize("seq, sp", [(8, 8), (257, 2), (1000, 3), (4097, 8)]) + def test_shards_cover_every_position_once(self, seq, sp): + bounds = [SPRowLayout(seq, 1, sp, r).bounds(r) for r in range(sp)] + assert bounds[0][0] == 0 and bounds[-1][1] == seq + assert all(a[1] == b[0] and a[0] < a[1] for a, b in zip(bounds, bounds[1:])) + + @pytest.mark.parametrize("seq, batch, sp, rank", [(3, 1, 4, 0), (8, 0, 2, 0), (8, 1, 2, 2)]) + def test_rejects(self, seq, batch, sp, rank): + with pytest.raises(ValueError): + sp_row_layout(seq, batch, sp, rank) + + +@requires_cuda +class TestBackwardRanks: + @pytest.mark.parametrize( + "sp, batch, seq, hidden, layout, dtype", + [ + (2, 1, 4097, 5376, "block", torch.bfloat16), + (2, 1, 4097, 5376, "interleaved", torch.bfloat16), + (3, 2, 1000, 256, "interleaved", torch.bfloat16), + (4, 2, 777, 512, "block", torch.float32), + (5, 1, 1301, 128, "interleaved", torch.bfloat16), + (8, 2, 2049, 256, "block", torch.bfloat16), + (8, 1, 40, 64, "interleaved", torch.float32), # shards far smaller than a tile + ], + ) + def test_ranks_reproduce_ws1(self, sp, batch, seq, hidden, layout, dtype): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_gate_residual_backward, + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(batch, seq, hidden, layout=layout, dtype=dtype, seed=sp * 100 + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] # shift/scale_mlp, gate_msa + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = sp_plan(lay, idx, rows) + norm = sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + gate_grads = sp_gate_residual_backward(local["grad"], local["y"], gate, plan, coll) + return lay, norm, gate_grads + + for lay, (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + part = slice(lay.lo, lay.hi) + assert torch.equal( + dx.view(batch, -1, hidden), ws1_norm[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dtype)) and torch.equal( + dsc, ws1_norm[3].to(dtype) + ) + assert torch.equal( + dy.view(batch, -1, hidden), ws1_gate[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dgate, ws1_gate[1]) + + @pytest.mark.parametrize("sp, batch, seq", [(4, 2, 777), (3, 2, 1000), (8, 1, 40)]) + def test_padding_of_the_exchange_and_partial_gathers_is_never_read( + self, monkeypatch, sp, batch, seq + ): + """Pad every send block and every partial gather with NaN: WS1 equality still holds. + + ``recv_global`` is zero in the padding slots, which map to a real table row, so the + padding is only harmless because no tile list refers to it; a NaN would reach the sums + of any tile that did. + """ + + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3 import sp_norm_adaln as module + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + + hidden = 256 + c = sp_case(batch, seq, hidden, layout="interleaved", dtype=torch.float32, seed=sp + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + nan = float("nan") + gather_rows = module.gather_rows + + def exchange(self, collective, *tensors): + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_full((self.max_send, *t.shape[1:]), nan) # padding is NaN + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family, partial): + fam = self.families[family] + padded = partial.new_full((fam["width"], *partial.shape[1:]), nan) # padding is NaN + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + monkeypatch.setattr(module._Plan, "exchange", exchange) + monkeypatch.setattr(module._Plan, "gather_partials", gather_partials) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = module.sp_plan(lay, idx, rows) + norm = module.sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + return norm, module.sp_gate_residual_backward( + local["grad"], local["y"], gate, plan, coll + ) + + for (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + assert torch.isfinite(dx).all() and torch.isfinite(dy).all() + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dsh.dtype)) + assert torch.equal(dsc, ws1_norm[3].to(dsc.dtype)) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + H3SPNormAdaLNCudaOp, + ) + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@requires_cuda +def test_inference_mode_index_is_planned_every_time_and_runs_both_ops(): + """An index made under ``torch.inference_mode()`` has no version counter (reading it raises).""" + + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + seq, hidden, rows = 4, 16, 2 + gen = torch.Generator(device="cuda").manual_seed(0) + + def bf16(*shape): + return torch.randn(*shape, device="cuda", generator=gen).bfloat16() + + x, residual, y = bf16(1, seq, hidden), bf16(1, seq, hidden), bf16(1, seq, hidden) + weight, shift, scale, gate = ( + bf16(hidden), + bf16(rows, hidden), + bf16(rows, hidden), + bf16(rows, hidden), + ) + values = [0, 1, 0, 1] + + op = H3SPNormAdaLNCudaOp(collective, seq, 1) + normal = torch.tensor(values, device="cuda") + want_norm = op.norm_modulated(x, weight, shift, scale, normal) + want_gate = op.gate_residual(residual, y, gate, normal) + + with torch.inference_mode(): + index = torch.tensor(values, device="cuda") + assert index.is_inference() + first = op.plan( + index, rows + ) # used to raise "Inference tensors do not track version counter" + got_norm = op.norm_modulated(x, weight, shift, scale, index) + got_gate = op.gate_residual(residual, y, gate, index) + index.copy_(torch.tensor([1, 1, 0, 0], device="cuda")) # in place; no version to notice it + second = op.plan(index, rows) + assert torch.equal(got_norm, want_norm) + assert torch.equal(got_gate, want_gate) + assert second is not first + assert torch.equal(second.local_index, torch.tensor([1, 1, 0, 0], device="cuda")) + assert torch.equal(first.local_index, torch.tensor(values, device="cuda")) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_plain_norm_uses_activation_device_and_separate_cache(): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + op = H3SPNormAdaLNCudaOp(collective, 5, 1) + with torch.cuda.device(0): + for device in ("cuda:1", "cuda:0"): + x = torch.randn(1, 5, 8, device=device, requires_grad=True) + weight = torch.randn(8, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) + result = op.norm(x, weight) + expected = torch.nn.functional.rms_norm(ref_x, (8,), ref_w, 1e-5) + result.sum().backward() + expected.sum().backward() + torch.testing.assert_close(result, expected) + torch.testing.assert_close(x.grad, ref_x.grad) + torch.testing.assert_close(weight.grad, ref_w.grad) + assert len(op._plans) == 2 + assert op.plan(None, None, device="cuda:1").plan.send_local.device.index == 1 + assert op.plan(None, None, device="cuda:0").plan.send_local.device.index == 0 + + +def _worker_exit_without_result(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + os._exit(kwargs["exit_code"]) + time.sleep(60) + + +def _worker_reports_error(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + queue.put({"rank": rank, "error": "intentional worker failure"}) + time.sleep(60) + + +def _worker_never_reports(rank, world, init, queue, path, target, kwargs): + import time + + time.sleep(60) + + +@pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) +def test_run_world_cleans_up_failed_workers(monkeypatch, mode): + import time + + import torch.multiprocessing as mp + + from rl_engine.validation.distributed import h3_ws2 + + workers = { + "clean_exit": _worker_exit_without_result, + "crash": _worker_exit_without_result, + "reported_error": _worker_reports_error, + "timeout": _worker_never_reports, + } + monkeypatch.setattr(h3_ws2, "_bootstrap", workers[mode]) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + expected = TimeoutError if mode == "timeout" else RuntimeError + with pytest.raises(expected): + run_world( + 2, + None, + {}, + timeout=1 if mode == "timeout" else 30, + exit_code=7 if mode == "crash" else 0, + ) + assert time.monotonic() - start < 20 + assert {p.pid for p in mp.active_children()} <= before + + +def _worker_reports_result(rank, world, init, queue, path, target, kwargs): + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + + +def test_run_world_returns_rank_ordered_results(monkeypatch): + from rl_engine.validation.distributed import h3_ws2 + + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker_reports_result) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + results = run_world(2, None, {}, timeout=30) + assert [result["rank"] for result in results] == [0, 1] + assert [result["value"].item() for result in results] == [0, 1] diff --git a/tests/distributed/test_h3_ws2_processes.py b/tests/distributed/test_h3_ws2_processes.py index 548886490..5150acf1a 100644 --- a/tests/distributed/test_h3_ws2_processes.py +++ b/tests/distributed/test_h3_ws2_processes.py @@ -37,7 +37,7 @@ def test_run_world_reaps_workers(monkeypatch, mode): assert [r["rank"] for r in ranks] == [0, 1] assert [r["value"].item() for r in ranks] == [0, 1] elif mode == "timeout": - with pytest.raises(TimeoutError, match="timed out waiting for ranks"): + with pytest.raises(TimeoutError, match="timed out"): h3_ws2.run_world(2, None, {}, timeout=0.2, mode=mode) else: message = "original rank traceback" if mode == "error" else "rank 0 exited with code 7" diff --git a/tests/distributed/tp/test_h3_tp_adaln.py b/tests/distributed/tp/test_h3_tp_adaln.py index 9219d5135..2ca5be32b 100644 --- a/tests/distributed/tp/test_h3_tp_adaln.py +++ b/tests/distributed/tp/test_h3_tp_adaln.py @@ -62,7 +62,7 @@ def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): class TestOwnership: - @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8, 9]) def test_shards_tile_the_table_once(self, tp): shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] @@ -234,3 +234,9 @@ def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): assert tp_matches_ws1(ws1, ranks) == ALL_TRUE assert [r["readback"]["rank"] for r in ranks] == list(range(world)) assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@pytest.mark.parametrize("n_total", [0, -1152]) +def test_rejects_nonpositive_projection(n_total): + with pytest.raises(ValueError, match="N must be positive"): + adaln_column_shard(n_total, 1, 0) diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py index 2a45ec657..8d3d985c1 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -445,7 +445,7 @@ def test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype( ).to(out_dtype) assert torch.equal(got, expected) - def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): + def test_backward_with_no_table_rows_or_no_boundaries_is_rejected(self): from rl_engine.backends.extension import _C _cuda_op() @@ -453,8 +453,8 @@ def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_bound zero = torch.zeros(0, dtype=torch.int64, device="cuda") pos = torch.arange(4, device="cuda") one = torch.zeros(1, dtype=torch.int64, device="cuda") # R = 0 - empty = _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) - assert empty.shape == (0, 48) + with pytest.raises(RuntimeError, match="must be positive"): + _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) with pytest.raises(RuntimeError, match="R \\+ 1 entries"): _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, zero, torch.float32) diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py index 0a9bf216c..bd3a44363 100644 --- a/tests/models/minimax_h3/test_h3_benchmark_timing.py +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -136,9 +136,7 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append( - [leaf._base if leaf._base is not None else leaf for leaf in leaves] - ) + prepared_leaves.append([leaf if leaf.is_leaf else leaf._base for leaf in leaves]) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -168,7 +166,7 @@ def forward(x, weight, shift, scale, index, **kwargs): def inputs(seq): return ( - torch.randn(3, 6 * 4), + torch.randn(3, 24), torch.randn(1, seq, 4), torch.randn(1, seq, 4), torch.arange(seq) % 3, @@ -177,7 +175,7 @@ def inputs(seq): provider = h3_report.provider_gate_residual def forward(residual, y, gate, index, **kwargs): - return track(lambda r, y_, g, i: provider(r, g, i, y_), [residual, y, gate], index) + return track(lambda r, y_, g: provider(r, g, index, y_), [residual, y, gate]) monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) monkeypatch.setattr(h3_report, "_gate_inputs", inputs) @@ -198,10 +196,14 @@ def inputs(seq): torch.arange(seq) % 3, ) - provider = h3_report.provider_final_adaln_out - def forward(x, nw, temb, w, b, ti): - return track(provider, [x, nw, temb, w, b], ti) + def compute(x, nw, temb, w, b): + rows = torch.nn.functional.linear(torch.nn.functional.silu(temb), w, b) + shift, scale = rows.chunk(2, dim=-1) + norm = torch.nn.functional.rms_norm(x, (4,), nw) + return norm * (1 + scale[ti]) + shift[ti] + + return track(compute, [x, nw, temb, w, b]) monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) monkeypatch.setattr(h3_report, "_final_inputs", inputs) diff --git a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py index 091001e36..d031f1d03 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -3,7 +3,7 @@ """A heatmap cell counts every forward output of the operator, not only the first. -Also: the TP evidence figure for a report that has no TP > 1 run. +Also: the TP and SP evidence figures for a report that has no TP or SP > 1 run. """ from __future__ import annotations @@ -126,3 +126,36 @@ def test_tp_figure_for_reports_with_and_without_tp_above_one(worlds): texts = [t.get_text() for ax in fig.axes for t in ax.texts] assert ("no run with TP > 1 in this report" in texts) == (worlds == [1]) plt.close(fig) + + +def _sp_report(worlds: list[int]) -> dict: + from plot_h3_evidence import SP_TENSORS + + case = { + "layout": "block", + "seq": 64, + "batch": 1, + "sp_forward_backward_us": 2000.0, + "ws1_forward_backward_us": 1900.0, + "equal": dict.fromkeys((key for key, _ in SP_TENSORS), True), + } + naive = {"d_weight": 0.5, "d_table": 0.5} + return { + "op": "sp_norm_adaln", + "environment": {"gpu": "gpu", "gpus": 8, "torch": "t", "cuda": "c"}, + "rl_kernel_commit": "0123456789", + "model_revision": "0123456789", + "runs": [{"sp": sp, "ranks": [{"rank": 0, "cases": [case]}]} for sp in worlds], + "naive_sp": {"sp": 2, "block": naive, "interleaved": naive}, + } + + +@pytest.mark.parametrize("worlds", [[1], [1, 2, 4], [2, 4]]) +def test_sp_figure_for_reports_with_and_without_sp_above_one(worlds): + import matplotlib.pyplot as plt + from plot_h3_evidence import plot_sp_norm + + fig = plot_sp_norm(_sp_report(worlds)) + texts = [t.get_text() for ax in fig.axes for t in ax.texts] + assert ("no run with SP > 1 in this report" in texts) == (worlds == [1]) + plt.close(fig) diff --git a/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py b/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py new file mode 100644 index 000000000..7d5e465e3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""A build that leaves the H3 sources out has no ``h3_*`` symbols; no CUDA op may then construct. + +``KERNEL_ALIGN_USE_FAST_MATH=1``, ROCm and relaxed ``NVCC_*_FLAGS`` builds omit the H3 kernels and +their bindings (build_tools/extensions.py). Every H3 CUDA op checks for its symbols when it is +constructed, so the registry falls back to the PyTorch reference instead of failing on the first +call with an AttributeError. +""" + +from __future__ import annotations + +import importlib + +import pytest + +PACKAGE = "rl_engine.backends.cuda.model_specific.minimax_h3" +# the modules that define the probes and own the ``_C`` they look in +PROBE_MODULES = ( + "timestep_sinusoid", + "det_linear", + "adaln_row_gather", + "rmsnorm", + "gate_residual", + "sp_norm_adaln", +) +# module, class, constructor arguments +OPS = ( + ("timestep_sinusoid", "H3TimestepSinusoidCudaOp", ()), + ("timestep_mlp", "H3TimestepMLPCudaOp", ()), + ("adaln_projection", "H3AdaLNProjectionCudaOp", ()), + ("adaln_row_gather", "H3AdaLNRowGatherCudaOp", ()), + ("adaln_modulation", "H3AdaLNModulationCudaOp", ()), + ("rmsnorm", "H3RMSNormCudaOp", ()), + ("gate_residual", "H3GateResidualCudaOp", ()), + ("final_adaln_out", "H3FinalAdaLNOutCudaOp", ()), + ("sp_norm_adaln", "H3SPNormAdaLNCudaOp", (None, 1, 1)), +) + + +class ExtensionWithoutH3: + """Stands in for a compiled ``rl_engine._C`` that has none of the H3 bindings.""" + + +@pytest.fixture +def extension_without_h3(monkeypatch): + for name in PROBE_MODULES: + module = importlib.import_module(f"{PACKAGE}.{name}") + monkeypatch.setattr(module, "_C", ExtensionWithoutH3()) + monkeypatch.setattr(module, "_EXT_AVAILABLE", True) + + +@pytest.mark.parametrize(("module", "op", "args"), OPS, ids=[o[1] for o in OPS]) +def test_every_cuda_op_refuses_to_construct(extension_without_h3, module, op, args): + cls = getattr(importlib.import_module(f"{PACKAGE}.{module}"), op) + with pytest.raises(RuntimeError): + cls(*args) + + +def test_the_sequence_parallel_op_checks_every_symbol_it_calls(monkeypatch): + sp = importlib.import_module(f"{PACKAGE}.sp_norm_adaln") + monkeypatch.setattr(sp, "_EXT_AVAILABLE", True) + + def extension(missing=None): + names = {n: staticmethod(lambda *a, **k: None) for n in sp._SYMBOLS if n != missing} + return type("Extension", (), names)() + + monkeypatch.setattr(sp, "_C", extension()) + assert sp.sp_norm_adaln_available() is True + for name in sp._SYMBOLS: # one missing symbol is enough to refuse: the op does not half-work + monkeypatch.setattr(sp, "_C", extension(missing=name)) + assert sp.sp_norm_adaln_available() is False, name diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index d04b16e3d..d16f07395 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -16,7 +16,6 @@ def native_extension(): from rl_engine.backends.extension import _C, _EXT_AVAILABLE names = ( - "h3_adaln_row_gather_backward", "h3_rmsnorm_forward", "h3_rmsnorm_backward", "h3_gate_residual_forward", @@ -50,10 +49,6 @@ def native_case(): def _native_call(extension, operation, boundary, case): x, index = case["x"], case["index"] tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) - if operation == "gather": - return extension.h3_adaln_row_gather_backward( - case["grad"].unsqueeze(0), *tiles, case["grad"].dtype - ) if operation == "rmsnorm": modulation = (case["shift"], case["scale"], index) if boundary == "forward": @@ -91,7 +86,7 @@ def test_native_rejects_invalid_row_index( _native_call(native_extension, operation, boundary, native_case) -@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual", "gather"]) +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) @pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) @pytest.mark.parametrize("device", ["cpu", "other_cuda"]) def test_native_rejects_mixed_device_tiles( @@ -129,45 +124,394 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat torch.testing.assert_close(actual, expected) -@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) -@pytest.mark.parametrize("boundary", ["forward", "backward"]) +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +@pytest.mark.cuda_only +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("entrypoint", ["partials", "dx"]) +def test_sp_backward_entry_points_reject_invalid_rstd(native, inputs, defect, entrypoint): + rstd = _malformed(inputs["rstd"], defect) + args = (inputs["grad"], inputs["x"], inputs["weight"], rstd) + modulation = (inputs["shift"], inputs["scale"], inputs["index"]) + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + if entrypoint == "dx": + native.h3_rmsnorm_backward_dx(*args, *modulation) + else: + tile = torch.zeros(1, dtype=torch.int32, device=inputs["x"].device) + native.h3_rmsnorm_backward_partials( + *args, *modulation, tile, tile, tile, None, None, None + ) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.cuda_only +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize( + "operation", ["gate_forward", "gate_dy", "norm_forward", "norm_dx", "norm_partials"] +) @pytest.mark.parametrize("bad_index", [-1, 2]) -def test_native_rejects_out_of_range_index_in_subprocess( - native_extension, operation, boundary, bad_index -): - # A device assertion poisons the CUDA context, so isolate each invalid launch. - import os +def test_native_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): import subprocess import sys + import textwrap - code = f""" -import torch -from rl_engine import _C -x = torch.ones(2, 8, device='cuda') -w = torch.ones(8, device='cuda') -index = torch.tensor([0, {bad_index}], device='cuda', dtype=torch.int64) -pos = torch.arange(2, device='cuda', dtype=torch.int64) -ends = pos + 1 -seg = torch.arange(3, device='cuda', dtype=torch.int64) -if {operation!r} == 'gate_residual': - if {boundary!r} == 'forward': - _C.h3_gate_residual_forward(x, x, x, index) - else: - _C.h3_gate_residual_backward(x, x, x, index, pos, pos, ends, seg) -else: - if {boundary!r} == 'forward': - _C.h3_rmsnorm_forward(x, w, 1e-5, x, x, index) - else: - rstd = torch.ones(2, device='cuda') - _C.h3_rmsnorm_backward(x, x, w, rstd, x, x, index, pos, pos, ends, seg) -torch.cuda.synchronize() -""" + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + op, bad = sys.argv[1], int(sys.argv[2]) + x = torch.ones(2, 8, device="cuda") + w = torch.ones(8, device="cuda") + table = torch.ones(2, 8, device="cuda") + index = torch.tensor([bad, 0], device="cuda") + rstd = torch.ones(2, device="cuda") + rows = torch.arange(2, device="cuda") + begin = torch.tensor([0], device="cuda") + end = torch.tensor([2], device="cuda") + try: + if op == "gate_forward": + _C.h3_gate_residual_forward(x, x, table, index) + elif op == "gate_dy": + _C.h3_gate_residual_backward_dy(x, table, index) + elif op == "norm_forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, table, table, index) + elif op == "norm_dx": + _C.h3_rmsnorm_backward_dx(x, x, w, rstd, table, table, index) + else: + _C.h3_rmsnorm_backward_partials( + x, x, w, rstd, table, table, index, rows, begin, end, rows, begin, end) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid index reached native table lookup without an assertion") + """) result = subprocess.run( - [sys.executable, "-c", code], + [sys.executable, "-c", script, operation, str(bad_index)], capture_output=True, text=True, timeout=60, - env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, ) - assert result.returncode != 0 - assert "device-side assert" in result.stderr + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +def test_row_gather_rejects_metadata_on_other_device(native_extension, metadata): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + grad = torch.ones(6, 2, 8, device="cuda:0") + tensors = dict( + sorted_pos=torch.tensor([0, 1], device="cuda:0"), + tile_begin=torch.tensor([0], device="cuda:0"), + tile_end=torch.tensor([2], device="cuda:0"), + seg_first_tile=torch.tensor([0, 1], device="cuda:0"), + ) + tensors[metadata] = tensors[metadata].to("cuda:1") + with pytest.raises(RuntimeError, match="grad device"): + native_extension.h3_adaln_row_gather_backward(grad, **tensors, out_dtype=torch.float32) + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("bound", ["negative", "past_end"]) +def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metadata, bound): + import subprocess + import sys + import textwrap + + script = textwrap.dedent(""" + import sys + import torch + from rl_engine import _C + name, bound = sys.argv[1:] + grad = torch.ones(6, 2, 8, device="cuda") + metadata = dict(sorted_pos=torch.tensor([0, 1], device="cuda"), + tile_begin=torch.tensor([0], device="cuda"), + tile_end=torch.tensor([2], device="cuda"), + seg_first_tile=torch.tensor([0, 1], device="cuda")) + metadata[name][0] = -1 if bound == "negative" else 3 + try: + _C.h3_adaln_row_gather_backward(grad, **metadata, out_dtype=torch.float32) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid metadata was not rejected") + """) + result = subprocess.run( + [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +def test_rmsnorm_rejects_modulation_rows_that_overlap(native_extension, native_case): + """A table of rows needs a row stride of at least N; an overlapping view is not one.""" + + table = torch.arange(32, dtype=torch.float32, device="cuda") + overlapping = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, overlapping, overlapping, + native_case["index"], + ) # fmt: skip + + +def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension, native_case): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + views = torch.randn(2, 16, generator=g).cuda().chunk(2, dim=-1) + copies = tuple(v.contiguous() for v in views) + assert views[0].stride(0) == 16 and copies[0].stride(0) == 8 + out = [ + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, shift, scale, native_case["index"] + )[0] + for shift, scale in (views, copies) + ] + assert torch.equal(out[0], out[1]) + + +_BACKWARD_CHILD = """ +import sys +import torch +from rl_engine.backends.extension import _C + +kind, mutation = sys.argv[1:] +x = torch.randn(4, 8, device="cuda") +weight = torch.ones(8, device="cuda") +gate = torch.ones(2, 8, device="cuda") +index = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") +i64 = dict(dtype=torch.int64, device="cuda") +sorted_pos = torch.tensor([0, 2, 1, 3], **i64) +begin, end = torch.tensor([0, 2], **i64), torch.tensor([2, 4], **i64) +first = torch.tensor([0, 1, 2], **i64) +if mutation == "tile_end_past_rows": + end = torch.tensor([2, 9], **i64) +elif mutation == "position_past_rows": + sorted_pos = torch.tensor([0, 2, 1, 4], **i64) +elif mutation == "segments_past_tiles": + first = torch.tensor([0, 1, 5], **i64) +elif mutation == "index_past_table": + index = torch.tensor([0, 1, 0, 7], **i64) +elif mutation == "index_negative": + index = torch.tensor([0, 1, 0, -1], **i64) +if kind == "rmsnorm_backward": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index, sorted_pos, + begin, end, first) +elif kind == "gate_backward": + _C.h3_gate_residual_backward(torch.ones_like(x), x, gate, index, sorted_pos, begin, end, first) +elif kind == "gate_partials": # sequence parallel: explicit row list and tiles + _C.h3_gate_grad_partials(torch.ones_like(x), x, sorted_pos, begin, end) +elif kind == "gate_fold": + _C.h3_gate_grad_fold(torch.ones(2, 8, device="cuda"), first, torch.float32) +elif kind == "rmsnorm_partials": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, None, None, None, + sorted_pos, begin, end, None, None, None) +elif kind == "rmsnorm_fold": + _C.h3_rmsnorm_fold_partials(torch.ones(2, 8, device="cuda"), weight, + torch.ones(2, 16, device="cuda"), first) +elif kind == "rmsnorm_partials_modulated": + # called without the dx entry point having run: the partials check the index values themselves + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, shift, scale, index, + sorted_pos, begin, end, sorted_pos, begin, end) +torch.cuda.synchronize() +""" + +_METADATA_CASES = { + "rmsnorm_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_partials": ("tile_end_past_rows", "position_past_rows"), + "gate_fold": ("segments_past_tiles",), + "rmsnorm_partials": ("tile_end_past_rows", "position_past_rows"), + "rmsnorm_partials_modulated": ("index_past_table", "index_negative"), + "rmsnorm_fold": ("segments_past_tiles",), +} + + +@pytest.mark.parametrize("kind", sorted(_METADATA_CASES)) +def test_segment_metadata_that_does_not_describe_the_rows_ends_in_a_device_assert( + native_extension, kind +): + """Every entry point that reads caller-built tile lists checks them before it reads.""" + + import os + import subprocess + import sys + + if not hasattr(native_extension, "h3_gate_grad_partials"): + pytest.skip("rl_engine._C lacks the H3 SP entry points") + + def run(mutation): + return subprocess.run( + [sys.executable, "-c", _BACKWARD_CHILD, kind, mutation], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + control = run("none") + assert control.returncode == 0, control.stderr + for mutation in _METADATA_CASES[kind]: + done = run(mutation) + assert done.returncode != 0, (kind, mutation) + assert "device-side assert triggered" in done.stderr, (kind, mutation, done.stderr) diff --git a/tools/migration/layout.json b/tools/migration/layout.json index b1e48a46f..58cfd113f 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -493,6 +493,8 @@ "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.png": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.png", "docs/usage/evidence/h3-rmsnorm-b200/figure.png": "reports/experiments/h3-rmsnorm-b200/figure.png", "docs/usage/evidence/h3-rmsnorm-b200/report.json": "reports/experiments/h3-rmsnorm-b200/report.json", + "docs/usage/evidence/h3-sp-norm-adaln-b200/figure.png": "reports/experiments/h3-sp-norm-adaln-b200/figure.png", + "docs/usage/evidence/h3-sp-norm-adaln-b200/report.json": "reports/experiments/h3-sp-norm-adaln-b200/report.json", "docs/usage/evidence/h3-timestep-mlp-b200/figure.png": "reports/experiments/h3-timestep-mlp-b200/figure.png", "docs/usage/evidence/h3-timestep-mlp-b200/report.json": "reports/experiments/h3-timestep-mlp-b200/report.json", "docs/usage/evidence/h3-timestep-sinusoid-b200/figure.png": "reports/experiments/h3-timestep-sinusoid-b200/figure.png", @@ -716,6 +718,7 @@ "rl_engine/kernels/ops/cuda/h3/final_adaln_out.py": "rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py", "rl_engine/kernels/ops/cuda/h3/gate_residual.py": "rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py", "rl_engine/kernels/ops/cuda/h3/rmsnorm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py": "rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py", "rl_engine/kernels/ops/cuda/h3/timestep_mlp.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py", "rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py", @@ -904,7 +907,7 @@ "tests/h3/test_h3_native_validation.py": "tests/models/minimax_h3/test_h3_native_validation.py", "tests/h3/test_h3_report.py": "tests/models/minimax_h3/test_h3_report.py", "tests/h3/test_h3_rmsnorm.py": "tests/models/minimax_h3/test_h3_rmsnorm.py", - "tests/h3/test_h3_rmsnorm_native_validation.py": "tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py", + "tests/h3/test_h3_sp_norm_adaln.py": "tests/distributed/sp/test_h3_sp_norm_adaln.py", "tests/h3/test_h3_timestep_mlp.py": "tests/models/minimax_h3/test_h3_timestep_mlp.py", "tests/h3/test_h3_timestep_sinusoid.py": "tests/models/minimax_h3/test_h3_timestep_sinusoid.py", "tests/h3/test_h3_tp_adaln.py": "tests/distributed/tp/test_h3_tp_adaln.py", diff --git a/tools/validation/models/h3_ws2_evidence.py b/tools/validation/models/h3_ws2_evidence.py index cf9efc466..f44bb4e5a 100644 --- a/tools/validation/models/h3_ws2_evidence.py +++ b/tools/validation/models/h3_ws2_evidence.py @@ -11,6 +11,8 @@ export RL_KERNEL_H3_WEIGHTS= python tools/validation/models/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ --out reports/experiments/h3-tp-adaln-3mod-b200/report.json + python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json """ from __future__ import annotations @@ -24,7 +26,13 @@ import torch # noqa: E402 -from rl_engine.validation.distributed.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 +from rl_engine.validation.distributed.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 from rl_engine.validation.models.h3_weights import ( # noqa: E402 load_h3_conditioning_weights, @@ -56,7 +64,37 @@ def _tp_adaln(worlds: list[int], max_t: int) -> dict: return {"weights": [WEIGHT, BIAS], "runs": runs} -OPS = {"tp_adaln_3mod": _tp_adaln} +SP_CASES = [ + {"batch": 1, "seq": 4097, "layout": "block", "seed": 1}, + {"batch": 1, "seq": 4097, "layout": "interleaved", "seed": 2}, + {"batch": 2, "seq": 4097, "layout": "block", "seed": 3}, + {"batch": 1, "seq": 32768, "layout": "block", "seed": 4}, + {"batch": 1, "seq": 32768, "layout": "interleaved", "seed": 5}, + {"batch": 1, "seq": 131072, "layout": "block", "seed": 6}, +] + + +def _sp_norm(worlds: list[int], _max_t: int) -> dict: + runs = [] + for world in worlds: + ranks = run_world(world, sp_region_sweep, {}, cases=SP_CASES) + runs.append({"sp": world, "ranks": ranks}) + equal = all(all(c["equal"].values()) for r in ranks for c in r["cases"]) + print(f"sp={world}: byte-equal to WS1 on every rank and case: {equal}") + naive = { + layout: naive_sp_mismatch(sp_case(1, 32768, layout=layout, seed=4), max(worlds)) + for layout in ("block", "interleaved") + } + region = "norm2(residual + gate_msa[row] * y) * (1 + scale_mlp[row]) + shift_mlp[row]" + return { + "region": region, + "hidden": 5376, + "runs": runs, + "naive_sp": {"sp": max(worlds), **naive}, + } + + +OPS = {"tp_adaln_3mod": _tp_adaln, "sp_norm_adaln": _sp_norm} def main() -> None: @@ -67,6 +105,10 @@ def main() -> None: parser.add_argument("--out", type=Path, required=True) args = parser.parse_args() worlds = [int(w) for w in args.worlds.split(",")] + if min(worlds) < 1: + parser.error(f"--worlds must be positive, got {args.worlds}") + if args.max_timesteps < 1: + parser.error(f"--max-timesteps must be positive, got {args.max_timesteps}") if torch.cuda.device_count() < max(worlds): raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index 35688f318..bafc137c3 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -454,6 +454,8 @@ def repeat(mode: str) -> str: def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + combined = all(row.get("backward_timing_scope") == "forward_and_backward" for row in perf) + backward_label = "fwd+bwd" if combined else "bwd" _style(ax, "Forward and backward time (log, lower is better)", "ms") _bars( ax, @@ -461,9 +463,14 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: [ ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), - ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), ( - f"{provider} bwd", + f"CUDA {backward_label}", + RL_KERNEL, + [row["candidate_backward_us"] / 1e3 for row in perf], + "////", + ), + ( + f"{provider} {backward_label}", PROVIDER, [row["provider_backward_us"] / 1e3 for row in perf], "////", @@ -473,6 +480,8 @@ def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: log=True, ) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") @@ -504,6 +513,7 @@ def label(name: str, mode: str) -> str: for text in ax.texts: if text.get_text() == f"{floor:.1e}": text.set_text("exact") + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=1) ax.set_ylim(floor / 3, 1e4) ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) @@ -553,43 +563,13 @@ def plot_tp_adaln(report: dict[str, Any]): report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" ) runs = report["runs"] - # Left: every rank, every T, every tensor byte-equal to WS1? - left.set_facecolor(SURFACE) - left.set_title( - "Byte-equal to WS1 (every rank, every T)", loc="left", fontsize=11, color=INK, pad=10 - ) - for y, run in enumerate(runs): - checks = [e for r in run["ranks"] for e in r["equality"]] - for x, (key, _) in enumerate(TP_TENSORS): - ok = sum(bool(e[key]) for e in checks) - good = ok == len(checks) - left.add_patch( - matplotlib.patches.FancyBboxPatch( - (x + 0.06, y + 0.08), - 0.88, - 0.84, - boxstyle="round,pad=0,rounding_size=0.06", - facecolor=RL_KERNEL if good else PROVIDER, - edgecolor=SURFACE, - linewidth=2, - ) - ) - left.text( - x + 0.5, - y + 0.5, - f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} rank x T", - ha="center", - va="center", - fontsize=8, - color="white", - ) - left.set_xlim(0, len(TP_TENSORS)) - left.set_ylim(len(runs), 0) - left.set_xticks([i + 0.5 for i in range(len(TP_TENSORS))], [label for _, label in TP_TENSORS]) - left.set_yticks([i + 0.5 for i in range(len(runs))], [f"TP{run['tp']}" for run in runs]) - left.tick_params(colors=MUTED, labelsize=8, length=0) - for side in left.spines.values(): - side.set_visible(False) + _equality_grid( + left, + "Byte-equal to WS1", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "ranks x T", + ) # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) @@ -644,6 +624,161 @@ def worst(run, key): return fig +def _equality_grid(ax, title, rows, tensors, unit) -> None: + """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" + + ax.set_facecolor(SURFACE) + ax.set_title(f"{title}; checks = {unit}", loc="left", fontsize=11, color=INK, pad=10) + for y, (_, checks) in enumerate(rows): + for x, (key, _) in enumerate(tensors): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + ax.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)}" + ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") + ax.set_xlim(0, len(tensors)) + ax.set_ylim(len(rows), 0) + ax.set_xticks([i + 0.5 for i in range(len(tensors))], [label for _, label in tensors]) + ax.set_yticks([i + 0.5 for i in range(len(rows))], [name for name, _ in rows]) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + for side in ax.spines.values(): + side.set_visible(False) + + +SP_TENSORS = ( + ("out", "output rows"), + ("d_residual", "d_residual"), + ("d_y", "d_sublayer"), + ("d_weight", "d_norm_w"), + ("d_table", "d_table"), +) + + +def plot_sp_norm(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(15, 3.9), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.35, 1, 0.9]} + ) + env = report["environment"] + fig.suptitle( + "sp_norm_adaln (sequence-parallel norm2(residual + gate * y) with AdaLN modulation, " + "H = 5376, NCCL)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.01, + f"{env['gpu']} x {env['gpus']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · bf16", + fontsize=7, + color=MUTED, + ) + runs = report["runs"] + _equality_grid( + axes[0], + "Byte-equal to WS1", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "ranks x cases", + ) + + # Middle: the largest block-layout case, slowest rank. + big = max( + (c for c in runs[0]["ranks"][0]["cases"] if c["layout"] == "block"), + key=lambda c: c["seq"] * c["batch"], + ) + + def worst(run, key): + return ( + max( + c[key] + for r in run["ranks"] + for c in r["cases"] + if (c["seq"], c["batch"], c["layout"]) == (big["seq"], big["batch"], big["layout"]) + ) + / 1e3 + ) + + sp_runs = [run for run in runs if run["sp"] > 1] + ax = axes[1] + _style(ax, f"Region fwd+bwd, S = {big['seq']}, block layout", "ms (slowest rank)") + if not sp_runs: + ax.text( + 0.5, + 0.5, + "no run with SP > 1 in this report", + transform=ax.transAxes, + ha="center", + va="center", + fontsize=10, + color=MUTED, + ) + else: + _bars( + ax, + [f"SP{run['sp']}" for run in sp_runs], + [("SP", RL_KERNEL, [worst(run, "sp_forward_backward_us") for run in sp_runs])], + "{:.2f}", + ) + ws1 = worst(runs[0], "ws1_forward_backward_us") + ax.axhline( + ws1, + color=PROVIDER, + linestyle="--", + linewidth=1.4, + zorder=1, + label=f"WS1, 1 GPU: {ws1:.2f}", + ) + ax.set_ylim(0, ws1 * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper right") + + # Right: what a naive SP reduction would do instead. + naive = report["naive_sp"] + ax = axes[2] + _style(ax, f"Naive SP{naive['sp']} vs WS1 (S = 32768)", "% of elements that differ") + _bars( + ax, + ["d_norm_w", "d_table"], + [ + ("naive, block", PROVIDER, [100 * naive["block"][k] for k in ("d_weight", "d_table")]), + ( + "naive, interleaved", + PROVIDER, + [100 * naive["interleaved"][k] for k in ("d_weight", "d_table")], + "////", + ), + ("RL-Kernel SP", RL_KERNEL, [0.0, 0.0]), + ], + "{:.1f}", + ) + ax.text( + 0.98, + 0.62, + "naive: each rank's WS1 backward,\nthen a rank-order sum", + transform=ax.transAxes, + ha="right", + fontsize=7, + color=MUTED, + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -653,6 +788,7 @@ def worst(run, key): "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, }