From e7bc53aa89f3c4b2c3d6045c0b2ccfec8c091391 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 02:51:41 +0000 Subject: [PATCH 01/62] [WS1][CUDA][MiniMax-H3] adaln_gate_residual: gated residual with in-kernel gate gather RFC #420 step 6, row `adaln_gate_residual` (stacked on h3_rmsnorm). - residual + gate[adaln_index] * sublayer_output in one pass, the gate row read from the AdaLN table view in the kernel; the two roundings sit where the eager expression rounds (__fmul_rn/__fadd_rn, no FMA contraction), so the forward is bitwise equal to diffusers in bf16/fp16/fp32. 1.7x faster forward at S = 131072. - Backward: d_residual = grad; d_sublayer bitwise equal to the eager VJP (exact product, one rounding); d_gate a deterministic FP32 segment sum (diffusers: non-deterministic BF16 atomics, 13-49x larger error, varying run to run). The PyTorch reference gets an equally deterministic gate gradient. - Chain stage residual + gate_msa * (stand-in attention output), covered by the end-to-end test; registry, gtest, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 242 ++++++++++++++++++ csrc/ops.cpp | 15 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-adaln-gate-residual.md | 91 +++++++ rl_engine/_C.pyi | 13 + rl_engine/kernels/gtest/operator_inputs.py | 18 ++ rl_engine/kernels/gtest/operator_specs.py | 11 + .../kernels/ops/cuda/h3/gate_residual.py | 83 ++++++ .../kernels/ops/pytorch/h3/gate_residual.py | 97 +++++++ rl_engine/kernels/registry.py | 11 + rl_engine/testing/h3_chain.py | 38 +++ rl_engine/testing/h3_provider.py | 8 + rl_engine/testing/h3_report.py | 107 ++++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 72 ++++++ setup.py | 1 + tests/h3/test_h3_adaln_gate_residual.py | 115 +++++++++ tests/test_ws1_gtest_gpu.py | 1 + 19 files changed, 926 insertions(+) create mode 100644 csrc/cuda/h3/gate_residual.cu create mode 100644 docs/operators/h3-adaln-gate-residual.md create mode 100644 rl_engine/kernels/ops/cuda/h3/gate_residual.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/gate_residual.py create mode 100644 tests/h3/test_h3_adaln_gate_residual.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..18ede9e72 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,242 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 gated residual (RFC #420 `adaln_gate_residual`). +// +// out[r, j] = cast(residual[r, j] + cast(gate[i, j] * y[r, j])) i = index[r % S] +// +// The gate row is read from a strided view of the AdaLN table (gate_msa or +// gate_mlp) inside the kernel, and the two roundings sit exactly where the +// eager `residual + gate.index_select(0, i) * y` rounds (__fmul_rn/__fadd_rn +// keep FP32 from contracting them into one FMA), so the result is +// bitwise equal to diffusers. Elementwise: every output depends only on its +// own inputs. +// +// Backward: d_residual = grad; d_y = cast(grad * gate) (the product of two +// 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. + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +struct Gate { + const void* rows; // (R, N) view, row stride `stride` + int64_t stride; + const int64_t* index; // (S,) + int64_t seq; +}; + +// One block per row (grid-strided); 16-byte vectors when the row allows it. +// `kVecElems` elements are loaded together; each is computed independently. +template +__global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t rows, int64_t n, + Gate gate, bool vectorized) { + constexpr int kVecElems = 16 / sizeof(T); + for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { + const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const T* ar = a + r * n; + const T* br = kDy ? nullptr : b + r * n; + T* orow = out + r * n; + if (vectorized) { + for (int64_t v = threadIdx.x; v < n / kVecElems; v += blockDim.x) { + const uint4 av = reinterpret_cast(ar)[v]; + const uint4 gv = reinterpret_cast(g)[v]; + uint4 bv = av; + if (!kDy) bv = reinterpret_cast(br)[v]; + const T* ae = reinterpret_cast(&av); + const T* ge = reinterpret_cast(&gv); + const T* be = reinterpret_cast(&bv); + uint4 ov; + T* oe = reinterpret_cast(&ov); +#pragma unroll + for (int k = 0; k < kVecElems; ++k) { + if (kDy) { + oe[k] = from_f(to_f(ae[k]) * to_f(ge[k])); // grad * gate + } else { + const float p = round_t(__fmul_rn(to_f(ge[k]), to_f(be[k]))); // gate * y + oe[k] = from_f(__fadd_rn(to_f(ae[k]), p)); // residual + p + } + } + reinterpret_cast(orow)[v] = ov; + } + } else { + for (int64_t j = threadIdx.x; j < n; j += blockDim.x) { + if (kDy) { + orow[j] = from_f(to_f(ar[j]) * to_f(g[j])); + } else { + const float p = round_t(__fmul_rn(to_f(g[j]), to_f(br[j]))); + orow[j] = from_f(__fadd_rn(to_f(ar[j]), p)); + } + } + } + } +} + +// partial[tile, j] = sum over sorted positions of the tile (ascending) of grad * y +template +__global__ void gate_grad_partial_kernel(const T* __restrict__ grad, const T* __restrict__ y, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t n) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t e = sorted_pos[p] * n + j; + acc = fmaf(to_f(grad[e]), to_f(y[e]), acc); + } + partial[tile * n + j] = acc; +} + +template +__global__ void gate_grad_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + T* __restrict__ out, int64_t n) { + const int64_t seg = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * n + j]; + out[seg * n + j] = from_f(acc); +} + +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torch::Tensor& x) { + TORCH_CHECK(gate.is_cuda() && gate.device() == x.device() && gate.scalar_type() == x.scalar_type(), + "gate must be a CUDA tensor with the activations' dtype"); + TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, + "gate must be (R, N) with unit column stride"); + TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && + index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; +} + +void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device() && t.scalar_type() == like.scalar_type(), + name, " must match the residual's device and dtype"); + TORCH_CHECK(t.sizes() == like.sizes() && t.is_contiguous(), name, + " must be contiguous with the residual's shape"); +} + +unsigned row_blocks(int64_t rows) { + return static_cast(std::min(rows, 1 << 30)); +} + +bool can_vectorize(const torch::Tensor& gate, int64_t n, std::initializer_list acts) { + const int64_t vec = 16 / gate.element_size(); + if (n % vec != 0 || (gate.stride(0) * gate.element_size()) % 16 != 0) return false; + if (reinterpret_cast(gate.data_ptr()) % 16 != 0) return false; + for (const auto* t : acts) { + if (reinterpret_cast(t->data_ptr()) % 16 != 0) return false; + } + return true; +} + +} // namespace + +// residual, y: (M, N); gate: (R, N) view; index: (S,), M % S == 0. +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && residual.is_contiguous() && + residual.numel() > 0, + "residual must be a non-empty contiguous (M, N) CUDA tensor"); + check_act(y, "y", residual); + const Gate g = make_gate(gate, index, residual); + const c10::cuda::CUDAGuard guard(residual.device()); + auto out = torch::empty_like(residual); + auto stream = at::cuda::getCurrentCUDAStream(); + const bool vec = can_vectorize(gate, residual.size(1), {&residual, &y, &out}); + H3_DISPATCH(residual.scalar_type(), "h3_gate_residual_forward", [&] { + gate_residual_rows_kernel<<>>( + reinterpret_cast(residual.data_ptr()), reinterpret_cast(y.data_ptr()), + reinterpret_cast(out.data_ptr()), residual.size(0), residual.size(1), g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Returns {d_y, d_gate (R, N) in the gate dtype}. d_residual is the incoming grad. +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile) { + 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); + const Gate g = make_gate(gate, index, grad); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors"); + } + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dy = torch::empty_like(y); + auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); + auto dgate = torch::empty({segments, n}, grad.options()); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { + const T* gp = reinterpret_cast(grad.data_ptr()); + const T* yp = reinterpret_cast(y.data_ptr()); + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + gp, nullptr, reinterpret_cast(dy.data_ptr()), rows, n, g, vec); + if (tiles > 0) { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, stream>>>( + gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), n); + } + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dy, dgate}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index f2288efe4..08085e721 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -527,6 +527,14 @@ std::vector h3_rmsnorm_backward( c10::optional index, c10::optional sorted_pos, c10::optional tile_begin, c10::optional tile_end, c10::optional seg_first_tile); +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index); +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile); #endif // PyBind11 Module Registration @@ -819,5 +827,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_residual_forward", &h3_gate_residual_forward, + "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)", + py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index")); + m.def("h3_gate_residual_backward", &h3_gate_residual_backward, + "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")); #endif } diff --git a/docs/.nav.yml b/docs/.nav.yml index af8c1863f..f6db1bc36 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -35,6 +35,7 @@ nav: - operators/h3-adaln-projection.md - operators/h3-adaln-row-gather.md - operators/h3-rmsnorm.md + - operators/h3-adaln-gate-residual.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index b687a65b5..7042d4561 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -37,4 +37,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) +- [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..7a3010185 --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,91 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and about 35× further from the FP64 golden. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python scripts/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op adaln_gate_residual \ + --out docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 01e6e5348..43fb55ff5 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -377,3 +377,16 @@ def h3_rmsnorm_backward( tile_end: torch.Tensor | None = None, seg_first_tile: torch.Tensor | None = None, ) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 02f7f2e37..6afafc3ae 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -51,6 +51,7 @@ def make_operator_inputs( "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, "adaln_row_gather": _make_adaln_row_gather_inputs, "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, } try: return builders[op_name](args, dtype, device) @@ -89,6 +90,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" f"->{batch * seq}", "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -446,6 +448,22 @@ def _make_h3_rmsnorm_inputs( } +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 197db8874..5254aba2e 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -311,6 +311,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("x", "weight"), ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp", + "cuda": "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp", + }, + grad_input_names=("residual", "y", "gate"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/gate_residual.py b/rl_engine/kernels/ops/cuda/h3/gate_residual.py new file mode 100644 index 000000000..2a0fc9e1d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + 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, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/gate_residual.py b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py new file mode 100644 index 000000000..d2be659d5 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.testing.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 8fa1d1453..bcfc897b8 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -223,6 +223,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) PYTORCH_H3_RMSNORM = "rl_engine.kernels.ops.pytorch.h3.rmsnorm.NativeH3RMSNormOp" CUDA_H3_RMSNORM = "rl_engine.kernels.ops.cuda.h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = ( + "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp" + ) + CUDA_H3_GATE_RESIDUAL = "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -650,6 +654,10 @@ def __init__(self): OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, ], "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], }, "rocm": { "logp": [ @@ -695,6 +703,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -744,6 +753,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -774,6 +784,7 @@ def __init__(self): "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index fe010c7e3..7b13fc60f 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,6 +33,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -146,6 +147,35 @@ def _norm1_modulated(ctx: dict[str, Any], modulation, call): ) ) +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + # The backward replay covers the conditioning chain up to the row gather. CONDITIONING_STAGES = ( "timestep_sinusoid_h3", @@ -199,6 +229,14 @@ def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dic ) .to(torch.bfloat16) .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), } diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 05c7ebec9..8c084afc7 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -81,3 +81,11 @@ def provider_norm_modulate( return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( 0, indices ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index dfecec419..021dab4b0 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_gate_residual, provider_norm_modulate, provider_time_embedder, provider_time_proj, @@ -467,12 +468,117 @@ def golden(x_, w_, sh_, sc_): } +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + def backward(fn, residual=residual, y=y, table=table, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": lambda b=backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward": lambda b=backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, "adaln_projection_3mod": _projection_perf, "adaln_row_gather": _gather_perf, "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, @@ -480,4 +586,5 @@ def golden(x_, w_, sh_, sc_): "adaln_projection_3mod": _projection_accuracy, "adaln_row_gather": _gather_accuracy, "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index 49d502f12..ee8e2c0ff 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -232,6 +232,7 @@ def fake_load_backend(backend): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index da8bc6774..2c62208b7 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -409,12 +409,84 @@ def repeat(mode: str) -> str: return fig +def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + _style(ax, "Forward and backward time (log, lower is better)", "ms") + _bars( + ax, + [row["case"] for row in perf], + [ + ("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", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + + +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") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + for text in ax.texts: + if text.get_text() == f"{floor:.1e}": + text.set_text("exact") + ax.set_ylim(floor / 3, 1e4) + ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, "adaln_projection_3mod": plot_projection, "adaln_row_gather": plot_gather, "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, } diff --git a/setup.py b/setup.py index fb671d0f0..3f38e1fbe 100644 --- a/setup.py +++ b/setup.py @@ -161,6 +161,7 @@ def get_extensions(): cuda_sources.append("csrc/cuda/h3/det_linear.cu") cuda_sources.append("csrc/cuda/h3/adaln_row_gather.cu") cuda_sources.append("csrc/cuda/h3/rmsnorm_modulate.cu") + cuda_sources.append("csrc/cuda/h3/gate_residual.cu") # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") diff --git a/tests/h3/test_h3_adaln_gate_residual.py b/tests/h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..d67210257 --- /dev/null +++ b/tests/h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.h3.gate_residual import NativeH3GateResidualOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_gate_residual + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 2b165c088..32c6322c3 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -54,6 +54,7 @@ def test_all_ws1_single_ops_are_registered(): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", } <= names From cd688b354823109b19b1abd24d8fdfeb5b7a7fc5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:17:29 +0000 Subject: [PATCH 02/62] [WS1][MiniMax-H3] adaln_gate_residual: B200 evidence report and figure Produced by scripts/h3_evidence.py at a26b41a on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 7a3010185..f55412449 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -51,7 +51,7 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) once. Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is - non-deterministic and about 35× further from the FP64 golden. + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). ## Performance Notes @@ -63,14 +63,22 @@ B200, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.21 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.40 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.31 ms | 10.6 ms | +| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | +| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | +| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), +written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + ## Tests ```bash From e7428b1e58ea71f0d036f9733faf31b76e0d4d44 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:02:12 +0000 Subject: [PATCH 03/62] [WS1][CUDA][MiniMax-H3] final_adaln_out: norm_out projection + final norm/modulation RFC #420 step 7, row `final_adaln_out` (stacked on adaln_gate_residual). - norm_out as one autograd node: FP32 SiLU -> one BF16 cast -> the deterministic tensor-core GEMV on norm_out.linear (shift first), then the h3_rmsnorm kernel indexed by timestep_indices. Forward equal to diffusers on 99.96% of elements (1-ULP projection ties); rows bitwise batch/position invariant; 2.4x faster forward at S = 131072. - Backward keeps the table gradient in FP32 into the projection backward: d_temb/dW/db 13-57x closer to FP64 than diffusers (whose error varies run to run), all gradients repeat-bitwise. - Golden rounds only at the declared boundaries (SiLU cast, BF16 table), straight-through for the gradient; gtest inputs at realistic scale. - Chain stage: the gated residual stream -> norm_out with the MLP stage's temb (end-to-end test now runs timestep -> ... -> norm_out); registry, report/figure hooks, doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-final-adaln-out.md | 89 ++++++++++++ rl_engine/kernels/gtest/operator_inputs.py | 27 ++++ rl_engine/kernels/gtest/operator_specs.py | 11 ++ .../kernels/ops/cuda/h3/final_adaln_out.py | 110 +++++++++++++++ .../kernels/ops/pytorch/h3/final_adaln_out.py | 95 +++++++++++++ rl_engine/kernels/registry.py | 11 ++ rl_engine/testing/h3_chain.py | 28 ++++ rl_engine/testing/h3_provider.py | 18 +++ rl_engine/testing/h3_report.py | 92 +++++++++++++ rl_engine/tests/test_dispatch.py | 1 + scripts/plot_h3_evidence.py | 17 +++ tests/h3/test_h3_final_adaln_out.py | 129 ++++++++++++++++++ tests/test_ws1_gtest_gpu.py | 1 + 15 files changed, 631 insertions(+) create mode 100644 docs/operators/h3-final-adaln-out.md create mode 100644 rl_engine/kernels/ops/cuda/h3/final_adaln_out.py create mode 100644 rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py create mode 100644 tests/h3/test_h3_final_adaln_out.py diff --git a/docs/.nav.yml b/docs/.nav.yml index f6db1bc36..1ef4fa017 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -36,6 +36,7 @@ nav: - operators/h3-adaln-row-gather.md - operators/h3-rmsnorm.md - operators/h3-adaln-gate-residual.md + - operators/h3-final-adaln-out.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 7042d4561..27ffce7d0 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -38,4 +38,5 @@ Every operator page should include: - [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) +- [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.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 new file mode 100644 index 000000000..e05f1d15c --- /dev/null +++ b/docs/operators/h3-final-adaln-out.md @@ -0,0 +1,89 @@ +# MiniMax-H3 Final AdaLN Output + +## Summary + +`final_adaln_out` is H3's `norm_out` (`MiniMaxH3AdaLayerNormOut`), which runs once after the 50 +blocks (RFC #420, WS1 step 7): + +```text +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] +``` + +The table has one row per distinct timestep and is indexed by `timestep_indices`, not by +`adaln_indices`. Diffusers then upcasts `out` to the FP32 output heads' dtype. That cast is exact +and is left to the caller. + +## Entry Point + +```python +op = kernel_registry.get_op("final_adaln_out", device="cuda") +out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timestep_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+) | `rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp` | `rl_engine._C.h3_det_linear_*`, `rl_engine._C.h3_rmsnorm_*` | One autograd node | +| PyTorch reference | `rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp` | n/a | `forward`: diffusers replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Numerics + +- **Projection.** The [AdaLN projection](h3-adaln-projection.md)'s deterministic tensor-core + GEMV (`h3-det-linear-bf16-mma-v1`), applied to `norm_out.linear`. The FP32 SiLU is rounded + once at the declared cast. BF16 `temb` is rejected (probe H7). +- **Norm and modulation.** The [`h3_rmsnorm`](h3-rmsnorm.md) kernel indexed by + `timestep_indices`. Given the same table, it is bitwise equal to diffusers. Overall, about + 0.04% of output elements differ from diffusers by 1 ULP, all traced to the projection's + summation tree versus cuBLAS. Rows are batch- and position-invariant. +- **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into + the projection backward without an extra BF16 rounding: + - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers (whose error varies from + run to run), which also rounds the + table gradient to BF16 and accumulates it with `index_select` atomics; + - `dx` and `d_norm_w` are at the BF16 rounding level for both; + - every gradient is repeat-bitwise. +- **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, + and `norm_out.linear`'s BF16 output table, which every position of a timestep shares. Both are + rounded straight-through for the gradient. + +## Performance Notes + +```bash +python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out +``` + +B200, B = 1, T = 3, pinned `norm_out` weights: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.13 ms | 0.15 ms | 1.62 ms | 1.09 ms | +| 32768 | 0.39 ms | 0.78 ms | 2.86 ms | 5.32 ms | +| 131072 | 1.24 ms | 2.93 ms | 7.78 ms | 20.7 ms | + +At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the +projection backward. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_final_adaln_out.py -v # operator +python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out +python scripts/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 64 --normalized-dim 5376 --check-grad +python scripts/h3_evidence.py --op final_adaln_out \ + --out docs/usage/evidence/h3-final-adaln-out-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/report.json +``` + +## Known Limitations + +- **gtest at larger S.** With `--seq 257`, a few gradient elements fail the per-element + contract tolerance (BF16 `d_norm_w`, FP32 `d_temb`). Those elements are what is left after + summing hundreds of rows into values around 1e2, which an absolute per-element atol can't + cover. Relative to the gradient's maximum the error is 4e-7 in FP32 and 5e-3 in BF16, and the + PyTorch reference fails the same elements, plus more. `--seq 64` passes in both dtypes. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 6afafc3ae..4abc17a80 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -52,6 +52,7 @@ def make_operator_inputs( "adaln_row_gather": _make_adaln_row_gather_inputs, "h3_rmsnorm": _make_h3_rmsnorm_inputs, "adaln_gate_residual": _make_adaln_gate_residual_inputs, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -91,6 +92,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: f"->{batch * seq}", "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -464,6 +466,31 @@ def _make_adaln_gate_residual_inputs( } +def _make_final_adaln_out_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract; norm_out.linear/norm/x use ``dtype`` (BF16 in H3). + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + num_timesteps = _h3_num_timesteps(args) + # Realistic scales (norm weight in [0.5, 1], unit activations): with unscaled + # inputs the parameter gradients sum hundreds of rows into values ~1e3, and + # their cancelling elements cannot meet a per-element FP32 atol. + linear = _floating_tensor((2 * hidden, H3_TIME_EMBED), args, torch.float32, device, 3) + norm = 0.5 + 0.5 * torch.sigmoid(_floating_tensor((hidden,), args, torch.float32, device, 4)) + generator = _generator(args, device, offset=37) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "norm_weight": norm.to(dtype), + "temb": _floating_tensor((num_timesteps, H3_TIME_EMBED), args, torch.float32, device, 1), + "weight": (linear * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (_floating_tensor((2 * hidden,), args, torch.float32, device, 2) * 0.1).to(dtype), + "timestep_indices": torch.randint( + 0, num_timesteps, (seq,), generator=generator, device=device + ), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 5254aba2e..d7f53b754 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -322,6 +322,17 @@ def _load_object(path: str) -> Any: }, grad_input_names=("residual", "y", "gate"), ), + "final_adaln_out": OperatorSpec( + name="final_adaln_out", + op_class="reduction", + gold_path="rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp"), + "cuda": "rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp", + }, + grad_input_names=("x", "norm_weight", "temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py b/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py new file mode 100644 index 000000000..f2700ef06 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/final_adaln_out.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 final AdaLN output (RFC #420 ``final_adaln_out``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.cuda.h3.rmsnorm import h3_rmsnorm_available +from rl_engine.kernels.ops.pytorch.h3.final_adaln_out import validate_h3_final_adaln_out +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import H3_NORM_EPS + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_rmsnorm_forward[modulated]" +) +BACKWARD_IMPL = "fp32_table_grad_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_rmsnorm_backward+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.kernels.ops.cuda.h3.det_linear.silu_backward_fp32" +) + + +class _H3FinalAdaLNOutCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): + hidden = x.shape[-1] + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) # declared boundary: one cast after the FP32 SiLU + (table,) = det_linear_forward(act, weight, bias) + shift, scale = table.chunk(2, dim=-1) + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward( + x2, norm_weight.contiguous(), float(eps), shift, scale, timestep_indices + ) + ctx.save_for_backward(x2, norm_weight, rstd, temb, act, weight, table, timestep_indices) + ctx.x_shape = x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + shift, scale = table.chunk(2, dim=-1) + positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + tiles = _segment_tiles(positions, table.shape[0]) + dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( + grad.contiguous().view_as(x2), + x2, + norm_weight.contiguous(), + rstd, + shift, + scale, + timestep_indices, + *tiles, + ) + # The table gradient stays FP32 into the projection backward. + d_table = torch.cat([d_shift, d_scale], dim=1) + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + record_backward( + "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + + +class H3FinalAdaLNOutCudaOp: + """CUDA candidate: tensor-core shift/scale projection + fused norm/modulation. + + The projection is ``adaln_projection_3mod``'s deterministic GEMV + (``h3-det-linear-bf16-mma-v1``) on ``norm_out.linear``; the norm and + modulation are ``h3_rmsnorm``'s kernel indexed by ``timestep_indices``. + One autograd node, so the table gradient reaches the projection in FP32. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and h3_rmsnorm_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + if weight.dtype != x.dtype: + raise TypeError("norm_out.linear must share the activations' dtype") + if not x.is_cuda: + raise ValueError("H3FinalAdaLNOutCudaOp needs CUDA tensors") + return _H3FinalAdaLNOutCuda.apply( + x, norm_weight, temb, weight, bias, timestep_indices.contiguous(), eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps).float() diff --git a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py new file mode 100644 index 000000000..94aaa28c1 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 final AdaLN output (RFC #420 row ``final_adaln_out``). + +``MiniMaxH3AdaLayerNormOut`` runs after the 50 blocks:: + + shift, scale = linear(silu(temb).to(bf16)).chunk(2) # (T, H) each, shift first + out = norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] + +``linear`` is ``norm_out.linear`` (``2688 -> 2 * 5376``, BF16), ``norm`` is +``norm_out.norm`` (``nn.RMSNorm(5376, eps=1e-5)``), and the table is indexed +by ``timestep_indices`` (one row per distinct timestep), not by +``adaln_indices``. Diffusers then casts the result to the FP32 output heads' +dtype, an exact upcast left to the caller. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.rmsnorm import ( + H3_NORM_EPS, + NativeH3RMSNormOp, + validate_h3_rmsnorm, +) + + +def validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, *, same_dtype: bool = True +) -> int: + hidden = x.shape[-1] + if same_dtype: + validate_h3_rmsnorm(x, norm_weight, eps) + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if temb.dim() != 2 or temb.shape[0] == 0: + raise ValueError(f"temb must be a non-empty (T, D) matrix, got {tuple(temb.shape)}") + if weight.shape != (2 * hidden, temb.shape[1]) or bias.shape != (2 * hidden,): + raise ValueError( + f"norm_out.linear must be ({2 * hidden}, {temb.shape[1]}) with a ({2 * hidden},) " + f"bias, got {tuple(weight.shape)} / {tuple(bias.shape)}" + ) + if bias.dtype != weight.dtype: + raise TypeError("norm_out.linear weight and bias must share a dtype") + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if timestep_indices.dtype != torch.int64 or timestep_indices.dim() != 1: + raise TypeError("timestep_indices must be a 1-D int64 tensor") + if x.dim() < 2 or x.shape[-2] != timestep_indices.shape[0]: + raise ValueError("x must be (..., S, H) with S = len(timestep_indices)") + if bool(((timestep_indices < 0) | (timestep_indices >= temb.shape[0])).any()): + raise IndexError(f"timestep_indices must lie in [0, {temb.shape[0]})") + return hidden + + +class NativeH3FinalAdaLNOutOp: + """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only at the + module boundaries the model declares (the SiLU cast and the BF16 table), + straight-through for the gradient.""" + + op_class = "reduction" + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + return NativeH3RMSNormOp().forward_modulated( + x, norm_weight, shift, scale, timestep_indices, eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, same_dtype=False + ) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + act = act + (act.to(weight.dtype).double() - act).detach() # declared cast, identity VJP + table = F.linear(act, weight.double(), bias.double()) + # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the + # weight dtype, and every position of a timestep shares that row, so the + # rounding is model semantics (straight-through for the gradient). + table = table + (table.to(weight.dtype).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() + out = n * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + return out.float() diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index bcfc897b8..39255a9d4 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -227,6 +227,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): "rl_engine.kernels.ops.pytorch.h3.gate_residual.NativeH3GateResidualOp" ) CUDA_H3_GATE_RESIDUAL = "rl_engine.kernels.ops.cuda.h3.gate_residual.H3GateResidualCudaOp" + PYTORCH_H3_FINAL_ADALN_OUT = ( + "rl_engine.kernels.ops.pytorch.h3.final_adaln_out.NativeH3FinalAdaLNOutOp" + ) + CUDA_H3_FINAL_ADALN_OUT = "rl_engine.kernels.ops.cuda.h3.final_adaln_out.H3FinalAdaLNOutCudaOp" def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: @@ -658,6 +662,10 @@ def __init__(self): OpBackend.CUDA_H3_GATE_RESIDUAL, OpBackend.PYTORCH_H3_GATE_RESIDUAL, ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -704,6 +712,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "musa": { "logp": [OpBackend.TRITON_LOGP, OpBackend.PYTORCH_NATIVE], @@ -754,6 +763,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -785,6 +795,7 @@ def __init__(self): "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 7b13fc60f..f4a9599f9 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -33,6 +33,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_final_adaln_out, provider_gate_residual, provider_norm_modulate, provider_time_embedder, @@ -176,6 +177,33 @@ def _norm1_modulated(ctx: dict[str, Any], modulation, call): ) ) + +def _final(ctx: dict[str, Any], call): + weights = ctx["weights"] + history = ctx["history"] + return call( + history["adaln_gate_residual"], # the residual stream after the gated sublayer + weights["norm_out.norm.weight"], + history["timestep_mlp_fp32"], # temb + weights["norm_out.linear.weight"], + weights["norm_out.linear.bias"], + ctx["timestep_indices"], + ) + + +STAGES.append( + Stage( + name="final_adaln_out", + op_type="final_adaln_out", + candidate=lambda op, ctx, _up: _final(ctx, op), + provider=lambda ctx, _up: _final(ctx, provider_final_adaln_out), + golden=lambda op, ctx, _up: _final(ctx, op.forward_fp32), + provider_bitwise_isolated=False, # shift/scale projection: tensor-core tree vs cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + # The backward replay covers the conditioning chain up to the row gather. CONDITIONING_STAGES = ( "timestep_sinusoid_h3", diff --git a/rl_engine/testing/h3_provider.py b/rl_engine/testing/h3_provider.py index 8c084afc7..c654611ec 100644 --- a/rl_engine/testing/h3_provider.py +++ b/rl_engine/testing/h3_provider.py @@ -89,3 +89,21 @@ def provider_gate_residual( """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" return residual + gate.index_select(0, indices) * sublayer_output + + +def provider_final_adaln_out( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + temb: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + timestep_indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``MiniMaxH3AdaLayerNormOut.forward``: shift first, indexed by timestep.""" + + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 021dab4b0..41f8d647d 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -20,6 +20,7 @@ from rl_engine.testing.h3_provider import ( provider_adaln_modulation, provider_adaln_row_gather, + provider_final_adaln_out, provider_gate_residual, provider_norm_modulate, provider_time_embedder, @@ -572,6 +573,95 @@ def grads(fn, dtype=None): } +# --------------------------------------------------------------------------- # +# final_adaln_out (norm_out) +# --------------------------------------------------------------------------- # + +FINAL_NAMES = ["norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias"] + + +def _final_inputs(seq: int, seed: int = 0): + norm_weight, weight, bias = h3_params(FINAL_NAMES, [(5376,), (10752, 2688), (10752,)]) + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, _ = h3_packed_layout(seq, 3, seed=seed) + return x, norm_weight.bfloat16(), temb, weight.bfloat16(), bias.bfloat16(), ti + + +def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("final_adaln_out", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, nw, temb, w, b, ti = _final_inputs(seq) + grad = torch.randn_like(x) + + def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + leaves = [t.detach().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + + cases.append( + { + "op": "final_adaln_out", + "case": f"S={seq}", + "backend": type(op).__name__, + # read x + norm_out.linear, write out + "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, + "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), + "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), + "provider_backward": lambda bw=backward, i=ti: bw( + lambda *t: provider_final_adaln_out(*t, i) + ), + } + ) + return cases + + +def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("final_adaln_out", device="cuda") + golden_op = registry._get_or_create_backend( + registry._priority_map["cpu"]["final_adaln_out"][-1] + ) + x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) + ours = op(x, nw, temb, w, b, ti) + theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) + golden = golden_op.forward_fp32(x, nw, temb, w, b, ti) + invariant = bool( + torch.equal(op(x[:, 100:160], nw, temb, w, b, ti[100:160])[0], ours[0, 100:160]) + ) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, nw, temb, w, b)] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda *t: golden_op.forward_fp32(*t, ti)) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(*t, ti)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), + "rel_error": { + key: float((g.double() - r.double()).abs().max() / r.double().abs().max()) + for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) + }, + } + return { + "equal_to_diffusers_fraction": float((ours == theirs).float().mean()), + "max_abs_vs_golden": float((ours.float() - golden).abs().max()), + "provider_max_abs_vs_golden": float((theirs.float() - golden).abs().max()), + "rows_batch_invariant": invariant, + "backward": backward, + } + + PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { "timestep_sinusoid_h3": _sinusoid_perf, "timestep_mlp_fp32": _mlp_perf, @@ -579,6 +669,7 @@ def grads(fn, dtype=None): "adaln_row_gather": _gather_perf, "h3_rmsnorm": _norm_perf, "adaln_gate_residual": _gate_perf, + "final_adaln_out": _final_perf, } ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { "timestep_sinusoid_h3": _sinusoid_accuracy, @@ -587,4 +678,5 @@ def grads(fn, dtype=None): "adaln_row_gather": _gather_accuracy, "h3_rmsnorm": _norm_accuracy, "adaln_gate_residual": _gate_accuracy, + "final_adaln_out": _final_accuracy, } diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index ee8e2c0ff..1e12e3dc5 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -233,6 +233,7 @@ def fake_load_backend(backend): "adaln_row_gather", "h3_rmsnorm", "adaln_gate_residual", + "final_adaln_out", ) diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 2c62208b7..af4496904 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -480,6 +480,22 @@ def plot_gate_residual(report: dict[str, Any]): return fig +def plot_final(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "final_adaln_out (norm_out: projection + final norm, T = 3)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + _backward_errors( + right, + acc["backward"], + ["dx", "d_norm_w", "d_temb", "dW", "db"], + f"forward equal to diffusers on {100 * acc['equal_to_diffusers_fraction']:.2f}% " + "of elements (projection 1-ULP ties)", + ) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -487,6 +503,7 @@ def plot_gate_residual(report: dict[str, Any]): "adaln_row_gather": plot_gather, "h3_rmsnorm": plot_rmsnorm, "adaln_gate_residual": plot_gate_residual, + "final_adaln_out": plot_final, } diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..f322cd93d --- /dev/null +++ b/tests/h3/test_h3_final_adaln_out.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``final_adaln_out``: norm_out shift/scale projection + final norm/modulation. + +* layout: ``shift`` is the first half of ``norm_out.linear``'s output, rows + are indexed by ``timestep_indices`` (not ``adaln_indices``); +* forward matches diffusers' ``norm_out`` to the projection's 1-ULP ties + (the norm and modulation are bitwise; the GEMV tree differs from cuBLAS); +* backward is deterministic, with the table gradient kept in FP32; +* rows are batch/position invariant; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.pytorch.h3.final_adaln_out import NativeH3FinalAdaLNOutOp +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_final_adaln_out + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.kernels.ops.cuda.h3.final_adaln_out import H3FinalAdaLNOutCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3FinalAdaLNOutCudaOp() + + +@pytest.fixture(scope="module") +def norm_out(h3_weights_cpu): + return [ + h3_weights_cpu[name].cuda() + for name in ("norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias") + ] + + +def _inputs(seq, num_timesteps=3, seed=0, batch=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(batch, seq, 5376, generator=g) * 2).bfloat16().cuda() + temb = (torch.randn(num_timesteps, 2688, generator=g) * 2).cuda() + ti, _ = h3_packed_layout(seq, num_timesteps, seed=seed) + return x, temb, ti + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3FinalAdaLNOutOp() + x, nw = torch.zeros(1, 4, 8), torch.ones(8) + temb, w, b = torch.zeros(2, 16), torch.zeros(16, 16), torch.zeros(16) + ti = torch.tensor([0, 1, 1, 0]) + with pytest.raises(TypeError): # BF16 temb (probe H7) + op(x, nw, temb.bfloat16(), w, b, ti) + with pytest.raises(ValueError): # not 2H rows + op(x, nw, temb, w[:10], b[:10], ti) + with pytest.raises(IndexError): + op(x, nw, temb, w, b, ti + 2) + + +@requires_cuda +class TestCuda: + def test_shift_first_layout_and_timestep_indexing(self): + """Zero weights: the bias halves are (shift, scale); rows follow timestep_indices.""" + + hidden, seq = 16, 6 + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + nw = torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(2 * hidden, 32, device="cuda", dtype=torch.bfloat16) + b = torch.cat([torch.full((hidden,), 1.0), torch.full((hidden,), -1.0)]).bfloat16().cuda() + ti = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda") + out = _cuda_op()(x, nw, torch.zeros(2, 32, device="cuda"), w, b, ti) + # scale = -1 -> norm(x) * 0 + shift = 1 everywhere + assert torch.equal(out, torch.ones_like(out)) + + @pytest.mark.parametrize("seq, num_timesteps", [(1, 1), (257, 2), (4097, 3)]) + def test_forward_against_diffusers_and_golden(self, norm_out, seq, num_timesteps): + x, temb, ti = _inputs(seq, num_timesteps, seed=seq) + ours = _cuda_op()(x, *norm_out[:1], temb, *norm_out[1:], ti) + theirs = provider_final_adaln_out(x, norm_out[0], temb, *norm_out[1:], ti) + golden = NativeH3FinalAdaLNOutOp().forward_fp32(x, norm_out[0], temb, *norm_out[1:], ti) + assert (ours == theirs).float().mean() > 0.999 + torch.testing.assert_close(ours.float(), golden, atol=5e-2, rtol=2e-2) + + def test_rows_are_batch_and_position_invariant(self, norm_out): + x, temb, ti = _inputs(600, seed=1, batch=2) + full = _cuda_op()(x, norm_out[0], temb, *norm_out[1:], ti) + part = _cuda_op()(x[1:2, 200:260], norm_out[0], temb, *norm_out[1:], ti[200:260]) + assert torch.equal(part[0], full[1, 200:260]) + + def test_backward_deterministic_and_fp32_table_gradient(self, norm_out): + x, temb, ti = _inputs(4097, seed=2) + grad = torch.randn_like(x) + + def grads(fn, dtype=None): + tensors = [x, norm_out[0], temb, *norm_out[1:]] + leaves = [ + (t if dtype is None else t.to(dtype)).detach().clone().requires_grad_(True) + for t in tensors + ] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, nw, t, w, b): + act = t * torch.sigmoid(t) + act = act + (act.to(torch.bfloat16).double() - act).detach() + shift, scale = F.linear(act, w, b).chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw + return n * (1 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ours = grads(lambda *t: _cuda_op()(*t, ti)) + again = grads(lambda *t: _cuda_op()(*t, ti)) + ref = grads(golden, torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): + rel = ((g.double() - r).abs().max() / r.abs().max()).item() + assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + + def test_registry_dispatches_cuda(self): + from rl_engine.kernels.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("final_adaln_out", device="cuda") + assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" diff --git a/tests/test_ws1_gtest_gpu.py b/tests/test_ws1_gtest_gpu.py index 32c6322c3..3679b409c 100644 --- a/tests/test_ws1_gtest_gpu.py +++ b/tests/test_ws1_gtest_gpu.py @@ -55,6 +55,7 @@ def test_all_ws1_single_ops_are_registered(): "adaln_row_gather", "h3_rmsnorm", "adaln_gate_residual", + "final_adaln_out", } <= names From 623bec886f679f168e17cea916ebbca47e868160 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 03:21:05 +0000 Subject: [PATCH 04/62] [WS1][MiniMax-H3] final_adaln_out: B200 evidence report, chain replay and figure Produced by scripts/h3_evidence.py and scripts/h3_chain_replay.py --backward at 38d575f on a clean tree; rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e05f1d15c..3e64efbef 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,9 +40,9 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers (whose error varies from - run to run), which also rounds the - table gradient to BF16 and accumulates it with `index_select` atomics; + - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers, whose error varies from + run to run because it rounds the table gradient to BF16 and accumulates it with + `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. - **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, @@ -59,13 +59,25 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.62 ms | 1.09 ms | -| 32768 | 0.39 ms | 0.78 ms | 2.86 ms | 5.32 ms | -| 131072 | 1.24 ms | 2.93 ms | 7.78 ms | 20.7 ms | +| 4097 | 0.13 ms | 0.15 ms | 1.64 ms | 1.10 ms | +| 32768 | 0.39 ms | 0.78 ms | 2.88 ms | 5.36 ms | +| 131072 | 1.23 ms | 2.92 ms | 7.80 ms | 20.7 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. +## Evidence + +![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) + +There are two data files, both written from a clean tree at commit `38d575f`: + +- [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the + forward-equality fraction, row invariance and backward accuracy. +- [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the + whole conditioning chain (timestep → … → norm_out) replayed stage by stage over + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. + ## Tests ```bash From cc867401e434b8e3d33daf28f3e8a64cd548dcef Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:27:16 +0000 Subject: [PATCH 05/62] fix(h3): reject mixed-device RMSNorm native inputs Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/rmsnorm_modulate.cu | 16 ++++++------ tests/h3/test_h3_rmsnorm.py | 43 ++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 7 deletions(-) diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..d42fe3525 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -276,8 +276,9 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous(), + "row index must be a contiguous int64 CUDA tensor on x's device"); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -291,9 +292,9 @@ Modulation make_modulation(const c10::optional& shift, void check_xw(const torch::Tensor& x, const torch::Tensor& w) { TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); - TORCH_CHECK(w.is_cuda() && w.dim() == 1 && w.size(0) == x.size(1) && w.is_contiguous() && - w.scalar_type() == x.scalar_type(), - "weight must be a contiguous (N,) CUDA tensor with x's dtype"); + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && w.size(0) == x.size(1) && + w.is_contiguous() && w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype on x's device"); TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, " (PyTorch's vectorized RMSNorm path)"); } @@ -344,8 +345,9 @@ std::vector h3_rmsnorm_backward( 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.scalar_type() == at::kFloat && rstd.numel() == x.size(0), - "rstd must be the forward's float32 (M,) statistics"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); diff --git a/tests/h3/test_h3_rmsnorm.py b/tests/h3/test_h3_rmsnorm.py index f0d013bfb..fffb7e3e0 100644 --- a/tests/h3/test_h3_rmsnorm.py +++ b/tests/h3/test_h3_rmsnorm.py @@ -143,6 +143,49 @@ def test_rejects_unsupported(self): _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +class TestCudaNativeDeviceValidation: + @pytest.mark.parametrize("argument", ["weight", "index"]) + def test_forward_rejects_different_device(self, argument): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + if argument == "weight": + weight = weight.to("cuda:1") + else: + index = index.to("cuda:1") + with pytest.raises(RuntimeError, match="on x's device"): + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @pytest.mark.parametrize("argument", ["weight", "index", "rstd"]) + def test_backward_rejects_different_device(self, argument): + from rl_engine.kernels.ops.base import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + if argument == "weight": + weight = weight.to("cuda:1") + elif argument == "index": + index = index.to("cuda:1") + else: + rstd = rstd.to("cuda:1") + with pytest.raises(RuntimeError, match="on x's device"): + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @requires_cuda class TestCudaBackward: def _grads(self, fn, tensors, grad): From 40a9b445bbdc46527bd5f8b1a99c7c1f08f0b77a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:27:16 +0000 Subject: [PATCH 06/62] fix(h3): retain FP64 final AdaLN golden gradients Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/testing/h3_report.py | 21 ++++++++--- tests/h3/test_h3_final_adaln_out.py | 58 +++++++++++++++++++++++++---- 2 files changed, 66 insertions(+), 13 deletions(-) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 41f8d647d..4b331e3fc 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -634,12 +634,23 @@ def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) ).to(x.dtype) - def grads(fn): - leaves = [t.detach().clone().requires_grad_(True) for t in (x, nw, temb, w, b)] - fn(*leaves).backward(grad) + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, nw, temb, w, b)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) return [leaf.grad for leaf in leaves] - ref = grads(lambda *t: golden_op.forward_fp32(*t, ti)) + def golden_backward(x_, nw_, t_, w_, b_): + # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. + act = t_ * torch.sigmoid(t_) + act = act + (act.to(torch.bfloat16).double() - act).detach() + table = torch.nn.functional.linear(act, w_, b_) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw_ + return n * (1.0 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ref = grads(golden_backward, torch.float64) backward = {} for name, fn in ( ("cuda", lambda *t: op(*t, ti)), @@ -649,7 +660,7 @@ def grads(fn): backward[name] = { "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), "rel_error": { - key: float((g.double() - r.double()).abs().max() / r.double().abs().max()) + key: float((g.double() - r).abs().max() / r.abs().max()) for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) }, } diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py index f322cd93d..17a3843b8 100644 --- a/tests/h3/test_h3_final_adaln_out.py +++ b/tests/h3/test_h3_final_adaln_out.py @@ -49,6 +49,18 @@ def _inputs(seq, num_timesteps=3, seed=0, batch=1): return x, temb, ti +def _golden(x, norm_weight, temb, weight, bias, timestep_indices): + act = temb * torch.sigmoid(temb) + act = act + (act.to(torch.bfloat16).double() - act).detach() + table = F.linear(act, weight, bias) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) * norm_weight + return n * (1 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + class TestReference: def test_rejects_bad_inputs(self): op = NativeH3FinalAdaLNOutOp() @@ -106,21 +118,51 @@ def grads(fn, dtype=None): fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) return [leaf.grad for leaf in leaves] - def golden(x_, nw, t, w, b): - act = t * torch.sigmoid(t) - act = act + (act.to(torch.bfloat16).double() - act).detach() - shift, scale = F.linear(act, w, b).chunk(2, dim=-1) - n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw - return n * (1 + scale.index_select(0, ti)) + shift.index_select(0, ti) - ours = grads(lambda *t: _cuda_op()(*t, ti)) again = grads(lambda *t: _cuda_op()(*t, ti)) - ref = grads(golden, torch.float64) + ref = grads(lambda *t: _golden(*t, ti), torch.float64) assert all(torch.equal(a, b) for a, b in zip(ours, again)) for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): rel = ((g.double() - r).abs().max() / r.abs().max()).item() assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + def test_report_backward_uses_fp64_golden(self, monkeypatch): + from types import SimpleNamespace + + from rl_engine.testing import h3_report + + rng = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn(1, 160, 8, device="cuda", generator=rng).bfloat16() + nw = torch.randn(8, device="cuda", generator=rng).bfloat16() + temb = torch.randn(3, 4, device="cuda", generator=rng) + w = torch.randn(16, 4, device="cuda", generator=rng).bfloat16() + b = torch.randn(16, device="cuda", generator=rng).bfloat16() + ti = torch.arange(160, device="cuda") % 3 + inputs = (x, nw, temb, w, b) + native = NativeH3FinalAdaLNOutOp() + registry = SimpleNamespace( + get_op=lambda *args, **kwargs: native, + _get_or_create_backend=lambda _: native, + _priority_map={"cpu": {"final_adaln_out": [None]}}, + ) + monkeypatch.setattr(h3_report, "_final_inputs", lambda *args, **kwargs: (*inputs, ti)) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", native) + report = h3_report._final_accuracy(registry) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + leaves = [t.detach().clone().requires_grad_(True) for t in inputs] + native(*leaves, ti).backward(grad) + ref = [t.double().detach().requires_grad_(True) for t in inputs] + _golden(*ref, ti).backward(grad.double()) + for name, leaf, reference in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), leaves, ref): + assert reference.grad.dtype == torch.float64 + expected = float( + (leaf.grad.double() - reference.grad).abs().max() / reference.grad.abs().max() + ) + for backend in ("cuda", "provider"): + assert report["backward"][backend]["rel_error"][name] == pytest.approx(expected) + def test_registry_dispatches_cuda(self): from rl_engine.kernels.registry import KernelRegistry From 7090f00fc986130391ebe91541f404586313df91 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:28:04 +0000 Subject: [PATCH 07/62] fix(h3): address native validation and evidence review comments Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 11 +- csrc/cuda/h3/rmsnorm_modulate.cu | 18 ++- docs/operators/h3-adaln-gate-residual.md | 2 + docs/operators/h3-rmsnorm.md | 2 + rl_engine/testing/h3_report.py | 61 +++++--- scripts/plot_h3_evidence.py | 14 ++ tests/h3/test_h3_native_validation.py | 124 ++++++++++++++++ tests/h3/test_h3_report.py | 177 +++++++++++++++++++++++ 8 files changed, 380 insertions(+), 29 deletions(-) create mode 100644 tests/h3/test_h3_native_validation.py create mode 100644 tests/h3/test_h3_report.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..c2bd7921a 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -144,9 +144,9 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc "gate must be a CUDA tensor with the activations' dtype"); TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, "gate must be (R, N) with unit column stride"); - TORCH_CHECK(index.is_cuda() && index.scalar_type() == at::kLong && index.dim() == 1 && - index.is_contiguous() && index.numel() > 0, - "index must be a non-empty contiguous int64 CUDA tensor"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && + index.dim() == 1 && index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; } @@ -207,8 +207,9 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: check_act(y, "y", grad); const Gate g = make_gate(gate, index, grad); for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { - TORCH_CHECK(t->is_cuda() && t->scalar_type() == at::kLong && t->dim() == 1 && t->is_contiguous(), - "tile metadata must be contiguous int64 CUDA tensors"); + 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 tensors on ", grad.device()); } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 5212324b1..2775572e6 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -276,8 +276,9 @@ Modulation make_modulation(const c10::optional& shift, } TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), "shift and scale must be views of the same table layout"); - TORCH_CHECK(ix.is_cuda() && ix.scalar_type() == at::kLong && ix.dim() == 1 && ix.is_contiguous(), - "row index must be a contiguous int64 CUDA tensor"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous() && ix.numel() > 0, + "row index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", ix.size(0), ")"); mod.shift = sh.data_ptr(); @@ -349,6 +350,16 @@ std::vector h3_rmsnorm_backward( const int64_t rows = x.size(0); const int64_t n = x.size(1); const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + for (const auto* t : {&*sorted_pos, &*tile_begin, &*tile_end, &*seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", x.device()); + } + } const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); auto dx = torch::empty_like(x); @@ -375,9 +386,6 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); if (mod.index == nullptr) return {dx, dweight}; - TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && - seg_first_tile.has_value(), - "modulated backward needs the sorted segment tiles"); const int64_t segments = seg_first_tile->numel() - 1; const int64_t seg_tiles = tile_begin->numel(); const int64_t width = 2 * n; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index f55412449..2ef31ba54 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -16,6 +16,8 @@ gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never m ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("adaln_gate_residual", device="cuda") hidden = op(residual, attn_output, gate_msa, adaln_indices) ``` diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 957670552..faaed48e2 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -21,6 +21,8 @@ gathered inside the kernel, so the `(S, H)` tensors that ## Entry Point ```python +from rl_engine.kernels.registry import kernel_registry + op = kernel_registry.get_op("h3_rmsnorm", device="cuda") n = op(x, weight) # plain RMSNorm out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 021dab4b0..489404bb6 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -25,7 +25,11 @@ provider_time_embedder, provider_time_proj, ) -from rl_engine.testing.h3_weights import h3_weights_dir, load_h3_conditioning_weights +from rl_engine.testing.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) # Keys of a perf case that hold a timed callable, in display order. TIMED_KEYS = ( @@ -247,7 +251,7 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, invariant = all( all( torch.equal(s, f[3 * i : 3 * i + 3]) - for s, f in zip(op(temb[i : i + 1], weight, bias), full) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) ) for i in range(9) ) @@ -304,7 +308,6 @@ def backward(fn, ti=ti, tags=tags, grads=grads): def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: from rl_engine.testing.h3_chain import run_backward_case - from rl_engine.testing.h3_weights import load_h3_conditioning_weights op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) @@ -314,7 +317,9 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: ti, tags = h3_packed_layout(seq, 3, seed=seq) ours = op(rows, ti, tags) theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) - forward_bitwise[str(seq)] = all(torch.equal(a, b) for a, b in zip(ours, theirs)) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. ti, tags = h3_packed_layout(4097, 3, seed=1) @@ -339,16 +344,22 @@ def grad_of(fn): "correctly_rounded_fraction": float((first == gold).float().mean()), } - weights = load_h3_conditioning_weights("cuda") - chain = [ - run_backward_case(registry, weights, num_timesteps=t, seq_len=s) - for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) - ] - return { + result: dict[str, Any] = { "forward_bitwise_vs_index_select": forward_bitwise, "op_backward": backward, - "chain_backward": chain, + "chain_backward": [], } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run scripts/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return result # --------------------------------------------------------------------------- # @@ -406,14 +417,19 @@ def backward(fn, x=x, weight=weight, shift=shift, scale=scale, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") - weights = load_h3_conditioning_weights("cuda") - plain = {} - for name in ( + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ "transformer_blocks.0.norm1.weight", "transformer_blocks.0.norm2.weight", "token_refiner.final_norm.weight", "norm_out.norm.weight", - ): + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + plain = {} + for name in names: x = torch.randn(2, 777, 5376, device="cuda").bfloat16() ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) plain[name] = bool(torch.equal(op(x, weights[name]), ref)) @@ -454,13 +470,16 @@ def golden(x_, w_, sh_, sc_): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) }, } return { + "weight_source": weight_source, "plain_bitwise_vs_nn_rmsnorm": plain, "modulated_bitwise_vs_diffusers": modulated, "rows_batch_invariant": invariant, @@ -559,10 +578,14 @@ def grads(fn, dtype=None): ): first, second = grads(fn), grads(fn) backward[name] = { - "repeat_bitwise_equal": all(torch.equal(a, b) for a, b in zip(first, second)), + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), "rel_error": { key: float((g.double() - r).abs().max() / r.abs().max()) - for key, g, r in zip(("d_residual", "d_sublayer", "d_gate"), first, ref) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) }, } return { diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 15c800b5f..d8bd5ea43 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -302,6 +302,20 @@ def plot_gather(report: dict[str, Any]): "Whole-chain grads of the FP32 time embedder vs FP64", "max abs error / golden max", ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig names = { "candidate": ("RL-Kernel, separate ops", THIRD, "o"), "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py new file mode 100644 index 000000000..8dd919028 --- /dev/null +++ b/tests/h3/test_h3_native_validation.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +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 == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@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( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) diff --git a/tests/h3/test_h3_report.py b/tests/h3/test_h3_report.py new file mode 100644 index 000000000..b6bb6ec3f --- /dev/null +++ b/tests/h3/test_h3_report.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.testing import h3_chain, h3_report +from rl_engine.testing.h3_cases import h3_packed_layout +from rl_engine.testing.h3_provider import provider_adaln_row_gather, provider_norm_modulate +from rl_engine.testing.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run scripts/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = Path(__file__).resolve().parents[2] / "scripts" / "plot_h3_evidence.py" + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) From cff41e857503e6ba3d7147c7d43dfd0d5a3cb93b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 16:28:41 +0000 Subject: [PATCH 08/62] docs(h3): regenerate final AdaLN evidence against FP64 gradients Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3e64efbef..59b4ab0a9 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,8 +40,8 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 13–57× closer to FP64 than diffusers, whose error varies from - run to run because it rounds the table gradient to BF16 and accumulates it with + - `d_temb`, `dW` and `db` are 22–43× closer to FP64 in the recorded run than diffusers, + whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. @@ -59,9 +59,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.64 ms | 1.10 ms | -| 32768 | 0.39 ms | 0.78 ms | 2.88 ms | 5.36 ms | -| 131072 | 1.23 ms | 2.92 ms | 7.80 ms | 20.7 ms | +| 4097 | 0.14 ms | 0.15 ms | 1.67 ms | 1.12 ms | +| 32768 | 0.40 ms | 0.78 ms | 2.97 ms | 5.39 ms | +| 131072 | 1.25 ms | 2.94 ms | 7.87 ms | 22.2 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -70,13 +70,15 @@ projection backward. ![final_adaln_out on B200: latency and backward accuracy](../usage/evidence/h3-final-adaln-out-b200/figure.png) -There are two data files, both written from a clean tree at commit `38d575f`: +There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the - forward-equality fraction, row invariance and backward accuracy. + forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean + tree at `372d6d7`, with FP64 leaves and upstream gradients and explicit BF16 forward boundaries. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over - T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean + tree at `38d575f`. ## Tests From 507c9ad86540dd0aaab9376996d85d8c1c5f60a8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:39:17 +0000 Subject: [PATCH 09/62] [WS2][CUDA][MiniMax-H3] tp_adaln_3mod: column-parallel AdaLN projection, byte-equal to WS1 RFC #420 WS2 row `tp_adaln_3mod` (stacked on final_adaln_out). - Ownership: rank r owns the contiguous output columns [r*N/tp, (r+1)*N/tp) of the 3 x 6 x 5376 table (the same weight/bias rows); N/tp must be a multiple of the 64-row d_input chunk, so a shard never splits a WS1 partial. Unsupported TP sizes fail closed. - Every floating-point op is the WS1 kernel: the local tensor-core GEMV gives the WS1 columns, a rank-ordered all-gather (a copy) rebuilds the table, dW/db are shard-local, and d_temb all-gathers the WS1 chunk partials and runs the WS1 ascending fold on every rank. No collective reduction arithmetic. - det_linear: d_input split into h3_det_linear_backward_input_partials and h3_det_linear_fold_chunks (same kernels; backward_input composes them). - ws2_comm.gather_rows: rank-ordered row gather that pads outputs of at most 256 KiB past DeterministicCollective's single-block byte-copy path (~4 us/KiB); keeps the T = 1 forward at ~0.12 ms. - Readback: tp, rank, owned columns and (modality, chunk) slots, collective backend, kernel ids, fallback. - Tests: ownership, single-GPU shard arithmetic (columns and chunk partials equal the WS1 slice, the fold equals WS1 d_input), fail-closed inputs, and real NCCL TP2/4/8 on the pinned block-0 weights. - scripts/h3_ws2_evidence.py + figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 41 +++- csrc/ops.cpp | 8 + docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-tp-adaln-3mod.md | 94 ++++++++ rl_engine/_C.pyi | 4 + rl_engine/kernels/ops/cuda/h3/det_linear.py | 17 ++ .../ops/cuda/h3/tp_adaln_projection.py | 211 ++++++++++++++++++ rl_engine/kernels/ops/cuda/h3/ws2_comm.py | 33 +++ rl_engine/testing/h3_ws2.py | 187 ++++++++++++++++ scripts/h3_ws2_evidence.py | 90 ++++++++ scripts/plot_h3_evidence.py | 96 +++++++- tests/h3/test_h3_tp_adaln.py | 148 ++++++++++++ 13 files changed, 923 insertions(+), 8 deletions(-) create mode 100644 docs/operators/h3-tp-adaln-3mod.md create mode 100644 rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py create mode 100644 rl_engine/kernels/ops/cuda/h3/ws2_comm.py create mode 100644 rl_engine/testing/h3_ws2.py create mode 100644 scripts/h3_ws2_evidence.py create mode 100644 tests/h3/test_h3_tp_adaln.py diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index 95e16b234..66d62931d 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -26,6 +26,9 @@ // * N is split into fixed chunks of 64 rows; each chunk is an ascending // fmaf chain into FP32 from 0; the chunk partials are then left-folded in // ascending chunk order and cast once. +// The partials and the fold are also bound separately: a contiguous N +// slice starting on a chunk boundary yields exactly the full call's +// partials for those chunks, which is what tensor-parallel shards gather. // d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] // * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one // cross-row reduction; its order is the logical row order. @@ -487,24 +490,24 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor return {out}; } -// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. -torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, - c10::ScalarType out_dtype) { +// grad [T, N] (float32), weight [N, K] -> partial [ceil(N / 64), T, K] float32: +// the per-chunk sums of contract h3-det-linear-v1, before the ascending fold. +// A contiguous slice of N that starts on a chunk boundary yields exactly the +// matching rows of the full call (tensor-parallel shards rely on this). +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight) { TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && grad.scalar_type() == at::kFloat, "grad must be a contiguous 2-D float32 CUDA tensor"); check_matrix(weight, "weight"); TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); - TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, - "out_dtype must be float32 or bfloat16"); const int64_t rows = grad.size(0); const int64_t n_out = weight.size(0); const int64_t k_in = weight.size(1); const c10::cuda::CUDAGuard device_guard(grad.device()); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; auto partial = torch::empty({chunks, rows, k_in}, grad.options()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); + if (partial.numel() == 0) return partial; auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); @@ -518,7 +521,23 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei n_out, k_in); } C10_CUDA_KERNEL_LAUNCH_CHECK(); - const int64_t elems = rows * k_in; + return partial; +} + +// partial [C, T, K] float32 -> left fold over C in ascending order, cast once. +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 3 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat && partial.size(0) > 0, + "partial must be a non-empty contiguous 3-D float32 CUDA tensor"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const c10::cuda::CUDAGuard device_guard(partial.device()); + const int64_t chunks = partial.size(0); + auto out = torch::empty({partial.size(1), partial.size(2)}, partial.options().dtype(out_dtype)); + const int64_t elems = out.numel(); + if (elems == 0) return out; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); if (out_dtype == at::kFloat) { det_fold_chunks_kernel<<>>( @@ -531,6 +550,14 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei return out; } +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + return h3_det_linear_fold_chunks(h3_det_linear_backward_input_partials(grad, weight), out_dtype); +} + // grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 52ece3f13..f3dc5cca2 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -507,6 +507,8 @@ std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor int64_t activation, bool save_pre_activation); torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, c10::ScalarType out_dtype); +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight); +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype); std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, c10::ScalarType w_dtype, bool with_bias); @@ -806,6 +808,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, "Deterministic grad @ weight with fixed 64-row N chunks folded in order", py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_input_partials", &h3_det_linear_backward_input_partials, + "Per-64-row-chunk FP32 partials of grad @ weight, before the ascending fold", + py::arg("grad"), py::arg("weight")); + m.def("h3_det_linear_fold_chunks", &h3_det_linear_fold_chunks, + "Ascending left fold of h3_det_linear_backward_input_partials, cast once", + py::arg("partial"), py::arg("out_dtype")); m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, "Deterministic dW/dbias as ascending-row FP32 folds", py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); diff --git a/docs/.nav.yml b/docs/.nav.yml index 1ef4fa017..48f12fe3b 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -37,6 +37,7 @@ nav: - operators/h3-rmsnorm.md - operators/h3-adaln-gate-residual.md - operators/h3-final-adaln-out.md + - operators/h3-tp-adaln-3mod.md - Developer Guide: - contributing/README.md - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 27ffce7d0..99e02cf0c 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -39,4 +39,5 @@ Every operator page should include: - [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) - [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) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..da7b0ae42 --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,94 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when +the op is built. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 89843f979..672bf6049 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -341,6 +341,10 @@ def h3_det_linear_forward( def h3_det_linear_backward_input( grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype ) -> torch.Tensor: ... +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: ... diff --git a/rl_engine/kernels/ops/cuda/h3/det_linear.py b/rl_engine/kernels/ops/cuda/h3/det_linear.py index f5dab854e..17e7bce0f 100644 --- a/rl_engine/kernels/ops/cuda/h3/det_linear.py +++ b/rl_engine/kernels/ops/cuda/h3/det_linear.py @@ -16,10 +16,13 @@ ACT_NONE = 0 ACT_SILU = 1 CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) _SYMBOLS = ( "h3_det_linear_forward", "h3_det_linear_backward_input", "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", ) @@ -62,6 +65,20 @@ def det_linear_backward_input( ) +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + def det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True ) -> list[torch.Tensor]: diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py new file mode 100644 index 000000000..0eb5a4709 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.cuda.h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.kernels.ops.pytorch.h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +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 % (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): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/kernels/ops/cuda/h3/ws2_comm.py b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py new file mode 100644 index 000000000..a0dc116f6 --- /dev/null +++ b/rl_engine/testing/h3_ws2.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS1 vs WS2 comparisons for the H3 conditioning path (tests and evidence). + +The per-rank helpers run inside real multi-process groups (tests and +``scripts/h3_ws2_evidence.py``); the WS1 side runs on one GPU. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def ws1_projection(temb, weight, bias, grad) -> dict[str, torch.Tensor]: + """The single-GPU ``adaln_projection_3mod`` table and its three gradients.""" + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + table = H3AdaLNProjectionCudaOp().forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight": leaves[1].grad, + "d_bias": leaves[2].grad, + } + + +def tp_projection_rank(collective, temb, weight, bias, grad) -> dict[str, Any]: + """One TP rank: shard the full weight, run the TP op, return its view of the results.""" + + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + op = H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, w_shard, b_shard)] + table = op.forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight_shard": leaves[1].grad, + "d_bias_shard": leaves[2].grad, + "readback": op.readback(), + } + + +def tp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Byte equality of every rank's replicated outputs and of the concatenated shards.""" + + return { + "table": all(torch.equal(r["table"], ws1["table"]) for r in ranks), + "d_temb": all(torch.equal(r["d_temb"], ws1["d_temb"]) for r in ranks), + "d_weight": torch.equal(torch.cat([r["d_weight_shard"] for r in ranks]), ws1["d_weight"]), + "d_bias": torch.equal(torch.cat([r["d_bias_shard"] for r in ranks]), ws1["d_bias"]), + } + + +def tp_projection_sweep(collective, temb, weight, bias, grad, iters: int = 100) -> dict[str, Any]: + """Evidence for one rank: per-T byte equality against WS1 on this GPU, timing, readback. + + ``temb``/``grad`` hold the largest T; every T in ``1..T`` uses their first + rows. Each rank compares its own replicated outputs and its own shard of + ``dW``/``db`` with the WS1 result computed here, so nothing large is shipped. + """ + + from rl_engine.kernels.ops.cuda.h3.adaln_projection import H3AdaLNProjectionCudaOp + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + adaln_column_shard, + ) + from rl_engine.testing.h3_report import time_us + + shard = adaln_column_shard(weight.shape[0], collective.world_size, collective.rank) + cols = slice(shard.begin, shard.end) + tp_op, ws1_op = ( + H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]), + H3AdaLNProjectionCudaOp(), + ) + equality, perf = [], [] + for num_t in range(1, temb.shape[0] + 1): + t, g = temb[:num_t], grad[:num_t] + ws1 = ws1_projection(t, weight, bias, g) + mine = tp_projection_rank(collective, t, weight, bias, g) + equality.append( + { + "num_timesteps": num_t, + "table": torch.equal(mine["table"], ws1["table"]), + "d_temb": torch.equal(mine["d_temb"], ws1["d_temb"]), + "d_weight_shard": torch.equal(mine["d_weight_shard"], ws1["d_weight"][cols]), + "d_bias_shard": torch.equal(mine["d_bias_shard"], ws1["d_bias"][cols]), + } + ) + leaves = [x.detach().clone().requires_grad_(True) for x in (t, weight[cols], bias[cols])] + full = [x.detach().clone().requires_grad_(True) for x in (t, weight, bias)] + + def step(op, args): + for x in args: + x.grad = None + op.forward_table(*args).backward(g) + + perf.append( + { + "num_timesteps": num_t, + "tp_forward_us": time_us( + lambda: tp_op.forward_table(t, weight[cols], bias[cols]), 30, iters + ), + "tp_forward_backward_us": time_us(lambda: step(tp_op, leaves), 10, iters), + "ws1_forward_us": time_us(lambda: ws1_op.forward_table(t, weight, bias), 30, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_op, full), 10, iters), + } + ) + return {"equality": equality, "perf": perf, "readback": tp_op.readback()} + + +# --- real multi-process groups ------------------------------------------------- + + +def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): + import traceback + + import torch.distributed as dist + + from rl_engine.distributed.collectives import DeterministicCollective + + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", init_method=init_method, rank=rank, world_size=world, device_id=rank + ) + inputs = {k: v.cuda() for k, v in torch.load(inputs_path).items()} + collective = DeterministicCollective(device=rank, max_size_bytes=256 << 20) + out = target(collective, **inputs, **kwargs) + # Results go through a file: tensors sent over the queue are shared by + # file descriptor and vanish if this process exits before the parent reads. + torch.save( + {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, + inputs_path.with_name(f"rank{rank}.pt"), + ) + queue.put({"rank": rank}) + collective.close() + dist.destroy_process_group() + except Exception: # forwarded to the parent + queue.put({"rank": rank, "error": traceback.format_exc()}) + + +def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: float = 900, **kwargs): + """Run ``target(collective, **inputs, **kwargs)`` on ranks ``0..world-1`` (one GPU each). + + ``target`` must be importable (it is pickled into spawned processes). Each + rank gets a NCCL group and a ``DeterministicCollective``; results come back + on the CPU, ordered by rank. Any rank's exception is re-raised here. + """ + + import tempfile + from pathlib import Path + + import torch.multiprocessing as mp + + if torch.cuda.device_count() < world: + raise RuntimeError(f"needs {world} GPUs, found {torch.cuda.device_count()}") + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "inputs.pt" + torch.save({k: v.cpu() for k, v in inputs.items()}, path) + queue = ctx.Queue() + init = (Path(tmp) / "init").as_uri() + procs = [ + ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) + for r in range(world) + ] + for proc in procs: + proc.start() + status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) + for proc in procs: + proc.join(timeout=120) + for result in status: + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py new file mode 100644 index 000000000..c27332f77 --- /dev/null +++ b/scripts/h3_ws2_evidence.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``scripts/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import torch # noqa: E402 + +from rl_engine.testing.h3_chain import environment, git_state # noqa: E402 +from rl_engine.testing.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) +from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +OPS = {"tp_adaln_3mod": _tp_adaln} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 3159c7379..ba4b3b4ef 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -2,7 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Render an RFC #420 operator report (``scripts/h3_evidence.py``) as a PNG. +"""Render an RFC #420 operator or WS2 report (``scripts/h3_evidence.py``, +``scripts/h3_ws2_evidence.py``) as a PNG. Needs only the report JSON and matplotlib (not torch or a GPU), so it can run anywhere: @@ -22,6 +23,7 @@ import matplotlib matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 import matplotlib.pyplot as plt # noqa: E402 # Fixed series identity across every H3 figure (categorical slots 1-3). @@ -496,6 +498,97 @@ def plot_final(report: dict[str, Any]): return fig +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + 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) + # 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"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + return fig + + PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { "timestep_sinusoid_h3": plot_sinusoid, "timestep_mlp_fp32": plot_mlp, @@ -504,6 +597,7 @@ def plot_final(report: dict[str, Any]): "h3_rmsnorm": plot_rmsnorm, "adaln_gate_residual": plot_gate_residual, "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, } diff --git a/tests/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py new file mode 100644 index 000000000..8b341c8d3 --- /dev/null +++ b/tests/h3/test_h3_tp_adaln.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import adaln_column_shard +from rl_engine.testing.h3_ws2 import run_world, tp_matches_ws1, tp_projection_rank, ws1_projection + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.kernels.ops.cuda.h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8]) + 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]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.kernels.ops.cuda.h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.kernels.ops.cuda.h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@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_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + 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"} From ecde5bb76a19bdc47457b67dad19e5423225fe55 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 10:58:54 +0000 Subject: [PATCH 10/62] [WS2][CUDA][MiniMax-H3] sp_norm_adaln: sequence-parallel norm/AdaLN/gated residual, byte-equal to WS1 RFC #420 WS2 row `sp_norm_adaln` (stacked on tp_adaln_3mod). - Ownership: rank r holds packed positions [r*S//sp, (r+1)*S//sp) of every batch item with the matching metadata; any S >= sp. The (S,) row index is replicated and every rank derives the same reduction plan from it. - Forward, dx, d_sublayer and d_residual: the WS1 kernels on local rows. - d_norm_weight, d_shift/d_scale, d_gate keep the WS1 two-level order (256-element tiles, then an ascending fold): each tile is computed by the rank holding its first row, after all-gathering the tile's rows from other ranks (only tiles crossing a shard boundary move rows); the partials are all-gathered into WS1 tile order and every rank runs the WS1 fold. Rank-ordered all-gathers only; no collective arithmetic. - New bindings around the unchanged WS1 kernels: h3_rmsnorm_backward_{dx,partials}, h3_rmsnorm_fold_partials, h3_gate_residual_backward_dy, h3_gate_grad_{partials,fold}; the WS1 dweight tile kernel also takes explicit row lists (same arithmetic). - The plan is built once per index object and shared by every norm and gated residual that uses it. - Tests: layout, per-rank backward vs WS1 for SP 2-8 on one GPU (odd S, B = 2, shards smaller than a tile, block and interleaved packing, bf16/fp32), fail-closed inputs, and the autograd block region over real NCCL SP2/4/8. - h3_block_layout (H3's block packing), evidence sweep with rows exchanged, timing and a naive-SP comparison; figure; operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 80 ++++ csrc/cuda/h3/rmsnorm_modulate.cu | 165 ++++++- csrc/ops.cpp | 42 ++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/h3-sp-norm-adaln.md | 96 +++++ rl_engine/_C.pyi | 43 ++ .../kernels/ops/cuda/h3/sp_norm_adaln.py | 404 ++++++++++++++++++ rl_engine/testing/h3_cases.py | 24 ++ rl_engine/testing/h3_ws2.py | 181 ++++++++ scripts/h3_ws2_evidence.py | 42 +- scripts/plot_h3_evidence.py | 183 ++++++-- tests/h3/test_h3_sp_norm_adaln.py | 210 +++++++++ 13 files changed, 1428 insertions(+), 44 deletions(-) create mode 100644 docs/operators/h3-sp-norm-adaln.md create mode 100644 rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py create mode 100644 tests/h3/test_h3_sp_norm_adaln.py diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 18ede9e72..9a03660f1 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 @@ -240,3 +242,81 @@ 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 unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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"); + 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)); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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 d42fe3525..0c7a1bbfe 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -20,6 +20,10 @@ // dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) // 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 @@ -180,18 +184,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; const T* scale = mod.index != nullptr ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride : nullptr; @@ -403,3 +414,147 @@ std::vector h3_rmsnorm_backward( C10_CUDA_KERNEL_LAUNCH_CHECK(); 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.numel() == x.size(0), + "rstd must be the forward's 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(); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + 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; + 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"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + 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.numel() == x.size(0), + "rstd must be the forward's 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/ops.cpp b/csrc/ops.cpp index f3dc5cca2..342d09516 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -537,6 +537,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 @@ -843,5 +862,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/docs/.nav.yml b/docs/.nav.yml index 48f12fe3b..cfc833fe2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -38,6 +38,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 - General: diff --git a/docs/operators/README.md b/docs/operators/README.md index 99e02cf0c..065d754b6 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -40,4 +40,5 @@ 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) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..167b63f7c --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,96 @@ +# 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.collectives import DeterministicCollective +from rl_engine.kernels.ops.cuda.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 – 174 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 +node, one GPU per rank. + +## Tests + +```bash +python -m pytest tests/h3/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json +python scripts/plot_h3_evidence.py docs/usage/evidence/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 672bf6049..d9948d98d 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -397,3 +397,46 @@ def h3_gate_residual_backward( tile_end: torch.Tensor, seg_first_tile: torch.Tensor, ) -> list[torch.Tensor]: ... +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]: ... +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]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +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: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py new file mode 100644 index 000000000..19a867b3e --- /dev/null +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -0,0 +1,404 @@ +# 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.kernels.ops.backward_runtime import record_backward +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.h3.adaln_row_gather import BACKWARD_TILE, _segment_tiles +from rl_engine.kernels.ops.cuda.h3.ws2_comm import gather_rows +from rl_engine.kernels.ops.pytorch.h3.gate_residual import validate_h3_gate_residual +from rl_engine.kernels.ops.pytorch.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) -> 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("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].contiguous() + 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, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + version = -1 if index is None else index._version + for obj, ver, rows, plan in self._plans: + if obj is index and ver == version and rows == num_rows: + return plan + plan = sp_plan(self.layout, index, num_rows) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, 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), 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/testing/h3_cases.py b/rl_engine/testing/h3_cases.py index 527a820ec..512d3ca4a 100644 --- a/rl_engine/testing/h3_cases.py +++ b/rl_engine/testing/h3_cases.py @@ -33,3 +33,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/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..b2e3b448f 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -120,6 +120,187 @@ 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.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.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.kernels.ops.cuda.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.testing.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.kernels.ops.cuda.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.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.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.kernels.ops.cuda.h3.gate_residual import H3GateResidualCudaOp + from rl_engine.kernels.ops.cuda.h3.rmsnorm import H3RMSNormCudaOp + from rl_engine.kernels.ops.cuda.h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + from rl_engine.testing.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 ------------------------------------------------- diff --git a/scripts/h3_ws2_evidence.py b/scripts/h3_ws2_evidence.py index c27332f77..021cfc1f8 100644 --- a/scripts/h3_ws2_evidence.py +++ b/scripts/h3_ws2_evidence.py @@ -11,6 +11,8 @@ export RL_KERNEL_H3_WEIGHTS= python scripts/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ --out docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json + python scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out docs/usage/evidence/h3-sp-norm-adaln-b200/report.json """ from __future__ import annotations @@ -29,7 +31,13 @@ load_h3_conditioning_weights, load_h3_manifest, ) -from rl_engine.testing.h3_ws2 import run_world, tp_projection_sweep # noqa: E402 +from rl_engine.testing.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" BIAS = "transformer_blocks.0.adaln_proj.linear.bias" @@ -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: diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index ba4b3b4ef..6ee3de114 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -511,43 +511,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 (every rank, every T)", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "rank 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"]) @@ -589,6 +559,144 @@ 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(title, 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)} {unit}" + 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 (every rank, every case)", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "rank x case", + ) + + # 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)") + _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, @@ -598,6 +706,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, } diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..04d58395c --- /dev/null +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,210 @@ +# 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.kernels.ops.cuda.h3.sp_norm_adaln import SPRowLayout, sp_row_layout +from rl_engine.testing.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.kernels.ops.cuda.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.kernels.ops.cuda.h3.adaln_row_gather import _segment_tiles + from rl_engine.kernels.ops.cuda.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]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.kernels.ops.cuda.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.kernels.ops.cuda.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"} From ef0c9ee24bf2924af0cbdea5e6d64022118a5366 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:01 +0000 Subject: [PATCH 11/62] [WS2][MiniMax-H3] tp_adaln_3mod: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --worlds 1,2,4,8 at 4947996 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-tp-adaln-3mod.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index da7b0ae42..6f525efe6 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -70,8 +70,24 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![tp_adaln_3mod on 8 x B200: byte equality and time](../usage/evidence/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. ## Tests From 63b8caa226e2af3aa49b7e9a2a913f8f235692be Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Wed, 7 Oct 2026 19:43:48 +0000 Subject: [PATCH 12/62] [WS2][MiniMax-H3] sp_norm_adaln: 8 x B200 NCCL evidence report and figure Produced by scripts/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 at 6c0d380 on a clean tree (real NCCL, one GPU per rank); rendered with scripts/plot_h3_evidence.py (equality-grid label moved to the panel title) and embedded in the operator doc. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-sp-norm-adaln.md | 25 +++++++++++++++++++++++-- scripts/plot_h3_evidence.py | 12 ++++++------ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 167b63f7c..8506060ea 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,8 +70,29 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -The report is written by `scripts/h3_ws2_evidence.py` with real NCCL processes on one 8 x B200 -node, one GPU per rank. +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../usage/evidence/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `scripts/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, 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 diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 6ee3de114..c4f7c05c6 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -513,10 +513,10 @@ def plot_tp_adaln(report: dict[str, Any]): runs = report["runs"] _equality_grid( left, - "Byte-equal to WS1 (every rank, every T)", + "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, - "rank x T", + "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"]) @@ -563,7 +563,7 @@ 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(title, loc="left", fontsize=11, color=INK, pad=10) + 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) @@ -579,7 +579,7 @@ def _equality_grid(ax, title, rows, tensors, unit) -> None: linewidth=2, ) ) - label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)} {unit}" + 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) @@ -624,13 +624,13 @@ def plot_sp_norm(report: dict[str, Any]): runs = report["runs"] _equality_grid( axes[0], - "Byte-equal to WS1 (every rank, every case)", + "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, - "rank x case", + "ranks x cases", ) # Middle: the largest block-layout case, slowest rank. From e4b8a62afa5980330fa73a0085e0c09b3cd58a92 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:50:47 +0000 Subject: [PATCH 13/62] fix(h3): store norm_out.norm and 1 + scale in BF16 in the final AdaLN golden The gtest gold skipped two roundings diffusers' norm_out performs, so the BF16 check missed d_norm_weight from S=257 and d_temb/dW from S=1024 for CUDA and the PyTorch reference alike. With them declared straight-through, the BF16 gtest passes at S=257/1024/4097; FP32 is unchanged and its long-reduction misses are documented. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 37 ++++++++++++++----- .../kernels/ops/pytorch/h3/final_adaln_out.py | 19 +++++++--- tests/h3/test_h3_final_adaln_out.py | 28 ++++++++++++++ 3 files changed, 70 insertions(+), 14 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 59b4ab0a9..7439e84a1 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -45,9 +45,12 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; - every gradient is repeat-bitwise. -- **Golden.** FP64, rounding only where the model declares a dtype boundary: the SiLU cast, - and `norm_out.linear`'s BF16 output table, which every position of a timestep shares. Both are - rounded straight-through for the gradient. +- **Golden.** FP64, rounding only where the model stores a value in its own dtype: the SiLU + cast, `norm_out.linear`'s BF16 output table, `norm_out.norm`'s BF16 output and the BF16 + `1 + scale`. All four are rounded straight-through for the gradient. The table and `1 + scale` + are shared by every position of a timestep. Without them the golden's `d_norm_w` drifts + systematically with S, and the gtest missed it from S = 257. `norm_out.norm`'s rounding + enters `d_scale`, `dW` and `d_temb` summed over S, and the gtest missed those from S = 1024. ## Performance Notes @@ -87,7 +90,7 @@ export RL_KERNEL_H3_WEIGHTS= python -m pytest tests/h3/test_h3_final_adaln_out.py -v # operator python -m pytest tests/h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out python scripts/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ - --dtype bf16 --batch 3 --seq 64 --normalized-dim 5376 --check-grad + --dtype bf16 --batch 3 --seq 4097 --normalized-dim 5376 --check-grad # also 257, 1024 python scripts/h3_evidence.py --op final_adaln_out \ --out docs/usage/evidence/h3-final-adaln-out-b200/report.json python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/report.json @@ -95,9 +98,25 @@ python scripts/plot_h3_evidence.py docs/usage/evidence/h3-final-adaln-out-b200/r ## Known Limitations -- **gtest at larger S.** With `--seq 257`, a few gradient elements fail the per-element - contract tolerance (BF16 `d_norm_w`, FP32 `d_temb`). Those elements are what is left after - summing hundreds of rows into values around 1e2, which an absolute per-element atol can't - cover. Relative to the gradient's maximum the error is 4e-7 in FP32 and 5e-3 in BF16, and the - PyTorch reference fails the same elements, plus more. `--seq 64` passes in both dtypes. +- **FP32 gtest at larger S.** H3 runs this op with BF16 activations and weights, and the BF16 + gtest passes at S = 257, 1024 and 4097. With `--dtype fp32` (every input FP32), the FP32 + reduction tolerance (atol = rtol = 1e-4) misses `d_temb` from S = 257 and `dW` from + S = 1024. Both gradients are FP32 sums over 3·S rows. Where large terms cancel, the + summation error exceeds the 1e-4 absolute floor. The PyTorch reference misses the same + gradients at the same S (`--batch 3`, seed 123, max abs error): + + | S | `d_temb` CUDA | `d_temb` PyTorch | `dW` CUDA | `dW` PyTorch | + | --- | --- | --- | --- | --- | + | 64 | 9.2e-5 | 1.5e-4 | 1.5e-4 | 9.2e-5 | + | 257 | 3.1e-4 ✗ | 5.9e-4 ✗ | 8.5e-4 | 2.4e-4 | + | 1024 | 1.5e-3 ✗ | 2.4e-3 ✗ | 2.0e-3 ✗ | 1.5e-3 ✗ | + | 4097 | 5.9e-3 ✗ | 9.3e-3 ✗ | 5.9e-3 ✗ | 3.9e-3 ✗ | + + ✗ fails the gtest. Unmarked values above 1e-4 pass through the rtol term. In a PyTorch + emulation, accumulating only the projection backward in FP64 clears S = 257 but not + S ≥ 1024, because the segment sums contribute as well. Making every reduction FP64 would + change the `h3_det_linear` kernels shared with `timestep_mlp_fp32` and + `adaln_projection_3mod`. The alternative is an FP32 gradient + tolerance that scales with reduction length. Which of the two to take is an open question + for the maintainers. - There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py index 94aaa28c1..b40924e12 100644 --- a/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py +++ b/rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py @@ -58,9 +58,10 @@ def validate_h3_final_adaln_out( class NativeH3FinalAdaLNOutOp: """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; - ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only at the - module boundaries the model declares (the SiLU cast and the BF16 table), - straight-through for the gradient.""" + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only where + the model itself stores a value in its own dtype (the SiLU cast, the BF16 + table, the ``norm_out.norm`` output and ``1 + scale``), straight-through for + the gradient, and computes everything else in FP64.""" op_class = "reduction" @@ -89,7 +90,15 @@ def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps shift, scale = table.chunk(2, dim=-1) x64 = x.double() n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() - out = n * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( - 0, timestep_indices + # norm_out.norm is an x.dtype module and diffusers forms 1 + scale in + # x.dtype, so both are stored rounded. 1 + scale is shared by every + # position of a timestep: left unrounded, its error is systematic and + # grows with S in d_norm_weight; n's rounding enters d_scale (and so + # dW, d_temb) summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.index_select(0, timestep_indices) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() ) + out = n * one_plus_scale + shift.index_select(0, timestep_indices) return out.float() diff --git a/tests/h3/test_h3_final_adaln_out.py b/tests/h3/test_h3_final_adaln_out.py index 17a3843b8..7a1f5b6d8 100644 --- a/tests/h3/test_h3_final_adaln_out.py +++ b/tests/h3/test_h3_final_adaln_out.py @@ -169,3 +169,31 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("final_adaln_out", device="cuda") assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_bf16_gradients_meet_gtest_contract(self, seq): + # The check_operator path; S=257 used to miss d_norm_weight before the + # golden stored norm(x) and 1 + scale in BF16 like the model does. + import argparse + + from rl_engine.kernels.gtest import run_operator_suite + from rl_engine.kernels.gtest.operator_specs import make_candidate, make_operator_case + + _cuda_op() + args = argparse.Namespace( + op="final_adaln_out", + candidate="cuda", + arch_key=None, + batch=3, + seq=seq, + normalized_dim=5376, + seed=123, + input_mode="random", + ) + report = run_operator_suite( + "final_adaln_out", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.bfloat16, torch.device("cuda"))], + check_grad=True, + ) + assert report.passed, report.candidates[0].cases[0].outputs From dcbd55bfd63cb067f407f3abdcc08fa674ea6638 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 09:57:30 +0000 Subject: [PATCH 14/62] docs(h3): regenerate final AdaLN evidence with the BF16 norm_out golden Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 7439e84a1..b7dd49d65 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,7 +40,7 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 22–43× closer to FP64 in the recorded run than diffusers, + - `d_temb`, `dW` and `db` are 22–48× closer to FP64 in the recorded run than diffusers, whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; @@ -62,9 +62,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.14 ms | 0.15 ms | 1.67 ms | 1.12 ms | -| 32768 | 0.40 ms | 0.78 ms | 2.97 ms | 5.39 ms | -| 131072 | 1.25 ms | 2.94 ms | 7.87 ms | 22.2 ms | +| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | +| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | +| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,11 +77,13 @@ There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `372d6d7`, with FP64 leaves and upstream gradients and explicit BF16 forward boundaries. + tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + backward reference keeps only the SiLU and table roundings, so the errors include the BF16 + rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `38d575f`. + tree at `001684d`. ## Tests From dd35c14482d7a7a5693b78848230acccd0e555ca Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:16:48 +0000 Subject: [PATCH 15/62] docs(h3): gate_residual vs existing implementations (B200, 4010854) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 2ef31ba54..6b923d740 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -81,6 +81,29 @@ The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/rep written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](../usage/evidence/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../usage/evidence/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python scripts/h3_prior_art.py --op gate_residual --out docs/usage/evidence/h3-prior-art-b200/gate_residual.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From 88053296fef5e325c8ed84c89f3434dd68738b85 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 8 Oct 2026 23:30:16 +0000 Subject: [PATCH 16/62] docs(h3): final_adaln_out vs existing implementations (B200, 03da729) Batch invariance (every row, full batch vs covering sub-batches, size sweep), accuracy against FP64 and latency, from scripts/h3_prior_art.py on a clean tree. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index b7dd49d65..e7bac8c59 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -85,6 +85,29 @@ There are two data files: T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `001684d`. +## Existing implementations (RFC #420 reuse rule) + +![final_adaln_out vs existing implementations](../usage/evidence/h3-prior-art-b200/final_adaln_out.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay) | **no** (param/table grads not repeatable) | 8.3e-03 / 7.3e-02 / 1021 µs | 8.9e-03 / 2.4e-01 / 4525 µs | +| rl-kernel H3FinalAdaLNOutCudaOp | yes | 8.3e-03 / 5.2e-03 / 1642 µs | 8.9e-03 / 5.7e-03 / 2959 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../usage/evidence/h3-prior-art-b200/final_adaln_out.json) +was written from a clean tree at `03da729` by + +```bash +python scripts/h3_prior_art.py --op final_adaln_out --out docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json +python scripts/plot_h3_prior_art.py docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + ## Tests ```bash From a9ed1267326f3efb50a9cf1d2a27188fdf52664a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:11:51 +0000 Subject: [PATCH 17/62] Fix H3 native validation, backward timing, and rank cleanup Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 4 +- csrc/cuda/h3/gate_residual.cu | 7 +++- csrc/cuda/h3/rmsnorm_modulate.cu | 23 +++++++---- docs/operators/h3-adaln-projection.md | 5 +++ docs/operators/h3-tp-adaln-3mod.md | 1 + rl_engine/testing/h3_report.py | 28 ++++++++++---- rl_engine/testing/h3_ws2.py | 55 ++++++++++++++++++++++----- tests/h3/test_h3_benchmark_timing.py | 50 +++++++++++++++++++++++- tests/h3/test_h3_native_validation.py | 51 ++++++++++++++++++++++++- tests/h3/test_h3_ws2_processes.py | 47 +++++++++++++++++++++++ 10 files changed, 239 insertions(+), 32 deletions(-) create mode 100644 tests/h3/test_h3_ws2_processes.py diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..9f3d00973 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -207,9 +207,9 @@ 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->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 1-D int64 CUDA tensors"); + "tile metadata must be contiguous int64 tensors on ", grad.device()); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index c2bd7921a..6e0b5a85b 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -50,6 +50,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -60,7 +61,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -148,7 +151,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index e2c05d2d0..19700d35a 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -106,6 +106,7 @@ struct Modulation { int64_t row_stride; 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 @@ -120,6 +121,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -156,9 +158,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -193,9 +198,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* const int64_t r1 = min(r0 + kRowTile, rows); float acc = 0.0f; for (int64_t r = r0; r < r1; ++r) { - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -260,7 +268,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + Modulation mod{nullptr, nullptr, 0, nullptr, 1, 0}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -287,6 +295,7 @@ Modulation make_modulation(const c10::optional& shift, mod.row_stride = sh.stride(0); mod.index = ix.data_ptr(); mod.seq = ix.size(0); + mod.num_rows = sh.size(0); return mod; } diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index 910165a3b..2cd2cbe0d 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -135,6 +135,11 @@ or in a batch of 9. ![adaln_projection vs existing implementations](../usage/evidence/h3-prior-art-b200/adaln_projection.png) +The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` +streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional +weight passes and are outside its intended T ≤ 8 range. This accounts for the +steep increase in the reported CUDA forward + backward times (77.4 and 627 ms). + | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| | diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..7a33f6d26 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -83,6 +83,7 @@ byte-equal to WS1 computed on the same GPU. | 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | | 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | | 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about 5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..3bd094314 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -606,10 +609,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +700,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 backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b)): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -708,8 +716,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index a0dc116f6..ad9041cf9 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -144,9 +144,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -160,7 +160,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -176,12 +178,45 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + try: + for proc in procs: + proc.start() + started.append(proc) + deadline = time.monotonic() + timeout + pending = set(range(world)) + while pending: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"timed out waiting for ranks {sorted(pending)}") + 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 + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + pending.remove(result["rank"]) + 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") + 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)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..47d5575b1 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,9 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(leaves) + prepared_leaves.append( + [leaf._base if leaf._base is not None else leaf for leaf in leaves] + ) cuda_clock.time += 1.0 outputs = fn(*leaves, *args) return track_backward(outputs) if leaves[0].requires_grad else outputs @@ -162,6 +164,50 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 6 * 4), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + 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) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + 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) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..8c71eba08 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -16,6 +16,7 @@ def native_extension(): from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE names = ( + "h3_adaln_row_gather_backward", "h3_rmsnorm_forward", "h3_rmsnorm_backward", "h3_gate_residual_forward", @@ -49,6 +50,10 @@ 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": @@ -86,7 +91,7 @@ def test_native_rejects_invalid_row_index( _native_call(native_extension, operation, boundary, native_case) -@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual", "gather"]) @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( @@ -122,3 +127,47 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@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 + import subprocess + import sys + + 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() +""" + result = subprocess.run( + [sys.executable, "-c", code], + 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 diff --git a/tests/h3/test_h3_ws2_processes.py b/tests/h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..fb69430cd --- /dev/null +++ b/tests/h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.testing import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=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"): + 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" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before From 72ed9702872be0ae4a5518ebe59579c264f54ead Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:09 +0000 Subject: [PATCH 18/62] Regenerate gate and final AdaLN backward-only B200 evidence Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 10 +++++----- docs/operators/h3-final-adaln-out.md | 13 +++++++------ 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..8d3b4b380 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -61,13 +61,13 @@ hidden = op(residual, attn_output, gate_msa, adaln_indices) python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual ``` -B200, B = 1, H = 5376, BF16: +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](../usage/evidence/h3-adaln-gate-residual-b200/report.json), -written by `scripts/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `scripts/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..2c2e215c6 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -40,7 +40,7 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste summation tree versus cuBLAS. Rows are batch- and position-invariant. - **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into the projection backward without an extra BF16 rounding: - - `d_temb`, `dW` and `db` are 22–48× closer to FP64 in the recorded run than diffusers, + - `d_temb`, `dW` and `db` are 21–37× closer to FP64 in the recorded run than diffusers, whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with `index_select` atomics; - `dx` and `d_norm_w` are at the BF16 rounding level for both; @@ -58,13 +58,14 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out ``` -B200, B = 1, T = 3, pinned `norm_out` weights: +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 | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | -| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | -| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | +| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | +| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | +| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,7 +78,7 @@ There are two data files: - [`report.json`](../usage/evidence/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../usage/evidence/h3-final-adaln-out-b200/chain_replay.json): the From d9600066ccd459f6f71e2daf4e071aee6d28854a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:14:40 +0000 Subject: [PATCH 19/62] test(h3): consolidate native validation tests for review file limit Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/h3/test_h3_native_validation.py | 153 +++++++++++++++++ tests/h3/test_h3_rmsnorm_native_validation.py | 158 ------------------ 2 files changed, 153 insertions(+), 158 deletions(-) delete mode 100644 tests/h3/test_h3_rmsnorm_native_validation.py diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 8dd919028..9f2a06a07 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -122,3 +122,156 @@ def test_native_accepts_same_device_inputs(native_extension, native_case, operat actual_grads = _native_call(native_extension, operation, "backward", native_case) for actual, expected in zip(actual_grads, expected_grads, strict=True): torch.testing.assert_close(actual, expected) + + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.kernels.ops.base 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("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) diff --git a/tests/h3/test_h3_rmsnorm_native_validation.py b/tests/h3/test_h3_rmsnorm_native_validation.py deleted file mode 100644 index d6f2a0d81..000000000 --- a/tests/h3/test_h3_rmsnorm_native_validation.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# Copyright (c) 2026 RL-Kernel Contributors - -"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" - -from __future__ import annotations - -import pytest -import torch - -pytestmark = [ - pytest.mark.cuda_only, - pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), -] - -TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") - - -@pytest.fixture -def native(): - from rl_engine.kernels.ops.base 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}") - - -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.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.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.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.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.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.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.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) From 86c168f2f5091b8f3b6b947dd3ef5a15ce30b19d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 9 Oct 2026 02:42:12 +0000 Subject: [PATCH 20/62] fix(h3): address automated review bounds, device, and benchmark findings Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 20 +++- csrc/cuda/h3/gate_residual.cu | 7 +- csrc/cuda/h3/rmsnorm_modulate.cu | 23 ++-- docs/operators/h3-adaln-gate-residual.md | 7 +- docs/operators/h3-final-adaln-out.md | 7 +- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 6 +- .../kernels/ops/cuda/h3/sp_norm_adaln.py | 27 +++-- .../ops/cuda/h3/tp_adaln_projection.py | 2 + rl_engine/testing/h3_chain.py | 4 +- rl_engine/testing/h3_report.py | 30 +++-- rl_engine/testing/h3_ws2.py | 62 +++++++++-- scripts/plot_h3_evidence.py | 14 ++- tests/h3/test_h3_benchmark_timing.py | 52 ++++++++- tests/h3/test_h3_native_validation.py | 103 ++++++++++++++++++ tests/h3/test_h3_sp_norm_adaln.py | 103 ++++++++++++++++++ tests/h3/test_h3_tp_adaln.py | 8 +- 17 files changed, 423 insertions(+), 54 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4b3a168a0..0de2d98d5 100644 --- a/csrc/cuda/h3/adaln_row_gather.cu +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -110,7 +110,10 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, const int64_t h = j - c * hidden; const g_t* g = grad + c * seq * hidden + h; float acc = 0.0f; - for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t begin = tile_begin[tile], end = tile_end[tile]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= seq); + for (int64_t p = begin; p < end; ++p) { + CUDA_KERNEL_ASSERT(sorted_pos[p] >= 0 && sorted_pos[p] < seq); acc += to_float(g[sorted_pos[p] * hidden]); } partial[tile * width + j] = acc; @@ -119,12 +122,14 @@ __global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, template __global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, const int64_t* __restrict__ seg_first_tile, - out_t* __restrict__ out, int64_t width) { + out_t* __restrict__ out, int64_t width, int64_t tiles) { const int64_t r = blockIdx.x; const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; if (j >= width) return; float acc = 0.0f; - for (int64_t tile = seg_first_tile[r]; tile < seg_first_tile[r + 1]; ++tile) { + const int64_t begin = seg_first_tile[r], end = seg_first_tile[r + 1]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= tiles); + for (int64_t tile = begin; tile < end; ++tile) { acc += partial[tile * width + j]; } out[r * width + j] = from_float(acc); @@ -207,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->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 1-D int64 CUDA tensors"); + "tile metadata must be contiguous 1-D int64 tensors on grad device"); } TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t chunks = grad.size(0); @@ -218,6 +224,8 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t width = chunks * hidden; const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); @@ -244,7 +252,7 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor adaln_row_gather_fold_kernel <<(num_rows), col_blocks), threads, 0, stream>>>( partial.data_ptr(), seg_first_tile.data_ptr(), - reinterpret_cast(out.data_ptr()), width); + reinterpret_cast(out.data_ptr()), width, tiles); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return out; diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 2d8aea209..eb0b03f50 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -52,6 +52,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -62,7 +63,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -150,7 +153,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 42e3ee18a..f9304ab13 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -108,6 +108,7 @@ 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] }; @@ -124,6 +125,7 @@ __global__ void __launch_bounds__(32 * kWarps) const T* scale = nullptr; if (mod.index != nullptr) { const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); shift = static_cast(mod.shift) + i * mod.row_stride; scale = static_cast(mod.scale) + i * mod.row_stride; } @@ -160,9 +162,12 @@ __global__ void __launch_bounds__(32 * kWarps) const int64_t r = blockIdx.x; const T* xr = x + r * n; const T* gr = g + r * n; - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } const int numx = blockDim.x * blockDim.y; const int thrx = threadIdx.x + threadIdx.y * blockDim.x; float s = 0.0f; @@ -204,9 +209,12 @@ __global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* float acc = 0.0f; for (int64_t p = p0; p < p1; ++p) { const int64_t r = row_list != nullptr ? row_list[p] : p; - const T* scale = mod.index != nullptr ? static_cast(mod.scale) + - mod.index[r % mod.seq] * mod.row_stride - : nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); } partial[static_cast(blockIdx.x) * n + j] = acc; @@ -271,7 +279,7 @@ Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { - Modulation mod{nullptr, nullptr, 0, nullptr, 1}; + 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; @@ -296,6 +304,7 @@ 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); return mod; diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6b923d740..9d5ee7e09 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -63,7 +63,7 @@ python benchmarks/benchmark_h3_conditioning.py --op adaln_gate_residual B200, B = 1, H = 5376, BF16: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | | 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | @@ -73,6 +73,11 @@ The forward is one pass over `residual` and `sublayer_output`, using 16-byte vec block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and tile setup. +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 ![adaln_gate_residual on B200: latency and backward accuracy](../usage/evidence/h3-adaln-gate-residual-b200/figure.png) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index e7bac8c59..6022af96a 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -60,7 +60,7 @@ python benchmarks/benchmark_h3_conditioning.py --op final_adaln_out B200, B = 1, T = 3, pinned `norm_out` weights: -| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | | --- | --- | --- | --- | --- | | 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | | 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | @@ -69,6 +69,11 @@ B200, B = 1, T = 3, pinned `norm_out` weights: 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](../usage/evidence/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 index 8506060ea..65b02d021 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -65,7 +65,7 @@ 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 – 174 | +| 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 diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 6f525efe6..0f6a64a35 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -28,8 +28,10 @@ op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the `N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. `N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard -never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when -the op is built. +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. | TP | Columns per rank | Rank 0 owns | | --- | --- | --- | diff --git a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py index 19a867b3e..7b5eb20d3 100644 --- a/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py +++ b/rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py @@ -211,16 +211,18 @@ class SPPlan: num_rows: int | None -def sp_plan(layout: SPRowLayout, index_full=None, num_rows: int | None = None) -> SPPlan: +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("cuda") + 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].contiguous() + 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) @@ -334,18 +336,21 @@ def __init__(self, collective, seq_len: int, batch: int) -> None: 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, SPPlan]] = [] + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] - def plan(self, index: torch.Tensor | None, num_rows: int | None) -> 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()) version = -1 if index is None else index._version - for obj, ver, rows, plan in self._plans: - if obj is index and ver == version and rows == num_rows: + 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) + 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, plan), *self._plans[:3]] + 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: @@ -368,7 +373,9 @@ def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | No 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), self.collective, 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]) diff --git a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py index 0eb5a4709..f5ec080a3 100644 --- a/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py +++ b/rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py @@ -84,6 +84,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/testing/h3_chain.py b/rl_engine/testing/h3_chain.py index 3af6ffd38..8f41be297 100644 --- a/rl_engine/testing/h3_chain.py +++ b/rl_engine/testing/h3_chain.py @@ -378,7 +378,9 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc else: value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) outputs = list(value) - torch.autograd.backward(outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs)]) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) return [leaves[name].grad for name in GRAD_LEAVES] diff --git a/rl_engine/testing/h3_report.py b/rl_engine/testing/h3_report.py index 407906a88..877a2534d 100644 --- a/rl_engine/testing/h3_report.py +++ b/rl_engine/testing/h3_report.py @@ -590,9 +590,13 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -606,10 +610,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } @@ -695,9 +701,13 @@ def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: x, nw, temb, w, b, ti = _final_inputs(seq) grad = torch.randn_like(x) - def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves), grad + + def backward(state): + output, grad = state + output.backward(grad) cases.append( { @@ -708,8 +718,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: provider_final_adaln_out(*t, i) ), } diff --git a/rl_engine/testing/h3_ws2.py b/rl_engine/testing/h3_ws2.py index b2e3b448f..13174f7f6 100644 --- a/rl_engine/testing/h3_ws2.py +++ b/rl_engine/testing/h3_ws2.py @@ -325,9 +325,9 @@ def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, inputs_path.with_name(f"rank{rank}.pt"), ) - queue.put({"rank": rank}) collective.close() dist.destroy_process_group() + queue.put({"rank": rank}) except Exception: # forwarded to the parent queue.put({"rank": rank, "error": traceback.format_exc()}) @@ -341,7 +341,9 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa """ import tempfile + import time from pathlib import Path + from queue import Empty import torch.multiprocessing as mp @@ -357,12 +359,52 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) for r in range(world) ] - for proc in procs: - proc.start() - status = sorted((queue.get(timeout=timeout) for _ in procs), key=lambda r: r["rank"]) - for proc in procs: - proc.join(timeout=120) - for result in status: - if "error" in result: - raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") - return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + started = [] + deadline = time.monotonic() + timeout + status = {} + try: + for proc in procs: + proc.start() + started.append(proc) + while len(status) < world or any(proc.is_alive() for proc in procs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"rank results timed out after {timeout}s") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + 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']}") + status[result["rank"]] = result + for rank, proc in enumerate(procs): + 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)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/scripts/plot_h3_evidence.py b/scripts/plot_h3_evidence.py index 7328207ad..bd2ada020 100644 --- a/scripts/plot_h3_evidence.py +++ b/scripts/plot_h3_evidence.py @@ -440,6 +440,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, @@ -447,9 +449,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], "////", @@ -459,6 +466,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") @@ -490,6 +499,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) diff --git a/tests/h3/test_h3_benchmark_timing.py b/tests/h3/test_h3_benchmark_timing.py index e6517bd70..d709468b3 100644 --- a/tests/h3/test_h3_benchmark_timing.py +++ b/tests/h3/test_h3_benchmark_timing.py @@ -114,7 +114,7 @@ def run(key=key, duration=position): assert not any(key.endswith("_setup") for key in row) -@pytest.mark.parametrize("operator", ["norm", "gather"]) +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( monkeypatch, cuda_clock, operator ): @@ -136,7 +136,7 @@ def backward_hook(grad): def track(fn, leaves, *args): if leaves[0].requires_grad: assert not cuda_clock.active - prepared_leaves.append(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 @@ -162,6 +162,54 @@ def forward(x, weight, shift, scale, index, **kwargs): monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) op = SimpleNamespace(forward_modulated=forward) build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + 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) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(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) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf else: randn = torch.randn diff --git a/tests/h3/test_h3_native_validation.py b/tests/h3/test_h3_native_validation.py index 9f2a06a07..4e59cfcce 100644 --- a/tests/h3/test_h3_native_validation.py +++ b/tests/h3/test_h3_native_validation.py @@ -275,3 +275,106 @@ def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name) 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_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + 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", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + 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 diff --git a/tests/h3/test_h3_sp_norm_adaln.py b/tests/h3/test_h3_sp_norm_adaln.py index 04d58395c..2b91aad64 100644 --- a/tests/h3/test_h3_sp_norm_adaln.py +++ b/tests/h3/test_h3_sp_norm_adaln.py @@ -208,3 +208,106 @@ def test_nccl_region_byte_equal_to_ws1(world): 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.kernels.ops.cuda.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) + + +@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.kernels.ops.cuda.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.testing 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.testing 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/h3/test_h3_tp_adaln.py b/tests/h3/test_h3_tp_adaln.py index 8b341c8d3..91079b691 100644 --- a/tests/h3/test_h3_tp_adaln.py +++ b/tests/h3/test_h3_tp_adaln.py @@ -48,7 +48,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]]] @@ -146,3 +146,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) From 3bbcca42f7e8a5a8a7b4d8753a46c108cd1bd653 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:30:51 +0000 Subject: [PATCH 21/62] evidence(h3): gate_residual direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 910d614a. There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e69de13cd..e003fd5b8 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -104,6 +104,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + ## Tests ```bash From 28099904ff6ed0197d9389c8ca456598bbbb9045 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 14:32:41 +0000 Subject: [PATCH 22/62] evidence(h3): final_adaln_out direct bitwise comparison on B200 Written by tools/validation/models/h3_bitwise_reuse.py from a clean tree at 4430ef16. There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index d3c938d55..3b571efec 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -108,6 +108,25 @@ python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior Libraries that do not import are skipped and recorded as unavailable in the report. +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) + +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. +[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png +``` + ## Tests ```bash From a4bc99f958003230273f05de48cd03ac796a4474 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:54:31 +0000 Subject: [PATCH 23/62] fix(h3): round the final_adaln_out goldens' activation from the FP32 SiLU Same mismatch as a0834552 for adaln_projection: forward_fp32, the test's _golden and the evidence report's golden_backward rounded an FP64 SiLU to BF16 for the declared cast, while the model and the kernel round the FP32 SiLU. Near a BF16 tie that picks the other neighbour and shifts a whole column of the norm_out table. All three now round F.silu(temb.float()); the VJP is unchanged. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/reference/minimax_h3/final_adaln_out.py | 5 ++++- rl_engine/validation/models/h3_report.py | 4 +++- tests/models/minimax_h3/test_h3_final_adaln_out.py | 4 +++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/rl_engine/reference/minimax_h3/final_adaln_out.py b/rl_engine/reference/minimax_h3/final_adaln_out.py index 0b19aaee3..a21a5cbdb 100644 --- a/rl_engine/reference/minimax_h3/final_adaln_out.py +++ b/rl_engine/reference/minimax_h3/final_adaln_out.py @@ -81,7 +81,10 @@ def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps ) t64 = temb.double() act = t64 * torch.sigmoid(t64) - act = act + (act.to(weight.dtype).double() - act).detach() # declared cast, identity VJP + # Declared cast of the FP32 SiLU, as the model computes it (an FP64 SiLU can + # round to the other neighbour near a tie); identity VJP. + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() table = F.linear(act, weight.double(), bias.double()) # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the # weight dtype, and every position of a timestep shares that row, so the diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 51407a030..107ee2ac2 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -742,7 +742,9 @@ def grads(fn, dtype=None): def golden_backward(x_, nw_, t_, w_, b_): # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. act = t_ * torch.sigmoid(t_) - act = act + (act.to(torch.bfloat16).double() - act).detach() + # The declared cast rounds the FP32 SiLU, as the model computes it. + declared = torch.nn.functional.silu(t_.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() table = torch.nn.functional.linear(act, w_, b_) table = table + (table.to(torch.bfloat16).double() - table).detach() shift, scale = table.chunk(2, dim=-1) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py index 5095d2c44..5fb964677 100644 --- a/tests/models/minimax_h3/test_h3_final_adaln_out.py +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -53,7 +53,9 @@ def _inputs(seq, num_timesteps=3, seed=0, batch=1): def _golden(x, norm_weight, temb, weight, bias, timestep_indices): act = temb * torch.sigmoid(temb) - act = act + (act.to(torch.bfloat16).double() - act).detach() + # Declared cast of the FP32 SiLU, as the model computes it; identity VJP. + declared = F.silu(temb.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() table = F.linear(act, weight, bias) table = table + (table.to(torch.bfloat16).double() - table).detach() shift, scale = table.chunk(2, dim=-1) From 7d716f227333a5265756f8f6fe3680ec5f38cf5e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:25:03 +0000 Subject: [PATCH 24/62] evidence(h3): regenerate final_adaln_out report and chain replay after the golden fix Written by tools/validation/models/h3_evidence.py and h3_chain_replay.py (--timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 --backward) from a clean tree at 50e913fa on an idle B200. The CUDA accuracy and every pass/fail flag are unchanged; the h3_rmsnorm stage's chained error against the golden drops in all 16 cases (worst 0.0796 -> 0.0671). The latency table in the doc follows the new report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3b571efec..44e08e507 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -62,9 +62,9 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.13 ms | 0.15 ms | 1.57 ms | 1.07 ms | -| 32768 | 0.38 ms | 0.78 ms | 2.80 ms | 5.34 ms | -| 131072 | 1.24 ms | 2.94 ms | 7.75 ms | 20.6 ms | +| 4097 | 0.18 ms | 0.19 ms | 1.50 ms | 1.07 ms | +| 32768 | 0.44 ms | 0.84 ms | 2.75 ms | 5.34 ms | +| 131072 | 1.29 ms | 2.99 ms | 7.76 ms | 20.8 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -77,13 +77,16 @@ There are two data files: - [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `001684d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `50e913fa` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `001684d`. + tree at `50e913fa`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` + stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, + 0.0671 now). ## Existing implementations (RFC #420 reuse rule) From 76594af5cb43cb73698faa7b1698cdb85ab9aff8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 13 Oct 2026 16:28:22 +0000 Subject: [PATCH 25/62] evidence(h3): regenerate this branch's final_adaln_out report after the golden fix This branch carries a newer final_adaln_out report.json than #489. Regenerated by tools/validation/models/h3_evidence.py from a clean tree at 0a2d5589 on an idle B200; chain_replay.json comes from #489's regeneration at 50e913fa. CUDA accuracy and every pass/fail flag are unchanged; the doc's latency table follows the new report. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index b95c906b0..3e884dcf2 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -63,9 +63,9 @@ and forward-graph setup: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.19 ms | 0.18 ms | 1.40 ms | 0.85 ms | -| 32768 | 0.46 ms | 0.80 ms | 2.59 ms | 4.49 ms | -| 131072 | 1.32 ms | 3.02 ms | 7.05 ms | 17.69 ms | +| 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 | +| 131072 | 1.28 ms | 2.97 ms | 7.02 ms | 17.61 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. @@ -78,13 +78,16 @@ There are two data files: - [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `bde1d8d` on an otherwise idle B200, with FP64 leaves and upstream gradients. The + tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `001684d`. + tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). ## Existing implementations (RFC #420 reuse rule) From a8e5901b5d194bf846c5ad670ea9780850087c04 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:18 +0000 Subject: [PATCH 26/62] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 6 files changed, 37 insertions(+), 37 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e003fd5b8..8847d6161 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash From 39fb03ddf1447c5ac8a23eaab3df1298476b6a1a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:19 +0000 Subject: [PATCH 27/62] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ 7 files changed, 44 insertions(+), 44 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index e003fd5b8..8847d6161 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a76cc0417..a9c8832e2 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 44e08e507..bb1194db8 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -71,17 +71,17 @@ projection backward. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `50e913fa` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` @@ -90,7 +90,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -101,7 +101,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -118,10 +118,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash From c469b72becec588256a436e7ed45b4609a529d01 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:20 +0000 Subject: [PATCH 28/62] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ docs/operators/h3-tp-adaln-3mod.md | 2 +- 8 files changed, 45 insertions(+), 45 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6c972994d..9a5d815a5 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a4cb57dd8..e8fc09537 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional @@ -157,7 +157,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -174,10 +174,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 3e884dcf2..7970b5c2d 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -72,17 +72,17 @@ projection backward. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![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) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the @@ -91,7 +91,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -102,7 +102,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -119,10 +119,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 614be59c5..8354dff8b 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -70,7 +70,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -![tp_adaln_3mod on 8 x B200: byte equality and time](../../reports/experiments/h3-tp-adaln-3mod-b200/figure.png) +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For From 0033b3a5a513604c72987abbc2cc08bd22d5a67d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:57:22 +0000 Subject: [PATCH 29/62] docs: move the measurement outputs out of the pull request The raw results (JSON, figures, audit files) now live on the fork's pr_assets branch; the documents link to them at a fixed commit. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 12 ++++++------ docs/operators/h3-adaln-projection.md | 12 ++++++------ docs/operators/h3-adaln-row-gather.md | 14 +++++++------- docs/operators/h3-final-adaln-out.md | 14 +++++++------- docs/operators/h3-rmsnorm.md | 12 ++++++------ docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 12 ++++++------ docs/operators/h3-timestep-sinusoid.md | 12 ++++++------ docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 46 insertions(+), 46 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 6c972994d..9a5d815a5 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -75,15 +75,15 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](../../reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-gate-residual-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) -![gate_residual vs existing implementations](../../reports/experiments/h3-prior-art-b200/gate_residual.png) +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -94,7 +94,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](../../reports/experiments/h3-prior-art-b200/gate_residual.json) +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) was written from a clean tree at `4010854` by ```bash @@ -111,10 +111,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![gate_residual bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](../../reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a4cb57dd8..e8fc09537 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -124,9 +124,9 @@ dependence on T. ## Evidence -![adaln_projection_3mod on B200: latency and correctly rounded outputs](../../reports/experiments/h3-adaln-projection-b200/figure.png) +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-adaln-projection-b200/report.json), +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also @@ -135,7 +135,7 @@ or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![adaln_projection vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_projection.png) +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional @@ -157,7 +157,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](../../reports/experiments/h3-prior-art-b200/adaln_projection.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) was written from a clean tree at `c88691c` by ```bash @@ -174,10 +174,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_projection bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 88fee5f3d..94e32b711 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -122,13 +122,13 @@ alternates each iteration and is recorded in the report. ## Evidence -![adaln_row_gather on B200: latency and whole-chain gradient accuracy](../../reports/experiments/h3-adaln-row-gather-b200/figure.png) +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) There are two data files: -- [`report.json`](../../reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward bitwise checks, op-level backward, and the chain-backward cases plotted above. -- [`chain_replay.json`](../../reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. `report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only @@ -141,7 +141,7 @@ model does; every CUDA pass/fail flag is unchanged. ## Existing implementations (RFC #420 reuse rule) -![adaln_row_gather vs existing implementations](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.png) +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -152,7 +152,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](../../reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) was written from a clean tree at `6c900ae` by ```bash @@ -169,10 +169,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![adaln_row_gather bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](../../reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 93143b858..d40a49f5f 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -77,17 +77,17 @@ with their original commit and corrected scope. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](../../reports/experiments/h3-final-adaln-out-b200/figure.png) +![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) There are two data files: -- [`report.json`](../../reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. -- [`chain_replay.json`](../../reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the @@ -96,7 +96,7 @@ There are two data files: ## Existing implementations (RFC #420 reuse rule) -![final_adaln_out vs existing implementations](../../reports/experiments/h3-prior-art-b200/final_adaln_out.png) +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -107,7 +107,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](../../reports/experiments/h3-prior-art-b200/final_adaln_out.json) +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) was written from a clean tree at `03da729` by ```bash @@ -124,10 +124,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![final_adaln_out bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](../../reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 740f16232..75dd3d477 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -92,9 +92,9 @@ each iteration and is recorded in the report. ## Evidence -![h3_rmsnorm on B200: latency and backward accuracy](../../reports/experiments/h3-rmsnorm-b200/figure.png) +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report.json), written by +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; @@ -103,7 +103,7 @@ The data is in [`report.json`](../../reports/experiments/h3-rmsnorm-b200/report. ## Existing implementations (RFC #420 reuse rule) -![norm_modulate vs existing implementations](../../reports/experiments/h3-prior-art-b200/norm_modulate.png) +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) | Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | |---|---|---|---| @@ -119,7 +119,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](../../reports/experiments/h3-prior-art-b200/norm_modulate.json) +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) was written from a clean tree at `ee83dec` by ```bash @@ -136,10 +136,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![norm_modulate bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](../../reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 1f19d0444..8eb91cc38 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -70,7 +70,7 @@ How many rows move depends on how the table rows interleave along the sequence: ## Evidence -![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](../../reports/experiments/h3-sp-norm-adaln-b200/figure.png) +![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 `6c0d380`, with real NCCL processes on one 8 x B200 node, one GPU per rank. The region is diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index 8049a87df..cc5809cee 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -109,15 +109,15 @@ kernel time. ## Evidence -![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](../../reports/experiments/h3-timestep-mlp-b200/figure.png) +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-mlp-b200/report.json), written +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) -![timestep_mlp vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_mlp.png) +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -134,7 +134,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](../../reports/experiments/h3-prior-art-b200/timestep_mlp.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) was written from a clean tree at `ffd958a` by ```bash @@ -151,10 +151,10 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by (each global batch-invariant mode in its own process) and compares every output and leaf gradient with diffusers' path and with this kernel. -![timestep_mlp bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index c5f69a15b..75edf172b 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -88,15 +88,15 @@ about 1 KB per timestep. ## Evidence -![timestep_sinusoid_h3 on B200: latency and error vs FP64](../../reports/experiments/h3-timestep-sinusoid-b200/figure.png) +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) -The data is in [`report.json`](../../reports/experiments/h3-timestep-sinusoid-b200/report.json). +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). `tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report records that commit and the environment. ## Existing implementations (RFC #420 reuse rule) -![timestep_sinusoid vs existing implementations](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) | Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | |---|---|---|---|---| @@ -108,7 +108,7 @@ Errors are max|err| / max|ref| against the same computation in FP64; latency is forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full 8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A -"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](../../reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) was written from a clean tree at `7a82917` by ```bash @@ -126,11 +126,11 @@ Equal FP64 errors do not show equal bits, so the candidates are also compared by compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). -![timestep_sinusoid bitwise comparison](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. -[`timestep_sinusoid.json`](../../reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) was written from a clean tree at `0eb5ddd5` by ```bash diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index c348bce35..318dbd2d0 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -72,7 +72,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ## Evidence -![tp_adaln_3mod on 8 x B200: byte equality and time](../../reports/experiments/h3-tp-adaln-3mod-b200/figure.png) +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For From 7e9475d187a68a68897fd5fbb12c9d7115cb6804 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 05:08:41 +0200 Subject: [PATCH 30/62] feat(h3): add projection rows (audio in/out, text context, QKV) on a fixed-order Triton GEMM Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 650e2148c6f0e5dcdef136537e0e62f745020273) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d099c0b11585477d5b56367551c3ee8987954196) --- docs/operators/README.md | 1 + docs/operators/h3-projection-rows.md | 104 +++++ .../model_specific/minimax_h3/projection.py | 193 ++++++++ .../shared/triton/gemm/h3_projection.py | 192 ++++++++ rl_engine/reference/minimax_h3/projection.py | 188 ++++++++ rl_engine/runtime/registry.py | 44 ++ rl_engine/validation/models/h3_block.py | 25 +- tests/models/minimax_h3/test_h3_one_block.py | 7 +- .../minimax_h3/test_h3_projection_rows.py | 421 ++++++++++++++++++ .../validation/models/h3_projection_bench.py | 96 ++++ 10 files changed, 1267 insertions(+), 4 deletions(-) create mode 100644 docs/operators/h3-projection-rows.md create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/projection.py create mode 100644 rl_engine/backends/shared/triton/gemm/h3_projection.py create mode 100644 rl_engine/reference/minimax_h3/projection.py create mode 100644 tests/models/minimax_h3/test_h3_projection_rows.py create mode 100644 tools/validation/models/h3_projection_bench.py diff --git a/docs/operators/README.md b/docs/operators/README.md index c99451553..79302ed6b 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -43,3 +43,4 @@ Every operator page should include: - [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) - [MiniMax-H3 One Block, Forward and Backward](h3-one-block.md) - [Operator Doc Template](../contributing/operator-doc-template.md) +- [MiniMax-H3 Projection Rows](h3-projection-rows.md) diff --git a/docs/operators/h3-projection-rows.md b/docs/operators/h3-projection-rows.md new file mode 100644 index 000000000..a64c5b349 --- /dev/null +++ b/docs/operators/h3-projection-rows.md @@ -0,0 +1,104 @@ +# MiniMax-H3 Projection Rows + +## Summary + +Four RFC #420 rows share one contract: `audio_input_projection_fp32`, +`audio_output_projection_fp32`, `text_context_projection` and `h3_qkv_gemm`. Each is a linear +projection pinned to checkpoint tensors (`MiniMaxAI/MiniMax-H3@42ed227`, see +[the architecture fingerprint](h3-arch-fingerprint.md)): + +| Row | in -> out | dtype | bias | tensors | +| --- | --- | --- | --- | --- | +| `audio_input_projection_fp32` | 32 -> 5376 | float32 | yes | `audio_proj_in` | +| `audio_output_projection_fp32` | 5376 -> 32 | float32 | yes | `audio_proj_out` | +| `text_context_projection` | 5120 -> 5376 | bfloat16 | yes | `context_embedder` | +| `h3_qkv_gemm` (per projection) | 5376 -> 7168 | bfloat16 | no | `attn.to_q/to_k/to_v` | + +`audio_*` modules are FP32 modules: a lower-precision argument is a contract violation, not +a variant. `h3_qkv_gemm` also covers the fused serving layout (`[Q; K; V]` rows concatenated, +as in vLLM-Omni's `QKVParallelLinear`) and tensor-parallel head shards (any multiple of +128 output columns, so TP2/4/8 over 56 heads). + +## Entry Points + +```python +from rl_engine.runtime.registry import KernelRegistry + +registry = KernelRegistry() +op = registry.get_op("h3_qkv_gemm", device="cuda") +q, k, v = op.forward(x, wq, wk, wv) # Diffusers layout +fused = op.forward_fused(x, torch.cat([wq, wk, wv])) # serving layout, == [q | k | v] bitwise +y = registry.get_op("text_context_projection", device="cuda")(x, weight, bias) +``` + +## Backends + +| Backend | Class | Kernel | Status | +| --- | --- | --- | --- | +| CUDA (Triton) | `rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3*CudaOp` | `rl_engine.backends.shared.triton.gemm.h3_projection._projection_kernel` | strict | +| PyTorch reference | `rl_engine.reference.minimax_h3.projection.NativeH3*Op` | n/a | `forward`: provider `F.linear`; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | falls back to the PyTorch reference | + +## Reduction Contract `h3-projection-gemm-v1` + +* one program per output tile; the reduction axis is walked in ascending `BLOCK_K` steps into an + FP32 accumulator; no split-K, stream-K, atomics, autotuning or M/N-dependent configuration; +* fixed tile configurations (part of the contract): BF16 `128x128x64`, 8 warps, 3 stages; + FP32 `64x64x32`, 4 warps, 3 stages; +* FP32 rows use `input_precision="ieee"` (no TF32); BF16 rows use the tensor-core dot. The + bias is added once in FP32 and the result rounded once to the output dtype; +* `d_input = g @ W` reduces the output axis ascending; `d_weight = g^T x` and `d_bias` + reduce tokens in logical row order (the one cross-token reduction); results are cast once; +* `h3_qkv_gemm` backward is the fused `[Q; K; V]` backward: `N = 7168` is a multiple of + `BLOCK_K`, so concatenating the three operands moves no block boundary and the separate and + fused layouts have identical gradient bytes. + +Consequences, all tested bitwise: a token's output does not depend on the batch size, the +row position or the other rows; a TP output shard equals the same columns of the full product; +fused equals separate. The accumulation order inside one MMA instruction is the architecture's, +so the contract is bitwise per architecture family (validated on Hopper/H100; Ampere/Ada +use `mma.sync` and Blackwell its own instruction, so bytes may differ across families). + +An earlier candidate, the warp-per-column GEMV of `h3-det-linear-v1` (also exact), was measured +at 50-250x cuBLAS time at thousands of tokens and rejected for these rows; it stays the right +kernel for the handful-of-rows conditioning layers. + +## Accuracy + +Against the FP64 golden (rounded once to the module dtype), measured as max error over max +|golden|: BF16 rows stay within 1.25x of the provider plus a 4e-3 one-ulp floor; FP32 rows +within 1.25x the provider plus 1e-6 (K = 32) or 4e-6 (K = 5376). A sequential FP32 chain over +K = 5376 lands at 2-3e-6 where cuBLAS's blocked sums reach about 5e-7: this is the price of a +fixed single-chain order, and the test names it. Gradients are checked against FP64 `g @ W`, +`g^T x`, `sum g`. Real checkpoint tensors (shard 1: audio, context embedder; block 0: Q/K/V) are +used when `RL_KERNEL_H3_WEIGHTS` points at them. + +## Performance (H100 SXM, torch 2.13, forward, provider = cuBLAS with TF32 off) + +Evidence: `tools/validation/models/h3_projection_bench.py`. + +| Row | T = 64 | T = 1024 | T = 4096 | +| --- | --- | --- | --- | +| `audio_input_projection_fp32` | x2.6 | x2.4 | x2.1 | +| `audio_output_projection_fp32` | x11.9 | x9.4 | x3.5 | +| `text_context_projection` | x2.8 | x2.1 | x1.4 | +| `h3_qkv_gemm` (3 projections) | x1.7 | x1.4 | x1.4 | + +Ratios are candidate time over provider time. `audio_output_projection_fp32` is 0.55 ms flat: with +`N = 32` the grid has only `T / 64` programs and one sequential FP32 chain over 5376, so it +cannot fill the GPU; it runs once per step on the audio tokens, and a per-row tile +configuration is the `cuda_perf_closeout` follow-up. + +## Fail-Closed Contract + +Wrong dtype (no promotion), missing or unexpected bias, wrong `K`, empty or non-2-D input, +widths that are neither full nor a whole-unit shard, mixed devices, CPU tensors for the CUDA +op, mismatched Q/K/V widths and a missing Triton all raise. Nothing falls back to cuBLAS. + +## Tests + +`tests/models/minimax_h3/test_h3_projection_rows.py`: reference/provider/golden agreement; +dtype/bias/shape rejection; forward and backward against FP64 on synthetic and pinned +weights; batch, position, permutation and padding invariance; TP shards; repeat and backward +determinism; fused vs separate QKV forward and backward bitwise; registry dispatch. +`h3_block` now binds Q/K/V to this row (`status="row"`). diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py new file mode 100644 index 000000000..f77af10c6 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 linear projection rows (RFC #420): audio in/out, text context, Q/K/V. + +Every projection runs on one fixed-order Triton GEMM (contract ``h3-projection-gemm-v1``, +``rl_engine/backends/shared/triton/gemm/h3_projection.py``): one program per output tile, +the reduction axis walked in ascending ``BLOCK_K`` steps into an FP32 accumulator, no +split-K / stream-K / atomics / autotuning, one rounding at the store. So + +* a token's output bytes do not depend on how many tokens share the call or where it sits; +* the fused ``[Q; K; V]`` projection equals the three separate ones column for column, + forward and backward; +* a tensor-parallel output shard (whole rows of the weight) equals the same columns of + the full product. + +FP32 rows use ``input_precision="ieee"`` (no TF32); BF16 rows use the tensor-core dot. +``d_input`` reduces the output axis ascending, ``d_weight`` / ``d_bias`` reduce tokens in +logical row order. Nothing falls back to cuBLAS: a missing Triton or an unsupported +dtype/shape raises. (An earlier warp-per-column GEMV kernel, ``h3-det-linear-v1``, is +exact too but 50-250x slower than cuBLAS at thousands of tokens, so it is not used here.) +""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.shared.triton.gemm import h3_projection as gemm +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import projection as ref + +CONTRACT = gemm.CONTRACT +FORWARD_KERNEL_ID = "rl_engine.backends.shared.triton.gemm.h3_projection._projection_kernel" +BACKWARD_IMPL = "h3_projection_gemm_v1_ascending_n_dinput_ascending_t_dweight" +BACKWARD_KERNEL_ID = FORWARD_KERNEL_ID + + +def _require_triton() -> None: + if not gemm.available(): + raise RuntimeError( + "triton + CUDA are required for the H3 projection rows (no cuBLAS fallback is allowed)" + ) + + +class _Projection(torch.autograd.Function): + @staticmethod + def forward(ctx, row, x, weight, bias): + x = x.contiguous() + ctx.row = row + ctx.has_bias = bias is not None + ctx.save_for_backward(x, weight) + return gemm.linear_forward(x, weight, bias) + + @staticmethod + def backward(ctx, grad_out): + x, weight = ctx.saved_tensors + need_x, need_w, need_b = ctx.needs_input_grad[1:] + dx = dw = db = None + if need_x: + dx = gemm.linear_backward_input(grad_out, weight) + if need_w or need_b: + dw, db = gemm.linear_backward_weight(grad_out, x, with_bias=ctx.has_bias and need_b) + dw = dw if need_w else None + record_backward(ctx.row, kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda") + return None, dx, dw, db + + +class _QKV(torch.autograd.Function): + """Q, K and V of one block; backward is the fused ``[Q; K; V]`` backward. + + ``d_input`` reduces over the concatenated ``3 * N`` output rows in ascending order, which + is what the fused projection does. ``N`` is a multiple of ``BLOCK_K`` so concatenating the + three operands does not move any block boundary: separate and fused layouts have the same + backward bytes. + """ + + @staticmethod + def forward(ctx, x, wq, wk, wv): + x = x.contiguous() + ctx.save_for_backward(x, wq, wk, wv) + return tuple(gemm.linear_forward(x, w, None) for w in (wq, wk, wv)) + + @staticmethod + def backward(ctx, gq, gk, gv): + x, wq, wk, wv = ctx.saved_tensors + need_x = ctx.needs_input_grad[0] + weights = (wq, wk, wv) + grads = [ + torch.zeros(x.shape[0], w.shape[0], device=x.device, dtype=x.dtype) if g is None else g + for g, w in zip((gq, gk, gv), weights) + ] + dx = None + if need_x: + dx = gemm.linear_backward_input( + torch.cat(grads, dim=1).contiguous(), torch.cat(weights, dim=0) + ) + dws = [ + gemm.linear_backward_weight(g, x, with_bias=False)[0] if need else None + for g, need in zip(grads, ctx.needs_input_grad[1:]) + ] + record_backward( + "h3_qkv_gemm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return (dx, *dws) + + +class _CudaProjection: + """Shared CUDA operator for one projection row.""" + + op_class = "reduction" + kernel_id = FORWARD_KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + spec: ref.ProjectionSpec + reference_cls: type + + def __init__(self) -> None: + _require_triton() + + def __call__(self, x, weight, bias=None): + return self.forward(x, weight, bias) + + def forward(self, x, weight, bias=None) -> torch.Tensor: + ref.validate_h3_projection(self.spec, x, weight, bias) + if not x.is_cuda: + raise ValueError(f"{self.spec.row}: needs CUDA tensors") + return _Projection.apply(self.spec.row, x, weight, bias) + + def forward_fp32(self, x, weight, bias=None) -> torch.Tensor: + """The CUDA projection already accumulates in FP32 and rounds once.""" + + return self.forward(x, weight, bias) + + +class H3AudioInputProjectionCudaOp(_CudaProjection): + spec = ref.AUDIO_INPUT + + +class H3AudioOutputProjectionCudaOp(_CudaProjection): + spec = ref.AUDIO_OUTPUT + + +class H3TextContextProjectionCudaOp(_CudaProjection): + spec = ref.TEXT_CONTEXT + + +class H3QKVGemmCudaOp: + """Q/K/V projections of a block (BF16, no bias); separate and fused layouts agree bitwise.""" + + spec = ref.QKV + op_class = "reduction" + kernel_id = FORWARD_KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + _require_triton() + + def __call__(self, x, weight_q, weight_k, weight_v): + return self.forward(x, weight_q, weight_k, weight_v) + + def forward(self, x, weight_q, weight_k, weight_v): + widths = { + ref.validate_h3_projection(self.spec, x, w, None) + for w in (weight_q, weight_k, weight_v) + } + if len(widths) != 1: + raise ValueError(f"h3_qkv_gemm: Q/K/V widths differ: {sorted(widths)}") + if not x.is_cuda: + raise ValueError("h3_qkv_gemm: needs CUDA tensors") + return _QKV.apply(x, weight_q, weight_k, weight_v) + + def project(self, x, weight) -> torch.Tensor: + """One of the three projections (what a graph node that owns only Q, K or V runs).""" + + ref.validate_h3_projection(self.spec, x, weight, None) + if not x.is_cuda: + raise ValueError("h3_qkv_gemm: needs CUDA tensors") + return _Projection.apply(self.spec.row, x, weight, None) + + def forward_fused(self, x, weight_qkv) -> torch.Tensor: + """``(T, 3 * N)`` product of the concatenated ``[Q; K; V]`` weight.""" + + n = weight_qkv.shape[0] // 3 if weight_qkv.dim() == 2 and weight_qkv.shape[0] % 3 == 0 else 0 + if n == 0: + raise ValueError("h3_qkv_gemm: fused weight must be (3 * N, K)") + ref.validate_h3_projection(self.spec, x, weight_qkv[:n], None) + if not x.is_cuda: + raise ValueError("h3_qkv_gemm: needs CUDA tensors") + return _Projection.apply(self.spec.row, x, weight_qkv, None) + + def forward_fp32(self, x, weight_q, weight_k, weight_v): + return self.forward(x, weight_q, weight_k, weight_v) diff --git a/rl_engine/backends/shared/triton/gemm/h3_projection.py b/rl_engine/backends/shared/triton/gemm/h3_projection.py new file mode 100644 index 000000000..a1e8a985e --- /dev/null +++ b/rl_engine/backends/shared/triton/gemm/h3_projection.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fixed-order Triton GEMM for the H3 projection rows (contract ``h3-projection-gemm-v1``). + +One program owns one output tile and walks the whole reduction axis in ascending +``BLOCK_K`` steps, accumulating into FP32 registers. There is no split-K, no +stream-K, no atomics, no autotuning and no dependence of the tile shape, the +instruction sequence or the order of the reduction on ``M``, ``N`` or the grid: + +* a token's output bytes are independent of how many tokens share the call and of + its position (``M`` only decides how many programs run, not what each computes); +* a tensor-parallel output shard (a block of whole rows of ``W``) equals the same columns of + the full product, because ``N`` only selects which tiles exist; +* the reduction axis is always ``K`` in ascending order (``T`` ascending for ``d_weight``). + +Arithmetic: BF16 operands use the tensor-core ``tl.dot`` with an FP32 accumulator; +FP32 operands use ``input_precision="ieee"`` (no TF32, no reduced-precision reductions). +The bias is added once in FP32 and the result is rounded once to the output dtype. +The accumulation order inside one MMA instruction is the architecture's, so the contract +is bitwise per architecture family (Hopper wgmma, Ampere/Ada mma.sync), not across them; +tile configurations below are fixed per dtype and therefore part of the contract. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - import guard mirrors the other Triton backends + import triton + import triton.language as tl + + HAS_TRITON = True +except ImportError: # pragma: no cover + HAS_TRITON = False + +CONTRACT = "h3-projection-gemm-v1" +# Fixed tile configurations (part of the contract; changing one changes the bytes). +CONFIG_BF16 = dict(BLOCK_M=128, BLOCK_N=128, BLOCK_K=64, num_warps=8, num_stages=3) +CONFIG_FP32 = dict(BLOCK_M=64, BLOCK_N=64, BLOCK_K=32, num_warps=4, num_stages=3) + +if HAS_TRITON: + + @triton.jit + def _projection_kernel( + a_ptr, + b_ptr, + bias_ptr, + c_ptr, + M, + N, + K, + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + HAS_BIAS: tl.constexpr, + IEEE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_K)): + k_rem = K - k * BLOCK_K + a = tl.load(a_ptrs, mask=(offs_m[:, None] < M) & (offs_k[None, :] < k_rem), other=0.0) + b = tl.load(b_ptrs, mask=(offs_k[:, None] < k_rem) & (offs_n[None, :] < N), other=0.0) + if IEEE: + acc = tl.dot(a, b, acc, input_precision="ieee") + else: + acc = tl.dot(a, b, acc) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + if HAS_BIAS: + bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0).to(tl.float32) + acc = acc + bias[None, :] + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store( + c_ptrs, + acc.to(c_ptr.dtype.element_ty), + mask=(offs_m[:, None] < M) & (offs_n[None, :] < N), + ) + + +def available() -> bool: + return bool(HAS_TRITON and torch.cuda.is_available()) + + +def _require(*tensors: torch.Tensor) -> torch.dtype: + if not HAS_TRITON: + raise RuntimeError("triton is required for the H3 projection rows (no cuBLAS fallback)") + dtype = tensors[0].dtype + if dtype not in (torch.float32, torch.bfloat16): + raise TypeError(f"h3 projection GEMM supports float32 and bfloat16, got {dtype}") + for t in tensors: + if not t.is_cuda or t.dtype != dtype or t.device != tensors[0].device: + raise ValueError("h3 projection GEMM needs same-dtype tensors on one CUDA device") + return dtype + + +def gemm( + a: torch.Tensor, + b: torch.Tensor, + *, + bias: torch.Tensor | None = None, + out_dtype: torch.dtype | None = None, + a_t: bool = False, + b_t: bool = False, +) -> torch.Tensor: + """``C = A @ B (+ bias)`` with reduction in ascending index order. + + ``a`` is ``(M, K)`` (``a_t``: given as ``(K, M)``), ``b`` is ``(K, N)`` (``b_t``: given as + ``(N, K)``). Operands may be strided; nothing is copied. FP32 accumulate, one rounding. + """ + + dtype = _require(a, b) if bias is None else _require(a, b, bias) + config = CONFIG_BF16 if dtype == torch.bfloat16 else CONFIG_FP32 + if a_t: + k, m = a.shape + sam, sak = a.stride(1), a.stride(0) + else: + m, k = a.shape + sam, sak = a.stride(0), a.stride(1) + if b_t: + n, kb = b.shape + sbk, sbn = b.stride(1), b.stride(0) + else: + kb, n = b.shape + sbk, sbn = b.stride(0), b.stride(1) + if k != kb: + raise ValueError(f"reduction axes differ: {k} vs {kb}") + if m == 0 or n == 0 or k == 0: + raise ValueError("h3 projection GEMM needs non-empty operands") + c = torch.empty((m, n), device=a.device, dtype=out_dtype or dtype) + grid = (triton.cdiv(m, config["BLOCK_M"]), triton.cdiv(n, config["BLOCK_N"])) + _projection_kernel[grid]( + a, + b, + bias if bias is not None else c, + c, + m, + n, + k, + sam, + sak, + sbk, + sbn, + c.stride(0), + c.stride(1), + HAS_BIAS=bias is not None, + IEEE=dtype == torch.float32, + BLOCK_M=config["BLOCK_M"], + BLOCK_N=config["BLOCK_N"], + BLOCK_K=config["BLOCK_K"], + num_warps=config["num_warps"], + num_stages=config["num_stages"], + ) + return c + + +def linear_forward(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: + """``y = x @ W^T (+ bias)`` for ``x (T, K)`` and ``W (N, K)``.""" + + return gemm(x, weight, bias=bias, b_t=True) + + +def linear_backward_input(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``dx = g @ W`` for ``g (T, N)`` and ``W (N, K)``; reduction over ``N`` ascending.""" + + return gemm(grad, weight) + + +def linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, *, with_bias: bool +) -> tuple[torch.Tensor, torch.Tensor | None]: + """``dW = g^T @ x`` and ``db = sum_t g`` with the token axis reduced in ascending order.""" + + dweight = gemm(grad, x, a_t=True) + if not with_bias: + return dweight, None + ones = torch.ones(grad.shape[0], 1, device=grad.device, dtype=grad.dtype) + dbias = gemm(grad, ones, a_t=True).reshape(-1) + return dweight, dbias diff --git a/rl_engine/reference/minimax_h3/projection.py b/rl_engine/reference/minimax_h3/projection.py new file mode 100644 index 000000000..53713d9d0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/projection.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 linear projection rows (RFC #420). + +Four rows share one contract, each pinned to its checkpoint tensors: + +=============================== ============ ============ ===== ================ +row in -> out dtype bias checkpoint tensors +=============================== ============ ============ ===== ================ +``audio_input_projection_fp32`` 32 -> 5376 float32 yes ``audio_proj_in`` +``audio_output_projection_fp32`` 5376 -> 32 float32 yes ``audio_proj_out`` +``text_context_projection`` 5120 -> 5376 bfloat16 yes ``context_embedder`` +``h3_qkv_gemm`` (per projection) 5376 -> 7168 bfloat16 no ``attn.to_{q,k,v}`` +=============================== ============ ============ ===== ================ + +``audio_*`` modules are in ``_keep_in_fp32_modules``: weights, bias, input and output +are FP32 and a lower-precision argument is a contract violation (RFC probe H7), not a +variant. The text projection and Q/K/V are BF16 modules (FP32 accumulate, one rounding). + +``h3_qkv_gemm`` also covers the fused layout used by serving engines (vLLM-Omni's +``QKVParallelLinear``): the rows of ``to_q``, ``to_k`` and ``to_v`` concatenated along +the output axis. ``forward_fused`` must equal the separate projections column for column. + +Output width may be a shard of the full width so tensor-parallel layouts reuse the +rows: Q/K/V shards are whole heads (a multiple of 128 output columns). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +HEAD_DIM = 128 + + +@dataclass(frozen=True) +class ProjectionSpec: + row: str + in_features: int + out_features: int # full width; shards must be whole multiples of ``shard_unit`` + dtype: torch.dtype + has_bias: bool + shard_unit: int # smallest legal output width (columns) + + +AUDIO_INPUT = ProjectionSpec("audio_input_projection_fp32", 32, 5376, torch.float32, True, 64) +AUDIO_OUTPUT = ProjectionSpec("audio_output_projection_fp32", 5376, 32, torch.float32, True, 32) +TEXT_CONTEXT = ProjectionSpec("text_context_projection", 5120, 5376, torch.bfloat16, True, 64) +QKV = ProjectionSpec("h3_qkv_gemm", 5376, 7168, torch.bfloat16, False, HEAD_DIM) +SPECS = {spec.row: spec for spec in (AUDIO_INPUT, AUDIO_OUTPUT, TEXT_CONTEXT, QKV)} + + +def validate_h3_projection( + spec: ProjectionSpec, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + allow_shard: bool = True, +) -> int: + """Check ``x (T, K)``, ``weight (N, K)``, ``bias (N,)`` against ``spec``; return ``N``. + + Dtype is exact (no promotion), tensors share a device, ``T > 0``, ``K`` is the + pinned input width, and ``N`` is the full width or (``allow_shard``) a whole-unit + shard of it. A bias is required when the module has one and forbidden when it has none. + """ + + for name, tensor in (("x", x), ("weight", weight), ("bias", bias)): + if tensor is None: + continue + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != spec.dtype: + raise TypeError( + f"{spec.row}: {name} must be {spec.dtype}, got {tensor.dtype} " + "(the module's declared precision is part of the contract)" + ) + if tensor.device != x.device: + raise ValueError(f"{spec.row}: {name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"{spec.row}: x must be a non-empty (T, {spec.in_features}) matrix") + if x.shape[1] != spec.in_features: + raise ValueError(f"{spec.row}: x width {x.shape[1]} != {spec.in_features}") + if weight.dim() != 2 or weight.shape[1] != spec.in_features: + raise ValueError(f"{spec.row}: weight must be (N, {spec.in_features})") + n = weight.shape[0] + legal = n == spec.out_features or ( + allow_shard and 0 < n < spec.out_features and n % spec.shard_unit == 0 + ) + if not legal: + raise ValueError( + f"{spec.row}: output width {n} must be {spec.out_features}" + + (f" or a multiple of {spec.shard_unit} below it" if allow_shard else "") + ) + if spec.has_bias: + if bias is None: + raise ValueError(f"{spec.row}: this module has a bias") + if bias.shape != (n,): + raise ValueError(f"{spec.row}: bias must be ({n},), got {tuple(bias.shape)}") + elif bias is not None: + raise ValueError(f"{spec.row}: this module has no bias") + return n + + +class _NativeProjection: + """PyTorch reference: ``forward`` is the provider (``F.linear``), ``forward_fp32`` the golden.""" + + spec: ProjectionSpec + op_class = "reduction" + + def __call__(self, x, weight, bias=None): + return self.forward(x, weight, bias) + + def forward(self, x, weight, bias=None) -> torch.Tensor: + """Provider path (``nn.Linear``); run with TF32 disabled on CUDA.""" + + validate_h3_projection(self.spec, x, weight, bias) + return F.linear(x, weight, bias) + + def forward_fp32(self, x, weight, bias=None) -> torch.Tensor: + """Independent golden: the product in FP64, rounded once to the module dtype.""" + + validate_h3_projection(self.spec, x, weight, bias) + out = F.linear(x.double(), weight.double(), None if bias is None else bias.double()) + return out.to(self.spec.dtype) + + +class NativeH3AudioInputProjectionOp(_NativeProjection): + spec = AUDIO_INPUT + + +class NativeH3AudioOutputProjectionOp(_NativeProjection): + spec = AUDIO_OUTPUT + + +class NativeH3TextContextProjectionOp(_NativeProjection): + spec = TEXT_CONTEXT + + +class NativeH3QKVGemmOp: + """Q, K and V projections of one attention block (no bias), separate or fused.""" + + spec = QKV + op_class = "reduction" + + def __call__(self, x, weight_q, weight_k, weight_v): + return self.forward(x, weight_q, weight_k, weight_v) + + def forward(self, x, weight_q, weight_k, weight_v): + """``(q, k, v)`` from three separate ``F.linear`` calls (the Diffusers layout).""" + + widths = { + validate_h3_projection(QKV, x, w, None) for w in (weight_q, weight_k, weight_v) + } + if len(widths) != 1: + raise ValueError(f"h3_qkv_gemm: Q/K/V widths differ: {sorted(widths)}") + return F.linear(x, weight_q), F.linear(x, weight_k), F.linear(x, weight_v) + + def forward_fused(self, x, weight_qkv): + """One ``(T, 3 * N)`` product of the concatenated ``[Q; K; V]`` weight (serving layout).""" + + n = _fused_width(weight_qkv) + validate_h3_projection(QKV, x, weight_qkv[:n], None) + return F.linear(x, weight_qkv) + + def forward_fp32(self, x, weight_q, weight_k, weight_v): + """FP64 golden of each projection, rounded once to BF16.""" + + outputs = [] + for w in (weight_q, weight_k, weight_v): + validate_h3_projection(QKV, x, w, None) + outputs.append(F.linear(x.double(), w.double()).to(torch.bfloat16)) + return tuple(outputs) + + +def _fused_width(weight_qkv: torch.Tensor) -> int: + if weight_qkv.dim() != 2 or weight_qkv.shape[0] % 3: + raise ValueError("h3_qkv_gemm: fused weight must be (3 * N, K)") + return weight_qkv.shape[0] // 3 + + +def fuse_qkv_weights(weight_q, weight_k, weight_v) -> torch.Tensor: + """Concatenate Q/K/V rows into the fused serving layout ``[Q; K; V]``.""" + + return torch.cat([weight_q, weight_k, weight_v], dim=0) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index a733e95ce..cf79bde32 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -216,6 +216,22 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_MLP = ( "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" ) + PYTORCH_H3_AUDIO_INPUT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3AudioInputProjectionOp" + CUDA_H3_AUDIO_INPUT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioInputProjectionCudaOp" + ) + PYTORCH_H3_AUDIO_OUTPUT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3AudioOutputProjectionOp" + CUDA_H3_AUDIO_OUTPUT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioOutputProjectionCudaOp" + ) + PYTORCH_H3_TEXT_CONTEXT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3TextContextProjectionOp" + CUDA_H3_TEXT_CONTEXT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3TextContextProjectionCudaOp" + ) + PYTORCH_H3_QKV_GEMM = "rl_engine.reference.minimax_h3.projection.NativeH3QKVGemmOp" + CUDA_H3_QKV_GEMM = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3QKVGemmCudaOp" + ) PYTORCH_H3_ADALN_PROJECTION = ( "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" ) @@ -719,6 +735,22 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_MLP, OpBackend.PYTORCH_H3_TIMESTEP_MLP, ], + "audio_input_projection_fp32": [ + OpBackend.CUDA_H3_AUDIO_INPUT_PROJECTION, + OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION, + ], + "audio_output_projection_fp32": [ + OpBackend.CUDA_H3_AUDIO_OUTPUT_PROJECTION, + OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION, + ], + "text_context_projection": [ + OpBackend.CUDA_H3_TEXT_CONTEXT_PROJECTION, + OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION, + ], + "h3_qkv_gemm": [ + OpBackend.CUDA_H3_QKV_GEMM, + OpBackend.PYTORCH_H3_QKV_GEMM, + ], "adaln_projection_3mod": [ OpBackend.CUDA_H3_ADALN_PROJECTION, OpBackend.PYTORCH_H3_ADALN_PROJECTION, @@ -779,6 +811,10 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], @@ -835,6 +871,10 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], @@ -868,6 +908,10 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py index b20e885a1..1bdeaba46 100644 --- a/rl_engine/validation/models/h3_block.py +++ b/rl_engine/validation/models/h3_block.py @@ -93,6 +93,7 @@ def __init__(self, registry: KernelRegistry | None = None) -> None: self.norm = registry.get_op("h3_rmsnorm", device="cuda") self.gate = registry.get_op("adaln_gate_residual", device="cuda") self.gemm = TritonDetGemmOp() + self.qkv = registry.get_op("h3_qkv_gemm", device="cuda") self.attention = DeterministicAttentionOp() self.swiglu = SwiGLUCudaOp() @@ -169,6 +170,24 @@ def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> ) +def _qkv_node(name: str, weight: str) -> Node: + """Q, K or V projection: the ``h3_qkv_gemm`` row (fixed-order Triton GEMM, BF16).""" + + def candidate(ops, ctx, x): + flat = x.reshape(-1, x.shape[-1]) + return ops.qkv.project(flat, _w(ctx, weight)).view(*x.shape[:-1], -1) + + return Node( + name=name, + inputs=("norm1",), + binding=Binding("h3_qkv_gemm", "row", "H3QKVGemmCudaOp.project"), + candidate=candidate, + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # fixed-order K loop vs cuBLAS + ) + + def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: return Node( name=name, @@ -240,9 +259,9 @@ def _adaln_params(ctx): provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS ), _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), - _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), - _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), - _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _qkv_node("q_proj", "attn.to_q.weight"), + _qkv_node("k_proj", "attn.to_k.weight"), + _qkv_node("v_proj", "attn.to_v.weight"), _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), _rope_node("rope_q", "q_norm"), diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py index 8f64a4508..b55e4ffcd 100644 --- a/tests/models/minimax_h3/test_h3_one_block.py +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -72,7 +72,12 @@ def test_graph_is_well_formed(): seen.add(node.name) assert h3_block.OUTPUT == "residual_mlp" rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} - assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + assert rows == { + "adaln_projection_3mod", + "h3_rmsnorm", + "adaln_gate_residual", + "h3_qkv_gemm", + } def test_partial_rope_rotates_96_of_128_channels(): diff --git a/tests/models/minimax_h3/test_h3_projection_rows.py b/tests/models/minimax_h3/test_h3_projection_rows.py new file mode 100644 index 000000000..d2134ab72 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_projection_rows.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 projection rows: ``audio_input_projection_fp32``, ``audio_output_projection_fp32``, +``text_context_projection`` and ``h3_qkv_gemm``. + +* accuracy: forward and every gradient against an FP64 golden, on synthetic data and on + the pinned checkpoint tensors (shard 1), with the provider (``F.linear``) as yardstick; +* invariance: a token's bytes do not depend on the batch size, the row position or the other + rows; a tensor-parallel output shard equals the same columns of the full product; +* QKV layouts: the fused ``[Q; K; V]`` product equals the separate projections column for + column, and their backwards agree bitwise; +* determinism and fail-closed contract (dtype, bias, widths, devices, missing extension). +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import projection as ref + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +SHARD1 = "diffusion_pytorch_model-00001-of-00014.safetensors" + +ROWS = { + "audio_input_projection_fp32": (ref.AUDIO_INPUT, "audio_proj_in"), + "audio_output_projection_fp32": (ref.AUDIO_OUTPUT, "audio_proj_out"), + "text_context_projection": (ref.TEXT_CONTEXT, "context_embedder"), +} +REFERENCE = { + "audio_input_projection_fp32": ref.NativeH3AudioInputProjectionOp, + "audio_output_projection_fp32": ref.NativeH3AudioOutputProjectionOp, + "text_context_projection": ref.NativeH3TextContextProjectionOp, +} + + +def _cuda_ops(): + from rl_engine.backends.cuda.model_specific.minimax_h3 import projection as cu + if not cu.gemm.available(): + pytest.skip("needs triton and a CUDA device") + return cu + + +def _cuda_op(row): + cu = _cuda_ops() + return { + "audio_input_projection_fp32": cu.H3AudioInputProjectionCudaOp, + "audio_output_projection_fp32": cu.H3AudioOutputProjectionCudaOp, + "text_context_projection": cu.H3TextContextProjectionCudaOp, + }[row]() + + +def _params(spec, n=None, seed=0, device="cuda"): + n = n or spec.out_features + g = torch.Generator(device="cpu").manual_seed(seed) + w = (torch.randn(n, spec.in_features, generator=g) / spec.in_features**0.5).to(spec.dtype) + b = (torch.randn(n, generator=g) * 0.1).to(spec.dtype) if spec.has_bias else None + return w.to(device), None if b is None else b.to(device) + + +def _x(spec, t, seed=1, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(t, spec.in_features, generator=g).to(spec.dtype).to(device) + + +def _real(names): + root = os.environ.get("RL_KERNEL_H3_WEIGHTS") + path = Path(root) / "hf" / "transformer" / SHARD1 if root else None + if path is None or not path.exists(): + pytest.skip("RL_KERNEL_H3_WEIGHTS/hf/transformer shard 1 not available") + from safetensors import safe_open + + with safe_open(str(path), framework="pt", device="cpu") as f: + return {n: f.get_tensor(n).cuda() for n in names} + + +def _floor(spec): + """Absolute slack (relative to max|golden|) on top of 1.25x the provider's error. + + A sequential FP32 chain over K = 5376 lands near 2-3e-6 (cuBLAS's blocked sums reach + ~5e-7); BF16 outputs round once at 2**-9 so 4e-3 covers a one-ulp flip at the maximum. + """ + + if spec.dtype == torch.float32: + return 4e-6 if spec.in_features > 1024 else 1e-6 + return 4e-3 + + +def _bits(t): + return t.contiguous().view(torch.uint8).cpu().numpy().tobytes() + + +def _golden_grads(x, w, grad): + xd, wd, gd = x.double(), w.double(), grad.double() + return gd @ wd, gd.t() @ xd, gd.sum(0) + + +def _rel(value, golden): + g = golden.double() + return float((value.double() - g).abs().max() / g.abs().max().clamp_min(1e-30)) + + +# --- CPU: reference and contract ------------------------------------------------------------ + + +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_reference_provider_and_golden_agree(row): + spec, _ = ROWS[row] + op = REFERENCE[row]() + w, b = _params(spec, device="cpu") + x = _x(spec, 9, device="cpu") + provider, golden = op.forward(x, w, b), op.forward_fp32(x, w, b) + assert provider.dtype == golden.dtype == spec.dtype + assert _rel(provider, golden) < (1e-5 if spec.dtype == torch.float32 else 2e-2) + + +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_reference_rejects_wrong_dtype_bias_and_shapes(row): + spec, _ = ROWS[row] + op = REFERENCE[row]() + w, b = _params(spec, device="cpu") + x = _x(spec, 4, device="cpu") + other = torch.bfloat16 if spec.dtype == torch.float32 else torch.float32 + for bad in ( + (x.to(other), w, b), + (x, w.to(other), b), + (x, w, b.to(other)), + (x, w, None), # these modules have a bias + (x[:0], w, b), + (x[:, :-1], w, b), + (x.flatten(), w, b), + (x, w[:-1], b[:-1]), # not a legal width + (x, w, b[:-1]), + ): + with pytest.raises((TypeError, ValueError)): + op.forward(*bad) + + +def test_qkv_reference_contract(): + spec = ref.QKV + op = ref.NativeH3QKVGemmOp() + ws = [_params(spec, seed=s, device="cpu")[0] for s in (1, 2, 3)] + x = _x(spec, 5, device="cpu") + q, k, v = op.forward(x, *ws) + assert all(t.shape == (5, 7168) and t.dtype == torch.bfloat16 for t in (q, k, v)) + fused = op.forward_fused(x, ref.fuse_qkv_weights(*ws)) + assert fused.shape == (5, 3 * 7168) + assert torch.equal(fused[:, :7168], q) or _rel(fused[:, :7168], q) < 1e-2 # CPU BLAS order + with pytest.raises(ValueError): + op.forward(x, ws[0], ws[1], ws[2][:-128]) # widths differ + with pytest.raises(ValueError): + op.forward(x, *[w[:100] for w in ws]) # not whole heads + with pytest.raises(ValueError): + op.forward_fused(x, ws[0][:-1]) + shard = op.forward(x, *[w[:1792] for w in ws]) # TP4 head shard + assert shard[0].shape == (5, 1792) + + +def test_registry_dispatches_reference_on_cpu(): + from rl_engine.runtime.registry import KernelRegistry + + registry = KernelRegistry() + for row, cls in REFERENCE.items(): + assert type(registry.get_op(row, device="cpu")) is cls + assert type(registry.get_op("h3_qkv_gemm", device="cpu")) is ref.NativeH3QKVGemmOp + + +# --- CUDA: accuracy ----------------------------------------------------------------------------- + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +@pytest.mark.parametrize("tokens", [1, 7, 8, 9, 257, 1025]) +def test_forward_against_fp64_golden(row, tokens): + spec, _ = ROWS[row] + op, native = _cuda_op(row), REFERENCE[row]() + w, b = _params(spec) + x = _x(spec, tokens) + out = op.forward(x, w, b) + golden = native.forward_fp32(x, w, b) + assert out.dtype == spec.dtype and out.shape == golden.shape + with torch.backends.cuda.matmul.flags(allow_tf32=False) if hasattr( + torch.backends.cuda.matmul, "flags" + ) else _noop(): + provider = native.forward(x, w, b) + floor = _floor(spec) + assert _rel(out, golden) <= 1.25 * _rel(provider, golden) + floor + + +class _noop: + def __enter__(self): + return None + + def __exit__(self, *a): + return False + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_forward_on_pinned_checkpoint_weights(row): + spec, module = ROWS[row] + t = _real([f"{module}.weight", f"{module}.bias"]) + w, b = t[f"{module}.weight"], t[f"{module}.bias"] + assert (w.dtype, tuple(w.shape)) == (spec.dtype, (spec.out_features, spec.in_features)) + x = _x(spec, 300, seed=5) + out = _cuda_op(row).forward(x, w, b) + golden = REFERENCE[row]().forward_fp32(x, w, b) + floor = _floor(spec) + provider = F.linear(x, w, b) + assert _rel(out, golden) <= 1.25 * _rel(provider, golden) + floor + + +@requires_cuda +def test_qkv_on_pinned_checkpoint_weights(): + names = [f"transformer_blocks.0.attn.to_{c}.weight" for c in "qkv"] + t = _real(names) + x = _x(ref.QKV, 130, seed=2) + cu = _cuda_ops() + outs = cu.H3QKVGemmCudaOp().forward(x, *[t[n] for n in names]) + golden = ref.NativeH3QKVGemmOp().forward_fp32(x, *[t[n] for n in names]) + for out, gold, n in zip(outs, golden, names): + assert _rel(out, gold) <= 1.25 * _rel(F.linear(x, t[n]), gold) + 4e-3 + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_backward_against_fp64_golden(row): + spec, _ = ROWS[row] + op = _cuda_op(row) + w, b = _params(spec) + w.requires_grad_(), b.requires_grad_() + x = _x(spec, 129).requires_grad_() + out = op.forward(x, w, b) + grad = torch.randn(out.shape, generator=torch.Generator("cpu").manual_seed(3)).to( + out.dtype + ).cuda() + out.backward(grad) + dx, dw, db = _golden_grads(x.detach(), w.detach(), grad) + floor = 1e-5 if spec.dtype == torch.float32 else 2e-2 + assert _rel(x.grad, dx) < floor + assert _rel(w.grad, dw) < floor + assert _rel(b.grad, db) < floor + assert x.grad.dtype == w.grad.dtype == b.grad.dtype == spec.dtype + + +# --- CUDA: invariance and determinism ---------------------------------------------------------- + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_rows_are_batch_and_position_invariant(row): + spec, _ = ROWS[row] + op = _cuda_op(row) + w, b = _params(spec) + x = _x(spec, 67) + full = op.forward(x, w, b) + for start, stop in ((0, 1), (3, 4), (60, 67), (1, 64), (8, 17)): + assert _bits(op.forward(x[start:stop].contiguous(), w, b)) == _bits(full[start:stop]) + perm = torch.randperm(67, generator=torch.Generator("cpu").manual_seed(0)).cuda() + assert _bits(op.forward(x[perm].contiguous(), w, b)) == _bits(full[perm]) + padded = torch.cat([x, _x(spec, 33, seed=9)]) + assert _bits(op.forward(padded, w, b)[:67]) == _bits(full) + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_output_shard_equals_columns_of_full_product(row): + spec, _ = ROWS[row] + if spec.out_features % (2 * spec.shard_unit): + pytest.skip("not splittable into two legal shards") + op = _cuda_op(row) + w, b = _params(spec) + x = _x(spec, 40) + full = op.forward(x, w, b) + half = spec.out_features // 2 + assert _bits(op.forward(x, w[:half].contiguous(), b[:half].contiguous())) == _bits( + full[:, :half] + ) + assert _bits(op.forward(x, w[half:].contiguous(), b[half:].contiguous())) == _bits( + full[:, half:] + ) + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_repeat_and_backward_are_deterministic(row): + spec, _ = ROWS[row] + op = _cuda_op(row) + w, b = _params(spec) + + def run(): + x = _x(spec, 70).requires_grad_() + wp, bp = w.clone().requires_grad_(), b.clone().requires_grad_() + out = op.forward(x, wp, bp) + out.backward(torch.ones_like(out)) + return [_bits(t) for t in (out, x.grad, wp.grad, bp.grad)] + + assert run() == run() + + +@requires_cuda +def test_dx_rows_are_batch_invariant(): + spec = ref.TEXT_CONTEXT + op = _cuda_op("text_context_projection") + w, b = _params(spec) + g = torch.randn(40, spec.out_features, generator=torch.Generator("cpu").manual_seed(7)).to( + spec.dtype + ).cuda() + + def dx(rows): + x = _x(spec, 40)[rows].requires_grad_() + out = op.forward(x, w, b) + out.backward(g[rows]) + return x.grad + + full = dx(slice(0, 40)) + assert _bits(dx(slice(5, 9))) == _bits(full[5:9]) + + +# --- CUDA: h3_qkv_gemm layouts ----------------------------------------------------------------- + + +@requires_cuda +@pytest.mark.parametrize("tokens", [1, 9, 130]) +def test_fused_qkv_equals_separate_columns_bitwise(tokens): + cu = _cuda_ops() + op = cu.H3QKVGemmCudaOp() + ws = [_params(ref.QKV, seed=s)[0] for s in (1, 2, 3)] + x = _x(ref.QKV, tokens) + q, k, v = op.forward(x, *ws) + fused = op.forward_fused(x, ref.fuse_qkv_weights(*ws)) + n = ref.QKV.out_features + assert _bits(fused[:, :n]) == _bits(q) + assert _bits(fused[:, n : 2 * n]) == _bits(k) + assert _bits(fused[:, 2 * n :]) == _bits(v) + + +@requires_cuda +def test_fused_and_separate_qkv_backward_are_bitwise_equal(): + cu = _cuda_ops() + op = cu.H3QKVGemmCudaOp() + n = ref.QKV.out_features + base = [_params(ref.QKV, seed=s)[0] for s in (1, 2, 3)] + gq, gk, gv = ( + torch.randn(66, n, generator=torch.Generator("cpu").manual_seed(s)).bfloat16().cuda() + for s in (4, 5, 6) + ) + + ws = [w.clone().requires_grad_() for w in base] + x1 = _x(ref.QKV, 66).requires_grad_() + q, k, v = op.forward(x1, *ws) + torch.autograd.backward([q, k, v], [gq, gk, gv]) + + wf = ref.fuse_qkv_weights(*base).clone().requires_grad_() + x2 = _x(ref.QKV, 66).requires_grad_() + fused = op.forward_fused(x2, wf) + fused.backward(torch.cat([gq, gk, gv], dim=1)) + + assert _bits(x1.grad) == _bits(x2.grad) + assert _bits(torch.cat([w.grad for w in ws])) == _bits(wf.grad) + + +@requires_cuda +def test_qkv_tp_head_shards_equal_full_columns_and_backward_matches_golden(): + cu = _cuda_ops() + op = cu.H3QKVGemmCudaOp() + ws = [_params(ref.QKV, seed=s)[0] for s in (1, 2, 3)] + x = _x(ref.QKV, 33) + full = op.forward(x, *ws) + for rank in range(4): # TP4: 14 heads per rank + sl = slice(rank * 1792, (rank + 1) * 1792) + shard = op.forward(x, *[w[sl].contiguous() for w in ws]) + for s, f in zip(shard, full): + assert _bits(s) == _bits(f[:, sl]) + # a missing output gradient is zero, not an error + wq = ws[0].clone().requires_grad_() + xg = x.clone().requires_grad_() + q, k, v = op.forward(xg, wq, ws[1], ws[2]) + q.sum().backward() + gold_dx = (torch.ones(33, 7168, device="cuda", dtype=torch.float64) @ ws[0].double()) + assert _rel(xg.grad, gold_dx) < 2e-2 + assert wq.grad.shape == ws[0].shape + + +# --- CUDA: fail closed ---------------------------------------------------------------------------- + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_cuda_rejects_cpu_mixed_devices_and_wrong_dtype(row): + spec, _ = ROWS[row] + op = _cuda_op(row) + w, b = _params(spec) + x = _x(spec, 4) + with pytest.raises(ValueError): + op.forward(x.cpu(), w.cpu(), b.cpu()) + with pytest.raises(ValueError): + op.forward(x, w.cpu(), b) + other = torch.bfloat16 if spec.dtype == torch.float32 else torch.float32 + with pytest.raises(TypeError): + op.forward(x.to(other), w, b) + with pytest.raises(ValueError): + op.forward(x, w, None) + + +@requires_cuda +def test_registry_dispatches_cuda(): + from rl_engine.backends.cuda.model_specific.minimax_h3 import projection as cu + from rl_engine.runtime.registry import KernelRegistry + + _cuda_ops() + registry = KernelRegistry() + assert type(registry.get_op("audio_input_projection_fp32", device="cuda")) is ( + cu.H3AudioInputProjectionCudaOp + ) + assert type(registry.get_op("h3_qkv_gemm", device="cuda")) is cu.H3QKVGemmCudaOp diff --git a/tools/validation/models/h3_projection_bench.py b/tools/validation/models/h3_projection_bench.py new file mode 100644 index 000000000..6736ed2f2 --- /dev/null +++ b/tools/validation/models/h3_projection_bench.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Time the H3 projection rows against the provider (cuBLAS ``F.linear``) and record backends. + + python tools/validation/models/h3_projection_bench.py --out reports/h3_projection_bench.json + +Per row and token count: median forward and forward+backward milliseconds of the +RL-Kernel operator and of the provider, their ratio, and the operator actually dispatched +(registry readback). Provider timings run with TF32 off. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.reference.minimax_h3 import projection as ref # noqa: E402 +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 + +ROWS = ( + ("audio_input_projection_fp32", ref.AUDIO_INPUT), + ("audio_output_projection_fp32", ref.AUDIO_OUTPUT), + ("text_context_projection", ref.TEXT_CONTEXT), +) + + +def _time(fn, repeats): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return statistics.median(samples) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="*", default=[64, 1024, 4096]) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + torch.backends.cuda.matmul.allow_tf32 = False + registry = KernelRegistry() + report = {"device": torch.cuda.get_device_name(0), "torch": torch.__version__, "rows": []} + for row, spec in [*ROWS, ("h3_qkv_gemm", ref.QKV)]: + op = registry.get_op(row, device="cuda") + for tokens in args.tokens: + gen = torch.Generator("cpu").manual_seed(0) + x = torch.randn(tokens, spec.in_features, generator=gen).to(spec.dtype).cuda() + n = spec.out_features + ws = [ + (torch.randn(n, spec.in_features, generator=gen) * 0.02).to(spec.dtype).cuda() + for _ in range(3 if row == "h3_qkv_gemm" else 1) + ] + b = torch.zeros(n, dtype=spec.dtype).cuda() if spec.has_bias else None + if row == "h3_qkv_gemm": + cand = lambda: op.forward(x, *ws) # noqa: E731 + prov = lambda: [F.linear(x, w) for w in ws] # noqa: E731 + else: + cand = lambda: op.forward(x, ws[0], b) # noqa: E731 + prov = lambda: F.linear(x, ws[0], b) # noqa: E731 + c, p = _time(cand, args.repeats), _time(prov, args.repeats) + report["rows"].append( + { + "row": row, + "tokens": tokens, + "backend": f"{type(op).__module__}.{type(op).__name__}", + "kernel_id": getattr(op, "kernel_id", ""), + "candidate_ms": c, + "provider_ms": p, + "ratio": c / p, + } + ) + print(f"{row:32s} T={tokens:5d} candidate {c:8.3f} ms provider {p:8.3f} ms x{c / p:.2f}") + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1) + "\n") + + +if __name__ == "__main__": + main() From a6533fa98df4599602ebd9f3bfb9d99d5a0f9847 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:58:00 +0200 Subject: [PATCH 31/62] feat(h3): projection rows Hopper evidence, charts and lint Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit b41f868b0220a6aa0aecb489dd76bd3544b0cd25) --- docs/operators/h3-projection-rows.md | 9 +++++ .../model_specific/minimax_h3/projection.py | 6 ++- .../shared/triton/gemm/h3_projection.py | 4 +- rl_engine/reference/minimax_h3/projection.py | 6 +-- rl_engine/runtime/registry.py | 12 ++++-- .../minimax_h3/test_h3_projection_rows.py | 25 +++++++----- .../validation/models/h3_projection_bench.py | 14 ++++++- .../validation/models/h3_projection_charts.py | 38 +++++++++++++++++++ 8 files changed, 94 insertions(+), 20 deletions(-) create mode 100644 tools/validation/models/h3_projection_charts.py diff --git a/docs/operators/h3-projection-rows.md b/docs/operators/h3-projection-rows.md index a64c5b349..a56eea810 100644 --- a/docs/operators/h3-projection-rows.md +++ b/docs/operators/h3-projection-rows.md @@ -102,3 +102,12 @@ dtype/bias/shape rejection; forward and backward against FP64 on synthetic and p weights; batch, position, permutation and padding invariance; TP shards; repeat and backward determinism; fused vs separate QKV forward and backward bitwise; registry dispatch. `h3_block` now binds Q/K/V to this row (`status="row"`). + +## Hopper evidence (H100) + +These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description). +`reports/experiments/h3-projection-rows-h100/bench.json` (from `tools/validation/models/h3_projection_bench.py`) and the chart below record +forward time and relative error vs an FP64 `F.linear`, per row and token count, for the RL-Kernel operator and the cuBLAS provider. +The fixed-order kernels are 1.3x to 4.3x slower than cuBLAS on the BF16 rows; `audio_output_projection_fp32` (FP32, K=5376, ieee +precision, 64x64x32 tile) is the slowest (3.5x to 12x) and is the first candidate for a per-row tile sweep. Strict mode forbids +split-K, so this is the cost of the fixed reduction order. diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py index f77af10c6..ac9ae4c56 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py @@ -181,7 +181,11 @@ def project(self, x, weight) -> torch.Tensor: def forward_fused(self, x, weight_qkv) -> torch.Tensor: """``(T, 3 * N)`` product of the concatenated ``[Q; K; V]`` weight.""" - n = weight_qkv.shape[0] // 3 if weight_qkv.dim() == 2 and weight_qkv.shape[0] % 3 == 0 else 0 + n = ( + weight_qkv.shape[0] // 3 + if weight_qkv.dim() == 2 and weight_qkv.shape[0] % 3 == 0 + else 0 + ) if n == 0: raise ValueError("h3_qkv_gemm: fused weight must be (3 * N, K)") ref.validate_h3_projection(self.spec, x, weight_qkv[:n], None) diff --git a/rl_engine/backends/shared/triton/gemm/h3_projection.py b/rl_engine/backends/shared/triton/gemm/h3_projection.py index a1e8a985e..4b1f548b5 100644 --- a/rl_engine/backends/shared/triton/gemm/h3_projection.py +++ b/rl_engine/backends/shared/triton/gemm/h3_projection.py @@ -167,7 +167,9 @@ def gemm( return c -def linear_forward(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None) -> torch.Tensor: +def linear_forward( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None +) -> torch.Tensor: """``y = x @ W^T (+ bias)`` for ``x (T, K)`` and ``W (N, K)``.""" return gemm(x, weight, bias=bias, b_t=True) diff --git a/rl_engine/reference/minimax_h3/projection.py b/rl_engine/reference/minimax_h3/projection.py index 53713d9d0..1ab933156 100644 --- a/rl_engine/reference/minimax_h3/projection.py +++ b/rl_engine/reference/minimax_h3/projection.py @@ -106,7 +106,7 @@ def validate_h3_projection( class _NativeProjection: - """PyTorch reference: ``forward`` is the provider (``F.linear``), ``forward_fp32`` the golden.""" + """PyTorch reference: ``forward`` is the provider (``F.linear``); ``forward_fp32`` is golden.""" spec: ProjectionSpec op_class = "reduction" @@ -152,9 +152,7 @@ def __call__(self, x, weight_q, weight_k, weight_v): def forward(self, x, weight_q, weight_k, weight_v): """``(q, k, v)`` from three separate ``F.linear`` calls (the Diffusers layout).""" - widths = { - validate_h3_projection(QKV, x, w, None) for w in (weight_q, weight_k, weight_v) - } + widths = {validate_h3_projection(QKV, x, w, None) for w in (weight_q, weight_k, weight_v)} if len(widths) != 1: raise ValueError(f"h3_qkv_gemm: Q/K/V widths differ: {sorted(widths)}") return F.linear(x, weight_q), F.linear(x, weight_k), F.linear(x, weight_v) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index cf79bde32..e84fc4ddb 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -216,15 +216,21 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_MLP = ( "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" ) - PYTORCH_H3_AUDIO_INPUT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3AudioInputProjectionOp" + PYTORCH_H3_AUDIO_INPUT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3AudioInputProjectionOp" + ) CUDA_H3_AUDIO_INPUT_PROJECTION = ( "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioInputProjectionCudaOp" ) - PYTORCH_H3_AUDIO_OUTPUT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3AudioOutputProjectionOp" + PYTORCH_H3_AUDIO_OUTPUT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3AudioOutputProjectionOp" + ) CUDA_H3_AUDIO_OUTPUT_PROJECTION = ( "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioOutputProjectionCudaOp" ) - PYTORCH_H3_TEXT_CONTEXT_PROJECTION = "rl_engine.reference.minimax_h3.projection.NativeH3TextContextProjectionOp" + PYTORCH_H3_TEXT_CONTEXT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3TextContextProjectionOp" + ) CUDA_H3_TEXT_CONTEXT_PROJECTION = ( "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3TextContextProjectionCudaOp" ) diff --git a/tests/models/minimax_h3/test_h3_projection_rows.py b/tests/models/minimax_h3/test_h3_projection_rows.py index d2134ab72..c3a352a51 100644 --- a/tests/models/minimax_h3/test_h3_projection_rows.py +++ b/tests/models/minimax_h3/test_h3_projection_rows.py @@ -41,6 +41,7 @@ def _cuda_ops(): from rl_engine.backends.cuda.model_specific.minimax_h3 import projection as cu + if not cu.gemm.available(): pytest.skip("needs triton and a CUDA device") return cu @@ -184,9 +185,11 @@ def test_forward_against_fp64_golden(row, tokens): out = op.forward(x, w, b) golden = native.forward_fp32(x, w, b) assert out.dtype == spec.dtype and out.shape == golden.shape - with torch.backends.cuda.matmul.flags(allow_tf32=False) if hasattr( - torch.backends.cuda.matmul, "flags" - ) else _noop(): + with ( + torch.backends.cuda.matmul.flags(allow_tf32=False) + if hasattr(torch.backends.cuda.matmul, "flags") + else _noop() + ): provider = native.forward(x, w, b) floor = _floor(spec) assert _rel(out, golden) <= 1.25 * _rel(provider, golden) + floor @@ -236,9 +239,9 @@ def test_backward_against_fp64_golden(row): w.requires_grad_(), b.requires_grad_() x = _x(spec, 129).requires_grad_() out = op.forward(x, w, b) - grad = torch.randn(out.shape, generator=torch.Generator("cpu").manual_seed(3)).to( - out.dtype - ).cuda() + grad = ( + torch.randn(out.shape, generator=torch.Generator("cpu").manual_seed(3)).to(out.dtype).cuda() + ) out.backward(grad) dx, dw, db = _golden_grads(x.detach(), w.detach(), grad) floor = 1e-5 if spec.dtype == torch.float32 else 2e-2 @@ -308,9 +311,11 @@ def test_dx_rows_are_batch_invariant(): spec = ref.TEXT_CONTEXT op = _cuda_op("text_context_projection") w, b = _params(spec) - g = torch.randn(40, spec.out_features, generator=torch.Generator("cpu").manual_seed(7)).to( - spec.dtype - ).cuda() + g = ( + torch.randn(40, spec.out_features, generator=torch.Generator("cpu").manual_seed(7)) + .to(spec.dtype) + .cuda() + ) def dx(rows): x = _x(spec, 40)[rows].requires_grad_() @@ -382,7 +387,7 @@ def test_qkv_tp_head_shards_equal_full_columns_and_backward_matches_golden(): xg = x.clone().requires_grad_() q, k, v = op.forward(xg, wq, ws[1], ws[2]) q.sum().backward() - gold_dx = (torch.ones(33, 7168, device="cuda", dtype=torch.float64) @ ws[0].double()) + gold_dx = torch.ones(33, 7168, device="cuda", dtype=torch.float64) @ ws[0].double() assert _rel(xg.grad, gold_dx) < 2e-2 assert wq.grad.shape == ws[0].shape diff --git a/tools/validation/models/h3_projection_bench.py b/tools/validation/models/h3_projection_bench.py index 6736ed2f2..a0cbef075 100644 --- a/tools/validation/models/h3_projection_bench.py +++ b/tools/validation/models/h3_projection_bench.py @@ -75,6 +75,14 @@ def main() -> None: cand = lambda: op.forward(x, ws[0], b) # noqa: E731 prov = lambda: F.linear(x, ws[0], b) # noqa: E731 c, p = _time(cand, args.repeats), _time(prov, args.repeats) + gold = F.linear(x.double(), ws[0].double(), None if b is None else b.double()) + got_c, got_p = cand(), prov() + got_c = got_c[0] if isinstance(got_c, (list, tuple)) else got_c + got_p = got_p[0] if isinstance(got_p, (list, tuple)) else got_p + + def rel(a): + return ((a.double() - gold).norm() / gold.norm()).item() + report["rows"].append( { "row": row, @@ -84,9 +92,13 @@ def main() -> None: "candidate_ms": c, "provider_ms": p, "ratio": c / p, + "rel_err_candidate": rel(got_c), + "rel_err_provider": rel(got_p), } ) - print(f"{row:32s} T={tokens:5d} candidate {c:8.3f} ms provider {p:8.3f} ms x{c / p:.2f}") + print( + f"{row:32s} T={tokens:5d} candidate {c:8.3f} ms provider {p:8.3f} ms x{c / p:.2f}" + ) if args.out: args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text(json.dumps(report, indent=1) + "\n") diff --git a/tools/validation/models/h3_projection_charts.py b/tools/validation/models/h3_projection_charts.py new file mode 100644 index 000000000..118b2bb31 --- /dev/null +++ b/tools/validation/models/h3_projection_charts.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Accuracy and performance charts of the H3 projection rows from ``h3_projection_bench`` output.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + + +def main(report: Path, out: Path) -> None: + rows = json.loads(report.read_text())["rows"] + names = sorted({r["row"] for r in rows}) + fig, ax = plt.subplots(1, 2, figsize=(12, 4)) + for name in names: + sel = [r for r in rows if r["row"] == name] + ax[0].plot([r["tokens"] for r in sel], [r["candidate_ms"] for r in sel], "o-", label=name) + ax[0].plot([r["tokens"] for r in sel], [r["provider_ms"] for r in sel], "x--", label=f"{name} provider") + ax[1].plot([r["tokens"] for r in sel], [r["rel_err_candidate"] for r in sel], "o-", label=name) + ax[1].plot([r["tokens"] for r in sel], [r["rel_err_provider"] for r in sel], "x--", label=f"{name} provider") + ax[0].set(xscale="log", yscale="log", xlabel="tokens", ylabel="forward ms (H100)", title="Performance") + ax[1].set(xscale="log", yscale="log", xlabel="tokens", ylabel="relative error vs FP64", title="Accuracy") + for a in ax: + a.grid(alpha=0.3) + a.legend(fontsize=6) + out.parent.mkdir(parents=True, exist_ok=True) + fig.tight_layout() + fig.savefig(out, dpi=120) + + +if __name__ == "__main__": + main(Path(sys.argv[1]), Path(sys.argv[2])) From e109e16c7ac3bcd1c85e956b4aee0ab2d22b3d72 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 13:31:23 +0200 Subject: [PATCH 32/62] test(h3): make the single-unit audio_output shard case an explicit fail-closed test instead of a skip Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 74e5b831adc922c11f8f7660a21be86d7fd93164) --- .../minimax_h3/test_h3_projection_rows.py | 27 +++++++++++++++-- .../validation/models/h3_projection_charts.py | 30 +++++++++++++++---- 2 files changed, 49 insertions(+), 8 deletions(-) diff --git a/tests/models/minimax_h3/test_h3_projection_rows.py b/tests/models/minimax_h3/test_h3_projection_rows.py index c3a352a51..04c92a806 100644 --- a/tests/models/minimax_h3/test_h3_projection_rows.py +++ b/tests/models/minimax_h3/test_h3_projection_rows.py @@ -270,12 +270,15 @@ def test_rows_are_batch_and_position_invariant(row): assert _bits(op.forward(padded, w, b)[:67]) == _bits(full) +def _splittable(row): + spec, _ = ROWS[row] + return spec.out_features % (2 * spec.shard_unit) == 0 + + @requires_cuda -@pytest.mark.parametrize("row", sorted(ROWS)) +@pytest.mark.parametrize("row", [r for r in sorted(ROWS) if _splittable(r)]) def test_output_shard_equals_columns_of_full_product(row): spec, _ = ROWS[row] - if spec.out_features % (2 * spec.shard_unit): - pytest.skip("not splittable into two legal shards") op = _cuda_op(row) w, b = _params(spec) x = _x(spec, 40) @@ -289,6 +292,24 @@ def test_output_shard_equals_columns_of_full_product(row): ) +@requires_cuda +@pytest.mark.parametrize("row", [r for r in sorted(ROWS) if not _splittable(r)]) +def test_single_unit_output_is_not_shardable_and_fails_closed(row): + """A row whose full width is one ``shard_unit`` (``audio_output_projection_fp32``: 32 columns) + has no legal output shard: tensor parallelism replicates it (or splits the reduction axis, which + the WS2 row-parallel path covers). Any narrower weight must be rejected, not silently computed. + """ + spec, _ = ROWS[row] + assert spec.out_features == spec.shard_unit + op = _cuda_op(row) + w, b = _params(spec) + x = _x(spec, 8) + assert op.forward(x, w, b).shape == (8, spec.out_features) + half = spec.out_features // 2 + with pytest.raises(ValueError): + op.forward(x, w[:half].contiguous(), b[:half].contiguous()) + + @requires_cuda @pytest.mark.parametrize("row", sorted(ROWS)) def test_repeat_and_backward_are_deterministic(row): diff --git a/tools/validation/models/h3_projection_charts.py b/tools/validation/models/h3_projection_charts.py index 118b2bb31..a47a0d8d0 100644 --- a/tools/validation/models/h3_projection_charts.py +++ b/tools/validation/models/h3_projection_charts.py @@ -21,11 +21,31 @@ def main(report: Path, out: Path) -> None: for name in names: sel = [r for r in rows if r["row"] == name] ax[0].plot([r["tokens"] for r in sel], [r["candidate_ms"] for r in sel], "o-", label=name) - ax[0].plot([r["tokens"] for r in sel], [r["provider_ms"] for r in sel], "x--", label=f"{name} provider") - ax[1].plot([r["tokens"] for r in sel], [r["rel_err_candidate"] for r in sel], "o-", label=name) - ax[1].plot([r["tokens"] for r in sel], [r["rel_err_provider"] for r in sel], "x--", label=f"{name} provider") - ax[0].set(xscale="log", yscale="log", xlabel="tokens", ylabel="forward ms (H100)", title="Performance") - ax[1].set(xscale="log", yscale="log", xlabel="tokens", ylabel="relative error vs FP64", title="Accuracy") + ax[0].plot( + [r["tokens"] for r in sel], + [r["provider_ms"] for r in sel], + "x--", + label=f"{name} provider", + ) + ax[1].plot( + [r["tokens"] for r in sel], [r["rel_err_candidate"] for r in sel], "o-", label=name + ) + ax[1].plot( + [r["tokens"] for r in sel], + [r["rel_err_provider"] for r in sel], + "x--", + label=f"{name} provider", + ) + ax[0].set( + xscale="log", yscale="log", xlabel="tokens", ylabel="forward ms (H100)", title="Performance" + ) + ax[1].set( + xscale="log", + yscale="log", + xlabel="tokens", + ylabel="relative error vs FP64", + title="Accuracy", + ) for a in ax: a.grid(alpha=0.3) a.legend(fontsize=6) From c8a9bae6325c9d6dc0d44bb88b8cb9614cb32b82 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:14:40 +0200 Subject: [PATCH 33/62] fix(h3): bitwise row-permutation invariant d_weight/d_bias via logical keys Review finding (#568): jointly permuting x and g changed dW/dbias because the token axis was reduced in call order. All projection ops (audio in/out, text context, Q/K/V) now take an optional logical_keys [T,2|3] (negative first column = padding); rows are reduced in ascending key order with padding dropped, active keys must be unique. Key format follows #562. Adds the reviewer's repro plus permutation/padding tests for every row and Q/K/V. Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com> Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit d2941f0dd14a86469aa60afcff031d7221948e07) --- docs/operators/h3-projection-rows.md | 8 +- .../model_specific/minimax_h3/projection.py | 126 +++++++++++++----- .../minimax_h3/test_h3_projection_rows.py | 105 +++++++++++++++ 3 files changed, 208 insertions(+), 31 deletions(-) diff --git a/docs/operators/h3-projection-rows.md b/docs/operators/h3-projection-rows.md index a56eea810..1ed31d15c 100644 --- a/docs/operators/h3-projection-rows.md +++ b/docs/operators/h3-projection-rows.md @@ -48,7 +48,13 @@ y = registry.get_op("text_context_projection", device="cuda")(x, weight, bias) * FP32 rows use `input_precision="ieee"` (no TF32); BF16 rows use the tensor-core dot. The bias is added once in FP32 and the result rounded once to the output dtype; * `d_input = g @ W` reduces the output axis ascending; `d_weight = g^T x` and `d_bias` - reduce tokens in logical row order (the one cross-token reduction); results are cast once; + reduce the token axis (the one cross-token reduction); results are cast once. Every op takes + an optional `logical_keys` (`[T, 2|3]` int64, a negative first component marks padding): + rows are then reduced in ascending key order with padding dropped, so `d_weight` / `d_bias` + are bitwise independent of how the caller permutes, pads or splits the rows. Active keys must + be unique. Without keys the reduction follows the call order and is only invariant for an + unchanged row order; the backward provenance records which of the two was used. The key + format follows #562; * `h3_qkv_gemm` backward is the fused `[Q; K; V]` backward: `N = 7168` is a multiple of `BLOCK_K`, so concatenating the three operands moves no block boundary and the separate and fused layouts have identical gradient bytes. diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py index ac9ae4c56..21c0b3971 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py @@ -15,10 +15,14 @@ the full product. FP32 rows use ``input_precision="ieee"`` (no TF32); BF16 rows use the tensor-core dot. -``d_input`` reduces the output axis ascending, ``d_weight`` / ``d_bias`` reduce tokens in -logical row order. Nothing falls back to cuBLAS: a missing Triton or an unsupported -dtype/shape raises. (An earlier warp-per-column GEMV kernel, ``h3-det-linear-v1``, is -exact too but 50-250x slower than cuBLAS at thousands of tokens, so it is not used here.) +``d_input`` reduces the output axis ascending. ``d_weight`` / ``d_bias`` reduce the token axis: +with ``logical_keys`` (``[T, 2|3]`` int64, a negative first component marks padding) the rows +are first put in ascending key order and padding is dropped, so the parameter gradients are +bitwise independent of how the caller permutes, pads or splits the rows; without keys they +reduce in call order, which is only invariant for an unchanged row order. Nothing falls back +to cuBLAS: a missing Triton or an unsupported dtype/shape raises. (An earlier warp-per-column +GEMV kernel, ``h3-det-linear-v1``, is exact too but 50-250x slower than cuBLAS at thousands of +tokens, so it is not used here.) """ from __future__ import annotations @@ -32,6 +36,7 @@ CONTRACT = gemm.CONTRACT FORWARD_KERNEL_ID = "rl_engine.backends.shared.triton.gemm.h3_projection._projection_kernel" BACKWARD_IMPL = "h3_projection_gemm_v1_ascending_n_dinput_ascending_t_dweight" +BACKWARD_IMPL_KEYED = "h3_projection_gemm_v1_ascending_n_dinput_logical_key_order_dweight" BACKWARD_KERNEL_ID = FORWARD_KERNEL_ID @@ -42,27 +47,77 @@ def _require_triton() -> None: ) +def canonical_row_order(logical_keys: torch.Tensor | None, x: torch.Tensor) -> torch.Tensor | None: + """Row indices of the active rows in ascending logical-key order (``None`` without keys). + + Active keys must be unique, otherwise the order of two equal keys would depend on the call. + """ + + if logical_keys is None: + return None + if not isinstance(logical_keys, torch.Tensor): + raise TypeError("logical_keys must be a torch.Tensor") + rows = x.shape[0] + if logical_keys.dtype != torch.int64 or logical_keys.device != x.device: + raise ValueError("logical_keys must be int64 on the projection device") + if ( + logical_keys.ndim != 2 + or logical_keys.shape[0] != rows + or logical_keys.shape[1] not in (2, 3) + ): + raise ValueError(f"logical_keys must have shape [{rows}, 2] or [{rows}, 3]") + active = (logical_keys[:, 0] >= 0).nonzero().flatten() + keys = logical_keys[active] + if torch.unique(keys, dim=0).shape[0] != keys.shape[0]: + raise ValueError("active logical keys must be unique") + order = torch.arange(keys.shape[0], device=x.device) + for column in reversed(range(keys.shape[1])): # stable passes, least significant key first + order = order[torch.argsort(keys[order, column], stable=True)] + return active[order] + + +def _keyed(logical_keys, x, *params): + """``(order, keyed)`` for ``Function.apply``; keys are validated whenever they are given.""" + + order = canonical_row_order(logical_keys, x) + wants_grad = torch.is_grad_enabled() and any(p is not None and p.requires_grad for p in params) + return (order if wants_grad else None), logical_keys is not None + + +def _weight_grads(grad_out, x, order, with_bias): + if order is not None: + grad_out, x = grad_out[order].contiguous(), x[order].contiguous() + return gemm.linear_backward_weight(grad_out, x, with_bias=with_bias) + + class _Projection(torch.autograd.Function): @staticmethod - def forward(ctx, row, x, weight, bias): + def forward(ctx, row, x, weight, bias, order=None, keyed=False): x = x.contiguous() ctx.row = row ctx.has_bias = bias is not None + ctx.order = order + ctx.keyed = keyed ctx.save_for_backward(x, weight) return gemm.linear_forward(x, weight, bias) @staticmethod def backward(ctx, grad_out): x, weight = ctx.saved_tensors - need_x, need_w, need_b = ctx.needs_input_grad[1:] + need_x, need_w, need_b = ctx.needs_input_grad[1:4] dx = dw = db = None if need_x: dx = gemm.linear_backward_input(grad_out, weight) if need_w or need_b: - dw, db = gemm.linear_backward_weight(grad_out, x, with_bias=ctx.has_bias and need_b) + dw, db = _weight_grads(grad_out, x, ctx.order, ctx.has_bias and need_b) dw = dw if need_w else None - record_backward(ctx.row, kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda") - return None, dx, dw, db + record_backward( + ctx.row, + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL_KEYED if ctx.keyed else BACKWARD_IMPL, + family="cuda", + ) + return None, dx, dw, db, None, None class _QKV(torch.autograd.Function): @@ -75,8 +130,10 @@ class _QKV(torch.autograd.Function): """ @staticmethod - def forward(ctx, x, wq, wk, wv): + def forward(ctx, x, wq, wk, wv, order=None, keyed=False): x = x.contiguous() + ctx.order = order + ctx.keyed = keyed ctx.save_for_backward(x, wq, wk, wv) return tuple(gemm.linear_forward(x, w, None) for w in (wq, wk, wv)) @@ -95,13 +152,16 @@ def backward(ctx, gq, gk, gv): torch.cat(grads, dim=1).contiguous(), torch.cat(weights, dim=0) ) dws = [ - gemm.linear_backward_weight(g, x, with_bias=False)[0] if need else None - for g, need in zip(grads, ctx.needs_input_grad[1:]) + _weight_grads(g, x, ctx.order, False)[0] if need else None + for g, need in zip(grads, ctx.needs_input_grad[1:4]) ] record_backward( - "h3_qkv_gemm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + "h3_qkv_gemm", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL_KEYED if ctx.keyed else BACKWARD_IMPL, + family="cuda", ) - return (dx, *dws) + return (dx, *dws, None, None) class _CudaProjection: @@ -117,19 +177,21 @@ class _CudaProjection: def __init__(self) -> None: _require_triton() - def __call__(self, x, weight, bias=None): - return self.forward(x, weight, bias) + def __call__(self, x, weight, bias=None, *, logical_keys=None): + return self.forward(x, weight, bias, logical_keys=logical_keys) - def forward(self, x, weight, bias=None) -> torch.Tensor: + def forward(self, x, weight, bias=None, *, logical_keys=None) -> torch.Tensor: ref.validate_h3_projection(self.spec, x, weight, bias) if not x.is_cuda: raise ValueError(f"{self.spec.row}: needs CUDA tensors") - return _Projection.apply(self.spec.row, x, weight, bias) + return _Projection.apply( + self.spec.row, x, weight, bias, *_keyed(logical_keys, x, weight, bias) + ) - def forward_fp32(self, x, weight, bias=None) -> torch.Tensor: + def forward_fp32(self, x, weight, bias=None, *, logical_keys=None) -> torch.Tensor: """The CUDA projection already accumulates in FP32 and rounds once.""" - return self.forward(x, weight, bias) + return self.forward(x, weight, bias, logical_keys=logical_keys) class H3AudioInputProjectionCudaOp(_CudaProjection): @@ -156,10 +218,10 @@ class H3QKVGemmCudaOp: def __init__(self) -> None: _require_triton() - def __call__(self, x, weight_q, weight_k, weight_v): - return self.forward(x, weight_q, weight_k, weight_v) + def __call__(self, x, weight_q, weight_k, weight_v, *, logical_keys=None): + return self.forward(x, weight_q, weight_k, weight_v, logical_keys=logical_keys) - def forward(self, x, weight_q, weight_k, weight_v): + def forward(self, x, weight_q, weight_k, weight_v, *, logical_keys=None): widths = { ref.validate_h3_projection(self.spec, x, w, None) for w in (weight_q, weight_k, weight_v) @@ -168,17 +230,19 @@ def forward(self, x, weight_q, weight_k, weight_v): raise ValueError(f"h3_qkv_gemm: Q/K/V widths differ: {sorted(widths)}") if not x.is_cuda: raise ValueError("h3_qkv_gemm: needs CUDA tensors") - return _QKV.apply(x, weight_q, weight_k, weight_v) + return _QKV.apply( + x, weight_q, weight_k, weight_v, *_keyed(logical_keys, x, weight_q, weight_k, weight_v) + ) - def project(self, x, weight) -> torch.Tensor: + def project(self, x, weight, *, logical_keys=None) -> torch.Tensor: """One of the three projections (what a graph node that owns only Q, K or V runs).""" ref.validate_h3_projection(self.spec, x, weight, None) if not x.is_cuda: raise ValueError("h3_qkv_gemm: needs CUDA tensors") - return _Projection.apply(self.spec.row, x, weight, None) + return _Projection.apply(self.spec.row, x, weight, None, *_keyed(logical_keys, x, weight)) - def forward_fused(self, x, weight_qkv) -> torch.Tensor: + def forward_fused(self, x, weight_qkv, *, logical_keys=None) -> torch.Tensor: """``(T, 3 * N)`` product of the concatenated ``[Q; K; V]`` weight.""" n = ( @@ -191,7 +255,9 @@ def forward_fused(self, x, weight_qkv) -> torch.Tensor: ref.validate_h3_projection(self.spec, x, weight_qkv[:n], None) if not x.is_cuda: raise ValueError("h3_qkv_gemm: needs CUDA tensors") - return _Projection.apply(self.spec.row, x, weight_qkv, None) + return _Projection.apply( + self.spec.row, x, weight_qkv, None, *_keyed(logical_keys, x, weight_qkv) + ) - def forward_fp32(self, x, weight_q, weight_k, weight_v): - return self.forward(x, weight_q, weight_k, weight_v) + def forward_fp32(self, x, weight_q, weight_k, weight_v, *, logical_keys=None): + return self.forward(x, weight_q, weight_k, weight_v, logical_keys=logical_keys) diff --git a/tests/models/minimax_h3/test_h3_projection_rows.py b/tests/models/minimax_h3/test_h3_projection_rows.py index 04c92a806..120729257 100644 --- a/tests/models/minimax_h3/test_h3_projection_rows.py +++ b/tests/models/minimax_h3/test_h3_projection_rows.py @@ -445,3 +445,108 @@ def test_registry_dispatches_cuda(): cu.H3AudioInputProjectionCudaOp ) assert type(registry.get_op("h3_qkv_gemm", device="cuda")) is cu.H3QKVGemmCudaOp + + +def _keys(tokens, device="cuda"): + t = torch.arange(tokens, device=device) + return torch.stack([t // 7, t % 7, t], dim=1).to(torch.int64) + + +@requires_cuda +def test_audio_input_param_grads_are_row_permutation_invariant_with_keys(): + """Review repro (#568): jointly permuting x and g must not change dW / dbias.""" + + op = _cuda_ops().H3AudioInputProjectionCudaOp() + x = torch.ones(3, 32, device="cuda") + g = torch.zeros(3, 5376, device="cuda") + g[:, 0] = torch.tensor([1.0, 2**-24, -1.0], device="cuda") + keys = torch.tensor([[0, 0, 0], [0, 0, 1], [0, 0, 2]], device="cuda") + + def grads(order, use_keys): + w = torch.zeros(5376, 32, device="cuda", requires_grad=True) + b = torch.zeros(5376, device="cuda", requires_grad=True) + out = op(x[order], w, b, logical_keys=keys[order] if use_keys else None) + return torch.autograd.grad(out, (w, b), g[order]) + + for i, j in (([0, 1, 2], [0, 2, 1]), ([0, 1, 2], [2, 1, 0]), ([0, 1, 2], [1, 0, 2])): + a, b_ = grads(i, True), grads(j, True) + assert torch.equal(a[0], b_[0]) and torch.equal(a[1], b_[1]) + # without keys the reduction follows the call order (documented), so it is not asserted here + + +@requires_cuda +@pytest.mark.parametrize("row", sorted(ROWS)) +def test_param_grads_are_invariant_to_permutation_and_padding(row): + cu = _cuda_ops() + spec, name = ROWS[row] + op = getattr( + cu, + { + "audio_input_projection_fp32": "H3AudioInputProjectionCudaOp", + "audio_output_projection_fp32": "H3AudioOutputProjectionCudaOp", + "text_context_projection": "H3TextContextProjectionCudaOp", + }[row], + )() + tokens = 37 + x = _x(spec, tokens) + gen = torch.Generator(device="cuda").manual_seed(5) + g = torch.randn(tokens, spec.out_features, device="cuda", generator=gen).to(spec.dtype) + w0, b0 = _params(spec) + keys = _keys(tokens) + + def run(order, pad): + w, b = w0.detach().clone().requires_grad_(), ( + None if b0 is None else b0.detach().clone().requires_grad_() + ) + xs, gs, ks = x[order], g[order], keys[order] + if pad: + xs = torch.cat([xs, torch.randn_like(xs[:pad])]) + gs = torch.cat([gs, torch.randn_like(gs[:pad])]) + ks = torch.cat([ks, torch.full((pad, 3), -1, device="cuda", dtype=torch.int64)]) + params = (w,) if b is None else (w, b) + return torch.autograd.grad(op(xs, w, b, logical_keys=ks), params, gs) + + ident = torch.arange(tokens, device="cuda") + base = run(ident, 0) + for order, pad in ( + (ident.flip(0), 0), + (torch.randperm(tokens, device="cuda", generator=gen), 0), + (ident, 5), + ): + got = run(order, pad) + assert all(torch.equal(a, b) for a, b in zip(base, got)) + + +@requires_cuda +def test_qkv_param_grads_are_invariant_to_permutation_with_keys(): + cu = _cuda_ops() + op = cu.H3QKVGemmCudaOp() + tokens = 41 + gen = torch.Generator(device="cuda").manual_seed(9) + x = torch.randn(tokens, 5376, device="cuda", generator=gen).bfloat16() + gs = [torch.randn(tokens, 7168, device="cuda", generator=gen).bfloat16() for _ in range(3)] + ws = [torch.randn(7168, 5376, device="cuda", generator=gen).bfloat16() for _ in range(3)] + keys = _keys(tokens) + + def run(order): + params = [w.detach().clone().requires_grad_() for w in ws] + outs = op(x[order], *params, logical_keys=keys[order]) + return torch.autograd.grad(outs, params, [g[order] for g in gs]) + + ident = torch.arange(tokens, device="cuda") + base = run(ident) + other = run(torch.randperm(tokens, device="cuda", generator=gen)) + assert all(torch.equal(a, b) for a, b in zip(base, other)) + + +@requires_cuda +def test_logical_keys_are_validated(): + op = _cuda_ops().H3AudioInputProjectionCudaOp() + x = torch.ones(2, 32, device="cuda") + w, b = torch.zeros(5376, 32, device="cuda"), torch.zeros(5376, device="cuda") + with pytest.raises(ValueError, match="unique"): + op(x, w, b, logical_keys=torch.tensor([[0, 0], [0, 0]], device="cuda")) + with pytest.raises(ValueError, match="int64"): + op(x, w, b, logical_keys=torch.zeros(2, 2, device="cuda", dtype=torch.int32)) + with pytest.raises(ValueError, match="shape"): + op(x, w, b, logical_keys=torch.zeros(3, 2, device="cuda", dtype=torch.int64)) From 5f7536c0efc09bed3560d3e1b96dc5153af6d600 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:19:22 +0200 Subject: [PATCH 34/62] feat(h3): block graph passes logical row keys; skip the gather for already ordered rows Follows the canonical-row design of #562: the Q/K/V node of the block graph supplies (batch row, position) keys so its dW is reduced in logical order, and rows that already arrive in canonical order are reduced in place (same bytes, no gather of x and g). Co-authored-by: Yitao Wang <75578057+walker-ai@users.noreply.github.com> Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> (cherry picked from commit 0da058326ae8fb5a9f97cbe02124f8f3b21b4f04) --- .../model_specific/minimax_h3/projection.py | 10 ++++++++-- rl_engine/validation/models/h3_block.py | 17 ++++++++++++++++- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py index 21c0b3971..a7d63e33c 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py @@ -48,7 +48,10 @@ def _require_triton() -> None: def canonical_row_order(logical_keys: torch.Tensor | None, x: torch.Tensor) -> torch.Tensor | None: - """Row indices of the active rows in ascending logical-key order (``None`` without keys). + """Row indices of the active rows in ascending logical-key order. + + ``None`` without keys, and also when the rows already are the active rows in canonical order + (the reduction then runs in place and gives the same bytes). Active keys must be unique, otherwise the order of two equal keys would depend on the call. """ @@ -73,7 +76,10 @@ def canonical_row_order(logical_keys: torch.Tensor | None, x: torch.Tensor) -> t order = torch.arange(keys.shape[0], device=x.device) for column in reversed(range(keys.shape[1])): # stable passes, least significant key first order = order[torch.argsort(keys[order, column], stable=True)] - return active[order] + order = active[order] + if order.numel() == rows and torch.equal(order, torch.arange(rows, device=x.device)): + return None # already in canonical order: reduce in place, no gather of x and g + return order def _keyed(logical_keys, x, *params): diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py index 1bdeaba46..0f895613e 100644 --- a/rl_engine/validation/models/h3_block.py +++ b/rl_engine/validation/models/h3_block.py @@ -170,12 +170,27 @@ def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> ) +def _row_keys(x: torch.Tensor) -> torch.Tensor: + """Logical identity ``(batch row, token position)`` of every flattened row of ``[B, S, C]``. + + The parameter-gradient reduction of a projection sorts its rows by this key, so dW does not + depend on the physical order the rows reach the GEMM in. + """ + + batch, seq = x.shape[0], x.shape[1] + rows = torch.arange(batch, device=x.device).repeat_interleave(seq) + positions = torch.arange(seq, device=x.device).repeat(batch) + return torch.stack([rows, positions], dim=1) + + def _qkv_node(name: str, weight: str) -> Node: """Q, K or V projection: the ``h3_qkv_gemm`` row (fixed-order Triton GEMM, BF16).""" def candidate(ops, ctx, x): flat = x.reshape(-1, x.shape[-1]) - return ops.qkv.project(flat, _w(ctx, weight)).view(*x.shape[:-1], -1) + return ops.qkv.project(flat, _w(ctx, weight), logical_keys=_row_keys(x)).view( + *x.shape[:-1], -1 + ) return Node( name=name, From 17d0e9e9b2aac7d30eee2b3ec79ec7dd36b4d0e0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:54:43 +0200 Subject: [PATCH 35/62] [WS1][CUDA][MiniMax-H3] ws1_one_h3_block + full-model qualification: block 0 node-by-node and 50-block/generation gates Squashes the one-block branch (its row-by-row history: timestep sinusoid, RMSNorm, AdaLN rows, sequence/tensor parallel norm rows, final AdaLN, one block, 50-block qualification) onto the current test-h3 so the pull request carries a reviewable, signed history. RFC #420. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .coderabbit.yaml | 6 + .github/workflows/ci.yml | 7 +- .../models/benchmark_h3_conditioning.py | 71 + build_tools/extensions.py | 22 +- csrc/bindings/ops.cpp | 133 + csrc/cuda/h3/adaln_row_gather.cu | 259 ++ csrc/cuda/h3/det_linear.cu | 620 ++++ csrc/cuda/h3/gate_residual.cu | 326 ++ csrc/cuda/h3/rmsnorm_modulate.cu | 593 ++++ csrc/cuda/h3/timestep_sinusoid.cu | 93 + docs/.nav.yml | 10 + docs/operators/README.md | 10 + docs/operators/h3-adaln-gate-residual.md | 143 + docs/operators/h3-adaln-projection.md | 207 ++ docs/operators/h3-adaln-row-gather.md | 204 ++ docs/operators/h3-final-adaln-out.md | 175 + docs/operators/h3-one-block.md | 183 + docs/operators/h3-rmsnorm.md | 169 + docs/operators/h3-sp-norm-adaln.md | 117 + docs/operators/h3-timestep-mlp.md | 184 + docs/operators/h3-timestep-sinusoid.md | 163 + docs/operators/h3-tp-adaln-3mod.md | 113 + examples/models/minimax_h3/README.md | 105 + examples/models/minimax_h3/candidate.sbatch | 59 + .../models/minimax_h3/diagnose-block.sbatch | 41 + .../minimax_h3/full-model-ensemble-rule.json | 2346 +++++++++++++ .../minimax_h3/full-model-thresholds.json | 3064 +++++++++++++++++ .../minimax_h3/generation-compare.sbatch | 26 + .../minimax_h3/generation-contract.json | 18 + .../models/minimax_h3/generation-inputs.json | 8 + .../models/minimax_h3/generation-prompt.txt | 7 + .../minimax_h3/generation-resume-check.sbatch | 77 + examples/models/minimax_h3/generation.sbatch | 56 + .../models/minimax_h3/reference-manifest.json | 1641 +++++++++ examples/models/minimax_h3/reference.sbatch | 51 + rl_engine/_C.pyi | 150 + .../backends/cuda/model_specific/__init__.py | 2 + .../model_specific/minimax_h3/__init__.py | 2 + .../minimax_h3/adaln_modulation.py | 133 + .../minimax_h3/adaln_projection.py | 132 + .../minimax_h3/adaln_row_gather.py | 113 + .../model_specific/minimax_h3/det_linear.py | 121 + .../minimax_h3/final_adaln_out.py | 110 + .../minimax_h3/gate_residual.py | 83 + .../cuda/model_specific/minimax_h3/rmsnorm.py | 103 + .../minimax_h3/sp_norm_adaln.py | 418 +++ .../model_specific/minimax_h3/timestep_mlp.py | 116 + .../minimax_h3/timestep_sinusoid.py | 117 + .../minimax_h3/tp_adaln_projection.py | 213 ++ .../model_specific/minimax_h3/ws2_comm.py | 33 + rl_engine/reference/minimax_h3/__init__.py | 30 + .../reference/minimax_h3/adaln_projection.py | 123 + .../reference/minimax_h3/adaln_row_gather.py | 141 + rl_engine/reference/minimax_h3/block.py | 98 + .../reference/minimax_h3/final_adaln_out.py | 107 + rl_engine/reference/minimax_h3/fixed_order.py | 46 + .../reference/minimax_h3/gate_residual.py | 97 + rl_engine/reference/minimax_h3/rmsnorm.py | 153 + .../reference/minimax_h3/timestep_mlp.py | 93 + .../reference/minimax_h3/timestep_sinusoid.py | 129 + rl_engine/runtime/registry.py | 84 + rl_engine/validation/models/h3_block.py | 718 ++++ rl_engine/validation/models/h3_cases.py | 188 + rl_engine/validation/models/h3_chain.py | 455 +++ .../validation/models/h3_full_manifest.json | 311 ++ rl_engine/validation/models/h3_full_model.py | 242 ++ .../validation/models/h3_generation_resume.py | 320 ++ rl_engine/validation/models/h3_manifest.json | 137 + rl_engine/validation/models/h3_provider.py | 192 ++ rl_engine/validation/models/h3_report.py | 831 +++++ rl_engine/validation/models/h3_weights.py | 108 + rl_engine/validation/models/h3_ws2.py | 418 +++ .../validation/operators/operator_inputs.py | 168 + .../validation/operators/operator_specs.py | 104 + .../build/test_build_platform_collectives.py | 32 +- tests/models/minimax_h3/conftest.py | 58 + .../minimax_h3/test_h3_adaln_gate_residual.py | 115 + .../minimax_h3/test_h3_adaln_modulation.py | 126 + .../minimax_h3/test_h3_adaln_projection.py | 298 ++ .../minimax_h3/test_h3_adaln_row_gather.py | 300 ++ .../minimax_h3/test_h3_benchmark_timing.py | 267 ++ .../minimax_h3/test_h3_chain_replay_cli.py | 83 + tests/models/minimax_h3/test_h3_cli.py | 174 + .../minimax_h3/test_h3_conditioning_e2e.py | 83 + tests/models/minimax_h3/test_h3_det_linear.py | 68 + .../minimax_h3/test_h3_final_adaln_out.py | 203 ++ tests/models/minimax_h3/test_h3_full_model.py | 187 + .../minimax_h3/test_h3_generation_resume.py | 384 +++ .../minimax_h3/test_h3_native_validation.py | 384 +++ tests/models/minimax_h3/test_h3_one_block.py | 139 + tests/models/minimax_h3/test_h3_report.py | 259 ++ tests/models/minimax_h3/test_h3_rmsnorm.py | 302 ++ .../minimax_h3/test_h3_sp_norm_adaln.py | 327 ++ .../models/minimax_h3/test_h3_timestep_mlp.py | 296 ++ .../minimax_h3/test_h3_timestep_sinusoid.py | 268 ++ tests/models/minimax_h3/test_h3_tp_adaln.py | 159 + tests/models/minimax_h3/test_h3_weights.py | 74 + .../minimax_h3/test_h3_ws2_processes.py | 47 + .../minimax_h3/test_prepare_h3_weights.py | 70 + tests/runtime/test_dispatch.py | 39 + .../operators/test_ws1_gtest_gpu.py | 9 + tools/migration/README.md | 3 + tools/migration/layout.json | 94 + tools/validation/models/h3_bitwise_reuse.py | 250 ++ .../validation/models/h3_block_diagnostic.py | 148 + tools/validation/models/h3_block_replay.py | 114 + tools/validation/models/h3_block_ulp_sweep.py | 96 + .../validation/models/h3_candidate_prepare.py | 90 + tools/validation/models/h3_chain_replay.py | 137 + tools/validation/models/h3_evidence.py | 74 + .../models/h3_evidence_integrity.py | 34 + .../models/h3_first_step_diagnostic.py | 117 + tools/validation/models/h3_full_compare.py | 152 + .../models/h3_full_ensemble_rule.py | 156 + tools/validation/models/h3_full_generation.py | 365 ++ tools/validation/models/h3_full_model.py | 297 ++ .../models/h3_full_provider_spread.py | 92 + .../validation/models/h3_generation_adopt.py | 63 + .../models/h3_generation_compare.py | 150 + tools/validation/models/h3_prior_art.py | 946 +++++ tools/validation/models/h3_ws2_evidence.py | 128 + .../models/plot_h3_bitwise_reuse.py | 111 + tools/validation/models/plot_h3_block.py | 172 + tools/validation/models/plot_h3_evidence.py | 768 +++++ tools/validation/models/plot_h3_prior_art.py | 139 + tools/weights/prepare_h3_weights.py | 125 + 126 files changed, 28014 insertions(+), 4 deletions(-) create mode 100644 .coderabbit.yaml create mode 100644 benchmarks/models/benchmark_h3_conditioning.py create mode 100644 csrc/cuda/h3/adaln_row_gather.cu create mode 100644 csrc/cuda/h3/det_linear.cu create mode 100644 csrc/cuda/h3/gate_residual.cu create mode 100644 csrc/cuda/h3/rmsnorm_modulate.cu create mode 100644 csrc/cuda/h3/timestep_sinusoid.cu create mode 100644 docs/operators/h3-adaln-gate-residual.md create mode 100644 docs/operators/h3-adaln-projection.md create mode 100644 docs/operators/h3-adaln-row-gather.md create mode 100644 docs/operators/h3-final-adaln-out.md create mode 100644 docs/operators/h3-one-block.md create mode 100644 docs/operators/h3-rmsnorm.md create mode 100644 docs/operators/h3-sp-norm-adaln.md create mode 100644 docs/operators/h3-timestep-mlp.md create mode 100644 docs/operators/h3-timestep-sinusoid.md create mode 100644 docs/operators/h3-tp-adaln-3mod.md create mode 100644 examples/models/minimax_h3/README.md create mode 100644 examples/models/minimax_h3/candidate.sbatch create mode 100644 examples/models/minimax_h3/diagnose-block.sbatch create mode 100644 examples/models/minimax_h3/full-model-ensemble-rule.json create mode 100644 examples/models/minimax_h3/full-model-thresholds.json create mode 100644 examples/models/minimax_h3/generation-compare.sbatch create mode 100644 examples/models/minimax_h3/generation-contract.json create mode 100644 examples/models/minimax_h3/generation-inputs.json create mode 100644 examples/models/minimax_h3/generation-prompt.txt create mode 100644 examples/models/minimax_h3/generation-resume-check.sbatch create mode 100644 examples/models/minimax_h3/generation.sbatch create mode 100644 examples/models/minimax_h3/reference-manifest.json create mode 100644 examples/models/minimax_h3/reference.sbatch create mode 100644 rl_engine/backends/cuda/model_specific/__init__.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py create mode 100644 rl_engine/reference/minimax_h3/__init__.py create mode 100644 rl_engine/reference/minimax_h3/adaln_projection.py create mode 100644 rl_engine/reference/minimax_h3/adaln_row_gather.py create mode 100644 rl_engine/reference/minimax_h3/block.py create mode 100644 rl_engine/reference/minimax_h3/final_adaln_out.py create mode 100644 rl_engine/reference/minimax_h3/fixed_order.py create mode 100644 rl_engine/reference/minimax_h3/gate_residual.py create mode 100644 rl_engine/reference/minimax_h3/rmsnorm.py create mode 100644 rl_engine/reference/minimax_h3/timestep_mlp.py create mode 100644 rl_engine/reference/minimax_h3/timestep_sinusoid.py create mode 100644 rl_engine/validation/models/h3_block.py create mode 100644 rl_engine/validation/models/h3_cases.py create mode 100644 rl_engine/validation/models/h3_chain.py create mode 100644 rl_engine/validation/models/h3_full_manifest.json create mode 100644 rl_engine/validation/models/h3_full_model.py create mode 100644 rl_engine/validation/models/h3_generation_resume.py create mode 100644 rl_engine/validation/models/h3_manifest.json create mode 100644 rl_engine/validation/models/h3_provider.py create mode 100644 rl_engine/validation/models/h3_report.py create mode 100644 rl_engine/validation/models/h3_weights.py create mode 100644 rl_engine/validation/models/h3_ws2.py create mode 100644 tests/models/minimax_h3/conftest.py create mode 100644 tests/models/minimax_h3/test_h3_adaln_gate_residual.py create mode 100644 tests/models/minimax_h3/test_h3_adaln_modulation.py create mode 100644 tests/models/minimax_h3/test_h3_adaln_projection.py create mode 100644 tests/models/minimax_h3/test_h3_adaln_row_gather.py create mode 100644 tests/models/minimax_h3/test_h3_benchmark_timing.py create mode 100644 tests/models/minimax_h3/test_h3_chain_replay_cli.py create mode 100644 tests/models/minimax_h3/test_h3_cli.py create mode 100644 tests/models/minimax_h3/test_h3_conditioning_e2e.py create mode 100644 tests/models/minimax_h3/test_h3_det_linear.py create mode 100644 tests/models/minimax_h3/test_h3_final_adaln_out.py create mode 100644 tests/models/minimax_h3/test_h3_full_model.py create mode 100644 tests/models/minimax_h3/test_h3_generation_resume.py create mode 100644 tests/models/minimax_h3/test_h3_native_validation.py create mode 100644 tests/models/minimax_h3/test_h3_one_block.py create mode 100644 tests/models/minimax_h3/test_h3_report.py create mode 100644 tests/models/minimax_h3/test_h3_rmsnorm.py create mode 100644 tests/models/minimax_h3/test_h3_sp_norm_adaln.py create mode 100644 tests/models/minimax_h3/test_h3_timestep_mlp.py create mode 100644 tests/models/minimax_h3/test_h3_timestep_sinusoid.py create mode 100644 tests/models/minimax_h3/test_h3_tp_adaln.py create mode 100644 tests/models/minimax_h3/test_h3_weights.py create mode 100644 tests/models/minimax_h3/test_h3_ws2_processes.py create mode 100644 tests/models/minimax_h3/test_prepare_h3_weights.py create mode 100644 tools/validation/models/h3_bitwise_reuse.py create mode 100644 tools/validation/models/h3_block_diagnostic.py create mode 100644 tools/validation/models/h3_block_replay.py create mode 100644 tools/validation/models/h3_block_ulp_sweep.py create mode 100644 tools/validation/models/h3_candidate_prepare.py create mode 100644 tools/validation/models/h3_chain_replay.py create mode 100644 tools/validation/models/h3_evidence.py create mode 100644 tools/validation/models/h3_evidence_integrity.py create mode 100644 tools/validation/models/h3_first_step_diagnostic.py create mode 100644 tools/validation/models/h3_full_compare.py create mode 100644 tools/validation/models/h3_full_ensemble_rule.py create mode 100644 tools/validation/models/h3_full_generation.py create mode 100644 tools/validation/models/h3_full_model.py create mode 100644 tools/validation/models/h3_full_provider_spread.py create mode 100644 tools/validation/models/h3_generation_adopt.py create mode 100644 tools/validation/models/h3_generation_compare.py create mode 100644 tools/validation/models/h3_prior_art.py create mode 100644 tools/validation/models/h3_ws2_evidence.py create mode 100644 tools/validation/models/plot_h3_bitwise_reuse.py create mode 100644 tools/validation/models/plot_h3_block.py create mode 100644 tools/validation/models/plot_h3_evidence.py create mode 100644 tools/validation/models/plot_h3_prior_art.py create mode 100644 tools/weights/prepare_h3_weights.py diff --git a/.coderabbit.yaml b/.coderabbit.yaml new file mode 100644 index 000000000..ebea58ae2 --- /dev/null +++ b/.coderabbit.yaml @@ -0,0 +1,6 @@ +# yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json +inheritance: true +reviews: + auto_review: + base_branches: + - "^test-h3$" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4634569cd..58d2dd175 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-h3 ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-h3 ] permissions: contents: read @@ -71,6 +71,9 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q + - name: Run full H3 binding contracts + run: python -m pytest -q tests/models/minimax_h3/test_h3_full_model.py + - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v diff --git a/benchmarks/models/benchmark_h3_conditioning.py b/benchmarks/models/benchmark_h3_conditioning.py new file mode 100644 index 000000000..261556809 --- /dev/null +++ b/benchmarks/models/benchmark_h3_conditioning.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. + +CUDA-event medians, bandwidth and peak memory per case, plus the backend +the registry dispatched. Timings alternate candidate/provider execution order; +backward timings exclude forward setup. Cases live in ``rl_engine/validation/models/h3_report.py``. + + python benchmarks/models/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 + python benchmarks/models/benchmark_h3_conditioning.py --op all --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment # noqa: E402 +from rl_engine.validation.models.h3_report import PERF_CASES, TIMED_KEYS, measure # noqa: E402 + + +def main() -> None: + """Benchmark selected CUDA operators, print summaries, and optionally save JSON.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--json", type=Path, default=None) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + + registry = KernelRegistry() + env = environment() + print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + print("timing order alternates each iteration; backward timings exclude forward setup") + results = [] + for name in list(PERF_CASES) if args.op == "all" else [args.op]: + for case in PERF_CASES[name](registry): + row = measure(case, args.warmup, args.iters) + results.append(row) + parts = [ + f"{key}={row[f'{key}_us']:.2f}us ({row[f'{key}_gbps']:.1f} GB/s, " + f"peak {row[f'{key}_peak_mib']:.2f} MiB)" + for key in TIMED_KEYS + if f"{key}_us" in row + ] + order = row["execution_order"] + phases = [" -> ".join(order[f"iteration_{i}"]) for i in (0, 1)] + print( + f"{name} {row['case']} [{row['backend']}]: " + + "; ".join(parts) + + "; alternating order: " + + " / ".join(phases) + ) + if args.json is not None: + args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 670fdb072..dbddf0c0a 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -14,6 +14,20 @@ _PROJECT_ROOT = Path(__file__).resolve().parents[1] +# MiniMax-H3 (RFC #420) kernels promise bitwise parity with diffusers' eager +# FP32 math. --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for +# approximations and flushes denormals, so these sources are left out of +# KERNEL_ALIGN_USE_FAST_MATH=1 builds; their bindings are gated on +# KERNEL_ALIGN_WITH_H3_KERNELS and dispatch falls back to the PyTorch reference. +H3_BITWISE_SOURCES = ( + "csrc/cuda/h3/timestep_sinusoid.cu", + "csrc/cuda/h3/det_linear.cu", + "csrc/cuda/h3/adaln_row_gather.cu", + "csrc/cuda/h3/rmsnorm_modulate.cu", + "csrc/cuda/h3/gate_residual.cu", +) +H3_KERNELS_DEFINE = "-DKERNEL_ALIGN_WITH_H3_KERNELS" + def _resolve_package_data_source( filename: str | os.PathLike[str], @@ -219,6 +233,7 @@ def get_extensions(): ) if is_rocm or torch.cuda.is_available(): + use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) cuda_sources = [ "csrc/bindings/ops.cpp", "csrc/cuda/logprob/fused_logp_kernel.cu", @@ -241,12 +256,14 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") + if not use_fast_math: + cuda_sources.extend(H3_BITWISE_SOURCES) # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + if use_fast_math: nvcc_flags.append("--use_fast_math") if not is_rocm: cc_major, cc_minor = torch.cuda.get_device_capability() @@ -333,6 +350,9 @@ def get_extensions(): if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") + if not is_rocm and not use_fast_math: + cxx_flags.append(H3_KERNELS_DEFINE) + nvcc_flags.append(H3_KERNELS_DEFINE) if not is_rocm: sm90_srcs = [ diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 0c7c3e7f4..0046fce21 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -500,6 +500,67 @@ at::Tensor prefix_shared_attention( #endif #endif +// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. +// Fast-math builds omit the H3 sources (bitwise contract), so gate on the build define. +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period, bool check_range); +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation); +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype); +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight); +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype); +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias); +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num); +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype); +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index); +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile); +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index); +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + 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 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; @@ -758,4 +819,76 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif + +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) + m.def("h3_timestep_sinusoid_forward", torch::wrap_pybind_function(h3_timestep_sinusoid_forward), + "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", + py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, + py::arg("check_range") = true); + m.def("h3_det_linear_forward", &h3_det_linear_forward, + "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU", + py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), + py::arg("activation") = 0, py::arg("save_pre_activation") = false); + m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, + "Deterministic grad @ weight with fixed 64-row N chunks folded in order", + py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_input_partials", &h3_det_linear_backward_input_partials, + "Per-64-row-chunk FP32 partials of grad @ weight, before the ascending fold", + py::arg("grad"), py::arg("weight")); + m.def("h3_det_linear_fold_chunks", &h3_det_linear_fold_chunks, + "Ascending left fold of h3_det_linear_backward_input_partials, cast once", + py::arg("partial"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, + "Deterministic dW/dbias as ascending-row FP32 folds", + py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); + m.def("h3_adaln_row_gather_forward", &h3_adaln_row_gather_forward, + "Fused six-way AdaLN row gather by timestep_index * 3 + token_tag (pure copy)", + py::arg("rows"), py::arg("timestep_indices"), py::arg("token_tags"), + py::arg("chunks") = 6, py::arg("modality_num") = 3); + m.def("h3_adaln_row_gather_backward", &h3_adaln_row_gather_backward, + "Deterministic segmented sum (sorted tiles folded in order) for the row gather", + py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), + py::arg("seg_first_tile"), py::arg("out_dtype")); + m.def("h3_rmsnorm_forward", &h3_rmsnorm_forward, + "RMSNorm replaying PyTorch's reduction order, with optional fused AdaLN modulation", + py::arg("x"), py::arg("weight"), py::arg("eps"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_rmsnorm_backward", &h3_rmsnorm_backward, + "Deterministic RMSNorm(+modulation) backward: row-local dx, tiled dweight, sorted table grads", + 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(), py::arg("sorted_pos") = py::none(), + py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_residual_forward", &h3_gate_residual_forward, + "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)", + py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index")); + m.def("h3_gate_residual_backward", &h3_gate_residual_backward, + "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 new file mode 100644 index 000000000..0de2d98d5 --- /dev/null +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 AdaLN row gather (RFC #420 `adaln_row_gather`). +// +// rows: (3T, C * H) modulation rows, row r = timestep * 3 + modality and +// column block c = shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, +// gate_mlp. For every packed-sequence position s: +// +// r[s] = timestep_indices[s] * modality_num + token_tags[s] +// out[c, s, :] = rows[r[s], c * H : (c + 1) * H] +// +// Forward is a pure copy: bitwise equal to C index_select calls, one launch, +// independent of S, packing order and repeats. +// +// Backward d_rows[r, j] = sum_{s : r[s] = r} grad[c(j), s, h(j)] is the one +// cross-row reduction. Positions are visited in a fixed order: the caller +// sorts them by (r, s) with a stable sort and cuts every segment into tiles of +// at most kTile positions. Each tile is an ascending FP32 chain from 0; the +// tiles of a segment are then left-folded in ascending order and cast once. +// No atomics: the result is a function of (r[s], grad) only. + +#include +#include +#include +#include +#include + +#include + +namespace { + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(double v) { return static_cast(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> +__device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> +__device__ __forceinline__ double from_float(float v) { return v; } + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; +template <> +struct CudaType { + using type = __half; +}; + +template +__device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, + int64_t modality_num, int64_t num_timesteps) { + const int64_t timestep = static_cast(ti[s]); + const int64_t tag = static_cast(tags[s]); + // Check each semantic index before multiplication, including values that overflow int64. + CUDA_KERNEL_ASSERT(timestep >= 0 && timestep < num_timesteps); + CUDA_KERNEL_ASSERT(tag >= 0 && tag < modality_num); + return timestep * modality_num + tag; +} + +// One block per (s, c); 16-byte copies when every row start is 16-byte aligned. +template +__global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_stride, + const idx_t* __restrict__ ti, + const idx_t* __restrict__ tags, T* __restrict__ out, + int64_t seq, int64_t hidden, int64_t chunks, + int64_t modality_num, int64_t num_timesteps) { + for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { + const int64_t c = blockIdx.y; + const int64_t r = adaln_row(ti, tags, s, modality_num, num_timesteps); + const T* src = rows + r * row_stride + c * hidden; + T* dst = out + (c * seq + s) * hidden; + if constexpr (kVector) { + constexpr int kVec = 16 / sizeof(T); + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t i = threadIdx.x; i < hidden / kVec; i += blockDim.x) dst4[i] = src4[i]; + } else { + for (int64_t i = threadIdx.x; i < hidden; i += blockDim.x) dst[i] = src[i]; + } + } +} + +// partial[tile, j] = sum over sorted positions p in tile (ascending) of grad[c, s_p, h] +template +__global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t seq, + int64_t hidden, int64_t width) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + const int64_t c = j / hidden; + const int64_t h = j - c * hidden; + const g_t* g = grad + c * seq * hidden + h; + float acc = 0.0f; + const int64_t begin = tile_begin[tile], end = tile_end[tile]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= seq); + for (int64_t p = begin; p < end; ++p) { + CUDA_KERNEL_ASSERT(sorted_pos[p] >= 0 && sorted_pos[p] < seq); + acc += to_float(g[sorted_pos[p] * hidden]); + } + partial[tile * width + j] = acc; +} + +template +__global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + out_t* __restrict__ out, int64_t width, int64_t tiles) { + const int64_t r = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + const int64_t begin = seg_first_tile[r], end = seg_first_tile[r + 1]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= tiles); + for (int64_t tile = begin; tile < end; ++tile) { + acc += partial[tile * width + j]; + } + out[r * width + j] = from_float(acc); +} + +void check_index(const torch::Tensor& t, const char* name, int64_t seq, + const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 1 && t.size(0) == seq, name, " must be (", seq, ",)"); + TORCH_CHECK(t.scalar_type() == at::kLong || t.scalar_type() == at::kInt, name, + " must be int64 or int32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +// rows (R, C * H) with unit column stride -> out (C, S, H) contiguous. +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num) { + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && rows.stride(1) == 1, + "rows must be a 2-D CUDA tensor with unit column stride"); + TORCH_CHECK(chunks > 0 && rows.size(1) % chunks == 0, "rows width ", rows.size(1), + " is not a multiple of ", chunks); + TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", + rows.size(0), " is not a multiple of ", modality_num); + const int64_t num_timesteps = rows.size(0) / modality_num; + const int64_t seq = timestep_indices.numel(); + TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); + check_index(timestep_indices, "timestep_indices", seq, rows); + check_index(token_tags, "token_tags", seq, rows); + TORCH_CHECK(timestep_indices.scalar_type() == token_tags.scalar_type(), + "timestep_indices and token_tags must share an integer dtype"); + const int64_t hidden = rows.size(1) / chunks; + const c10::cuda::CUDAGuard device_guard(rows.device()); + auto out = torch::empty({chunks, seq, hidden}, rows.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(std::min(seq, 1 << 20)), + static_cast(chunks)); + const int threads = 128; + const int64_t elem = rows.element_size(); + const bool vector = (hidden * elem) % 16 == 0 && (rows.stride(0) * elem) % 16 == 0 && + reinterpret_cast(rows.data_ptr()) % 16 == 0; + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, rows.scalar_type(), "h3_adaln_row_gather_forward", [&] { + using T = typename CudaType::type; + const T* src = reinterpret_cast(rows.data_ptr()); + T* dst = reinterpret_cast(out.data_ptr()); + auto launch = [&](auto idx_tag) { + using idx_t = decltype(idx_tag); + const idx_t* ti = timestep_indices.data_ptr(); + const idx_t* tags = token_tags.data_ptr(); + if (vector) { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); + } else { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); + } + }; + if (timestep_indices.scalar_type() == at::kLong) { + launch(int64_t{0}); + } else { + launch(int32_t{0}); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad (C, S, H) contiguous; sorted_pos (S,) positions sorted stably by row; +// tiles [tile_begin, tile_end) in sorted order; seg_first_tile (R + 1,). +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype) { + 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() && + t->scalar_type() == at::kLong, + "tile metadata must be contiguous 1-D int64 tensors on grad device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); + const int64_t chunks = grad.size(0); + const int64_t seq = grad.size(1); + const int64_t hidden = grad.size(2); + const int64_t width = chunks * hidden; + const int64_t num_rows = seg_first_tile.numel() - 1; + const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); + TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); + auto partial = torch::empty({std::max(tiles, 1), width}, + grad.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((width + threads - 1) / threads); + if (tiles > 0) { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { + using G = typename CudaType::type; + adaln_row_gather_partial_kernel + <<(tiles), col_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), + sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), seq, hidden, width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, out_dtype, "h3_adaln_row_gather_fold", [&] { + using O = typename CudaType::type; + adaln_row_gather_fold_kernel + <<(num_rows), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), width, tiles); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu new file mode 100644 index 000000000..96b954b0e --- /dev/null +++ b/csrc/cuda/h3/det_linear.cu @@ -0,0 +1,620 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic, batch-invariant row-wise linear layers for the MiniMax-H3 +// conditioning path (RFC #420: `timestep_mlp_fp32`, `adaln_projection_3mod`). +// +// These layers run on a handful of rows (one per distinct timestep) against +// large weights, so they are GEMVs bound by weight bandwidth. The reduction +// order is fixed and depends only on K (and N for the input gradient), never +// on the number of rows or their position. +// +// Reduction contract h3-det-linear-v1 +// ----------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * one warp per output column n; lane l owns the 16-byte K chunks +// c = l, l + 32, l + 64, ... (VEC = 16 / sizeof(w) elements each); +// * each lane accumulates its chunks in ascending order and the elements of +// a chunk in ascending order with fmaf into an FP32 accumulator from 0; +// * the 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1), +// whose result is identical on every lane (FP add is commutative); +// * the bias is added once after the butterfly; the activation runs in FP32 +// (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// * this FP32 forward is used for FP32 weights; BF16 weights take the +// tensor-core forward of contract h3-det-linear-bf16-mma-v1 below. +// d_input dx[t, k] = sum_n g[t, n] * W[n, k] +// * N is split into fixed chunks of 64 rows; each chunk is an ascending +// fmaf chain into FP32 from 0; the chunk partials are then left-folded in +// ascending chunk order and cast once. +// The partials and the fold are also bound separately: a contiguous N +// slice starting on a chunk boundary yields exactly the full call's +// partials for those chunks, which is what tensor-parallel shards gather. +// d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] +// * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one +// cross-row reduction; its order is the logical row order. +// +// Reduction contract h3-det-linear-bf16-mma-v1 (BF16 x and W, FP32 accumulate) +// ----------------------------------------------------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * a warp owns 16 output columns (MMA rows) and 8 input rows (MMA columns; +// rows past T are zero), so every launch runs the same instruction +// sequence for every T <= 8, and an MMA column never sees another's data; +// * K is visited in groups of 16 in ascending order. Within a 32-wide step, +// quad q supplies k = 32s + 8q + {0..3} to the first mma.sync m16n8k16 and +// k = 32s + 8q + {4..7} to the second (the same k map for W and x); +// * every mma.sync starts from a zero accumulator, so the tensor core only +// sums 16 products; its FP32 result is added to the running FP32 sum with +// an IEEE add, group by group in ascending k; +// * then bias, FP32 activation and one cast, as above. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +#ifndef H3_FWD_F32_COLS +#define H3_FWD_F32_COLS 2 +#define H3_FWD_F32_AHEAD 6 +#endif +#ifndef H3_FWD_CONFIGS +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) +#endif + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 8; +constexpr int kRowTile = 4; +constexpr int kDInputChunk = 64; + +enum Activation : int64_t { kActNone = 0, kActSilu = 1 }; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); // round to nearest even, like at::BFloat16 +} + +// 16-byte vector load converted to FP32. +template +struct Vec16 { + static constexpr int kN = 16 / sizeof(T); + __device__ __forceinline__ static void load(const T* ptr, float (&out)[kN]) { + unpack(load_raw(ptr), out); + } + __device__ __forceinline__ static uint4 load_raw(const T* ptr) { + return *reinterpret_cast(ptr); + } + // Streaming weights are read once: bypass L1 allocation. + __device__ __forceinline__ static uint4 load_stream(const T* ptr) { + return __ldcs(reinterpret_cast(ptr)); + } + __device__ __forceinline__ static void unpack(const uint4& raw, float (&out)[kN]) { + const T* vals = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kN; ++i) out[i] = to_float(vals[i]); + } +}; + +__device__ __forceinline__ float warp_butterfly_sum(float v) { +#pragma unroll + for (int offset = kWarp / 2; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + +// Each warp owns kCols adjacent output columns so that kLoadAhead x kCols +// 16-byte weight loads are in flight per lane. Grouping columns and chunks +// changes only when loads are issued: every column still accumulates its own +// chunks in ascending order, exactly as a one-column-per-warp kernel would. +template +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_kernel(const x_t* __restrict__ x, const w_t* __restrict__ w, + const w_t* __restrict__ bias, out_t* __restrict__ out, + float* __restrict__ pre_act, int64_t rows, int64_t n_out, + int64_t k_in, int64_t activation) { + using WV = Vec16; + using XV = Vec16; + constexpr int kVec = WV::kN; + static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); + static_assert(XV::kN == kVec, "x and weight share a dtype"); + constexpr int64_t kStride = static_cast(kWarp) * kVec; + + const int lane = threadIdx.x % kWarp; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + + for (int64_t n0 = warp_global * kCols; n0 < n_out; n0 += warp_stride * kCols) { + for (int64_t t0 = 0; t0 < rows; t0 += kRows) { + float acc[kCols][kRows]; +#pragma unroll + for (int c = 0; c < kCols; ++c) { +#pragma unroll + for (int r = 0; r < kRows; ++r) acc[c][r] = 0.0f; + } + for (int64_t g0 = static_cast(lane) * kVec; g0 < k_in; + g0 += kStride * kLoadAhead) { + uint4 raw[kLoadAhead][kCols]; +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; +#pragma unroll + for (int c = 0; c < kCols; ++c) { + if (k0 < k_in && n0 + c < n_out) raw[u][c] = WV::load_stream(w + (n0 + c) * k_in + k0); + } + } +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; + if (k0 >= k_in) break; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + if (t0 + r < rows) { + float xv[kVec]; + XV::load(x + (t0 + r) * k_in + k0, xv); +#pragma unroll + for (int c = 0; c < kCols; ++c) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); +#pragma unroll + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } + } + } + } + } +#pragma unroll + for (int c = 0; c < kCols; ++c) { + const int64_t n = n0 + c; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + const int64_t t = t0 + r; + const float sum = warp_butterfly_sum(acc[c][r]); + if (lane == c * kRows + r && t < rows && n < n_out) { + float v = sum + (bias != nullptr ? to_float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = from_float(v); + } + } + } + } + } +} + +__device__ __forceinline__ void mma_m16n8k16_bf16(float (&c)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // BF16 mma.sync needs SM80+; the host refuses to launch below that. + __trap(); +#else + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +#endif +} + +constexpr int kMmaCols = 16; // output columns per warp (MMA M) +constexpr int kMmaRows = 8; // input rows per pass (MMA N) +constexpr int kMmaAhead = 4; // 32-wide K steps loaded before use + +__device__ __forceinline__ void mma_group_accumulate(float (&acc)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { + float part[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_m16n8k16_bf16(part, a, b0, b1); +#pragma unroll + for (int i = 0; i < 4; ++i) acc[i] += part[i]; +} + +// Contract h3-det-linear-bf16-mma-v1. Lane = 4 * g + q: g picks the MMA row +// (output columns n0 + g and n0 + g + 8) and the MMA column (input row t0 + g). +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_bf16_mma_kernel(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ pre_act, + int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation) { + const int lane = threadIdx.x % kWarp; + const int g = lane >> 2; + const int q = lane & 3; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + + for (int64_t n0 = warp_global * kMmaCols; n0 < n_out; n0 += warp_stride * kMmaCols) { + const bool lo_ok = n0 + g < n_out; + const bool hi_ok = n0 + g + 8 < n_out; + const __nv_bfloat16* w_lo = w + (n0 + g) * k_in; + const __nv_bfloat16* w_hi = w + (n0 + g + 8) * k_in; + for (int64_t t0 = 0; t0 < rows; t0 += kMmaRows) { + const bool x_ok = t0 + g < rows; + const __nv_bfloat16* x_row = x + (t0 + g) * k_in; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t kb = 0; kb < k_in; kb += 32 * kMmaAhead) { + uint4 a_lo[kMmaAhead], a_hi[kMmaAhead], bx[kMmaAhead]; +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + const int64_t k = kb + 32 * u + 8 * q; + const bool k_ok = k < k_in; + a_lo[u] = (k_ok && lo_ok) ? __ldcs(reinterpret_cast(w_lo + k)) : zero; + a_hi[u] = (k_ok && hi_ok) ? __ldcs(reinterpret_cast(w_hi + k)) : zero; + bx[u] = (k_ok && x_ok) ? *reinterpret_cast(x_row + k) : zero; + } +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + if (kb + 32 * u >= k_in) break; + const uint32_t first[4] = {a_lo[u].x, a_hi[u].x, a_lo[u].y, a_hi[u].y}; + const uint32_t second[4] = {a_lo[u].z, a_hi[u].z, a_lo[u].w, a_hi[u].w}; + mma_group_accumulate(acc, first, bx[u].x, bx[u].y); + mma_group_accumulate(acc, second, bx[u].z, bx[u].w); + } + } + // acc = {(n0+g, t0+2q), (n0+g, t0+2q+1), (n0+g+8, t0+2q), (n0+g+8, t0+2q+1)} +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int64_t n = n0 + g + (i >= 2 ? 8 : 0); + const int64_t t = t0 + 2 * q + (i & 1); + if (n < n_out && t < rows) { + float v = acc[i] + (bias != nullptr ? __bfloat162float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = __float2bfloat16(v); + } + } + } + } +} + +// partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] +template +__global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, + const w_t* __restrict__ w, + float* __restrict__ partial, int64_t rows, + int64_t n_out, int64_t k_in) { + const int64_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = blockIdx.y; + if (k >= k_in) return; + const int64_t n_begin = chunk * kDInputChunk; + const int64_t n_end = min(n_begin + kDInputChunk, n_out); + for (int64_t t0 = 0; t0 < rows; t0 += kRowTile) { + float acc[kRowTile]; +#pragma unroll + for (int r = 0; r < kRowTile; ++r) acc[r] = 0.0f; + for (int64_t n = n_begin; n < n_end; ++n) { + const float wv = to_float(w[n * k_in + k]); +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) acc[r] = fmaf(to_float(grad[t * n_out + n]), wv, acc[r]); + } + } +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) partial[(chunk * rows + t) * k_in + k] = acc[r]; + } + } +} + +template +__global__ void det_fold_chunks_kernel(const float* __restrict__ partial, out_t* __restrict__ out, + int64_t chunks, int64_t elems) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elems) return; + float acc = partial[i]; + for (int64_t c = 1; c < chunks; ++c) acc += partial[c * elems + i]; + out[i] = from_float(acc); +} + +// dW[n, k] = sum_t g[t, n] * x[t, k] ; ascending t, FP32 from 0, one cast. +template +__global__ void det_linear_dweight_kernel(const g_t* __restrict__ grad, const x_t* __restrict__ x, + w_t* __restrict__ dw, int64_t rows, int64_t n_out, + int64_t k_in) { + const int64_t total = n_out * k_in; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = i / k_in; + const int64_t k = i - n * k_in; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + acc = fmaf(to_float(grad[t * n_out + n]), to_float(x[t * k_in + k]), acc); + } + dw[i] = from_float(acc); + } +} + +template +__global__ void det_linear_dbias_kernel(const g_t* __restrict__ grad, w_t* __restrict__ db, + int64_t rows, int64_t n_out) { + const int64_t n = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= n_out) return; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) acc += to_float(grad[t * n_out + n]); + db[n] = from_float(acc); +} + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; + +template +const typename CudaType::type* cptr(const torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} +template +typename CudaType::type* mptr(torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} + +void check_matrix(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, name, + " must be float32 or bfloat16, got ", t.scalar_type()); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +int64_t grid_for(int64_t work, int64_t per_block) { + const int64_t blocks = (work + per_block - 1) / per_block; + return std::max(1, std::min(blocks, 1 << 20)); +} + +// Launch shape only: the row tile, the columns per warp and the load depth +// change which loads are in flight, never a column's accumulation order. +struct ForwardConfig { + int rows; + int cols; + int ahead; +}; + +ForwardConfig pick_forward_config(int64_t rows) { + const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); + return {row_tile, H3_FWD_F32_COLS, H3_FWD_F32_AHEAD}; +} + +template +void launch_forward_impl(const T* x, const T* w, const T* bias, T* out, float* pre, + int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + const int64_t blocks = grid_for((n_out + kCols - 1) / kCols, kWarpsPerBlock); + det_linear_forward_kernel + <<(blocks), kWarpsPerBlock * kWarp, 0, stream>>>( + x, w, bias, out, pre, rows, n_out, k_in, activation); +} + +template +void launch_forward_rows(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation, cudaStream_t stream) { +#define H3_FWD_CASE(C, A) \ + if (cfg.cols == C && cfg.ahead == A) { \ + launch_forward_impl(x, w, bias, out, pre, rows, n_out, k_in, activation, \ + stream); \ + return; \ + } + H3_FWD_CONFIGS +#undef H3_FWD_CASE + TORCH_CHECK(false, "unsupported det_linear forward config cols=", cfg.cols, + " ahead=", cfg.ahead); +} + +template +void launch_forward(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + if (cfg.rows == 1) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else if (cfg.rows == 2) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } +} + +} // namespace + +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation) { + check_matrix(x, "x"); + check_matrix(weight, "weight"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must be on the same device"); + TORCH_CHECK(x.scalar_type() == weight.scalar_type(), + "x and weight must share a dtype (cast at the declared boundary first), got ", + x.scalar_type(), " and ", weight.scalar_type()); + TORCH_CHECK(activation == kActNone || activation == kActSilu, "unknown activation ", + activation); + const int64_t rows = x.size(0); + const int64_t k_in = x.size(1); + const int64_t n_out = weight.size(0); + TORCH_CHECK(rows > 0, "x must have at least one row"); + TORCH_CHECK(weight.size(1) == k_in, "weight is [", n_out, ", ", weight.size(1), + "] but x has K=", k_in); + const int64_t vec = 16 / x.element_size(); + TORCH_CHECK(k_in % vec == 0, "K=", k_in, " must be a multiple of ", vec, + " for 16-byte vector loads"); + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && + bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), + "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + TORCH_CHECK(bias->device() == weight.device(), + "bias and weight must be on the same device"); + } + + const c10::cuda::CUDAGuard device_guard(x.device()); + auto out = torch::empty({rows, n_out}, x.options()); + torch::Tensor pre; + if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + if (x.scalar_type() == at::kFloat) { + const ForwardConfig cfg = pick_forward_config(rows); + launch_forward(cfg, cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, + k_in, activation, stream); + } else { + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "the BF16 h3 det_linear forward uses mma.sync and needs SM80 or newer"); + const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); + det_linear_forward_bf16_mma_kernel<<(blocks), kWarpsPerBlock * kWarp, + 0, stream>>>( + cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (save_pre_activation) return {out, pre}; + return {out}; +} + +// grad [T, N] (float32), weight [N, K] -> partial [ceil(N / 64), T, K] float32: +// the per-chunk sums of contract h3-det-linear-v1, before the ascending fold. +// A contiguous slice of N that starts on a chunk boundary yields exactly the +// matching rows of the full call (tensor-parallel shards rely on this). +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(weight, "weight"); + TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); + TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = weight.size(0); + const int64_t k_in = weight.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; + auto partial = torch::empty({chunks, rows, k_in}, grad.options()); + if (partial.numel() == 0) return partial; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); + if (weight.scalar_type() == at::kFloat) { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, n_out, + k_in); + } else { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, + n_out, k_in); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// partial [C, T, K] float32 -> left fold over C in ascending order, cast once. +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 3 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat && partial.size(0) > 0, + "partial must be a non-empty contiguous 3-D float32 CUDA tensor"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const c10::cuda::CUDAGuard device_guard(partial.device()); + const int64_t chunks = partial.size(0); + auto out = torch::empty({partial.size(1), partial.size(2)}, partial.options().dtype(out_dtype)); + const int64_t elems = out.numel(); + if (elems == 0) return out; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); + if (out_dtype == at::kFloat) { + det_fold_chunks_kernel<<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } else { + det_fold_chunks_kernel<__nv_bfloat16><<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + auto partial = h3_det_linear_backward_input_partials(grad, weight); + if (partial.size(0) == 0) { // N == 0: no output column contributes + return torch::zeros({partial.size(1), partial.size(2)}, grad.options().dtype(out_dtype)); + } + return h3_det_linear_fold_chunks(partial, out_dtype); +} + +// grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(x, "x"); + TORCH_CHECK(grad.device() == x.device(), "grad and x must be on the same device"); + TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); + TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, + "w_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = grad.size(1); + const int64_t k_in = x.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + if (n_out == 0) { + if (!with_bias) return {dw}; + return {dw, torch::empty({n_out}, grad.options().dtype(w_dtype))}; + } + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const int64_t total = n_out * k_in; + const unsigned blocks = static_cast(grid_for(total, threads)); + const float* g = grad.data_ptr(); +#define H3_DW_LAUNCH(XT, XC, WT, WC) \ + det_linear_dweight_kernel<<>>( \ + g, cptr(x), mptr(dw), rows, n_out, k_in) + if (x.scalar_type() == at::kFloat && w_dtype == at::kFloat) { + H3_DW_LAUNCH(float, float, float, float); + } else if (x.scalar_type() == at::kFloat) { + H3_DW_LAUNCH(float, float, at::BFloat16, __nv_bfloat16); + } else if (w_dtype == at::kFloat) { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, float, float); + } else { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, at::BFloat16, __nv_bfloat16); + } +#undef H3_DW_LAUNCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!with_bias) return {dw}; + auto db = torch::empty({n_out}, grad.options().dtype(w_dtype)); + const unsigned bias_blocks = static_cast((n_out + threads - 1) / threads); + if (w_dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } else { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dw, db}; +} diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..eb0b03f50 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 gated residual (RFC #420 `adaln_gate_residual`). +// +// out[r, j] = cast(residual[r, j] + cast(gate[i, j] * y[r, j])) i = index[r % S] +// +// The gate row is read from a strided view of the AdaLN table (gate_msa or +// gate_mlp) inside the kernel, and the two roundings sit exactly where the +// eager `residual + gate.index_select(0, i) * y` rounds (__fmul_rn/__fadd_rn +// keep FP32 from contracting them into one FMA), so the result is +// bitwise equal to diffusers. Elementwise: every output depends only on its +// own inputs. +// +// Backward: d_residual = grad; d_y = cast(grad * gate) (the product of two +// 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 +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +struct Gate { + const void* rows; // (R, N) view, row stride `stride` + int64_t stride; + const int64_t* index; // (S,) + int64_t seq; + int64_t num_rows; +}; + +// One block per row (grid-strided); 16-byte vectors when the row allows it. +// `kVecElems` elements are loaded together; each is computed independently. +template +__global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t rows, int64_t n, + Gate gate, bool vectorized) { + constexpr int kVecElems = 16 / sizeof(T); + for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; + const T* ar = a + r * n; + const T* br = kDy ? nullptr : b + r * n; + T* orow = out + r * n; + if (vectorized) { + for (int64_t v = threadIdx.x; v < n / kVecElems; v += blockDim.x) { + const uint4 av = reinterpret_cast(ar)[v]; + const uint4 gv = reinterpret_cast(g)[v]; + uint4 bv = av; + if (!kDy) bv = reinterpret_cast(br)[v]; + const T* ae = reinterpret_cast(&av); + const T* ge = reinterpret_cast(&gv); + const T* be = reinterpret_cast(&bv); + uint4 ov; + T* oe = reinterpret_cast(&ov); +#pragma unroll + for (int k = 0; k < kVecElems; ++k) { + if (kDy) { + oe[k] = from_f(to_f(ae[k]) * to_f(ge[k])); // grad * gate + } else { + const float p = round_t(__fmul_rn(to_f(ge[k]), to_f(be[k]))); // gate * y + oe[k] = from_f(__fadd_rn(to_f(ae[k]), p)); // residual + p + } + } + reinterpret_cast(orow)[v] = ov; + } + } else { + for (int64_t j = threadIdx.x; j < n; j += blockDim.x) { + if (kDy) { + orow[j] = from_f(to_f(ar[j]) * to_f(g[j])); + } else { + const float p = round_t(__fmul_rn(to_f(g[j]), to_f(br[j]))); + orow[j] = from_f(__fadd_rn(to_f(ar[j]), p)); + } + } + } + } +} + +// partial[tile, j] = sum over sorted positions of the tile (ascending) of grad * y +template +__global__ void gate_grad_partial_kernel(const T* __restrict__ grad, const T* __restrict__ y, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t n) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t e = sorted_pos[p] * n + j; + acc = fmaf(to_f(grad[e]), to_f(y[e]), acc); + } + partial[tile * n + j] = acc; +} + +template +__global__ void gate_grad_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + T* __restrict__ out, int64_t n) { + const int64_t seg = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * n + j]; + out[seg * n + j] = from_f(acc); +} + +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torch::Tensor& x) { + TORCH_CHECK(gate.is_cuda() && gate.device() == x.device() && gate.scalar_type() == x.scalar_type(), + "gate must be a CUDA tensor with the activations' dtype"); + TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, + "gate must be (R, N) with unit column stride"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && + index.dim() == 1 && index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 tensor on ", x.device()); + TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; +} + +void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device() && t.scalar_type() == like.scalar_type(), + name, " must match the residual's device and dtype"); + TORCH_CHECK(t.sizes() == like.sizes() && t.is_contiguous(), name, + " must be contiguous with the residual's shape"); +} + +unsigned row_blocks(int64_t rows) { + return static_cast(std::min(rows, 1 << 30)); +} + +bool can_vectorize(const torch::Tensor& gate, int64_t n, std::initializer_list acts) { + const int64_t vec = 16 / gate.element_size(); + if (n % vec != 0 || (gate.stride(0) * gate.element_size()) % 16 != 0) return false; + if (reinterpret_cast(gate.data_ptr()) % 16 != 0) return false; + for (const auto* t : acts) { + if (reinterpret_cast(t->data_ptr()) % 16 != 0) return false; + } + return true; +} + +} // namespace + +// residual, y: (M, N); gate: (R, N) view; index: (S,), M % S == 0. +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && residual.is_contiguous() && + residual.numel() > 0, + "residual must be a non-empty contiguous (M, N) CUDA tensor"); + check_act(y, "y", residual); + const Gate g = make_gate(gate, index, residual); + const c10::cuda::CUDAGuard guard(residual.device()); + auto out = torch::empty_like(residual); + auto stream = at::cuda::getCurrentCUDAStream(); + const bool vec = can_vectorize(gate, residual.size(1), {&residual, &y, &out}); + H3_DISPATCH(residual.scalar_type(), "h3_gate_residual_forward", [&] { + gate_residual_rows_kernel<<>>( + reinterpret_cast(residual.data_ptr()), reinterpret_cast(y.data_ptr()), + reinterpret_cast(out.data_ptr()), residual.size(0), residual.size(1), g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Returns {d_y, d_gate (R, N) in the gate dtype}. d_residual is the incoming grad. +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile) { + 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); + const Gate g = make_gate(gate, index, grad); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + 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 tensors on ", grad.device()); + } + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dy = torch::empty_like(y); + auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); + auto dgate = torch::empty({segments, n}, grad.options()); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { + const T* gp = reinterpret_cast(grad.data_ptr()); + const T* yp = reinterpret_cast(y.data_ptr()); + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + gp, nullptr, reinterpret_cast(dy.data_ptr()), rows, n, g, vec); + if (tiles > 0) { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, stream>>>( + gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), n); + } + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + }); + 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 unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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"); + 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)); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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 new file mode 100644 index 000000000..f9304ab13 --- /dev/null +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -0,0 +1,593 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 RMSNorm and fused AdaLN modulation (RFC #420 `h3_rmsnorm`). +// +// n[r] = cast(w * (rstd[r] * x[r])) rstd = rsqrtf(sum(x^2) / N + eps) +// out[r] = cast(cast(n[r] * cast(1 + scale[i])) + shift[i]) i = row_index[r % S] +// +// Forward contract h3-rmsnorm-v1: the statistics replay PyTorch's +// vectorized_layer_norm_kernel (torch cf30153): one +// (32, 4) block per row, 4-element vectors, thread t sums vectors t, t+128, ... +// in order, a shuffle-down tree (16..1), a cross-warp tree, sum / N, rsqrtf, +// then w * (rstd * x) and one cast. nn.RMSNorm therefore matches bitwise, and +// each modulation step rounds to the tensor dtype exactly where the eager +// `n * (1.0 + scale) + shift` expression does. Rows are independent: the result +// does not depend on batch size, position or the other rows. +// +// Backward (FP32, one cast per output, no atomics): +// d_n = g * cast(1 + scale) (g without modulation) +// dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) +// 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 +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kVec = 4; // PyTorch's layer-norm vec_size +constexpr int kWarps = 4; // num_threads() / warp_size = 128 / 32 +constexpr int kRowTile = 256; // dweight row tile + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +// Round to T and back: the value an eager T-dtype op would produce. +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +// PyTorch's rms statistics: sequential per thread, shuffle-down, cross-warp tree. +// Returns sum(v^2) / N on every thread. +template +__device__ __forceinline__ float row_mean_square(const T* __restrict__ row, int n, float* buf) { + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const float v = to_f(row[i * kVec + k]); + s = s + v * v; + } + } + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + if (thrx == 0) buf[0] = s / float(n); + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +// Same tree for an arbitrary per-element FP32 term; returns the row sum. +__device__ __forceinline__ float row_sum_tree(float s, float* buf) { + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + if (thrx == 0) buf[0] = s; + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +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] +}; + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, + float* __restrict__ rstd_out, int n, float eps, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const float rstd = rsqrtf(row_mean_square(xr, n, buf) + eps); + const T* shift = nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + shift = static_cast(mod.shift) + i * mod.row_stride; + scale = static_cast(mod.scale) + i * mod.row_stride; + } + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + float v = round_t(__fmul_rn(to_f(w[j]), __fmul_rn(rstd, to_f(xr[j])))); + if (shift != nullptr) { + // Separately rounded ops (no FMA contraction), as the eager expression. + const float t1 = round_t(__fadd_rn(1.0f, to_f(scale[j]))); + v = round_t(__fadd_rn(round_t(__fmul_rn(v, t1)), to_f(shift[j]))); + } + y[r * n + j] = from_f(v); + } + } + if (thrx == 0) rstd_out[r] = rstd; +} + +template +__device__ __forceinline__ float d_norm_out(const T* g_row, const T* scale, int j) { + const float g = to_f(g_row[j]); + return scale == nullptr ? g : g * round_t(1.0f + to_f(scale[j])); +} + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_dx_kernel(const T* __restrict__ g, const T* __restrict__ x, + const T* __restrict__ w, const float* __restrict__ rstd, + T* __restrict__ dx, int n, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const T* gr = g + r * n; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + s = fmaf(to_f(w[j]) * d_norm_out(gr, scale, j), to_f(xr[j]), s); + } + } + const float dot = row_sum_tree(s, buf); + const float rs = rstd[r]; + const float coef = rs * rs * rs * dot / float(n); + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + dx[r * n + j] = from_f(rs * to_f(w[j]) * d_norm_out(gr, scale, j) - to_f(xr[j]) * coef); + } + } +} + +// 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, + 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 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 p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } + acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); + } + partial[static_cast(blockIdx.x) * n + j] = acc; +} + +template +__global__ void fold_tiles_kernel(const float* __restrict__ partial, T* __restrict__ out, + int64_t tiles, int n) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = 0; t < tiles; ++t) acc += partial[t * n + j]; + out[j] = from_f(acc); +} + +// Table gradient tiles: positions sorted by (index, position); columns j < n +// accumulate g (shift), columns j >= n accumulate g * n_out (scale). +template +__global__ void modulation_grad_partial_kernel( + const T* __restrict__ g, const T* __restrict__ x, const T* __restrict__ w, + const float* __restrict__ rstd, const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int n) { + const int64_t tile = blockIdx.x; + const int j2 = blockIdx.y * blockDim.x + threadIdx.x; + if (j2 >= 2 * n) return; + const bool is_scale = j2 >= n; + const int j = is_scale ? j2 - n : j2; + const float wj = to_f(w[j]); + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t r = sorted_pos[p]; + const float gv = to_f(g[r * n + j]); + if (is_scale) { + const float nv = round_t(wj * (rstd[r] * to_f(x[r * n + j]))); + acc = fmaf(gv, nv, acc); + } else { + acc += gv; + } + } + partial[tile * 2 * n + j2] = acc; +} + +__global__ void fold_segments_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + float* __restrict__ out, int width) { + const int64_t seg = blockIdx.x; + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * width + j]; + out[seg * width + j] = acc; +} + +void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 2 && t.is_contiguous(), name, " must be a contiguous 2-D tensor"); + TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); +} + +Modulation make_modulation(const c10::optional& shift, + const c10::optional& scale, + const c10::optional& index, const torch::Tensor& x, + int64_t n) { + 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; + } + TORCH_CHECK(shift.has_value() && scale.has_value(), "modulation needs both shift and scale"); + const auto& sh = *shift; + const auto& sc = *scale; + const auto& ix = *index; + for (const auto* t : {&sh, &sc}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == x.scalar_type(), + "shift/scale must be CUDA tensors with x's dtype"); + TORCH_CHECK(t->dim() == 2 && t->size(1) == n && t->stride(1) == 1, + "shift/scale must be (R, N) with unit column stride"); + } + TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), + "shift and scale must be views of the same table layout"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous() && ix.numel() > 0, + "row index must be a non-empty contiguous int64 tensor on ", x.device()); + TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", + ix.size(0), ")"); + 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); + return mod; +} + +void check_xw(const torch::Tensor& x, const torch::Tensor& w) { + TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); + TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && + w.size(0) == x.size(1) && w.is_contiguous() && + w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype and device"); + TORCH_CHECK(x.size(1) > 0, "x must have at least one column"); + TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, + " (PyTorch's vectorized RMSNorm path)"); +} + +// float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +} // namespace + +// x (M, N); optional shift/scale (R, N) strided views and index (S,), M % S == 0. +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + TORCH_CHECK(x.scalar_type() != at::kDouble, "float64 is not supported"); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto y = torch::empty_like(x); + auto rstd = torch::empty({x.size(0)}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_forward", [&] { + rmsnorm_modulate_fwd_kernel<<(x.size(0)), dim3(32, kWarps), 0, stream>>>( + reinterpret_cast(x.data_ptr()), reinterpret_cast(weight.data_ptr()), + reinterpret_cast(y.data_ptr()), rstd.data_ptr(), static_cast(n), + static_cast(eps), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, rstd}; +} + +// Returns {dx, dweight} and, with modulation, {d_shift, d_scale} as FP32 (R, N). +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile) { + 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 int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + const std::pair tiles[] = {{&*sorted_pos, "sorted_pos"}, + {&*tile_begin, "tile_begin"}, + {&*tile_end, "tile_end"}, + {&*seg_first_tile, "seg_first_tile"}}; + for (const auto& [t, name] : tiles) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + name, ": tile metadata must be contiguous int64 tensors on ", x.device()); + } + TORCH_CHECK(sorted_pos->size(0) == rows, "sorted_pos must have M entries"); + TORCH_CHECK(seg_first_tile->size(0) == shift->size(0) + 1, + "seg_first_tile must have R + 1 entries"); + TORCH_CHECK(tile_begin->size(0) == tile_end->size(0), + "tile_begin and tile_end must have the same number of entries"); + } else { + TORCH_CHECK(!sorted_pos.has_value() && !tile_begin.has_value() && !tile_end.has_value() && + !seg_first_tile.has_value(), + "sorted segment tiles require modulation"); + } + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dx = torch::empty_like(x); + const int64_t tiles = (rows + kRowTile - 1) / kRowTile; + auto partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + auto dweight = torch::empty_like(weight); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward", [&] { + const T* g = reinterpret_cast(grad.data_ptr()); + const T* xp = reinterpret_cast(x.data_ptr()); + const T* wp = reinterpret_cast(weight.data_ptr()); + rmsnorm_modulate_dx_kernel<<(rows), dim3(32, kWarps), 0, stream>>>( + g, xp, wp, rstd.data_ptr(), reinterpret_cast(dx.data_ptr()), + static_cast(n), mod); + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, 0, + stream>>>(g, xp, rstd.data_ptr(), + partial.data_ptr(), rows, + static_cast(n), mod); + fold_tiles_kernel<<>>( + partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), tiles, + static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (mod.index == nullptr) return {dx, dweight}; + + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t seg_tiles = tile_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); + auto table_grad = torch::empty({segments, 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_grad", [&] { + 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(), + sorted_pos->data_ptr(), tile_begin->data_ptr(), + tile_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + 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 {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.numel() == x.size(0), + "rstd must be the forward's 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(); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + 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; + 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"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + 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.numel() == x.size(0), + "rstd must be the forward's 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/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu new file mode 100644 index 000000000..2acfd516c --- /dev/null +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 sinusoidal timestep features (RFC #420 `timestep_sinusoid_h3`). +// +// freq[k] = expf((c * k) * (1 / half)) c = (float)(-ln(max_period)) +// arg[t, k] = t[t] * freq[k] +// out[t] = [cosf(arg) | sinf(arg)] (T, 2 * half) FP32 +// +// Every value is produced by the same FP32 operation sequence as diffusers' +// get_timestep_embedding on CUDA (scalar * arange, multiply by the reciprocal +// of the CPU-scalar divisor, exp, multiply, cos/sin), so the output is meant +// to be bitwise equal to that path. Each element depends only on (t, k): +// batch-, position- and repeat-invariant by construction. Build without +// --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. +// nvcc defines no macro for fast math, so build_tools/extensions.py lists this +// file in H3_BITWISE_SOURCES, which KERNEL_ALIGN_USE_FAST_MATH=1 builds leave +// out; the op then falls back to the PyTorch reference. + +#include +#include +#include +#include + +#include +#include + +namespace { + +__global__ void h3_timestep_sinusoid_kernel( + const float* __restrict__ timestep, + float* __restrict__ out, + int64_t num_timesteps, + int64_t half, + float neg_log_max_period, + float inv_half) { + const int64_t total = num_timesteps * half; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t row = idx / half; + const int64_t k = idx - row * half; + // __fmul_rn keeps each product a separately rounded FP32 multiply (no FMA + // contraction), matching the two elementwise torch kernels it replays. + const float exponent = __fmul_rn(__fmul_rn(neg_log_max_period, static_cast(k)), inv_half); + const float freq = expf(exponent); + const float arg = __fmul_rn(timestep[row], freq); + float* out_row = out + row * 2 * half; + out_row[k] = cosf(arg); + out_row[half + k] = sinf(arg); + } +} + +} // namespace + +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period, bool check_range) { + TORCH_CHECK(timestep.is_cuda(), "timestep must be a CUDA tensor"); + TORCH_CHECK(timestep.scalar_type() == at::kFloat, "timestep must be float32, got ", + timestep.scalar_type()); + TORCH_CHECK(timestep.dim() == 1, "timestep must be 1-D (num_timesteps,)"); + TORCH_CHECK(timestep.numel() > 0, "timestep must hold at least one timestep"); + TORCH_CHECK(num_channels > 0 && num_channels % 2 == 0, + "num_channels must be a positive even number, got ", num_channels); + TORCH_CHECK(max_period > 0.0, "max_period must be positive"); + + const c10::cuda::CUDAGuard device_guard(timestep.device()); + auto t = timestep.contiguous(); + // Check at the native boundary, including direct extension calls. The explicit + // opt-out is for already-validated inputs and kernel-only profiling. + if (check_range) { + TORCH_CHECK_VALUE(((t >= 0) & (t <= 1)).all().item(), + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + "unscaled"); + } + const int64_t num_timesteps = t.size(0); + const int64_t half = num_channels / 2; + auto out = torch::empty({num_timesteps, num_channels}, t.options()); + + // torch evaluates `-math.log(max_period) * arange` with the Python double + // rounded to the FP32 opmath type, and `x / half` as x * (1 / half) in FP32. + const float neg_log_max_period = static_cast(-std::log(max_period)); + const float inv_half = 1.0f / static_cast(half); + + const int threads = 256; + const int64_t total = num_timesteps * half; + const int64_t blocks = std::min((total + threads - 1) / threads, 65535); + auto stream = at::cuda::getCurrentCUDAStream(); + h3_timestep_sinusoid_kernel<<(blocks), threads, 0, stream>>>( + t.data_ptr(), out.data_ptr(), num_timesteps, half, neg_log_max_period, + inv_half); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 1037edce5..05e5aba93 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -30,6 +30,16 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - MiniMax-H3: + - operators/h3-timestep-sinusoid.md + - operators/h3-timestep-mlp.md + - operators/h3-adaln-projection.md + - operators/h3-adaln-row-gather.md + - operators/h3-rmsnorm.md + - 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 00f4cbb45..c99451553 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,4 +32,14 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) +- [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) +- [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) +- [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) +- [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) +- [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) +- [MiniMax-H3 One Block, Forward and Backward](h3-one-block.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..9a5d815a5 --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,143 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python tools/validation/models/h3_prior_art.py --op gate_residual --out reports/experiments/h3-prior-art-b200/gate_residual.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python tools/validation/operators/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_gate_residual \ + --out reports/experiments/h3-adaln-gate-residual-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md new file mode 100644 index 000000000..e8fc09537 --- /dev/null +++ b/docs/operators/h3-adaln-projection.md @@ -0,0 +1,207 @@ +# MiniMax-H3 Three-Modality AdaLN Projection + +## Summary + +`adaln_projection_3mod` is the per-block `MiniMaxH3AdaLayerNormModulation` +(RFC #420, WS1 step 3). It turns the shared FP32 timestep embedding into the six +modulation tensors of one transformer block, for all three modalities: + +```text +act = silu(temb).to(bf16) SiLU in FP32, one declared cast +table = act @ W.T + b (T, 96768) BF16, W: (96768, 2688) +rows = table.view(3T, 32256) row t * 3 + m, m: 0 video, 1 text, 2 audio +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) # (3T, 5376) each +``` + +Projection output channel `o = m * 32256 + c * 5376 + h` is chunk `c`, hidden index `h` +of modality `m`. The table rows are laid out `[t0m0, t0m1, t0m2, t1m0, ...]`. That is the +layout `adaln_row_gather` addresses with `timestep_indices * 3 + token_tags`. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`, `transformer_blocks.0.adaln_proj.linear` +(BF16 weight and bias). All 50 blocks have the same shape. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_projection_3mod", device="cuda") +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, weight, bias) +table = op.forward_table(temb, weight, bias) # raw (T, 96768) for a fused gather +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp` | `rl_engine._C.h3_det_linear_*` | BF16 weights: contract `h3-det-linear-bf16-mma-v1`; FP32 weights: `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp` | n/a | `forward`: provider path; `forward_fp32`: declared-cast FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `temb` | `(T, D)`, `T >= 1` | float32 | H3: `D = 2688`. BF16 is rejected (RFC probe H7) | +| `weight` | `(6 * H * 3, D)` | bfloat16 (checkpoint) or float32 | First dim must be a multiple of 18. H3: `H = 5376` | +| `bias` | `(6 * H * 3,)` | same as `weight` | | +| outputs | six `(3T, H)` views | weight dtype | Views of one `(T, 18H)` table, as in diffusers | + +The CUDA kernel needs `D` to be a multiple of 8 for BF16 or 4 for FP32, and BF16 needs SM80+. These inputs fail +closed: a BF16 `temb`, mismatched dtypes or shapes, an empty `temb`, and mixed devices. + +## Numerics + +- **Mixed-precision boundary.** SiLU runs in FP32 at `temb`'s precision, and the result is + rounded to BF16 exactly once. These are the provider's own elementwise ops, so the + activation is bitwise equal to diffusers. Casting before the SiLU (probe H7) is + rejected at the API. It is also numerically visible: on block 0 it changes about 52% of + the BF16 outputs. +- **Projection (BF16 weights).** It follows contract `h3-det-linear-bf16-mma-v1` in + `csrc/cuda/h3/det_linear.cu`: + - Each warp owns 16 output columns and 8 input rows. Rows past T are zero, so every + T <= 8 runs the same instruction sequence. A tensor-core column never sees another + column's data. + - K is visited in groups of 16 in ascending order. Each `mma.sync m16n8k16` (BF16 in, + FP32 out) starts from zero, so the tensor core only sums 16 products. Its result is + added to an FP32 running sum with an IEEE add. + - Then the bias is added and the result is rounded to BF16 once. + + A timestep's 18 modulation rows are bitwise independent of the other timesteps in the + call and of their order. Restarting the accumulator for every 16-wide group is what keeps + accuracy at the level of an FP32 FMA chain. Letting the tensor core accumulate across + the whole of K, as cuBLAS does, loses about 0.15% of correctly rounded outputs. FP32 + weights use the FMA path of `h3-det-linear-v1` instead, because FP32 on tensor cores + would mean TF32, which the contract forbids. +- **Golden.** `forward_fp32` computes the SiLU in FP64, rounds it to BF16 at the declared + boundary (the cast is model semantics), and runs the projection in FP64. It returns + FP32 without the final rounding. Its gradient stays in FP64 end to end, because the + cast is applied straight-through (identity VJP). Writing the cast as + `.to(bf16).double()` would make autograd round the golden's own `d_temb` to BF16. +- **Backward.** No cuBLAS and no atomics. `d_act` is computed with the chunked + deterministic `grad @ W`, kept in FP32 through the cast (whose VJP is the identity), + and passed through the FP32 SiLU VJP; it is row-local and batch-invariant. + `dW` and `db` are ascending-row FP32 folds, rounded to BF16 once. Unlike the provider, + `d_temb` is not rounded to BF16 on its way through the cast. + +Measured on a B200 with the pinned block-0 weights and torch 2.13.0+cu130. Fractions are +medians over 20 random draws of `temb`, with T = 4: + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| outputs equal to the correctly rounded golden | 99.982% (worst draw 99.976%) | 99.842% (worst draw 99.787%) | +| CUDA equal to provider | 99.84% | | +| early-cast golden (probe H7) equal to CUDA | 43% | | +| `d_temb` vs FP64 golden (gtest, T = 3) | max abs 2.3e-5 | 6.1e-2 (`d_act` rounded to BF16) | +| `dW`, `db` vs FP64 golden rounded to BF16 | > 99.9% bitwise equal | | + +The contract tolerance is `reduction` / `bfloat16`: atol 5e-2, rtol 2e-2, and atol 1e-1 +for gradients. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_projection_3mod +``` + +B200, pinned BF16 weights (520 MB per block, streamed once per call): + +| T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | +| --- | --- | --- | --- | --- | +| 1 | 100 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 100 µs | 100 µs | 78.4 µs | 79 µs | +| 4 | 101 µs | 101 µs | 79.0 µs | 79 µs | + +The kernel takes the same time for every T <= 8, because it always computes 8 rows. It +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 22 µs of the op time is +the SiLU and cast kernels plus the Python wrapper. With those included, the op is within +about 1% of the provider for T >= 2. + +An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time +grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA +(`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the +dependence on T. + +## Evidence + +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 +SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also +records that a timestep's 18 modulation rows are bitwise identical whether it runs alone +or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) + +The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` +streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional +weight passes and are outside its intended T ≤ 8 range. This accounts for the +steep increase in the reported CUDA forward + backward times (77.4 and 627 ms). + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | +| rl-kernel H3AdaLNProjectionCudaOp | yes | 3.6e-03 / 3.9e-03 / 2276 µs | 3.4e-03 / 3.2e-03 / 77446 µs | 3.2e-03 / 3.5e-03 / 626992 µs | +| diffusers AdaLN projection, BF16 F.linear [vllm] | yes | 3.6e-03 / 3.9e-03 / 681 µs | 3.4e-03 / 3.2e-03 / 722 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang] | yes | 3.6e-03 / 3.9e-03 / 3096 µs | 3.4e-03 / 3.2e-03 / 3690 µs | 3.2e-03 / 3.5e-03 / 14296 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang_ieee] | yes | 3.6e-03 / 3.9e-03 / 3138 µs | 3.4e-03 / 3.2e-03 / 3740 µs | 3.2e-03 / 3.5e-03 / 14325 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_te_native] | **no** (7107 rows, 2056 sub-batches, 1104 sweep cases) | 3.6e-03 / 3.9e-03 / 702 µs | 3.4e-03 / 3.2e-03 / 706 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton] | yes | 3.6e-03 / 3.9e-03 / 2532 µs | 3.4e-03 / 3.2e-03 / 3204 µs | 3.2e-03 / 3.5e-03 / 13771 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton_ieee] | yes | 3.6e-03 / 3.9e-03 / 2533 µs | 3.4e-03 / 3.2e-03 / 3206 µs | 3.2e-03 / 3.5e-03 / 13773 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) +was written from a clean tree at `c88691c` by + +```bash +python tools/validation/models/h3_prior_art.py --op adaln_projection --out reports/experiments/h3-prior-art-b200/adaln_projection.json --megatron-src +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/adaln_projection.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) + +This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_projection.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP -> projection +python tools/validation/operators/check_operator.py --op adaln_projection_3mod --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_projection_3mod \ + --out reports/experiments/h3-adaln-projection-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-projection-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- The BF16 path needs SM80 or newer (`mma.sync`). It is rejected at launch on older GPUs. +- T > 8 runs one more weight pass per 8 rows. +- Sharding the 96768-wide projection (`tp_adaln_3mod`) is a separate WS2 row. diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md new file mode 100644 index 000000000..94e32b711 --- /dev/null +++ b/docs/operators/h3-adaln-row-gather.md @@ -0,0 +1,204 @@ +# MiniMax-H3 AdaLN Row Gather + +## Summary + +`adaln_row_gather` selects every packed-sequence position's six modulation vectors from +one block's AdaLN table (RFC #420, WS1 step 3): + +```text +adaln_indices = timestep_indices * 3 + token_tags (S,) 0 video, 1 text, 2 audio +shift_msa[adaln_indices], scale_msa[...], gate_msa[...], +shift_mlp[...], scale_mlp[...], gate_mlp[...] six (S, 5376) +``` + +`rows` is the `(3T, 6H)` view of the [AdaLN projection](h3-adaln-projection.md) table, and +its six column blocks are the six modulation tensors. `timestep_indices` and `token_tags` +are semantic inputs (RFC #420 §4): they are validated, never clamped. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_row_gather", device="cuda") +outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) +outs = op.gather_chunks((shift_msa, ..., gate_mlp), timestep_indices, token_tags) # drop-in +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather.H3AdaLNRowGatherCudaOp` | `rl_engine._C.h3_adaln_row_gather_{forward,backward}` | Forward bitwise equal to `index_select`; deterministic backward | +| PyTorch reference | `rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp` | n/a | `index_select` forward; deterministic per-row FP32 backward | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +The raw diffusers path, including its atomic backward, is kept as +`rl_engine.validation.models.h3_provider.provider_adaln_row_gather` for evidence. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `rows` | `(3T, 6H)` | bf16 / fp16 / fp32 | Unit column stride. Any row stride (16-byte copies when aligned) | +| `timestep_indices` | `(S,)`, `S >= 1` | int64 or int32 | In `[0, T)` | +| `token_tags` | `(S,)` | same as `timestep_indices` | In `{0, 1, 2}` | +| outputs | six `(S, H)` | `rows.dtype` | Contiguous; slices of one `(6, S, H)` buffer | + +These inputs fail closed: out-of-range tags (a fourth modality, probe H2) or timesteps +(an offset past `T`, probe H3), length or dtype mismatches, empty sequences, and tables +that are not 3 rows per timestep or 6 blocks wide. The range check is one host +read-back; pass `check_range=False` when the indices are already validated. + +## Numerics + +- **Forward** is a byte copy: one launch, one block per `(position, chunk)`, using + 16-byte vector copies. It is bitwise equal to diffusers' six `index_select` calls for + every packing, length and index dtype tested (S up to 131072). +- **Backward**: `d_rows[r] = sum over {s : r[s] = r} of grad[s]` is the op's only + reduction. Positions are sorted stably by `(row, position)` and cut into tiles of 256. + Each tile is an ascending FP32 chain from 0, and a row's tiles are left-folded in order + and rounded once. There are no atomics, so the result depends only on the indices and + the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient + dtype (BF16), which is neither deterministic nor accurate. +- **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as + `reduction`. The forward is asserted bitwise in `tests/models/minimax_h3`. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison (T = 3, S = 4097, H = 5376, BF16) | CUDA | provider (`index_select`) | +| --- | --- | --- | +| Forward vs `index_select` | bitwise equal | | +| `d_rows` repeat-bitwise | yes | no | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.9% | +| `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | + +### Fused modulation: `H3AdaLNModulationCudaOp` + +`rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation.H3AdaLNModulationCudaOp(temb, weight, bias, +timestep_indices, token_tags)` runs [`adaln_projection_3mod`](h3-adaln-projection.md) and +this gather as one autograd node. Its forward uses the same kernels and is bitwise equal +to calling the two ops in sequence. + +The difference is in the backward. When the ops are called separately, autograd returns +the gather's FP32 segment sum to the projection in the table's BF16 dtype, which rounds it +once more. The fused node feeds the FP32 table gradient straight into the projection +backward. Use it wherever both ops run back to back, as they do in an H3 block. + +### Whole-chain backward + +`tools/validation/models/h3_chain_replay.py --backward` runs on the pinned weights for +T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. It checks six parameter gradients in each +of the 16 cases (the time embedder's `linear_{1,2}.{weight,bias}` and block 0's AdaLN +weight and bias) against an FP64 golden in which every declared cast is straight-through: + +| S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | +| --- | --- | --- | --- | +| repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 3.8e-3 to 1.3e-1, varying run to run | +| AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | + +At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the +time embedder, against 3.2e-3 for diffusers. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_row_gather +``` + +B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table stays in L2. + +| S | CUDA forward | provider forward | CUDA backward | provider backward | +| --- | --- | --- | --- | --- | +| 4097 | 0.07 ms | 0.11 ms | 0.88 ms | 1.28 ms | +| 32768 | 0.33 ms | 0.62 ms | 1.75 ms | 8.04 ms | +| 131072 | 1.17 ms | 2.42 ms | 5.62 ms | 32.45 ms | + +Backward timings and peak memory exclude leaf creation and the forward pass. The CUDA +backward stacks the six output gradients and materialises per-tile FP32 partials, adding +8130 MiB at S = 131072; the provider adds under 2 MiB. Candidate/provider execution order +alternates each iteration and is recorded in the report. + +## Evidence + +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) + +There are two data files: + +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward + bitwise checks, op-level backward, and the chain-backward cases plotted above. +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full + stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. + +`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +timings and alternating execution order, and with the op-backward check's rows and gradients +drawn from a seeded generator. Before, they came from the process's default CUDA generator, +so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). +`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the +model does; every CUDA pass/fail flag is unchanged. + +## Existing implementations (RFC #420 reuse rule) + +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers six index_select calls | **no** (param/table grads not repeatable) | 0.0e+00 / 6.0e-02 / 1381 µs | 0.0e+00 / 2.1e-01 / 9753 µs | +| rl-kernel H3AdaLNRowGatherCudaOp | yes | 0.0e+00 / 2.5e-03 / 1106 µs | 0.0e+00 / 3.0e-03 / 2776 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +was written from a clean tree at `6c900ae` by + +```bash +python tools/validation/models/h3_prior_art.py --op adaln_row_gather --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/adaln_row_gather.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) + +There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_row_gather.py tests/models/minimax_h3/test_h3_adaln_modulation.py -v +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # whole chain, forward + backward +python tools/validation/operators/check_operator.py --op adaln_row_gather --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_row_gather \ + --out reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/h3_chain_replay.py --timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 \ + --backward --out reports/experiments/h3-adaln-row-gather-b200/chain_replay.json +``` + +## Known Limitations + +- Like diffusers, the op materialises six `(S, H)` tensors, which is 8 GB at + S = 131072. Fusing the gather into the norm/modulate and gate-residual consumers + (`h3_rmsnorm`, `adaln_gate_residual`) avoids this. Those are separate rows. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md new file mode 100644 index 000000000..d40a49f5f --- /dev/null +++ b/docs/operators/h3-final-adaln-out.md @@ -0,0 +1,175 @@ +# MiniMax-H3 Final AdaLN Output + +## Summary + +`final_adaln_out` is H3's `norm_out` (`MiniMaxH3AdaLayerNormOut`), which runs once after the 50 +blocks (RFC #420, WS1 step 7): + +```text +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] +``` + +The table has one row per distinct timestep and is indexed by `timestep_indices`, not by +`adaln_indices`. Diffusers then upcasts `out` to the FP32 output heads' dtype. That cast is exact +and is left to the caller. + +## Entry Point + +```python +op = kernel_registry.get_op("final_adaln_out", device="cuda") +out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timestep_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+) | `rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out.H3FinalAdaLNOutCudaOp` | `rl_engine._C.h3_det_linear_*`, `rl_engine._C.h3_rmsnorm_*` | One autograd node | +| PyTorch reference | `rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp` | n/a | `forward`: diffusers replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Numerics + +- **Projection.** The [AdaLN projection](h3-adaln-projection.md)'s deterministic tensor-core + GEMV (`h3-det-linear-bf16-mma-v1`), applied to `norm_out.linear`. The FP32 SiLU is rounded + once at the declared cast. BF16 `temb` is rejected (probe H7). +- **Norm and modulation.** The [`h3_rmsnorm`](h3-rmsnorm.md) kernel indexed by + `timestep_indices`. Given the same table, it is bitwise equal to diffusers. Overall, about + 0.04% of output elements differ from diffusers by 1 ULP, all traced to the projection's + summation tree versus cuBLAS. Rows are batch- and position-invariant. +- **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into + the projection backward without an extra BF16 rounding: + - `d_temb`, `dW` and `db` are 21–37× closer to FP64 in the recorded run than diffusers, + whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with + `index_select` atomics; + - `dx` and `d_norm_w` are at the BF16 rounding level for both; + - every gradient is repeat-bitwise. +- **Golden.** FP64, rounding only where the model stores a value in its own dtype: the SiLU + cast, `norm_out.linear`'s BF16 output table, `norm_out.norm`'s BF16 output and the BF16 + `1 + scale`. All four are rounded straight-through for the gradient. The table and `1 + scale` + are shared by every position of a timestep. Without them the golden's `d_norm_w` drifts + systematically with S, and the gtest missed it from S = 257. `norm_out.norm`'s rounding + enters `d_scale`, `dW` and `d_temb` summed over S, and the gtest missed those from S = 1024. + +## Performance Notes + +```bash +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 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 | +| 131072 | 1.28 ms | 2.97 ms | 7.02 ms | 17.61 ms | + +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) + +There are two data files: + +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the + forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean + tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The + backward reference keeps only the SiLU and table roundings, so the errors include the BF16 + rounding of `norm(x)` and `1 + scale`. +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the + whole conditioning chain (timestep → … → norm_out) replayed stage by stage over + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean + tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). + +## Existing implementations (RFC #420 reuse rule) + +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay) | **no** (param/table grads not repeatable) | 8.3e-03 / 7.3e-02 / 1021 µs | 8.9e-03 / 2.4e-01 / 4525 µs | +| rl-kernel H3FinalAdaLNOutCudaOp | yes | 8.3e-03 / 5.2e-03 / 1642 µs | 8.9e-03 / 5.7e-03 / 2959 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) +was written from a clean tree at `03da729` by + +```bash +python tools/validation/models/h3_prior_art.py --op final_adaln_out --out reports/experiments/h3-prior-art-b200/final_adaln_out.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/final_adaln_out.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) + +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_final_adaln_out.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out +python tools/validation/operators/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 4097 --normalized-dim 5376 --check-grad # also 257, 1024 +python tools/validation/models/h3_evidence.py --op final_adaln_out \ + --out reports/experiments/h3-final-adaln-out-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-final-adaln-out-b200/report.json +``` + +## Known Limitations + +- **FP32 gtest at larger S.** H3 runs this op with BF16 activations and weights, and the BF16 + gtest passes at S = 257, 1024 and 4097. With `--dtype fp32` (every input FP32), the FP32 + reduction tolerance (atol = rtol = 1e-4) misses `d_temb` from S = 257 and `dW` from + S = 1024. Both gradients are FP32 sums over 3·S rows. Where large terms cancel, the + summation error exceeds the 1e-4 absolute floor. The PyTorch reference misses the same + gradients at the same S (`--batch 3`, seed 123, max abs error): + + | S | `d_temb` CUDA | `d_temb` PyTorch | `dW` CUDA | `dW` PyTorch | + | --- | --- | --- | --- | --- | + | 64 | 9.2e-5 | 1.5e-4 | 1.5e-4 | 9.2e-5 | + | 257 | 3.1e-4 ✗ | 5.9e-4 ✗ | 8.5e-4 | 2.4e-4 | + | 1024 | 1.5e-3 ✗ | 2.4e-3 ✗ | 2.0e-3 ✗ | 1.5e-3 ✗ | + | 4097 | 5.9e-3 ✗ | 9.3e-3 ✗ | 5.9e-3 ✗ | 3.9e-3 ✗ | + + ✗ fails the gtest. Unmarked values above 1e-4 pass through the rtol term. In a PyTorch + emulation, accumulating only the projection backward in FP64 clears S = 257 but not + S ≥ 1024, because the segment sums contribute as well. Making every reduction FP64 would + change the `h3_det_linear` kernels shared with `timestep_mlp_fp32` and + `adaln_projection_3mod`. The alternative is an FP32 gradient + tolerance that scales with reduction length. Which of the two to take is an open question + for the maintainers. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-one-block.md b/docs/operators/h3-one-block.md new file mode 100644 index 000000000..b2f1d4ce0 --- /dev/null +++ b/docs/operators/h3-one-block.md @@ -0,0 +1,183 @@ +# MiniMax-H3 One Block, Forward and Backward + +## Summary + +`ws1_one_h3_block` (RFC #420, WS1 closeout) runs **block 0 of the pinned MiniMax-H3 checkpoint** +node by node, forward and backward, and reports where it first departs from diffusers. The block +is a fixed graph of 17 nodes, from the AdaLN projection to the second gated residual. It runs on +packed FL2VA layouts that are **bitwise equal to the pinned diffusers pipeline's** `position_ids` +and modality tags. + +On one B200, at S = 264, 925 and 3160: + +- **Determinism.** Every node, every gradient reaching a node, and all 14 leaf gradients are + bitwise repeatable. +- **Batch invariance.** Batch row 0 of a two-row batch has the single-row bytes at every node, + forward and backward. +- **Agreement with diffusers.** Every node that promises it is bitwise equal to diffusers on + diffusers' inputs. The first drift from diffusers is `adaln_projection`, the first reduction + whose tree differs from cuBLAS. +- **Accuracy.** The block output is exactly as close to the FP32 golden as diffusers' (relative + error 5.2e-3 / 7.9e-3 / 9.0e-3). + +Only three of the 17 nodes are bound to this stack's row operators. The others use **interim** +deterministic operators until their own rows land, or the provider replay where no kernel exists +yet (see Bindings). The block-level promises above hold with whatever a node is bound to. Speed +and the exact rounding of interim nodes belong to those operators, not to this row. + +## Entry Point + +```python +from rl_engine.validation.models.h3_block import CandidateOps, run_block_backward, run_block_case +from rl_engine.validation.models.h3_weights import load_h3_block_weights + +weights = load_h3_block_weights("cuda") # conditioning + block-0 tensors, sha256-checked +ops = CandidateOps() +forward = run_block_case(weights, layout="small", ops=ops) # per node, first_*_drift +backward = run_block_backward(weights, layout="small", ops=ops) +``` + +The CLI writes the JSON evidence and the plot script renders it: + +```bash +python tools/weights/prepare_h3_weights.py --out --block # shard 1 of 14, ~4.8 GB +export RL_KERNEL_H3_WEIGHTS= +python tools/validation/models/h3_block_replay.py --layouts tiny,small,medium --out report.json +python tools/validation/models/plot_h3_block.py report.json +``` + +## Graph and Bindings + +`rl_engine/validation/models/h3_block.py` `NODES` (graph order): + +| Node | Owning RFC row | Binding today | +|---|---|---| +| `adaln_projection` | `adaln_projection_3mod` | **row**: `H3AdaLNProjectionCudaOp` (registry) | +| `norm1`, `norm2` | `h3_rmsnorm` | **row**: `H3RMSNormCudaOp.forward_modulated` (registry) | +| `q_proj`, `k_proj`, `v_proj` | `h3_qkv_gemm` | interim: Triton FP32-accumulate GEMM, one BF16 rounding | +| `q_norm`, `k_norm` | `h3_qk_rmsnorm_d128` | interim: `H3RMSNormCudaOp.forward` on the 128-wide heads | +| `rope_q`, `rope_k` | `h3_mm_rope_3axis_partial` | reference: the provider replay (no kernel exists yet) | +| `attention` | `h3_full_attention` | interim: `DeterministicAttentionOp(causal=False)` | +| `o_proj` | `h3_attention_o_gemm` | interim: Triton FP32-accumulate GEMM | +| `residual_attn`, `residual_mlp` | `adaln_gate_residual` | **row**: `H3GateResidualCudaOp` (registry) | +| `ffn_gate_up` | `h3_ffn_gate_up_gemm` | interim: Triton FP32-accumulate GEMM | +| `swiglu` | `h3_swiglu` | interim: `SwiGLUCudaOp` (FP32 math, one rounding) | +| `ffn_down` | `h3_ffn_down_gemm` | interim: Triton FP32-accumulate GEMM | + +When a row lands, replacing its node's `candidate` callable is the whole integration. The node's +report entries then become that operator's block-level evidence. + +The interim GEMMs use `_triton_gemm_fp32`, one fixed-order FP32 K loop with no split-K, and round +to BF16 once. They do not use `TritonDetGemmOp.__call__`, whose K tree has BF16 nodes. Measured on +B200: + +| Shape | Tree path err | FP32 K loop → BF16 | cuBLAS | Tree | FP32 K loop | +|---|---|---|---|---|---| +| 264 × 5376 → 7168 | 7.0e-3 | 2.4e-3 | 2.4e-3 | 10.4 ms | 1.2 ms | +| 925 × 14336 → 5376 | 7.5e-3 | 2.6e-3 | 2.6e-3 | 47 ms | 2.7 ms | + +Both paths are bitwise repeatable and row-invariant. Error is max|err| / max|ref| against FP64. + +## What Each Node Reports + +Every node runs three ways: the provider replay (`h3_provider`, diffusers@f53d552 op for op), the +candidate (chained, and *isolated* on the provider's inputs), and the FP32 golden +(`rl_engine/reference/minimax_h3/block.py`: FP32 math, TF32 off, no BF16 storage between nodes). + +- **Forward:** `repeat_bitwise_equal` and `batch_row_bitwise_equal` (row 0 of a B = 2 batch whose + row 1 has a different input). Also chained and isolated agreement with the provider (bitwise, + max abs, first differing packed row), the candidate's and provider's `max|err| / max|golden|`, + and the sha256 of the node's output. +- **Backward:** for the block input, `temb` and the 12 block parameters: repeat bitwise, and the + candidate's and provider's error against the golden. For the gradient reaching each node's + output, in reverse graph order: repeat bitwise and batch-row bitwise. Parameter gradients sum + over the batch, so only activation gradients take part in the batch check. +- **First drift:** `first_drift`, `first_isolated_drift`, `first_repeat_drift`, + `first_batch_drift`, `first_grad_repeat_drift` and `first_grad_batch_drift` each name the first + node where that comparison stops being bitwise. +- **Permutation probe (control H-C3, reported, not promised):** permute the packed rows together + with positions and AdaLN indices, and compare the output to the permuted original. + +Inputs: `temb` comes from the provider's conditioning path (the conditioning rows are checked by +`h3_chain`). The block input is a seeded BF16 stand-in, because the packed input projections are +other rows. Text rows use AdaLN table row 0 and media rows use row 1. + +## Evidence + +[`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/report.json) was generated from a clean tree +at `41e4293` on an otherwise idle B200 (torch 2.13.0+cu130): + +![one-block evidence](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/figure.png) + +| Layout | S | First drift / first isolated drift | Repeat / batch drift (fwd, grad) | Output err: RL-Kernel / diffusers | Fwd ms: RL-Kernel / diffusers | Fwd+bwd ms | +|---|---|---|---|---|---|---| +| tiny | 264 | `adaln_projection` / `adaln_projection` | none | 5.21e-3 / 5.21e-3 | 7.4 / 0.61 | 22.7 / 2.8 | +| small | 925 | `adaln_projection` / `adaln_projection` | none | 7.87e-3 / 7.87e-3 | 30.1 / 1.12 | 81.7 / 5.0 | +| medium | 3160 | `adaln_projection` / `adaln_projection` | none | 8.97e-3 / 8.97e-3 | 225 / 3.8 | 577 / 11.5 | + +- **Forward accuracy.** Per node, the candidate's error is within 0.93–1.20× diffusers'. The + largest is `attention` at S = 3160, which is the interim attention operator's. +- **Gradient accuracy.** The candidate is better than diffusers on `adaln_proj.weight`, + `adaln_proj.bias` and `temb`, by up to 7×: diffusers' `index_select` backward accumulates the + table gradient in BF16 atomics, while the row operators use FP32 segment sums. The candidate is + within 1.16× on the other leaves, except `attn.norm_q.weight` at S = 3160 (9.5e-3 vs 5.5e-3). + At S = 925 the same leaf is 2× *more* accurate than diffusers (1.3e-2 vs 2.4e-2), so this + 128-element gradient's error moves with BF16 rounding rather than favouring one side. +- **Permutation probe (H-C3).** The output changes in 79 / 450 / 2199 of the rows, by at most 8. + Attention sums over keys in packed order, so a consistent permutation is not byte-preserving with + the interim attention operator. This is for `h3_full_attention` to settle. +- **Speed.** The candidate is 12–60× slower than diffusers forward. Almost all of it is the + interim GEMMs and attention (FP32 scores materialised). The three row operators are each + faster than or equal to diffusers in their own evidence. + +## Existing implementations (RFC #420 reuse rule) + +- **Is there an existing block-level harness?** diffusers' `MiniMaxH3TransformerBlock` is the + executable reference. It is replayed op for op as the provider, so diffusers is not imported and + the replay pins the reference's dtypes and ordering. The Qwen3-8B chain gate from #315 + (`rl_engine/validation/models/chain_gate.py`) has per-node token digests and first-drift + localisation, but its graph, shapes and acceptance are Qwen3-specific (causal attention, no + AdaLN). This row reuses its method (per-node digests, graph-ordered first drift) on the H3 + `h3_chain` stage pattern. +- **Interim operators.** These are reused from RL-Kernel and are not new code. The only choice + made here is the FP32-accumulate entry of the Triton GEMM instead of its BF16 tree, for the + accuracy and speed measured above. + +## Tests + +`tests/models/minimax_h3/test_h3_one_block.py`: + +- **CPU:** + - the three layouts hash to the pinned diffusers `build_packed_sequence` output; + - the graph is well formed; + - the partial MM-RoPE leaves channels 96–127 unchanged. +- **GPU, `tiny` and `small` forward:** + - every node repeats and is batch-row bitwise; + - promised nodes are isolated-bitwise to diffusers; + - node and output errors are at most 1.25 × diffusers' + 1e-3. +- **GPU, `tiny` backward:** + - every leaf repeats; + - dX is batch-row bitwise; + - no gradient drift; + - leaf errors are at most 1.5 × diffusers' + 5e-3, or, for a BF16 leaf, at most 6 ulps + (of its largest golden element) above diffusers'. A stored gradient's worst element moves + in whole ulps, so the relative rule alone is seed- and torch-version-sensitive; + `tools/validation/models/h3_block_ulp_sweep.py` measured at most +4.15 ulps over + tiny/small × 10 seeds on torch 2.13 and 2.9.1, and +5.15 ulps over 34 single-node + candidate/provider swaps (reports: `reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json`). + +```bash +RL_KERNEL_H3_WEIGHTS= CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/models/minimax_h3 -q +``` + +## Known Limitations + +- **Single platform.** CUDA only. `DeterministicAttentionOp` and the Triton GEMM also build for + ROCm, but the block has not been run there. +- **Interim nodes.** Eleven of the 17 nodes are interim or reference bindings. Their speed and + rounding are not the owning rows' final ones, and the MM-RoPE nodes are the provider replay + itself, so they carry no RL-Kernel evidence until `h3_mm_rope_3axis_partial` lands. +- **Stand-in inputs.** The block input is a seeded stand-in and `temb` is the provider's. + Real packed projections arrive with the input-projection rows. +- **Not covered here.** Block 0 only (`ws1_full_h3_transformer` covers 50 blocks and the refiner), + and one sequence per batch row (H3's batch is a replication axis). diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md new file mode 100644 index 000000000..75dd3d477 --- /dev/null +++ b/docs/operators/h3-rmsnorm.md @@ -0,0 +1,169 @@ +# MiniMax-H3 RMSNorm and AdaLN Modulation + +## Summary + +`h3_rmsnorm` covers every RMSNorm in MiniMax-H3 (RFC #420, WS1 step 4). That is the block +`norm1`/`norm2`, the token-refiner norms, the refiner `final_norm` and `norm_out.norm`. All of +them are `nn.RMSNorm(5376, eps=1e-5)` with a BF16 affine weight. In a transformer block and in +`norm_out`, the normalised rows are immediately modulated by per-row AdaLN parameters, and the +op fuses that step: + +```text +n = rms_norm(x, weight, eps) +out = n * (1.0 + scale[index]) + shift[index] +``` + +`index` is `adaln_indices` in a block (rows of the [projection](h3-adaln-projection.md) table) +and `timestep_indices` in `norm_out`. `shift`/`scale` are `(R, H)` row views of the AdaLN table, +gathered inside the kernel, so the `(S, H)` tensors that +[`adaln_row_gather`](h3-adaln-row-gather.md) materialises are never needed. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("h3_rmsnorm", device="cuda") +n = op(x, weight) # plain RMSNorm +out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | +| PyTorch reference | `rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden (the modulated one stores `norm(x)` and `1 + scale` in BF16, straight-through, as the model does) | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(..., S, N)` | bf16 / fp16 / fp32 | `N % 4 == 0`; H3: `N = 5376` | +| `weight` | `(N,)` | `x.dtype` | | +| `shift`, `scale` | `(R, N)` row views, same stride | `x.dtype` | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; one entry per position, shared across the batch | + +Out-of-range indices, mismatched dtypes or shapes, a non-positive eps, and `N % 4 != 0` all fail +closed. + +## Numerics + +- **Statistics (contract `h3-rmsnorm-v1`).** The kernel replays PyTorch's own + `vectorized_layer_norm_kernel` (torch `cf30153`): + - one `(32, 4)` block per row, reading 4-element vectors; + - thread `t` sums vectors `t, t+128, …` in order; + - a shuffle-down tree (16 → 1), then a cross-warp tree; + - `rsqrtf(Σx²/N + eps)`, then `w * (rstd * x)` with one cast. + + The plain norm is therefore **bitwise equal to `nn.RMSNorm`**. Rows are independent of batch + size and position. +- **Modulation.** `1 + scale`, `n * (…)` and `+ shift` are each rounded to the tensor dtype, + exactly where the eager expression rounds. The fused output is **bitwise equal to diffusers' + `norm(x) * (1.0 + scale.index_select(0, i)) + shift.index_select(0, i)`**. +- **Backward.** Everything is in FP32 with one cast per output and no atomics: + - `dx` is row-local, using the same reduction tree; + - `dweight` is folded over fixed 256-row tiles in ascending order; + - `dshift`/`dscale` are segmented sums over positions sorted stably by table row, in the + same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). + + Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is + non-deterministic and about 15–20× less accurate (its error varies from run to run). + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op h3_rmsnorm +``` + +B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.08 ms | 0.12 ms | 1.23 ms | 0.85 ms | +| 32768 | 0.35 ms | 0.76 ms | 2.18 ms | 4.40 ms | +| 131072 | 1.21 ms | 2.90 ms | 6.17 ms | 17.55 ms | + +The forward is a single pass that never materialises the gathered rows. At small S the backward +is dominated by the fixed cost of the stable sort and tile setup. Backward timings and peak +memory exclude leaf creation and the forward pass. Candidate/provider execution order alternates +each iteration and is recorded in the report. + +## Evidence + +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by +`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: + +- bitwise equality with `nn.RMSNorm` for all four pinned norm weights; +- bitwise equality with diffusers for the modulation; +- row invariance. + +## Existing implementations (RFC #420 reuse rule) + +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers composition (F.rms_norm + index_select modulation) | **no** (param/table grads not repeatable) | 8.6e-03 / 4.5e-02 / 872 µs | 7.5e-03 / 1.9e-01 / 4434 µs | +| torch F.rms_norm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 331 µs | 2.0e-03 / 2.8e-03 / 850 µs | +| TE 2.20.2 RMSNorm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 471 µs | 2.0e-03 / 2.8e-03 / 1597 µs | +| Liger 0.8.4 modulated RMSNorm + index_select | **no** (param/table grads not repeatable) | 6.6e-03 / 4.7e-02 / 810 µs | 6.9e-03 / 2.0e-01 / 3911 µs | +| Liger 0.8.4 modulated RMSNorm + rl-kernel row gather | yes | 6.6e-03 / 6.5e-03 / 1824 µs | 6.9e-03 / 8.0e-03 / 4763 µs | +| SGLang 0.5.21 fused_norm_scale_shift (forward only) | yes | 4.4e-03 / — / 85 µs (fwd only) | 4.1e-03 / — / 436 µs (fwd only) | +| rl-kernel H3RMSNormCudaOp.forward_modulated | yes | 8.6e-03 / 4.2e-03 / 1093 µs | 7.5e-03 / 5.0e-03 / 2653 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) +was written from a clean tree at `ee83dec` by + +```bash +python tools/validation/models/h3_prior_art.py --op norm_modulate --out reports/experiments/h3-prior-art-b200/norm_modulate.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/norm_modulate.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) + +This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_rmsnorm.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, incl. block norm1 + modulation +python tools/validation/operators/check_operator.py --op h3_rmsnorm --candidate cuda --device cuda \ + --dtype bf16 --batch 2 --seq 257 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op h3_rmsnorm --out reports/experiments/h3-rmsnorm-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-rmsnorm-b200/report.json +``` + +## Known Limitations + +- Bitwise parity with `nn.RMSNorm` is tied to PyTorch's vectorized layer-norm path (`N % 4 == 0`, + aligned tensors). Other shapes are rejected rather than approximated. +- The token-refiner blocks' own norm weights live in shard 14, which is not pinned. They are the + same module (`nn.RMSNorm(5376, eps=1e-5)`) as the four pinned norms tested here. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..8eb91cc38 --- /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 `6c0d380`, 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/models/minimax_h3/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/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md new file mode 100644 index 000000000..cc5809cee --- /dev/null +++ b/docs/operators/h3-timestep-mlp.md @@ -0,0 +1,184 @@ +# MiniMax-H3 FP32 Timestep MLP + +## Summary + +`timestep_mlp_fp32` is H3's `TimestepEmbedding(in_channels=256, time_embed_dim=5376, +out_dim=2688)`. It runs on the sinusoidal features of the distinct timesteps +(RFC #420, WS1 step 3): + +```text +temb = linear_2(silu(linear_1(features))) (T, 256) -> (T, 5376) -> (T, 2688) +``` + +`time_embedder` is a `_keep_in_fp32_modules` module. Its weights, biases, +activations and output are all FP32. `temb` stays FP32 because every AdaLN +projection applies its own SiLU before casting to BF16. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`. The tensors are +`time_embedder.linear_{1,2}.{weight,bias}` (see `rl_engine/validation/models/h3_manifest.json`). + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("timestep_mlp_fp32", device="cuda") +temb = op(features, w1, b1, w2, b2) # all float32 +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp` | `rl_engine._C.h3_det_linear_{forward,backward_input,backward_weight}` | Contract `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp` | n/a | `forward`: provider path (`F.linear`/`F.silu`); `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(T, K)`, `T >= 1` | float32 | H3: `K = 256` | +| `w1`, `b1` | `(H, K)`, `(H,)` | float32 | H3: `H = 5376` | +| `w2`, `b2` | `(D, H)`, `(D,)` | float32 | H3: `D = 2688` | +| output | `(T, D)` | float32 | | + +The CUDA kernel needs `K` and `H` to be multiples of 4, for 16-byte vector loads. +These inputs fail closed: BF16 in any argument (RFC probe H7), mismatched shapes, +an empty `x`, mixed devices, and unaligned `K`. + +## Numerics: contract `h3-det-linear-v1` + +Forward, for each output `y[t, n]`: + +- One warp owns a group of output columns. Lane `l` reads the 16-byte K chunks + `l, l + 32, l + 64, ...`. +- Each lane accumulates its chunks in ascending order, and the elements inside a chunk + in ascending order, using `fmaf` into an FP32 accumulator that starts at 0. +- The 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1). +- The bias is added once, then SiLU runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula). The result is stored as FP32. + +How many columns a warp owns, the row tile (1, 2 or 4), and how many chunks are +loaded ahead change only when loads are issued, never the order in which a column +is summed. A timestep's `temb` is therefore bitwise identical however many other +timesteps share the call, and wherever it sits among them. + +Backward is deterministic. There is no cuBLAS and there are no atomics: + +- `d_hidden = g @ W2` and `dx = d_pre @ W1`: N is split into fixed 64-row chunks. + Each chunk is an ascending `fmaf` chain, and the chunks are then left-folded in + ascending order. These gradients are row-local and batch-invariant. +- `d_pre = d_hidden * s * (1 + z * (1 - s))`, elementwise in FP32, using the saved + pre-activation `z`. +- `dW[n, k] = sum_t g[t, n] * x[t, k]` and `db[n] = sum_t g[t, n]`: an ascending-`t` + FP32 fold that starts from 0, in logical row order. + +Accuracy, measured on a B200 with the pinned checkpoint weights, TF32 disabled and 200 +draws of 4 timesteps (each draw includes t = 0 and t = 1): + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| max abs error vs FP64 golden, median over draws | 1.1e-6 | 5.4e-7 | +| gradients vs FP64 autograd (`x`, `w1`, `b1`, `w2`, `b2`) | <= 1.2e-5 | | + +Both outputs are about 100× inside the contract (`reduction` / `float32`, atol and rtol +1e-4). Neither is more accurate in general: cuBLAS's tree happens to do slightly better +on these weights. What the CUDA kernel adds is a fixed summation order, which makes each +row batch- and position-invariant and every run repeat-bitwise, and it is 2× faster +for T >= 2. The two outputs are not bitwise equal, so `tools/validation/models/h3_chain_replay.py` +reports this stage as the chain's `first_drift`, which is expected for a reduction. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op timestep_mlp_fp32 +``` + +B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrapper: + +| T | CUDA op | provider | kernels only (CUDA) | +| --- | --- | --- | --- | +| 1 | 35.5 µs | 40.8 µs | 3.8 + 11.6 µs | +| 2 | 40.4 µs | 84.9 µs | 3.9 + 13.7 µs | +| 4 | 44.9 µs | 88.6 µs | 5.0 + 15.8 µs | + +The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the +provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of +kernel time. + +## Evidence + +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +timestep's row is bitwise identical whether it runs alone or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers TimestepEmbedding, FP32 F.linear [plain] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 5.7e-07 / 4.1e-07 / 777 µs | 5.7e-07 / 1.0e-06 / 832 µs | 3.7e-06 / 2.0e-06 / 3609 µs | +| rl-kernel H3TimestepMLPCudaOp | yes | 2.6e-07 / 4.2e-07 / 748 µs | 2.3e-07 / 7.6e-07 / 4595 µs | 2.6e-07 / 2.3e-06 / 42993 µs | +| diffusers TimestepEmbedding, FP32 F.linear [vllm] | **no** (14214 rows, 4096 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.5e-07 / 752 µs | 3.7e-06 / 2.8e-06 / 921 µs | 3.7e-06 / 3.2e-06 / 3826 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang] | yes | 1.6e-03 / 1.6e-03 / 1036 µs | 1.5e-03 / 1.5e-03 / 1140 µs | 1.5e-03 / 1.7e-03 / 2636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang_ieee] | yes | 3.3e-06 / 2.3e-06 / 2234 µs | 3.7e-06 / 2.8e-06 / 2428 µs | 3.7e-06 / 3.2e-06 / 7329 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_te_native] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.1e-07 / 729 µs | 3.7e-06 / 1.0e-06 / 814 µs | 3.7e-06 / 2.0e-06 / 3636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton] | yes | 1.6e-03 / 1.6e-03 / 1027 µs | 1.5e-03 / 1.5e-03 / 1137 µs | 1.5e-03 / 1.7e-03 / 2634 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton_ieee] | yes | 3.3e-06 / 2.3e-06 / 2309 µs | 3.7e-06 / 2.8e-06 / 2435 µs | 3.7e-06 / 3.2e-06 / 7349 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) +was written from a clean tree at `ffd958a` by + +```bash +python tools/validation/models/h3_prior_art.py --op timestep_mlp --out reports/experiments/h3-prior-art-b200/timestep_mlp.json --megatron-src +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/timestep_mlp.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) + +No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_mlp.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP +python tools/validation/operators/check_operator.py --op timestep_mlp_fp32 --candidate cuda --device cuda \ + --dtype fp32 --batch 3 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_mlp_fp32 \ + --out reports/experiments/h3-timestep-mlp-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timestep-mlp-b200/report.json +``` + +## Known Limitations + +- FP32 only, by contract. The BF16 and FP16 gtest rows do not apply to this op. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- Weight gradients are reductions across timesteps, so they follow logical row + order. They are deterministic, but not batch-invariant (no cross-row reduction can be). diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md new file mode 100644 index 000000000..75edf172b --- /dev/null +++ b/docs/operators/h3-timestep-sinusoid.md @@ -0,0 +1,163 @@ +# MiniMax-H3 Timestep Sinusoid + +## Summary + +`timestep_sinusoid_h3` produces the FP32 sinusoidal timestep features that feed the +MiniMax-H3 timestep MLP (RFC #420, WS1 step 3). H3 builds +`Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)` and calls it +on the *distinct* timesteps of the packed sequence, unscaled in `[0, 1]` +(`t = 1 - sigma`): + +```text +freq[k] = exp(-ln(10000) * k / 128) k = 0..127 +arg[t, k] = t * freq[k] +out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256), float32 +``` + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227` (`rl_engine/validation/models/h3_manifest.json`). +Provider reference: `huggingface/diffusers@f53d552`, `get_timestep_embedding`. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("timestep_sinusoid_h3", device="cuda") +features = op(timestep) # timestep: (T,) in [0, 1] -> (T, 256) float32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid.H3TimestepSinusoidCudaOp` | `rl_engine._C.h3_timestep_sinusoid_forward` | Bitwise equal to the provider path | +| PyTorch reference | `rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp` | n/a | `forward`: provider replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `timestep` | `(T,)`, `T >= 1` | fp32 / bf16 / fp16 | Finite, in `[0, 1]`; any stride; CUDA for the CUDA backend | +| `num_channels` | scalar | int | Positive and even; H3 uses 256 | +| output | `(T, num_channels)` | float32 | `[cos \| sin]` order | + +Inputs that break the contract fail closed: non-1-D, empty, integer or out-of-range +timesteps raise. `t * 1000` callers (RFC probe H10) therefore fail instead of +silently producing a different embedding. The range check costs one host read-back. +The native CUDA entrypoint performs the value check, so the wrapper does not repeat +the host synchronization. Pass `check_range=False` to `forward` or +`rl_engine._C.h3_timestep_sinusoid_forward` only when the caller has already validated +`t`; this trusted path skips value validation and is used for kernel-only profiling. + +## Numerics + +- **Operation order.** The kernel evaluates exactly the FP32 operations the provider + runs on CUDA: `(float)(-ln 10000) * k`, multiplied by the FP32 reciprocal of `128` + (PyTorch divides by a CPU scalar this way), `expf`, `t * freq`, `cosf` and `sinf`. + The multiplies use `__fmul_rn`, so they are never fused into an FMA. The extension + is built without `--use_fast_math`, so `expf`, `sinf` and `cosf` are the precise + libdevice functions. +- **Cast points.** bf16 and fp16 timesteps are upcast to FP32 before anything else, + matching `timesteps[:, None].float()`. The output is never rounded below FP32. +- **Invariance.** Each output element depends only on `(t, k)`. Results are + bitwise identical across batch size, position, permutation and repeated runs. +- **Backward.** `d/dt` is the analytic row-local VJP + (`-f * sin(t f)` for the cos half, `f * cos(t f)` for the sin half). It is summed over + the 256 channels with a fixed pairwise tree (`tree_sum_lastdim_fp32`), so it does not + depend on how many timesteps share the launch. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison | Result | +| --- | --- | +| CUDA vs provider path, T in {1, 2, 3, 7, 64, 1000, 4097} | bitwise equal | +| CUDA vs FP64 golden | max abs 1.19e-7 (contract: FP32 elementwise atol 1e-5) | +| gradient vs FP64 golden (gtest, T = 7) | max abs 2.4e-7 | + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 +``` + +On a B200 this is a single launch of about 15 µs, independent of `T` for `T <= 64`. With +the range check it takes about 49 µs. The provider path takes about 62–70 µs, because it is +six eager kernels plus two concatenations. Either way the op is launch-bound: it moves +about 1 KB per timestep. + +## Evidence + +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). +`tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report +records that commit and the environment. + +## Existing implementations (RFC #420 reuse rule) + +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers get_timestep_embedding (op-for-op replay) | yes | 6.7e-08 / — / 64 µs (fwd only) | 8.6e-08 / — / 64 µs (fwd only) | 9.9e-08 / — / 65 µs (fwd only) | +| SGLang 0.5.21 timestep_embedding | yes | 6.7e-08 / — / 13 µs (fwd only) | 8.6e-08 / — / 12 µs (fwd only) | 9.9e-08 / — / 13 µs (fwd only) | +| rl-kernel H3TimestepSinusoidCudaOp (check_range=False) | yes | 6.7e-08 / — / 16 µs (fwd only) | 8.6e-08 / — / 16 µs (fwd only) | 9.9e-08 / — / 16 µs (fwd only) | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +was written from a clean tree at `7a82917` by + +```bash +python tools/validation/models/h3_prior_art.py --op timestep_sinusoid --out reports/experiments/h3-prior-art-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/timestep_sinusoid.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs: +2^20 uniform timesteps, every 997th FP32 bit pattern in [0, 1], and a 100001-point grid. It +compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this +kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). + +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) + +SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all +2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +was written from a clean tree at `0eb5ddd5` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_bitwise_reuse.py \ + reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python tools/validation/operators/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timestep-sinusoid-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- `max_period` is fixed at 10000, the only value H3 uses. +- `KERNEL_ALIGN_USE_FAST_MATH=1` builds leave out every H3 kernel (`H3_BITWISE_SOURCES` in + `build_tools/extensions.py`): fast math replaces `expf`/`sinf`/`cosf` with approximations and + breaks bitwise parity. Such builds dispatch the PyTorch reference. diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..318dbd2d0 --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,113 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +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/plot_h3_evidence.py reports/experiments/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/examples/models/minimax_h3/README.md b/examples/models/minimax_h3/README.md new file mode 100644 index 000000000..e39ca714f --- /dev/null +++ b/examples/models/minimax_h3/README.md @@ -0,0 +1,105 @@ +# Full H3 qualification (in progress) + +The existing block-0 evidence remains a single-block result. This extension has +not yet qualified the complete model or generation quality. + +`tools.validation.models.h3_full_model` runs the real 50 distinct transformer +layers, two text-refiner layers and input/conditioning/output components. Each +candidate block binds its own parameters to the existing graph. Row/interim/ +reference labels are retained; boundary components stay provider/reference. +The runner checks every parameter/input gradient, checkpoints activations, +offloads saved tensors and persists/reclaims completed parameter gradients. + +The `.sbatch` scripts request one GPU node and leave the partition, account and +job name to the submitting site, for example +`sbatch --partition= --account= reference.sbatch ...`. + +`reference.sbatch` first runs independent FP32 and BF16 provider passes, then freezes per-tensor error +bounds using the committed formula. It does not run a candidate. Candidate +acceptance compares immutable reference artifacts and records the first failing +tensor without changing bounds. Existing block-0 golden tolerances are unchanged. + +Generation uses the pinned public reference image identified in +`generation-inputs.json`, the accompanying prompt and the three seeds/configuration +in `generation-contract.json`. The preflight is explicitly lower resolution. +Full generation keeps the official 124 frames, 1344×768, 24 fps and joint audio. +Every initial packed input and predicted velocity is retained. The comparator +checks identical initial noise, video/audio shape and quality gates and reports +PSNR/SSIM, temporal RMSE and audio differences. Query chunking bounds candidate +attention workspace while retaining all keys for each query; it is inference only. + +An independent reference run from `2681cadb` completed the FP32/BF16 full-model +passes (638 parameter gradients each), three-seed generation and repeat comparisons. +`reference-manifest.json` anchors its 1,635 evidence files; `full-model-thresholds.json` +contains the unchanged 764 tensor bounds. These provider results do not establish +candidate acceptance. + +**Acceptance rule (2026-10-14).** The single-configuration bounds in +`full-model-thresholds.json` (1.5 × one provider pass + floor) accept only the provider +configuration they were frozen from: provider passes that differ only in SDPA backend +(flash, memory-efficient) or BLAS library (cuBLASLt) fail 36–66 of them. The first +candidate failed 54, all on single-element `max_abs` except two `relative_l2`. Acceptance +is therefore `full-model-ensemble-rule.json`, calibrated from four distinct provider +configurations with no candidate data (`tools/validation/models/h3_full_ensemble_rule.py`): +every tensor's `relative_l2` must be at most k × the largest provider error + floor, +with k = 2.6 the smallest factor at which every provider configuration held out passes +against the others; `max_abs` is reported but not gated, because held-out providers +exceed the others' maximum by up to 4.1× on it. A full-model node-swap diagnostic +(`--provider-nodes`) found no single node responsible for the candidate's high +tensors: swapping any group moves which tensors land high rather than removing them. +The single-configuration file is kept and still compared, as a record. + +**Regenerated reference (2026-10-16).** The artifacts of the first reference run (about 200 GB) +were deleted to save disk space, so the reference was regenerated with the original commands at +the same source commit (`2681cadb`). 1,614 of the 1,635 pinned files are byte-identical to the +first run's: every parameter gradient, both `tensors.pt` and `gradients.json`, every initial +input, every per-step velocity file and the `video.npy` / `audio.npy` that the comparison +reads. The 21 that differ are the eight `report.json` files (timings, job metadata), the +thresholds file that embeds their hashes, and the twelve `video.mp4` / `audio.wav` containers, +which also differ between two runs of one job (the two WAVs differ in two header bytes). +The calibrated bounds are identical to the first run's, and the ensemble rule re-derived from +the regenerated reference (six provider runs compared with it) is identical in every numeric +field. `reference-manifest.json`, `full-model-thresholds.json` and `full-model-ensemble-rule.json` +are re-pinned to the regenerated run; the first run's SHA-256 were: manifest `1640d5c7fccb`, +thresholds `3064c7eb8924`, rule `2e2a4ac58471` and `fp32-reference/report.json` `4fd74feb00ad` +(first 12 hex digits). + +`candidate.sbatch` builds its own H3 CUDA extension, records its identity, verifies +all reference/checkpoint hashes, and runs block-0 golden tests without skips before +full-model forward/backward. It requires calls to every candidate graph node in +every block and rejects CUDA backend fallback. The comparator orders first drift +by actual forward/backward execution. Only a passing numerical comparison permits +the generation preflight (the end of `candidate.sbatch`) and full generation. + +Full candidate generation is its own job. With the candidate attention it takes +about 20 hours per seed (about 1500 s per sampling step), longer than the three +seeds fit in one 24-hour job, and a preemptible partition can evict a job every few +hours. `generation.sbatch` therefore runs one seed per array task (`--array=0-2`) +and `h3_full_generation --resume` continues after an eviction: every saved step is +replayed from disk instead of recomputed and each replayed call's inputs are checked +against their saved digest (`h3_generation_resume.py`). Recomputing the last saved step +and requiring a bitwise match (`--verify-last-step`) costs a full step per restart, which +on a heavily preempted partition is most of an attempt, so it is opt-in. A run is tied to +the git trees of the model paths, not to a commit (`h3_generation_adopt.py` moves an older +commit-based run over, only if those paths are identical). The task refuses +to start unless the gates of `candidate.sbatch` passed (`` argument) at a commit whose +model paths are byte-identical to this checkout's. After all seeds finish, `generation-compare.sbatch` +runs the quality comparison. Each GPU job uses one B200 and about 512 GB of host +memory. + +Status of the full-model qualification (evidence in `reports/experiments/h3-full-model-b200/`): +numerical acceptance (50-block forward/backward under the provider-ensemble rule) has +passed; three-seed candidate generation has run, and its quality comparison has **not** +passed the bounds fixed before generation (seed 44 passes, seeds 42 and 43 do not). The +same bounds are also failed by the provider itself when one attention backend is forced +(flash and memory-efficient: 0 of 3 seeds each), and the first denoising step differs +from the provider's by the same amount (about 0.02 relative L2 on video) for the +candidate as for those provider variants, so the bounds, which had only been checked +against a bitwise-identical provider repeat, were never calibrated for a change of kernel. +They are unchanged here. No passed candidate result is claimed by CPU tests. + +The generation environment must provide torchvision matching PyTorch (the B200 +reference environment uses torch 2.9.1+cu130 and torchvision 0.24.1+cu130). +Qwen3VLProcessor needs its video processor even for the fixed image conditioning. +The runner loads this processor locally before model weights and rejects missing +required pipeline components instead of proceeding with a partially loaded pipeline. diff --git a/examples/models/minimax_h3/candidate.sbatch b/examples/models/minimax_h3/candidate.sbatch new file mode 100644 index 000000000..d88d5fe9b --- /dev/null +++ b/examples/models/minimax_h3/candidate.sbatch @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Numerical gates and the generation preflight of the full-model candidate. Full generation is +# generation.sbatch (one task per seed), allowed only after this job passed for the same commit. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +reference_root=$4 +result_root=$5 +reference_image=$6 +provider_source=$7 +extra_dependencies=$8 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 MAX_JOBS=8 +export PATH="$CUDA_HOME/bin:$PATH" +export TORCH_KERNEL_CACHE_PATH="$result_dir/torch-cache" TRITON_CACHE_DIR="$result_dir/triton-cache" +mkdir -p "$TORCH_KERNEL_CACHE_PATH" "$TRITON_CACHE_DIR" +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +"$python_bin" -m pip list --path "$extra_dependencies" --format=json > "$result_dir/extra-dependencies.json" +"$python_bin" setup.py build_ext --inplace > "$result_dir/build.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_candidate_prepare --checkpoint "$checkpoint" \ + --reference-root "$reference_root" --output "$result_dir/preflight" > "$result_dir/preflight.log" 2>&1 +export RL_KERNEL_H3_WEIGHTS="$result_dir/preflight/weights" +"$python_bin" -m pytest -q tests/models/minimax_h3/test_h3_one_block.py \ + --junitxml="$result_dir/block0.xml" > "$result_dir/block0.log" 2>&1 +"$python_bin" -c 'import sys,xml.etree.ElementTree as E; r=E.parse(sys.argv[1]); assert all(int(s.get("skipped",0)) == 0 for s in r.iter("testsuite")), "block0 skipped"' "$result_dir/block0.xml" +"$python_bin" -m tools.validation.models.h3_full_model --checkpoint "$checkpoint" \ + --manifest rl_engine/validation/models/h3_full_manifest.json --mode candidate \ + --thresholds examples/models/minimax_h3/full-model-thresholds.json \ + --output "$result_dir/candidate" > "$result_dir/candidate.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_compare --reference "$reference_root/fp32-reference" \ + --other "$result_dir/candidate" --thresholds examples/models/minimax_h3/full-model-thresholds.json \ + --reference-root "$reference_root" --reference-manifest examples/models/minimax_h3/reference-manifest.json \ + --output "$result_dir/full-model-comparison.json" > "$result_dir/comparison.log" 2>&1 || true +# The single-configuration bounds above are kept as a record; acceptance is the +# provider-ensemble rule (README), which exits non-zero on any failure. +"$python_bin" -m tools.validation.models.h3_full_ensemble_rule check \ + --rule examples/models/minimax_h3/full-model-ensemble-rule.json \ + --candidate "$result_dir/full-model-comparison.json" \ + --output "$result_dir/full-model-acceptance.json" > "$result_dir/acceptance.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --preflight --output "$result_dir/generation-preflight" \ + > "$result_dir/generation-preflight.log" 2>&1 diff --git a/examples/models/minimax_h3/diagnose-block.sbatch b/examples/models/minimax_h3/diagnose-block.sbatch new file mode 100644 index 000000000..128ea1955 --- /dev/null +++ b/examples/models/minimax_h3/diagnose-block.sbatch @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +failed_result=$3 +result_root=$4 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir" RL_KERNEL_H3_WEIGHTS="$failed_result/preflight/weights" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +# Preserve the exact native binary from the failed attempt; record and recheck +# its identity instead of changing both the kernel and diagnostic at once. +"$python_bin" - "$failed_result" "$result_dir" <<'PY' +import json +import shutil +import sys +from pathlib import Path +from tools.validation.models.h3_full_model import digest +failed, output = map(Path, sys.argv[1:]) +identity = json.loads((failed / "preflight/identity.json").read_text()) +source = Path(identity["extension"]) +assert digest(source) == identity["extension_sha256"] +shutil.copy2(source, Path("rl_engine") / source.name) +(output / "native-origin.json").write_text(json.dumps(identity, indent=2) + "\n") +(output / "retry-of.txt").write_text(str(failed) + "\n") +PY +"$python_bin" -m tools.validation.models.h3_block_diagnostic --failed-job "$failed_result" --causal-sweep \ + --output "$result_dir/diagnostic" > "$result_dir/diagnostic.log" 2>&1 diff --git a/examples/models/minimax_h3/full-model-ensemble-rule.json b/examples/models/minimax_h3/full-model-ensemble-rule.json new file mode 100644 index 000000000..ee4a86d3c --- /dev/null +++ b/examples/models/minimax_h3/full-model-ensemble-rule.json @@ -0,0 +1,2346 @@ +{ + "status": "frozen_ensemble_rule", + "rule": "relative_l2 <= k * max over provider configurations + floor; max_abs reported", + "factor_k": 2.6, + "leave_one_out_worst_ratio": { + "blas-cublaslt": 2.20842743779203, + "default": 2.5483393774198175, + "sdpa-efficient": 1.5394680908162333, + "sdpa-flash": 1.6790097485497042 + }, + "providers": [ + "blas-cublaslt", + "default", + "sdpa-efficient", + "sdpa-flash" + ], + "duplicates_dropped": [ + "reference-run", + "sdpa-cudnn" + ], + "sources": { + "blas-cublaslt": { + "sha256": "8aa7346ef2a6f9774513ecbafedb2974bad9ad62599490334a3d250024fce379", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "default": { + "sha256": "0e66bc8f799ddcf4a61acbf265fb3ef024131d2cbfd0c8ef6946d5441eeba4e1", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "reference-run": { + "sha256": "19f2e39b8280755ddf7f4a553f147f0216b756d50ba690e84a9ca09684ab21fe", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-cudnn": { + "sha256": "adc1509415f6f21aee49ec9d1b3899de62671e087f859f23ffed9db318afdda9", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-efficient": { + "sha256": "f16950d3ef7913bce1a2461f777c76f57ff82a52dd3023e1a40ae15c73c98280", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-flash": { + "sha256": "42a975d807dfe6ecec7ce81e8e4b5caadb04c9bf9f74e0aa397c76e11070531a", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + } + }, + "reference_report_sha256": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "ensemble_relative_l2": { + "audio": 0.02110560096918055, + "input_gradients.audio_hidden_states": 0.20277222521992394, + "input_gradients.encoder_hidden_states": 0.31350433311250625, + "input_gradients.hidden_states": 0.21767735582914077, + "nodes.audio_proj_in": 0.0, + "nodes.audio_proj_in.grad": 0.21529843009419838, + "nodes.audio_proj_out": 0.021090187694741094, + "nodes.audio_proj_out.grad": 0.02110560096918055, + "nodes.block.0": 0.007637028339181236, + "nodes.block.0.grad": 0.8528404577370075, + "nodes.block.1": 0.007167718706812642, + "nodes.block.1.grad": 0.83313024074139, + "nodes.block.10": 0.009470979132713336, + "nodes.block.10.grad": 0.9145230710597959, + "nodes.block.11": 0.009609765952726198, + "nodes.block.11.grad": 0.9151679613375059, + "nodes.block.12": 0.009811757785775614, + "nodes.block.12.grad": 0.9458007065819563, + "nodes.block.13": 0.010102301900968155, + "nodes.block.13.grad": 0.9579435857285441, + "nodes.block.14": 0.01029801834941549, + "nodes.block.14.grad": 0.9632512267797664, + "nodes.block.15": 0.010673081941103698, + "nodes.block.15.grad": 0.969776464301216, + "nodes.block.16": 0.010944323125835815, + "nodes.block.16.grad": 0.9900424517939425, + "nodes.block.17": 0.01116479582191621, + "nodes.block.17.grad": 0.9971977671738045, + "nodes.block.18": 0.012193206051473853, + "nodes.block.18.grad": 0.9965371368019056, + "nodes.block.19": 0.013078463243021731, + "nodes.block.19.grad": 1.0162687635352703, + "nodes.block.2": 0.007404503527807438, + "nodes.block.2.grad": 0.9834410772840688, + "nodes.block.20": 0.01347476355078683, + "nodes.block.20.grad": 1.0190441077729735, + "nodes.block.21": 0.013788720833893198, + "nodes.block.21.grad": 1.0289967632293302, + "nodes.block.22": 0.013559447098486279, + "nodes.block.22.grad": 1.0586827213971899, + "nodes.block.23": 0.012802505292154518, + "nodes.block.23.grad": 1.0800576612954356, + "nodes.block.24": 0.010981097081068994, + "nodes.block.24.grad": 1.1046818558335174, + "nodes.block.25": 0.00993079603502247, + "nodes.block.25.grad": 1.1126083625809122, + "nodes.block.26": 0.009286344548814976, + "nodes.block.26.grad": 1.113991113025653, + "nodes.block.27": 0.01933541921659969, + "nodes.block.27.grad": 1.1180943367539482, + "nodes.block.28": 0.01645995510288913, + "nodes.block.28.grad": 1.1277243036800457, + "nodes.block.29": 0.016225393322503424, + "nodes.block.29.grad": 1.0763523804698858, + "nodes.block.3": 0.008075193149480336, + "nodes.block.3.grad": 0.9060932909313238, + "nodes.block.30": 0.015821176889678927, + "nodes.block.30.grad": 1.0156486154560846, + "nodes.block.31": 0.01627388190675362, + "nodes.block.31.grad": 1.0538874207349838, + "nodes.block.32": 0.015005661062321091, + "nodes.block.32.grad": 1.0038473074352139, + "nodes.block.33": 0.015151115277153682, + "nodes.block.33.grad": 1.001759230958048, + "nodes.block.34": 0.015260938453177957, + "nodes.block.34.grad": 1.0882704655567403, + "nodes.block.35": 0.01361960971320334, + "nodes.block.35.grad": 1.1198217711779717, + "nodes.block.36": 0.012297703371840298, + "nodes.block.36.grad": 1.0856430212191899, + "nodes.block.37": 0.011772510842582027, + "nodes.block.37.grad": 1.0654626726976346, + "nodes.block.38": 0.012050407362486405, + "nodes.block.38.grad": 1.0858962868337536, + "nodes.block.39": 0.011985727655582418, + "nodes.block.39.grad": 1.095775280759746, + "nodes.block.4": 0.00868913976043236, + "nodes.block.4.grad": 0.924952776102672, + "nodes.block.40": 0.011512896942200267, + "nodes.block.40.grad": 1.1018309295987487, + "nodes.block.41": 0.011611925677531472, + "nodes.block.41.grad": 1.1067364188449467, + "nodes.block.42": 0.009726727784226567, + "nodes.block.42.grad": 1.1876454639541187, + "nodes.block.43": 0.007361845810865862, + "nodes.block.43.grad": 1.174099900135046, + "nodes.block.44": 0.0071179542127801695, + "nodes.block.44.grad": 1.1571739376426546, + "nodes.block.45": 0.0077513342359457945, + "nodes.block.45.grad": 1.1530878877550301, + "nodes.block.46": 0.007713059002009989, + "nodes.block.46.grad": 1.1393752324442692, + "nodes.block.47": 0.00577020720572875, + "nodes.block.47.grad": 1.141632154523112, + "nodes.block.48": 0.006725423713444281, + "nodes.block.48.grad": 1.1713012568585635, + "nodes.block.49": 0.007460893197348834, + "nodes.block.49.grad": 0.02028429422458945, + "nodes.block.5": 0.008896497807929678, + "nodes.block.5.grad": 0.8714454280768207, + "nodes.block.6": 0.008949304485577259, + "nodes.block.6.grad": 0.9109083971517853, + "nodes.block.7": 0.008971055142252412, + "nodes.block.7.grad": 0.905662113323111, + "nodes.block.8": 0.009235740680008939, + "nodes.block.8.grad": 0.9036578329629903, + "nodes.block.9": 0.009412928731248495, + "nodes.block.9.grad": 0.9007381090490874, + "nodes.context_embedder": 0.002384571186068664, + "nodes.context_embedder.grad": 0.3146938130796209, + "nodes.norm_out": 0.01728418962285463, + "nodes.norm_out.grad": 0.020587306859649657, + "nodes.proj_in": 0.0, + "nodes.proj_in.grad": 0.23277925686464554, + "nodes.proj_out": 0.018702757639433267, + "nodes.proj_out.grad": 0.006463551732884611, + "nodes.refiner.0": 0.0052852776939599985, + "nodes.refiner.0.grad": 0.3185379392032242, + "nodes.refiner.1": 0.005989274309705795, + "nodes.refiner.1.grad": 0.31814344423924806, + "nodes.time_embedder": 0.0, + "nodes.time_embedder.grad": 0.5162050269731083, + "nodes.time_proj": 0.0, + "nodes.token_refiner": 0.006084313328960681, + "nodes.token_refiner.grad": 0.3160029449176521, + "parameter.audio_proj_in.bias": 0.1874771619328801, + "parameter.audio_proj_in.weight": 0.21519850916146643, + "parameter.audio_proj_out.bias": 0.013905871137514934, + "parameter.audio_proj_out.weight": 0.017141809562141865, + "parameter.context_embedder.bias": 0.34122486124431056, + "parameter.context_embedder.weight": 0.31444085786700243, + "parameter.norm_out.linear.bias": 0.015900100954407113, + "parameter.norm_out.linear.weight": 0.01586639338847783, + "parameter.norm_out.norm.weight": 0.011882855956542197, + "parameter.proj_in.bias": 0.2177238654494255, + "parameter.proj_in.weight": 0.23583015398086374, + "parameter.proj_out.bias": 0.005525584430620126, + "parameter.proj_out.weight": 0.010896080050018181, + "parameter.time_embedder.linear_1.bias": 0.03383056500771813, + "parameter.time_embedder.linear_1.weight": 0.03108368457958572, + "parameter.time_embedder.linear_2.bias": 0.49613847877224526, + "parameter.time_embedder.linear_2.weight": 0.14648554956113208, + "parameter.token_refiner.final_norm.weight": 0.32675350633917283, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 0.26645555754500405, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 0.26593696921549076, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 0.3163188188666285, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 0.32687503983930444, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 0.3045241605999762, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 0.3270055880199932, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 0.3219667791477918, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 0.32660888280066347, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 0.32353553759927595, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 0.3315986685912315, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 0.43792017421068, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 0.438433276061005, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 0.35206328838048234, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 0.3352967680856234, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 0.3583973153812709, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 0.3410664838137415, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 0.36494682444311155, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 0.335625810567975, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 0.26516590237840665, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 0.36832092767828667, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 0.31899451766122605, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 0.17867782631623222, + "parameter.transformer_blocks.0.attn.norm_k.weight": 0.28960468817995333, + "parameter.transformer_blocks.0.attn.norm_q.weight": 0.31245708376305836, + "parameter.transformer_blocks.0.attn.to_k.weight": 0.22809571887143615, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 0.1752762846442529, + "parameter.transformer_blocks.0.attn.to_q.weight": 0.2325037713076897, + "parameter.transformer_blocks.0.attn.to_v.weight": 0.17366440926942864, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 0.18026215904611417, + "parameter.transformer_blocks.0.ff.net.2.weight": 0.17771809513331524, + "parameter.transformer_blocks.0.norm1.weight": 0.21693076268469982, + "parameter.transformer_blocks.0.norm2.weight": 0.17937090662880234, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 0.18509642602034504, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 0.1774306170524224, + "parameter.transformer_blocks.1.attn.norm_k.weight": 0.6107248821324249, + "parameter.transformer_blocks.1.attn.norm_q.weight": 0.6443177983145605, + "parameter.transformer_blocks.1.attn.to_k.weight": 0.6697545238218464, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 0.5362465973157898, + "parameter.transformer_blocks.1.attn.to_q.weight": 0.5868832674821828, + "parameter.transformer_blocks.1.attn.to_v.weight": 0.578016937573896, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 0.18332852162379282, + "parameter.transformer_blocks.1.ff.net.2.weight": 0.17509047422265817, + "parameter.transformer_blocks.1.norm1.weight": 0.6435346837382664, + "parameter.transformer_blocks.1.norm2.weight": 0.17115076008546634, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 0.17941840962409109, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 0.1623874068552092, + "parameter.transformer_blocks.10.attn.norm_k.weight": 0.6152553123197887, + "parameter.transformer_blocks.10.attn.norm_q.weight": 0.6176182711966747, + "parameter.transformer_blocks.10.attn.to_k.weight": 0.5929514096870396, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 0.506602926058383, + "parameter.transformer_blocks.10.attn.to_q.weight": 0.5846389710107828, + "parameter.transformer_blocks.10.attn.to_v.weight": 0.5517895970800015, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 0.17987433666892472, + "parameter.transformer_blocks.10.ff.net.2.weight": 0.15051711004439988, + "parameter.transformer_blocks.10.norm1.weight": 0.6431711828853883, + "parameter.transformer_blocks.10.norm2.weight": 0.22894892387570479, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 0.16100573162632656, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 0.15062333205170744, + "parameter.transformer_blocks.11.attn.norm_k.weight": 0.7059893490695992, + "parameter.transformer_blocks.11.attn.norm_q.weight": 0.736105549346304, + "parameter.transformer_blocks.11.attn.to_k.weight": 0.6789991211550914, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 0.49334014494012723, + "parameter.transformer_blocks.11.attn.to_q.weight": 0.6664320046003854, + "parameter.transformer_blocks.11.attn.to_v.weight": 0.5363533879985244, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 0.15950074480074164, + "parameter.transformer_blocks.11.ff.net.2.weight": 0.1427400350951154, + "parameter.transformer_blocks.11.norm1.weight": 0.6165197387097995, + "parameter.transformer_blocks.11.norm2.weight": 0.1743280989426046, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 0.16669251305544974, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 0.15450476555580314, + "parameter.transformer_blocks.12.attn.norm_k.weight": 0.6707093903946136, + "parameter.transformer_blocks.12.attn.norm_q.weight": 0.6762641775318454, + "parameter.transformer_blocks.12.attn.to_k.weight": 0.5035471662679206, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 0.43496434592746525, + "parameter.transformer_blocks.12.attn.to_q.weight": 0.5499982670952236, + "parameter.transformer_blocks.12.attn.to_v.weight": 0.48818185711813206, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 0.1722070930596971, + "parameter.transformer_blocks.12.ff.net.2.weight": 0.1460392638774976, + "parameter.transformer_blocks.12.norm1.weight": 0.5046925581248843, + "parameter.transformer_blocks.12.norm2.weight": 0.17744799383842155, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 0.1696187851756362, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 0.15451788367374297, + "parameter.transformer_blocks.13.attn.norm_k.weight": 0.6826694775939346, + "parameter.transformer_blocks.13.attn.norm_q.weight": 0.6739919946570334, + "parameter.transformer_blocks.13.attn.to_k.weight": 0.616337247011177, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 0.6185542328063899, + "parameter.transformer_blocks.13.attn.to_q.weight": 0.613351106050686, + "parameter.transformer_blocks.13.attn.to_v.weight": 0.6313333975892937, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 0.168929578806616, + "parameter.transformer_blocks.13.ff.net.2.weight": 0.14457795218884922, + "parameter.transformer_blocks.13.norm1.weight": 0.655974573599716, + "parameter.transformer_blocks.13.norm2.weight": 0.17869929032611745, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 0.17394482005841622, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 0.1590707776637609, + "parameter.transformer_blocks.14.attn.norm_k.weight": 0.5634564709394162, + "parameter.transformer_blocks.14.attn.norm_q.weight": 0.6043468403497236, + "parameter.transformer_blocks.14.attn.to_k.weight": 0.5561731527009006, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 0.49695488228434914, + "parameter.transformer_blocks.14.attn.to_q.weight": 0.5806546639369049, + "parameter.transformer_blocks.14.attn.to_v.weight": 0.5676983678537382, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 0.1637974398672475, + "parameter.transformer_blocks.14.ff.net.2.weight": 0.14413908284561155, + "parameter.transformer_blocks.14.norm1.weight": 0.6548673998875053, + "parameter.transformer_blocks.14.norm2.weight": 0.1932632933697487, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 0.16735203350888342, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 0.15188667677679074, + "parameter.transformer_blocks.15.attn.norm_k.weight": 0.6675522437962607, + "parameter.transformer_blocks.15.attn.norm_q.weight": 0.69642257060271, + "parameter.transformer_blocks.15.attn.to_k.weight": 0.6497632709849944, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 0.4787558426271283, + "parameter.transformer_blocks.15.attn.to_q.weight": 0.6026810255486286, + "parameter.transformer_blocks.15.attn.to_v.weight": 0.5504704747572478, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 0.162270031903906, + "parameter.transformer_blocks.15.ff.net.2.weight": 0.1453468746058194, + "parameter.transformer_blocks.15.norm1.weight": 0.6601224321425502, + "parameter.transformer_blocks.15.norm2.weight": 0.20789091452427774, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 0.17192920023271227, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 0.15172464758144125, + "parameter.transformer_blocks.16.attn.norm_k.weight": 0.5411303219731743, + "parameter.transformer_blocks.16.attn.norm_q.weight": 0.5627860173783555, + "parameter.transformer_blocks.16.attn.to_k.weight": 0.5147295853848322, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.37626756060993605, + "parameter.transformer_blocks.16.attn.to_q.weight": 0.5211155316562834, + "parameter.transformer_blocks.16.attn.to_v.weight": 0.4601833083803108, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 0.1611149175734667, + "parameter.transformer_blocks.16.ff.net.2.weight": 0.13647756377949197, + "parameter.transformer_blocks.16.norm1.weight": 0.5305065646401627, + "parameter.transformer_blocks.16.norm2.weight": 0.16299974442833884, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 0.15750520234871465, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 0.1442080584059882, + "parameter.transformer_blocks.17.attn.norm_k.weight": 0.5064534732130231, + "parameter.transformer_blocks.17.attn.norm_q.weight": 0.5585420282834562, + "parameter.transformer_blocks.17.attn.to_k.weight": 0.5945243422467011, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 0.5920735852348725, + "parameter.transformer_blocks.17.attn.to_q.weight": 0.5748802118736813, + "parameter.transformer_blocks.17.attn.to_v.weight": 0.6715329239199194, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 0.15044139611347626, + "parameter.transformer_blocks.17.ff.net.2.weight": 0.13287263448993833, + "parameter.transformer_blocks.17.norm1.weight": 0.6910566752579711, + "parameter.transformer_blocks.17.norm2.weight": 0.15128633761209442, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 0.1434070672190233, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 0.13647082108089864, + "parameter.transformer_blocks.18.attn.norm_k.weight": 0.7149893895321376, + "parameter.transformer_blocks.18.attn.norm_q.weight": 0.6497792945045617, + "parameter.transformer_blocks.18.attn.to_k.weight": 0.6257317399231312, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 0.4735823683215769, + "parameter.transformer_blocks.18.attn.to_q.weight": 0.6227941989918219, + "parameter.transformer_blocks.18.attn.to_v.weight": 0.5951698295070317, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 0.141487463750178, + "parameter.transformer_blocks.18.ff.net.2.weight": 0.12124054956470803, + "parameter.transformer_blocks.18.norm1.weight": 0.6196138920784167, + "parameter.transformer_blocks.18.norm2.weight": 0.14495111923251633, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 0.1414025275323694, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 0.13174059786195905, + "parameter.transformer_blocks.19.attn.norm_k.weight": 0.6492313403519856, + "parameter.transformer_blocks.19.attn.norm_q.weight": 0.6377195588380845, + "parameter.transformer_blocks.19.attn.to_k.weight": 0.6360903900367645, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 0.4765003925776794, + "parameter.transformer_blocks.19.attn.to_q.weight": 0.629942788876498, + "parameter.transformer_blocks.19.attn.to_v.weight": 0.5691853563022023, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 0.12873415655769765, + "parameter.transformer_blocks.19.ff.net.2.weight": 0.10139029534587113, + "parameter.transformer_blocks.19.norm1.weight": 0.6353850059551215, + "parameter.transformer_blocks.19.norm2.weight": 0.1508616045251074, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 0.17733209220723198, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 0.1706727714082375, + "parameter.transformer_blocks.2.attn.norm_k.weight": 0.5551056072295623, + "parameter.transformer_blocks.2.attn.norm_q.weight": 0.5369920381607626, + "parameter.transformer_blocks.2.attn.to_k.weight": 0.5302745024408893, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 0.4763106748686571, + "parameter.transformer_blocks.2.attn.to_q.weight": 0.4913420936131699, + "parameter.transformer_blocks.2.attn.to_v.weight": 0.6139169897083216, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 0.17577450548046156, + "parameter.transformer_blocks.2.ff.net.2.weight": 0.16919691639079415, + "parameter.transformer_blocks.2.norm1.weight": 0.5438326294406551, + "parameter.transformer_blocks.2.norm2.weight": 0.17908548259702073, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 0.15516000757003306, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 0.13473269748319175, + "parameter.transformer_blocks.20.attn.norm_k.weight": 0.8717416419892866, + "parameter.transformer_blocks.20.attn.norm_q.weight": 0.9198161269361697, + "parameter.transformer_blocks.20.attn.to_k.weight": 0.8572434742105496, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 0.3892794559046116, + "parameter.transformer_blocks.20.attn.to_q.weight": 0.791981844639221, + "parameter.transformer_blocks.20.attn.to_v.weight": 0.478350420897483, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 0.14013424458483464, + "parameter.transformer_blocks.20.ff.net.2.weight": 0.11631561270226247, + "parameter.transformer_blocks.20.norm1.weight": 0.6788613775770197, + "parameter.transformer_blocks.20.norm2.weight": 0.13949718891009422, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 0.1482196449124827, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 0.1300058383040362, + "parameter.transformer_blocks.21.attn.norm_k.weight": 0.7481017080404762, + "parameter.transformer_blocks.21.attn.norm_q.weight": 0.7592458857742853, + "parameter.transformer_blocks.21.attn.to_k.weight": 0.7171547011740366, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 0.5013809935656205, + "parameter.transformer_blocks.21.attn.to_q.weight": 0.676030423741261, + "parameter.transformer_blocks.21.attn.to_v.weight": 0.6191617303664513, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 0.1420040837772908, + "parameter.transformer_blocks.21.ff.net.2.weight": 0.1180330600974897, + "parameter.transformer_blocks.21.norm1.weight": 0.640965339782901, + "parameter.transformer_blocks.21.norm2.weight": 0.1434696448275675, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 0.14936310971176525, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 0.132431273579337, + "parameter.transformer_blocks.22.attn.norm_k.weight": 0.6627742980495887, + "parameter.transformer_blocks.22.attn.norm_q.weight": 0.6658411201577026, + "parameter.transformer_blocks.22.attn.to_k.weight": 0.6216034582065753, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 0.6723760585981233, + "parameter.transformer_blocks.22.attn.to_q.weight": 0.6059599015397963, + "parameter.transformer_blocks.22.attn.to_v.weight": 0.772219315272921, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 0.1426257615377658, + "parameter.transformer_blocks.22.ff.net.2.weight": 0.11360215573011313, + "parameter.transformer_blocks.22.norm1.weight": 0.6697341123206519, + "parameter.transformer_blocks.22.norm2.weight": 0.14625720523831526, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 0.13884457252088103, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 0.12478032787063244, + "parameter.transformer_blocks.23.attn.norm_k.weight": 0.6814540343533161, + "parameter.transformer_blocks.23.attn.norm_q.weight": 0.684038690122064, + "parameter.transformer_blocks.23.attn.to_k.weight": 0.6908737360594973, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 0.49140031547262497, + "parameter.transformer_blocks.23.attn.to_q.weight": 0.6379720333772122, + "parameter.transformer_blocks.23.attn.to_v.weight": 0.6241045529727879, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 0.13312269435414717, + "parameter.transformer_blocks.23.ff.net.2.weight": 0.1143442047795657, + "parameter.transformer_blocks.23.norm1.weight": 0.7184202157238492, + "parameter.transformer_blocks.23.norm2.weight": 0.14232311638461428, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 0.14042778599981806, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 0.12344647130721911, + "parameter.transformer_blocks.24.attn.norm_k.weight": 0.5957750471379901, + "parameter.transformer_blocks.24.attn.norm_q.weight": 0.5858490156955045, + "parameter.transformer_blocks.24.attn.to_k.weight": 0.5651326224307139, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 0.3861928792070328, + "parameter.transformer_blocks.24.attn.to_q.weight": 0.5318470735889794, + "parameter.transformer_blocks.24.attn.to_v.weight": 0.5294161705020743, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 0.13253318691665164, + "parameter.transformer_blocks.24.ff.net.2.weight": 0.08400841861636524, + "parameter.transformer_blocks.24.norm1.weight": 0.5916204321295466, + "parameter.transformer_blocks.24.norm2.weight": 0.13317067128474694, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 0.10912880779097324, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 0.1021035135517177, + "parameter.transformer_blocks.25.attn.norm_k.weight": 0.6487957220130574, + "parameter.transformer_blocks.25.attn.norm_q.weight": 0.6523428868096084, + "parameter.transformer_blocks.25.attn.to_k.weight": 0.6660812607231732, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 0.5174420452991362, + "parameter.transformer_blocks.25.attn.to_q.weight": 0.6397013876867854, + "parameter.transformer_blocks.25.attn.to_v.weight": 0.660878895059845, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 0.10628172036996353, + "parameter.transformer_blocks.25.ff.net.2.weight": 0.07574625326211328, + "parameter.transformer_blocks.25.norm1.weight": 0.5410480441062008, + "parameter.transformer_blocks.25.norm2.weight": 0.10128470936739821, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 0.1280437196985238, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 0.10888463268784916, + "parameter.transformer_blocks.26.attn.norm_k.weight": 0.6441648947168217, + "parameter.transformer_blocks.26.attn.norm_q.weight": 0.6484796721945846, + "parameter.transformer_blocks.26.attn.to_k.weight": 0.6491471682152515, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.21242931299844153, + "parameter.transformer_blocks.26.attn.to_q.weight": 0.6467800475358854, + "parameter.transformer_blocks.26.attn.to_v.weight": 0.32522441779907785, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 0.11659499821369074, + "parameter.transformer_blocks.26.ff.net.2.weight": 0.08279338541473577, + "parameter.transformer_blocks.26.norm1.weight": 0.5092815994108456, + "parameter.transformer_blocks.26.norm2.weight": 0.12215671300627264, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 0.1390507418408379, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 0.11275655533933447, + "parameter.transformer_blocks.27.attn.norm_k.weight": 0.7734440691053224, + "parameter.transformer_blocks.27.attn.norm_q.weight": 0.7953313142746422, + "parameter.transformer_blocks.27.attn.to_k.weight": 0.8133527123298837, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 0.6348785766665711, + "parameter.transformer_blocks.27.attn.to_q.weight": 0.7918463015928954, + "parameter.transformer_blocks.27.attn.to_v.weight": 0.7301223125733811, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 0.1263956489072167, + "parameter.transformer_blocks.27.ff.net.2.weight": 0.09920134398712835, + "parameter.transformer_blocks.27.norm1.weight": 0.8697917621100352, + "parameter.transformer_blocks.27.norm2.weight": 0.11998925710626789, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 0.1332879678030465, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 0.10131284826629, + "parameter.transformer_blocks.28.attn.norm_k.weight": 0.9210698288645108, + "parameter.transformer_blocks.28.attn.norm_q.weight": 0.9192904857217888, + "parameter.transformer_blocks.28.attn.to_k.weight": 0.9299335039349791, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 0.8426652469633431, + "parameter.transformer_blocks.28.attn.to_q.weight": 0.9482478007716462, + "parameter.transformer_blocks.28.attn.to_v.weight": 0.8578499632178054, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 0.12237747003347534, + "parameter.transformer_blocks.28.ff.net.2.weight": 0.08989724385841778, + "parameter.transformer_blocks.28.norm1.weight": 0.9234679202995794, + "parameter.transformer_blocks.28.norm2.weight": 0.13418668037996598, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 0.12769710071712426, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 0.09735381316720622, + "parameter.transformer_blocks.29.attn.norm_k.weight": 1.086104090000951, + "parameter.transformer_blocks.29.attn.norm_q.weight": 1.0901827753128521, + "parameter.transformer_blocks.29.attn.to_k.weight": 0.9437641228092345, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 0.8368198991044917, + "parameter.transformer_blocks.29.attn.to_q.weight": 0.951135692714322, + "parameter.transformer_blocks.29.attn.to_v.weight": 0.8757337926823521, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 0.11756568410754052, + "parameter.transformer_blocks.29.ff.net.2.weight": 0.07980869084435108, + "parameter.transformer_blocks.29.norm1.weight": 0.855527361817073, + "parameter.transformer_blocks.29.norm2.weight": 0.11290706902206783, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 0.1736136707981538, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 0.16752519826321585, + "parameter.transformer_blocks.3.attn.norm_k.weight": 0.7638730850842748, + "parameter.transformer_blocks.3.attn.norm_q.weight": 0.7465127871209892, + "parameter.transformer_blocks.3.attn.to_k.weight": 0.8239210843040061, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 0.5116743271863021, + "parameter.transformer_blocks.3.attn.to_q.weight": 0.7078958561704025, + "parameter.transformer_blocks.3.attn.to_v.weight": 0.584106712366063, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 0.1734745564727508, + "parameter.transformer_blocks.3.ff.net.2.weight": 0.1647559335686392, + "parameter.transformer_blocks.3.norm1.weight": 0.7427319268410297, + "parameter.transformer_blocks.3.norm2.weight": 0.17633769659016266, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 0.1614848620480023, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 0.09714668064766162, + "parameter.transformer_blocks.30.attn.norm_k.weight": 0.9262721419142987, + "parameter.transformer_blocks.30.attn.norm_q.weight": 0.9425131749841288, + "parameter.transformer_blocks.30.attn.to_k.weight": 1.0460792057173103, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 0.6686605200736664, + "parameter.transformer_blocks.30.attn.to_q.weight": 0.9757513962929028, + "parameter.transformer_blocks.30.attn.to_v.weight": 0.7299885484468922, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 0.12726474297280962, + "parameter.transformer_blocks.30.ff.net.2.weight": 0.08961859178187792, + "parameter.transformer_blocks.30.norm1.weight": 0.9618845175797996, + "parameter.transformer_blocks.30.norm2.weight": 0.12280940291215998, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 0.1450064713873809, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 0.10459082237081521, + "parameter.transformer_blocks.31.attn.norm_k.weight": 0.5653750877917905, + "parameter.transformer_blocks.31.attn.norm_q.weight": 0.5560376981215728, + "parameter.transformer_blocks.31.attn.to_k.weight": 0.6695545906672978, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 0.6623203254172305, + "parameter.transformer_blocks.31.attn.to_q.weight": 0.6356325720384296, + "parameter.transformer_blocks.31.attn.to_v.weight": 0.7361700881562175, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 0.14358311859226364, + "parameter.transformer_blocks.31.ff.net.2.weight": 0.09965486066250556, + "parameter.transformer_blocks.31.norm1.weight": 0.6471378646598047, + "parameter.transformer_blocks.31.norm2.weight": 0.1389725319732575, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 0.152713623242756, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 0.0902834080472529, + "parameter.transformer_blocks.32.attn.norm_k.weight": 1.1410089440835625, + "parameter.transformer_blocks.32.attn.norm_q.weight": 1.1158782942368046, + "parameter.transformer_blocks.32.attn.to_k.weight": 1.0004746487190592, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 0.889554343452924, + "parameter.transformer_blocks.32.attn.to_q.weight": 0.9849043486731729, + "parameter.transformer_blocks.32.attn.to_v.weight": 0.9132571793910841, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 0.13161416846598228, + "parameter.transformer_blocks.32.ff.net.2.weight": 0.09298675626882538, + "parameter.transformer_blocks.32.norm1.weight": 0.9349206141687386, + "parameter.transformer_blocks.32.norm2.weight": 0.13212903942288529, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 0.13979614110149272, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 0.09553208066293542, + "parameter.transformer_blocks.33.attn.norm_k.weight": 0.9155705517529243, + "parameter.transformer_blocks.33.attn.norm_q.weight": 0.9154534423974686, + "parameter.transformer_blocks.33.attn.to_k.weight": 0.8314888895641861, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 0.8453351563331559, + "parameter.transformer_blocks.33.attn.to_q.weight": 0.8044708403916625, + "parameter.transformer_blocks.33.attn.to_v.weight": 0.8490922798253322, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 0.13045382202277842, + "parameter.transformer_blocks.33.ff.net.2.weight": 0.0840068196187977, + "parameter.transformer_blocks.33.norm1.weight": 0.7570142897318359, + "parameter.transformer_blocks.33.norm2.weight": 0.14674158754893832, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 0.13477742308433482, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 0.08323276800712222, + "parameter.transformer_blocks.34.attn.norm_k.weight": 0.9791420574656059, + "parameter.transformer_blocks.34.attn.norm_q.weight": 0.9652674685185889, + "parameter.transformer_blocks.34.attn.to_k.weight": 0.8292524820974579, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 0.7930937938479834, + "parameter.transformer_blocks.34.attn.to_q.weight": 0.7589795609253875, + "parameter.transformer_blocks.34.attn.to_v.weight": 0.8549580578578369, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 0.11995605351038738, + "parameter.transformer_blocks.34.ff.net.2.weight": 0.0838291183465568, + "parameter.transformer_blocks.34.norm1.weight": 0.9403059130060623, + "parameter.transformer_blocks.34.norm2.weight": 0.12189935037058316, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 0.10189507844898495, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 0.07609252844677573, + "parameter.transformer_blocks.35.attn.norm_k.weight": 0.8454882627709907, + "parameter.transformer_blocks.35.attn.norm_q.weight": 0.8218902481390318, + "parameter.transformer_blocks.35.attn.to_k.weight": 0.8230795944844709, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 0.7995565691361847, + "parameter.transformer_blocks.35.attn.to_q.weight": 0.8884296214201013, + "parameter.transformer_blocks.35.attn.to_v.weight": 0.8713868733028493, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 0.11097041565052795, + "parameter.transformer_blocks.35.ff.net.2.weight": 0.07438589686382978, + "parameter.transformer_blocks.35.norm1.weight": 0.7699673515094488, + "parameter.transformer_blocks.35.norm2.weight": 0.130018302136186, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 0.1425909917060937, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 0.10349439246805126, + "parameter.transformer_blocks.36.attn.norm_k.weight": 0.9145770498663861, + "parameter.transformer_blocks.36.attn.norm_q.weight": 0.9484815620003966, + "parameter.transformer_blocks.36.attn.to_k.weight": 1.1298876570819247, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 0.5106671007649235, + "parameter.transformer_blocks.36.attn.to_q.weight": 1.1055395676246789, + "parameter.transformer_blocks.36.attn.to_v.weight": 0.5525522622895881, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 0.13211208971682759, + "parameter.transformer_blocks.36.ff.net.2.weight": 0.07943197232542645, + "parameter.transformer_blocks.36.norm1.weight": 0.9174376768312953, + "parameter.transformer_blocks.36.norm2.weight": 0.3773931626682203, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 0.12218262090280432, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 0.09287378428596499, + "parameter.transformer_blocks.37.attn.norm_k.weight": 0.9123803015450559, + "parameter.transformer_blocks.37.attn.norm_q.weight": 0.9302709725598957, + "parameter.transformer_blocks.37.attn.to_k.weight": 1.000961027885407, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 0.7427503354176709, + "parameter.transformer_blocks.37.attn.to_q.weight": 0.9484356139329904, + "parameter.transformer_blocks.37.attn.to_v.weight": 0.7748231151004865, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 0.1467101168691995, + "parameter.transformer_blocks.37.ff.net.2.weight": 0.06782686777903887, + "parameter.transformer_blocks.37.norm1.weight": 0.8829208798747256, + "parameter.transformer_blocks.37.norm2.weight": 0.17577249861019198, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 0.1254128039706359, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 0.0690761047640983, + "parameter.transformer_blocks.38.attn.norm_k.weight": 0.8418897031323198, + "parameter.transformer_blocks.38.attn.norm_q.weight": 0.8231596183694648, + "parameter.transformer_blocks.38.attn.to_k.weight": 0.8308298748231153, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 0.826364197591171, + "parameter.transformer_blocks.38.attn.to_q.weight": 0.8747986913890288, + "parameter.transformer_blocks.38.attn.to_v.weight": 0.8249675556685392, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 0.11340375170577656, + "parameter.transformer_blocks.38.ff.net.2.weight": 0.06541392236676424, + "parameter.transformer_blocks.38.norm1.weight": 0.7646574747887432, + "parameter.transformer_blocks.38.norm2.weight": 0.10926704894255639, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 0.10287534823972645, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 0.06472699315158703, + "parameter.transformer_blocks.39.attn.norm_k.weight": 0.9668497714880203, + "parameter.transformer_blocks.39.attn.norm_q.weight": 0.9802806969412944, + "parameter.transformer_blocks.39.attn.to_k.weight": 0.9821426346333626, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 0.7643879929788783, + "parameter.transformer_blocks.39.attn.to_q.weight": 1.023908223876388, + "parameter.transformer_blocks.39.attn.to_v.weight": 0.775839462822879, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 0.09925381460058395, + "parameter.transformer_blocks.39.ff.net.2.weight": 0.06097360985984232, + "parameter.transformer_blocks.39.norm1.weight": 0.7778455401258393, + "parameter.transformer_blocks.39.norm2.weight": 0.16601035712608037, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 0.17159093153329363, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 0.1638525474502011, + "parameter.transformer_blocks.4.attn.norm_k.weight": 0.8085543754044254, + "parameter.transformer_blocks.4.attn.norm_q.weight": 0.7691899487095933, + "parameter.transformer_blocks.4.attn.to_k.weight": 0.6670061012138586, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 0.6291277678966494, + "parameter.transformer_blocks.4.attn.to_q.weight": 0.6630847996693151, + "parameter.transformer_blocks.4.attn.to_v.weight": 0.7004136365823667, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 0.17499912855333324, + "parameter.transformer_blocks.4.ff.net.2.weight": 0.15735854396481597, + "parameter.transformer_blocks.4.norm1.weight": 0.7134663928601898, + "parameter.transformer_blocks.4.norm2.weight": 0.16745667286470833, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 0.0995681843335852, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 0.059305223345764504, + "parameter.transformer_blocks.40.attn.norm_k.weight": 0.9435521245892633, + "parameter.transformer_blocks.40.attn.norm_q.weight": 0.9891083149347497, + "parameter.transformer_blocks.40.attn.to_k.weight": 0.9628497990950426, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 0.7036645759214321, + "parameter.transformer_blocks.40.attn.to_q.weight": 1.0394286869298817, + "parameter.transformer_blocks.40.attn.to_v.weight": 0.7389363281243745, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 0.10314438485333412, + "parameter.transformer_blocks.40.ff.net.2.weight": 0.05791788845953324, + "parameter.transformer_blocks.40.norm1.weight": 0.9204294948298577, + "parameter.transformer_blocks.40.norm2.weight": 0.08970946823680186, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 0.07167722151177615, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 0.05385167994476892, + "parameter.transformer_blocks.41.attn.norm_k.weight": 0.7415739676265102, + "parameter.transformer_blocks.41.attn.norm_q.weight": 0.7214113518977477, + "parameter.transformer_blocks.41.attn.to_k.weight": 0.8375918144818602, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 0.7181446985424611, + "parameter.transformer_blocks.41.attn.to_q.weight": 0.8633605626921834, + "parameter.transformer_blocks.41.attn.to_v.weight": 0.7340670197934589, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 0.08418384433323474, + "parameter.transformer_blocks.41.ff.net.2.weight": 0.055784632945377714, + "parameter.transformer_blocks.41.norm1.weight": 0.7416192510696097, + "parameter.transformer_blocks.41.norm2.weight": 0.07178941417708969, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 0.06136013817900501, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 0.04869513884491029, + "parameter.transformer_blocks.42.attn.norm_k.weight": 0.5782331610559092, + "parameter.transformer_blocks.42.attn.norm_q.weight": 0.5610783050313803, + "parameter.transformer_blocks.42.attn.to_k.weight": 0.5961723816303608, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 0.4970380312375704, + "parameter.transformer_blocks.42.attn.to_q.weight": 0.6264928024228196, + "parameter.transformer_blocks.42.attn.to_v.weight": 0.5601926606479264, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 0.08107068095989993, + "parameter.transformer_blocks.42.ff.net.2.weight": 0.052568672636339825, + "parameter.transformer_blocks.42.norm1.weight": 0.611284124711771, + "parameter.transformer_blocks.42.norm2.weight": 0.08594602357157691, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 0.08455479479594155, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 0.06298888911994822, + "parameter.transformer_blocks.43.attn.norm_k.weight": 0.8367560596212958, + "parameter.transformer_blocks.43.attn.norm_q.weight": 0.8293012138736153, + "parameter.transformer_blocks.43.attn.to_k.weight": 0.8355610570416689, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 0.5048074999452841, + "parameter.transformer_blocks.43.attn.to_q.weight": 0.8589702799195638, + "parameter.transformer_blocks.43.attn.to_v.weight": 0.5237629013042355, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 0.08180036987950778, + "parameter.transformer_blocks.43.ff.net.2.weight": 0.05492635350283063, + "parameter.transformer_blocks.43.norm1.weight": 0.625254111632119, + "parameter.transformer_blocks.43.norm2.weight": 0.08185618169640825, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 0.0747035755085786, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 0.06660390438644587, + "parameter.transformer_blocks.44.attn.norm_k.weight": 0.7054325330091511, + "parameter.transformer_blocks.44.attn.norm_q.weight": 0.7153772458747647, + "parameter.transformer_blocks.44.attn.to_k.weight": 0.8151165065352909, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 0.4896220425182629, + "parameter.transformer_blocks.44.attn.to_q.weight": 0.8899012780688226, + "parameter.transformer_blocks.44.attn.to_v.weight": 0.5503025423219586, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 0.08021322348273459, + "parameter.transformer_blocks.44.ff.net.2.weight": 0.05349598808605537, + "parameter.transformer_blocks.44.norm1.weight": 0.6243989599839498, + "parameter.transformer_blocks.44.norm2.weight": 0.11799748672207468, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 0.0497453175382903, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 0.04457923582338126, + "parameter.transformer_blocks.45.attn.norm_k.weight": 0.7810134190748194, + "parameter.transformer_blocks.45.attn.norm_q.weight": 0.7955276983624541, + "parameter.transformer_blocks.45.attn.to_k.weight": 0.7142341183040185, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.3006170712170958, + "parameter.transformer_blocks.45.attn.to_q.weight": 0.5358664533020208, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.3333600839476919, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 0.06397931304789603, + "parameter.transformer_blocks.45.ff.net.2.weight": 0.04759898484212081, + "parameter.transformer_blocks.45.norm1.weight": 0.5074765325765653, + "parameter.transformer_blocks.45.norm2.weight": 0.07059256845556657, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 0.03973402481862122, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 0.039342269700501455, + "parameter.transformer_blocks.46.attn.norm_k.weight": 0.5281689931014756, + "parameter.transformer_blocks.46.attn.norm_q.weight": 0.5023460667545797, + "parameter.transformer_blocks.46.attn.to_k.weight": 0.5710493648291428, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.20458083566418112, + "parameter.transformer_blocks.46.attn.to_q.weight": 0.6043269611697017, + "parameter.transformer_blocks.46.attn.to_v.weight": 0.21222939602038893, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 0.053632445429609044, + "parameter.transformer_blocks.46.ff.net.2.weight": 0.043544126943884376, + "parameter.transformer_blocks.46.norm1.weight": 0.20378245936065947, + "parameter.transformer_blocks.46.norm2.weight": 0.12094285973244451, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 0.05022102495701803, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 0.04967724417388883, + "parameter.transformer_blocks.47.attn.norm_k.weight": 0.40581739556096796, + "parameter.transformer_blocks.47.attn.norm_q.weight": 0.40086424817021526, + "parameter.transformer_blocks.47.attn.to_k.weight": 0.30190780592594835, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 0.1607475406502512, + "parameter.transformer_blocks.47.attn.to_q.weight": 0.32646706522384894, + "parameter.transformer_blocks.47.attn.to_v.weight": 0.17795703897959772, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 0.05300780342409638, + "parameter.transformer_blocks.47.ff.net.2.weight": 0.03804820836236104, + "parameter.transformer_blocks.47.norm1.weight": 0.2770958821336741, + "parameter.transformer_blocks.47.norm2.weight": 0.05101560480227945, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 0.03333339913016972, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 0.03236733810896702, + "parameter.transformer_blocks.48.attn.norm_k.weight": 0.19457433853619419, + "parameter.transformer_blocks.48.attn.norm_q.weight": 0.2039308465959488, + "parameter.transformer_blocks.48.attn.to_k.weight": 0.20109957601376408, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 0.06825592517444747, + "parameter.transformer_blocks.48.attn.to_q.weight": 0.23789554363993, + "parameter.transformer_blocks.48.attn.to_v.weight": 0.06767046317004026, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 0.04273345975348288, + "parameter.transformer_blocks.48.ff.net.2.weight": 0.030467613849704606, + "parameter.transformer_blocks.48.norm1.weight": 0.1027296515699527, + "parameter.transformer_blocks.48.norm2.weight": 0.036571067769884724, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 0.08135452791594873, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 0.026706518747267025, + "parameter.transformer_blocks.49.attn.norm_k.weight": 0.21441411702278998, + "parameter.transformer_blocks.49.attn.norm_q.weight": 0.18738898861526473, + "parameter.transformer_blocks.49.attn.to_k.weight": 0.2729530254337735, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 0.05813700787592316, + "parameter.transformer_blocks.49.attn.to_q.weight": 0.7771356836911435, + "parameter.transformer_blocks.49.attn.to_v.weight": 0.0805187373024994, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 0.030523150592913086, + "parameter.transformer_blocks.49.ff.net.2.weight": 0.01776529192249012, + "parameter.transformer_blocks.49.norm1.weight": 0.7904321091518686, + "parameter.transformer_blocks.49.norm2.weight": 0.03715688505236718, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 0.17641330345085307, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 0.1601695998366226, + "parameter.transformer_blocks.5.attn.norm_k.weight": 0.9930043052829239, + "parameter.transformer_blocks.5.attn.norm_q.weight": 0.9196082022003691, + "parameter.transformer_blocks.5.attn.to_k.weight": 0.9654403710082802, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 0.3941884282252417, + "parameter.transformer_blocks.5.attn.to_q.weight": 1.0187611646153703, + "parameter.transformer_blocks.5.attn.to_v.weight": 0.42468143004203623, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 0.16811237874287088, + "parameter.transformer_blocks.5.ff.net.2.weight": 0.1576003582577836, + "parameter.transformer_blocks.5.norm1.weight": 0.76708667521147, + "parameter.transformer_blocks.5.norm2.weight": 0.37211491427211013, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 0.19024775578310005, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 0.15817875829841233, + "parameter.transformer_blocks.6.attn.norm_k.weight": 0.4710895866015602, + "parameter.transformer_blocks.6.attn.norm_q.weight": 0.43767176224906923, + "parameter.transformer_blocks.6.attn.to_k.weight": 0.49508575795553283, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 0.5608999448263835, + "parameter.transformer_blocks.6.attn.to_q.weight": 0.49082194900600334, + "parameter.transformer_blocks.6.attn.to_v.weight": 0.5702360994896027, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 0.1721629763762543, + "parameter.transformer_blocks.6.ff.net.2.weight": 0.15903645990316964, + "parameter.transformer_blocks.6.norm1.weight": 0.5200554294016089, + "parameter.transformer_blocks.6.norm2.weight": 0.16227311821904442, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 0.20821190039464113, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 0.1750668113552889, + "parameter.transformer_blocks.7.attn.norm_k.weight": 0.767506548561085, + "parameter.transformer_blocks.7.attn.norm_q.weight": 0.724381651371519, + "parameter.transformer_blocks.7.attn.to_k.weight": 0.7161623747587325, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 0.6452027625697531, + "parameter.transformer_blocks.7.attn.to_q.weight": 0.7234735440770655, + "parameter.transformer_blocks.7.attn.to_v.weight": 0.6694764858944942, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 0.18568970462110013, + "parameter.transformer_blocks.7.ff.net.2.weight": 0.1577253307671931, + "parameter.transformer_blocks.7.norm1.weight": 0.7940419976295847, + "parameter.transformer_blocks.7.norm2.weight": 0.1862207062091121, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 0.177743452577233, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 0.16363789239079013, + "parameter.transformer_blocks.8.attn.norm_k.weight": 0.7775983592302588, + "parameter.transformer_blocks.8.attn.norm_q.weight": 0.8258994911211479, + "parameter.transformer_blocks.8.attn.to_k.weight": 0.796482110330491, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 0.49648199918847563, + "parameter.transformer_blocks.8.attn.to_q.weight": 0.7766663436028602, + "parameter.transformer_blocks.8.attn.to_v.weight": 0.5437818585307683, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 0.17455968335439612, + "parameter.transformer_blocks.8.ff.net.2.weight": 0.15671060212477078, + "parameter.transformer_blocks.8.norm1.weight": 0.7127991715379555, + "parameter.transformer_blocks.8.norm2.weight": 0.1771072697254998, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 0.18711840743085115, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 0.1618146031833099, + "parameter.transformer_blocks.9.attn.norm_k.weight": 0.8153320809159224, + "parameter.transformer_blocks.9.attn.norm_q.weight": 0.8138568088872633, + "parameter.transformer_blocks.9.attn.to_k.weight": 0.7583346203665146, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 0.4967006237582202, + "parameter.transformer_blocks.9.attn.to_q.weight": 0.7777404983883319, + "parameter.transformer_blocks.9.attn.to_v.weight": 0.5323776028893205, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 0.1668908492336362, + "parameter.transformer_blocks.9.ff.net.2.weight": 0.15348178168709498, + "parameter.transformer_blocks.9.norm1.weight": 0.704262611952027, + "parameter.transformer_blocks.9.norm2.weight": 0.16444916255759165, + "video": 0.006463553178251248 + }, + "ensemble_max_abs": { + "audio": 25.881103515625, + "input_gradients.audio_hidden_states": 1659.3463134765625, + "input_gradients.encoder_hidden_states": 84.62007141113281, + "input_gradients.hidden_states": 1499.3985595703125, + "nodes.audio_proj_in": 0.0, + "nodes.audio_proj_in.grad": 459.2862548828125, + "nodes.audio_proj_out": 25.881103515625, + "nodes.audio_proj_out.grad": 0.4043922424316406, + "nodes.block.0": 140320768.0, + "nodes.block.0.grad": 0.0034225070849061012, + "nodes.block.1": 143748096.0, + "nodes.block.1.grad": 0.0011117004323750734, + "nodes.block.10": 162256896.0, + "nodes.block.10.grad": 0.00033802539110183716, + "nodes.block.11": 164428800.0, + "nodes.block.11.grad": 0.000311677431454882, + "nodes.block.12": 200382464.0, + "nodes.block.12.grad": 0.00032012685551308095, + "nodes.block.13": 205912064.0, + "nodes.block.13.grad": 0.0003194731834810227, + "nodes.block.14": 181512192.0, + "nodes.block.14.grad": 0.000302475003991276, + "nodes.block.15": 229742592.0, + "nodes.block.15.grad": 0.0002795343170873821, + "nodes.block.16": 291952640.0, + "nodes.block.16.grad": 0.00028924597427248955, + "nodes.block.17": 237170688.0, + "nodes.block.17.grad": 0.0002339307829970494, + "nodes.block.18": 187920384.0, + "nodes.block.18.grad": 0.0002275904844282195, + "nodes.block.19": 216483840.0, + "nodes.block.19.grad": 0.00020624657918233424, + "nodes.block.2": 149746688.0, + "nodes.block.2.grad": 0.0010053855367004871, + "nodes.block.20": 209188864.0, + "nodes.block.20.grad": 0.00018948683282360435, + "nodes.block.21": 199888896.0, + "nodes.block.21.grad": 0.00018563814228400588, + "nodes.block.22": 248029184.0, + "nodes.block.22.grad": 0.00018029732746072114, + "nodes.block.23": 290775040.0, + "nodes.block.23.grad": 0.00017719886091072112, + "nodes.block.24": 463339520.0, + "nodes.block.24.grad": 0.0001746931520756334, + "nodes.block.25": 753213440.0, + "nodes.block.25.grad": 0.00016678865358699113, + "nodes.block.26": 1228570624.0, + "nodes.block.26.grad": 0.0001614233769942075, + "nodes.block.27": 1964654592.0, + "nodes.block.27.grad": 0.00018509774236008525, + "nodes.block.28": 928874496.0, + "nodes.block.28.grad": 0.00013470835983753204, + "nodes.block.29": 952860672.0, + "nodes.block.29.grad": 0.00011252994590904564, + "nodes.block.3": 175917056.0, + "nodes.block.3.grad": 0.0008070977637544274, + "nodes.block.30": 1398390784.0, + "nodes.block.30.grad": 8.35541941341944e-05, + "nodes.block.31": 1785257984.0, + "nodes.block.31.grad": 8.580036228522658e-05, + "nodes.block.32": 1542389760.0, + "nodes.block.32.grad": 6.19091370026581e-05, + "nodes.block.33": 1992736768.0, + "nodes.block.33.grad": 5.435024286271073e-05, + "nodes.block.34": 2048032768.0, + "nodes.block.34.grad": 5.132899968884885e-05, + "nodes.block.35": 2313650176.0, + "nodes.block.35.grad": 4.8147980123758316e-05, + "nodes.block.36": 2989260800.0, + "nodes.block.36.grad": 3.9090195059543476e-05, + "nodes.block.37": 3730833408.0, + "nodes.block.37.grad": 3.7519050238188356e-05, + "nodes.block.38": 4208328704.0, + "nodes.block.38.grad": 3.63895705959294e-05, + "nodes.block.39": 5305368576.0, + "nodes.block.39.grad": 3.337875386932865e-05, + "nodes.block.4": 149050368.0, + "nodes.block.4.grad": 0.0005667342338711023, + "nodes.block.40": 5325160448.0, + "nodes.block.40.grad": 3.524516796460375e-05, + "nodes.block.41": 5474811904.0, + "nodes.block.41.grad": 3.4466203942429274e-05, + "nodes.block.42": 5105451008.0, + "nodes.block.42.grad": 3.930190359824337e-05, + "nodes.block.43": 10013900800.0, + "nodes.block.43.grad": 3.9428363379556686e-05, + "nodes.block.44": 14327480320.0, + "nodes.block.44.grad": 4.173372872173786e-05, + "nodes.block.45": 15591800832.0, + "nodes.block.45.grad": 4.137665018788539e-05, + "nodes.block.46": 16392912896.0, + "nodes.block.46.grad": 4.1100756789091974e-05, + "nodes.block.47": 24343740416.0, + "nodes.block.47.grad": 4.002563582616858e-05, + "nodes.block.48": 19063373824.0, + "nodes.block.48.grad": 3.985674993600696e-05, + "nodes.block.49": 23810801664.0, + "nodes.block.49.grad": 4.638408768187219e-09, + "nodes.block.5": 139992064.0, + "nodes.block.5.grad": 0.0005438930820673704, + "nodes.block.6": 236392448.0, + "nodes.block.6.grad": 0.0005069520557299256, + "nodes.block.7": 190355456.0, + "nodes.block.7.grad": 0.00044696839177049696, + "nodes.block.8": 188913664.0, + "nodes.block.8.grad": 0.0004436581220943481, + "nodes.block.9": 189696000.0, + "nodes.block.9.grad": 0.0004370134265627712, + "nodes.context_embedder": 0.12600326538085938, + "nodes.context_embedder.grad": 3.906439781188965, + "nodes.norm_out": 111.63232421875, + "nodes.norm_out.grad": 1.3662357330322266, + "nodes.proj_in": 0.0, + "nodes.proj_in.grad": 177.72679138183594, + "nodes.proj_out": 9.693748474121094, + "nodes.proj_out.grad": 0.004043877124786377, + "nodes.refiner.0": 46.240234375, + "nodes.refiner.0.grad": 0.05538301542401314, + "nodes.refiner.1": 43.7119140625, + "nodes.refiner.1.grad": 0.056103259325027466, + "nodes.time_embedder": 0.0, + "nodes.time_embedder.grad": 1092.195556640625, + "nodes.time_proj": 0.0, + "nodes.token_refiner": 0.13956832885742188, + "nodes.token_refiner.grad": 29.5113582611084, + "parameter.audio_proj_in.bias": 959.9775390625, + "parameter.audio_proj_in.weight": 1504.7664794921875, + "parameter.audio_proj_out.bias": 0.6827297210693359, + "parameter.audio_proj_out.weight": 3522.9375, + "parameter.context_embedder.bias": 6.652046203613281, + "parameter.context_embedder.weight": 15.703447341918945, + "parameter.norm_out.linear.bias": 3.800384521484375, + "parameter.norm_out.linear.weight": 967.58203125, + "parameter.norm_out.norm.weight": 1044.5390625, + "parameter.proj_in.bias": 283.2672119140625, + "parameter.proj_in.weight": 538.5294189453125, + "parameter.proj_out.bias": 0.009851455688476562, + "parameter.proj_out.weight": 8.053604125976562, + "parameter.time_embedder.linear_1.bias": 588.97265625, + "parameter.time_embedder.linear_1.weight": 588.4208984375, + "parameter.time_embedder.linear_2.bias": 1201.2255859375, + "parameter.time_embedder.linear_2.weight": 7405.4541015625, + "parameter.token_refiner.final_norm.weight": 89.69309997558594, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 83.28971862792969, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 83.48487854003906, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 15.937400817871094, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 21.675384521484375, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 14.408926010131836, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 14.414284706115723, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 27.380104064941406, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 46.86830520629883, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 79.89678955078125, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 116.41938781738281, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 8.864188194274902, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 8.92668342590332, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 1.5547832250595093, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 3.674546718597412, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 2.011857509613037, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 2.641394853591919, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 31.77730941772461, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 54.34901809692383, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 45.30865478515625, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 29.382301330566406, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 1196.275390625, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 26491.515625, + "parameter.transformer_blocks.0.attn.norm_k.weight": 4221.015625, + "parameter.transformer_blocks.0.attn.norm_q.weight": 3929.76025390625, + "parameter.transformer_blocks.0.attn.to_k.weight": 2417.5927734375, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 1356.8145751953125, + "parameter.transformer_blocks.0.attn.to_q.weight": 8516.501953125, + "parameter.transformer_blocks.0.attn.to_v.weight": 747.5717163085938, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 3595.708740234375, + "parameter.transformer_blocks.0.ff.net.2.weight": 708.46337890625, + "parameter.transformer_blocks.0.norm1.weight": 2407.3369140625, + "parameter.transformer_blocks.0.norm2.weight": 7685.876953125, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 26.141883850097656, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 6241.2802734375, + "parameter.transformer_blocks.1.attn.norm_k.weight": 7.931436538696289, + "parameter.transformer_blocks.1.attn.norm_q.weight": 7.464145660400391, + "parameter.transformer_blocks.1.attn.to_k.weight": 7.599443435668945, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 3.394611358642578, + "parameter.transformer_blocks.1.attn.to_q.weight": 6.851078033447266, + "parameter.transformer_blocks.1.attn.to_v.weight": 4.701122760772705, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 846.48876953125, + "parameter.transformer_blocks.1.ff.net.2.weight": 481.535400390625, + "parameter.transformer_blocks.1.norm1.weight": 18.98525047302246, + "parameter.transformer_blocks.1.norm2.weight": 1647.931640625, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 5.168201446533203, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 1230.390625, + "parameter.transformer_blocks.10.attn.norm_k.weight": 3.2422327995300293, + "parameter.transformer_blocks.10.attn.norm_q.weight": 2.38580060005188, + "parameter.transformer_blocks.10.attn.to_k.weight": 2.1780014038085938, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 0.6007388830184937, + "parameter.transformer_blocks.10.attn.to_q.weight": 1.3468154668807983, + "parameter.transformer_blocks.10.attn.to_v.weight": 1.3952274322509766, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 643.40576171875, + "parameter.transformer_blocks.10.ff.net.2.weight": 187.44125366210938, + "parameter.transformer_blocks.10.norm1.weight": 5.480815410614014, + "parameter.transformer_blocks.10.norm2.weight": 2093.111328125, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 17.13475799560547, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 3995.078125, + "parameter.transformer_blocks.11.attn.norm_k.weight": 2.1234676837921143, + "parameter.transformer_blocks.11.attn.norm_q.weight": 2.16512393951416, + "parameter.transformer_blocks.11.attn.to_k.weight": 1.8319860696792603, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 0.587449848651886, + "parameter.transformer_blocks.11.attn.to_q.weight": 1.8382810354232788, + "parameter.transformer_blocks.11.attn.to_v.weight": 0.7373056411743164, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 535.1204833984375, + "parameter.transformer_blocks.11.ff.net.2.weight": 328.0040283203125, + "parameter.transformer_blocks.11.norm1.weight": 4.554431438446045, + "parameter.transformer_blocks.11.norm2.weight": 3220.455078125, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 5.411174774169922, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 1297.1611328125, + "parameter.transformer_blocks.12.attn.norm_k.weight": 3.307208776473999, + "parameter.transformer_blocks.12.attn.norm_q.weight": 3.3316569328308105, + "parameter.transformer_blocks.12.attn.to_k.weight": 1.5812640190124512, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 0.367980033159256, + "parameter.transformer_blocks.12.attn.to_q.weight": 2.3092703819274902, + "parameter.transformer_blocks.12.attn.to_v.weight": 0.45317959785461426, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 244.15122985839844, + "parameter.transformer_blocks.12.ff.net.2.weight": 103.07870483398438, + "parameter.transformer_blocks.12.norm1.weight": 2.0407357215881348, + "parameter.transformer_blocks.12.norm2.weight": 526.4371337890625, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 7.499958038330078, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 1779.9697265625, + "parameter.transformer_blocks.13.attn.norm_k.weight": 2.817554473876953, + "parameter.transformer_blocks.13.attn.norm_q.weight": 2.7233643531799316, + "parameter.transformer_blocks.13.attn.to_k.weight": 0.8722900152206421, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 0.5390692353248596, + "parameter.transformer_blocks.13.attn.to_q.weight": 1.7306736707687378, + "parameter.transformer_blocks.13.attn.to_v.weight": 0.88811194896698, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 151.831787109375, + "parameter.transformer_blocks.13.ff.net.2.weight": 205.3848876953125, + "parameter.transformer_blocks.13.norm1.weight": 5.726235866546631, + "parameter.transformer_blocks.13.norm2.weight": 912.68212890625, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 5.064319610595703, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 1209.17236328125, + "parameter.transformer_blocks.14.attn.norm_k.weight": 1.3951852321624756, + "parameter.transformer_blocks.14.attn.norm_q.weight": 1.3705897331237793, + "parameter.transformer_blocks.14.attn.to_k.weight": 1.9639244079589844, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 0.41162002086639404, + "parameter.transformer_blocks.14.attn.to_q.weight": 1.3147220611572266, + "parameter.transformer_blocks.14.attn.to_v.weight": 1.1094074249267578, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 241.68624877929688, + "parameter.transformer_blocks.14.ff.net.2.weight": 93.27825927734375, + "parameter.transformer_blocks.14.norm1.weight": 3.7792177200317383, + "parameter.transformer_blocks.14.norm2.weight": 551.478759765625, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 6.347064971923828, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 1483.4775390625, + "parameter.transformer_blocks.15.attn.norm_k.weight": 2.051187038421631, + "parameter.transformer_blocks.15.attn.norm_q.weight": 2.2277870178222656, + "parameter.transformer_blocks.15.attn.to_k.weight": 0.8541905879974365, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 0.7474973201751709, + "parameter.transformer_blocks.15.attn.to_q.weight": 1.0169093608856201, + "parameter.transformer_blocks.15.attn.to_v.weight": 0.6167980432510376, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 696.7670288085938, + "parameter.transformer_blocks.15.ff.net.2.weight": 140.95977783203125, + "parameter.transformer_blocks.15.norm1.weight": 3.8490424156188965, + "parameter.transformer_blocks.15.norm2.weight": 1405.282958984375, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 5.253047466278076, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 1252.7545166015625, + "parameter.transformer_blocks.16.attn.norm_k.weight": 1.5553314685821533, + "parameter.transformer_blocks.16.attn.norm_q.weight": 1.941348671913147, + "parameter.transformer_blocks.16.attn.to_k.weight": 1.1562013626098633, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.4928184151649475, + "parameter.transformer_blocks.16.attn.to_q.weight": 0.8413535356521606, + "parameter.transformer_blocks.16.attn.to_v.weight": 0.7481282949447632, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 268.7306823730469, + "parameter.transformer_blocks.16.ff.net.2.weight": 95.50224304199219, + "parameter.transformer_blocks.16.norm1.weight": 3.8275866508483887, + "parameter.transformer_blocks.16.norm2.weight": 1096.85888671875, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 9.43035888671875, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 2227.4814453125, + "parameter.transformer_blocks.17.attn.norm_k.weight": 1.7146778106689453, + "parameter.transformer_blocks.17.attn.norm_q.weight": 1.5827388763427734, + "parameter.transformer_blocks.17.attn.to_k.weight": 1.1171233654022217, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 0.5116302967071533, + "parameter.transformer_blocks.17.attn.to_q.weight": 1.061017394065857, + "parameter.transformer_blocks.17.attn.to_v.weight": 0.7643781900405884, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 312.38629150390625, + "parameter.transformer_blocks.17.ff.net.2.weight": 247.89013671875, + "parameter.transformer_blocks.17.norm1.weight": 4.868851661682129, + "parameter.transformer_blocks.17.norm2.weight": 607.411865234375, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 7.408510208129883, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 1765.047607421875, + "parameter.transformer_blocks.18.attn.norm_k.weight": 1.7119410037994385, + "parameter.transformer_blocks.18.attn.norm_q.weight": 1.1268473863601685, + "parameter.transformer_blocks.18.attn.to_k.weight": 0.8185585141181946, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 0.5142297744750977, + "parameter.transformer_blocks.18.attn.to_q.weight": 0.8801079988479614, + "parameter.transformer_blocks.18.attn.to_v.weight": 0.8254555463790894, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 580.1936645507812, + "parameter.transformer_blocks.18.ff.net.2.weight": 428.6776123046875, + "parameter.transformer_blocks.18.norm1.weight": 2.4846653938293457, + "parameter.transformer_blocks.18.norm2.weight": 1136.5263671875, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 6.499603271484375, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 1539.50634765625, + "parameter.transformer_blocks.19.attn.norm_k.weight": 2.237013578414917, + "parameter.transformer_blocks.19.attn.norm_q.weight": 2.2605738639831543, + "parameter.transformer_blocks.19.attn.to_k.weight": 1.7174041271209717, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 0.6810566186904907, + "parameter.transformer_blocks.19.attn.to_q.weight": 1.4289164543151855, + "parameter.transformer_blocks.19.attn.to_v.weight": 1.0487899780273438, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 294.0188293457031, + "parameter.transformer_blocks.19.ff.net.2.weight": 261.71728515625, + "parameter.transformer_blocks.19.norm1.weight": 6.750887870788574, + "parameter.transformer_blocks.19.norm2.weight": 1107.74365234375, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 33.29774475097656, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 8001.673828125, + "parameter.transformer_blocks.2.attn.norm_k.weight": 10.58828353881836, + "parameter.transformer_blocks.2.attn.norm_q.weight": 9.941483497619629, + "parameter.transformer_blocks.2.attn.to_k.weight": 6.533878803253174, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 2.6041641235351562, + "parameter.transformer_blocks.2.attn.to_q.weight": 5.902656078338623, + "parameter.transformer_blocks.2.attn.to_v.weight": 5.091973304748535, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 561.8043212890625, + "parameter.transformer_blocks.2.ff.net.2.weight": 381.8466796875, + "parameter.transformer_blocks.2.norm1.weight": 10.919822692871094, + "parameter.transformer_blocks.2.norm2.weight": 1716.290283203125, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 7.788921356201172, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 1854.8935546875, + "parameter.transformer_blocks.20.attn.norm_k.weight": 2.3954052925109863, + "parameter.transformer_blocks.20.attn.norm_q.weight": 2.0874135494232178, + "parameter.transformer_blocks.20.attn.to_k.weight": 1.894758939743042, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 0.5271312594413757, + "parameter.transformer_blocks.20.attn.to_q.weight": 1.0881155729293823, + "parameter.transformer_blocks.20.attn.to_v.weight": 0.8003813028335571, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 155.64996337890625, + "parameter.transformer_blocks.20.ff.net.2.weight": 190.9737548828125, + "parameter.transformer_blocks.20.norm1.weight": 5.4173431396484375, + "parameter.transformer_blocks.20.norm2.weight": 334.54669189453125, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 4.380756378173828, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 1040.6484375, + "parameter.transformer_blocks.21.attn.norm_k.weight": 2.262387990951538, + "parameter.transformer_blocks.21.attn.norm_q.weight": 2.1462185382843018, + "parameter.transformer_blocks.21.attn.to_k.weight": 1.493720531463623, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 0.5406997203826904, + "parameter.transformer_blocks.21.attn.to_q.weight": 0.84881591796875, + "parameter.transformer_blocks.21.attn.to_v.weight": 0.9486402869224548, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 208.00839233398438, + "parameter.transformer_blocks.21.ff.net.2.weight": 202.0222930908203, + "parameter.transformer_blocks.21.norm1.weight": 4.693154335021973, + "parameter.transformer_blocks.21.norm2.weight": 287.727294921875, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 5.165304183959961, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 1219.8623046875, + "parameter.transformer_blocks.22.attn.norm_k.weight": 1.7254773378372192, + "parameter.transformer_blocks.22.attn.norm_q.weight": 1.5732178688049316, + "parameter.transformer_blocks.22.attn.to_k.weight": 1.8890459537506104, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 0.524356484413147, + "parameter.transformer_blocks.22.attn.to_q.weight": 0.6948188543319702, + "parameter.transformer_blocks.22.attn.to_v.weight": 1.7596900463104248, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 253.32228088378906, + "parameter.transformer_blocks.22.ff.net.2.weight": 169.9898681640625, + "parameter.transformer_blocks.22.norm1.weight": 7.527772426605225, + "parameter.transformer_blocks.22.norm2.weight": 492.393798828125, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 5.38870906829834, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 1281.716796875, + "parameter.transformer_blocks.23.attn.norm_k.weight": 1.782730221748352, + "parameter.transformer_blocks.23.attn.norm_q.weight": 1.8004608154296875, + "parameter.transformer_blocks.23.attn.to_k.weight": 2.6039834022521973, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 0.3485468327999115, + "parameter.transformer_blocks.23.attn.to_q.weight": 0.8478580117225647, + "parameter.transformer_blocks.23.attn.to_v.weight": 1.2505159378051758, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 173.40414428710938, + "parameter.transformer_blocks.23.ff.net.2.weight": 195.6923828125, + "parameter.transformer_blocks.23.norm1.weight": 5.559892654418945, + "parameter.transformer_blocks.23.norm2.weight": 351.5206298828125, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 4.866369247436523, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 1158.40771484375, + "parameter.transformer_blocks.24.attn.norm_k.weight": 1.8183107376098633, + "parameter.transformer_blocks.24.attn.norm_q.weight": 1.4892802238464355, + "parameter.transformer_blocks.24.attn.to_k.weight": 1.4131574630737305, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 0.6507848501205444, + "parameter.transformer_blocks.24.attn.to_q.weight": 0.6354339122772217, + "parameter.transformer_blocks.24.attn.to_v.weight": 0.9285950660705566, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 348.76458740234375, + "parameter.transformer_blocks.24.ff.net.2.weight": 370.705078125, + "parameter.transformer_blocks.24.norm1.weight": 6.474706649780273, + "parameter.transformer_blocks.24.norm2.weight": 533.75048828125, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 5.577777862548828, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 1323.35205078125, + "parameter.transformer_blocks.25.attn.norm_k.weight": 2.0326733589172363, + "parameter.transformer_blocks.25.attn.norm_q.weight": 1.8265665769577026, + "parameter.transformer_blocks.25.attn.to_k.weight": 2.347825050354004, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 0.5998079776763916, + "parameter.transformer_blocks.25.attn.to_q.weight": 1.2458218336105347, + "parameter.transformer_blocks.25.attn.to_v.weight": 1.4591829776763916, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 409.814697265625, + "parameter.transformer_blocks.25.ff.net.2.weight": 540.3759765625, + "parameter.transformer_blocks.25.norm1.weight": 2.727937698364258, + "parameter.transformer_blocks.25.norm2.weight": 473.1929931640625, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 4.299077033996582, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 1016.748291015625, + "parameter.transformer_blocks.26.attn.norm_k.weight": 2.916896343231201, + "parameter.transformer_blocks.26.attn.norm_q.weight": 2.9030942916870117, + "parameter.transformer_blocks.26.attn.to_k.weight": 2.8463592529296875, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.7373838424682617, + "parameter.transformer_blocks.26.attn.to_q.weight": 1.6248605251312256, + "parameter.transformer_blocks.26.attn.to_v.weight": 1.5674855709075928, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 259.17681884765625, + "parameter.transformer_blocks.26.ff.net.2.weight": 319.775390625, + "parameter.transformer_blocks.26.norm1.weight": 5.834235668182373, + "parameter.transformer_blocks.26.norm2.weight": 694.09619140625, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 4.972381114959717, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 723.10498046875, + "parameter.transformer_blocks.27.attn.norm_k.weight": 2.396589994430542, + "parameter.transformer_blocks.27.attn.norm_q.weight": 2.8167920112609863, + "parameter.transformer_blocks.27.attn.to_k.weight": 4.00289249420166, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 0.48014622926712036, + "parameter.transformer_blocks.27.attn.to_q.weight": 2.6316230297088623, + "parameter.transformer_blocks.27.attn.to_v.weight": 3.5429189205169678, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 211.07672119140625, + "parameter.transformer_blocks.27.ff.net.2.weight": 83.690185546875, + "parameter.transformer_blocks.27.norm1.weight": 17.755285263061523, + "parameter.transformer_blocks.27.norm2.weight": 436.00537109375, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 4.6777191162109375, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 1108.4140625, + "parameter.transformer_blocks.28.attn.norm_k.weight": 2.4593915939331055, + "parameter.transformer_blocks.28.attn.norm_q.weight": 2.287010669708252, + "parameter.transformer_blocks.28.attn.to_k.weight": 2.3165788650512695, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 0.8154343366622925, + "parameter.transformer_blocks.28.attn.to_q.weight": 2.055948257446289, + "parameter.transformer_blocks.28.attn.to_v.weight": 3.6044247150421143, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 163.99996948242188, + "parameter.transformer_blocks.28.ff.net.2.weight": 396.37841796875, + "parameter.transformer_blocks.28.norm1.weight": 7.464785575866699, + "parameter.transformer_blocks.28.norm2.weight": 983.402099609375, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 4.8386688232421875, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 1127.87890625, + "parameter.transformer_blocks.29.attn.norm_k.weight": 6.578639507293701, + "parameter.transformer_blocks.29.attn.norm_q.weight": 5.110691547393799, + "parameter.transformer_blocks.29.attn.to_k.weight": 4.1849684715271, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 0.9435173869132996, + "parameter.transformer_blocks.29.attn.to_q.weight": 2.8432507514953613, + "parameter.transformer_blocks.29.attn.to_v.weight": 3.2650294303894043, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 344.19805908203125, + "parameter.transformer_blocks.29.ff.net.2.weight": 282.1435546875, + "parameter.transformer_blocks.29.norm1.weight": 8.832938194274902, + "parameter.transformer_blocks.29.norm2.weight": 376.4471435546875, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 31.928050994873047, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 7585.13330078125, + "parameter.transformer_blocks.3.attn.norm_k.weight": 7.919702053070068, + "parameter.transformer_blocks.3.attn.norm_q.weight": 7.012223720550537, + "parameter.transformer_blocks.3.attn.to_k.weight": 4.903183460235596, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 3.3478007316589355, + "parameter.transformer_blocks.3.attn.to_q.weight": 11.044124603271484, + "parameter.transformer_blocks.3.attn.to_v.weight": 1.4261471033096313, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 801.4122924804688, + "parameter.transformer_blocks.3.ff.net.2.weight": 260.0557861328125, + "parameter.transformer_blocks.3.norm1.weight": 9.486372947692871, + "parameter.transformer_blocks.3.norm2.weight": 2734.337890625, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 5.7413330078125, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 1350.337890625, + "parameter.transformer_blocks.30.attn.norm_k.weight": 5.465303897857666, + "parameter.transformer_blocks.30.attn.norm_q.weight": 5.46435546875, + "parameter.transformer_blocks.30.attn.to_k.weight": 5.119855880737305, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 0.9348153471946716, + "parameter.transformer_blocks.30.attn.to_q.weight": 3.1151809692382812, + "parameter.transformer_blocks.30.attn.to_v.weight": 2.243969440460205, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 159.632080078125, + "parameter.transformer_blocks.30.ff.net.2.weight": 159.099365234375, + "parameter.transformer_blocks.30.norm1.weight": 11.394719123840332, + "parameter.transformer_blocks.30.norm2.weight": 301.135498046875, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 5.110774993896484, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 739.34765625, + "parameter.transformer_blocks.31.attn.norm_k.weight": 3.8568966388702393, + "parameter.transformer_blocks.31.attn.norm_q.weight": 3.906276226043701, + "parameter.transformer_blocks.31.attn.to_k.weight": 1.2982009649276733, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 0.7056084275245667, + "parameter.transformer_blocks.31.attn.to_q.weight": 1.1121916770935059, + "parameter.transformer_blocks.31.attn.to_v.weight": 1.6242420673370361, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 139.44873046875, + "parameter.transformer_blocks.31.ff.net.2.weight": 104.65277099609375, + "parameter.transformer_blocks.31.norm1.weight": 4.170856475830078, + "parameter.transformer_blocks.31.norm2.weight": 417.41650390625, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 5.7159857749938965, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 770.952392578125, + "parameter.transformer_blocks.32.attn.norm_k.weight": 5.170241832733154, + "parameter.transformer_blocks.32.attn.norm_q.weight": 5.211503982543945, + "parameter.transformer_blocks.32.attn.to_k.weight": 3.110686779022217, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 1.5664355754852295, + "parameter.transformer_blocks.32.attn.to_q.weight": 2.448791027069092, + "parameter.transformer_blocks.32.attn.to_v.weight": 1.7731201648712158, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 140.49258422851562, + "parameter.transformer_blocks.32.ff.net.2.weight": 68.61331176757812, + "parameter.transformer_blocks.32.norm1.weight": 9.787178039550781, + "parameter.transformer_blocks.32.norm2.weight": 297.69940185546875, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 4.766090393066406, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 774.0791015625, + "parameter.transformer_blocks.33.attn.norm_k.weight": 4.196714878082275, + "parameter.transformer_blocks.33.attn.norm_q.weight": 4.320964813232422, + "parameter.transformer_blocks.33.attn.to_k.weight": 2.6738462448120117, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 2.137948513031006, + "parameter.transformer_blocks.33.attn.to_q.weight": 1.688585877418518, + "parameter.transformer_blocks.33.attn.to_v.weight": 1.9710156917572021, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 206.24911499023438, + "parameter.transformer_blocks.33.ff.net.2.weight": 71.080810546875, + "parameter.transformer_blocks.33.norm1.weight": 10.011297225952148, + "parameter.transformer_blocks.33.norm2.weight": 667.790771484375, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 7.7424468994140625, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 1036.16064453125, + "parameter.transformer_blocks.34.attn.norm_k.weight": 8.191692352294922, + "parameter.transformer_blocks.34.attn.norm_q.weight": 8.252612113952637, + "parameter.transformer_blocks.34.attn.to_k.weight": 4.370151519775391, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 1.28891122341156, + "parameter.transformer_blocks.34.attn.to_q.weight": 2.719578742980957, + "parameter.transformer_blocks.34.attn.to_v.weight": 2.4563608169555664, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 129.66920471191406, + "parameter.transformer_blocks.34.ff.net.2.weight": 223.3360595703125, + "parameter.transformer_blocks.34.norm1.weight": 25.90180778503418, + "parameter.transformer_blocks.34.norm2.weight": 483.9815673828125, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 4.1862640380859375, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 971.1640625, + "parameter.transformer_blocks.35.attn.norm_k.weight": 3.796952247619629, + "parameter.transformer_blocks.35.attn.norm_q.weight": 3.838136672973633, + "parameter.transformer_blocks.35.attn.to_k.weight": 2.2202110290527344, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 1.6180669069290161, + "parameter.transformer_blocks.35.attn.to_q.weight": 2.895798683166504, + "parameter.transformer_blocks.35.attn.to_v.weight": 1.8437247276306152, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 138.53948974609375, + "parameter.transformer_blocks.35.ff.net.2.weight": 124.335205078125, + "parameter.transformer_blocks.35.norm1.weight": 4.438076972961426, + "parameter.transformer_blocks.35.norm2.weight": 495.2100830078125, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 6.160249710083008, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 1463.8109130859375, + "parameter.transformer_blocks.36.attn.norm_k.weight": 3.277141809463501, + "parameter.transformer_blocks.36.attn.norm_q.weight": 3.317197799682617, + "parameter.transformer_blocks.36.attn.to_k.weight": 5.186769962310791, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 0.866863489151001, + "parameter.transformer_blocks.36.attn.to_q.weight": 3.173929214477539, + "parameter.transformer_blocks.36.attn.to_v.weight": 1.6129945516586304, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 206.4375, + "parameter.transformer_blocks.36.ff.net.2.weight": 71.85504150390625, + "parameter.transformer_blocks.36.norm1.weight": 15.505533218383789, + "parameter.transformer_blocks.36.norm2.weight": 1674.610595703125, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 5.213738441467285, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 1062.2938232421875, + "parameter.transformer_blocks.37.attn.norm_k.weight": 2.9631030559539795, + "parameter.transformer_blocks.37.attn.norm_q.weight": 3.029543876647949, + "parameter.transformer_blocks.37.attn.to_k.weight": 3.477022409439087, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 0.7805371284484863, + "parameter.transformer_blocks.37.attn.to_q.weight": 1.8869693279266357, + "parameter.transformer_blocks.37.attn.to_v.weight": 1.9407142400741577, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 447.45355224609375, + "parameter.transformer_blocks.37.ff.net.2.weight": 49.300811767578125, + "parameter.transformer_blocks.37.norm1.weight": 7.276589870452881, + "parameter.transformer_blocks.37.norm2.weight": 364.8778076171875, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 4.798489093780518, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 683.1319580078125, + "parameter.transformer_blocks.38.attn.norm_k.weight": 4.893096923828125, + "parameter.transformer_blocks.38.attn.norm_q.weight": 4.971759796142578, + "parameter.transformer_blocks.38.attn.to_k.weight": 2.918351173400879, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 0.5770251154899597, + "parameter.transformer_blocks.38.attn.to_q.weight": 1.906826376914978, + "parameter.transformer_blocks.38.attn.to_v.weight": 1.2777994871139526, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 97.0000228881836, + "parameter.transformer_blocks.38.ff.net.2.weight": 127.248291015625, + "parameter.transformer_blocks.38.norm1.weight": 9.897371292114258, + "parameter.transformer_blocks.38.norm2.weight": 468.759765625, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 7.547196388244629, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 1791.341552734375, + "parameter.transformer_blocks.39.attn.norm_k.weight": 3.645379066467285, + "parameter.transformer_blocks.39.attn.norm_q.weight": 3.6634902954101562, + "parameter.transformer_blocks.39.attn.to_k.weight": 2.178628444671631, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 0.7987926602363586, + "parameter.transformer_blocks.39.attn.to_q.weight": 1.5225403308868408, + "parameter.transformer_blocks.39.attn.to_v.weight": 1.4502114057540894, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 83.56243896484375, + "parameter.transformer_blocks.39.ff.net.2.weight": 121.65234375, + "parameter.transformer_blocks.39.norm1.weight": 5.666716575622559, + "parameter.transformer_blocks.39.norm2.weight": 720.703857421875, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 21.32648468017578, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 5028.822265625, + "parameter.transformer_blocks.4.attn.norm_k.weight": 4.957897186279297, + "parameter.transformer_blocks.4.attn.norm_q.weight": 4.433900833129883, + "parameter.transformer_blocks.4.attn.to_k.weight": 6.1132001876831055, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 1.9584068059921265, + "parameter.transformer_blocks.4.attn.to_q.weight": 5.3016510009765625, + "parameter.transformer_blocks.4.attn.to_v.weight": 1.5592217445373535, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 773.1954345703125, + "parameter.transformer_blocks.4.ff.net.2.weight": 268.1669921875, + "parameter.transformer_blocks.4.norm1.weight": 14.660944938659668, + "parameter.transformer_blocks.4.norm2.weight": 887.4386596679688, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 4.005348205566406, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 854.2412109375, + "parameter.transformer_blocks.40.attn.norm_k.weight": 2.949293613433838, + "parameter.transformer_blocks.40.attn.norm_q.weight": 4.571164608001709, + "parameter.transformer_blocks.40.attn.to_k.weight": 3.0293209552764893, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 0.6123585104942322, + "parameter.transformer_blocks.40.attn.to_q.weight": 3.568161964416504, + "parameter.transformer_blocks.40.attn.to_v.weight": 1.8074942827224731, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 140.12953186035156, + "parameter.transformer_blocks.40.ff.net.2.weight": 152.8375244140625, + "parameter.transformer_blocks.40.norm1.weight": 5.113841533660889, + "parameter.transformer_blocks.40.norm2.weight": 129.9615478515625, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 4.217306137084961, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 987.54638671875, + "parameter.transformer_blocks.41.attn.norm_k.weight": 2.3821609020233154, + "parameter.transformer_blocks.41.attn.norm_q.weight": 2.412914752960205, + "parameter.transformer_blocks.41.attn.to_k.weight": 1.237536072731018, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 0.7576071619987488, + "parameter.transformer_blocks.41.attn.to_q.weight": 1.2888585329055786, + "parameter.transformer_blocks.41.attn.to_v.weight": 1.4355392456054688, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 127.60647583007812, + "parameter.transformer_blocks.41.ff.net.2.weight": 152.33447265625, + "parameter.transformer_blocks.41.norm1.weight": 3.5250301361083984, + "parameter.transformer_blocks.41.norm2.weight": 445.77001953125, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 3.4761009216308594, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 846.796875, + "parameter.transformer_blocks.42.attn.norm_k.weight": 2.687297821044922, + "parameter.transformer_blocks.42.attn.norm_q.weight": 2.7009596824645996, + "parameter.transformer_blocks.42.attn.to_k.weight": 1.5777289867401123, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 0.6879916191101074, + "parameter.transformer_blocks.42.attn.to_q.weight": 0.9663207530975342, + "parameter.transformer_blocks.42.attn.to_v.weight": 1.3828082084655762, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 321.3841552734375, + "parameter.transformer_blocks.42.ff.net.2.weight": 144.8499755859375, + "parameter.transformer_blocks.42.norm1.weight": 4.310583114624023, + "parameter.transformer_blocks.42.norm2.weight": 276.15570068359375, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 5.307106018066406, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 1281.81640625, + "parameter.transformer_blocks.43.attn.norm_k.weight": 6.124272346496582, + "parameter.transformer_blocks.43.attn.norm_q.weight": 6.2348408699035645, + "parameter.transformer_blocks.43.attn.to_k.weight": 2.432459831237793, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 0.7861360311508179, + "parameter.transformer_blocks.43.attn.to_q.weight": 1.5843608379364014, + "parameter.transformer_blocks.43.attn.to_v.weight": 1.4354814291000366, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 281.88043212890625, + "parameter.transformer_blocks.43.ff.net.2.weight": 68.20425415039062, + "parameter.transformer_blocks.43.norm1.weight": 8.903609275817871, + "parameter.transformer_blocks.43.norm2.weight": 188.57647705078125, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 11.816520690917969, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 2805.9755859375, + "parameter.transformer_blocks.44.attn.norm_k.weight": 2.603389263153076, + "parameter.transformer_blocks.44.attn.norm_q.weight": 2.5335259437561035, + "parameter.transformer_blocks.44.attn.to_k.weight": 1.9004267454147339, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 0.9370132684707642, + "parameter.transformer_blocks.44.attn.to_q.weight": 1.4908075332641602, + "parameter.transformer_blocks.44.attn.to_v.weight": 1.3552401065826416, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 266.840576171875, + "parameter.transformer_blocks.44.ff.net.2.weight": 86.6273193359375, + "parameter.transformer_blocks.44.norm1.weight": 7.830251693725586, + "parameter.transformer_blocks.44.norm2.weight": 565.53564453125, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 5.412269592285156, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 1297.621337890625, + "parameter.transformer_blocks.45.attn.norm_k.weight": 2.0145363807678223, + "parameter.transformer_blocks.45.attn.norm_q.weight": 2.071850299835205, + "parameter.transformer_blocks.45.attn.to_k.weight": 2.3817930221557617, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.45184069871902466, + "parameter.transformer_blocks.45.attn.to_q.weight": 30.458595275878906, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.9223560690879822, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 193.89520263671875, + "parameter.transformer_blocks.45.ff.net.2.weight": 166.10183715820312, + "parameter.transformer_blocks.45.norm1.weight": 8.281604766845703, + "parameter.transformer_blocks.45.norm2.weight": 309.4776611328125, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 6.8365631103515625, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 1614.33984375, + "parameter.transformer_blocks.46.attn.norm_k.weight": 2.728790283203125, + "parameter.transformer_blocks.46.attn.norm_q.weight": 2.4782252311706543, + "parameter.transformer_blocks.46.attn.to_k.weight": 2.4199938774108887, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.3956280052661896, + "parameter.transformer_blocks.46.attn.to_q.weight": 2.4051148891448975, + "parameter.transformer_blocks.46.attn.to_v.weight": 1.3095606565475464, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 296.1636047363281, + "parameter.transformer_blocks.46.ff.net.2.weight": 853.3876953125, + "parameter.transformer_blocks.46.norm1.weight": 9.261405944824219, + "parameter.transformer_blocks.46.norm2.weight": 2051.6435546875, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 17.834823608398438, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 4247.4453125, + "parameter.transformer_blocks.47.attn.norm_k.weight": 5.551021575927734, + "parameter.transformer_blocks.47.attn.norm_q.weight": 5.684200286865234, + "parameter.transformer_blocks.47.attn.to_k.weight": 12.534589767456055, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 1.4007599353790283, + "parameter.transformer_blocks.47.attn.to_q.weight": 4.060039520263672, + "parameter.transformer_blocks.47.attn.to_v.weight": 10.387222290039062, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 795.94580078125, + "parameter.transformer_blocks.47.ff.net.2.weight": 1438.3203125, + "parameter.transformer_blocks.47.norm1.weight": 10.565017700195312, + "parameter.transformer_blocks.47.norm2.weight": 1186.939453125, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 21.9805908203125, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 4901.0, + "parameter.transformer_blocks.48.attn.norm_k.weight": 2.7005364894866943, + "parameter.transformer_blocks.48.attn.norm_q.weight": 2.77207088470459, + "parameter.transformer_blocks.48.attn.to_k.weight": 4.525178909301758, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 1.0088205337524414, + "parameter.transformer_blocks.48.attn.to_q.weight": 6.6492791175842285, + "parameter.transformer_blocks.48.attn.to_v.weight": 2.435667037963867, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 954.7305297851562, + "parameter.transformer_blocks.48.ff.net.2.weight": 712.043212890625, + "parameter.transformer_blocks.48.norm1.weight": 5.328737258911133, + "parameter.transformer_blocks.48.norm2.weight": 1848.45703125, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 6.877416610717773, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 960.73828125, + "parameter.transformer_blocks.49.attn.norm_k.weight": 5.781886100769043, + "parameter.transformer_blocks.49.attn.norm_q.weight": 5.775971412658691, + "parameter.transformer_blocks.49.attn.to_k.weight": 120.58187866210938, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 1.635695219039917, + "parameter.transformer_blocks.49.attn.to_q.weight": 5951.6728515625, + "parameter.transformer_blocks.49.attn.to_v.weight": 12.527933120727539, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 1403.05126953125, + "parameter.transformer_blocks.49.ff.net.2.weight": 1077.126953125, + "parameter.transformer_blocks.49.norm1.weight": 88.3868179321289, + "parameter.transformer_blocks.49.norm2.weight": 731.34423828125, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 22.666481018066406, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 5374.4609375, + "parameter.transformer_blocks.5.attn.norm_k.weight": 3.1488566398620605, + "parameter.transformer_blocks.5.attn.norm_q.weight": 3.0029892921447754, + "parameter.transformer_blocks.5.attn.to_k.weight": 2.35954213142395, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 2.4877138137817383, + "parameter.transformer_blocks.5.attn.to_q.weight": 3.6635231971740723, + "parameter.transformer_blocks.5.attn.to_v.weight": 0.8009382486343384, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 417.8459167480469, + "parameter.transformer_blocks.5.ff.net.2.weight": 180.81723022460938, + "parameter.transformer_blocks.5.norm1.weight": 8.364011764526367, + "parameter.transformer_blocks.5.norm2.weight": 7164.46484375, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 6.457424163818359, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 1539.517333984375, + "parameter.transformer_blocks.6.attn.norm_k.weight": 3.098749876022339, + "parameter.transformer_blocks.6.attn.norm_q.weight": 3.6933846473693848, + "parameter.transformer_blocks.6.attn.to_k.weight": 2.2403249740600586, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 0.9580168724060059, + "parameter.transformer_blocks.6.attn.to_q.weight": 4.970716953277588, + "parameter.transformer_blocks.6.attn.to_v.weight": 1.5979580879211426, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 856.0050659179688, + "parameter.transformer_blocks.6.ff.net.2.weight": 134.00628662109375, + "parameter.transformer_blocks.6.norm1.weight": 5.737216949462891, + "parameter.transformer_blocks.6.norm2.weight": 756.50048828125, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 6.233695030212402, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 1038.2841796875, + "parameter.transformer_blocks.7.attn.norm_k.weight": 3.7008655071258545, + "parameter.transformer_blocks.7.attn.norm_q.weight": 3.2430977821350098, + "parameter.transformer_blocks.7.attn.to_k.weight": 1.7681100368499756, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 0.9606469869613647, + "parameter.transformer_blocks.7.attn.to_q.weight": 1.8657087087631226, + "parameter.transformer_blocks.7.attn.to_v.weight": 1.2660351991653442, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 1445.8360595703125, + "parameter.transformer_blocks.7.ff.net.2.weight": 104.09483337402344, + "parameter.transformer_blocks.7.norm1.weight": 11.101558685302734, + "parameter.transformer_blocks.7.norm2.weight": 1281.765869140625, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 9.229780197143555, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 2193.487548828125, + "parameter.transformer_blocks.8.attn.norm_k.weight": 1.500385046005249, + "parameter.transformer_blocks.8.attn.norm_q.weight": 1.459054946899414, + "parameter.transformer_blocks.8.attn.to_k.weight": 1.144398808479309, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 0.9251453876495361, + "parameter.transformer_blocks.8.attn.to_q.weight": 1.5949926376342773, + "parameter.transformer_blocks.8.attn.to_v.weight": 0.6618691086769104, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 718.069580078125, + "parameter.transformer_blocks.8.ff.net.2.weight": 114.80775451660156, + "parameter.transformer_blocks.8.norm1.weight": 3.614859104156494, + "parameter.transformer_blocks.8.norm2.weight": 2606.68359375, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 5.06758975982666, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 1203.7950439453125, + "parameter.transformer_blocks.9.attn.norm_k.weight": 2.6946144104003906, + "parameter.transformer_blocks.9.attn.norm_q.weight": 2.358290195465088, + "parameter.transformer_blocks.9.attn.to_k.weight": 1.6546334028244019, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 0.5754650235176086, + "parameter.transformer_blocks.9.attn.to_q.weight": 3.5623562335968018, + "parameter.transformer_blocks.9.attn.to_v.weight": 1.1003837585449219, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 407.1868896484375, + "parameter.transformer_blocks.9.ff.net.2.weight": 90.68083190917969, + "parameter.transformer_blocks.9.norm1.weight": 6.272876739501953, + "parameter.transformer_blocks.9.norm2.weight": 448.25537109375, + "video": 0.77642822265625 + }, + "bounds_relative_l2": { + "audio": 0.055874562519869435, + "input_gradients.audio_hidden_states": 0.5322077855718023, + "input_gradients.encoder_hidden_states": 0.8201112660925163, + "input_gradients.hidden_states": 0.570961125155766, + "nodes.audio_proj_in": 0.001, + "nodes.audio_proj_in.grad": 0.5647759182449158, + "nodes.audio_proj_out": 0.05583448800632685, + "nodes.audio_proj_out.grad": 0.05987456251986943, + "nodes.block.0": 0.020856273681871214, + "nodes.block.0.grad": 2.2223851901162193, + "nodes.block.1": 0.01963606863771287, + "nodes.block.1.grad": 2.1711386259276138, + "nodes.block.10": 0.025624545745054676, + "nodes.block.10.grad": 2.3827599847554692, + "nodes.block.11": 0.025985391477088116, + "nodes.block.11.grad": 2.3844366994775155, + "nodes.block.12": 0.026510570243016597, + "nodes.block.12.grad": 2.464081837113086, + "nodes.block.13": 0.027265984942517207, + "nodes.block.13.grad": 2.495653322894215, + "nodes.block.14": 0.027774847708480276, + "nodes.block.14.grad": 2.5094531896273926, + "nodes.block.15": 0.028750013046869617, + "nodes.block.15.grad": 2.5264188071831617, + "nodes.block.16": 0.02945524012717312, + "nodes.block.16.grad": 2.5791103746642503, + "nodes.block.17": 0.03002846913698215, + "nodes.block.17.grad": 2.597714194651892, + "nodes.block.18": 0.03270233573383202, + "nodes.block.18.grad": 2.5959965556849545, + "nodes.block.19": 0.03500400443185651, + "nodes.block.19.grad": 2.647298785191703, + "nodes.block.2": 0.02025170917229934, + "nodes.block.2.grad": 2.561946800938579, + "nodes.block.20": 0.03603438523204576, + "nodes.block.20.grad": 2.654514680209731, + "nodes.block.21": 0.036850674168122315, + "nodes.block.21.grad": 2.6803915843962587, + "nodes.block.22": 0.03625456245606433, + "nodes.block.22.grad": 2.7575750756326936, + "nodes.block.23": 0.03428651375960175, + "nodes.block.23.grad": 2.8131499193681324, + "nodes.block.24": 0.029550852410779385, + "nodes.block.24.grad": 2.877172825167145, + "nodes.block.25": 0.026820069691058424, + "nodes.block.25.grad": 2.8977817427103716, + "nodes.block.26": 0.02514449582691894, + "nodes.block.26.grad": 2.901376893866698, + "nodes.block.27": 0.05127208996315919, + "nodes.block.27.grad": 2.9120452755602653, + "nodes.block.28": 0.04379588326751174, + "nodes.block.28.grad": 2.937083189568119, + "nodes.block.29": 0.0431860226385089, + "nodes.block.29.grad": 2.803516189221703, + "nodes.block.3": 0.021995502188648874, + "nodes.block.3.grad": 2.360842556421442, + "nodes.block.30": 0.04213505991316521, + "nodes.block.30.grad": 2.64568640018582, + "nodes.block.31": 0.04331209295755942, + "nodes.block.31.grad": 2.7451072939109578, + "nodes.block.32": 0.040014718762034836, + "nodes.block.32.grad": 2.615002999331556, + "nodes.block.33": 0.040392899720599576, + "nodes.block.33.grad": 2.609574000490925, + "nodes.block.34": 0.04067843997826269, + "nodes.block.34.grad": 2.834503210447525, + "nodes.block.35": 0.03641098525432868, + "nodes.block.35.grad": 2.9165366050627264, + "nodes.block.36": 0.03297402876678478, + "nodes.block.36.grad": 2.8276718551698936, + "nodes.block.37": 0.03160852819071327, + "nodes.block.37.grad": 2.77520294901385, + "nodes.block.38": 0.03233105914246465, + "nodes.block.38.grad": 2.8283303457677595, + "nodes.block.39": 0.03216289190451429, + "nodes.block.39.grad": 2.8540157299753397, + "nodes.block.4": 0.023591763377124136, + "nodes.block.4.grad": 2.4098772178669474, + "nodes.block.40": 0.030933532049720696, + "nodes.block.40.grad": 2.8697604169567468, + "nodes.block.41": 0.03119100676158183, + "nodes.block.41.grad": 2.8825146889968614, + "nodes.block.42": 0.026289492238989076, + "nodes.block.42.grad": 3.0928782062807088, + "nodes.block.43": 0.020140799108251242, + "nodes.block.43.grad": 3.0576597403511196, + "nodes.block.44": 0.019506680953228443, + "nodes.block.44.grad": 3.013652237870902, + "nodes.block.45": 0.02115346901345907, + "nodes.block.45.grad": 3.0030285081630783, + "nodes.block.46": 0.02105395340522597, + "nodes.block.46.grad": 2.9673756043551, + "nodes.block.47": 0.016002538734894753, + "nodes.block.47.grad": 2.973243601760091, + "nodes.block.48": 0.018486101654955133, + "nodes.block.48.grad": 3.050383267832265, + "nodes.block.49": 0.02039832231310697, + "nodes.block.49.grad": 0.05773916498393257, + "nodes.block.5": 0.024130894300617163, + "nodes.block.5.grad": 2.270758112999734, + "nodes.block.6": 0.024268191662500875, + "nodes.block.6.grad": 2.3733618325946417, + "nodes.block.7": 0.024324743369856273, + "nodes.block.7.grad": 2.3597214946400884, + "nodes.block.8": 0.025012925768023244, + "nodes.block.8.grad": 2.3545103657037747, + "nodes.block.9": 0.025473614701246088, + "nodes.block.9.grad": 2.3469190835276272, + "nodes.context_embedder": 0.007199885083778527, + "nodes.context_embedder.grad": 0.8232039140070144, + "nodes.norm_out": 0.04593889301942204, + "nodes.norm_out.grad": 0.058526997835089106, + "nodes.proj_in": 0.001, + "nodes.proj_in.grad": 0.6102260678480784, + "nodes.proj_out": 0.0496271698625265, + "nodes.proj_out.grad": 0.021805234505499992, + "nodes.refiner.0": 0.014741722004295996, + "nodes.refiner.0.grad": 0.833198641928383, + "nodes.refiner.1": 0.016572113205235068, + "nodes.refiner.1.grad": 0.832172955022045, + "nodes.time_embedder": 0.001, + "nodes.time_embedder.grad": 1.3471330701300817, + "nodes.time_proj": 0.001, + "nodes.token_refiner": 0.01681921465529777, + "nodes.token_refiner.grad": 0.8266076567858954, + "parameter.audio_proj_in.bias": 0.4924406210254883, + "parameter.audio_proj_in.weight": 0.5645161238198128, + "parameter.audio_proj_out.bias": 0.04115526495753882, + "parameter.audio_proj_out.weight": 0.04956870486156885, + "parameter.context_embedder.bias": 0.8921846392352075, + "parameter.context_embedder.weight": 0.8225462304542064, + "parameter.norm_out.linear.bias": 0.04634026248145849, + "parameter.norm_out.linear.weight": 0.04625262281004235, + "parameter.norm_out.norm.weight": 0.03589542548700971, + "parameter.proj_in.bias": 0.5710820501685063, + "parameter.proj_in.weight": 0.6181584003502457, + "parameter.proj_out.bias": 0.019366519519612328, + "parameter.proj_out.weight": 0.033329808130047274, + "parameter.time_embedder.linear_1.bias": 0.09295946902006715, + "parameter.time_embedder.linear_1.weight": 0.08581757990692288, + "parameter.time_embedder.linear_2.bias": 1.2949600448078376, + "parameter.time_embedder.linear_2.weight": 0.3858624288589434, + "parameter.token_refiner.final_norm.weight": 0.8545591164818493, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 0.6977844496170106, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 0.696436119960276, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 0.8274289290532341, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 0.8548751035821915, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 0.7967628175599382, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 0.8552145288519823, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 0.8421136257842587, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 0.8541830952817251, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 0.8461923977581175, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 0.867156538337202, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 1.143592452947768, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 1.144926517758613, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 0.9203645497892541, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 0.876771597022621, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 0.9368330199913044, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 0.8917728579157279, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 0.9538617435520901, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 0.8776271074767351, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 0.6944313461838573, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 0.9626344119635454, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 0.8343857459191878, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 0.4695623484222038, + "parameter.transformer_blocks.0.attn.norm_k.weight": 0.7579721892678787, + "parameter.transformer_blocks.0.attn.norm_q.weight": 0.8173884177839518, + "parameter.transformer_blocks.0.attn.to_k.weight": 0.598048869065734, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 0.4607183400750576, + "parameter.transformer_blocks.0.attn.to_q.weight": 0.6095098053999932, + "parameter.transformer_blocks.0.attn.to_v.weight": 0.4565274641005145, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 0.47368161351989685, + "parameter.transformer_blocks.0.ff.net.2.weight": 0.46706704734661963, + "parameter.transformer_blocks.0.norm1.weight": 0.5690199829802196, + "parameter.transformer_blocks.0.norm2.weight": 0.4713643572348861, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 0.48625070765289713, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 0.46631960433629827, + "parameter.transformer_blocks.1.attn.norm_k.weight": 1.5928846935443048, + "parameter.transformer_blocks.1.attn.norm_q.weight": 1.6802262756178572, + "parameter.transformer_blocks.1.attn.to_k.weight": 1.7463617619368006, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 1.3992411530210536, + "parameter.transformer_blocks.1.attn.to_q.weight": 1.5308964954536752, + "parameter.transformer_blocks.1.attn.to_v.weight": 1.5078440376921296, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 0.48165415622186136, + "parameter.transformer_blocks.1.ff.net.2.weight": 0.4602352329789112, + "parameter.transformer_blocks.1.norm1.weight": 1.6781901777194925, + "parameter.transformer_blocks.1.norm2.weight": 0.4499919762222125, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 0.47148786502263684, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 0.4272072578235439, + "parameter.transformer_blocks.10.attn.norm_k.weight": 1.6046638120314505, + "parameter.transformer_blocks.10.attn.norm_q.weight": 1.6108075051113542, + "parameter.transformer_blocks.10.attn.to_k.weight": 1.546673665186303, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 1.3221676077517956, + "parameter.transformer_blocks.10.attn.to_q.weight": 1.5250613246280351, + "parameter.transformer_blocks.10.attn.to_v.weight": 1.439652952408004, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 0.4726732753392043, + "parameter.transformer_blocks.10.ff.net.2.weight": 0.3963444861154397, + "parameter.transformer_blocks.10.norm1.weight": 1.6772450755020096, + "parameter.transformer_blocks.10.norm2.weight": 0.6002672020768325, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 0.42361490222844905, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 0.39662066333443935, + "parameter.transformer_blocks.11.attn.norm_k.weight": 1.840572307580958, + "parameter.transformer_blocks.11.attn.norm_q.weight": 1.9188744283003905, + "parameter.transformer_blocks.11.attn.to_k.weight": 1.7703977150032377, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 1.2876843768443307, + "parameter.transformer_blocks.11.attn.to_q.weight": 1.737723211961002, + "parameter.transformer_blocks.11.attn.to_v.weight": 1.3995188087961634, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 0.4197019364819283, + "parameter.transformer_blocks.11.ff.net.2.weight": 0.37612409124730006, + "parameter.transformer_blocks.11.norm1.weight": 1.6079513206454785, + "parameter.transformer_blocks.11.norm2.weight": 0.458253057250772, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 0.43840053394416933, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 0.40671239044508817, + "parameter.transformer_blocks.12.attn.norm_k.weight": 1.7488444150259954, + "parameter.transformer_blocks.12.attn.norm_q.weight": 1.7632868615827981, + "parameter.transformer_blocks.12.attn.to_k.weight": 1.3142226322965935, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 1.1359072994114097, + "parameter.transformer_blocks.12.attn.to_q.weight": 1.4349954944475813, + "parameter.transformer_blocks.12.attn.to_v.weight": 1.2742728285071432, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 0.45273844195521246, + "parameter.transformer_blocks.12.ff.net.2.weight": 0.3847020860814938, + "parameter.transformer_blocks.12.norm1.weight": 1.3172006511246992, + "parameter.transformer_blocks.12.norm2.weight": 0.46636478397989606, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 0.44600884145665415, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 0.40674649755173176, + "parameter.transformer_blocks.13.attn.norm_k.weight": 1.7799406417442298, + "parameter.transformer_blocks.13.attn.norm_q.weight": 1.7573791861082868, + "parameter.transformer_blocks.13.attn.to_k.weight": 1.6074768422290602, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 1.6132410052966137, + "parameter.transformer_blocks.13.attn.to_q.weight": 1.5997128757317836, + "parameter.transformer_blocks.13.attn.to_v.weight": 1.6464668337321637, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 0.4442169048972016, + "parameter.transformer_blocks.13.ff.net.2.weight": 0.380902675691008, + "parameter.transformer_blocks.13.norm1.weight": 1.7105338913592616, + "parameter.transformer_blocks.13.norm2.weight": 0.46961815484790537, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 0.45725653215188217, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 0.41858402192577837, + "parameter.transformer_blocks.14.attn.norm_k.weight": 1.469986824442482, + "parameter.transformer_blocks.14.attn.norm_q.weight": 1.5763017849092813, + "parameter.transformer_blocks.14.attn.to_k.weight": 1.4510501970223415, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 1.2970826939393076, + "parameter.transformer_blocks.14.attn.to_q.weight": 1.5147021262359528, + "parameter.transformer_blocks.14.attn.to_v.weight": 1.4810157564197193, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 0.43087334365484353, + "parameter.transformer_blocks.14.ff.net.2.weight": 0.37976161539859005, + "parameter.transformer_blocks.14.norm1.weight": 1.7076552397075138, + "parameter.transformer_blocks.14.norm2.weight": 0.5074845627613467, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 0.4401152871230969, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 0.39990535961965595, + "parameter.transformer_blocks.15.attn.norm_k.weight": 1.7406358338702776, + "parameter.transformer_blocks.15.attn.norm_q.weight": 1.815698683567046, + "parameter.transformer_blocks.15.attn.to_k.weight": 1.6943845045609855, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 1.2497651908305334, + "parameter.transformer_blocks.15.attn.to_q.weight": 1.5719706664264341, + "parameter.transformer_blocks.15.attn.to_v.weight": 1.4362232343688441, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 0.4269020829501556, + "parameter.transformer_blocks.15.ff.net.2.weight": 0.3829018739751305, + "parameter.transformer_blocks.15.norm1.weight": 1.7213183235706304, + "parameter.transformer_blocks.15.norm2.weight": 0.5455163777631221, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 0.4520159206050519, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 0.39948408371174726, + "parameter.transformer_blocks.16.attn.norm_k.weight": 1.4119388371302533, + "parameter.transformer_blocks.16.attn.norm_q.weight": 1.4682436451837244, + "parameter.transformer_blocks.16.attn.to_k.weight": 1.3432969220005635, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.9832956575858338, + "parameter.transformer_blocks.16.attn.to_q.weight": 1.359900382306337, + "parameter.transformer_blocks.16.attn.to_v.weight": 1.201476601788808, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 0.4238987856910134, + "parameter.transformer_blocks.16.ff.net.2.weight": 0.35984166582667915, + "parameter.transformer_blocks.16.norm1.weight": 1.3843170680644228, + "parameter.transformer_blocks.16.norm2.weight": 0.428799335513681, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 0.4145135261066581, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 0.3799409518555693, + "parameter.transformer_blocks.17.attn.norm_k.weight": 1.3217790303538601, + "parameter.transformer_blocks.17.attn.norm_q.weight": 1.4572092735369861, + "parameter.transformer_blocks.17.attn.to_k.weight": 1.550763289841423, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 1.5443913216106686, + "parameter.transformer_blocks.17.attn.to_q.weight": 1.4996885508715714, + "parameter.transformer_blocks.17.attn.to_v.weight": 1.7509856021917902, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 0.3961476298950383, + "parameter.transformer_blocks.17.ff.net.2.weight": 0.3504688496738397, + "parameter.transformer_blocks.17.norm1.weight": 1.8017473556707249, + "parameter.transformer_blocks.17.norm2.weight": 0.3983444777914455, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 0.3778583747694606, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 0.3598241348103365, + "parameter.transformer_blocks.18.attn.norm_k.weight": 1.8639724127835577, + "parameter.transformer_blocks.18.attn.norm_q.weight": 1.6944261657118604, + "parameter.transformer_blocks.18.attn.to_k.weight": 1.6319025238001412, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 1.2363141576360999, + "parameter.transformer_blocks.18.attn.to_q.weight": 1.624264917378737, + "parameter.transformer_blocks.18.attn.to_v.weight": 1.5524415567182823, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 0.37286740575046284, + "parameter.transformer_blocks.18.ff.net.2.weight": 0.3202254288682409, + "parameter.transformer_blocks.18.norm1.weight": 1.6159961194038832, + "parameter.transformer_blocks.18.norm2.weight": 0.3818729100045425, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 0.3726465715841605, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 0.34752555444109356, + "parameter.transformer_blocks.19.attn.norm_k.weight": 1.6930014849151624, + "parameter.transformer_blocks.19.attn.norm_q.weight": 1.6630708529790197, + "parameter.transformer_blocks.19.attn.to_k.weight": 1.6588350140955876, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 1.2439010207019663, + "parameter.transformer_blocks.19.attn.to_q.weight": 1.6428512510788948, + "parameter.transformer_blocks.19.attn.to_v.weight": 1.484881926385726, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 0.33970880705001394, + "parameter.transformer_blocks.19.ff.net.2.weight": 0.268614767899265, + "parameter.transformer_blocks.19.norm1.weight": 1.657001015483316, + "parameter.transformer_blocks.19.norm2.weight": 0.3972401717652793, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 0.46606343973880315, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 0.4487492056614175, + "parameter.transformer_blocks.2.attn.norm_k.weight": 1.4482745787968618, + "parameter.transformer_blocks.2.attn.norm_q.weight": 1.4011792992179826, + "parameter.transformer_blocks.2.attn.to_k.weight": 1.383713706346312, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 1.2434077546585083, + "parameter.transformer_blocks.2.attn.to_q.weight": 1.2824894433942415, + "parameter.transformer_blocks.2.attn.to_v.weight": 1.6011841732416363, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 0.46201371424920007, + "parameter.transformer_blocks.2.ff.net.2.weight": 0.4449119826160648, + "parameter.transformer_blocks.2.norm1.weight": 1.4189648365457033, + "parameter.transformer_blocks.2.norm2.weight": 0.4706222547522539, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 0.40841601968208596, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 0.35530501345629856, + "parameter.transformer_blocks.20.attn.norm_k.weight": 2.271528269172145, + "parameter.transformer_blocks.20.attn.norm_q.weight": 2.3965219300340412, + "parameter.transformer_blocks.20.attn.to_k.weight": 2.233833032947429, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 1.0171265853519902, + "parameter.transformer_blocks.20.attn.to_q.weight": 2.0641527960619745, + "parameter.transformer_blocks.20.attn.to_v.weight": 1.2487110943334556, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 0.3693490359205701, + "parameter.transformer_blocks.20.ff.net.2.weight": 0.3074205930258824, + "parameter.transformer_blocks.20.norm1.weight": 1.7700395817002512, + "parameter.transformer_blocks.20.norm2.weight": 0.367692691166245, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 0.3903710767724551, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 0.34301517959049416, + "parameter.transformer_blocks.21.attn.norm_k.weight": 1.950064440905238, + "parameter.transformer_blocks.21.attn.norm_q.weight": 1.9790393030131417, + "parameter.transformer_blocks.21.attn.to_k.weight": 1.8696022230524951, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 1.3085905832706133, + "parameter.transformer_blocks.21.attn.to_q.weight": 1.7626791017272785, + "parameter.transformer_blocks.21.attn.to_v.weight": 1.6148204989527732, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 0.37421061782095605, + "parameter.transformer_blocks.21.ff.net.2.weight": 0.31188595625347326, + "parameter.transformer_blocks.21.norm1.weight": 1.6715098834355426, + "parameter.transformer_blocks.21.norm2.weight": 0.3780210765516755, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 0.39334408525058967, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 0.34932131130627625, + "parameter.transformer_blocks.22.attn.norm_k.weight": 1.7282131749289305, + "parameter.transformer_blocks.22.attn.norm_q.weight": 1.7361869124100267, + "parameter.transformer_blocks.22.attn.to_k.weight": 1.6211689913370957, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 1.7531777523551206, + "parameter.transformer_blocks.22.attn.to_q.weight": 1.5804957440034704, + "parameter.transformer_blocks.22.attn.to_v.weight": 2.0127702197095947, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 0.3758269799981911, + "parameter.transformer_blocks.22.ff.net.2.weight": 0.30036560489829417, + "parameter.transformer_blocks.22.norm1.weight": 1.746308692033695, + "parameter.transformer_blocks.22.norm2.weight": 0.38526873361961966, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 0.3659958885542907, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 0.32942885246364434, + "parameter.transformer_blocks.23.attn.norm_k.weight": 1.7767804893186219, + "parameter.transformer_blocks.23.attn.norm_q.weight": 1.7835005943173663, + "parameter.transformer_blocks.23.attn.to_k.weight": 1.801271713754693, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 1.2826408202288249, + "parameter.transformer_blocks.23.attn.to_q.weight": 1.6637272867807518, + "parameter.transformer_blocks.23.attn.to_v.weight": 1.6276718377292485, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 0.35111900532078266, + "parameter.transformer_blocks.23.ff.net.2.weight": 0.30229493242687083, + "parameter.transformer_blocks.23.norm1.weight": 1.8728925608820077, + "parameter.transformer_blocks.23.norm2.weight": 0.37504010259999715, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 0.370112243599527, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 0.3259608253987697, + "parameter.transformer_blocks.24.attn.norm_k.weight": 1.5540151225587742, + "parameter.transformer_blocks.24.attn.norm_q.weight": 1.5282074408083117, + "parameter.transformer_blocks.24.attn.to_k.weight": 1.4743448183198562, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 1.0091014859382852, + "parameter.transformer_blocks.24.attn.to_q.weight": 1.3878023913313464, + "parameter.transformer_blocks.24.attn.to_v.weight": 1.3814820433053931, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 0.3495862859832943, + "parameter.transformer_blocks.24.ff.net.2.weight": 0.22342188840254962, + "parameter.transformer_blocks.24.norm1.weight": 1.5432131235368212, + "parameter.transformer_blocks.24.norm2.weight": 0.3512437453403421, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 0.2887349002565304, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 0.27046913523446603, + "parameter.transformer_blocks.25.attn.norm_k.weight": 1.691868877233949, + "parameter.transformer_blocks.25.attn.norm_q.weight": 1.7010915057049818, + "parameter.transformer_blocks.25.attn.to_k.weight": 1.7368112778802502, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 1.350349317777754, + "parameter.transformer_blocks.25.attn.to_q.weight": 1.668223607985642, + "parameter.transformer_blocks.25.attn.to_v.weight": 1.723285127155597, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 0.28133247296190517, + "parameter.transformer_blocks.25.ff.net.2.weight": 0.20194025848149452, + "parameter.transformer_blocks.25.norm1.weight": 1.411724914676122, + "parameter.transformer_blocks.25.norm2.weight": 0.26834024435523535, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 0.3379136712161619, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 0.28810004498840786, + "parameter.transformer_blocks.26.attn.norm_k.weight": 1.6798287262637364, + "parameter.transformer_blocks.26.attn.norm_q.weight": 1.69104714770592, + "parameter.transformer_blocks.26.attn.to_k.weight": 1.692782637359654, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.557316213795948, + "parameter.transformer_blocks.26.attn.to_q.weight": 1.6866281235933018, + "parameter.transformer_blocks.26.attn.to_v.weight": 0.8505834862776024, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 0.3081469953555959, + "parameter.transformer_blocks.26.ff.net.2.weight": 0.220262802078313, + "parameter.transformer_blocks.26.norm1.weight": 1.3291321584681985, + "parameter.transformer_blocks.26.norm2.weight": 0.32260745381630884, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 0.3665319287861786, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 0.29816704388226967, + "parameter.transformer_blocks.27.attn.norm_k.weight": 2.015954579673838, + "parameter.transformer_blocks.27.attn.norm_q.weight": 2.0728614171140696, + "parameter.transformer_blocks.27.attn.to_k.weight": 2.119717052057698, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 1.655684299333085, + "parameter.transformer_blocks.27.attn.to_q.weight": 2.063800384141528, + "parameter.transformer_blocks.27.attn.to_v.weight": 1.903318012690791, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 0.3336286871587634, + "parameter.transformer_blocks.27.ff.net.2.weight": 0.2629234943665337, + "parameter.transformer_blocks.27.norm1.weight": 2.2664585814860914, + "parameter.transformer_blocks.27.norm2.weight": 0.3169720684762965, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 0.3515487162879209, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 0.26841340549235404, + "parameter.transformer_blocks.28.attn.norm_k.weight": 2.399781555047728, + "parameter.transformer_blocks.28.attn.norm_q.weight": 2.395155262876651, + "parameter.transformer_blocks.28.attn.to_k.weight": 2.4228271102309455, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 2.195929642104692, + "parameter.transformer_blocks.28.attn.to_q.weight": 2.47044428200628, + "parameter.transformer_blocks.28.attn.to_v.weight": 2.235409904366294, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 0.3231814220870359, + "parameter.transformer_blocks.28.ff.net.2.weight": 0.23873283403188625, + "parameter.transformer_blocks.28.norm1.weight": 2.406016592778906, + "parameter.transformer_blocks.28.norm2.weight": 0.35388536898791156, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 0.3370124618645231, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 0.2581199142347362, + "parameter.transformer_blocks.29.attn.norm_k.weight": 2.828870634002473, + "parameter.transformer_blocks.29.attn.norm_q.weight": 2.8394752158134153, + "parameter.transformer_blocks.29.attn.to_k.weight": 2.4587867193040096, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 2.1807317376716786, + "parameter.transformer_blocks.29.attn.to_q.weight": 2.477952801057237, + "parameter.transformer_blocks.29.attn.to_v.weight": 2.2819078609741155, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 0.31067077867960535, + "parameter.transformer_blocks.29.ff.net.2.weight": 0.2125025961953128, + "parameter.transformer_blocks.29.norm1.weight": 2.22937114072439, + "parameter.transformer_blocks.29.norm2.weight": 0.2985583794573764, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 0.4563955440751999, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 0.44056551548436124, + "parameter.transformer_blocks.3.attn.norm_k.weight": 1.9910700212191144, + "parameter.transformer_blocks.3.attn.norm_q.weight": 1.945933246514572, + "parameter.transformer_blocks.3.attn.to_k.weight": 2.1471948191904158, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 1.3353532506843855, + "parameter.transformer_blocks.3.attn.to_q.weight": 1.8455292260430465, + "parameter.transformer_blocks.3.attn.to_v.weight": 1.5236774521517638, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 0.45603384682915216, + "parameter.transformer_blocks.3.ff.net.2.weight": 0.43336542727846195, + "parameter.transformer_blocks.3.norm1.weight": 1.9361030097866772, + "parameter.transformer_blocks.3.norm2.weight": 0.46347801113442294, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 0.424860641324806, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 0.2575813696839202, + "parameter.transformer_blocks.30.attn.norm_k.weight": 2.4133075689771766, + "parameter.transformer_blocks.30.attn.norm_q.weight": 2.4555342549587347, + "parameter.transformer_blocks.30.attn.to_k.weight": 2.7248059348650067, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 1.7435173521915326, + "parameter.transformer_blocks.30.attn.to_q.weight": 2.5419536303615473, + "parameter.transformer_blocks.30.attn.to_v.weight": 1.9029702259619197, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 0.33588833172930505, + "parameter.transformer_blocks.30.ff.net.2.weight": 0.2380083386328826, + "parameter.transformer_blocks.30.norm1.weight": 2.505899745707479, + "parameter.transformer_blocks.30.norm2.weight": 0.324304447571616, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 0.38201682560719036, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 0.27693613816411955, + "parameter.transformer_blocks.31.attn.norm_k.weight": 1.4749752282586552, + "parameter.transformer_blocks.31.attn.norm_q.weight": 1.450698015116089, + "parameter.transformer_blocks.31.attn.to_k.weight": 1.7458419357349744, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 1.7270328460847992, + "parameter.transformer_blocks.31.attn.to_q.weight": 1.6576446872999169, + "parameter.transformer_blocks.31.attn.to_v.weight": 1.9190422292061655, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 0.3783161083398855, + "parameter.transformer_blocks.31.ff.net.2.weight": 0.26410263772251447, + "parameter.transformer_blocks.31.norm1.weight": 1.6875584481154922, + "parameter.transformer_blocks.31.norm2.weight": 0.36632858313046956, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 0.4020554204311656, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 0.23973686092285756, + "parameter.transformer_blocks.32.attn.norm_k.weight": 2.9716232546172625, + "parameter.transformer_blocks.32.attn.norm_q.weight": 2.906283565015692, + "parameter.transformer_blocks.32.attn.to_k.weight": 2.606234086669554, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 2.317841292977602, + "parameter.transformer_blocks.32.attn.to_q.weight": 2.5657513065502493, + "parameter.transformer_blocks.32.attn.to_v.weight": 2.3794686664168188, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 0.34719683801155393, + "parameter.transformer_blocks.32.ff.net.2.weight": 0.246765566298946, + "parameter.transformer_blocks.32.norm1.weight": 2.4357935968387205, + "parameter.transformer_blocks.32.norm2.weight": 0.34853550249950177, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 0.36846996686388106, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 0.2533834097236321, + "parameter.transformer_blocks.33.attn.norm_k.weight": 2.3854834345576035, + "parameter.transformer_blocks.33.attn.norm_q.weight": 2.3851789502334184, + "parameter.transformer_blocks.33.attn.to_k.weight": 2.166871112866884, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 2.202871406466205, + "parameter.transformer_blocks.33.attn.to_q.weight": 2.0966241850183223, + "parameter.transformer_blocks.33.attn.to_v.weight": 2.2126399275458635, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 0.3441799372592239, + "parameter.transformer_blocks.33.ff.net.2.weight": 0.22341773100887405, + "parameter.transformer_blocks.33.norm1.weight": 1.9732371533027735, + "parameter.transformer_blocks.33.norm2.weight": 0.38652812762723965, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 0.3554213000192706, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 0.22140519681851778, + "parameter.transformer_blocks.34.attn.norm_k.weight": 2.5507693494105754, + "parameter.transformer_blocks.34.attn.norm_q.weight": 2.514695418148331, + "parameter.transformer_blocks.34.attn.to_k.weight": 2.1610564534533907, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 2.0670438640047566, + "parameter.transformer_blocks.34.attn.to_q.weight": 1.9783468584060075, + "parameter.transformer_blocks.34.attn.to_v.weight": 2.2278909504303757, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 0.3168857391270072, + "parameter.transformer_blocks.34.ff.net.2.weight": 0.2229557077010477, + "parameter.transformer_blocks.34.norm1.weight": 2.449795373815762, + "parameter.transformer_blocks.34.norm2.weight": 0.3219383109635162, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 0.26992720396736086, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 0.20284057396161692, + "parameter.transformer_blocks.35.attn.norm_k.weight": 2.2032694832045756, + "parameter.transformer_blocks.35.attn.norm_q.weight": 2.1419146451614828, + "parameter.transformer_blocks.35.attn.to_k.weight": 2.1450069456596244, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 2.0838470797540802, + "parameter.transformer_blocks.35.attn.to_q.weight": 2.3149170156922634, + "parameter.transformer_blocks.35.attn.to_v.weight": 2.270605870587408, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 0.2935230806913727, + "parameter.transformer_blocks.35.ff.net.2.weight": 0.19840333184595743, + "parameter.transformer_blocks.35.norm1.weight": 2.006915113924567, + "parameter.transformer_blocks.35.norm2.weight": 0.34304758555408366, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 0.3757365784358437, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 0.2740854204169333, + "parameter.transformer_blocks.36.attn.norm_k.weight": 2.3829003296526037, + "parameter.transformer_blocks.36.attn.norm_q.weight": 2.471052061201031, + "parameter.transformer_blocks.36.attn.to_k.weight": 2.9427079084130043, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 1.332734461988801, + "parameter.transformer_blocks.36.attn.to_q.weight": 2.879402875824165, + "parameter.transformer_blocks.36.attn.to_v.weight": 1.441635881952929, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 0.34849143326375176, + "parameter.transformer_blocks.36.ff.net.2.weight": 0.2115231280461088, + "parameter.transformer_blocks.36.norm1.weight": 2.390337959761368, + "parameter.transformer_blocks.36.norm2.weight": 0.9862222229373728, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 0.32267481434729123, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 0.24647183914350898, + "parameter.transformer_blocks.37.attn.norm_k.weight": 2.3771887840171453, + "parameter.transformer_blocks.37.attn.norm_q.weight": 2.423704528655729, + "parameter.transformer_blocks.37.attn.to_k.weight": 2.6074986725020586, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 1.9361508720859442, + "parameter.transformer_blocks.37.attn.to_q.weight": 2.470932596225775, + "parameter.transformer_blocks.37.attn.to_v.weight": 2.019540099261265, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 0.3864463038599187, + "parameter.transformer_blocks.37.ff.net.2.weight": 0.1813498562255011, + "parameter.transformer_blocks.37.norm1.weight": 2.3005942876742864, + "parameter.transformer_blocks.37.norm2.weight": 0.46200849638649916, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 0.33107329032365335, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 0.1845978723866556, + "parameter.transformer_blocks.38.attn.norm_k.weight": 2.1939132281440314, + "parameter.transformer_blocks.38.attn.norm_q.weight": 2.1452150077606085, + "parameter.transformer_blocks.38.attn.to_k.weight": 2.1651576745401, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 2.1535469137370447, + "parameter.transformer_blocks.38.attn.to_q.weight": 2.279476597611475, + "parameter.transformer_blocks.38.attn.to_v.weight": 2.149915644738202, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 0.29984975443501904, + "parameter.transformer_blocks.38.ff.net.2.weight": 0.17507619815358705, + "parameter.transformer_blocks.38.norm1.weight": 1.9931094344507323, + "parameter.transformer_blocks.38.norm2.weight": 0.2890943272506466, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 0.2724759054232888, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 0.17329018219412629, + "parameter.transformer_blocks.39.attn.norm_k.weight": 2.518809405868853, + "parameter.transformer_blocks.39.attn.norm_q.weight": 2.5537298120473655, + "parameter.transformer_blocks.39.attn.to_k.weight": 2.558570850046743, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 1.9924087817450835, + "parameter.transformer_blocks.39.attn.to_q.weight": 2.6671613820786084, + "parameter.transformer_blocks.39.attn.to_v.weight": 2.0221826033394854, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 0.2630599179615183, + "parameter.transformer_blocks.39.ff.net.2.weight": 0.16353138563559005, + "parameter.transformer_blocks.39.norm1.weight": 2.027398404327182, + "parameter.transformer_blocks.39.norm2.weight": 0.436626928527809, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 0.45113642198656345, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 0.43101662337052293, + "parameter.transformer_blocks.4.attn.norm_k.weight": 2.107241376051506, + "parameter.transformer_blocks.4.attn.norm_q.weight": 2.004893866644943, + "parameter.transformer_blocks.4.attn.to_k.weight": 1.7392158631560324, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 1.6407321965312882, + "parameter.transformer_blocks.4.attn.to_q.weight": 1.7290204791402193, + "parameter.transformer_blocks.4.attn.to_v.weight": 1.8260754551141534, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 0.45999773423866647, + "parameter.transformer_blocks.4.ff.net.2.weight": 0.41413221430852154, + "parameter.transformer_blocks.4.norm1.weight": 1.8600126214364934, + "parameter.transformer_blocks.4.norm2.weight": 0.4403873494482417, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 0.2638772792673216, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 0.15919358069898773, + "parameter.transformer_blocks.40.attn.norm_k.weight": 2.4582355239320846, + "parameter.transformer_blocks.40.attn.norm_q.weight": 2.5766816188303494, + "parameter.transformer_blocks.40.attn.to_k.weight": 2.5084094776471106, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 1.8345278973957233, + "parameter.transformer_blocks.40.attn.to_q.weight": 2.7075145860176923, + "parameter.transformer_blocks.40.attn.to_v.weight": 1.9262344531233737, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 0.27317540061866874, + "parameter.transformer_blocks.40.ff.net.2.weight": 0.15558650999478643, + "parameter.transformer_blocks.40.norm1.weight": 2.3981166865576298, + "parameter.transformer_blocks.40.norm2.weight": 0.23824461741568484, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 0.191360775930618, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 0.1450143678563992, + "parameter.transformer_blocks.41.attn.norm_k.weight": 1.9330923158289264, + "parameter.transformer_blocks.41.attn.norm_q.weight": 1.8806695149341441, + "parameter.transformer_blocks.41.attn.to_k.weight": 2.1827387176528363, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 1.8721762162103988, + "parameter.transformer_blocks.41.attn.to_q.weight": 2.249737462999677, + "parameter.transformer_blocks.41.attn.to_v.weight": 1.913574251462993, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 0.22387799526641033, + "parameter.transformer_blocks.41.ff.net.2.weight": 0.15004004565798207, + "parameter.transformer_blocks.41.norm1.weight": 1.9332100527809852, + "parameter.transformer_blocks.41.norm2.weight": 0.19165247686043319, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 0.16453635926541302, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 0.13160736099676676, + "parameter.transformer_blocks.42.attn.norm_k.weight": 1.5084062187453637, + "parameter.transformer_blocks.42.attn.norm_q.weight": 1.463803593081589, + "parameter.transformer_blocks.42.attn.to_k.weight": 1.555048192238938, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 1.2972988812176829, + "parameter.transformer_blocks.42.attn.to_q.weight": 1.6338812862993308, + "parameter.transformer_blocks.42.attn.to_v.weight": 1.4615009176846085, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 0.21578377049573982, + "parameter.transformer_blocks.42.ff.net.2.weight": 0.14167854885448355, + "parameter.transformer_blocks.42.norm1.weight": 1.5943387242506044, + "parameter.transformer_blocks.42.norm2.weight": 0.22845966128609996, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 0.22484246646944805, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 0.1687711117118654, + "parameter.transformer_blocks.43.attn.norm_k.weight": 2.180565755015369, + "parameter.transformer_blocks.43.attn.norm_q.weight": 2.1611831560713997, + "parameter.transformer_blocks.43.attn.to_k.weight": 2.1774587483083394, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 1.3174994998577385, + "parameter.transformer_blocks.43.attn.to_q.weight": 2.238322727790866, + "parameter.transformer_blocks.43.attn.to_v.weight": 1.3667835433910123, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 0.21768096168672024, + "parameter.transformer_blocks.43.ff.net.2.weight": 0.14780851910735965, + "parameter.transformer_blocks.43.norm1.weight": 1.6306606902435092, + "parameter.transformer_blocks.43.norm2.weight": 0.21782607241066146, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 0.19922929632230438, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 0.17817015140475925, + "parameter.transformer_blocks.44.attn.norm_k.weight": 1.8391245858237928, + "parameter.transformer_blocks.44.attn.norm_q.weight": 1.864980839274388, + "parameter.transformer_blocks.44.attn.to_k.weight": 2.1243029169917564, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 1.2780173105474835, + "parameter.transformer_blocks.44.attn.to_q.weight": 2.3187433229789387, + "parameter.transformer_blocks.44.attn.to_v.weight": 1.4357866100370924, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 0.21355438105510993, + "parameter.transformer_blocks.44.ff.net.2.weight": 0.14408956902374398, + "parameter.transformer_blocks.44.norm1.weight": 1.6284372959582694, + "parameter.transformer_blocks.44.norm2.weight": 0.31179346547739417, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 0.1343378255995548, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 0.12090601314079129, + "parameter.transformer_blocks.45.attn.norm_k.weight": 2.0356348895945304, + "parameter.transformer_blocks.45.attn.norm_q.weight": 2.0733720157423807, + "parameter.transformer_blocks.45.attn.to_k.weight": 1.8620087075904481, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.7866043851644491, + "parameter.transformer_blocks.45.attn.to_q.weight": 1.3982527785852539, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.8717362182639989, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 0.1713462139245297, + "parameter.transformer_blocks.45.ff.net.2.weight": 0.1287573605895141, + "parameter.transformer_blocks.45.norm1.weight": 1.3244389846990698, + "parameter.transformer_blocks.45.norm2.weight": 0.18854067798447308, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 0.10830846452841518, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 0.10728990122130379, + "parameter.transformer_blocks.46.attn.norm_k.weight": 1.3782393820638366, + "parameter.transformer_blocks.46.attn.norm_q.weight": 1.3110997735619072, + "parameter.transformer_blocks.46.attn.to_k.weight": 1.4897283485557713, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.5369101727268709, + "parameter.transformer_blocks.46.attn.to_q.weight": 1.5762500990412245, + "parameter.transformer_blocks.46.attn.to_v.weight": 0.5567964296530112, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 0.14444435811698353, + "parameter.transformer_blocks.46.ff.net.2.weight": 0.11821473005409938, + "parameter.transformer_blocks.46.norm1.weight": 0.5348343943377146, + "parameter.transformer_blocks.46.norm2.weight": 0.3194514353043557, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 0.13557466488824688, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 0.13416083485211097, + "parameter.transformer_blocks.47.attn.norm_k.weight": 1.0601252284585165, + "parameter.transformer_blocks.47.attn.norm_q.weight": 1.0472470452425595, + "parameter.transformer_blocks.47.attn.to_k.weight": 0.7899602954074657, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 0.42294360569065315, + "parameter.transformer_blocks.47.attn.to_q.weight": 0.8538143695820073, + "parameter.transformer_blocks.47.attn.to_v.weight": 0.4676883013469541, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 0.1428202889026506, + "parameter.transformer_blocks.47.ff.net.2.weight": 0.1039253417421387, + "parameter.transformer_blocks.47.norm1.weight": 0.7254492935475526, + "parameter.transformer_blocks.47.norm2.weight": 0.1376405724859266, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 0.09166683773844128, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 0.08915507908331426, + "parameter.transformer_blocks.48.attn.norm_k.weight": 0.510893280194105, + "parameter.transformer_blocks.48.attn.norm_q.weight": 0.5352202011494669, + "parameter.transformer_blocks.48.attn.to_k.weight": 0.5278588976357866, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 0.18246540545356343, + "parameter.transformer_blocks.48.attn.to_q.weight": 0.623528413463818, + "parameter.transformer_blocks.48.attn.to_v.weight": 0.1809432042421047, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 0.1161069953590555, + "parameter.transformer_blocks.48.ff.net.2.weight": 0.08421579600923199, + "parameter.transformer_blocks.48.norm1.weight": 0.27209709408187704, + "parameter.transformer_blocks.48.norm2.weight": 0.1000847762017003, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 0.2165217725814667, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 0.07443694874289428, + "parameter.transformer_blocks.49.attn.norm_k.weight": 0.562476704259254, + "parameter.transformer_blocks.49.attn.norm_q.weight": 0.4922113703996883, + "parameter.transformer_blocks.49.attn.to_k.weight": 0.714677866127811, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 0.15615622047740021, + "parameter.transformer_blocks.49.attn.to_q.weight": 2.0255527775969733, + "parameter.transformer_blocks.49.attn.to_v.weight": 0.21434871698649843, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 0.08436019154157402, + "parameter.transformer_blocks.49.ff.net.2.weight": 0.05118975899847431, + "parameter.transformer_blocks.49.norm1.weight": 2.0601234837948583, + "parameter.transformer_blocks.49.norm2.weight": 0.10160790113615467, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 0.463674588972218, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 0.42144095957521877, + "parameter.transformer_blocks.5.attn.norm_k.weight": 2.586811193735602, + "parameter.transformer_blocks.5.attn.norm_q.weight": 2.3959813257209595, + "parameter.transformer_blocks.5.attn.to_k.weight": 2.5151449646215283, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 1.0298899133856283, + "parameter.transformer_blocks.5.attn.to_q.weight": 2.6537790279999625, + "parameter.transformer_blocks.5.attn.to_v.weight": 1.1091717181092942, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 0.44209218473146433, + "parameter.transformer_blocks.5.ff.net.2.weight": 0.41476093147023735, + "parameter.transformer_blocks.5.norm1.weight": 1.999425355549822, + "parameter.transformer_blocks.5.norm2.weight": 0.9724987771074863, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 0.49964416503606013, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 0.41626477157587205, + "parameter.transformer_blocks.6.attn.norm_k.weight": 1.2298329251640565, + "parameter.transformer_blocks.6.attn.norm_q.weight": 1.14294658184758, + "parameter.transformer_blocks.6.attn.to_k.weight": 1.2922229706843853, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 1.463339856548597, + "parameter.transformer_blocks.6.attn.to_q.weight": 1.2811370674156086, + "parameter.transformer_blocks.6.attn.to_v.weight": 1.487613858672967, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 0.4526237385782612, + "parameter.transformer_blocks.6.ff.net.2.weight": 0.4184947957482411, + "parameter.transformer_blocks.6.norm1.weight": 1.357144116444183, + "parameter.transformer_blocks.6.norm2.weight": 0.4269101073695155, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 0.546350941026067, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 0.46017370952375114, + "parameter.transformer_blocks.7.attn.norm_k.weight": 2.000517026258821, + "parameter.transformer_blocks.7.attn.norm_q.weight": 1.8883922935659494, + "parameter.transformer_blocks.7.attn.to_k.weight": 1.8670221743727042, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 1.682527182681358, + "parameter.transformer_blocks.7.attn.to_q.weight": 1.8860312146003704, + "parameter.transformer_blocks.7.attn.to_v.weight": 1.7456388633256847, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 0.48779323201486036, + "parameter.transformer_blocks.7.ff.net.2.weight": 0.4150858599947021, + "parameter.transformer_blocks.7.norm1.weight": 2.0695091938369203, + "parameter.transformer_blocks.7.norm2.weight": 0.4891738361436915, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 0.46713297670080584, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 0.43045852021605435, + "parameter.transformer_blocks.8.attn.norm_k.weight": 2.0267557339986726, + "parameter.transformer_blocks.8.attn.norm_q.weight": 2.1523386769149844, + "parameter.transformer_blocks.8.attn.to_k.weight": 2.0758534868592764, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 1.2958531978900365, + "parameter.transformer_blocks.8.attn.to_q.weight": 2.0243324933674365, + "parameter.transformer_blocks.8.attn.to_v.weight": 1.4188328321799977, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 0.45885517672142995, + "parameter.transformer_blocks.8.ff.net.2.weight": 0.41244756552440404, + "parameter.transformer_blocks.8.norm1.weight": 1.8582778459986842, + "parameter.transformer_blocks.8.norm2.weight": 0.46547890128629954, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 0.491507859320213, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 0.42571796827660574, + "parameter.transformer_blocks.9.attn.norm_k.weight": 2.124863410381398, + "parameter.transformer_blocks.9.attn.norm_q.weight": 2.1210277031068845, + "parameter.transformer_blocks.9.attn.to_k.weight": 1.9766700129529378, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 1.2964216217713724, + "parameter.transformer_blocks.9.attn.to_q.weight": 2.027125295809663, + "parameter.transformer_blocks.9.attn.to_v.weight": 1.389181767512233, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 0.4389162080074541, + "parameter.transformer_blocks.9.ff.net.2.weight": 0.40405263238644695, + "parameter.transformer_blocks.9.norm1.weight": 1.8360827910752702, + "parameter.transformer_blocks.9.norm2.weight": 0.4325678226497383, + "video": 0.017805238263453246 + } +} diff --git a/examples/models/minimax_h3/full-model-thresholds.json b/examples/models/minimax_h3/full-model-thresholds.json new file mode 100644 index 000000000..82b29d58c --- /dev/null +++ b/examples/models/minimax_h3/full-model-thresholds.json @@ -0,0 +1,3064 @@ +{ + "reference_report_sha256": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "other_report_sha256": "d49a80d6e438eb1f7d91193ad3f3ef5b5b3692b2c86a7532b86f23359ad33d79", + "status": "frozen_reference_bounds", + "rule": "1.5 * independent BF16 provider error + fixed floor", + "bounds": { + "video": { + "max_abs": 1.0716748962402343, + "relative_l2": 0.010294890439757714 + }, + "audio": { + "max_abs": 34.65894563293457, + "relative_l2": 0.030129260422720988 + }, + "nodes.audio_proj_in": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.audio_proj_in.grad": { + "max_abs": 656.3094281005859, + "relative_l2": 0.2961049755980539 + }, + "nodes.audio_proj_out": { + "max_abs": 34.65894563293457, + "relative_l2": 0.03022506434096451 + }, + "nodes.audio_proj_out.grad": { + "max_abs": 0.5465304005146027, + "relative_l2": 0.034129260422720985 + }, + "nodes.block.0": { + "max_abs": 86476800.001, + "relative_l2": 0.009254473682794001 + }, + "nodes.block.0.grad": { + "max_abs": 0.007892671618610621, + "relative_l2": 1.0282748009575065 + }, + "nodes.block.1": { + "max_abs": 135257088.001, + "relative_l2": 0.00974318315515435 + }, + "nodes.block.1.grad": { + "max_abs": 0.006624704105779529, + "relative_l2": 1.0305674729163048 + }, + "nodes.block.10": { + "max_abs": 243385344.001, + "relative_l2": 0.014183544640477025 + }, + "nodes.block.10.grad": { + "max_abs": 0.005376999407308176, + "relative_l2": 1.1788970789964635 + }, + "nodes.block.11": { + "max_abs": 246643200.001, + "relative_l2": 0.014389428832832524 + }, + "nodes.block.11.grad": { + "max_abs": 0.0053723131376318635, + "relative_l2": 1.1852776533292766 + }, + "nodes.block.12": { + "max_abs": 291520512.001, + "relative_l2": 0.014761840861898093 + }, + "nodes.block.12.grad": { + "max_abs": 0.005355812138877809, + "relative_l2": 1.2346247101218673 + }, + "nodes.block.13": { + "max_abs": 308868096.001, + "relative_l2": 0.015281297232727049 + }, + "nodes.block.13.grad": { + "max_abs": 0.005372488284483552, + "relative_l2": 1.2471660330417302 + }, + "nodes.block.14": { + "max_abs": 272268288.001, + "relative_l2": 0.015601342139096724 + }, + "nodes.block.14.grad": { + "max_abs": 0.005347652858181391, + "relative_l2": 1.2606747758455823 + }, + "nodes.block.15": { + "max_abs": 269202432.001, + "relative_l2": 0.016103171225580006 + }, + "nodes.block.15.grad": { + "max_abs": 0.005317293059197255, + "relative_l2": 1.2776078942880231 + }, + "nodes.block.16": { + "max_abs": 268137984.001, + "relative_l2": 0.016473838830060608 + }, + "nodes.block.16.grad": { + "max_abs": 0.005337768455210608, + "relative_l2": 1.2982750317435832 + }, + "nodes.block.17": { + "max_abs": 279490560.001, + "relative_l2": 0.01700976352461589 + }, + "nodes.block.17.grad": { + "max_abs": 0.005300611849525012, + "relative_l2": 1.304552845914669 + }, + "nodes.block.18": { + "max_abs": 281880576.001, + "relative_l2": 0.018818761997117264 + }, + "nodes.block.18.grad": { + "max_abs": 0.0053020529038622045, + "relative_l2": 1.3132094053007688 + }, + "nodes.block.19": { + "max_abs": 324725760.001, + "relative_l2": 0.020049400375340646 + }, + "nodes.block.19.grad": { + "max_abs": 0.005287827766442206, + "relative_l2": 1.331032060282588 + }, + "nodes.block.2": { + "max_abs": 123956736.001, + "relative_l2": 0.010425392138777783 + }, + "nodes.block.2.grad": { + "max_abs": 0.00622760480735451, + "relative_l2": 1.2071082994824671 + }, + "nodes.block.20": { + "max_abs": 313783296.001, + "relative_l2": 0.02069248717055676 + }, + "nodes.block.20.grad": { + "max_abs": 0.005284230249235407, + "relative_l2": 1.3333540488485576 + }, + "nodes.block.21": { + "max_abs": 299833344.001, + "relative_l2": 0.0211042005101494 + }, + "nodes.block.21.grad": { + "max_abs": 0.005257708161370829, + "relative_l2": 1.3444240875437083 + }, + "nodes.block.22": { + "max_abs": 372043776.001, + "relative_l2": 0.020841782421864578 + }, + "nodes.block.22.grad": { + "max_abs": 0.005248349679168314, + "relative_l2": 1.371844846626805 + }, + "nodes.block.23": { + "max_abs": 436162560.001, + "relative_l2": 0.019816861965668092 + }, + "nodes.block.23.grad": { + "max_abs": 0.005236955092987046, + "relative_l2": 1.3894633432736114 + }, + "nodes.block.24": { + "max_abs": 695009280.001, + "relative_l2": 0.016932971507712787 + }, + "nodes.block.24.grad": { + "max_abs": 0.005229643585043959, + "relative_l2": 1.418867550586571 + }, + "nodes.block.25": { + "max_abs": 1129820160.001, + "relative_l2": 0.015896194052533705 + }, + "nodes.block.25.grad": { + "max_abs": 0.005218857559957542, + "relative_l2": 1.4337955436014975 + }, + "nodes.block.26": { + "max_abs": 1842855936.001, + "relative_l2": 0.014708202945811977 + }, + "nodes.block.26.grad": { + "max_abs": 0.005217536799027585, + "relative_l2": 1.4267937146443628 + }, + "nodes.block.27": { + "max_abs": 2946981888.001, + "relative_l2": 0.029825333550548658 + }, + "nodes.block.27.grad": { + "max_abs": 0.005218868648516946, + "relative_l2": 1.42950586952009 + }, + "nodes.block.28": { + "max_abs": 1393311744.001, + "relative_l2": 0.02539093504291072 + }, + "nodes.block.28.grad": { + "max_abs": 0.005193499094166327, + "relative_l2": 1.4388679605626724 + }, + "nodes.block.29": { + "max_abs": 1097785344.001, + "relative_l2": 0.024924610281992207 + }, + "nodes.block.29.grad": { + "max_abs": 0.0051151480319094845, + "relative_l2": 1.2236330922921828 + }, + "nodes.block.3": { + "max_abs": 131919360.001, + "relative_l2": 0.011420138219926868 + }, + "nodes.block.3.grad": { + "max_abs": 0.005938520999625325, + "relative_l2": 1.136913714467664 + }, + "nodes.block.30": { + "max_abs": 1956667392.001, + "relative_l2": 0.023872450055785956 + }, + "nodes.block.30.grad": { + "max_abs": 0.005070673440859537, + "relative_l2": 1.16461148905023 + }, + "nodes.block.31": { + "max_abs": 2265833472.001, + "relative_l2": 0.024456562630709992 + }, + "nodes.block.31.grad": { + "max_abs": 0.005072068880399456, + "relative_l2": 1.194078650490841 + }, + "nodes.block.32": { + "max_abs": 2067637248.001, + "relative_l2": 0.023105217172228312 + }, + "nodes.block.32.grad": { + "max_abs": 0.005050663617003011, + "relative_l2": 1.1115885666094312 + }, + "nodes.block.33": { + "max_abs": 2405548032.001, + "relative_l2": 0.023271942704352414 + }, + "nodes.block.33.grad": { + "max_abs": 0.005049794474762166, + "relative_l2": 1.1311506019932958 + }, + "nodes.block.34": { + "max_abs": 2844155904.001, + "relative_l2": 0.02325442521193387 + }, + "nodes.block.34.grad": { + "max_abs": 0.005045931737986393, + "relative_l2": 1.2013872772100338 + }, + "nodes.block.35": { + "max_abs": 3470475264.001, + "relative_l2": 0.02134088292154418 + }, + "nodes.block.35.grad": { + "max_abs": 0.005041465970753052, + "relative_l2": 1.233510966039407 + }, + "nodes.block.36": { + "max_abs": 4340367360.001, + "relative_l2": 0.0180463405343924 + }, + "nodes.block.36.grad": { + "max_abs": 0.005039195159843075, + "relative_l2": 1.2096277002335567 + }, + "nodes.block.37": { + "max_abs": 5596250112.001, + "relative_l2": 0.017395995963329733 + }, + "nodes.block.37.grad": { + "max_abs": 0.005037690586395911, + "relative_l2": 1.1923039041300327 + }, + "nodes.block.38": { + "max_abs": 6312493056.001, + "relative_l2": 0.01715249891038743 + }, + "nodes.block.38.grad": { + "max_abs": 0.005036132002605882, + "relative_l2": 1.218786553128028 + }, + "nodes.block.39": { + "max_abs": 7958052864.001, + "relative_l2": 0.01719502785204246 + }, + "nodes.block.39.grad": { + "max_abs": 0.005035766980545304, + "relative_l2": 1.2318844221994263 + }, + "nodes.block.4": { + "max_abs": 169340928.001, + "relative_l2": 0.012555142596216557 + }, + "nodes.block.4.grad": { + "max_abs": 0.0056757682422176004, + "relative_l2": 1.17920320434519 + }, + "nodes.block.40": { + "max_abs": 7987740672.001, + "relative_l2": 0.017254775666558798 + }, + "nodes.block.40.grad": { + "max_abs": 0.0050362005121132826, + "relative_l2": 1.2320456657478307 + }, + "nodes.block.41": { + "max_abs": 8212217856.001, + "relative_l2": 0.017141120518213926 + }, + "nodes.block.41.grad": { + "max_abs": 0.005036102641388425, + "relative_l2": 1.2386607608755411 + }, + "nodes.block.42": { + "max_abs": 7658176512.001, + "relative_l2": 0.015161639630964174 + }, + "nodes.block.42.grad": { + "max_abs": 0.005034739564398478, + "relative_l2": 1.2944377183229072 + }, + "nodes.block.43": { + "max_abs": 15020851200.001, + "relative_l2": 0.010976331445410721 + }, + "nodes.block.43.grad": { + "max_abs": 0.005036297476981417, + "relative_l2": 1.2822040421243264 + }, + "nodes.block.44": { + "max_abs": 21491220480.001, + "relative_l2": 0.010692049224040517 + }, + "nodes.block.44.grad": { + "max_abs": 0.005036446064768825, + "relative_l2": 1.274834343770201 + }, + "nodes.block.45": { + "max_abs": 23387701248.001, + "relative_l2": 0.011531374588086778 + }, + "nodes.block.45.grad": { + "max_abs": 0.005037532945498242, + "relative_l2": 1.287014693377177 + }, + "nodes.block.46": { + "max_abs": 24589369344.001, + "relative_l2": 0.011653401897591616 + }, + "nodes.block.46.grad": { + "max_abs": 0.005035186672903365, + "relative_l2": 1.2755247924762254 + }, + "nodes.block.47": { + "max_abs": 36515610624.001, + "relative_l2": 0.009354598313891512 + }, + "nodes.block.47.grad": { + "max_abs": 0.005034289247196284, + "relative_l2": 1.2837355978859173 + }, + "nodes.block.48": { + "max_abs": 28595060736.001, + "relative_l2": 0.00991133297587967 + }, + "nodes.block.48.grad": { + "max_abs": 0.005034035918361042, + "relative_l2": 1.3231800171294437 + }, + "nodes.block.49": { + "max_abs": 35716202496.001, + "relative_l2": 0.01111432926598677 + }, + "nodes.block.49.grad": { + "max_abs": 0.005000006957613152, + "relative_l2": 0.03542644133688417 + }, + "nodes.block.5": { + "max_abs": 209988096.001, + "relative_l2": 0.01310967836419382 + }, + "nodes.block.5.grad": { + "max_abs": 0.005638774537947029, + "relative_l2": 1.1153482427434613 + }, + "nodes.block.6": { + "max_abs": 286685184.001, + "relative_l2": 0.013200938568209863 + }, + "nodes.block.6.grad": { + "max_abs": 0.005473614592920058, + "relative_l2": 1.1466965710007662 + }, + "nodes.block.7": { + "max_abs": 223693824.001, + "relative_l2": 0.013329474464451103 + }, + "nodes.block.7.grad": { + "max_abs": 0.005471000399556942, + "relative_l2": 1.1582630616146892 + }, + "nodes.block.8": { + "max_abs": 283370496.001, + "relative_l2": 0.013902013824406895 + }, + "nodes.block.8.grad": { + "max_abs": 0.005416834795032628, + "relative_l2": 1.1586090943783933 + }, + "nodes.block.9": { + "max_abs": 284544000.001, + "relative_l2": 0.014139204849605763 + }, + "nodes.block.9.grad": { + "max_abs": 0.0054162817203905434, + "relative_l2": 1.160017287530927 + }, + "nodes.context_embedder": { + "max_abs": 0.19000489807128906, + "relative_l2": 0.004576856779102996 + }, + "nodes.context_embedder.grad": { + "max_abs": 4.159599070549011, + "relative_l2": 0.4770407196194314 + }, + "nodes.norm_out": { + "max_abs": 79.44240625, + "relative_l2": 0.0241334814018383 + }, + "nodes.norm_out.grad": { + "max_abs": 1.6565254974365233, + "relative_l2": 0.034355922337465906 + }, + "nodes.proj_in": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.proj_in.grad": { + "max_abs": 192.61858642578125, + "relative_l2": 0.3541688852969683 + }, + "nodes.proj_out": { + "max_abs": 14.54162271118164, + "relative_l2": 0.0290541364591499 + }, + "nodes.proj_out.grad": { + "max_abs": 0.010576401948928834, + "relative_l2": 0.014294886490478394 + }, + "nodes.refiner.0": { + "max_abs": 69.3613515625, + "relative_l2": 0.008927916540939999 + }, + "nodes.refiner.0.grad": { + "max_abs": 0.08660967938601971, + "relative_l2": 0.48280690880483634 + }, + "nodes.refiner.1": { + "max_abs": 65.56887109375, + "relative_l2": 0.009971913328415988 + }, + "nodes.refiner.1.grad": { + "max_abs": 0.0854927796125412, + "relative_l2": 0.48221516635887207 + }, + "nodes.time_embedder": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.time_embedder.grad": { + "max_abs": 1082.1651867675782, + "relative_l2": 0.3660479013861043 + }, + "nodes.time_proj": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.token_refiner": { + "max_abs": 0.2103524932861328, + "relative_l2": 0.010126469993441023 + }, + "nodes.token_refiner.grad": { + "max_abs": 42.15805328369141, + "relative_l2": 0.47900441737647814 + }, + "input_gradients.audio_hidden_states": { + "max_abs": 2186.4831494140626, + "relative_l2": 0.2894651121905773 + }, + "input_gradients.encoder_hidden_states": { + "max_abs": 126.93510711669921, + "relative_l2": 0.4752564996687594 + }, + "input_gradients.hidden_states": { + "max_abs": 1468.9695385742189, + "relative_l2": 0.28346112575648574 + }, + "parameter.audio_proj_in.bias": { + "max_abs": 1127.2670849609376, + "relative_l2": 0.27396454330724807 + }, + "parameter.audio_proj_in.weight": { + "max_abs": 2011.3793896484376, + "relative_l2": 0.30171097338711617 + }, + "parameter.audio_proj_out.bias": { + "max_abs": 1.0290945816040038, + "relative_l2": 0.025858806706272403 + }, + "parameter.audio_proj_out.weight": { + "max_abs": 5284.41125, + "relative_l2": 0.030712714343212797 + }, + "parameter.context_embedder.bias": { + "max_abs": 9.983069305419923, + "relative_l2": 0.5168372918664659 + }, + "parameter.context_embedder.weight": { + "max_abs": 23.560171012878417, + "relative_l2": 0.4766612868005037 + }, + "parameter.norm_out.linear.bias": { + "max_abs": 5.705576782226562, + "relative_l2": 0.028850151431610672 + }, + "parameter.norm_out.linear.weight": { + "max_abs": 1451.378046875, + "relative_l2": 0.028799590082716747 + }, + "parameter.norm_out.norm.weight": { + "max_abs": 1272.362421875, + "relative_l2": 0.02272375101128014 + }, + "parameter.proj_in.bias": { + "max_abs": 388.4746044921875, + "relative_l2": 0.33158579817413825 + }, + "parameter.proj_in.weight": { + "max_abs": 785.5183972167969, + "relative_l2": 0.3587452309712956 + }, + "parameter.proj_out.bias": { + "max_abs": 0.01868856430053711, + "relative_l2": 0.012918536528168385 + }, + "parameter.proj_out.weight": { + "max_abs": 12.085406188964845, + "relative_l2": 0.021091651937679996 + }, + "parameter.time_embedder.linear_1.bias": { + "max_abs": 783.52501953125, + "relative_l2": 0.04770286067239733 + }, + "parameter.time_embedder.linear_1.weight": { + "max_abs": 853.1036328125, + "relative_l2": 0.04727435712437821 + }, + "parameter.time_embedder.linear_2.bias": { + "max_abs": 1232.4355419921876, + "relative_l2": 0.33858927966178637 + }, + "parameter.time_embedder.linear_2.weight": { + "max_abs": 11108.18615234375, + "relative_l2": 0.2193347288362636 + }, + "parameter.token_refiner.final_norm.weight": { + "max_abs": 90.47127807617187, + "relative_l2": 0.4951302595087592 + }, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": { + "max_abs": 124.09522140502929, + "relative_l2": 0.4046833363175061 + }, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": { + "max_abs": 125.23231781005859, + "relative_l2": 0.40390545382323617 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": { + "max_abs": 16.892093544006347, + "relative_l2": 0.47947822829994274 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": { + "max_abs": 25.671869163513183, + "relative_l2": 0.49531255975895666 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": { + "max_abs": 20.153124694824218, + "relative_l2": 0.46178624089996434 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": { + "max_abs": 21.626427059173583, + "relative_l2": 0.4955083820299898 + }, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": { + "max_abs": 33.706654434204104, + "relative_l2": 0.4879501687216877 + }, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": { + "max_abs": 70.30745780944824, + "relative_l2": 0.4949133242009952 + }, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": { + "max_abs": 119.85018432617187, + "relative_l2": 0.49030330639891395 + }, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": { + "max_abs": 174.6340817260742, + "relative_l2": 0.5023980028868472 + }, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": { + "max_abs": 13.301282291412354, + "relative_l2": 0.6618802613160201 + }, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": { + "max_abs": 13.395025138854981, + "relative_l2": 0.6626499140915075 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": { + "max_abs": 2.337174837589264, + "relative_l2": 0.5330949325707235 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": { + "max_abs": 5.516820077896118, + "relative_l2": 0.5079451521284352 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": { + "max_abs": 3.0227862644195556, + "relative_l2": 0.5425959730719063 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": { + "max_abs": 3.9670922803878783, + "relative_l2": 0.5165997257206122 + }, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": { + "max_abs": 34.41474044799805, + "relative_l2": 0.5385256887991842 + }, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": { + "max_abs": 81.52852714538574, + "relative_l2": 0.5084387158519625 + }, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": { + "max_abs": 43.96798217773438, + "relative_l2": 0.3892203088479711 + }, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": { + "max_abs": 41.61128318786621, + "relative_l2": 0.55748139151743 + }, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": { + "max_abs": 618.4180859375, + "relative_l2": 0.47075492675148795 + }, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": { + "max_abs": 29961.772578125, + "relative_l2": 0.27301673947434835 + }, + "parameter.transformer_blocks.0.attn.norm_k.weight": { + "max_abs": 4459.5284375, + "relative_l2": 0.37966807565997 + }, + "parameter.transformer_blocks.0.attn.norm_q.weight": { + "max_abs": 4694.645380859375, + "relative_l2": 0.4193497891611958 + }, + "parameter.transformer_blocks.0.attn.to_k.weight": { + "max_abs": 2666.39416015625, + "relative_l2": 0.3421879560960825 + }, + "parameter.transformer_blocks.0.attn.to_out.0.weight": { + "max_abs": 1295.931025390625, + "relative_l2": 0.26791442696637935 + }, + "parameter.transformer_blocks.0.attn.to_q.weight": { + "max_abs": 3227.72814453125, + "relative_l2": 0.3076763356611271 + }, + "parameter.transformer_blocks.0.attn.to_v.weight": { + "max_abs": 967.9125012207031, + "relative_l2": 0.265496613904143 + }, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": { + "max_abs": 5272.87560546875, + "relative_l2": 0.27539323856917125 + }, + "parameter.transformer_blocks.0.ff.net.2.weight": { + "max_abs": 931.70568359375, + "relative_l2": 0.27157714269997285 + }, + "parameter.transformer_blocks.0.norm1.weight": { + "max_abs": 2536.1047314453126, + "relative_l2": 0.30649315358287554 + }, + "parameter.transformer_blocks.0.norm2.weight": { + "max_abs": 8055.1895703125, + "relative_l2": 0.25356405562355117 + }, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": { + "max_abs": 34.71782577514649, + "relative_l2": 0.28264463903051756 + }, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": { + "max_abs": 8285.483515625, + "relative_l2": 0.2711459255786336 + }, + "parameter.transformer_blocks.1.attn.norm_k.weight": { + "max_abs": 9.839654808044434, + "relative_l2": 0.6866799937412663 + }, + "parameter.transformer_blocks.1.attn.norm_q.weight": { + "max_abs": 8.86622465133667, + "relative_l2": 0.6603229434845339 + }, + "parameter.transformer_blocks.1.attn.to_k.weight": { + "max_abs": 6.760801439285278, + "relative_l2": 0.8091230188837336 + }, + "parameter.transformer_blocks.1.attn.to_out.0.weight": { + "max_abs": 5.096917037963867, + "relative_l2": 0.716292400958965 + }, + "parameter.transformer_blocks.1.attn.to_q.weight": { + "max_abs": 6.934056406021118, + "relative_l2": 0.7213361749357194 + }, + "parameter.transformer_blocks.1.attn.to_v.weight": { + "max_abs": 5.85525155544281, + "relative_l2": 0.7914735134494816 + }, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": { + "max_abs": 1124.0305126953126, + "relative_l2": 0.27999278243568926 + }, + "parameter.transformer_blocks.1.ff.net.2.weight": { + "max_abs": 722.3081005859375, + "relative_l2": 0.2676357113339872 + }, + "parameter.transformer_blocks.1.norm1.weight": { + "max_abs": 19.10787570953369, + "relative_l2": 0.7871692466073358 + }, + "parameter.transformer_blocks.1.norm2.weight": { + "max_abs": 1245.1071728515626, + "relative_l2": 0.2617261401281995 + }, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": { + "max_abs": 7.491905097961426, + "relative_l2": 0.27412761443613665 + }, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": { + "max_abs": 1787.8017529296876, + "relative_l2": 0.2485811102828138 + }, + "parameter.transformer_blocks.10.attn.norm_k.weight": { + "max_abs": 3.919461612701416, + "relative_l2": 0.8939340974928423 + }, + "parameter.transformer_blocks.10.attn.norm_q.weight": { + "max_abs": 3.357305293083191, + "relative_l2": 0.91772600843942 + }, + "parameter.transformer_blocks.10.attn.to_k.weight": { + "max_abs": 1.9735518741607665, + "relative_l2": 0.7451501197444428 + }, + "parameter.transformer_blocks.10.attn.to_out.0.weight": { + "max_abs": 0.7348604547977448, + "relative_l2": 0.6738693588687075 + }, + "parameter.transformer_blocks.10.attn.to_q.weight": { + "max_abs": 1.4673158574104308, + "relative_l2": 0.7542932763495229 + }, + "parameter.transformer_blocks.10.attn.to_v.weight": { + "max_abs": 1.669773106575012, + "relative_l2": 0.7198321350834086 + }, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": { + "max_abs": 946.703486328125, + "relative_l2": 0.2748115050033871 + }, + "parameter.transformer_blocks.10.ff.net.2.weight": { + "max_abs": 281.16688049316406, + "relative_l2": 0.23077566506659983 + }, + "parameter.transformer_blocks.10.norm1.weight": { + "max_abs": 6.122771863937378, + "relative_l2": 0.8132083951553788 + }, + "parameter.transformer_blocks.10.norm2.weight": { + "max_abs": 2005.6719921875, + "relative_l2": 0.2858285655459195 + }, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": { + "max_abs": 10.026488189697266, + "relative_l2": 0.24650859743948983 + }, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": { + "max_abs": 2381.72375, + "relative_l2": 0.23093499807756118 + }, + "parameter.transformer_blocks.11.attn.norm_k.weight": { + "max_abs": 2.3551847982406615, + "relative_l2": 0.9131311070090182 + }, + "parameter.transformer_blocks.11.attn.norm_q.weight": { + "max_abs": 2.4871248054504393, + "relative_l2": 0.9289690749562688 + }, + "parameter.transformer_blocks.11.attn.to_k.weight": { + "max_abs": 1.9326666045188903, + "relative_l2": 0.831162001950036 + }, + "parameter.transformer_blocks.11.attn.to_out.0.weight": { + "max_abs": 0.6217981028556824, + "relative_l2": 0.6472011831701678 + }, + "parameter.transformer_blocks.11.attn.to_q.weight": { + "max_abs": 2.5200601863861083, + "relative_l2": 0.8275675088169366 + }, + "parameter.transformer_blocks.11.attn.to_v.weight": { + "max_abs": 0.9879325973987579, + "relative_l2": 0.7010794076362376 + }, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": { + "max_abs": 681.3028210449219, + "relative_l2": 0.24425111720111248 + }, + "parameter.transformer_blocks.11.ff.net.2.weight": { + "max_abs": 300.01104248046875, + "relative_l2": 0.2191100526426731 + }, + "parameter.transformer_blocks.11.norm1.weight": { + "max_abs": 4.620777015686035, + "relative_l2": 0.7356413076782026 + }, + "parameter.transformer_blocks.11.norm2.weight": { + "max_abs": 596.8696850585937, + "relative_l2": 0.14912470039405729 + }, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": { + "max_abs": 8.121762161254884, + "relative_l2": 0.2550387695831746 + }, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": { + "max_abs": 1945.74669921875, + "relative_l2": 0.2367571483337047 + }, + "parameter.transformer_blocks.12.attn.norm_k.weight": { + "max_abs": 2.9853444147109984, + "relative_l2": 0.7001627545316077 + }, + "parameter.transformer_blocks.12.attn.norm_q.weight": { + "max_abs": 3.0688916492462157, + "relative_l2": 0.7122141032106739 + }, + "parameter.transformer_blocks.12.attn.to_k.weight": { + "max_abs": 1.6972085285186767, + "relative_l2": 0.5879689607748688 + }, + "parameter.transformer_blocks.12.attn.to_out.0.weight": { + "max_abs": 0.49494848132133484, + "relative_l2": 0.5773987091429704 + }, + "parameter.transformer_blocks.12.attn.to_q.weight": { + "max_abs": 2.5782805728912352, + "relative_l2": 0.6454071223125697 + }, + "parameter.transformer_blocks.12.attn.to_v.weight": { + "max_abs": 0.6551486897468567, + "relative_l2": 0.6474215209980447 + }, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": { + "max_abs": 354.23184478759765, + "relative_l2": 0.26331063958954565 + }, + "parameter.transformer_blocks.12.ff.net.2.weight": { + "max_abs": 130.85853088378906, + "relative_l2": 0.2240588958162464 + }, + "parameter.transformer_blocks.12.norm1.weight": { + "max_abs": 2.7988663363456725, + "relative_l2": 0.6063861929230794 + }, + "parameter.transformer_blocks.12.norm2.weight": { + "max_abs": 789.6607006835937, + "relative_l2": 0.2711719907576323 + }, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": { + "max_abs": 9.051812057495118, + "relative_l2": 0.2594281777634543 + }, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": { + "max_abs": 2153.95958984375, + "relative_l2": 0.23677682551061446 + }, + "parameter.transformer_blocks.13.attn.norm_k.weight": { + "max_abs": 3.3524254608154296, + "relative_l2": 0.8954921306191933 + }, + "parameter.transformer_blocks.13.attn.norm_q.weight": { + "max_abs": 2.6193434047698974, + "relative_l2": 0.8582240172254604 + }, + "parameter.transformer_blocks.13.attn.to_k.weight": { + "max_abs": 1.079060022830963, + "relative_l2": 0.8219867581710362 + }, + "parameter.transformer_blocks.13.attn.to_out.0.weight": { + "max_abs": 0.7083306360244751, + "relative_l2": 0.8192653910512632 + }, + "parameter.transformer_blocks.13.attn.to_q.weight": { + "max_abs": 1.8803073811531066, + "relative_l2": 0.7905141412197559 + }, + "parameter.transformer_blocks.13.attn.to_v.weight": { + "max_abs": 1.1077276515960692, + "relative_l2": 0.8386404244712579 + }, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": { + "max_abs": 227.7526806640625, + "relative_l2": 0.2558986494449096 + }, + "parameter.transformer_blocks.13.ff.net.2.weight": { + "max_abs": 218.08233154296875, + "relative_l2": 0.22186692828327384 + }, + "parameter.transformer_blocks.13.norm1.weight": { + "max_abs": 3.649041299819946, + "relative_l2": 0.8200788291571931 + }, + "parameter.transformer_blocks.13.norm2.weight": { + "max_abs": 574.981806640625, + "relative_l2": 0.24275651736493498 + }, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": { + "max_abs": 6.945229415893555, + "relative_l2": 0.26591723008762436 + }, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": { + "max_abs": 1645.763544921875, + "relative_l2": 0.24360616649564137 + }, + "parameter.transformer_blocks.14.attn.norm_k.weight": { + "max_abs": 2.0977778482437133, + "relative_l2": 0.6688106360712017 + }, + "parameter.transformer_blocks.14.attn.norm_q.weight": { + "max_abs": 1.3061669516563414, + "relative_l2": 0.6895363279211654 + }, + "parameter.transformer_blocks.14.attn.to_k.weight": { + "max_abs": 1.5680741119384765, + "relative_l2": 0.6876301501287961 + }, + "parameter.transformer_blocks.14.attn.to_out.0.weight": { + "max_abs": 0.5789981383085251, + "relative_l2": 0.6718812940143374 + }, + "parameter.transformer_blocks.14.attn.to_q.weight": { + "max_abs": 1.2793809819221496, + "relative_l2": 0.7055059914862861 + }, + "parameter.transformer_blocks.14.attn.to_v.weight": { + "max_abs": 1.4136269736289977, + "relative_l2": 0.758456532499533 + }, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": { + "max_abs": 358.0343731689453, + "relative_l2": 0.2506961598008713 + }, + "parameter.transformer_blocks.14.ff.net.2.weight": { + "max_abs": 138.15642822265625, + "relative_l2": 0.22120862426841734 + }, + "parameter.transformer_blocks.14.norm1.weight": { + "max_abs": 3.1570816087722777, + "relative_l2": 0.7692865880722793 + }, + "parameter.transformer_blocks.14.norm2.weight": { + "max_abs": 827.2231396484375, + "relative_l2": 0.2948949400546231 + }, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": { + "max_abs": 8.299082641601563, + "relative_l2": 0.25602805026332515 + }, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": { + "max_abs": 2009.92736328125, + "relative_l2": 0.23283001516518612 + }, + "parameter.transformer_blocks.15.attn.norm_k.weight": { + "max_abs": 2.2380307364463805, + "relative_l2": 0.7711324359701877 + }, + "parameter.transformer_blocks.15.attn.norm_q.weight": { + "max_abs": 2.4794930267333983, + "relative_l2": 0.7798992025704999 + }, + "parameter.transformer_blocks.15.attn.to_k.weight": { + "max_abs": 1.2511296319961547, + "relative_l2": 0.7908391612072826 + }, + "parameter.transformer_blocks.15.attn.to_out.0.weight": { + "max_abs": 0.8765870380401611, + "relative_l2": 0.6449978432370598 + }, + "parameter.transformer_blocks.15.attn.to_q.weight": { + "max_abs": 1.296251540184021, + "relative_l2": 0.7386447155926827 + }, + "parameter.transformer_blocks.15.attn.to_v.weight": { + "max_abs": 0.8481322872638702, + "relative_l2": 0.7339279795477406 + }, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": { + "max_abs": 995.5617932128906, + "relative_l2": 0.24840504785585898 + }, + "parameter.transformer_blocks.15.ff.net.2.weight": { + "max_abs": 154.35751998901367, + "relative_l2": 0.22302031190872912 + }, + "parameter.transformer_blocks.15.norm1.weight": { + "max_abs": 4.56264889717102, + "relative_l2": 0.8064254345561096 + }, + "parameter.transformer_blocks.15.norm2.weight": { + "max_abs": 2107.9294384765626, + "relative_l2": 0.3168363717864166 + }, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": { + "max_abs": 7.884571199417114, + "relative_l2": 0.2628938003490684 + }, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": { + "max_abs": 1879.1367749023439, + "relative_l2": 0.2325869713721619 + }, + "parameter.transformer_blocks.16.attn.norm_k.weight": { + "max_abs": 2.1271226453781127, + "relative_l2": 0.7322144551851629 + }, + "parameter.transformer_blocks.16.attn.norm_q.weight": { + "max_abs": 2.4217547225952147, + "relative_l2": 0.7826353295239833 + }, + "parameter.transformer_blocks.16.attn.to_k.weight": { + "max_abs": 1.7393020439147948, + "relative_l2": 0.7223069267374671 + }, + "parameter.transformer_blocks.16.attn.to_out.0.weight": { + "max_abs": 0.5886289674043655, + "relative_l2": 0.5151741110535427 + }, + "parameter.transformer_blocks.16.attn.to_q.weight": { + "max_abs": 1.2670303034782409, + "relative_l2": 0.7348671488129176 + }, + "parameter.transformer_blocks.16.attn.to_v.weight": { + "max_abs": 0.9935271489620209, + "relative_l2": 0.6170055938519416 + }, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": { + "max_abs": 384.8454083251953, + "relative_l2": 0.24667237636020004 + }, + "parameter.transformer_blocks.16.ff.net.2.weight": { + "max_abs": 133.7273129272461, + "relative_l2": 0.20971634566923797 + }, + "parameter.transformer_blocks.16.norm1.weight": { + "max_abs": 4.996379976272583, + "relative_l2": 0.7024999699709787 + }, + "parameter.transformer_blocks.16.norm2.weight": { + "max_abs": 1165.293330078125, + "relative_l2": 0.24776021583071195 + }, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": { + "max_abs": 13.588038330078126, + "relative_l2": 0.24125780352307197 + }, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": { + "max_abs": 3197.22716796875, + "relative_l2": 0.2213120876089823 + }, + "parameter.transformer_blocks.17.attn.norm_k.weight": { + "max_abs": 2.577016716003418, + "relative_l2": 0.7274001496154391 + }, + "parameter.transformer_blocks.17.attn.norm_q.weight": { + "max_abs": 2.37910831451416, + "relative_l2": 0.7732131521152253 + }, + "parameter.transformer_blocks.17.attn.to_k.weight": { + "max_abs": 1.0713100481033324, + "relative_l2": 0.7684038527415521 + }, + "parameter.transformer_blocks.17.attn.to_out.0.weight": { + "max_abs": 0.77244544506073, + "relative_l2": 0.7915539672303255 + }, + "parameter.transformer_blocks.17.attn.to_q.weight": { + "max_abs": 1.3973073410987853, + "relative_l2": 0.7598057776059467 + }, + "parameter.transformer_blocks.17.attn.to_v.weight": { + "max_abs": 0.9992235350608826, + "relative_l2": 0.8915697210644591 + }, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": { + "max_abs": 182.5413426208496, + "relative_l2": 0.23066209417021438 + }, + "parameter.transformer_blocks.17.ff.net.2.weight": { + "max_abs": 367.0468701171875, + "relative_l2": 0.20430895173490748 + }, + "parameter.transformer_blocks.17.norm1.weight": { + "max_abs": 3.7926521348953246, + "relative_l2": 0.8373703936130231 + }, + "parameter.transformer_blocks.17.norm2.weight": { + "max_abs": 863.1227978515625, + "relative_l2": 0.23192950641814164 + }, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": { + "max_abs": 11.117765312194825, + "relative_l2": 0.22011060082853495 + }, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": { + "max_abs": 2647.5764111328126, + "relative_l2": 0.20970623162134797 + }, + "parameter.transformer_blocks.18.attn.norm_k.weight": { + "max_abs": 1.8815052556991576, + "relative_l2": 0.7812328618762753 + }, + "parameter.transformer_blocks.18.attn.norm_q.weight": { + "max_abs": 0.9511458325386047, + "relative_l2": 0.7175421337378723 + }, + "parameter.transformer_blocks.18.attn.to_k.weight": { + "max_abs": 0.9281016337871552, + "relative_l2": 0.7242657368447106 + }, + "parameter.transformer_blocks.18.attn.to_out.0.weight": { + "max_abs": 0.6100253510475159, + "relative_l2": 0.6399585470232062 + }, + "parameter.transformer_blocks.18.attn.to_q.weight": { + "max_abs": 0.8241854357719421, + "relative_l2": 0.7026953919302968 + }, + "parameter.transformer_blocks.18.attn.to_v.weight": { + "max_abs": 1.0181871104240416, + "relative_l2": 0.7925326835764924 + }, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": { + "max_abs": 870.2954968261719, + "relative_l2": 0.21723119562526702 + }, + "parameter.transformer_blocks.18.ff.net.2.weight": { + "max_abs": 535.0214184570312, + "relative_l2": 0.18686082434706205 + }, + "parameter.transformer_blocks.18.norm1.weight": { + "max_abs": 1.7391082096099852, + "relative_l2": 0.7427907656180031 + }, + "parameter.transformer_blocks.18.norm2.weight": { + "max_abs": 760.1420849609375, + "relative_l2": 0.20643528437354 + }, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": { + "max_abs": 9.754404907226563, + "relative_l2": 0.21710379129855412 + }, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": { + "max_abs": 2309.264521484375, + "relative_l2": 0.2026108967929386 + }, + "parameter.transformer_blocks.19.attn.norm_k.weight": { + "max_abs": 3.3605203676223754, + "relative_l2": 0.9235690859452362 + }, + "parameter.transformer_blocks.19.attn.norm_q.weight": { + "max_abs": 3.3958607959747313, + "relative_l2": 0.9027156051345979 + }, + "parameter.transformer_blocks.19.attn.to_k.weight": { + "max_abs": 2.5811061906814574, + "relative_l2": 0.9278696870441093 + }, + "parameter.transformer_blocks.19.attn.to_out.0.weight": { + "max_abs": 0.9226974296569824, + "relative_l2": 0.6308925415130978 + }, + "parameter.transformer_blocks.19.attn.to_q.weight": { + "max_abs": 1.8905621814727782, + "relative_l2": 0.9197084146310966 + }, + "parameter.transformer_blocks.19.attn.to_v.weight": { + "max_abs": 1.4341714429855346, + "relative_l2": 0.7554498717910688 + }, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": { + "max_abs": 428.02108276367187, + "relative_l2": 0.19810123483654649 + }, + "parameter.transformer_blocks.19.ff.net.2.weight": { + "max_abs": 348.52154052734375, + "relative_l2": 0.1570854430188067 + }, + "parameter.transformer_blocks.19.norm1.weight": { + "max_abs": 10.131331806182862, + "relative_l2": 0.9411943421518384 + }, + "parameter.transformer_blocks.19.norm2.weight": { + "max_abs": 1661.620478515625, + "relative_l2": 0.2312924067876611 + }, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": { + "max_abs": 38.34929931640625, + "relative_l2": 0.27099813831084796 + }, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": { + "max_abs": 9139.78625, + "relative_l2": 0.2610091571123563 + }, + "parameter.transformer_blocks.2.attn.norm_k.weight": { + "max_abs": 11.38742530822754, + "relative_l2": 0.7203443894744488 + }, + "parameter.transformer_blocks.2.attn.norm_q.weight": { + "max_abs": 10.698475246429444, + "relative_l2": 0.7308512325109128 + }, + "parameter.transformer_blocks.2.attn.to_k.weight": { + "max_abs": 7.837068204879761, + "relative_l2": 0.7116781769614272 + }, + "parameter.transformer_blocks.2.attn.to_out.0.weight": { + "max_abs": 3.5596836853027343, + "relative_l2": 0.5998331316236671 + }, + "parameter.transformer_blocks.2.attn.to_q.weight": { + "max_abs": 7.5656751441955565, + "relative_l2": 0.6356281954464494 + }, + "parameter.transformer_blocks.2.attn.to_v.weight": { + "max_abs": 7.642959957122803, + "relative_l2": 0.7641174143769899 + }, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": { + "max_abs": 835.4142712402344, + "relative_l2": 0.2686617582206923 + }, + "parameter.transformer_blocks.2.ff.net.2.weight": { + "max_abs": 552.6587780761719, + "relative_l2": 0.25879537458619123 + }, + "parameter.transformer_blocks.2.norm1.weight": { + "max_abs": 13.009734039306641, + "relative_l2": 0.6925190518757222 + }, + "parameter.transformer_blocks.2.norm2.weight": { + "max_abs": 1900.908564453125, + "relative_l2": 0.2736282238955311 + }, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": { + "max_abs": 11.688382034301759, + "relative_l2": 0.2377400113550496 + }, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": { + "max_abs": 2782.34533203125, + "relative_l2": 0.20709904622478764 + }, + "parameter.transformer_blocks.20.attn.norm_k.weight": { + "max_abs": 2.9652954387664794, + "relative_l2": 1.1392256242682777 + }, + "parameter.transformer_blocks.20.attn.norm_q.weight": { + "max_abs": 2.5787290382385253, + "relative_l2": 1.1756615296147424 + }, + "parameter.transformer_blocks.20.attn.to_k.weight": { + "max_abs": 2.683075909614563, + "relative_l2": 1.1153968264775866 + }, + "parameter.transformer_blocks.20.attn.to_out.0.weight": { + "max_abs": 0.7956968891620636, + "relative_l2": 0.528389423061085 + }, + "parameter.transformer_blocks.20.attn.to_q.weight": { + "max_abs": 1.4907193827629088, + "relative_l2": 1.0275720720510009 + }, + "parameter.transformer_blocks.20.attn.to_v.weight": { + "max_abs": 1.0080752718448638, + "relative_l2": 0.6380884034094756 + }, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": { + "max_abs": 220.31978881835937, + "relative_l2": 0.21520136687725197 + }, + "parameter.transformer_blocks.20.ff.net.2.weight": { + "max_abs": 286.46563232421875, + "relative_l2": 0.1794734190533937 + }, + "parameter.transformer_blocks.20.norm1.weight": { + "max_abs": 6.677889709472656, + "relative_l2": 0.9076549223343746 + }, + "parameter.transformer_blocks.20.norm2.weight": { + "max_abs": 384.90670288085937, + "relative_l2": 0.21424578336514133 + }, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": { + "max_abs": 6.576134567260742, + "relative_l2": 0.22732946736872406 + }, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": { + "max_abs": 1560.97765625, + "relative_l2": 0.2000087574560543 + }, + "parameter.transformer_blocks.21.attn.norm_k.weight": { + "max_abs": 3.070456986427307, + "relative_l2": 0.9745551614882625 + }, + "parameter.transformer_blocks.21.attn.norm_q.weight": { + "max_abs": 2.8962028074264525, + "relative_l2": 0.9897367111919141 + }, + "parameter.transformer_blocks.21.attn.to_k.weight": { + "max_abs": 1.5696979212760924, + "relative_l2": 0.9636766966378697 + }, + "parameter.transformer_blocks.21.attn.to_out.0.weight": { + "max_abs": 0.8160495805740356, + "relative_l2": 0.6789384150673038 + }, + "parameter.transformer_blocks.21.attn.to_q.weight": { + "max_abs": 1.102442626953125, + "relative_l2": 0.9350146184006228 + }, + "parameter.transformer_blocks.21.attn.to_v.weight": { + "max_abs": 1.402775650024414, + "relative_l2": 0.8394066497398459 + }, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": { + "max_abs": 172.8049038696289, + "relative_l2": 0.2180061256659362 + }, + "parameter.transformer_blocks.21.ff.net.2.weight": { + "max_abs": 207.50718963623046, + "relative_l2": 0.18204959014623456 + }, + "parameter.transformer_blocks.21.norm1.weight": { + "max_abs": 5.1228568601608275, + "relative_l2": 0.8691599430647023 + }, + "parameter.transformer_blocks.21.norm2.weight": { + "max_abs": 395.5959423828125, + "relative_l2": 0.22020446724135123 + }, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": { + "max_abs": 7.377956275939941, + "relative_l2": 0.22904466456764788 + }, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": { + "max_abs": 1733.79845703125, + "relative_l2": 0.2036469103690055 + }, + "parameter.transformer_blocks.22.attn.norm_k.weight": { + "max_abs": 2.5229035067558288, + "relative_l2": 0.9044992781713185 + }, + "parameter.transformer_blocks.22.attn.norm_q.weight": { + "max_abs": 2.3648268032073974, + "relative_l2": 0.8756975510736572 + }, + "parameter.transformer_blocks.22.attn.to_k.weight": { + "max_abs": 1.9956800985336303, + "relative_l2": 0.85522739174773 + }, + "parameter.transformer_blocks.22.attn.to_out.0.weight": { + "max_abs": 0.7181864130496979, + "relative_l2": 0.8913080694256287 + }, + "parameter.transformer_blocks.22.attn.to_q.weight": { + "max_abs": 0.8015691387653351, + "relative_l2": 0.8314708081082037 + }, + "parameter.transformer_blocks.22.attn.to_v.weight": { + "max_abs": 2.644535069465637, + "relative_l2": 1.0197672534558853 + }, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": { + "max_abs": 357.3243588256836, + "relative_l2": 0.2189386423066487 + }, + "parameter.transformer_blocks.22.ff.net.2.weight": { + "max_abs": 207.0422848510742, + "relative_l2": 0.1754032335951697 + }, + "parameter.transformer_blocks.22.norm1.weight": { + "max_abs": 3.6996399211883544, + "relative_l2": 0.9034981506161427 + }, + "parameter.transformer_blocks.22.norm2.weight": { + "max_abs": 522.5956982421875, + "relative_l2": 0.21523293735280408 + }, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": { + "max_abs": 8.08806360244751, + "relative_l2": 0.21326685878132157 + }, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": { + "max_abs": 1922.5801953125, + "relative_l2": 0.19217049180594867 + }, + "parameter.transformer_blocks.23.attn.norm_k.weight": { + "max_abs": 2.515032832622528, + "relative_l2": 1.0121865478435368 + }, + "parameter.transformer_blocks.23.attn.norm_q.weight": { + "max_abs": 2.541628723144531, + "relative_l2": 1.0178988868215497 + }, + "parameter.transformer_blocks.23.attn.to_k.weight": { + "max_abs": 3.535975103378296, + "relative_l2": 0.9932593402245732 + }, + "parameter.transformer_blocks.23.attn.to_out.0.weight": { + "max_abs": 0.48513928532600403, + "relative_l2": 0.644548323931213 + }, + "parameter.transformer_blocks.23.attn.to_q.weight": { + "max_abs": 1.0787883150577544, + "relative_l2": 0.9093563448523482 + }, + "parameter.transformer_blocks.23.attn.to_v.weight": { + "max_abs": 1.7548005628585814, + "relative_l2": 0.8057142720129602 + }, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": { + "max_abs": 239.49857604980468, + "relative_l2": 0.20468404153122077 + }, + "parameter.transformer_blocks.23.ff.net.2.weight": { + "max_abs": 293.54357421875, + "relative_l2": 0.17651630716934855 + }, + "parameter.transformer_blocks.23.norm1.weight": { + "max_abs": 6.985463981628418, + "relative_l2": 1.0053685364073084 + }, + "parameter.transformer_blocks.23.norm2.weight": { + "max_abs": 456.5930432128906, + "relative_l2": 0.21848467457692144 + }, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": { + "max_abs": 7.304553871154785, + "relative_l2": 0.21564167899972708 + }, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": { + "max_abs": 1737.616572265625, + "relative_l2": 0.1901697069608287 + }, + "parameter.transformer_blocks.24.attn.norm_k.weight": { + "max_abs": 2.216841106414795, + "relative_l2": 0.7926188934928867 + }, + "parameter.transformer_blocks.24.attn.norm_q.weight": { + "max_abs": 1.5476918268203734, + "relative_l2": 0.7777840222904054 + }, + "parameter.transformer_blocks.24.attn.to_k.weight": { + "max_abs": 1.8396434831619262, + "relative_l2": 0.7579427336702298 + }, + "parameter.transformer_blocks.24.attn.to_out.0.weight": { + "max_abs": 0.9694585251808167, + "relative_l2": 0.5167980950331458 + }, + "parameter.transformer_blocks.24.attn.to_q.weight": { + "max_abs": 0.9235194373130798, + "relative_l2": 0.7017912116896795 + }, + "parameter.transformer_blocks.24.attn.to_v.weight": { + "max_abs": 1.2338300991058349, + "relative_l2": 0.6999683243661816 + }, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": { + "max_abs": 487.1518811035156, + "relative_l2": 0.20379978037497748 + }, + "parameter.transformer_blocks.24.ff.net.2.weight": { + "max_abs": 556.0626171875, + "relative_l2": 0.13101262792454788 + }, + "parameter.transformer_blocks.24.norm1.weight": { + "max_abs": 6.81080997467041, + "relative_l2": 0.7571361191504765 + }, + "parameter.transformer_blocks.24.norm2.weight": { + "max_abs": 800.630732421875, + "relative_l2": 0.20475600692712043 + }, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": { + "max_abs": 8.371666793823243, + "relative_l2": 0.16869321168645987 + }, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": { + "max_abs": 1985.033076171875, + "relative_l2": 0.15815527032757656 + }, + "parameter.transformer_blocks.25.attn.norm_k.weight": { + "max_abs": 2.4211975383758544, + "relative_l2": 0.8235665728543179 + }, + "parameter.transformer_blocks.25.attn.norm_q.weight": { + "max_abs": 2.205787365436554, + "relative_l2": 0.8263968452846001 + }, + "parameter.transformer_blocks.25.attn.to_k.weight": { + "max_abs": 2.5189250755310058, + "relative_l2": 0.8722942908154744 + }, + "parameter.transformer_blocks.25.attn.to_out.0.weight": { + "max_abs": 0.9047119665145874, + "relative_l2": 0.6978460086808806 + }, + "parameter.transformer_blocks.25.attn.to_q.weight": { + "max_abs": 1.404982750415802, + "relative_l2": 0.8480601970016735 + }, + "parameter.transformer_blocks.25.attn.to_v.weight": { + "max_abs": 2.1768286275863646, + "relative_l2": 0.8885601899867497 + }, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": { + "max_abs": 408.89319885253906, + "relative_l2": 0.1644225805549453 + }, + "parameter.transformer_blocks.25.ff.net.2.weight": { + "max_abs": 753.0723828125, + "relative_l2": 0.11861937989316992 + }, + "parameter.transformer_blocks.25.norm1.weight": { + "max_abs": 3.4406565475463866, + "relative_l2": 0.7214860130902033 + }, + "parameter.transformer_blocks.25.norm2.weight": { + "max_abs": 537.69640625, + "relative_l2": 0.15529167748626888 + }, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": { + "max_abs": 4.974030380249023, + "relative_l2": 0.1970655795477857 + }, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": { + "max_abs": 1166.8911083984376, + "relative_l2": 0.16832694903177375 + }, + "parameter.transformer_blocks.26.attn.norm_k.weight": { + "max_abs": 3.5365945148468017, + "relative_l2": 0.8607390883814103 + }, + "parameter.transformer_blocks.26.attn.norm_q.weight": { + "max_abs": 3.5158914375305175, + "relative_l2": 0.863374440509692 + }, + "parameter.transformer_blocks.26.attn.to_k.weight": { + "max_abs": 3.6291722106933593, + "relative_l2": 0.8664953205185234 + }, + "parameter.transformer_blocks.26.attn.to_out.0.weight": { + "max_abs": 1.0129927444458007, + "relative_l2": 0.2929568074262903 + }, + "parameter.transformer_blocks.26.attn.to_q.weight": { + "max_abs": 1.9735407876968383, + "relative_l2": 0.8663906879530491 + }, + "parameter.transformer_blocks.26.attn.to_v.weight": { + "max_abs": 2.1177463579177855, + "relative_l2": 0.4406044410223554 + }, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": { + "max_abs": 371.87398803710937, + "relative_l2": 0.1798924973205361 + }, + "parameter.transformer_blocks.26.ff.net.2.weight": { + "max_abs": 457.3574169921875, + "relative_l2": 0.12919007812210365 + }, + "parameter.transformer_blocks.26.norm1.weight": { + "max_abs": 7.477498416900635, + "relative_l2": 0.689315960809784 + }, + "parameter.transformer_blocks.26.norm2.weight": { + "max_abs": 594.1813305664062, + "relative_l2": 0.18823506950940896 + }, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": { + "max_abs": 5.869821672439575, + "relative_l2": 0.21336485205109756 + }, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": { + "max_abs": 1084.662470703125, + "relative_l2": 0.1741348330090017 + }, + "parameter.transformer_blocks.27.attn.norm_k.weight": { + "max_abs": 3.150988345146179, + "relative_l2": 1.062841235494389 + }, + "parameter.transformer_blocks.27.attn.norm_q.weight": { + "max_abs": 3.6676880168914794, + "relative_l2": 1.0967533088519172 + }, + "parameter.transformer_blocks.27.attn.to_k.weight": { + "max_abs": 4.32183874130249, + "relative_l2": 1.0950772065469172 + }, + "parameter.transformer_blocks.27.attn.to_out.0.weight": { + "max_abs": 0.6107107448577881, + "relative_l2": 0.8264677710738436 + }, + "parameter.transformer_blocks.27.attn.to_q.weight": { + "max_abs": 3.3824980306625365, + "relative_l2": 1.070568942653313 + }, + "parameter.transformer_blocks.27.attn.to_v.weight": { + "max_abs": 4.458062772750854, + "relative_l2": 0.9492867035892784 + }, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": { + "max_abs": 259.40442932128906, + "relative_l2": 0.19459347336082505 + }, + "parameter.transformer_blocks.27.ff.net.2.weight": { + "max_abs": 125.5402783203125, + "relative_l2": 0.15380201598069254 + }, + "parameter.transformer_blocks.27.norm1.weight": { + "max_abs": 11.696521644592286, + "relative_l2": 0.962568903689285 + }, + "parameter.transformer_blocks.27.norm2.weight": { + "max_abs": 630.013056640625, + "relative_l2": 0.18498388565940183 + }, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": { + "max_abs": 6.271578674316406, + "relative_l2": 0.20026362707009437 + }, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": { + "max_abs": 1470.62609375, + "relative_l2": 0.156969272399435 + }, + "parameter.transformer_blocks.28.attn.norm_k.weight": { + "max_abs": 3.694087390899658, + "relative_l2": 1.1393344702953399 + }, + "parameter.transformer_blocks.28.attn.norm_q.weight": { + "max_abs": 3.2554990100860595, + "relative_l2": 1.132576919349636 + }, + "parameter.transformer_blocks.28.attn.to_k.weight": { + "max_abs": 2.87049311876297, + "relative_l2": 1.190224288047722 + }, + "parameter.transformer_blocks.28.attn.to_out.0.weight": { + "max_abs": 1.1344015049934386, + "relative_l2": 1.1170925582000946 + }, + "parameter.transformer_blocks.28.attn.to_q.weight": { + "max_abs": 2.7373598861694335, + "relative_l2": 1.2229476895309177 + }, + "parameter.transformer_blocks.28.attn.to_v.weight": { + "max_abs": 4.708512072563171, + "relative_l2": 1.1407878598746384 + }, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": { + "max_abs": 246.0049542236328, + "relative_l2": 0.18856620505021301 + }, + "parameter.transformer_blocks.28.ff.net.2.weight": { + "max_abs": 281.19640625, + "relative_l2": 0.13984586578762667 + }, + "parameter.transformer_blocks.28.norm1.weight": { + "max_abs": 11.20217836380005, + "relative_l2": 1.2122775780156614 + }, + "parameter.transformer_blocks.28.norm2.weight": { + "max_abs": 359.1081494140625, + "relative_l2": 0.16797781699942133 + }, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": { + "max_abs": 6.513003234863281, + "relative_l2": 0.19296473938499076 + }, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": { + "max_abs": 1499.823359375, + "relative_l2": 0.15103071975080934 + }, + "parameter.transformer_blocks.29.attn.norm_k.weight": { + "max_abs": 9.076084260940553, + "relative_l2": 1.6341561350014264 + }, + "parameter.transformer_blocks.29.attn.norm_q.weight": { + "max_abs": 7.577287321090698, + "relative_l2": 1.640274162969278 + }, + "parameter.transformer_blocks.29.attn.to_k.weight": { + "max_abs": 6.094952707290649, + "relative_l2": 1.4206461842138516 + }, + "parameter.transformer_blocks.29.attn.to_out.0.weight": { + "max_abs": 1.0908178436756133, + "relative_l2": 0.9675961856620527 + }, + "parameter.transformer_blocks.29.attn.to_q.weight": { + "max_abs": 4.269876127243042, + "relative_l2": 1.4317035390714827 + }, + "parameter.transformer_blocks.29.attn.to_v.weight": { + "max_abs": 3.974481587409973, + "relative_l2": 1.0206224707618594 + }, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": { + "max_abs": 516.3020886230469, + "relative_l2": 0.18134852616131078 + }, + "parameter.transformer_blocks.29.ff.net.2.weight": { + "max_abs": 341.305537109375, + "relative_l2": 0.12471303626652662 + }, + "parameter.transformer_blocks.29.norm1.weight": { + "max_abs": 13.254407291412354, + "relative_l2": 1.201718957896938 + }, + "parameter.transformer_blocks.29.norm2.weight": { + "max_abs": 412.28252685546875, + "relative_l2": 0.17436060353310176 + }, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": { + "max_abs": 26.373812561035155, + "relative_l2": 0.2654205061972307 + }, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": { + "max_abs": 6267.35216796875, + "relative_l2": 0.25628779739482377 + }, + "parameter.transformer_blocks.3.attn.norm_k.weight": { + "max_abs": 10.009553079605103, + "relative_l2": 1.0049191909026862 + }, + "parameter.transformer_blocks.3.attn.norm_q.weight": { + "max_abs": 8.742085580825806, + "relative_l2": 0.9629340394722943 + }, + "parameter.transformer_blocks.3.attn.to_k.weight": { + "max_abs": 7.2191501903533934, + "relative_l2": 1.0015157745442331 + }, + "parameter.transformer_blocks.3.attn.to_out.0.weight": { + "max_abs": 3.831388597488403, + "relative_l2": 0.6430799687149178 + }, + "parameter.transformer_blocks.3.attn.to_q.weight": { + "max_abs": 9.674795513153077, + "relative_l2": 0.8449424927333038 + }, + "parameter.transformer_blocks.3.attn.to_v.weight": { + "max_abs": 1.722345654964447, + "relative_l2": 0.7180312690570306 + }, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": { + "max_abs": 1106.1234387207032, + "relative_l2": 0.26521183470912624 + }, + "parameter.transformer_blocks.3.ff.net.2.weight": { + "max_abs": 289.41797912597656, + "relative_l2": 0.2521339003529588 + }, + "parameter.transformer_blocks.3.norm1.weight": { + "max_abs": 11.262598876953126, + "relative_l2": 0.9202157109443916 + }, + "parameter.transformer_blocks.3.norm2.weight": { + "max_abs": 4101.5118359375, + "relative_l2": 0.269506544885244 + }, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": { + "max_abs": 4.232561950683594, + "relative_l2": 0.21645874428585316 + }, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": { + "max_abs": 1016.190546875, + "relative_l2": 0.15072002097149242 + }, + "parameter.transformer_blocks.30.attn.norm_k.weight": { + "max_abs": 6.234205846786499, + "relative_l2": 1.024429056236177 + }, + "parameter.transformer_blocks.30.attn.norm_q.weight": { + "max_abs": 6.326533203125, + "relative_l2": 1.0369492200530552 + }, + "parameter.transformer_blocks.30.attn.to_k.weight": { + "max_abs": 4.29625440120697, + "relative_l2": 1.1695632920193195 + }, + "parameter.transformer_blocks.30.attn.to_out.0.weight": { + "max_abs": 0.9467710602283478, + "relative_l2": 0.7817796444989354 + }, + "parameter.transformer_blocks.30.attn.to_q.weight": { + "max_abs": 3.7871464538574218, + "relative_l2": 1.0976680529287743 + }, + "parameter.transformer_blocks.30.attn.to_v.weight": { + "max_abs": 2.525110249519348, + "relative_l2": 0.8613141581962568 + }, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": { + "max_abs": 239.4531201171875, + "relative_l2": 0.19589711445921443 + }, + "parameter.transformer_blocks.30.ff.net.2.weight": { + "max_abs": 238.6540478515625, + "relative_l2": 0.13942788767281689 + }, + "parameter.transformer_blocks.30.norm1.weight": { + "max_abs": 8.512052898406983, + "relative_l2": 1.0641770769293457 + }, + "parameter.transformer_blocks.30.norm2.weight": { + "max_abs": 451.7082470703125, + "relative_l2": 0.18921410436824 + }, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": { + "max_abs": 7.6711624908447265, + "relative_l2": 0.2123018243287658 + }, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": { + "max_abs": 799.74376953125, + "relative_l2": 0.16188623355622284 + }, + "parameter.transformer_blocks.31.attn.norm_k.weight": { + "max_abs": 4.266907458305359, + "relative_l2": 0.7253633062716792 + }, + "parameter.transformer_blocks.31.attn.norm_q.weight": { + "max_abs": 4.317539339065552, + "relative_l2": 0.7171574495366576 + }, + "parameter.transformer_blocks.31.attn.to_k.weight": { + "max_abs": 1.8372212696075438, + "relative_l2": 0.8527959018711727 + }, + "parameter.transformer_blocks.31.attn.to_out.0.weight": { + "max_abs": 0.7584056901931763, + "relative_l2": 0.7532493180744085 + }, + "parameter.transformer_blocks.31.attn.to_q.weight": { + "max_abs": 1.5440465259552, + "relative_l2": 0.8164946551799815 + }, + "parameter.transformer_blocks.31.attn.to_v.weight": { + "max_abs": 1.8497110652923583, + "relative_l2": 0.8529554209495858 + }, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": { + "max_abs": 154.2845066833496, + "relative_l2": 0.22037467788839546 + }, + "parameter.transformer_blocks.31.ff.net.2.weight": { + "max_abs": 112.23086822509765, + "relative_l2": 0.15448229099375835 + }, + "parameter.transformer_blocks.31.norm1.weight": { + "max_abs": 5.710070734024048, + "relative_l2": 0.8192065322889842 + }, + "parameter.transformer_blocks.31.norm2.weight": { + "max_abs": 349.50445068359375, + "relative_l2": 0.21345879795988626 + }, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": { + "max_abs": 4.077955131530762, + "relative_l2": 0.20541480550039418 + }, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": { + "max_abs": 980.7512158203125, + "relative_l2": 0.14042511207087938 + }, + "parameter.transformer_blocks.32.attn.norm_k.weight": { + "max_abs": 6.775987749099731, + "relative_l2": 1.288531562899753 + }, + "parameter.transformer_blocks.32.attn.norm_q.weight": { + "max_abs": 6.837880973815918, + "relative_l2": 1.2434342562895009 + }, + "parameter.transformer_blocks.32.attn.to_k.weight": { + "max_abs": 4.671030168533325, + "relative_l2": 1.1509488580181235 + }, + "parameter.transformer_blocks.32.attn.to_out.0.weight": { + "max_abs": 1.6398096132278441, + "relative_l2": 0.992625717903801 + }, + "parameter.transformer_blocks.32.attn.to_q.weight": { + "max_abs": 2.183147554397583, + "relative_l2": 1.1746264337260195 + }, + "parameter.transformer_blocks.32.attn.to_v.weight": { + "max_abs": 1.9204409170150756, + "relative_l2": 1.0286533770205062 + }, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": { + "max_abs": 209.76043975830078, + "relative_l2": 0.20242125269897343 + }, + "parameter.transformer_blocks.32.ff.net.2.weight": { + "max_abs": 89.70086181640625, + "relative_l2": 0.1444801344032381 + }, + "parameter.transformer_blocks.32.norm1.weight": { + "max_abs": 9.582114820480347, + "relative_l2": 1.1194710214665151 + }, + "parameter.transformer_blocks.32.norm2.weight": { + "max_abs": 336.048212890625, + "relative_l2": 0.20319355913432793 + }, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": { + "max_abs": 7.154135589599609, + "relative_l2": 0.19469886900881594 + }, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": { + "max_abs": 966.8830822753906, + "relative_l2": 0.14829812099440315 + }, + "parameter.transformer_blocks.33.attn.norm_k.weight": { + "max_abs": 5.550072317123413, + "relative_l2": 1.1274773312509463 + }, + "parameter.transformer_blocks.33.attn.norm_q.weight": { + "max_abs": 5.736447219848633, + "relative_l2": 1.1253863704143687 + }, + "parameter.transformer_blocks.33.attn.to_k.weight": { + "max_abs": 2.679831688404083, + "relative_l2": 0.9850919152514763 + }, + "parameter.transformer_blocks.33.attn.to_out.0.weight": { + "max_abs": 2.063485448360443, + "relative_l2": 0.9390572904823536 + }, + "parameter.transformer_blocks.33.attn.to_q.weight": { + "max_abs": 1.764441316127777, + "relative_l2": 0.9817237812251518 + }, + "parameter.transformer_blocks.33.attn.to_v.weight": { + "max_abs": 2.399023537635803, + "relative_l2": 0.9501950616159477 + }, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": { + "max_abs": 309.37867248535156, + "relative_l2": 0.2006807330341676 + }, + "parameter.transformer_blocks.33.ff.net.2.weight": { + "max_abs": 80.20376480102539, + "relative_l2": 0.13101022942819657 + }, + "parameter.transformer_blocks.33.norm1.weight": { + "max_abs": 15.021945838928223, + "relative_l2": 0.9127326214831956 + }, + "parameter.transformer_blocks.33.norm2.weight": { + "max_abs": 1001.6911572265625, + "relative_l2": 0.2251123813234075 + }, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": { + "max_abs": 8.243670349121095, + "relative_l2": 0.1806048426281073 + }, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": { + "max_abs": 1148.6772412109376, + "relative_l2": 0.12984915201068334 + }, + "parameter.transformer_blocks.34.attn.norm_k.weight": { + "max_abs": 7.219762449264526, + "relative_l2": 1.032384395529538 + }, + "parameter.transformer_blocks.34.attn.norm_q.weight": { + "max_abs": 7.133918170928955, + "relative_l2": 1.0149967125277624 + }, + "parameter.transformer_blocks.34.attn.to_k.weight": { + "max_abs": 4.497727279663086, + "relative_l2": 0.9591314698602402 + }, + "parameter.transformer_blocks.34.attn.to_out.0.weight": { + "max_abs": 1.03602308511734, + "relative_l2": 0.8783822880153688 + }, + "parameter.transformer_blocks.34.attn.to_q.weight": { + "max_abs": 2.7484306144714354, + "relative_l2": 0.8864441081344899 + }, + "parameter.transformer_blocks.34.attn.to_v.weight": { + "max_abs": 3.29366970539093, + "relative_l2": 0.9451329266505385 + }, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": { + "max_abs": 194.5088070678711, + "relative_l2": 0.18493408026558106 + }, + "parameter.transformer_blocks.34.ff.net.2.weight": { + "max_abs": 245.620478515625, + "relative_l2": 0.1307436775198352 + }, + "parameter.transformer_blocks.34.norm1.weight": { + "max_abs": 24.60771167755127, + "relative_l2": 1.0233308374050827 + }, + "parameter.transformer_blocks.34.norm2.weight": { + "max_abs": 319.7883786010742, + "relative_l2": 0.15874418206460217 + }, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": { + "max_abs": 5.534396057128906, + "relative_l2": 0.1467721804613339 + }, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": { + "max_abs": 1264.75109375, + "relative_l2": 0.11913879267016361 + }, + "parameter.transformer_blocks.35.attn.norm_k.weight": { + "max_abs": 4.387928371429443, + "relative_l2": 0.942640852962107 + }, + "parameter.transformer_blocks.35.attn.norm_q.weight": { + "max_abs": 4.449705009460449, + "relative_l2": 0.919191599683251 + }, + "parameter.transformer_blocks.35.attn.to_k.weight": { + "max_abs": 1.7174570608139037, + "relative_l2": 0.9179817867494432 + }, + "parameter.transformer_blocks.35.attn.to_out.0.weight": { + "max_abs": 1.459444110393524, + "relative_l2": 0.8814012609524116 + }, + "parameter.transformer_blocks.35.attn.to_q.weight": { + "max_abs": 2.8721355247497558, + "relative_l2": 1.0034159715694124 + }, + "parameter.transformer_blocks.35.attn.to_v.weight": { + "max_abs": 1.7862120914459227, + "relative_l2": 0.9437613436042724 + }, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": { + "max_abs": 191.31423461914062, + "relative_l2": 0.17145562347579194 + }, + "parameter.transformer_blocks.35.ff.net.2.weight": { + "max_abs": 186.5078076171875, + "relative_l2": 0.11657884529574467 + }, + "parameter.transformer_blocks.35.norm1.weight": { + "max_abs": 6.005865459442139, + "relative_l2": 0.8703644599721103 + }, + "parameter.transformer_blocks.35.norm2.weight": { + "max_abs": 679.1898754882812, + "relative_l2": 0.200027453204279 + }, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": { + "max_abs": 9.245374565124513, + "relative_l2": 0.2087435271818816 + }, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": { + "max_abs": 2195.7213696289064, + "relative_l2": 0.1602415887020769 + }, + "parameter.transformer_blocks.36.attn.norm_k.weight": { + "max_abs": 4.618183856010437, + "relative_l2": 1.0325610439726254 + }, + "parameter.transformer_blocks.36.attn.norm_q.weight": { + "max_abs": 4.725456480979919, + "relative_l2": 1.0538600913481053 + }, + "parameter.transformer_blocks.36.attn.to_k.weight": { + "max_abs": 5.769529585838318, + "relative_l2": 1.221919776036673 + }, + "parameter.transformer_blocks.36.attn.to_out.0.weight": { + "max_abs": 1.0138367462158202, + "relative_l2": 0.5783897208725631 + }, + "parameter.transformer_blocks.36.attn.to_q.weight": { + "max_abs": 3.6643313217163085, + "relative_l2": 1.2180755207186575 + }, + "parameter.transformer_blocks.36.attn.to_v.weight": { + "max_abs": 1.8370685625076293, + "relative_l2": 0.6205732856999983 + }, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": { + "max_abs": 309.66125, + "relative_l2": 0.20316813457524138 + }, + "parameter.transformer_blocks.36.ff.net.2.weight": { + "max_abs": 107.78756225585937, + "relative_l2": 0.12236753561577896 + }, + "parameter.transformer_blocks.36.norm1.weight": { + "max_abs": 22.513299827575683, + "relative_l2": 1.0778564356317475 + }, + "parameter.transformer_blocks.36.norm2.weight": { + "max_abs": 2511.9208935546876, + "relative_l2": 0.5710897440023305 + }, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": { + "max_abs": 6.676037197113037, + "relative_l2": 0.1808701004652349 + }, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": { + "max_abs": 1593.4457348632814, + "relative_l2": 0.1443106764289475 + }, + "parameter.transformer_blocks.37.attn.norm_k.weight": { + "max_abs": 2.832467083930969, + "relative_l2": 1.0101240952318908 + }, + "parameter.transformer_blocks.37.attn.norm_q.weight": { + "max_abs": 2.861815993785858, + "relative_l2": 1.0127782878012312 + }, + "parameter.transformer_blocks.37.attn.to_k.weight": { + "max_abs": 4.04865861415863, + "relative_l2": 1.1374777364018036 + }, + "parameter.transformer_blocks.37.attn.to_out.0.weight": { + "max_abs": 0.854398285150528, + "relative_l2": 0.8419456640459979 + }, + "parameter.transformer_blocks.37.attn.to_q.weight": { + "max_abs": 2.1792039918899535, + "relative_l2": 1.1004063862861377 + }, + "parameter.transformer_blocks.37.attn.to_v.weight": { + "max_abs": 2.1426338601112365, + "relative_l2": 0.8787250645274481 + }, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": { + "max_abs": 656.2790783691406, + "relative_l2": 0.2233693593556396 + }, + "parameter.transformer_blocks.37.ff.net.2.weight": { + "max_abs": 50.611208801269534, + "relative_l2": 0.10374361375682521 + }, + "parameter.transformer_blocks.37.norm1.weight": { + "max_abs": 8.482384805679322, + "relative_l2": 1.0256669036930872 + }, + "parameter.transformer_blocks.37.norm2.weight": { + "max_abs": 547.3217114257812, + "relative_l2": 0.268658747915288 + }, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": { + "max_abs": 5.421483640670776, + "relative_l2": 0.161460090239007 + }, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": { + "max_abs": 771.6396435546875, + "relative_l2": 0.10861415714614744 + }, + "parameter.transformer_blocks.38.attn.norm_k.weight": { + "max_abs": 5.844645385742187, + "relative_l2": 0.9242552679150525 + }, + "parameter.transformer_blocks.38.attn.norm_q.weight": { + "max_abs": 5.868889694213867, + "relative_l2": 0.9179216618472709 + }, + "parameter.transformer_blocks.38.attn.to_k.weight": { + "max_abs": 2.9294017601013183, + "relative_l2": 0.9094237527007583 + }, + "parameter.transformer_blocks.38.attn.to_out.0.weight": { + "max_abs": 0.668610428571701, + "relative_l2": 0.9234331916237591 + }, + "parameter.transformer_blocks.38.attn.to_q.weight": { + "max_abs": 1.974614565372467, + "relative_l2": 0.9598699375006258 + }, + "parameter.transformer_blocks.38.attn.to_v.weight": { + "max_abs": 1.4428296136856078, + "relative_l2": 0.9253544945282148 + }, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": { + "max_abs": 145.5050343322754, + "relative_l2": 0.17510562755866485 + }, + "parameter.transformer_blocks.38.ff.net.2.weight": { + "max_abs": 65.21968353271484, + "relative_l2": 0.09506300458020124 + }, + "parameter.transformer_blocks.38.norm1.weight": { + "max_abs": 13.351056938171388, + "relative_l2": 0.8662618342953593 + }, + "parameter.transformer_blocks.38.norm2.weight": { + "max_abs": 703.1446484375, + "relative_l2": 0.16890057341383458 + }, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": { + "max_abs": 6.127394561767578, + "relative_l2": 0.13168527159369647 + }, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": { + "max_abs": 1456.505, + "relative_l2": 0.08965412934922393 + }, + "parameter.transformer_blocks.39.attn.norm_k.weight": { + "max_abs": 4.7785280990600585, + "relative_l2": 1.0694236943687239 + }, + "parameter.transformer_blocks.39.attn.norm_q.weight": { + "max_abs": 4.7358480739593505, + "relative_l2": 1.0685993433522698 + }, + "parameter.transformer_blocks.39.attn.to_k.weight": { + "max_abs": 2.62364173412323, + "relative_l2": 1.124200114084775 + }, + "parameter.transformer_blocks.39.attn.to_out.0.weight": { + "max_abs": 0.757876490354538, + "relative_l2": 0.8481142640103078 + }, + "parameter.transformer_blocks.39.attn.to_q.weight": { + "max_abs": 1.915063087940216, + "relative_l2": 1.1942952625059793 + }, + "parameter.transformer_blocks.39.attn.to_v.weight": { + "max_abs": 1.7589626359939574, + "relative_l2": 0.8567941748841181 + }, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": { + "max_abs": 117.97527587890625, + "relative_l2": 0.14981088979566995 + }, + "parameter.transformer_blocks.39.ff.net.2.weight": { + "max_abs": 182.483515625, + "relative_l2": 0.0897690232787083 + }, + "parameter.transformer_blocks.39.norm1.weight": { + "max_abs": 8.069322471618653, + "relative_l2": 0.9237021151344867 + }, + "parameter.transformer_blocks.39.norm2.weight": { + "max_abs": 1081.0607861328126, + "relative_l2": 0.25401553568912055 + }, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": { + "max_abs": 27.046805267333983, + "relative_l2": 0.26238639729994045 + }, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": { + "max_abs": 6488.14513671875, + "relative_l2": 0.25077882117530165 + }, + "parameter.transformer_blocks.4.attn.norm_k.weight": { + "max_abs": 4.5355954217910766, + "relative_l2": 0.9585719035587216 + }, + "parameter.transformer_blocks.4.attn.norm_q.weight": { + "max_abs": 3.9605391073226928, + "relative_l2": 0.9581588861708147 + }, + "parameter.transformer_blocks.4.attn.to_k.weight": { + "max_abs": 3.6364445734024047, + "relative_l2": 0.8205567058445037 + }, + "parameter.transformer_blocks.4.attn.to_out.0.weight": { + "max_abs": 2.0168289589881896, + "relative_l2": 0.789058322187198 + }, + "parameter.transformer_blocks.4.attn.to_q.weight": { + "max_abs": 5.848101501464844, + "relative_l2": 0.8094523151615062 + }, + "parameter.transformer_blocks.4.attn.to_v.weight": { + "max_abs": 1.6423358964920043, + "relative_l2": 0.8591711598148015 + }, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": { + "max_abs": 1159.7981518554689, + "relative_l2": 0.2674986928299999 + }, + "parameter.transformer_blocks.4.ff.net.2.weight": { + "max_abs": 402.25548828125, + "relative_l2": 0.24103781594722395 + }, + "parameter.transformer_blocks.4.norm1.weight": { + "max_abs": 17.1214174079895, + "relative_l2": 0.8361116592439933 + }, + "parameter.transformer_blocks.4.norm2.weight": { + "max_abs": 801.8566235351562, + "relative_l2": 0.2561850092970625 + }, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": { + "max_abs": 3.2005223083496093, + "relative_l2": 0.12521540524849173 + }, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": { + "max_abs": 705.36681640625, + "relative_l2": 0.08374239247167095 + }, + "parameter.transformer_blocks.40.attn.norm_k.weight": { + "max_abs": 3.4289771366119384, + "relative_l2": 1.0478765270951127 + }, + "parameter.transformer_blocks.40.attn.norm_q.weight": { + "max_abs": 4.846121912002563, + "relative_l2": 1.0795186674130992 + }, + "parameter.transformer_blocks.40.attn.to_k.weight": { + "max_abs": 3.5802519369125365, + "relative_l2": 1.140589196399972 + }, + "parameter.transformer_blocks.40.attn.to_out.0.weight": { + "max_abs": 0.6302498775720596, + "relative_l2": 0.7935126763845027 + }, + "parameter.transformer_blocks.40.attn.to_q.weight": { + "max_abs": 5.073410515785217, + "relative_l2": 1.3077308914794323 + }, + "parameter.transformer_blocks.40.attn.to_v.weight": { + "max_abs": 2.0222979831695556, + "relative_l2": 0.8269094116055061 + }, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": { + "max_abs": 210.19929779052734, + "relative_l2": 0.15512173606582044 + }, + "parameter.transformer_blocks.40.ff.net.2.weight": { + "max_abs": 229.26128662109375, + "relative_l2": 0.0881041806052548 + }, + "parameter.transformer_blocks.40.norm1.weight": { + "max_abs": 4.467885737419128, + "relative_l2": 0.9969333401675976 + }, + "parameter.transformer_blocks.40.norm2.weight": { + "max_abs": 194.94732177734375, + "relative_l2": 0.1290674790515727 + }, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": { + "max_abs": 4.080959205627441, + "relative_l2": 0.09656603372527352 + }, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": { + "max_abs": 953.324580078125, + "relative_l2": 0.08195072423360418 + }, + "parameter.transformer_blocks.41.attn.norm_k.weight": { + "max_abs": 2.757928853034973, + "relative_l2": 0.7563508971808889 + }, + "parameter.transformer_blocks.41.attn.norm_q.weight": { + "max_abs": 2.8040596294403075, + "relative_l2": 0.7409621277067486 + }, + "parameter.transformer_blocks.41.attn.to_k.weight": { + "max_abs": 1.3457395482063292, + "relative_l2": 0.8211947853401178 + }, + "parameter.transformer_blocks.41.attn.to_out.0.weight": { + "max_abs": 0.731994201540947, + "relative_l2": 0.8026639137656787 + }, + "parameter.transformer_blocks.41.attn.to_q.weight": { + "max_abs": 0.8410365927219391, + "relative_l2": 0.8438698268641297 + }, + "parameter.transformer_blocks.41.attn.to_v.weight": { + "max_abs": 1.642683868408203, + "relative_l2": 0.8220985051578446 + }, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": { + "max_abs": 154.27525985717773, + "relative_l2": 0.12454783262895078 + }, + "parameter.transformer_blocks.41.ff.net.2.weight": { + "max_abs": 228.506708984375, + "relative_l2": 0.0842293967140266 + }, + "parameter.transformer_blocks.41.norm1.weight": { + "max_abs": 5.2925452041625975, + "relative_l2": 0.7890362572700109 + }, + "parameter.transformer_blocks.41.norm2.weight": { + "max_abs": 668.660029296875, + "relative_l2": 0.11268412126563454 + }, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": { + "max_abs": 4.917359237670898, + "relative_l2": 0.09542605967411237 + }, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": { + "max_abs": 1169.406611328125, + "relative_l2": 0.07712454346483907 + }, + "parameter.transformer_blocks.42.attn.norm_k.weight": { + "max_abs": 4.035946731567383, + "relative_l2": 0.8238593889325051 + }, + "parameter.transformer_blocks.42.attn.norm_q.weight": { + "max_abs": 4.056439523696899, + "relative_l2": 0.7996155756600228 + }, + "parameter.transformer_blocks.42.attn.to_k.weight": { + "max_abs": 1.6770926761627196, + "relative_l2": 0.7641442861987374 + }, + "parameter.transformer_blocks.42.attn.to_out.0.weight": { + "max_abs": 0.5470972406864166, + "relative_l2": 0.5442167760593228 + }, + "parameter.transformer_blocks.42.attn.to_q.weight": { + "max_abs": 1.1088188636302947, + "relative_l2": 0.8081196804260867 + }, + "parameter.transformer_blocks.42.attn.to_v.weight": { + "max_abs": 1.2766135680675506, + "relative_l2": 0.6157095221288996 + }, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": { + "max_abs": 155.3827389526367, + "relative_l2": 0.11766137098501636 + }, + "parameter.transformer_blocks.42.ff.net.2.weight": { + "max_abs": 129.46139038085937, + "relative_l2": 0.07740887728938908 + }, + "parameter.transformer_blocks.42.norm1.weight": { + "max_abs": 4.736500029563904, + "relative_l2": 0.7671526716518661 + }, + "parameter.transformer_blocks.42.norm2.weight": { + "max_abs": 270.3913067626953, + "relative_l2": 0.10944300490817432 + }, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": { + "max_abs": 5.132880096435547, + "relative_l2": 0.11273468825966387 + }, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": { + "max_abs": 1240.3321484375, + "relative_l2": 0.08187268981591711 + }, + "parameter.transformer_blocks.43.attn.norm_k.weight": { + "max_abs": 6.753908519744873, + "relative_l2": 0.9205244132931526 + }, + "parameter.transformer_blocks.43.attn.norm_q.weight": { + "max_abs": 6.826011304855347, + "relative_l2": 0.9200605195684711 + }, + "parameter.transformer_blocks.43.attn.to_k.weight": { + "max_abs": 2.645546441078186, + "relative_l2": 0.9111346814717839 + }, + "parameter.transformer_blocks.43.attn.to_out.0.weight": { + "max_abs": 0.7447509664297104, + "relative_l2": 0.571198163788551 + }, + "parameter.transformer_blocks.43.attn.to_q.weight": { + "max_abs": 1.889353756904602, + "relative_l2": 0.9313925335595282 + }, + "parameter.transformer_blocks.43.attn.to_v.weight": { + "max_abs": 1.4668158936500548, + "relative_l2": 0.599385356863508 + }, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": { + "max_abs": 300.6612957763672, + "relative_l2": 0.1136452440596051 + }, + "parameter.transformer_blocks.43.ff.net.2.weight": { + "max_abs": 66.31138122558593, + "relative_l2": 0.08056046834619426 + }, + "parameter.transformer_blocks.43.norm1.weight": { + "max_abs": 13.360413913726807, + "relative_l2": 0.7114983351779064 + }, + "parameter.transformer_blocks.43.norm2.weight": { + "max_abs": 156.96438110351562, + "relative_l2": 0.09389628263383776 + }, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": { + "max_abs": 17.729781036376952, + "relative_l2": 0.11705536326286792 + }, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": { + "max_abs": 4208.96837890625, + "relative_l2": 0.1049058565796688 + }, + "parameter.transformer_blocks.44.attn.norm_k.weight": { + "max_abs": 2.269458894729614, + "relative_l2": 0.7091757295194187 + }, + "parameter.transformer_blocks.44.attn.norm_q.weight": { + "max_abs": 2.164663915634155, + "relative_l2": 0.7054317049993656 + }, + "parameter.transformer_blocks.44.attn.to_k.weight": { + "max_abs": 1.7072026181221007, + "relative_l2": 0.8324279544785277 + }, + "parameter.transformer_blocks.44.attn.to_out.0.weight": { + "max_abs": 0.8187230277061462, + "relative_l2": 0.5460598839539066 + }, + "parameter.transformer_blocks.44.attn.to_q.weight": { + "max_abs": 1.233398710489273, + "relative_l2": 0.8872870872880635 + }, + "parameter.transformer_blocks.44.attn.to_v.weight": { + "max_abs": 1.7586665201187133, + "relative_l2": 0.6282548434192357 + }, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": { + "max_abs": 149.37920272827148, + "relative_l2": 0.10937401822764103 + }, + "parameter.transformer_blocks.44.ff.net.2.weight": { + "max_abs": 92.94004333496093, + "relative_l2": 0.07597955498059568 + }, + "parameter.transformer_blocks.44.norm1.weight": { + "max_abs": 8.42225254058838, + "relative_l2": 0.6669851115364491 + }, + "parameter.transformer_blocks.44.norm2.weight": { + "max_abs": 848.308466796875, + "relative_l2": 0.18199623008311203 + }, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": { + "max_abs": 6.1834706115722655, + "relative_l2": 0.07386063824844913 + }, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": { + "max_abs": 1467.5729931640626, + "relative_l2": 0.06511389998485793 + }, + "parameter.transformer_blocks.45.attn.norm_k.weight": { + "max_abs": 2.2533670711517333, + "relative_l2": 0.771607820457134 + }, + "parameter.transformer_blocks.45.attn.norm_q.weight": { + "max_abs": 2.2690254497528075, + "relative_l2": 0.7876834677150717 + }, + "parameter.transformer_blocks.45.attn.to_k.weight": { + "max_abs": 2.8169800090789794, + "relative_l2": 0.768766108578039 + }, + "parameter.transformer_blocks.45.attn.to_out.0.weight": { + "max_abs": 0.510720004439354, + "relative_l2": 0.34819230093551706 + }, + "parameter.transformer_blocks.45.attn.to_q.weight": { + "max_abs": 25.44289291381836, + "relative_l2": 0.4792674547547203 + }, + "parameter.transformer_blocks.45.attn.to_v.weight": { + "max_abs": 1.0604091036319732, + "relative_l2": 0.3754161340434127 + }, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": { + "max_abs": 201.54845703125, + "relative_l2": 0.09227611490400561 + }, + "parameter.transformer_blocks.45.ff.net.2.weight": { + "max_abs": 119.0656689453125, + "relative_l2": 0.07023417264246236 + }, + "parameter.transformer_blocks.45.norm1.weight": { + "max_abs": 9.380532150268555, + "relative_l2": 0.569506798252519 + }, + "parameter.transformer_blocks.45.norm2.weight": { + "max_abs": 368.22149169921875, + "relative_l2": 0.1011082779913207 + }, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": { + "max_abs": 8.651469116210938, + "relative_l2": 0.06329771186354746 + }, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": { + "max_abs": 1905.06359375, + "relative_l2": 0.05911674769537031 + }, + "parameter.transformer_blocks.46.attn.norm_k.weight": { + "max_abs": 2.558462266921997, + "relative_l2": 0.585320050556141 + }, + "parameter.transformer_blocks.46.attn.norm_q.weight": { + "max_abs": 2.6533019590377807, + "relative_l2": 0.5598057701695375 + }, + "parameter.transformer_blocks.46.attn.to_k.weight": { + "max_abs": 2.455328290462494, + "relative_l2": 0.6239648127845449 + }, + "parameter.transformer_blocks.46.attn.to_out.0.weight": { + "max_abs": 0.549632688164711, + "relative_l2": 0.24239122366787058 + }, + "parameter.transformer_blocks.46.attn.to_q.weight": { + "max_abs": 2.604859833717346, + "relative_l2": 0.6585910041072844 + }, + "parameter.transformer_blocks.46.attn.to_v.weight": { + "max_abs": 1.3638730692863463, + "relative_l2": 0.24932692934413414 + }, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": { + "max_abs": 410.6479443359375, + "relative_l2": 0.07948161125168707 + }, + "parameter.transformer_blocks.46.ff.net.2.weight": { + "max_abs": 971.273310546875, + "relative_l2": 0.06522679281801508 + }, + "parameter.transformer_blocks.46.norm1.weight": { + "max_abs": 10.147108917236329, + "relative_l2": 0.2227172172388825 + }, + "parameter.transformer_blocks.46.norm2.weight": { + "max_abs": 2693.47033203125, + "relative_l2": 0.16338661750553896 + }, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": { + "max_abs": 16.257235412597655, + "relative_l2": 0.07627798102860456 + }, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": { + "max_abs": 3683.17296875, + "relative_l2": 0.06780493854220315 + }, + "parameter.transformer_blocks.47.attn.norm_k.weight": { + "max_abs": 5.5190323638916015, + "relative_l2": 0.4231322946674436 + }, + "parameter.transformer_blocks.47.attn.norm_q.weight": { + "max_abs": 5.7188004302978515, + "relative_l2": 0.4193198161245315 + }, + "parameter.transformer_blocks.47.attn.to_k.weight": { + "max_abs": 11.963134651184083, + "relative_l2": 0.330995632161804 + }, + "parameter.transformer_blocks.47.attn.to_out.0.weight": { + "max_abs": 1.202230279445648, + "relative_l2": 0.18708588305373947 + }, + "parameter.transformer_blocks.47.attn.to_q.weight": { + "max_abs": 4.0040234375, + "relative_l2": 0.35734895930583077 + }, + "parameter.transformer_blocks.47.attn.to_v.weight": { + "max_abs": 11.460832719802857, + "relative_l2": 0.20559483258091166 + }, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": { + "max_abs": 1193.923701171875, + "relative_l2": 0.07835797547764027 + }, + "parameter.transformer_blocks.47.ff.net.2.weight": { + "max_abs": 1576.885859375, + "relative_l2": 0.051879736665692025 + }, + "parameter.transformer_blocks.47.norm1.weight": { + "max_abs": 15.85252655029297, + "relative_l2": 0.34502870028341526 + }, + "parameter.transformer_blocks.47.norm2.weight": { + "max_abs": 1780.4141796875, + "relative_l2": 0.08152340720341918 + }, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": { + "max_abs": 26.97588623046875, + "relative_l2": 0.04928975797603363 + }, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": { + "max_abs": 7351.505, + "relative_l2": 0.049891999286638504 + }, + "parameter.transformer_blocks.48.attn.norm_k.weight": { + "max_abs": 2.1515078592300414, + "relative_l2": 0.19186781900909602 + }, + "parameter.transformer_blocks.48.attn.norm_q.weight": { + "max_abs": 2.2295125770568847, + "relative_l2": 0.1988307599059829 + }, + "parameter.transformer_blocks.48.attn.to_k.weight": { + "max_abs": 4.589643363952637, + "relative_l2": 0.2042588278502384 + }, + "parameter.transformer_blocks.48.attn.to_out.0.weight": { + "max_abs": 1.014324550628662, + "relative_l2": 0.096045697959378 + }, + "parameter.transformer_blocks.48.attn.to_q.weight": { + "max_abs": 7.781969909667969, + "relative_l2": 0.2424748187686756 + }, + "parameter.transformer_blocks.48.attn.to_v.weight": { + "max_abs": 2.8530453491210936, + "relative_l2": 0.0977908986332009 + }, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": { + "max_abs": 1301.2202099609376, + "relative_l2": 0.06038765587406777 + }, + "parameter.transformer_blocks.48.ff.net.2.weight": { + "max_abs": 1068.0698193359376, + "relative_l2": 0.045891303085466316 + }, + "parameter.transformer_blocks.48.norm1.weight": { + "max_abs": 6.685605888366699, + "relative_l2": 0.12466980346726744 + }, + "parameter.transformer_blocks.48.norm2.weight": { + "max_abs": 2388.690546875, + "relative_l2": 0.05236776060520865 + }, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": { + "max_abs": 8.821124916076661, + "relative_l2": 0.10311497979458151 + }, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": { + "max_abs": 1441.112421875, + "relative_l2": 0.043180249258132096 + }, + "parameter.transformer_blocks.49.attn.norm_k.weight": { + "max_abs": 4.704728012084961, + "relative_l2": 0.22726517993617218 + }, + "parameter.transformer_blocks.49.attn.norm_q.weight": { + "max_abs": 3.2715772438049315, + "relative_l2": 0.22292634803167102 + }, + "parameter.transformer_blocks.49.attn.to_k.weight": { + "max_abs": 180.87781799316406, + "relative_l2": 0.3738389842182632 + }, + "parameter.transformer_blocks.49.attn.to_out.0.weight": { + "max_abs": 2.1655740785598754, + "relative_l2": 0.07821954800563145 + }, + "parameter.transformer_blocks.49.attn.to_q.weight": { + "max_abs": 7475.51427734375, + "relative_l2": 0.9976771172686144 + }, + "parameter.transformer_blocks.49.attn.to_v.weight": { + "max_abs": 3.403834466934204, + "relative_l2": 0.07990269277842786 + }, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": { + "max_abs": 1048.5288037109376, + "relative_l2": 0.04472210669022003 + }, + "parameter.transformer_blocks.49.ff.net.2.weight": { + "max_abs": 1158.47814453125, + "relative_l2": 0.028978917566814043 + }, + "parameter.transformer_blocks.49.norm1.weight": { + "max_abs": 79.21369064331054, + "relative_l2": 0.8639755692388044 + }, + "parameter.transformer_blocks.49.norm2.weight": { + "max_abs": 558.84435546875, + "relative_l2": 0.04921952358353799 + }, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": { + "max_abs": 34.00472152709961, + "relative_l2": 0.26746757834431667 + }, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": { + "max_abs": 8061.69640625, + "relative_l2": 0.2452543997549339 + }, + "parameter.transformer_blocks.5.attn.norm_k.weight": { + "max_abs": 3.9417589950561522, + "relative_l2": 1.2664483396718582 + }, + "parameter.transformer_blocks.5.attn.norm_q.weight": { + "max_abs": 3.946983938217163, + "relative_l2": 1.2224533507336455 + }, + "parameter.transformer_blocks.5.attn.to_k.weight": { + "max_abs": 2.9919197607040404, + "relative_l2": 1.1727527541387843 + }, + "parameter.transformer_blocks.5.attn.to_out.0.weight": { + "max_abs": 2.8459457206726073, + "relative_l2": 0.510954297011832 + }, + "parameter.transformer_blocks.5.attn.to_q.weight": { + "max_abs": 5.031534795761108, + "relative_l2": 1.2450305679568294 + }, + "parameter.transformer_blocks.5.attn.to_v.weight": { + "max_abs": 1.0378571796417235, + "relative_l2": 0.539951216146229 + }, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": { + "max_abs": 493.2738751220703, + "relative_l2": 0.25716856811430633 + }, + "parameter.transformer_blocks.5.ff.net.2.weight": { + "max_abs": 234.38091552734375, + "relative_l2": 0.24140053738667538 + }, + "parameter.transformer_blocks.5.norm1.weight": { + "max_abs": 10.863516931533814, + "relative_l2": 0.946200087408264 + }, + "parameter.transformer_blocks.5.norm2.weight": { + "max_abs": 10746.702265625, + "relative_l2": 0.5631723714081652 + }, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": { + "max_abs": 8.736480836868287, + "relative_l2": 0.28376338375463944 + }, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": { + "max_abs": 2082.391901855469, + "relative_l2": 0.2422681374476185 + }, + "parameter.transformer_blocks.6.attn.norm_k.weight": { + "max_abs": 3.949999814033508, + "relative_l2": 0.6616773576672742 + }, + "parameter.transformer_blocks.6.attn.norm_q.weight": { + "max_abs": 4.279451971054077, + "relative_l2": 0.6088499948792981 + }, + "parameter.transformer_blocks.6.attn.to_k.weight": { + "max_abs": 3.365487461090088, + "relative_l2": 0.671770525013 + }, + "parameter.transformer_blocks.6.attn.to_out.0.weight": { + "max_abs": 1.4420253086090087, + "relative_l2": 0.7280617994287809 + }, + "parameter.transformer_blocks.6.attn.to_q.weight": { + "max_abs": 7.461075429916382, + "relative_l2": 0.671384648994924 + }, + "parameter.transformer_blocks.6.attn.to_v.weight": { + "max_abs": 1.8169648694992064, + "relative_l2": 0.7366241094623122 + }, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": { + "max_abs": 1212.0285290527345, + "relative_l2": 0.26324446456438144 + }, + "parameter.transformer_blocks.6.ff.net.2.weight": { + "max_abs": 201.01442993164062, + "relative_l2": 0.24355468985475448 + }, + "parameter.transformer_blocks.6.norm1.weight": { + "max_abs": 8.35041940689087, + "relative_l2": 0.7145662697849823 + }, + "parameter.transformer_blocks.6.norm2.weight": { + "max_abs": 942.755732421875, + "relative_l2": 0.23473965726291585 + }, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": { + "max_abs": 7.761792545318603, + "relative_l2": 0.3173178505919617 + }, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": { + "max_abs": 1557.43126953125, + "relative_l2": 0.26760021703293335 + }, + "parameter.transformer_blocks.7.attn.norm_k.weight": { + "max_abs": 4.056298260688782, + "relative_l2": 0.8996356894661431 + }, + "parameter.transformer_blocks.7.attn.norm_q.weight": { + "max_abs": 3.4165216732025145, + "relative_l2": 0.8404564580857118 + }, + "parameter.transformer_blocks.7.attn.to_k.weight": { + "max_abs": 1.8837275552749633, + "relative_l2": 0.8492421136620586 + }, + "parameter.transformer_blocks.7.attn.to_out.0.weight": { + "max_abs": 1.1330161440372466, + "relative_l2": 0.83062026638771 + }, + "parameter.transformer_blocks.7.attn.to_q.weight": { + "max_abs": 2.040867094993591, + "relative_l2": 0.8312702507157053 + }, + "parameter.transformer_blocks.7.attn.to_v.weight": { + "max_abs": 1.3269965398311614, + "relative_l2": 0.8473707200751844 + }, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": { + "max_abs": 2168.759089355469, + "relative_l2": 0.2835345569316502 + }, + "parameter.transformer_blocks.7.ff.net.2.weight": { + "max_abs": 156.14725006103515, + "relative_l2": 0.24158799615078969 + }, + "parameter.transformer_blocks.7.norm1.weight": { + "max_abs": 16.6573380279541, + "relative_l2": 0.9899422775170169 + }, + "parameter.transformer_blocks.7.norm2.weight": { + "max_abs": 1497.3561962890626, + "relative_l2": 0.2729074793630908 + }, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": { + "max_abs": 13.849670295715333, + "relative_l2": 0.2716151788658495 + }, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": { + "max_abs": 3290.2363232421876, + "relative_l2": 0.24767352098462966 + }, + "parameter.transformer_blocks.8.attn.norm_k.weight": { + "max_abs": 1.5559553551673888, + "relative_l2": 0.8947608029700306 + }, + "parameter.transformer_blocks.8.attn.norm_q.weight": { + "max_abs": 1.506643121242523, + "relative_l2": 0.9384623469999616 + }, + "parameter.transformer_blocks.8.attn.to_k.weight": { + "max_abs": 1.2076490569114684, + "relative_l2": 0.9773064344087591 + }, + "parameter.transformer_blocks.8.attn.to_out.0.weight": { + "max_abs": 1.310686831474304, + "relative_l2": 0.633518587303856 + }, + "parameter.transformer_blocks.8.attn.to_q.weight": { + "max_abs": 2.280301456451416, + "relative_l2": 0.9573002642277444 + }, + "parameter.transformer_blocks.8.attn.to_v.weight": { + "max_abs": 0.8934161114692688, + "relative_l2": 0.6871977728427118 + }, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": { + "max_abs": 832.8469494628906, + "relative_l2": 0.2663712599558485 + }, + "parameter.transformer_blocks.8.ff.net.2.weight": { + "max_abs": 145.30633056640625, + "relative_l2": 0.24006590318715618 + }, + "parameter.transformer_blocks.8.norm1.weight": { + "max_abs": 4.724163656234741, + "relative_l2": 0.8735804939347517 + }, + "parameter.transformer_blocks.8.norm2.weight": { + "max_abs": 3142.030390625, + "relative_l2": 0.23299972377987604 + }, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": { + "max_abs": 6.839249973297119, + "relative_l2": 0.2846913416140406 + }, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": { + "max_abs": 1627.18126953125, + "relative_l2": 0.24772190477496486 + }, + "parameter.transformer_blocks.9.attn.norm_k.weight": { + "max_abs": 3.0111261653900145, + "relative_l2": 0.9238515507972674 + }, + "parameter.transformer_blocks.9.attn.norm_q.weight": { + "max_abs": 3.0268102931976317, + "relative_l2": 0.932745375224013 + }, + "parameter.transformer_blocks.9.attn.to_k.weight": { + "max_abs": 2.0182001042366027, + "relative_l2": 0.9300877917722118 + }, + "parameter.transformer_blocks.9.attn.to_out.0.weight": { + "max_abs": 0.7163272845745087, + "relative_l2": 0.6560424752082701 + }, + "parameter.transformer_blocks.9.attn.to_q.weight": { + "max_abs": 4.336956386566162, + "relative_l2": 0.9655093562800879 + }, + "parameter.transformer_blocks.9.attn.to_v.weight": { + "max_abs": 1.6555756378173827, + "relative_l2": 0.6974471762705327 + }, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": { + "max_abs": 450.18311584472656, + "relative_l2": 0.2535606699307462 + }, + "parameter.transformer_blocks.9.ff.net.2.weight": { + "max_abs": 119.15561950683593, + "relative_l2": 0.23522267253064247 + }, + "parameter.transformer_blocks.9.norm1.weight": { + "max_abs": 6.78931510925293, + "relative_l2": 0.8447450481755553 + }, + "parameter.transformer_blocks.9.norm2.weight": { + "max_abs": 476.0451611328125, + "relative_l2": 0.23699570586893662 + } + } +} diff --git a/examples/models/minimax_h3/generation-compare.sbatch b/examples/models/minimax_h3/generation-compare.sbatch new file mode 100644 index 000000000..c188e3c4a --- /dev/null +++ b/examples/models/minimax_h3/generation-compare.sbatch @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Quality comparison of the three candidate seeds with the provider reference (CPU only). +# Submit after the generation array finished successfully (--dependency=afterok:). +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=128G +#SBATCH --time=06:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +reference_root=$3 +result_root=$4 +extra_dependencies=$5 # overlay with scikit-image / scipy (the reference job used it too) +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +commit=$(git rev-parse HEAD) +export PYTHONPATH="$source_dir:$extra_dependencies" +"$python_bin" -m tools.validation.models.h3_generation_compare \ + --reference "$reference_root/generation-reference" \ + --candidate "$result_root/generation-${commit:0:8}" \ + --contract examples/models/minimax_h3/generation-contract.json \ + --reference-root "$reference_root" \ + --reference-manifest examples/models/minimax_h3/reference-manifest.json \ + --output "$result_root/generation-comparison-${commit:0:8}.json" diff --git a/examples/models/minimax_h3/generation-contract.json b/examples/models/minimax_h3/generation-contract.json new file mode 100644 index 000000000..7c29221c6 --- /dev/null +++ b/examples/models/minimax_h3/generation-contract.json @@ -0,0 +1,18 @@ +{ + "schema": 1, + "model_revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "provider_revision": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "workflow": "fl2va", + "seeds": [42, 43, 44], + "num_frames": 124, + "height": 768, + "width": 1344, + "fps": 24, + "sampling_rate": 32000, + "num_inference_steps": 50, + "preflight": {"num_frames": 124, "height": 128, "width": 128, "num_inference_steps": 2}, + "quality_bounds": {"minimum_psnr_db": 30, "minimum_ssim": 0.99, "maximum_temporal_rmse": 0.01, "maximum_audio_relative_l2": 0.05}, + "quality_bounds_policy": "Conservative absolute gates fixed before candidate generation; a failure remains a failure. Reference repeats must first satisfy these same bounds.", + "attention_query_chunk": 256, + "status": "unqualified_until_full_reference_and_candidate_runs_pass" +} diff --git a/examples/models/minimax_h3/generation-inputs.json b/examples/models/minimax_h3/generation-inputs.json new file mode 100644 index 000000000..ea62adc23 --- /dev/null +++ b/examples/models/minimax_h3/generation-inputs.json @@ -0,0 +1,8 @@ +{ + "image_sha256": "352e959e08b886406d3766cee25ebe7bbf679e6d53eeffbd1c607b8dc93e2068", + "image_url": "https://cdn.hailuoai.com/prod/hailuo_demo/testsets/H3_AA_I2VA/gallery/sr_v17_variants_seed42_43_20260724/inputs/4a3a90bf9100_KDmcbkhzYo5sjjxr9FqcVmWVnzb.png", + "prompt_reused_verbatim": true, + "prompt_sha256": "fba3868d31e9911a81d60c191261ec2192c827bcaa25c937660bf35aeb0b1673", + "source_revision": "bbee488da3baf65817f4f8826b1f6e89cc28fda7", + "source_suite": "tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py" +} diff --git a/examples/models/minimax_h3/generation-prompt.txt b/examples/models/minimax_h3/generation-prompt.txt new file mode 100644 index 000000000..ba5be1ceb --- /dev/null +++ b/examples/models/minimax_h3/generation-prompt.txt @@ -0,0 +1,7 @@ +For the target video, at 0.00 seconds into the target video, (from [Shot 1]) is fully referenced. + +integrated_multimodal_description: [Shot 1] This is a live-action, cinematic shot with a shallow depth of field. The camera holds a perfectly static shot throughout the entire eight-second duration, capturing a cozy family gathering in a traditional Japanese dining room. The scene opens with a large, intricately patterned blue and white ceramic bowl of ramen in the immediate foreground, rendered in crisp, sharp focus. The bowl sits on a smooth, polished long wooden table. Inside the bowl, a rich, oily golden-brown broth surrounds yellow wavy noodles, topped with two thick, round slices of chashu pork featuring visible fat marbling and a distinct spiral meat pattern. A generous mound of freshly chopped, bright green scallions rests in the center, and a crisp, dark green rectangular sheet of nori seaweed is tucked into the right edge. To the left of the bowl, a pair of light brown wooden chopsticks rests horizontally on a small, dark rectangular chopstick rest, near a small cylindrical ceramic teacup with blue painted patterns. On the right side of the table, a spherical paper lantern with a ribbed bamboo frame sits on a black wooden base. In the background, a large family of seven is gathered around the table, initially appearing as a soft, blurred presence. Behind them, traditional Japanese sliding shoji screens with wooden lattice frames are open, revealing a bright outdoor scene with lush green trees. Early in the clip, the thick, white steam rising from the hot ramen broth immediately intensifies, billowing upwards in thick, swirling clouds that dance continuously above the bowl. As the clip progresses into the middle seconds, the camera maintains its static position while the focus begins a deliberate, smooth shift deeper into the room. The foreground ramen bowl, its vibrant ingredients, and the rising steam gradually soften into a hazy, out-of-focus blur. Simultaneously, the family members in the background come into sharp, detailed clarity. The heavy steam continues to rise from the foreground, creating a dynamic, translucent veil between the camera and the family. With the focus now firmly locked on the background, the vibrant family dinner comes alive. The man in the dark navy blue long-sleeved shirt on the left leans forward, his mouth moving animatedly in a silent exchange. The young girl in the crisp white short-sleeved t-shirt beside him smiles brightly, looking toward the center of the table. The woman on the far left, wearing a soft light blue long-sleeved blouse, turns her head slightly, smiling gently. Across the table, the woman in the light grey button-down shirt smiles broadly, her eyes crinkling, as she rests her hands near her plate. The woman in the dark grey top further back uses her wooden chopsticks to pick up a small piece of food from a central ceramic dish filled with bright red pickled vegetables. The woman in the center back in the light grey sweater smiles gently, her hands clasped softly in front of her, observing the interaction. Throughout the remainder of the clip, the family continues their lively physical interaction, their mouths moving in continuous, silent cadences of conversation, while the thick, white steam from the blurred ramen bowl in the foreground never stops rising, adding a comforting atmosphere to the warm gathering. + +overall_soundscape: The soundscape begins with a quiet room tone mixed with the faint, airy rustle of the thick steam billowing from the hot ramen bowl in the foreground, accompanied by the subtle, continuous hissing and bubbling of the rich broth. As the visual focus shifts deeper into the room, the physical sounds of the bustling family dinner become dominant in the foreground. The clear, sharp clinking of ceramic bowls and wooden chopsticks touching plates is clearly heard as the family members reach for food. This is followed by the faint, muffled thud of a cup being set down on the smooth wooden table, and the subtle, rhythmic rustle of cotton and wool clothing as the family members lean forward and gesture, perfectly capturing the lively, physical atmosphere of the shared meal. + +non_diegetic_music: A gentle, heartwarming acoustic guitar melody plays softly in the background, accompanied by the subtle, resonant notes of a traditional Japanese koto. The music maintains a slow, comforting tempo that enhances the cozy, nostalgic, and joyful atmosphere of the family gathering. diff --git a/examples/models/minimax_h3/generation-resume-check.sbatch b/examples/models/minimax_h3/generation-resume-check.sbatch new file mode 100644 index 000000000..be544b38f --- /dev/null +++ b/examples/models/minimax_h3/generation-resume-check.sbatch @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Resume check on the real pipeline at preflight size (128x128, 6 sampling steps, 5 transformer +# calls): one run without interruption, and two preempted runs that are resumed -- "cut" +# after 2, 4 and 5 calls (the last one after the final step, before decoding, so the final +# attempt has nothing left to compute) and "cutverify" after 2 calls with --verify-last-step. +# Every step file, the initial inputs and the decoded media must be byte-equal to "whole". +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=04:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +reference_image=$4 +provider_source=$5 +extra_dependencies=$6 +result_root=$7 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 +export PATH="$CUDA_HOME/bin:$PATH" +out="$result_root/resume-check-$(git rev-parse --short=8 HEAD)" +mkdir "$out" +gen() { + "$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --preflight --preflight-steps 6 "$@" +} +gen --output "$out/whole" > "$out/whole.log" 2>&1 +gen --output "$out/cut" --stop-after-step 2 > "$out/cut-1.log" 2>&1 || true +gen --output "$out/cut" --resume --stop-after-step 4 > "$out/cut-2.log" 2>&1 || true +gen --output "$out/cut" --resume --stop-after-step 5 > "$out/cut-3.log" 2>&1 || true +gen --output "$out/cut" --resume > "$out/cut-4.log" 2>&1 +gen --output "$out/cutverify" --stop-after-step 2 > "$out/cutverify-1.log" 2>&1 || true +gen --output "$out/cutverify" --resume --verify-last-step > "$out/cutverify-2.log" 2>&1 +"$python_bin" - "$out" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +whole = root / "whole" / "42" +whole_report = json.loads((whole / "report.json").read_text()) +names = sorted( + p.name + for p in whole.iterdir() + if p.name.startswith("step-") and p.suffix == ".pt" + or p.name in ("initial.pt", "video.npy", "audio.npy") +) +result = {"compared_files": len(names), "steps": whole_report["steps"], "runs": {}} +for name in ("cut", "cutverify"): + directory = root / name / "42" + report = json.loads((directory / "report.json").read_text()) + result["runs"][name] = { + "differing": [n for n in names if (whole / n).read_bytes() != (directory / n).read_bytes()], + "resume": {k: v for k, v in report["resume"].items() if k != "attempts"}, + "attempts": len(report["resume"]["attempts"]), + "trace_calls_match": report["trace"]["calls"].keys() == whole_report["trace"]["calls"].keys() + and report["trace"]["candidate_calls"].keys() + == whole_report["trace"]["candidate_calls"].keys(), + "failure_markers": sorted(p.name for p in (root / name).glob("failure*.json")), + } +(root / "resume-check.json").write_text(json.dumps(result, indent=2) + "\n") +print(json.dumps(result, indent=2)) +cut, verify = result["runs"]["cut"], result["runs"]["cutverify"] +assert not cut["differing"] and not verify["differing"], result +assert cut["resume"]["real_forward_calls"] == 0, "the final cut attempt should replay everything" +assert cut["resume"]["recomputed_and_verified_step"] is None +assert verify["resume"]["recomputed_and_verified_step"] is not None +assert cut["trace_calls_match"] and verify["trace_calls_match"] +PY diff --git a/examples/models/minimax_h3/generation.sbatch b/examples/models/minimax_h3/generation.sbatch new file mode 100644 index 000000000..ca1a97e0e --- /dev/null +++ b/examples/models/minimax_h3/generation.sbatch @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Full-resolution candidate generation, one seed per array task (--array=0-2). Preemption +# requeues the task; the generation continues from its saved steps (--resume). +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +#SBATCH --requeue +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +gate_job=$4 # job id of the candidate.sbatch run whose gates passed for this commit +result_root=$5 +reference_image=$6 +provider_source=$7 +extra_dependencies=$8 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +commit=$(git rev-parse HEAD) +# Only a passing numerical comparison of the same model code permits full generation: the gates +# passed at a commit whose model paths (h3_generation_resume.MODEL_PATHS) are byte-identical to +# this checkout's. The same commit qualifies trivially. +gate_dir= +for attempt in "$result_root"/job-"$gate_job"-attempt-*; do + if [[ -f $attempt/commit.txt ]] && "$python_bin" - "$attempt/commit.txt" "$attempt/full-model-acceptance.json" \ + 2> /dev/null <<'PY' +import json +import sys + +from rl_engine.validation.models.h3_generation_resume import model_trees + +gate_commit = open(sys.argv[1]).read().strip() +assert json.load(open(sys.argv[2]))["status"] == "passed" +assert model_trees(gate_commit) == model_trees("HEAD") +PY + then + gate_dir=$attempt + fi +done +[[ -n $gate_dir ]] || { echo "no passing numerical gate for these model paths in job $gate_job" >&2; exit 1; } +seeds=($("$python_bin" -c 'import json; print(*json.load(open("examples/models/minimax_h3/generation-contract.json"))["seeds"])')) +seed=${seeds[$SLURM_ARRAY_TASK_ID]} +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 +export PATH="$CUDA_HOME/bin:$PATH" +out="$result_root/generation-${commit:0:8}" +mkdir -p "$out" +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --seeds "$seed" --resume --output "$out" \ + >> "$out/seed-$seed.log" 2>&1 diff --git a/examples/models/minimax_h3/reference-manifest.json b/examples/models/minimax_h3/reference-manifest.json new file mode 100644 index 000000000..acfe57486 --- /dev/null +++ b/examples/models/minimax_h3/reference-manifest.json @@ -0,0 +1,1641 @@ +{ + "source_commit": "2681cadb4dd14c8980f6d760ba11751be657f261", + "note": "Regenerated on 2026-10-16 with the original commands at source_commit after the first run's artifacts were deleted. 1614 of the first run's 1635 files are byte-identical; the others are metadata (report.json, the thresholds file that embeds their hashes) and the video.mp4 / audio.wav containers, which differ even between two runs of one job.", + "files": { + "fp32-reference/report.json": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "fp32-reference/gradients.json": "1e933436ca157876aa7dd5b6f1f4b17d1a51329ad52ac9c20107033be16292e2", + "fp32-reference/tensors.pt": "3944d33e8331efcd6b6ef3a0aa7a7d5528bde3cb7a51cf3bbe27d47bbf5264c6", + "fp32-reference/gradients/0637.pt": "44f5bbaed42fff5acf99e8c204882a0c574de0b29d65fef631383e2dd5098d74", + "fp32-reference/gradients/0636.pt": "254731e038755b142a49f7751856177c062266a3a6c3b7bd59efd4187f0f1158", + "fp32-reference/gradients/0635.pt": "a48608cbefc9a5aa90a1981238b8aacef860f7225f0ed6b6e0b217e2fa89effa", + "fp32-reference/gradients/0634.pt": "a7909e90edad95fc8c1721f7e8dd5dff6b326bfdd4c525eb93e373b4dc1efbc7", + "fp32-reference/gradients/0631.pt": "0d75a0b1311b45f7d2d74622a61bfcbcdbd59b32f455fb16a45cb5d9a444f2aa", + "fp32-reference/gradients/0633.pt": "665c07fb50f760eaf86cb0de99a9b0550fff06a4f0133c34991ad08d99d7fb4d", + "fp32-reference/gradients/0632.pt": "54ca30b8ff8c43b8b41e49e145d526c8b9a678ac2d80b4cb645293fd4f4e8c2b", + "fp32-reference/gradients/0628.pt": "e853596f30bf518fccd1508de3682c15b89fe823f7a2916ede958b54efd6c2e8", + "fp32-reference/gradients/0627.pt": "90f3f9cb0982f472ccc91fd6e696c42bb204c1faef5de058bdf01ebdfdff2976", + "fp32-reference/gradients/0626.pt": "bc42369ec3de7ea1166281989736b960ee1bd92dfc34d28b3e5b8bc1e9c22570", + "fp32-reference/gradients/0625.pt": "10beec67d5f8a74e03451f12e1ce7867c3b9f8740287d2ee6a64962bb290570d", + "fp32-reference/gradients/0624.pt": "decd792208d9969a670b04bf357c0bf7a08061a8d44fbc401a378d0f2665be49", + "fp32-reference/gradients/0623.pt": "b70e894832b85566ff63c954bc817f2d22ec8edb0a1591f89fbd9986c194ee7b", + "fp32-reference/gradients/0622.pt": "b74e7fe0e27483e4bc02030c72f9bd8d9cf13fe3e9b36981d5e7c766520a2f72", + "fp32-reference/gradients/0621.pt": "3a21a0af227127a43604e0dc12ec4019d34ce597481ea019e85f5d2382c1bb2c", + "fp32-reference/gradients/0620.pt": "31679f9ddaebeecc04ed14f0566dabb117ee185e8665500232a691e95decb1e9", + "fp32-reference/gradients/0619.pt": "e90a3f84a10a51357abf17e880cf5f26ec4875423cfab369acc9837d5950a24c", + "fp32-reference/gradients/0630.pt": "00c5f09323896833fdb4e8f7e472ccc92a87e2c8d77ba635f5e3e9655fe1f67f", + "fp32-reference/gradients/0629.pt": "05658b8175118bdd51cdf1703f585e62381593314a06a9bcf296ce698dbf593a", + "fp32-reference/gradients/0616.pt": "dc2d855b1c2a8c0689a31f26dbcc49b52f59342d8abbfd02c53555653c5a7a99", + "fp32-reference/gradients/0615.pt": "273f5f9e8f5f285402b951ceaf434d44ba8fa594b32e57b6af586880c33b929d", + "fp32-reference/gradients/0614.pt": "8efca892a52bdfffa7abc41f2bdee17397e73e0c4bcdc28be888dd0959579c28", + "fp32-reference/gradients/0613.pt": "37e41946f8b7dc1de33eed0275f37562249fe0d55ce037e13f1c35a5683d211c", + "fp32-reference/gradients/0612.pt": "3dc925c6fdba3cd5033a637ff8ca873f9ae655eb0e58ab1360a864c22c122995", + "fp32-reference/gradients/0611.pt": "fb09ec9e84b92033dc0d4644f6ec72e0868830ecc362670782f198f3795dfc55", + "fp32-reference/gradients/0610.pt": "7529a771ef1c37b3415ca8de67e454c88c83952047457f6334b0d8ea0124d3ed", + "fp32-reference/gradients/0609.pt": "22e05312b01f606c9481481686665ec98bd68dd23bb522098769726c0081c0a3", + "fp32-reference/gradients/0608.pt": "180abc034fbed4b13c7558c20e5ba3d2a5ccca4cf1647e3e0d2addc25a01accb", + "fp32-reference/gradients/0607.pt": "1a34a0d1bb63396ee46373f80034e76cd06dcb7c56582f3eea15715eeaa681ca", + "fp32-reference/gradients/0618.pt": "bf8e30cfd617e225b5aa4beaa44a44be00e89e3510a0f46d3ed946b0817fb9ae", + "fp32-reference/gradients/0617.pt": "b2814b4164c79a314e50b652d97c7dde7bd18033a9e7af86ad3a6ce1be6bc839", + "fp32-reference/gradients/0604.pt": "33b5102eceb55ce3ef632e08e5f0d258d060833e9d34ab7aced4e1a4513f804f", + "fp32-reference/gradients/0603.pt": "bf302df2dc6a5fcdd4c94587eafdf84bec36dd716d902a0c35d01449a474ca45", + "fp32-reference/gradients/0602.pt": "3fe3279733a24f54602fc0b76e0413dbc967cbcf4a6890f67b3df16bc27ef631", + "fp32-reference/gradients/0601.pt": "a720456c6e189915151fe83a96c991c4dfec8327c83ff97c9f52627e54077273", + "fp32-reference/gradients/0600.pt": "baf9bd4336b597dcf9775959c63be60eb094687945e88a5702ede50f9c226a42", + "fp32-reference/gradients/0599.pt": "f9536b7ade5e21e340c122636b7761964d863e5abb25c97bf26ca26b37221087", + "fp32-reference/gradients/0598.pt": "d14af50b4470ca9245f1a2998949e09531af4d06f0d56b71b65fcff925ea3a93", + "fp32-reference/gradients/0597.pt": "bb9016a2949975d81fe4cbe92f4b3600078c360ad170715054596fc9876fbf06", + "fp32-reference/gradients/0596.pt": "33e3f6636ce02508301b8891be3551908e82da23db45b48b2ac461deb51b0abc", + "fp32-reference/gradients/0595.pt": "197e55d14ca2f0b0a59c16a8532add4336471be3efdeb16d67608cf512fed9c3", + "fp32-reference/gradients/0606.pt": "ca5e63333d5e4c48af48fb4c2aefa3420701d67b9d7c5cfde98aa306fc90faf7", + "fp32-reference/gradients/0605.pt": "646f515364bfc1ab4372208d51d405066c62e541a86f8e6021d8d07f30a2fe56", + "fp32-reference/gradients/0592.pt": "41f2515a72db0c0cbfc31691bdcfc2f683d67ab5877618c83480afa36a40fcfe", + "fp32-reference/gradients/0591.pt": "f453e824a82d73817f087d9a18dc6702d18075212b2280ede5268c49fe91c0df", + "fp32-reference/gradients/0590.pt": "2535ff23d17149d142ad967251ec4a659a281f0624488cf514f9ed0e658c98ce", + "fp32-reference/gradients/0589.pt": "8a3d3f193b56a3ede7ea543ea59a890f923a4fd91927667a02672744c249ebdf", + "fp32-reference/gradients/0588.pt": "c780e2182cc311a799ce3dd6399fff3a03fc1a4e3efa807ce22c25a5ba1508d6", + "fp32-reference/gradients/0587.pt": "562fc28ed2dce596e281aec9ed03f5a0757cbd8ddba265940bfb723866113bed", + "fp32-reference/gradients/0586.pt": "921043d6bf593d325aa59acf2d54e6e1a68c3effdf9b129293945f3477366794", + "fp32-reference/gradients/0585.pt": "c972070823d14a2f3a93b508204d9b69dd34b974a4bc408b5bc704de971b0553", + "fp32-reference/gradients/0584.pt": "d583d795b1478b779ade820ad6d5d578afd3c57d225237a7022926ae90ee0b53", + "fp32-reference/gradients/0583.pt": "6193692f06308374c4b957e22eac7903d1b3ed2ef8ea3bc9ede70db7440b23d5", + "fp32-reference/gradients/0594.pt": "74e34ec8619044dda13074a19b9142e24154a8172b8f0b942f005f29f2c7be0d", + "fp32-reference/gradients/0593.pt": "bae9313909674635cf9f33352f58f1b67bd568aca424ae9b25e6c5db2aecf2a9", + "fp32-reference/gradients/0580.pt": "c04fe8669c0ecac74a1b733fe1fd7debb0ee0d046528ad533913b5e3089963f0", + "fp32-reference/gradients/0579.pt": "fa351a8a3fe71204357f9e802709beba5fd9efffe0d5e2b849faa8da300e99b1", + "fp32-reference/gradients/0578.pt": "1f40e8332afa6a18a9ca57add6205fbf08e8d1bb7c6eeedcf62b3f89e77a4dcb", + "fp32-reference/gradients/0577.pt": "aea66f63e383a5d071ce9a20061b17fbeb9807604529431f6522f9dedbf07d65", + "fp32-reference/gradients/0576.pt": "3075eb7cd56feab500821f11a8097809949095c2a31e3941d189a0375507731e", + "fp32-reference/gradients/0575.pt": "e47d9e80455259ebe88c8991d38b07b0a7d8617625a62a536874e42dd9edb5e2", + "fp32-reference/gradients/0574.pt": "b8e196fff622d71b28ff84c19e70c78bea5b803c4eb43c7f5e2f1dce6a6278e7", + "fp32-reference/gradients/0573.pt": "dac0a85d7bc8ae9b0c08ad76374ed7881285b1d552552ddf92bc426127a02d3b", + "fp32-reference/gradients/0572.pt": "6aa08bc98cbcca883cb53fa12f80f5f0bdc5be204d33a46734697709c40d288c", + "fp32-reference/gradients/0571.pt": "88eef77acf516a8cc66a709b738746bb732f33ee9d20e5a98035ce1b21b8ee62", + "fp32-reference/gradients/0582.pt": "f0a6f8cf798c92ce501a2ecb43727bbc0cb656dcf67af11fc0bc17f96518e576", + "fp32-reference/gradients/0581.pt": "c9f1cc11bbd50411948b1785c09c7a1b0b2bc0f43dcc39ae0f6a968f80d1a0b2", + "fp32-reference/gradients/0568.pt": "25eba65a0b9ab0c565008aafebc65f3e0ad5ca8e7d73a39c504f88ef92d30bee", + "fp32-reference/gradients/0567.pt": "c0a0a485d1e25c2c275bc4614b5b1251ad04bc07a0e1fc682fcce8a0a33ec056", + "fp32-reference/gradients/0566.pt": "12c60adc2b6d12959d1d7d6bc6e959ccdf47a31d5fb657203736b2721474b382", + "fp32-reference/gradients/0565.pt": "abe6e3af6bccd98f5418ff0994a99f38f55b71b849a6a6bd48300162a092cc16", + "fp32-reference/gradients/0564.pt": "f6c11d204ae2f6a1f8cf53b3d59f0cf8618ec0f518650784ac12367e9aaf785b", + "fp32-reference/gradients/0563.pt": "13a275ed3a508c43e5e8116a75c6db44ea27c26fa3f68bc8b65bf1e43f389659", + "fp32-reference/gradients/0562.pt": "9b3839bbe302539a908f053a0b96662f657313a74bea71f56460411f32a1c1bb", + "fp32-reference/gradients/0561.pt": "c1614ca2877742c4bc056e33e1b7e47e7f525ece5ec2f9f045875ef2688ef2a2", + "fp32-reference/gradients/0560.pt": "bd1f4e9f86449803a63fde68633487f426d2503208fa6a2a74354440138254c8", + "fp32-reference/gradients/0559.pt": "069e800a99f054f11ffb3709c942e76049e7b05314cf2a1cf370ac62c17ae99f", + "fp32-reference/gradients/0570.pt": "ae867dd8f4e994bc4275508ae2c1caf8f1f8a35c6485a12b02121257471d3586", + "fp32-reference/gradients/0569.pt": "efd950800dd04a89d271131c81a0715c53b2fc98240e84948eb7c996ac14a952", + "fp32-reference/gradients/0556.pt": "f5096f499f47380174c2b551409907df82f929c2225cb37c1466344b9ecb7e3a", + "fp32-reference/gradients/0555.pt": "54d1f6946665e11ae7c8a9f6a8f24c34f53da005f58bbe1a49e83325b2b473de", + "fp32-reference/gradients/0554.pt": "e266c4857acb712b3f34394d24982838318358afcf3422e6638270677fa7b3c7", + "fp32-reference/gradients/0553.pt": "de2138d1665e1e5fb59a830d82602010e22f9ecc70958bc51028e330d3dfa92f", + "fp32-reference/gradients/0552.pt": "0d478c3f601c1d04769663777ec496a96d095e936006d5995e348b1fd74993fc", + "fp32-reference/gradients/0551.pt": "cf98d824f0f80bf32329cb91290685c5246a590a477477157942a3ba9f8444dc", + "fp32-reference/gradients/0550.pt": "ca7a07a9ad353f2b5f334333c661501c782f35c24a7bb2ee1bbf24b2849740f2", + "fp32-reference/gradients/0549.pt": "942b6b6d0fd0e9c81d84705bf7aa9c9da6a63f269f2a0fdeb466a77742b7b6df", + "fp32-reference/gradients/0548.pt": "96f2fcddc5d2ce4bf5199432f70f08dbeab0fc5f8910bd0667fa7a6e1cf1d841", + "fp32-reference/gradients/0547.pt": "1250c5a774a0837114d2927fe25f9abcbd29f2541a06cb6fc711862a50f86c06", + "fp32-reference/gradients/0558.pt": "98d6c6c7b9cf59fc506ee12a825da8d9818091c3e2ac3dbc2571372b77223ada", + "fp32-reference/gradients/0557.pt": "b10177d042eded9e6ab9905cdb666be6741a462bbc68983f98bc9482d2a580f7", + "fp32-reference/gradients/0544.pt": "50edc8355df398821f71866d565d5315410bb7e315c38793a2bba56a4777362d", + "fp32-reference/gradients/0543.pt": "8fa3752eff89183e99bb8e2cad6580752007ec066a50c0befdae0cf56edd80a5", + "fp32-reference/gradients/0542.pt": "2f3305e9de3502d8c4eb39f6ec7edd88d095bc89a3c7564e653e11aac36478bb", + "fp32-reference/gradients/0541.pt": "fbd3c1d13b2b3c77c669423deb4b6bab16f8533b83ab21157fa8e7faf4cbf23c", + "fp32-reference/gradients/0540.pt": "c3345fc3b40045674a6056adb5a8b950c3d68809169f4b9f869b107c2be20d11", + "fp32-reference/gradients/0539.pt": "6a2034e9f9fc7f97545623164c6294dc5d8ad144dd47496757e9fa2c9a4b1754", + "fp32-reference/gradients/0538.pt": "9b69cd5698548300b19f15a6f47bf2da4c67cfe7fdcb1c8cd98ffff7237e4311", + "fp32-reference/gradients/0537.pt": "c501a28230fff172649e2582d00bd0045245be2579cba72a10c1bafc8ee71a60", + "fp32-reference/gradients/0536.pt": "e608e623f2f61c554f61f4ea3dba5d3a19652190609623150e30a0555d6b6540", + "fp32-reference/gradients/0535.pt": "f46efdb2f5aefd044af14259bc8aabb2401f805a7ee0040880133c128d260770", + "fp32-reference/gradients/0546.pt": "85c068aacfed9e8192ff677e7737d061da7d856223270bcc1c03f289a1304585", + "fp32-reference/gradients/0545.pt": "85f48b3f6304874c95a9809f6985ba2945a5a6030347c6dae41372001a1c1bad", + "fp32-reference/gradients/0532.pt": "943695429503883526618ca2bf73208139d79471fa5bb741f91ee921723af1a5", + "fp32-reference/gradients/0531.pt": "77a97794ff01fa1888a3d06512e4103da118d821fb255d81876455a0ea058fb7", + "fp32-reference/gradients/0530.pt": "28a5f17c058e33274d209ca9ad4ea9b0d9156804a99dc3efe19641800217f54c", + "fp32-reference/gradients/0529.pt": "90a288c366c9cdf0aa20c4ac71820ecf9c97ddc5ea2a3a73cba8bf61786eb85b", + "fp32-reference/gradients/0528.pt": "e90a01fb080473bebb60405512d74810ea034960d9e48c14a1b2375ccbeb5011", + "fp32-reference/gradients/0527.pt": "70d05a36b918f5413b63b8a09034a64d7b7c6fb223f6aaa6abe171724ae5ee96", + "fp32-reference/gradients/0526.pt": "ce091cef2919bcc951a8ef1cfd8128fdd2c46f628dd4dd1aff13bb2141e5bf2f", + "fp32-reference/gradients/0525.pt": "1207b3d5a9a08f6fd525365a405529c67a9f86f959e0c361881310e38037be2a", + "fp32-reference/gradients/0524.pt": "1c5058f32db55c0615d936c8ff4ee8b2da5e745fac42e88dcd01e8f77a0358b2", + "fp32-reference/gradients/0523.pt": "1796a50450a530a31f47ed133de263736b0034c0c8ced0981ed572fd882f6f58", + "fp32-reference/gradients/0534.pt": "4d7bf480f2b928a86e909594ae851461d42bdef139e1cf9a62cbdad51dad8ca4", + "fp32-reference/gradients/0533.pt": "670ea78d38faca647db06d2c2c6d1e53a19280e4d55bc1f5f851ed38d26a5537", + "fp32-reference/gradients/0520.pt": "43cc5f2f0e18e1f9ec61a750c87bd9db3cf097f453bfe6f45fa9c67a65edb647", + "fp32-reference/gradients/0519.pt": "2e72ba1a0ea33725378e093d98bf853cdc68886ce67cfe0d03169709b83e7197", + "fp32-reference/gradients/0518.pt": "da36f02ee00005f0702d5a0a01687bc4bfe87ee617ed4e7ede9ebfe1370cfa34", + "fp32-reference/gradients/0517.pt": "7ec01591947add93c475af44b6e1d8728e0d0fea292f998b2e92001da5a8ff22", + "fp32-reference/gradients/0516.pt": "8b36ed32494a658f7d6037e57bb90f5b0f2c44d10c86c63d3188cbe2701970e4", + "fp32-reference/gradients/0515.pt": "73794c97a169cec101c036e22b867b1f80c665eaae74454119b08a6cabf228cb", + "fp32-reference/gradients/0514.pt": "21ac37ba854658a299d18635b581b1a731fe7a873090b6bf2527609744c2f706", + "fp32-reference/gradients/0513.pt": "c45611708b5dcfd3f37fc561e29706b20552566dfe07ec802639a9fb53dc76cc", + "fp32-reference/gradients/0512.pt": "f17e67901af8cee77e500cd7ccc816658254d0e9f0e7f97ac966e37eec414066", + "fp32-reference/gradients/0511.pt": "72ec8e7ff122c89b9ad5f081bb38780a05798baf5db199ec65b0bb9531861a03", + "fp32-reference/gradients/0522.pt": "006b34aa59e38dd230e0deb54437a4d865f21dd5973904e65180fce39e10d9ac", + "fp32-reference/gradients/0521.pt": "6568995b15407cbd88fb3acf9d8d53b66a73ece10977a45a130a9e63f61a8fe1", + "fp32-reference/gradients/0508.pt": "b06e0d4f4ec0c738ffa9aec689ab50ccfb9f1f5b21347ed722779a9a9270c6c8", + "fp32-reference/gradients/0507.pt": "543df7123d3a2a58123ab577238b08f4ae9a996550ac427f9b4e67b0c410beae", + "fp32-reference/gradients/0506.pt": "c1281ee12c79c65278789b3f3088a3878cd180bee4b356eee859f1e15f521662", + "fp32-reference/gradients/0505.pt": "28d4541906dca31a95663cd8871ad0d1bd8f2f9a9b0ec4fec2fbd131e929678c", + "fp32-reference/gradients/0504.pt": "3e0a93c47d2e1aeecf7b88523df493ed221452dca89ad7a7303350ac7e97b49b", + "fp32-reference/gradients/0503.pt": "6e5574a0ac9f3ebf84943be5e0ea709dfff1d14736162babd9dc6822a64cf236", + "fp32-reference/gradients/0502.pt": "2a58c4eccb5676c4a62438925c86a5b4cddb1fcac574428979dbbe0ccc7e6e49", + "fp32-reference/gradients/0501.pt": "13d5eac783b4d9fd547ac31888e1a81187e810d13de9d4edb7f3a74a9579c490", + "fp32-reference/gradients/0500.pt": "9db2150d5ffa5de04cec3cea3c918955b56ebcd2abdccf012553b47b1e09c52a", + "fp32-reference/gradients/0499.pt": "a6ce0ef449017393dd6d0f193f43a390c1694d50cfa234c5ba87575db6c255d0", + "fp32-reference/gradients/0510.pt": "dc6b7b5902b4fdd6cc8f95a4664ef737148c4f658a37e672f599aa49aa878800", + "fp32-reference/gradients/0509.pt": "f06a43d2844738a53cc4465b536eda78a51363c42c4eb301191451ddb72649bd", + "fp32-reference/gradients/0496.pt": "40e8e219f53a6fd32cd13f25f016b5fee910167ea3693d277088df2e30aae283", + "fp32-reference/gradients/0495.pt": "fa09566b9cc7aa8282c5c633a6e020818489d151dc860771a7aeba5130569182", + "fp32-reference/gradients/0494.pt": "a8e1bacaf41297dd7a02e4c0694683410db0cc83f74289f48dc4ef54d3881c20", + "fp32-reference/gradients/0493.pt": "38826546a90bfe5288484fbbc7e16fbfcdb6f99b73afda65f0ae2319b1cf62e8", + "fp32-reference/gradients/0492.pt": "206d571b87c88493d5e302d424c1e21c9b4f8fefb1c5d5b676deafa9dab5c943", + "fp32-reference/gradients/0491.pt": "2ff197a7f8d707fece3e2471469cbbe1e82936f3ee36822c6f95a87ab2cdd726", + "fp32-reference/gradients/0490.pt": "0390cd25a8adf59249cda2d30c3d06aad0d7bade0eb0f76ae049481fd8d42532", + "fp32-reference/gradients/0489.pt": "f4ea4a3182afcff0da048f60ecd1bca38def8ed141ae3e0f38ce36c63c95aa91", + "fp32-reference/gradients/0488.pt": "400662f9763b4ae3850e2d38ab235b2ed26041215462c2c2d8908312465e7aa0", + "fp32-reference/gradients/0487.pt": "30a3ed098eb7eb5f14732fd40665cf8688b246cd8c0f473b6e714525b2c15fdc", + "fp32-reference/gradients/0498.pt": "d825eee41db8392a5a5732534a5ad28a235d0de217b3f3e53b5a752d2cfd2a92", + "fp32-reference/gradients/0497.pt": "704279a36d508ded75e63d94e52be769c79853cc9dd890a72b644962413ccede", + "fp32-reference/gradients/0484.pt": "9e0cb0aa6de8230b78bf2e6c3ce1e545b697ae5953c5fd40497f5cd291f10673", + "fp32-reference/gradients/0483.pt": "e661fe7eb319d398ccf230f331f638c6ec766e1a9da286450ab52cdc0a320658", + "fp32-reference/gradients/0482.pt": "dcb47d66beb78c80a357abdbc490ab890fdf10a1c3ba5585a8f07637c42ab6d3", + "fp32-reference/gradients/0481.pt": "326ecce695eb2154fc11dcb8183c6d6d3f2c539492c7dda9cface28f4a830f10", + "fp32-reference/gradients/0480.pt": "6a1bc745f3205e625de2982a281dbfa8fe9e64a0720dda995b39a4f79159ea05", + "fp32-reference/gradients/0479.pt": "c261a24a15163c2363e7a3a22b8a25dbe3612ac588c1ad8881a9a6f80365a323", + "fp32-reference/gradients/0478.pt": "14b50f0bb5050805c84ad92dc420254d1ec031c4e8574f37f4a40e4f80f76603", + "fp32-reference/gradients/0477.pt": "4c3a42f151a26de43655da95b7987019eff5a4f041e2da6dcc67fac8d3fae685", + "fp32-reference/gradients/0476.pt": "713271643a14adf133f189ebbb69ba97aea5761a975572a8be93d1dd8dc4d9ba", + "fp32-reference/gradients/0475.pt": "5a399f516b0e67d766153961a6f4031f82d1ad223bc169a1d851c7e5321c5f6b", + "fp32-reference/gradients/0486.pt": "46fdef3e713423cb93df30ee7285674bc26d14e5e7060fb3407987a763be2583", + "fp32-reference/gradients/0485.pt": "d8975b3bd49a354ed8a856863bb38cce8090082bca012f7bd6238958567f2b1d", + "fp32-reference/gradients/0472.pt": "cbbb64c0c541c893a28a09427a0d1ebfc3892f4d48dd648fa2ed48d16b39dbe2", + "fp32-reference/gradients/0471.pt": "322be960b040310dc6037828a1ab58439059daf66fe1a36a1b5436e710624cd0", + "fp32-reference/gradients/0470.pt": "ef81d5a07a2ba847257b5cbef97233b346cdce42417a4df78748802ad0251fa8", + "fp32-reference/gradients/0469.pt": "59a1240c5f65c6f87a658123d7c25db803a29a30839503649fa5f75bbb7382c8", + "fp32-reference/gradients/0468.pt": "957a78f6d9061c3f7f4a0e2c8dce6cf958425c63178dec09a9789fd20944211f", + "fp32-reference/gradients/0467.pt": "4aee1927a6078a296a97226231b13dabc98d6e79ed5584f462dbe60e1187e4db", + "fp32-reference/gradients/0466.pt": "8dd58690d2b120c1f9a6519ec7f030351206f11abc01b49e1c513faf3c9d7a22", + "fp32-reference/gradients/0465.pt": "f21a5c5a9cc17a30f1d095e64d9e2d0e6442c80ab3f4c6fe185fe4b129ac8019", + "fp32-reference/gradients/0464.pt": "f425133d2df900b0fdb4c0308fcb0bb624fd8b4e5d82c8bfd0e6e2a85cf53b6a", + "fp32-reference/gradients/0463.pt": "3e2b380eeeaec92c66c2042d35c5ac491dc0870025eb0f578d7070db0324a80d", + "fp32-reference/gradients/0474.pt": "930be3a0d6fa35fe355628d8a3e203604c5174baf0f56e91de547cd36ebb9595", + "fp32-reference/gradients/0473.pt": "5e9f2affab05ae19018a73d6317635b76bca52196226d4b36f50081a3acbe1e4", + "fp32-reference/gradients/0460.pt": "c6a71ba438f79c7e7cc15f581aceebab346bf2b3ea3aa97bbd28cff471c28205", + "fp32-reference/gradients/0459.pt": "35267c9a1f22e7ea83e08a7e1b47760de3a119f173a941ea58b5e03d5fa2564e", + "fp32-reference/gradients/0458.pt": "7b3cc3bfc0c6c03c9694830532949aba765d6c8eac45843532328973f7a9e878", + "fp32-reference/gradients/0457.pt": "c6ff9b6a3ad24ba2ed0dcbf80a8c42bbe3d23e96d2f64daef69f8cfc07b6ffe9", + "fp32-reference/gradients/0456.pt": "603609358dbbb1edf4d134e4edf8f8e2eafb7980a6d9ec4a1847e3c92c1abb6c", + "fp32-reference/gradients/0455.pt": "8c8805ecad1a9281fb9aff6f25c8239ac38bd35fc10e7afa5055d607dd18bcf2", + "fp32-reference/gradients/0454.pt": "984b2d3a175b67d0b2b3ee91dc6a72c0e2074faa0958386c2f9dfab2ed65bf7f", + "fp32-reference/gradients/0453.pt": "deab52b241bfbe5a72a47bfd8115fa91df25f0d3732cb8d8de8db17afdd42345", + "fp32-reference/gradients/0452.pt": "f2c89000d816ae526c707d1806862160a80f69cc9932a131f28b07453c10a302", + "fp32-reference/gradients/0451.pt": "dd33693e1fc442b2779b7192aa95e6607b396fd8eb34133fb1259dc0adc4bdfa", + "fp32-reference/gradients/0462.pt": "98cfb654e544a673ca011163aed6f1e9d9f9c0e551ba96147ddd8729470f3d06", + "fp32-reference/gradients/0461.pt": "2c03f7c8e551207b401d849529c1861edad03f1c18a7087f9b0dd80116348054", + "fp32-reference/gradients/0448.pt": "81cd65ae67a59af052232c49299b08cd30d669c216602207bab48f49182d8828", + "fp32-reference/gradients/0447.pt": "af8e51ba5a0d5e4bbecb1dc7b107a77a14a1ac487da2ee16185573f7be5e3455", + "fp32-reference/gradients/0446.pt": "677a8910044036fce0c62f5bcb84b7a372da4c3d7181250c0ddb0b8a63b13e09", + "fp32-reference/gradients/0445.pt": "487db7dcc022a52162cd98cf26dd8147355c608a21d13c6af16155bb82318070", + "fp32-reference/gradients/0444.pt": "df3201fb9a082cdb9d93c8066e74a5318681a60f5f924ab7e064e61030ee8974", + "fp32-reference/gradients/0443.pt": "97af11e81de3f569e077cb031cd98f894e6c5d3d4588867302e776def0f5d1c2", + "fp32-reference/gradients/0442.pt": "127a82048329e8f6784e0774a03c1f19397b9a69feaeab3e6b7b3fc244c84435", + "fp32-reference/gradients/0441.pt": "6a76887e80fc3135b73a551ce908e64e23e78b67dfaa5f4f4988b7a8a4d3ba83", + "fp32-reference/gradients/0440.pt": "7a6edef7f1ee390b8c52a06161c7fe053acc7e2f38c178f28203643715627021", + "fp32-reference/gradients/0439.pt": "993197a391f39470dd4d0f5868ac165767b54b7abec0465e77ef4a2e62a7f6f8", + "fp32-reference/gradients/0450.pt": "973638a3a1b1673a22a62cbdabec4617f0aee24a103978c2cd6aba5ee672047d", + "fp32-reference/gradients/0449.pt": "2d54af03139f375dede9804624777a461dcbb4e65be05c0e4e34fa53c0c475ad", + "fp32-reference/gradients/0436.pt": "8ccb65336ec154b5ceec97e7c7e77771959094c970304cbc9e669eedafbc984b", + "fp32-reference/gradients/0435.pt": "a06840af610d2d695f155cb579496efc71a48eec470900a9367e4b4d2d6d0589", + "fp32-reference/gradients/0434.pt": "498768b33a2ddfe59e2d05643197b65d88532805e5ff2424bb13ea34183b335c", + "fp32-reference/gradients/0433.pt": "402b8c894be2381a1b3effb145f8d7c2b315fe7d861a2b554171f1dd5c5932f5", + "fp32-reference/gradients/0432.pt": "281948db17d5cba6b3147be99d40d83a61a0eb584afb64a5a3559931cbb1f1d0", + "fp32-reference/gradients/0431.pt": "45f30f12ed0ed990450967f11f708e01131c7010ea8a911cc30dfceac68c54ed", + "fp32-reference/gradients/0430.pt": "aba14a507adb1457d5bf2f1bd78252a402195b9b79eca8daa88b5ed5b8af64bf", + "fp32-reference/gradients/0429.pt": "c9da0de7771e8ad19613ce08bb8de95adc5b4c9577ba57946f13d7c1ff3e8832", + "fp32-reference/gradients/0428.pt": "4979d33eb8e70d2c8951a180cab912154ef015b4d7808aeddf061043ee5c4331", + "fp32-reference/gradients/0427.pt": "1e4030f3ca19efea712ac4f1939f3522fb3dc630cc79766b86f0fbd9ef543b23", + "fp32-reference/gradients/0438.pt": "81fd57d56eb2af15ce2d8e45c56250cbc5e9b2ece6227bf71ecb12c73277b7be", + "fp32-reference/gradients/0437.pt": "eb3313558184ac8a8dc339736ff5c5658a2ed3d865d43cd585072c66fa8ff2fa", + "fp32-reference/gradients/0424.pt": "0cc8f805a9744b814159f4e6b272e70a1b92662f85d604db8d46acca1128c38a", + "fp32-reference/gradients/0423.pt": "5530b03cf876b7fb8a5450f1f0709eb6e91ab4b5e23ac76d5fed093714375eaa", + "fp32-reference/gradients/0422.pt": "690554a5d80a0b98ba6378f3c6ed6d41357e4899bf9bc88cebb32b8bae9abfa7", + "fp32-reference/gradients/0421.pt": "9879ec70748ad0f7c274d709b1d8a24335398aaef2db6e09a2cd7d8f912ae3d1", + "fp32-reference/gradients/0420.pt": "c2e67c2a62f16d8f19068d3b9bad00bf1dd9ca62cd2ae6e54232f6dfb71e3b66", + "fp32-reference/gradients/0419.pt": "31be6eff12f5ad14cdcbb43dabe3c493911a953f32dc1c60f06306fa6310e2c7", + "fp32-reference/gradients/0418.pt": "ed3d0feb5ed1eb58bdb4e516aee3ee1caa866d0bf9d9c04b37f25e464317e108", + "fp32-reference/gradients/0417.pt": "0740fdd7326f0ea2c1bb4ea7140359b1f329341d3f20db1bc1f8c3d0b5d927d8", + "fp32-reference/gradients/0416.pt": "556ac4f1e52dcdb880c4288242a7d1761021b7cfe7d3f9c3c7671c49d0ab1a65", + "fp32-reference/gradients/0415.pt": "46a96495be56bb52abdfd7267938fd55044658cbb2894767f4f9ca014f698fb3", + "fp32-reference/gradients/0426.pt": "e070b09cf801ed3a8a65e5b2ec63ecc1642b2389bbf8f7486a9e792af4562c3e", + "fp32-reference/gradients/0425.pt": "55c1d8439a2cd53dc49bc35c43f76ca4a3f76a2d2054cd85e025bea41cd3991b", + "fp32-reference/gradients/0412.pt": "61dcac766920bc765df7fc69fd30e77fa8a07f0874bfa9ed4889164a60e37cd3", + "fp32-reference/gradients/0411.pt": "99c45965ccd1495cc6e7821ff2030354791c435be19706609dffc55f52197429", + "fp32-reference/gradients/0410.pt": "81b4d92fbcef872dc767575bd04fc04d0a99e7d73c9bd7b534753277bb4acf20", + "fp32-reference/gradients/0409.pt": "8304957c813471f601c71ea518fe5a49249638d0413a3f14a20e20df00a17ae1", + "fp32-reference/gradients/0408.pt": "ca69b37d38be66fafa21ef99922f3e6ce737701ff17c862db68779f15d9eda15", + "fp32-reference/gradients/0407.pt": "b927fc05062ecbae856da3b7bb845d5f46988a2b11763d02720a60abffa2de35", + "fp32-reference/gradients/0406.pt": "b2188a96cd251ab54b15c7f05be6168823cdec771bc8a5563cc2ed0ee713c57c", + "fp32-reference/gradients/0405.pt": "cd009b42b8316e84b06670df68c3f1056058ce1b3c6ea5eebcf2e39ab7fdc7a4", + "fp32-reference/gradients/0404.pt": "0c092e2448b5f7c815bbd67c06052019551ce387d77000e192841c8aab18878f", + "fp32-reference/gradients/0403.pt": "1c10e4c4aee8e5232206d8b8268d2e98f768ea5075122e50d99f7273d13c282c", + "fp32-reference/gradients/0414.pt": "a1a5c98c122831855aaadffe0f55a746acb3562b5cbacbe983640a880edf74ab", + "fp32-reference/gradients/0413.pt": "d226b6558e5ef7bb333da5318dfc2c998674634c47c59333fb4812bff8df1565", + "fp32-reference/gradients/0400.pt": "babdf8d72aa4ebadf25177e8755a63e216a639ca599c868583d71158aa4b2ad1", + "fp32-reference/gradients/0399.pt": "f4c1f48734bfc6a92d095bc358b0bca49e29058ec8bb0bb312d4662e8fe4a6e0", + "fp32-reference/gradients/0398.pt": "993c0beb924ca34bd796c74ff2ace0e3b6b7b63189e6822fed76ae022e40c151", + "fp32-reference/gradients/0397.pt": "cb1b248337dd8869f80f077de96241ee6a3e0f1ad89ee70ca32797bcebd7d2b1", + "fp32-reference/gradients/0396.pt": "479869b785c2e46f81e047701031624a34ddd0daa1fce350a4c4c13f5708f046", + "fp32-reference/gradients/0395.pt": "02bf140d6ca2c7b05b627c671771d58a12525b556fea24d0e6b3cb1da8a19894", + "fp32-reference/gradients/0394.pt": "6ea8d33e05f192a3904d0df8238b4d3dc06d60885170d677504109b0b2f42f7b", + "fp32-reference/gradients/0393.pt": "497dd32db340a8602a5667f0967854fd47efd8b0eedb0d78a1b1892dae69b0c2", + "fp32-reference/gradients/0392.pt": "e241fdb37146d5df78d6a5cd2a386af505a314af8cce700d4b31b858412aa07e", + "fp32-reference/gradients/0391.pt": "4cd2bc752abfee65b34d656a719e3d41ccfaaca69b83252dbdf441b0ced575cc", + "fp32-reference/gradients/0402.pt": "b269ddd15d6b3d1c7eb93abaa14eefb096f11a07ae2e1a7b40b00f8479969364", + "fp32-reference/gradients/0401.pt": "c3f791af8aa2f3dd31b82faf4e24be613ee1af339fb3dd0b72307e21d905b3f6", + "fp32-reference/gradients/0388.pt": "055d0f42b1d00d33caafb3b36e0a191d77e664ac9fa6d1622fb73de6ed29cf8f", + "fp32-reference/gradients/0387.pt": "09eadfc79f8d76fc9fc1527980e87fb32cdabbf47820fba6a2bc9a96e417e112", + "fp32-reference/gradients/0386.pt": "bcbac108b3e79a5920c78755848acc74d3ea88ab83c7d4cfb66a3ddae5e49701", + "fp32-reference/gradients/0385.pt": "dab9514a1d728bed815212d68f49ccd086f65685636de4f6428d3a4b62aa25ca", + "fp32-reference/gradients/0384.pt": "61c28bf78e73158fa599ae5e4fca1fff10f631fcb242fbb19603dc063f42e5b6", + "fp32-reference/gradients/0383.pt": "f98993334332d6b3fb603bff01157ef4121f68e83ac4ffe866d64ce7e7f79e74", + "fp32-reference/gradients/0382.pt": "da0dfb211b031fac433ddc0f22bc5837e944ea9724e8912866b8085f4337c6f1", + "fp32-reference/gradients/0381.pt": "5c4b234b4780b7aaaa601c4cfe152bba2369655900fdb39f4989c8e8f3899c2e", + "fp32-reference/gradients/0380.pt": "0bcafb5cf2c1c7521b15dfaadd32d254675adb96e9b0bdf2c2350f60ebf1fda2", + "fp32-reference/gradients/0379.pt": "fc1adaad6c208cf4ed78501d6ca918624f9b5abe395267703a9b449c3a391d89", + "fp32-reference/gradients/0390.pt": "926fc65ea61633bc44e4d42e1e32eec564dce7b3043d617512e0fabb5f90b087", + "fp32-reference/gradients/0389.pt": "d7f9ef3c7b1aab9cf97eba892e19793aaca7b4270d2a04411eed323c608b6fd4", + "fp32-reference/gradients/0376.pt": "fca479e9c367975ec12aef6bcdcbbff639faca4441503a7c4a7388ddbd1e81cf", + "fp32-reference/gradients/0375.pt": "8fedbec498e53baadf9b7921de73e53fc3b4ef06ac7a39a7717b09511909eacf", + "fp32-reference/gradients/0374.pt": "cb14cd46f5564f1f558ca0ed65cdb0782157ba931887de4efdc973d0c4d79f94", + "fp32-reference/gradients/0373.pt": "6deadc0364719fe6f8fc6c5dcf3c308bc540ed97a46da3b17e2304eb9a197dfc", + "fp32-reference/gradients/0372.pt": "6d87e6164df3fb9e066120461671f47be986b399d2a97890fb86e1ccc250dcc4", + "fp32-reference/gradients/0371.pt": "26acf2f7ba06424123a879fc2c08251fc1afcc72fe71ce13764fb14f21a9e247", + "fp32-reference/gradients/0370.pt": "21a95af2c977ede2b79e6a14e83e70f44b5eea58fae021bb3ba189505823f6d5", + "fp32-reference/gradients/0369.pt": "5040b554bd3a9e027cdb0c3a53b609b6eb2b90636932c575a00b5f9867244a29", + "fp32-reference/gradients/0368.pt": "c5f6a9a097c9fd6aa9c978ffc1d9b5d223370be963f293e2d9ebdff9acffe1ef", + "fp32-reference/gradients/0367.pt": "141fc2a9740cca0e795932932f082b5f0a6fd2f4bd235f5ba88b3d3550703767", + "fp32-reference/gradients/0378.pt": "2a5160f53f07d9545f97166fbc156e0a998557d251201100b5c8ad04e72c3216", + "fp32-reference/gradients/0377.pt": "4b891cea1186ea29963536e51bff3f61f4b131e9902396104eab8fe0883a6301", + "fp32-reference/gradients/0364.pt": "297e433afd8df2f4e363d2deb099d38e41c13ddc4acc0c5ac4d5beae92353886", + "fp32-reference/gradients/0363.pt": "830c8fa1f764f5d3c1a89351070f1ad76e4340273e46b38fdcc068b2f1a4fe04", + "fp32-reference/gradients/0362.pt": "b71b38fbcc78f6d31641241633e001f20db6f57529c570c8071e1147169fa42f", + "fp32-reference/gradients/0361.pt": "cfc0ea845a88b63b1e22c922564d39fe64a6d7e9eb0f049905bf513245625f5d", + "fp32-reference/gradients/0360.pt": "c1b8a40a0859cdf37efd0133c492c043a5231f12f616cb7acf46601ac3af6899", + "fp32-reference/gradients/0359.pt": "6da7d1c92ea12e9a8562f985157967e36f840e07f5868d41f1d9873941d289ca", + "fp32-reference/gradients/0358.pt": "98381516e776803417d7441f825aeef0cce21a357310943c3851a2bb914790c1", + "fp32-reference/gradients/0357.pt": "d230e6f3b0f30ad25069564c7f6fe583acf3ca88840c86bac1b404c30f590caa", + "fp32-reference/gradients/0356.pt": "f142bbc7005fd0daad1cf4f2b9e5e698752eda421b75b18bee14a1f3fb0f0ae3", + "fp32-reference/gradients/0355.pt": "5622fee6dcaeae422325f614b9feb9d809d1cb0d86dc35f90735bb015baf51ef", + "fp32-reference/gradients/0366.pt": "be69db3dcc33b8749c8b21bcf9da4afcccfb68fb671884772afb623c57b2fd99", + "fp32-reference/gradients/0365.pt": "4c254dd353ebc1b09b7eab21170765e489b8d9ddd798dd99addf772285fe0cdc", + "fp32-reference/gradients/0352.pt": "d0ea48a2b67f353d9f3f8148e93069b85182cb3f306c5c7749e989f9776be585", + "fp32-reference/gradients/0351.pt": "664941612ea2d708d020617616bdecc811efc94c6baf4b3d8e8fa060a3f322f2", + "fp32-reference/gradients/0350.pt": "a46c006861b58b0959738f1a7e4af5385b9e9ff72bf8f6e1085856c1cef6adcf", + "fp32-reference/gradients/0349.pt": "01646aee7e5ec1f0c3e0b270a89dd93955c977981b0a835d88018e771f191078", + "fp32-reference/gradients/0348.pt": "dad5413edd983c8a8bfde79146d99dfe938564fa8d0029cf1279e28b48690186", + "fp32-reference/gradients/0347.pt": "8e67e8e8331059794a4da7d91441617ebaeadde4879a1acfd23d61c1fc2703f7", + "fp32-reference/gradients/0346.pt": "7c44095544782f10c97dbee924dd647e6297553676ab13543a3d832f2c32aadf", + "fp32-reference/gradients/0345.pt": "3b91220e343c52b71eac538752c6e2661b4756ecdb3a5d314b8546e5d0c7a77f", + "fp32-reference/gradients/0344.pt": "18529e38c196dac0c5fe86c3f7bbfbfb02db2875d9fd89d0a722e6c699e12573", + "fp32-reference/gradients/0343.pt": "0ebb81d928088e9a98ae13c93db5964545e382beee8cbf9c58d366c59e19f97e", + "fp32-reference/gradients/0354.pt": "7f28f9f54e3efc7e43c30f4beb199b0784eaf365e86cb343f8579f2cb31cceef", + "fp32-reference/gradients/0353.pt": "e9bf6f6ec4046508264dc20f914f629f77e9d6e3010a26c47334cf22cf22a0a7", + "fp32-reference/gradients/0340.pt": "30c1ca0449204c49165ee0ae647516732b457697e4a0c46dbbeff6bb9bfe4103", + "fp32-reference/gradients/0339.pt": "bbb21d4990a263e4bd00bc6a685d963beb4e8ead1397e6710678804a78a697a3", + "fp32-reference/gradients/0338.pt": "1ec28f042a16d01964213fe873cd8f4358fa8b84341763307d0393fd7ab7f20c", + "fp32-reference/gradients/0337.pt": "31cbfb86edcd646f8e67050b218451fa2696bb22ba4599a6a3b4e32849c57236", + "fp32-reference/gradients/0336.pt": "26e3d7b7d6ea02900d252c202288afdb36b34b580393ddfd90f877ac1d158792", + "fp32-reference/gradients/0335.pt": "ad309b703b42366a9bd5df22967d2663240da28c9056a32f7a039f164b439c32", + "fp32-reference/gradients/0334.pt": "e6140fd83b826824ebe2dce3131947772fb0d95210a110b101b54ad8d99d32d5", + "fp32-reference/gradients/0333.pt": "91061cdce387adfbadecd2b52e31be797d093ccff436cf7c337483d26df30f9b", + "fp32-reference/gradients/0332.pt": "6784885842c58905df6382ba505a76d4271d5faae11d8486203ad93db3b8e086", + "fp32-reference/gradients/0331.pt": "f87214d58ec9e88064e24e472391264823ae025a939c9c0cff9408b63d72798b", + "fp32-reference/gradients/0342.pt": "c80472d05577d563baeeb3da54f51511a4d588b29e95aac30872ee23d617265d", + "fp32-reference/gradients/0341.pt": "1b97fe915bab23a8f2acdcb9a4e14e12c1bde2b495c19cfad61a5d19dd52a0f0", + "fp32-reference/gradients/0328.pt": "a46f533054a001ff28f32b915cf85a05d8b95bd4cd13ea4205842b611b1e6f9b", + "fp32-reference/gradients/0327.pt": "9e85dcbfb51f3d492c3c93b87d5cae0bc961110b3be6c94a6281fbb19608d151", + "fp32-reference/gradients/0326.pt": "61446f3b0ab7612c75cd61a325ec49977d827985ef78c0c9d4407073599eb586", + "fp32-reference/gradients/0325.pt": "f30573b2637ee6c35c97042a913a2d48112169bdcd6b31ab5155aabe4febdd0b", + "fp32-reference/gradients/0324.pt": "9ab6259674209f759dc294e731a10a80205973b98c201e8a2cec04a63a51bfbf", + "fp32-reference/gradients/0323.pt": "e7bc8e825b3b558a0740fd26425d103647650801e3cc1103fd0d39c4cf3fbfda", + "fp32-reference/gradients/0322.pt": "cac33dfe30e43034a3d814b379cd9ec6390e3307941ad5bed24a1c0a79d0f2f7", + "fp32-reference/gradients/0321.pt": "932fa66b71a1bdf1e8ababcfd9ecae9ce9ed1bd8589791b7342140ecb4cc5bc0", + "fp32-reference/gradients/0320.pt": "94d781055abb8fa924731009a5ac4c6f23b27d69364b92a367eecc5917ac9536", + "fp32-reference/gradients/0319.pt": "6047d1b068cf0252ef47033e1111f070a19919ae4deba4e95e6486b2fdd1367d", + "fp32-reference/gradients/0330.pt": "43b519b8c6d8a59c735918b465fb461b8026356941cfa3fc0fb777f4f8439f4d", + "fp32-reference/gradients/0329.pt": "c31311bb28b3f451096ee4017f227c457f0f9f4d9af8044ee906c20af5e37516", + "fp32-reference/gradients/0316.pt": "aa535cfd1f5e50eb830a4451df49b5b4a017b4a5ffa0a9cc3bd0ce7b39c7af06", + "fp32-reference/gradients/0315.pt": "0e3a47131d9328e1f4d9eb6d202ef20decab26ff352d84b3bd8c16c232cb998e", + "fp32-reference/gradients/0314.pt": "c900945c0b6af837c551a71df7edbafed618c66d15f8ac8be9340608c3bcfa49", + "fp32-reference/gradients/0313.pt": "c98d8cc2c560a21ca866f8d4a0b3ad3b52ec9997e8329e61f18db9aa8af1de09", + "fp32-reference/gradients/0312.pt": "529b655af8096c58d58fb1b3c3aff067b97a68729ea84e9604748c4ac54a7844", + "fp32-reference/gradients/0311.pt": "07369812478dbcca4562706fa2523f4b0262f8c2b8ca5e1ea5f6b10649ca599e", + "fp32-reference/gradients/0310.pt": "7d810f9a97fb1a1a5f3f0ad9c4919b6cc05ce2439b1989a2cd5ff28a90aad19f", + "fp32-reference/gradients/0309.pt": "8aae11e69cea48219c2994364b21b5939b6d72a082393881cb50ddbb4658fc4d", + "fp32-reference/gradients/0308.pt": "89bec3cbe3662565b5a735087e2468b2284b3be9bcd6efa646dfc783803f0747", + "fp32-reference/gradients/0307.pt": "3e5769ec13b31cc8dfb80caa425e24defb8c245815f3b04bcb2a021bba2356b9", + "fp32-reference/gradients/0318.pt": "6a58c2cd5662cf5f600b467b827730026bb8e7fefaf1ed0dd59dc31ac1b50b03", + "fp32-reference/gradients/0317.pt": "8349ec6d6d454cbee6161a21fe93381a31d99f041f545c0420df84511e3ed59e", + "fp32-reference/gradients/0304.pt": "56927f97e433dc8c4a4ac52ff00747a03e7c8f4c9e3014540ff256fdaaccfc2c", + "fp32-reference/gradients/0303.pt": "7eca12e84511596c21ffb80a2b25d80cded57e5c8ad6b0de1e785a6b9501f590", + "fp32-reference/gradients/0302.pt": "86bc7606eca062317ec2bcfefa4eda0cd3e53c5f2b5d4175281b4ad55a53672e", + "fp32-reference/gradients/0301.pt": "e826563d9b100cd3e3e22b7538d4961642fbfd1381b3973664dd96e5d7e50755", + "fp32-reference/gradients/0300.pt": "f1a98109766ef78c2eecb238dec287d4ff647fb8dce0c2caad229f853aaff82d", + "fp32-reference/gradients/0299.pt": "f948e050467b4d540dc4a9339f551c317e8ac12f0cdcc7b1ddac35bb50778d71", + "fp32-reference/gradients/0298.pt": "36d0016c7ca4e22235e2e676b5ad3b82678375744e2a3da7f86a63dc39506099", + "fp32-reference/gradients/0297.pt": "2e7e293b196e88cfcdda3313ee7cb9f2586af250b2dfd64894c52a71380636e6", + "fp32-reference/gradients/0296.pt": "44fc8f38aad56d152da5931db0c48a1d1adb09cf93da010feed7e718a6542d53", + "fp32-reference/gradients/0295.pt": "d428d1ae88d44f78d842e2f86c83a55c8614dd98d9146615c5bb86c76c5839da", + "fp32-reference/gradients/0306.pt": "47977dee08ee2be3dc9b3457e67fbc668b0748879f61f320e37e3f4802a0a60f", + "fp32-reference/gradients/0305.pt": "83a717de444b9e8c67f092a1486764954c44df1c5a3169994f02491ae3e2febb", + "fp32-reference/gradients/0292.pt": "e6f26acfa525135c68d7aa2a621385489211f29234ffaa1391f71d9136ba1abd", + "fp32-reference/gradients/0291.pt": "6568c1c9332be4f04d93f7c9df5403705244be1680065951f8d3111c10a76311", + "fp32-reference/gradients/0290.pt": "ac9d5e5ee0f40b4da80cc7b250c479e98d16fd24080344313298ff263978c342", + "fp32-reference/gradients/0289.pt": "9ccf2de8ca240f9ae55c7d54b517c71c39c9eec1e6e328cb0309c72255878cd5", + "fp32-reference/gradients/0288.pt": "b44208f51754dcf8dc83e82a48b3bc9390c5fbbe8ddb8d20c05c556ddcc0fd1c", + "fp32-reference/gradients/0287.pt": "4e9ee7340584f74e2769eaf5397214455a126302c5c9c769dda55c5429a9fdf1", + "fp32-reference/gradients/0286.pt": "e958e029137c0eaf4321f4e2becd71fcc591d7e41fc56f003019ab6af5cbea44", + "fp32-reference/gradients/0285.pt": "2eb3c1cf558bd3efbf1ba135be3309ea7f4625a69ba53d8ef7b7f7d540e5ab77", + "fp32-reference/gradients/0284.pt": "c9f18d0ba332d302fe1db2ae201c3159759650f697db7583df26fde6f1450460", + "fp32-reference/gradients/0283.pt": "57d52095ba5b541d7bf1e96ab47b16d12071ac63ad48386c120ff458492f9ca4", + "fp32-reference/gradients/0294.pt": "7f54a9b76f96e3c5a15d276f057d7215db701fbaff25104b7461a6a4e534d992", + "fp32-reference/gradients/0293.pt": "445dab90ea8620c59f2ede7de5ee4a79a9da227aab31767d45ab433116e2ceaa", + "fp32-reference/gradients/0280.pt": "bb173c55f854dc97199e915da49e6c61f763decf02e3af4fd2a9229932f600ca", + "fp32-reference/gradients/0279.pt": "09918cb3e216463cfe2229411b587edf64e184db29ac4449d1db6924241a93c8", + "fp32-reference/gradients/0278.pt": "24ead138531983bdc29613177e2a12d1e7c5a99521100e444930ffd84779ee23", + "fp32-reference/gradients/0277.pt": "0791e037b3685357f850875d4eaa239127eeb3170ff72985fd509950c7b1926e", + "fp32-reference/gradients/0276.pt": "bb7689d72d84a0d2ad58677ca28a9747c33fcb1f5e50a94843aa0425f5380270", + "fp32-reference/gradients/0275.pt": "f8b425462d1239aa42a7d33410233d39ac02a9df052d3d0cc1d4aca8d4594ab1", + "fp32-reference/gradients/0274.pt": "9e3a878665c09d337a3a349a2536f4e2dfceb3dc68f44dc9be4e2863a925d55f", + "fp32-reference/gradients/0273.pt": "c3f20c91bc0b816d38ec0c46219e0776fdb1041b59b5c3df7c3708052974ed7c", + "fp32-reference/gradients/0272.pt": "f9f149bc940f7cd388b1d2303857732ca510494c392a303d2404ef598e61a408", + "fp32-reference/gradients/0271.pt": "89c8029ad82790e44abbdbe41c813e520967b96cdcebb45e98362cb419c03d4d", + "fp32-reference/gradients/0282.pt": "01461d30708557a310882c612cdaea64dcf2f4db11372f6c7c932ff110bb8236", + "fp32-reference/gradients/0281.pt": "83a5000db45c95e38d4be269495bc82fc456ad2ce7294fe6728ce65e49bcabc9", + "fp32-reference/gradients/0268.pt": "74b1618f5397c50676856eb6ec11be0dfd0934599535f89fc2a2709791247620", + "fp32-reference/gradients/0267.pt": "48aa3a0fb63758a5ab02cceb6a0b2a3d6d339b8420085091b5a05e572819ae72", + "fp32-reference/gradients/0266.pt": "a06bf21815dbd6ef3b6b7f08022ad523ea583af27ccdd82b6a58221274a57f16", + "fp32-reference/gradients/0265.pt": "87d805f43b18dac430d0d63ca05dadd55a5ca32e4852856f0310381a281dcab8", + "fp32-reference/gradients/0264.pt": "a469ad817abaeec10adeebaffb9aea07c8befb7279634d2d788b29c40821a6a3", + "fp32-reference/gradients/0263.pt": "08a25e62a6f8680731ae589381bae0429081c03cfa255cafbe06d04b4ea08675", + "fp32-reference/gradients/0262.pt": "3904581f9ef042d39c459e0c62070846962d3d5b7200920080d815af59307c94", + "fp32-reference/gradients/0261.pt": "015443a3746e4b298cea82ff383735370802948a008fa478d831ae47fde942de", + "fp32-reference/gradients/0260.pt": "62672784add9ad400444f1ce37f8cdc5c05a2c04275920a2408fd4818a7aaa99", + "fp32-reference/gradients/0259.pt": "16cfb35a77918631f4affa6bf494d7b7ef28de3e3cb3ba078112bb380117f008", + "fp32-reference/gradients/0270.pt": "a798d90c7b1e37acfc75eb3c8fb9609ec51a607cf02350ddce86fdd1f96c1c65", + "fp32-reference/gradients/0269.pt": "4042dde7e9a5f8a2137dd9f3b21adf3444cbbe0eee4094dde3532a31e2616bbf", + "fp32-reference/gradients/0256.pt": "5d9ba1422fc50d1ced3a8e7845cf09dfae1e2034f62cb1646887000a2aa2278c", + "fp32-reference/gradients/0255.pt": "0f6ed31fe4b00d96ce35c3754d1439f08f4201a23869164840ad921fc7aa523a", + "fp32-reference/gradients/0254.pt": "13bfc45f4ff6d72e810ace53e61656914ea406da30913629b7dce9e6b7fcdd22", + "fp32-reference/gradients/0253.pt": "957610f17e94a964b3ce98682830f9d5f033338888d431e0fcdf4636192478ae", + "fp32-reference/gradients/0252.pt": "cdefa2367ac2f76cb315427623e7751731247cfc463fb6961e830209a8f36565", + "fp32-reference/gradients/0251.pt": "ff6c7251f77ffa0380a5b216f474e5cb89b2d0c65e78df0fadb9b782a2a8a7b6", + "fp32-reference/gradients/0250.pt": "331a77e539eb03f25188dbbba071e4890eeb5ee0353efa2d6e04fc3a0e4791fc", + "fp32-reference/gradients/0249.pt": "58cb63a6e35cb65fbb897bd7fb45cfd59a5a1fde8be6ea089e76f7428f4937c8", + "fp32-reference/gradients/0248.pt": "6573b84ca394825d338b9b3d7b3f6e3a329ca505ebd81480ada17009071343da", + "fp32-reference/gradients/0247.pt": "53f565fc6d5264df52f5b0947e017d072cbd57f578932041e2506e6290a15820", + "fp32-reference/gradients/0258.pt": "9710d5632b54a5317ff9563a004544de71fd583e2b328293e3977d42d46a333e", + "fp32-reference/gradients/0257.pt": "805d4b08afd97ace82c902e0aaee865d417b1999cf52f445219deeafd8b70d7d", + "fp32-reference/gradients/0244.pt": "15da11791244a0b7a82e9c4b8fd4a7ca450d999b8296ff93767bb52f13c09321", + "fp32-reference/gradients/0243.pt": "f52c438fc875bf0fa379b8f89a49bf45a3871a967290445066ff196a9005eba4", + "fp32-reference/gradients/0242.pt": "68fbe79a87f5151b3d6aac46b7915788eb0f06e8df15ee7d4d98ab7a97c473fa", + "fp32-reference/gradients/0241.pt": "32018a1c61616fe55d4eb061f4fe97691dd20384c47cc91ae78c9874e9c8c540", + "fp32-reference/gradients/0240.pt": "fe17142d8523a3e82c108c252aee0ef93968eeebf2b39109651cd22b11a2d1c7", + "fp32-reference/gradients/0239.pt": "52eeb688c6bde3a35c304e8b3723e80ca13ec8bdf0fe9a650c564ddc23016b5a", + "fp32-reference/gradients/0238.pt": "90523ef79ff4d4822ec45016d7d033d848220f9a215b54a2ac376aaad2ee73d2", + "fp32-reference/gradients/0237.pt": "4c844c2ddf60ebc6711e4f4df72bcfcdb2985016ff30affa67f1226c92f028b1", + "fp32-reference/gradients/0236.pt": "3cce77166f9b8a42cbcb57bca033688f6b6bff865884449454437acb361b1859", + "fp32-reference/gradients/0235.pt": "63dd5be6524df313c429b7c64049a8a710b5e0b8e3e7c67e456e99ea110e96b9", + "fp32-reference/gradients/0246.pt": "5c74a456b051370097290c6aae8c654804ef4f5d819ff5f36f21a0758855b6ef", + "fp32-reference/gradients/0245.pt": "f34ddd4fdc299207f52de538234b5071ff28a6f0f7ffe817916af075dd1f0f12", + "fp32-reference/gradients/0232.pt": "63a832e049d3a280a339d9667adc479fa409493e08f7911244114a128cd3e854", + "fp32-reference/gradients/0231.pt": "53a28778723343eeec78750c57c7ff4d69d610ff906a4d2915a50cb3b16d9af2", + "fp32-reference/gradients/0230.pt": "2a01daadbf34e70ecc1b5cc1fa8d44f6add72ad9a7f80c0e7f6310810afba0b3", + "fp32-reference/gradients/0229.pt": "ec95345a5bb65c8b09f7bfa56a6f18f10b693b81878dae1b13bf157b27f5f9f9", + "fp32-reference/gradients/0228.pt": "88893e25ac13fbb7f61549868f5efd41e0a310fad2ef92efe3a67f6a9bbd0f50", + "fp32-reference/gradients/0227.pt": "ece7264aeb17f5cbccfebe102f64ca822a9d36d786958126da75454f61d3a8ca", + "fp32-reference/gradients/0226.pt": "737db907ea2d2bc477b0f713340f6b6a3968ee2c4c29e199bd6343d96cf74ada", + "fp32-reference/gradients/0225.pt": "780d7021878199e7385d612bc2736ef6afa30bab555c9a5bd1d3551751a0094b", + "fp32-reference/gradients/0224.pt": "f31570780d9507fe089fc08ea55913db84390b42eb433ab42e0b95f771d2fe0b", + "fp32-reference/gradients/0223.pt": "2ffdf58e282f2776a4e13b1fea9aa737dc7dc7ea0486bc14f0489a5f6c34fa81", + "fp32-reference/gradients/0234.pt": "e16320be47fef92d457a1282ee5bacbd32d8406e99840016852c942dcb9191e7", + "fp32-reference/gradients/0233.pt": "ab8ffb2e37b4fd738831b1e0d14532175235cc9563ca59e6127fbf8816e4a011", + "fp32-reference/gradients/0220.pt": "b79aaeaef0a4df18ab2e8363ae21d175185f78b088b79232a8b1ee4a11b6f566", + "fp32-reference/gradients/0219.pt": "f7f8ed60e67a53edbaf38770c80f4dac3db2683381c7df9d2ce346283645cf67", + "fp32-reference/gradients/0218.pt": "d1919d831be54c10c1fff2a78f604774d280f7bc6e33d42a4e38a3a18ffd975d", + "fp32-reference/gradients/0217.pt": "a62dc9c713e64abf9433f7d2baa7bd0209649f44009d6a3a3d114808bbe1beef", + "fp32-reference/gradients/0216.pt": "1b0d2e81da8cc84a5297d622b0ce1760f83a125319b32ca04461d092608946e6", + "fp32-reference/gradients/0215.pt": "4181323351f413d3b8d96352c742893d072f220d5dec5c6cc2894fa049dc6e76", + "fp32-reference/gradients/0214.pt": "52aec816316351f1cf09e3d2f8a1fac7ddf60c2be166f442a344958ec465e915", + "fp32-reference/gradients/0213.pt": "cc75dcf14ed94d45a8f993f1e36e73124ec946134ff88e28d02ee11ee08087c9", + "fp32-reference/gradients/0212.pt": "7a0ab54bfd179305ef69f6555cdd724c2c0e017be27385fa6966088717c32757", + "fp32-reference/gradients/0211.pt": "1a4525da9aa3cc3ebbb07f10edb4c03a22eecb299868e253a4df142c64f49a87", + "fp32-reference/gradients/0222.pt": "7034b0dccf92c1fea7caba281b139c41527a853db8a60e560616b149a92a686e", + "fp32-reference/gradients/0221.pt": "e422fc8c1978ab439c4920978ec8148784135544af2d9abda46c2e317e8a4da9", + "fp32-reference/gradients/0208.pt": "17afd90ee3fd66673308cbca0d91138d6f63637f73bceb1e0e734b74bfb61961", + "fp32-reference/gradients/0207.pt": "73284b75bfa35fb6f4453632ad0021175c0ffb2a8cfb0192f9598d67d4b7abcc", + "fp32-reference/gradients/0206.pt": "be3853417df32c614a41562ba6ce11e2a3d12ba341b0113acae8fb27a7dd25e3", + "fp32-reference/gradients/0205.pt": "392c3545647a01f50144be8c0908bb5974cc995d8f416521c269fce3b24fe426", + "fp32-reference/gradients/0204.pt": "5c76c7f71014a0abeaf47753763483186bb4ee1d9b83db15a886fe9b56a4d128", + "fp32-reference/gradients/0203.pt": "ed5588d737e6ca4f5d5c71708516fe43ab90afd49f966308201cef6db624a4d6", + "fp32-reference/gradients/0202.pt": "0b1b89065e42478eb33ce7c71f580f528dda416b65c997b75202e4fe27e913b9", + "fp32-reference/gradients/0201.pt": "e95e304dd8bbf323ba3a644314d8901dfa3f5cf6f65ac14ee1b6446463029a85", + "fp32-reference/gradients/0200.pt": "21fe5ea00c1eb8648dcefc6a352d04c56d47b3d3975ad4f8339dd1091b570c54", + "fp32-reference/gradients/0199.pt": "1f22dac9532faa8d265e9678f8e63f759b4162ebe3fd236aa72a01e78889535a", + "fp32-reference/gradients/0210.pt": "3962095170b0967eeee1c86e6fc27256225160abf987133fcba67ee5e0bb88ba", + "fp32-reference/gradients/0209.pt": "77231e7a1d4b5b3a9c99ed04d022864ccd5a321e290515c3111764ec9408fdde", + "fp32-reference/gradients/0196.pt": "9973b017c9345214248b04363544ef1d3ece268ffd383d505ff510a6c8dc824b", + "fp32-reference/gradients/0195.pt": "9a960ae7c115ffc2bdc2458017ed9c33d2e1746f0a615d725b08f5f72e920388", + "fp32-reference/gradients/0194.pt": "e14d1ba9ab68ede0a055ac665761f62bfaa5acfb12fba96fec91fcfaf165c3b2", + "fp32-reference/gradients/0193.pt": "c3394f2d71d8bf6572ae31239fdfc7a7d925a93894911efbd778fbdee4b5889f", + "fp32-reference/gradients/0192.pt": "0191597afc0d3ad11d02adc401028a9967b5d3cf1984f2137161b8c1219b4b90", + "fp32-reference/gradients/0191.pt": "50a83e66da7b52aa46618325411680efdf7a008df939e82ece17aa7d4eff36a0", + "fp32-reference/gradients/0190.pt": "05e90464879175abba5316f3e852091d11748e6c1bedc53c102ab7e4b9721b16", + "fp32-reference/gradients/0189.pt": "7ac70df33e3b5a715eeebed46db148f3c12060ba3453dc97426fba2b8afd3fc0", + "fp32-reference/gradients/0188.pt": "5158f69cd770cf5db251a96b2b1634fd36dc781ed1249e59880caad14f1a107f", + "fp32-reference/gradients/0187.pt": "2bf8c64a232aa5c0259287f76214eac70cba6282dfb06db4b361cd02e1064adc", + "fp32-reference/gradients/0198.pt": "78ef8746e59eab1471ef933617d8aec8c4dea3803a132e047b1d2fc44509753b", + "fp32-reference/gradients/0197.pt": "408b3bdbb4f384d2459866122b893d3810f0b07191b790a6d1cef3041b76f372", + "fp32-reference/gradients/0184.pt": "b3e02e5674aa74e6371679edc8485ee573ee3d861b205ab8c6db88858ba23c68", + "fp32-reference/gradients/0183.pt": "dde7a8bc4595275db1c6fb26bfb62a802c9dddde2807a3bf0981d08de2febf2b", + "fp32-reference/gradients/0182.pt": "be56c87888d39287928f45af5d828976a34ec73a4a9437eef6d4a881b61b51cd", + "fp32-reference/gradients/0181.pt": "7777e4eb88e1e2312c8dbe1fdd7210d2fc9086b687300526ef0fe1f5418c5681", + "fp32-reference/gradients/0180.pt": "f4532dbd854c7b95ad09dc80841dcf11410b6e495f6710f5126cf04bb7a91434", + "fp32-reference/gradients/0179.pt": "54be67e63f3159371053dc664a051b60fb52bc42ed8a883087cd3b60cace98e3", + "fp32-reference/gradients/0178.pt": "09360ac4602a57af14b53ed431f50f7dac1d41bd65c2d396a4c991fdd2c8b6c0", + "fp32-reference/gradients/0177.pt": "330cd7a890d573598367fa63e37f56ac86c1b416eb26cbf7125137ef6b4518e0", + "fp32-reference/gradients/0176.pt": "9b71236e3685061d9a4bf032ba06ec5faef50b82ece1b7255208a37578c5e5b1", + "fp32-reference/gradients/0175.pt": "520dfc2f5c73201a764a36b927e88cc259708d7c791a84974b94f26875f9bfc4", + "fp32-reference/gradients/0186.pt": "4437758b852edc34f8cfad3abfd5a9a0fcca6092a3d6cee9bec375c189d969c0", + "fp32-reference/gradients/0185.pt": "020c5c3262063b7de91284e2c0cea160ac1ea2aa2f8cd49ddec81fe720d5713f", + "fp32-reference/gradients/0172.pt": "882862c9b0142712b1c4df5601ea8c6f169cc68bc9c4500901d84715ceb3fd19", + "fp32-reference/gradients/0171.pt": "096ceea720f53c78afff13b465a9640bfac01d1c6a3b5b361ccc07b033c299d0", + "fp32-reference/gradients/0170.pt": "d2e0784fc06524ffecdb4a3482e986fd07621ccebe1d6cbf56a25c98351d9255", + "fp32-reference/gradients/0169.pt": "f3dddfc6ae5b77c9a1d39cb902cc856f3ce8ab85f512061046c0e8f146f69806", + "fp32-reference/gradients/0168.pt": "7e6fbf962e9166dbd5c271eafdddc8071a40a446cac930b77c676afb4731fbe3", + "fp32-reference/gradients/0167.pt": "ee003a6246d66822acca80d7e4a6135cac6e6e61215738a978ac6fa23a48474b", + "fp32-reference/gradients/0166.pt": "dad720f418df89b6a9076ccfda4c4454f0b921a48589043a4932318cfab7ae6c", + "fp32-reference/gradients/0165.pt": "d3c36343db437f2ec2d69fc6182256271d153a68f98a25ce01a7d09dd82555bd", + "fp32-reference/gradients/0164.pt": "c4a8a78c37e9f71bc4bf273d8af0b428ad3cc400142dec2cefa8c1f064a4ebc1", + "fp32-reference/gradients/0163.pt": "4a99b33fb4182893bea399da2f290b312b3cb252683292aae11d2e96ebddb581", + "fp32-reference/gradients/0174.pt": "cc2ff6280cae269b3d0f817c88f8a0a715b4a50d5fb17d764c79b3382325a30d", + "fp32-reference/gradients/0173.pt": "4591792c167fccfa8630c57547f1c0f01ba2182e4edbe93f488a6018dc537289", + "fp32-reference/gradients/0160.pt": "7f1f4106d8097ee5c5531c8e38e2b7bd16c81d81aebf60a9ff056889fc11bf31", + "fp32-reference/gradients/0159.pt": "befc5346742535b4b009d4800bc69403bbd9e14b21a640b7f58c9267e23f5c14", + "fp32-reference/gradients/0158.pt": "19df487803b9e7534d28c91fcd2aae4abd47eacc59c3bf7398645d273cba5c53", + "fp32-reference/gradients/0157.pt": "c5d25f9369eb7b457e4b708dd8703adacf76a15b206004d1b34dd7345421f1a5", + "fp32-reference/gradients/0156.pt": "3166f23d79b4ed7aadd48d6532ca02e8090521d14e35159f11f60ccca493db3e", + "fp32-reference/gradients/0155.pt": "077f1bd6d24a041db7f4ae9032d897821608eff5a61db7a53f468bf2eac1eedc", + "fp32-reference/gradients/0154.pt": "4c4575ac1020a8efb8ea4925507ac2876423c053ab8823271fac7bd4eb2a10d7", + "fp32-reference/gradients/0153.pt": "0c9efab1fe6e2eb94d425d1c575803bc7225dac499b9d1210b8e7abae56913bf", + "fp32-reference/gradients/0152.pt": "f8be0b6b44bfe1b67bebeeca2bbc8f1ab8145b857136115982572b680d77eed0", + "fp32-reference/gradients/0151.pt": "1c8be24f986eefed4dbe1959588edf1422ad1756ffb2283125c5cb6548b0f766", + "fp32-reference/gradients/0162.pt": "d4ac31763f94cefac0e7ffdf139a9b1b0a1f834ba8041a2b2111ae9ae8c15afd", + "fp32-reference/gradients/0161.pt": "d2ff019717053fc34e5f0206174a770866f699ba21962a03cbb924a7ba1cbb6d", + "fp32-reference/gradients/0148.pt": "f597968fbe2d77bee43d5559d1431d5b7da81b18874f932559393406eaaf7177", + "fp32-reference/gradients/0147.pt": "fe3fc60bc9005807ff5c21cdc3ceea0973edc448ac5b44dce4604e9d18b4043c", + "fp32-reference/gradients/0146.pt": "20e4bfe11ba93df308bc986e6c4a45c36b7ae8813fd5b2e1006e7c4820af1dfe", + "fp32-reference/gradients/0145.pt": "4ab42890c9082364010b4df131170f989a2538c486d89907de1104bba2fe4ef6", + "fp32-reference/gradients/0144.pt": "cd60b2a8446b40d6a7a6f0a889ca5c421bbd73bf539cf73fc3ce2041699b63f9", + "fp32-reference/gradients/0143.pt": "4f495c888e6cf02d397f37339c50aa0ab1d098d57f2068aa100c5ecd3541e30c", + "fp32-reference/gradients/0142.pt": "577149715278de3eb1e73c6875e6840ef7d8329baf33ffc424c6006a703221f5", + "fp32-reference/gradients/0141.pt": "c90e0e8ede5f868e097ac2fb04e8215cb4b34659c3114205f2a96db9ee400e34", + "fp32-reference/gradients/0140.pt": "505f0ce6e6462c38703163cf6733dbf28990f0985ebf6e9c1dfac35aeefc069c", + "fp32-reference/gradients/0139.pt": "dcb2c21276b9bfaa597605352caa86b66019441b19785058d0d8b9f3ae6f49a6", + "fp32-reference/gradients/0150.pt": "91ecd6deb6402a6b9654562418a8dd8f231d19735576982188ce886a0173bfc2", + "fp32-reference/gradients/0149.pt": "f4b515d08ef82923144e9ea06e82ddfbae76d7e017277746f7307912af6f90f3", + "fp32-reference/gradients/0136.pt": "4db5a0dc8f6127148d007e28d163968f820508625e2d8cba0c30adb245b367df", + "fp32-reference/gradients/0135.pt": "ab09687d9b1df75194c431a43f041784e1da0dc3e0279b9001f208a0269e90a0", + "fp32-reference/gradients/0134.pt": "e82046c52150ba940b6fb1a17bef72518e96923a5e358920ad0d3d5b99bccc19", + "fp32-reference/gradients/0133.pt": "b96f4bfcf59e9660377e7f4a09c3d716c6a310bac91c3108e439c7a716b374a4", + "fp32-reference/gradients/0132.pt": "9b8a47fa6fda0b78fc3230473f5e3397c66d286508faf91121e33aa1b8eb0a0d", + "fp32-reference/gradients/0131.pt": "197e9a7ab040df42415317b624ae6a8179d2cb897a048e8ec8b1585253d987d1", + "fp32-reference/gradients/0130.pt": "a82998d9cb3f69b29f02e973d4f6b05d7fd5b71a8b3274c6e01d22d7d083283d", + "fp32-reference/gradients/0129.pt": "c603922d13f5d498e416addfd24cdc692f0c88826463857985e91f51b6242d61", + "fp32-reference/gradients/0128.pt": "eac06846dde403b093c414e46211ae5420c95e6676028cc0c19e0fb820221a2f", + "fp32-reference/gradients/0127.pt": "584dd5c4dafe5526ada168b79564d2d5a2d6ac9885475235c79b7c6db5f7dc17", + "fp32-reference/gradients/0138.pt": "31b9f1e938dd2b239e872c6d83599349f210bb52c9236e84765c04ea4ceda70a", + "fp32-reference/gradients/0137.pt": "6af6a3eaaa857bd1b34a075c77bdecd0bc232448bb762c49c95202a15b771219", + "fp32-reference/gradients/0124.pt": "b5846f779c74725d02a6cedcd676eb43734d4d20bcde782f3a3012fd25daf113", + "fp32-reference/gradients/0123.pt": "1e130cfd1e73fdf877088f9c998f8f5754ecdda5ecdf6ed42ba96faee3bf2dac", + "fp32-reference/gradients/0122.pt": "78ca0aa4d199db71f297f54c2b52a0aac683760322cfdc61476857a5661e85fd", + "fp32-reference/gradients/0121.pt": "20231c6a4f9db313eb50f34a38583e245fd0c32f52257bbce19b3483fde2569b", + "fp32-reference/gradients/0120.pt": "cf26d994ef0dd0b6077a9f79699287d032d04e196d7ac942433d78af6ee1b7b6", + "fp32-reference/gradients/0119.pt": "3ff84f127c5bf5e89ff6a817d4b3be0b4f457ca915ac666f38d1f9aeba2ecb76", + "fp32-reference/gradients/0118.pt": "d744872d12bbba24546dd2ba89c4109e921367a974937d6c8628e4adfc58065d", + "fp32-reference/gradients/0117.pt": "e2f4c477bd9e8d41f05b433baa278277b1888797e4d40ab7eb6cce8814649e29", + "fp32-reference/gradients/0116.pt": "4e0a4f1866f621f004de2b20a08eb631f572833b5232bc0acfad693cda1a67f4", + "fp32-reference/gradients/0115.pt": "44fe298e7d2d8f0a9ac3013172486e1c7db7862e5510c4e278f4e7c8f0c3c42a", + "fp32-reference/gradients/0126.pt": "bb5b91b0fc8f5505dd13c59279fb36545020ac46543afba7d1029eabff7d5a18", + "fp32-reference/gradients/0125.pt": "4e605ef9b9b5a6dd2cf3c0d097f07264e8c86fb5315173386a52907d1444a283", + "fp32-reference/gradients/0112.pt": "1d81d90143fdee2c5bbceba8b10e18f5460c91c79892f70f7bd452ebe853b4d0", + "fp32-reference/gradients/0111.pt": "1f4cf67047321d713c447c5182377ac7e70cd792e2185b2ba8f498affa30caa1", + "fp32-reference/gradients/0110.pt": "af8b9e1ecbc48d297314e2886a80be8196c3ad50b62780891f47ba833c9f0d2f", + "fp32-reference/gradients/0109.pt": "f28d836783666399312c65d077a44478f890459bdb8b7f62bba025b05a0f791c", + "fp32-reference/gradients/0108.pt": "ad1506ed4e9d06d43381af842d2e603771e39d8ea35cbffcc169f7fff6584226", + "fp32-reference/gradients/0107.pt": "ee3522f5f906c621463224feb56f7f75021586c4c78d02aee28defc866954c19", + "fp32-reference/gradients/0106.pt": "8ed67af3aba1a42c0173f5fee9da3d65fd63683568ec679bf7ff0c556b7d6da1", + "fp32-reference/gradients/0105.pt": "fe1228cda80642b2d3e541909b26e3e53e99bb3ad5254aa9b80d79c8f32af793", + "fp32-reference/gradients/0104.pt": "c60e6980543b98ebd0dd271687cbd4bebd75f5dfb2244267b4c5c7ed7bbe0aea", + "fp32-reference/gradients/0103.pt": "11dac88c8b802a30dcf3e77a08400e7fbe3eecf3564538966df03c6ede4f9c34", + "fp32-reference/gradients/0114.pt": "493169ce6544ee958d39d54f7ae2c8b68aea1b998987794c6a2fd9d6d5231535", + "fp32-reference/gradients/0113.pt": "98ad3557a69c4218a09241a3476ff8b5049d6056ee29ad0fb574d504facaf968", + "fp32-reference/gradients/0100.pt": "70f66af21bb9f5d5f868eb3c064c7b810377123da2d57c253e70f9dd2505c44d", + "fp32-reference/gradients/0099.pt": "7a4e1dd9733c55b88d5275e40325b8ffefbf9c2c70ecd19d0d8d9b0ecd3f02c4", + "fp32-reference/gradients/0098.pt": "4dbc44dc7e03ad611ca70fb20179749afc75dc296809b9ece2a762729411032b", + "fp32-reference/gradients/0097.pt": "b769dd57b82e763c726e908f275998bb3854901ae6278695ef797668c2ebc8b4", + "fp32-reference/gradients/0096.pt": "a7f51e770c17c1a83fb39ca6561268efe315421e7aa3a1da7be8ca1740857487", + "fp32-reference/gradients/0095.pt": "edb1849e30227917f2c234c0e808499493d25866b14cb523379f62285458c349", + "fp32-reference/gradients/0094.pt": "eee1eb0387492d2eeef00eb8224c9a524e55b3fd555264e3d1034291988194d2", + "fp32-reference/gradients/0093.pt": "f21d6b20698c4daf711e87917064e7ce2291d13cb775f90b4bdf0a852d4a5fdb", + "fp32-reference/gradients/0092.pt": "82000f2b130e0585f1609129b09d8ecb5cef98a3412ce47b7107c2b48bb6cfd5", + "fp32-reference/gradients/0091.pt": "75b0f1d8c206663516eb65d47f94b1fbc5d20c91b4853ea3544e02449077d70a", + "fp32-reference/gradients/0102.pt": "b625032731794390d28521187feac7c92745d2ed56789d821540b73b64f854e2", + "fp32-reference/gradients/0101.pt": "fbaa0c0ff79a307daa72da00b9f3d0d16f60e17a872621d2c0931f4813085816", + "fp32-reference/gradients/0088.pt": "6630823b8faa6cda492942db91e73524e2dc9d5760175e9858b7bca512b77690", + "fp32-reference/gradients/0087.pt": "7001edfd5dc9854b2c6c55031590d9cacb1db146b3d148b6b62f334ab2c5b81c", + "fp32-reference/gradients/0086.pt": "a5c226ce3763a07fb29e3b2656bbb69f6ee4505b232830b0ccfc14772e2a6be2", + "fp32-reference/gradients/0085.pt": "2b66738ad8814a1ba47dcd9776029b3c27b19cfdec8d018720622e764e223708", + "fp32-reference/gradients/0084.pt": "83bb85db5b6579a0bb7a7b4af6f276b8130991a32805051755c91a2e69f8e308", + "fp32-reference/gradients/0083.pt": "9722fb5d0e3560b5cd1dee15e382ae36c85077219aefe1d6d9cb86faad7a6079", + "fp32-reference/gradients/0082.pt": "40d597f6fe90c7f9c0482008f02b3d3674ff62e24fc1a09d825d8737cdaabbea", + "fp32-reference/gradients/0081.pt": "5152cc9686464f57f751f8f50c9ec3a28dab0869ba3450e33c3ab3188319587e", + "fp32-reference/gradients/0080.pt": "1c410fb3553d74495ae5621ad82060807a463ce82732701c5e7919345a4a1cc5", + "fp32-reference/gradients/0079.pt": "2a5b84b3868ae17c3fa593118ebb50f3010aa0c80366e76e8dc638f503eecc50", + "fp32-reference/gradients/0090.pt": "dfa763e291bd9fc454ab24d0afb6bb554d485cd309eb691459484e3ae5d0d71f", + "fp32-reference/gradients/0089.pt": "cdddaa734b91a56d3982037ddc06b31b3de1b127da0c18d24d418c4a22e5aba5", + "fp32-reference/gradients/0076.pt": "728eaaf85f2fe00fca2709fc058ec4251489905d0b80e0c125d7b5f9aa91cfe5", + "fp32-reference/gradients/0075.pt": "20013f7b225182f41902427568763ed16c57a37ee8c6f3fccbcb6dc2db0a0240", + "fp32-reference/gradients/0074.pt": "4bd48f9b26a284f88094ee2d1b5e25d7a8551dd0adcd2a1022e61e7033cc1e2e", + "fp32-reference/gradients/0073.pt": "2d8c58d640fff1c73e8ff384be388e516ff7c141707ca5fabb63058783d51eaa", + "fp32-reference/gradients/0072.pt": "8f465620b1bc98917d2afa8afa9df1d7b598b620b9b43ab3bab980b672ba59d5", + "fp32-reference/gradients/0071.pt": "7c3b0806b92db3c8b232d47bf69b780c79d818d7a18a34d59547cc933ed52369", + "fp32-reference/gradients/0070.pt": "61ef4c0077c8eeb8425fc99afc26fbb80bd305e508528770434aec19477a4424", + "fp32-reference/gradients/0069.pt": "f3cccc0dd4318165e0a57c5ccc6e350396f175e88fb225535041cc57b6766814", + "fp32-reference/gradients/0068.pt": "6d2ad0be502f304db305602b70e10ff016717a25cff20bee9c21cf7d6870d11a", + "fp32-reference/gradients/0067.pt": "97522b486c15336bbcad70a081771d91bd6b0c496b9e7b68580b180338f3c32d", + "fp32-reference/gradients/0078.pt": "1aba7ec1b18b858cab2a82553b1b5b22ad8a1ec8d5ac619c46fc2e652330e4b8", + "fp32-reference/gradients/0077.pt": "7a8a4c1eb4380d436f070aeb4b062c56e3e389ec20b8e0ec75c777252140e729", + "fp32-reference/gradients/0064.pt": "9d487ed2abd411d54da33a569dd4ab64f37f126f76c5327b72794c0d0b77ad2c", + "fp32-reference/gradients/0063.pt": "53f2f6cd0c8985d68e2ba59c4374ee61ee308047cc41e6c70d9e515eb6a5fc9b", + "fp32-reference/gradients/0062.pt": "06c1a9ad1b28e22cee769d513b771784c6ace0130ac97497c1c2a8070b36ba67", + "fp32-reference/gradients/0061.pt": "7071b184165b47a92d5f415a47db3ee8705b9187af41bc5b57a49a188961f1ce", + "fp32-reference/gradients/0060.pt": "4ef21ba0a0a59600d8a33c4b8bfad87463d0988343a30add0827b7715945e7c5", + "fp32-reference/gradients/0059.pt": "fe2eca1fe0316fd0917bd5abf263e19d1324513f53b11ee94f1372c089fecf94", + "fp32-reference/gradients/0058.pt": "c4332f8a32817da839859cfdba35217a3d48ee77ef99bf776a807f6d1c92479e", + "fp32-reference/gradients/0057.pt": "d788d2c6d6d2136cde78c9b2a87328f7f1f473d6b8170414c579c40cc7a44bdc", + "fp32-reference/gradients/0056.pt": "b18fe46471fa08b76922c5c8469b45e32f4cce5b667fb140b5c9ee13850e2a62", + "fp32-reference/gradients/0055.pt": "7f5f5d80d2d0413c175d5d7f1fbb4e979a4a5e75d66eddaf269e7a5a15406c18", + "fp32-reference/gradients/0066.pt": "674e9a06e7d0a44fe1c680cbadb8b694714cf8988c68e625ce82a2cf1078e3e5", + "fp32-reference/gradients/0065.pt": "2170eb1d748979d92954df19908340928198ad5637faadc399ffe4356b0789c3", + "fp32-reference/gradients/0052.pt": "65eb7545e27ff9452acde808b36b588a2f9ebdda83c9cb34fa29ef2c55a72438", + "fp32-reference/gradients/0051.pt": "1b1345f6ac31172a92562a5f0ba8b82d94ceb7f476bfc8c679e203799b1a660a", + "fp32-reference/gradients/0050.pt": "55e84dd9a515859978ae9a5a96ce178baf057da1cc14cc49864dc56897539727", + "fp32-reference/gradients/0049.pt": "10a7c34b3b1997aee50246b7fb5a51b87866153f57e54d3faa75cac826cca248", + "fp32-reference/gradients/0048.pt": "d4c84c0b1e38e077e69dfbb2eaef2dbfffb022d38bc6284514aeecfd36444b83", + "fp32-reference/gradients/0047.pt": "ec1a717e2a51887118ab8f4c62af44ce5176a16cffff35e25af6a75b955b8845", + "fp32-reference/gradients/0046.pt": "d6670b2182ffd8f77f953328dcfea550cf00ee03122ab2e804b629b516302814", + "fp32-reference/gradients/0045.pt": "a42872fcc9b0203c79329e0f731c62e2cffaddfffc3dc105ee136157d0f777d1", + "fp32-reference/gradients/0044.pt": "e68139314f366b7c18709f828e37f7c063c5d44d198808f59c1dcd042c711f00", + "fp32-reference/gradients/0043.pt": "8ac7a0cab906637797d6919c6a2e26a536babee044e564deeddb16e18cee45cc", + "fp32-reference/gradients/0054.pt": "4dc8058a898db51a5310c51838d3500feb29e00a833ce8c244b714c92691bdab", + "fp32-reference/gradients/0053.pt": "5148c5396672518eec5c4fc641fa5301231971fa70d3c15db300ac230bfa7158", + "fp32-reference/gradients/0040.pt": "50f4495cde40013a908e8110c47fc15e6de4c11254649b154573187b3131d8da", + "fp32-reference/gradients/0039.pt": "8a0b2d7962117ea6f70ca44abd2d23748c71bad596388e397a780780c318b0e6", + "fp32-reference/gradients/0038.pt": "e816f0674c99299f65154c3a43fc21ff40b0fb05b7a3a5f49462526b2fed2fab", + "fp32-reference/gradients/0037.pt": "40605f89ea7f0dd0fdae830b978d06dd216240beb9d93bdb0306fe9f7e3271b3", + "fp32-reference/gradients/0036.pt": "99ccc3a8e9739279978cf612bf1703f9ea64874fd0a83f9a39fb8a289d43361c", + "fp32-reference/gradients/0035.pt": "733a0981c95af7b1e5a4f14f05bb576f905796e9d09ba065193fcafc4b91c7aa", + "fp32-reference/gradients/0034.pt": "2ad4201fe5b25d1ee3b13a32c794f94a85e73962d614ec4f9b91a6e2cf5cdfb5", + "fp32-reference/gradients/0033.pt": "3906ebb90fd6a62739df341571e2c64375bdfeedb47536513c4e52433c9c6aa6", + "fp32-reference/gradients/0032.pt": "06fc4c4d064a07037b709e41059bac406daa7cf468f371f1ce5c9b635482e400", + "fp32-reference/gradients/0031.pt": "c0aae5701c42c167c4b65ac48f0f55086bd3afe390a2fcce405655ae2be7ed39", + "fp32-reference/gradients/0042.pt": "3984142de083c8d2c12c2ff72e672a3ef2ab16ac182c80b2e4765c444f29c1fb", + "fp32-reference/gradients/0041.pt": "f0e5722180894e058da84b2bec560b937e35e59378ca095d53a4d6745935649d", + "fp32-reference/gradients/0009.pt": "97d4a4c63a4d19a156195d96cd0489d2c94b8fc05b9ba5b07f5cecd083263924", + "fp32-reference/gradients/0008.pt": "ff74c7140174f036a208d51310f75fcd7c7aa8d10b03d0c172d58eb245b94157", + "fp32-reference/gradients/0007.pt": "a9ed52ea232dc4559e4fbcf024bad48a5aaafb79ee6e399d68b5844931c3b24e", + "fp32-reference/gradients/0006.pt": "ea30e4874686fd68cdeb741a2f0fc7dab4704a250b3b46b808da46fdaf1ca392", + "fp32-reference/gradients/0030.pt": "2559c36a9f9532b818bfa734baefdaa5bd6ab465a7456f3306153377715011b5", + "fp32-reference/gradients/0029.pt": "77d6954cfd6fb5b4dce2fdb75f75bc24b6e20bbe61f61eba6d381929450438e1", + "fp32-reference/gradients/0028.pt": "3a4bc8d1864ce3048424bb7760800882aa41aff44acfa94b405d1d053fac9ba5", + "fp32-reference/gradients/0027.pt": "0831a649430c7fccdfe0d09f9c2836f151ed1ad9a2d812440bea3dbc4877e4a4", + "fp32-reference/gradients/0026.pt": "539679efb96862b5792975b15de270da58ac14a45d6dac71c824dda9fa1d1723", + "fp32-reference/gradients/0025.pt": "e69675232c8ac02fb7874276df98c44fe1b570ff3dd217ad92c39c4a0cde15de", + "fp32-reference/gradients/0024.pt": "27a8fc9823a9d73b3f2fda8a39fa87af6d7368b9ca4f8a036a23967257f3b035", + "fp32-reference/gradients/0023.pt": "cdbee4034ed3595a82c19c2e128019a9b54624e1b904e9f8b98cecd8f1f422eb", + "fp32-reference/gradients/0022.pt": "ba98182fec10ffb5faf61fe825fe24348d47524a6afb9e6b73861164dda17e44", + "fp32-reference/gradients/0021.pt": "3ac846501b778b7ac829a350173d68e24c4655901027d9d99858c00126e5543d", + "fp32-reference/gradients/0020.pt": "6373d0f0a778e4c4e76ae7910ac98ec5cad79716d870a593f8956fca3d22f690", + "fp32-reference/gradients/0019.pt": "8d5a5f8d5d4f6721c6a4447ea9d91b57c5a16e57d0f8c333a9cfe359202cdb68", + "fp32-reference/gradients/0018.pt": "260dba92837138bfd2c58a4dd4c2e23f224a32c1b607146d88d0c9770dedbc03", + "fp32-reference/gradients/0017.pt": "f4e2784dc14cc048acc47cd0e1448c87af3a2b94656e555b68124ff6edc59c60", + "fp32-reference/gradients/0016.pt": "ee3e74c0568d935973f196a7db97c83437d9c923713f84665ad2187017dbba8f", + "fp32-reference/gradients/0015.pt": "11b82267677ec4b6ec327c5c34b578fec10e55e73d50b0d16e51bf7f9ad43422", + "fp32-reference/gradients/0014.pt": "a06fe168ab2e23c279b8aa1d3facde62a7f764423d3195e0366f71884f33b4df", + "fp32-reference/gradients/0013.pt": "bb5e1e290216b6ff0747caf7e5f9dc44f7ab386d71d7a5cbd439d13a5358c11a", + "fp32-reference/gradients/0012.pt": "630b5c9bf756301a9a4a14ee045202b463fdbc9bc8c297d90a61dabdbae536fd", + "fp32-reference/gradients/0011.pt": "60fcc18f64a16e2f0b29fefd7d2c248abe2ef3587a0b3ff6b82a8323f856d66c", + "fp32-reference/gradients/0010.pt": "319cc63c3a9858022c03b06f513070c6ace8615a5d144a58bd0be3c9cc386a9a", + "fp32-reference/gradients/0005.pt": "d454aff83dcd6e14cbf9685128797a913c73539ddb470fd286c47e7ad5e14c3f", + "fp32-reference/gradients/0004.pt": "2d20102884ed602eb8ec075554844fc5840bf9464a4bc353be5f993ad741074f", + "fp32-reference/gradients/0003.pt": "53c91edfe9637cc99b14ff5e1503e0e6a21996c73fdb94f3149099a52a384137", + "fp32-reference/gradients/0002.pt": "0419ee6d6ad48307830cc52cdeb3421584ce0e43175cdf5ef832b13f10ac02ab", + "fp32-reference/gradients/0001.pt": "a13fa20784d73f4073e29ccc80a63d6cd8d1d3e99127c939d1f3e0532d74aadc", + "fp32-reference/gradients/0000.pt": "82f4f19f10c9540f0f9e7daf898b793723970846f0e2af0f19e5cb6137ddc335", + "bf16-provider/report.json": "d49a80d6e438eb1f7d91193ad3f3ef5b5b3692b2c86a7532b86f23359ad33d79", + "bf16-provider/gradients.json": "4aa8c3e141a8649ba4384ef1d7852cef2d1ea1aabcc3a06f64acd17efed61e9f", + "bf16-provider/tensors.pt": "0946fbd60bef1996693d2f97d69f13e01165c9094fca9fe988cfc95ec64773db", + "bf16-provider/gradients/0637.pt": "e3844a24c49f10a55cb8cd4fcd97f71d814bccebad3edd81988a49502bdf3b86", + "bf16-provider/gradients/0636.pt": "4f0171f5de3f271bdbab29b9250b2e7b9415209be2c43fde44314b75be9082ff", + "bf16-provider/gradients/0635.pt": "0a93c6cebf6871f3a8872fca367c5b26a23800fc924d9ed5543df2f2a335f0bf", + "bf16-provider/gradients/0634.pt": "bd513cac9f098d35c3b83e59d5dba5400e32865958783fd2e21e151ecad331c1", + "bf16-provider/gradients/0631.pt": "51598eb239d21a6e2cd3079821ea65660515f1468b8df864c3a5e566b93c6214", + "bf16-provider/gradients/0633.pt": "33806dce327b8335d65d3b898fa90056564828a353b823b625f11593e73d5514", + "bf16-provider/gradients/0632.pt": "0e60adf736145369074cf76f434d7cf9eb221d6f49a67da2f4978e24485a92bc", + "bf16-provider/gradients/0628.pt": "4ef0834886c1f5916a6f6305766e1024b1203483797a21460c28b06bb5beba6a", + "bf16-provider/gradients/0627.pt": "913665b9659bdb29a70bc3d02eb2d445c85a5009c83169a4147330676b21e264", + "bf16-provider/gradients/0626.pt": "28c56ce68138d83acaa74e7b2a26e4f7bbcf5f76551df90ef6c01ccd1bb6f113", + "bf16-provider/gradients/0625.pt": "ff8c9ec07cec56bb2d12787fffbea708f8b734291f62730e0019dd0baf15b1fa", + "bf16-provider/gradients/0624.pt": "3fc9857d0ddac770e071e5d998908b0988d5ef7405eec8138f8f45c462082ef2", + "bf16-provider/gradients/0623.pt": "a0c593df280bde9d01eeef67decdd64a51828b36d51046c114cf9bd5bb590a1d", + "bf16-provider/gradients/0622.pt": "5cb69d6bc4671abe1a18e59031c26a746e6dd585c6bbe3cc2be15aa813cd54d8", + "bf16-provider/gradients/0621.pt": "851a411f41d4b9e9a278c97a2a92043aab6366a7a4bc4a9cbcb364979dcb18fb", + "bf16-provider/gradients/0620.pt": "9db3e1ffb4b9ae818bdf8bfe0f06ecf09cca5dcaea67e7d901d5db79e53c0b10", + "bf16-provider/gradients/0619.pt": "d1babab5f9512b93122dd571e60da19cb76da15dfadf504509d67d03a94f68d8", + "bf16-provider/gradients/0630.pt": "ca7f24a16c16afeda5d876e29ef3b7e3293ec77645bb33cc7b1b8fda6925cc64", + "bf16-provider/gradients/0629.pt": "38ec65bd40e3f255dbbf49da0cadc63becbe53a5d5c6c89f2770fca50bfbbc8f", + "bf16-provider/gradients/0616.pt": "fa914d3a6b6cbecc3a38cccbcddbff1fa34aca0cbc2775cb91055e43a76d1b65", + "bf16-provider/gradients/0615.pt": "2b3075e68c3d6f92b541d642810ad7b11f505bb5d33b67637a44968688093ac0", + "bf16-provider/gradients/0614.pt": "ea0d5c550b383dab6225d6a7790239a77c0a90b6e5fbe4a5f3df8f680bccdf9a", + "bf16-provider/gradients/0613.pt": "62d8bca5994bd2ca25a853951ff7449554c7c1389644deef17da1b72ae5d17c1", + "bf16-provider/gradients/0612.pt": "7a7229e523859a3d0582b2ca62d54dcfaea25d6c4757be3bd8895515b724261c", + "bf16-provider/gradients/0611.pt": "bdf44f033071dd300e9e55b20e5db52b7838b67e2db3561532b642b3c3765f8a", + "bf16-provider/gradients/0610.pt": "d3fdd2b2236b3bc9f013f9efa2fa36226b6aefe2c9f2c6701ed99070b7d7605a", + "bf16-provider/gradients/0609.pt": "e8d0bfbd34d5ebe63d08c85c2eb61f040425e8885dea6dbe89282861358fac32", + "bf16-provider/gradients/0608.pt": "2b2fc55e2d8f80a9a81354b259bc49cbaaa109b1e509eb75521d7ac68fe00760", + "bf16-provider/gradients/0607.pt": "7fb3904db359e9aade77abd89a6adceccc3bbe3c781bc1c21b99b0efe5ae6245", + "bf16-provider/gradients/0618.pt": "018aa39a33eb4c1b40782f074212f339aae5b9fa0b3e892838d90abdef0a1bc5", + "bf16-provider/gradients/0617.pt": "016dc3451f14408615f3ea626beb19d7572df1c340c76fd12e839a9a5f69bb6e", + "bf16-provider/gradients/0604.pt": "bcefb03c9741e6625375bbe3bf23a4e183e15b94e04ec9f17937b490b6913643", + "bf16-provider/gradients/0603.pt": "0ed996bb8bf4df3dd76ec08056f76f03aab6eedd219f13a2f160cc782222a0f6", + "bf16-provider/gradients/0602.pt": "54e929bb4912a7ea5c0df1d971c6d7f2592608e03882d34337e2d440627b1434", + "bf16-provider/gradients/0601.pt": "6a9e96d568953e43e0cfdad1bf21f0c7af94ab58ec20d9f7455e4e0897b7f39f", + "bf16-provider/gradients/0600.pt": "fbdbc4a352ba654f2c3a9fcf102e2a4b297876e9e5bdacc0ebd2cf9b4535e83d", + "bf16-provider/gradients/0599.pt": "88dd012e28d652906053a614309260f8dff041b49f5bd25acc735694016ee9a2", + "bf16-provider/gradients/0598.pt": "221d3d71a7f2ae8df65643da20c1598506dd329486e48bfcefd2ae90fa15a1d2", + "bf16-provider/gradients/0597.pt": "8ae935a939083e830c030362f857d70dc3b403f31aa60b069c8502fe4c78d92c", + "bf16-provider/gradients/0596.pt": "9888c82184122e0b7f1b24eb98d886545a11871e6f765be63c4250d58d25b6a9", + "bf16-provider/gradients/0595.pt": "890446a2799cd772077e1b3a96b9c28b44bd67a238835915352c72d69bbc9b62", + "bf16-provider/gradients/0606.pt": "d878219a9c322bdc610b840e657905ebece4bcf489d4f60e20ecb8a634cc9b27", + "bf16-provider/gradients/0605.pt": "ad169fbd40159df0b2a04665eccf1ede575c60d63c1c7837fd484f0aa38de5a3", + "bf16-provider/gradients/0592.pt": "8bbdd22256b837d712b1ee07cd7638feacebdb19cb05d123a0b14aef46f7dda0", + "bf16-provider/gradients/0591.pt": "fd3c0839a61f3c8e8b3387b0d64aa57c0cdc8314c1cab47b2af6fcf0de18e6ea", + "bf16-provider/gradients/0590.pt": "689f0c9b67e7e3fcd0fa372b130626ec3b3912940caf34b39c2f35a79ae54d1a", + "bf16-provider/gradients/0589.pt": "0a8199f7a5c05ac343dec7fff537a821690c2eab374e9bb4b72e985920637188", + "bf16-provider/gradients/0588.pt": "bd960d7750933208a3a0deaf760806e8f41d1b809fc952ca4721efb30112e196", + "bf16-provider/gradients/0587.pt": "0ecedec069314be2c5c804ac9c080aed436a26bb617e00e16143482e17abb9ca", + "bf16-provider/gradients/0586.pt": "c68f6b73c56a0364d8726b63b5014cb85ab0d44fe4a6b97d28cb881acb77c11e", + "bf16-provider/gradients/0585.pt": "8cb2b4527d51586788fce5328f2046fdd02fb3dff61b3d4db1bca938797cbb82", + "bf16-provider/gradients/0584.pt": "fd9e2170bfdaee78cbad2af809c172ed1cc011c39fe5ce5e81126b5e21e2d485", + "bf16-provider/gradients/0583.pt": "945e8f576364faffaaab6ddde5e24d26b159369d124c04510da98221f40f82af", + "bf16-provider/gradients/0594.pt": "8602f38b3b8d337dff24593156a4adecf274da9a80cfeb66c477c131547f240e", + "bf16-provider/gradients/0593.pt": "46dbd687badab03b9018c139d00adc6bb4b4e8dbb577bf1eebeec2458ab871fe", + "bf16-provider/gradients/0580.pt": "1491b47c7978d9ea8359a56292fc44c30b89a89ea59f208715c78daf9d6ac2d6", + "bf16-provider/gradients/0579.pt": "c63604ee2ffc3ac441857498a5cc8e563ae0a1a68d79d378ebdbb654726c9985", + "bf16-provider/gradients/0578.pt": "8573eed9b76c6d67e917a9568684fbf073e4967f1b47aa2f1b3ab6f330fe08cc", + "bf16-provider/gradients/0577.pt": "c5222d4c103c9860ea6fddcfb2bb08cf1dd0986784e3abeab51e0fedba6e54ec", + "bf16-provider/gradients/0576.pt": "23aef65cf4f34badb6f54f79937f785907a60361f8ff8f376d38ec4f9325ff3c", + "bf16-provider/gradients/0575.pt": "a5a3e293f026a3ccfc1f3a9f69eaf9c2b48fb36b307e4e5c63d4bcb273150402", + "bf16-provider/gradients/0574.pt": "ad03e712eae857ae84a4390c896367ade6323952d7ced6eed7e4a8d7eb40da31", + "bf16-provider/gradients/0573.pt": "57ee0afb6f94214756c9c54a490823b27abdb9005a849b0fdc5b640747a32a89", + "bf16-provider/gradients/0572.pt": "8768cade03c13d12225068b4809d7cb13d89a387055fc3ece71bd1cbd01a04d1", + "bf16-provider/gradients/0571.pt": "7be2bc4784f7e9950a73628a98feb51c287665f55c8d471a6e63f27c73df57db", + "bf16-provider/gradients/0582.pt": "54baeabca7ef89c73e80f7380f99f0de89601beafe54afde037852fb0eb6f7ab", + "bf16-provider/gradients/0581.pt": "9359cbd8a00cb77ccd6dcd98303949a3d661bdd1f5194be1c3302565d0a67ab7", + "bf16-provider/gradients/0568.pt": "be46d5c4351a616510ab17d76b5e806bb29dbcd48b00554ca91d91effa60ecb4", + "bf16-provider/gradients/0567.pt": "5770498c485f2adba56c536b11ec8e6b310b125f35e3d161c9114b50e48c08fd", + "bf16-provider/gradients/0566.pt": "984de7b29f1ec4a997c47ac9cde37713c5318e38a03974f225e435980322e5c5", + "bf16-provider/gradients/0565.pt": "0edd6067054df8df147fab513af9d769bd5ffcbcae67417b1fd14a4cf7d3af2f", + "bf16-provider/gradients/0564.pt": "a7a9edb3ab3d9d41d05df54033c91f211eb4ff312ef936f7cb55c83f679cd673", + "bf16-provider/gradients/0563.pt": "4a025606f6cfe00c2de7eaf2b80cd472133eebf6108e7d64f71588453b4339fd", + "bf16-provider/gradients/0562.pt": "9fecc5de4630537cf580421569821b71bafcfe65dcc38ad81633c08134045ed2", + "bf16-provider/gradients/0561.pt": "533684836475cff6d2bc5d95a923ebc2ba5e0a313f91b98d597f5f6a7c80c47b", + "bf16-provider/gradients/0560.pt": "27ed55d0e7e8f8ccb0d3e8d5cd88f31df7cba976ee71537919c6106cf3fdf908", + "bf16-provider/gradients/0559.pt": "53a2f2a9f5ff38cd615ea1d0b716554685ccc76b1c1c0ed574a419e3a61ab48b", + "bf16-provider/gradients/0570.pt": "17ac192b08997134047773026351be1d1db636cc7f6f8b7932db24e40031ccfb", + "bf16-provider/gradients/0569.pt": "86c813be3a9da399fc2b9ef799f5362f2113a34cb8dfa09759004d819adc70de", + "bf16-provider/gradients/0556.pt": "5d38387ad6a82948ce726466261d5080e0be20ba316cdd2a733063555bc6dbad", + "bf16-provider/gradients/0555.pt": "22cc0d5842e8410830f2c4bfa8aece9afd10cf946cb8bec33c76fdf8fac329a7", + "bf16-provider/gradients/0554.pt": "aefb9d245e82e0489b09a907b22cbce33049b23a6d2eae2288a4ce8d6d791414", + "bf16-provider/gradients/0553.pt": "a7a46058da0b6bbba09abc1eb9ba6476c8e6b79b1a1b9c44a46640337ac7a0d8", + "bf16-provider/gradients/0552.pt": "4f5217e95931f275f7f8f123e11bc785e993ae3e1666097fcea3e3d9a3d8b04e", + "bf16-provider/gradients/0551.pt": "a4316761352fa318b00679b1ac45ac02e367d12e1e3cb684d76fd9f06fccede3", + "bf16-provider/gradients/0550.pt": "5146dc1f6469a40d648a442cc26722c6a936ec0d61ab9f68d5a4582ae5c0dfc8", + "bf16-provider/gradients/0549.pt": "fac434ba5e865bd6692d8a6547e2884ef72aa51784ccfa5b2534cc30599f8a1f", + "bf16-provider/gradients/0548.pt": "5c102395455ce8168e4be4c20b397924957ed1d5eed6fb190ed706188f57c7cc", + "bf16-provider/gradients/0547.pt": "ffc8141ac3d86477b39d0866968032db4e834753e6280a38439383705ba296e2", + "bf16-provider/gradients/0558.pt": "ef70685243de893d4f8a4ca105ee66d81ff4f65b791a4ad971cca9b1069689f7", + "bf16-provider/gradients/0557.pt": "eb10aeab6ce010ee9836d3fa25e99d0847483c30be9b5215e01712306a1081cf", + "bf16-provider/gradients/0544.pt": "f5812b13181808ded4de9c07802f9b30b5181d1a1d3724177a4e90803bff9c60", + "bf16-provider/gradients/0543.pt": "ccc434434e34eccaf5744e962d07bc25e461f57f8e1e862f2da4bd5b9eca6c16", + "bf16-provider/gradients/0542.pt": "f065482e4865b0847c0a4c677177a9db80c52fbde3315460a52b82d2b01cae32", + "bf16-provider/gradients/0541.pt": "c87bd58dc250a8d960ab75703713325c8d64d944ac80cf59cb6185625bf54e7d", + "bf16-provider/gradients/0540.pt": "9d8a6ab3d9f4f746760ea4e590fb645062a13d0da36bf49ebb90ee254aa33328", + "bf16-provider/gradients/0539.pt": "1db2f753557e53896f7a3b55594405472f91dc51389e3fc77f1f74dadd9a2ffe", + "bf16-provider/gradients/0538.pt": "a37e2b4389b788567f0748fa937358f1d92decdcbbee2aa30f5be376081f227d", + "bf16-provider/gradients/0537.pt": "63e82dcd38eb8f7e9bfd752299bd89642ee4650111b083b048127d8c276dbd2f", + "bf16-provider/gradients/0536.pt": "178f0d0b39cecb680d755ebba612815be2bca21b2b31adf43ab52a3ba8e9e382", + "bf16-provider/gradients/0535.pt": "b1cbc58016470fb40996db3f6ee8e4d49c6872edf3cf1c0ec748d1d9986385dc", + "bf16-provider/gradients/0546.pt": "8aba3a6f188e499801e293ea99e6a24872a0b3da322a0f2686472fbadb97bcd0", + "bf16-provider/gradients/0545.pt": "13cbdc90c566f913377702a0579a5ba75c5517ce23a68e3120fb9d5a4ccacb3e", + "bf16-provider/gradients/0532.pt": "2c615efd2923d609bf1cf6f77d44dcb3e351cc7dbf0283fe843be0198c04c77c", + "bf16-provider/gradients/0531.pt": "6e045c469beefbc41e97702d42c908ccb777c3682e278ade0ba47bef890580c6", + "bf16-provider/gradients/0530.pt": "b703f285f960fcf2f4e0984ee64f3119ce26820ce9204db657d0c21043ec4fd0", + "bf16-provider/gradients/0529.pt": "eaef32f9eeb96ba92f46432c8114feb21dd6d1d4fc3bde96d6300d84bd8665ed", + "bf16-provider/gradients/0528.pt": "e0592d062f6a01396ccb01cc53deed7a84edafa914af63241a8d06c16cc3b60d", + "bf16-provider/gradients/0527.pt": "8e0ff55072cf1798c0963f4a00b0b9ec25e20052cc8cd79b3ed633c826e24d96", + "bf16-provider/gradients/0526.pt": "18d9af8ebcd68e06a3dc9a0cea36ca6d6efd2231d7289a3e385e0fbe88d4cb22", + "bf16-provider/gradients/0525.pt": "e441d5ec143c790b19762e2f217028ab7266b8d8b92864e61c07749499a19158", + "bf16-provider/gradients/0524.pt": "3ae5eaca1a440c8373b5b61584644db4d650c30a024d446c06d631f310ef19f1", + "bf16-provider/gradients/0523.pt": "92f7e448dd792f3cf7ae0b006ccb190edc0310fd08028cd388f194ae2f98f9a2", + "bf16-provider/gradients/0534.pt": "b0b7832cfecec6c81a0f1910d5141a8a8b5b13bef75599abdf588ebd80bc39df", + "bf16-provider/gradients/0533.pt": "5cb82cd6cd1f7709dd1402d3029f28703e7eb599dec844e8e3b422bbbfbc3c04", + "bf16-provider/gradients/0520.pt": "dc6193ebf4b7e1f9acf742f6c049b44914dee467f09b4ec704de28547f773962", + "bf16-provider/gradients/0519.pt": "fde3ad4471b0504fe2fbf6795e385394831e9eb98822ecc390922e8f3cc43b1b", + "bf16-provider/gradients/0518.pt": "306f556e6b7df14938884265422e149b638a6ed52888432e9c7a598ef35f7a97", + "bf16-provider/gradients/0517.pt": "0e6ffec908b0b43cada76470fea711bddac5c6104a3db49412e357eec175769c", + "bf16-provider/gradients/0516.pt": "835066eb37267c91e1fa65874e5700976430603a6e7a8fbe99699c402642196a", + "bf16-provider/gradients/0515.pt": "90043d814566b6a61e2ea7888f15d681101165a8833a51f503d58a2a6d90ca1b", + "bf16-provider/gradients/0514.pt": "f3f633bdeabcf26df4bb70dcde9f5e10fc50c4479f190cf09ceb4395e569e67e", + "bf16-provider/gradients/0513.pt": "211944e3666590a18abef957222a5999fe7c9895cde6f34d02ba5d1a20f0c8e9", + "bf16-provider/gradients/0512.pt": "f7d95606e6c293746562cda0a7697ccc9a73b20cb2d2062681ae73802730513d", + "bf16-provider/gradients/0511.pt": "7dcbb1587b8a88bdf4df3016df63b7db60ec82e350c01d8d1b60324b76f879f2", + "bf16-provider/gradients/0522.pt": "3c623cfae38c555d279e4ebc55dce1580241c46d9fd80e7d79b9a7f3bb3aceba", + "bf16-provider/gradients/0521.pt": "c7fef0cd677c2c2a5c0906590d3fe37f28459f1b1a9124ae8813a60260523b39", + "bf16-provider/gradients/0508.pt": "b9962e684b42a46ea85472f2677b753424c202c4fc47dfc402fb7bedb36db7c1", + "bf16-provider/gradients/0507.pt": "a4ab90249ece77aaf07ba8ada94c057904c1892d54873aa3461aed852724c6f7", + "bf16-provider/gradients/0506.pt": "cc059e7e82af46b0fb5fca3377ac0ec774411bf8fe79224d86350023b49fcf8d", + "bf16-provider/gradients/0505.pt": "35e5ad1f88d24fbda6cb315a64c991646c22a9101cd84bee0866b730e25261cf", + "bf16-provider/gradients/0504.pt": "01f2cf395c8e0a322301ab737620f025ff4c51241497bfb5b156f48b39fd0f0b", + "bf16-provider/gradients/0503.pt": "77cff9445ba8f1019aa7b9828f7b8d38c3bcdb1225fa28c06b296f822c29e64f", + "bf16-provider/gradients/0502.pt": "7130564a5ed1b02ea68e72aee594c2bff4302db8ab5e0f45fbaad2b75c19c9bf", + "bf16-provider/gradients/0501.pt": "8ff81f867922902edb9453e61fd21d7896b12c63b1b226c691bdfd7171b0f9b3", + "bf16-provider/gradients/0500.pt": "94263e7b0e65c6b51ee84fe45c61873723e89d320b08d5a4e1ca78af04e04756", + "bf16-provider/gradients/0499.pt": "663aaff4d99ec8b2f58085cc3bd5cf6f4062467b980aabdb0099d0c04bad9b07", + "bf16-provider/gradients/0510.pt": "5658660b4191fe7a0b7b9ae7c3ecc5ceef8b0576be6a83c270785dd129094a46", + "bf16-provider/gradients/0509.pt": "817adb51b2e64d38622d83fd59e360243862be28863311a019792583f3f7f5c2", + "bf16-provider/gradients/0496.pt": "1ae60be582155b48fe58a92a2a9a1e25c1c97b0d58b07be1334b70882795f0af", + "bf16-provider/gradients/0495.pt": "b7f691c0beae28cc99ecb050dae1154ca27d5a83ebf759bdf27b785e69593514", + "bf16-provider/gradients/0494.pt": "2f441517cc9479f5cf538ac4e0dcc86f37df201b80e9bdeab69b1daa7a64b615", + "bf16-provider/gradients/0493.pt": "f20240e0b40e95b910eb4a262948fe8ca21d4d2dab6f0236d50662619d94dcd0", + "bf16-provider/gradients/0492.pt": "f56d7184fa1fc74eab5e8ffe910151b7df772170d0beba7ff099ee9e4e6cd9d0", + "bf16-provider/gradients/0491.pt": "03ebd4d2a58ebdbc22ebf2e6bdcdee6533c256c7c3c1184cdb0e8a2fef1feeb6", + "bf16-provider/gradients/0490.pt": "4d796d39f7b8bb74e311c777d5ba00ba259c4556fb8fb0d54e917e6901cdcd2a", + "bf16-provider/gradients/0489.pt": "84d87af311c8b91664283ab5b2aae4299313b999d188df6798b309a9cd1c1fae", + "bf16-provider/gradients/0488.pt": "78b4e4da3bb7a0b2b0c68d203f70631143beb55af9e38e55eb73731426dc1915", + "bf16-provider/gradients/0487.pt": "6c52d0b39d7d7d8521743459744aac609ad2310464333b7635e2307bf96bbce0", + "bf16-provider/gradients/0498.pt": "ac4a9f164643ddb499717218c034112daf3419c7d5355fd160b8522843c9eb6b", + "bf16-provider/gradients/0497.pt": "ed7fd638cc5f89b9d794f9851ac983683ba48f54d1bd4a6fca55c09e93d10fbe", + "bf16-provider/gradients/0484.pt": "d677b348266831ebbc4c523e3d84ec83885bb851a00e266b0a51fa2ab5493621", + "bf16-provider/gradients/0483.pt": "22a12b0ef8de6043639f920c524e8c930d2a02b1d20451109a8b6fca0a919529", + "bf16-provider/gradients/0482.pt": "c5bced8939f0d1005c917c2d4d3f1e42eb115b4cf95f6e9ec086e9103f5d3a37", + "bf16-provider/gradients/0481.pt": "9c6cb5e1b11bf8f01a34aa7955cbd31b0b02d4eabc29250f11edb94eff21f0fc", + "bf16-provider/gradients/0480.pt": "5da555815a3fe105b85057c507eb8ff2498f3102d13028f65969daf0ffa90195", + "bf16-provider/gradients/0479.pt": "c59623111a24369ef542337593b4244453ec11121f932831d845a09fa10767eb", + "bf16-provider/gradients/0478.pt": "51b3a8f563657f5fb55bec5774d9e3b431a6fca3be71a95c789e0838ad4bd431", + "bf16-provider/gradients/0477.pt": "891d737107ca018bb54c89fb5c919af92709baf5e55281108cb6d1b62f5b577e", + "bf16-provider/gradients/0476.pt": "03c78db510bb19dae85acb75d11e5f5ad65d6e54d7aa6ede0c5f5ae1e598ff29", + "bf16-provider/gradients/0475.pt": "bf623529e4399a3b5f45d58fbaa3e994e101e35d4e7ef4bf174d11038216eb03", + "bf16-provider/gradients/0486.pt": "844838a0d4a67972838a0fa63c0d19c52924a10bcf7e64a6a16631b501325c84", + "bf16-provider/gradients/0485.pt": "d34cf4a12d9f02908d36a7f7775411de7e3a1e3e6539ee98c6cf88e6b812078d", + "bf16-provider/gradients/0472.pt": "cffca0566f6916795e9a852648b53c5a12ecc297141bb2d28dbcb1cea1f94140", + "bf16-provider/gradients/0471.pt": "47789d466380fa6c27a12fbed584875b4adfbecdb92be3ab8bf86996097d0fe8", + "bf16-provider/gradients/0470.pt": "335519ee69f6d2b2732de1f686500bb6cf1c0e1989f0b13e6f05aee7107f84dc", + "bf16-provider/gradients/0469.pt": "852f5894d5bd0e2e876b2dbcbc818348aa4b8828cbecf8992bedc826460d56ba", + "bf16-provider/gradients/0468.pt": "3b8e415537faefd522f822b0def895db00fda3fa56eb1a9e099b084ceb8a76cd", + "bf16-provider/gradients/0467.pt": "00de339693e9e42a50d2d68053f506fb74dd2edd0cbcb5e36c600c05386516f7", + "bf16-provider/gradients/0466.pt": "7c4d30defcd08033f135e5b18b8755561b05cf5d483b96459a5cbe571245a065", + "bf16-provider/gradients/0465.pt": "3b78037be58521e544113886f171d23c1373f2737a67045e19718f3d2419199e", + "bf16-provider/gradients/0464.pt": "b051ca08bef62117ec3460c5401114293d18e85e7425bcf4865998c3846105aa", + "bf16-provider/gradients/0463.pt": "05e46077a1beaca7d3424fdc71b7caa74f7819351b07c567a994374a4b01bc13", + "bf16-provider/gradients/0474.pt": "71be156216ba221a97e1423a03217a7a7167c88f150f603efd2d8c3a11f38503", + "bf16-provider/gradients/0473.pt": "bf55e7d09d5c5b856d6feca13155c7c4e82ecca6a4520856365e706ab00f0d5b", + "bf16-provider/gradients/0460.pt": "7fabc4c35ccc0907d2ecd5b49a57c1efb64cfef04c8828dc96de9c22529cbde9", + "bf16-provider/gradients/0459.pt": "459398bfa18a8c7021ea269175b9d639fdf138ee51dd84cc9d1101f52bbb6298", + "bf16-provider/gradients/0458.pt": "c1c19be6ddf5f7d84700f5bb13f9154b39cfd949f51a730d946bb66a8677c2bd", + "bf16-provider/gradients/0457.pt": "30dd1380984026f13f99f8d9a3df21ceeba1b3dc526684db5fd5acee94c64715", + "bf16-provider/gradients/0456.pt": "830b6ce8c09e1e0d1d51acd72d2ad4ad8595d23072ceae4efc3970bed5570cb3", + "bf16-provider/gradients/0455.pt": "d243fbca25f5f6aa29ad48552f38fec86f8a01c8534eae18f5fefa3187210d99", + "bf16-provider/gradients/0454.pt": "8d6507b3747f7224ea8d168b034911ba1a7f1837a52fc070faa696186b8c9c02", + "bf16-provider/gradients/0453.pt": "4f1524673d3334f5b0884719e71620c95df2719d8d68e990bb8fbe89c99827fe", + "bf16-provider/gradients/0452.pt": "81a00bb1487bdc2dd595e0e3f7ee75e7d9148eb943f566d7bf3777c97c5e915b", + "bf16-provider/gradients/0451.pt": "d31dd660f1c1e8ae3fe121978d12fe12d006a3fe4ea06b43b0ca948c43c56261", + "bf16-provider/gradients/0462.pt": "99cee9ccb61645d1dbdfb9abdc91a677fda2c0c9d67c654f4ee265e6d3df614d", + "bf16-provider/gradients/0461.pt": "0c0bcf12ef370674906bdf31844ee8cbc460516270da6b924faa8eb22354a7b4", + "bf16-provider/gradients/0448.pt": "d669605d10e1954f66a509a33d4cbe8f01820a0aae89accb562a946699bd0a43", + "bf16-provider/gradients/0447.pt": "f0991a36dc1b3f2de055077396835da8cd918b592291f189b11b21968fedfe6f", + "bf16-provider/gradients/0446.pt": "8fee8797705a82f026c1d4369af14836a07113a3000396c9afdec448c149e539", + "bf16-provider/gradients/0445.pt": "fead9a7e8c04fc2624058097546e874eb151025d806206c7b8c0b78583111fa0", + "bf16-provider/gradients/0444.pt": "f94353bd4cdfac7d0d14be7193cf401a5bf920d0790deaa7841459878614bfff", + "bf16-provider/gradients/0443.pt": "a9e90619d6e5521aeb2334f3f6be6c803274580eec082d91c209219b5d5d51b3", + "bf16-provider/gradients/0442.pt": "8e21e37aaf689d5ac7a935c6a00b271ae54ac323d891970c8f0c7e2f3503b6b8", + "bf16-provider/gradients/0441.pt": "5a0ab2c9b9163f3a568e2bfdffa89227dcd6986f41f295cc03a570f24e748d41", + "bf16-provider/gradients/0440.pt": "36a3c16a220274d2bc149cf95b75770ffb0660b21f893dae6c75e20d2ab5c9d5", + "bf16-provider/gradients/0439.pt": "b171415dc97e7661edf1ad12a2674373d5dd9f12f1a7ae45c6fb2c7143464d1e", + "bf16-provider/gradients/0450.pt": "8e6f07c28eae7a00bc37f8d6f9d92a8c0b06d6de8231b0361bf4bb7b3d9bdee5", + "bf16-provider/gradients/0449.pt": "7b2a0f5393f7c4b6342495cae101f5d0fa69194bf5c190bed72bf28fc75003e6", + "bf16-provider/gradients/0436.pt": "834ec333f78bd28821895823aca1b30ed8fe08da07b718b05dce0ac5e6dc8bf0", + "bf16-provider/gradients/0435.pt": "34cb3acae0d595bca5bdafa674db93e07af5f2b622d1d522c51891886cc6baca", + "bf16-provider/gradients/0434.pt": "73068876b071591f76196e1c2c8ec8056a40b8c76e47a1791cfbb6a164ddecf8", + "bf16-provider/gradients/0433.pt": "dac1afd8bcd9a2aff4726ab8ae02fc31a549f990e957584b2f1f3f943f355c67", + "bf16-provider/gradients/0432.pt": "eda8de51540f5547b3bf15019803b443f0c53a34bfae2c69c7895bf0b603fb72", + "bf16-provider/gradients/0431.pt": "9595b4baac45a409602faa397f33dcc17748636ce674617fb5fc4e7334c1bce9", + "bf16-provider/gradients/0430.pt": "940334deb170ddfed4b1a31e8c229628aa009f745d708ab6104bea92c8fc1e0f", + "bf16-provider/gradients/0429.pt": "e4d4a1c4cc701cd55c69c00c508ac3bf76925cf2b31d66d4e07a5ee929485a97", + "bf16-provider/gradients/0428.pt": "14d3e7302550ae94a58643aed0a2511a3bfa63daecb3cede0d62eaa682a8377c", + "bf16-provider/gradients/0427.pt": "3d104208990f1ea302885f316d2b08c2774a3d5e028ddae08c94eda048d88f32", + "bf16-provider/gradients/0438.pt": "1baed369a2ce360a0de3eb5863270a5f257117221db024f0f98932a86fc26405", + "bf16-provider/gradients/0437.pt": "d6dc7e416eb2fdb75eaac86bd4f1c8b78efd500cfafae9a8f065429d56ad530b", + "bf16-provider/gradients/0424.pt": "1b898b8bc3b14d42334488f9df4165690876762702441029edf32b6e753d9035", + "bf16-provider/gradients/0423.pt": "f87f47ba3ad6649007e9020f8706aa331b98513294ea61f62b3420f382d7aff1", + "bf16-provider/gradients/0422.pt": "10288c7ee2c446c385058a64c8c095fdf56bda47f389999a513c382a63748e12", + "bf16-provider/gradients/0421.pt": "50b3c863396334d9f8ea937c0c716052a4edefaa2bfd31f70e314d61b5d9e04a", + "bf16-provider/gradients/0420.pt": "e00669d8f3e71b772e5a0f7fa8911c7cda6e59195f36df868659c5e2bf1bfd8f", + "bf16-provider/gradients/0419.pt": "2c2f4bead308123ce434bdb8ba3ad24a06e4b45770b50066885d02484b470006", + "bf16-provider/gradients/0418.pt": "dc69f98edaa1b6b26ea2c0df259f60f097129f3cba6ff943ea8a8277ff3d05b7", + "bf16-provider/gradients/0417.pt": "e723311caaf808dc9c0459b41f94618b7c2e2bcf806cd1691a14875eb239f41b", + "bf16-provider/gradients/0416.pt": "68d79df4f2a59dbb98ea8758ea1fe1a5b9a649909531df08e3ef5b19bee21790", + "bf16-provider/gradients/0415.pt": "c657bfb1ec98bd4b6e56e8a8e6dae63e716cf4b91688bde1f39066ca70222c64", + "bf16-provider/gradients/0426.pt": "7cf0e8f7f7b12fbc639cef97ace25ca674d93f1d88d4bf2597bf64653679f483", + "bf16-provider/gradients/0425.pt": "0500fb0d766b281637ba24e2050de2e3fc10601d564c674a8640dfe3c6f373fc", + "bf16-provider/gradients/0412.pt": "4bfc5c032d26964efa3fd2a84e1d8968b769d40eec8bbe9718d392855c626f71", + "bf16-provider/gradients/0411.pt": "2f4ae0f3c0e4096397e7786b3ea0259099187206e5c1b7b1b6b3c5fa86b25d69", + "bf16-provider/gradients/0410.pt": "42f419cbd8544eda7be7fa2f667097d8bd88e006538e9d67d1bce3961b11225a", + "bf16-provider/gradients/0409.pt": "eef39b3b6b1d834cf0e506ca09c0de4f3843fc891e2995ebf09174704b4371dd", + "bf16-provider/gradients/0408.pt": "b3adae34f2641f63b9d9429a286631d22ce93a8b119572b47594f3f04439e6ab", + "bf16-provider/gradients/0407.pt": "ce7cf574a586123e9cc0f8d0117e409b1411663269383a346af027ce19db29cf", + "bf16-provider/gradients/0406.pt": "a499999820e53a0f239336712c66ee76ead36582c149c9d3d6f905f65f6653ec", + "bf16-provider/gradients/0405.pt": "3111bb92fd28d64f382b4aa0767397cd2514d5602b0a2de635f47009b7133861", + "bf16-provider/gradients/0404.pt": "e751f5cc45343da810a307d273962c735c8edbe8d28ec19b1957e80c9c392c96", + "bf16-provider/gradients/0403.pt": "6336ce3fbc6f62af1a59f40650229c6a72791e15c6a78c730a5445029e9edc23", + "bf16-provider/gradients/0414.pt": "983932ac4f61e0a7c2a592a58ddeea75dea31f7aba0ced0b922e213064e72ed3", + "bf16-provider/gradients/0413.pt": "3c3e33a6f32818988c23262eed7ef1a226d69f5a5f68ef5c43582bae43f99420", + "bf16-provider/gradients/0400.pt": "6e20fa0e1e3756137131e6744ffd2d8ee67d285accc58638ed14b56f641daa2b", + "bf16-provider/gradients/0399.pt": "199826583b91b8a1e4ab99510e120815e327d477876a48cd06436f80568061ab", + "bf16-provider/gradients/0398.pt": "ad1c0c75231bb2ccf4dddbf7252ba64c7f85011718730af27f15ab91866f42ae", + "bf16-provider/gradients/0397.pt": "6126e25b6ba7c4f39aa4169f82efb2edd8866a0f7ebab22901fe4806280bc99d", + "bf16-provider/gradients/0396.pt": "b9401ae1183b64c57979e10df3fd738439960d17687bd10fcfbe23fd74c72d6d", + "bf16-provider/gradients/0395.pt": "00cb1f8b781f7df14ccbeaccdfe13d0f184258680d09e5b121ed4f62d4b50b43", + "bf16-provider/gradients/0394.pt": "7c41ba4078643c2c007c57a88417a1d82ba8159083c979b518f4ff440c48ce05", + "bf16-provider/gradients/0393.pt": "dee0364f82594a2d9d3880ea80a46c85322f38155f6b45b6fa46b4a8728da63c", + "bf16-provider/gradients/0392.pt": "7d2210cf08a9090efe94c510c1ea16cb1f547511f5dad337d55cbdfa0c72d10d", + "bf16-provider/gradients/0391.pt": "a8ad0998896ee8ffe5972b3f214746319e4e52733984e8e1df0424bf5ee799a4", + "bf16-provider/gradients/0402.pt": "3f8dbe7fadc5ea145906611363bf22173a47368175f76ab266384018ea60017a", + "bf16-provider/gradients/0401.pt": "de62139743fb85546f4a2c5a71e9c1ecdeb53c48a49b7d97bc257748623203e4", + "bf16-provider/gradients/0388.pt": "c01a779947fb30a631e357fd880decdd520c4caa90f31473d4b58b546b596eef", + "bf16-provider/gradients/0387.pt": "1af9a5d52305c25c7ef28239e6b73cc6b69aed344f9c2c677d8580c189405401", + "bf16-provider/gradients/0386.pt": "83f94e71e1ac221cb1e56132ade7704b2856cf3d126227a8b217c0300e5144ca", + "bf16-provider/gradients/0385.pt": "93daf48ecb6e04cf0dd96764571546a3d8b0cbc08eebc39978a87dbbc360eb3f", + "bf16-provider/gradients/0384.pt": "4c889512e4ba10b671cfd286a629b3d19f965584a50b7f5916593109eb05dd2c", + "bf16-provider/gradients/0383.pt": "83a4d89b8d6fbdd88a4f3224565faf09180735e7da060fd6a44395b78b4131f4", + "bf16-provider/gradients/0382.pt": "736fdc79e54f720023033cf411b5a56eed176d259cda852d6fac8e4fe0070962", + "bf16-provider/gradients/0381.pt": "e1650693f27a890330a91b847ff669f2845d962df460ff9f18e7706b0b8f0ced", + "bf16-provider/gradients/0380.pt": "6af3e26fcaf59b4ca8d600dffdded8191cfa7d31773b4bad11dd689d8fe1891c", + "bf16-provider/gradients/0379.pt": "7a9026ef060cb6d7ef16718bac6f3aab381f38bb4afb1bed63cf75213cd45ae5", + "bf16-provider/gradients/0390.pt": "9d01f008020aa5f7bab6bd45d2321186ae97ecb3f747f21c38d56d13cd08563a", + "bf16-provider/gradients/0389.pt": "fed763e3838a95b12a4a4683b1a863eaf4f8a17e1ca7ff596b8f1f1824e80fb8", + "bf16-provider/gradients/0376.pt": "80d991324b2a369e92703b4d330cf9f95a2625c05a6db6e2a1b33cad19ef6cc8", + "bf16-provider/gradients/0375.pt": "9a2548d35c9168c4b76fa30fae9a755e8b2401130261a645fab732f9ba58b631", + "bf16-provider/gradients/0374.pt": "366e7a501ca2023befd3a1f1a6599a72f59f49f9af9cd52895a9e1d245b67d47", + "bf16-provider/gradients/0373.pt": "2f7113d23b7e40fd5f5c3940bb0b54b8caa829e8b278ea720ba99ccbd4533640", + "bf16-provider/gradients/0372.pt": "fb22006855623f04e5823a46ce0ba4e0ac7ae939bfe5e81efbb25a812ec196c1", + "bf16-provider/gradients/0371.pt": "c04413582df24961b0f65b52f7a0480f922105e877c465f574dd9ec4bdb66df8", + "bf16-provider/gradients/0370.pt": "df53702eb87cd51d31af9d714012b7fdfc031b0b6f334903c9d02507324e5db5", + "bf16-provider/gradients/0369.pt": "5cd145903aeaddfd0523565eabc48ed2ec3a9f9e2de099ca02d3085d76387672", + "bf16-provider/gradients/0368.pt": "442b10d48479cd3bb6637e77028c7f37275c0498e0b12c657461af6f9e479dfd", + "bf16-provider/gradients/0367.pt": "2820a3ddd747ee37c3f92ef04921cfb014f1bbcafbd47304f48aa4cd64557401", + "bf16-provider/gradients/0378.pt": "7bc6ac4c284eef24a85062aeab27d2db52d9f52294f8249473186a2b9619b406", + "bf16-provider/gradients/0377.pt": "0fb26d3d09283c9c93ee5ff99e8d015058afa2759f59678297a05465d586d941", + "bf16-provider/gradients/0364.pt": "7086c5889af19c994a10c9fee802da9a0d3505be3ca936a39bc819379635b3c1", + "bf16-provider/gradients/0363.pt": "b10767c3b65bd0440f5374265e67bdf2998868413214e1d5348b7c7a6496af7b", + "bf16-provider/gradients/0362.pt": "7540995f2696f4fd1fd660595c81d78d2ca64ae36553ac4eb0cac167c3f3fd36", + "bf16-provider/gradients/0361.pt": "cfd05b67644a4c309d6a72c22f399f36f977156a150d66ffa6e2586d9a49cbd9", + "bf16-provider/gradients/0360.pt": "c587df6077f08da4b7e0619960cca589f5248db0c32bb82cdb781602cebc59c5", + "bf16-provider/gradients/0359.pt": "92b58b50957eaa4c9bde642ae32635a31f1fab21f5979b1a4b98ba09edbde855", + "bf16-provider/gradients/0358.pt": "10c4b56975dd96a39c64a0fc1296f89e0c2ba0b63d363bdb65d333eed53648c7", + "bf16-provider/gradients/0357.pt": "16f230ee1e7ea6f01259b29c1747d9126043a8bf5939d8a44dacd673cd2ef072", + "bf16-provider/gradients/0356.pt": "e8ec4e58e9d1df18453191536147588efc7a9d7734c81cb748997e6234ee26a5", + "bf16-provider/gradients/0355.pt": "215feb49ffa57c65ead59544bf40b64cddbc2b408f5650684218810e9f9428ab", + "bf16-provider/gradients/0366.pt": "8dc4a8cb1ef2ef03b7261b9057a2d1b3415925b8e102dd713d7952324bac9906", + "bf16-provider/gradients/0365.pt": "048dfe529664b2f752583875838bf01c6382bd1356ba5c6ad4181b533aba765c", + "bf16-provider/gradients/0352.pt": "73b885e92b77e1a87a2ef06a0b9ba5b6106294192b8bfa925fd1a7ceb0b0bd46", + "bf16-provider/gradients/0351.pt": "8ada9649ccd6e091746424be7c0e9f3c4d7f4b5de9d0bf1624d6c94616638c61", + "bf16-provider/gradients/0350.pt": "b897d88524957c46b4b5c829385b1ded1c97a5c60995b1de494e047fa542c019", + "bf16-provider/gradients/0349.pt": "d4f9dda2f1686a3daf3cb6303d07bf76b14f2d162a6ec68edd2826c26bfffab3", + "bf16-provider/gradients/0348.pt": "114a6555be1fea1a2411da97bd39ee547da87ff8401f7c884b38111d285769af", + "bf16-provider/gradients/0347.pt": "d99e95c384ac1e1a79269ab365e7a641fb51cec7e53f48f19023b591162ed321", + "bf16-provider/gradients/0346.pt": "c8196dfd759a551f9073919e1184d9ae2c1e7df7f679b53ff72e5fb323828c90", + "bf16-provider/gradients/0345.pt": "061fbb5691890cb52bb3d4c4d523c30f19441631617f27f52a2156ab1e713d41", + "bf16-provider/gradients/0344.pt": "bc7afb0a9f9b5bf678de90d4e32760c9dd2133b2ff7665542f68bd10c73ca3ac", + "bf16-provider/gradients/0343.pt": "c89979544fadc5881e4db713ef85ba71fa97f3875b750ab6cfa1f61d6c8cbda5", + "bf16-provider/gradients/0354.pt": "9d28f231db5552c04e92baeeed5c1039326d284d896354b39151fb13059d9acd", + "bf16-provider/gradients/0353.pt": "2cdfa89c6fa01d30190c3561c3dcfa1a99cf345d252352b3d36a1d85d3b40f7c", + "bf16-provider/gradients/0340.pt": "16e736c422313236fa0687931f559902adbbc030a71a6a9265f0e0d48a744ca7", + "bf16-provider/gradients/0339.pt": "c44c4f8a6e6b690c9a676688a102a57a557d6b0fadaa2e5ab3914c37831d514e", + "bf16-provider/gradients/0338.pt": "432eb47149d2699aa0938b44e8311cf02bd578b2e4e3d9026e1b1ccecd078af7", + "bf16-provider/gradients/0337.pt": "9d2996d3e10e01646574bb35f9e4ea39d32041dda99842ba88443f6b1e90a187", + "bf16-provider/gradients/0336.pt": "5ac82f7000cebb3952880b1624f7c822661dd1e80e36272205c130eaf2c7de9a", + "bf16-provider/gradients/0335.pt": "5f96493dd608cb99d6d6b55638f872230e4074eefcaeb2ad19e1f74ccd52f775", + "bf16-provider/gradients/0334.pt": "ab438df3ff6b41104fbb062d14be1d32db9acd9bcce90586eff6d28af0952db7", + "bf16-provider/gradients/0333.pt": "c0b604625ceea9e2be57f6712581231b1c8a7fb19a7466a57203e1ef9b55fc95", + "bf16-provider/gradients/0332.pt": "89fa2ab19f9530652a5589e5dec63dbc5658d6292728abfceb6e971a7361cd68", + "bf16-provider/gradients/0331.pt": "a28d779aa3c7ea69749cfe92f72150ab4253b3620a25fd26157bc50eba51249f", + "bf16-provider/gradients/0342.pt": "dbe29d903f6f33fb03d6e57fc57066a77783484ac332a7001aec997aac07f2d3", + "bf16-provider/gradients/0341.pt": "15db439a8cc038776a700c98ff1af3dd5535d5b9a81f1f14d33102c3697266d3", + "bf16-provider/gradients/0328.pt": "9c2096499920a1f2ab43be7737404f8c7ca77df77ee8da8c4d175ffe79b1b8e7", + "bf16-provider/gradients/0327.pt": "6c209b6e47c1dfec4b212471dcc803ae83376d08853117a860d4d4800481ee4d", + "bf16-provider/gradients/0326.pt": "9c65c872e4526ca10382a1a58cf2f917aaf21e12ab9e23ac78a6f711b1c83c79", + "bf16-provider/gradients/0325.pt": "7f2fca93366b019fa9cf2399d3614722bf4026d2c62be5998cfa4f4ae8436e71", + "bf16-provider/gradients/0324.pt": "a383beb4c6dc0a9b81668c4ebd111b5a6e527c576253ccbbcc574967277fb386", + "bf16-provider/gradients/0323.pt": "bd611186afd49b0f5ba6c530c7c2307f30bd03e0a1fbfeeb4d15ad01818a6c27", + "bf16-provider/gradients/0322.pt": "9f7a464f463566fc87aa8800106925523a6e651f42595abc6969936574bb9845", + "bf16-provider/gradients/0321.pt": "07f5921da45d7d5a3d10203bb49465d56b92106aa9799baa0bf5f30947a287cf", + "bf16-provider/gradients/0320.pt": "dab1df649ff666c8e37eee09779a8b39eefd171a8f99da0102af1590bf4cb0d8", + "bf16-provider/gradients/0319.pt": "a706aca7f87492191c7c7434ee5d185764eb1f59d94ae7c4299d316aaeef51d3", + "bf16-provider/gradients/0330.pt": "134317c0c9e4b21682d825d0c35c9436e269c9517e42dff6440f73aba83c7867", + "bf16-provider/gradients/0329.pt": "46b37f67a291d05e35ededf73764f5763d68dab2ffb115ffe5d85afb8b8f7fe4", + "bf16-provider/gradients/0316.pt": "9390fd6efd83be43ee17922a778e5be65806fc0a88fd127c250d3139bf9a4eb6", + "bf16-provider/gradients/0315.pt": "d56e52357f154ae94ce733262f7284c38bbcb2cb4ac06079bc8c7008a485055a", + "bf16-provider/gradients/0314.pt": "bf43b0a5eab07b921a7c7d587f90fae7abca00158cce1e089ed0acc182ccec13", + "bf16-provider/gradients/0313.pt": "6136f702dce8910065120cd350ebf5cf3cef47ed4fe158f87d8421cfb3da669d", + "bf16-provider/gradients/0312.pt": "d9a8b18e267e976a7346bf311cfd79fa4625f091d03c104b74f0724a011231d1", + "bf16-provider/gradients/0311.pt": "5b295970ff06425a7b72ba8021c74ce779e14c49b4d3fe2be355e05d5655ff6b", + "bf16-provider/gradients/0310.pt": "0d653d6145f8ced6c4fa61a187bb1773503d188b51729343434c153d0c954877", + "bf16-provider/gradients/0309.pt": "e010c4c69dadd88e4b434a3c39f4fd87b84ae264659ce400b6343e935aadf4a0", + "bf16-provider/gradients/0308.pt": "da3c2b653c9a47bd73805db2bd693b463e8ef20beb16273b11484df3c76b19b4", + "bf16-provider/gradients/0307.pt": "5023bfd386b924ee3b89c869ef73022bc9a77fc7bdbe3d1c90a636874ddf071f", + "bf16-provider/gradients/0318.pt": "cfccb4ffc0d753db41273a3fc30c11773a5b36a3c9ad58e575a2f9f3e2d6328c", + "bf16-provider/gradients/0317.pt": "47aed8bf8daa5dcd663f8e5bdfa0ccf9942fdc886baf5b374b383adb86500462", + "bf16-provider/gradients/0304.pt": "97739c64dced790b37ae7c068e9c6d53f3e20d30e6e54c6a179a1b9f84453d71", + "bf16-provider/gradients/0303.pt": "3304aa31f19d4af1973f91e58484b35ebeed42f1ccbfa5979be754e9f2e11149", + "bf16-provider/gradients/0302.pt": "ca8c7d81f786b9ce66cb5889306f513d6e03024376bdb3d44acdacae6dcbbfeb", + "bf16-provider/gradients/0301.pt": "787675e431dc8cbea94c09decc3149fca55cb1e70293365770e2214615b47451", + "bf16-provider/gradients/0300.pt": "668ef60e15794443875978942371c0e340578c7991eceb3ee832830ec0c82130", + "bf16-provider/gradients/0299.pt": "6eef743f79d8f93aa8e010987b292b0f4a72c7d4a5203cb0c74f37da09c88ef4", + "bf16-provider/gradients/0298.pt": "233c25284295be84a203fe10833fa9801ab197a468a9b7bc1153e7aa4bc3d8ab", + "bf16-provider/gradients/0297.pt": "7f2bb4b77e5380e5caa61cb8781b66a7b68e20679497e7cc334dbd4e48af6ce5", + "bf16-provider/gradients/0296.pt": "49b573ab96c800002052bf1d470d39ab94810c9ee7a7135d7b03e9f1ce62f4bf", + "bf16-provider/gradients/0295.pt": "f88f6163f48a80473a21d129f619ab42a4d88334f1ecfa74005d1627e4057333", + "bf16-provider/gradients/0306.pt": "796f1ecc30c0895b2645a17d2820d9c0c3c437f0c1629f086869d764236758d3", + "bf16-provider/gradients/0305.pt": "fb7e07b6d951996018c55a750b4de1a7587e50a9b1e8b1ace03d381423b5166d", + "bf16-provider/gradients/0292.pt": "d774e15a385a6e2affe56fd4bf7565b8a62f90a7984dc1121cace9520140d5c2", + "bf16-provider/gradients/0291.pt": "0f295a53ea57aa75fa7607fb78bec5b95247494f29651269b54385f9f3d63f31", + "bf16-provider/gradients/0290.pt": "5e4ae29bfeb14e558df98e786d1a10bc280afa49d5f2ccdb89ded768f3f2bec8", + "bf16-provider/gradients/0289.pt": "0e50704ebdab957ad87686cf78306b329410a668bd177568538388e5cf8c9e38", + "bf16-provider/gradients/0288.pt": "3f5034461985f9f4a3d3c03d67dca680b3693b7bedc4e98abcb661baf72c9762", + "bf16-provider/gradients/0287.pt": "86c0f75adb9d77407f34f834b3123a66113a0e4433efad21dc6d644919897d02", + "bf16-provider/gradients/0286.pt": "35913b01eb81cf78be88754a54e9d3001561f33363145b6df93492f39bef5486", + "bf16-provider/gradients/0285.pt": "afc40a230f35770796db2500ae780c5f3be23bb0ae7f25f77fc0bc64cc807665", + "bf16-provider/gradients/0284.pt": "26e2bec40c56e38cc2ff769a5668677cfc914242de5d996fb799d0bf525dbbb7", + "bf16-provider/gradients/0283.pt": "4417139e2f7fa76dbc6bd9234fab69923b8ec6dfb2ab8006e2d4e97d5c0225eb", + "bf16-provider/gradients/0294.pt": "808a1412bd7e2daef88ea1c267c0135b823be7e1038336a884721c56cabc4bb6", + "bf16-provider/gradients/0293.pt": "76ba8b80c0d1b68afd1d6b058511a7b61d2672d1fb1f265d53d1d750f43d64d4", + "bf16-provider/gradients/0280.pt": "dc6b8ea09a1e08f2ced13ff7112a2ed17c3c137b50a6cf8dbbd962eb0650c421", + "bf16-provider/gradients/0279.pt": "3f29a1ebd7a98e252bad3ff33801d197d9e955445217e878cd5dbe032336cb90", + "bf16-provider/gradients/0278.pt": "9f0d2d5267f3011ac21138c731cbc0099ba70eccb81d11a5c8e2c45e73021c4c", + "bf16-provider/gradients/0277.pt": "8d9315411992cf348e28707747169acf8084122ef7e78839a51aa615ab33dcde", + "bf16-provider/gradients/0276.pt": "2bc86872c0427eef27c8d41828ed65185cd20d43c97faa02faa276b435112a6d", + "bf16-provider/gradients/0275.pt": "d18373120bfc7561ba43364cc7ef2952c1412a934b11bf19aaa43446b6319b6e", + "bf16-provider/gradients/0274.pt": "d03638a9a7cdc8097cdb71710c5f6fde6a4cb0c72e9f904ed29bcc3be379b993", + "bf16-provider/gradients/0273.pt": "d2a9a7ea8eaffffe8c698ca8eec6c4e3ca404eb00ec77d808fc3834c62c3b331", + "bf16-provider/gradients/0272.pt": "6db399f8af1c9925068b79ddd174c9a3a8707db1f6d33b6c89e53ff2d43d673f", + "bf16-provider/gradients/0271.pt": "a186a972fad6db0a157ccb39a88f0abc5926f4129b582090077aafd84f22d74f", + "bf16-provider/gradients/0282.pt": "9ee552757a1496ec847ba7c77445854cf51d5b857f5a869e7c779883e0154605", + "bf16-provider/gradients/0281.pt": "0d0252a452b5a900e0e1cb3211fc88b44f2b9f6aa38b183d55e19a37f3a2b9ad", + "bf16-provider/gradients/0268.pt": "eda4537d480a47f9b8c854ad5da8c51c14e11e9ef79445db161bb87ead56034c", + "bf16-provider/gradients/0267.pt": "a543f6e2eef7495a8ed588339eabc7df314634ab816e08e8b1e430a9dcba66e5", + "bf16-provider/gradients/0266.pt": "e759cff02bd6196fe5de0c1028b9ad2668d467354bdb16cd416708cdb15ae185", + "bf16-provider/gradients/0265.pt": "042bb32b5a2a98e22fb7f4e7c9b27867c70d2eedb1198a637dc703c5cb3703f0", + "bf16-provider/gradients/0264.pt": "35b9ad420c48646419a247e113f0319c3b5b28dd85ce64c009b83943e254db70", + "bf16-provider/gradients/0263.pt": "c7964735b2b3bc3b030a62306d2f99c3ffa91487c92a8fee0dcc4832a7173faa", + "bf16-provider/gradients/0262.pt": "f9ff815ad5e7325b80892fb584703cd6e8735e17efce777cff2e1144fc65bbdf", + "bf16-provider/gradients/0261.pt": "4ad2f2b3328ddfa1828d3ebd55a5f3605231fe5b919e0c5d17594aa0151e5210", + "bf16-provider/gradients/0260.pt": "53f7d3922bbea80523b20f065de6c4d58161a69404bda4010828e119bc14c5a3", + "bf16-provider/gradients/0259.pt": "deb98987e5cdc89e76beb9fc92c52409bf8beb9e3fd66f29b5c557ff944a898c", + "bf16-provider/gradients/0270.pt": "331af297b425fca0d9f996e2afe83661f98533818922f3bc7ef35ecd90c7458f", + "bf16-provider/gradients/0269.pt": "21e1849f0dc4290caa005a924f36dee7aa4a5ca43718f7b2aafa8fd1652f976e", + "bf16-provider/gradients/0256.pt": "d9a3f3ad49d2ec14f3e1c98939daac81d748bd5ce3223209fd0f79c50e4ccb40", + "bf16-provider/gradients/0255.pt": "8827e510654d2727d8d9a24316e2332c30c43d31db69dd0b988df016fcfc5824", + "bf16-provider/gradients/0254.pt": "0e3bfcea4d4857a0affab4319659a6de26be790c102e76a1c142de78a8e877b3", + "bf16-provider/gradients/0253.pt": "e182a9865b8f41da2e50b5c625a502aa4b69afd3a7f3e027f90bc7ac22b6b95d", + "bf16-provider/gradients/0252.pt": "64044c1e4986c3f95008f294d1dd807b27e1aece6ed48745422c13cbd3de32fc", + "bf16-provider/gradients/0251.pt": "995d73480e764bfeca642a3e98aad010e02aeb1f382bb6a43b7968f23becd61d", + "bf16-provider/gradients/0250.pt": "e4d24a090de0e3b44a7dc0a0e7b37eb776cf600ad4f7fbf47f0e56a31ca10e92", + "bf16-provider/gradients/0249.pt": "2b47350687d676b10fe5b477f8448b396276b40618645b866c5292ea907d28d9", + "bf16-provider/gradients/0248.pt": "3eee607af2409ff8291636a53d8653616c2155633d69f07f35a72192721bfb88", + "bf16-provider/gradients/0247.pt": "7b4921e404de60678f39763f3f01dc9aa0525b3700f6a666868175d4ed6d2a70", + "bf16-provider/gradients/0258.pt": "7fa9e2361f768b3c11a05912ffad5db282a8c92627cb5c1328015bfa1ed66b16", + "bf16-provider/gradients/0257.pt": "35a0be6b96d695c4bace48f372f6d6da79d529917241a0fc408f542a92da8f79", + "bf16-provider/gradients/0244.pt": "2edda58756a5909e00df680ba20935aaf4b2560219d6795271c7f07cf63fff32", + "bf16-provider/gradients/0243.pt": "0d5dd5c1bc41d326425e58aa87814b63c524da5202202f65ed38e5f955a6e236", + "bf16-provider/gradients/0242.pt": "5e0d121d7d055e4c666de5ea54194efff27fb38124dc46efdf88379fbc952282", + "bf16-provider/gradients/0241.pt": "97ff757508a1f23ea0a5037c65db9afc86d6438abe1b613ebd036f9d0e674918", + "bf16-provider/gradients/0240.pt": "4c994d82c228bf5a5c39d3eaba7fb8290a4593176396979cd05732a34b02ef2b", + "bf16-provider/gradients/0239.pt": "81390d41da237daaf46382ae3d9281a258228ef0d7fa88b2f1bb90abe1f6a63b", + "bf16-provider/gradients/0238.pt": "cce6e8691c4e078d1bc2f6a660fb3b68cdbbad9cd605a062c43f1f4735021dba", + "bf16-provider/gradients/0237.pt": "ababa08f47c4071fa17302e9029a220d69c1d076dd62680ea60aa78ed560724a", + "bf16-provider/gradients/0236.pt": "5c2b4c74861056484b085664f3be32a3dfeb088f4a792876bdc6474c9faa35fa", + "bf16-provider/gradients/0235.pt": "e0818c2664fe961bda4170d740112c554245340796bfc88af79a2465e0aaa1ab", + "bf16-provider/gradients/0246.pt": "96769ebad77b2012bcf50570d35720984269ad4886376203150a0315b2e77b03", + "bf16-provider/gradients/0245.pt": "856296f7d8d4c371d7e6496c781b73b753d90fd1bb93933a308d9e0c3c426969", + "bf16-provider/gradients/0232.pt": "138e34674452dacd239e811f4d0fd95a03ee33c4e2e4ffc82ddc3ab81fb0747a", + "bf16-provider/gradients/0231.pt": "ef62f36d3acd46a543240d3e4eb2457c078041e0d6d7c247266495c32688423d", + "bf16-provider/gradients/0230.pt": "028f6ec9ac2acb4596391ab04bed596121ba3590e8ea542025b3a7b36ba76233", + "bf16-provider/gradients/0229.pt": "f6a26bdf2356aecce69b55a10ed5f757376a572b4c1bf9adf1ea5b48c60889e5", + "bf16-provider/gradients/0228.pt": "3337413bb7b2476214add605a710ddad5b52b0b884c0909625e6aed1fd2fa903", + "bf16-provider/gradients/0227.pt": "f8e336ec4317ab0959462b1514b00488be782431a251f8226d5aeb327bb95ddd", + "bf16-provider/gradients/0226.pt": "e5acb44c93e7eba97d8766da911f7f2cec1d9a81e8494bc0fcee36ac39628392", + "bf16-provider/gradients/0225.pt": "fd46ba7fe63dd918a34eb099b720e3a2cbea2fe8d2a32ce710487dd7e8825249", + "bf16-provider/gradients/0224.pt": "66feef305f8e73b239261d70d32282cc6a4862fc6389284e667b2f9cb1bcdc5c", + "bf16-provider/gradients/0223.pt": "4c00f3c929248a54acb02ebcb2691f3c98eea5dba89b836f641bfe87deeddd4e", + "bf16-provider/gradients/0234.pt": "3a00acaf15200bda3868ebf8657e48a66d743fe15b4a84ad46b8bbbbe6c3ce1e", + "bf16-provider/gradients/0233.pt": "f3713c4b0601d73011409a8a7786a3a2c5005575040d7b7b67e64c3b9db2d951", + "bf16-provider/gradients/0220.pt": "9b9a49fe27a7d4885a41c31084cda29cbfe820861ceec6931c9fb9f272cc2900", + "bf16-provider/gradients/0219.pt": "43a94aa6cda2a92ac3d0b458f1c6b6aab9f8f94149f08fce0d6290179160492c", + "bf16-provider/gradients/0218.pt": "231c88e71b80ba4b5ee9658309e540e75dbae979fd203d5f91f6570bb8c01545", + "bf16-provider/gradients/0217.pt": "02ea797c01ff789685a3daaaccd55b4e4b3da1bc69afc9d015732890d27d2d51", + "bf16-provider/gradients/0216.pt": "a68db9b9da5c22cc5015d4db7f5eff44ea8048ce7aa74a3d5c1948a3d23c4de8", + "bf16-provider/gradients/0215.pt": "e6a8c4f881f31c148bf2424cdcd3f62a89492385659db4cf638d1768f6c3e7d4", + "bf16-provider/gradients/0214.pt": "6268b4b6c92da5faadfeaf8a5268f20f0783fb127d64f25a10916fab0d296e93", + "bf16-provider/gradients/0213.pt": "d1ab32df4a95afc81a053f3087a05f0c6d8dd88eca14ebac174d415bbbb1aaec", + "bf16-provider/gradients/0212.pt": "accce9d5f93aa2cea3c290073756588ae67a72e875031660a5d711ce8744669d", + "bf16-provider/gradients/0211.pt": "c7bb5e8a8094d7addf1260486a3768ce296f18fa7a3790b44afb734d9dc4c6eb", + "bf16-provider/gradients/0222.pt": "bb262babfe633986374c150d32a6c2c4918008707e38d79118cfb3138d0a4a8a", + "bf16-provider/gradients/0221.pt": "367d15082c9c3443e1d302044f0ba89d48a6c4607c1e75a7e00f2bc66af6f081", + "bf16-provider/gradients/0208.pt": "ec69a37acf4fd5ec8b48896c0984b2f36ff580d3c847f0d3f821bc4ddd399d5f", + "bf16-provider/gradients/0207.pt": "ef7b8ad9b92f94a04799bf5f0e470499ad35e382baa15f4c8fa7e9ef402ac66d", + "bf16-provider/gradients/0206.pt": "7ea4af6d3be7edf6d195cacaf5360ab95b4256f4f41fb4cc10a8ed4bb1771d6f", + "bf16-provider/gradients/0205.pt": "3da5315785f276d61cef3108d748074a5a16cc9c6a94e625a059755113e523a4", + "bf16-provider/gradients/0204.pt": "c860184bf7f61c1b423b7a76c496923b20eaa48104265b04aa3da36f4bcd4662", + "bf16-provider/gradients/0203.pt": "142f97eeb0d5376046d69e5e325a0b8a4797adca669687e6f393a64e6bd168fb", + "bf16-provider/gradients/0202.pt": "1f070a3c00ea16a81f7bb03fdc123837f9a6f72eca06c4bed0e7d74d05fe5ea5", + "bf16-provider/gradients/0201.pt": "583c07b39f496470dabd088eee80f0b2e4b0dcb0f1c5da9a796025bd80303955", + "bf16-provider/gradients/0200.pt": "8d2b63e23ee19dd7f9cbfeafcd61d4d5a4052dc7dcbe1344762c7e4473d7b05e", + "bf16-provider/gradients/0199.pt": "a68321b3801cf41e9fa06723164feb62e8d285986cbb7bb3422c85bb4b4cce1f", + "bf16-provider/gradients/0210.pt": "3d1b6ef3a7257378e4c0745b0e8140ab68b1a83b513acf1507417c74429aaf08", + "bf16-provider/gradients/0209.pt": "8da497487256c7cebbeb3acc10d59bf7c438108fa2d0b75fede9c3c6d482e922", + "bf16-provider/gradients/0196.pt": "d56d202cdd76c0302322836571ab2463c93479fe288bfed1857826694fa42742", + "bf16-provider/gradients/0195.pt": "57b41605efeadb40d6aa9108db0e7a5aa7e1f46e5b0f41551fdff305bf2a7f07", + "bf16-provider/gradients/0194.pt": "5d23277607ee99ce66886822f8914ce125a695c1c4a507d9b94d91ce529f9aa0", + "bf16-provider/gradients/0193.pt": "8b7104e81471686c7793303c31404cc359890fb701afedfdc1f0825e8ff03b4d", + "bf16-provider/gradients/0192.pt": "a972ea28abae2e9606fa3ec4e30bf39052e18aff339b66c2de785f6104e8f4ef", + "bf16-provider/gradients/0191.pt": "cad419f09a97deb1912989597ca98c29148705ebb4768f6ede89a5de7d0a4f20", + "bf16-provider/gradients/0190.pt": "dc61442a3ae246cf85e03efed12725f4ae033237e3b92760b81bff1a61b40b05", + "bf16-provider/gradients/0189.pt": "22bc489135f9d4c1ff0f56d8c52977857b32922900555ec1f5bd097ebd9c1a3e", + "bf16-provider/gradients/0188.pt": "b2614dc807f92b54e0ea8589b3fc30241d572e90b8351d5d1db9e467bfc9cab3", + "bf16-provider/gradients/0187.pt": "daba5072a2576fce5fb57cf32e3ee225e4c323ecb5ca67890b7f0e59351658d2", + "bf16-provider/gradients/0198.pt": "45f6987de952c6c13e8379d76188cd83bc648c7a2a05f4292f2f206338ba2211", + "bf16-provider/gradients/0197.pt": "dd4d032f2330d6d263bb9e8ec402ccf5b8103e5ca453286637f67cf373009625", + "bf16-provider/gradients/0184.pt": "7ffd4d1b9fe71b0f514abe5b564a7bc60bdc08551c5fd76ca643589f40279e6e", + "bf16-provider/gradients/0183.pt": "9975bac895e20c1ddc9937b17c851960c2b35a9e5975e9529a36fd246dfe366d", + "bf16-provider/gradients/0182.pt": "a9c98f1544a7c4d2ec3c412b01250ebb273779d6e7d1cfcd76c51533a2bf1b1f", + "bf16-provider/gradients/0181.pt": "6bc9d0c23aae1a4377728a5f80fe0159283d0cc5ed475e8524e2c6427ab4f1c6", + "bf16-provider/gradients/0180.pt": "c5739199f54a40ce008e8500c1dab2b3bd12b7e4c467796152957c652eb3b800", + "bf16-provider/gradients/0179.pt": "e2c76c1bbe74302a29de725df3c8b2e0f14611df60c94a0a2dd6d5c7d8089309", + "bf16-provider/gradients/0178.pt": "cb668f84ff41f5caab5b717d5ca608228d770e208f842de229d71d4c8b550e69", + "bf16-provider/gradients/0177.pt": "7b3dc7301357f3702cd3fd5694d2e5cd456f07b43c439c9e22a2548d30a2de8c", + "bf16-provider/gradients/0176.pt": "6d7ebff37525d991c1879f82dee3e5fa2dbe4dc6373a0c709108c661b2768701", + "bf16-provider/gradients/0175.pt": "a4eeab0695a22cf78fd777301257a2589261390e1cc6f9a14c4d87f6a883b559", + "bf16-provider/gradients/0186.pt": "8434eeded31e0c7ea8b047c94ff5f164c7b50062b574dc0939dd18cbd248d8a8", + "bf16-provider/gradients/0185.pt": "1c1cc9022678617e00a2e806b8e18dafa8b85e6165fcbf78e0590be822441d43", + "bf16-provider/gradients/0172.pt": "016a36a9d3567779fe23d91f99ff855e1a9af8066e373eef9f23299bfffc5837", + "bf16-provider/gradients/0171.pt": "b0c8b6b6aa50924a8139da67dee1efce799e4a8cba77c994ee4087a15ceeca04", + "bf16-provider/gradients/0170.pt": "dc0221a72663a02aa2c447e624e9f5a81d3c8324d1c0918329b72fa5b9ea5682", + "bf16-provider/gradients/0169.pt": "ccc96b95774d18294008e74868b40e27006ccc8c2bd914b5386346524080a97d", + "bf16-provider/gradients/0168.pt": "505b0a5d3af2787f104dda0b8ef704250f737816488c52fbea7d5d7d18a53330", + "bf16-provider/gradients/0167.pt": "711d2fd41c74af125cec317c7cd3b2fda5722f7572124486e2b4913b94591225", + "bf16-provider/gradients/0166.pt": "9c74031094116729122f0f311dbc62f65f56a36e4c2df75ab6a96afff189d3a6", + "bf16-provider/gradients/0165.pt": "58804c807ce3e933f23acc19a9f7298fe75a921a9ea41bd148db05e1bdcf94bc", + "bf16-provider/gradients/0164.pt": "69db65b8d18eeed754cc5896693ad5f42ef0a1b8187416cb761de299cce77ee0", + "bf16-provider/gradients/0163.pt": "ef62de4a59ed2b7e58a66dd933affab1263387e3798e07cfa3490b84f6e11b83", + "bf16-provider/gradients/0174.pt": "09ede8866a958b2ee9e8126369fef6ce7e4323409fe3c00563d8360d9d42101f", + "bf16-provider/gradients/0173.pt": "6a1096389ef27a462216c6b69f687496093ac17f42981176afb6dc45fb6413a0", + "bf16-provider/gradients/0160.pt": "9bde85c82bffb4798e008201297dfd29dbfa0724eb98d499ee1ade1e9f43a549", + "bf16-provider/gradients/0159.pt": "1c559ae77b1bd082da5e9c5e54f2b50c4e58f9e1b49fc37f2ecc303642d7ddba", + "bf16-provider/gradients/0158.pt": "018d969699505bdc6bf01a63d9d49a24004ca43d4e00558fbfede39ad7f7731a", + "bf16-provider/gradients/0157.pt": "9f18ab851c0c1d9c73813dc72daec46bc826635fc0818bdf48df0295b8386e81", + "bf16-provider/gradients/0156.pt": "cd2a439b79d958254592d7f5960001b3f3e6ac4bd9930f7d7774a57f06b957f8", + "bf16-provider/gradients/0155.pt": "c4e09428b8557906d650939be1935fa41c04757640720cd1e8d82351d458821d", + "bf16-provider/gradients/0154.pt": "13b9ecfc3c2af43fde4b48e02654633c6fe0ddd7423190ee23fe785d46c90e56", + "bf16-provider/gradients/0153.pt": "29c671db274b7200b9a57226cc85ab5241cd7c1a30a2ce5c9461470c27484e17", + "bf16-provider/gradients/0152.pt": "3a7377c5348ea6165e9b7b82683e7a84d73e3047bfa84275598c7c2135abd8a4", + "bf16-provider/gradients/0151.pt": "593708635544f6403cf94e36c59324a516cc8cfed11530c16da1dc54ab90e5d4", + "bf16-provider/gradients/0162.pt": "afae5321af111190bbe3a27f8633ca05f7a95b7c195178682c5b3d18062eda0f", + "bf16-provider/gradients/0161.pt": "15bb8d7e4ef8e2b4bff6ab0ac4641aec9faf6f9e6bcac66c62cba59b851315b0", + "bf16-provider/gradients/0148.pt": "e23cd4f3ae6d0f9b19affbd1c48dc80a47d16ef05c624ca126577e0d0495cb7b", + "bf16-provider/gradients/0147.pt": "85f26343f8242e0ff98a6ea9faf60cffaafe9076d4e415476121ae9c28547a68", + "bf16-provider/gradients/0146.pt": "815508d3a5721fb5850dc275553142796cc47f9fcbcc3c6d97a003c66d95df28", + "bf16-provider/gradients/0145.pt": "a759e15ec38a04639058443c47a2a1fa5660eb0d439524f6b986ee652f6936fc", + "bf16-provider/gradients/0144.pt": "acb1925dc915c8cc874d66abd51c5e944b98c629dd334639ff34cc0a97ec2245", + "bf16-provider/gradients/0143.pt": "10da5b0274b175a6223cf48b88ddf37ebf16e0083abfcf7c1a7e96b948f26283", + "bf16-provider/gradients/0142.pt": "162739213b2cd276a285823e42b40e49b737108a5802cce1132cb3f5b812222d", + "bf16-provider/gradients/0141.pt": "cd4c8c4f39356449b76bfa1ca2e7126f6b5e843345ce221cc34cc590524dea3b", + "bf16-provider/gradients/0140.pt": "ec04f8fe998018da5fbd1820a5856d53c690fa17206212057a59e33f321d935e", + "bf16-provider/gradients/0139.pt": "53a5cd74f51b7995c2c0418ae70d94c27d24a05ef83e5dcd1a27bebe5471770a", + "bf16-provider/gradients/0150.pt": "07932e7ee2dbc633b8ab0f8e213bc5f793058f7b59a4b8ebd8c6d5f98485e2d3", + "bf16-provider/gradients/0149.pt": "4f66852763363e3cc77cdc62cfe13ded71b03f9b607849f32443b50b5b520931", + "bf16-provider/gradients/0136.pt": "1d712a653fcf0981214b1924d30d3e8c8c04c50f57e075d70c39aaf65df3eaba", + "bf16-provider/gradients/0135.pt": "c75dc1ed5222c4fa46a564d6032b151bdd36c46cd789f9d75593bbcd14ffb6a2", + "bf16-provider/gradients/0134.pt": "826cc623be29f797ecf4546c34870dbe86bcad0a20b6a61043c7771851153ede", + "bf16-provider/gradients/0133.pt": "1eae2273a5afb2b3acdae724666600f103af57b3d072cd4967bfd0107442d840", + "bf16-provider/gradients/0132.pt": "534b45c8e194ec24845c3ab685390f5d6a8194b2058c85f1e2f91f4bd5210312", + "bf16-provider/gradients/0131.pt": "087678ebe2ed7fb82735eda24a4631aa309a2f40546753fda03d6d691e1b6d00", + "bf16-provider/gradients/0130.pt": "b33178c0f5a3921fba3d9adf920f222a4d20c33f1aeeb206975bdc258c92dc98", + "bf16-provider/gradients/0129.pt": "b701fbc13ec95e85222495fec6c6028b1a1e93aff213eca68c3e7105e450b6c4", + "bf16-provider/gradients/0128.pt": "4c7621976840a1f8459727e71edd47c42ca19f20475c49597355c25344e4f842", + "bf16-provider/gradients/0127.pt": "274dc09a775da8f4a61b32173d2178e8df21ea1b83bbe765c99baa2714f64870", + "bf16-provider/gradients/0138.pt": "877c7a18047c4e3ceadaf73b508d90877081685c75ba1102249a6939ec1d4c25", + "bf16-provider/gradients/0137.pt": "18b570db82f4deeecbf75fb8dc029cca837e6a8d0da23a5fe45f73b663e41892", + "bf16-provider/gradients/0124.pt": "cf19e165eaf709bfe0359d8e486e78497e2b0e00db7c024df03fadf42ebea5f5", + "bf16-provider/gradients/0123.pt": "ac8722337c50c2fb9798afb3afb098e71fc5bf036dc081b6479525d12481959d", + "bf16-provider/gradients/0122.pt": "774e8af6593a0dea7dd67593f6485c293625ff5645fcc251f25d841ea572d7a7", + "bf16-provider/gradients/0121.pt": "9161e0f787bfcc9cbf4e1108de217bc77d25d37811eb4b641c95fe4443e047de", + "bf16-provider/gradients/0120.pt": "89e5bbe7aa9f2dd6dbc0bf76c5e405fb784ced9c68b1de06133a43cc59a65fd3", + "bf16-provider/gradients/0119.pt": "55b11262ad988c34f26974d3cd40a86ae59543f911c8c591517ee27f85efd640", + "bf16-provider/gradients/0118.pt": "0a098c417551c7e109535459167c92d26294628cf3fcbca6f3f17cc22442c2e3", + "bf16-provider/gradients/0117.pt": "fe8c8212c7b9673111ec777f46013df5fd477d8bf00b86c7e4bdab7bfc5e2c93", + "bf16-provider/gradients/0116.pt": "8c3033828275a2771b6bf640047642bebc6c05c0f7369723e8e9e2195dca89bb", + "bf16-provider/gradients/0115.pt": "31cc3f5422a20c671f2cce549e93cc8c9ca1da510f16e46a17a3b4e5aa32bc63", + "bf16-provider/gradients/0126.pt": "3d5e9096ac850b1f195a13943969b7f9ee28a241943c1f70850473ca4672d4d0", + "bf16-provider/gradients/0125.pt": "144130e8520ffbcca2a2ab0f41f6f84807f181c1d1077c3435206965b75b53ab", + "bf16-provider/gradients/0112.pt": "5c7af8d04a634bededab66333074d7626ad411ac306d693dc95ed6def0459cd8", + "bf16-provider/gradients/0111.pt": "f59db7383925bd7ea207b917235a06ede25592c747ee02a38c721677a6003314", + "bf16-provider/gradients/0110.pt": "ac80f3091f2725f92ba02eeb582c806eb863da016ff92af423bf777814181f0d", + "bf16-provider/gradients/0109.pt": "f0a2afc6c0fc3844b3529e2be8ffb7255a2ca864a8ae2d0a9e3f669a20c14ca8", + "bf16-provider/gradients/0108.pt": "9bb5c3eb397301d96fd1490756031a1f8d689f9b6242e5f1d6de836a34292d1c", + "bf16-provider/gradients/0107.pt": "7c613e45624c62ae2e08c0c477923fd46fdb9092cb733edd403976c4a04b4b6b", + "bf16-provider/gradients/0106.pt": "63c7b0a3b11ccf9e9d6011884d9cb36b6284bbe67d160f14028fbfb1de61bbac", + "bf16-provider/gradients/0105.pt": "aecae0cef955e9e1e1fb66d901716569c8e5d74dd3db01eee052e81ea3c6dc42", + "bf16-provider/gradients/0104.pt": "d6d8d40daeac243f6fd9975d4ccac10bcf30354f35173a562ba6e05f6ad6a7a1", + "bf16-provider/gradients/0103.pt": "0db7dbf32f62f6659ee1205f91b81ca1805b529ae75e725eb5ca7b1909ff44cd", + "bf16-provider/gradients/0114.pt": "6da9b3527255cd82ac91ea421bb0802a9474902783191eb0f5d65b0ac062dde4", + "bf16-provider/gradients/0113.pt": "10630f0a39d2d0abfb49aa3ac08d173f34ab6eb11ce15bf3672f207a887ead14", + "bf16-provider/gradients/0100.pt": "2832a9cc77db09443da100feb1553606e3ea625218f5ef4574724c909cc54a9e", + "bf16-provider/gradients/0099.pt": "5e37cd9e586a650413176939381063b26530930fa26daabb00e60d50b1353d2c", + "bf16-provider/gradients/0098.pt": "8bbf8b390b4441065205c70a7e5e8b011bd1b8f35f5f2da2a827562190c3d9c7", + "bf16-provider/gradients/0097.pt": "6ebc32a32feb456765d1e3ad7bd96987e6ff8cad906c1b7d3e352e08e410f711", + "bf16-provider/gradients/0096.pt": "4bee3136156a423a43b4ff236a72d3a2a56932233ab25e89a984bbaaf2c8e4f4", + "bf16-provider/gradients/0095.pt": "bb34ac908f20a239d98c1f5c0416d91cefa95f5c105ecd241f24de8bebacfa3a", + "bf16-provider/gradients/0094.pt": "17d4322dd958c3e096b66230312b58cf14646e24b500234e3c6f302b70f2568d", + "bf16-provider/gradients/0093.pt": "dd70dcf21fb8ccf93a51991e89415380bf421b2f550245076f7fd4c877d786b0", + "bf16-provider/gradients/0092.pt": "904a0e7949f730a8257cbb4fb7afe68e05bf491ad40bf36e51e55aacf2438eee", + "bf16-provider/gradients/0091.pt": "15ce62a31c2871e75845a26cb32f8668c13c6d544e0ee941e090eddf48af1f41", + "bf16-provider/gradients/0102.pt": "2105ddf33f8b5d864df4e9c8f005099fda2d241cea290e0c8c0edc7f7f9d1f44", + "bf16-provider/gradients/0101.pt": "3f6b9693e70bc787745771556135dfed5ccf4265de5eef6da488622de0517577", + "bf16-provider/gradients/0088.pt": "2cb9cf8ce9dd39c7f7f57f4ad225f9a5bc99e9113b243095905f5edcedde9cbd", + "bf16-provider/gradients/0087.pt": "435fa661238b6c1b4551c5491d18099d7468149a4d48a1da7a45b3cdd8a116e8", + "bf16-provider/gradients/0086.pt": "b9d949c5cadedcfe6d46f8192de6f307fff9f7b05c9faa0d73b165f349f8c05f", + "bf16-provider/gradients/0085.pt": "4c2a4273bd9defa879591794a590eab1242047245433c618cf23e95f41cb87e8", + "bf16-provider/gradients/0084.pt": "f35f309d595786a6e929928b5b50c9a93bf0a91d8b9a2e744836ab53978eea6c", + "bf16-provider/gradients/0083.pt": "166cfdaf62b8893a83d92dcf8bbdd162edfa88bc48f0fdfd4d2e5aa8fab1d8d0", + "bf16-provider/gradients/0082.pt": "143aae86d17c6a1935980f0b2d1eff88d1aab8c942d28fe982a8091f3f9a0031", + "bf16-provider/gradients/0081.pt": "5f473c39fc3374181c611d1eb265679ddcdda92ca1fab207f75a4cf5762bb459", + "bf16-provider/gradients/0080.pt": "4d42a16f74698f5bcee7ae76b980878506cb23b4b8f3801bd1f4bd0cbc014ba4", + "bf16-provider/gradients/0079.pt": "c61dd2aa8d7b9c462fbb35803ca2232ebafe460a83147eb4aa8db48922035b2b", + "bf16-provider/gradients/0090.pt": "0141bc5cf23698c2f891fe0824f2f7d1bea7400b01f1ffb38c908574260fe5b9", + "bf16-provider/gradients/0089.pt": "938b32ab210aa1ce7cff95866361c0230371d44cf54c67dd945324613926972a", + "bf16-provider/gradients/0076.pt": "293d60dffa0f358385c13b7e3d456b5598dd9966e3fd121c6ff82455c318f374", + "bf16-provider/gradients/0075.pt": "f36571750844e76325093b7c37ac3bc916ce5374afec6f260b7d046c03c3a70a", + "bf16-provider/gradients/0074.pt": "9038d73341d28c607d92439ff89ca0bda9c372953ad4c6ea5251e7e748c9cb09", + "bf16-provider/gradients/0073.pt": "339cd2fcee85b12299e2cb9548f33ceaff90765e4eedfb086edb6c13a9ecd65f", + "bf16-provider/gradients/0072.pt": "1ab591630cbed9bdc50b6efc206608f81ff91db1398c80174d6b941f181ecdea", + "bf16-provider/gradients/0071.pt": "d09d2cdc8fa64016333be8b54e51b169944e4a4a408e216152760d79f2ada1dd", + "bf16-provider/gradients/0070.pt": "5d91967d801a1adc4cfcc5a317fad0c0fce6aed02f7b406ba30883f493b9d33f", + "bf16-provider/gradients/0069.pt": "d793217a0b139aa0130f525a81bf5ca4d70b91ccf920782d2f6d0ed49e199640", + "bf16-provider/gradients/0068.pt": "f4d4e7601a3d80129447decf647ead76d36587b33ffa68cf135b54c68395a202", + "bf16-provider/gradients/0067.pt": "2d646468bfd531007d63d06c684387e9a86b9c634746ba8c20790a9e348f58c5", + "bf16-provider/gradients/0078.pt": "1b8c659f0574cf4febdf7f8fa8df027d81d411244705639760530716780647de", + "bf16-provider/gradients/0077.pt": "e48b2388258704879784d4b0bd4e589e85f84285c333fb04d0d23e3c1af332b6", + "bf16-provider/gradients/0064.pt": "b7d99f27a3c4b724fc6b7cc74903657b04465c0432e52c92121ff65ffef82eb7", + "bf16-provider/gradients/0063.pt": "4133fb7012619a578275b9011dfc5eb37c3c1619d22a18379d69d7808f7d3e08", + "bf16-provider/gradients/0062.pt": "0fa4df1da173e0fbc15ede70c252eed208b1b3e6c1b4a6ef77f064bc160be685", + "bf16-provider/gradients/0061.pt": "39b847e7d283acdd0acada980458ea5272dacf5d9a89902228ae58106f596de2", + "bf16-provider/gradients/0060.pt": "97872df7510332413337eca50c394132423e2c26c231339d1328261332892216", + "bf16-provider/gradients/0059.pt": "152c079726dd0ae673badbc4c878db1ba144bbdffa781947a9018dc3a901530d", + "bf16-provider/gradients/0058.pt": "17996bdb44c46cdba536edcb3271ac25a6e914c11f6f8b9392702475012e19da", + "bf16-provider/gradients/0057.pt": "e6a2d5f495f9dfa98402fcfe5292a85b31927351a00fb0aa0c01e139f21ed7c8", + "bf16-provider/gradients/0056.pt": "6c7eb067b1ef06be780b031196b052f2065f7646798962af6ee5910251ab814a", + "bf16-provider/gradients/0055.pt": "ee511ef6481d96e4178843075ea5d3c65d5fdb5a0640b3fb80b83215a01d79b0", + "bf16-provider/gradients/0066.pt": "71fa3630d746d844ef8332249afde4763182fb108b603d90d32807e60ac085ea", + "bf16-provider/gradients/0065.pt": "3a5690dfac45adec64f184ee74a2e514bd435c2ab6bb7eea078e924e7b3b9975", + "bf16-provider/gradients/0052.pt": "55f1a640d777a587dae8ecba98331694099dcee1ec05f2801e3e34a890c675ec", + "bf16-provider/gradients/0051.pt": "10d6da1d756568758347c08247e50184ada2aafaf5c3ddbdf73d661d672f71e6", + "bf16-provider/gradients/0050.pt": "56bed646e4f450fa1d6f8a8bd5ec45484a69c874efb7ba9b2fd7d6149e1b1318", + "bf16-provider/gradients/0049.pt": "dbbd40845e50516541f796388f91efaaf805d1594880bdb63292f97bd6a9f025", + "bf16-provider/gradients/0048.pt": "2bc9cb41dd36ffe8d4681a45e69969a29b97a30ecd717cf1fe6f776dc239a76e", + "bf16-provider/gradients/0047.pt": "4c14231efb114fea51cd04115a861cb9d1e3acc56257c84cac2c3a8d3ff63a95", + "bf16-provider/gradients/0046.pt": "1f6dfc8143744f1bdfc3b6dd2a42e041ec1dcf863369a02b4431da4c2b0e0b46", + "bf16-provider/gradients/0045.pt": "e42525b6185b5070da925ba8ecea153090ba1d08edf3f9d89f0bd7dddd105bf7", + "bf16-provider/gradients/0044.pt": "911f5953d0c2a5f4746549c1e9703f1f596c02496198d2831797a3e37c4f5d71", + "bf16-provider/gradients/0043.pt": "44d0fda8dc1ce5e177bbd028228a60f3c71fe96689c0811ad700476ebec370b4", + "bf16-provider/gradients/0054.pt": "1ffec37aa9359996c99b32f7d9c785d89bb107806e8fa1ba3bbb1f5d6bf60402", + "bf16-provider/gradients/0053.pt": "d874a0151cfe590905adcd4b089cb0cf8d0193ec4073da7dcac82b3f3999eb6d", + "bf16-provider/gradients/0040.pt": "4895de0e001cb9602018260946f4b3c5f508d526aaac4659b8b18b30e677c441", + "bf16-provider/gradients/0039.pt": "12eebd64cd38985f12a7f1d39ca002b35206c2d37a00c94846209fe96b2b9811", + "bf16-provider/gradients/0038.pt": "44d110c789407ada35b4b984fed865ade732f26d79c78762afcba528897ed0e1", + "bf16-provider/gradients/0037.pt": "24fb83adcaa5f42f51ea86c546221a85ab7937605150460652459d44c6dd29e4", + "bf16-provider/gradients/0036.pt": "ba28120bda4a87b324e20c4e03e9497e9df18e489e1b89df82eb5b9233ec5537", + "bf16-provider/gradients/0035.pt": "c5f967b3ef117e01c047e6b578deaf86154b44619055e82f10378833356c571e", + "bf16-provider/gradients/0034.pt": "fcf2313216dd2b8fdc8ed31a389b459d8704e281685f83b8be8893cbbdf8236a", + "bf16-provider/gradients/0033.pt": "48ed901ebc30a8eb15c8aa00bec6bee7fd589c70031fa84227a34bd830dc51c9", + "bf16-provider/gradients/0032.pt": "9f23a16e0d67d0d5781e55fe6d90230f2c7a69f668867d084dc064843f861e6a", + "bf16-provider/gradients/0031.pt": "89b10eef54ff52932a865f92174f3d97306ced96759aed3a069e8c11a6eee8ac", + "bf16-provider/gradients/0042.pt": "080dafd3d8a7c0c8077931a5a7f3a160e4b1add36f84cd70ce7d359b0313853b", + "bf16-provider/gradients/0041.pt": "07d1ba8a4a7d26df9da7ca67280d374596a320739659b21500814befcfce439d", + "bf16-provider/gradients/0009.pt": "de47e82f2f3500d93f56e4c028f0a4fa8d2790f4e62479754dd2412a99fa7b7d", + "bf16-provider/gradients/0008.pt": "1dfc2144fa6d768523cb6898d6a20b99cca90b46e3515ea4fb1ec738310a5b74", + "bf16-provider/gradients/0007.pt": "17f995c829b58ad1bb6a584222727be244251458f564b7fef5c7dd88d4c198dd", + "bf16-provider/gradients/0006.pt": "62bcb914fc74eff342f30a65eb2cf3d1509d550912f1bc58bf684250b7017534", + "bf16-provider/gradients/0030.pt": "e88add9ff07e2e340f831ddeb0d6cdb7234da718dbeacf483e96fcddb340b3dc", + "bf16-provider/gradients/0029.pt": "565007293ed7d6ccd577c0879757e489f63aa81ee5846f268a55831610b14adf", + "bf16-provider/gradients/0028.pt": "ae8bf5675c319d904ad3a47466ac56616eea4cd76dcad9f21a860eaea6702cbe", + "bf16-provider/gradients/0027.pt": "5b3b2dba09b7e4811283d4797c97eef57378dfee7c4f0bb0bc89b3732f3c3d79", + "bf16-provider/gradients/0026.pt": "ed74b7bca7380650b552056563196958929d6d130ebff126550df717eef15123", + "bf16-provider/gradients/0025.pt": "d15c5a45bb7858ce7a2033573d3bd8d15b3d502db00ddc9335ac2880d5f20ef6", + "bf16-provider/gradients/0024.pt": "6970eb0e2285784f4c5495009a87dfbb36ac612c5570e04d220c7484669e727d", + "bf16-provider/gradients/0023.pt": "de6c1a0c38b17d5054a85664dd4587a942e036305754bf3796bead97ce2b9957", + "bf16-provider/gradients/0022.pt": "55761c3189e84cc1bc2de51c4f270e679ffbd532145994deaf253930117d6e7d", + "bf16-provider/gradients/0021.pt": "b232701911cda69913bff55f8c305b3d6da1b08f3ddd1eef63c1b3883a4e2313", + "bf16-provider/gradients/0020.pt": "b113fe0092d38ff7c209ca0cae8bb438e437e1d9e80cc59b7f1933a84fdf6be1", + "bf16-provider/gradients/0019.pt": "73c9019f4869b3233e55b8d3ddcf641f27a2c398e6d843cf50040b6c0215ffa8", + "bf16-provider/gradients/0018.pt": "94a2d5447adf7806a9b1c012b4ba95e180dec535f9b658b8fea5c160dc108e0f", + "bf16-provider/gradients/0017.pt": "d721ce7e8891d57d9a5e34c06211d2e3e5aafeb610fb4b00c23ed21d61a17591", + "bf16-provider/gradients/0016.pt": "6ffaf1248e2c88315943bdfbece1abc8e24dc655634d801a6203ea6016572e85", + "bf16-provider/gradients/0015.pt": "d1d638a00f6a41657f8073129fd3f7c3d57cf58915aab47a24d79f66f69a6aec", + "bf16-provider/gradients/0014.pt": "b12ab91a2aa7392166ad5d2090a26a3ca1f62b98ccdabab9594d685806ecfa42", + "bf16-provider/gradients/0013.pt": "ba45fd0ff6f5600ecb2a9ec6babda8836c863225430c5430d5b78c16b0ccb75a", + "bf16-provider/gradients/0012.pt": "816a2593b6597196ee735b7909ba63ff33f8003e98e0411a834daad804333047", + "bf16-provider/gradients/0011.pt": "a0b24512b80d5a7de9c1c5d636fc977c263d9507d30352e537f4c4af04d5b8ac", + "bf16-provider/gradients/0010.pt": "aa8e3825d6987a4ae2ebe5af3e3d7326d2102f1a98ac4fc4516327897e5b6a7b", + "bf16-provider/gradients/0005.pt": "5a661a8f6aa437196011babdafe14d8c754ad2e4be75900786e6507cb134b700", + "bf16-provider/gradients/0004.pt": "4beae74ac56b4eaed53c2b17a7d31f98d045e64e01d067d5067851304a0efcb3", + "bf16-provider/gradients/0003.pt": "b3427ad172c897d9ca3f684d70286d98f99f9074b37a369e7aad5e2fb69eb3df", + "bf16-provider/gradients/0002.pt": "1c56aa9e0a604d21e52062270e0c6ce5e319b79694986338da13dbaa8dd81f08", + "bf16-provider/gradients/0001.pt": "f575932d4d780dcf5da597987286d0c9cccb1bc589113a38d0a0b010b56a5e87", + "bf16-provider/gradients/0000.pt": "369cf6e3c0a16b175a5eb8d95e43b339b96a1f7f01b1b765f05eb8a38048897a", + "generation-reference/42/report.json": "e55e877570fa3f089b158be76a4423f8b25443af5c85fca45ba3dd378974c123", + "generation-reference/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference/42/audio.wav": "373efac29bd3073f1dfebf0cba30c51848866351fb6d2c8e55921047c82cd296", + "generation-reference/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference/42/video.mp4": "560d134d198e725880f308be6b8ebafd71718d01f323ba7b5f3fbee379302d01", + "generation-reference/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference/43/report.json": "7e181902ac5984cd04c8773a2c0bae77ba10618b63b0e5b2a3b56599c4ac5525", + "generation-reference/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference/43/audio.wav": "eb404d2f89fbf3eb9795ef2704c245618e47e552a3b4d782b52602762569de6a", + "generation-reference/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference/43/video.mp4": "14dc8307aab50e41849f02e3726612a5e2b5d197dc04987ae9196d2b63c2e5c7", + "generation-reference/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference/44/report.json": "50f571441bad710594cb53fb861cc73fd6b8b9e81a2fbe82de1f9dea4feaddfc", + "generation-reference/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference/44/audio.wav": "0cab0190275bb143bbd997f50928cea3e2c0208f90cd71b9633046ec6e4ddf0c", + "generation-reference/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference/44/video.mp4": "e68e18a0fd4a1ea9c5da84969c70299d9d21db6ce9aa8592bd649f84c6b9b595", + "generation-reference/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "generation-reference-repeat/42/report.json": "e1f254991fa8c8e72f2dffc4e30a056d38e9e302fd2f684d8c593895afc020e4", + "generation-reference-repeat/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference-repeat/42/audio.wav": "b9891b3a34e097361f06e63a9e890dc4551600a9ec575a08dc07ef9493c19fe8", + "generation-reference-repeat/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference-repeat/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference-repeat/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference-repeat/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference-repeat/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference-repeat/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference-repeat/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference-repeat/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference-repeat/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference-repeat/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference-repeat/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference-repeat/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference-repeat/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference-repeat/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference-repeat/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference-repeat/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference-repeat/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference-repeat/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference-repeat/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference-repeat/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference-repeat/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference-repeat/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference-repeat/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference-repeat/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference-repeat/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference-repeat/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference-repeat/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference-repeat/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference-repeat/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference-repeat/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference-repeat/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference-repeat/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference-repeat/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference-repeat/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference-repeat/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference-repeat/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference-repeat/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference-repeat/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference-repeat/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference-repeat/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference-repeat/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference-repeat/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference-repeat/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference-repeat/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference-repeat/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference-repeat/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference-repeat/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference-repeat/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference-repeat/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference-repeat/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference-repeat/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference-repeat/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference-repeat/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference-repeat/42/video.mp4": "f97c122eada187d8805bb088a2326774af2524e598091446c1302924ee043d9e", + "generation-reference-repeat/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference-repeat/43/report.json": "914a6e564097b23305d6bce5ba4b06a30603d5dd3c0866e3688f3a8a8bd76987", + "generation-reference-repeat/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference-repeat/43/audio.wav": "1cc055d49792fa9760701ad092f1f9e9dba691503ca7a4205702fc3e0c6d63e1", + "generation-reference-repeat/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference-repeat/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference-repeat/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference-repeat/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference-repeat/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference-repeat/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference-repeat/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference-repeat/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference-repeat/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference-repeat/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference-repeat/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference-repeat/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference-repeat/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference-repeat/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference-repeat/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference-repeat/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference-repeat/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference-repeat/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference-repeat/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference-repeat/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference-repeat/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference-repeat/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference-repeat/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference-repeat/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference-repeat/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference-repeat/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference-repeat/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference-repeat/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference-repeat/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference-repeat/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference-repeat/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference-repeat/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference-repeat/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference-repeat/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference-repeat/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference-repeat/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference-repeat/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference-repeat/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference-repeat/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference-repeat/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference-repeat/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference-repeat/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference-repeat/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference-repeat/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference-repeat/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference-repeat/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference-repeat/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference-repeat/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference-repeat/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference-repeat/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference-repeat/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference-repeat/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference-repeat/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference-repeat/43/video.mp4": "51af10f6386bbe0b12db43d7c5d4727fe5e6c1aa91e5ee7502ff4b9e71c96f3a", + "generation-reference-repeat/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference-repeat/44/report.json": "d72c8d8cdc680c7af21481cd1fd35e821fc231020caa8c310497e25833433100", + "generation-reference-repeat/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference-repeat/44/audio.wav": "edda8f495c4288cd872d2a585bfa17d9a4f7e6d0df79893a0cf1b08f3adaff74", + "generation-reference-repeat/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference-repeat/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference-repeat/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference-repeat/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference-repeat/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference-repeat/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference-repeat/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference-repeat/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference-repeat/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference-repeat/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference-repeat/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference-repeat/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference-repeat/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference-repeat/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference-repeat/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference-repeat/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference-repeat/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference-repeat/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference-repeat/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference-repeat/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference-repeat/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference-repeat/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference-repeat/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference-repeat/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference-repeat/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference-repeat/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference-repeat/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference-repeat/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference-repeat/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference-repeat/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference-repeat/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference-repeat/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference-repeat/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference-repeat/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference-repeat/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference-repeat/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference-repeat/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference-repeat/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference-repeat/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference-repeat/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference-repeat/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference-repeat/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference-repeat/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference-repeat/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference-repeat/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference-repeat/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference-repeat/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference-repeat/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference-repeat/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference-repeat/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference-repeat/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference-repeat/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference-repeat/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference-repeat/44/video.mp4": "929285c63024073b07eab1f6b7c3fbf5d2e49cc0f07b2e426dff5fb006d09196", + "generation-reference-repeat/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "full-model-thresholds.json": "9542bce25d394e72f24c2cc02012fc47dc16100a37d5f86feda59c132132eb30", + "reference-repeat-quality.json": "3817569746ac3a1c99cffe32bdf183fbe12de2817c3e60fd6c21c17711bdd691", + "commit.txt": "7c9a380b0ba89a2288ca661bd2fda1e102a02606521681ab0bcf3077e5d2d600", + "dependencies.txt": "6139443eb38b47a805548b0d892808f5b6a26c7bc3b3796ff96ef6b9b0185cbf", + "extra-dependencies.json": "d2655f41325039e28fd4aae62d5637de0df49d08ea8dc9545295cdbbf1f47ce6" + } +} diff --git a/examples/models/minimax_h3/reference.sbatch b/examples/models/minimax_h3/reference.sbatch new file mode 100644 index 000000000..ce7429710 --- /dev/null +++ b/examples/models/minimax_h3/reference.sbatch @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +result_root=$4 +reference_image=$5 +provider_source=$6 +extra_dependencies=$7 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +"$python_bin" -m pip list --path "$extra_dependencies" --format=json > "$result_dir/extra-dependencies.json" +for mode in fp32-reference bf16-provider; do + "$python_bin" -m tools.validation.models.h3_full_model \ + --checkpoint "$checkpoint" --manifest rl_engine/validation/models/h3_full_manifest.json \ + --mode "$mode" --output "$result_dir/$mode" > "$result_dir/$mode.log" 2>&1 +done +"$python_bin" -m tools.validation.models.h3_full_compare \ + --reference "$result_dir/fp32-reference" --other "$result_dir/bf16-provider" \ + --output "$result_dir/full-model-thresholds.json" > "$result_dir/calibration.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --preflight --output "$result_dir/generation-preflight" \ + > "$result_dir/generation-preflight.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference" \ + > "$result_dir/generation-reference.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference-repeat" \ + > "$result_dir/generation-reference-repeat.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_generation_compare --reference "$result_dir/generation-reference" \ + --candidate "$result_dir/generation-reference-repeat" --provider-repeat \ + --contract examples/models/minimax_h3/generation-contract.json --output "$result_dir/reference-repeat-quality.json" \ + > "$result_dir/reference-repeat-quality.log" 2>&1 +# Candidate jobs are submitted only after these immutable reference artifacts +# are reviewed and the numerical thresholds have been committed. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..40fb56213 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,153 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... +def h3_timestep_sinusoid_forward( + timestep: torch.Tensor, + num_channels: int = 256, + max_period: float = 10000.0, + check_range: bool = True, +) -> torch.Tensor: + """Return FP32 CUDA cosine-then-sine features of shape ``(T, num_channels)``. + + Require a nonempty 1-D FP32 CUDA timestep tensor, positive even channels + and positive ``max_period``. Check finite values in ``[0, 1]`` when enabled; + the Python operator supplies the analytic backward for this native forward. + """ + ... + +def h3_det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: int = 0, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Return a deterministic CUDA projection and optional FP32 pre-activation. + + Require same-device/dtype FP32 or BF16 contiguous, 16-byte-aligned matrices + ``x (T, K)`` and ``weight (N, K)``, ``T > 0`` and ``K`` divisible by the + 16-byte vector width. Optional contiguous bias is ``(N,)`` in that dtype. + Activation 0 is identity and 1 is SiLU; BF16 forward requires SM80 or newer. + Return ``[out]`` in the input dtype or ``[out, pre]`` when saving activation. + """ + ... + +def h3_det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Return a deterministic ``(T, K)`` CUDA input gradient in ``out_dtype``. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``weight (N, K)``. Output dtype + must be FP32 or BF16; output-column chunks accumulate in fixed order. + """ + ... + +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... +def h3_det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True +) -> list[torch.Tensor]: + """Return CUDA weight and optional bias gradients using an ascending row fold. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``x (T, K)``. Return ``[dweight]`` + with shape ``(N, K)``, or ``[dweight, dbias]`` with ``dbias (N,)`` when + requested, in FP32 or BF16 as specified by ``w_dtype``. + """ + ... + +def h3_adaln_row_gather_forward( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + chunks: int = 6, + modality_num: int = 3, +) -> torch.Tensor: ... +def h3_adaln_row_gather_backward( + grad: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... +def h3_rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward( + 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, + sorted_pos: torch.Tensor | None = None, + tile_begin: torch.Tensor | None = None, + tile_end: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... +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]: ... +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]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +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: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/__init__.py b/rl_engine/backends/cuda/model_specific/__init__.py new file mode 100644 index 000000000..3b09d3c11 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA implementations with model-specific arithmetic contracts.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py new file mode 100644 index 000000000..17a3a3884 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused H3 AdaLN modulation: ``adaln_projection_3mod`` followed by ``adaln_row_gather``. + +The forward is exactly the two RFC #420 ops back to back (same kernels, same +bits). The difference is the backward. With two separate ops, autograd hands +the gather's table gradient to the projection in the table's dtype (BF16), +which rounds a long FP32 segment sum once more before it reaches the +projection weights and ``temb``. Here the segment sum stays FP32 into the +projection backward, so the parameter gradients carry no extra BF16 rounding. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + adaln_row_gather_available, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import validate_h3_adaln_projection +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_adaln_row_gather_forward" +) +BACKWARD_IMPL = "fp32_segment_sum_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_adaln_row_gather_backward[fp32]" + "+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNModulationCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias) + rows = table.view(-1, H3_ADALN_CHUNKS * hidden) + packed = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(temb, act, weight, timestep_indices, token_tags) + ctx.table_shape = table.shape + return packed + + @staticmethod + def backward(ctx, grad_packed): + temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] + num_timesteps, width = ctx.table_shape + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) + d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) + d_table = d_rows.view(num_timesteps, width) + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) + record_backward( + "adaln_modulation_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return ( + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + None, + ) + + +class H3AdaLNModulationCudaOp: + """``adaln_projection_3mod`` + ``adaln_row_gather`` with an FP32 table gradient. + + Returns the six ``(S, H)`` modulation tensors for every packed position, + bitwise equal to running the two ops separately. + """ + + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and adaln_row_gather_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + + def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") + _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + packed = _H3AdaLNModulationCuda.apply( + temb, + weight, + bias, + timestep_indices.contiguous(), + token_tags.contiguous(), + hidden, + ) + return tuple(packed.unbind(0)) + + +def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): + # Reuse the gather's checks against a stand-in of the right row count. + stand_in = torch.empty( + (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device + ) + validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..f02431903 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 three-modality AdaLN projection (RFC #420 ``adaln_projection_3mod``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.adaln_projection import ( + split_adaln_table, + validate_h3_adaln_projection, +) + +KERNEL_ID = "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNProjectionCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + """Project FP32 ``(T, D)`` embeddings to a weight-dtype ``(T, 18H)`` table. + + Save the embedding, rounded SiLU activation and weight for the VJP. + """ + + temb = temb.contiguous() + # Declared mixed-precision boundary: SiLU at temb's FP32 precision, + # then exactly one rounding to the projection dtype. These are the + # provider's own elementwise ops, so the activation is bitwise equal. + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + def backward(ctx, grad_table): + """Return requested embedding, weight and bias gradients in their input dtypes. + + Accumulate projection gradients deterministically and evaluate the + straight-through cast and SiLU VJP in FP32 for the embedding gradient. + """ + + temb, act, weight = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad + grad_table = grad_table.float().contiguous() + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + if need_temb: + # The cast's VJP is the identity; the SiLU VJP runs in FP32. + d_act = det_linear_backward_input(grad_table, weight, torch.float32) + d_temb = silu_backward_fp32(d_act, temb) + record_backward( + "adaln_projection_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return d_temb, dw if need_weight else None, db if need_bias else None + + +class H3AdaLNProjectionCudaOp: + """CUDA candidate: FP32 SiLU, one BF16 cast, deterministic GEMV, six views. + + The projection weight (96768 x 2688 BF16, 520 MB per block) is streamed + once per call; every output element follows contract ``h3-det-linear-v1``, + so a timestep's 3 x 6 modulation rows do not depend on the other + timesteps in the call. The six outputs are views of one ``(T, 6H*3)`` + table, laid out exactly like diffusers' ``view(-1, 6H).chunk(6)``. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six ``(3T, H)`` views after FP32 SiLU and a deterministic projection. + + Require CUDA FP32 ``temb`` of shape ``(T, D)`` with ``T > 0`` and + same-device FP32 or BF16 weight/bias of shapes ``(18H, D)``/``(18H,)``. + Outputs share one table in the weight dtype; autograd covers all inputs. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + table = _H3AdaLNProjectionCuda.apply(temb, weight, bias) + return split_adaln_table(table, hidden) + + def forward_table(self, temb, weight, bias) -> torch.Tensor: + """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" + + validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + return _H3AdaLNProjectionCuda.apply(temb, weight, bias) + + def forward_fp32(self, temb, weight, bias): + """Run the CUDA projection with its declared weight-dtype output boundary.""" + + return self.forward(temb, weight, bias) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..ee681454d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 AdaLN row gather (RFC #420 ``adaln_row_gather``).""" + +from __future__ import annotations + +import torch + +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 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = "rl_engine._C.h3_adaln_row_gather_forward" +BACKWARD_IMPL = "stable_sorted_segment_tiles_fp32_ordered_fold" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_adaln_row_gather_backward" +# Positions per backward tile; part of the fixed summation order. +BACKWARD_TILE = 256 + + +def adaln_row_gather_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_adaln_row_gather_forward") + and hasattr(_C, "h3_adaln_row_gather_backward") + ) + + +def _segment_tiles(index: torch.Tensor, num_rows: int): + """Sorted positions and fixed-size tiles per row segment (all integer ops).""" + + sorted_pos = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows) + seg_start = torch.cumsum(counts, 0) - counts + tiles_per_seg = (counts + BACKWARD_TILE - 1) // BACKWARD_TILE + seg_first_tile = torch.zeros(num_rows + 1, dtype=torch.int64, device=index.device) + seg_first_tile[1:] = torch.cumsum(tiles_per_seg, 0) + tile_seg = torch.repeat_interleave(torch.arange(num_rows, device=index.device), tiles_per_seg) + tile_local = torch.arange(tile_seg.numel(), device=index.device) - seg_first_tile[tile_seg] + tile_begin = seg_start[tile_seg] + tile_local * BACKWARD_TILE + tile_end = torch.minimum(tile_begin + BACKWARD_TILE, seg_start[tile_seg] + counts[tile_seg]) + return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile + + +class _H3AdaLNRowGatherCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, rows, timestep_indices, token_tags): + out = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(timestep_indices, token_tags) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return out + + @staticmethod + def backward(ctx, grad_out): + timestep_indices, token_tags = ctx.saved_tensors + num_rows, rows_dtype = ctx.rows_meta + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_rows) + grad_rows = _C.h3_adaln_row_gather_backward(grad_out.contiguous(), *tiles, rows_dtype) + record_backward( + "adaln_row_gather", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad_rows, None, None + + +class H3AdaLNRowGatherCudaOp: + """CUDA candidate: one launch gathers all six modulation tensors. + + The forward is a byte copy (bitwise equal to six ``index_select`` calls). + The backward replaces ``index_select``'s atomic BF16 scatter-add with a + deterministic FP32 segmented sum in stable sorted order, cast once. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not adaln_row_gather_available(): + raise RuntimeError( + "rl_engine._C lacks h3_adaln_row_gather_*; rebuild with " + "csrc/cuda/h3/adaln_row_gather.cu" + ) + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + if not rows.is_cuda: + raise ValueError("H3AdaLNRowGatherCudaOp needs CUDA tensors") + if rows.stride(1) != 1: + rows = rows.contiguous() + packed = _H3AdaLNRowGatherCuda.apply( + rows, timestep_indices.contiguous(), token_tags.contiguous() + ) + return tuple(packed.unbind(0)) + + def gather_chunks(self, chunks, timestep_indices, token_tags, **kwargs): + """Drop-in for diffusers' six ``(3T, H)`` tensors: concatenate, then gather.""" + + if len(chunks) != H3_ADALN_CHUNKS: + raise ValueError(f"expected {H3_ADALN_CHUNKS} modulation tensors, got {len(chunks)}") + return self.forward(torch.cat(list(chunks), dim=1), timestep_indices, token_tags, **kwargs) + + def forward_fp32(self, rows, timestep_indices, token_tags): + return tuple(out.float() for out in self.forward(rows, timestep_indices, token_tags)) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py new file mode 100644 index 000000000..1274364f2 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python surface of the H3 deterministic row-wise linear (contract h3-det-linear-v1). + +See ``csrc/cuda/h3/det_linear.cu`` for the reduction order. Every helper +raises instead of falling back to cuBLAS when the extension is missing. +""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE + +ACT_NONE = 0 +ACT_SILU = 1 +CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) +_SYMBOLS = ( + "h3_det_linear_forward", + "h3_det_linear_backward_input", + "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", +) + + +def det_linear_available() -> bool: + """Return whether the extension exposes all three H3 deterministic linear APIs.""" + + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +def _require() -> None: + """Raise when a deterministic linear entry point is unavailable.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " + "with csrc/cuda/h3/det_linear.cu (no cuBLAS fallback is allowed)" + ) + + +def det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + activation: int = ACT_NONE, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Project CUDA ``x`` of shape ``(T, K)`` with ``(N, K)`` weights and optional bias. + + Inputs share an FP32 or BF16 dtype/device, ``T > 0``, and ``K`` supports + 16-byte vector loads; BF16 forward needs SM80 or newer. Return ``[out]`` + in the input dtype, plus an FP32 pre-activation when requested. Activation + is identity or SiLU; contiguous copies preserve the reduction contract. + """ + + _require() + return _C.h3_det_linear_forward( + x.contiguous(), + weight.contiguous(), + None if bias is None else bias.contiguous(), + int(activation), + bool(save_pre_activation), + ) + + +def det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Compute a deterministic ``(T, K)`` input VJP on the inputs' CUDA device. + + Convert nonempty ``(T, N)`` gradients to FP32 and combine them with FP32 + or BF16 ``(N, K)`` weights; return FP32 or BF16 as specified by ``out_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_input( + grad.float().contiguous(), weight.contiguous(), out_dtype + ) + + +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + +def det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True +) -> list[torch.Tensor]: + """Fold nonempty CUDA ``(T, N)`` gradients and ``(T, K)`` inputs in row order. + + Convert gradients to FP32; ``x`` is FP32 or BF16 on the same device. + Return ``[dweight]`` of shape ``(N, K)``, optionally followed by ``dbias`` + of shape ``(N,)``, both in the requested FP32 or BF16 ``w_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_weight( + grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) + ) + + +def silu_backward_fp32(grad: torch.Tensor, pre_activation: torch.Tensor) -> torch.Tensor: + """Elementwise SiLU VJP in FP32: g * s * (1 + z * (1 - s)), s = sigmoid(z).""" + + z = pre_activation.float() + s = torch.sigmoid(z) + return grad.float() * (s * (1.0 + z * (1.0 - s))) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..741a0a85e --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 final AdaLN output (RFC #420 ``final_adaln_out``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out +from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_rmsnorm_forward[modulated]" +) +BACKWARD_IMPL = "fp32_table_grad_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_rmsnorm_backward+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3FinalAdaLNOutCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): + hidden = x.shape[-1] + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) # declared boundary: one cast after the FP32 SiLU + (table,) = det_linear_forward(act, weight, bias) + shift, scale = table.chunk(2, dim=-1) + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward( + x2, norm_weight.contiguous(), float(eps), shift, scale, timestep_indices + ) + ctx.save_for_backward(x2, norm_weight, rstd, temb, act, weight, table, timestep_indices) + ctx.x_shape = x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + shift, scale = table.chunk(2, dim=-1) + positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + tiles = _segment_tiles(positions, table.shape[0]) + dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( + grad.contiguous().view_as(x2), + x2, + norm_weight.contiguous(), + rstd, + shift, + scale, + timestep_indices, + *tiles, + ) + # The table gradient stays FP32 into the projection backward. + d_table = torch.cat([d_shift, d_scale], dim=1) + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + record_backward( + "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + + +class H3FinalAdaLNOutCudaOp: + """CUDA candidate: tensor-core shift/scale projection + fused norm/modulation. + + The projection is ``adaln_projection_3mod``'s deterministic GEMV + (``h3-det-linear-bf16-mma-v1``) on ``norm_out.linear``; the norm and + modulation are ``h3_rmsnorm``'s kernel indexed by ``timestep_indices``. + One autograd node, so the table gradient reaches the projection in FP32. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and h3_rmsnorm_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + if weight.dtype != x.dtype: + raise TypeError("norm_out.linear must share the activations' dtype") + if not x.is_cuda: + raise ValueError("H3FinalAdaLNOutCudaOp needs CUDA tensors") + return _H3FinalAdaLNOutCuda.apply( + x, norm_weight, temb, weight, bias, timestep_indices.contiguous(), eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py new file mode 100644 index 000000000..3d040e9e8 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +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 + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + 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, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..d9266e759 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 RMSNorm with fused AdaLN modulation (RFC #420 ``h3_rmsnorm``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +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.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +KERNEL_ID = "rl_engine._C.h3_rmsnorm_forward" +BACKWARD_IMPL = "row_local_dx_tiled_dweight_sorted_table_grads" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_rmsnorm_backward" + + +def h3_rmsnorm_available() -> bool: + return bool( + _EXT_AVAILABLE and hasattr(_C, "h3_rmsnorm_forward") and hasattr(_C, "h3_rmsnorm_backward") + ) + + +class _H3RMSNormCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, index, eps): + hidden = x.shape[-1] + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward(x2, weight.contiguous(), float(eps), shift, scale, index) + ctx.save_for_backward(x2, weight, rstd, shift, scale, index) + ctx.x_shape = x.shape + ctx.modulated = index is not None + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + if not ctx.modulated: + dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) + d_shift = d_scale = None + else: + positions = index.repeat(x2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, shift.shape[0]) + dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( + g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles + ) + d_shift, d_scale = d_shift.to(shift.dtype), d_scale.to(scale.dtype) + record_backward( + "h3_rmsnorm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None + + +class H3RMSNormCudaOp: + """CUDA candidate. + + The statistics replay PyTorch's own RMSNorm reduction order, so the plain + norm is bitwise equal to ``nn.RMSNorm``; the modulated form gathers its + ``shift``/``scale`` rows in the kernel (no ``(S, H)`` intermediates) and + rounds each step where the eager expression does, so it is bitwise equal to + diffusers' ``n * (1.0 + scale[i]) + shift[i]`` as well. Rows are + independent of batch size and position. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_rmsnorm_available(): + raise RuntimeError( + "rl_engine._C lacks h3_rmsnorm_*; rebuild with csrc/cuda/h3/rmsnorm_modulate.cu" + ) + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + return _H3RMSNormCuda.apply(x, weight, None, None, None, eps) + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + return self.forward(x, weight, eps).float() 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..37cea620b --- /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) + if num_rows is None: + raise ValueError("indexed SP plans require num_rows") + *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()) + 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 num_rows is None: + raise ValueError("indexed SP validation requires num_rows") + 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/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..4d58f6d9f --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 FP32 timestep MLP (RFC #420 ``timestep_mlp_fp32``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + ACT_SILU, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3TimestepMLPCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2): + """Evaluate the FP32 CUDA linear-SiLU-linear graph and save VJP inputs.""" + + x = x.contiguous() + hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) + (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) + ctx.save_for_backward(x, w1, w2, pre, hidden) + return out + + @staticmethod + def backward(ctx, grad_out): + """Return requested FP32 input and parameter VJPs with fixed reduction order.""" + + x, w1, w2, pre, hidden = ctx.saved_tensors + need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad + grad_out = grad_out.float().contiguous() + dw2 = db2 = dw1 = db1 = dx = None + if need_w2 or need_b2: + dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + d_pre = None + if need_x or need_w1 or need_b1: + d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) + d_pre = silu_backward_fp32(d_hidden, pre).contiguous() + if need_w1 or need_b1: + dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + if need_x: + dx = det_linear_backward_input(d_pre, w1, torch.float32) + record_backward( + "timestep_mlp_fp32", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return ( + dx, + dw1 if need_w1 else None, + db1 if need_b1 else None, + dw2 if need_w2 else None, + db2 if need_b2 else None, + ) + + +class H3TimestepMLPCudaOp: + """CUDA candidate: two warp-per-column GEMVs, SiLU fused into the first. + + The weights dominate the traffic (5.5 MB + 57.8 MB FP32), the rows are the + handful of distinct timesteps, and every output element follows contract + ``h3-det-linear-v1``, so a timestep's ``temb`` bytes do not depend on which + other timesteps share the call. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` timestep embeddings for same-device CUDA inputs.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Apply two deterministic FP32 CUDA projections with an intervening SiLU. + + Require same-device FP32 tensors: ``x`` is nonempty ``(T, K)``, + weights are ``(H, K)`` and ``(D, H)``, and biases are ``(H,)`` and + ``(D,)``. Return ``(T, D)`` with autograd support for all five inputs. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + if not x.is_cuda: + raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") + return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Use the standard CUDA MLP, whose inputs and output are already FP32.""" + + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..65934e1e9 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 sinusoidal timestep features (RFC #420 ``timestep_sinusoid_h3``).""" + +from __future__ import annotations + +import torch + +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 import H3_FREQ_DIM, H3_MAX_PERIOD +from rl_engine.reference.minimax_h3.fixed_order import tree_sum_lastdim_fp32 +from rl_engine.reference.minimax_h3.timestep_sinusoid import ( + NativeH3TimestepSinusoidOp, + validate_h3_timesteps, +) + +KERNEL_ID = "rl_engine._C.h3_timestep_sinusoid_forward" +BACKWARD_IMPL = "row_local_fp32_analytic_tree_sum" + + +def h3_sinusoid_cuda_available() -> bool: + """Return whether the extension provides the CUDA H3 sinusoid entry point.""" + + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +class _H3TimestepSinusoidCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): + """Compute FP32 CUDA features and save them for the analytic timestep VJP.""" + + t32 = timestep.detach().float().contiguous() + out = _C.h3_timestep_sinusoid_forward( + t32, int(num_channels), float(H3_MAX_PERIOD), check_range + ) + ctx.save_for_backward(out) + ctx.num_channels = int(num_channels) + ctx.timestep_dtype = timestep.dtype + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + """Reduce channel derivatives with a fixed FP32 tree into the timestep dtype.""" + + (out,) = ctx.saved_tensors + half = ctx.num_channels // 2 + freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) + cos_part, sin_part = out[:, :half], out[:, half:] + g = grad_out.float() + # d cos(t f)/dt = -f sin(t f); d sin(t f)/dt = f cos(t f). Row-local, + # summed over channels with a fixed pairwise tree. + contrib = torch.cat( + [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 + ) + grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) + record_backward( + "timestep_sinusoid_h3", + kernel_id="rl_engine.reference.minimax_h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None, None + + +class H3TimestepSinusoidCudaOp: + """CUDA candidate: one thread per (timestep, frequency), FP32 output. + + Construction fails when the extension lacks the symbol, so the registry + falls back to the PyTorch reference instead of silently mis-dispatching. + """ + + op_class = "elementwise" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + """Require the native H3 CUDA sinusoid symbol before creating this operator.""" + + if not h3_sinusoid_cuda_available(): + raise RuntimeError( + "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu (it is omitted " + "from KERNEL_ALIGN_USE_FAST_MATH=1 builds to keep bitwise parity)" + ) + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 CUDA features for timesteps in ``[0, 1]`` with gradients.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on CUDA. + + Require nonempty 1-D FP32, FP16 or BF16 timesteps and a positive even + channel count. Native validation rejects nonfinite/out-of-range values + when ``check_range`` is true; backward returns the original input dtype. + """ + + # Native code owns value validation, avoiding a second host sync here. + validate_h3_timesteps(timestep, num_channels, check_range=False) + if not timestep.is_cuda: + raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") + return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) + + def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Use the standard CUDA sinusoid path, which already returns FP32 features.""" + + # The op is FP32 end to end; the FP32 path is the op itself. + return self.forward(timestep, num_channels=num_channels) 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 new file mode 100644 index 000000000..c4ceed334 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +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): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/reference/minimax_h3/__init__.py b/rl_engine/reference/minimax_h3/__init__.py new file mode 100644 index 000000000..a01f9218b --- /dev/null +++ b/rl_engine/reference/minimax_h3/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 (RFC #420) PyTorch goldens. + +Shape constants are pinned to ``MiniMaxAI/MiniMax-H3@42ed227`` (see +``rl_engine/validation/models/h3_manifest.json``). The ops accept other sizes so that +tiny synthetic shapes can be tested, but the H3 layout rules (three modality +rows per timestep, six modulation chunks) are fixed. +""" + +from __future__ import annotations + +H3_FREQ_DIM = 256 +H3_MAX_PERIOD = 10000 +H3_TIME_EMBED_HIDDEN_DIM = 5376 +H3_TIME_EMBED_DIM = 2688 +H3_HIDDEN_SIZE = 5376 +# Modality tags of the packed sequence: 0 video, 1 text, 2 audio. +H3_MODALITY_NUM = 3 +# shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp +H3_ADALN_CHUNKS = 6 +H3_ADALN_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..be83239c5 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 three-modality AdaLN projection (RFC #420 row ``adaln_projection_3mod``). + +One ``MiniMaxH3AdaLayerNormModulation`` per transformer block:: + + act = silu(temb).to(weight.dtype) SiLU in FP32, one cast to BF16 + table = act @ W.T + b (T, 6 * H * 3), BF16 + rows = table.view(3 * T, 6 * H) row t * 3 + m, m in {video, text, audio} + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) + +Output channel ``o = m * 6H + c * H + h`` of the projection is chunk ``c``, hidden +index ``h`` of modality ``m``. Each of the six outputs has shape ``(3T, H)``, +and row ``t * 3 + m`` is what ``timestep_indices * 3 + token_tags`` addresses. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +def adaln_hidden_size(weight: torch.Tensor) -> int: + """Infer ``H`` from a 2-D projection weight whose row count is divisible by 18.""" + + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM + if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: + raise ValueError( + f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " + f"dim is not a multiple of {rows_per_hidden}" + ) + return weight.shape[0] // rows_per_hidden + + +def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """Check the AdaLN input shapes, shared device and precision boundary; return ``H``. + + Require nonempty FP32 ``temb`` of shape ``(T, D)`` and FP32 or BF16 weight + of shape ``(18H, D)`` with matching-dtype bias of shape ``(18H,)``. + """ + + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != temb.device: + raise ValueError(f"{name} is on {tensor.device}, temb is on {temb.device}") + if temb.dtype != torch.float32: + # RFC probe H7: the activation must run at the timestep-embedding + # precision; a BF16 temb means the cast happened before the SiLU. + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in _WEIGHT_DTYPES or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" + ) + hidden = adaln_hidden_size(weight) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + if bias.shape != (weight.shape[0],): + raise ValueError(f"bias must be ({weight.shape[0]},), got {tuple(bias.shape)}") + return hidden + + +def split_adaln_table(table: torch.Tensor, hidden: int) -> tuple[torch.Tensor, ...]: + """(T, 6H*3) projection output -> six (3T, H) views, diffusers order.""" + + return table.view(-1, H3_ADALN_CHUNKS * hidden).chunk(H3_ADALN_CHUNKS, dim=-1) + + +class NativeH3AdaLNProjectionOp: + """PyTorch reference for the H3 AdaLN projection. + + ``forward`` is the provider path (``F.silu`` in FP32, cast, ``F.linear``). + ``forward_fp32`` is the golden: the SiLU in FP64, rounded once to the + weight dtype at the declared boundary (the cast is model semantics, not + an implementation detail), then the projection in FP64. The six outputs are + returned in FP32 without the final BF16 rounding. Its gradient is FP64 + end to end: the cast is applied straight-through (identity VJP). + """ + + op_class = "reduction" + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` views using the PyTorch provider path.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Apply FP32 SiLU, cast to the weight dtype and project on the input device. + + Validate ``(T, D)`` embeddings and ``(18H, D)``/``(18H,)`` parameters; + return six weight-dtype ``(3T, H)`` views with PyTorch autograd support. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + return split_adaln_table(table, hidden) + + def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six FP32 golden outputs from FP64 SiLU and projection arithmetic. + + Preserve the declared activation rounding to the weight dtype while + using an identity VJP at that cast boundary for the FP64 golden graph. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # The declared cast rounds the FP32 SiLU, as the model (and the kernel) + # computes it: an FP64 SiLU can round to the other weight-dtype neighbour + # near a tie and shift a whole output column. Its VJP is the identity, + # written straight-through so autograd does not round the FP64 gradient + # to the weight dtype on its way back (``.to(bf16)`` would). + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/rl_engine/reference/minimax_h3/adaln_row_gather.py b/rl_engine/reference/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..2cd713d45 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_row_gather.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 AdaLN row gather (RFC #420 row ``adaln_row_gather``). + +H3 computes one row index per packed-sequence position and uses it to pick +every position's six modulation vectors:: + + adaln_indices = timestep_indices * 3 + token_tags (S,) + shift_msa[adaln_indices], scale_msa[adaln_indices], ... six (S, H) + +``rows`` is the ``(3T, 6H)`` view of one block's AdaLN projection +(``table.view(-1, 6H)``); its six column blocks are the six modulation +tensors. ``timestep_indices`` and ``token_tags`` are semantic inputs +(RFC #420 §4): out-of-range values are an error, never clamped. +""" + +from __future__ import annotations + +import torch + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_INDEX_DTYPES = (torch.int64, torch.int32) + + +def validate_h3_adaln_row_gather( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + *, + check_range: bool = True, +) -> int: + """Return the hidden size; raise on anything H3 does not define.""" + + if not isinstance(rows, torch.Tensor) or rows.dim() != 2: + raise ValueError("rows must be a 2-D (3T, 6H) tensor") + if not rows.is_floating_point(): + raise TypeError(f"rows must be floating point, got {rows.dtype}") + num_rows, width = rows.shape + if num_rows == 0 or num_rows % H3_MODALITY_NUM != 0: + raise ValueError(f"rows must hold 3 rows per timestep, got {num_rows}") + if width == 0 or width % H3_ADALN_CHUNKS != 0: + raise ValueError(f"rows width must be 6 * H, got {width}") + for name, index in (("timestep_indices", timestep_indices), ("token_tags", token_tags)): + if not isinstance(index, torch.Tensor) or index.dim() != 1: + raise ValueError(f"{name} must be a 1-D (S,) tensor") + if index.dtype not in _INDEX_DTYPES: + raise TypeError(f"{name} must be int64 or int32, got {index.dtype}") + if index.device != rows.device: + raise ValueError(f"{name} is on {index.device}, rows are on {rows.device}") + if timestep_indices.shape != token_tags.shape: + raise ValueError( + f"timestep_indices {tuple(timestep_indices.shape)} and token_tags " + f"{tuple(token_tags.shape)} must have the same length" + ) + if timestep_indices.numel() == 0: + raise ValueError("the packed sequence must not be empty") + if timestep_indices.dtype != token_tags.dtype: + raise TypeError("timestep_indices and token_tags must share an integer dtype") + if check_range: + num_timesteps = num_rows // H3_MODALITY_NUM + bad = ( + (timestep_indices < 0) + | (timestep_indices >= num_timesteps) + | (token_tags < 0) + | (token_tags >= H3_MODALITY_NUM) + ) + if bool(bad.any()): # one host sync + raise IndexError( + f"timestep_indices must lie in [0, {num_timesteps}) and token_tags in " + f"[0, {H3_MODALITY_NUM}) (0 video, 1 text, 2 audio)" + ) + return width // H3_ADALN_CHUNKS + + +def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) -> torch.Tensor: + return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() + + +class _DeterministicRowGather(torch.autograd.Function): + """``index_select`` forward; FP32 per-row segment sums for the backward. + + ``index_select``'s own backward is an atomic scatter-add in the grad dtype + (BF16 for H3), which is neither deterministic nor accurate. There are + only 3T table rows, so the reference sums each row's positions with one + FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. + """ + + @staticmethod + def forward(ctx, rows, index): + ctx.save_for_backward(index) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return torch.stack( + [chunk.index_select(0, index) for chunk in rows.chunk(H3_ADALN_CHUNKS, dim=-1)] + ) + + @staticmethod + def backward(ctx, grad): + (index,) = ctx.saved_tensors + num_rows, dtype = ctx.rows_meta + grad32 = grad.float() # (6, S, H) + out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) + for row in range(num_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + out[row] = grad32.index_select(1, positions).sum(dim=1) + return out.reshape(num_rows, -1).to(dtype), None + + +class NativeH3AdaLNRowGatherOp: + """PyTorch reference: ``index_select`` on each of the six column blocks. + + The forward is byte-identical to diffusers; the backward is deterministic + (see ``_DeterministicRowGather``). The raw diffusers path, including its + atomic backward, is ``rl_engine.validation.models.h3_provider.provider_adaln_row_gather``. + ``forward_fp32`` returns the same rows in FP32 with an FP64 gradient. + + The forward is a copy, but the VJP is a segmented reduction over the + packed sequence, so the op is judged with the ``reduction`` tolerance + class; forward bitwise equality is asserted separately. + """ + + op_class = "reduction" + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple(_DeterministicRowGather.apply(rows, index).unbind(0)) + + def forward_fp32(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple( + chunk.index_select(0, index).float() + for chunk in rows.double().chunk(H3_ADALN_CHUNKS, dim=-1) + ) diff --git a/rl_engine/reference/minimax_h3/block.py b/rl_engine/reference/minimax_h3/block.py new file mode 100644 index 000000000..427d7de2e --- /dev/null +++ b/rl_engine/reference/minimax_h3/block.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 golden of one MiniMax-H3 transformer block (RFC #420 ``ws1_one_h3_block``). + +Every node takes FP32 inputs (BF16 weights and activations upcast exactly), +computes in FP32 with TF32 off, and returns FP32 without rounding back to +BF16, so the golden is the block's arithmetic without the BF16 storage the +provider and the candidate both have. It is the yardstick their errors are +measured against, not a bitwise target. The one declared cast kept from the +checkpoint contract is the AdaLN input: ``silu(temb)`` is rounded to BF16 +before the projection, as every implementation must do it. + +The partial MM-RoPE uses exact FP32 ``cos``/``sin`` tables on the leading 96 +of 128 head channels (the provider rounds the tables to BF16 first). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +H3_HEADS = 56 +H3_HEAD_DIM = 128 +H3_ROTARY_DIM = 96 +H3_NORM_EPS = 1e-5 + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +def golden_adaln_tables(temb, weight, bias, hidden_size: int = 5376) -> tuple[torch.Tensor, ...]: + """Six ``(3T, H)`` FP32 tables; ``silu(temb)`` rounded to the weight dtype first.""" + + act = F.silu(_f32(temb)).to(weight.dtype) + table = F.linear(_f32(act), _f32(weight), _f32(bias)) + return table.view(-1, 6 * hidden_size).chunk(6, dim=-1) + + +def golden_rmsnorm(x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + x = _f32(x) + rstd = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + return x * rstd * _f32(weight) + + +def golden_norm_modulate(x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + n = golden_rmsnorm(x, weight, eps) + return n * (1.0 + _f32(scale).index_select(0, index)) + _f32(shift).index_select(0, index) + + +def golden_linear(x, weight) -> torch.Tensor: + return F.linear(_f32(x), _f32(weight)) + + +def golden_rope_tables(position_ids, rope_freq_dim: int = 16, rope_theta: float = 10000.0): + """FP32 ``(S, 96)`` cos/sin, the same angles as the provider's rotary module.""" + + exponent = torch.arange( + 0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device + ) + inv_freq = 1.0 / (rope_theta ** (exponent / (2 * rope_freq_dim))) + freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs = torch.cat(freqs.unbind(dim=1), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def golden_apply_rotary(x, cos, sin) -> torch.Tensor: + """Rotate-half on the leading ``cos.shape[-1]`` channels of ``(B, S, H, D)``; FP32 tables.""" + + x = _f32(x) + rot = cos.shape[-1] + head, tail = x[..., :rot], x[..., rot:] + x1, x2 = head.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + head = head * cos[None, :, None, :] + rotated * sin[None, :, None, :] + return torch.cat((head, tail), dim=-1) + + +def golden_attention(q, k, v) -> torch.Tensor: + """Non-causal full attention over one packed document, ``(B, S, H, D)`` -> ``(B, S, H*D)``.""" + + q, k, v = (_f32(t).permute(0, 2, 1, 3) for t in (q, k, v)) + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.shape[-1]) + out = torch.matmul(torch.softmax(scores, dim=-1), v) + return out.permute(0, 2, 1, 3).flatten(2, 3) + + +def golden_swiglu(projected) -> torch.Tensor: + hidden, gate = _f32(projected).chunk(2, dim=-1) + return hidden * F.silu(gate) + + +def golden_gate_residual(residual, gate, index, y) -> torch.Tensor: + return _f32(residual) + _f32(gate).index_select(0, index) * _f32(y) diff --git a/rl_engine/reference/minimax_h3/final_adaln_out.py b/rl_engine/reference/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..a21a5cbdb --- /dev/null +++ b/rl_engine/reference/minimax_h3/final_adaln_out.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 final AdaLN output (RFC #420 row ``final_adaln_out``). + +``MiniMaxH3AdaLayerNormOut`` runs after the 50 blocks:: + + shift, scale = linear(silu(temb).to(bf16)).chunk(2) # (T, H) each, shift first + out = norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] + +``linear`` is ``norm_out.linear`` (``2688 -> 2 * 5376``, BF16), ``norm`` is +``norm_out.norm`` (``nn.RMSNorm(5376, eps=1e-5)``), and the table is indexed +by ``timestep_indices`` (one row per distinct timestep), not by +``adaln_indices``. Diffusers then casts the result to the FP32 output heads' +dtype, an exact upcast left to the caller. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + NativeH3RMSNormOp, + validate_h3_rmsnorm, +) + + +def validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, *, same_dtype: bool = True +) -> int: + hidden = x.shape[-1] + if same_dtype: + validate_h3_rmsnorm(x, norm_weight, eps) + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if temb.dim() != 2 or temb.shape[0] == 0: + raise ValueError(f"temb must be a non-empty (T, D) matrix, got {tuple(temb.shape)}") + if weight.shape != (2 * hidden, temb.shape[1]) or bias.shape != (2 * hidden,): + raise ValueError( + f"norm_out.linear must be ({2 * hidden}, {temb.shape[1]}) with a ({2 * hidden},) " + f"bias, got {tuple(weight.shape)} / {tuple(bias.shape)}" + ) + if bias.dtype != weight.dtype: + raise TypeError("norm_out.linear weight and bias must share a dtype") + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if timestep_indices.dtype != torch.int64 or timestep_indices.dim() != 1: + raise TypeError("timestep_indices must be a 1-D int64 tensor") + if x.dim() < 2 or x.shape[-2] != timestep_indices.shape[0]: + raise ValueError("x must be (..., S, H) with S = len(timestep_indices)") + if bool(((timestep_indices < 0) | (timestep_indices >= temb.shape[0])).any()): + raise IndexError(f"timestep_indices must lie in [0, {temb.shape[0]})") + return hidden + + +class NativeH3FinalAdaLNOutOp: + """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only where + the model itself stores a value in its own dtype (the SiLU cast, the BF16 + table, the ``norm_out.norm`` output and ``1 + scale``), straight-through for + the gradient, and computes everything else in FP64.""" + + op_class = "reduction" + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + return NativeH3RMSNormOp().forward_modulated( + x, norm_weight, shift, scale, timestep_indices, eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, same_dtype=False + ) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # Declared cast of the FP32 SiLU, as the model computes it (an FP64 SiLU can + # round to the other neighbour near a tie); identity VJP. + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the + # weight dtype, and every position of a timestep shares that row, so the + # rounding is model semantics (straight-through for the gradient). + table = table + (table.to(weight.dtype).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() + # norm_out.norm is an x.dtype module and diffusers forms 1 + scale in + # x.dtype, so both are stored rounded. 1 + scale is shared by every + # position of a timestep: left unrounded, its error is systematic and + # grows with S in d_norm_weight; n's rounding enters d_scale (and so + # dW, d_temb) summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.index_select(0, timestep_indices) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.index_select(0, timestep_indices) + return out.float() diff --git a/rl_engine/reference/minimax_h3/fixed_order.py b/rl_engine/reference/minimax_h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/reference/minimax_h3/fixed_order.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shape-independent FP32 reductions shared by the H3 goldens and backwards. + +``torch.sum`` picks its reduction tree from the tensor shape, so a row's +result can change with the number of rows next to it. These helpers spell +the order out with elementwise ops only, which keeps every row's bytes a +function of that row alone. +""" + +from __future__ import annotations + +import torch + + +def tree_sum_lastdim_fp32(x: torch.Tensor) -> torch.Tensor: + """Pairwise tree over the last dim: pad to a power of two, halve until 1. + + Level ``l`` adds element ``i`` and ``i + width/2`` for every ``i``. The + tree depends only on the last-dim size. + """ + + x = x.float() + width = x.shape[-1] + if width == 0: + return x.new_zeros(x.shape[:-1]) + padded = 1 << (width - 1).bit_length() + if padded != width: + x = torch.nn.functional.pad(x, (0, padded - width)) + while x.shape[-1] > 1: + half = x.shape[-1] // 2 + x = x[..., :half] + x[..., half:] + return x[..., 0] + + +def fold_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left fold over dim 0 in ascending row order: ((r0 + r1) + r2) + ...""" + + rows = rows.float() + if rows.shape[0] == 0: + return rows.new_zeros(rows.shape[1:]) + acc = rows[0].clone() + for index in range(1, rows.shape[0]): + acc = acc + rows[index] + return acc diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py new file mode 100644 index 000000000..f68a0fdf4 --- /dev/null +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.validation.models.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..1727fd1d0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 RMSNorm and its fused AdaLN modulation (RFC #420 row ``h3_rmsnorm``). + +Every H3 RMSNorm is ``nn.RMSNorm(5376, eps=1e-5)`` with an affine weight in the +block dtype: the block ``norm1``/``norm2``, the token-refiner norms and the +final ``norm_out.norm``. Inside a transformer block (and in ``norm_out``) the +normalised rows are immediately modulated by per-row AdaLN parameters:: + + n = rms_norm(x, weight, eps) + out = n * (1.0 + scale[index]) + shift[index] # eager, rounded after every op + +``index`` is ``adaln_indices`` in a block and ``timestep_indices`` in +``norm_out``; ``shift``/``scale`` are ``(R, H)`` rows (views of the AdaLN +table). ``x`` is ``(..., S, H)`` and ``index`` has one entry per position ``S``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +H3_NORM_EPS = 1e-5 +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> int: + for name, tensor in (("x", x), ("weight", weight)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype not in _DTYPES: + raise TypeError(f"{name} must be fp32, bf16 or fp16, got {tensor.dtype}") + if weight.dtype != x.dtype: + raise TypeError(f"weight dtype {weight.dtype} must match x dtype {x.dtype}") + if weight.device != x.device: + raise ValueError(f"weight is on {weight.device}, x is on {x.device}") + if x.dim() < 2 or x.numel() == 0: + raise ValueError(f"x must be a non-empty (..., N) tensor, got {tuple(x.shape)}") + if weight.shape != (x.shape[-1],): + raise ValueError(f"weight must be ({x.shape[-1]},), got {tuple(weight.shape)}") + if not eps > 0: + raise ValueError(f"eps must be positive, got {eps}") + return x.shape[-1] + + +def validate_h3_modulation( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + hidden = x.shape[-1] + for name, tensor in (("shift", shift), ("scale", scale)): + if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): + raise TypeError(f"{name} must match x's dtype and device") + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") + if shift.shape != scale.shape: + raise ValueError("shift and scale must have the same shape") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != x.device: + raise TypeError("index must be a 1-D int64 tensor on x's device") + if x.dim() < 2 or x.shape[-2] != index.shape[0]: + raise ValueError( + f"x must be (..., S, {hidden}) with S = len(index) = {index.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if check_range and bool(((index < 0) | (index >= shift.shape[0])).any()): + raise IndexError(f"index must lie in [0, {shift.shape[0]})") + + +class _DeterministicModulation(torch.autograd.Function): + """Eager ``n * (1.0 + scale[i]) + shift[i]``; the ``shift``/``scale`` + gradients as FP32 per-row sums instead of ``index_select``'s atomic + scatter-add (``d_n`` is the eager VJP).""" + + @staticmethod + def forward(ctx, n, shift, scale, index): + t1 = 1.0 + scale.index_select(0, index) + ctx.save_for_backward(n, t1, index) + ctx.rows = shift.shape[0] + return n * t1 + shift.index_select(0, index) + + @staticmethod + def backward(ctx, grad): + n, t1, index = ctx.saved_tensors + hidden = n.shape[-1] + g = grad.float().reshape(-1, index.shape[0], hidden) + per_shift = g.sum(dim=0) + per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) + d_shift = per_shift.new_zeros((ctx.rows, hidden)) + d_scale = per_scale.new_zeros((ctx.rows, hidden)) + for row in range(ctx.rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + return grad * t1, d_shift.to(t1.dtype), d_scale.to(t1.dtype), None + + +class NativeH3RMSNormOp: + """PyTorch reference. + + ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and + the eager modulation expression, with a deterministic shift/scale gradient; + the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / + ``forward_modulated_fp32`` are the golden: the same math in FP64, returned in + FP32. The modulated golden rounds, straight-through, only the two values the + model stores in ``x``'s dtype: ``norm(x)`` and ``1 + scale``. + """ + + op_class = "reduction" + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + hidden = validate_h3_rmsnorm(x, weight, eps) + return F.rms_norm(x, (hidden,), weight, eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) + return (x64 * rstd * weight.double()).float() + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + n = self.forward(x, weight, eps) + return _DeterministicModulation.apply(n, shift, scale, index) + + def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + # The golden may be fed FP32 golden modulation rows for a BF16 x. + validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() + # The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype. + # 1 + scale is shared by every position of a row, so leaving it unrounded + # skews d_norm_weight systematically with S; n's rounding enters d_scale + # summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.double().index_select(0, index) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.double().index_select(0, index) + return out.float() diff --git a/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..51ef48db7 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_mlp.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 timestep MLP (RFC #420 row ``timestep_mlp_fp32``). + +``TimestepEmbedding(in_channels=256, time_embed_dim=5376, out_dim=2688)``:: + + temb = linear_2(silu(linear_1(features))) 256 -> 5376 -> 2688 + +``time_embedder`` is in ``_keep_in_fp32_modules``: weights, biases, +activations and the output are FP32, and ``temb`` stays FP32 because every +AdaLN projection applies its own SiLU to it before casting (RFC #420 §4). +A BF16 argument is a contract violation (RFC probe H7), not a lower-precision +variant, so it is rejected. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def validate_h3_timestep_mlp( + x: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> None: + """Check a same-device FP32 ``(T, K) -> (T, H) -> (T, D)`` MLP contract. + + Require ``T > 0``, weights of shapes ``(H, K)``/``(D, H)`` and biases + of shapes ``(H,)``/``(D,)``; reject lower-precision parameters or inputs. + """ + + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} + for name, tensor in named.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError( + f"{name} must be float32: the H3 time_embedder is an FP32 module " + f"(got {tensor.dtype})" + ) + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") + hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if k_in != x.shape[1]: + raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") + if b1.shape != (hidden,): + raise ValueError(f"b1 must be ({hidden},), got {tuple(b1.shape)}") + if w2.dim() != 2 or w2.shape[1] != hidden: + raise ValueError(f"w2 must be (D, {hidden}), got {tuple(w2.shape)}") + if b2.shape != (w2.shape[0],): + raise ValueError(f"b2 must be ({w2.shape[0]},), got {tuple(b2.shape)}") + + +class NativeH3TimestepMLPOp: + """PyTorch reference for the H3 timestep MLP. + + ``forward`` is the provider path (``nn.Linear`` / ``F.silu`` in FP32; run + it with TF32 disabled, as the RL-Kernel contract requires). + ``forward_fp32`` is the independent golden: the same graph in FP64, + rounded once to FP32, so its own rounding error (~1e-16 relative) is far + below the FP32 contract and its reduction order does not matter. + """ + + op_class = "reduction" + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` embeddings using the PyTorch provider graph.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Validate and evaluate linear-SiLU-linear in FP32 on the input device. + + Inputs follow the ``(T, K) -> (T, H) -> (T, D)`` contract and retain + PyTorch autograd support; CUDA callers must disable TF32 for this path. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Evaluate the validated MLP in FP64 and round its ``(T, D)`` output to FP32.""" + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + z = F.linear(x.double(), w1.double(), b1.double()) + h = z * torch.sigmoid(z) + return F.linear(h, w2.double(), b2.double()).float() diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..0d95e58f0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 sinusoidal timestep features (RFC #420 row ``timestep_sinusoid_h3``). + +H3 builds ``Timesteps(num_channels=256, flip_sin_to_cos=True, +downscale_freq_shift=0)`` and feeds it the *distinct* timestep values of the +packed sequence, unscaled in ``[0, 1]`` (``t = 1 - sigma``):: + + freq[k] = exp(-ln(10000) * k / 128) k = 0..127, FP32 + arg[t, k] = t * freq[k] FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + +The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` +module in the checkpoint, so the features never pass through BF16. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD + +_TIMESTEP_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_timesteps( + timestep: torch.Tensor, num_channels: int, *, check_range: bool = True +) -> None: + """Fail closed on inputs the H3 convention does not define. + + ``check_range`` rejects values outside ``[0, 1]``: H3 consumes + ``t = 1 - sigma`` unscaled, so a ``t * 1000`` caller (RFC probe H10) is a + convention error, not a different valid input. It reads the tensor back + to the host, which is cheap for the handful of distinct timesteps. + """ + + if not isinstance(timestep, torch.Tensor): + raise TypeError("timestep must be a torch.Tensor") + if timestep.dim() != 1: + raise ValueError(f"timestep must be 1-D (num_timesteps,), got {tuple(timestep.shape)}") + if timestep.numel() == 0: + raise ValueError("timestep must hold at least one timestep") + if timestep.dtype not in _TIMESTEP_DTYPES: + raise TypeError(f"timestep must be fp32, bf16 or fp16, got {timestep.dtype}") + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") + if check_range: + t32 = timestep.detach().float() + # One host sync: NaN fails both comparisons, so this also rejects it. + if not bool(((t32 >= 0) & (t32 <= 1)).all()): + raise ValueError( + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + f"unscaled (got {t32.cpu().tolist()[:8]})" + ) + + +class NativeH3TimestepSinusoidOp: + """PyTorch reference for the H3 sinusoidal timestep features. + + ``forward`` replays diffusers' ``get_timestep_embedding`` op for op on the + input's device (the provider path). ``forward_fp32`` is the independent + golden: the same formula evaluated in FP64 and rounded once to FP32. + """ + + op_class = "elementwise" + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 sinusoidal features for nonempty timesteps in ``[0, 1]``.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on the input device. + + Require nonempty 1-D FP32, BF16 or FP16 timesteps and positive even + channels. Validate finite values in ``[0, 1]`` when ``check_range`` is + true and preserve PyTorch autograd through the provider formula. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + # Same op sequence and dtypes as diffusers get_timestep_embedding with + # flip_sin_to_cos=True, downscale_freq_shift=0, scale=1. + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / half + freq = torch.exp(exponent) + arg = timestep[:, None].float() * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1) + + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Evaluate the validated sinusoid formula in FP64 and round once to FP32. + + Return ``(T, num_channels)`` on the timestep device with cosine channels + followed by sine channels and gradients through the FP64 golden graph. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + k = torch.arange(half, dtype=torch.float64, device=timestep.device) + freq = torch.exp(-math.log(H3_MAX_PERIOD) * k / half) + arg = timestep.double()[:, None] * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1).float() + + @staticmethod + def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: + """The FP32 frequency table the provider path multiplies by.""" + + half = num_channels // 2 + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=device + ) + return torch.exp(exponent / half) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 528c77b98..a733e95ce 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -204,6 +204,43 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + # MiniMax-H3 (RFC #420) conditioning path + PYTORCH_H3_TIMESTEP_SINUSOID = ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + PYTORCH_H3_TIMESTEP_MLP = "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp" + CUDA_H3_TIMESTEP_MLP = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" + ) + PYTORCH_H3_ADALN_PROJECTION = ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) + PYTORCH_H3_ADALN_ROW_GATHER = ( + "rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ) + CUDA_H3_ADALN_ROW_GATHER = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" + ) + PYTORCH_H3_RMSNORM = "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp" + CUDA_H3_RMSNORM = "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp" + CUDA_H3_GATE_RESIDUAL = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp" + ) + PYTORCH_H3_FINAL_ADALN_OUT = ( + "rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp" + ) + CUDA_H3_FINAL_ADALN_OUT = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out.H3FinalAdaLNOutCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -502,6 +539,7 @@ class KernelRegistry: """ def __init__(self): + """Initialize backend caches, capability contracts, and device-specific priorities.""" self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) @@ -673,6 +711,31 @@ def __init__(self): OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], + "timestep_sinusoid_h3": [ + OpBackend.CUDA_H3_TIMESTEP_SINUSOID, + OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, + ], + "timestep_mlp_fp32": [ + OpBackend.CUDA_H3_TIMESTEP_MLP, + OpBackend.PYTORCH_H3_TIMESTEP_MLP, + ], + "adaln_projection_3mod": [ + OpBackend.CUDA_H3_ADALN_PROJECTION, + OpBackend.PYTORCH_H3_ADALN_PROJECTION, + ], + "adaln_row_gather": [ + OpBackend.CUDA_H3_ADALN_ROW_GATHER, + OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, + ], + "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -714,6 +777,13 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "musa": { "logp": [ @@ -763,6 +833,13 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -789,6 +866,13 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py new file mode 100644 index 000000000..b20e885a1 --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,718 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 transformer block, node by node (RFC #420 ``ws1_one_h3_block``). + +The block is a fixed graph of 17 nodes (``NODES``), from the AdaLN projection +through attention and the FFN to the second gated residual, on block 0's +pinned weights and a packed FL2VA layout. Every node runs three ways: + +* provider: the diffusers op sequence (``h3_provider``), fed provider inputs; +* candidate: the RL-Kernel binding of the node, fed candidate inputs + (chained) and, separately, provider inputs (isolated); +* golden: the FP32 reference (``rl_engine.reference.minimax_h3.block``). + +Each node also names the RFC row that owns its arithmetic and how it is bound +today (``Binding.status``): + +* ``row`` -- that row's operator, in this PR stack; +* ``interim`` -- an existing deterministic RL-Kernel operator standing in + until the owning row lands (its arithmetic, speed and rounding are that + operator's, not the row's); +* ``reference`` -- no kernel exists yet; the node runs the provider replay + itself, so it is trivially bitwise and contributes no RL-Kernel evidence. + +The report gives, per node: repeat and batch-row bitwise checks, chained and +isolated agreement with the provider, error against the golden for both the +candidate and the provider, and the backend that ran. ``first_*_drift`` name +the first node (graph order) where a comparison stops being bitwise; the +backward report does the same for the gradients reaching each node's output, +in reverse graph order. +""" + +from __future__ import annotations + +import hashlib +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable + +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout, h3_timesteps + +HIDDEN = 5376 +HEADS = 56 +HEAD_DIM = 128 +BLOCK = "transformer_blocks.0" +NUM_TIMESTEPS = 2 + + +@dataclass(frozen=True) +class Binding: + rfc_row: str + status: str # "row" | "interim" | "reference" + backend: str + + +@dataclass(frozen=True) +class Node: + name: str + inputs: tuple[str, ...] + binding: Binding + # (ops, ctx, *inputs) -> output + candidate: Callable[..., Any] + # (ctx, *inputs) -> output + provider: Callable[..., Any] + golden: Callable[..., Any] + # Promise checked by the tests: the isolated candidate equals the provider bitwise. + provider_bitwise_isolated: bool + + +class CandidateOps: + """The operators bound to the block's nodes, resolved once per run. + + The three rows of this stack are resolved through the registry, so the + report records what dispatch actually returned; the interim operators are + named directly because their registry entries (``det_gemm``, + ``attention``, ``swiglu``) prefer backends that are not available or not + deterministic on every GPU this replay runs on. + """ + + def __init__(self, registry: KernelRegistry | None = None) -> None: + from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + registry = registry or KernelRegistry() + self.projection = registry.get_op("adaln_projection_3mod", device="cuda") + self.norm = registry.get_op("h3_rmsnorm", device="cuda") + self.gate = registry.get_op("adaln_gate_residual", device="cuda") + self.gemm = TritonDetGemmOp() + self.attention = DeterministicAttentionOp() + self.swiglu = SwiGLUCudaOp() + + def backends(self) -> dict[str, str]: + return { + name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items() + } + + +def _w(ctx: dict[str, Any], name: str) -> torch.Tensor: + return ctx["weights"][f"{BLOCK}.{name}"] + + +def _heads(x: torch.Tensor) -> torch.Tensor: + return x.unflatten(-1, (HEADS, HEAD_DIM)) + + +class _InterimLinear(torch.autograd.Function): + """``y = x @ W^T`` on the Triton FP32-accumulate kernel, rounded to BF16 once. + + Interim binding for the block's GEMM rows. ``TritonDetGemmOp.__call__`` + reduces K as a tree with BF16 nodes (about 3x cuBLAS's error at these + shapes); ``_triton_gemm_fp32`` runs one fixed-order FP32 K loop (no + split-K), so each output row depends only on its own input row. The + backward uses the same kernel for ``dX`` (reduction over the output + features, row-local) and ``dW`` (reduction over tokens). + """ + + @staticmethod + def forward(ctx, x2, weight): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + ctx.save_for_backward(x2, weight) + return _triton_gemm_fp32(x2, weight.t()).to(x2.dtype) + + @staticmethod + def backward(ctx, grad): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x2, weight = ctx.saved_tensors + grad = grad.contiguous() + dx = _triton_gemm_fp32(grad, weight).to(x2.dtype) + dw = _triton_gemm_fp32(grad.t(), x2).to(weight.dtype) + return dx, dw + + +def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat = x.reshape(-1, x.shape[-1]) + return _InterimLinear.apply(flat, weight).view(*x.shape[:-1], weight.shape[0]) + + +def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor: + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v))) + out = ops.attention(q, k, v, causal=False) + return out.transpose(1, 2).flatten(2, 3) + + +def _cand_swiglu(ops: CandidateOps, projected: torch.Tensor) -> torch.Tensor: + hidden, gate = projected.chunk(2, dim=-1) + return ops.swiglu(gate.contiguous(), hidden.contiguous()) + + +def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding( + row, status, "triton _det_gemm_fp32_kernel (FP32 K loop, one BF16 rounding)" + ), + candidate=lambda ops, ctx, x: _cand_linear(ops, x, _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # tree reduction vs cuBLAS + ) + + +def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: + return Node( + name=name, + inputs=(source, "adaln_projection"), + binding=Binding("h3_rmsnorm", "row", "H3RMSNormCudaOp.forward_modulated"), + candidate=lambda ops, ctx, x, t: ops.norm.forward_modulated( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider=lambda ctx, x, t: prov.provider_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + golden=lambda ctx, x, t: gold.golden_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider_bitwise_isolated=True, + ) + + +def _head_norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + # The row is someone else's; until it lands, this stack's plain H3 RMSNorm + # (bitwise equal to nn.RMSNorm at any width) runs on the 128-wide heads. + binding=Binding("h3_qk_rmsnorm_d128", "interim", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(_heads(x), _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(_heads(x), _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(_heads(x), _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _rope_node(name: str, source: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_mm_rope_3axis_partial", "reference", "provider_apply_rotary"), + candidate=lambda ops, ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + provider=lambda ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + golden=lambda ctx, x: gold.golden_apply_rotary(x, *ctx["rope_fp32"]), + provider_bitwise_isolated=True, + ) + + +def _gate_node(name: str, residual: str, sublayer: str, gate: int) -> Node: + return Node( + name=name, + inputs=(residual, sublayer, "adaln_projection"), + binding=Binding("adaln_gate_residual", "row", "H3GateResidualCudaOp"), + candidate=lambda ops, ctx, r, y, t: ops.gate(r, y, t[gate], ctx["index"]), + provider=lambda ctx, r, y, t: prov.provider_gate_residual(r, t[gate], ctx["index"], y), + golden=lambda ctx, r, y, t: gold.golden_gate_residual(r, t[gate], ctx["index"], y), + provider_bitwise_isolated=True, + ) + + +def _adaln_params(ctx): + return _w(ctx, "adaln_proj.linear.weight"), _w(ctx, "adaln_proj.linear.bias") + + +NODES: list[Node] = [ + Node( + name="adaln_projection", + inputs=("temb",), + binding=Binding("adaln_projection_3mod", "row", "H3AdaLNProjectionCudaOp"), + candidate=lambda ops, ctx, temb: ops.projection(temb, *_adaln_params(ctx)), + provider=lambda ctx, temb: prov.provider_adaln_modulation(temb, *_adaln_params(ctx)), + golden=lambda ctx, temb: gold.golden_adaln_tables(temb, *_adaln_params(ctx)), + provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS + ), + _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), + _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), + _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), + _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), + _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), + _rope_node("rope_q", "q_norm"), + _rope_node("rope_k", "k_norm"), + Node( + name="attention", + inputs=("rope_q", "rope_k", "v_proj"), + binding=Binding("h3_full_attention", "interim", "DeterministicAttentionOp(causal=False)"), + candidate=lambda ops, ctx, q, k, v: _cand_attention(ops, q, k, v), + provider=lambda ctx, q, k, v: prov.provider_attention(q, k, _heads(v)), + golden=lambda ctx, q, k, v: gold.golden_attention(q, k, _heads(v)), + provider_bitwise_isolated=False, + ), + _linear_node("o_proj", "attention", "attn.to_out.0.weight", "h3_attention_o_gemm", "interim"), + _gate_node("residual_attn", "hidden", "o_proj", gate=2), + _norm_modulate_node("norm2", "residual_attn", "norm2.weight", shift=3, scale=4), + _linear_node("ffn_gate_up", "norm2", "ff.net.0.proj.weight", "h3_ffn_gate_up_gemm", "interim"), + Node( + name="swiglu", + inputs=("ffn_gate_up",), + binding=Binding("h3_swiglu", "interim", "SwiGLUCudaOp"), + candidate=lambda ops, ctx, x: _cand_swiglu(ops, x), + provider=lambda ctx, x: prov.provider_swiglu(x), + golden=lambda ctx, x: gold.golden_swiglu(x), + provider_bitwise_isolated=False, # one FP32 rounding vs the eager two + ), + _linear_node("ffn_down", "swiglu", "ff.net.2.weight", "h3_ffn_down_gemm", "interim"), + _gate_node("residual_mlp", "residual_attn", "ffn_down", gate=5), +] +NODE_NAMES = tuple(node.name for node in NODES) +OUTPUT = NODES[-1].name + +# Gradient leaves of the backward replay: the block input, temb and every block parameter. +PARAM_LEAVES = ( + "adaln_proj.linear.weight", + "adaln_proj.linear.bias", + "norm1.weight", + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + "attn.to_out.0.weight", + "norm2.weight", + "ff.net.0.proj.weight", + "ff.net.2.weight", +) + + +# --- Inputs --------------------------------------------------------------------------------- + + +def make_block_context( + weights: dict[str, torch.Tensor], *, layout: str, seed: int = 0, batch: int = 1 +) -> dict[str, Any]: + """Block-0 weights, a packed FL2VA layout, the timestep embedding and a seeded input. + + ``temb`` comes from the provider's conditioning path (``h3_chain`` covers + those rows); the block input is a seeded BF16 stand-in for the packed + projections, which are other rows. Batch rows ``b > 0`` get their own + seeded inputs; row 0 is the same tensor at every ``batch``. + """ + + spec = H3_BLOCK_LAYOUTS[layout] + packed = h3_fl2va_layout( + num_text_tokens=spec["num_text_tokens"], + num_latent_frames=spec["num_latent_frames"], + latent_height=spec["latent_height"], + latent_width=spec["latent_width"], + num_audio_latents=spec["num_audio_latents"], + ) + seq_len = packed["position_ids"].shape[0] + timestep = h3_timesteps(NUM_TIMESTEPS, seed=seed) + w = weights + with torch.no_grad(): + features = prov.provider_time_proj(timestep) + temb = prov.provider_time_embedder( + features, + w["time_embedder.linear_1.weight"], + w["time_embedder.linear_1.bias"], + w["time_embedder.linear_2.weight"], + w["time_embedder.linear_2.bias"], + ) + rows = [ + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2 + b)) + for b in range(batch) + ] + hidden = torch.cat(rows).to(torch.bfloat16).cuda() + return { + "layout": layout, + "seq_len": seq_len, + "batch": batch, + "weights": weights, + "temb": temb, + "hidden": hidden, + "index": packed["timestep_indices"] * 3 + packed["token_tags"], + "position_ids": packed["position_ids"], + "rope": prov.provider_rope(packed["position_ids"]), + "rope_fp32": gold.golden_rope_tables(packed["position_ids"]), + } + + +@contextmanager +def strict_fp32_matmul(): + """Golden matmuls in full FP32: no TF32, no reduced-precision reductions.""" + + saved = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = saved + + +# --- Graph execution ------------------------------------------------------------------------ + + +def run_graph(mode: str, ctx: dict[str, Any], ops: CandidateOps | None = None) -> dict[str, Any]: + """Every node's output for one mode ("candidate", "provider" or "golden").""" + + values: dict[str, Any] = {"hidden": ctx["hidden"], "temb": ctx["temb"]} + # Diagnostic hybrids only: these candidate nodes run their provider replay instead. + provider_nodes = ctx.get("provider_nodes", ()) + with strict_fp32_matmul(): + for node in NODES: + args = [values[name] for name in node.inputs] + if mode == "candidate" and node.name in provider_nodes: + values[node.name] = node.provider(ctx, *args) + elif mode == "candidate": + values[node.name] = node.candidate(ops, ctx, *args) + elif mode == "provider": + values[node.name] = node.provider(ctx, *args) + else: + values[node.name] = node.golden(ctx, *args) + observer = ctx.get("node_observer") + if observer is not None: + observer(node.name) + return values + + +def run_isolated( + ctx: dict[str, Any], ops: CandidateOps, provider: dict[str, Any] +) -> dict[str, Any]: + """Each candidate node fed the provider's inputs to that node.""" + + return { + node.name: node.candidate(ops, ctx, *[provider[name] for name in node.inputs]) + for node in NODES + } + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [t for item in value for t in _flat(item)] + + +def _batch_row(value: Any, row: int, batch: int) -> Any: + """Row ``row`` of a batched activation; AdaLN tables have no batch axis.""" + + if isinstance(value, torch.Tensor): + return value[row : row + 1] if value.dim() >= 3 and value.shape[0] == batch else value + return tuple(_batch_row(v, row, batch) for v in value) + + +def bitwise(lhs: Any, rhs: Any) -> bool: + return all(torch.equal(a, b) for a, b in zip(_flat(lhs), _flat(rhs), strict=True)) + + +def max_abs(lhs: Any, rhs: Any) -> float: + return max( + float((a.float() - b.float()).abs().max()) if a.numel() else 0.0 + for a, b in zip(_flat(lhs), _flat(rhs), strict=True) + ) + + +def first_differing_row(lhs: Any, rhs: Any) -> int | None: + """First packed row (sequence position) where two node outputs differ, or None.""" + + for a, b in zip(_flat(lhs), _flat(rhs), strict=True): + if a.dim() < 3: + if not torch.equal(a, b): + return -1 # a table row, not a sequence position + continue + diff = (a != b).reshape(a.shape[0], a.shape[1], -1).any(dim=-1).any(dim=0) + rows = diff.nonzero() + if rows.numel(): + return int(rows[0]) + return None + + +def rel_err(value: Any, golden: Any) -> float: + """``max|value - golden| / max|golden|`` over a node's tensors.""" + + num = den = 0.0 + for a, g in zip(_flat(value), _flat(golden), strict=True): + g32 = g.float() + num = max(num, float((a.float() - g32).abs().max())) + den = max(den, float(g32.abs().max())) + return num / max(den, 1e-30) + + +def ulp_err(value: torch.Tensor, golden: torch.Tensor) -> float: + """``max|value - golden|`` in units of ``value.dtype``'s ulp at ``max|golden|``. + + ``rel_err`` of a stored gradient moves in steps of this ulp: its worst + element is a single rounded number, so two equally valid reduction orders + can differ by whole ulps there while every other element agrees. + """ + + g32 = golden.float() + top = float(g32.abs().max()) + if top == 0.0: + return 0.0 + ulp = torch.finfo(value.dtype).eps * 2.0 ** torch.floor(torch.log2(torch.tensor(top))).item() + return float((value.float() - g32).abs().max()) / ulp + + +def digest(value: Any) -> str: + """sha256 over a node's tensors' bytes, for comparing runs across processes or hosts.""" + + h = hashlib.sha256() + for t in _flat(value): + h.update(t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return h.hexdigest() + + +def _time_ms(fn: Callable[[], Any], repeats: int) -> float: + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return sorted(samples)[len(samples) // 2] + + +# --- Forward report ------------------------------------------------------------------------- + + +def run_block_case( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Forward replay of block 0: per-node agreement, invariance, accuracy and first drift.""" + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + with torch.no_grad(): + provider = run_graph("provider", ctx) + golden = run_graph("golden", ctx) + chained = run_graph("candidate", ctx, ops) + repeat = run_graph("candidate", ctx, ops) + batched = run_graph("candidate", ctx2, ops) + isolated = run_isolated(ctx, ops, provider) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "backends": ops.backends(), + "nodes": [], + } + firsts: dict[str, str | None] = {"first_drift": None, "first_isolated_drift": None} + firsts.update({"first_repeat_drift": None, "first_batch_drift": None}) + for node in NODES: + name = node.name + row0 = _batch_row(batched[name], 0, 2) + entry: dict[str, Any] = { + "node": name, + "rfc_row": node.binding.rfc_row, + "status": node.binding.status, + "backend": node.binding.backend, + "provider_bitwise_isolated_promised": node.provider_bitwise_isolated, + "repeat_bitwise_equal": bitwise(repeat[name], chained[name]), + "batch_row_bitwise_equal": bitwise(row0, chained[name]), + "chained_vs_provider": { + "bitwise_equal": bitwise(chained[name], provider[name]), + "max_abs": max_abs(chained[name], provider[name]), + "first_differing_row": first_differing_row(chained[name], provider[name]), + }, + "isolated_vs_provider": { + "bitwise_equal": bitwise(isolated[name], provider[name]), + "max_abs": max_abs(isolated[name], provider[name]), + }, + "candidate_rel_err_vs_golden": rel_err(chained[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "candidate_sha256": digest(chained[name]), + } + checks = { + "first_drift": entry["chained_vs_provider"]["bitwise_equal"], + "first_isolated_drift": entry["isolated_vs_provider"]["bitwise_equal"], + "first_repeat_drift": entry["repeat_bitwise_equal"], + "first_batch_drift": entry["batch_row_bitwise_equal"], + } + for key, ok in checks.items(): + if firsts[key] is None and not ok: + firsts[key] = name + report["nodes"].append(entry) + report.update(firsts) + report["output"] = { + "candidate_rel_err_vs_golden": rel_err(chained[OUTPUT], golden[OUTPUT]), + "provider_rel_err_vs_golden": rel_err(provider[OUTPUT], golden[OUTPUT]), + "candidate_vs_provider_max_abs": max_abs(chained[OUTPUT], provider[OUTPUT]), + } + report["permutation_probe"] = permutation_probe(ctx, ops, chained[OUTPUT], seed) + if timing_repeats: + with torch.no_grad(): + report["timing_ms"] = { + "candidate_forward": _time_ms( + lambda: run_graph("candidate", ctx, ops), timing_repeats + ), + "provider_forward": _time_ms(lambda: run_graph("provider", ctx), timing_repeats), + } + return report + + +def permutation_probe( + ctx: dict[str, Any], ops: CandidateOps, base_output: torch.Tensor, seed: int +) -> dict[str, Any]: + """RFC #420 control H-C3, reported and not promised. + + Permute the packed rows together with every per-row tensor (positions, + AdaLN indices); the block output should be the same permutation of the + original. Row-local nodes keep each row's bytes, but attention sums over + keys in packed order, so a different key order may round differently. + """ + + perm = torch.randperm(ctx["seq_len"], generator=torch.Generator().manual_seed(seed + 7)).cuda() + permuted = { + **ctx, + "hidden": ctx["hidden"][:, perm], + "index": ctx["index"][perm], + "position_ids": ctx["position_ids"][perm], + "rope": tuple(t[perm] for t in ctx["rope"]), + "rope_fp32": tuple(t[perm] for t in ctx["rope_fp32"]), + } + with torch.no_grad(): + out = run_graph("candidate", permuted, ops)[OUTPUT] + expected = base_output[:, perm] + return { + "bitwise_equal_after_permutation": bool(torch.equal(out, expected)), + "max_abs": max_abs(out, expected), + "differing_rows": int((out != expected).reshape(out.shape[1], -1).any(dim=-1).sum()), + } + + +# --- Backward report ------------------------------------------------------------------------ + + +def _backward( + mode: str, ctx: dict[str, Any], ops: CandidateOps | None, upstream: torch.Tensor +) -> tuple[dict[str, torch.Tensor], dict[str, list[torch.Tensor]]]: + """Gradients of every leaf, and the gradient reaching each node's output, for one mode. + + The golden runs on FP32 copies of the leaves, so its gradients are not + rounded to the BF16 parameter dtype. + """ + + def leaf(t: torch.Tensor) -> torch.Tensor: + t = t.detach().float() if mode == "golden" else t.detach() + return t.clone().requires_grad_(True) + + leaves = {name: leaf(ctx["weights"][f"{BLOCK}.{name}"]) for name in PARAM_LEAVES} + leaves["hidden"] = leaf(ctx["hidden"]) + leaves["temb"] = leaf(ctx["temb"]) + weights = {**ctx["weights"], **{f"{BLOCK}.{k}": leaves[k] for k in PARAM_LEAVES}} + run_ctx = {**ctx, "weights": weights, "hidden": leaves["hidden"], "temb": leaves["temb"]} + values = run_graph(mode, run_ctx, ops) + for node in NODES: + for t in _flat(values[node.name]): + if t.requires_grad: + t.retain_grad() + out = values[OUTPUT] + torch.autograd.backward(out, upstream.to(out.dtype)) + node_grads = { + node.name: [ + t.grad if t.grad is not None else torch.zeros_like(t) for t in _flat(values[node.name]) + ] + for node in NODES + } + return {name: t.grad for name, t in leaves.items()}, node_grads + + +def run_block_backward( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Backward replay of block 0. + + Per leaf (block input, ``temb`` and the 12 block parameters): repeat + bitwise, and the candidate's and provider's error against the FP32 + golden. Per node, in reverse graph order: whether the gradient reaching + the node's output repeats bitwise, and whether batch row 0 of a two-row + batch gets the same bytes as the batch of one (rows are separate packed + documents, so every activation gradient is row-local). Parameter + gradients sum over the batch and are not part of the batch check. + """ + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + generator = torch.Generator().manual_seed(seed + 11) + up_row0 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + up_row1 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + upstream = up_row0.bfloat16().cuda() + upstream2 = torch.cat([up_row0, up_row1]).bfloat16().cuda() + + golden, _ = _backward("golden", ctx, None, upstream) + provider, _ = _backward("provider", ctx, None, upstream) + first, first_nodes = _backward("candidate", ctx, ops, upstream) + second, second_nodes = _backward("candidate", ctx, ops, upstream) + batched, batched_nodes = _backward("candidate", ctx2, ops, upstream2) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "leaves": {}, + "nodes": [], + } + for name in (*PARAM_LEAVES, "hidden", "temb"): + report["leaves"][name] = { + "repeat_bitwise_equal": bool(torch.equal(first[name], second[name])), + "candidate_rel_err_vs_golden": rel_err(first[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "storage_dtype": str(first[name].dtype).removeprefix("torch."), + "candidate_ulps_vs_golden": ulp_err(first[name], golden[name]), + "provider_ulps_vs_golden": ulp_err(provider[name], golden[name]), + } + report["leaves"]["hidden"]["batch_row_bitwise_equal"] = bool( + torch.equal(batched["hidden"][:1], first["hidden"]) + ) + first_repeat = first_batch = None + for node in reversed(NODES): + name = node.name + row0 = [_batch_row(g, 0, 2) for g in batched_nodes[name]] + entry: dict[str, Any] = { + "node": name, + "grad_repeat_bitwise_equal": bitwise(first_nodes[name], second_nodes[name]), + "grad_batch_row_bitwise_equal": ( + bitwise(row0, first_nodes[name]) if node.name != "adaln_projection" else None + ), + } + if first_repeat is None and not entry["grad_repeat_bitwise_equal"]: + first_repeat = name + if first_batch is None and entry["grad_batch_row_bitwise_equal"] is False: + first_batch = name + report["nodes"].append(entry) + report["first_grad_repeat_drift"] = first_repeat + report["first_grad_batch_drift"] = first_batch + if timing_repeats: + report["timing_ms"] = { + "candidate_forward_backward": _time_ms( + lambda: _backward("candidate", ctx, ops, upstream), timing_repeats + ), + "provider_forward_backward": _time_ms( + lambda: _backward("provider", ctx, None, upstream), timing_repeats + ), + } + return report diff --git a/rl_engine/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py new file mode 100644 index 000000000..92455bae0 --- /dev/null +++ b/rl_engine/validation/models/h3_cases.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic MiniMax-H3 (RFC #420) test inputs shared by tests and scripts.""" + +from __future__ import annotations + +import torch + + +def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: + """Return seeded FP32 timesteps of shape ``(num,)`` on ``device``. + + Requires positive ``num``; the first value is 0 and the last is 1 when + ``num > 1``, with all remaining values sampled uniformly in [0, 1]. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + t = torch.rand(num, generator=generator) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(device) + + +def h3_packed_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """Return seeded int64 timestep indices and modality tags, each shape ``(seq_len,)``. + + Each row independently samples a timestep in ``[0, num_timesteps)`` and + a video/text/audio tag in ``[0, 3)``. The first three rows cover all tags + when available; both tensors are returned on ``device``. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + token_tags = torch.randint(0, 3, (seq_len,), generator=generator) + 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) + + +# diffusers@f53d552 modular_pipelines/minimax_h3/before_denoise.py constants. +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +def _spatial_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + """One aspect-normalised spatial axis, as numpy's ``linspace(endpoint=False)`` builds it.""" + + import numpy as np + + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def h3_fl2va_layout( + *, + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + audio_channels: int = 2, + patch: tuple[int, int] = (2, 2), + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Packed text/audio/video rows of a text-only FL2VA request, laid out as diffusers packs them. + + Replays ``position_ids``/``token_tags`` of the pinned pipeline's FL2VA + layout step (no keyframe anchors): text rows on the time axis at their row + index, channel-major audio rows pinned to the width extremes, then the video + frames on the non-uniform ``5/3 * (1, 4, 4, 4, 4)`` clock. ``position_ids`` + is ``(S, 3)`` float64 as the pipeline makes it; the rotary module casts it. + ``timestep_indices`` puts text on table row 0 and the generated media on + row 1, a two-timestep stand-in for the conditioning rows. + """ + + import numpy as np + + patch_h, patch_w = patch + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_audio_rows = num_audio_latents * audio_channels + num_video_rows = num_latent_frames * rows_per_frame + seq_len = num_text_tokens + num_audio_rows + num_video_rows + audio_start = num_text_tokens + video_start = audio_start + num_audio_rows + + position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + frame_grid = torch.stack([grid.reshape(-1) for grid in grids], dim=-1) + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(audio_channels) + position_ids[audio_start:video_start, 2] = torch.cat( + [ + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full( + (num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 + ), + ] + ) + spans = torch.tensor( + [ + _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] + for i in range(num_latent_frames) + ], + dtype=torch.float64, + ) + frame_time = float(num_text_tokens) + torch.cat( + [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] + ) + video = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video[:, :, 0] = frame_time[:, None] + video[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video.reshape(-1, 3) + + token_tags = torch.full((seq_len,), VIDEO_TAG, dtype=torch.long) + token_tags[:num_text_tokens] = TEXT_TAG + token_tags[audio_start:video_start] = AUDIO_TAG + timestep_indices = torch.ones(seq_len, dtype=torch.long) + timestep_indices[:num_text_tokens] = 0 + return { + "position_ids": position_ids.to(device), + "token_tags": token_tags.to(device), + "timestep_indices": timestep_indices.to(device), + } + + +# Named FL2VA layouts for the one-block replay: (text, latent frames, latent H, latent W, audio). +H3_BLOCK_LAYOUTS = { + # 32 + 2*20 + 2*(8*12) = 264 rows: the CI-sized case. + "tiny": dict( + num_text_tokens=32, + num_latent_frames=2, + latent_height=16, + latent_width=24, + num_audio_latents=20, + ), + # 77 + 2*40 + 3*(16*16) = 925 rows: a non-power-of-two, three-frame case. + "small": dict( + num_text_tokens=77, + num_latent_frames=3, + latent_height=32, + latent_width=32, + num_audio_latents=40, + ), + # 120 + 2*80 + 5*(18*32) = 3160 rows: a 16:9 five-frame case. + "medium": dict( + num_text_tokens=120, + num_latent_frames=5, + latent_height=36, + latent_width=64, + num_audio_latents=80, + ), +} diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py new file mode 100644 index 000000000..3aa1aa191 --- /dev/null +++ b/rl_engine/validation/models/h3_chain.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stage-by-stage replay of the MiniMax-H3 conditioning chain (RFC #420). + +Each stage runs three ways on the same device: + +* provider: the diffusers op sequence (``rl_engine.validation.models.h3_provider``); +* candidate: the dispatched RL-Kernel op, fed the previous candidate output + (chained) and, separately, the previous provider output (isolated); +* golden: the FP32/FP64 reference, fed the previous golden output. + +A stage also declares what it promises, so the end-to-end test and the +evidence JSON check the same thing: whether its isolated output is bitwise +equal to the provider, and its tolerance against the golden (rows of +``tolerance_contract.json``). ``first_drift`` is the first stage whose chained +output differs from the provider; ``first_isolated_drift`` the first stage +that differs even on provider inputs. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +HIDDEN = 5376 + + +@dataclass(frozen=True) +class Stage: + name: str + op_type: str + # (op, ctx, upstream) -> outputs; ``upstream`` is the previous stage's output. + candidate: Callable[[Any, dict, Any], Any] + provider: Callable[[dict, Any], Any] + golden: Callable[[Any, dict, Any], Any] + # Promises checked by tests/models/minimax_h3/test_h3_conditioning_e2e.py. + provider_bitwise_isolated: bool + golden_atol: float + golden_rtol: float + + +def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + """Return the context's MLP weights and biases in linear_1, linear_2 call order.""" + + weights = ctx["weights"] + return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] + + +def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the context's block-0 AdaLN projection weight and bias.""" + + prefix = "transformer_blocks.0.adaln_proj.linear" + return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] + + +STAGES: list[Stage] = [ + Stage( + name="timestep_sinusoid_h3", + op_type="timestep_sinusoid_h3", + candidate=lambda op, ctx, _up: op(ctx["timestep"]), + provider=lambda ctx, _up: provider_time_proj(ctx["timestep"]), + golden=lambda op, ctx, _up: op.forward_fp32(ctx["timestep"]), + provider_bitwise_isolated=True, + golden_atol=1e-5, # elementwise / float32 + golden_rtol=1e-5, + ), + Stage( + name="timestep_mlp_fp32", + op_type="timestep_mlp_fp32", + candidate=lambda op, ctx, up: op(up, *mlp_params(ctx)), + provider=lambda ctx, up: provider_time_embedder(up, *mlp_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *mlp_params(ctx)), + provider_bitwise_isolated=False, # reduction: different tree from cuBLAS + golden_atol=1e-4, # reduction / float32 + golden_rtol=1e-4, + ), + Stage( + name="adaln_projection_3mod", + op_type="adaln_projection_3mod", + candidate=lambda op, ctx, up: op(up, *adaln_params(ctx)), + provider=lambda ctx, up: provider_adaln_modulation(up, *adaln_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *adaln_params(ctx)), + provider_bitwise_isolated=False, # reduction: tensor-core tree differs from cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ), + Stage( + name="adaln_row_gather", + op_type="adaln_row_gather", + candidate=lambda op, ctx, up: op.gather_chunks( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + provider=lambda ctx, up: provider_adaln_row_gather( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + golden=lambda op, ctx, up: op.forward_fp32( + torch.cat(list(up), dim=1), ctx["timestep_indices"], ctx["token_tags"] + ), + provider_bitwise_isolated=True, # a copy + golden_atol=5e-2, # carries the projection's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ), +] + + +def adaln_indices(ctx: dict[str, Any]) -> torch.Tensor: + return ctx["timestep_indices"] * 3 + ctx["token_tags"] + + +def _norm1_modulated(ctx: dict[str, Any], modulation, call): + shift_msa, scale_msa = modulation[0], modulation[1] + weight = ctx["weights"]["transformer_blocks.0.norm1.weight"] + return call(ctx["hidden"], weight, shift_msa, scale_msa, adaln_indices(ctx)) + + +STAGES.append( + Stage( + name="h3_rmsnorm", + op_type="h3_rmsnorm", + # Block norm1 followed by the MSA shift/scale of the projection stage. + candidate=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated + ), + provider=lambda ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], provider_norm_modulate + ), + golden=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated_fp32 + ), + provider_bitwise_isolated=True, # replays nn.RMSNorm's reduction order + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + + +def _final(ctx: dict[str, Any], call): + weights = ctx["weights"] + history = ctx["history"] + return call( + history["adaln_gate_residual"], # the residual stream after the gated sublayer + weights["norm_out.norm.weight"], + history["timestep_mlp_fp32"], # temb + weights["norm_out.linear.weight"], + weights["norm_out.linear.bias"], + ctx["timestep_indices"], + ) + + +STAGES.append( + Stage( + name="final_adaln_out", + op_type="final_adaln_out", + candidate=lambda op, ctx, _up: _final(ctx, op), + provider=lambda ctx, _up: _final(ctx, provider_final_adaln_out), + golden=lambda op, ctx, _up: _final(ctx, op.forward_fp32), + provider_bitwise_isolated=False, # shift/scale projection: tensor-core tree vs cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +# The backward replay covers the conditioning chain up to the row gather. +CONDITIONING_STAGES = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) + +# Parameters whose gradients the backward replay reports, in chain order. +GRAD_LEAVES = ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +) +BACKWARD_MODES = ("candidate", "candidate_fused", "provider") + + +def _flat(value: Any) -> list[torch.Tensor]: + """Collect tensors from nested sequences in their original order.""" + + if isinstance(value, torch.Tensor): + return [value] + return [tensor for item in value for tensor in _flat(item)] + + +def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + """Report tensor equality, maximum FP32 absolute error, and tolerance checks.""" + + lhs_t, rhs_t = _flat(lhs), _flat(rhs) + bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) + max_abs = 0.0 + within = True + for a, b in zip(lhs_t, rhs_t, strict=True): + diff = (a.float() - b.float()).abs() + if diff.numel(): + max_abs = max(max_abs, float(diff.max())) + within = within and bool((diff <= atol + rtol * b.float().abs()).all()) + return {"bitwise_equal": bitwise, "max_abs": max_abs, "within_tolerance": within} + + +def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + """Build seeded CUDA timesteps and packed row labels with the supplied weights.""" + + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) + return { + "timestep": h3_timesteps(num_timesteps, seed=seed), + "timestep_indices": timestep_indices, + "token_tags": token_tags, + "weights": weights, + # Stand-in for the packed hidden states (the patch/text projections are + # other RFC rows): BF16 (1, S, H), seeded. + "hidden": torch.randn( + (1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2) + ) + .to(torch.bfloat16) + .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), + } + + +def golden_op(registry: KernelRegistry, op_type: str): + """Return the PyTorch reference at the end of the operator's CPU priority list.""" + + return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) + + +def run_case( + registry: KernelRegistry, + weights, + *, + num_timesteps: int, + seq_len: int, + seed: int = 0, + stages: list[Stage] | None = None, +) -> dict[str, Any]: + """Replay CUDA candidate, provider, and golden paths and report each stage's drift. + + Supplied stages must include their prerequisites in execution order. + Omitting ``stages`` runs the full registered chain without gradients. + """ + + stages = STAGES if stages is None else stages + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + report: dict[str, Any] = { + "num_timesteps": num_timesteps, + "seq_len": seq_len, + "seed": seed, + "stages": [], + } + chained: Any = None + provider: Any = None + golden: Any = None + first_drift = first_isolated = None + # Each mode sees its own earlier outputs (a stage may read any earlier stage). + history: dict[str, dict[str, Any]] = {"candidate": {}, "provider": {}, "golden": {}} + ctx_c = {**ctx, "history": history["candidate"]} + ctx_p = {**ctx, "history": history["provider"]} + ctx_g = {**ctx, "history": history["golden"]} + for stage in stages: + op = registry.get_op(stage.op_type, device="cuda") + gold = golden_op(registry, stage.op_type) + with torch.no_grad(): + isolated = stage.candidate(op, ctx_p, provider) + chained_input = chained + chained = stage.candidate(op, ctx_c, chained_input) + repeat = stage.candidate(op, ctx_c, chained_input) + golden = stage.golden(gold, ctx_g, golden) + provider = stage.provider(ctx_p, provider) + history["candidate"][stage.name] = chained + history["provider"][stage.name] = provider + history["golden"][stage.name] = golden + entry = { + "stage": stage.name, + "backend": type(op).__name__, + "kernel_id": getattr(op, "kernel_id", type(op).__name__), + "repeat_bitwise_equal": compare(repeat, chained)["bitwise_equal"], + "chained_vs_provider": compare(chained, provider), + "isolated_vs_provider": compare(isolated, provider), + "chained_vs_golden": compare(chained, golden, stage.golden_atol, stage.golden_rtol), + "provider_vs_golden": compare(provider, golden, stage.golden_atol, stage.golden_rtol), + } + if first_drift is None and not entry["chained_vs_provider"]["bitwise_equal"]: + first_drift = stage.name + if first_isolated is None and not entry["isolated_vs_provider"]["bitwise_equal"]: + first_isolated = stage.name + report["stages"].append(entry) + report["first_drift"] = first_drift + report["first_isolated_drift"] = first_isolated + return report + + +def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torch.Tensor]: + """Parameter gradients of the full chain for one execution mode. + + ``candidate`` runs the four RL-Kernel ops separately, ``candidate_fused`` + runs projection + gather as ``H3AdaLNModulationCudaOp``, ``provider`` the + diffusers path, ``golden`` the FP64 references. + """ + + leaves = { + name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES + } + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}, "history": {}} + value = None + if mode == "candidate_fused": + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + + for stage in STAGES[:2]: + value = stage.candidate(registry.get_op(stage.op_type, device="cuda"), run_ctx, value) + value = H3AdaLNModulationCudaOp()( + value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] + ) + else: + for stage in (s for s in STAGES if s.name in CONDITIONING_STAGES): + if mode == "provider": + value = stage.provider(run_ctx, value) + elif mode == "candidate": + op = registry.get_op(stage.op_type, device="cuda") + value = stage.candidate(op, run_ctx, value) + else: + value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + if value is None: + raise ValueError("conditioning chain produced no output") + outputs = list(value) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) + return [leaves[name].grad for name in GRAD_LEAVES] + + +def run_backward_case( + registry: KernelRegistry, weights, *, num_timesteps: int, seq_len: int, seed: int = 0 +) -> dict[str, Any]: + """Determinism and accuracy of the chain's parameter gradients, per mode.""" + + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + generator = torch.Generator(device="cuda").manual_seed(seed + 1) + hidden = weights[GRAD_LEAVES[4]].shape[0] // 18 + upstream = [ + torch.randn(seq_len, hidden, device="cuda", generator=generator).bfloat16() + for _ in range(6) + ] + golden = chain_grads("golden", registry, ctx, upstream) + runs = { + mode: [chain_grads(mode, registry, ctx, upstream) for _ in range(2)] + for mode in BACKWARD_MODES + } + report: dict[str, Any] = {"num_timesteps": num_timesteps, "seq_len": seq_len, "leaves": {}} + for index, name in enumerate(GRAD_LEAVES): + gold = golden[index].float() + entry: dict[str, Any] = {"golden_absmax": float(gold.abs().max())} + for mode, (first, second) in runs.items(): + grad = first[index] + err = (grad.float() - gold).abs().max() + entry[mode] = { + "repeat_bitwise_equal": bool(torch.equal(first[index], second[index])), + "max_abs_vs_golden": float(err), + "max_abs_vs_golden_over_absmax": float(err / gold.abs().max().clamp_min(1e-30)), + "correctly_rounded_fraction": float( + (grad == golden[index].to(grad.dtype)).float().mean() + ), + } + report["leaves"][name] = entry + return report + + +def git_state() -> dict[str, Any]: + """Return HEAD and tracked-tree dirtiness, defaulting to unknown/dirty if Git fails.""" + + def run(*args: str) -> str: + """Run Git in the repository root and return stripped stdout.""" + + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() + + try: + sha = run("rev-parse", "HEAD") + dirty = bool(run("status", "--porcelain", "--untracked-files=no")) + except Exception: # noqa: BLE001 + sha, dirty = "unknown", True + return {"rl_kernel_commit": sha, "tracked_tree_dirty": dirty} + + +def environment() -> dict[str, Any]: + """Record the active CUDA device, software versions, and matmul TF32 setting.""" + + return { + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": bool(torch.backends.cuda.matmul.allow_tf32), + } diff --git a/rl_engine/validation/models/h3_full_manifest.json b/rl_engine/validation/models/h3_full_manifest.json new file mode 100644 index 000000000..6cc2e5b5a --- /dev/null +++ b/rl_engine/validation/models/h3_full_manifest.json @@ -0,0 +1,311 @@ +{ + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "files": [ + { + "path": "audio_scheduler/scheduler_config.json", + "size": 96, + "sha256": "804780f7133477067bd6bbfbc02dc8b3cf9feeb400f97c08f5b1d5f6cbab3840" + }, + { + "path": "audio_vae/config.json", + "size": 2271, + "sha256": "9a3c645ff892b376c6f5f4c8685964cd75474731af594ff058492a0000caabb6" + }, + { + "path": "audio_vae/diffusion_pytorch_model.safetensors", + "size": 605429340, + "sha256": "52c59e67ba8de5477c81bfbced0327aabf500f1bfdeefd5ee754529241cb26cb" + }, + { + "path": "model_index.json", + "size": 2936, + "sha256": "5a587fe13b2371427415ac892463142683aefcd8d322e274a3a095eac37ac7d2" + }, + { + "path": "modular_model_index.json", + "size": 2935, + "sha256": "a2b6a210e482ffb78e613b553f570c44e101afce6741bd4ed91429d0559af031" + }, + { + "path": "processor/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "processor/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "processor/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "processor/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "processor/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "processor/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "processor/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "scheduler/scheduler_config.json", + "size": 97, + "sha256": "8fa6c3aa70dc9e691e1a6df899fd1b6f75f70481a27cee6e18a303817075c304" + }, + { + "path": "text_encoder/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "text_encoder/config.json", + "size": 1474, + "sha256": "d2dd0c60d01b9e195d9447c52da61c7302d28828524914c044d9c6e1b81d0427" + }, + { + "path": "text_encoder/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "text_encoder/model-00001-of-00014.safetensors", + "size": 4932328944, + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb" + }, + { + "path": "text_encoder/model-00002-of-00014.safetensors", + "size": 4875990528, + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa" + }, + { + "path": "text_encoder/model-00003-of-00014.safetensors", + "size": 4875990552, + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940" + }, + { + "path": "text_encoder/model-00004-of-00014.safetensors", + "size": 4875990584, + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f" + }, + { + "path": "text_encoder/model-00005-of-00014.safetensors", + "size": 4875990584, + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2" + }, + { + "path": "text_encoder/model-00006-of-00014.safetensors", + "size": 4875990584, + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974" + }, + { + "path": "text_encoder/model-00007-of-00014.safetensors", + "size": 4875990584, + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934" + }, + { + "path": "text_encoder/model-00008-of-00014.safetensors", + "size": 4875990584, + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f" + }, + { + "path": "text_encoder/model-00009-of-00014.safetensors", + "size": 4875990584, + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47" + }, + { + "path": "text_encoder/model-00010-of-00014.safetensors", + "size": 4875990584, + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f" + }, + { + "path": "text_encoder/model-00011-of-00014.safetensors", + "size": 4875990584, + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f" + }, + { + "path": "text_encoder/model-00012-of-00014.safetensors", + "size": 4875990584, + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e" + }, + { + "path": "text_encoder/model-00013-of-00014.safetensors", + "size": 4875990584, + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c" + }, + { + "path": "text_encoder/model-00014-of-00014.safetensors", + "size": 3270697008, + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b" + }, + { + "path": "text_encoder/model.safetensors.index.json", + "size": 97831, + "sha256": "06c952c569285870b811989b794b9766493e280fb77fbcb957fc4e5fcf25403a" + }, + { + "path": "text_encoder/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "text_encoder/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "text_encoder/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "text_encoder/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "text_encoder/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "tokenizer/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "tokenizer/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "tokenizer/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "tokenizer/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "transformer/config.json", + "size": 546, + "sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00014.safetensors", + "size": 4825958704, + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00014.safetensors", + "size": 4702158032, + "sha256": "2c4d362eddd2802180ac9c744849eb9ba8d9c8b984bdf9822cb02ed004b29184" + }, + { + "path": "transformer/diffusion_pytorch_model-00003-of-00014.safetensors", + "size": 4933368192, + "sha256": "949c5aafbbfa5654da730a6a7fafd75adb164d0857b095a30e8bb6d390887d69" + }, + { + "path": "transformer/diffusion_pytorch_model-00004-of-00014.safetensors", + "size": 4567069608, + "sha256": "eef7616790105ee839766bb2027203bf2c0d87c6aa038dca84145a8675f5ce28" + }, + { + "path": "transformer/diffusion_pytorch_model-00005-of-00014.safetensors", + "size": 4702158080, + "sha256": "43fdf42d638e8bc6745f713fae80c93bb301807a1a5ae7249344ce28e202a494" + }, + { + "path": "transformer/diffusion_pytorch_model-00006-of-00014.safetensors", + "size": 4933368232, + "sha256": "6442510b34d173653f0cce5c964b935395a8f7accf0b9cc0aa31aec59805239d" + }, + { + "path": "transformer/diffusion_pytorch_model-00007-of-00014.safetensors", + "size": 4567069608, + "sha256": "29f48f535c91dac76496ca821eeb16ca24bc4caf3f0cae8b920a89b1f966da6d" + }, + { + "path": "transformer/diffusion_pytorch_model-00008-of-00014.safetensors", + "size": 4702158080, + "sha256": "c711b096c764bd60f0b8b6ad49518bfab6d614fb788c725add8741c0674a4cd8" + }, + { + "path": "transformer/diffusion_pytorch_model-00009-of-00014.safetensors", + "size": 4933368232, + "sha256": "44428defe3976cbb87635ad200b958199e739986697cd29fdf27aeb7294b5944" + }, + { + "path": "transformer/diffusion_pytorch_model-00010-of-00014.safetensors", + "size": 4567069608, + "sha256": "3d44939c374c9da382e9c6877e1946adf7b84e08c7a881c068f228d6849411c9" + }, + { + "path": "transformer/diffusion_pytorch_model-00011-of-00014.safetensors", + "size": 4702158080, + "sha256": "224d24430b58127a5577721084e0e704a0e74ec96dd7c35bc6fc0994ebd87c33" + }, + { + "path": "transformer/diffusion_pytorch_model-00012-of-00014.safetensors", + "size": 4933368232, + "sha256": "48fa2bd8fe134eef565ab2464f1c2589a6657cba0d14283dfc06b532f8961f3c" + }, + { + "path": "transformer/diffusion_pytorch_model-00013-of-00014.safetensors", + "size": 4567069608, + "sha256": "be5b4b1809f9d546ffd4b3fcf41e5c1e02b819125caa6bc105c109b04c051bd3" + }, + { + "path": "transformer/diffusion_pytorch_model-00014-of-00014.safetensors", + "size": 4644161920, + "sha256": "8fbd5e6c1fb1df7ce988ca90f3d59e7610e465c7517e4b344eda4a214ba4b97d" + }, + { + "path": "transformer/diffusion_pytorch_model.safetensors.index.json", + "size": 64488, + "sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e" + }, + { + "path": "vae/config.json", + "size": 2011, + "sha256": "78f67deec3d63aae807f2bfe7154bc1e26f6372cb20b63265fcbae1b62bb5745" + }, + { + "path": "vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "size": 5061033024, + "sha256": "72f4c6be84ac0674f27398cde991dd9d719762f3952c4921aa66b2ce542f6374" + }, + { + "path": "vae/diffusion_pytorch_model-00002-of-00003.safetensors", + "size": 4955986528, + "sha256": "2e05e8bc23fa4071043e17fd242be8acd0685e781a43987432b2eae925be4198" + }, + { + "path": "vae/diffusion_pytorch_model-00003-of-00003.safetensors", + "size": 398539336, + "sha256": "c05d6ac4b1a33de372799d708531da6320f6a3ce6d1ce6d895e770988e004a39" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors.index.json", + "size": 74228, + "sha256": "15f6d44553c3c616b0dc999920aa784f92ecee7e4201f1f99ac405cfbf3061ca" + } + ] +} diff --git a/rl_engine/validation/models/h3_full_model.py b/rl_engine/validation/models/h3_full_model.py new file mode 100644 index 000000000..1219ef89f --- /dev/null +++ b/rl_engine/validation/models/h3_full_model.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Instrument the full pinned H3 provider with the actual RL-Kernel block graph. + +Input/output projections and both text-refiner layers retain provider arithmetic. +Each of the 50 distinct blocks executes the existing candidate graph with its own +parameters. No weights are copied from block zero. This is integration evidence, +not qualification of the reference/interim bindings as final RFC kernels. +""" + +from __future__ import annotations + +from collections import Counter +from contextlib import contextmanager +from types import MethodType + +import torch + +from rl_engine.validation.models import h3_block + +MODEL_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe" +PROVIDER_REVISION = "f53d552036a0d1bd5570782a39cd40cfabf112bc" + + +def validate_full_structure(model): + expected = dict( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + for key, value in expected.items(): + if getattr(model.config, key) != value: + raise ValueError(f"full H3 requires {key}={value}") + blocks = model.transformer_blocks + if len(blocks) != 50 or len({id(b) for b in blocks}) != 50: + raise ValueError("50 distinct transformer blocks are required") + owned = set() + for i, block in enumerate(blocks): + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError(f"block {i} has missing or shared parameters") + owned.update(id(p) for p in params) + required = ( + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ) + if any(not hasattr(model, name) for name in required): + raise ValueError("full-model boundary components are missing") + refiners = model.token_refiner.refiner_blocks + if len(refiners) != 2 or len({id(b) for b in refiners}) != 2: + raise ValueError("two distinct real refiner blocks are required") + for block in refiners: + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError("refiner parameters are missing or shared") + owned.update(id(p) for p in params) + return required + + +class FullModelTrace: + def __init__(self): + self.calls = Counter() + self.backward_calls = Counter() + self.snapshots = {} + self.capture = False + self.actual_backends = {} + self.candidate = False + self.candidate_calls = Counter() + self.forward_order = [] + self.backward_order = [] + self.execution_order = [] + self._events = set() + + def record_event(self, name): + if name not in self._events: + self._events.add(name) + self.execution_order.append(name) + + def observe(self, name, value): + self.calls[name] += 1 + self.record_event("nodes." + name) + if name not in self.forward_order: + self.forward_order.append(name) + if not isinstance(value, torch.Tensor): + return + if self.capture: + self.snapshots[name] = value.detach().cpu() + if value.requires_grad: + + def record(grad): + self.backward_calls[name] += 1 + self.record_event("nodes." + name + ".grad") + if name not in self.backward_order: + self.backward_order.append(name) + if self.capture: + self.snapshots[name + ".grad"] = grad.detach().cpu() + + value.register_hook(record) + + def validate(self, *, backward=False): + required = ( + [f"block.{i}" for i in range(50)] + + ["refiner.0", "refiner.1"] + + [ + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ] + ) + missing = [name for name in required if not self.calls[name]] + if backward: + missing += [name + ".backward" for name in required if not self.backward_calls[name]] + if missing: + raise ValueError(f"incomplete full-model execution: {missing}") + if self.candidate: + missing_nodes = [ + f"block.{i}.{node.name}" + for i in range(50) + for node in h3_block.NODES + if not self.candidate_calls[f"block.{i}.{node.name}"] + ] + if missing_nodes: + raise ValueError(f"candidate nodes did not execute: {missing_nodes}") + return dict( + candidate_calls=dict(self.candidate_calls), + execution_order=self.execution_order, + forward_order=self.forward_order, + backward_order=self.backward_order, + actual_backends=self.actual_backends, + calls=dict(self.calls), + backward_calls=dict(self.backward_calls), + bindings={ + n.name: dict( + rfc_row=n.binding.rfc_row, status=n.binding.status, backend=n.binding.backend + ) + for n in h3_block.NODES + }, + boundary_binding="provider/reference", + full_transformer_blocks=50, + refiner_layers=2, + ) + + +def require_candidate_cuda(ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + if ".cuda." not in backends[name]: + raise ValueError(f"candidate {name} fell back from its CUDA implementation") + return backends + + +@contextmanager +def full_model_execution( + model, *, candidate: bool, trace: FullModelTrace, ops=None, provider_nodes=() +): + """Temporarily bind each real block, preserving module parameters and restoring on error. + + ``provider_nodes`` (diagnostic hybrids only) names candidate graph nodes that run + their provider replay in every block; such a run localises drift and is never a + candidate result. + """ + unknown = set(provider_nodes) - {n.name for n in h3_block.NODES} + if unknown: + raise ValueError(f"unknown block nodes: {sorted(unknown)}") + boundaries = validate_full_structure(model) + originals, handles = [], [] + trace.candidate = candidate + if candidate: + ops = ops or h3_block.CandidateOps() + if hasattr(ops, "backends"): + trace.actual_backends = require_candidate_cuda(ops) + try: + for i, block in enumerate(model.transformer_blocks): + if candidate: + original = block.forward + originals.append((block, original)) + + def forward( + self, + hidden_states, + temb, + adaln_indices, + rotary_emb, + attention_mask=None, + _block_index=i, + ): + if attention_mask is not None: + raise ValueError( + "the qualified H3 candidate requires an unmasked packed sequence" + ) + weights = {h3_block.BLOCK + "." + n: p for n, p in self.named_parameters()} + ctx = dict( + weights=weights, + hidden=hidden_states, + temb=temb, + index=adaln_indices, + rope=rotary_emb, + provider_nodes=frozenset(provider_nodes), + node_observer=lambda name: trace.candidate_calls.update( + [f"block.{_block_index}.{name}"] + ), + ) + return h3_block.run_graph("candidate", ctx, ops)[h3_block.OUTPUT] + + block.forward = MethodType(forward, block) + + def hook(module, args, output, name=f"block.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for i, block in enumerate(model.token_refiner.refiner_blocks): + + def hook(module, args, output, name=f"refiner.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for name in boundaries: + + def hook(module, args, output, name=name): + trace.observe(name, output) + + handles.append(getattr(model, name).register_forward_hook(hook)) + yield trace + finally: + for handle in handles: + handle.remove() + for block, original in originals: + block.forward = original diff --git a/rl_engine/validation/models/h3_generation_resume.py b/rl_engine/validation/models/h3_generation_resume.py new file mode 100644 index 000000000..4d1cfc428 --- /dev/null +++ b/rl_engine/validation/models/h3_generation_resume.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable denoising for the full-H3 candidate generation (RFC #420). + +A candidate generation takes about 20 hours per seed and its cluster partition +preempts jobs every few hours, so a restart must not redo finished steps. The +generation already keeps every step's predicted velocities (``step-NNN.pt``); +``StepStore`` adds, next to each, a sidecar with the digest of the transformer's +inputs for that step and the digest of the saved file, both written atomically +(the sidecar last, so a step counts only when both exist). + +On resume every saved transformer call is *replayed*: the saved velocities are +returned instead of recomputed, and the pipeline's own scheduler update turns them +into the next latents exactly as before. The input digest of every replayed call is +checked against the one saved, so the replay is proven to have rebuilt the trajectory +the earlier attempt had. Recomputing the last saved step (``verify_last_step``) +additionally checks that the forward pass reproduces its stored output bitwise; it +costs one full step per restart, which on a heavily preempted partition is most of an +attempt, so it is opt-in. The execution trace is saved after the first real step of +an attempt, so a run whose remaining work is only decoding keeps a complete trace. + +A run is tied to the numerics it started with (``model_trees``: the git tree of every +path in ``MODEL_PATHS``), not to a commit, so a commit that changes only orchestration +code can continue it. ``adopt_identity`` moves a run started under a commit-based +identity onto this one and refuses unless those paths are byte-identical. +""" + +from __future__ import annotations + +import datetime +import functools +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Callable + +import torch + +RUN_FILE = "run.json" +TRACE_FILE = "trace.json" + +# Everything that determines the generated numbers. The resume and orchestration code is +# deliberately absent: changing it must not invalidate saved steps. +MODEL_PATHS = ( + "rl_engine/backends", + "rl_engine/contracts", + "rl_engine/ops", + "rl_engine/reference", + "rl_engine/runtime", + "rl_engine/validation/models/h3_block.py", + "rl_engine/validation/models/h3_cases.py", + "rl_engine/validation/models/h3_full_manifest.json", + "rl_engine/validation/models/h3_full_model.py", + "rl_engine/validation/models/h3_provider.py", + "rl_engine/validation/models/h3_weights.py", + "csrc", + "build_tools", + "examples/models/minimax_h3/generation-contract.json", + "examples/models/minimax_h3/generation-inputs.json", + "examples/models/minimax_h3/generation-prompt.txt", +) + + +class ResumeError(ValueError): + """A saved run cannot be continued, or a replay disagrees with what was saved.""" + + +def tensor_digest(named: dict[str, torch.Tensor]) -> str: + """sha256 over tensor names, dtypes, shapes and raw bytes, in name order.""" + + digest = hashlib.sha256() + for name in sorted(named): + tensor = named[name].detach().cpu().contiguous() + digest.update(f"{name}|{tensor.dtype}|{tuple(tensor.shape)}|".encode()) + if tensor.numel(): + digest.update(memoryview(tensor.reshape(-1).view(torch.uint8).numpy())) + return digest.hexdigest() + + +def inputs_digest(args: tuple, kwargs: dict[str, Any]) -> str: + """Digest of every tensor passed to one transformer call (non-tensors are ignored).""" + + named = {f"arg{i}": v for i, v in enumerate(args) if isinstance(v, torch.Tensor)} + named.update({k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)}) + return tensor_digest(named) + + +def model_trees(revision: str = "HEAD", cwd: Path | None = None) -> dict[str, str]: + """The git tree (or blob) id of every model path at ``revision``.""" + + import subprocess + + trees = {} + for path in MODEL_PATHS: + trees[path] = subprocess.check_output( + ["git", "rev-parse", f"{revision}:{path}"], text=True, cwd=cwd + ).strip() + return trees + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1 << 24): + digest.update(block) + return digest.hexdigest() + + +def _atomic_write(path: Path, write: Callable[[Path], object]) -> None: + temporary = path.with_name(path.name + ".tmp") + write(temporary) + os.replace(temporary, path) + + +class StepStore: + """The per-step velocities and input digests of one seed, in one directory.""" + + def __init__(self, directory: Path): + self.directory = directory + + def _step(self, step: int) -> Path: + return self.directory / f"step-{step:03}.pt" + + def _sidecar(self, step: int) -> Path: + return self.directory / f"step-{step:03}.in.json" + + def open(self, identity: dict[str, Any], *, resume: bool) -> int: + """Create the directory, or on ``resume`` validate it; return the saved step count. + + ``identity`` (commit, contract, seed, flags, ...) must equal the one stored, + so steps from a different source or setting can never be mixed in. + """ + + run_file = self.directory / RUN_FILE + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + if not self.directory.exists(): + self.directory.mkdir(parents=True) + elif not resume: + raise ResumeError(f"{self.directory} exists; pass resume to continue it") + if run_file.exists(): + stored = json.loads(run_file.read_text()) + if stored["identity"] != identity: + raise ResumeError("saved run has a different identity; refusing to continue it") + else: + stored = {"identity": identity, "attempts": []} + saved = self._complete_steps() + stored["attempts"].append({"started": now, "saved_steps": saved}) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) + return saved + + def attempts(self) -> list[dict[str, Any]]: + return json.loads((self.directory / RUN_FILE).read_text())["attempts"] + + def _complete_steps(self) -> int: + count = 0 + while self._step(count).exists() and self._sidecar(count).exists(): + sidecar = json.loads(self._sidecar(count).read_text()) + if _file_sha256(self._step(count)) != sidecar["output_sha256"]: + break + count += 1 + if count and not (self.directory / "initial.pt").exists(): + return 0 + return count + + def save(self, step: int, values: tuple[torch.Tensor, ...], digest: str) -> None: + """Persist one step's outputs, then its sidecar; both atomically.""" + + _atomic_write(self._step(step), lambda p: torch.save(tuple(values), p)) + sidecar = {"inputs_sha256": digest, "output_sha256": _file_sha256(self._step(step))} + _atomic_write( + self._sidecar(step), lambda p: p.write_text(json.dumps(sidecar, indent=2) + "\n") + ) + + def save_trace(self, trace: dict[str, Any]) -> None: + """Keep the validated execution trace of this seed's real forward passes.""" + + _atomic_write( + self.directory / TRACE_FILE, lambda p: p.write_text(json.dumps(trace, indent=2) + "\n") + ) + + def load_trace(self) -> dict[str, Any] | None: + path = self.directory / TRACE_FILE + return json.loads(path.read_text()) if path.exists() else None + + def saved_digest(self, step: int) -> str: + return json.loads(self._sidecar(step).read_text())["inputs_sha256"] + + def load(self, step: int) -> tuple[torch.Tensor, ...]: + sidecar = json.loads(self._sidecar(step).read_text()) + if _file_sha256(self._step(step)) != sidecar["output_sha256"]: + raise ResumeError(f"saved step {step} changed on disk") + return tuple(torch.load(self._step(step), map_location="cpu", weights_only=True)) + + +class ResumableRun: + """Replays saved steps and records new ones for one seed's transformer calls.""" + + def __init__( + self, + store: StepStore, + identity: dict[str, Any], + *, + resume: bool, + verify_last_step: bool = False, + ): + self.store = store + self.saved = store.open(identity, resume=resume) + self.verify_last_step = verify_last_step + self.replay = self.saved - 1 if verify_last_step and self.saved else self.saved + self.verified_step: int | None = None + self._forward_calls = 0 + self.recorded = 0 + self.real_calls = 0 + + def install(self, module: torch.nn.Module, make_output: Callable[[tuple, dict], Any]): + """Serve the first ``replay`` calls of ``module`` from disk; return an uninstaller.""" + + original = module.forward + + # The pipeline picks the transformer's layout inputs from inspect.signature(forward), + # so the wrapper must keep the original signature (functools.wraps sets __wrapped__). + @functools.wraps(original) + def forward(*args, **kwargs): + index = self._forward_calls + self._forward_calls += 1 + if index >= self.replay: + return original(*args, **kwargs) + if inputs_digest(args, kwargs) != self.store.saved_digest(index): + raise ResumeError(f"replayed step {index} has different inputs than when saved") + device = next( + v.device for v in (*args, *kwargs.values()) if isinstance(v, torch.Tensor) + ) + values = tuple(v.to(device) for v in self.store.load(index)) + return make_output(values, kwargs) + + module.forward = forward + + def uninstall(): + module.forward = original + + return uninstall + + def record(self, args: tuple, kwargs: dict[str, Any], values: tuple[torch.Tensor, ...]): + """Handle one finished call (forward hook): save it, or verify a recomputed step.""" + + index = self.recorded + self.recorded += 1 + if index < self.replay: + return + self.real_calls += 1 + digest = inputs_digest(args, kwargs) + values = tuple(v.detach().cpu() for v in values) + if index < self.saved: + if digest != self.store.saved_digest(index): + raise ResumeError(f"recomputed step {index} has different inputs than saved") + stored = self.store.load(index) + if len(stored) != len(values) or not all( + torch.equal(a, b) for a, b in zip(stored, values, strict=True) + ): + raise ResumeError(f"recomputed step {index} differs from the saved output") + self.verified_step = index + return + self.store.save(index, values, digest) + + def summary(self) -> dict[str, Any]: + return dict( + saved_steps_at_start=self.saved, + replayed_steps=self.replay, + verify_last_step=self.verify_last_step, + recomputed_and_verified_step=self.verified_step, + real_forward_calls=self.real_calls, + attempts=self.store.attempts(), + ) + + +def adopt_identity( + directory: Path, + new_identity: dict[str, Any], + from_commit: str, + cwd: Path | None = None, + *, + dry_run: bool = False, +) -> None: + """Move a run started under a commit-based identity onto a ``model_trees`` identity. + + Refuses unless everything but the commit matches and the model paths at + ``from_commit`` are byte-identical to the ones in ``new_identity``; records the + adoption in the run's attempt list. ``dry_run`` performs every check and writes nothing. + """ + + run_file = directory / RUN_FILE + stored = json.loads(run_file.read_text()) + old = stored["identity"] + if "model_trees" in old: + raise ResumeError("run already has a model-tree identity") + if old.get("commit", "").startswith(from_commit) is False: + raise ResumeError(f"run was started at {old.get('commit')}, not {from_commit}") + rest_old = {k: v for k, v in old.items() if k != "commit"} + rest_new = {k: v for k, v in new_identity.items() if k != "model_trees"} + if rest_old != rest_new: + raise ResumeError("run differs from the new identity in more than the commit") + then = model_trees(old["commit"], cwd=cwd) + if then != new_identity["model_trees"]: + changed = sorted(k for k in then if then[k] != new_identity["model_trees"][k]) + raise ResumeError(f"model paths changed since {from_commit}: {changed}") + if dry_run: + return + stored["identity"] = new_identity + stored["attempts"].append( + { + "adopted": { + "from_commit": old["commit"], + "model_paths_identical": True, + "when": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + } + ) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) diff --git a/rl_engine/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json new file mode 100644 index 000000000..e590a8473 --- /dev/null +++ b/rl_engine/validation/models/h3_manifest.json @@ -0,0 +1,137 @@ +{ + "version": "h3-conditioning-v1", + "rfc": "RL-Align/RL-Kernel#420", + "model_identity": { + "hf_repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "subfolder": "transformer", + "architecture": "MiniMaxH3Transformer3DModel", + "config_sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae", + "index_file": "diffusion_pytorch_model.safetensors.index.json", + "index_sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e", + "config_fingerprint": { + "hidden_size": 5376, + "num_attention_heads": 56, + "attention_head_dim": 128, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "norm_eps": 1e-05, + "final_norm_eps": 1e-05 + } + }, + "weight_shards": { + "diffusion_pytorch_model-00001-of-00014.safetensors": { + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size_bytes": 4825958704 + } + }, + "tensors": { + "time_embedder.linear_1.weight": { + "dtype": "float32", "shape": [5376, 256], + "sha256": "4c50eabe3c434dc05422e6501cdf7baa5cecbb1a7706ae9599c7a5b0e34e2913" + }, + "time_embedder.linear_1.bias": { + "dtype": "float32", "shape": [5376], + "sha256": "64763a5e39cb078eb0d8c5be5d56ad36cac3dec37be93a87e106d82ea6c071b7" + }, + "time_embedder.linear_2.weight": { + "dtype": "float32", "shape": [2688, 5376], + "sha256": "697819dd9413e0fec3a50987a964f1fd5db568ca545267414238f452294a8154" + }, + "time_embedder.linear_2.bias": { + "dtype": "float32", "shape": [2688], + "sha256": "096a3c7539f0248698cc65376c8da260e7365bf4ccef847c09bb9a36782a743c" + }, + "transformer_blocks.0.adaln_proj.linear.weight": { + "dtype": "bfloat16", "shape": [96768, 2688], + "sha256": "a0000193b30ea36a1996f53fc0d6c4c37ce5fd4585284f0f9bae570d25c6b8cf" + }, + "transformer_blocks.0.adaln_proj.linear.bias": { + "dtype": "bfloat16", "shape": [96768], + "sha256": "d59f9de9e07e6b4e3b57009627e276a26971ce379555fc99dbb06d3a237c8fc0" + }, + "transformer_blocks.0.norm1.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "5e32e4d7379a9f9302755af395ed50b5b998352a221dc2c99204a85dd9f9b49f" + }, + "transformer_blocks.0.norm2.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "a2c513ec459414e211a3d37f77c1f556f11117f56516b22eb5912f8ff5f98872" + }, + "token_refiner.final_norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "cf10a56a11370a216a4ac3afe9e8b8a3133206bccd6011a92b9aa791a9196b93" + }, + "norm_out.norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "91ac17792929e0a84533b28caf71ac831ec3d666c07a473393639be20392f5ce" + }, + "norm_out.linear.weight": { + "dtype": "bfloat16", "shape": [10752, 2688], + "sha256": "38063ae4b492aacd320e71d4a38a32325c0d47b29f59978f0f4979ca080fe673" + }, + "norm_out.linear.bias": { + "dtype": "bfloat16", "shape": [10752], + "sha256": "3b12cdc7bbda772e7e31539ffa9a32f529dba7f4148273d2128bdc621fd27438" + } + }, + "block_tensors": { + "transformer_blocks.0.attn.to_q.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "196f9120ad018784bcefd4224f0e08fae88ba1adf079ecdca18410a414a6a2a4" + }, + "transformer_blocks.0.attn.to_k.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "8e9171f3b8cf894904d1ddbd96510739d7291d2992310de6a52e1d733ec895cd" + }, + "transformer_blocks.0.attn.to_v.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "275bbc993b1ba738831e04d86f290c9988108d96fe8c97c3b576a043c7ff3fe1" + }, + "transformer_blocks.0.attn.norm_q.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "34a3f13557a2c0a9af260e48d431041e381779a0f61616cfe9800ec839693519" + }, + "transformer_blocks.0.attn.norm_k.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "8991949ddea3bc2136a0fcf0b980fa63bd0df760207e3116f7ce96df4e6ca058" + }, + "transformer_blocks.0.attn.to_out.0.weight": { + "dtype": "bfloat16", "shape": [5376, 7168], + "sha256": "822cb426aaebe57d3d219b3c4f6855ca5c13821cb84ebe0a7474fe3a87869fd6" + }, + "transformer_blocks.0.ff.net.0.proj.weight": { + "dtype": "bfloat16", "shape": [28672, 5376], + "sha256": "49587044f325aae73f046c4cda9637052670227ec3c2c289359d9d1ebf29f95d" + }, + "transformer_blocks.0.ff.net.2.weight": { + "dtype": "bfloat16", "shape": [5376, 14336], + "sha256": "e19d4e1649e4b55b1931842e5cc4bb46892c2e635b9315b2996630a86284e88f" + } + }, + "reference_implementation": { + "repo": "huggingface/diffusers", + "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "file": "src/diffusers/models/transformers/transformer_minimax_h3.py", + "file_sha256": "435873d9bc26b8770342961e18cb20324aeeaa4d045e85e0941638d4857d7db3", + "timestep_embedding_file": "src/diffusers/models/embeddings.py" + }, + "conventions": { + "timestep": "distinct timesteps t = 1 - sigma, unscaled in [0, 1], shape (T,)", + "modality_tags": {"video": 0, "text": 1, "audio": 2}, + "adaln_row": "timestep_index * 3 + token_tag", + "adaln_table_rows": "[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]", + "adaln_chunk_order": ["shift_msa", "scale_msa", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"], + "precision": { + "timestep_sinusoid": "float32", + "timestep_mlp": "float32 weights, float32 activations", + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output", + "rmsnorm": "nn.RMSNorm(5376, eps=1e-5), bfloat16 affine weight, statistics in float32", + "final_modulation": "norm_out.linear(silu(temb).bfloat16) -> (shift, scale), indexed by timestep_indices" + } + } +} diff --git a/rl_engine/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py new file mode 100644 index 000000000..3eec1818f --- /dev/null +++ b/rl_engine/validation/models/h3_provider.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The diffusers MiniMax-H3 conditioning path, op for op, without diffusers. + +Each function replays the tensor ops of the pinned reference +(``huggingface/diffusers@f53d552``, see ``h3_manifest.json``) in the same +order and dtypes, so it runs the same kernels the provider does on a given +device. This is the *provider* side of the RFC #420 comparisons; the FP32/FP64 +goldens live in ``rl_engine.reference.minimax_h3``. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + + +def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: + """Replay provider FP32 cosine/sine features for ``(T,)`` timesteps. + + Matches ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)`` + with the default channel count, returning shape ``(T, 256)``. + """ + + half_dim = num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / (half_dim - 0) + emb = torch.exp(exponent) + emb = timestep[:, None].float() * emb[None, :] + emb = 1 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb + + +def provider_time_embedder( + features: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> torch.Tensor: + """Replay linear_1, SiLU, and linear_2 using the checkpoint weight dtype. + + H3's FP32 weights map features of shape ``(T, 256)`` to ``(T, 2688)``. + """ + + sample = F.linear(features.to(w1.dtype), w1, b1) + sample = F.silu(sample) + return F.linear(sample, w2, b2) + + +def provider_adaln_modulation( + temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 +) -> tuple[torch.Tensor, ...]: + """Replay provider SiLU, the weight-dtype cast, and the AdaLN projection. + + Return six ``(3T, hidden_size)`` views in timestep-major modality order + for checkpoint weights shaped ``(18 * hidden_size, 2688)``. + """ + + temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + temb = temb.view(-1, 6 * hidden_size) + return temb.chunk(6, dim=-1) + + +def provider_adaln_row_gather( + modulation: tuple[torch.Tensor, ...], timestep_indices: torch.Tensor, token_tags: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """``adaln_indices = timestep_indices * 3 + token_tags`` then six ``index_select``.""" + + adaln_indices = timestep_indices * 3 + token_tags + return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) + + +def provider_norm_modulate( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``norm(x) * (1.0 + scale[i]) + shift[i]`` with ``index_select``; block and norm_out.""" + + norm_hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( + 0, indices + ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output + + +def provider_final_adaln_out( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + temb: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + timestep_indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``MiniMaxH3AdaLayerNormOut.forward``: shift first, indexed by timestep.""" + + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +# --- One transformer block (RFC #420 ``ws1_one_h3_block``) ------------------------------------ + + +def provider_rope( + position_ids: torch.Tensor, rope_freq_dim: int = 16, rope_theta: float = 10000.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """``MiniMaxH3RotaryPosEmbed``: FP32 ``(S, 96)`` cos/sin over the ``(t, h, w)`` axes.""" + + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device) + / (2 * rope_freq_dim) + ) + ) + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def provider_apply_rotary( + hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """``_apply_rotary_emb``: rotate-half on the leading 96 of 128 channels of ``(B, S, H, D)``. + + The FP32 tables are cast to the activation dtype (BF16) before the rotation. + """ + + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + x1, x2 = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-x2, x1), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + +def provider_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A bias-free ``nn.Linear`` of the block (Q/K/V, O and both FFN projections).""" + + return F.linear(x, weight) + + +def provider_head_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + """``norm_q``/``norm_k``: ``nn.RMSNorm(128)`` on ``(B, S, H, D)``.""" + + return F.rms_norm(x, (x.shape[-1],), weight, eps) + + +def provider_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """The native diffusers backend: SDPA on ``(B, H, S, D)`` views, non-causal, no mask. + + Returns ``(B, S, H*D)`` after ``flatten(2, 3).type_as(query)``. + """ + + q, k, v = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + return out.permute(0, 2, 1, 3).flatten(2, 3).type_as(query) + + +def provider_swiglu(projected: torch.Tensor) -> torch.Tensor: + """diffusers ``SwiGLU`` after its projection: ``hidden * silu(gate)`` on the two halves.""" + + hidden_states, gate = projected.chunk(2, dim=-1) + return hidden_states * F.silu(gate) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py new file mode 100644 index 000000000..16b212e6d --- /dev/null +++ b/rl_engine/validation/models/h3_report.py @@ -0,0 +1,831 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Per-operator performance and accuracy measurements for the H3 evidence. + +``PERF_CASES[op]`` builds the timed cases that ``benchmarks/models/benchmark_h3_conditioning.py`` +prints and ``tools/validation/models/h3_evidence.py`` stores. ``ACCURACY[op]`` measures the +op against the provider path and its golden. Each RFC #420 row registers both. +""" + +from __future__ import annotations + +import statistics +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) + +# Keys of a perf case that hold a timed callable, in display order. +TIMED_KEYS = ( + "candidate", + "candidate_checked", + "provider", + "candidate_backward", + "provider_backward", +) + + +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ + + for _ in range(warmup): + _prepared_call(fn, setup)() + torch.cuda.synchronize() + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + run() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + row["timing_samples_us"] = samples + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) + return row + + +# --------------------------------------------------------------------------- # +# timestep_sinusoid_h3 +# --------------------------------------------------------------------------- # + + +def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA sinusoid timing cases with checked, unchecked, and provider calls.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + cases = [] + for num in (1, 2, 4, 64): + t = h3_timesteps(num) + cases.append( + { + "op": "timestep_sinusoid_h3", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": num * 4 + num * 256 * 4, + "candidate": lambda t=t: op.forward(t, check_range=False), + "candidate_checked": lambda t=t: op.forward(t), + "provider": lambda t=t: provider_time_proj(t), + } + ) + return cases + + +def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + """Compare sinusoid outputs to provider and FP64 golden across timestep counts.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + rows = [] + for num in (1, 2, 3, 7, 64, 1000, 4097): + t = h3_timesteps(num, seed=num) + out, ref, gold = op(t), provider_time_proj(t), golden.forward_fp32(t) + rows.append( + { + "num_timesteps": num, + "bitwise_equal_to_provider": bool(torch.equal(out, ref)), + "max_abs_vs_fp64": float((out - gold).abs().max()), + "provider_max_abs_vs_fp64": float((ref - gold).abs().max()), + } + ) + return {"contract_atol": 1e-5, "cases": rows} + + +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Return CUDA checkpoint parameters, or seeded random tensors when weights are unset.""" + + if h3_weights_dir() is not None: + weights = load_h3_conditioning_weights("cuda", names) + return [weights[name] for name in names] + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 + for shape in shapes + ] + + +# --------------------------------------------------------------------------- # +# timestep_mlp_fp32 +# --------------------------------------------------------------------------- # + +MLP_NAMES = [f"time_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")] +MLP_SHAPES = [(5376, 256), (5376,), (2688, 5376), (2688,)] + + +def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA MLP candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + weight_bytes = sum(p.numel() * p.element_size() for p in params) + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + x = torch.rand(num, 256, device="cuda", generator=gen) * 2 - 1 + cases.append( + { + "op": "timestep_mlp_fp32", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight_bytes, + "candidate": lambda x=x: op(x, *params), + "provider": lambda x=x: provider_time_embedder(x, *params), + } + ) + return cases + + +def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + """Measure per-draw MLP error against FP64 and test single-row/batched equality.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + cuda_err, provider_err = [], [] + for seed in range(draws): + x = sinusoid(h3_timesteps(4, seed=seed)) + gold = golden.forward_fp32(x, *params) + cuda_err.append(float((op(x, *params) - gold).abs().max())) + provider_err.append(float((provider_time_embedder(x, *params) - gold).abs().max())) + # Row invariance: each timestep alone equals its row in the batch of 9. + x = sinusoid(h3_timesteps(9, seed=1)) + full = op(x, *params) + invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) + return { + "contract_atol": 1e-4, + "draws": draws, + "num_timesteps": 4, + "cuda_max_abs_vs_fp64": cuda_err, + "provider_max_abs_vs_fp64": provider_err, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_projection_3mod +# --------------------------------------------------------------------------- # + +ADALN_NAMES = [ + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +] + + +def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + """Return block-0 AdaLN projection weights and bias as BF16 CUDA tensors.""" + + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) + return weight.bfloat16(), bias.bfloat16() + + +def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA AdaLN candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + weight, bias = _adaln_params() + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) + cases.append( + { + "op": "adaln_projection_3mod", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight.numel() * weight.element_size(), + "candidate": lambda temb=temb: op(temb, weight, bias), + "provider": lambda temb=temb: provider_adaln_modulation(temb, weight, bias), + } + ) + return cases + + +def _flat_cat(outputs) -> torch.Tensor: + """Concatenate flattened modulation tensors in their output order.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: + """Report per-draw BF16 equality to FP64 goldens and single-row/batch invariance.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["adaln_projection_3mod"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + weight, bias = _adaln_params() + cuda_frac, provider_frac, early_frac = [], [], [] + for seed in range(draws): + g = torch.Generator(device="cuda").manual_seed(seed) + temb = torch.randn(4, 2688, device="cuda", generator=g) * 2 + gold = _flat_cat(golden.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat_cat(golden.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + ours = _flat_cat(op(temb, weight, bias)) + theirs = _flat_cat(provider_adaln_modulation(temb, weight, bias)) + cuda_frac.append(float((ours == gold).float().mean())) + provider_frac.append(float((theirs == gold).float().mean())) + early_frac.append(float((ours == early).float().mean())) + gen = torch.Generator(device="cuda").manual_seed(draws) + temb = torch.randn(9, 2688, device="cuda", generator=gen) + full = op(temb, weight, bias) + invariant = all( + all( + torch.equal(s, f[3 * i : 3 * i + 3]) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) + ) + for i in range(9) + ) + return { + "draws": draws, + "num_timesteps": 4, + "cuda_correctly_rounded": cuda_frac, + "provider_correctly_rounded": provider_frac, + "early_cast_golden_match": early_frac, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_row_gather +# --------------------------------------------------------------------------- # + +GATHER_SEQ_LENS = (4097, 32768, 131072) + + +def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_row_gather", device="cuda") + num_timesteps, hidden = 3, 5376 + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda", generator=gen).bfloat16() + chunks = rows.chunk(6, dim=-1) + cases = [] + for seq in GATHER_SEQ_LENS: + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + grads = [ + torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) + ] + + def prepare_backward(fn, ti=ti, tags=tags): + leaf = rows.detach().requires_grad_(True) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) + + cases.append( + { + "op": "adaln_row_gather", + "case": f"S={seq}", + "backend": type(op).__name__, + # bytes written (six (S, H) BF16 outputs); the 3T-row table stays in L2 + "bytes": 6 * seq * hidden * 2, + "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), + "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) + ), + } + ) + return cases + + +def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: + from rl_engine.validation.models.h3_chain import run_backward_case + + op = registry.get_op("adaln_row_gather", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + # Seeded so the op-backward fractions below are reproducible from this commit. + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() + forward_bitwise = {} + for seq in (1, 257, 4097, 32768): + ti, tags = h3_packed_layout(seq, 3, seed=seq) + ours = op(rows, ti, tags) + theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) + + # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = [torch.randn(4097, 5376, device="cuda", generator=gen).bfloat16() for _ in range(6)] + + def grad_of(fn): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(fn(leaf)), grads) + return leaf.grad + + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward(list(golden.forward_fp32(ref, ti, tags)), [g.float() for g in grads]) + gold = ref.grad.bfloat16() + backward = {} + for name, fn in ( + ("cuda", lambda leaf: op(leaf, ti, tags)), + ("provider", lambda leaf: provider_adaln_row_gather(leaf.chunk(6, dim=-1), ti, tags)), + ): + first, second = grad_of(fn), grad_of(fn) + backward[name] = { + "repeat_bitwise_equal": bool(torch.equal(first, second)), + "correctly_rounded_fraction": float((first == gold).float().mean()), + } + + result: dict[str, Any] = { + "forward_bitwise_vs_index_select": forward_bitwise, + "op_backward": backward, + "chain_backward": [], + } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return result + + +# --------------------------------------------------------------------------- # +# h3_rmsnorm (block norm1 + MSA modulation) +# --------------------------------------------------------------------------- # + +NORM_SEQ_LENS = (4097, 32768, 131072) + + +def _norm_inputs(seq: int, seed: int = 0): + weight = h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0].bfloat16() + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).bfloat16() + shift, scale = table.view(9, 6, 5376)[:, 0], table.view(9, 6, 5376)[:, 1] + ti, tags = h3_packed_layout(seq, 3, seed=seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + return x, weight, shift, scale, ti * 3 + tags + + +def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("h3_rmsnorm", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, weight, shift, scale, index = _norm_inputs(seq) + grad = torch.randn_like(x) + + def prepare_backward(fn, x=x, weight=weight, shift=shift, scale=scale): + leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) + + cases.append( + { + "op": "h3_rmsnorm", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 2 * seq * 5376 * 2, # read x, write y + "candidate": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + op.forward_modulated(x, w, sh, sc, i, check_range=False) + ), + "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + provider_norm_modulate(x, w, sh, sc, i) + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: op.forward_modulated( + x_, w_, sh_, sc_, i, check_range=False + ) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) + ), + } + ) + return cases + + +def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("h3_rmsnorm", device="cuda") + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + gen = torch.Generator(device="cuda").manual_seed(0) + plain = {} + for name in names: + x = torch.randn(2, 777, 5376, device="cuda", generator=gen).bfloat16() + ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) + plain[name] = bool(torch.equal(op(x, weights[name]), ref)) + x, weight, shift, scale, index = _norm_inputs(4097, seed=1) + modulated = bool( + torch.equal( + op.forward_modulated(x, weight, shift, scale, index), + provider_norm_modulate(x, weight, shift, scale, index), + ) + ) + invariant = bool( + torch.equal( + op.forward_modulated(x[:, 100:140], weight, shift, scale, index[100:140])[0], + op.forward_modulated(x, weight, shift, scale, index)[0, 100:140], + ) + ) + + # Backward against FP64, for the CUDA op and the diffusers expression. + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, weight, shift, scale)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, w_, sh_, sc_): + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * w_ + return n * (1 + sc_.index_select(0, index)) + sh_.index_select(0, index) + + ref = grads(golden, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op.forward_modulated(*t, index)), + ("provider", lambda *t: provider_norm_modulate(t[0], t[1], t[2], t[3], index)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) + }, + } + return { + "weight_source": weight_source, + "plain_bitwise_vs_nn_rmsnorm": plain, + "modulated_bitwise_vs_diffusers": modulated, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + 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])), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# final_adaln_out (norm_out) +# --------------------------------------------------------------------------- # + +FINAL_NAMES = ["norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias"] + + +def _final_inputs(seq: int, seed: int = 0): + norm_weight, weight, bias = h3_params(FINAL_NAMES, [(5376,), (10752, 2688), (10752,)]) + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, _ = h3_packed_layout(seq, 3, seed=seed) + return x, norm_weight.bfloat16(), temb, weight.bfloat16(), bias.bfloat16(), ti + + +def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("final_adaln_out", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, nw, temb, w, b, ti = _final_inputs(seq) + grad = torch.randn_like(x) + + 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), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "final_adaln_out", + "case": f"S={seq}", + "backend": type(op).__name__, + # read x + norm_out.linear, write out + "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, + "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out( + t[0], t[1], t[2], t[3], t[4], i + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], i) + ), + } + ) + return cases + + +def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("final_adaln_out", device="cuda") + golden_op = registry._get_or_create_backend( + registry._priority_map["cpu"]["final_adaln_out"][-1] + ) + if golden_op is None: + raise RuntimeError("required H3 golden backend is unavailable") + x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) + ours = op(x, nw, temb, w, b, ti) + theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) + golden = golden_op.forward_fp32(x, nw, temb, w, b, ti) + invariant = bool( + torch.equal(op(x[:, 100:160], nw, temb, w, b, ti[100:160])[0], ours[0, 100:160]) + ) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, nw, temb, w, b)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden_backward(x_, nw_, t_, w_, b_): + # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. + act = t_ * torch.sigmoid(t_) + # The declared cast rounds the FP32 SiLU, as the model computes it. + declared = torch.nn.functional.silu(t_.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = torch.nn.functional.linear(act, w_, b_) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw_ + return n * (1.0 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ref = grads(golden_backward, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], ti)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) + }, + } + return { + "equal_to_diffusers_fraction": float((ours == theirs).float().mean()), + "max_abs_vs_golden": float((ours.float() - golden).abs().max()), + "provider_max_abs_vs_golden": float((theirs.float() - golden).abs().max()), + "rows_batch_invariant": invariant, + "backward": backward, + } + + +PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { + "timestep_sinusoid_h3": _sinusoid_perf, + "timestep_mlp_fp32": _mlp_perf, + "adaln_projection_3mod": _projection_perf, + "adaln_row_gather": _gather_perf, + "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, + "final_adaln_out": _final_perf, +} +ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { + "timestep_sinusoid_h3": _sinusoid_accuracy, + "timestep_mlp_fp32": _mlp_accuracy, + "adaln_projection_3mod": _projection_accuracy, + "adaln_row_gather": _gather_accuracy, + "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, + "final_adaln_out": _final_accuracy, +} diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py new file mode 100644 index 000000000..f587f2726 --- /dev/null +++ b/rl_engine/validation/models/h3_weights.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pinned MiniMax-H3 conditioning weights for the RFC #420 tests. + +Only the tensors listed in ``h3_manifest.json`` are used. They all live in +shard 1 of 14, so ``tools/weights/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. With ``--block`` it also +extracts block 0's attention and FFN weights (manifest ``block_tensors``, +about 1.1 GB) into ``h3_block0.safetensors`` for ``ws1_one_h3_block``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch + +MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") +WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" +EXTRACTED_FILE = "h3_conditioning.safetensors" +BLOCK_FILE = "h3_block0.safetensors" + + +def load_h3_manifest() -> dict[str, Any]: + """Read the bundled checkpoint identity and pinned tensor contracts.""" + + return json.loads(MANIFEST_PATH.read_text()) + + +def h3_weights_dir() -> Path | None: + """Expand the configured extraction directory, or return None when unset or blank.""" + + value = os.environ.get(WEIGHTS_ENV, "").strip() + return Path(value).expanduser() if value else None + + +def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + """Return a file SHA256 while reading bounded chunks.""" + + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def sha256_tensor(tensor: torch.Tensor) -> str: + """Hash a CPU tensor's contiguous bytes without safetensors metadata.""" + + tensor_bytes = tensor.contiguous().view(torch.uint8).numpy() + return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() + + +def _load_pinned( + filename: str, section: str, device: torch.device | str, names: list[str] | None +) -> dict[str, torch.Tensor]: + """Load one extracted file and check every tensor of a manifest section against it.""" + + from safetensors.torch import load_file + + root = h3_weights_dir() + if root is None: + raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") + path = root / filename + if not path.is_file(): + flag = " --block" if section == "block_tensors" else "" + raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py{flag}") + specs = load_h3_manifest()[section] + tensors = load_file(str(path), device="cpu") + wanted = names if names is not None else list(specs) + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + out: dict[str, torch.Tensor] = {} + for name in wanted: + if name not in specs: + raise KeyError(name) + out[name] = tensors[name].to(device=device) + return out + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + return _load_pinned(EXTRACTED_FILE, "tensors", device, names) + + +def load_h3_block_weights(device: torch.device | str = "cpu") -> dict[str, torch.Tensor]: + """Conditioning tensors plus block 0's attention/FFN weights, every one sha-checked.""" + + return { + **load_h3_conditioning_weights(device), + **_load_pinned(BLOCK_FILE, "block_tensors", device, None), + } diff --git a/rl_engine/validation/models/h3_ws2.py b/rl_engine/validation/models/h3_ws2.py new file mode 100644 index 000000000..9ec7b22c3 --- /dev/null +++ b/rl_engine/validation/models/h3_ws2.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS1 vs WS2 comparisons for the H3 conditioning path (tests and evidence). + +The per-rank helpers run inside real multi-process groups (tests and +``tools/validation/models/h3_ws2_evidence.py``); the WS1 side runs on one GPU. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def ws1_projection(temb, weight, bias, grad) -> dict[str, torch.Tensor]: + """The single-GPU ``adaln_projection_3mod`` table and its three gradients.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + table = H3AdaLNProjectionCudaOp().forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight": leaves[1].grad, + "d_bias": leaves[2].grad, + } + + +def tp_projection_rank(collective, temb, weight, bias, grad) -> dict[str, Any]: + """One TP rank: shard the full weight, run the TP op, return its view of the results.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + op = H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, w_shard, b_shard)] + table = op.forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight_shard": leaves[1].grad, + "d_bias_shard": leaves[2].grad, + "readback": op.readback(), + } + + +def tp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Byte equality of every rank's replicated outputs and of the concatenated shards.""" + + return { + "table": all(torch.equal(r["table"], ws1["table"]) for r in ranks), + "d_temb": all(torch.equal(r["d_temb"], ws1["d_temb"]) for r in ranks), + "d_weight": torch.equal(torch.cat([r["d_weight_shard"] for r in ranks]), ws1["d_weight"]), + "d_bias": torch.equal(torch.cat([r["d_bias_shard"] for r in ranks]), ws1["d_bias"]), + } + + +def tp_projection_sweep(collective, temb, weight, bias, grad, iters: int = 100) -> dict[str, Any]: + """Evidence for one rank: per-T byte equality against WS1 on this GPU, timing, readback. + + ``temb``/``grad`` hold the largest T; every T in ``1..T`` uses their first + rows. Each rank compares its own replicated outputs and its own shard of + ``dW``/``db`` with the WS1 result computed here, so nothing large is shipped. + """ + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + adaln_column_shard, + ) + from rl_engine.validation.models.h3_report import time_us + + shard = adaln_column_shard(weight.shape[0], collective.world_size, collective.rank) + cols = slice(shard.begin, shard.end) + tp_op, ws1_op = ( + H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]), + H3AdaLNProjectionCudaOp(), + ) + equality, perf = [], [] + for num_t in range(1, temb.shape[0] + 1): + t, g = temb[:num_t], grad[:num_t] + ws1 = ws1_projection(t, weight, bias, g) + mine = tp_projection_rank(collective, t, weight, bias, g) + equality.append( + { + "num_timesteps": num_t, + "table": torch.equal(mine["table"], ws1["table"]), + "d_temb": torch.equal(mine["d_temb"], ws1["d_temb"]), + "d_weight_shard": torch.equal(mine["d_weight_shard"], ws1["d_weight"][cols]), + "d_bias_shard": torch.equal(mine["d_bias_shard"], ws1["d_bias"][cols]), + } + ) + leaves = [x.detach().clone().requires_grad_(True) for x in (t, weight[cols], bias[cols])] + full = [x.detach().clone().requires_grad_(True) for x in (t, weight, bias)] + + def step(op, args): + for x in args: + x.grad = None + op.forward_table(*args).backward(g) + + perf.append( + { + "num_timesteps": num_t, + "tp_forward_us": time_us( + lambda: tp_op.forward_table(t, weight[cols], bias[cols]), 30, iters + ), + "tp_forward_backward_us": time_us(lambda: step(tp_op, leaves), 10, iters), + "ws1_forward_us": time_us(lambda: ws1_op.forward_table(t, weight, bias), 30, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_op, full), 10, iters), + } + ) + 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 ------------------------------------------------- + + +def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): + import traceback + + import torch.distributed as dist + + from rl_engine.distributed.algorithms.collectives import DeterministicCollective + + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", init_method=init_method, rank=rank, world_size=world, device_id=rank + ) + inputs = {k: v.cuda() for k, v in torch.load(inputs_path).items()} + collective = DeterministicCollective(device=rank, max_size_bytes=256 << 20) + out = target(collective, **inputs, **kwargs) + # Results go through a file: tensors sent over the queue are shared by + # file descriptor and vanish if this process exits before the parent reads. + torch.save( + {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, + inputs_path.with_name(f"rank{rank}.pt"), + ) + collective.close() + dist.destroy_process_group() + queue.put({"rank": rank}) + except Exception: # forwarded to the parent + queue.put({"rank": rank, "error": traceback.format_exc()}) + + +def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: float = 900, **kwargs): + """Run ``target(collective, **inputs, **kwargs)`` on ranks ``0..world-1`` (one GPU each). + + ``target`` must be importable (it is pickled into spawned processes). Each + rank gets a NCCL group and a ``DeterministicCollective``; results come back + on the CPU, ordered by rank. Any rank's exception is re-raised here. + """ + + import tempfile + import time + from pathlib import Path + from queue import Empty + + import torch.multiprocessing as mp + + if torch.cuda.device_count() < world: + raise RuntimeError(f"needs {world} GPUs, found {torch.cuda.device_count()}") + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "inputs.pt" + torch.save({k: v.cpu() for k, v in inputs.items()}, path) + queue = ctx.Queue() + init = (Path(tmp) / "init").as_uri() + procs = [ + ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) + for r in range(world) + ] + started = [] + deadline = time.monotonic() + timeout + status: dict[int, object] = {} + try: + for proc in procs: + proc.start() + started.append(proc) + while len(status) < world or any(proc.is_alive() for proc in procs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"rank results timed out after {timeout}s") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + 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']}") + status[result["rank"]] = result + for rank, proc in enumerate(procs): + 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)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index aaebc279f..15d7adc91 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -16,6 +16,10 @@ DEFAULT_VOCAB = 151936 DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 +H3_FREQ_DIM = 256 +H3_TIME_HIDDEN = 5376 +H3_TIME_EMBED = 2688 +H3_HIDDEN = 5376 def make_operator_inputs( @@ -24,6 +28,7 @@ def make_operator_inputs( dtype: torch.dtype, device: torch.device, ) -> dict[str, Any]: + """Build operator keyword inputs on the requested device from CLI shape and seed options.""" builders = { "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, @@ -44,6 +49,13 @@ def make_operator_inputs( "embedding": _make_embedding_inputs, "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, + "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, + "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, + "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, + "adaln_row_gather": _make_adaln_row_gather_inputs, + "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -52,6 +64,7 @@ def make_operator_inputs( def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: + """Return the operator's dimension label for benchmark and evidence reports.""" if op_name == "video_patch_pack_unpack": shape = "x".join(str(dim) for dim in _video_patch_shape(args)) return f"{_arg_str(args, 'patch_direction', 'pack')}-{shape}" @@ -78,6 +91,16 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", + "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", + "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" + f"x{H3_TIME_EMBED}", + "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" + f"x{6 * 3 * _h3_hidden(args)}", + "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" + f"->{batch * seq}", + "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -370,6 +393,151 @@ def _make_kv_cache_attention_inputs( } +def _h3_num_timesteps(args: argparse.Namespace) -> int: + """Read the packed H3 timestep count, falling back to the CLI batch size.""" + # H3 packs a handful of distinct timesteps; reuse --batch as their count. + return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) + + +def _h3_timesteps( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """Build seeded timesteps in [0, 1] with endpoint cases in the requested storage dtype.""" + num = _h3_num_timesteps(args) + mode = _arg_str(args, "input_mode", "random") + if mode == "constant": + t = torch.full((num,), 0.5, device=device) + else: + t = torch.rand((num,), generator=_generator(args, device, offset=7), device=device) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(dtype) + + +def _make_timestep_sinusoid_h3_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Supply a one-dimensional packed timestep tensor to the H3 sinusoid operator.""" + return {"timestep": _h3_timesteps(args, dtype, device)} + + +def _make_timestep_mlp_fp32_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build [T, 256] sinusoid features and seeded FP32 256→5376→2688 MLP parameters.""" + # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. + del dtype + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + features = NativeH3TimestepSinusoidOp().forward(_h3_timesteps(args, torch.float32, device)) + scale1, scale2 = H3_FREQ_DIM**-0.5, H3_TIME_HIDDEN**-0.5 + return { + "x": features, + "w1": _floating_tensor((H3_TIME_HIDDEN, H3_FREQ_DIM), args, torch.float32, device, 1) + * scale1, + "b1": _floating_tensor((H3_TIME_HIDDEN,), args, torch.float32, device, 2) * 0.1, + "w2": _floating_tensor((H3_TIME_EMBED, H3_TIME_HIDDEN), args, torch.float32, device, 3) + * scale2, + "b2": _floating_tensor((H3_TIME_EMBED,), args, torch.float32, device, 4) * 0.1, + } + + +def _h3_hidden(args: argparse.Namespace) -> int: + """Read the AdaLN channel width, defaulting to the checkpoint's 5376 channels.""" + return _arg_int(args, "normalized_dim", H3_HIDDEN) + + +def _make_adaln_projection_3mod_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build FP32 [T, 2688] embeddings and 18H projection parameters in the chosen dtype.""" + # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the + # projection's weight dtype, BF16 in the checkpoint. + n_out = 6 * 3 * _h3_hidden(args) + temb = _floating_tensor( + (_h3_num_timesteps(args), H3_TIME_EMBED), args, torch.float32, device, 0 + ) + weight = _floating_tensor((n_out, H3_TIME_EMBED), args, torch.float32, device, 1) + bias = _floating_tensor((n_out,), args, torch.float32, device, 2) + return { + "temb": temb * 2.0, + "weight": (weight * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (bias * 0.1).to(dtype), + } + + +def _make_adaln_row_gather_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # T distinct timesteps (--batch) and a packed sequence of batch * seq rows + # mixing all three modalities. + batch, seq = _batch_seq(args) + num_timesteps = _h3_num_timesteps(args) + generator = _generator(args, device, offset=29) + packed = batch * seq + token_tags = torch.randint(0, 3, (packed,), generator=generator, device=device) + timestep_indices = torch.randint( + 0, num_timesteps, (packed,), generator=generator, device=device + ) + rows = _floating_tensor((3 * num_timesteps, 6 * _h3_hidden(args)), args, dtype, device, 0) + return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} + + +def _make_h3_rmsnorm_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + weight = _floating_tensor((hidden,), args, torch.float32, device, 1).abs() + 0.5 + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "weight": weight.to(dtype), + "eps": 1e-5, + } + + +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + +def _make_final_adaln_out_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract; norm_out.linear/norm/x use ``dtype`` (BF16 in H3). + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + num_timesteps = _h3_num_timesteps(args) + # Realistic scales (norm weight in [0.5, 1], unit activations): with unscaled + # inputs the parameter gradients sum hundreds of rows into values ~1e3, and + # their cancelling elements cannot meet a per-element FP32 atol. + linear = _floating_tensor((2 * hidden, H3_TIME_EMBED), args, torch.float32, device, 3) + norm = 0.5 + 0.5 * torch.sigmoid(_floating_tensor((hidden,), args, torch.float32, device, 4)) + generator = _generator(args, device, offset=37) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "norm_weight": norm.to(dtype), + "temb": _floating_tensor((num_timesteps, H3_TIME_EMBED), args, torch.float32, device, 1), + "weight": (linear * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (_floating_tensor((2 * hidden,), args, torch.float32, device, 2) * 0.1).to(dtype), + "timestep_indices": torch.randint( + 0, num_timesteps, (seq,), generator=generator, device=device + ), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 92995b938..d5e6a7c85 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -272,6 +272,110 @@ def _load_object(path: str) -> Any: }, grad_input_names=("logits",), ), + # MiniMax-H3 (RFC #420) conditioning path. The op is FP32 end to end, so + # its FP32 rows are the declared contract; other dtypes only change the + # timestep input's storage dtype. + "timestep_sinusoid_h3": OperatorSpec( + name="timestep_sinusoid_h3", + op_class="elementwise", + gold_path="rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ), + }, + grad_input_names=("timestep",), + ), + "timestep_mlp_fp32": OperatorSpec( + name="timestep_mlp_fp32", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_mlp.H3TimestepMLPCudaOp" + ), + }, + grad_input_names=("x", "w1", "b1", "w2", "b2"), + ), + "adaln_projection_3mod": OperatorSpec( + name="adaln_projection_3mod", + op_class="reduction", + gold_path=("rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_projection.H3AdaLNProjectionCudaOp" + ), + }, + grad_input_names=("temb", "weight", "bias"), + ), + # The gather's forward is a copy (asserted bitwise in tests/models/minimax_h3), but its VJP + # sums every packed position of a table row, so it is judged as a reduction. + "adaln_row_gather": OperatorSpec( + name="adaln_row_gather", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_row_gather.H3AdaLNRowGatherCudaOp" + ), + }, + grad_input_names=("rows",), + ), + "h3_rmsnorm": OperatorSpec( + name="h3_rmsnorm", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + "cuda": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "gate_residual.H3GateResidualCudaOp" + ), + }, + grad_input_names=("residual", "y", "gate"), + ), + "final_adaln_out": OperatorSpec( + name="final_adaln_out", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "final_adaln_out.H3FinalAdaLNOutCudaOp" + ), + }, + grad_input_names=("x", "norm_weight", "temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..820347673 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -10,7 +10,12 @@ import torch from torch.utils import cpp_extension -from build_tools.extensions import _resolve_package_data_source, get_cmdclass +from build_tools.extensions import ( + H3_BITWISE_SOURCES, + H3_KERNELS_DEFINE, + _resolve_package_data_source, + get_cmdclass, +) def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: @@ -83,3 +88,28 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert _resolve_package_data_source(source, project_root=tmp_path) == source assert "build_py" in get_cmdclass() + + +def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + extension = _load_extension_config(monkeypatch, hip=None) + + assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["nvcc"] + assert "--use_fast_math" not in extension["extra_compile_args"]["nvcc"] + + +def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: + # --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for approximations + # and enables FTZ, which breaks the H3 kernels' bitwise parity promise; nvcc + # exposes no macro to reject it in-source, so the build drops the kernels + # and their bindings. + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + extension = _load_extension_config(monkeypatch, hip=None) + + assert "--use_fast_math" in extension["extra_compile_args"]["nvcc"] + assert not set(H3_BITWISE_SOURCES) & set(extension["sources"]) + assert not [s for s in extension["sources"] if s.startswith("csrc/cuda/h3/")] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py new file mode 100644 index 000000000..f9c2b000f --- /dev/null +++ b/tests/models/minimax_h3/conftest.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared fixtures for the MiniMax-H3 (RFC #420) operator tests. + +Real-weight tests read the pinned tensors from ``$RL_KERNEL_H3_WEIGHTS`` +(written by ``tools/weights/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_block_weights, + load_h3_conditioning_weights, +) + + +@pytest.fixture(scope="session") +def h3_weights_cpu() -> dict[str, torch.Tensor]: + """Load pinned CPU weights once per session, skipping unavailable artifacts.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") + try: + return load_h3_conditioning_weights("cpu") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(scope="session") +def h3_block_weights_cuda() -> dict[str, torch.Tensor]: + """Conditioning plus block-0 tensors on CUDA, skipping when they were not extracted.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py --block") + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + try: + return load_h3_block_weights("cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +def h3_kernels_built() -> bool: + """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +@pytest.fixture +def require_h3_kernels() -> None: + if not h3_kernels_built(): + pytest.skip("rl_engine._C has no H3 kernels (not built, or a fast-math build)") diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..e4c303e10 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.gate_residual import NativeH3GateResidualOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_gate_residual + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py new file mode 100644 index 000000000..066445fc0 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused AdaLN modulation (projection + row gather) with an FP32 table gradient. + +The forward must be bitwise equal to the two RFC #420 ops run separately; the +backward must be deterministic and must not re-round the table gradient to +BF16 (so ``d_temb`` reaches FP64-golden accuracy). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models.h3_cases import h3_packed_layout + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _ops(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks the H3 kernels") + return H3AdaLNModulationCudaOp(), H3AdaLNProjectionCudaOp(), H3AdaLNRowGatherCudaOp() + + +def _params(hidden, dim, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = 18 * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).bfloat16().cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).bfloat16().cuda() + temb = (torch.randn(3, dim, generator=g) * 2).cuda() + return temb, weight, bias + + +def _grads(seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + +def _leaf_grads(fn, temb, weight, bias, grads): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(fn(*leaves)), grads) + return [leaf.grad for leaf in leaves] + + +@requires_cuda +class TestFusedModulation: + def test_forward_bitwise_equal_to_separate_ops(self): + fused, proj, gather = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + ours = fused(temb, weight, bias, ti, tags) + separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) + assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + + def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): + fused, proj, gather = _ops() + weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() + bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() + g = torch.Generator(device="cuda").manual_seed(0) + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = _grads(4097, 5376, seed=1) + ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + separate = _leaf_grads( + lambda t, w, b: gather.gather_chunks(proj(t, w, b), ti, tags), + temb, + weight, + bias, + grads, + ) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column (about 1 temb draw in 30 dropped below 0.9999 that way). + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # identity-VJP cast + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) + index = ti * 3 + tags + torch.autograd.backward( + [chunk.index_select(0, index) for chunk in table.chunk(6, dim=-1)], + [g.double() for g in grads], + ) + + def rel(grad, golden): + return ((grad.double() - golden).abs().max() / golden.abs().max()).item() + + # d_temb: no BF16 rounding anywhere on its path, so it reaches FP32 accuracy. + assert rel(ours[0], ref[0].grad) < 1e-5 + assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 + # dW, db: only their own final BF16 rounding remains. + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + + def test_backward_repeat_bitwise(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=3) + ti, tags = h3_packed_layout(3000, 3, seed=3) + grads = _grads(3000, 64, seed=3) + first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + assert all(torch.equal(a, b) for a, b in zip(first, again)) + + def test_rejects_bad_indices(self): + fused, _, _ = _ops() + temb, weight, bias = _params(8, 16) + ti, tags = h3_packed_layout(10, 3) + with pytest.raises(IndexError): + fused(temb, weight, bias, ti + 3, tags) + with pytest.raises(TypeError): + fused(temb.bfloat16(), weight, bias, ti, tags) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..5fb10a73f --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_projection_3mod``: SiLU + 2688 -> 96768 projection, 6 x 3 table. + +* layout: the six outputs and the three modality rows per timestep land + exactly where diffusers puts them (RFC probe H4); +* mixed precision: SiLU in FP32, one cast to BF16; an early cast (probe H7) + is rejected at the API and is measurably different numerically; +* accuracy: forward and gradients against the declared-cast golden on the + pinned block-0 weights; +* invariance and determinism: a timestep's 18 modulation rows do not depend + on the other timesteps, and repeated runs are bitwise equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.validation.models.h3_provider import provider_adaln_modulation + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3AdaLNProjectionCudaOp() + + +def _temb(num, dim=2688, seed=0, device="cuda"): + """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(num, dim, generator=g) * 2).to(device) + + +def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + """Create seeded projection weights for six chunks and three modalities per channel.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +def _flat(outputs): + """Flatten and concatenate all modulation chunks for elementwise comparisons.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +class TestReference: + def test_provider_layout_matches_diffusers(self): + """Match the provider values and timestep-major modality layout on CPU.""" + + temb = _temb(2, dim=16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + for a, b in zip(ours, theirs, strict=True): + assert a.shape == (6, 8) and torch.equal(a, b) + + def test_rejects_bf16_temb(self): + # RFC probe H7: casting temb before the SiLU. + """Reject an early BF16 cast so SiLU receives the declared FP32 embeddings.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(TypeError, match="float32"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) + + def test_rejects_bad_shapes(self): + """Reject invalid projection dimensions, empty batches, and mismatched bias dtype.""" + + weight, bias = _synthetic(8, 16, device="cpu") + op = NativeH3AdaLNProjectionOp() + with pytest.raises(ValueError): # not a multiple of 6 * 3 rows + op.forward(_temb(2, 16, device="cpu"), weight[:-1], bias[:-1]) + with pytest.raises(ValueError): + op.forward(_temb(2, 15, device="cpu"), weight, bias) + with pytest.raises(ValueError): + op.forward(_temb(0, 16, device="cpu"), weight, bias) + with pytest.raises(TypeError): + op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + + +@requires_cuda +class TestLayout: + def test_channel_identity_lands_in_its_chunk_and_modality_row(self): + """Probe H4: bias o = m*6H + c*H + h tagged (m, c); zero weight passes it through.""" + + hidden, dim, num = 16, 32, 3 + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + o = torch.arange(n_out) + modality, chunk = o // (H3_ADALN_CHUNKS * hidden), (o // hidden) % H3_ADALN_CHUNKS + bias = (modality * 10 + chunk).to(torch.bfloat16).cuda() + weight = torch.zeros(n_out, dim, dtype=torch.bfloat16, device="cuda") + outputs = _cuda_op()(_temb(num, dim), weight, bias) + assert len(outputs) == 6 + for c, out in enumerate(outputs): + assert out.shape == (num * 3, hidden) + for t in range(num): + for m in range(3): + expected = torch.full((hidden,), float(m * 10 + c), device="cuda") + assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) + + def test_outputs_are_views_of_one_table_like_diffusers(self): + """Require all six modulation chunks to share storage with the expected offsets.""" + + weight, bias = _synthetic(16, 32) + outputs = _cuda_op()(_temb(2, 32), weight, bias) + base = outputs[0].untyped_storage().data_ptr() + assert all(out.untyped_storage().data_ptr() == base for out in outputs) + assert outputs[1].data_ptr() - outputs[0].data_ptr() == 16 * outputs[0].element_size() + + +@pytest.fixture(scope="module") +def block0(h3_weights_cpu): + """Move the pinned block-zero projection weight and bias to CUDA for this module.""" + + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4]) + def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + """Check checkpoint accuracy and near-exact BF16 rounding of the declared golden.""" + + weight, bias = block0 + temb = _temb(num, seed=num) + outputs = _cuda_op()(temb, weight, bias) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + for out, gold in zip(outputs, golden, strict=True): + assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) + # tolerance_contract.json forward_accuracy / reduction / bfloat16. + torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) + # Beyond the contract: almost every element is the correctly rounded + # golden value. The rest come from the FP32 tensor-core accumulation over + # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact + # equality is not reachable: over 300 temb draws per T the worst match was + # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). + declared = _flat(golden).bfloat16() + assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 + + def test_early_cast_is_numerically_detectable(self, block0): + """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" + + weight, bias = block0 + temb = _temb(3, seed=7) + op = NativeH3AdaLNProjectionOp() + ours = _flat(_cuda_op()(temb, weight, bias)) + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + assert (ours == declared).float().mean() >= 0.9995 + assert (ours == early).float().mean() < 0.8 + + def test_rows_are_batch_and_position_invariant(self, block0): + """Keep each timestep modulation row identical alone and under batch reordering.""" + + weight, bias = block0 + op = _cuda_op() + temb = _temb(5, seed=3) + full = op(temb, weight, bias) + for i in range(5): + single = op(temb[i : i + 1], weight, bias) + for s, f in zip(single, full, strict=True): + assert torch.equal(s, f[3 * i : 3 * i + 3]) + swapped = op(temb.flip(0), weight, bias) + for s, f in zip(swapped, full, strict=True): + assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) + + def test_backward_against_golden(self, block0): + """Check embedding and parameter gradients against a high-precision declared-cast VJP.""" + + weight, bias = block0 + temb = _temb(2, seed=5) + g = torch.Generator(device="cuda").manual_seed(0) + grads = [torch.randn(6, 5376, device="cuda", generator=g).bfloat16() for _ in range(6)] + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(_cuda_op()(*leaves)), grads) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column. + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # declared cast, identity VJP + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) + torch.autograd.backward(list(table), [gr.double() for gr in grads]) + # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad + # and BF16 weight/bias grads (rounded once from an FP32 accumulation). + torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) + for leaf, r in zip(leaves[1:], ref[1:], strict=True): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + # With T = 2 each dW/db element sums two products, and the one BF16 cast + # of the FP32 sum is the correctly rounded value of every element + # (checked on 500 of 500 temb draws). + assert torch.equal(leaf.grad, r.grad.to(torch.bfloat16)) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) + def test_other_sizes(self, dtype, hidden, dim): + """Match the projection golden at non-checkpoint dimensions in FP32 and BF16.""" + + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) + temb = _temb(3, dim, seed=hidden) + tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) + for out, gold in zip( + _cuda_op()(temb, weight, bias), + NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + strict=True, + ): + torch.testing.assert_close(out.float(), gold, **tol) + + def test_repeat_forward_backward_bitwise(self): + """Require repeated projection outputs and all input gradients to match bitwise.""" + + weight, bias = _synthetic(64, 2688, seed=9) + temb = _temb(3, seed=9) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + outs = _cuda_op()(*leaves) + torch.autograd.backward(list(outs), grads) + runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_temb_grad_rows_are_batch_invariant(self): + """Keep embedding gradients bitwise equal in full batches and single-row calls.""" + + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(4, seed=11) + grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] + + def d_temb(rows): + """Compute selected embedding-row gradients with their matching modality gradients.""" + + leaf = temb[rows].detach().clone().requires_grad_(True) + outs = _cuda_op()(leaf, weight, bias) + torch.autograd.backward( + list(outs), [g.view(4, 3, -1)[rows].reshape(-1, 64) for g in grads] + ) + return leaf.grad + + full = d_temb(slice(0, 4)) + for i in range(4): + assert torch.equal(d_temb(slice(i, i + 1))[0], full[i]) + + @pytest.mark.parametrize("num", [8, 9, 17]) + def test_rows_invariant_across_tensor_core_row_tiles(self, num): + """The BF16 path computes rows in tiles of 8; a row's bits must not depend on its tile.""" + + weight, bias = _synthetic(64, 2688, seed=13) + temb = _temb(num, seed=num) + full = _cuda_op()(temb, weight, bias) + for i in (0, 7, num - 1): + single = _cuda_op()(temb[i : i + 1], weight, bias) + for s_out, f_out in zip(single, full, strict=True): + assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + + def test_rejects_cpu(self): + """Reject CPU tensors at the CUDA projection boundary.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError): + _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) + + def test_registry_dispatches_cuda(self): + """Resolve the projection registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py new file mode 100644 index 000000000..0f9fec6c3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_row_gather``: ``timestep_index * 3 + token_tag`` lookup of six tensors. + +* forward is a copy: bitwise equal to diffusers' six ``index_select`` calls, + for every packing, length and index dtype; +* semantic indices fail closed (out-of-range tags or timesteps, RFC probes + H2/H3), and flipping a tag or offsetting a timestep picks a different row; +* backward is a deterministic FP32 segmented sum: repeat-bitwise and + correctly rounded, unlike the atomic BF16 ``index_select`` backward; +* batch invariance: a position's output does not depend on which other + positions share the call, and a row's gradient does not depend on the + positions that reference other rows. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +from rl_engine.reference.minimax_h3.adaln_row_gather import NativeH3AdaLNRowGatherOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_adaln_row_gather + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks h3_adaln_row_gather_*") + return H3AdaLNRowGatherCudaOp() + + +def _rows(num_timesteps, hidden=5376, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(3 * num_timesteps, 6 * hidden, generator=g).to(dtype).to(device) + + +class TestReference: + def test_matches_diffusers_index_select(self): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = h3_packed_layout(11, 2, device="cpu") + index = ti * 3 + tags + chunks = rows.chunk(6, dim=-1) + for ours, chunk in zip(NativeH3AdaLNRowGatherOp().forward(rows, ti, tags), chunks): + assert torch.equal(ours, chunk.index_select(0, index)) + + @pytest.mark.parametrize( + "mutate, error", + [ + (lambda ti, tags: (ti, tags.clone().fill_(3)), IndexError), # H2: no 4th modality + (lambda ti, tags: (ti, tags - 1), IndexError), + (lambda ti, tags: (ti + 2, tags), IndexError), # H3: timestep offset past T + (lambda ti, tags: (ti, tags[:-1]), ValueError), + (lambda ti, tags: (ti[:0], tags[:0]), ValueError), + (lambda ti, tags: (ti.float(), tags), TypeError), + (lambda ti, tags: (ti.int(), tags), TypeError), # mixed index dtypes + ], + ) + def test_rejects_invalid_indices(self, mutate, error): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = mutate(*h3_packed_layout(9, 2, device="cpu")) + with pytest.raises(error): + NativeH3AdaLNRowGatherOp().forward(rows, ti, tags) + + def test_rejects_bad_rows(self): + ti, tags = h3_packed_layout(4, 1, device="cpu") + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(4, 24), ti, tags) # not 3 per t + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) + @pytest.mark.parametrize("num_timesteps", [1, 3]) + def test_bitwise_equal_to_index_select(self, seq, num_timesteps): + rows = _rows(num_timesteps, seed=seq) + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + ours = _cuda_op()(rows, ti, tags) + ref = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + assert len(ours) == 6 + for a, b in zip(ours, ref): + assert a.shape == (seq, 5376) and a.is_contiguous() and torch.equal(a, b) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [24, 5, 5376]) + def test_dtypes_and_unaligned_hidden(self, dtype, hidden): + rows = _rows(2, hidden=hidden, dtype=dtype) + ti, tags = h3_packed_layout(33, 2) + for a, b in zip( + _cuda_op()(rows, ti.int(), tags.int()), NativeH3AdaLNRowGatherOp()(rows, ti, tags) + ): + assert torch.equal(a, b) + + def test_strided_rows_view(self): + table = _rows(4, hidden=16) + rows = table[::2] # row stride 2 * 6H, still 3 rows per timestep after slicing 12 -> 6 + ti, tags = h3_packed_layout(20, 2) + for a, b in zip(_cuda_op()(rows, ti, tags), NativeH3AdaLNRowGatherOp()(rows, ti, tags)): + assert torch.equal(a, b) + + def test_tag_flip_and_timestep_offset_select_other_rows(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(30, 2) + base = _cuda_op()(rows, ti, tags) + flipped = _cuda_op()(rows, ti, (tags + 1) % 3) # H2 + shifted = _cuda_op()(rows, 1 - ti, tags) # H3 within range + assert not torch.equal(base[0], flipped[0]) + assert not torch.equal(base[0], shifted[0]) + # Every output row is exactly the addressed table row, nothing else. + index = ti * 3 + tags + for c, out in enumerate(base): + assert torch.equal(out, rows[index, c * 16 : (c + 1) * 16]) + + def test_rejects_out_of_range_on_device(self): + rows = _rows(1, hidden=16) + ti, tags = h3_packed_layout(8, 1) + with pytest.raises(IndexError): + _cuda_op()(rows, ti, tags + 3) + + @pytest.mark.parametrize("entrypoint", ["native", "unchecked_wrapper"]) + @pytest.mark.parametrize( + "timestep, tag, index_dtype", + [ + pytest.param(-1, 0, "int32", id="negative-timestep"), + pytest.param(2, 0, "int64", id="timestep-past-table"), + pytest.param(0, 3, "int32", id="tag-past-modality-valid-row"), + pytest.param(1, -1, "int64", id="negative-tag-valid-row"), + # Multiplication by 3 wraps this value to row 2 in signed int64. + pytest.param((2**64 + 2) // 3, 0, "int64", id="timestep-overflow-valid-row"), + pytest.param(0, 2**63 - 1, "int64", id="int64-max-tag"), + ], + ) + def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): + """Reject semantic index violations without poisoning pytest's CUDA context.""" + + _cuda_op() + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + entrypoint, timestep, tag, index_dtype = sys.argv[1:] + rows = torch.zeros((6, 6), dtype=torch.float32, device="cuda") + dtype = getattr(torch, index_dtype) + ti = torch.tensor([int(timestep)], dtype=dtype, device="cuda") + tags = torch.tensor([int(tag)], dtype=dtype, device="cuda") + if entrypoint == "native": + _C.h3_adaln_row_gather_forward(rows, ti, tags, 6, 3) + else: + H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "native gather accepted invalid semantic indices" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + def test_gather_chunks_drop_in(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(10, 2) + chunks = rows.chunk(6, dim=-1) + for a, b in zip(_cuda_op().gather_chunks(chunks, ti, tags), _cuda_op()(rows, ti, tags)): + assert torch.equal(a, b) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + def _cuda_grad(self, rows, ti, tags, grads): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(_cuda_op()(leaf, ti, tags)), grads) + return leaf.grad + + def test_correctly_rounded_fp32_segment_sum(self): + rows = _rows(3, seed=1) + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = self._grads(4097, 5376) + ours = self._cuda_grad(rows, ti, tags, grads) + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward( + list(NativeH3AdaLNRowGatherOp().forward_fp32(ref, ti, tags)), [g.float() for g in grads] + ) + assert ours.dtype == torch.bfloat16 + # tolerance_contract.json gradient_accuracy / elementwise / bfloat16. + torch.testing.assert_close(ours.double(), ref.grad, atol=2e-2, rtol=1.6e-2) + assert (ours == ref.grad.bfloat16()).float().mean() > 0.9999 + + def test_repeat_bitwise(self): + rows = _rows(2, hidden=512, seed=2) + ti, tags = h3_packed_layout(5000, 2, seed=2) + grads = self._grads(5000, 512, seed=2) + first = self._cuda_grad(rows, ti, tags, grads) + for _ in range(3): + assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + + def test_forward_rows_invariant_to_batch_size_and_position(self): + rows = _rows(3, hidden=512, seed=5) + ti, tags = h3_packed_layout(4097, 3, seed=5) + full = _cuda_op()(rows, ti, tags) + perm = torch.randperm(4097, generator=torch.Generator().manual_seed(5)).cuda() + for pick in ( + torch.tensor([0]), + torch.tensor([2048]), + torch.arange(100, 357), + torch.arange(4000, 4097), + perm, + ): + pick = pick.cuda() + part = _cuda_op()(rows, ti[pick], tags[pick]) + for a, b in zip(part, full): + assert torch.equal(a, b[pick]) + + def test_backward_row_gradient_independent_of_other_rows_tokens(self): + # Each row's segment is tiled on its own, so a row's gradient from the + # full packing equals the one from a packing of only its own positions + # (same order), even though ~455 positions per row cross tile boundaries. + rows = _rows(3, hidden=512, seed=6) + ti, tags = h3_packed_layout(4097, 3, seed=6) + grads = self._grads(4097, 512, seed=6) + full = self._cuda_grad(rows, ti, tags, grads) + flat = ti * 3 + tags + for r in range(rows.shape[0]): + own = (flat == r).nonzero().squeeze(1) + assert own.numel() > 256 + alone = self._cuda_grad(rows, ti[own], tags[own], [g[own] for g in grads]) + assert torch.equal(alone[r], full[r]) + assert torch.count_nonzero(alone[torch.arange(rows.shape[0], device="cuda") != r]) == 0 + + def test_unreferenced_rows_get_zero_grad(self): + rows = _rows(3, hidden=16) + ti = torch.zeros(7, dtype=torch.long, device="cuda") + tags = torch.zeros(7, dtype=torch.long, device="cuda") # only row 0 is used + grad = self._cuda_grad(rows, ti, tags, self._grads(7, 16)) + assert torch.count_nonzero(grad[1:]) == 0 + assert torch.count_nonzero(grad[0]) > 0 + + def test_segments_longer_than_one_tile(self): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import BACKWARD_TILE + + seq = 3 * BACKWARD_TILE + 17 + rows = _rows(1, hidden=8, dtype=torch.float32) + ti = torch.zeros(seq, dtype=torch.long, device="cuda") + tags = torch.zeros(seq, dtype=torch.long, device="cuda") + grads = [torch.ones(seq, 8, device="cuda") for _ in range(6)] + grad = self._cuda_grad(rows, ti, tags, grads) + assert torch.equal(grad[0], torch.full((48,), float(seq), device="cuda")) + + def test_reference_backward_is_deterministic_and_accurate(self): + rows = _rows(2, hidden=256, seed=4) + ti, tags = h3_packed_layout(3000, 2, seed=4) + grads = self._grads(3000, 256, seed=4) + + def ref_grad(): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(NativeH3AdaLNRowGatherOp()(leaf, ti, tags)), grads) + return leaf.grad + + first = ref_grad() + assert torch.equal(ref_grad(), first) + torch.testing.assert_close( + first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 + ) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_row_gather", device="cuda") + assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py new file mode 100644 index 000000000..bd3a44363 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU regressions for benchmark ordering and backward timing boundaries.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_report + + +@pytest.fixture +def cuda_clock(monkeypatch): + clock = SimpleNamespace(time=0.0, active=False, log=[], allocated=0, peak=0, events=0) + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + self.start = clock.events % 2 == 0 + clock.events += 1 + + def record(self): + clock.active = self.start + clock.log.append("start" if self.start else "end") + self.time = clock.time + + def synchronize(self): + assert not clock.active + + def elapsed_time(self, end): + return end.time - self.time + + def reset_peak(): + clock.log.append("reset_peak") + clock.peak = clock.allocated + + def allocated(): + clock.log.append("memory_baseline") + return clock.allocated + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: clock.log.append("synchronize")) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset_peak) + monkeypatch.setattr(torch.cuda, "memory_allocated", allocated) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: clock.peak) + return clock + + +def test_setup_is_outside_backward_events_and_memory_baseline(cuda_clock): + states = [] + + def setup(): + assert not cuda_clock.active + cuda_clock.log.append("setup") + cuda_clock.time += 50.0 + cuda_clock.allocated += 64 * 2**20 + return object() + + def backward(state): + states.append(state) + cuda_clock.log.append("backward") + cuda_clock.time += 2.0 + cuda_clock.allocated += 3 * 2**20 + cuda_clock.peak = cuda_clock.allocated + + assert h3_report.time_us(backward, warmup=1, iters=2, setup=setup) == 2000.0 + assert len({id(state) for state in states}) == 3 + assert cuda_clock.log.count("start") == 2 + for index, entry in enumerate(cuda_clock.log): + if entry == "start": + assert cuda_clock.log[index - 1] == "setup" + assert cuda_clock.log[index + 1 : index + 3] == ["backward", "end"] + + cuda_clock.log.clear() + assert h3_report.peak_mib(backward, setup=setup) == 3.0 + assert cuda_clock.log == [ + "setup", + "synchronize", + "reset_peak", + "memory_baseline", + "backward", + "synchronize", + ] + + +@pytest.mark.parametrize("keys", [("candidate", "provider"), h3_report.TIMED_KEYS]) +def test_measure_interleaves_and_records_both_orders(cuda_clock, keys): + calls = [] + case = {"op": "test", "case": "tiny", "backend": "cpu", "bytes": 4096} + for position, key in enumerate(keys, start=1): + + def run(key=key, duration=position): + calls.append(key) + cuda_clock.time += duration + + case[key] = run + + row = h3_report.measure(case, warmup=2, iters=4) + first, second = list(keys), list(reversed(keys)) + assert calls == (first + second) * 3 + first # warmup, samples, then memory calls + assert row["execution_order"] == { + "policy": "alternating", + "iteration_0": first, + "iteration_1": second, + } + for position, key in enumerate(keys, start=1): + assert row[f"{key}_us"] == position * 1000.0 + assert row[f"{key}_gbps"] == pytest.approx(4096 / (position * 1e-3) / 1e9) + assert row[f"{key}_peak_mib"] == 0.0 + assert not any(key.endswith("_setup") for key in row) + + +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) +def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( + monkeypatch, cuda_clock, operator +): + prepared_leaves = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def track_backward(outputs): + tensors = outputs if isinstance(outputs, tuple) else (outputs,) + + def backward_hook(grad): + cuda_clock.time += 2.0 / len(tensors) + return grad + + for tensor in tensors: + tensor.register_hook(backward_hook) + return outputs + + def track(fn, leaves, *args): + if leaves[0].requires_grad: + assert not cuda_clock.active + 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 + + if operator == "norm": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 4), + torch.randn(3, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_norm_modulate + + def forward(x, weight, shift, scale, index, **kwargs): + return track(provider, [x, weight, shift, scale], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_norm_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) + op = SimpleNamespace(forward_modulated=forward) + build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + 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) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(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) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf + else: + randn = torch.randn + cpu_generators = {} + + def cpu_randn(*args, **kwargs): + kwargs["device"] = "cpu" + # A seeded CUDA generator cannot draw a CPU tensor: continue the same seed on + # one CPU generator per CUDA generator. + gen = kwargs.get("generator") + if gen is not None and gen.device.type != "cpu": + if id(gen) not in cpu_generators: + cpu_generators[id(gen)] = torch.Generator().manual_seed(gen.initial_seed()) + kwargs["generator"] = cpu_generators[id(gen)] + return randn(*args, **kwargs) + + provider = h3_report.provider_adaln_row_gather + + def forward(rows, ti, tags, **kwargs): + def gather(leaf, ti, tags): + return provider(leaf.chunk(6, dim=-1), ti, tags) + + return track(gather, [rows], ti, tags) + + def provider_forward(chunks, ti, tags): + if chunks[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([chunks[0]._base]) + cuda_clock.time += 1.0 + outputs = provider(chunks, ti, tags) + return track_backward(outputs) if chunks[0].requires_grad else outputs + + monkeypatch.setattr(torch, "randn", cpu_randn) + monkeypatch.setattr(h3_report, "GATHER_SEQ_LENS", (2, 3)) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: (torch.arange(seq) % num, torch.arange(seq) % 3), + ) + monkeypatch.setattr(h3_report, "provider_adaln_row_gather", provider_forward) + op = SimpleNamespace(forward=forward) + build = h3_report._gather_perf + + try: + registry = SimpleNamespace(get_op=lambda *args, **kwargs: op) + for case in build(registry): + row = h3_report.measure(case, warmup=1, iters=2) + assert row["backward_timing_scope"] == "backward_only" + assert row["candidate_backward_us"] == pytest.approx(2000.0) + assert row["provider_backward_us"] == pytest.approx(2000.0) + finally: + torch.set_num_threads(previous_threads) + + assert prepared_leaves + assert len({id(leaves[0]) for leaves in prepared_leaves}) == len(prepared_leaves) + assert all(leaf.grad is not None for leaves in prepared_leaves for leaf in leaves) diff --git a/tests/models/minimax_h3/test_h3_chain_replay_cli.py b/tests/models/minimax_h3/test_h3_chain_replay_cli.py new file mode 100644 index 000000000..b9baea84d --- /dev/null +++ b/tests/models/minimax_h3/test_h3_chain_replay_cli.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Reject incomplete chain dependencies before checking CUDA or loading weights.""" + +from __future__ import annotations + +import sys + +import pytest + +from tools.validation.models import h3_chain_replay + +STAGE_NAMES = [stage.name for stage in h3_chain_replay.h3_chain.STAGES] + + +@pytest.fixture +def no_cuda_check(monkeypatch): + """Fail if invalid CLI arguments reach a CUDA availability check.""" + + def unexpected_check(): + """Report an environment check performed before argument rejection.""" + + pytest.fail("invalid arguments must be rejected before checking CUDA") + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", unexpected_check) + + +@pytest.mark.parametrize("stages, message", [("", "unknown stages"), ("unknown", "unknown stages")]) +def test_reject_invalid_stages(monkeypatch, capsys, no_cuda_check, stages, message): + """Reject unknown stages before any device check.""" + + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES))) +def test_backward_requires_all_stages(monkeypatch, capsys, no_cuda_check, count): + """Require the full chain for parameter-gradient replay.""" + + monkeypatch.setattr( + sys, + "argv", + ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count]), "--backward"], + ) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert "--backward requires every stage" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "stages", + [ + *(",".join(STAGE_NAMES[:count]) for count in range(1, len(STAGE_NAMES) + 1)), + STAGE_NAMES[1], + STAGE_NAMES[-1], + ",".join(STAGE_NAMES[:1] + STAGE_NAMES[2:]), + ",".join(reversed(STAGE_NAMES)), + ",".join([STAGE_NAMES[0], STAGE_NAMES[0]]), + ], +) +def test_accept_any_known_selection(monkeypatch, stages): + """Any selection of known stages reaches the device requirement; the replay + runs its prerequisites (tests/models/minimax_h3/test_h3_cli.py).""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() + + +@pytest.mark.parametrize("options", [[], ["--backward"]]) +def test_accept_all_stages_by_default(monkeypatch, options): + """Keep the complete chain as the default for forward and backward replay.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", *options]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py new file mode 100644 index 000000000..8e97fc6a2 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 CLIs preserve chain dependencies and require pinned report weights.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_chain +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV +from tools.validation.models import h3_chain_replay, h3_evidence + + +@pytest.fixture +def cpu_chain_cli(monkeypatch): + """Replace CUDA chain stages with CPU stubs that track dependencies and provider drift.""" + + executed = [] + state = {"provider_shift": 0} + + class CpuOp: + def __call__(self, value): + """Advance the synthetic chain value by one for candidate and golden replay.""" + + return value + 1 + + forward_fp32 = __call__ + + stages = [] + for index, stage in enumerate(h3_chain.STAGES): + + def stage_input(ctx, upstream, index=index): + """Use raw timesteps for stage zero and require predecessor outputs thereafter.""" + + if index == 0: + return ctx["timestep"] + assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" + return upstream + + stages.append( + replace( + stage, + candidate=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + golden=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + provider=lambda ctx, up, stage_input=stage_input, index=index: ( + stage_input(ctx, up) + 1 + (state["provider_shift"] if index == 0 else 0) + ), + ) + ) + + def get_op(op_type, *, device): + """Record each dispatched stage and return its deterministic CPU stub.""" + + executed.append(op_type) + return CpuOp() + + monkeypatch.setattr(h3_chain, "STAGES", stages) + monkeypatch.setattr(h3_chain, "make_context", lambda *a, **kw: {"timestep": torch.zeros(1)}) + monkeypatch.setattr(h3_chain, "golden_op", lambda *a: CpuOp()) + monkeypatch.setattr(h3_chain_replay, "KernelRegistry", lambda: SimpleNamespace(get_op=get_op)) + monkeypatch.setattr(h3_chain_replay, "load_h3_conditioning_weights", lambda *a: {}) + monkeypatch.setattr(h3_chain, "environment", lambda: {}) + monkeypatch.setattr(h3_chain, "git_state", lambda: {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + return executed, state + + +@pytest.mark.parametrize("provider_shift", [0, 1]) +@pytest.mark.parametrize( + "requested, execution_count", + [ + ("timestep_mlp_fp32", 2), + ("adaln_projection_3mod", 3), + ("timestep_sinusoid_h3,adaln_projection_3mod", 3), + ("timestep_mlp_fp32,timestep_sinusoid_h3", 2), + ], +) +def test_chain_selection_runs_dependencies( + cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift +): + """Execute prerequisite stages but report requested stages and preserve first drift.""" + + executed, state = cpu_chain_cli + state["provider_shift"] = provider_shift + out = tmp_path / "chain.json" + monkeypatch.setattr( + sys, + "argv", + [ + "h3_chain_replay.py", + "--stages", + requested, + "--timesteps", + "1", + "--seq-lens", + "3", + "--out", + str(out), + ], + ) + h3_chain_replay.main() + + report = json.loads(out.read_text()) + stage_names = [stage.name for stage in h3_chain.STAGES] + expected_report = [name for name in stage_names if name in requested.split(",")] + assert executed == stage_names[:execution_count] + assert report["executed_stages"] == executed + assert report["stages"] == expected_report + case = report["cases"][0] + assert [entry["stage"] for entry in case["stages"]] == expected_report + assert all(entry["chained_vs_golden"]["bitwise_equal"] for entry in case["stages"]) + assert case["first_drift"] == (stage_names[0] if provider_shift else None) + assert case["first_isolated_drift"] == (stage_names[0] if provider_shift else None) + summary = capsys.readouterr().out.split("; first_drift=", 1)[0] + for name in stage_names: + assert (f"{name}=" in summary) == (name in expected_report) + + +@pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) +@pytest.mark.parametrize("weights_env", [None, " "]) +def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + + if weights_env is None: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, weights_env) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + with pytest.raises(SystemExit, match=WEIGHTS_ENV): + h3_evidence.main() + assert not out.exists() + + +@pytest.mark.parametrize( + "op, weight_source", + [ + ("timestep_sinusoid_h3", "not_applicable"), + ("timestep_mlp_fp32", "pinned_checkpoint"), + ("adaln_projection_3mod", "pinned_checkpoint"), + ], +) +def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + + if weight_source == "not_applicable": + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + monkeypatch.setattr(h3_evidence, "KernelRegistry", lambda: object()) + monkeypatch.setattr( + h3_evidence, + "git_state", + lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, + ) + monkeypatch.setattr(h3_evidence, "environment", lambda: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py new file mode 100644 index 000000000..07cd49776 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End to end: the H3 conditioning chain on the pinned checkpoint, through the registry. + +Runs every stage registered in ``rl_engine.validation.models.h3_chain.STAGES`` (each +RFC #420 row adds its own), chained exactly as the model calls them, and +checks every stage's promises: + +* the registry dispatched the CUDA backend; +* a repeated run is bitwise equal; +* the chained output is within the stage's contract tolerance of the golden; +* stages that promise it are bitwise equal to diffusers on diffusers' inputs, + so the first isolated drift can only be a stage that declares a reduction. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_chain + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. +CASES = [(1, 3), (3, 257), (4, 4097)] + + +@pytest.fixture(scope="module") +def chain_weights(h3_weights_cpu): + """Move all pinned conditioning tensors to CUDA once for the end-to-end module.""" + + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} + + +@pytest.fixture(scope="module") +def registry(): + """Provide the production registry used to dispatch every conditioning stage.""" + + return KernelRegistry() + + +@pytest.mark.parametrize("num_timesteps, seq_len", CASES) +def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Check CUDA dispatch, golden tolerance, repeatability, and declared provider parity.""" + + report = h3_chain.run_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + stages = {stage.name: stage for stage in h3_chain.STAGES} + for entry in report["stages"]: + stage = stages[entry["stage"]] + assert entry["backend"].endswith("CudaOp"), entry + assert entry["repeat_bitwise_equal"], entry["stage"] + assert entry["chained_vs_golden"]["within_tolerance"], entry + if stage.provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] + drift = report["first_isolated_drift"] + assert drift is None or not stages[drift].provider_bitwise_isolated + + +@pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) +def test_chain_backward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" + + report = h3_chain.run_backward_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + for name, entry in report["leaves"].items(): + for mode in ("candidate", "candidate_fused"): + assert entry[mode]["repeat_bitwise_equal"], (mode, name) + fused = entry["candidate_fused"] + if name.startswith("time_embedder"): + # FP32 parameters: no BF16 rounding anywhere on the fused path. + assert fused["max_abs_vs_golden_over_absmax"] < 1e-5, (name, fused) + else: + # BF16 parameters: only their own final rounding remains. + assert fused["correctly_rounded_fraction"] > 0.99, (name, fused) diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py new file mode 100644 index 000000000..ca11c4e53 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Native input validation for the H3 deterministic linear kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available +from rl_engine.backends.extension import _C + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), + pytest.mark.skipif(not det_linear_available(), reason="rl_engine._C lacks h3_det_linear_*"), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_input_rejects_zero_rows(dtype): + """Reject empty gradient batches in native input-gradient kernels for both dtypes.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + weight = torch.empty(16, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_input(grad, weight, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_rejects_zero_rows(dtype, with_bias): + """Reject empty gradient batches before parameter-gradient or bias-gradient reduction.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + x = torch.empty(0, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +def test_backward_input_handles_empty_columns(dtype, out_dtype, k_in, n_out): + """Match the input VJP of empty linear dimensions in the requested output dtype.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + weight = torch.empty(n_out, k_in, device="cuda", dtype=dtype) + expected = (grad @ weight.float()).to(out_dtype) + actual = _C.h3_det_linear_backward_input(grad, weight, out_dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with_bias): + """Return empty weight gradients while preserving any nonempty bias reduction.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + x = torch.arange(3 * k_in, device="cuda", dtype=dtype).reshape(3, k_in) + expected_weight = (grad.T @ x.float()).to(w_dtype) + actual = _C.h3_det_linear_backward_weight(grad, x, w_dtype, with_bias) + assert len(actual) == (2 if with_bias else 1) + torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) + if with_bias: + torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..5fb964677 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``final_adaln_out``: norm_out shift/scale projection + final norm/modulation. + +* layout: ``shift`` is the first half of ``norm_out.linear``'s output, rows + are indexed by ``timestep_indices`` (not ``adaln_indices``); +* forward matches diffusers' ``norm_out`` to the projection's 1-ULP ties + (the norm and modulation are bitwise; the GEMV tree differs from cuBLAS); +* backward is deterministic, with the table gradient kept in FP32; +* rows are batch/position invariant; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.final_adaln_out import NativeH3FinalAdaLNOutOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_final_adaln_out + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3FinalAdaLNOutCudaOp() + + +@pytest.fixture(scope="module") +def norm_out(h3_weights_cpu): + return [ + h3_weights_cpu[name].cuda() + for name in ("norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias") + ] + + +def _inputs(seq, num_timesteps=3, seed=0, batch=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(batch, seq, 5376, generator=g) * 2).bfloat16().cuda() + temb = (torch.randn(num_timesteps, 2688, generator=g) * 2).cuda() + ti, _ = h3_packed_layout(seq, num_timesteps, seed=seed) + return x, temb, ti + + +def _golden(x, norm_weight, temb, weight, bias, timestep_indices): + act = temb * torch.sigmoid(temb) + # Declared cast of the FP32 SiLU, as the model computes it; identity VJP. + declared = F.silu(temb.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = F.linear(act, weight, bias) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) * norm_weight + return n * (1 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3FinalAdaLNOutOp() + x, nw = torch.zeros(1, 4, 8), torch.ones(8) + temb, w, b = torch.zeros(2, 16), torch.zeros(16, 16), torch.zeros(16) + ti = torch.tensor([0, 1, 1, 0]) + with pytest.raises(TypeError): # BF16 temb (probe H7) + op(x, nw, temb.bfloat16(), w, b, ti) + with pytest.raises(ValueError): # not 2H rows + op(x, nw, temb, w[:10], b[:10], ti) + with pytest.raises(IndexError): + op(x, nw, temb, w, b, ti + 2) + + +@requires_cuda +class TestCuda: + def test_shift_first_layout_and_timestep_indexing(self): + """Zero weights: the bias halves are (shift, scale); rows follow timestep_indices.""" + + hidden, seq = 16, 6 + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + nw = torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(2 * hidden, 32, device="cuda", dtype=torch.bfloat16) + b = torch.cat([torch.full((hidden,), 1.0), torch.full((hidden,), -1.0)]).bfloat16().cuda() + ti = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda") + out = _cuda_op()(x, nw, torch.zeros(2, 32, device="cuda"), w, b, ti) + # scale = -1 -> norm(x) * 0 + shift = 1 everywhere + assert torch.equal(out, torch.ones_like(out)) + + @pytest.mark.parametrize("seq, num_timesteps", [(1, 1), (257, 2), (4097, 3)]) + def test_forward_against_diffusers_and_golden(self, norm_out, seq, num_timesteps): + x, temb, ti = _inputs(seq, num_timesteps, seed=seq) + ours = _cuda_op()(x, *norm_out[:1], temb, *norm_out[1:], ti) + theirs = provider_final_adaln_out(x, norm_out[0], temb, *norm_out[1:], ti) + golden = NativeH3FinalAdaLNOutOp().forward_fp32(x, norm_out[0], temb, *norm_out[1:], ti) + assert (ours == theirs).float().mean() > 0.999 + torch.testing.assert_close(ours.float(), golden, atol=5e-2, rtol=2e-2) + + def test_rows_are_batch_and_position_invariant(self, norm_out): + x, temb, ti = _inputs(600, seed=1, batch=2) + full = _cuda_op()(x, norm_out[0], temb, *norm_out[1:], ti) + part = _cuda_op()(x[1:2, 200:260], norm_out[0], temb, *norm_out[1:], ti[200:260]) + assert torch.equal(part[0], full[1, 200:260]) + + def test_backward_deterministic_and_fp32_table_gradient(self, norm_out): + x, temb, ti = _inputs(4097, seed=2) + grad = torch.randn_like(x) + + def grads(fn, dtype=None): + tensors = [x, norm_out[0], temb, *norm_out[1:]] + leaves = [ + (t if dtype is None else t.to(dtype)).detach().clone().requires_grad_(True) + for t in tensors + ] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda *t: _cuda_op()(*t, ti)) + again = grads(lambda *t: _cuda_op()(*t, ti)) + ref = grads(lambda *t: _golden(*t, ti), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): + rel = ((g.double() - r).abs().max() / r.abs().max()).item() + assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + + def test_report_backward_uses_fp64_golden(self, monkeypatch): + from types import SimpleNamespace + + from rl_engine.validation.models import h3_report + + rng = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn(1, 160, 8, device="cuda", generator=rng).bfloat16() + nw = torch.randn(8, device="cuda", generator=rng).bfloat16() + temb = torch.randn(3, 4, device="cuda", generator=rng) + w = torch.randn(16, 4, device="cuda", generator=rng).bfloat16() + b = torch.randn(16, device="cuda", generator=rng).bfloat16() + ti = torch.arange(160, device="cuda") % 3 + inputs = (x, nw, temb, w, b) + native = NativeH3FinalAdaLNOutOp() + registry = SimpleNamespace( + get_op=lambda *args, **kwargs: native, + _get_or_create_backend=lambda _: native, + _priority_map={"cpu": {"final_adaln_out": [None]}}, + ) + monkeypatch.setattr(h3_report, "_final_inputs", lambda *args, **kwargs: (*inputs, ti)) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", native) + report = h3_report._final_accuracy(registry) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + leaves = [t.detach().clone().requires_grad_(True) for t in inputs] + native(*leaves, ti).backward(grad) + ref = [t.double().detach().requires_grad_(True) for t in inputs] + _golden(*ref, ti).backward(grad.double()) + for name, leaf, reference in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), leaves, ref): + assert reference.grad.dtype == torch.float64 + expected = float( + (leaf.grad.double() - reference.grad).abs().max() / reference.grad.abs().max() + ) + for backend in ("cuda", "provider"): + assert report["backward"][backend]["rel_error"][name] == pytest.approx(expected) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("final_adaln_out", device="cuda") + assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_bf16_gradients_meet_gtest_contract(self, seq): + # The check_operator path; S=257 used to miss d_norm_weight before the + # golden stored norm(x) and 1 + scale in BF16 like the model does. + import argparse + + from rl_engine.validation.operators import run_operator_suite + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + _cuda_op() + args = argparse.Namespace( + op="final_adaln_out", + candidate="cuda", + arch_key=None, + batch=3, + seq=seq, + normalized_dim=5376, + seed=123, + input_mode="random", + ) + report = run_operator_suite( + "final_adaln_out", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.bfloat16, torch.device("cuda"))], + check_grad=True, + ) + assert report.passed, report.candidates[0].cases[0].outputs diff --git a/tests/models/minimax_h3/test_h3_full_model.py b/tests/models/minimax_h3/test_h3_full_model.py new file mode 100644 index 000000000..b10898105 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_full_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_block +from rl_engine.validation.models.h3_full_model import ( + FullModelTrace, + full_model_execution, + validate_full_structure, +) + + +class Block(torch.nn.Module): + def __init__(self, value): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(float(value))) + + def forward(self, x, *args): + return x + self.weight + + +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + self.transformer_blocks = torch.nn.ModuleList([Block(i) for i in range(50)]) + for name in [ + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + setattr(self, name, torch.nn.Identity()) + self.token_refiner.refiner_blocks = torch.nn.ModuleList([Block(50), Block(51)]) + + +def test_full_stack_uses_each_layers_own_parameters_and_keeps_gradients(monkeypatch): + model = Model() + seen = [] + + def run(mode, ctx, ops): + p = ctx["weights"][h3_block.BLOCK + ".weight"] + seen.append(p.item()) + return {h3_block.OUTPUT: ctx["hidden"] + p} + + monkeypatch.setattr(h3_block, "run_graph", run) + original = model.transformer_blocks[0].forward + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + x = torch.ones(1, requires_grad=True) + for block in model.transformer_blocks: + x = block(x, None, None, None) + x.sum().backward() + assert x.item() == 1 + sum(range(50)) + assert seen == list(range(50)) + assert all(b.weight.grad.item() == 1 for b in model.transformer_blocks) + assert model.transformer_blocks[0].forward == original + + +def test_repeated_block_zero_is_not_full_model(): + model = Model() + model.transformer_blocks[49] = model.transformer_blocks[0] + with pytest.raises(ValueError, match="distinct"): + validate_full_structure(model) + + +def test_partial_execution_cannot_publish_full_model_pass(): + trace = FullModelTrace() + trace.observe("block.0", torch.ones(1)) + with pytest.raises(ValueError, match="incomplete"): + trace.validate() + + +def test_binding_restored_after_failure(): + model = Model() + original = model.transformer_blocks[12].forward + with pytest.raises(RuntimeError, match="injected"): + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + raise RuntimeError("injected") + assert model.transformer_blocks[12].forward == original + + +def test_candidate_execution_requires_every_node(): + trace = FullModelTrace() + trace.candidate = True + for name in [f"block.{i}" for i in range(50)] + [ + "refiner.0", + "refiner.1", + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + trace.calls[name] = 1 + with pytest.raises(ValueError, match="candidate nodes did not execute"): + trace.validate() + + +def test_drift_order_follows_execution_instead_of_lexical_names(): + from tools.validation.models.h3_full_compare import drift_order + + report = { + "trace": { + "calls": {"block.2": 1, "block.10": 1}, + "backward_calls": {"block.10": 1, "block.2": 1}, + } + } + names = [ + "video", + "nodes.block.10", + "nodes.block.2", + "nodes.block.2.grad", + "nodes.block.10.grad", + ] + assert drift_order(report, names) == [ + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "nodes.block.2.grad", + "video", + ] + + +def test_reference_inventory_rejects_tampered_tensor(tmp_path): + import hashlib + import json + + from tools.validation.models.h3_evidence_integrity import verify_inventory + + tensor = tmp_path / "tensor.pt" + tensor.write_bytes(b"original") + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps({"files": {"tensor.pt": hashlib.sha256(b"original").hexdigest()}}) + ) + verify_inventory(tmp_path, manifest) + tensor.write_bytes(b"changed") + with pytest.raises(ValueError, match="changed"): + verify_inventory(tmp_path, manifest) + + +def test_cuda_guard_rejects_one_fallback_among_native_bindings(): + from rl_engine.validation.models.h3_full_model import require_candidate_cuda + + native = "rl_engine.backends.cuda.model_specific.minimax_h3.NativeOp" + ops = SimpleNamespace( + backends=lambda: {"projection": native, "norm": "torch.nn.RMSNorm", "gate": native} + ) + with pytest.raises(ValueError, match="candidate norm fell back"): + require_candidate_cuda(ops) + + +def test_drift_order_interleaves_parameter_and_activation_backward_hooks(): + from tools.validation.models.h3_full_compare import drift_order + + trace = FullModelTrace() + for name in ( + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "parameter.block.10.weight", + "nodes.block.2.grad", + ): + trace.record_event(name) + trace.record_event("nodes.block.2") # checkpoint recomputation is not a new first event + names = list(reversed(trace.execution_order)) + assert drift_order({"trace": {"execution_order": trace.execution_order}}, names) == ( + trace.execution_order + ) diff --git a/tests/models/minimax_h3/test_h3_generation_resume.py b/tests/models/minimax_h3/test_h3_generation_resume.py new file mode 100644 index 000000000..5c0d3d56b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_generation_resume.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable candidate generation: a preempted run continues to the same result. + +A fake transformer and denoising loop stand in for the pipeline, so these run on CPU. +The loop feeds each call's output back into the next call's inputs (as the scheduler +does), so any wrong replay would change every later step. +""" + +from __future__ import annotations + +import inspect +import subprocess +import types + +import pytest +import torch + +from rl_engine.validation.models import h3_generation_resume as resume_module +from rl_engine.validation.models.h3_generation_resume import ( + ResumableRun, + ResumeError, + StepStore, + adopt_identity, + inputs_digest, + model_trees, +) + +STEPS = 6 +IDENTITY = {"commit": "abc", "seed": 42, "candidate": True, "preflight": False} + + +class FakeTransformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.computed = 0 + + def forward(self, hidden_states, audio_hidden_states, *, return_dict=True): + self.computed += 1 + video = torch.tanh(hidden_states) * 0.5 + audio = torch.sin(audio_hidden_states) * 0.25 + return _output((video, audio), {"return_dict": return_dict}) + + +def _output(values, kwargs): + if kwargs.get("return_dict", True): + return types.SimpleNamespace(sample=values[0], audio_sample=values[1]) + return values + + +def _initial(seed=0): + generator = torch.Generator().manual_seed(seed) + return torch.randn(8, 4, generator=generator), torch.randn(3, 5, generator=generator) + + +def denoise( + directory, *, resume, stop_after=None, initial=None, identity=IDENTITY, verify_last_step=False +): + """Run the fake loop through a ResumableRun; return (latents, transformer, run).""" + + video, audio = initial if initial is not None else _initial() + store = StepStore(directory) + run = ResumableRun(store, identity, resume=resume, verify_last_step=verify_last_step) + if not (directory / "initial.pt").exists(): + torch.save({"hidden_states": video, "audio_hidden_states": audio}, directory / "initial.pt") + module = FakeTransformer() + run.install(module, _output) + + def hook(mod, args, kwargs, output): + values = (output.sample, output.audio_sample) + run.record(args, kwargs, values) + if stop_after is not None and run.recorded >= stop_after: + raise KeyboardInterrupt("simulated preemption") + + module.register_forward_hook(hook, with_kwargs=True) + for _ in range(STEPS): + out = module(hidden_states=video, audio_hidden_states=audio) + # The scheduler's update: the next inputs depend on this call's outputs. + video = video + out.sample + audio = audio - 0.5 * out.audio_sample + return (video, audio), module, run + + +def _step_bytes(directory): + return { + p.name: p.read_bytes() + for p in sorted(directory.iterdir()) + if p.name.startswith("step-") and p.suffix in (".pt", ".json") + } + + +def test_uninterrupted_run_saves_every_step(tmp_path): + (video, _), module, run = denoise(tmp_path / "a", resume=False) + assert module.computed == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.pt")]) == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.in.json")]) == STEPS + assert run.summary()["replayed_steps"] == 0 + + +def test_resumed_run_with_verification_matches_an_uninterrupted_one(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True, verify_last_step=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert _step_bytes(tmp_path / "whole").keys() == _step_bytes(tmp_path / "cut").keys() + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + # Four steps were saved: three replayed, the fourth recomputed and verified, two new. + summary = run.summary() + assert summary["saved_steps_at_start"] == 4 + assert summary["replayed_steps"] == 3 + assert summary["recomputed_and_verified_step"] == 3 + assert module.computed == STEPS - 3 + assert [a["saved_steps"] for a in summary["attempts"]] == [0, 4] + + +def test_repeated_preemptions_still_converge(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + for stop in (2, 3, 5): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=stop > 2, stop_after=stop) + resumed, _, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert len(run.summary()["attempts"]) == 4 + + +def test_with_verification_a_complete_run_recomputes_only_its_last_step(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True, verify_last_step=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 1 + assert run.summary()["recomputed_and_verified_step"] == STEPS - 1 + + +def test_different_inputs_are_rejected_during_replay(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + # Rejected at the first replayed call, not only at the recomputed last saved step. + with pytest.raises(ResumeError, match="replayed step 0 has different inputs"): + denoise(tmp_path / "a", resume=True, initial=_initial(seed=1)) + + +def test_a_nondeterministic_recompute_is_rejected_when_verifying(tmp_path, monkeypatch): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + + original = FakeTransformer.forward + + def noisy(self, hidden_states, audio_hidden_states, *, return_dict=True): + out = original(self, hidden_states, audio_hidden_states, return_dict=return_dict) + return _output((out.sample + 1e-6, out.audio_sample), {"return_dict": return_dict}) + + monkeypatch.setattr(FakeTransformer, "forward", noisy) + with pytest.raises(ResumeError, match="differs from the saved output"): + denoise(tmp_path / "a", resume=True, verify_last_step=True) + # Without verification nothing recomputes the saved steps, so the drift goes unnoticed. + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "b", resume=False, stop_after=3) + denoise(tmp_path / "b", resume=True) + + +def test_a_different_identity_cannot_be_continued(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + with pytest.raises(ResumeError, match="different identity"): + denoise(tmp_path / "a", resume=True, identity={**IDENTITY, "commit": "other"}) + + +def test_an_existing_directory_needs_resume(tmp_path): + denoise(tmp_path / "a", resume=False, stop_after=None) + with pytest.raises(ResumeError, match="resume"): + denoise(tmp_path / "a", resume=False) + + +def test_a_corrupted_step_ends_the_valid_prefix(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=4) + (tmp_path / "a" / "step-002.pt").write_bytes(b"damaged") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 2 + + +def test_leftover_temporary_files_are_ignored(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "step-003.pt.tmp").write_bytes(b"half written") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 3 + + +def test_a_run_without_initial_inputs_is_not_resumable(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "initial.pt").unlink() + assert ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True).saved == 0 + + +def test_inputs_digest_tracks_value_dtype_and_shape_and_ignores_non_tensors(): + x = torch.arange(6.0).reshape(2, 3) + base = inputs_digest((), {"a": x, "flag": True}) + assert base == inputs_digest((), {"a": x.clone(), "flag": False}) + assert base != inputs_digest((), {"a": x + 1}) + assert base != inputs_digest((), {"a": x.double()}) + assert base != inputs_digest((), {"a": x.reshape(3, 2)}) + assert base != inputs_digest((), {"b": x}) + assert inputs_digest((x,), {}) != inputs_digest((), {"a": x}) + assert inputs_digest((), {"s": torch.tensor(3.0)}) # zero-dim tensors hash + + +def test_replay_wrapper_keeps_the_forward_signature(tmp_path): + """The pipeline selects the transformer's layout inputs from inspect.signature(forward).""" + + module = FakeTransformer() + before = list(inspect.signature(module.forward).parameters) + (tmp_path / "a").mkdir() + run = ResumableRun(StepStore(tmp_path / "a" / "x"), IDENTITY, resume=False) + uninstall = run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + uninstall() + assert list(inspect.signature(module.forward).parameters) == before + + +def test_replay_wrapper_keeps_a_wrapped_forward_signature(tmp_path): + """Same through an accelerate-style partial that itself carries __wrapped__.""" + + import functools + + module = FakeTransformer() + plain = module.forward + + def new_forward(m, *args, **kwargs): + return plain(*args, **kwargs) + + hooked = functools.update_wrapper(functools.partial(new_forward, module), plain) + module.forward = hooked + before = list(inspect.signature(module.forward).parameters) + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + run = ResumableRun(StepStore(tmp_path / "b"), IDENTITY, resume=False) + run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + + +# --- default resume: replay everything saved, recompute nothing -------------------------- + + +def test_default_resume_replays_every_saved_step_and_computes_only_the_rest(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + summary = run.summary() + assert (summary["saved_steps_at_start"], summary["replayed_steps"]) == (4, 4) + assert summary["recomputed_and_verified_step"] is None + assert summary["real_forward_calls"] == STEPS - 4 == module.computed + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + + +def test_a_complete_run_resumes_with_no_forward_pass_at_all(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 0 + assert run.real_calls == 0 + + +def test_each_attempt_only_needs_a_few_steps_to_make_progress(tmp_path): + """Attempts that each finish one new step still add up to the whole run.""" + + whole, _, _ = denoise(tmp_path / "whole", resume=False) + saved = 0 + result = run = None + for attempt in range(STEPS): + stop = saved + 1 if saved + 1 < STEPS else None # the last attempt is left to finish + try: + result, _, run = denoise(tmp_path / "cut", resume=attempt > 0, stop_after=stop) + except KeyboardInterrupt: + saved += 1 + continue + break + assert result is not None and saved == STEPS - 1 + assert all(torch.equal(a, b) for a, b in zip(whole, result, strict=True)) + assert run.real_calls == 1 and len(run.summary()["attempts"]) == STEPS + + +def test_the_validated_trace_survives_for_a_later_replay_only_attempt(tmp_path): + store = StepStore(tmp_path / "a") + assert store.load_trace() is None + (tmp_path / "a").mkdir() + trace = {"calls": {"block.0": 3}, "execution_order": ["block.0", "block.1"]} + store.save_trace(trace) + assert StepStore(tmp_path / "a").load_trace() == trace + assert not list((tmp_path / "a").glob("*.tmp")) + + +# --- identity tied to the model paths, and adopting a commit-based run ------------------- + + +def _git(repo, *args): + return subprocess.check_output(["git", *args], cwd=repo, text=True).strip() + + +@pytest.fixture +def repo(tmp_path, monkeypatch): + root = tmp_path / "repo" + (root / "model").mkdir(parents=True) + (root / "orchestration").mkdir() + (root / "model" / "kernel.py").write_text("v1\n") + (root / "orchestration" / "resume.py").write_text("v1\n") + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@example.com") + _git(root, "config", "user.name", "t") + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "a") + monkeypatch.setattr(resume_module, "MODEL_PATHS", ("model",)) + return root + + +def _commit(repo, path, text): + (repo / path).write_text(text) + _git(repo, "commit", "-q", "-am", f"change {path}") + return _git(repo, "rev-parse", "HEAD") + + +def _old_run(directory, commit, trees): + base = {k: v for k, v in IDENTITY.items() if k != "commit"} + store = StepStore(directory) + store.open({"commit": commit, **base}, resume=False) + return {**base, "model_trees": trees} + + +def test_model_trees_ignore_orchestration_changes_but_not_model_changes(repo): + first = model_trees("HEAD", cwd=repo) + _commit(repo, "orchestration/resume.py", "v2\n") + assert model_trees("HEAD", cwd=repo) == first + _commit(repo, "model/kernel.py", "v2\n") + assert model_trees("HEAD", cwd=repo) != first + + +def test_adopting_a_commit_based_run_lets_a_later_commit_continue_it(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + new_identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "orchestration/resume.py", "v2\n") + new_identity["model_trees"] = model_trees("HEAD", cwd=repo) + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert StepStore(tmp_path / "run").open(new_identity, resume=True) == 0 + attempts = StepStore(tmp_path / "run").attempts() + assert attempts[-2]["adopted"]["from_commit"] == old + assert attempts[-2]["adopted"]["model_paths_identical"] is True + + +def test_adoption_is_refused_when_a_model_path_changed(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "model/kernel.py", "v2\n") + new_identity = { + **{k: v for k, v in IDENTITY.items() if k != "commit"}, + "model_trees": model_trees("HEAD", cwd=repo), + } + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="model paths changed"): + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert (tmp_path / "run" / "run.json").read_text() == before + + +def test_adoption_refuses_other_differences_and_dry_run_writes_nothing(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="more than the commit"): + adopt_identity(tmp_path / "run", {**identity, "seed": 7}, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="not abc"): + adopt_identity(tmp_path / "run", identity, "abc", cwd=repo) + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo, dry_run=True) + assert (tmp_path / "run" / "run.json").read_text() == before + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="already"): + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py new file mode 100644 index 000000000..03a70f1a3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +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 == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@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( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) + + +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("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_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + 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", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + 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 diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py new file mode 100644 index 000000000..8f64a4508 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 block, forward and backward (RFC #420 ``ws1_one_h3_block``). + +GPU tests run block 0 of the pinned checkpoint on packed FL2VA layouts and +check the block-level promises of ``rl_engine.validation.models.h3_block``: + +* every node is bitwise repeatable, and batch row 0 of a two-row batch has the + bytes of the single-row run (forward, and the gradient reaching every node); +* nodes that promise it are bitwise equal to diffusers on diffusers' inputs; +* every node, the block output and every gradient are about as close to the + FP32 golden as diffusers itself. + +CPU tests pin the packed layouts to the pinned diffusers pipeline and check +the graph's structure. +""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.validation.models import h3_block +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# sha256(position_ids float64 bytes + token_tags int64 bytes) of diffusers@f53d552's +# ``MiniMaxH3PrepareLayoutStep.build_packed_sequence`` for each layout (text tags 1, patch +# (1, 2, 2), two audio channels, audio tag 2, video tag 0, no keyframe anchors). +PINNED_LAYOUTS = { + "tiny": "5329a0a13a9996c4fea7a0e676b65aa24302d65bea2e2248ba68d12184cc7c37", + "small": "64e4e79d8fb91d5a68ae7d15bb29049e08add38da2639030e089c6ef80968931", + "medium": "16e448800695eedfb85e0aacbca9bd5415ed53fbbe25f701bcb60df5f4754a31", +} + +# Accuracy: the candidate's error against the FP32 golden may exceed diffusers' by this +# factor plus an absolute floor (relative to the golden's max magnitude). +FORWARD_FACTOR, FORWARD_FLOOR = 1.25, 1e-3 +BACKWARD_FACTOR, BACKWARD_FLOOR = 1.5, 5e-3 +# A BF16 leaf gradient may instead exceed diffusers' error by this many ulps of its largest +# golden element. Its max|err| / max|golden| is set by one rounded element, so it moves in +# whole ulps (0.004 relative for attn.norm_q.weight), and the relative rule alone failed on +# seed 2 with torch 2.13 and on seed 0 with torch 2.9.1. Sized from +# tools/validation/models/h3_block_ulp_sweep.py over tiny/small x seeds 0-9 on both torch +# versions (max excess +4.15 ulp) and 34 single-node candidate/provider swaps (max +5.15 ulp); +# reports in reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json. +BACKWARD_BF16_EXCESS_ULPS = 6.0 + + +@pytest.mark.parametrize("layout", sorted(PINNED_LAYOUTS)) +def test_layout_matches_pinned_pipeline(layout): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS[layout], device="cpu") + data = packed["position_ids"].numpy().tobytes() + packed["token_tags"].numpy().tobytes() + assert hashlib.sha256(data).hexdigest() == PINNED_LAYOUTS[layout] + + +def test_graph_is_well_formed(): + seen = {"hidden", "temb"} + for node in h3_block.NODES: + assert set(node.inputs) <= seen, node.name + assert node.name not in seen, node.name + assert node.binding.status in ("row", "interim", "reference"), node.name + seen.add(node.name) + assert h3_block.OUTPUT == "residual_mlp" + rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} + assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + + +def test_partial_rope_rotates_96_of_128_channels(): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS["tiny"], device="cpu") + cos, sin = prov.provider_rope(packed["position_ids"]) + assert cos.shape == (packed["position_ids"].shape[0], 96) and cos.dtype == torch.float32 + gold_cos, gold_sin = gold.golden_rope_tables(packed["position_ids"]) + assert torch.equal(cos, gold_cos) and torch.equal(sin, gold_sin) + x = torch.randn(1, cos.shape[0], 2, 128).bfloat16() + for out in (prov.provider_apply_rotary(x, cos, sin), gold.golden_apply_rotary(x, cos, sin)): + assert torch.equal(out[..., 96:].float(), x[..., 96:].float()) + assert not torch.equal(out[..., :96].float(), x[..., :96].float()) + + +@pytest.fixture(scope="module") +def ops(): + return h3_block.CandidateOps() + + +@needs_gpu +def test_stack_rows_dispatch_to_cuda(require_h3_kernels, ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + assert backends[name].endswith("CudaOp"), backends + + +@needs_gpu +@pytest.mark.parametrize("layout", ["tiny", "small"]) +def test_block_forward(h3_block_weights_cuda, ops, layout): + report = h3_block.run_block_case(h3_block_weights_cuda, layout=layout, ops=ops) + nodes = {node.name: node for node in h3_block.NODES} + for entry in report["nodes"]: + name = entry["node"] + assert entry["repeat_bitwise_equal"], name + assert entry["batch_row_bitwise_equal"], name + if nodes[name].provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], name + bound = FORWARD_FACTOR * entry["provider_rel_err_vs_golden"] + FORWARD_FLOOR + assert entry["candidate_rel_err_vs_golden"] <= bound, entry + assert report["first_repeat_drift"] is None + assert report["first_batch_drift"] is None + drift = report["first_isolated_drift"] + assert drift is None or not nodes[drift].provider_bitwise_isolated + out = report["output"] + assert out["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * out["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ) + + +@needs_gpu +def test_block_backward(h3_block_weights_cuda, ops): + report = h3_block.run_block_backward(h3_block_weights_cuda, layout="tiny", ops=ops) + for name, entry in report["leaves"].items(): + assert entry["repeat_bitwise_equal"], name + bound = BACKWARD_FACTOR * entry["provider_rel_err_vs_golden"] + BACKWARD_FLOOR + within_rel = entry["candidate_rel_err_vs_golden"] <= bound + within_ulps = entry["storage_dtype"] == "bfloat16" and ( + entry["candidate_ulps_vs_golden"] + <= entry["provider_ulps_vs_golden"] + BACKWARD_BF16_EXCESS_ULPS + ) + assert within_rel or within_ulps, (name, entry) + assert report["leaves"]["hidden"]["batch_row_bitwise_equal"] + assert report["first_grad_repeat_drift"] is None + assert report["first_grad_batch_drift"] is None diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..ed74d1901 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements, and +comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models import h3_chain, h3_report +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import ( + provider_adaln_row_gather, + provider_norm_modulate, +) +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run tools/weights/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = ( + Path(__file__).resolve().parents[3] + / "tools" + / "validation" + / "models" + / "plot_h3_evidence.py" + ) + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + + assert enable_timing + + def record(self): + """Capture the synthetic elapsed clock at this event.""" + + self.timestamp = elapsed_us + + def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + + pass + + def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py new file mode 100644 index 000000000..ced9a9377 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_rmsnorm``: RMSNorm (eps 1e-5) and its fused AdaLN modulation. + +* forward is bitwise equal to ``nn.RMSNorm`` and to diffusers' + ``norm(x) * (1.0 + scale[i]) + shift[i]`` (block: ``adaln_indices``; + ``norm_out``: ``timestep_indices``) on the pinned norm weights; +* rows are independent of batch size and position; +* backward is deterministic and checked against an FP64 golden; +* malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.rmsnorm import NativeH3RMSNormOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_norm_modulate + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NORMS = ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", +) + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import ( + H3RMSNormCudaOp, + h3_rmsnorm_available, + ) + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3RMSNormCudaOp() + + +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(shape, generator=g) * scale).to(device=device, dtype=dtype) + + +def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(rows, chunks * hidden, generator=g) * 0.5).to(dtype).cuda() + return table.chunk(chunks, dim=-1) # strided (R, H) views, like the AdaLN outputs + + +def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * weight.double() + return n * (1 + scale.double().index_select(0, index)) + shift.double().index_select(0, index) + + +class TestReference: + def test_golden_close_to_provider(self): + x, w = _x((3, 64), torch.float32, device="cpu"), torch.rand(64) + 0.5 + op = NativeH3RMSNormOp() + torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + + def test_modulated_with_prevalidated_indices(self): + x, w = _x((2, 8), torch.float32, device="cpu"), torch.ones(8) + shift, scale = torch.randn(3, 8), torch.randn(3, 8) + index = torch.tensor([2, 0]) + op = NativeH3RMSNormOp() + expected = op.forward(x, w) * (1 + scale[index]) + shift[index] + torch.testing.assert_close( + op.forward_modulated(x, w, shift, scale, index, check_range=False), + expected, + ) + + def test_rejects_bad_inputs(self): + op = NativeH3RMSNormOp() + x, w = torch.randn(2, 8), torch.ones(8) + with pytest.raises(TypeError): + op.forward(x, w.double()) + with pytest.raises(TypeError): + op.forward(x, w.bfloat16()) + with pytest.raises(ValueError): + op.forward(x, torch.ones(7)) + with pytest.raises(ValueError): + op.forward(x, w, eps=0.0) + shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(ValueError): # one index per position S + op.forward_modulated(x, w, shift, scale, torch.tensor([0])) + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("name", NORMS) + def test_bitwise_equal_to_nn_rmsnorm_on_pinned_weights(self, h3_weights_cpu, name): + weight = h3_weights_cpu[name].cuda() + module = torch.nn.RMSNorm(5376, eps=1e-5).cuda().bfloat16() + module.weight.data.copy_(weight) + x = _x((2, 777, 5376), seed=1) + assert torch.equal(_cuda_op()(x, weight), module(x)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [12, 256, 4096, 5376]) + def test_bitwise_equal_to_f_rms_norm(self, dtype, hidden): + x = _x((5, 33, hidden), dtype, seed=hidden, scale=30.0) + w = (torch.rand(hidden, device="cuda") * 2).to(dtype) + ref = torch.nn.functional.rms_norm(x, (hidden,), w, 1e-5) + assert torch.equal(_cuda_op()(x, w), ref) + + def test_block_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["transformer_blocks.0.norm1.weight"].cuda() + shift, scale, *_ = _table(9, 5376, 6, seed=2) # (3T, H) rows, T = 3 + ti, tags = h3_packed_layout(4097, 3, seed=2) + index = ti * 3 + tags + x = _x((2, 4097, 5376), seed=2) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, index)) + + def test_norm_out_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["norm_out.norm.weight"].cuda() + shift, scale = _table(3, 5376, 2, seed=3) # (T, H) rows indexed by timestep + ti, _ = h3_packed_layout(1000, 3, seed=3) + x = _x((1, 1000, 5376), seed=3) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, ti) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, ti)) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_modulation_bitwise_in_every_dtype(self, dtype): + """FP32 must not contract the separately rounded eager ops into an FMA.""" + + w = (torch.rand(5376, device="cuda") + 0.5).to(dtype) + shift, scale, *_ = _table(9, 5376, 6, dtype=dtype, seed=9) + ti, tags = h3_packed_layout(333, 3, seed=9) + index = ti * 3 + tags + x = _x((2, 333, 5376), dtype, seed=9) + ours = _cuda_op().forward_modulated(x, w, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + + def test_rows_are_batch_and_position_invariant(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + shift, scale, *_ = _table(9, 5376, 6, seed=4) + ti, tags = h3_packed_layout(300, 3, seed=4) + index = ti * 3 + tags + x = _x((3, 300, 5376), seed=4) + full = _cuda_op().forward_modulated(x, w, shift, scale, index) + part = _cuda_op().forward_modulated(x[2:3, 50:90], w, shift, scale, index[50:90]) + assert torch.equal(part[0], full[2, 50:90]) + + def test_rejects_unsupported(self): + w = torch.ones(6, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(_x((2, 6)), w) + with pytest.raises(ValueError): + _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + + +_DEVICE_MESSAGES = { + "weight": "with x's dtype and device", + "index": r"row index must be a non-empty contiguous int64 tensor on cuda:0", + "rstd": "on x's device", +} + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +class TestCudaNativeDeviceValidation: + @pytest.mark.parametrize("argument", ["weight", "index"]) + def test_forward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + if argument == "weight": + weight = weight.to("cuda:1") + else: + index = index.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @pytest.mark.parametrize("argument", ["weight", "index", "rstd"]) + def test_backward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + if argument == "weight": + weight = weight.to("cuda:1") + elif argument == "index": + index = index.to("cuda:1") + else: + rstd = rstd.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, fn, tensors, grad): + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + def test_modulated_backward_against_fp64(self): + g = torch.Generator(device="cpu").manual_seed(5) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.5).bfloat16().cuda() + ti, tags = h3_packed_layout(2049, 3, seed=5) + index = ti * 3 + tags + x = _x((2, 2049, 5376), seed=5) + grad = _x((2, 2049, 5376), seed=6, scale=1.0) + op = _cuda_op() + + def ours(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return op.forward_modulated(x_, w_, shift, scale, index) + + def golden(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return _fp64_modulated(x_, w_, shift, scale, index) + + got = self._grads(ours, (x, w, table), grad) + ref = self._grads(golden, (x.double(), w.double(), table.double()), grad.double()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + rel = ((g_.double() - r).abs().max() / r.abs().max()).item() + # Two BF16 roundings remain: the eager VJP uses round(1 + scale), and + # every gradient is rounded once to BF16 (2^-9 relative each). + assert rel < 1e-2, (name, rel) + again = self._grads(ours, (x, w, table), grad) + assert all(torch.equal(a, b) for a, b in zip(got, again)) + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_modulated_gradients_meet_contract_against_golden(self, seq): + # Against forward_modulated_fp32 before it stored norm(x) and 1 + scale + # in BF16, d_weight missed from S=1024 and d_scale at S=4097. + from rl_engine.contracts.numerical import load_contract, resolve_tolerance + + spec = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=torch.bfloat16, + ) + g = torch.Generator(device="cpu").manual_seed(seq) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.1).bfloat16().cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seq) + index = ti * 3 + tags + x = _x((1, seq, 5376), seed=seq) + grad = _x((1, seq, 5376), seed=seq + 1, scale=1.0) + op, native = _cuda_op(), NativeH3RMSNormOp() + + def run(fn, upstream): + def call(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return fn(x_, w_, shift, scale, index) + + return self._grads(call, (x, w, table), upstream) + + got = run(op.forward_modulated, grad) + ref = run(native.forward_modulated_fp32, grad.float()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + torch.testing.assert_close( + g_.float(), r.float(), atol=spec.atol, rtol=spec.rtol, msg=lambda m: f"{name}: {m}" + ) + + def test_plain_backward_against_fp64_and_row_local_dx(self): + w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() + x = _x((4, 64, 5376), seed=7) + grad = _x((4, 64, 5376), seed=8, scale=1.0) + op = _cuda_op() + got = self._grads(lambda a, b: op(a, b), (x, w), grad) + + def golden(a, b): + return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + 1e-5) * b + + ref = self._grads(golden, (x.double(), w.double()), grad.double()) + for g, r in zip(got, ref): + assert ((g.double() - r).abs().max() / r.abs().max()).item() < 5e-3 + part = self._grads(lambda a, b: op(a, b), (x[1:2, 10:20], w), grad[1:2, 10:20]) + assert torch.equal(part[0][0], got[0][1, 10:20]) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + assert ( + type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" + ) diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..a60489a04 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,327 @@ +# 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.models.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]) + + 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) + + +@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.models 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.models 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/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py new file mode 100644 index 000000000..555a28ebe --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_mlp_fp32``: FP32 256 -> 5376 -> 2688 timestep MLP. + +* accuracy: forward and every gradient against an FP64 golden, on the pinned + checkpoint weights and on synthetic non-H3 sizes; +* invariance: a timestep's ``temb`` bytes (and its row-local ``dx``) do not + depend on how many timesteps share the call or where it sits; +* determinism: repeated forward/backward runs are bitwise equal; +* contract: BF16 inputs (RFC probe H7) and malformed shapes fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.timestep_mlp import NativeH3TimestepMLPOp +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NAMES = ("x", "w1", "b1", "w2", "b2") + + +def _cuda_op(): + """Construct the CUDA MLP operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import H3TimestepMLPCudaOp + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3TimestepMLPCudaOp() + + +def _features(num, device="cuda", seed=0): + """Generate seeded H3 sinusoidal features on the requested test device.""" + + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) + + +def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + """Create seeded FP32 parameters for a two-layer SiLU MLP of the requested sizes.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 + b1 = torch.randn(hidden, generator=g) * 0.1 + w2 = torch.randn(out, hidden, generator=g) / hidden**0.5 + b2 = torch.randn(out, generator=g) * 0.1 + return [t.to(device) for t in (w1, b1, w2, b2)] + + +def _real_params(weights): + """Return pinned timestep MLP weights and biases on CUDA in layer order.""" + + return [ + weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") + ] + + +def _fp64_grads(x, params, grad_out): + """Compute gradients of both linear layers and SiLU using FP64 autograd.""" + + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] + z = F.linear(leaves[0], leaves[1], leaves[2]) + F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) + return [leaf.grad for leaf in leaves] + + +class TestReference: + def test_provider_and_golden_agree(self): + """Keep the FP32 provider close to the high-precision MLP replay on CPU.""" + + x = _features(5, device="cpu") + params = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + torch.testing.assert_close( + op.forward(x, *params), op.forward_fp32(x, *params), atol=1e-5, rtol=1e-5 + ) + + @pytest.mark.parametrize("which", range(5)) + def test_rejects_bf16_anywhere(self, which): + # RFC probe H7: casting a declared FP32 path to BF16 early. + """Reject BF16 in every argument of the declared FP32 MLP path.""" + + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] + args[which] = args[which].bfloat16() + with pytest.raises(TypeError, match="float32"): + NativeH3TimestepMLPOp().forward(*args) + + def test_rejects_shape_mismatches(self): + """Reject incompatible layer dimensions, malformed biases, and empty feature batches.""" + + x = _features(2, device="cpu") + w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + with pytest.raises(ValueError): + op.forward(x[:, :128], w1, b1, w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1[:10], w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1, w2[:, :10], b2) + with pytest.raises(ValueError): + op.forward(x[:0], w1, b1, w2, b2) + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) + def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + """Check checkpoint forward accuracy, FP32 dtype, and the timestep output shape.""" + + params = _real_params(h3_weights_cpu) + x = _features(num, seed=num) + out = _cuda_op()(x, *params) + gold = NativeH3TimestepMLPOp().forward_fp32(x, *params) + assert out.dtype == torch.float32 and out.shape == (num, 2688) + # tolerance_contract.json forward_accuracy / reduction / float32. + torch.testing.assert_close(out, gold, atol=1e-4, rtol=1e-4) + assert (out - gold).abs().max().item() < 1e-5 + + def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + """Keep each checkpoint embedding row bitwise unchanged across batch sizes and positions.""" + + op = _cuda_op() + params = _real_params(h3_weights_cpu) + x_all = _features(9, seed=1) + full = op(x_all, *params) + for i in range(9): + assert torch.equal(op(x_all[i : i + 1], *params)[0], full[i]) + for num in (2, 3, 4, 5, 8): + for pos in (0, num - 1): + x = _features(num, seed=10 + num) + x[pos] = x_all[4] + assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) + + def test_backward_against_fp64_golden(self, h3_weights_cpu): + """Match all five checkpoint input gradients to the FP64 autograd golden.""" + + params = _real_params(h3_weights_cpu) + x = _features(3, seed=2) + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + grad_out = torch.randn( + 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) + ) + _cuda_op()(*leaves).backward(grad_out) + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out), strict=True): + # tolerance_contract.json gradient_accuracy / reduction / float32. + torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) + def test_other_sizes(self, k_in, hidden, out): + """Match the high-precision forward golden for small and checkpoint-sized MLPs.""" + + params = _synthetic_params(k_in, hidden, out, seed=k_in) + x = torch.randn(5, k_in, device="cuda") + torch.testing.assert_close( + _cuda_op()(x, *params), + NativeH3TimestepMLPOp().forward_fp32(x, *params), + atol=1e-4, + rtol=1e-4, + ) + + def test_zero_hidden_size_matches_cpu_forward_backward(self): + """Match CPU forward values and all input gradients with a zero-width hidden layer.""" + + generator = torch.Generator(device="cpu").manual_seed(0) + cpu_inputs = [ + torch.randn(3, 8, generator=generator), + torch.empty(0, 8), + torch.empty(0), + torch.empty(16, 0), + torch.linspace(-1.0, 1.0, 16), + ] + cpu_inputs = [tensor.requires_grad_(True) for tensor in cpu_inputs] + cuda_inputs = [tensor.detach().cuda().requires_grad_(True) for tensor in cpu_inputs] + cpu_out = NativeH3TimestepMLPOp()(*cpu_inputs) + cuda_out = _cuda_op()(*cuda_inputs) + torch.testing.assert_close(cuda_out.cpu(), cpu_out, atol=0, rtol=0) + grad = torch.randn(3, 16, generator=generator) + cpu_out.backward(grad) + cuda_out.backward(grad.cuda()) + for name, cpu_input, cuda_input in zip(NAMES, cpu_inputs, cuda_inputs, strict=True): + torch.testing.assert_close( + cuda_input.grad.cpu(), cpu_input.grad, atol=0, rtol=0, msg=name + ) + + def test_rejects_unaligned_k(self): + """Reject inner dimensions that violate the native FP32 four-element alignment.""" + + params = _synthetic_params(6, 8, 8) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(torch.randn(2, 6, device="cuda"), *params) + + def test_rejects_cpu_and_mixed_devices(self): + """Reject CPU-only and mixed-device inputs at the CUDA MLP boundary.""" + + params = _synthetic_params(8, 24, 16) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *params) + + def test_repeat_and_backward_are_deterministic(self): + """Require repeated MLP outputs and all input gradients to be bitwise equal.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=3) + x = _features(4, seed=3) + grad_out = torch.randn(4, 2688, device="cuda") + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + out = op(*leaves) + out.backward(grad_out) + runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_dx_rows_are_batch_invariant(self): + """Keep feature gradients bitwise unchanged when rows are evaluated independently.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=4) + x = _features(6, seed=4) + grad_out = torch.randn(6, 2688, device="cuda") + + def dx(rows): + """Compute selected feature-row gradients using matching upstream rows.""" + + leaf = x[rows].detach().clone().requires_grad_(True) + op(leaf, *params).backward(grad_out[rows]) + return leaf.grad + + full = dx(slice(0, 6)) + for i in range(6): + assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + + def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") + assert type(op).__name__ == "H3TimestepMLPCudaOp" + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +class TestCudaDetLinearDeviceContract: + def test_forward_rejects_bias_on_another_device(self, dtype): + """Reject a projection bias on a different CUDA device from its weight.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_forward + + _cuda_op() + x = torch.randn(2, 8, device="cuda:0", dtype=dtype) + weight = torch.randn(16, 8, device="cuda:0", dtype=dtype) + bias = torch.randn(16, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="bias and weight must be on the same device"): + det_linear_forward(x, weight, bias) + + def test_backward_input_rejects_weight_on_another_device(self, dtype): + """Reject native input-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_input, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + weight = torch.randn(16, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and weight must be on the same device"): + det_linear_backward_input(grad, weight, dtype) + + def test_backward_weight_rejects_input_on_another_device(self, dtype): + """Reject native parameter-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_weight, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + x = torch.randn(2, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and x must be on the same device"): + det_linear_backward_weight(grad, x, dtype) diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..277beadea --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_sinusoid_h3``: FP32 [cos | sin] timestep features. + +Three properties, each checked separately: + +* provider parity: the CUDA kernel is bitwise equal to the diffusers + ``get_timestep_embedding`` replay on the same device (an elementwise op + with the same FP32 operation sequence); +* accuracy: within the FP32 elementwise contract of an FP64 golden; +* invariance: a row's bytes do not depend on how many timesteps share the + launch, their order, or repeated runs. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps +from rl_engine.validation.models.h3_provider import provider_time_proj + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + """Construct the CUDA sinusoid operator, skipping a missing native forward entrypoint.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + h3_sinusoid_cuda_available, + ) + + if not h3_sinusoid_cuda_available(): + pytest.skip("rl_engine._C lacks h3_timestep_sinusoid_forward") + return H3TimestepSinusoidCudaOp() + + +def _dense_unit_sweep() -> torch.Tensor: + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + edges = torch.tensor([0x00000001, 0x007FFFFF, 0x00800000, 0x3F800000], dtype=torch.int32) + return torch.cat([bits, edges.cuda()]).view(torch.float32) + + +class TestReference: + def test_layout_is_cos_then_sin(self): + """Check FP32 cosine-then-sine layout and the unit-frequency boundary values.""" + + t = torch.tensor([0.0, 1.0]) + out = NativeH3TimestepSinusoidOp().forward(t) + assert out.shape == (2, 256) and out.dtype == torch.float32 + assert torch.equal(out[0, :128], torch.ones(128)) + assert torch.equal(out[0, 128:], torch.zeros(128)) + # k = 0 has frequency 1: cos(1), sin(1) at t = 1. + assert out[1, 0].item() == pytest.approx(math.cos(1.0), abs=1e-7) + assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) + + def test_provider_path_matches_diffusers_cpu(self): + """Require CPU sinusoid values to match the diffusers provider bitwise.""" + + t = h3_timesteps(9, device="cpu") + ours = NativeH3TimestepSinusoidOp().forward(t) + assert torch.equal(ours, provider_time_proj(t, 256)) + + def test_golden_close_to_provider_path(self): + """Keep the CPU provider within FP32 error of the high-precision sinusoid golden.""" + + t = h3_timesteps(33, device="cpu") + op = NativeH3TimestepSinusoidOp() + torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) + + @pytest.mark.parametrize( + "bad, error", + [ + (torch.tensor([[0.5]]), ValueError), # not 1-D + (torch.empty(0), ValueError), # empty + (torch.tensor([1, 0]), TypeError), # integer + (torch.tensor([0.5, 1.5]), ValueError), # H10: t outside [0, 1] + (torch.tensor([500.0]), ValueError), # H10: t * 1000 convention + (torch.tensor([float("nan")]), ValueError), + ], + ) + def test_rejects_invalid_timesteps(self, bad, error): + """Reject empty, malformed, non-floating, non-finite, or out-of-range timesteps.""" + + with pytest.raises(error): + NativeH3TimestepSinusoidOp().forward(bad) + + def test_rejects_odd_channels(self): + """Reject odd feature counts that cannot split into equal cosine and sine halves.""" + + with pytest.raises(ValueError): + NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) + def test_bitwise_equal_to_diffusers_cuda_path(self, num): + """Require CUDA sinusoid values to match both provider replays for varied batch sizes.""" + + t = h3_timesteps(num, seed=num) + out = _cuda_op()(t) + assert out.dtype == torch.float32 and out.shape == (num, 256) + assert torch.equal(out, provider_time_proj(t, 256)) + assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + + def test_bitwise_over_dense_fp32_sweep_of_unit_interval(self): + # Every 997th FP32 bit pattern in [0, 1] plus the subnormals' edge: a + # fast-math build (approximate __expf/__sinf/__cosf, FTZ) fails here. + t = _dense_unit_sweep() + assert torch.equal(_cuda_op().forward(t, check_range=False), provider_time_proj(t, 256)) + + def test_sglang_timestep_embedding_is_bitwise_equal(self): + # RFC #420 reuse rule: parity with SGLang is checked bit for bit, not + # through matching FP64 error summaries. + jit = pytest.importorskip("sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit") + t = torch.cat([_dense_unit_sweep(), h3_timesteps(4097, seed=3)]) + theirs = jit.timestep_embedding(t, 256, flip_sin_to_cos=True, downscale_freq_shift=0.0) + assert torch.equal(theirs, provider_time_proj(t, 256)) + assert torch.equal(theirs, _cuda_op().forward(t, check_range=False)) + + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) + def test_other_channel_counts_match_reference(self, num_channels): + """Preserve bitwise reference parity across supported even channel counts.""" + + t = h3_timesteps(5) + out = _cuda_op()(t, num_channels=num_channels) + ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) + assert torch.equal(out, ref) + + def test_fp32_contract_against_fp64_golden(self): + """Bound CUDA sinusoid error against the high-precision golden by FP32 precision.""" + + t = h3_timesteps(257) + out = _cuda_op()(t) + gold = NativeH3TimestepSinusoidOp().forward_fp32(t) + # tolerance_contract.json forward_accuracy / elementwise / float32. + torch.testing.assert_close(out, gold, atol=1e-5, rtol=1e-5) + assert (out - gold).abs().max().item() <= 2.0**-23 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_low_precision_timesteps_are_upcast_first(self, dtype): + """Match provider FP32 upcasting for BF16 and FP16 timesteps.""" + + t = h3_timesteps(6).to(dtype) + assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) + + def test_batch_position_and_repeat_invariance(self): + """Keep row bits unchanged across batch positions, repeated calls, and permutations.""" + + op = _cuda_op() + base = torch.tensor(0.3141592, device="cuda") + single = op(base.view(1)) + for num in (2, 5, 64, 1023): + for pos in (0, num // 2, num - 1): + t = h3_timesteps(num, seed=num) + t[pos] = base + assert torch.equal(op(t)[pos], single[0]), (num, pos) + t = h3_timesteps(129, seed=11) + first = op(t) + for _ in range(3): + assert torch.equal(op(t), first) + perm = torch.randperm(129, generator=torch.Generator().manual_seed(0)).cuda() + assert torch.equal(op(t[perm]), first[perm]) + + def test_non_contiguous_input(self): + """Treat strided timesteps identically to their contiguous copies.""" + + t = h3_timesteps(20) + strided = t[::2] + assert not strided.is_contiguous() + assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) + + def test_rejects_cpu_input(self): + """Reject CPU timesteps at the CUDA sinusoid boundary.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.5])) + + def test_rejects_out_of_range(self): + """Reject scaled timesteps outside the declared unit interval.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + + @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) + def test_native_entrypoint_rejects_invalid_values(self, bad): + """Reject invalid values in native and wrapper calls even when neighbors are valid.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + # Valid neighbors must not hide an invalid timestep in the native call. + t = torch.tensor([0.0, bad, 1.0], device="cuda") + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _C.h3_timestep_sinusoid_forward(t) + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _cuda_op()(t) + + def test_native_entrypoint_accepts_boundaries(self): + """Accept both interval endpoints and match provider values in the native call.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + t = torch.tensor([0.0, 0.5, 1.0], device="cuda") + assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) + + def test_trusted_path_supports_cuda_graph(self): + """Capture and replay the prevalidated sinusoid path without changing its output.""" + + op = _cuda_op() + t = h3_timesteps(7) + expected = op(t) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op.forward(t, check_range=False) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op.forward(t, check_range=False) + graph.replay() + assert torch.equal(out, expected) + + def test_backward_matches_autograd_reference(self): + """Match timestep gradients to the high-precision autograd sinusoid replay.""" + + t = h3_timesteps(17).requires_grad_(True) + t_ref = t.detach().clone().requires_grad_(True) + grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) + torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + + def test_backward_is_row_invariant(self): + """Keep a timestep gradient bitwise unchanged when unrelated rows join the batch.""" + + op = _cuda_op() + grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) + + def grad_of(t, pos): + """Differentiate one selected timestep using a fixed upstream feature-gradient row.""" + + t = t.clone().requires_grad_(True) + grad = torch.zeros(t.shape[0], 256, device="cuda") + grad[pos] = grad_row + op(t).backward(grad) + return t.grad[pos] + + single = grad_of(torch.tensor([0.7], device="cuda"), 0) + t = h3_timesteps(40) + t[13] = 0.7 + assert torch.equal(grad_of(t, 13), single) + + def test_registry_dispatches_cuda(self): + """Resolve the sinusoid registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") + assert type(op).__name__ == "H3TimestepSinusoidCudaOp" diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py new file mode 100644 index 000000000..c9b511153 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import adaln_column_shard +from rl_engine.validation.models.h3_ws2 import ( + run_world, + tp_matches_ws1, + tp_projection_rank, + ws1_projection, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @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]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@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_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + 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_weights.py b/tests/models/minimax_h3/test_h3_weights.py new file mode 100644 index 000000000..13b9b2489 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weights.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Extracted H3 weights must retain the pinned tensor contents.""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.validation.models import h3_weights + + +@pytest.mark.parametrize( + "dtype,raw_bytes", + [ + (torch.float32, bytes.fromhex("000000000000803f")), + (torch.bfloat16, bytes.fromhex("0000803f")), + ], +) +def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + """Hash the exact FP32 and BF16 bit patterns without numerical conversion.""" + + tensor = torch.tensor([0.0, 1.0], dtype=dtype) + assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +@pytest.mark.parametrize("names", [None, ["bias"]]) +def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + + from safetensors.torch import save_file + + tensors = { + "weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=dtype), + "bias": torch.tensor([0.0, 1.0], dtype=dtype), + } + manifest = { + "tensors": { + name: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "sha256": h3_weights.sha256_tensor(tensor), + } + for name, tensor in tensors.items() + } + } + if mismatch == "dtype": + tensors["weight"] = tensors["weight"].to(torch.float64) + elif mismatch == "shape": + tensors["weight"] = tensors["weight"].reshape(4) + elif mismatch == "sha256": + tensors["weight"][0, 0] = 1.0 + # Tensor checksums do not depend on serializer metadata or tensor key order. + save_file( + dict(reversed(list(tensors.items()))), + str(tmp_path / h3_weights.EXTRACTED_FILE), + metadata={"revision": "a different serializer metadata value"}, + ) + monkeypatch.setenv(h3_weights.WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(h3_weights, "load_h3_manifest", lambda: manifest) + + if mismatch: + with pytest.raises(ValueError, match=f"weight: {mismatch} .* != manifest"): + h3_weights.load_h3_conditioning_weights("cpu", names=names) + else: + actual = h3_weights.load_h3_conditioning_weights("cpu", names=names) + assert list(actual) == (names if names is not None else list(tensors)) + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/tests/models/minimax_h3/test_h3_ws2_processes.py b/tests/models/minimax_h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..f448ba580 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.validation.models import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=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"): + 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" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before diff --git a/tests/models/minimax_h3/test_prepare_h3_weights.py b/tests/models/minimax_h3/test_prepare_h3_weights.py new file mode 100644 index 000000000..5c9a966b9 --- /dev/null +++ b/tests/models/minimax_h3/test_prepare_h3_weights.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Weight extraction rejects off-manifest tensors even under Python optimization.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.validation.models.h3_weights import EXTRACTED_FILE, sha256_file, sha256_tensor + + +@pytest.mark.parametrize("optimize", [0, 1]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + + import huggingface_hub + from safetensors.torch import load_file, save_file + + tensor = torch.tensor([0.0, 1.0]) + config = tmp_path / "config.json" + config.write_text("{}") + index = tmp_path / "index.json" + index.write_text(json.dumps({"weight_map": {"weight": "shard.safetensors"}})) + shard = tmp_path / "shard.safetensors" + save_file({"weight": tensor}, str(shard)) + manifest = { + "model_identity": { + "hf_repo": "test/h3", + "subfolder": "dit", + "revision": "pinned", + "config_sha256": sha256_file(config), + "index_file": index.name, + "index_sha256": sha256_file(index), + }, + "weight_shards": {shard.name: {"sha256": sha256_file(shard)}}, + "tensors": { + "weight": { + "dtype": "bfloat16" if mismatch == "dtype" else "float32", + "shape": [3] if mismatch == "shape" else [2], + "sha256": "0" * 64 if mismatch == "sha256" else sha256_tensor(tensor), + } + }, + } + monkeypatch.setattr( + huggingface_hub, + "hf_hub_download", + lambda repo, filename, **kwargs: str(tmp_path / Path(filename).name), + ) + out_dir = tmp_path / "extracted" + monkeypatch.setattr(sys, "argv", ["prepare_h3_weights.py", "--out", str(out_dir)]) + script = Path(__file__).resolve().parents[3] / "tools" / "weights" / "prepare_h3_weights.py" + namespace = {"__name__": "prepare_h3_weights_test", "__file__": str(script)} + # Compiling the actual script at optimize=1 reproduces python -O's assert removal. + exec(compile(script.read_text(), str(script), "exec", optimize=optimize), namespace) + namespace["load_h3_manifest"] = lambda: manifest + target = out_dir / EXTRACTED_FILE + if mismatch: + with pytest.raises(SystemExit, match=f"weight: {mismatch} .* != manifest"): + namespace["main"]() + assert not target.exists() + else: + namespace["main"]() + assert torch.equal(load_file(str(target))["weight"], tensor) diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 3c47bbbcf..647053ce1 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -4,6 +4,7 @@ import sys from types import ModuleType +import pytest import torch import rl_engine.platforms.device as device_module @@ -231,6 +232,44 @@ def fake_load_backend(backend): ] +H3_OPS = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", +) + + +@pytest.mark.parametrize("op_name", H3_OPS) +def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + """Prefer H3 CUDA backends on CUDA and use their PyTorch fallbacks elsewhere.""" + + registry = KernelRegistry() + cuda_chain = registry._priority_map["cuda"][op_name] + assert cuda_chain[0].name.startswith("CUDA_H3_") + assert cuda_chain[-1].name.startswith("PYTORCH_H3_") + for platform in ("rocm", "musa", "cpu", "npu"): + chain = registry._priority_map[platform][op_name] + assert [backend.name for backend in chain] == [cuda_chain[-1].name] + + +def test_h3_timestep_falls_back_when_fast_math_build_omits_kernel(monkeypatch): + # KERNEL_ALIGN_USE_FAST_MATH=1 builds leave out the bitwise timestep kernel + # and its binding; dispatch must then pick the PyTorch reference. + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_sinusoid + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + monkeypatch.setattr(timestep_sinusoid, "h3_sinusoid_cuda_available", lambda: False) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda:0") + + assert isinstance(op, NativeH3TimestepSinusoidOp) + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod diff --git a/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index bcaebc2f5..e7d1e5275 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/test_ws1_gtest_gpu.py @@ -33,6 +33,8 @@ def _run(script: str, *args: str, timeout: int = 300) -> None: def test_all_ws1_single_ops_are_registered(): + """Require gtest specs for all WS1 single operators and the H3 conditioning stages.""" + names = set(operator_names()) assert { "rms_norm", @@ -49,6 +51,13 @@ def test_all_ws1_single_ops_are_registered(): "pack", "linear_logp", "prefix_shared_attention", + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", } <= names diff --git a/tools/migration/README.md b/tools/migration/README.md index 4a6cee300..48491e383 100644 --- a/tools/migration/README.md +++ b/tools/migration/README.md @@ -8,3 +8,6 @@ Historical workload/operator identifiers and archived evidence intentionally retain their original values. New source imports should use the mapped canonical module. Remove compatibility paths only in a separately announced migration once downstream launchers, patches, imports and saved objects have migrated. + +The MiniMax-H3 entries also record the open RFC #420 PR stack migration to +`test-h3`; those H3 files were not part of the original main baseline. diff --git a/tools/migration/layout.json b/tools/migration/layout.json index cb8fe469a..6c7fc9898 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -347,6 +347,7 @@ "benchmarks/benchmark_deterministic_attention.py": "benchmarks/operators/attention/benchmark_deterministic_attention.py", "benchmarks/benchmark_grpo_loss.py": "benchmarks/operators/loss/benchmark_grpo_loss.py", "benchmarks/benchmark_grpo_op.py": "benchmarks/operators/loss/benchmark_grpo_op.py", + "benchmarks/benchmark_h3_conditioning.py": "benchmarks/models/benchmark_h3_conditioning.py", "benchmarks/benchmark_linear_logp.py": "benchmarks/operators/logprob/benchmark_linear_logp.py", "benchmarks/benchmark_pack.py": "benchmarks/operators/packing/benchmark_pack.py", "benchmarks/benchmark_qwen_ffn_layout.py": "benchmarks/layers/benchmark_qwen_ffn_layout.py", @@ -466,6 +467,42 @@ "docs/design/ws2-attention-single-gpu-harness.md": "docs/archive/design/ws2-attention-single-gpu-harness.md", "docs/design/ws2-cp-attention-contract.md": "docs/contracts/ws2-cp-attention-contract.md", "docs/design/ws2-module-debug-matrix.md": "docs/archive/design/ws2-module-debug-matrix.md", + "docs/usage/evidence/h3-adaln-gate-residual-b200/figure.png": "reports/experiments/h3-adaln-gate-residual-b200/figure.png", + "docs/usage/evidence/h3-adaln-gate-residual-b200/report.json": "reports/experiments/h3-adaln-gate-residual-b200/report.json", + "docs/usage/evidence/h3-adaln-projection-b200/figure.png": "reports/experiments/h3-adaln-projection-b200/figure.png", + "docs/usage/evidence/h3-adaln-projection-b200/report.json": "reports/experiments/h3-adaln-projection-b200/report.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/chain_replay.json": "reports/experiments/h3-adaln-row-gather-b200/chain_replay.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/figure.png": "reports/experiments/h3-adaln-row-gather-b200/figure.png", + "docs/usage/evidence/h3-adaln-row-gather-b200/report.json": "reports/experiments/h3-adaln-row-gather-b200/report.json", + "docs/usage/evidence/h3-final-adaln-out-b200/chain_replay.json": "reports/experiments/h3-final-adaln-out-b200/chain_replay.json", + "docs/usage/evidence/h3-final-adaln-out-b200/figure.png": "reports/experiments/h3-final-adaln-out-b200/figure.png", + "docs/usage/evidence/h3-final-adaln-out-b200/report.json": "reports/experiments/h3-final-adaln-out-b200/report.json", + "docs/usage/evidence/h3-one-block-b200/figure.png": "reports/experiments/h3-one-block-b200/figure.png", + "docs/usage/evidence/h3-one-block-b200/report.json": "reports/experiments/h3-one-block-b200/report.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json": "reports/experiments/h3-prior-art-b200/adaln_row_gather.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.png": "reports/experiments/h3-prior-art-b200/adaln_row_gather.png", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json": "reports/experiments/h3-prior-art-b200/final_adaln_out.json", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.png": "reports/experiments/h3-prior-art-b200/final_adaln_out.png", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.json": "reports/experiments/h3-prior-art-b200/gate_residual.json", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.png": "reports/experiments/h3-prior-art-b200/gate_residual.png", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.json": "reports/experiments/h3-prior-art-b200/norm_modulate.json", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.png": "reports/experiments/h3-prior-art-b200/norm_modulate.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json": "reports/experiments/h3-prior-art-b200/timestep_mlp.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.png": "reports/experiments/h3-prior-art-b200/timestep_mlp.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.json", + "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", + "docs/usage/evidence/h3-timestep-sinusoid-b200/report.json": "reports/experiments/h3-timestep-sinusoid-b200/report.json", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/figure.png": "reports/experiments/h3-tp-adaln-3mod-b200/figure.png", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json": "reports/experiments/h3-tp-adaln-3mod-b200/report.json", "envs.py": "build_tools/envs.py", "examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json": "configs/experiments/cross_config/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json", "examples/cross_config_s0_cpu_smoke.json": "configs/experiments/cross_config/cross_config_s0_cpu_smoke.json", @@ -675,6 +712,19 @@ "rl_engine/kernels/ops/cuda/attention/prefix_shared_attn.py": "rl_engine/backends/cuda/attention/prefix_shared_attn.py", "rl_engine/kernels/ops/cuda/attention/strict_runtime.py": "rl_engine/backends/cuda/attention/strict_runtime.py", "rl_engine/kernels/ops/cuda/ffn.py": "rl_engine/backends/cuda/ffn/ffn.py", + "rl_engine/kernels/ops/cuda/h3/__init__.py": "rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py", + "rl_engine/kernels/ops/cuda/h3/adaln_modulation.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py", + "rl_engine/kernels/ops/cuda/h3/adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/cuda/h3/det_linear.py": "rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py", + "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", + "rl_engine/kernels/ops/cuda/h3/ws2_comm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py", "rl_engine/kernels/ops/cuda/linear/__init__.py": "rl_engine/backends/cuda/embedding/__init__.py", "rl_engine/kernels/ops/cuda/linear/embedding.py": "rl_engine/backends/cuda/embedding/embedding.py", "rl_engine/kernels/ops/cuda/linear/lm_head.py": "rl_engine/backends/cuda/embedding/lm_head.py", @@ -705,6 +755,15 @@ "rl_engine/kernels/ops/pytorch/attention/stateful_kv.py": "rl_engine/reference/attention/stateful_kv.py", "rl_engine/kernels/ops/pytorch/ffn/__init__.py": "rl_engine/ops/ffn/__init__.py", "rl_engine/kernels/ops/pytorch/ffn/ffn.py": "rl_engine/ops/ffn/qwen3.py", + "rl_engine/kernels/ops/pytorch/h3/__init__.py": "rl_engine/reference/minimax_h3/__init__.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_projection.py": "rl_engine/reference/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py": "rl_engine/reference/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py": "rl_engine/reference/minimax_h3/final_adaln_out.py", + "rl_engine/kernels/ops/pytorch/h3/fixed_order.py": "rl_engine/reference/minimax_h3/fixed_order.py", + "rl_engine/kernels/ops/pytorch/h3/gate_residual.py": "rl_engine/reference/minimax_h3/gate_residual.py", + "rl_engine/kernels/ops/pytorch/h3/rmsnorm.py": "rl_engine/reference/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py": "rl_engine/reference/minimax_h3/timestep_mlp.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py": "rl_engine/reference/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/pytorch/linear/__init__.py": "rl_engine/reference/embedding/__init__.py", "rl_engine/kernels/ops/pytorch/linear/embedding.py": "rl_engine/reference/embedding/embedding.py", "rl_engine/kernels/ops/pytorch/linear/lm_head.py": "rl_engine/reference/embedding/lm_head.py", @@ -776,6 +835,13 @@ "rl_engine/testing/__init__.py": "rl_engine/validation/common/__init__.py", "rl_engine/testing/attention_comparison.py": "rl_engine/validation/common/attention_comparison.py", "rl_engine/testing/distributed_logprob_comparison.py": "rl_engine/validation/common/distributed_logprob_comparison.py", + "rl_engine/testing/h3_cases.py": "rl_engine/validation/models/h3_cases.py", + "rl_engine/testing/h3_chain.py": "rl_engine/validation/models/h3_chain.py", + "rl_engine/testing/h3_manifest.json": "rl_engine/validation/models/h3_manifest.json", + "rl_engine/testing/h3_provider.py": "rl_engine/validation/models/h3_provider.py", + "rl_engine/testing/h3_report.py": "rl_engine/validation/models/h3_report.py", + "rl_engine/testing/h3_weights.py": "rl_engine/validation/models/h3_weights.py", + "rl_engine/testing/h3_ws2.py": "rl_engine/validation/models/h3_ws2.py", "rl_engine/testing/logprob_comparison.py": "rl_engine/validation/common/logprob_comparison.py", "rl_engine/testing/logprob_drift.py": "rl_engine/validation/common/logprob_drift.py", "rl_engine/testing/reference_ops.py": "rl_engine/reference/definitions.py", @@ -790,6 +856,13 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", + "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", + "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/h3_ws2_evidence.py": "tools/validation/models/h3_ws2_evidence.py", + "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", + "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", + "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", "scripts/prepare_ws1_weights.py": "tools/weights/prepare_ws1_weights.py", "scripts/run_perf.py": "tools/benchmarking/run_perf.py", "scripts/run_profile_suite.py": "tools/benchmarking/run_profile_suite.py", @@ -820,6 +893,27 @@ "tests/distributed/test_rocm_attention_transport.py": "tests/distributed/cp/test_rocm_attention_transport.py", "tests/distributed/test_rocm_strict_attention_cp.py": "tests/distributed/cp/test_rocm_strict_attention_cp.py", "tests/distributed/test_transport_deterministic_collective.py": "tests/distributed/collectives/test_transport_deterministic_collective.py", + "tests/h3/conftest.py": "tests/models/minimax_h3/conftest.py", + "tests/h3/test_h3_adaln_gate_residual.py": "tests/models/minimax_h3/test_h3_adaln_gate_residual.py", + "tests/h3/test_h3_adaln_modulation.py": "tests/models/minimax_h3/test_h3_adaln_modulation.py", + "tests/h3/test_h3_adaln_projection.py": "tests/models/minimax_h3/test_h3_adaln_projection.py", + "tests/h3/test_h3_adaln_row_gather.py": "tests/models/minimax_h3/test_h3_adaln_row_gather.py", + "tests/h3/test_h3_benchmark_timing.py": "tests/models/minimax_h3/test_h3_benchmark_timing.py", + "tests/h3/test_h3_chain_replay_cli.py": "tests/models/minimax_h3/test_h3_chain_replay_cli.py", + "tests/h3/test_h3_cli.py": "tests/models/minimax_h3/test_h3_cli.py", + "tests/h3/test_h3_conditioning_e2e.py": "tests/models/minimax_h3/test_h3_conditioning_e2e.py", + "tests/h3/test_h3_det_linear.py": "tests/models/minimax_h3/test_h3_det_linear.py", + "tests/h3/test_h3_final_adaln_out.py": "tests/models/minimax_h3/test_h3_final_adaln_out.py", + "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_sp_norm_adaln.py": "tests/models/minimax_h3/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/models/minimax_h3/test_h3_tp_adaln.py", + "tests/h3/test_h3_weights.py": "tests/models/minimax_h3/test_h3_weights.py", + "tests/h3/test_h3_ws2_processes.py": "tests/models/minimax_h3/test_h3_ws2_processes.py", + "tests/h3/test_prepare_h3_weights.py": "tests/models/minimax_h3/test_prepare_h3_weights.py", "tests/linear_logp_tp.py": "tests/distributed/tp/linear_logp_tp.py", "tests/test_alignment_model_wrappers.py": "tests/models/qwen3/test_alignment_model_wrappers.py", "tests/test_alignment_wrapper_interfaces.py": "tests/models/qwen3/test_alignment_wrapper_interfaces.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py new file mode 100644 index 000000000..d43e5c70e --- /dev/null +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct bitwise comparison of one RFC #420 operator with its reuse candidates. + +``h3_prior_art.py`` reports each implementation's batch invariance and its error against FP64. +Matching error summaries do not show that two implementations produce the same bits, so before +a candidate is adopted this runner compares outputs directly. Every candidate that +``h3_prior_art.candidates()`` defines runs on the same inputs (each global batch-invariant mode +in its own process, as in ``h3_prior_art.py``), and every forward output and leaf gradient is +compared byte for byte with diffusers' path and with the rl-kernel op: + + python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \\ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \\ + [--megatron-src /path/to/Megatron-LM] [--fast-math-variant] + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png + +For ``timestep_sinusoid`` the inputs are three timestep sets: 2^20 uniform draws, every 997th +FP32 bit pattern in [0, 1] (a dense sweep, mostly tiny values), and a 100001-point grid. +``--fast-math-variant`` also compiles ``csrc/cuda/h3/timestep_sinusoid.cu`` with +``--use_fast_math`` and compares it, to show what the fast-math build guard prevents. Other +operators use each candidate's own input generator at four sizes. + +Intermediate tensors are written next to the report (not to /tmp, which cluster job epilogs +can clear) and removed afterwards. Run on a clean tree; the report records the commit, the tree +state and every library version. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import h3_prior_art as PA +import torch + +SEED = 7 +SIZES = { + "timestep_mlp": (1, 4, 64, 2048), + "adaln_projection": (1, 4, 64, 2048), + "adaln_row_gather": (1, 257, 4097, 32768), + "norm_modulate": (1, 257, 4097, 32768), + "gate_residual": (1, 257, 4097, 32768), + "final_adaln_out": (1, 257, 4097, 32768), +} +SINUSOID_SETS = ("random_2^20", "dense_fp32_sweep", "grid_100001") +_FAST_MATH_BINDING = """ +#include +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor, int64_t, double, bool); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &h3_timestep_sinusoid_forward); } +""" + + +def _sinusoid_inputs(name: str) -> dict[str, torch.Tensor]: + if name == "random_2^20": + t = torch.rand(1 << 20, device="cuda", generator=torch.Generator("cuda").manual_seed(0)) + elif name == "dense_fp32_sweep": + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + t = bits.view(torch.float32) + else: + t = torch.linspace(0, 1, 100001, device="cuda") + return {"t": t} + + +def _fast_math_sinusoid(build_dir: Path): + from torch.utils.cpp_extension import load + + build_dir.mkdir(parents=True, exist_ok=True) + binding = build_dir / "binding.cpp" + binding.write_text(_FAST_MATH_BINDING) + root = Path(__file__).resolve().parents[3] + mod = load( + name="h3_timestep_sinusoid_fast_math", + sources=[str(root / "csrc/cuda/h3/timestep_sinusoid.cu"), str(binding)], + build_directory=str(build_dir), + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + return { + "source": "rl-kernel timestep_sinusoid.cu compiled with --use_fast_math", + "grads": [], + "fn": lambda t: mod.forward(t, PA.F1, 10000.0, False), + } + + +def _run(cand: dict[str, Any], inputs: dict[str, torch.Tensor]): + args = {} + for k, v in inputs.items(): + args[k] = v.detach().clone().requires_grad_(True) if k in cand["grads"] else v + out = cand["fn"](**args) + outs = list(out) if isinstance(out, (tuple, list)) else [out] + live = [o for o in outs if o.requires_grad] + if live: + gen = torch.Generator("cuda").manual_seed(SEED) + ups = [torch.randn(o.shape, device="cuda", generator=gen).to(o.dtype) for o in live] + torch.autograd.backward(live, ups) + grads = {k: args[k].grad.cpu() if args[k].grad is not None else None for k in cand["grads"]} + return [o.detach().cpu() for o in outs], grads + + +def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: bool) -> None: + """Child process: run every candidate of one mode and save its outputs.""" + + try: + PA.enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + (parts / f"{mode}.unavailable").write_text(f"{type(exc).__name__}: {exc}"[:300]) + return + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cands = PA.candidates(op, mode) + if fast_math and op == "timestep_sinusoid" and mode == "plain": + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + for cand in cands: + name = cand.pop("name") + if "unavailable" in cand: + (parts / f"{name}.unavailable").write_text(cand["unavailable"][:300]) + continue + cases = SINUSOID_SETS if op == "timestep_sinusoid" else SIZES[op] + for case in cases: + inputs = ( + _sinusoid_inputs(case) if op == "timestep_sinusoid" else cand["make"](SEED, case) + ) + outs, grads = _run(cand, inputs) + torch.save( + {"source": cand["source"], "out": outs, "grads": grads}, + parts / f"{name}__{case}.pt", + ) + print(f" [{mode}] {name}: done", flush=True) + torch.cuda.empty_cache() + + +def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | None: + if a is None or b is None: + return None + if a.shape != b.shape or a.dtype != b.dtype: + return {"equal": False, "note": f"{a.dtype}{tuple(a.shape)} vs {b.dtype}{tuple(b.shape)}"} + width = a.element_size() + differ = ( + a.contiguous().view(torch.uint8).view(-1, width) + != b.contiguous().view(torch.uint8).view(-1, width) + ).any(1) + return { + "equal": not bool(differ.any()), + "differ": int(differ.sum()), + "numel": a.numel(), + "max_abs": float((a.double() - b.double()).abs().max()), + } + + +def compare(parts: Path) -> dict[str, Any]: + runs: dict[str, dict[str, Any]] = {} + for f in sorted(parts.glob("*.pt")): + name, case = f.stem.rsplit("__", 1) + runs.setdefault(name, {})[case] = torch.load(f) + refs = [r for r in ("diffusers[plain]", "diffusers", "rl_kernel") if r in runs] + out: dict[str, Any] = {} + for name, by_case in runs.items(): + entry: dict[str, Any] = {"source": next(iter(by_case.values()))["source"]} + for ref in refs: + if ref == name: + continue + entry[f"vs {ref}"] = { + case: { + "out": [_same(a, b) for a, b in zip(run["out"], runs[ref][case]["out"])], + "grads": { + g: _same(run["grads"].get(g), runs[ref][case]["grads"].get(g)) + for g in run["grads"] + }, + } + for case, run in by_case.items() + } + out[name] = entry + for f in parts.glob("*.unavailable"): + out[f.stem] = {"unavailable": f.read_text()} + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(PA.OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument( + "--fast-math-variant", + action="store_true", + help="timestep_sinusoid only: also compare a --use_fast_math build of the kernel", + ) + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(PA.OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {PA.OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + parts = args.out.with_name(f".{args.out.stem}.parts") + + if args.mode is not None: # child process: one global mode + dump(args.op, args.mode, parts, args.megatron_src, args.fast_math_variant) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + shutil.rmtree(parts, ignore_errors=True) + parts.mkdir(parents=True) + modes = list(PA.MODES) if args.op in PA.GEMM_OPS else ["plain"] + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(args.out), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.fast_math_variant: + cmd.append("--fast-math-variant") + proc = subprocess.run(cmd, env={**os.environ, **PA.MODES[mode]}) + if proc.returncode != 0: + (parts / f"{mode}.unavailable").write_text(f"subprocess exited with {proc.returncode}") + results = compare(parts) + shutil.rmtree(parts, ignore_errors=True) + report = { + "kind": "h3_bitwise_reuse", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: PA._version(lib) for lib in PA.LIBS}, + "mode_environment": {m: PA.MODES[m] for m in modes}, + }, + "cases": list(SINUSOID_SETS if args.op == "timestep_sinusoid" else SIZES[args.op]), + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_diagnostic.py b/tools/validation/models/h3_block_diagnostic.py new file mode 100644 index 000000000..19e5b111a --- /dev/null +++ b/tools/validation/models/h3_block_diagnostic.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Localize a failed block-0 gradient gate without changing its numerical contract.""" +import argparse +import json +import os +import subprocess +from dataclasses import replace +from pathlib import Path + +import torch + +from rl_engine.backends.extension import _C +from rl_engine.validation.models import h3_block as block +from rl_engine.validation.models.h3_weights import load_h3_block_weights +from tools.validation.models.h3_full_model import digest + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--failed-job", required=True) + parser.add_argument("--causal-sweep", action="store_true") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=False) + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("diagnostic requires clean source") + ops = block.CandidateOps() + weights = load_h3_block_weights("cuda") + ctx = block.make_block_context(weights, layout="tiny", seed=0, batch=1) + upstream = ( + torch.randn((1, ctx["seq_len"], block.HIDDEN), generator=torch.Generator().manual_seed(11)) + .bfloat16() + .cuda() + ) + report = dict( + status="diagnostic_only_not_acceptance", + failed_job=args.failed_job, + job_id=os.environ.get("SLURM_JOB_ID"), + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + extension_sha256=digest(Path(_C.__file__)), + torch=torch.__version__, + device=torch.cuda.get_device_name(), + backends=ops.backends(), + default_matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision_reduction=( + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + ), + leaves={}, + nodes=[], + ) + # Keep FP32 arithmetic strict through backward as well as forward. + with block.strict_fp32_matmul(): + golden, golden_nodes = block._backward("golden", ctx, None, upstream) + provider, provider_nodes = block._backward("provider", ctx, None, upstream) + candidate, candidate_nodes = block._backward("candidate", ctx, ops, upstream) + for name in golden: + report["leaves"][name] = dict( + candidate=block.rel_err(candidate[name], golden[name]), + provider=block.rel_err(provider[name], golden[name]), + ) + for node in reversed(block.NODES): + report["nodes"].append( + dict( + node=node.name, + errors=[ + dict(candidate=block.rel_err(c, g), provider=block.rel_err(p, g)) + for g, p, c in zip( + golden_nodes[node.name], + provider_nodes[node.name], + candidate_nodes[node.name], + strict=True, + ) + ], + ) + ) + # Same input and upstream gradient isolate the RMSNorm implementation from + # error accumulated through other candidate nodes. + with torch.no_grad(): + values = block.run_graph("provider", ctx) + isolated = {} + for name, source, weight_name in ( + ("q_norm", "q_proj", "attn.norm_q.weight"), + ("k_norm", "k_proj", "attn.norm_k.weight"), + ): + grads = {} + node = next(n for n in block.NODES if n.name == name) + for mode in ("golden", "provider", "candidate"): + x = values[source].detach().clone() + w = weights[block.BLOCK + "." + weight_name].detach().clone() + if mode == "golden": + x, w = x.float(), w.float() + x.requires_grad_() + w.requires_grad_() + local = {**ctx, "weights": {**weights, block.BLOCK + "." + weight_name: w}} + y = ( + node.candidate(ops, local, x) + if mode == "candidate" + else getattr(node, mode)(local, x) + ) + torch.autograd.backward(y, provider_nodes[name][0].to(y.dtype)) + grads[mode] = w.grad + isolated[name] = dict( + candidate=block.rel_err(grads["candidate"], grads["golden"]), + provider=block.rel_err(grads["provider"], grads["golden"]), + ) + report["isolated_norm_weight_gradients"] = isolated + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + if args.causal_sweep: + # These hybrids intentionally use provider nodes for localization only. + # They cannot be submitted to the candidate acceptance comparator. + report["causal_sweep"] = [] + original_nodes = block.NODES + try: + for node in original_nodes: + for baseline in ("candidate", "provider"): + mixed = [] + for item in original_nodes: + use_provider = (item.name == node.name) == (baseline == "candidate") + mixed.append( + replace( + item, + candidate=lambda ops, ctx, *xs, item=item: item.provider(ctx, *xs), + ) + if use_provider + else item + ) + block.NODES = tuple(mixed) + with block.strict_fp32_matmul(): + leaves, _ = block._backward("candidate", ctx, ops, upstream) + report["causal_sweep"].append( + dict( + baseline=baseline, + replaced_node=node.name, + errors={ + name: block.rel_err(leaves[name], golden[name]) for name in golden + }, + ) + ) + del leaves + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + finally: + block.NODES = original_nodes + report["causal_sweep_complete"] = True + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_replay.py b/tools/validation/models/h3_block_replay.py new file mode 100644 index 000000000..7d8230294 --- /dev/null +++ b/tools/validation/models/h3_block_replay.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay one real MiniMax-H3 block forward and backward and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_block.py`` for what is compared. + + python tools/weights/prepare_h3_weights.py --out --block + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_replay.py --out block_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_block_weights, + load_h3_manifest, +) + + +def main() -> None: + """Run the forward (and optionally backward) block replay per layout and save the report.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--layouts", + default="tiny,small,medium", + help=f"comma list of packed FL2VA layouts: {', '.join(H3_BLOCK_LAYOUTS)}", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-backward", action="store_true", help="forward replay only") + parser.add_argument("--timing-repeats", type=int, default=5) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + layouts = args.layouts.split(",") + unknown = sorted(set(layouts) - set(H3_BLOCK_LAYOUTS)) + if unknown: + parser.error(f"unknown layouts: {unknown}") + if not torch.cuda.is_available(): + raise SystemExit("the block replay needs a CUDA device") + + weights = load_h3_block_weights("cuda") + manifest = load_h3_manifest() + ops = h3_block.CandidateOps() + + forward, backward = [], [] + for layout in layouts: + case = h3_block.run_block_case( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + forward.append(case) + print( + f"{layout} S={case['seq_len']}: first_drift={case['first_drift']} " + f"first_isolated_drift={case['first_isolated_drift']} " + f"repeat_drift={case['first_repeat_drift']} batch_drift={case['first_batch_drift']} " + f"out err cand {case['output']['candidate_rel_err_vs_golden']:.2e} " + f"prov {case['output']['provider_rel_err_vs_golden']:.2e}" + ) + if not args.no_backward: + case = h3_block.run_block_backward( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + backward.append(case) + det = all(v["repeat_bitwise_equal"] for v in case["leaves"].values()) + print( + f"{layout} backward: leaves {'det' if det else 'NONDET'}, " + f"grad_repeat_drift={case['first_grad_repeat_drift']} " + f"grad_batch_drift={case['first_grad_batch_drift']} " + f"dX batch-row {case['leaves']['hidden']['batch_row_bitwise_equal']}" + ) + torch.cuda.empty_cache() + + evidence = { + "kind": "h3_one_block_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "nodes": [ + { + "node": n.name, + "inputs": list(n.inputs), + "rfc_row": n.binding.rfc_row, + "status": n.binding.status, + "backend": n.binding.backend, + } + for n in h3_block.NODES + ], + "forward_cases": forward, + "backward_cases": backward, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_ulp_sweep.py b/tools/validation/models/h3_block_ulp_sweep.py new file mode 100644 index 000000000..7fe66646b --- /dev/null +++ b/tools/validation/models/h3_block_ulp_sweep.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sweep the one-block backward over layouts and seeds and record leaf errors in ulps. + +Backs the ulp allowance of ``tests/models/minimax_h3/test_h3_one_block.py``: +for every leaf gradient it records the candidate's and diffusers' error against +the FP32 golden both as ``max|err| / max|golden|`` and in ulps of the leaf's +storage dtype at ``max|golden|``, and the per-leaf excess +``candidate_ulps - provider_ulps``. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_ulp_sweep.py --layouts tiny,small --seeds 0,1,2,3,4 \\ + --out ulp_sweep.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_block_weights # noqa: E402 + + +def main() -> None: + """Run the backward replay for every (layout, seed) and write per-leaf ulp errors.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layouts", default="tiny,small") + parser.add_argument("--seeds", default="0,1,2,3,4") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("the sweep needs a CUDA device") + + weights = load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + cases = [] + for layout in args.layouts.split(","): + for seed in (int(s) for s in args.seeds.split(",")): + report = h3_block.run_block_backward(weights, layout=layout, seed=seed, ops=ops) + leaves = { + name: { + key: entry[key] + for key in ( + "storage_dtype", + "candidate_rel_err_vs_golden", + "provider_rel_err_vs_golden", + "candidate_ulps_vs_golden", + "provider_ulps_vs_golden", + ) + } + | { + "excess_ulps": entry["candidate_ulps_vs_golden"] + - entry["provider_ulps_vs_golden"] + } + for name, entry in report["leaves"].items() + } + worst = max(leaves, key=lambda n: leaves[n]["excess_ulps"]) + print( + f"{layout} S={report['seq_len']} seed={seed}: worst excess " + f"{leaves[worst]['excess_ulps']:+.2f} ulp ({worst})" + ) + cases.append( + {"layout": layout, "seq_len": report["seq_len"], "seed": seed, "leaves": leaves} + ) + torch.cuda.empty_cache() + + bf16 = [ + leaf["excess_ulps"] + for case in cases + for leaf in case["leaves"].values() + if leaf["storage_dtype"] == "bfloat16" + ] + evidence = { + "kind": "h3_one_block_backward_ulp_sweep", + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "cases": cases, + "max_bf16_excess_ulps": max(bf16), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"max BF16 excess: {max(bf16):+.2f} ulp; wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_candidate_prepare.py b/tools/validation/models/h3_candidate_prepare.py new file mode 100644 index 000000000..ae17b6600 --- /dev/null +++ b/tools/validation/models/h3_candidate_prepare.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fail closed on missing native H3 kernels and pin the candidate/reference identities.""" +import argparse +import json +import os +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from rl_engine.backends.extension import _C +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import require_candidate_cuda +from rl_engine.validation.models.h3_weights import load_h3_manifest, sha256_tensor +from tools.validation.models.h3_evidence_integrity import verify_inventory +from tools.validation.models.h3_full_model import digest, verify_checkpoint + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + a = p.parse_args() + a.output.mkdir() + manifest = Path("examples/models/minimax_h3/reference-manifest.json") + verify_inventory(a.reference_root, manifest) + verify_checkpoint( + a.checkpoint, Path("rl_engine/validation/models/h3_full_manifest.json"), all_files=True + ) + if _C is None or not hasattr(_C, "h3_timestep_sinusoid_forward"): + raise ValueError("candidate extension lacks H3 kernels") + ops = CandidateOps() + backends = require_candidate_cuda(ops) + # Inject a real eager CUDA fallback and require the same guard used by the + # full-model path to reject it before it can count as candidate execution. + original = ops.norm + try: + ops.norm = torch.nn.RMSNorm(128).cuda() + try: + require_candidate_cuda(ops) + except ValueError as exc: + if "candidate norm fell back" not in str(exc): + raise + else: + raise AssertionError("eager CUDA norm fallback was accepted") + finally: + ops.norm = original + require_candidate_cuda(ops) + extension = Path(_C.__file__).resolve() + if not extension.is_relative_to(Path.cwd().resolve()): + raise ValueError("candidate loaded an extension outside its clean source snapshot") + spec = load_h3_manifest() + index = json.loads( + (a.checkpoint / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + )["weight_map"] + weights = a.output / "weights" + weights.mkdir() + for section, filename in [ + ("tensors", "h3_conditioning.safetensors"), + ("block_tensors", "h3_block0.safetensors"), + ]: + values = {} + for shard in sorted({index[name] for name in spec[section]}): + with safe_open(str(a.checkpoint / "transformer" / shard), framework="pt") as handle: + for name, row in spec[section].items(): + if index[name] == shard: + value = handle.get_tensor(name).contiguous() + if sha256_tensor(value) != row["sha256"]: + raise ValueError(f"block weight identity changed: {name}") + values[name] = value + save_file(values, str(weights / filename)) + report = dict( + extension=str(extension), + extension_sha256=digest(extension), + torch=torch.__version__, + backends=backends, + negative_controls={"eager_cuda_norm_fallback_rejected": True}, + reference_manifest_sha256=digest(manifest), + build={ + key: os.environ.get(key) + for key in ("CUDA_HOME", "MAX_JOBS", "KERNEL_ALIGN_USE_FAST_MATH") + }, + ) + (a.output / "identity.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_chain_replay.py b/tools/validation/models/h3_chain_replay.py new file mode 100644 index 000000000..11068259c --- /dev/null +++ b/tools/validation/models/h3_chain_replay.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the MiniMax-H3 conditioning chain and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_chain_replay.py --out chain_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models import h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) + + +def _ints(text: str) -> list[int]: + """Parse a comma-separated list of integer case sizes.""" + + return [int(v) for v in text.split(",")] + + +def main() -> None: + """Replay requested CUDA stages with prerequisites and optionally save JSON evidence.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages to report; preceding stages run automatically", + ) + parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") + parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None) + parser.add_argument( + "--backward", + action="store_true", + help="also replay the chain backward (needs every stage): parameter-gradient " + "determinism and accuracy for the separate-op, fused and diffusers chains", + ) + args = parser.parse_args() + + wanted = args.stages.split(",") + stage_names = [s.name for s in h3_chain.STAGES] + unknown = sorted(set(wanted) - set(stage_names)) + if unknown: + parser.error(f"unknown stages: {unknown}") + stages = [s for s in h3_chain.STAGES if s.name in wanted] + if args.backward and len(stages) != len(h3_chain.STAGES): + parser.error("--backward requires every stage of the chain") + last_stage = max(i for i, stage in enumerate(h3_chain.STAGES) if stage.name in wanted) + execution_stages = h3_chain.STAGES[: last_stage + 1] + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + weights = load_h3_conditioning_weights("cuda") + registry = KernelRegistry() + manifest = load_h3_manifest() + + cases = [] + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_case( + registry, + weights, + num_timesteps=num_timesteps, + seq_len=seq_len, + seed=args.seed, + stages=execution_stages, + ) + case["stages"] = [entry for entry in case["stages"] if entry["stage"] in wanted] + cases.append(case) + summary = ", ".join( + f"{e['stage']}=" + f"{'eq' if e['chained_vs_provider']['bitwise_equal'] else 'drift'}" + f"(gold {e['chained_vs_golden']['max_abs']:.3e})" + for e in case["stages"] + ) + print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + + backward_cases = [] + if args.backward: + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_backward_case( + registry, weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed + ) + backward_cases.append(case) + parts = [] + for mode in h3_chain.BACKWARD_MODES: + entries = case["leaves"].values() + det = all(e[mode]["repeat_bitwise_equal"] for e in entries) + worst = max(e[mode]["max_abs_vs_golden_over_absmax"] for e in entries) + parts.append(f"{mode}: {'det' if det else 'NONDET'} {worst:.1e}") + summary = " | ".join(parts) + print(f"backward T={num_timesteps} S={seq_len}: {summary}") + + evidence = { + "kind": "h3_conditioning_chain_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "stages": [s.name for s in stages], + "executed_stages": [s.name for s in execution_stages], + "cases": cases, + "backward_cases": backward_cases, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py new file mode 100644 index 000000000..213d5beb3 --- /dev/null +++ b/tools/validation/models/h3_evidence.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write one RFC #420 operator's evidence report (performance + accuracy) as JSON. + +Run it on a clean tree; the report records the commit and whether the tree +was dirty. ``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) + + +def main() -> None: + """Collect CUDA accuracy and performance evidence, requiring weights for weighted ops.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + uses_weights = args.op != "timestep_sinusoid_h3" + if uses_weights and h3_weights_dir() is None: + raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + registry = KernelRegistry() + manifest = load_h3_manifest() + report = { + "kind": "h3_operator_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weight_source": "pinned_checkpoint" if uses_weights else "not_applicable", + "reference_commit": manifest["reference_implementation"]["commit"], + **git_state(), + "environment": environment(), + "accuracy": ACCURACY[args.op](registry), + "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence_integrity.py b/tools/validation/models/h3_evidence_integrity.py new file mode 100644 index 000000000..42d9b383d --- /dev/null +++ b/tools/validation/models/h3_evidence_integrity.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate immutable reference inventories before consuming any tensor evidence.""" +from __future__ import annotations + +import json +from pathlib import Path + +from tools.validation.models.h3_full_model import digest + + +def verify_inventory(root: Path, manifest: Path): + records = json.loads(manifest.read_text())["files"] + for name, expected in records.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError("evidence path escapes root") + if digest(path) != expected: + raise ValueError(f"reference evidence changed: {name}") + + +def verify_generation(root): + report = json.loads((root / "report.json").read_text()) + if report["status"] != "generation_complete_comparison_pending": + raise ValueError("generation did not complete") + required = {"initial.pt", "video.npy", "audio.npy"} | { + f"step-{i:03}.pt" for i in range(report["steps"]) + } + if not required.issubset(report["artifacts"]): + raise ValueError("incomplete generation artifact inventory") + for name, expected in report["artifacts"].items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()) or digest(path) != expected: + raise ValueError(f"generation artifact changed: {name}") + return report diff --git a/tools/validation/models/h3_first_step_diagnostic.py b/tools/validation/models/h3_first_step_diagnostic.py new file mode 100644 index 000000000..e433b8481 --- /dev/null +++ b/tools/validation/models/h3_first_step_diagnostic.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the first denoising step of a generation under one execution configuration. + +The first transformer call of a generation sees the saved packed inputs (``initial.pt``) +and nothing has fed back yet, so the difference to the provider's saved ``step-000.pt`` +is the model's own error on a real generation input. A configuration is the provider +(optionally with one forced SDPA backend or cuBLASLt), the candidate, or the candidate +with chosen block nodes run as their provider replay (``--provider-nodes``). Replaying +the provider default must reproduce ``step-000.pt`` bitwise; ``baseline_bitwise`` records +whether it did, which validates the replay itself. Diagnostic only: nothing here accepts +or rejects a candidate. + + python -m tools.validation.models.h3_first_step_diagnostic --checkpoint \\ + --generation --seed 43 --mode candidate \\ + --provider-nodes attention --output out.json +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_generation import QueryChunkAttention +from tools.validation.models.h3_full_model import PROVIDER_VARIANTS, provider_variant_context + + +def errors(value: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: + """max_abs and relative_l2 of ``value`` against ``reference`` (float64 norms).""" + + delta = value.double() - reference.double() + return dict( + max_abs=float(delta.abs().max()), + relative_l2=float(delta.norm() / reference.double().norm().clamp_min(1e-12)), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--contract", type=Path, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--mode", choices=["provider", "candidate"], required=True) + parser.add_argument("--provider-variant", choices=PROVIDER_VARIANTS, default="default") + parser.add_argument("--provider-nodes", default="") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.mode == "candidate" and args.provider_variant != "default": + parser.error("--provider-variant applies to --mode provider only") + if args.mode == "provider" and args.provider_nodes: + parser.error("--provider-nodes applies to --mode candidate only") + nodes = tuple(n for n in args.provider_nodes.split(",") if n) + + from diffusers import MiniMaxH3Transformer3DModel + + directory = args.generation / str(args.seed) + inputs = torch.load(directory / "initial.pt", map_location="cpu", weights_only=True) + saved = torch.load(directory / "step-000.pt", map_location="cpu", weights_only=True) + contract = json.loads(args.contract.read_text()) + + started = time.perf_counter() + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), torch_dtype=torch.bfloat16, local_files_only=True + ) + model = model.to("cuda").eval() + loaded = time.perf_counter() - started + kwargs = {k: v.to("cuda") for k, v in inputs.items()} + ops = None + if args.mode == "candidate": + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + + trace = FullModelTrace() + started = time.perf_counter() + with ( + torch.inference_mode(), + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + ops=ops, + provider_nodes=nodes, + ), + provider_variant_context(args.provider_variant), + ): + video, audio = model(**kwargs, attention_kwargs=None, return_dict=False) + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - started + video, audio = video.float().cpu(), audio.float().cpu() + result = dict( + seed=args.seed, + mode=args.mode, + provider_variant=args.provider_variant, + provider_nodes=list(nodes), + video=errors(video, saved[0]), + audio=errors(audio, saved[1]), + baseline_bitwise=bool(torch.equal(video, saved[0]) and torch.equal(audio, saved[1])), + load_seconds=loaded, + forward_seconds=forward_seconds, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_compare.py b/tools/validation/models/h3_full_compare.py new file mode 100644 index 000000000..6f1645e67 --- /dev/null +++ b/tools/validation/models/h3_full_compare.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Freeze full-H3 reference error bounds, or evaluate a candidate without changing them.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from tools.validation.models.h3_full_model import digest + + +def tensors(root): + report = json.loads((root / "report.json").read_text()) + if digest(root / "tensors.pt") != report["tensors_sha256"]: + raise ValueError("full-model tensor evidence changed") + payload = torch.load(root / "tensors.pt", map_location="cpu", weights_only=True) + yield "video", payload["video"] + yield "audio", payload["audio"] + for section in ("nodes", "input_gradients"): + for name, value in sorted(payload[section].items()): + yield section + "." + name, value + for name, row in sorted(json.loads((root / "gradients.json").read_text()).items()): + path = root / row["path"] + if digest(path) != row["sha256"]: + raise ValueError(f"gradient artifact changed: {path}") + yield "parameter." + name, torch.load(path, map_location="cpu", weights_only=True) + + +def compare(left, right): + a, b = tensors(left), tensors(right) + from itertools import zip_longest + + for x, y in zip_longest(a, b): + if x is None or y is None or x[0] != y[0]: + raise ValueError("different full-model tensor inventories") + name, value = x + other = y[1] + if ( + value.shape != other.shape + or not torch.isfinite(value).all() + or not torch.isfinite(other).all() + ): + raise ValueError(f"nonfinite/shape mismatch: {name}") + delta = value.float() - other.float() + yield name, dict( + max_abs=delta.abs().max().item(), + relative_l2=(delta.double().norm() / value.double().norm().clamp_min(1e-12)).item(), + ) + + +def drift_order(report, names): + trace = report["trace"] + if "execution_order" in trace: + observed = trace["execution_order"] + return [name for name in observed if name in names] + [ + name for name in names if name not in observed + ] + forward = trace.get("forward_order", list(trace["calls"])) + backward = trace.get("backward_order", list(trace["backward_calls"])) + ordered = ["nodes." + name for name in forward] + ordered += ["nodes." + name + ".grad" for name in backward] + ordered += [name for name in names if name not in ordered] + return [name for name in ordered if name in names] + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--reference", type=Path, required=True) + p.add_argument("--other", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--thresholds", type=Path) + p.add_argument("--reference-manifest", type=Path) + p.add_argument("--reference-root", type=Path) + args = p.parse_args() + if args.output.exists(): + raise ValueError("never overwrite calibration or failed evidence") + from tools.validation.models.h3_evidence_integrity import verify_inventory + + if args.thresholds and not (args.reference_root and args.reference_manifest): + raise ValueError("candidate comparison requires pinned reference inventory") + if args.reference_manifest: + verify_inventory(args.reference_root, args.reference_manifest) + ref = json.loads((args.reference / "report.json").read_text()) + other = json.loads((args.other / "report.json").read_text()) + if ref["mode"] != "fp32-reference" or ref["status"] != "execution_complete_acceptance_pending": + raise ValueError("independent completed FP32 reference required") + if other["status"] != "execution_complete_acceptance_pending": + raise ValueError("other execution is incomplete") + reference_sha = digest(args.reference / "report.json") + if args.thresholds: + contract = json.loads(args.thresholds.read_text()) + if other["mode"] != "candidate" or contract["reference_report_sha256"] != reference_sha: + raise ValueError("wrong candidate/reference contract") + if other.get("thresholds_sha256") != digest(args.thresholds): + raise ValueError("thresholds were not frozen before candidate execution") + elif other["mode"] != "bf16-provider": + raise ValueError("only independent provider can establish thresholds") + rows = dict(compare(args.reference, args.other)) + report = dict( + reference_report_sha256=reference_sha, + other_report_sha256=digest(args.other / "report.json"), + ) + if args.thresholds: + if rows.keys() != contract["bounds"].keys(): + raise ValueError("threshold tensor inventory differs") + failures = [ + name + for name in drift_order(other if "execution_order" in other["trace"] else ref, rows) + if any(rows[name][k] > contract["bounds"][name][k] for k in rows[name]) + ] + hybrid = other.get("diagnostic_provider_nodes") or [] + report.update( + status="diagnostic_only" if hybrid else ("failed" if failures else "passed"), + diagnostic_provider_nodes=hybrid, + errors=rows, + failures=failures, + first_drift=failures[0] if failures else None, + thresholds_sha256=digest(args.thresholds), + ) + else: + # Committed before candidate execution. Forward and gradient tensors use + # separate floors; the existing block-0 golden contract is unchanged. + bounds = { + name: { + key: 1.5 * value + + ( + 5e-3 + if ".grad" in name + or name.startswith("parameter.") + or name.startswith("input_gradients.") + else 1e-3 + ) + for key, value in row.items() + } + for name, row in rows.items() + } + report.update( + status="frozen_reference_bounds", + rule="1.5 * independent BF16 provider error + fixed floor", + bounds=bounds, + errors=rows, + ) + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.thresholds and failures and not report.get("diagnostic_provider_nodes"): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_ensemble_rule.py b/tools/validation/models/h3_full_ensemble_rule.py new file mode 100644 index 000000000..72c397533 --- /dev/null +++ b/tools/validation/models/h3_full_ensemble_rule.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Full-model acceptance from an ensemble of independent BF16 provider configurations. + +The single-configuration rule (``1.5 * one provider's error + floor``) accepts only the +configuration it was frozen from: providers that differ only in their SDPA backend or +BLAS library fail it on tens of tensors. This rule is calibrated on several provider +configurations, compared with the same frozen FP32 reference, and contains no candidate +data: + +* E[t] is the largest ``relative_l2`` error any configuration has on tensor ``t``; +* the bound is ``k * E[t] + floor`` (floors as in the single-configuration rule); +* ``k`` is the smallest factor, rounded up to 0.1, for which every configuration held + out passes against the maximum of the others (leave-one-out); +* ``max_abs``, set by one element, is reported but not gated: held-out configurations + exceed the others' maximum by several times on it. + + h3_full_ensemble_rule.py calibrate --provider NAME=comparison.json ... \\ + --output rule.json + h3_full_ensemble_rule.py check --rule rule.json --candidate comparison.json \\ + --output out.json +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path + + +def _floor(name: str) -> float: + gradient = ( + ".grad" in name or name.startswith("parameter.") or name.startswith("input_gradients.") + ) + return 5e-3 if gradient else 1e-3 + + +def _digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def calibrate(args) -> None: + """Build the ensemble bounds and the leave-one-out factor from provider comparisons.""" + + providers, sources = {}, {} + for item in args.provider: + name, _, path = item.partition("=") + report = json.loads(Path(path).read_text()) + if report.get("diagnostic_provider_nodes"): + raise ValueError(f"{name} is a hybrid diagnostic, not a provider configuration") + providers[name] = report["errors"] + sources[name] = { + "sha256": _digest(Path(path)), + "reference": report["reference_report_sha256"], + } + if len({s["reference"] for s in sources.values()}) != 1: + raise ValueError("provider comparisons use different FP32 references") + # Configurations whose errors are identical to an earlier one are one configuration. + distinct = {} + for name in sorted(providers): + if not any(providers[name] == providers[kept] for kept in distinct): + distinct[name] = providers[name] + if len(distinct) < 3: + raise ValueError("need at least three distinct provider configurations") + names = sorted(next(iter(distinct.values()))) + loo = {} + for held in distinct: + others = [distinct[o] for o in distinct if o != held] + worst = 0.0 + for t in names: + top = max(o[t]["relative_l2"] for o in others) + if top > 0: + worst = max(worst, (distinct[held][t]["relative_l2"] - _floor(t)) / top) + loo[held] = worst + factor = math.ceil(max(loo.values()) * 10) / 10 + ensemble = {t: max(d[t]["relative_l2"] for d in distinct.values()) for t in names} + rule = { + "status": "frozen_ensemble_rule", + "rule": "relative_l2 <= k * max over provider configurations + floor; max_abs reported", + "factor_k": factor, + "leave_one_out_worst_ratio": loo, + "providers": sorted(distinct), + "duplicates_dropped": sorted(set(providers) - set(distinct)), + "sources": sources, + "reference_report_sha256": next(iter(sources.values()))["reference"], + "ensemble_relative_l2": ensemble, + "ensemble_max_abs": {t: max(d[t]["max_abs"] for d in distinct.values()) for t in names}, + "bounds_relative_l2": {t: factor * ensemble[t] + _floor(t) for t in names}, + } + args.output.write_text(json.dumps(rule, indent=2) + "\n") + print( + f"k = {factor} from leave-one-out {loo}; " + f"{len(distinct)} configurations, {len(names)} tensors" + ) + + +def check(args) -> None: + """Apply a frozen ensemble rule to one candidate comparison; exit 1 on any failure.""" + + rule = json.loads(args.rule.read_text()) + report = json.loads(args.candidate.read_text()) + if report["reference_report_sha256"] != rule["reference_report_sha256"]: + raise ValueError("candidate was compared with a different FP32 reference") + if report.get("diagnostic_provider_nodes"): + raise ValueError("a hybrid diagnostic is not a candidate result") + errors = report["errors"] + if errors.keys() != rule["bounds_relative_l2"].keys(): + raise ValueError("tensor inventory differs from the rule") + failures = [ + t for t in sorted(errors) if errors[t]["relative_l2"] > rule["bounds_relative_l2"][t] + ] + ratios = {t: errors[t]["relative_l2"] / rule["bounds_relative_l2"][t] for t in errors} + max_abs_ratio = { + t: errors[t]["max_abs"] / rule["ensemble_max_abs"][t] + for t in errors + if rule["ensemble_max_abs"][t] + } + out = { + "status": "failed" if failures else "passed", + "rule_sha256": _digest(args.rule), + "candidate_comparison_sha256": _digest(args.candidate), + "failures": failures, + "worst_relative_l2_over_bound": sorted(ratios.items(), key=lambda x: -x[1])[:10], + "outputs": {t: errors[t] for t in ("video", "audio") if t in errors}, + "max_abs_over_ensemble_max_reported": sorted(max_abs_ratio.items(), key=lambda x: -x[1])[ + :10 + ], + } + args.output.write_text(json.dumps(out, indent=2) + "\n") + print( + f"{out['status']}: {len(failures)} failures; worst {out['worst_relative_l2_over_bound'][0]}" + ) + if failures: + raise SystemExit(1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="cmd", required=True) + cal = sub.add_parser("calibrate") + cal.add_argument("--provider", action="append", required=True, help="name=comparison.json") + cal.add_argument("--output", type=Path, required=True) + chk = sub.add_parser("check") + chk.add_argument("--rule", type=Path, required=True) + chk.add_argument("--candidate", type=Path, required=True) + chk.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + calibrate(args) if args.cmd == "calibrate" else check(args) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_generation.py b/tools/validation/models/h3_full_generation.py new file mode 100644 index 000000000..64b215e9f --- /dev/null +++ b/tools/validation/models/h3_full_generation.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Pinned three-seed FL2VA generation evidence for the full H3 candidate stack. + +A candidate seed takes about 20 hours and its partition preempts jobs, so a seed can be +run on its own (``--seeds``) and continued after a preemption (``--resume``): saved steps +are replayed from disk and verified (see ``h3_generation_resume``), never recomputed. +""" +from __future__ import annotations + +import argparse +import datetime +import functools +import json +import os +import subprocess +import time +from pathlib import Path + +import numpy as np +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from rl_engine.validation.models.h3_generation_resume import ResumableRun, StepStore, model_trees +from tools.validation.models.h3_full_model import ( + PROVIDER_VARIANTS, + digest, + provider_variant_context, + verify_checkpoint, +) + + +class QueryChunkAttention: + """Bound attention workspace; every query still attends to every key.""" + + def __init__(self, op, chunk): + self.op, self.chunk = op, chunk + + def __call__(self, q, k, v, **kwargs): + if torch.is_grad_enabled(): + raise ValueError("generation query chunking is inference only") + return torch.cat( + [self.op(part.contiguous(), k, v, **kwargs) for part in q.split(self.chunk, dim=2)], + dim=2, + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-image", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--candidate", action="store_true") + p.add_argument("--preflight", action="store_true") + p.add_argument( + "--seeds", help="comma list of contract seeds to run (default: all; preflight: first)" + ) + p.add_argument( + "--resume", + action="store_true", + help="continue a preempted run in --output: skip finished seeds, replay saved steps", + ) + p.add_argument( + "--verify-last-step", + action="store_true", + help="on resume, recompute the last saved step and require it to match bitwise " + "(costs one step per restart)", + ) + p.add_argument( + "--provider-variant", + choices=PROVIDER_VARIANTS, + default="default", + help="provider only, a variability study: force one SDPA backend or cuBLASLt", + ) + p.add_argument("--preflight-steps", type=int, help="test only: preflight num_inference_steps") + p.add_argument( + "--stop-after-step", type=int, help="test only: raise after N calls (preemption)" + ) + args = p.parse_args() + if args.provider_variant != "default" and args.candidate: + p.error("--provider-variant applies to provider generation only") + if args.preflight_steps is not None and not args.preflight: + p.error("--preflight-steps applies to --preflight only") + args.output.mkdir(parents=True, exist_ok=args.resume) + stale = args.output / failure_name(args.seeds) + if args.resume and stale.exists(): + # A failure of an earlier attempt is kept, not overwritten or hidden. + stale.rename(args.output / f"{stale.stem}-{int(time.time())}.json") + source = Path(__file__).resolve().parents[3] + config_root = source / "examples/models/minimax_h3" + contract = json.loads((config_root / "generation-contract.json").read_text()) + identity_inputs = json.loads((config_root / "generation-inputs.json").read_text()) + prompt = config_root / "generation-prompt.txt" + if ( + digest(prompt) != identity_inputs["prompt_sha256"] + or digest(args.reference_image) != identity_inputs["image_sha256"] + ): + raise ValueError("fixed prompt/reference input identity changed") + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("generation requires clean source") + verify_checkpoint( + args.checkpoint, + source / "rl_engine/validation/models/h3_full_manifest.json", + all_files=True, + ) + import diffusers + + provider_root = Path(diffusers.__file__).resolve().parents[2] + if ( + subprocess.check_output( + ["git", "-C", str(provider_root), "rev-parse", "HEAD"], text=True + ).strip() + != contract["provider_revision"] + ): + raise ValueError("generation provider revision differs") + if subprocess.check_output( + ["git", "-C", str(provider_root), "status", "--porcelain"], text=True + ).strip(): + raise ValueError("generation provider is dirty") + import soundfile as sf + from diffusers import ComponentsManager, ModularPipeline + from diffusers.utils.export_utils import encode_video + from PIL import Image + from transformers import Qwen3VLProcessor + + # Load before the large models: ModularPipeline otherwise catches processor + # dependency failures and leaves None, failing much later in conditioning. + processor = Qwen3VLProcessor.from_pretrained( + str(args.checkpoint / "processor"), local_files_only=True + ) + manager = ComponentsManager() + manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="20GB") + pipe = ModularPipeline.from_pretrained( + str(args.checkpoint), workflow="fl2va", components_manager=manager + ) + pipe.update_components(processor=processor) + pipe.load_components( + dtype=torch.bfloat16, + pretrained_model_name_or_path=str(args.checkpoint), + local_files_only=True, + ) + for name in ( + "processor", + "tokenizer", + "text_encoder", + "transformer", + "vae", + "audio_vae", + "scheduler", + "audio_scheduler", + ): + if getattr(pipe, name, None) is None: + raise RuntimeError(f"Required FL2VA component failed to load: {name}") + ops = None + if args.candidate: + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + settings = contract["preflight"] if args.preflight else contract + if args.preflight_steps is not None: + settings = dict(settings, num_inference_steps=args.preflight_steps) + seeds = contract["seeds"][:1] if args.preflight else contract["seeds"] + if args.seeds: + seeds = [int(s) for s in args.seeds.split(",")] + if not set(seeds) <= set(contract["seeds"]): + raise ValueError(f"seeds {seeds} are not all in the contract {contract['seeds']}") + from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput + + def make_output(values, kwargs): + if kwargs.get("return_dict", True): + return MiniMaxH3TransformerOutput(sample=values[0], audio_sample=values[1]) + return values + + for seed in seeds: + directory = args.output / str(seed) + if args.resume and (directory / "report.json").exists(): + print(f"seed {seed} already complete; skipping", flush=True) + continue + identity = run_identity( + config_root, + contract, + identity_inputs, + candidate=args.candidate, + preflight=args.preflight, + seed=seed, + settings=settings, + provider_variant=args.provider_variant, + ) + run = ResumableRun( + StepStore(directory), + identity, + resume=args.resume, + verify_last_step=args.verify_last_step, + ) + trace = FullModelTrace() + initial = {} + + def capture(module, positional, kwargs, output, run=run, directory=directory): + # Preserve the complete first packed denoiser input to prove identical + # noise and conditioning, then retain every step's predicted velocities. + values = (output.sample, output.audio_sample) if hasattr(output, "sample") else output + if run.saved == 0 and run.recorded == 0: + initial.update( + {k: v.detach().cpu() for k, v in kwargs.items() if isinstance(v, torch.Tensor)} + ) + tmp = directory / "initial.pt.tmp" + torch.save(initial, tmp) + os.replace(tmp, directory / "initial.pt") + run.record(positional, kwargs, tuple(values)) + if run.real_calls == 1: + # One real forward pass has run every block and boundary: keep that trace, so a + # later attempt that only replays and decodes still reports a complete one. + run.store.save_trace(trace.validate()) + if args.stop_after_step is not None and run.recorded >= args.stop_after_step: + raise KeyboardInterrupt(f"simulated preemption after {run.recorded} calls") + + if args.provider_variant != "default": + # Only the transformer runs under the forced kernel: the text encoder and VAE + # (which have their own attention masks) and the conditioning stay the provider's. + plain_forward = pipe.transformer.forward + + @functools.wraps(plain_forward) + def variant_forward(*a, _plain=plain_forward, **kw): + with provider_variant_context(args.provider_variant): + return _plain(*a, **kw) + + pipe.transformer.forward = variant_forward + uninstall = run.install(pipe.transformer, make_output) + handle = pipe.transformer.register_forward_hook(capture, with_kwargs=True) + started = time.perf_counter() + try: + with ( + torch.inference_mode(), + full_model_execution( + pipe.transformer, candidate=args.candidate, trace=trace, ops=ops + ), + ): + result = pipe( + prompt=prompt.read_text(), + image=Image.open(args.reference_image).convert("RGB"), + num_frames=settings["num_frames"], + height=settings["height"], + width=settings["width"], + num_inference_steps=settings["num_inference_steps"], + generator=torch.Generator().manual_seed(seed), + output_type="np", + output=["videos", "audio", "sampling_rate"], + ) + finally: + handle.remove() + uninstall() + if args.provider_variant != "default": + pipe.transformer.forward = plain_forward + video = np.asarray(result["videos"][0]) + audio = result["audio"][0].detach().cpu().numpy() + if video.shape != ( + settings["num_frames"], + settings["height"], + settings["width"], + 3, + ): + raise ValueError(f"wrong video dimensions {video.shape}") + if not np.isfinite(video).all() or not np.isfinite(audio).all(): + raise ValueError("non-finite generated media") + if ( + result["sampling_rate"] != contract["sampling_rate"] + or abs( + audio.shape[-1] / result["sampling_rate"] - settings["num_frames"] / contract["fps"] + ) + > 0.05 + ): + raise ValueError("incorrect audio sample rate/duration") + np.save(directory / "video.npy", video) + np.save(directory / "audio.npy", audio) + sf.write(directory / "audio.wav", audio.T, result["sampling_rate"], subtype="FLOAT") + encode_video( + video, + fps=contract["fps"], + output_path=str(directory / "video.mp4"), + audio=result["audio"][0], + audio_sample_rate=result["sampling_rate"], + ) + for index in (0, len(video) // 2, len(video) - 1): + Image.fromarray(np.round(video[index] * 255).clip(0, 255).astype("uint8")).save( + directory / f"frame-{index:03}.png" + ) + final_trace = trace.validate() if run.real_calls else run.store.load_trace() + if final_trace is None: + raise ValueError("no execution trace: this seed never ran a real forward pass") + report = dict( + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + model_revision=contract["model_revision"], + provider_revision=contract["provider_revision"], + status="generation_complete_comparison_pending", + candidate=args.candidate, + provider_variant=args.provider_variant, + seed=seed, + preflight=args.preflight, + trace=final_trace, + steps=run.recorded, + resume=run.summary(), + seconds=time.perf_counter() - started, + peak_allocated_bytes=torch.cuda.max_memory_allocated(), + contract_sha256=digest(config_root / "generation-contract.json"), + ) + report["artifacts"] = { + path.name: digest(path) for path in sorted(directory.iterdir()) if path.is_file() + } + (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +def run_identity( + config_root, + contract, + inputs, + *, + candidate, + preflight, + seed, + settings, + provider_variant="default", +): + """What a saved run is tied to: the numerics (model-path git trees), inputs and settings.""" + + return dict( + model_trees=model_trees(), + contract_sha256=digest(config_root / "generation-contract.json"), + image_sha256=inputs["image_sha256"], + prompt_sha256=inputs["prompt_sha256"], + candidate=candidate, + preflight=preflight, + seed=seed, + **({"provider_variant": provider_variant} if provider_variant != "default" else {}), + settings={k: settings[k] for k in ("num_frames", "height", "width")} + | {"num_inference_steps": settings["num_inference_steps"]}, + ) + + +def failure_name(seeds): + """Failure marker name: per seed selection, so parallel seed tasks do not collide.""" + + return "failure.json" if not seeds else f"failure-seeds-{seeds.replace(',', '-')}.json" + + +if __name__ == "__main__": + try: + main() + except BaseException as exc: + import sys + + if "--output" in sys.argv: + directory = Path(sys.argv[sys.argv.index("--output") + 1]) + seeds = sys.argv[sys.argv.index("--seeds") + 1] if "--seeds" in sys.argv else None + marker = directory / failure_name(seeds) + if directory.exists() and not marker.exists(): + marker.write_text( + json.dumps( + { + "status": "failed", + "error": f"{type(exc).__name__}: {exc}", + "when": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + ) + ) + raise diff --git a/tools/validation/models/h3_full_model.py b/tools/validation/models/h3_full_model.py new file mode 100644 index 000000000..09a8b07dc --- /dev/null +++ b/tools/validation/models/h3_full_model.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Full pinned H3 forward/backward evidence; reference calibration precedes candidates. + +The small packed input checks the entire training graph, not generation quality. +FP32 reference and BF16 provider runs never instantiate candidate operators. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_cases import h3_fl2va_layout +from rl_engine.validation.models.h3_full_model import ( + MODEL_REVISION, + PROVIDER_REVISION, + FullModelTrace, + full_model_execution, +) + + +def digest(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024**2), b""): + h.update(block) + return h.hexdigest() + + +def verify_checkpoint(root, manifest, *, all_files=False): + spec = json.loads(manifest.read_text()) + if spec["revision"] != MODEL_REVISION: + raise ValueError("wrong H3 revision") + needed = [r for r in spec["files"] if (all_files or r["path"].startswith("transformer/"))] + if not needed: + raise ValueError("no pinned transformer shards") + for record in needed: + path = root / record["path"] + if path.stat().st_size != record["size"] or digest(path) != record["sha256"]: + raise ValueError(f"checkpoint identity mismatch: {path}") + return spec + + +def packed_inputs(device): + layout = h3_fl2va_layout( + num_text_tokens=4, + num_latent_frames=1, + latent_height=4, + latent_width=4, + num_audio_latents=2, + device=device, + ) + generator = torch.Generator().manual_seed(420539) + values = dict( + hidden_states=(1, 4, 96), audio_hidden_states=(1, 4, 32), encoder_hidden_states=(1, 4, 5120) + ) + result = { + name: torch.randn(shape, generator=generator).to(device).requires_grad_() + for name, shape in values.items() + } + tags = layout["token_tags"] + return dict( + **result, + **layout, + timestep=torch.tensor([0.0, 500.0], device=device), + video_indices=(tags == 0).nonzero().flatten(), + text_indices=(tags == 1).nonzero().flatten(), + audio_indices=(tags == 2).nonzero().flatten(), + return_dict=False, + ) + + +# Provider passes that change only which kernels compute the same diffusers graph, to +# measure how much an independent BF16 provider's errors move on their own. Each forces +# one choice; an unavailable kernel raises instead of falling back. +PROVIDER_VARIANTS = ("default", "sdpa-flash", "sdpa-efficient", "sdpa-cudnn", "blas-cublaslt") + + +def provider_variant_context(variant: str): + """Context manager forcing one SDPA backend, or a no-op (BLAS choice is set globally).""" + + from contextlib import nullcontext + + from torch.nn.attention import SDPBackend, sdpa_kernel + + backends = { + "sdpa-flash": SDPBackend.FLASH_ATTENTION, + "sdpa-efficient": SDPBackend.EFFICIENT_ATTENTION, + "sdpa-cudnn": SDPBackend.CUDNN_ATTENTION, + } + if variant == "blas-cublaslt": + torch.backends.cuda.preferred_blas_library("cublaslt") + return sdpa_kernel([backends[variant]]) if variant in backends else nullcontext() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--mode", choices=["fp32-reference", "bf16-provider", "candidate"], required=True + ) + parser.add_argument("--thresholds", type=Path) + parser.add_argument( + "--provider-variant", + choices=PROVIDER_VARIANTS, + default="default", + help="bf16-provider only: force one SDPA backend or cuBLASLt (variability study)", + ) + parser.add_argument( + "--provider-nodes", + default="", + help="candidate only, diagnostic: comma list of block nodes that run the provider " + "replay in every block (the result is labelled diagnostic, never acceptance)", + ) + args = parser.parse_args() + provider_nodes = tuple(n for n in args.provider_nodes.split(",") if n) + if provider_nodes and args.mode != "candidate": + parser.error("--provider-nodes applies to --mode candidate only") + if args.mode == "candidate" and not args.thresholds: + parser.error("candidate execution requires previously frozen calibration thresholds") + if args.provider_variant != "default" and args.mode != "bf16-provider": + parser.error("--provider-variant applies to --mode bf16-provider only") + if args.thresholds: + thresholds = json.loads(args.thresholds.read_text()) + if thresholds.get("status") != "frozen_reference_bounds" or not thresholds.get( + "reference_report_sha256" + ): + raise ValueError("candidate requires immutable independent FP32 calibration") + args.output.mkdir(parents=True, exist_ok=False) + report = dict( + status="running", + mode=args.mode, + provider_variant=args.provider_variant, + diagnostic_provider_nodes=list(provider_nodes), + job_id=os.environ.get("SLURM_JOB_ID"), + model_revision=MODEL_REVISION, + command=list(os.sys.argv), + ) + report_path = args.output / "report.json" + started = time.perf_counter() + try: + commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("evidence requires a clean source commit") + report["commit"] = commit + verify_checkpoint(args.checkpoint, args.manifest) + import diffusers + from diffusers import MiniMaxH3Transformer3DModel + + provider_root = Path(diffusers.__file__).resolve().parents[2] + provider_commit = subprocess.check_output( + ["git", "-C", str(provider_root), "rev-parse", "HEAD"], text=True + ).strip() + if provider_commit != PROVIDER_REVISION: + raise ValueError("diffusers provider revision changed") + if subprocess.check_output( + ["git", "-C", str(provider_root), "status", "--porcelain"], text=True + ).strip(): + raise ValueError("diffusers provider source is dirty") + report.update( + provider_revision=provider_commit, + torch=torch.__version__, + diffusers=diffusers.__version__, + device=torch.cuda.get_device_name(), + checkpoint_manifest_sha256=digest(args.manifest), + ) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.manual_seed(420539) + dtype = torch.float32 if args.mode == "fp32-reference" else torch.bfloat16 + model, loading = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=dtype, + local_files_only=True, + low_cpu_mem_usage=True, + output_loading_info=True, + ) + report["checkpoint_loading"] = loading + if any( + loading.get(key) + for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") + ): + raise ValueError("full H3 checkpoint did not load every tensor exactly") + model = model.to("cuda").eval() + model.enable_gradient_checkpointing() + inputs = packed_inputs("cuda") + trace = FullModelTrace() + trace.capture = True + gradients = {} + handles = [] + # Persist and release each completed parameter gradient. This preserves + # the complete backward chain while bounding resident FP32 gradient memory. + directory = args.output / "gradients" + directory.mkdir() + for i, (name, parameter) in enumerate(model.named_parameters()): + + def save_grad(p, name=name, i=i): + trace.record_event("parameter." + name) + grad = p.grad.detach().cpu() + if not torch.isfinite(grad).all(): + raise ValueError(f"non-finite parameter gradient: {name}") + path = directory / f"{i:04}.pt" + torch.save(grad, path) + gradients[name] = dict( + path=str(path.relative_to(args.output)), + shape=list(grad.shape), + dtype=str(grad.dtype), + sha256=digest(path), + ) + p.grad = None + + handles.append(parameter.register_post_accumulate_grad_hook(save_grad)) + for name in ("hidden_states", "audio_hidden_states", "encoder_hidden_states"): + + def record_input(grad, name=name): + trace.record_event("input_gradients." + name) + return grad + + handles.append(inputs[name].register_hook(record_input)) + torch.cuda.reset_peak_memory_stats() + begin = time.perf_counter() + report["preferred_blas_library"] = None + with ( + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + provider_nodes=provider_nodes, + ), + provider_variant_context(args.provider_variant), + ): + report["preferred_blas_library"] = str(torch.backends.cuda.preferred_blas_library()) + with torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = model(**inputs) + if not torch.isfinite(video).all() or not torch.isfinite(audio).all(): + raise ValueError("non-finite full-model prediction") + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - begin + begin = time.perf_counter() + loss = video.float().square().mean() + audio.float().square().mean() + loss.backward() + torch.cuda.synchronize() + backward_seconds = time.perf_counter() - begin + for h in handles: + h.remove() + if set(gradients) != set(dict(model.named_parameters())): + raise ValueError("missing full-model parameter gradients") + input_grads = {} + for name in ("hidden_states", "audio_hidden_states", "encoder_hidden_states"): + grad = inputs[name].grad + if grad is None or not torch.isfinite(grad).all(): + raise ValueError(f"missing/non-finite input gradient: {name}") + input_grads[name] = grad.cpu() + torch.save( + dict( + video=video.detach().cpu(), + audio=audio.detach().cpu(), + input_gradients=input_grads, + nodes=trace.snapshots, + ), + args.output / "tensors.pt", + ) + (args.output / "gradients.json").write_text(json.dumps(gradients, indent=2) + "\n") + report.update( + trace=trace.validate(backward=True), + forward_seconds=forward_seconds, + backward_seconds=backward_seconds, + peak_allocated_bytes=torch.cuda.max_memory_allocated(), + peak_reserved_bytes=torch.cuda.max_memory_reserved(), + loss=loss.item(), + parameter_gradients=len(gradients), + tensors_sha256=digest(args.output / "tensors.pt"), + ) + # Numerical acceptance is a separate comparison against immutable reference + # artifacts; merely finishing all layers is never a numerical PASS. + report["status"] = "execution_complete_acceptance_pending" + if args.thresholds: + report["thresholds_sha256"] = digest(args.thresholds) + except BaseException as exc: + report.update(status="failed", error=f"{type(exc).__name__}: {exc}") + raise + finally: + report["total_seconds"] = time.perf_counter() - started + report_path.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_provider_spread.py b/tools/validation/models/h3_full_provider_spread.py new file mode 100644 index 000000000..733fc2d2b --- /dev/null +++ b/tools/validation/models/h3_full_provider_spread.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""How far independent BF16 provider passes spread, and where the candidate falls. + +Inputs are ``h3_full_compare`` outputs against the same frozen FP32 reference: one +for the candidate and one per provider variant (``h3_full_model --mode bf16-provider +--provider-variant ...``). For every tensor and metric (``max_abs``, ``relative_l2``) +it reports the providers' minimum and maximum error, the candidate's error, and +whether the candidate lies above the providers' maximum. It also applies the frozen +bounds to each provider variant, which shows whether the bounds hold for providers +themselves. Diagnostic only: it changes no bound and accepts nothing. + + python tools/validation/models/h3_full_provider_spread.py \\ + --thresholds examples/models/minimax_h3/full-model-thresholds.json \\ + --candidate candidate-comparison.json \\ + --provider default=default.json --provider sdpa-flash=flash.json ... \\ + --output spread.json +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +METRICS = ("max_abs", "relative_l2") + + +def main() -> None: + """Load the comparisons, compute per-tensor spreads, and write a summary.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--thresholds", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--provider", action="append", required=True, help="name=comparison.json") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + bounds = json.loads(args.thresholds.read_text())["bounds"] + candidate = json.loads(args.candidate.read_text())["errors"] + providers = {} + for item in args.provider: + name, _, path = item.partition("=") + providers[name] = json.loads(Path(path).read_text())["errors"] + names = sorted(bounds) + for label, errors in [("candidate", candidate), *providers.items()]: + missing = set(names) - set(errors) + if missing: + raise ValueError(f"{label} lacks {len(missing)} tensors, e.g. {sorted(missing)[:3]}") + + tensors = {} + summary = { + "tensors": len(names), + "providers": sorted(providers), + "candidate_failures": {m: 0 for m in METRICS}, + "candidate_above_provider_max": {m: 0 for m in METRICS}, + "provider_failures": {p: {m: 0 for m in METRICS} for p in providers}, + } + for name in names: + entry = {} + for m in METRICS: + values = [providers[p][name][m] for p in providers] + cand = candidate[name][m] + bound = bounds[name][m] + entry[m] = { + "bound": bound, + "candidate": cand, + "provider_min": min(values), + "provider_max": max(values), + "candidate_over_provider_max": cand / max(values) if max(values) else None, + } + summary["candidate_failures"][m] += cand > bound + summary["candidate_above_provider_max"][m] += cand > max(values) + for p in providers: + summary["provider_failures"][p][m] += providers[p][name][m] > bound + tensors[name] = entry + worst = sorted( + names, + key=lambda n: tensors[n]["max_abs"]["candidate_over_provider_max"] or 0.0, + reverse=True, + )[:20] + summary["worst_candidate_over_provider_max"] = [ + {"tensor": n, **tensors[n]["max_abs"]} for n in worst + ] + args.output.write_text(json.dumps({"summary": summary, "tensors": tensors}, indent=2) + "\n") + print(json.dumps(summary, indent=2)[:4000]) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_generation_adopt.py b/tools/validation/models/h3_generation_adopt.py new file mode 100644 index 000000000..b87876878 --- /dev/null +++ b/tools/validation/models/h3_generation_adopt.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Move saved candidate-generation runs onto the model-tree identity (RFC #420, #539). + +Runs started before the identity moved from a commit to the git trees of the model paths +(``h3_generation_resume.MODEL_PATHS``) carry the commit they started at. This re-labels +them for the current checkout, once, and only if every model path is byte-identical +between that commit and ``HEAD``; anything else is refused and nothing is written. + + python tools/validation/models/h3_generation_adopt.py --output \\ + --from-commit 4792c8a9 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from rl_engine.validation.models.h3_generation_resume import adopt_identity # noqa: E402 +from tools.validation.models.h3_full_generation import run_identity # noqa: E402 + + +def main() -> None: + """Adopt every seed directory of ``--output`` (all contract seeds by default).""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--from-commit", required=True) + parser.add_argument("--seeds", help="comma list; default: the contract's seeds") + args = parser.parse_args() + source = Path(__file__).resolve().parents[3] + config_root = source / "examples/models/minimax_h3" + contract = json.loads((config_root / "generation-contract.json").read_text()) + inputs = json.loads((config_root / "generation-inputs.json").read_text()) + seeds = [int(s) for s in args.seeds.split(",")] if args.seeds else contract["seeds"] + # Check every seed before changing any, so a refusal leaves all runs untouched. + plans = [] + for seed in seeds: + identity = run_identity( + config_root, + contract, + inputs, + candidate=True, + preflight=False, + seed=seed, + settings=contract, + ) + plans.append((args.output / str(seed), identity)) + for directory, identity in plans: + adopt_identity(directory, identity, args.from_commit, cwd=source, dry_run=True) + for directory, identity in plans: + adopt_identity(directory, identity, args.from_commit, cwd=source) + print(f"adopted {directory.name} (from {args.from_commit})") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_generation_compare.py b/tools/validation/models/h3_generation_compare.py new file mode 100644 index 000000000..27510767a --- /dev/null +++ b/tools/validation/models/h3_generation_compare.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Compare fixed full-generation runs, preserving failed quality gates.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from skimage.metrics import structural_similarity + +from tools.validation.models.h3_full_model import digest + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--reference", type=Path, required=True) + p.add_argument("--candidate", type=Path, required=True) + p.add_argument("--contract", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--provider-repeat", action="store_true") + p.add_argument("--reference-manifest", type=Path) + p.add_argument("--reference-root", type=Path) + a = p.parse_args() + contract = json.loads(a.contract.read_text()) + from tools.validation.models.h3_evidence_integrity import verify_generation, verify_inventory + + if not a.provider_repeat and not (a.reference_root and a.reference_manifest): + raise ValueError("candidate comparison requires pinned reference inventory") + if a.reference_manifest: + verify_inventory(a.reference_root, a.reference_manifest) + rows = [] + if a.output.exists(): + raise ValueError("never overwrite a prior comparison") + for seed in contract["seeds"]: + left, right = a.reference / str(seed), a.candidate / str(seed) + reports = [verify_generation(root) for root in (left, right)] + if any( + r["preflight"] or r["seed"] != seed or r["contract_sha256"] != digest(a.contract) + for r in reports + ): + raise ValueError("wrong generation evidence/contract") + if reports[0]["candidate"] or reports[1]["candidate"] == a.provider_repeat: + raise ValueError("reference/candidate path mismatch") + initial = [ + torch.load(root / "initial.pt", map_location="cpu", weights_only=True) + for root in (left, right) + ] + if initial[0].keys() != initial[1].keys() or not initial[0]: + raise ValueError("missing packed initial noise evidence") + if any(not torch.equal(value, initial[1][name]) for name, value in initial[0].items()): + raise ValueError("different initial noise/conditioning") + x = np.load(left / "video.npy", mmap_mode="r") + y = np.load(right / "video.npy", mmap_mode="r") + if x.shape != y.shape or x.shape != ( + contract["num_frames"], + contract["height"], + contract["width"], + 3, + ): + raise ValueError("video geometry differs") + mse = 0.0 + ssim = 0.0 + temporal = 0.0 + previous = None + for frame in range(len(x)): + if not np.isfinite(x[frame]).all() or not np.isfinite(y[frame]).all(): + raise ValueError("nonfinite video frame") + delta = x[frame].astype("float64") - y[frame] + mse += float(np.mean(delta**2)) + ssim += float( + structural_similarity(x[frame], y[frame], channel_axis=-1, data_range=1.0) + ) + if previous is not None: + temporal += float(np.mean((delta - previous) ** 2)) + previous = delta + audio = np.load(left / "audio.npy") + other = np.load(right / "audio.npy") + if not np.isfinite(audio).all() or not np.isfinite(other).all(): + raise ValueError("nonfinite audio") + if audio.shape != other.shape: + raise ValueError("audio dimensions differ") + audio_error = float( + np.linalg.norm((audio - other).astype("float64")) + / max(np.linalg.norm(audio.astype("float64")), 1e-12) + ) + row = dict( + seed=seed, + psnr_db=float(-10 * np.log10(max(mse / len(x), 1e-30))), + ssim=ssim / len(x), + temporal_rmse=float(np.sqrt(temporal / (len(x) - 1))), + audio_relative_l2=audio_error, + audio_max_abs=float(np.max(np.abs(audio - other))), + ) + predictions = [] + if reports[0]["steps"] != reports[1]["steps"]: + raise ValueError("different scheduler step counts") + for step in range(reports[0]["steps"]): + v, w = [ + torch.load(root / f"step-{step:03}.pt", map_location="cpu", weights_only=True) + for root in (left, right) + ] + if len(v) != 2 or len(w) != 2: + raise ValueError("missing video/audio prediction") + if any( + b.shape != c.shape or not torch.isfinite(b).all() or not torch.isfinite(c).all() + for b, c in zip(v, w, strict=True) + ): + raise ValueError("nonfinite or malformed step prediction") + predictions.append( + [ + dict( + max_abs=(b.float() - c.float()).abs().max().item(), + relative_l2=( + (b.float() - c.float()).double().norm() + / b.double().norm().clamp_min(1e-12) + ).item(), + ) + for b, c in zip(v, w, strict=True) + ] + ) + bound = contract["quality_bounds"] + row["passed"] = ( + row["psnr_db"] >= bound["minimum_psnr_db"] + and row["ssim"] >= bound["minimum_ssim"] + and row["temporal_rmse"] <= bound["maximum_temporal_rmse"] + and audio_error <= bound["maximum_audio_relative_l2"] + ) + row["step_errors"] = predictions + rows.append(row) + passed = all(r["passed"] for r in rows) + a.output.write_text( + json.dumps( + dict( + status="passed" if passed else "failed", + contract_sha256=digest(a.contract), + seeds=rows, + ), + indent=2, + ) + + "\n" + ) + if not passed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_prior_art.py b/tools/validation/models/h3_prior_art.py new file mode 100644 index 000000000..9012542e2 --- /dev/null +++ b/tools/validation/models/h3_prior_art.py @@ -0,0 +1,946 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compare one RFC #420 operator with existing implementations: batch invariance, accuracy, latency. + +The RFC #420 reuse rule asks, for every operator, whether an existing implementation is +batch-invariant (then reuse it) and for accuracy and performance comparisons. This runner +measures the operator's rl-kernel CUDA op next to the existing implementations that import in +the current environment, and writes one JSON report per operator: + + python tools/validation/models/h3_prior_art.py --op adaln_row_gather \\ + --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json \\ + [--megatron-src /path/to/Megatron-LM] [--quick] + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/adaln_row_gather.json + +Only operators whose rl-kernel op exists on the current branch can be selected. Optional +libraries (SGLang, Liger, Transformer Engine, vLLM, Megatron-LM from a source tree) are skipped +and recorded as unavailable when they do not import. Global batch-invariant modes (vLLM, +SGLang, Megatron) patch aten for the whole process, so each runs in its own subprocess, with +the environment variables it relies on set before CUDA initialises. + +Batch invariance is bitwise, for the forward, the per-row gradients, and repeatability of every +parameter and table gradient, with three checks: + +1. every row computed alone vs inside full batches of 64, 257 and 2048 rows, seeds 3-5; +2. the full workload-size batch vs sub-batches that together cover every row; +3. eight probe rows at the front, middle and back of batches of every size 1..9 and + 2^k - 1, 2^k, 2^k + 1, plus the whole batch reversed. + +Sparse probes (a few rows, a few batch sizes) miss row-specific or large-batch dependence, which +is why every row is checked. Accuracy is max|err| / max|ref| against the same computation in +FP64; latency is the median of 100 CUDA-event samples after 20 warm-ups. Run on a clean tree on +an otherwise idle GPU; the report records the commit, the tree state and every library version. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as md +import importlib.util +import json +import os +import statistics +import subprocess +import sys +from pathlib import Path +from typing import Any, Callable + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +import torch # noqa: E402 +import torch.nn.functional as F # noqa: E402 + +H, D, F1, T3, EPS = 5376, 2688, 256, 3, 1e-5 + +#: operator -> rl-kernel module that must exist on the branch +OPS = { + "timestep_sinusoid": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid", + "timestep_mlp": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp", + "adaln_projection": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection", + "adaln_row_gather": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather", + "norm_modulate": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm", + "gate_residual": "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual", + "final_adaln_out": "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out", +} +#: global batch-invariant modes compared for the GEMM operators, with their environment +MODES = { + "plain": {}, + "vllm": {"CUBLAS_WORKSPACE_CONFIG": ":16:8", "CUBLASLT_WORKSPACE_SIZE": "1"}, + "sglang": {}, + "sglang_ieee": {"TRITON_F32_DEFAULT": "ieee"}, + "megatron_te_native": {"CUBLASLT_WORKSPACE_SIZE": "0"}, + "megatron_triton": {}, + "megatron_triton_ieee": {"TRITON_F32_DEFAULT": "ieee"}, +} +GEMM_OPS = ("timestep_mlp", "adaln_projection") +LIBS = ("torch", "triton", "diffusers", "vllm", "flashinfer-python", "sglang", "liger-kernel") +LIBS += ("transformer_engine", "megatron-core") + +DEV = "cuda" + + +# --------------------------------------------------------------------------- # +# Batch-invariance checks +# --------------------------------------------------------------------------- # + + +class _Case: + """``fn(**inputs)`` with row arguments sliced together; fixed upstream gradients per seed.""" + + def __init__(self, fn, inputs, rows, grads, out_dims, seed): + self.fn, self.inputs, self.rows, self.grads, self.out_dims = ( + fn, + inputs, + rows, + grads, + out_dims, + ) + self.n = next(inputs[k].shape[d] for k, d in rows.items()) + outs = self._flat(fn(**inputs)) + gen = torch.Generator(device=DEV).manual_seed(seed) + self.ups = [torch.randn(o.shape, device=DEV, generator=gen).to(o.dtype) for o in outs] + + @staticmethod + def _flat(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + def run(self, idx=None): + args = {} + for k, v in self.inputs.items(): + if idx is not None and k in self.rows: + v = v.index_select(self.rows[k], idx) + if k in self.grads: + v = v.detach().clone().requires_grad_(True) + args[k] = v + outs = self._flat(self.fn(**args)) + ups = ( + self.ups + if idx is None + else [u.index_select(d, idx) for u, d in zip(self.ups, self.out_dims)] + ) + live = [(o, u) for o, u in zip(outs, ups) if o.requires_grad] + if live: + torch.autograd.backward([o for o, _ in live], [u for _, u in live]) + grads = {k: args[k].grad for k in self.grads} + return [o.detach() for o in outs], grads + + +def _sweep_sizes(n: int) -> list[int]: + sizes = set(range(1, min(n, 9) + 1)) + k = 16 + while k <= n: + sizes.update(v for v in (k - 1, k, k + 1) if v <= n) + k *= 2 + return sorted(sizes | {n}) + + +def batch_invariance(cand: dict[str, Any], quick: bool) -> dict[str, Any]: + fn, make, rows, grads = cand["fn"], cand["make"], cand["rows"], cand["grads"] + out_dims = cand.get("out_dims", (0,)) + params = [k for k in grads if k not in rows] + full_sizes, big, subs = cand["bi_sizes"] + if quick: + full_sizes, big, subs = (17, 64), 256, (1, 7, 64) + seeds = (3, 4, 5) + result: dict[str, Any] = {"all_rows": {}, "params_repeatable": None} + ok = True + + def same_row(a, b, d, i): + return torch.equal(a.select(d, 0), b.select(d, i)) + + for n in full_sizes: + fwd = grd = 0 + repeat = True + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + fo, fg = case.run() + for i in range(n): + o, g = case.run(torch.tensor([i], device=DEV)) + fwd += not all(same_row(a, b, d, i) for a, b, d in zip(o, fo, out_dims)) + grd += not all( + same_row(g[k], fg[k], rows[k], i) + for k in rows + if k in grads and g[k] is not None + ) + if params: + _, again = case.run() + repeat &= all(torch.equal(again[k], fg[k]) for k in params if fg[k] is not None) + result["all_rows"][str(n)] = { + "rows": len(seeds) * n, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + if params: + result["params_repeatable"] = repeat and result["params_repeatable"] is not False + ok &= fwd == 0 and grd == 0 and repeat + + fwd = grd = checked = 0 + for seed in seeds[:2]: + case = _Case(fn, make(seed, big), rows, grads, out_dims, seed) + fo, fg = case.run() + for sb in subs: + step = sb if sb >= 1024 else max(sb, big // 512) + for start in range(0, big, step): + idx = torch.arange(start, min(start + sb, big), device=DEV) + o, g = case.run(idx) + checked += 1 + fwd += not all( + torch.equal(a, b.index_select(d, idx)) for a, b, d in zip(o, fo, out_dims) + ) + grd += not all( + torch.equal(g[k], fg[k].index_select(rows[k], idx)) + for k in rows + if k in grads and g[k] is not None + ) + del case, fo, fg + torch.cuda.empty_cache() + result["full_vs_sub_batches"] = { + "full_batch": big, + "sub_batch_sizes": list(subs), + "sub_batches": checked, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + ok &= fwd == 0 and grd == 0 + + n = full_sizes[-1] + sizes = _sweep_sizes(n) + probes = sorted({0, n - 1, *list(range(n // 8 + 3, n - 1, n // 8))[:6]}) + compared = failed = 0 + first: list[dict[str, int]] = [] + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + alone = {r: case.run(torch.tensor([r], device=DEV)) for r in probes} + gen = torch.Generator().manual_seed(seed) + + def compare(idx, pos, r, seed=seed, case=case, alone=alone): + nonlocal compared, failed + o, g = case.run(idx.to(DEV)) + ao, ag = alone[r] + good = all( + torch.equal(a.select(d, pos), b.select(d, 0)) for a, b, d in zip(o, ao, out_dims) + ) + good &= all( + torch.equal(g[k].select(rows[k], pos), ag[k].select(rows[k], 0)) + for k in rows + if k in grads and g[k] is not None + ) + compared += 1 + if not good: + failed += 1 + if len(first) < 8: + first.append( + {"seed": seed, "batch": int(idx.numel()), "position": pos, "row": r} + ) + + for m in sizes: + for r in probes: + others = torch.randperm(n, generator=gen) + others = others[others != r][: m - 1] + for pos in sorted({0, (m - 1) // 2, m - 1}): + compare(torch.cat([others[:pos], torch.tensor([r]), others[pos:]]), pos, r) + for r in probes: + compare(torch.arange(n - 1, -1, -1), n - 1 - r, r) + result["batch_size_sweep"] = { + "batch_sizes": sizes, + "probe_rows": probes, + "comparisons": compared, + "differ": failed, + "first_failures": first, + } + ok &= failed == 0 + result["batch_invariant"] = bool(ok) + return result + + +# --------------------------------------------------------------------------- # +# Accuracy and latency +# --------------------------------------------------------------------------- # + + +def _time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 100) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def _rel(a: torch.Tensor, ref: torch.Tensor) -> float: + return ((a.double() - ref).abs().max() / ref.abs().max().clamp_min(1e-300)).item() + + +def accuracy_latency(cand: dict[str, Any], mode_off, quick: bool) -> dict[str, Any]: + out: dict[str, Any] = {} + for size in cand["perf_sizes"][:2] if quick else cand["perf_sizes"]: + inputs = cand["make"](11, size) + names = list(inputs) + grads = [k for k in cand["grads"] if k in names] + leaves = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + y = cand["fn"](**leaves) + y = torch.cat([t.reshape(-1) for t in y]) if isinstance(y, (tuple, list)) else y + with mode_off(): + ref_in = { + k: (v.detach().double().requires_grad_(k in grads) if v.is_floating_point() else v) + for k, v in inputs.items() + } + ref = cand["ref"](**ref_in) + ref = torch.cat([t.reshape(-1) for t in ref]) if isinstance(ref, (tuple, list)) else ref + row: dict[str, Any] = {"fwd_rel_err": _rel(y, ref.detach())} + if grads: + gen = torch.Generator(device=DEV).manual_seed(9) + up = torch.randn(y.shape, device=DEV, generator=gen).to(y.dtype) + y.backward(up) + with mode_off(): + ref.backward(up.double()) + row["grad_rel_err"] = {k: _rel(leaves[k].grad, ref_in[k].grad) for k in grads} + row["fwd_us"] = _time_us(lambda inputs=inputs: cand["fn"](**inputs)) + if grads: + + def step(inputs=inputs): + lv = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + o = cand["fn"](**lv) + outs = list(o) if isinstance(o, (tuple, list)) else [o] + torch.autograd.backward(outs, [torch.ones_like(t) for t in outs]) + + row["fwd_bwd_us"] = _time_us(step) + out[str(size)] = row + del inputs, leaves, y, ref + torch.cuda.empty_cache() + return out + + +# --------------------------------------------------------------------------- # +# Candidates +# --------------------------------------------------------------------------- # + + +def _gen(seed: int, salt: int) -> torch.Generator: + return torch.Generator(device=DEV).manual_seed(seed * 1000 + salt) + + +def _rn(seed, salt, *shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, device=DEV, generator=_gen(seed, salt)) * scale).to(dtype) + + +def _layout(n: int, seed: int): + from rl_engine.validation.models.h3_cases import h3_packed_layout + + ti, tags = h3_packed_layout(n, T3, seed=seed) + return ti.contiguous(), tags.contiguous() + + +def _version(name: str) -> str | None: + try: + return md.version(name) + except md.PackageNotFoundError: + return None + + +T_BI = ((64, 257, 2048), 4096, (1, 7, 1024, 2048)) +S_BI = ((64, 257, 2048), 131072, (1, 7, 4097, 32768, 65536)) +T_PERF = (1, 3, 4, 64, 256, 2048) +S_PERF = (4097, 32768) + + +def _sinusoid_ref(t): + half = F1 // 2 + exponent = -torch.log(torch.tensor(10000.0, dtype=torch.float64, device=DEV)) + freqs = torch.exp(exponent * torch.arange(half, device=DEV, dtype=torch.float64) / half) + e = t.double()[:, None] * freqs[None] + return torch.cat([torch.cos(e), torch.sin(e)], -1) + + +def candidates(op: str, mode: str) -> list[dict[str, Any]]: + """Candidate dicts; a factory that raises (missing library) becomes an unavailable entry.""" + + from rl_engine.validation.models import h3_provider as P + + plain = mode == "plain" + found: list[tuple[str, Callable[[], dict[str, Any]]]] = [] + + if op == "timestep_sinusoid": + base = { + "make": lambda s, n: {"t": torch.rand(n, device=DEV, generator=_gen(s, 1))}, + "rows": {"t": 0}, + "grads": [], + "ref": _sinusoid_ref, + "bi_sizes": ((64, 257, 2048), 8192, (1, 7, 2048, 4097)), + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers get_timestep_embedding (op-for-op replay)", + "fn": lambda t: P.provider_time_proj(t), + } + + def sglang(): + from sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit import ( + timestep_embedding, + ) + + return { + **base, + "source": f"SGLang {_version('sglang')} timestep_embedding", + "fn": lambda t: timestep_embedding( + t, F1, flip_sin_to_cos=True, downscale_freq_shift=0.0 + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + ) + + op_ = H3TimestepSinusoidCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepSinusoidCudaOp (check_range=False)", + "fn": lambda t: op_.forward(t, check_range=False), + } + + found = [("diffusers", diffusers), ("sglang", sglang), ("rl_kernel", ours)] + + elif op == "timestep_mlp": + + def make(s, n): + return { + "x": _rn(s, 2, n, F1), + "w1": _rn(1, 3, H, F1, scale=F1**-0.5), + "b1": _rn(1, 4, H, scale=0.02), + "w2": _rn(1, 5, D, H, scale=H**-0.5), + "b2": _rn(1, 6, D, scale=0.02), + } + + base = { + "make": make, + "rows": {"x": 0}, + "grads": ["x", "w1", "b1", "w2", "b2"], + "ref": lambda x, w1, b1, w2, b2: F.linear(F.silu(F.linear(x, w1, b1)), w2, b2), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers TimestepEmbedding, FP32 F.linear [{mode}]", + "fn": lambda x, w1, b1, w2, b2: P.provider_time_embedder(x, w1, b1, w2, b2), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import ( + H3TimestepMLPCudaOp, + ) + + op_ = H3TimestepMLPCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepMLPCudaOp", + "fn": lambda x, w1, b1, w2, b2: op_(x, w1, b1, w2, b2), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_projection": + + def make(s, n): + return { + "temb": _rn(s, 7, n, D), + "w": _rn(1, 8, 18 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 9, 18 * H, dtype=torch.bfloat16, scale=0.02), + } + + base = { + "make": make, + "rows": {"temb": 0}, + "grads": ["temb", "w", "b"], + "ref": lambda temb, w, b: F.linear(F.silu(temb), w, b), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers AdaLN projection, BF16 F.linear [{mode}]", + "fn": lambda temb, w, b: F.linear(F.silu(temb).to(w.dtype), w, b), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + + op_ = H3AdaLNProjectionCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNProjectionCudaOp", + "fn": lambda temb, w, b: op_.forward_table(temb, w, b), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_row_gather": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "table": _rn(1, 14, 3 * T3, 6 * H, dtype=torch.bfloat16, scale=0.1), + "ti": ti, + "tags": tags, + } + + def ref(table, ti, tags): + return tuple(table.index_select(0, ti * 3 + tags).chunk(6, dim=1)) + + base = { + "make": make, + "rows": {"ti": 0, "tags": 0}, + "grads": ["table"], + "out_dims": (0,) * 6, + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return {**base, "source": "diffusers six index_select calls", "fn": ref} + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + op_ = H3AdaLNRowGatherCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNRowGatherCudaOp", + "fn": lambda table, ti, tags: op_(table, ti, tags), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "norm_modulate": + + def make(s, n): + ti, tags = _layout(n, s) + w = (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 11))).bfloat16() + return { + "x": _rn(s, 10, n, H, dtype=torch.bfloat16, scale=2.0), + "w": w, + "sh": _rn(1, 12, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "sc": _rn(1, 13, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(x, w, sh, sc, idx): + n = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + return n * (1 + sc.index_select(0, idx)) + sh.index_select(0, idx) + + base = { + "make": make, + "rows": {"x": 0, "idx": 0}, + "grads": ["x", "w", "sh", "sc"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def per_token(fn): + # Liger / SGLang take per-token shift/scale rows: gather them first (index_select). + return lambda x, w, sh, sc, idx: fn( + x, w, sh.index_select(0, idx), sc.index_select(0, idx) + ) + + def diffusers(): + return { + **base, + "source": "diffusers composition (F.rms_norm + index_select modulation)", + "fn": lambda x, w, sh, sc, idx: P.provider_norm_modulate(x, w, sh, sc, idx), + } + + def liger(): + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + return { + **base, + "source": f"Liger {_version('liger-kernel')} modulated RMSNorm + index_select", + "fn": per_token( + lambda x, w, shr, scr: LigerModulatedRMSNormFunction.apply( + x, w, scr, shr, EPS, 0.0, "llama", False + ) + ), + } + + def liger_rlk_gather(): + # Same Liger kernel, with the rows gathered by rl-kernel's deterministic row gather + # (adaln_row_gather) instead of index_select: separates Liger from the gather. + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + gather = H3AdaLNRowGatherCudaOp() + + def fn(x, w, sh, sc, idx): + pad = torch.zeros_like(sh) + table = torch.cat([sh, sc, pad, pad, pad, pad], dim=1) + rows = gather(table, torch.div(idx, 3, rounding_mode="floor"), idx % 3) + return LigerModulatedRMSNormFunction.apply( + x, w, rows[1], rows[0], EPS, 0.0, "llama", False + ) + + version = _version("liger-kernel") + return { + **base, + "source": f"Liger {version} modulated RMSNorm + rl-kernel row gather", + "fn": fn, + } + + def sglang(): + from sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl import ( + fused_norm_scale_shift, + ) + + return { + **base, + "grads": [], + "source": f"SGLang {_version('sglang')} fused_norm_scale_shift (forward only)", + "fn": per_token( + lambda x, w, shr, scr: fused_norm_scale_shift( + x[None], w, None, scr[None], shr[None], "rms", EPS + )[0] + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + + op_ = H3RMSNormCudaOp() + return { + **base, + "source": "rl-kernel H3RMSNormCudaOp.forward_modulated", + "fn": lambda x, w, sh, sc, idx: op_.forward_modulated(x, w, sh, sc, idx), + } + + def norm_only(make): + def m(s, n): + d = make(s, n) + return {"x": d["x"], "w": d["w"]} + + return m + + def norm_ref(x, w): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + + def torch_norm(): + # diffusers' norm without the modulation, to show which half is batch-invariant + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x", "w"], + "ref": norm_ref, + "source": "torch F.rms_norm (no modulation)", + "fn": lambda x, w: F.rms_norm(x, (H,), w, EPS), + } + + def te_norm(): + import transformer_engine.pytorch as te + + mod = te.RMSNorm(H, eps=EPS, params_dtype=torch.bfloat16, device=DEV) + + def fn(x, w): + # TE's backward reads its own parameter, so copy w in (functional_call + # gives a wrong dx); dweight is then TE's parameter gradient, not checked. + with torch.no_grad(): + mod.weight.copy_(w) + return mod(x) + + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x"], + "ref": norm_ref, + "source": f"TE {_version('transformer_engine')} RMSNorm (no modulation)", + "fn": fn, + } + + found = [ + ("diffusers", diffusers), + ("torch_rms_norm", torch_norm), + ("transformer_engine", te_norm), + ("liger", liger), + ("liger_rlk_gather", liger_rlk_gather), + ("sglang", sglang), + ("rl_kernel", ours), + ] + + elif op == "gate_residual": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "res": _rn(s, 15, n, H, dtype=torch.bfloat16), + "y": _rn(s, 16, n, H, dtype=torch.bfloat16), + "gate": _rn(1, 17, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(res, y, gate, idx): + return res + gate.index_select(0, idx) * y + + base = { + "make": make, + "rows": {"res": 0, "y": 0, "idx": 0}, + "grads": ["res", "y", "gate"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers residual + gate.index_select(...) * y", + "fn": lambda res, y, gate, idx: P.provider_gate_residual(res, gate, idx, y), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + ) + + op_ = H3GateResidualCudaOp() + return { + **base, + "source": "rl-kernel H3GateResidualCudaOp", + "fn": lambda res, y, gate, idx: op_(res, y, gate, idx), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "final_adaln_out": + + def make(s, n): + ti, _ = _layout(n, s) + return { + "x": _rn(s, 18, n, H, dtype=torch.bfloat16, scale=2.0), + "nw": (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 19))).bfloat16(), + "temb": _rn(1, 20, T3, D), + "w": _rn(1, 21, 2 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 22, 2 * H, dtype=torch.bfloat16, scale=0.02), + "ti": ti, + } + + def ref(x, nw, temb, w, b, ti): + return P.provider_final_adaln_out(x, nw, temb, w, b, ti) + + base = { + "make": make, + "rows": {"x": 0, "ti": 0}, + "grads": ["x", "nw", "temb", "w", "b"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay)", + "fn": ref, + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + + op_ = H3FinalAdaLNOutCudaOp() + return { + **base, + "source": "rl-kernel H3FinalAdaLNOutCudaOp", + "fn": lambda x, nw, temb, w, b, ti: op_(x, nw, temb, w, b, ti), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + result = [] + for name, factory in found: + try: + result.append({"name": name, **factory()}) + except Exception as exc: # noqa: BLE001 - optional library missing or unusable + result.append({"name": name, "unavailable": f"{type(exc).__name__}: {exc}"[:300]}) + return result + + +# --------------------------------------------------------------------------- # +# Modes and driver +# --------------------------------------------------------------------------- # + + +def enable_mode(mode: str, megatron_src: str | None): + """Switch on a global batch-invariant mode; returns a context manager that turns it off.""" + + base = mode[:-5] if mode.endswith("_ieee") else mode + if base == "plain": + return contextlib.nullcontext + if base == "vllm": + from vllm.model_executor.determinism.batch_invariant import enable_batch_invariant_mode + + enable_batch_invariant_mode() + return contextlib.nullcontext + if base == "sglang": + from sglang.srt.batch_invariant_ops import batch_invariant_ops as mod + + mod.enable_batch_invariant_mode() + kwargs: dict[str, str] = {} + else: + if megatron_src: + sys.path.insert(0, megatron_src) + from megatron.core.transformer.custom_layers import batch_invariant_kernels as mod + + kwargs = {"backend": base.split("_", 1)[1]} + mod.enable_batch_invariant_mode(**kwargs) + if kwargs["backend"] != "triton": + return contextlib.nullcontext + + @contextlib.contextmanager + def off(): + # The Triton matmuls have no FP64 configuration: FP64 references run with the mode off. + mod.disable_batch_invariant_mode() + try: + yield + finally: + mod.enable_batch_invariant_mode(**kwargs) + + return off + + +def run_mode(op: str, mode: str, megatron_src: str | None, quick: bool) -> dict[str, Any]: + try: + mode_off = enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + return {"unavailable": f"{type(exc).__name__}: {exc}"[:300]} + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + out: dict[str, Any] = {} + for cand in candidates(op, mode): + name = cand.pop("name") + if "unavailable" in cand: + out[name] = cand + print(f" [{mode}] {name}: unavailable ({cand['unavailable'][:80]})", flush=True) + continue + entry: dict[str, Any] = {"source": cand["source"], "has_backward": bool(cand["grads"])} + entry["batch_invariance"] = batch_invariance(cand, quick) + entry["accuracy_latency"] = accuracy_latency(cand, mode_off, quick) + out[name] = entry + print( + f" [{mode}] {name}: batch_invariant={entry['batch_invariance']['batch_invariant']}", + flush=True, + ) + torch.cuda.empty_cache() + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument("--quick", action="store_true", help="small sizes, for a smoke test only") + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + + if args.mode is not None: # child process: one global mode + args.out.write_text(json.dumps(run_mode(args.op, args.mode, args.megatron_src, args.quick))) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + modes = list(MODES) if args.op in GEMM_OPS else ["plain"] + results: dict[str, Any] = {} + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + # The part file sits next to the report, not in /tmp: cluster job epilogs can clear a + # user's /tmp while another of their jobs on the same node is still running. + part = args.out.with_name(f".{args.out.stem}.{mode}.part.json") + part.parent.mkdir(parents=True, exist_ok=True) + part.unlink(missing_ok=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.quick: + cmd.append("--quick") + proc = subprocess.run(cmd, env={**os.environ, **MODES[mode]}) + results[mode] = ( + json.loads(part.read_text()) + if proc.returncode == 0 and part.exists() + else {"unavailable": f"subprocess exited with {proc.returncode}"} + ) + part.unlink(missing_ok=True) + megatron_commit = None + if args.megatron_src: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + report = { + "kind": "h3_prior_art", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: _version(lib) for lib in LIBS}, + "megatron_source_commit": megatron_commit, + "mode_environment": {m: MODES[m] for m in modes}, + }, + "quick": args.quick, + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_ws2_evidence.py b/tools/validation/models/h3_ws2_evidence.py new file mode 100644 index 000000000..8ef5976fb --- /dev/null +++ b/tools/validation/models/h3_ws2_evidence.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + 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 + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +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, + load_h3_manifest, +) +from rl_engine.validation.models.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +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: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py new file mode 100644 index 000000000..85983fc0a --- /dev/null +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot ``h3_bitwise_reuse.py`` reports: share of forward outputs that are not bitwise equal. + +One heatmap per report. Rows are implementations, columns are the cases compared against +diffusers' path and against the rl-kernel op; 0 means bitwise identical. + + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +_REFS = ("diffusers[plain]", "diffusers", "rl_kernel") + + +def _cell(entry: dict | None) -> float | None: + if not entry: + return None + out = entry["out"][0] + if out is None or "differ" not in out: + return None + return 0.0 if out["equal"] else out["differ"] / out["numel"] + + +def _text(frac: float | None) -> str: + if frac is None: + return "—" + if frac == 0: + return "0" + return f"{frac:.1%}" if frac >= 0.001 else f"{frac:.2%}" + + +def _panel(ax, report: dict) -> None: + results = {k: v for k, v in report["results"].items() if "unavailable" not in v} + cases = report["cases"] + diffusers = next((r for r in _REFS[:2] if r in results), None) + refs = [r for r in (diffusers, "rl_kernel") if r] + rows = ["rl_kernel"] + sorted(r for r in results if r not in _REFS) + cols = [(ref, case) for ref in refs for case in cases] + grid = np.full((len(rows), len(cols)), np.nan) + for i, row in enumerate(rows): + for j, (ref, case) in enumerate(cols): + frac = _cell(results.get(row, {}).get(f"vs {ref}", {}).get(str(case))) + if frac is not None: + grid[i, j] = frac + ax.text( + j, + i, + "—" if row == ref else _text(frac), + ha="center", + va="center", + fontsize=7.5, + color="white" if frac is not None and frac > 0.3 else "black", + fontweight="bold" if frac == 0 else None, + ) + ax.imshow(np.sqrt(np.nan_to_num(grid)), cmap="Reds", vmin=0, vmax=1, aspect="auto") + ax.set_yticks(range(len(rows)), rows, fontsize=8) + ax.set_xticks(range(len(cols)), [str(c) for _, c in cols], fontsize=7, rotation=45) + transform = ax.get_xaxis_transform() + for k, ref in enumerate(refs): + if k: + ax.axvline(k * len(cases) - 0.5, color="k", lw=1.5) + label = "vs rl-kernel" if ref == "rl_kernel" else "vs diffusers" + ax.text( + k * len(cases) + (len(cases) - 1) / 2, + 1.01, + label, + ha="center", + va="bottom", + fontsize=8.5, + transform=transform, + ) + ax.set_title(f"{report['op']} @ {report['rl_kernel_commit'][:8]}", fontsize=10, pad=18) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reports", nargs="+", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + reports = [json.loads(p.read_text()) for p in args.reports] + fig, axes = plt.subplots(1, len(reports), figsize=(6.4 * len(reports), 5.8), squeeze=False) + for ax, report in zip(axes[0], reports): + _panel(ax, report) + gpu = reports[0]["environment"].get("gpu", "") + fig.suptitle( + "Forward outputs that are not bitwise equal\n" + f"(share of elements; 0 = bitwise identical), {gpu}", + fontsize=10, + ) + fig.tight_layout() + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=125) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_block.py b/tools/validation/models/plot_h3_block.py new file mode 100644 index 000000000..228f3a289 --- /dev/null +++ b/tools/validation/models/plot_h3_block.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render the RFC #420 one-block replay (``tools/validation/models/h3_block_replay.py``) as a PNG. + +Needs only the report JSON and matplotlib: + + python tools/validation/models/plot_h3_block.py \ + reports/experiments/h3-one-block-b200/report.json + +writes ``figure.png`` next to the report. Panels: the per-node invariance and +provider checks over every layout; the per-node error against the FP32 golden +for the candidate and diffusers; the same for every backward leaf. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import ListedColormap # noqa: E402 + +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +PASS = "#1baf7a" +FAIL = "#d64545" +NA = "#e4e3dd" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" +STATUS_MARK = {"row": "", "interim": "*", "reference": "†"} + + +def _style(ax, title: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=10.5, color=INK, pad=8) + ax.tick_params(colors=MUTED, labelsize=7.5, length=0) + for side in ("top", "right", "left", "bottom"): + ax.spines[side].set_visible(False) + + +def _checks_panel(ax, report: dict[str, Any]) -> None: + """Rows: nodes. Columns: forward/backward checks, ANDed over every layout.""" + + nodes = [n["node"] for n in report["nodes"]] + status = {n["node"]: n["status"] for n in report["nodes"]} + cols = [ + ("repeat", "forward", "repeat_bitwise_equal"), + ("batch row", "forward", "batch_row_bitwise_equal"), + ("= diffusers\n(isolated)", "forward", None), + ("grad\nrepeat", "backward", "grad_repeat_bitwise_equal"), + ("grad\nbatch row", "backward", "grad_batch_row_bitwise_equal"), + ] + grid = [] + for node in nodes: + row = [] + for _label, kind, key in cols: + cases = report["forward_cases" if kind == "forward" else "backward_cases"] + entries = [e for case in cases for e in case["nodes"] if e["node"] == node] + if key is None: + promised = entries[0]["provider_bitwise_isolated_promised"] + ok = all(e["isolated_vs_provider"]["bitwise_equal"] for e in entries) + row.append(2 if not promised and not ok else (1 if ok else 0)) + else: + values = [e[key] for e in entries] + row.append(2 if not values or None in values else (1 if all(values) else 0)) + grid.append(row) + ax.imshow(grid, cmap=ListedColormap([FAIL, PASS, NA]), vmin=0, vmax=2, aspect="auto") + ax.set_xticks(range(len(cols)), [c[0] for c in cols], fontsize=7.5, color=INK) + ax.set_yticks( + range(len(nodes)), [f"{n}{STATUS_MARK[status[n]]}" for n in nodes], fontsize=7.5, color=INK + ) + ax.xaxis.tick_top() + _style(ax, "Bitwise checks, all layouts") + ax.set_xlabel( + "green = holds grey = not promised (reduction differs from diffusers)\n" + "or not applicable (AdaLN tables have no batch axis)\n" + "* interim operator † provider replay (no kernel yet)", + fontsize=7, + color=MUTED, + ) + + +def _err_panel(ax, labels, cand, prov, title: str) -> None: + xs = range(len(labels)) + ax.plot( + xs, [math.log10(v) for v in cand], "o-", color=RL_KERNEL, lw=1.8, ms=5, label="RL-Kernel" + ) + ax.plot( + xs, [math.log10(v) for v in prov], "s--", color=PROVIDER, lw=1.4, ms=4, label="diffusers" + ) + ax.set_xticks(list(xs), labels, rotation=55, ha="right", fontsize=7, color=INK) + lo = math.floor(min(math.log10(v) for v in cand + prov)) + hi = math.ceil(max(math.log10(v) for v in cand + prov)) + ax.set_yticks(range(lo, hi + 1), [f"1e{v}" for v in range(lo, hi + 1)]) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_ylabel("max|err| / max|golden|", color=MUTED, fontsize=8) + ax.legend(frameon=False, fontsize=7.5, labelcolor=INK, loc="upper left") + _style(ax, title) + + +def plot(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(16, 5.6), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.0, 1.6, 1.3]} + ) + case = max(report["forward_cases"], key=lambda c: c["seq_len"]) + back = ( + max(report["backward_cases"], key=lambda c: c["seq_len"]) + if report["backward_cases"] + else None + ) + _checks_panel(axes[0], report) + _err_panel( + axes[1], + [e["node"] for e in case["nodes"]], + [e["candidate_rel_err_vs_golden"] for e in case["nodes"]], + [e["provider_rel_err_vs_golden"] for e in case["nodes"]], + f"Forward error vs FP32 golden, S = {case['seq_len']}", + ) + if back is not None: + leaves = list(back["leaves"]) + _err_panel( + axes[2], + [f"d {name}" for name in leaves], + [back["leaves"][n]["candidate_rel_err_vs_golden"] for n in leaves], + [back["leaves"][n]["provider_rel_err_vs_golden"] for n in leaves], + f"Gradient error vs FP32 golden, S = {back['seq_len']}", + ) + env = report["environment"] + sizes = ", ".join(str(c["seq_len"]) for c in report["forward_cases"]) + fig.suptitle( + "MiniMax-H3 block 0, forward and backward (RFC #420 ws1_one_h3_block)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.905, + f"{env['gpu']} | torch {env['torch']} | layouts S = {sizes} | " + f"first drift vs diffusers: {case['first_drift']} | " + f"rl-kernel {report['rl_kernel_commit'][:7]}", + fontsize=8, + color=MUTED, + ) + fig.tight_layout(rect=(0, 0, 1, 0.89)) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + out = args.report.with_name("figure.png") + plot(report).savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py new file mode 100644 index 000000000..467cb5b7d --- /dev/null +++ b/tools/validation/models/plot_h3_evidence.py @@ -0,0 +1,768 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator or WS2 report (``tools/validation/models/h3_evidence.py``, +``tools/validation/models/h3_ws2_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python tools/validation/models/plot_h3_evidence.py \\ + reports/experiments/h3-timestep-sinusoid-b200/report.json + +writes ``figure.png`` next to the report. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Callable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 +import matplotlib.pyplot as plt # noqa: E402 + +# Fixed series identity across every H3 figure (categorical slots 1-3). +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +THIRD = "#1baf7a" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" + + +def _style(ax, title: str, ylabel: str) -> None: + """Apply the shared report axis styling, title, and vertical-axis label.""" + + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_ylabel(ylabel, color=MUTED, fontsize=9) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + for side in ("top", "right", "left"): + ax.spines[side].set_visible(False) + ax.spines["bottom"].set_color(GRID) + + +def _bars(ax, labels: list[str], series: list[tuple], fmt: str, log: bool = False) -> None: + """Draw grouped series with formatted value labels and a legend; a 4th tuple item is a hatch.""" + + n = len(series) + width = 0.8 / n + for i, (name, color, values, *hatch) in enumerate(series): + xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] + bars = ax.bar( + xs, + values, + width * 0.92, + color=color if not hatch else SURFACE, + edgecolor=color, + hatch=hatch[0] if hatch else None, + linewidth=1.2, + label=name, + zorder=2, + ) + for bar, value in zip(bars, values, strict=True): + ax.annotate( + fmt.format(value), + (bar.get_x() + bar.get_width() / 2, bar.get_height()), + xytext=(0, 2), + textcoords="offset points", + ha="center", + fontsize=7, + color=INK, + ) + ax.set_xticks(range(len(labels)), labels) + top = max(v for _, _, values, *_ in series for v in values) + if log: + ax.set_yscale("log") + ax.set_ylim(min(v for _, _, values, *_ in series for v in values) / 2, top * 8) + else: + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=min(n, 4)) + + +# Exact zeros cannot sit on a log axis; draw them on this floor and label them. +LOG_FLOOR = 1e-9 + + +def _log_values(values: list[float]) -> list[float]: + """Replace nonpositive values with the plotting floor for logarithmic axes.""" + + return [v if v > 0 else LOG_FLOOR for v in values] + + +def _figure(report: dict[str, Any], suptitle: str): + """Create two report axes with a title and GPU, software, and revision provenance.""" + + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) + env = report["environment"] + fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") + fig.text( + 0.01, + 0.01, + f"{env['gpu']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · MiniMax-H3@{report['model_revision'][:7]}", + fontsize=7, + color=MUTED, + ) + return fig, axes + + +def plot_sinusoid(report: dict[str, Any]): + """Return a sinusoid report figure showing latency and errors against FP64.""" + + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") + perf = report["perf"] + _style(left, "Latency per call (lower is better)", "µs") + _bars( + left, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + ("CUDA + range check", THIRD, [row["candidate_checked_us"] for row in perf]), + ("diffusers", PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + cases = report["accuracy"]["cases"] + xs = [c["num_timesteps"] for c in cases] + _style(right, "Max |error| vs FP64 golden (log)", "max abs error") + cuda = [c["max_abs_vs_fp64"] for c in cases] + provider = [c["provider_max_abs_vs_fp64"] for c in cases] + right.plot(xs, _log_values(cuda), "o-", color=RL_KERNEL, lw=2, ms=8, label="RL-Kernel CUDA") + right.plot(xs, _log_values(provider), "s--", color=PROVIDER, lw=2, ms=5, label="diffusers") + atol = report["accuracy"]["contract_atol"] + right.axhline(atol, color=MUTED, lw=1, ls=":") + right.annotate( + "contract atol 1e-5", + (xs[-1], atol), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + for x, value in zip(xs, cuda, strict=True): + if value == 0: + right.annotate( + "exact", + (x, LOG_FLOOR), + xytext=(0, 8), + textcoords="offset points", + fontsize=7, + color=INK, + ha="center", + ) + right.set_xscale("log") + right.set_yscale("log") + right.set_ylim(LOG_FLOOR / 3, atol * 30) + right.set_xlabel("number of timesteps T", color=MUTED, fontsize=9) + equal = sum(c["bitwise_equal_to_provider"] for c in cases) + right.legend( + frameon=False, + fontsize=8, + labelcolor=INK, + loc="upper left", + ncol=2, + title=f"bitwise equal to diffusers at {equal}/{len(cases)} sizes", + title_fontsize=8, + alignment="left", + ) + return fig + + +def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + """Plot candidate and provider latency summaries in microseconds for each case.""" + + _style(ax, "Latency per call (lower is better)", "µs") + _bars( + ax, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + (provider_name, PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + + +def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], contract: float): + """Per-draw max |error| as boxes with every draw overlaid, log scale.""" + + _style(ax, title, "max abs error per draw") + for i, (_name, color, values) in enumerate(series): + ax.boxplot( + [values], + positions=[i], + widths=0.45, + showfliers=False, + patch_artist=True, + boxprops={"facecolor": color, "alpha": 0.25, "edgecolor": color}, + medianprops={"color": color, "linewidth": 2}, + whiskerprops={"color": color}, + capprops={"color": color}, + ) + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax.scatter(jitter, values, s=8, color=color, alpha=0.6, linewidths=0, zorder=3) + median = sorted(values)[len(values) // 2] + ax.annotate( + f"median {median:.1e}", + (i + 0.28, median), + fontsize=8, + color=INK, + va="center", + ) + ax.axhline(contract, color=MUTED, lw=1, ls=":") + ax.annotate( + f"contract atol {contract:g}", + (len(series) - 0.5, contract), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + ax.set_yscale("log") + ax.set_xticks(range(len(series)), [name for name, _, _ in series]) + ax.set_xlim(-0.6, len(series) - 0.1) + + +def plot_mlp(report: dict[str, Any]): + """Return an MLP report figure showing latency and per-draw FP64 error distributions.""" + + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _error_boxes( + right, + f"Error vs FP64 golden, {acc['draws']} draws of T={acc['num_timesteps']}", + [ + ("RL-Kernel CUDA", RL_KERNEL, acc["cuda_max_abs_vs_fp64"]), + ("diffusers (cuBLAS)", PROVIDER, acc["provider_max_abs_vs_fp64"]), + ], + acc["contract_atol"], + ) + return fig + + +def plot_projection(report: dict[str, Any]): + """Return an AdaLN figure with latency, BF16 rounding accuracy, and the early-cast probe.""" + + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _style( + right, + f"Outputs equal to the correctly rounded FP64 golden, {acc['draws']} draws", + "% of BF16 outputs", + ) + series = [ + ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), + ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), + ] + for i, (_name, color, values) in enumerate(series): + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax_vals = sorted(values) + median = ax_vals[len(ax_vals) // 2] + right.scatter(jitter, values, s=36, color=color, alpha=0.8, linewidths=0, zorder=3) + right.hlines(median, i - 0.25, i + 0.25, color=color, lw=2.5, zorder=4) + right.annotate( + f"median {median:.3f}%", (i + 0.28, median), fontsize=8, color=INK, va="center" + ) + right.set_xticks(range(len(series)), [name for name, _, _ in series]) + right.set_xlim(-0.6, len(series) - 0.1) + early = sorted(acc["early_cast_golden_match"])[len(acc["early_cast_golden_match"]) // 2] + right.text( + 0.98, + 0.62, + f"an early BF16 cast (probe H7)\nwould match only {100 * early:.0f}%", + transform=right.transAxes, + fontsize=8, + color=MUTED, + ha="right", + ) + return fig + + +def plot_gather(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_row_gather (T = 3, H = 5376, BF16)") + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("index_select fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "index_select bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + chain = report["accuracy"]["chain_backward"] + _style( + right, + "Whole-chain grads of the FP32 time embedder vs FP64", + "max abs error / golden max", + ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig + names = { + "candidate": ("RL-Kernel, separate ops", THIRD, "o"), + "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), + "provider": ("diffusers", PROVIDER, "s"), + } + labels = [f"T={c['num_timesteps']}\nS={c['seq_len']}" for c in chain] + for j, (mode, (name, color, marker)) in enumerate(names.items()): + leaves = [e for c in chain for e in c["leaves"].values()] + det = sum(e[mode]["repeat_bitwise_equal"] for e in leaves) + adaln = min( + e[mode]["correctly_rounded_fraction"] + for c in chain + for n, e in c["leaves"].items() + if n.startswith("transformer_blocks") + ) + label = ( + f"{name}\n repeat-bitwise {det}/{len(leaves)} · " + f"AdaLN BF16 grads ≥{100 * adaln:.1f}% correctly rounded" + ) + for i, case in enumerate(chain): + values = [ + e[mode]["max_abs_vs_golden_over_absmax"] + for n, e in case["leaves"].items() + if n.startswith("time_embedder") + ] + right.scatter( + [i + (j - 1) * 0.22] * len(values), + values, + s=30, + color=color, + marker=marker, + alpha=0.85, + linewidths=0, + zorder=3, + label=label if i == 0 else None, + ) + right.set_yscale("log") + right.set_ylim(1e-8, 1e4) + right.set_xticks(range(len(labels)), labels, fontsize=7) + right.legend(frameon=False, fontsize=7, labelcolor=INK, loc="upper left") + return fig + + +def plot_rmsnorm(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "h3_rmsnorm (block norm1 + MSA modulation, H = 5376, BF16)" + ) + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("diffusers fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "diffusers bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + acc = report["accuracy"] + bwd = acc["backward"] + keys = ["dx", "dweight", "dshift", "dscale"] + fwd_ok = acc["modulated_bitwise_vs_diffusers"] and all( + acc["plain_bitwise_vs_nn_rmsnorm"].values() + ) + _style(right, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def repeat(mode: str) -> str: + return "yes" if bwd[mode]["repeat_bitwise_equal"] else "no" + + _bars( + right, + keys, + [ + ( + f"RL-Kernel CUDA (repeat-bitwise: {repeat('cuda')})", + RL_KERNEL, + [bwd["cuda"]["rel_error"][k] for k in keys], + ), + ( + f"diffusers (repeat-bitwise: {repeat('provider')})", + PROVIDER, + [bwd["provider"]["rel_error"][k] for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + right.text( + 0.98, + 0.80, + f"forward bitwise equal to diffusers: {'yes' if fwd_ok else 'NO'}", + transform=right.transAxes, + ha="right", + fontsize=8, + color=INK, + ) + return fig + + +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, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ( + 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], + "////", + ), + ], + "{:.2f}", + 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") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + 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) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + +def plot_final(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "final_adaln_out (norm_out: projection + final norm, T = 3)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + _backward_errors( + right, + acc["backward"], + ["dx", "d_norm_w", "d_temb", "dW", "db"], + f"forward equal to diffusers on {100 * acc['equal_to_diffusers_fraction']:.2f}% " + "of elements (projection 1-ULP ties)", + ) + return fig + + +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" + ) + runs = report["runs"] + _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"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + 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)") + _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, + "adaln_projection_3mod": plot_projection, + "adaln_row_gather": plot_gather, + "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, + "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, +} + + +def main() -> None: + """Read an operator report and save its figure as a PNG beside it or at ``--out``.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + fig = PLOTS[report["op"]](report) + fig.tight_layout(rect=(0, 0.04, 1, 0.95)) + out = args.out or args.report.with_name("figure.png") + fig.savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_prior_art.py b/tools/validation/models/plot_h3_prior_art.py new file mode 100644 index 000000000..893612e87 --- /dev/null +++ b/tools/validation/models/plot_h3_prior_art.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot H3 prior-art latency, accuracy, and batch invariance. + + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/.json + +Writes ``.png`` next to the report. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib import transforms # noqa: E402 + + +def _entries(report): + for mode, cands in report["results"].items(): + if "unavailable" in cands: + continue + for name, entry in cands.items(): + if "unavailable" not in entry: + yield (name if name.endswith("]") or mode == "plain" else f"{name}[{mode}]"), entry + + +def _differ(bi) -> int: + rows = sum(v["fwd_differ"] + v["rowgrad_differ"] for v in bi["all_rows"].values()) + sub = bi["full_vs_sub_batches"] + return rows + sub["fwd_differ"] + sub["rowgrad_differ"] + bi["batch_size_sweep"]["differ"] + + +def main() -> None: + path = Path(sys.argv[1]) + report = json.loads(path.read_text()) + entries = list(_entries(report)) + names = [n for n, _ in entries] + fig, axes = plt.subplots(1, 3, figsize=(20, 6.5), layout="constrained") + fig.suptitle( + f"{report['op']} vs existing implementations — {report['environment']['gpu']}, " + f"commit {report['rl_kernel_commit'][:7]}", + fontsize=13, + ) + + ax = axes[0] + for name, entry in entries: + perf = entry["accuracy_latency"] + key = "fwd_bwd_us" if entry["has_backward"] else "fwd_us" + sizes = sorted(perf, key=int) + ax.plot( + [int(s) for s in sizes], + [perf[s][key] for s in sizes], + marker="o", + ls="-" if entry["has_backward"] else "--", + label=name + ("" if entry["has_backward"] else " (forward only)"), + ) + ax.set_xscale("log", base=2) + ax.set_yscale("log") + ax.set_xlabel("rows (timesteps or tokens)") + ax.set_ylabel("µs, median") + ax.set_title("forward + backward latency (dashed: forward only)") + ax.grid(True, which="both", alpha=0.3) + ax.legend(fontsize=7) + + ax = axes[1] + first = entries[0][1]["accuracy_latency"] + size = "3" if "3" in first else sorted(first, key=int)[-1] + ys = range(len(entries)) + fwd = [e["accuracy_latency"][size]["fwd_rel_err"] for _, e in entries] + grad = [ + max(e["accuracy_latency"][size].get("grad_rel_err", {}).values(), default=0) + for _, e in entries + ] + ax.barh([y - 0.2 for y in ys], fwd, 0.4, label="forward") + ax.barh([y + 0.2 for y in ys], [g or float("nan") for g in grad], 0.4, label="worst gradient") + values = [v for v in fwd + grad if v] + if values and max(values) / min(values) > 10: + ax.set_xscale("log") + inside = transforms.blended_transform_factory(ax.transAxes, ax.transData) + for y, a, b in zip(ys, fwd, grad): + ax.text( + 0.98, + y, + f"fwd {a:.1e} / grad {b:.1e}" if b else f"fwd {a:.1e}", + transform=inside, + ha="right", + va="center", + fontsize=7, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.8, "pad": 1}, + ) + ax.set_yticks(list(ys), names, fontsize=7) + ax.invert_yaxis() + ax.set_title(f"error vs FP64, max|err| / max|ref| (size {size})") + ax.legend(fontsize=7, loc="lower right") + ax.grid(True, axis="x", alpha=0.3) + + ax = axes[2] + for y, (_, e) in zip(ys, entries): + bi = e["batch_invariance"] + rep = bi["params_repeatable"] + lines = ["batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant"] + lines.append(f"{_differ(bi)} row comparisons differ") + if rep is not None: + lines.append("parameter/table grads " + ("repeatable" if rep else "NOT repeatable")) + ax.text( + 0.02, + y, + " | ".join(lines), + va="center", + fontsize=8, + color="#3aa676" if bi["batch_invariant"] else "#d14b4b", + ) + ax.set_ylim(len(entries) - 0.5, -0.5) + ax.set_xlim(0, 1) + ax.axis("off") + ax.set_title("batch invariance (bitwise)") + ax.text( + 0.02, + len(entries) - 0.3, + "every row alone vs full batches + full batch vs covering sub-batches + size sweep", + fontsize=7, + color="#555555", + ) + + out = path.with_suffix(".png") + fig.savefig(out, dpi=120) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py new file mode 100644 index 000000000..f6e32bcb5 --- /dev/null +++ b/tools/weights/prepare_h3_weights.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fetch and pin the MiniMax-H3 conditioning tensors used by the RFC #420 tests. + +Downloads only the config, the shard index and the shards that hold the +manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies +every sha256 against ``rl_engine/validation/models/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. ``--block`` also writes block 0's +attention/FFN weights (manifest ``block_tensors``, same shard) to +``/h3_block0.safetensors`` for the ``ws1_one_h3_block`` replay. + + python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 [--block] +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + BLOCK_FILE, + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, + sha256_tensor, +) + + +def _check(path: Path, expected: str) -> None: + """Verify a downloaded file's SHA-256, exiting on mismatch or printing success.""" + + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") + print(f"ok {path.name} {actual}") + + +def main() -> None: + """Download pinned model artifacts, verify tensor identities, and write extracted weights.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="output directory") + parser.add_argument( + "--download-dir", + type=Path, + default=None, + help="where the raw shard goes (default: /hf)", + ) + parser.add_argument( + "--block", action="store_true", help="also extract block 0 (manifest block_tensors)" + ) + args = parser.parse_args() + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + manifest = load_h3_manifest() + identity = manifest["model_identity"] + out_dir: Path = args.out.expanduser() + download_dir = (args.download_dir or out_dir / "hf").expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + + def fetch(filename: str) -> Path: + """Download one file from the manifest's model subfolder and pinned revision.""" + + return Path( + hf_hub_download( + identity["hf_repo"], + f"{identity['subfolder']}/{filename}", + revision=identity["revision"], + local_dir=str(download_dir), + ) + ) + + _check(fetch("config.json"), identity["config_sha256"]) + index_path = fetch(identity["index_file"]) + _check(index_path, identity["index_sha256"]) + weight_map = json.loads(index_path.read_text())["weight_map"] + + def extract(section: str, filename: str) -> None: + """Pull one manifest section's tensors from their pinned shards, verify, and save.""" + + specs = manifest[section] + tensors = {} + for shard in sorted({weight_map[name] for name in specs}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in specs: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + + target = out_dir / filename + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + + extract("tensors", EXTRACTED_FILE) + if args.block: + extract("block_tensors", BLOCK_FILE) + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main() From 3cf51bf9851493646b61c9aed2a26514f2bebc0b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:54:43 +0200 Subject: [PATCH 36/62] docs(h3): do not link to the fingerprint page, it lives in a sibling pull request Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-projection-rows.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/operators/h3-projection-rows.md b/docs/operators/h3-projection-rows.md index 1ed31d15c..241244744 100644 --- a/docs/operators/h3-projection-rows.md +++ b/docs/operators/h3-projection-rows.md @@ -5,7 +5,7 @@ Four RFC #420 rows share one contract: `audio_input_projection_fp32`, `audio_output_projection_fp32`, `text_context_projection` and `h3_qkv_gemm`. Each is a linear projection pinned to checkpoint tensors (`MiniMaxAI/MiniMax-H3@42ed227`, see -[the architecture fingerprint](h3-arch-fingerprint.md)): +the architecture fingerprint (`h3-arch-fingerprint`, #567)): | Row | in -> out | dtype | bias | tensors | | --- | --- | --- | --- | --- | From 38018081392ba4c1222a1bb32e0261a6122cf804 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:05:58 +0000 Subject: [PATCH 37/62] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 4 ++-- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 8847d6161..b35bc41c1 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ From 428baa1139daa781927a800be85ee5e705ab2ce7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:00 +0000 Subject: [PATCH 38/62] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 4 ++-- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- 7 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 8847d6161..b35bc41c1 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -78,7 +78,7 @@ tile setup. ![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a26b41a`. It also records +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index a9c8832e2..3f79732cd 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index bb1194db8..73eac42d6 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -77,14 +77,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `50e913fa` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `a4bc99f9` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` + tree at `a4bc99f9`. Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -121,7 +121,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ From 4530173e0ccd2f02dbc05535f6c21e98734d3613 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:02 +0000 Subject: [PATCH 39/62] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 2 +- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 2 +- 8 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 9a5d815a5..ffd200b0a 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e8fc09537..526e2b4cf 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -177,7 +177,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 7970b5c2d..c4c45ebb5 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -78,14 +78,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -122,7 +122,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 8354dff8b..f1b1a2409 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -72,7 +72,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are byte-equal to WS1 computed on the same GPU. From 4ef511c9e05e26073dc4baff6ec0075e5048be24 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 22:06:03 +0000 Subject: [PATCH 40/62] docs: point commit hashes at the rewritten history The history of this branch was rewritten on 2026-10-23 (evidence directories removed from every commit, sign-offs added). Commit hashes quoted in the reports and documents now name the rewritten commits; the evidence archives still record the pre-rewrite hashes. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 2 +- docs/operators/h3-adaln-projection.md | 4 ++-- docs/operators/h3-adaln-row-gather.md | 8 ++++---- docs/operators/h3-final-adaln-out.md | 6 +++--- docs/operators/h3-rmsnorm.md | 4 ++-- docs/operators/h3-sp-norm-adaln.md | 2 +- docs/operators/h3-timestep-mlp.md | 4 ++-- docs/operators/h3-timestep-sinusoid.md | 2 +- docs/operators/h3-tp-adaln-3mod.md | 2 +- 9 files changed, 17 insertions(+), 17 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index 9a5d815a5..ffd200b0a 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -114,7 +114,7 @@ gradient with diffusers' path and with this kernel. ![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. -[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md index e8fc09537..526e2b4cf 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -127,7 +127,7 @@ dependence on T. ![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `579775f1` on an otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also records that a timestep's 18 modulation rows are bitwise identical whether it runs alone @@ -177,7 +177,7 @@ gradient with diffusers' path and with this kernel. ![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. -[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `50d1e924` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md index 94e32b711..923ad2f77 100644 --- a/docs/operators/h3-adaln-row-gather.md +++ b/docs/operators/h3-adaln-row-gather.md @@ -131,11 +131,11 @@ There are two data files: - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. -`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +`report.json` was regenerated from a clean tree at commit `29b07c5d` with backward-only timings and alternating execution order, and with the op-backward check's rows and gradients drawn from a seeded generator. Before, they came from the process's default CUDA generator, -so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). -`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the model does; every CUDA pass/fail flag is unchanged. @@ -172,7 +172,7 @@ gradient with diffusers' path and with this kernel. ![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. -[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `62f8009e` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index d40a49f5f..a46954f06 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -83,14 +83,14 @@ There are two data files: - [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. - [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the whole conditioning chain (timestep → … → norm_out) replayed stage by stage over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean - tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 before, 0.0671 now). @@ -127,7 +127,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. -[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 75dd3d477..ac625ae35 100644 --- a/docs/operators/h3-rmsnorm.md +++ b/docs/operators/h3-rmsnorm.md @@ -95,7 +95,7 @@ each iteration and is recorded in the report. ![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by -`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: +`tools/validation/models/h3_evidence.py` from a clean tree at commit `d2a6546`. It also records: - bitwise equality with `nn.RMSNorm` for all four pinned norm weights; - bitwise equality with diffusers for the modulation; @@ -139,7 +139,7 @@ gradient with diffusers' path and with this kernel. ![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. -[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `83797833` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md index 8eb91cc38..eeb90f72f 100644 --- a/docs/operators/h3-sp-norm-adaln.md +++ b/docs/operators/h3-sp-norm-adaln.md @@ -72,7 +72,7 @@ How many rows move depends on how the table rows interleave along the sequence: ![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 `6c0d380`, with +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 diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md index cc5809cee..b65b19fca 100644 --- a/docs/operators/h3-timestep-mlp.md +++ b/docs/operators/h3-timestep-mlp.md @@ -112,7 +112,7 @@ kernel time. ![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written -by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `418998c`. It also records that a timestep's row is bitwise identical whether it runs alone or in a batch of 9. ## Existing implementations (RFC #420 reuse rule) @@ -154,7 +154,7 @@ gradient with diffusers' path and with this kernel. ![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. -[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `73939f2d` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md index 75edf172b..ff6579cd8 100644 --- a/docs/operators/h3-timestep-sinusoid.md +++ b/docs/operators/h3-timestep-sinusoid.md @@ -131,7 +131,7 @@ kernel compiled with `--use_fast_math`, which the build excludes (see Known Limi SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all 2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) -was written from a clean tree at `0eb5ddd5` by +was written from a clean tree at `be6a94ae` by ```bash python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index 318dbd2d0..868df445a 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -74,7 +74,7 @@ size, where the multi-block path takes about 45 µs. This matters for the forwar ![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) -The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are byte-equal to WS1 computed on the same GPU. From 193f079ead3e51d8834c2f8cdd7ab222a04a7bd1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:16:26 +0000 Subject: [PATCH 41/62] [WS1][CUDA][MiniMax-H3] ws1_one_h3_block: block 0 node-by-node replay and diagnostics (RFC #420, #539 split 1/3) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .github/workflows/ci.yml | 7 +- .../benchmark_deterministic_attention.py | 175 ++++- build_tools/envs.py | 26 + build_tools/extensions.py | 27 + csrc/bindings/h3_ops.cpp | 35 + .../cuda/attention/deterministic_attention.cu | 341 ++++++++- docs/operators/README.md | 1 + docs/operators/attention.md | 57 ++ docs/operators/h3-one-block.md | 183 +++++ .../models/minimax_h3/diagnose-block.sbatch | 41 + .../cuda/attention/h3_full_attention.py | 150 ++++ .../minimax_h3/adaln_modulation.py | 4 +- .../minimax_h3/adaln_projection.py | 2 +- .../minimax_h3/adaln_row_gather.py | 2 +- .../minimax_h3/final_adaln_out.py | 4 +- .../minimax_h3/gate_residual.py | 4 +- .../cuda/model_specific/minimax_h3/rmsnorm.py | 4 +- .../minimax_h3/sp_norm_adaln.py | 8 +- .../model_specific/minimax_h3/timestep_mlp.py | 2 +- .../minimax_h3/tp_adaln_projection.py | 2 +- .../reference/attention/standard_attn.py | 35 + rl_engine/reference/minimax_h3/block.py | 98 +++ rl_engine/runtime/registry.py | 32 + rl_engine/validation/models/h3_block.py | 718 ++++++++++++++++++ rl_engine/validation/models/h3_cases.py | 124 +++ rl_engine/validation/models/h3_chain.py | 8 +- rl_engine/validation/models/h3_manifest.json | 34 + rl_engine/validation/models/h3_provider.py | 72 ++ rl_engine/validation/models/h3_report.py | 26 +- rl_engine/validation/models/h3_weights.py | 41 +- rl_engine/validation/models/h3_ws2.py | 2 +- .../validation/operators/operator_inputs.py | 10 + .../validation/operators/operator_specs.py | 10 + tests/backends/cuda/test_h3_full_attention.py | 344 +++++++++ tests/models/minimax_h3/conftest.py | 15 + .../minimax_h3/test_h3_adaln_row_gather.py | 6 +- tests/models/minimax_h3/test_h3_det_linear.py | 2 +- .../minimax_h3/test_h3_native_validation.py | 12 +- tests/models/minimax_h3/test_h3_one_block.py | 139 ++++ .../minimax_h3/test_h3_sp_norm_adaln.py | 2 + tools/migration/layout.json | 2 + .../validation/models/h3_block_diagnostic.py | 148 ++++ tools/validation/models/h3_block_replay.py | 114 +++ tools/validation/models/h3_block_ulp_sweep.py | 96 +++ .../models/h3_evidence_integrity.py | 34 + .../models/h3_first_step_diagnostic.py | 117 +++ tools/validation/models/plot_h3_block.py | 172 +++++ tools/weights/prepare_h3_weights.py | 66 +- 48 files changed, 3455 insertions(+), 99 deletions(-) create mode 100644 csrc/bindings/h3_ops.cpp create mode 100644 docs/operators/h3-one-block.md create mode 100644 examples/models/minimax_h3/diagnose-block.sbatch create mode 100644 rl_engine/backends/cuda/attention/h3_full_attention.py create mode 100644 rl_engine/reference/minimax_h3/block.py create mode 100644 rl_engine/validation/models/h3_block.py create mode 100644 tests/backends/cuda/test_h3_full_attention.py create mode 100644 tests/models/minimax_h3/test_h3_one_block.py create mode 100644 tools/validation/models/h3_block_diagnostic.py create mode 100644 tools/validation/models/h3_block_replay.py create mode 100644 tools/validation/models/h3_block_ulp_sweep.py create mode 100644 tools/validation/models/h3_evidence_integrity.py create mode 100644 tools/validation/models/h3_first_step_diagnostic.py create mode 100644 tools/validation/models/plot_h3_block.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4634569cd..58d2dd175 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-h3 ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-h3 ] permissions: contents: read @@ -71,6 +71,9 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q + - name: Run full H3 binding contracts + run: python -m pytest -q tests/models/minimax_h3/test_h3_full_model.py + - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v diff --git a/benchmarks/operators/attention/benchmark_deterministic_attention.py b/benchmarks/operators/attention/benchmark_deterministic_attention.py index b41f98a00..b830c6393 100644 --- a/benchmarks/operators/attention/benchmark_deterministic_attention.py +++ b/benchmarks/operators/attention/benchmark_deterministic_attention.py @@ -5,12 +5,17 @@ Reports latency and peak memory for Qwen3-8B representative shapes, including the cost of full scores/P materialization. + +Use the ``h3`` subcommand for the strict-native/CUDA all-gradient matrix. """ import argparse +import math +import sys import time import torch +from torch.autograd.function import once_differentiable def benchmark_attention( @@ -93,5 +98,173 @@ def main(): torch.cuda.empty_cache() +def h3_value_grad(op, xs, dy): + out = op(*xs) + return [out.detach(), *[g.detach() for g in torch.autograd.grad(out, xs, dy)]] + + +def h3_bits_equal(a, b): + return torch.equal(a.contiguous().view(torch.uint8), b.contiguous().view(torch.uint8)) + + +class H3StrictNative(torch.autograd.Function): + """Unoptimized native baseline, compiled with the same strict arithmetic flags.""" + + @staticmethod + def forward(ctx, q, k, v): + from rl_engine.backends.cuda.attention.h3_full_attention import _extension + + q, k, v = (x.transpose(1, 2).contiguous() for x in (q, k, v)) + out, lse, p = _extension().deterministic_attention_forward( + q, k, v, False, 1.0 / math.sqrt(128), None + ) + ctx.save_for_backward(q, k, v, p) + lse = lse.transpose(1, 2) + ctx.mark_non_differentiable(lse) + return out.transpose(1, 2), lse + + @staticmethod + @once_differentiable + def backward(ctx, grad, unused_lse): + from rl_engine.backends.cuda.attention.h3_full_attention import _extension + + grads = _extension().deterministic_attention_backward( + grad.transpose(1, 2).contiguous(), *ctx.saved_tensors, False, 1.0 / math.sqrt(128), None + ) + return tuple(g.transpose(1, 2) for g in grads) + + +def h3_main(argv=None): + """Paired strict native/CUDA forward and all-QKV-gradient matrix.""" + import hashlib + import json + import statistics + from pathlib import Path + + ROOT = Path(__file__).resolve().parents[3] + sys.path.insert(0, str(ROOT)) + from rl_engine.backends.cuda.attention.h3_full_attention import H3FullAttentionOp, _extension + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path) + parser.add_argument( + "--shapes", + nargs="+", + default=["1x129", "1x512", "1x1025", "1x4096", "2x1024", "4x1024"], + ) + parser.add_argument("--chunks", nargs="+", type=int, default=[64, 128, 256]) + parser.add_argument("--samples", type=int, default=6) + parser.add_argument("--warmups", type=int, default=2) + args = parser.parse_args(argv) + try: + shapes = [tuple(map(int, x.split("x"))) for x in args.shapes] + except ValueError: + parser.error("shapes must be positive BxS pairs") + if any(len(s) != 2 or min(s) <= 0 for s in shapes): + parser.error("shapes must be positive BxS pairs") + if args.samples < 3 or args.warmups < 1 or min(args.chunks) <= 0: + parser.error("use >=3 samples, >=1 warmup and positive chunks") + args.output.mkdir(parents=True, exist_ok=True) + ops = {"strict_native": lambda *xs: H3StrictNative.apply(*xs)[0]} + for chunk in dict.fromkeys(args.chunks): + ops[f"cuda-c{chunk}"] = H3FullAttentionOp(chunk_rows=chunk) + sources = [Path(__file__), ROOT / "csrc/cuda/attention/deterministic_attention.cu"] + sources += [ + ROOT / "rl_engine/backends/cuda/attention/h3_full_attention.py", + ROOT / "rl_engine/reference/attention/standard_attn.py", + ] + report = { + "status": "running", + "gpu": torch.cuda.get_device_name(), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "extension_sha256": hashlib.sha256(Path(_extension().__file__).read_bytes()).hexdigest(), + "source_sha256": { + str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest() for p in sources + }, + "provenance": {k: v.provenance() for k, v in ops.items() if k != "strict_native"}, + "measurement": ( + "CUDA events; eager forward + backward, including layout copies and allocations" + ), + "accuracy": "output/dQ/dK/dV raw-bit equality against strict native on identical inputs", + "warmups": args.warmups, + "samples": args.samples, + "cases": [], + } + + def save(): + temporary = args.output / "report.tmp" + temporary.write_text(json.dumps(report, indent=2) + "\n") + temporary.replace(args.output / "report.json") + + save() + for batch, seq in shapes: + # The native baseline materializes quadratic FP32 intermediates. + free, _ = torch.cuda.mem_get_info() + estimate = batch * 56 * seq * seq * 4 * 3 + batch * seq * 56 * 128 * 2 * 20 + if estimate > free * 0.75: + raise RuntimeError(f"{batch}x{seq}: insufficient free memory for safe paired baseline") + torch.manual_seed(42) + xs = [ + torch.randn(batch, seq, 56, 128, device="cuda", dtype=torch.bfloat16).requires_grad_() + for _ in range(3) + ] + dy = torch.randn_like(xs[0]) + golden = h3_value_grad(ops["strict_native"], xs, dy) + for name, op in ops.items(): + actual = h3_value_grad(op, xs, dy) + assert all(h3_bits_equal(a, g) for a, g in zip(actual, golden)), (batch, seq, name) + del actual + del golden + for op in ops.values(): + for _ in range(args.warmups): + h3_value_grad(op, xs, dy) + torch.cuda.synchronize() + rows = {name: [] for name in ops} + names = list(ops) + for sample in range(args.samples): + order = names[sample % len(names) :] + names[: sample % len(names)] + if sample % 2: + order = list(reversed(order)) + for name in order: + resident = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + result = h3_value_grad(ops[name], xs, dy) + end.record() + end.synchronize() + rows[name].append( + { + "ms": start.elapsed_time(end), + "extra_peak_bytes": torch.cuda.max_memory_allocated() - resident, + } + ) + del result + case = dict( + shape=[batch, seq, 56, 128], + seed=42, + backends={ + name: dict( + median_ms=statistics.median(s["ms"] for s in samples), + peak_extra_bytes=max(s["extra_peak_bytes"] for s in samples), + raw_samples=samples, + ) + for name, samples in rows.items() + }, + ) + report["cases"].append(case) + save() + print(json.dumps(case), flush=True) + del xs, dy + report["status"] = "passed" + save() + print("PAIRED_MATRIX_PASSED", flush=True) + + if __name__ == "__main__": - main() + if sys.argv[1:2] == ["h3"]: + h3_main(sys.argv[2:]) + else: + main() diff --git a/build_tools/envs.py b/build_tools/envs.py index 77dcb8e35..b4c901fd1 100644 --- a/build_tools/envs.py +++ b/build_tools/envs.py @@ -8,6 +8,7 @@ """ import os +import shlex def env_flag(name: str, default: bool = False) -> bool: @@ -27,6 +28,31 @@ def env_flag(name: str, default: bool = False) -> bool: KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC = "KERNEL_ALIGN_ALLOW_UNSUPPORTED_MSVC" KERNEL_ALIGN_FORCE_SM90 = "KERNEL_ALIGN_FORCE_SM90" RL_KERNEL_REQUIRE_EXT = "RL_KERNEL_REQUIRE_EXT" +RL_KERNEL_H3_STRICT = "RL_KERNEL_H3_STRICT" + + +def h3_strict_cuda_flags() -> list[str]: + """Opt-in, cache-visible arithmetic profile for the H3 CUDA entry points. + + External nvcc flags are rejected in this profile: append flags and response + files can silently override these options after build_tools/extensions.py constructs them. + Other operators keep the existing build defaults when this opt-in is absent. + """ + if not env_flag(RL_KERNEL_H3_STRICT): + return [] + if env_flag(KERNEL_ALIGN_USE_FAST_MATH): + raise ValueError("RL_KERNEL_H3_STRICT is incompatible with KERNEL_ALIGN_USE_FAST_MATH") + for name in ("NVCC_PREPEND_FLAGS", "NVCC_APPEND_FLAGS"): + if shlex.split(os.environ.get(name, "")): + raise ValueError(f"Unset {name} for the audited RL_KERNEL_H3_STRICT build") + return [ + "-DRL_KERNEL_H3_STRICT=1", + "--fmad=false", + "--ftz=false", + "--prec-div=true", + "--prec-sqrt=true", + ] + KERNEL_ALIGN_FORCE_ASCEND = "KERNEL_ALIGN_FORCE_ASCEND" KERNEL_ALIGN_ASCEND_ARCH = "KERNEL_ALIGN_ASCEND_ARCH" diff --git a/build_tools/extensions.py b/build_tools/extensions.py index dd135beff..dbddf0c0a 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -119,6 +119,7 @@ def _native_extension_required() -> bool: """Whether the caller explicitly requested a native extension build.""" return ( envs.env_flag(envs.RL_KERNEL_REQUIRE_EXT) + or envs.env_flag(envs.RL_KERNEL_H3_STRICT) or bool(os.environ.get("PYTORCH_ROCM_ARCH", "").strip()) or bool(os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()) or envs.env_flag("FORCE_CUDA") @@ -172,6 +173,7 @@ def _filter_rocm_incompatible_nvcc_flags(flags: list[str]) -> list[str]: def get_extensions(): + h3_flags = envs.h3_strict_cuda_flags() torch, _, GpuExtension = _load_torch_extension_tools() if torch is None: message = ( @@ -196,6 +198,9 @@ def get_extensions(): torch_rpath.append(f"-Wl,-rpath,{torch_lib_dir}") is_rocm = getattr(torch.version, "hip", None) is not None + if h3_flags and (is_rocm or _musa_build_available(torch)): + raise RuntimeError("RL_KERNEL_H3_STRICT currently qualifies NVIDIA CUDA only") + if _musa_build_available(torch): extensions.append( GpuExtension( @@ -398,6 +403,28 @@ def get_extensions(): ) ) + if h3_flags: + # Separate DSO: strict arithmetic never changes other operators. + extensions.append( + GpuExtension( + name="rl_engine._H3_C", + sources=[ + "csrc/bindings/h3_ops.cpp", + "csrc/cuda/attention/deterministic_attention.cu", + ], + extra_compile_args={ + "cxx": ["-O3"], + "nvcc": [ + "-O3", + *h3_flags, + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + ], + }, + extra_link_args=torch_rpath, + ) + ) + extensions.extend(_ascend_extensions()) if _native_extension_required() and not extensions: diff --git a/csrc/bindings/h3_ops.cpp b/csrc/bindings/h3_ops.cpp new file mode 100644 index 000000000..f26859760 --- /dev/null +++ b/csrc/bindings/h3_ops.cpp @@ -0,0 +1,35 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +#include + +#include +#include +#include + +std::map h3_attention_build_info(); +std::vector h3_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v); +torch::Tensor h3_attention_probabilities( + torch::Tensor q, torch::Tensor k, torch::Tensor v); +std::vector h3_attention_backward_accumulate( + torch::Tensor grad_output, torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor probabilities, + torch::Tensor dk, torch::Tensor dv); + +// Expose the original core in the same strict build for numerical comparisons. +std::vector deterministic_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v, bool causal, double scale, + torch::optional key_padding_mask); +std::vector deterministic_attention_backward( + torch::Tensor grad_output, torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor probabilities, bool causal, double scale, + torch::optional key_padding_mask); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("h3_attention_build_info", &h3_attention_build_info); + m.def("h3_attention_forward", &h3_attention_forward); + m.def("h3_attention_probabilities", &h3_attention_probabilities); + m.def("h3_attention_backward_accumulate", &h3_attention_backward_accumulate); + m.def("deterministic_attention_forward", &deterministic_attention_forward); + m.def("deterministic_attention_backward", &deterministic_attention_backward); +} diff --git a/csrc/cuda/attention/deterministic_attention.cu b/csrc/cuda/attention/deterministic_attention.cu index 82ec824db..a3c7974dc 100644 --- a/csrc/cuda/attention/deterministic_attention.cu +++ b/csrc/cuda/attention/deterministic_attention.cu @@ -6,13 +6,17 @@ // All reductions use fixed order; no split-KV or dynamic dispatch. #include +#include #include #include #include #include #include +#include #include +#include +#include #include namespace { @@ -60,6 +64,117 @@ __global__ void qk_kernel( scores[out_idx] = scale * acc; } +// H3 fixed-D dot product. Cooperative, contiguous global loads convert each +// BF16 input once per CTA. Padding/transposition keeps shared-memory reads +// conflict-free for the (16 key, 16 query) thread layout. Each score retains +// the original ascending-d, separate-multiply/add FP32 reduction. +template +__global__ void h3_shared_qk_kernel( + const at::BFloat16* __restrict__ Q, + const at::BFloat16* __restrict__ K, + float* __restrict__ scores, + int64_t Sq, int64_t Skv, float scale) { + __shared__ float sq[16][129]; + __shared__ float sk[128][17]; + const int tid = threadIdx.y * 16 + threadIdx.x; + const int64_t q_start = (int64_t)blockIdx.y * 16; + const int64_t k_start = (int64_t)blockIdx.x * 16; + const int64_t bh = blockIdx.z; + for (int i = tid; i < 16 * 128; i += 256) { + const int row = i / 128; + const int d = i % 128; + sq[row][d] = q_start + row < Sq + ? (float)Q[(bh * Sq + q_start + row) * 128 + d] : 0.0f; + sk[d][row] = k_start + row < Skv + ? (float)K[(bh * Skv + k_start + row) * 128 + d] : 0.0f; + } + __syncthreads(); + const int64_t qi = q_start + threadIdx.y; + const int64_t ki = k_start + threadIdx.x; + if (qi >= Sq || ki >= Skv) return; + float acc = 0.0f; + #pragma unroll 8 + for (int d = 0; d < 128; ++d) { + acc += sq[threadIdx.y][d] * sk[d][threadIdx.x]; + } + scores[(bh * Sq + qi) * Skv + ki] = ApplyScale ? scale * acc : acc; +} + +// H3 PV/dQ and fused dK/dV reuse staged operands in shared memory. +// Each output retains the native scalar ascending reduction order. +// All threads take every barrier, including tail rows/channels. +template +__global__ void h3_staged_row_product( + const float* __restrict__ A, const at::BFloat16* __restrict__ X, + at::BFloat16* __restrict__ Y, int64_t Rows, int64_t Inner, float scale) { + __shared__ float a[8][33]; + __shared__ float x[32][33]; + const int tid = threadIdx.y * 32 + threadIdx.x; + const int64_t bh = blockIdx.z; + const int64_t row = (int64_t)blockIdx.y * 8 + threadIdx.y; + const int d = blockIdx.x * 32 + threadIdx.x; + float acc = 0.0f; + for (int64_t start = 0; start < Inner; start += 32) { + const int ar = tid / 32, ak = tid % 32; + const int64_t global_row = (int64_t)blockIdx.y * 8 + ar; + a[ar][ak] = global_row < Rows && start + ak < Inner + ? A[(bh * Rows + global_row) * Inner + start + ak] : 0.0f; + for (int i = tid; i < 32 * 32; i += 256) { + const int k = i / 32, c = i % 32; + x[k][c] = start + k < Inner + ? (float)X[(bh * Inner + start + k) * 128 + blockIdx.x * 32 + c] : 0.0f; + } + __syncthreads(); + // No padded contributions: even signed zero keeps the original operation count. + const int n = (int)min((int64_t)32, Inner - start); + #pragma unroll 8 + for (int j = 0; j < n; ++j) { + acc += a[threadIdx.y][j] * x[j][threadIdx.x]; + } + __syncthreads(); + } + if (row < Rows) Y[(bh * Rows + row) * 128 + d] = + (at::BFloat16)(Scale ? scale * acc : acc); +} + +__global__ void h3_staged_dkv( + const float* __restrict__ DS, const float* __restrict__ P, + const at::BFloat16* __restrict__ Q, const at::BFloat16* __restrict__ DO, + float* __restrict__ DK, float* __restrict__ DV, int64_t Sq, int64_t Skv) { + __shared__ float ds[32][9], p[32][9]; + __shared__ float q[32][33], dout[32][33]; + const int tid = threadIdx.y * 32 + threadIdx.x; + const int64_t bh = blockIdx.z; + const int64_t ki = (int64_t)blockIdx.y * 8 + threadIdx.y; + const int d = blockIdx.x * 32 + threadIdx.x; + const int64_t index = (bh * Skv + ki) * 128 + d; + float dk = ki < Skv ? DK[index] : 0.0f; + float dv = ki < Skv ? DV[index] : 0.0f; + for (int64_t start = 0; start < Sq; start += 32) { + const int r = tid / 8, c = tid % 8; + const int64_t k = (int64_t)blockIdx.y * 8 + c; + const bool valid = start + r < Sq && k < Skv; + ds[r][c] = valid ? DS[(bh * Sq + start + r) * Skv + k] : 0.0f; + p[r][c] = valid ? P[(bh * Sq + start + r) * Skv + k] : 0.0f; + for (int i = tid; i < 32 * 32; i += 256) { + const int row = i / 32, dim = i % 32; + const bool valid_q = start + row < Sq; + const int64_t offset = (bh * Sq + start + row) * 128 + blockIdx.x * 32 + dim; + q[row][dim] = valid_q ? (float)Q[offset] : 0.0f; + dout[row][dim] = valid_q ? (float)DO[offset] : 0.0f; + } + __syncthreads(); + const int n = (int)min((int64_t)32, Sq - start); + #pragma unroll 8 + for (int j = 0; j < n; ++j) { + dk += ds[j][threadIdx.y] * q[j][threadIdx.x]; + dv += p[j][threadIdx.y] * dout[j][threadIdx.x]; + } + __syncthreads(); + } + if (ki < Skv) { DK[index] = dk; DV[index] = dv; } +} + // --------------------------------------------------------------------------- // Masked Softmax + LSE Kernel // One CTA per (b, hq, q) row. Fixed 256 threads. @@ -263,7 +378,9 @@ std::vector deterministic_attention_forward_impl( bool causal, double scale, torch::optional key_padding_mask, - bool output_fp32) { + bool output_fp32, + bool compute_output = true, + bool h3_shared_qk = false) { check_deterministic_attention_inputs(q, k, v, key_padding_mask); const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); @@ -295,17 +412,24 @@ std::vector deterministic_attention_forward_impl( (Skv + kQKTileK - 1) / kQKTileK, (Sq + kQKTileQ - 1) / kQKTileQ, B * Hq); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, - q_contig.scalar_type(), "qk_kernel", [&] { - qk_kernel<<>>( - q_contig.data_ptr(), - k_contig.data_ptr(), - scores.data_ptr(), - B, Hq, Hkv, Sq, Skv, D, - (float)scale); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); + if (h3_shared_qk) { + h3_shared_qk_kernel<<>>( + q_contig.data_ptr(), k_contig.data_ptr(), + scores.data_ptr(), Sq, Skv, (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "qk_kernel", [&] { + qk_kernel<<>>( + q_contig.data_ptr(), + k_contig.data_ptr(), + scores.data_ptr(), + B, Hq, Hkv, Sq, Skv, D, + (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } } // --- Launch Masked Softmax + LSE kernel --- @@ -326,6 +450,10 @@ std::vector deterministic_attention_forward_impl( C10_CUDA_KERNEL_LAUNCH_CHECK(); } + if (!compute_output) { + return {torch::Tensor(), lse, scores}; + } + // --- Launch PV kernel --- auto out = output_fp32 ? torch::empty(q_contig.sizes(), q_contig.options().dtype(at::kFloat)) @@ -336,24 +464,33 @@ std::vector deterministic_attention_forward_impl( (D + kPVTileD - 1) / kPVTileD, (Sq + kPVTileQ - 1) / kPVTileQ, B * Hq); - AT_DISPATCH_FLOATING_TYPES_AND2( - at::ScalarType::Half, at::ScalarType::BFloat16, - q_contig.scalar_type(), "pv_kernel", [&] { - if (output_fp32) { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); - } else { - pv_kernel<<>>( - scores.data_ptr(), - v_contig.data_ptr(), - out.data_ptr(), - B, Hq, Hkv, Sq, Skv, D); - } - C10_CUDA_KERNEL_LAUNCH_CHECK(); - }); + if (h3_shared_qk && !output_fp32) { + dim3 h3block(32, 8); + dim3 h3grid(4, (Sq + 7) / 8, B * Hq); + h3_staged_row_product<<>>( + scores.data_ptr(), v_contig.data_ptr(), + out.data_ptr(), Sq, Skv, (float)scale); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } else { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::ScalarType::Half, at::ScalarType::BFloat16, + q_contig.scalar_type(), "pv_kernel", [&] { + if (output_fp32) { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } else { + pv_kernel<<>>( + scores.data_ptr(), + v_contig.data_ptr(), + out.data_ptr(), + B, Hq, Hkv, Sq, Skv, D); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + }); + } } return {out, lse, scores}; @@ -373,6 +510,63 @@ std::vector deterministic_attention_forward_fp32( q, k, v, causal, scale, key_padding_mask, true); } +namespace { +void check_h3_build() { +#if !defined(RL_KERNEL_H3_STRICT) || !RL_KERNEL_H3_STRICT || defined(USE_ROCM) + TORCH_CHECK(false, + "H3 requires an audited strict CUDA build; rebuild with RL_KERNEL_H3_STRICT=1"); +#endif +} + +void check_h3_inputs(const torch::Tensor& q, const torch::Tensor& k, + const torch::Tensor& v) { + check_h3_build(); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, + "H3 requires 4-D head-major Q/K/V"); + TORCH_CHECK(q.size(1) == 56 && k.size(1) == 56 && v.size(1) == 56, + "H3 requires 56 MHA heads"); + TORCH_CHECK(q.scalar_type() == at::kBFloat16, "H3 requires BF16"); + TORCH_CHECK(q.size(0) > 0 && q.size(0) <= 65535 / 56, + "H3 batch exceeds CUDA grid bounds"); + // H3's staged PV/dQ/dK/dV kernels use eight rows per grid block. Keep + // their y-dimension within CUDA's 16-bit grid limit as well as the QK + // kernel's sixteen-row limit. + TORCH_CHECK(q.size(2) <= 65535 * 8 && k.size(2) <= 65535 * 8, + "H3 sequence exceeds CUDA grid bounds"); + TORCH_CHECK(q.size(0) * 56 * q.size(2) <= 2147483647LL, + "H3 query chunk exceeds softmax row index bounds"); + check_deterministic_attention_inputs(q, k, v, torch::nullopt); + TORCH_CHECK(at::cuda::getDeviceProperties(q.get_device())->major >= 8, + "H3 BF16 CUDA profile requires SM80 or newer"); +} +} // namespace + +std::map h3_attention_build_info() { +#if defined(RL_KERNEL_H3_STRICT) && RL_KERNEL_H3_STRICT && !defined(USE_ROCM) + return {{"profile", "h3-cuda-strict-v1"}, {"fmad", "false"}, {"ftz", "false"}, + {"prec_div", "true"}, {"prec_sqrt", "true"}, {"fast_math", "false"}, + {"cuda_compiler", std::to_string(__CUDACC_VER_MAJOR__) + "." + + std::to_string(__CUDACC_VER_MINOR__) + "." + + std::to_string(__CUDACC_VER_BUILD__)}}; +#else + return {{"profile", "unqualified"}}; +#endif +} + +std::vector h3_attention_forward( + torch::Tensor q, torch::Tensor k, torch::Tensor v) { + check_h3_inputs(q, k, v); + return deterministic_attention_forward_impl(q, k, v, false, 1.0 / std::sqrt(128.0), + torch::nullopt, false, true, true); +} + +// Recompute only probabilities during H3 backward; do not execute an unused PV. +torch::Tensor h3_attention_probabilities(torch::Tensor q, torch::Tensor k, torch::Tensor v) { + check_h3_inputs(q, k, v); + return deterministic_attention_forward_impl(q, k, v, false, 1.0 / std::sqrt(128.0), + torch::nullopt, false, false, true)[2]; +} + // =========================================================================== // BACKWARD // =========================================================================== @@ -686,6 +880,93 @@ std::vector deterministic_attention_backward( return {dQ, dK, dV}; } +// H3 backward continues FP32 accumulators across query chunks. +std::vector h3_attention_backward_accumulate( + torch::Tensor grad_output, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor P, // recomputed [B, 56, Sq, Skv] FP32 + torch::Tensor dK, torch::Tensor dV) { + check_h3_inputs(q, k, v); + TORCH_CHECK(grad_output.sizes() == q.sizes() && grad_output.device() == q.device() + && grad_output.scalar_type() == q.scalar_type(), "invalid H3 grad_output"); + TORCH_CHECK(P.dim() == 4 && P.size(0) == q.size(0) && P.size(1) == 56 + && P.size(2) == q.size(2) && P.size(3) == k.size(2) + && P.device() == q.device() && P.scalar_type() == at::kFloat, "invalid H3 P"); + for (const auto& acc : {dK, dV}) { + TORCH_CHECK(acc.device() == q.device() && acc.scalar_type() == at::kFloat + && acc.sizes() == k.sizes() && acc.is_contiguous(), "invalid H3 accumulator"); + } + + at::assert_no_overlap(dK, dV); + for (const auto& input : {q, k, v, P, grad_output}) { + at::assert_no_overlap(dK, input); + at::assert_no_overlap(dV, input); + } + + const at::cuda::OptionalCUDAGuard device_guard(at::device_of(q)); + + auto dO = grad_output.contiguous(); + auto q_c = q.contiguous(); + auto k_c = k.contiguous(); + auto v_c = v.contiguous(); + auto P_c = P.contiguous(); + + const int64_t B = q_c.size(0); + const int64_t Hq = q_c.size(1); + const int64_t Sq = q_c.size(2); + const int64_t Skv = k_c.size(2); + + auto stream = at::cuda::getCurrentCUDAStream(); + + // dP = dO @ V^T [B, Hq, Sq, Skv] + auto dP = torch::empty({B, Hq, Sq, Skv}, q_c.options().dtype(at::kFloat)); + { + dim3 block(kQKTileK, kQKTileQ); + dim3 grid( + (Skv + kQKTileK - 1) / kQKTileK, + (Sq + kQKTileQ - 1) / kQKTileQ, + B * Hq); + h3_shared_qk_kernel<<>>( + dO.data_ptr(), v_c.data_ptr(), + dP.data_ptr(), Sq, Skv, 1.0f); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + // Softmax backward: dS = P * (dP - delta), writes in-place over dP + { + const int64_t num_rows = B * Hq * Sq; + softmax_backward_kernel<<>>( + dP.data_ptr(), + P_c.data_ptr(), + Skv); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + // dP buffer now contains dS + + auto dQ = torch::empty_like(q_c); + { + dim3 block(32, 8); + dim3 grid(4, (Sq + 7) / 8, B * Hq); + h3_staged_row_product<<>>( + dP.data_ptr(), k_c.data_ptr(), + dQ.data_ptr(), Sq, Skv, (float)(1.0 / std::sqrt(128.0))); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + { + dim3 block(32, 8); + dim3 grid(4, (Skv + 7) / 8, B * Hq); + h3_staged_dkv<<>>( + dP.data_ptr(), P_c.data_ptr(), + q_c.data_ptr(), dO.data_ptr(), + dK.data_ptr(), dV.data_ptr(), Sq, Skv); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + + return {dQ, dK, dV}; +} + #if defined(USE_ROCM) namespace { diff --git a/docs/operators/README.md b/docs/operators/README.md index 065d754b6..c99451553 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -41,4 +41,5 @@ Every operator page should include: - [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) +- [MiniMax-H3 One Block, Forward and Backward](h3-one-block.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/attention.md b/docs/operators/attention.md index 52f8517ea..886e91935 100644 --- a/docs/operators/attention.md +++ b/docs/operators/attention.md @@ -271,3 +271,60 @@ for measured peak memory at representative shapes. - `Hq` must be divisible by `Hkv` (raises `ValueError` otherwise). - CUDA KV-cache op wrapper is not in scope (caller does cat + calls this op). - No FP8, no multi-GPU / sequence-parallel. + +## MiniMax-H3 full attention + +This implements the CUDA `h3_full_attention` work item in [RFC #420](https://github.com/RL-Align/RL-Kernel/issues/420), with an independent PyTorch reference. It covers the attention core after Q/K normalization and RoPE; it does not close the complete model roadmap. + +### Interface and scope + +```python +from rl_engine.backends.cuda.attention.h3_full_attention import H3FullAttentionOp + +out = H3FullAttentionOp(chunk_rows=128)(q, k, v) +``` + +Q/K/V must be matching contiguous BF16 `[B,S,56,128]` tensors. They must already have the checkpoint's QK RMSNorm and partial three-axis RoPE applied. Output has the same shape and dtype. `forward_with_lse` additionally returns diagnostic FP32 `[B,S,56]` LSE without gradients. + +Only non-causal MHA, single rank, TP1/CP1 and first derivatives are supported. No masks, GQA, KV cache, CP or projections are provided. CUDA requires NVIDIA SM80 or newer and an audited strict extension; unsupported inputs/builds fail closed. H20 is the measured device. ROCm is outside this PR. + +`kernel_registry.get_op("h3_full_attention", device="cuda")` selects CUDA. The semantic registry explicitly declares dtype/layout/topology and rejects unsupported requests. + +### Arithmetic and workspace + +The CUDA path retains #240's ascending-D scalar QK reduction, 256-thread softmax tree and ascending-key PV/dQ reductions. Shared-memory staging reuses converted Q/K and dO/V operands and rows of probabilities/dS. Fused dK/dV loading shares operands while maintaining two separate ascending-query FP32 accumulation chains. Each gradient element has one writer; there are no cross-CTA atomics or split-K partial-sum merges. + +Queries are processed in fixed-size chunks, while each query consumes all keys. Forward saves Q/K/V; backward recomputes probabilities using the same QK/softmax arithmetic, skipping unused PV. dK/dV retain FP32 accumulators across chunks, dK is scaled once, and both gradients are cast once at the final boundary. Fixed-chunk probability workspace is `[B,56,min(chunk_rows,S),S]` FP32; backward additionally needs dP/dS and two `[B,56,S,128]` FP32 accumulators. Workspace grows linearly with S at a fixed chunk size; compute remains quadratic. + +Strict compilation uses `--fmad=false --ftz=false --prec-div=true --prec-sqrt=true`. Fast math and external NVCC prepend/append flags are rejected. An isolated `_H3_C` extension enforces this profile without changing ordinary `_C` source/compiler/linker settings. Native H3 entry points and the wrapper verify build identity. Raw-bit promises apply within a qualified device/build, not across platforms or compilers. Independent mathematical accuracy is a separate tolerance check. + +### Build and validation + +Use an isolated checkout and environment. The ordinary build intentionally omits the strict extension. + +```bash +RL_KERNEL_H3_STRICT=1 RL_KERNEL_REQUIRE_EXT=1 TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=2 \ + python setup.py build_ext --inplace +python -m pytest -q tests/backends/cuda/test_h3_full_attention.py +python tools/validation/operators/check_operator.py --op h3_full_attention \ + --candidate cuda --device cuda --dtype bf16 --batch 2 --seq 129 --check-grad --json +``` + +Tests cover independent FP32/FP64 output/all-gradient accuracy, strict-native raw bits, +repeat/batch/chunk invariance, LSE, invalid inputs, checkpointing and streams. CUDA tests require the optional strict `_H3_C`; +CPU-only success or skipped GPU tests do not establish GPU qualification. + +### Performance reproduction + +```bash +python benchmarks/operators/attention/benchmark_deterministic_attention.py h3 \ + --output /path/to/local-evidence/matrix +``` + +The matrix compares strict native and CUDA chunks 64/128/256 on six shapes. +Output/all-gradient raw-bit checks precede timing; two warmups and six rotating samples +include public-wrapper copies, allocations and forward plus all QKV gradients. JSON records +raw samples, extra peak allocation, environment, backend readback and source/binary hashes. +At fixed chunk size, probability workspace grows linearly with sequence length while compute +remains quadratic. These are operator measurements, without +full-model, multi-GPU or ROCm hardware qualification. diff --git a/docs/operators/h3-one-block.md b/docs/operators/h3-one-block.md new file mode 100644 index 000000000..b2f1d4ce0 --- /dev/null +++ b/docs/operators/h3-one-block.md @@ -0,0 +1,183 @@ +# MiniMax-H3 One Block, Forward and Backward + +## Summary + +`ws1_one_h3_block` (RFC #420, WS1 closeout) runs **block 0 of the pinned MiniMax-H3 checkpoint** +node by node, forward and backward, and reports where it first departs from diffusers. The block +is a fixed graph of 17 nodes, from the AdaLN projection to the second gated residual. It runs on +packed FL2VA layouts that are **bitwise equal to the pinned diffusers pipeline's** `position_ids` +and modality tags. + +On one B200, at S = 264, 925 and 3160: + +- **Determinism.** Every node, every gradient reaching a node, and all 14 leaf gradients are + bitwise repeatable. +- **Batch invariance.** Batch row 0 of a two-row batch has the single-row bytes at every node, + forward and backward. +- **Agreement with diffusers.** Every node that promises it is bitwise equal to diffusers on + diffusers' inputs. The first drift from diffusers is `adaln_projection`, the first reduction + whose tree differs from cuBLAS. +- **Accuracy.** The block output is exactly as close to the FP32 golden as diffusers' (relative + error 5.2e-3 / 7.9e-3 / 9.0e-3). + +Only three of the 17 nodes are bound to this stack's row operators. The others use **interim** +deterministic operators until their own rows land, or the provider replay where no kernel exists +yet (see Bindings). The block-level promises above hold with whatever a node is bound to. Speed +and the exact rounding of interim nodes belong to those operators, not to this row. + +## Entry Point + +```python +from rl_engine.validation.models.h3_block import CandidateOps, run_block_backward, run_block_case +from rl_engine.validation.models.h3_weights import load_h3_block_weights + +weights = load_h3_block_weights("cuda") # conditioning + block-0 tensors, sha256-checked +ops = CandidateOps() +forward = run_block_case(weights, layout="small", ops=ops) # per node, first_*_drift +backward = run_block_backward(weights, layout="small", ops=ops) +``` + +The CLI writes the JSON evidence and the plot script renders it: + +```bash +python tools/weights/prepare_h3_weights.py --out --block # shard 1 of 14, ~4.8 GB +export RL_KERNEL_H3_WEIGHTS= +python tools/validation/models/h3_block_replay.py --layouts tiny,small,medium --out report.json +python tools/validation/models/plot_h3_block.py report.json +``` + +## Graph and Bindings + +`rl_engine/validation/models/h3_block.py` `NODES` (graph order): + +| Node | Owning RFC row | Binding today | +|---|---|---| +| `adaln_projection` | `adaln_projection_3mod` | **row**: `H3AdaLNProjectionCudaOp` (registry) | +| `norm1`, `norm2` | `h3_rmsnorm` | **row**: `H3RMSNormCudaOp.forward_modulated` (registry) | +| `q_proj`, `k_proj`, `v_proj` | `h3_qkv_gemm` | interim: Triton FP32-accumulate GEMM, one BF16 rounding | +| `q_norm`, `k_norm` | `h3_qk_rmsnorm_d128` | interim: `H3RMSNormCudaOp.forward` on the 128-wide heads | +| `rope_q`, `rope_k` | `h3_mm_rope_3axis_partial` | reference: the provider replay (no kernel exists yet) | +| `attention` | `h3_full_attention` | interim: `DeterministicAttentionOp(causal=False)` | +| `o_proj` | `h3_attention_o_gemm` | interim: Triton FP32-accumulate GEMM | +| `residual_attn`, `residual_mlp` | `adaln_gate_residual` | **row**: `H3GateResidualCudaOp` (registry) | +| `ffn_gate_up` | `h3_ffn_gate_up_gemm` | interim: Triton FP32-accumulate GEMM | +| `swiglu` | `h3_swiglu` | interim: `SwiGLUCudaOp` (FP32 math, one rounding) | +| `ffn_down` | `h3_ffn_down_gemm` | interim: Triton FP32-accumulate GEMM | + +When a row lands, replacing its node's `candidate` callable is the whole integration. The node's +report entries then become that operator's block-level evidence. + +The interim GEMMs use `_triton_gemm_fp32`, one fixed-order FP32 K loop with no split-K, and round +to BF16 once. They do not use `TritonDetGemmOp.__call__`, whose K tree has BF16 nodes. Measured on +B200: + +| Shape | Tree path err | FP32 K loop → BF16 | cuBLAS | Tree | FP32 K loop | +|---|---|---|---|---|---| +| 264 × 5376 → 7168 | 7.0e-3 | 2.4e-3 | 2.4e-3 | 10.4 ms | 1.2 ms | +| 925 × 14336 → 5376 | 7.5e-3 | 2.6e-3 | 2.6e-3 | 47 ms | 2.7 ms | + +Both paths are bitwise repeatable and row-invariant. Error is max|err| / max|ref| against FP64. + +## What Each Node Reports + +Every node runs three ways: the provider replay (`h3_provider`, diffusers@f53d552 op for op), the +candidate (chained, and *isolated* on the provider's inputs), and the FP32 golden +(`rl_engine/reference/minimax_h3/block.py`: FP32 math, TF32 off, no BF16 storage between nodes). + +- **Forward:** `repeat_bitwise_equal` and `batch_row_bitwise_equal` (row 0 of a B = 2 batch whose + row 1 has a different input). Also chained and isolated agreement with the provider (bitwise, + max abs, first differing packed row), the candidate's and provider's `max|err| / max|golden|`, + and the sha256 of the node's output. +- **Backward:** for the block input, `temb` and the 12 block parameters: repeat bitwise, and the + candidate's and provider's error against the golden. For the gradient reaching each node's + output, in reverse graph order: repeat bitwise and batch-row bitwise. Parameter gradients sum + over the batch, so only activation gradients take part in the batch check. +- **First drift:** `first_drift`, `first_isolated_drift`, `first_repeat_drift`, + `first_batch_drift`, `first_grad_repeat_drift` and `first_grad_batch_drift` each name the first + node where that comparison stops being bitwise. +- **Permutation probe (control H-C3, reported, not promised):** permute the packed rows together + with positions and AdaLN indices, and compare the output to the permuted original. + +Inputs: `temb` comes from the provider's conditioning path (the conditioning rows are checked by +`h3_chain`). The block input is a seeded BF16 stand-in, because the packed input projections are +other rows. Text rows use AdaLN table row 0 and media rows use row 1. + +## Evidence + +[`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/report.json) was generated from a clean tree +at `41e4293` on an otherwise idle B200 (torch 2.13.0+cu130): + +![one-block evidence](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/figure.png) + +| Layout | S | First drift / first isolated drift | Repeat / batch drift (fwd, grad) | Output err: RL-Kernel / diffusers | Fwd ms: RL-Kernel / diffusers | Fwd+bwd ms | +|---|---|---|---|---|---|---| +| tiny | 264 | `adaln_projection` / `adaln_projection` | none | 5.21e-3 / 5.21e-3 | 7.4 / 0.61 | 22.7 / 2.8 | +| small | 925 | `adaln_projection` / `adaln_projection` | none | 7.87e-3 / 7.87e-3 | 30.1 / 1.12 | 81.7 / 5.0 | +| medium | 3160 | `adaln_projection` / `adaln_projection` | none | 8.97e-3 / 8.97e-3 | 225 / 3.8 | 577 / 11.5 | + +- **Forward accuracy.** Per node, the candidate's error is within 0.93–1.20× diffusers'. The + largest is `attention` at S = 3160, which is the interim attention operator's. +- **Gradient accuracy.** The candidate is better than diffusers on `adaln_proj.weight`, + `adaln_proj.bias` and `temb`, by up to 7×: diffusers' `index_select` backward accumulates the + table gradient in BF16 atomics, while the row operators use FP32 segment sums. The candidate is + within 1.16× on the other leaves, except `attn.norm_q.weight` at S = 3160 (9.5e-3 vs 5.5e-3). + At S = 925 the same leaf is 2× *more* accurate than diffusers (1.3e-2 vs 2.4e-2), so this + 128-element gradient's error moves with BF16 rounding rather than favouring one side. +- **Permutation probe (H-C3).** The output changes in 79 / 450 / 2199 of the rows, by at most 8. + Attention sums over keys in packed order, so a consistent permutation is not byte-preserving with + the interim attention operator. This is for `h3_full_attention` to settle. +- **Speed.** The candidate is 12–60× slower than diffusers forward. Almost all of it is the + interim GEMMs and attention (FP32 scores materialised). The three row operators are each + faster than or equal to diffusers in their own evidence. + +## Existing implementations (RFC #420 reuse rule) + +- **Is there an existing block-level harness?** diffusers' `MiniMaxH3TransformerBlock` is the + executable reference. It is replayed op for op as the provider, so diffusers is not imported and + the replay pins the reference's dtypes and ordering. The Qwen3-8B chain gate from #315 + (`rl_engine/validation/models/chain_gate.py`) has per-node token digests and first-drift + localisation, but its graph, shapes and acceptance are Qwen3-specific (causal attention, no + AdaLN). This row reuses its method (per-node digests, graph-ordered first drift) on the H3 + `h3_chain` stage pattern. +- **Interim operators.** These are reused from RL-Kernel and are not new code. The only choice + made here is the FP32-accumulate entry of the Triton GEMM instead of its BF16 tree, for the + accuracy and speed measured above. + +## Tests + +`tests/models/minimax_h3/test_h3_one_block.py`: + +- **CPU:** + - the three layouts hash to the pinned diffusers `build_packed_sequence` output; + - the graph is well formed; + - the partial MM-RoPE leaves channels 96–127 unchanged. +- **GPU, `tiny` and `small` forward:** + - every node repeats and is batch-row bitwise; + - promised nodes are isolated-bitwise to diffusers; + - node and output errors are at most 1.25 × diffusers' + 1e-3. +- **GPU, `tiny` backward:** + - every leaf repeats; + - dX is batch-row bitwise; + - no gradient drift; + - leaf errors are at most 1.5 × diffusers' + 5e-3, or, for a BF16 leaf, at most 6 ulps + (of its largest golden element) above diffusers'. A stored gradient's worst element moves + in whole ulps, so the relative rule alone is seed- and torch-version-sensitive; + `tools/validation/models/h3_block_ulp_sweep.py` measured at most +4.15 ulps over + tiny/small × 10 seeds on torch 2.13 and 2.9.1, and +5.15 ulps over 34 single-node + candidate/provider swaps (reports: `reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json`). + +```bash +RL_KERNEL_H3_WEIGHTS= CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/models/minimax_h3 -q +``` + +## Known Limitations + +- **Single platform.** CUDA only. `DeterministicAttentionOp` and the Triton GEMM also build for + ROCm, but the block has not been run there. +- **Interim nodes.** Eleven of the 17 nodes are interim or reference bindings. Their speed and + rounding are not the owning rows' final ones, and the MM-RoPE nodes are the provider replay + itself, so they carry no RL-Kernel evidence until `h3_mm_rope_3axis_partial` lands. +- **Stand-in inputs.** The block input is a seeded stand-in and `temb` is the provider's. + Real packed projections arrive with the input-projection rows. +- **Not covered here.** Block 0 only (`ws1_full_h3_transformer` covers 50 blocks and the refiner), + and one sequence per batch row (H3's batch is a replication axis). diff --git a/examples/models/minimax_h3/diagnose-block.sbatch b/examples/models/minimax_h3/diagnose-block.sbatch new file mode 100644 index 000000000..128ea1955 --- /dev/null +++ b/examples/models/minimax_h3/diagnose-block.sbatch @@ -0,0 +1,41 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +failed_result=$3 +result_root=$4 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir" RL_KERNEL_H3_WEIGHTS="$failed_result/preflight/weights" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +# Preserve the exact native binary from the failed attempt; record and recheck +# its identity instead of changing both the kernel and diagnostic at once. +"$python_bin" - "$failed_result" "$result_dir" <<'PY' +import json +import shutil +import sys +from pathlib import Path +from tools.validation.models.h3_full_model import digest +failed, output = map(Path, sys.argv[1:]) +identity = json.loads((failed / "preflight/identity.json").read_text()) +source = Path(identity["extension"]) +assert digest(source) == identity["extension_sha256"] +shutil.copy2(source, Path("rl_engine") / source.name) +(output / "native-origin.json").write_text(json.dumps(identity, indent=2) + "\n") +(output / "retry-of.txt").write_text(str(failed) + "\n") +PY +"$python_bin" -m tools.validation.models.h3_block_diagnostic --failed-job "$failed_result" --causal-sweep \ + --output "$result_dir/diagnostic" > "$result_dir/diagnostic.log" 2>&1 diff --git a/rl_engine/backends/cuda/attention/h3_full_attention.py b/rl_engine/backends/cuda/attention/h3_full_attention.py new file mode 100644 index 000000000..0a30fa04b --- /dev/null +++ b/rl_engine/backends/cuda/attention/h3_full_attention.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""H3 non-causal MHA with query-chunk recomputation, retaining #240 arithmetic.""" +from __future__ import annotations + +import importlib +import math +from types import ModuleType + +import torch +from torch.autograd import Function +from torch.autograd.function import once_differentiable + +_C: ModuleType | None +try: + _C = importlib.import_module("rl_engine._H3_C") +except ImportError: + _C = None + +_SCALE = 1.0 / math.sqrt(128) + + +def _extension(): + required = ( + "h3_attention_forward", + "h3_attention_build_info", + "h3_attention_backward_accumulate", + "h3_attention_probabilities", + ) + if _C is None or any(not hasattr(_C, name) for name in required): + raise RuntimeError("H3 CUDA extension unavailable; rebuild RL-Kernel (no fallback)") + if _C.h3_attention_build_info().get("profile") != "h3-cuda-strict-v1": + raise RuntimeError("H3 requires a strict CUDA build; rebuild with RL_KERNEL_H3_STRICT=1") + return _C + + +class _H3AttentionFn(Function): + @staticmethod + def forward(ctx, q, k, v, chunk_rows): + ext = _extension() + q, k, v = (x.contiguous() for x in (q, k, v)) + out = torch.empty_like(q) + lse = torch.empty(q.shape[:-1], device=q.device, dtype=torch.float32) + for start in range(0, q.shape[2], chunk_rows): + end = min(start + chunk_rows, q.shape[2]) + block_out, block_lse, p = ext.h3_attention_forward( + q[:, :, start:end].contiguous(), k, v + ) + out[:, :, start:end] = block_out + lse[:, :, start:end] = block_lse + del block_out, block_lse, p + ctx.save_for_backward(q, k, v) + ctx.chunk_rows = chunk_rows + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + @once_differentiable + def backward(ctx, grad_out, grad_lse): + ext = _extension() + q, k, v = ctx.saved_tensors + dq = torch.empty_like(q) + dk = torch.zeros_like(k, dtype=torch.float32) + dv = torch.zeros_like(v, dtype=torch.float32) + for start in range(0, q.shape[2], ctx.chunk_rows): + end = min(start + ctx.chunk_rows, q.shape[2]) + qt = q[:, :, start:end].contiguous() + # Deliberately retain the original softmax (not an LSE-based alternative). + p = ext.h3_attention_probabilities(qt, k, v) + dqt, _, _ = ext.h3_attention_backward_accumulate( + grad_out[:, :, start:end].contiguous(), + qt, + k, + v, + p, + dk, + dv, + ) + dq[:, :, start:end] = dqt + del qt, p, dqt + # No per-chunk rounding or partial-sum merge: scale/cast exactly once. + dk.mul_(_SCALE) + return dq, dk.to(k.dtype), dv.to(v.dtype), None + + +class H3FullAttentionOp: + """BF16 [B,S,56,128] -> [B,S,56,128], Q/K already normalized/rotated. + + Physical sequence order defines the fixed reduction order. + """ + + _function: type[Function] = _H3AttentionFn + + backend_id = "rlkernel.cuda.h3_full_attention.query_chunk.v1" + + def __init__(self, chunk_rows=128): + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, int) or chunk_rows <= 0: + raise ValueError("chunk_rows must be a positive integer") + self.chunk_rows = chunk_rows + + @staticmethod + def _validate(q, k, v): + if any(x.ndim != 4 for x in (q, k, v)): + raise ValueError("H3 Q/K/V must have shape [B,S,56,128]") + if q.shape != k.shape or q.shape != v.shape or q.shape[2:] != (56, 128): + raise ValueError("H3 requires equal Q/K/V shapes [B,S,56,128]") + if not 0 < q.shape[0] <= 65535 // 56 or q.shape[1] <= 0: + raise ValueError("H3 requires nonempty tensors and batch <= 1170") + # Native H3 staging launches eight query/key rows per grid block. + if q.shape[1] > 65535 * 8: + raise ValueError("H3 sequence exceeds CUDA grid bounds") + if any(x.dtype != torch.bfloat16 for x in (q, k, v)): + raise ValueError("H3 strict checkpoint profile requires BF16 Q/K/V") + if any(not x.is_contiguous() for x in (q, k, v)): + raise ValueError("H3 requires contiguous BSHD Q/K/V; call contiguous() explicitly") + if any(x.device != q.device for x in (k, v)): + raise ValueError("H3 Q/K/V must share a device") + if not q.is_cuda or torch.version.hip is not None: + raise ValueError("H3 CUDA implementation requires NVIDIA CUDA") + + def forward_with_lse(self, q, k, v): + self._validate(q, k, v) + out, lse = self._function.apply( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), self.chunk_rows + ) + return out.transpose(1, 2), lse.transpose(1, 2) + + def forward(self, q, k, v): + return self.forward_with_lse(q, k, v)[0] + + __call__ = forward + + def provenance(self): + return { + "backend": self.backend_id, + "build": dict(_extension().h3_attention_build_info()), + "fallback": False, + "layout": "BSHD", + "heads": 56, + "head_dim": 128, + "dtype": "bf16", + "causal": False, + "qk_preprocessing": "caller-owned RMSNorm + partial 3-axis RoPE", + "chunk_rows": self.chunk_rows, + "accumulator": "fp32", + "reduction": "#240: ascending D/K/Q, softmax 256-thread tree", + "backward": "recompute P; continuous FP32 dK/dV accumulators; final cast", + "supports_cp": False, + "supports_double_backward": False, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py index 763dba546..17a3a3884 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -16,8 +16,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( _segment_tiles, adaln_row_gather_available, @@ -29,6 +27,8 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import validate_h3_adaln_projection from rl_engine.reference.minimax_h3.adaln_row_gather import ( diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py index 6c9f5bf49..f02431903 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -8,7 +8,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, CONTRACT, @@ -18,6 +17,7 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.adaln_projection import ( split_adaln_table, validate_h3_adaln_projection, diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py index 92835f0cc..ee681454d 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -7,8 +7,8 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward 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 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_row_gather import ( h3_adaln_indices, diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py index 2ff1fd16a..741a0a85e 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -8,8 +8,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension 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.det_linear import ( det_linear_available, @@ -19,6 +17,8 @@ silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 2c6c6be55..3d040e9e8 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -7,9 +7,9 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +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 KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py index a6ed3a28b..d9266e759 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -7,9 +7,9 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +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.rmsnorm import ( H3_NORM_EPS, validate_h3_modulation, 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 index adf01094e..37cea620b 100644 --- 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 @@ -26,13 +26,13 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension import _C, _EXT_AVAILABLE 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, @@ -223,6 +223,8 @@ def sp_plan( 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) + if num_rows is None: + raise ValueError("indexed SP plans require num_rows") *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) @@ -370,6 +372,8 @@ def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | No f"index must be the full ({self.layout.seq_len},) row index, " f"got {tuple(index.shape)}" ) + if num_rows is None: + raise ValueError("indexed SP validation requires num_rows") if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): raise IndexError(f"index must be in [0, {num_rows})") diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py index aff81eaa8..4d58f6d9f 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py @@ -7,7 +7,6 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( ACT_NONE, ACT_SILU, @@ -18,6 +17,7 @@ det_linear_forward, silu_backward_fp32, ) +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" 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 177ce97f1..c4ceed334 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 @@ -31,7 +31,6 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( CONTRACT, DINPUT_CHUNK, @@ -43,6 +42,7 @@ silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table diff --git a/rl_engine/reference/attention/standard_attn.py b/rl_engine/reference/attention/standard_attn.py index 804efb6ed..b0a84c225 100644 --- a/rl_engine/reference/attention/standard_attn.py +++ b/rl_engine/reference/attention/standard_attn.py @@ -189,3 +189,38 @@ def _strict_fp32_math(device_type: str): yield finally: torch.backends.cuda.matmul.allow_tf32 = prev_tf32 + + +class H3FullAttentionReference: + def forward_fp32(self, q, k, v): + with NativeAttentionOp._strict_fp32_math(q.device.type): + q, k, v = (x.float().transpose(1, 2) for x in (q, k, v)) + return ((q @ k.transpose(-1, -2) / 128**0.5).softmax(-1) @ v).transpose(1, 2) + + __call__ = forward_fp32 + + @torch.no_grad() + def value_and_grad(self, q, k, v, grad_output, chunk_rows=128): + """FP32 explicit VJP: dS=P*(dO V^T-sum(P*dO V^T)). + + Query chunks bound temporary memory; dK/dV merge GEMM partial sums. + This intentionally differs from the candidate's scalar reduction order. + """ + if isinstance(chunk_rows, bool) or not isinstance(chunk_rows, int) or chunk_rows < 1: + raise ValueError("chunk_rows must be a positive integer") + with NativeAttentionOp._strict_fp32_math(q.device.type): + q, k, v, dy = (x.float().transpose(1, 2) for x in (q, k, v, grad_output)) + out, dq = torch.empty_like(q), torch.empty_like(q) + dk, dv = torch.zeros_like(k), torch.zeros_like(v) + for start in range(0, q.shape[2], chunk_rows): + stop = min(start + chunk_rows, q.shape[2]) + qt, dyt = q[:, :, start:stop], dy[:, :, start:stop] + p = (qt @ k.transpose(-1, -2) / 128**0.5).softmax(-1) + out[:, :, start:stop] = p @ v + dv.add_(p.transpose(-1, -2) @ dyt) + dp = dyt @ v.transpose(-1, -2) + ds = p * (dp - (p * dp).sum(-1, keepdim=True)) + dq[:, :, start:stop] = (ds @ k) / 128**0.5 + dk.add_((ds.transpose(-1, -2) @ qt) / 128**0.5) + del p, dp, ds + return [x.transpose(1, 2) for x in (out, dq, dk, dv)] diff --git a/rl_engine/reference/minimax_h3/block.py b/rl_engine/reference/minimax_h3/block.py new file mode 100644 index 000000000..427d7de2e --- /dev/null +++ b/rl_engine/reference/minimax_h3/block.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 golden of one MiniMax-H3 transformer block (RFC #420 ``ws1_one_h3_block``). + +Every node takes FP32 inputs (BF16 weights and activations upcast exactly), +computes in FP32 with TF32 off, and returns FP32 without rounding back to +BF16, so the golden is the block's arithmetic without the BF16 storage the +provider and the candidate both have. It is the yardstick their errors are +measured against, not a bitwise target. The one declared cast kept from the +checkpoint contract is the AdaLN input: ``silu(temb)`` is rounded to BF16 +before the projection, as every implementation must do it. + +The partial MM-RoPE uses exact FP32 ``cos``/``sin`` tables on the leading 96 +of 128 head channels (the provider rounds the tables to BF16 first). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +H3_HEADS = 56 +H3_HEAD_DIM = 128 +H3_ROTARY_DIM = 96 +H3_NORM_EPS = 1e-5 + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +def golden_adaln_tables(temb, weight, bias, hidden_size: int = 5376) -> tuple[torch.Tensor, ...]: + """Six ``(3T, H)`` FP32 tables; ``silu(temb)`` rounded to the weight dtype first.""" + + act = F.silu(_f32(temb)).to(weight.dtype) + table = F.linear(_f32(act), _f32(weight), _f32(bias)) + return table.view(-1, 6 * hidden_size).chunk(6, dim=-1) + + +def golden_rmsnorm(x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + x = _f32(x) + rstd = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + return x * rstd * _f32(weight) + + +def golden_norm_modulate(x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + n = golden_rmsnorm(x, weight, eps) + return n * (1.0 + _f32(scale).index_select(0, index)) + _f32(shift).index_select(0, index) + + +def golden_linear(x, weight) -> torch.Tensor: + return F.linear(_f32(x), _f32(weight)) + + +def golden_rope_tables(position_ids, rope_freq_dim: int = 16, rope_theta: float = 10000.0): + """FP32 ``(S, 96)`` cos/sin, the same angles as the provider's rotary module.""" + + exponent = torch.arange( + 0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device + ) + inv_freq = 1.0 / (rope_theta ** (exponent / (2 * rope_freq_dim))) + freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs = torch.cat(freqs.unbind(dim=1), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def golden_apply_rotary(x, cos, sin) -> torch.Tensor: + """Rotate-half on the leading ``cos.shape[-1]`` channels of ``(B, S, H, D)``; FP32 tables.""" + + x = _f32(x) + rot = cos.shape[-1] + head, tail = x[..., :rot], x[..., rot:] + x1, x2 = head.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + head = head * cos[None, :, None, :] + rotated * sin[None, :, None, :] + return torch.cat((head, tail), dim=-1) + + +def golden_attention(q, k, v) -> torch.Tensor: + """Non-causal full attention over one packed document, ``(B, S, H, D)`` -> ``(B, S, H*D)``.""" + + q, k, v = (_f32(t).permute(0, 2, 1, 3) for t in (q, k, v)) + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.shape[-1]) + out = torch.matmul(torch.softmax(scores, dim=-1), v) + return out.permute(0, 2, 1, 3).flatten(2, 3) + + +def golden_swiglu(projected) -> torch.Tensor: + hidden, gate = _f32(projected).chunk(2, dim=-1) + return hidden * F.silu(gate) + + +def golden_gate_residual(residual, gate, index, y) -> torch.Tensor: + return _f32(residual) + _f32(gate).index_select(0, index) * _f32(y) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 292ba453c..a733e95ce 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -76,6 +76,7 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): MUSA_FUSED_LOGP_GENERIC = "rl_engine.backends.musa.logprob.logp.MusaFusedLogpOp" CUDA_DETERMINISTIC_LOGP = "rl_engine.kernels.ops.cuda.loss.logp.DeterministicLogpCUDAOp" # Deterministic standard-softmax attention (issue #147); not FlashAttention. + CUDA_H3_FULL_ATTENTION = "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp" CUDA_DETERMINISTIC_ATTENTION = ( "rl_engine.kernels.ops.cuda.attention.deterministic_attn.DeterministicAttentionOp" ) @@ -243,6 +244,34 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( + OperatorBackendDescriptor( + semantic_op="h3_full_attention", + backend_id="rlkernel.cuda.h3_full_attention.query_chunk.v1", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cuda"}), + supported_dtypes=frozenset({"bfloat16"}), + supported_topologies={ + "world_size": (1,), + "tensor_parallel_size": (1,), + "context_parallel_size": (1,), + }, + determinism_or_alignment_properties={ + "deterministic": True, + "batch_invariant": True, + "strict_observable": True, + "arithmetic_profile": "h3-cuda-strict-v1", + "layout": "BSHD", + "heads": 56, + "head_dim": 128, + "causal": False, + }, + lifecycle=OperatorLifecycle.ENGINE_CONSTRUCTION, + implementation_class_or_factory=( + "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp" + ), + fallback_policy=OperatorFallbackPolicy.ERROR, + version_or_build_fingerprint="h3-query-chunk-v1; binary profile checked at invocation", + ), OperatorBackendDescriptor( semantic_op="video_patch_pack_unpack", backend_id="pytorch-video-patch-pack-unpack-v1", @@ -637,6 +666,7 @@ def __init__(self): OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_ATTN, ], + "h3_full_attention": [OpBackend.CUDA_H3_FULL_ATTENTION], "attention": [ OpBackend.CUDA_DETERMINISTIC_ATTENTION, OpBackend.PYTORCH_NATIVE_ATTENTION, @@ -1159,6 +1189,8 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: """Select the best legacy operator for the requested device.""" platform = self._platform_for_device(device) + if op_type == "h3_full_attention" and platform != "cuda": + raise RuntimeError("h3_full_attention requires NVIDIA CUDA; no fallback") candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py new file mode 100644 index 000000000..b20e885a1 --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,718 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 transformer block, node by node (RFC #420 ``ws1_one_h3_block``). + +The block is a fixed graph of 17 nodes (``NODES``), from the AdaLN projection +through attention and the FFN to the second gated residual, on block 0's +pinned weights and a packed FL2VA layout. Every node runs three ways: + +* provider: the diffusers op sequence (``h3_provider``), fed provider inputs; +* candidate: the RL-Kernel binding of the node, fed candidate inputs + (chained) and, separately, provider inputs (isolated); +* golden: the FP32 reference (``rl_engine.reference.minimax_h3.block``). + +Each node also names the RFC row that owns its arithmetic and how it is bound +today (``Binding.status``): + +* ``row`` -- that row's operator, in this PR stack; +* ``interim`` -- an existing deterministic RL-Kernel operator standing in + until the owning row lands (its arithmetic, speed and rounding are that + operator's, not the row's); +* ``reference`` -- no kernel exists yet; the node runs the provider replay + itself, so it is trivially bitwise and contributes no RL-Kernel evidence. + +The report gives, per node: repeat and batch-row bitwise checks, chained and +isolated agreement with the provider, error against the golden for both the +candidate and the provider, and the backend that ran. ``first_*_drift`` name +the first node (graph order) where a comparison stops being bitwise; the +backward report does the same for the gradients reaching each node's output, +in reverse graph order. +""" + +from __future__ import annotations + +import hashlib +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable + +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout, h3_timesteps + +HIDDEN = 5376 +HEADS = 56 +HEAD_DIM = 128 +BLOCK = "transformer_blocks.0" +NUM_TIMESTEPS = 2 + + +@dataclass(frozen=True) +class Binding: + rfc_row: str + status: str # "row" | "interim" | "reference" + backend: str + + +@dataclass(frozen=True) +class Node: + name: str + inputs: tuple[str, ...] + binding: Binding + # (ops, ctx, *inputs) -> output + candidate: Callable[..., Any] + # (ctx, *inputs) -> output + provider: Callable[..., Any] + golden: Callable[..., Any] + # Promise checked by the tests: the isolated candidate equals the provider bitwise. + provider_bitwise_isolated: bool + + +class CandidateOps: + """The operators bound to the block's nodes, resolved once per run. + + The three rows of this stack are resolved through the registry, so the + report records what dispatch actually returned; the interim operators are + named directly because their registry entries (``det_gemm``, + ``attention``, ``swiglu``) prefer backends that are not available or not + deterministic on every GPU this replay runs on. + """ + + def __init__(self, registry: KernelRegistry | None = None) -> None: + from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + registry = registry or KernelRegistry() + self.projection = registry.get_op("adaln_projection_3mod", device="cuda") + self.norm = registry.get_op("h3_rmsnorm", device="cuda") + self.gate = registry.get_op("adaln_gate_residual", device="cuda") + self.gemm = TritonDetGemmOp() + self.attention = DeterministicAttentionOp() + self.swiglu = SwiGLUCudaOp() + + def backends(self) -> dict[str, str]: + return { + name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items() + } + + +def _w(ctx: dict[str, Any], name: str) -> torch.Tensor: + return ctx["weights"][f"{BLOCK}.{name}"] + + +def _heads(x: torch.Tensor) -> torch.Tensor: + return x.unflatten(-1, (HEADS, HEAD_DIM)) + + +class _InterimLinear(torch.autograd.Function): + """``y = x @ W^T`` on the Triton FP32-accumulate kernel, rounded to BF16 once. + + Interim binding for the block's GEMM rows. ``TritonDetGemmOp.__call__`` + reduces K as a tree with BF16 nodes (about 3x cuBLAS's error at these + shapes); ``_triton_gemm_fp32`` runs one fixed-order FP32 K loop (no + split-K), so each output row depends only on its own input row. The + backward uses the same kernel for ``dX`` (reduction over the output + features, row-local) and ``dW`` (reduction over tokens). + """ + + @staticmethod + def forward(ctx, x2, weight): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + ctx.save_for_backward(x2, weight) + return _triton_gemm_fp32(x2, weight.t()).to(x2.dtype) + + @staticmethod + def backward(ctx, grad): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x2, weight = ctx.saved_tensors + grad = grad.contiguous() + dx = _triton_gemm_fp32(grad, weight).to(x2.dtype) + dw = _triton_gemm_fp32(grad.t(), x2).to(weight.dtype) + return dx, dw + + +def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat = x.reshape(-1, x.shape[-1]) + return _InterimLinear.apply(flat, weight).view(*x.shape[:-1], weight.shape[0]) + + +def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor: + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v))) + out = ops.attention(q, k, v, causal=False) + return out.transpose(1, 2).flatten(2, 3) + + +def _cand_swiglu(ops: CandidateOps, projected: torch.Tensor) -> torch.Tensor: + hidden, gate = projected.chunk(2, dim=-1) + return ops.swiglu(gate.contiguous(), hidden.contiguous()) + + +def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding( + row, status, "triton _det_gemm_fp32_kernel (FP32 K loop, one BF16 rounding)" + ), + candidate=lambda ops, ctx, x: _cand_linear(ops, x, _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # tree reduction vs cuBLAS + ) + + +def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: + return Node( + name=name, + inputs=(source, "adaln_projection"), + binding=Binding("h3_rmsnorm", "row", "H3RMSNormCudaOp.forward_modulated"), + candidate=lambda ops, ctx, x, t: ops.norm.forward_modulated( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider=lambda ctx, x, t: prov.provider_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + golden=lambda ctx, x, t: gold.golden_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider_bitwise_isolated=True, + ) + + +def _head_norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + # The row is someone else's; until it lands, this stack's plain H3 RMSNorm + # (bitwise equal to nn.RMSNorm at any width) runs on the 128-wide heads. + binding=Binding("h3_qk_rmsnorm_d128", "interim", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(_heads(x), _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(_heads(x), _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(_heads(x), _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _rope_node(name: str, source: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_mm_rope_3axis_partial", "reference", "provider_apply_rotary"), + candidate=lambda ops, ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + provider=lambda ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + golden=lambda ctx, x: gold.golden_apply_rotary(x, *ctx["rope_fp32"]), + provider_bitwise_isolated=True, + ) + + +def _gate_node(name: str, residual: str, sublayer: str, gate: int) -> Node: + return Node( + name=name, + inputs=(residual, sublayer, "adaln_projection"), + binding=Binding("adaln_gate_residual", "row", "H3GateResidualCudaOp"), + candidate=lambda ops, ctx, r, y, t: ops.gate(r, y, t[gate], ctx["index"]), + provider=lambda ctx, r, y, t: prov.provider_gate_residual(r, t[gate], ctx["index"], y), + golden=lambda ctx, r, y, t: gold.golden_gate_residual(r, t[gate], ctx["index"], y), + provider_bitwise_isolated=True, + ) + + +def _adaln_params(ctx): + return _w(ctx, "adaln_proj.linear.weight"), _w(ctx, "adaln_proj.linear.bias") + + +NODES: list[Node] = [ + Node( + name="adaln_projection", + inputs=("temb",), + binding=Binding("adaln_projection_3mod", "row", "H3AdaLNProjectionCudaOp"), + candidate=lambda ops, ctx, temb: ops.projection(temb, *_adaln_params(ctx)), + provider=lambda ctx, temb: prov.provider_adaln_modulation(temb, *_adaln_params(ctx)), + golden=lambda ctx, temb: gold.golden_adaln_tables(temb, *_adaln_params(ctx)), + provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS + ), + _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), + _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), + _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), + _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), + _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), + _rope_node("rope_q", "q_norm"), + _rope_node("rope_k", "k_norm"), + Node( + name="attention", + inputs=("rope_q", "rope_k", "v_proj"), + binding=Binding("h3_full_attention", "interim", "DeterministicAttentionOp(causal=False)"), + candidate=lambda ops, ctx, q, k, v: _cand_attention(ops, q, k, v), + provider=lambda ctx, q, k, v: prov.provider_attention(q, k, _heads(v)), + golden=lambda ctx, q, k, v: gold.golden_attention(q, k, _heads(v)), + provider_bitwise_isolated=False, + ), + _linear_node("o_proj", "attention", "attn.to_out.0.weight", "h3_attention_o_gemm", "interim"), + _gate_node("residual_attn", "hidden", "o_proj", gate=2), + _norm_modulate_node("norm2", "residual_attn", "norm2.weight", shift=3, scale=4), + _linear_node("ffn_gate_up", "norm2", "ff.net.0.proj.weight", "h3_ffn_gate_up_gemm", "interim"), + Node( + name="swiglu", + inputs=("ffn_gate_up",), + binding=Binding("h3_swiglu", "interim", "SwiGLUCudaOp"), + candidate=lambda ops, ctx, x: _cand_swiglu(ops, x), + provider=lambda ctx, x: prov.provider_swiglu(x), + golden=lambda ctx, x: gold.golden_swiglu(x), + provider_bitwise_isolated=False, # one FP32 rounding vs the eager two + ), + _linear_node("ffn_down", "swiglu", "ff.net.2.weight", "h3_ffn_down_gemm", "interim"), + _gate_node("residual_mlp", "residual_attn", "ffn_down", gate=5), +] +NODE_NAMES = tuple(node.name for node in NODES) +OUTPUT = NODES[-1].name + +# Gradient leaves of the backward replay: the block input, temb and every block parameter. +PARAM_LEAVES = ( + "adaln_proj.linear.weight", + "adaln_proj.linear.bias", + "norm1.weight", + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + "attn.to_out.0.weight", + "norm2.weight", + "ff.net.0.proj.weight", + "ff.net.2.weight", +) + + +# --- Inputs --------------------------------------------------------------------------------- + + +def make_block_context( + weights: dict[str, torch.Tensor], *, layout: str, seed: int = 0, batch: int = 1 +) -> dict[str, Any]: + """Block-0 weights, a packed FL2VA layout, the timestep embedding and a seeded input. + + ``temb`` comes from the provider's conditioning path (``h3_chain`` covers + those rows); the block input is a seeded BF16 stand-in for the packed + projections, which are other rows. Batch rows ``b > 0`` get their own + seeded inputs; row 0 is the same tensor at every ``batch``. + """ + + spec = H3_BLOCK_LAYOUTS[layout] + packed = h3_fl2va_layout( + num_text_tokens=spec["num_text_tokens"], + num_latent_frames=spec["num_latent_frames"], + latent_height=spec["latent_height"], + latent_width=spec["latent_width"], + num_audio_latents=spec["num_audio_latents"], + ) + seq_len = packed["position_ids"].shape[0] + timestep = h3_timesteps(NUM_TIMESTEPS, seed=seed) + w = weights + with torch.no_grad(): + features = prov.provider_time_proj(timestep) + temb = prov.provider_time_embedder( + features, + w["time_embedder.linear_1.weight"], + w["time_embedder.linear_1.bias"], + w["time_embedder.linear_2.weight"], + w["time_embedder.linear_2.bias"], + ) + rows = [ + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2 + b)) + for b in range(batch) + ] + hidden = torch.cat(rows).to(torch.bfloat16).cuda() + return { + "layout": layout, + "seq_len": seq_len, + "batch": batch, + "weights": weights, + "temb": temb, + "hidden": hidden, + "index": packed["timestep_indices"] * 3 + packed["token_tags"], + "position_ids": packed["position_ids"], + "rope": prov.provider_rope(packed["position_ids"]), + "rope_fp32": gold.golden_rope_tables(packed["position_ids"]), + } + + +@contextmanager +def strict_fp32_matmul(): + """Golden matmuls in full FP32: no TF32, no reduced-precision reductions.""" + + saved = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = saved + + +# --- Graph execution ------------------------------------------------------------------------ + + +def run_graph(mode: str, ctx: dict[str, Any], ops: CandidateOps | None = None) -> dict[str, Any]: + """Every node's output for one mode ("candidate", "provider" or "golden").""" + + values: dict[str, Any] = {"hidden": ctx["hidden"], "temb": ctx["temb"]} + # Diagnostic hybrids only: these candidate nodes run their provider replay instead. + provider_nodes = ctx.get("provider_nodes", ()) + with strict_fp32_matmul(): + for node in NODES: + args = [values[name] for name in node.inputs] + if mode == "candidate" and node.name in provider_nodes: + values[node.name] = node.provider(ctx, *args) + elif mode == "candidate": + values[node.name] = node.candidate(ops, ctx, *args) + elif mode == "provider": + values[node.name] = node.provider(ctx, *args) + else: + values[node.name] = node.golden(ctx, *args) + observer = ctx.get("node_observer") + if observer is not None: + observer(node.name) + return values + + +def run_isolated( + ctx: dict[str, Any], ops: CandidateOps, provider: dict[str, Any] +) -> dict[str, Any]: + """Each candidate node fed the provider's inputs to that node.""" + + return { + node.name: node.candidate(ops, ctx, *[provider[name] for name in node.inputs]) + for node in NODES + } + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [t for item in value for t in _flat(item)] + + +def _batch_row(value: Any, row: int, batch: int) -> Any: + """Row ``row`` of a batched activation; AdaLN tables have no batch axis.""" + + if isinstance(value, torch.Tensor): + return value[row : row + 1] if value.dim() >= 3 and value.shape[0] == batch else value + return tuple(_batch_row(v, row, batch) for v in value) + + +def bitwise(lhs: Any, rhs: Any) -> bool: + return all(torch.equal(a, b) for a, b in zip(_flat(lhs), _flat(rhs), strict=True)) + + +def max_abs(lhs: Any, rhs: Any) -> float: + return max( + float((a.float() - b.float()).abs().max()) if a.numel() else 0.0 + for a, b in zip(_flat(lhs), _flat(rhs), strict=True) + ) + + +def first_differing_row(lhs: Any, rhs: Any) -> int | None: + """First packed row (sequence position) where two node outputs differ, or None.""" + + for a, b in zip(_flat(lhs), _flat(rhs), strict=True): + if a.dim() < 3: + if not torch.equal(a, b): + return -1 # a table row, not a sequence position + continue + diff = (a != b).reshape(a.shape[0], a.shape[1], -1).any(dim=-1).any(dim=0) + rows = diff.nonzero() + if rows.numel(): + return int(rows[0]) + return None + + +def rel_err(value: Any, golden: Any) -> float: + """``max|value - golden| / max|golden|`` over a node's tensors.""" + + num = den = 0.0 + for a, g in zip(_flat(value), _flat(golden), strict=True): + g32 = g.float() + num = max(num, float((a.float() - g32).abs().max())) + den = max(den, float(g32.abs().max())) + return num / max(den, 1e-30) + + +def ulp_err(value: torch.Tensor, golden: torch.Tensor) -> float: + """``max|value - golden|`` in units of ``value.dtype``'s ulp at ``max|golden|``. + + ``rel_err`` of a stored gradient moves in steps of this ulp: its worst + element is a single rounded number, so two equally valid reduction orders + can differ by whole ulps there while every other element agrees. + """ + + g32 = golden.float() + top = float(g32.abs().max()) + if top == 0.0: + return 0.0 + ulp = torch.finfo(value.dtype).eps * 2.0 ** torch.floor(torch.log2(torch.tensor(top))).item() + return float((value.float() - g32).abs().max()) / ulp + + +def digest(value: Any) -> str: + """sha256 over a node's tensors' bytes, for comparing runs across processes or hosts.""" + + h = hashlib.sha256() + for t in _flat(value): + h.update(t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return h.hexdigest() + + +def _time_ms(fn: Callable[[], Any], repeats: int) -> float: + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return sorted(samples)[len(samples) // 2] + + +# --- Forward report ------------------------------------------------------------------------- + + +def run_block_case( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Forward replay of block 0: per-node agreement, invariance, accuracy and first drift.""" + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + with torch.no_grad(): + provider = run_graph("provider", ctx) + golden = run_graph("golden", ctx) + chained = run_graph("candidate", ctx, ops) + repeat = run_graph("candidate", ctx, ops) + batched = run_graph("candidate", ctx2, ops) + isolated = run_isolated(ctx, ops, provider) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "backends": ops.backends(), + "nodes": [], + } + firsts: dict[str, str | None] = {"first_drift": None, "first_isolated_drift": None} + firsts.update({"first_repeat_drift": None, "first_batch_drift": None}) + for node in NODES: + name = node.name + row0 = _batch_row(batched[name], 0, 2) + entry: dict[str, Any] = { + "node": name, + "rfc_row": node.binding.rfc_row, + "status": node.binding.status, + "backend": node.binding.backend, + "provider_bitwise_isolated_promised": node.provider_bitwise_isolated, + "repeat_bitwise_equal": bitwise(repeat[name], chained[name]), + "batch_row_bitwise_equal": bitwise(row0, chained[name]), + "chained_vs_provider": { + "bitwise_equal": bitwise(chained[name], provider[name]), + "max_abs": max_abs(chained[name], provider[name]), + "first_differing_row": first_differing_row(chained[name], provider[name]), + }, + "isolated_vs_provider": { + "bitwise_equal": bitwise(isolated[name], provider[name]), + "max_abs": max_abs(isolated[name], provider[name]), + }, + "candidate_rel_err_vs_golden": rel_err(chained[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "candidate_sha256": digest(chained[name]), + } + checks = { + "first_drift": entry["chained_vs_provider"]["bitwise_equal"], + "first_isolated_drift": entry["isolated_vs_provider"]["bitwise_equal"], + "first_repeat_drift": entry["repeat_bitwise_equal"], + "first_batch_drift": entry["batch_row_bitwise_equal"], + } + for key, ok in checks.items(): + if firsts[key] is None and not ok: + firsts[key] = name + report["nodes"].append(entry) + report.update(firsts) + report["output"] = { + "candidate_rel_err_vs_golden": rel_err(chained[OUTPUT], golden[OUTPUT]), + "provider_rel_err_vs_golden": rel_err(provider[OUTPUT], golden[OUTPUT]), + "candidate_vs_provider_max_abs": max_abs(chained[OUTPUT], provider[OUTPUT]), + } + report["permutation_probe"] = permutation_probe(ctx, ops, chained[OUTPUT], seed) + if timing_repeats: + with torch.no_grad(): + report["timing_ms"] = { + "candidate_forward": _time_ms( + lambda: run_graph("candidate", ctx, ops), timing_repeats + ), + "provider_forward": _time_ms(lambda: run_graph("provider", ctx), timing_repeats), + } + return report + + +def permutation_probe( + ctx: dict[str, Any], ops: CandidateOps, base_output: torch.Tensor, seed: int +) -> dict[str, Any]: + """RFC #420 control H-C3, reported and not promised. + + Permute the packed rows together with every per-row tensor (positions, + AdaLN indices); the block output should be the same permutation of the + original. Row-local nodes keep each row's bytes, but attention sums over + keys in packed order, so a different key order may round differently. + """ + + perm = torch.randperm(ctx["seq_len"], generator=torch.Generator().manual_seed(seed + 7)).cuda() + permuted = { + **ctx, + "hidden": ctx["hidden"][:, perm], + "index": ctx["index"][perm], + "position_ids": ctx["position_ids"][perm], + "rope": tuple(t[perm] for t in ctx["rope"]), + "rope_fp32": tuple(t[perm] for t in ctx["rope_fp32"]), + } + with torch.no_grad(): + out = run_graph("candidate", permuted, ops)[OUTPUT] + expected = base_output[:, perm] + return { + "bitwise_equal_after_permutation": bool(torch.equal(out, expected)), + "max_abs": max_abs(out, expected), + "differing_rows": int((out != expected).reshape(out.shape[1], -1).any(dim=-1).sum()), + } + + +# --- Backward report ------------------------------------------------------------------------ + + +def _backward( + mode: str, ctx: dict[str, Any], ops: CandidateOps | None, upstream: torch.Tensor +) -> tuple[dict[str, torch.Tensor], dict[str, list[torch.Tensor]]]: + """Gradients of every leaf, and the gradient reaching each node's output, for one mode. + + The golden runs on FP32 copies of the leaves, so its gradients are not + rounded to the BF16 parameter dtype. + """ + + def leaf(t: torch.Tensor) -> torch.Tensor: + t = t.detach().float() if mode == "golden" else t.detach() + return t.clone().requires_grad_(True) + + leaves = {name: leaf(ctx["weights"][f"{BLOCK}.{name}"]) for name in PARAM_LEAVES} + leaves["hidden"] = leaf(ctx["hidden"]) + leaves["temb"] = leaf(ctx["temb"]) + weights = {**ctx["weights"], **{f"{BLOCK}.{k}": leaves[k] for k in PARAM_LEAVES}} + run_ctx = {**ctx, "weights": weights, "hidden": leaves["hidden"], "temb": leaves["temb"]} + values = run_graph(mode, run_ctx, ops) + for node in NODES: + for t in _flat(values[node.name]): + if t.requires_grad: + t.retain_grad() + out = values[OUTPUT] + torch.autograd.backward(out, upstream.to(out.dtype)) + node_grads = { + node.name: [ + t.grad if t.grad is not None else torch.zeros_like(t) for t in _flat(values[node.name]) + ] + for node in NODES + } + return {name: t.grad for name, t in leaves.items()}, node_grads + + +def run_block_backward( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Backward replay of block 0. + + Per leaf (block input, ``temb`` and the 12 block parameters): repeat + bitwise, and the candidate's and provider's error against the FP32 + golden. Per node, in reverse graph order: whether the gradient reaching + the node's output repeats bitwise, and whether batch row 0 of a two-row + batch gets the same bytes as the batch of one (rows are separate packed + documents, so every activation gradient is row-local). Parameter + gradients sum over the batch and are not part of the batch check. + """ + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + generator = torch.Generator().manual_seed(seed + 11) + up_row0 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + up_row1 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + upstream = up_row0.bfloat16().cuda() + upstream2 = torch.cat([up_row0, up_row1]).bfloat16().cuda() + + golden, _ = _backward("golden", ctx, None, upstream) + provider, _ = _backward("provider", ctx, None, upstream) + first, first_nodes = _backward("candidate", ctx, ops, upstream) + second, second_nodes = _backward("candidate", ctx, ops, upstream) + batched, batched_nodes = _backward("candidate", ctx2, ops, upstream2) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "leaves": {}, + "nodes": [], + } + for name in (*PARAM_LEAVES, "hidden", "temb"): + report["leaves"][name] = { + "repeat_bitwise_equal": bool(torch.equal(first[name], second[name])), + "candidate_rel_err_vs_golden": rel_err(first[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "storage_dtype": str(first[name].dtype).removeprefix("torch."), + "candidate_ulps_vs_golden": ulp_err(first[name], golden[name]), + "provider_ulps_vs_golden": ulp_err(provider[name], golden[name]), + } + report["leaves"]["hidden"]["batch_row_bitwise_equal"] = bool( + torch.equal(batched["hidden"][:1], first["hidden"]) + ) + first_repeat = first_batch = None + for node in reversed(NODES): + name = node.name + row0 = [_batch_row(g, 0, 2) for g in batched_nodes[name]] + entry: dict[str, Any] = { + "node": name, + "grad_repeat_bitwise_equal": bitwise(first_nodes[name], second_nodes[name]), + "grad_batch_row_bitwise_equal": ( + bitwise(row0, first_nodes[name]) if node.name != "adaln_projection" else None + ), + } + if first_repeat is None and not entry["grad_repeat_bitwise_equal"]: + first_repeat = name + if first_batch is None and entry["grad_batch_row_bitwise_equal"] is False: + first_batch = name + report["nodes"].append(entry) + report["first_grad_repeat_drift"] = first_repeat + report["first_grad_batch_drift"] = first_batch + if timing_repeats: + report["timing_ms"] = { + "candidate_forward_backward": _time_ms( + lambda: _backward("candidate", ctx, ops, upstream), timing_repeats + ), + "provider_forward_backward": _time_ms( + lambda: _backward("provider", ctx, None, upstream), timing_repeats + ), + } + return report diff --git a/rl_engine/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py index 62184cc75..92455bae0 100644 --- a/rl_engine/validation/models/h3_cases.py +++ b/rl_engine/validation/models/h3_cases.py @@ -62,3 +62,127 @@ def h3_block_layout( token_tags[lo : lo + text] = 1 token_tags[hi - audio : hi] = 2 return timestep_indices.to(device), token_tags.to(device) + + +# diffusers@f53d552 modular_pipelines/minimax_h3/before_denoise.py constants. +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +def _spatial_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + """One aspect-normalised spatial axis, as numpy's ``linspace(endpoint=False)`` builds it.""" + + import numpy as np + + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def h3_fl2va_layout( + *, + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + audio_channels: int = 2, + patch: tuple[int, int] = (2, 2), + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Packed text/audio/video rows of a text-only FL2VA request, laid out as diffusers packs them. + + Replays ``position_ids``/``token_tags`` of the pinned pipeline's FL2VA + layout step (no keyframe anchors): text rows on the time axis at their row + index, channel-major audio rows pinned to the width extremes, then the video + frames on the non-uniform ``5/3 * (1, 4, 4, 4, 4)`` clock. ``position_ids`` + is ``(S, 3)`` float64 as the pipeline makes it; the rotary module casts it. + ``timestep_indices`` puts text on table row 0 and the generated media on + row 1, a two-timestep stand-in for the conditioning rows. + """ + + import numpy as np + + patch_h, patch_w = patch + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_audio_rows = num_audio_latents * audio_channels + num_video_rows = num_latent_frames * rows_per_frame + seq_len = num_text_tokens + num_audio_rows + num_video_rows + audio_start = num_text_tokens + video_start = audio_start + num_audio_rows + + position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + frame_grid = torch.stack([grid.reshape(-1) for grid in grids], dim=-1) + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(audio_channels) + position_ids[audio_start:video_start, 2] = torch.cat( + [ + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full( + (num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 + ), + ] + ) + spans = torch.tensor( + [ + _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] + for i in range(num_latent_frames) + ], + dtype=torch.float64, + ) + frame_time = float(num_text_tokens) + torch.cat( + [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] + ) + video = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video[:, :, 0] = frame_time[:, None] + video[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video.reshape(-1, 3) + + token_tags = torch.full((seq_len,), VIDEO_TAG, dtype=torch.long) + token_tags[:num_text_tokens] = TEXT_TAG + token_tags[audio_start:video_start] = AUDIO_TAG + timestep_indices = torch.ones(seq_len, dtype=torch.long) + timestep_indices[:num_text_tokens] = 0 + return { + "position_ids": position_ids.to(device), + "token_tags": token_tags.to(device), + "timestep_indices": timestep_indices.to(device), + } + + +# Named FL2VA layouts for the one-block replay: (text, latent frames, latent H, latent W, audio). +H3_BLOCK_LAYOUTS = { + # 32 + 2*20 + 2*(8*12) = 264 rows: the CI-sized case. + "tiny": dict( + num_text_tokens=32, + num_latent_frames=2, + latent_height=16, + latent_width=24, + num_audio_latents=20, + ), + # 77 + 2*40 + 3*(16*16) = 925 rows: a non-power-of-two, three-frame case. + "small": dict( + num_text_tokens=77, + num_latent_frames=3, + latent_height=32, + latent_width=32, + num_audio_latents=40, + ), + # 120 + 2*80 + 5*(18*32) = 3160 rows: a 16:9 five-frame case. + "medium": dict( + num_text_tokens=120, + num_latent_frames=5, + latent_height=36, + latent_width=64, + num_audio_latents=80, + ), +} diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py index 72cbc760a..3aa1aa191 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -307,10 +307,12 @@ def run_case( "seed": seed, "stages": [], } - chained = provider = golden = None + chained: Any = None + provider: Any = None + golden: Any = None first_drift = first_isolated = None # Each mode sees its own earlier outputs (a stage may read any earlier stage). - history = {"candidate": {}, "provider": {}, "golden": {}} + history: dict[str, dict[str, Any]] = {"candidate": {}, "provider": {}, "golden": {}} ctx_c = {**ctx, "history": history["candidate"]} ctx_p = {**ctx, "history": history["provider"]} ctx_g = {**ctx, "history": history["golden"]} @@ -379,6 +381,8 @@ def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torc value = stage.candidate(op, run_ctx, value) else: value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + if value is None: + raise ValueError("conditioning chain produced no output") outputs = list(value) torch.autograd.backward( outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] diff --git a/rl_engine/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json index 55f7e5220..e590a8473 100644 --- a/rl_engine/validation/models/h3_manifest.json +++ b/rl_engine/validation/models/h3_manifest.json @@ -79,6 +79,40 @@ "sha256": "3b12cdc7bbda772e7e31539ffa9a32f529dba7f4148273d2128bdc621fd27438" } }, + "block_tensors": { + "transformer_blocks.0.attn.to_q.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "196f9120ad018784bcefd4224f0e08fae88ba1adf079ecdca18410a414a6a2a4" + }, + "transformer_blocks.0.attn.to_k.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "8e9171f3b8cf894904d1ddbd96510739d7291d2992310de6a52e1d733ec895cd" + }, + "transformer_blocks.0.attn.to_v.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "275bbc993b1ba738831e04d86f290c9988108d96fe8c97c3b576a043c7ff3fe1" + }, + "transformer_blocks.0.attn.norm_q.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "34a3f13557a2c0a9af260e48d431041e381779a0f61616cfe9800ec839693519" + }, + "transformer_blocks.0.attn.norm_k.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "8991949ddea3bc2136a0fcf0b980fa63bd0df760207e3116f7ce96df4e6ca058" + }, + "transformer_blocks.0.attn.to_out.0.weight": { + "dtype": "bfloat16", "shape": [5376, 7168], + "sha256": "822cb426aaebe57d3d219b3c4f6855ca5c13821cb84ebe0a7474fe3a87869fd6" + }, + "transformer_blocks.0.ff.net.0.proj.weight": { + "dtype": "bfloat16", "shape": [28672, 5376], + "sha256": "49587044f325aae73f046c4cda9637052670227ec3c2c289359d9d1ebf29f95d" + }, + "transformer_blocks.0.ff.net.2.weight": { + "dtype": "bfloat16", "shape": [5376, 14336], + "sha256": "e19d4e1649e4b55b1931842e5cc4bb46892c2e635b9315b2996630a86284e88f" + } + }, "reference_implementation": { "repo": "huggingface/diffusers", "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", diff --git a/rl_engine/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py index 19c2b5034..3eec1818f 100644 --- a/rl_engine/validation/models/h3_provider.py +++ b/rl_engine/validation/models/h3_provider.py @@ -118,3 +118,75 @@ def provider_final_adaln_out( return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( 0, timestep_indices ) + + +# --- One transformer block (RFC #420 ``ws1_one_h3_block``) ------------------------------------ + + +def provider_rope( + position_ids: torch.Tensor, rope_freq_dim: int = 16, rope_theta: float = 10000.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """``MiniMaxH3RotaryPosEmbed``: FP32 ``(S, 96)`` cos/sin over the ``(t, h, w)`` axes.""" + + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device) + / (2 * rope_freq_dim) + ) + ) + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def provider_apply_rotary( + hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """``_apply_rotary_emb``: rotate-half on the leading 96 of 128 channels of ``(B, S, H, D)``. + + The FP32 tables are cast to the activation dtype (BF16) before the rotation. + """ + + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + x1, x2 = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-x2, x1), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + +def provider_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A bias-free ``nn.Linear`` of the block (Q/K/V, O and both FFN projections).""" + + return F.linear(x, weight) + + +def provider_head_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + """``norm_q``/``norm_k``: ``nn.RMSNorm(128)`` on ``(B, S, H, D)``.""" + + return F.rms_norm(x, (x.shape[-1],), weight, eps) + + +def provider_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """The native diffusers backend: SDPA on ``(B, H, S, D)`` views, non-causal, no mask. + + Returns ``(B, S, H*D)`` after ``flatten(2, 3).type_as(query)``. + """ + + q, k, v = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + return out.permute(0, 2, 1, 3).flatten(2, 3).type_as(query) + + +def provider_swiglu(projected: torch.Tensor) -> torch.Tensor: + """diffusers ``SwiGLU`` after its projection: ``hidden * silu(gate)`` on the two halves.""" + + hidden_states, gate = projected.chunk(2, dim=-1) + return hidden_states * F.silu(gate) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 9acee0c32..16b212e6d 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -100,7 +100,9 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. """ - row = {key: case[key] for key in ("op", "case", "backend", "bytes") if key in case} + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } keys = [key for key in TIMED_KEYS if key in case] if not keys: return row @@ -116,7 +118,7 @@ def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[st for key in orders[iteration % 2]: _prepared_call(case[key], case.get(f"{key}_setup"))() torch.cuda.synchronize() - samples = {key: [] for key in keys} + samples: dict[str, list[float]] = {key: [] for key in keys} for iteration in range(iters): for key in orders[iteration % 2]: samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) @@ -163,6 +165,8 @@ def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: golden = registry._get_or_create_backend( registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") rows = [] for num in (1, 2, 3, 7, 64, 1000, 4097): t = h3_timesteps(num, seed=num) @@ -227,6 +231,8 @@ def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: op = registry.get_op("timestep_mlp_fp32", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") params = h3_params(MLP_NAMES, MLP_SHAPES) cuda_err, provider_err = [], [] @@ -301,6 +307,8 @@ def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, golden = registry._get_or_create_backend( registry._priority_map["cpu"]["adaln_projection_3mod"][-1] ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") weight, bias = _adaln_params() cuda_frac, provider_frac, early_frac = [], [], [] for seed in range(draws): @@ -389,6 +397,8 @@ def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("adaln_row_gather", device="cuda") golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") # Seeded so the op-backward fractions below are reproducible from this commit. gen = torch.Generator(device="cuda").manual_seed(0) rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() @@ -552,7 +562,7 @@ def golden(x_, w_, sh_, sc_): backward = {} for name, fn in ( ("cuda", lambda *t: op.forward_modulated(*t, index)), - ("provider", lambda *t: provider_norm_modulate(*t, index)), + ("provider", lambda *t: provider_norm_modulate(t[0], t[1], t[2], t[3], index)), ): first, second = grads(fn), grads(fn) backward[name] = { @@ -726,14 +736,16 @@ def backward(state): # read x + norm_out.linear, write out "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), - "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out( + t[0], t[1], t[2], t[3], t[4], i + ), "candidate_backward": backward, "provider_backward": backward, "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( lambda *t: op(*t, i) ), "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( - lambda *t: provider_final_adaln_out(*t, i) + lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], i) ), } ) @@ -745,6 +757,8 @@ def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: golden_op = registry._get_or_create_backend( registry._priority_map["cpu"]["final_adaln_out"][-1] ) + if golden_op is None: + raise RuntimeError("required H3 golden backend is unavailable") x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) ours = op(x, nw, temb, w, b, ti) theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) @@ -778,7 +792,7 @@ def golden_backward(x_, nw_, t_, w_, b_): backward = {} for name, fn in ( ("cuda", lambda *t: op(*t, ti)), - ("provider", lambda *t: provider_final_adaln_out(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], ti)), ): first, second = grads(fn), grads(fn) backward[name] = { diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py index c6e5942c2..f587f2726 100644 --- a/rl_engine/validation/models/h3_weights.py +++ b/rl_engine/validation/models/h3_weights.py @@ -6,7 +6,9 @@ Only the tensors listed in ``h3_manifest.json`` are used. They all live in shard 1 of 14, so ``tools/weights/prepare_h3_weights.py`` downloads that shard, checks it against the manifest and extracts the tensors into -``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. With ``--block`` it also +extracts block 0's attention and FFN weights (manifest ``block_tensors``, +about 1.1 GB) into ``h3_block0.safetensors`` for ``ws1_one_h3_block``. """ from __future__ import annotations @@ -22,6 +24,7 @@ MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" EXTRACTED_FILE = "h3_conditioning.safetensors" +BLOCK_FILE = "h3_block0.safetensors" def load_h3_manifest() -> dict[str, Any]: @@ -54,23 +57,24 @@ def sha256_tensor(tensor: torch.Tensor) -> str: return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() -def load_h3_conditioning_weights( - device: torch.device | str = "cpu", names: list[str] | None = None +def _load_pinned( + filename: str, section: str, device: torch.device | str, names: list[str] | None ) -> dict[str, torch.Tensor]: - """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + """Load one extracted file and check every tensor of a manifest section against it.""" from safetensors.torch import load_file root = h3_weights_dir() if root is None: raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") - path = root / EXTRACTED_FILE + path = root / filename if not path.is_file(): - raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py") - manifest = load_h3_manifest() + flag = " --block" if section == "block_tensors" else "" + raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py{flag}") + specs = load_h3_manifest()[section] tensors = load_file(str(path), device="cpu") - wanted = names if names is not None else list(manifest["tensors"]) - for name, spec in manifest["tensors"].items(): + wanted = names if names is not None else list(specs) + for name, spec in specs.items(): tensor = tensors[name] if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") @@ -81,7 +85,24 @@ def load_h3_conditioning_weights( raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") out: dict[str, torch.Tensor] = {} for name in wanted: - if name not in manifest["tensors"]: + if name not in specs: raise KeyError(name) out[name] = tensors[name].to(device=device) return out + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + return _load_pinned(EXTRACTED_FILE, "tensors", device, names) + + +def load_h3_block_weights(device: torch.device | str = "cpu") -> dict[str, torch.Tensor]: + """Conditioning tensors plus block 0's attention/FFN weights, every one sha-checked.""" + + return { + **load_h3_conditioning_weights(device), + **_load_pinned(BLOCK_FILE, "block_tensors", device, None), + } diff --git a/rl_engine/validation/models/h3_ws2.py b/rl_engine/validation/models/h3_ws2.py index 7342774a0..9ec7b22c3 100644 --- a/rl_engine/validation/models/h3_ws2.py +++ b/rl_engine/validation/models/h3_ws2.py @@ -369,7 +369,7 @@ def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: floa ] started = [] deadline = time.monotonic() + timeout - status = {} + status: dict[int, object] = {} try: for proc in procs: proc.start() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index 16e282514..15d7adc91 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -37,6 +37,7 @@ def make_operator_inputs( "matmul": _make_matmul_inputs, "det_gemm": _make_det_gemm_inputs, "attention": _make_attention_inputs, + "h3_full_attention": _make_h3_attention_inputs, "prefix_shared_attention": _make_prefix_shared_attention_inputs, "cp_attention": _make_cp_attention_inputs, "logp": _make_logp_inputs, @@ -77,6 +78,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "matmul": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "det_gemm": f"{batch}x{seq}x{_matmul_k(args)}x{_matmul_n(args)}", "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", + "h3_full_attention": f"{batch}x{seq}x56x128", "prefix_shared_attention": f"{batch}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}" f"x{seq}x{DEFAULT_HEAD_DIM}", "cp_attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}xcp2", @@ -600,3 +602,11 @@ def _arg_int(args: argparse.Namespace, name: str, default: int) -> int: def _arg_str(args: argparse.Namespace, name: str, default: str) -> str: return str(getattr(args, name, default)) + + +def _make_h3_attention_inputs(args, dtype, device): + batch, seq = _batch_seq(args) + return { + name: _floating_tensor((batch, seq, 56, 128), args, dtype, device, offset) + for offset, name in enumerate(("q", "k", "v")) + } diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index c0f174982..d5e6a7c85 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -32,6 +32,16 @@ def _load_object(path: str) -> Any: OP_SPECS = { + "h3_full_attention": OperatorSpec( + name="h3_full_attention", + op_class="attention", + gold_path=("rl_engine.reference.attention.standard_attn.H3FullAttentionReference"), + gold_method="forward_fp32", + candidate_paths={ + "cuda": "rl_engine.backends.cuda.attention.h3_full_attention.H3FullAttentionOp", + }, + grad_input_names=("q", "k", "v"), + ), "rms_norm": OperatorSpec( name="rms_norm", op_class="reduction", diff --git a/tests/backends/cuda/test_h3_full_attention.py b/tests/backends/cuda/test_h3_full_attention.py new file mode 100644 index 000000000..5ee28e8f0 --- /dev/null +++ b/tests/backends/cuda/test_h3_full_attention.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +import pytest +import torch + +from benchmarks.operators.attention.benchmark_deterministic_attention import H3StrictNative +from benchmarks.operators.attention.benchmark_deterministic_attention import ( + h3_bits_equal as bits_equal, +) +from benchmarks.operators.attention.benchmark_deterministic_attention import ( + h3_value_grad as value_grad, +) +from rl_engine.backends.cuda.attention.h3_full_attention import _C, H3FullAttentionOp +from rl_engine.contracts.numerical import load_contract, resolve_tolerance +from rl_engine.reference.attention.standard_attn import H3FullAttentionReference + + +def case_data(b, s, seed): + torch.manual_seed(seed) + return [ + torch.randn(b, s, 56, 128, device="cuda", dtype=torch.bfloat16).requires_grad_() + for _ in range(3) + ] + + +def assert_accuracy(actual, reference): + contract = load_contract() + for i, (a, g) in enumerate(zip(actual, reference)): + tol = resolve_tolerance( + contract, + op_class="attention", + dtype="bfloat16", + judgment="forward_accuracy" if i == 0 else "gradient_accuracy", + backend_profile="cuda_bf16", + ) + torch.testing.assert_close(a.to(g.dtype), g, atol=tol.atol, rtol=tol.rtol) + + +@pytest.fixture +def gpu_op(): + if not torch.cuda.is_available(): + pytest.skip("requires a GPU") + if torch.version.hip is not None or _C is None: + pytest.skip("requires the optional strict NVIDIA CUDA extension") + assert _C.h3_attention_build_info()["profile"] == "h3-cuda-strict-v1" + return H3FullAttentionOp + + +@pytest.mark.parametrize("value", [0, -1, True, 1.5]) +def test_invalid_chunk(value): + with pytest.raises(ValueError, match="positive integer"): + H3FullAttentionOp(value) + + +@pytest.mark.parametrize( + "shape,dtype", + [ + ((1, 2, 55, 128), torch.bfloat16), + ((1, 2, 56, 128), torch.float32), + ((1, 0, 56, 128), torch.bfloat16), + ], +) +def test_invalid_profile(shape, dtype): + q = torch.empty(shape, dtype=dtype) + with pytest.raises(ValueError): + H3FullAttentionOp()(q, q, q) + + +def test_cpu_fails_closed(): + q = torch.empty(1, 2, 56, 128, dtype=torch.bfloat16) + with pytest.raises(ValueError, match="NVIDIA CUDA"): + H3FullAttentionOp()(q, q, q) + + +@pytest.mark.parametrize("platform", ["cpu", "rocm", "musa", "ascend"]) +def test_registry_rejects_unsupported_platform(monkeypatch, platform): + from rl_engine.runtime.registry import kernel_registry + + monkeypatch.setattr(kernel_registry, "_platform_for_device", lambda _: platform) + with pytest.raises(RuntimeError, match="no fallback"): + kernel_registry.get_op("h3_full_attention", device="cpu") + + +@pytest.mark.parametrize( + "seq,chunk,amplitude", + [ + (1, 1, 1), + (17, 16, 1), + (65, 64, 1), + (129, 128, 1), + (257, 128, 1), + (1024, 128, 1), + (1025, 128, 4), + (4096, 128, 1), + ], +) +def test_accuracy_and_chunk_invariance(gpu_op, seq, chunk, amplitude): + xs = [(x.detach() * amplitude).requires_grad_() for x in case_data(1, seq, 100 + seq)] + dy = torch.randn_like(xs[0]) + op = gpu_op(chunk) + actual = value_grad(op, xs, dy) + golden = H3FullAttentionReference().value_and_grad(*xs, dy) + references = [golden] + if seq <= 257: + + def fp64(q, k, v): + q, k, v = (x.double().transpose(1, 2) for x in (q, k, v)) + return ((q @ k.transpose(-1, -2) / 128**0.5).softmax(-1) @ v).transpose(1, 2) + + references.append( + value_grad(fp64, [x.detach().double().requires_grad_() for x in xs], dy.double()) + ) + for reference in references: + assert_accuracy(actual, reference) + other = value_grad(gpu_op(257), xs, dy) + assert all(bits_equal(a, b) for a, b in zip(actual, other)) + if _C is not None and torch.version.hip is None: + native = value_grad(lambda *x: H3StrictNative.apply(*x)[0], xs, dy) + assert all(bits_equal(a, b) for a, b in zip(actual, native)) + _, lse = op.forward_with_lse(*xs) + _, native_lse = H3StrictNative.apply(*xs) + assert bits_equal(lse, native_lse) + + +def test_batch_checkpoint_and_stream_invariance(gpu_op): + from torch.utils.checkpoint import checkpoint + + op = gpu_op(64) + xs = case_data(3, 129, 781) + dy = torch.randn_like(xs[0]) + expected = value_grad(op, xs, dy) + for inputs, grad, golden in ( + (xs, dy, expected), + ([x[1:2].detach().requires_grad_() for x in xs], dy[1:2], [x[1:2] for x in expected]), + ): + actual = value_grad(op, inputs, grad) + assert all(bits_equal(a, b) for a, b in zip(golden, actual)) + perm = torch.tensor([2, 0, 1], device="cuda") + actual = value_grad(op, [x[perm].detach().requires_grad_() for x in xs], dy[perm]) + assert all(bits_equal(a[perm], b) for a, b in zip(expected, actual)) + unrelated = [x.detach().clone().requires_grad_() for x in xs] + with torch.no_grad(): + for x in unrelated: + x[0].fill_(7) + x[2].fill_(-3) + changed = value_grad(op, unrelated, dy) + assert all(bits_equal(a[1], b[1]) for a, b in zip(expected, changed)) + # Real strided views, not merely a logically permuted contiguous copy. + strided = [] + for x in xs: + backing = torch.empty(3, 258, 56, 128, device="cuda", dtype=torch.bfloat16) + view = backing[:, ::2] + view.copy_(x.detach()) + strided.append(view.requires_grad_()) + with pytest.raises(ValueError, match="contiguous"): + op(*strided) + actual = value_grad(lambda *args: checkpoint(op, *args, use_reentrant=False), xs, dy) + assert all(bits_equal(a, b) for a, b in zip(expected, actual)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + actual = value_grad(op, xs, dy) + torch.cuda.current_stream().wait_stream(stream) + assert all(bits_equal(a, b) for a, b in zip(expected, actual)) + + +@pytest.mark.parametrize("seq,chunk", [(1, 1), (7, 3), (17, 16)]) +def test_bounded_reference_vjp(seq, chunk): + torch.manual_seed(241) + xs = [torch.randn(1, seq, 56, 128, requires_grad=True) for _ in range(3)] + dy = torch.randn_like(xs[0]) + ref = H3FullAttentionReference() + out = ref(*xs) + expected = [out, *torch.autograd.grad(out, xs, dy)] + for actual, golden in zip(ref.value_and_grad(*xs, dy, chunk), expected): + torch.testing.assert_close(actual, golden, atol=2e-6, rtol=2e-5) + + +def test_gtest_registration(): + from argparse import Namespace + + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + args = Namespace(op="h3_full_attention", candidate="cuda", arch_key="sm90", batch=1, seq=7) + case = make_operator_case(args, torch.bfloat16, torch.device("cpu")) + assert case.inputs["q"].shape == (1, 7, 56, 128) + assert isinstance(make_candidate(args).fn, H3FullAttentionOp) + assert case.gold_fn(**case.inputs).dtype == torch.float32 + + +def test_strict_build_flags(monkeypatch): + from build_tools import envs + + for name in ( + "RL_KERNEL_H3_STRICT", + "KERNEL_ALIGN_USE_FAST_MATH", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + ): + monkeypatch.delenv(name, raising=False) + assert envs.h3_strict_cuda_flags() == [] + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + assert set(envs.h3_strict_cuda_flags()) == { + "-DRL_KERNEL_H3_STRICT=1", + "--fmad=false", + "--ftz=false", + "--prec-div=true", + "--prec-sqrt=true", + } + for name, value in ( + ("KERNEL_ALIGN_USE_FAST_MATH", "1"), + ("NVCC_PREPEND_FLAGS", "--use_fast_math"), + ("NVCC_APPEND_FLAGS", "--fmad=true"), + ): + monkeypatch.setenv(name, value) + with pytest.raises(ValueError): + envs.h3_strict_cuda_flags() + monkeypatch.delenv(name) + + +@pytest.fixture +def extension_config(monkeypatch): + from build_tools import extensions + + for name in ( + "RL_KERNEL_H3_STRICT", + "KERNEL_ALIGN_USE_FAST_MATH", + "NVCC_PREPEND_FLAGS", + "NVCC_APPEND_FLAGS", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(extensions, "_load_torch_extension_tools", lambda: (torch, None, dict)) + monkeypatch.setattr(extensions, "_musa_build_available", lambda _: False) + monkeypatch.setattr(extensions, "_ascend_extensions", lambda: []) + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (8, 0)) + return extensions.get_extensions + + +def test_strict_extension_isolation(monkeypatch, extension_config): + baseline = extension_config() + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + ordinary, h3 = extension_config() + assert baseline == [ordinary] + assert ordinary["name"] == "rl_engine._C" + assert h3["name"] == "rl_engine._H3_C" + assert "-DRL_KERNEL_H3_STRICT=1" not in ordinary["extra_compile_args"]["nvcc"] + flags = h3["extra_compile_args"]["nvcc"] + assert "-DRL_KERNEL_H3_STRICT=1" in flags + assert "--fmad=false" in flags and "--ftz=false" in flags + assert "--use_fast_math" not in flags + + +def test_strict_extension_rejects_rocm(monkeypatch, extension_config): + monkeypatch.setenv("RL_KERNEL_H3_STRICT", "1") + monkeypatch.setattr(torch.version, "hip", "test") + with pytest.raises(RuntimeError, match="NVIDIA CUDA only"): + extension_config() + + +def test_unqualified_extension_rejected(monkeypatch): + from types import SimpleNamespace + + from rl_engine.backends.cuda.attention import h3_full_attention as h3 + + fake = SimpleNamespace( + h3_attention_forward=lambda: None, + h3_attention_backward_accumulate=lambda: None, + h3_attention_probabilities=lambda: None, + h3_attention_build_info=lambda: {"profile": "unqualified"}, + ) + monkeypatch.setattr(h3, "_C", fake) + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + h3._extension() + + +@pytest.mark.parametrize( + "device,dtype,world,cp", + [ + ("cuda", "bfloat16", 1, 1), + ("cpu", "bfloat16", 1, 1), + ("rocm", "bfloat16", 1, 1), + ("cuda", "float16", 1, 1), + ("cuda", "bfloat16", 2, 1), + ("cuda", "bfloat16", 1, 2), + ], +) +def test_semantic_dispatch_contract(device, dtype, world, cp): + from rl_engine.runtime.registry import kernel_registry + from rl_engine.runtime.semantic_registry import OperatorRequirements, OperatorResolutionError + + session = kernel_registry.semantic.session() + kwargs = dict( + semantic_op="h3_full_attention", + requested_backend=H3FullAttentionOp.backend_id, + target="training", + requirements=OperatorRequirements( + device=device, + dtype=dtype, + topology={"world_size": world, "tensor_parallel_size": 1, "context_parallel_size": cp}, + alignment_properties={"deterministic": True}, + ), + ) + if (device, dtype, world, cp) == ("cuda", "bfloat16", 1, 1): + resolution = session.resolve(**kwargs) + instance = session.instantiate(resolution) + assert isinstance(instance, H3FullAttentionOp) + assert ( + session.instance_provenance(resolution, instance).backend_id + == H3FullAttentionOp.backend_id + ) + else: + with pytest.raises(OperatorResolutionError): + session.resolve(**kwargs) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, reason="requires NVIDIA CUDA" +) +def test_cuda_build_and_invalid_inputs(): + if _C is None: + with pytest.raises(RuntimeError, match="extension unavailable"): + H3FullAttentionOp()(*case_data(1, 17, 9)) + return + assert hasattr(_C, "h3_attention_build_info") + xs = case_data(1, 17, 9) + op = H3FullAttentionOp() + if _C.h3_attention_build_info()["profile"] != "h3-cuda-strict-v1": + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + op(*xs) + with pytest.raises(RuntimeError, match="RL_KERNEL_H3_STRICT"): + _C.h3_attention_forward(*(x.transpose(1, 2) for x in xs)) + return + for bad in (xs[1].cpu(), xs[1].float(), xs[1][..., :64], xs[1][:, :0]): + with pytest.raises(ValueError): + op(xs[0], bad, xs[2]) + q, k, v = [x.transpose(1, 2).contiguous() for x in xs] + with pytest.raises(RuntimeError): + _C.h3_attention_probabilities(q, k[:, :0], v[:, :0]) + p = _C.h3_attention_probabilities(q, k, v) + acc = torch.zeros_like(k, dtype=torch.float32) + with pytest.raises(RuntimeError): + _C.h3_attention_backward_accumulate(q, q, k, v, p, acc, acc) diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py index 8bc18d0df..f9c2b000f 100644 --- a/tests/models/minimax_h3/conftest.py +++ b/tests/models/minimax_h3/conftest.py @@ -16,6 +16,7 @@ from rl_engine.validation.models.h3_weights import ( WEIGHTS_ENV, h3_weights_dir, + load_h3_block_weights, load_h3_conditioning_weights, ) @@ -32,6 +33,20 @@ def h3_weights_cpu() -> dict[str, torch.Tensor]: pytest.skip(str(exc)) +@pytest.fixture(scope="session") +def h3_block_weights_cuda() -> dict[str, torch.Tensor]: + """Conditioning plus block-0 tensors on CUDA, skipping when they were not extracted.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py --block") + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + try: + return load_h3_block_weights("cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + def h3_kernels_built() -> bool: """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) 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 2b7df3803..0f9fec6c3 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -149,7 +149,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): """Reject semantic index violations without poisoning pytest's CUDA context.""" _cuda_op() - probe = textwrap.dedent(""" + probe = textwrap.dedent( + """ import sys import torch from rl_engine.backends.extension import _C @@ -167,7 +168,8 @@ def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): else: H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) torch.cuda.synchronize() - """) + """ + ) completed = subprocess.run( [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], cwd=Path(__file__).resolve().parents[3], diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py index f9bc79829..ca11c4e53 100644 --- a/tests/models/minimax_h3/test_h3_det_linear.py +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -8,8 +8,8 @@ import pytest import torch -from rl_engine.backends.extension import _C from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available +from rl_engine.backends.extension import _C pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index bcb06f1f2..03a70f1a3 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -286,7 +286,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -318,7 +319,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio sys.exit(0) raise raise AssertionError("invalid index reached native table lookup without an assertion") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, operation, str(bad_index)], capture_output=True, @@ -352,7 +354,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -372,7 +375,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada 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 ) diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py new file mode 100644 index 000000000..8f64a4508 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 block, forward and backward (RFC #420 ``ws1_one_h3_block``). + +GPU tests run block 0 of the pinned checkpoint on packed FL2VA layouts and +check the block-level promises of ``rl_engine.validation.models.h3_block``: + +* every node is bitwise repeatable, and batch row 0 of a two-row batch has the + bytes of the single-row run (forward, and the gradient reaching every node); +* nodes that promise it are bitwise equal to diffusers on diffusers' inputs; +* every node, the block output and every gradient are about as close to the + FP32 golden as diffusers itself. + +CPU tests pin the packed layouts to the pinned diffusers pipeline and check +the graph's structure. +""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.validation.models import h3_block +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# sha256(position_ids float64 bytes + token_tags int64 bytes) of diffusers@f53d552's +# ``MiniMaxH3PrepareLayoutStep.build_packed_sequence`` for each layout (text tags 1, patch +# (1, 2, 2), two audio channels, audio tag 2, video tag 0, no keyframe anchors). +PINNED_LAYOUTS = { + "tiny": "5329a0a13a9996c4fea7a0e676b65aa24302d65bea2e2248ba68d12184cc7c37", + "small": "64e4e79d8fb91d5a68ae7d15bb29049e08add38da2639030e089c6ef80968931", + "medium": "16e448800695eedfb85e0aacbca9bd5415ed53fbbe25f701bcb60df5f4754a31", +} + +# Accuracy: the candidate's error against the FP32 golden may exceed diffusers' by this +# factor plus an absolute floor (relative to the golden's max magnitude). +FORWARD_FACTOR, FORWARD_FLOOR = 1.25, 1e-3 +BACKWARD_FACTOR, BACKWARD_FLOOR = 1.5, 5e-3 +# A BF16 leaf gradient may instead exceed diffusers' error by this many ulps of its largest +# golden element. Its max|err| / max|golden| is set by one rounded element, so it moves in +# whole ulps (0.004 relative for attn.norm_q.weight), and the relative rule alone failed on +# seed 2 with torch 2.13 and on seed 0 with torch 2.9.1. Sized from +# tools/validation/models/h3_block_ulp_sweep.py over tiny/small x seeds 0-9 on both torch +# versions (max excess +4.15 ulp) and 34 single-node candidate/provider swaps (max +5.15 ulp); +# reports in reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json. +BACKWARD_BF16_EXCESS_ULPS = 6.0 + + +@pytest.mark.parametrize("layout", sorted(PINNED_LAYOUTS)) +def test_layout_matches_pinned_pipeline(layout): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS[layout], device="cpu") + data = packed["position_ids"].numpy().tobytes() + packed["token_tags"].numpy().tobytes() + assert hashlib.sha256(data).hexdigest() == PINNED_LAYOUTS[layout] + + +def test_graph_is_well_formed(): + seen = {"hidden", "temb"} + for node in h3_block.NODES: + assert set(node.inputs) <= seen, node.name + assert node.name not in seen, node.name + assert node.binding.status in ("row", "interim", "reference"), node.name + seen.add(node.name) + assert h3_block.OUTPUT == "residual_mlp" + rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} + assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + + +def test_partial_rope_rotates_96_of_128_channels(): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS["tiny"], device="cpu") + cos, sin = prov.provider_rope(packed["position_ids"]) + assert cos.shape == (packed["position_ids"].shape[0], 96) and cos.dtype == torch.float32 + gold_cos, gold_sin = gold.golden_rope_tables(packed["position_ids"]) + assert torch.equal(cos, gold_cos) and torch.equal(sin, gold_sin) + x = torch.randn(1, cos.shape[0], 2, 128).bfloat16() + for out in (prov.provider_apply_rotary(x, cos, sin), gold.golden_apply_rotary(x, cos, sin)): + assert torch.equal(out[..., 96:].float(), x[..., 96:].float()) + assert not torch.equal(out[..., :96].float(), x[..., :96].float()) + + +@pytest.fixture(scope="module") +def ops(): + return h3_block.CandidateOps() + + +@needs_gpu +def test_stack_rows_dispatch_to_cuda(require_h3_kernels, ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + assert backends[name].endswith("CudaOp"), backends + + +@needs_gpu +@pytest.mark.parametrize("layout", ["tiny", "small"]) +def test_block_forward(h3_block_weights_cuda, ops, layout): + report = h3_block.run_block_case(h3_block_weights_cuda, layout=layout, ops=ops) + nodes = {node.name: node for node in h3_block.NODES} + for entry in report["nodes"]: + name = entry["node"] + assert entry["repeat_bitwise_equal"], name + assert entry["batch_row_bitwise_equal"], name + if nodes[name].provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], name + bound = FORWARD_FACTOR * entry["provider_rel_err_vs_golden"] + FORWARD_FLOOR + assert entry["candidate_rel_err_vs_golden"] <= bound, entry + assert report["first_repeat_drift"] is None + assert report["first_batch_drift"] is None + drift = report["first_isolated_drift"] + assert drift is None or not nodes[drift].provider_bitwise_isolated + out = report["output"] + assert out["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * out["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ) + + +@needs_gpu +def test_block_backward(h3_block_weights_cuda, ops): + report = h3_block.run_block_backward(h3_block_weights_cuda, layout="tiny", ops=ops) + for name, entry in report["leaves"].items(): + assert entry["repeat_bitwise_equal"], name + bound = BACKWARD_FACTOR * entry["provider_rel_err_vs_golden"] + BACKWARD_FLOOR + within_rel = entry["candidate_rel_err_vs_golden"] <= bound + within_ulps = entry["storage_dtype"] == "bfloat16" and ( + entry["candidate_ulps_vs_golden"] + <= entry["provider_ulps_vs_golden"] + BACKWARD_BF16_EXCESS_ULPS + ) + assert within_rel or within_ulps, (name, entry) + assert report["leaves"]["hidden"]["batch_row_bitwise_equal"] + assert report["first_grad_repeat_drift"] is None + assert report["first_grad_batch_drift"] is None diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py index 4866e68f3..a60489a04 100644 --- a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -284,7 +284,9 @@ def _worker_never_reports(rank, world, init, queue, path, target, kwargs): @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.models import h3_ws2 workers = { diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 8d469480c..6c7fc9898 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -477,6 +477,8 @@ "docs/usage/evidence/h3-final-adaln-out-b200/chain_replay.json": "reports/experiments/h3-final-adaln-out-b200/chain_replay.json", "docs/usage/evidence/h3-final-adaln-out-b200/figure.png": "reports/experiments/h3-final-adaln-out-b200/figure.png", "docs/usage/evidence/h3-final-adaln-out-b200/report.json": "reports/experiments/h3-final-adaln-out-b200/report.json", + "docs/usage/evidence/h3-one-block-b200/figure.png": "reports/experiments/h3-one-block-b200/figure.png", + "docs/usage/evidence/h3-one-block-b200/report.json": "reports/experiments/h3-one-block-b200/report.json", "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json": "reports/experiments/h3-prior-art-b200/adaln_row_gather.json", diff --git a/tools/validation/models/h3_block_diagnostic.py b/tools/validation/models/h3_block_diagnostic.py new file mode 100644 index 000000000..19e5b111a --- /dev/null +++ b/tools/validation/models/h3_block_diagnostic.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Localize a failed block-0 gradient gate without changing its numerical contract.""" +import argparse +import json +import os +import subprocess +from dataclasses import replace +from pathlib import Path + +import torch + +from rl_engine.backends.extension import _C +from rl_engine.validation.models import h3_block as block +from rl_engine.validation.models.h3_weights import load_h3_block_weights +from tools.validation.models.h3_full_model import digest + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--failed-job", required=True) + parser.add_argument("--causal-sweep", action="store_true") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=False) + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("diagnostic requires clean source") + ops = block.CandidateOps() + weights = load_h3_block_weights("cuda") + ctx = block.make_block_context(weights, layout="tiny", seed=0, batch=1) + upstream = ( + torch.randn((1, ctx["seq_len"], block.HIDDEN), generator=torch.Generator().manual_seed(11)) + .bfloat16() + .cuda() + ) + report = dict( + status="diagnostic_only_not_acceptance", + failed_job=args.failed_job, + job_id=os.environ.get("SLURM_JOB_ID"), + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + extension_sha256=digest(Path(_C.__file__)), + torch=torch.__version__, + device=torch.cuda.get_device_name(), + backends=ops.backends(), + default_matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision_reduction=( + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + ), + leaves={}, + nodes=[], + ) + # Keep FP32 arithmetic strict through backward as well as forward. + with block.strict_fp32_matmul(): + golden, golden_nodes = block._backward("golden", ctx, None, upstream) + provider, provider_nodes = block._backward("provider", ctx, None, upstream) + candidate, candidate_nodes = block._backward("candidate", ctx, ops, upstream) + for name in golden: + report["leaves"][name] = dict( + candidate=block.rel_err(candidate[name], golden[name]), + provider=block.rel_err(provider[name], golden[name]), + ) + for node in reversed(block.NODES): + report["nodes"].append( + dict( + node=node.name, + errors=[ + dict(candidate=block.rel_err(c, g), provider=block.rel_err(p, g)) + for g, p, c in zip( + golden_nodes[node.name], + provider_nodes[node.name], + candidate_nodes[node.name], + strict=True, + ) + ], + ) + ) + # Same input and upstream gradient isolate the RMSNorm implementation from + # error accumulated through other candidate nodes. + with torch.no_grad(): + values = block.run_graph("provider", ctx) + isolated = {} + for name, source, weight_name in ( + ("q_norm", "q_proj", "attn.norm_q.weight"), + ("k_norm", "k_proj", "attn.norm_k.weight"), + ): + grads = {} + node = next(n for n in block.NODES if n.name == name) + for mode in ("golden", "provider", "candidate"): + x = values[source].detach().clone() + w = weights[block.BLOCK + "." + weight_name].detach().clone() + if mode == "golden": + x, w = x.float(), w.float() + x.requires_grad_() + w.requires_grad_() + local = {**ctx, "weights": {**weights, block.BLOCK + "." + weight_name: w}} + y = ( + node.candidate(ops, local, x) + if mode == "candidate" + else getattr(node, mode)(local, x) + ) + torch.autograd.backward(y, provider_nodes[name][0].to(y.dtype)) + grads[mode] = w.grad + isolated[name] = dict( + candidate=block.rel_err(grads["candidate"], grads["golden"]), + provider=block.rel_err(grads["provider"], grads["golden"]), + ) + report["isolated_norm_weight_gradients"] = isolated + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + if args.causal_sweep: + # These hybrids intentionally use provider nodes for localization only. + # They cannot be submitted to the candidate acceptance comparator. + report["causal_sweep"] = [] + original_nodes = block.NODES + try: + for node in original_nodes: + for baseline in ("candidate", "provider"): + mixed = [] + for item in original_nodes: + use_provider = (item.name == node.name) == (baseline == "candidate") + mixed.append( + replace( + item, + candidate=lambda ops, ctx, *xs, item=item: item.provider(ctx, *xs), + ) + if use_provider + else item + ) + block.NODES = tuple(mixed) + with block.strict_fp32_matmul(): + leaves, _ = block._backward("candidate", ctx, ops, upstream) + report["causal_sweep"].append( + dict( + baseline=baseline, + replaced_node=node.name, + errors={ + name: block.rel_err(leaves[name], golden[name]) for name in golden + }, + ) + ) + del leaves + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + finally: + block.NODES = original_nodes + report["causal_sweep_complete"] = True + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_replay.py b/tools/validation/models/h3_block_replay.py new file mode 100644 index 000000000..7d8230294 --- /dev/null +++ b/tools/validation/models/h3_block_replay.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay one real MiniMax-H3 block forward and backward and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_block.py`` for what is compared. + + python tools/weights/prepare_h3_weights.py --out --block + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_replay.py --out block_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_block_weights, + load_h3_manifest, +) + + +def main() -> None: + """Run the forward (and optionally backward) block replay per layout and save the report.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--layouts", + default="tiny,small,medium", + help=f"comma list of packed FL2VA layouts: {', '.join(H3_BLOCK_LAYOUTS)}", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-backward", action="store_true", help="forward replay only") + parser.add_argument("--timing-repeats", type=int, default=5) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + layouts = args.layouts.split(",") + unknown = sorted(set(layouts) - set(H3_BLOCK_LAYOUTS)) + if unknown: + parser.error(f"unknown layouts: {unknown}") + if not torch.cuda.is_available(): + raise SystemExit("the block replay needs a CUDA device") + + weights = load_h3_block_weights("cuda") + manifest = load_h3_manifest() + ops = h3_block.CandidateOps() + + forward, backward = [], [] + for layout in layouts: + case = h3_block.run_block_case( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + forward.append(case) + print( + f"{layout} S={case['seq_len']}: first_drift={case['first_drift']} " + f"first_isolated_drift={case['first_isolated_drift']} " + f"repeat_drift={case['first_repeat_drift']} batch_drift={case['first_batch_drift']} " + f"out err cand {case['output']['candidate_rel_err_vs_golden']:.2e} " + f"prov {case['output']['provider_rel_err_vs_golden']:.2e}" + ) + if not args.no_backward: + case = h3_block.run_block_backward( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + backward.append(case) + det = all(v["repeat_bitwise_equal"] for v in case["leaves"].values()) + print( + f"{layout} backward: leaves {'det' if det else 'NONDET'}, " + f"grad_repeat_drift={case['first_grad_repeat_drift']} " + f"grad_batch_drift={case['first_grad_batch_drift']} " + f"dX batch-row {case['leaves']['hidden']['batch_row_bitwise_equal']}" + ) + torch.cuda.empty_cache() + + evidence = { + "kind": "h3_one_block_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "nodes": [ + { + "node": n.name, + "inputs": list(n.inputs), + "rfc_row": n.binding.rfc_row, + "status": n.binding.status, + "backend": n.binding.backend, + } + for n in h3_block.NODES + ], + "forward_cases": forward, + "backward_cases": backward, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_ulp_sweep.py b/tools/validation/models/h3_block_ulp_sweep.py new file mode 100644 index 000000000..7fe66646b --- /dev/null +++ b/tools/validation/models/h3_block_ulp_sweep.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sweep the one-block backward over layouts and seeds and record leaf errors in ulps. + +Backs the ulp allowance of ``tests/models/minimax_h3/test_h3_one_block.py``: +for every leaf gradient it records the candidate's and diffusers' error against +the FP32 golden both as ``max|err| / max|golden|`` and in ulps of the leaf's +storage dtype at ``max|golden|``, and the per-leaf excess +``candidate_ulps - provider_ulps``. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_ulp_sweep.py --layouts tiny,small --seeds 0,1,2,3,4 \\ + --out ulp_sweep.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_block_weights # noqa: E402 + + +def main() -> None: + """Run the backward replay for every (layout, seed) and write per-leaf ulp errors.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layouts", default="tiny,small") + parser.add_argument("--seeds", default="0,1,2,3,4") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("the sweep needs a CUDA device") + + weights = load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + cases = [] + for layout in args.layouts.split(","): + for seed in (int(s) for s in args.seeds.split(",")): + report = h3_block.run_block_backward(weights, layout=layout, seed=seed, ops=ops) + leaves = { + name: { + key: entry[key] + for key in ( + "storage_dtype", + "candidate_rel_err_vs_golden", + "provider_rel_err_vs_golden", + "candidate_ulps_vs_golden", + "provider_ulps_vs_golden", + ) + } + | { + "excess_ulps": entry["candidate_ulps_vs_golden"] + - entry["provider_ulps_vs_golden"] + } + for name, entry in report["leaves"].items() + } + worst = max(leaves, key=lambda n: leaves[n]["excess_ulps"]) + print( + f"{layout} S={report['seq_len']} seed={seed}: worst excess " + f"{leaves[worst]['excess_ulps']:+.2f} ulp ({worst})" + ) + cases.append( + {"layout": layout, "seq_len": report["seq_len"], "seed": seed, "leaves": leaves} + ) + torch.cuda.empty_cache() + + bf16 = [ + leaf["excess_ulps"] + for case in cases + for leaf in case["leaves"].values() + if leaf["storage_dtype"] == "bfloat16" + ] + evidence = { + "kind": "h3_one_block_backward_ulp_sweep", + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "cases": cases, + "max_bf16_excess_ulps": max(bf16), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"max BF16 excess: {max(bf16):+.2f} ulp; wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence_integrity.py b/tools/validation/models/h3_evidence_integrity.py new file mode 100644 index 000000000..42d9b383d --- /dev/null +++ b/tools/validation/models/h3_evidence_integrity.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate immutable reference inventories before consuming any tensor evidence.""" +from __future__ import annotations + +import json +from pathlib import Path + +from tools.validation.models.h3_full_model import digest + + +def verify_inventory(root: Path, manifest: Path): + records = json.loads(manifest.read_text())["files"] + for name, expected in records.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError("evidence path escapes root") + if digest(path) != expected: + raise ValueError(f"reference evidence changed: {name}") + + +def verify_generation(root): + report = json.loads((root / "report.json").read_text()) + if report["status"] != "generation_complete_comparison_pending": + raise ValueError("generation did not complete") + required = {"initial.pt", "video.npy", "audio.npy"} | { + f"step-{i:03}.pt" for i in range(report["steps"]) + } + if not required.issubset(report["artifacts"]): + raise ValueError("incomplete generation artifact inventory") + for name, expected in report["artifacts"].items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()) or digest(path) != expected: + raise ValueError(f"generation artifact changed: {name}") + return report diff --git a/tools/validation/models/h3_first_step_diagnostic.py b/tools/validation/models/h3_first_step_diagnostic.py new file mode 100644 index 000000000..e433b8481 --- /dev/null +++ b/tools/validation/models/h3_first_step_diagnostic.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the first denoising step of a generation under one execution configuration. + +The first transformer call of a generation sees the saved packed inputs (``initial.pt``) +and nothing has fed back yet, so the difference to the provider's saved ``step-000.pt`` +is the model's own error on a real generation input. A configuration is the provider +(optionally with one forced SDPA backend or cuBLASLt), the candidate, or the candidate +with chosen block nodes run as their provider replay (``--provider-nodes``). Replaying +the provider default must reproduce ``step-000.pt`` bitwise; ``baseline_bitwise`` records +whether it did, which validates the replay itself. Diagnostic only: nothing here accepts +or rejects a candidate. + + python -m tools.validation.models.h3_first_step_diagnostic --checkpoint \\ + --generation --seed 43 --mode candidate \\ + --provider-nodes attention --output out.json +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_generation import QueryChunkAttention +from tools.validation.models.h3_full_model import PROVIDER_VARIANTS, provider_variant_context + + +def errors(value: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: + """max_abs and relative_l2 of ``value`` against ``reference`` (float64 norms).""" + + delta = value.double() - reference.double() + return dict( + max_abs=float(delta.abs().max()), + relative_l2=float(delta.norm() / reference.double().norm().clamp_min(1e-12)), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--contract", type=Path, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--mode", choices=["provider", "candidate"], required=True) + parser.add_argument("--provider-variant", choices=PROVIDER_VARIANTS, default="default") + parser.add_argument("--provider-nodes", default="") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.mode == "candidate" and args.provider_variant != "default": + parser.error("--provider-variant applies to --mode provider only") + if args.mode == "provider" and args.provider_nodes: + parser.error("--provider-nodes applies to --mode candidate only") + nodes = tuple(n for n in args.provider_nodes.split(",") if n) + + from diffusers import MiniMaxH3Transformer3DModel + + directory = args.generation / str(args.seed) + inputs = torch.load(directory / "initial.pt", map_location="cpu", weights_only=True) + saved = torch.load(directory / "step-000.pt", map_location="cpu", weights_only=True) + contract = json.loads(args.contract.read_text()) + + started = time.perf_counter() + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), torch_dtype=torch.bfloat16, local_files_only=True + ) + model = model.to("cuda").eval() + loaded = time.perf_counter() - started + kwargs = {k: v.to("cuda") for k, v in inputs.items()} + ops = None + if args.mode == "candidate": + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + + trace = FullModelTrace() + started = time.perf_counter() + with ( + torch.inference_mode(), + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + ops=ops, + provider_nodes=nodes, + ), + provider_variant_context(args.provider_variant), + ): + video, audio = model(**kwargs, attention_kwargs=None, return_dict=False) + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - started + video, audio = video.float().cpu(), audio.float().cpu() + result = dict( + seed=args.seed, + mode=args.mode, + provider_variant=args.provider_variant, + provider_nodes=list(nodes), + video=errors(video, saved[0]), + audio=errors(audio, saved[1]), + baseline_bitwise=bool(torch.equal(video, saved[0]) and torch.equal(audio, saved[1])), + load_seconds=loaded, + forward_seconds=forward_seconds, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_block.py b/tools/validation/models/plot_h3_block.py new file mode 100644 index 000000000..228f3a289 --- /dev/null +++ b/tools/validation/models/plot_h3_block.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render the RFC #420 one-block replay (``tools/validation/models/h3_block_replay.py``) as a PNG. + +Needs only the report JSON and matplotlib: + + python tools/validation/models/plot_h3_block.py \ + reports/experiments/h3-one-block-b200/report.json + +writes ``figure.png`` next to the report. Panels: the per-node invariance and +provider checks over every layout; the per-node error against the FP32 golden +for the candidate and diffusers; the same for every backward leaf. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import ListedColormap # noqa: E402 + +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +PASS = "#1baf7a" +FAIL = "#d64545" +NA = "#e4e3dd" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" +STATUS_MARK = {"row": "", "interim": "*", "reference": "†"} + + +def _style(ax, title: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=10.5, color=INK, pad=8) + ax.tick_params(colors=MUTED, labelsize=7.5, length=0) + for side in ("top", "right", "left", "bottom"): + ax.spines[side].set_visible(False) + + +def _checks_panel(ax, report: dict[str, Any]) -> None: + """Rows: nodes. Columns: forward/backward checks, ANDed over every layout.""" + + nodes = [n["node"] for n in report["nodes"]] + status = {n["node"]: n["status"] for n in report["nodes"]} + cols = [ + ("repeat", "forward", "repeat_bitwise_equal"), + ("batch row", "forward", "batch_row_bitwise_equal"), + ("= diffusers\n(isolated)", "forward", None), + ("grad\nrepeat", "backward", "grad_repeat_bitwise_equal"), + ("grad\nbatch row", "backward", "grad_batch_row_bitwise_equal"), + ] + grid = [] + for node in nodes: + row = [] + for _label, kind, key in cols: + cases = report["forward_cases" if kind == "forward" else "backward_cases"] + entries = [e for case in cases for e in case["nodes"] if e["node"] == node] + if key is None: + promised = entries[0]["provider_bitwise_isolated_promised"] + ok = all(e["isolated_vs_provider"]["bitwise_equal"] for e in entries) + row.append(2 if not promised and not ok else (1 if ok else 0)) + else: + values = [e[key] for e in entries] + row.append(2 if not values or None in values else (1 if all(values) else 0)) + grid.append(row) + ax.imshow(grid, cmap=ListedColormap([FAIL, PASS, NA]), vmin=0, vmax=2, aspect="auto") + ax.set_xticks(range(len(cols)), [c[0] for c in cols], fontsize=7.5, color=INK) + ax.set_yticks( + range(len(nodes)), [f"{n}{STATUS_MARK[status[n]]}" for n in nodes], fontsize=7.5, color=INK + ) + ax.xaxis.tick_top() + _style(ax, "Bitwise checks, all layouts") + ax.set_xlabel( + "green = holds grey = not promised (reduction differs from diffusers)\n" + "or not applicable (AdaLN tables have no batch axis)\n" + "* interim operator † provider replay (no kernel yet)", + fontsize=7, + color=MUTED, + ) + + +def _err_panel(ax, labels, cand, prov, title: str) -> None: + xs = range(len(labels)) + ax.plot( + xs, [math.log10(v) for v in cand], "o-", color=RL_KERNEL, lw=1.8, ms=5, label="RL-Kernel" + ) + ax.plot( + xs, [math.log10(v) for v in prov], "s--", color=PROVIDER, lw=1.4, ms=4, label="diffusers" + ) + ax.set_xticks(list(xs), labels, rotation=55, ha="right", fontsize=7, color=INK) + lo = math.floor(min(math.log10(v) for v in cand + prov)) + hi = math.ceil(max(math.log10(v) for v in cand + prov)) + ax.set_yticks(range(lo, hi + 1), [f"1e{v}" for v in range(lo, hi + 1)]) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_ylabel("max|err| / max|golden|", color=MUTED, fontsize=8) + ax.legend(frameon=False, fontsize=7.5, labelcolor=INK, loc="upper left") + _style(ax, title) + + +def plot(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(16, 5.6), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.0, 1.6, 1.3]} + ) + case = max(report["forward_cases"], key=lambda c: c["seq_len"]) + back = ( + max(report["backward_cases"], key=lambda c: c["seq_len"]) + if report["backward_cases"] + else None + ) + _checks_panel(axes[0], report) + _err_panel( + axes[1], + [e["node"] for e in case["nodes"]], + [e["candidate_rel_err_vs_golden"] for e in case["nodes"]], + [e["provider_rel_err_vs_golden"] for e in case["nodes"]], + f"Forward error vs FP32 golden, S = {case['seq_len']}", + ) + if back is not None: + leaves = list(back["leaves"]) + _err_panel( + axes[2], + [f"d {name}" for name in leaves], + [back["leaves"][n]["candidate_rel_err_vs_golden"] for n in leaves], + [back["leaves"][n]["provider_rel_err_vs_golden"] for n in leaves], + f"Gradient error vs FP32 golden, S = {back['seq_len']}", + ) + env = report["environment"] + sizes = ", ".join(str(c["seq_len"]) for c in report["forward_cases"]) + fig.suptitle( + "MiniMax-H3 block 0, forward and backward (RFC #420 ws1_one_h3_block)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.905, + f"{env['gpu']} | torch {env['torch']} | layouts S = {sizes} | " + f"first drift vs diffusers: {case['first_drift']} | " + f"rl-kernel {report['rl_kernel_commit'][:7]}", + fontsize=8, + color=MUTED, + ) + fig.tight_layout(rect=(0, 0, 1, 0.89)) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + out = args.report.with_name("figure.png") + plot(report).savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py index 5c6f163f6..f6e32bcb5 100644 --- a/tools/weights/prepare_h3_weights.py +++ b/tools/weights/prepare_h3_weights.py @@ -8,9 +8,11 @@ manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies every sha256 against ``rl_engine/validation/models/h3_manifest.json`` and writes the tensors to ``/h3_conditioning.safetensors``. Point the tests at it with -``export RL_KERNEL_H3_WEIGHTS=``. +``export RL_KERNEL_H3_WEIGHTS=``. ``--block`` also writes block 0's +attention/FFN weights (manifest ``block_tensors``, same shard) to +``/h3_block0.safetensors`` for the ``ws1_one_h3_block`` replay. - python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 + python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 [--block] """ from __future__ import annotations @@ -24,6 +26,7 @@ sys.path.insert(0, str(REPO_ROOT)) from rl_engine.validation.models.h3_weights import ( # noqa: E402 + BLOCK_FILE, EXTRACTED_FILE, load_h3_manifest, sha256_file, @@ -51,6 +54,9 @@ def main() -> None: default=None, help="where the raw shard goes (default: /hf)", ) + parser.add_argument( + "--block", action="store_true", help="also extract block 0 (manifest block_tensors)" + ) args = parser.parse_args() from huggingface_hub import hf_hub_download @@ -80,30 +86,38 @@ def fetch(filename: str) -> Path: _check(index_path, identity["index_sha256"]) weight_map = json.loads(index_path.read_text())["weight_map"] - tensors = {} - for shard in sorted({weight_map[name] for name in manifest["tensors"]}): - if shard not in manifest["weight_shards"]: - raise SystemExit(f"shard {shard} is not pinned in the manifest") - shard_path = fetch(shard) - _check(shard_path, manifest["weight_shards"][shard]["sha256"]) - with safe_open(str(shard_path), framework="pt") as handle: - for name in manifest["tensors"]: - if weight_map[name] == shard: - tensors[name] = handle.get_tensor(name).contiguous() - - for name, spec in manifest["tensors"].items(): - tensor = tensors[name] - if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: - raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") - if list(tensor.shape) != spec["shape"]: - raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") - actual = sha256_tensor(tensor) - if actual != spec["sha256"]: - raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") - - target = out_dir / EXTRACTED_FILE - save_file(tensors, str(target), metadata={"revision": identity["revision"]}) - print(f"wrote {len(tensors)} tensors to {target}") + def extract(section: str, filename: str) -> None: + """Pull one manifest section's tensors from their pinned shards, verify, and save.""" + + specs = manifest[section] + tensors = {} + for shard in sorted({weight_map[name] for name in specs}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in specs: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + + target = out_dir / filename + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + + extract("tensors", EXTRACTED_FILE) + if args.block: + extract("block_tensors", BLOCK_FILE) print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") From 5ecb92decfc1181061a2a427e4a6774a1a183ad1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:16:27 +0000 Subject: [PATCH 42/62] [WS1][CUDA][MiniMax-H3] full-model qualification: thresholds, ensemble rule and reference manifest (RFC #420, #539 split 2/3) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- examples/models/minimax_h3/README.md | 105 + examples/models/minimax_h3/candidate.sbatch | 59 + .../minimax_h3/full-model-ensemble-rule.json | 2346 +++++++++++++ .../minimax_h3/full-model-thresholds.json | 3064 +++++++++++++++++ .../models/minimax_h3/reference-manifest.json | 1641 +++++++++ examples/models/minimax_h3/reference.sbatch | 51 + .../validation/models/h3_full_manifest.json | 311 ++ rl_engine/validation/models/h3_full_model.py | 242 ++ tests/models/minimax_h3/test_h3_full_model.py | 187 + .../validation/models/h3_candidate_prepare.py | 90 + tools/validation/models/h3_full_compare.py | 152 + .../models/h3_full_ensemble_rule.py | 156 + tools/validation/models/h3_full_model.py | 297 ++ .../models/h3_full_provider_spread.py | 92 + 14 files changed, 8793 insertions(+) create mode 100644 examples/models/minimax_h3/README.md create mode 100644 examples/models/minimax_h3/candidate.sbatch create mode 100644 examples/models/minimax_h3/full-model-ensemble-rule.json create mode 100644 examples/models/minimax_h3/full-model-thresholds.json create mode 100644 examples/models/minimax_h3/reference-manifest.json create mode 100644 examples/models/minimax_h3/reference.sbatch create mode 100644 rl_engine/validation/models/h3_full_manifest.json create mode 100644 rl_engine/validation/models/h3_full_model.py create mode 100644 tests/models/minimax_h3/test_h3_full_model.py create mode 100644 tools/validation/models/h3_candidate_prepare.py create mode 100644 tools/validation/models/h3_full_compare.py create mode 100644 tools/validation/models/h3_full_ensemble_rule.py create mode 100644 tools/validation/models/h3_full_model.py create mode 100644 tools/validation/models/h3_full_provider_spread.py diff --git a/examples/models/minimax_h3/README.md b/examples/models/minimax_h3/README.md new file mode 100644 index 000000000..e39ca714f --- /dev/null +++ b/examples/models/minimax_h3/README.md @@ -0,0 +1,105 @@ +# Full H3 qualification (in progress) + +The existing block-0 evidence remains a single-block result. This extension has +not yet qualified the complete model or generation quality. + +`tools.validation.models.h3_full_model` runs the real 50 distinct transformer +layers, two text-refiner layers and input/conditioning/output components. Each +candidate block binds its own parameters to the existing graph. Row/interim/ +reference labels are retained; boundary components stay provider/reference. +The runner checks every parameter/input gradient, checkpoints activations, +offloads saved tensors and persists/reclaims completed parameter gradients. + +The `.sbatch` scripts request one GPU node and leave the partition, account and +job name to the submitting site, for example +`sbatch --partition= --account= reference.sbatch ...`. + +`reference.sbatch` first runs independent FP32 and BF16 provider passes, then freezes per-tensor error +bounds using the committed formula. It does not run a candidate. Candidate +acceptance compares immutable reference artifacts and records the first failing +tensor without changing bounds. Existing block-0 golden tolerances are unchanged. + +Generation uses the pinned public reference image identified in +`generation-inputs.json`, the accompanying prompt and the three seeds/configuration +in `generation-contract.json`. The preflight is explicitly lower resolution. +Full generation keeps the official 124 frames, 1344×768, 24 fps and joint audio. +Every initial packed input and predicted velocity is retained. The comparator +checks identical initial noise, video/audio shape and quality gates and reports +PSNR/SSIM, temporal RMSE and audio differences. Query chunking bounds candidate +attention workspace while retaining all keys for each query; it is inference only. + +An independent reference run from `2681cadb` completed the FP32/BF16 full-model +passes (638 parameter gradients each), three-seed generation and repeat comparisons. +`reference-manifest.json` anchors its 1,635 evidence files; `full-model-thresholds.json` +contains the unchanged 764 tensor bounds. These provider results do not establish +candidate acceptance. + +**Acceptance rule (2026-10-14).** The single-configuration bounds in +`full-model-thresholds.json` (1.5 × one provider pass + floor) accept only the provider +configuration they were frozen from: provider passes that differ only in SDPA backend +(flash, memory-efficient) or BLAS library (cuBLASLt) fail 36–66 of them. The first +candidate failed 54, all on single-element `max_abs` except two `relative_l2`. Acceptance +is therefore `full-model-ensemble-rule.json`, calibrated from four distinct provider +configurations with no candidate data (`tools/validation/models/h3_full_ensemble_rule.py`): +every tensor's `relative_l2` must be at most k × the largest provider error + floor, +with k = 2.6 the smallest factor at which every provider configuration held out passes +against the others; `max_abs` is reported but not gated, because held-out providers +exceed the others' maximum by up to 4.1× on it. A full-model node-swap diagnostic +(`--provider-nodes`) found no single node responsible for the candidate's high +tensors: swapping any group moves which tensors land high rather than removing them. +The single-configuration file is kept and still compared, as a record. + +**Regenerated reference (2026-10-16).** The artifacts of the first reference run (about 200 GB) +were deleted to save disk space, so the reference was regenerated with the original commands at +the same source commit (`2681cadb`). 1,614 of the 1,635 pinned files are byte-identical to the +first run's: every parameter gradient, both `tensors.pt` and `gradients.json`, every initial +input, every per-step velocity file and the `video.npy` / `audio.npy` that the comparison +reads. The 21 that differ are the eight `report.json` files (timings, job metadata), the +thresholds file that embeds their hashes, and the twelve `video.mp4` / `audio.wav` containers, +which also differ between two runs of one job (the two WAVs differ in two header bytes). +The calibrated bounds are identical to the first run's, and the ensemble rule re-derived from +the regenerated reference (six provider runs compared with it) is identical in every numeric +field. `reference-manifest.json`, `full-model-thresholds.json` and `full-model-ensemble-rule.json` +are re-pinned to the regenerated run; the first run's SHA-256 were: manifest `1640d5c7fccb`, +thresholds `3064c7eb8924`, rule `2e2a4ac58471` and `fp32-reference/report.json` `4fd74feb00ad` +(first 12 hex digits). + +`candidate.sbatch` builds its own H3 CUDA extension, records its identity, verifies +all reference/checkpoint hashes, and runs block-0 golden tests without skips before +full-model forward/backward. It requires calls to every candidate graph node in +every block and rejects CUDA backend fallback. The comparator orders first drift +by actual forward/backward execution. Only a passing numerical comparison permits +the generation preflight (the end of `candidate.sbatch`) and full generation. + +Full candidate generation is its own job. With the candidate attention it takes +about 20 hours per seed (about 1500 s per sampling step), longer than the three +seeds fit in one 24-hour job, and a preemptible partition can evict a job every few +hours. `generation.sbatch` therefore runs one seed per array task (`--array=0-2`) +and `h3_full_generation --resume` continues after an eviction: every saved step is +replayed from disk instead of recomputed and each replayed call's inputs are checked +against their saved digest (`h3_generation_resume.py`). Recomputing the last saved step +and requiring a bitwise match (`--verify-last-step`) costs a full step per restart, which +on a heavily preempted partition is most of an attempt, so it is opt-in. A run is tied to +the git trees of the model paths, not to a commit (`h3_generation_adopt.py` moves an older +commit-based run over, only if those paths are identical). The task refuses +to start unless the gates of `candidate.sbatch` passed (`` argument) at a commit whose +model paths are byte-identical to this checkout's. After all seeds finish, `generation-compare.sbatch` +runs the quality comparison. Each GPU job uses one B200 and about 512 GB of host +memory. + +Status of the full-model qualification (evidence in `reports/experiments/h3-full-model-b200/`): +numerical acceptance (50-block forward/backward under the provider-ensemble rule) has +passed; three-seed candidate generation has run, and its quality comparison has **not** +passed the bounds fixed before generation (seed 44 passes, seeds 42 and 43 do not). The +same bounds are also failed by the provider itself when one attention backend is forced +(flash and memory-efficient: 0 of 3 seeds each), and the first denoising step differs +from the provider's by the same amount (about 0.02 relative L2 on video) for the +candidate as for those provider variants, so the bounds, which had only been checked +against a bitwise-identical provider repeat, were never calibrated for a change of kernel. +They are unchanged here. No passed candidate result is claimed by CPU tests. + +The generation environment must provide torchvision matching PyTorch (the B200 +reference environment uses torch 2.9.1+cu130 and torchvision 0.24.1+cu130). +Qwen3VLProcessor needs its video processor even for the fixed image conditioning. +The runner loads this processor locally before model weights and rejects missing +required pipeline components instead of proceeding with a partially loaded pipeline. diff --git a/examples/models/minimax_h3/candidate.sbatch b/examples/models/minimax_h3/candidate.sbatch new file mode 100644 index 000000000..d88d5fe9b --- /dev/null +++ b/examples/models/minimax_h3/candidate.sbatch @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Numerical gates and the generation preflight of the full-model candidate. Full generation is +# generation.sbatch (one task per seed), allowed only after this job passed for the same commit. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +reference_root=$4 +result_root=$5 +reference_image=$6 +provider_source=$7 +extra_dependencies=$8 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 MAX_JOBS=8 +export PATH="$CUDA_HOME/bin:$PATH" +export TORCH_KERNEL_CACHE_PATH="$result_dir/torch-cache" TRITON_CACHE_DIR="$result_dir/triton-cache" +mkdir -p "$TORCH_KERNEL_CACHE_PATH" "$TRITON_CACHE_DIR" +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +"$python_bin" -m pip list --path "$extra_dependencies" --format=json > "$result_dir/extra-dependencies.json" +"$python_bin" setup.py build_ext --inplace > "$result_dir/build.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_candidate_prepare --checkpoint "$checkpoint" \ + --reference-root "$reference_root" --output "$result_dir/preflight" > "$result_dir/preflight.log" 2>&1 +export RL_KERNEL_H3_WEIGHTS="$result_dir/preflight/weights" +"$python_bin" -m pytest -q tests/models/minimax_h3/test_h3_one_block.py \ + --junitxml="$result_dir/block0.xml" > "$result_dir/block0.log" 2>&1 +"$python_bin" -c 'import sys,xml.etree.ElementTree as E; r=E.parse(sys.argv[1]); assert all(int(s.get("skipped",0)) == 0 for s in r.iter("testsuite")), "block0 skipped"' "$result_dir/block0.xml" +"$python_bin" -m tools.validation.models.h3_full_model --checkpoint "$checkpoint" \ + --manifest rl_engine/validation/models/h3_full_manifest.json --mode candidate \ + --thresholds examples/models/minimax_h3/full-model-thresholds.json \ + --output "$result_dir/candidate" > "$result_dir/candidate.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_compare --reference "$reference_root/fp32-reference" \ + --other "$result_dir/candidate" --thresholds examples/models/minimax_h3/full-model-thresholds.json \ + --reference-root "$reference_root" --reference-manifest examples/models/minimax_h3/reference-manifest.json \ + --output "$result_dir/full-model-comparison.json" > "$result_dir/comparison.log" 2>&1 || true +# The single-configuration bounds above are kept as a record; acceptance is the +# provider-ensemble rule (README), which exits non-zero on any failure. +"$python_bin" -m tools.validation.models.h3_full_ensemble_rule check \ + --rule examples/models/minimax_h3/full-model-ensemble-rule.json \ + --candidate "$result_dir/full-model-comparison.json" \ + --output "$result_dir/full-model-acceptance.json" > "$result_dir/acceptance.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --preflight --output "$result_dir/generation-preflight" \ + > "$result_dir/generation-preflight.log" 2>&1 diff --git a/examples/models/minimax_h3/full-model-ensemble-rule.json b/examples/models/minimax_h3/full-model-ensemble-rule.json new file mode 100644 index 000000000..ee4a86d3c --- /dev/null +++ b/examples/models/minimax_h3/full-model-ensemble-rule.json @@ -0,0 +1,2346 @@ +{ + "status": "frozen_ensemble_rule", + "rule": "relative_l2 <= k * max over provider configurations + floor; max_abs reported", + "factor_k": 2.6, + "leave_one_out_worst_ratio": { + "blas-cublaslt": 2.20842743779203, + "default": 2.5483393774198175, + "sdpa-efficient": 1.5394680908162333, + "sdpa-flash": 1.6790097485497042 + }, + "providers": [ + "blas-cublaslt", + "default", + "sdpa-efficient", + "sdpa-flash" + ], + "duplicates_dropped": [ + "reference-run", + "sdpa-cudnn" + ], + "sources": { + "blas-cublaslt": { + "sha256": "8aa7346ef2a6f9774513ecbafedb2974bad9ad62599490334a3d250024fce379", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "default": { + "sha256": "0e66bc8f799ddcf4a61acbf265fb3ef024131d2cbfd0c8ef6946d5441eeba4e1", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "reference-run": { + "sha256": "19f2e39b8280755ddf7f4a553f147f0216b756d50ba690e84a9ca09684ab21fe", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-cudnn": { + "sha256": "adc1509415f6f21aee49ec9d1b3899de62671e087f859f23ffed9db318afdda9", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-efficient": { + "sha256": "f16950d3ef7913bce1a2461f777c76f57ff82a52dd3023e1a40ae15c73c98280", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + }, + "sdpa-flash": { + "sha256": "42a975d807dfe6ecec7ce81e8e4b5caadb04c9bf9f74e0aa397c76e11070531a", + "reference": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d" + } + }, + "reference_report_sha256": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "ensemble_relative_l2": { + "audio": 0.02110560096918055, + "input_gradients.audio_hidden_states": 0.20277222521992394, + "input_gradients.encoder_hidden_states": 0.31350433311250625, + "input_gradients.hidden_states": 0.21767735582914077, + "nodes.audio_proj_in": 0.0, + "nodes.audio_proj_in.grad": 0.21529843009419838, + "nodes.audio_proj_out": 0.021090187694741094, + "nodes.audio_proj_out.grad": 0.02110560096918055, + "nodes.block.0": 0.007637028339181236, + "nodes.block.0.grad": 0.8528404577370075, + "nodes.block.1": 0.007167718706812642, + "nodes.block.1.grad": 0.83313024074139, + "nodes.block.10": 0.009470979132713336, + "nodes.block.10.grad": 0.9145230710597959, + "nodes.block.11": 0.009609765952726198, + "nodes.block.11.grad": 0.9151679613375059, + "nodes.block.12": 0.009811757785775614, + "nodes.block.12.grad": 0.9458007065819563, + "nodes.block.13": 0.010102301900968155, + "nodes.block.13.grad": 0.9579435857285441, + "nodes.block.14": 0.01029801834941549, + "nodes.block.14.grad": 0.9632512267797664, + "nodes.block.15": 0.010673081941103698, + "nodes.block.15.grad": 0.969776464301216, + "nodes.block.16": 0.010944323125835815, + "nodes.block.16.grad": 0.9900424517939425, + "nodes.block.17": 0.01116479582191621, + "nodes.block.17.grad": 0.9971977671738045, + "nodes.block.18": 0.012193206051473853, + "nodes.block.18.grad": 0.9965371368019056, + "nodes.block.19": 0.013078463243021731, + "nodes.block.19.grad": 1.0162687635352703, + "nodes.block.2": 0.007404503527807438, + "nodes.block.2.grad": 0.9834410772840688, + "nodes.block.20": 0.01347476355078683, + "nodes.block.20.grad": 1.0190441077729735, + "nodes.block.21": 0.013788720833893198, + "nodes.block.21.grad": 1.0289967632293302, + "nodes.block.22": 0.013559447098486279, + "nodes.block.22.grad": 1.0586827213971899, + "nodes.block.23": 0.012802505292154518, + "nodes.block.23.grad": 1.0800576612954356, + "nodes.block.24": 0.010981097081068994, + "nodes.block.24.grad": 1.1046818558335174, + "nodes.block.25": 0.00993079603502247, + "nodes.block.25.grad": 1.1126083625809122, + "nodes.block.26": 0.009286344548814976, + "nodes.block.26.grad": 1.113991113025653, + "nodes.block.27": 0.01933541921659969, + "nodes.block.27.grad": 1.1180943367539482, + "nodes.block.28": 0.01645995510288913, + "nodes.block.28.grad": 1.1277243036800457, + "nodes.block.29": 0.016225393322503424, + "nodes.block.29.grad": 1.0763523804698858, + "nodes.block.3": 0.008075193149480336, + "nodes.block.3.grad": 0.9060932909313238, + "nodes.block.30": 0.015821176889678927, + "nodes.block.30.grad": 1.0156486154560846, + "nodes.block.31": 0.01627388190675362, + "nodes.block.31.grad": 1.0538874207349838, + "nodes.block.32": 0.015005661062321091, + "nodes.block.32.grad": 1.0038473074352139, + "nodes.block.33": 0.015151115277153682, + "nodes.block.33.grad": 1.001759230958048, + "nodes.block.34": 0.015260938453177957, + "nodes.block.34.grad": 1.0882704655567403, + "nodes.block.35": 0.01361960971320334, + "nodes.block.35.grad": 1.1198217711779717, + "nodes.block.36": 0.012297703371840298, + "nodes.block.36.grad": 1.0856430212191899, + "nodes.block.37": 0.011772510842582027, + "nodes.block.37.grad": 1.0654626726976346, + "nodes.block.38": 0.012050407362486405, + "nodes.block.38.grad": 1.0858962868337536, + "nodes.block.39": 0.011985727655582418, + "nodes.block.39.grad": 1.095775280759746, + "nodes.block.4": 0.00868913976043236, + "nodes.block.4.grad": 0.924952776102672, + "nodes.block.40": 0.011512896942200267, + "nodes.block.40.grad": 1.1018309295987487, + "nodes.block.41": 0.011611925677531472, + "nodes.block.41.grad": 1.1067364188449467, + "nodes.block.42": 0.009726727784226567, + "nodes.block.42.grad": 1.1876454639541187, + "nodes.block.43": 0.007361845810865862, + "nodes.block.43.grad": 1.174099900135046, + "nodes.block.44": 0.0071179542127801695, + "nodes.block.44.grad": 1.1571739376426546, + "nodes.block.45": 0.0077513342359457945, + "nodes.block.45.grad": 1.1530878877550301, + "nodes.block.46": 0.007713059002009989, + "nodes.block.46.grad": 1.1393752324442692, + "nodes.block.47": 0.00577020720572875, + "nodes.block.47.grad": 1.141632154523112, + "nodes.block.48": 0.006725423713444281, + "nodes.block.48.grad": 1.1713012568585635, + "nodes.block.49": 0.007460893197348834, + "nodes.block.49.grad": 0.02028429422458945, + "nodes.block.5": 0.008896497807929678, + "nodes.block.5.grad": 0.8714454280768207, + "nodes.block.6": 0.008949304485577259, + "nodes.block.6.grad": 0.9109083971517853, + "nodes.block.7": 0.008971055142252412, + "nodes.block.7.grad": 0.905662113323111, + "nodes.block.8": 0.009235740680008939, + "nodes.block.8.grad": 0.9036578329629903, + "nodes.block.9": 0.009412928731248495, + "nodes.block.9.grad": 0.9007381090490874, + "nodes.context_embedder": 0.002384571186068664, + "nodes.context_embedder.grad": 0.3146938130796209, + "nodes.norm_out": 0.01728418962285463, + "nodes.norm_out.grad": 0.020587306859649657, + "nodes.proj_in": 0.0, + "nodes.proj_in.grad": 0.23277925686464554, + "nodes.proj_out": 0.018702757639433267, + "nodes.proj_out.grad": 0.006463551732884611, + "nodes.refiner.0": 0.0052852776939599985, + "nodes.refiner.0.grad": 0.3185379392032242, + "nodes.refiner.1": 0.005989274309705795, + "nodes.refiner.1.grad": 0.31814344423924806, + "nodes.time_embedder": 0.0, + "nodes.time_embedder.grad": 0.5162050269731083, + "nodes.time_proj": 0.0, + "nodes.token_refiner": 0.006084313328960681, + "nodes.token_refiner.grad": 0.3160029449176521, + "parameter.audio_proj_in.bias": 0.1874771619328801, + "parameter.audio_proj_in.weight": 0.21519850916146643, + "parameter.audio_proj_out.bias": 0.013905871137514934, + "parameter.audio_proj_out.weight": 0.017141809562141865, + "parameter.context_embedder.bias": 0.34122486124431056, + "parameter.context_embedder.weight": 0.31444085786700243, + "parameter.norm_out.linear.bias": 0.015900100954407113, + "parameter.norm_out.linear.weight": 0.01586639338847783, + "parameter.norm_out.norm.weight": 0.011882855956542197, + "parameter.proj_in.bias": 0.2177238654494255, + "parameter.proj_in.weight": 0.23583015398086374, + "parameter.proj_out.bias": 0.005525584430620126, + "parameter.proj_out.weight": 0.010896080050018181, + "parameter.time_embedder.linear_1.bias": 0.03383056500771813, + "parameter.time_embedder.linear_1.weight": 0.03108368457958572, + "parameter.time_embedder.linear_2.bias": 0.49613847877224526, + "parameter.time_embedder.linear_2.weight": 0.14648554956113208, + "parameter.token_refiner.final_norm.weight": 0.32675350633917283, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 0.26645555754500405, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 0.26593696921549076, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 0.3163188188666285, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 0.32687503983930444, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 0.3045241605999762, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 0.3270055880199932, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 0.3219667791477918, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 0.32660888280066347, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 0.32353553759927595, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 0.3315986685912315, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 0.43792017421068, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 0.438433276061005, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 0.35206328838048234, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 0.3352967680856234, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 0.3583973153812709, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 0.3410664838137415, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 0.36494682444311155, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 0.335625810567975, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 0.26516590237840665, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 0.36832092767828667, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 0.31899451766122605, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 0.17867782631623222, + "parameter.transformer_blocks.0.attn.norm_k.weight": 0.28960468817995333, + "parameter.transformer_blocks.0.attn.norm_q.weight": 0.31245708376305836, + "parameter.transformer_blocks.0.attn.to_k.weight": 0.22809571887143615, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 0.1752762846442529, + "parameter.transformer_blocks.0.attn.to_q.weight": 0.2325037713076897, + "parameter.transformer_blocks.0.attn.to_v.weight": 0.17366440926942864, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 0.18026215904611417, + "parameter.transformer_blocks.0.ff.net.2.weight": 0.17771809513331524, + "parameter.transformer_blocks.0.norm1.weight": 0.21693076268469982, + "parameter.transformer_blocks.0.norm2.weight": 0.17937090662880234, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 0.18509642602034504, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 0.1774306170524224, + "parameter.transformer_blocks.1.attn.norm_k.weight": 0.6107248821324249, + "parameter.transformer_blocks.1.attn.norm_q.weight": 0.6443177983145605, + "parameter.transformer_blocks.1.attn.to_k.weight": 0.6697545238218464, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 0.5362465973157898, + "parameter.transformer_blocks.1.attn.to_q.weight": 0.5868832674821828, + "parameter.transformer_blocks.1.attn.to_v.weight": 0.578016937573896, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 0.18332852162379282, + "parameter.transformer_blocks.1.ff.net.2.weight": 0.17509047422265817, + "parameter.transformer_blocks.1.norm1.weight": 0.6435346837382664, + "parameter.transformer_blocks.1.norm2.weight": 0.17115076008546634, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 0.17941840962409109, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 0.1623874068552092, + "parameter.transformer_blocks.10.attn.norm_k.weight": 0.6152553123197887, + "parameter.transformer_blocks.10.attn.norm_q.weight": 0.6176182711966747, + "parameter.transformer_blocks.10.attn.to_k.weight": 0.5929514096870396, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 0.506602926058383, + "parameter.transformer_blocks.10.attn.to_q.weight": 0.5846389710107828, + "parameter.transformer_blocks.10.attn.to_v.weight": 0.5517895970800015, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 0.17987433666892472, + "parameter.transformer_blocks.10.ff.net.2.weight": 0.15051711004439988, + "parameter.transformer_blocks.10.norm1.weight": 0.6431711828853883, + "parameter.transformer_blocks.10.norm2.weight": 0.22894892387570479, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 0.16100573162632656, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 0.15062333205170744, + "parameter.transformer_blocks.11.attn.norm_k.weight": 0.7059893490695992, + "parameter.transformer_blocks.11.attn.norm_q.weight": 0.736105549346304, + "parameter.transformer_blocks.11.attn.to_k.weight": 0.6789991211550914, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 0.49334014494012723, + "parameter.transformer_blocks.11.attn.to_q.weight": 0.6664320046003854, + "parameter.transformer_blocks.11.attn.to_v.weight": 0.5363533879985244, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 0.15950074480074164, + "parameter.transformer_blocks.11.ff.net.2.weight": 0.1427400350951154, + "parameter.transformer_blocks.11.norm1.weight": 0.6165197387097995, + "parameter.transformer_blocks.11.norm2.weight": 0.1743280989426046, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 0.16669251305544974, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 0.15450476555580314, + "parameter.transformer_blocks.12.attn.norm_k.weight": 0.6707093903946136, + "parameter.transformer_blocks.12.attn.norm_q.weight": 0.6762641775318454, + "parameter.transformer_blocks.12.attn.to_k.weight": 0.5035471662679206, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 0.43496434592746525, + "parameter.transformer_blocks.12.attn.to_q.weight": 0.5499982670952236, + "parameter.transformer_blocks.12.attn.to_v.weight": 0.48818185711813206, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 0.1722070930596971, + "parameter.transformer_blocks.12.ff.net.2.weight": 0.1460392638774976, + "parameter.transformer_blocks.12.norm1.weight": 0.5046925581248843, + "parameter.transformer_blocks.12.norm2.weight": 0.17744799383842155, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 0.1696187851756362, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 0.15451788367374297, + "parameter.transformer_blocks.13.attn.norm_k.weight": 0.6826694775939346, + "parameter.transformer_blocks.13.attn.norm_q.weight": 0.6739919946570334, + "parameter.transformer_blocks.13.attn.to_k.weight": 0.616337247011177, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 0.6185542328063899, + "parameter.transformer_blocks.13.attn.to_q.weight": 0.613351106050686, + "parameter.transformer_blocks.13.attn.to_v.weight": 0.6313333975892937, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 0.168929578806616, + "parameter.transformer_blocks.13.ff.net.2.weight": 0.14457795218884922, + "parameter.transformer_blocks.13.norm1.weight": 0.655974573599716, + "parameter.transformer_blocks.13.norm2.weight": 0.17869929032611745, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 0.17394482005841622, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 0.1590707776637609, + "parameter.transformer_blocks.14.attn.norm_k.weight": 0.5634564709394162, + "parameter.transformer_blocks.14.attn.norm_q.weight": 0.6043468403497236, + "parameter.transformer_blocks.14.attn.to_k.weight": 0.5561731527009006, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 0.49695488228434914, + "parameter.transformer_blocks.14.attn.to_q.weight": 0.5806546639369049, + "parameter.transformer_blocks.14.attn.to_v.weight": 0.5676983678537382, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 0.1637974398672475, + "parameter.transformer_blocks.14.ff.net.2.weight": 0.14413908284561155, + "parameter.transformer_blocks.14.norm1.weight": 0.6548673998875053, + "parameter.transformer_blocks.14.norm2.weight": 0.1932632933697487, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 0.16735203350888342, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 0.15188667677679074, + "parameter.transformer_blocks.15.attn.norm_k.weight": 0.6675522437962607, + "parameter.transformer_blocks.15.attn.norm_q.weight": 0.69642257060271, + "parameter.transformer_blocks.15.attn.to_k.weight": 0.6497632709849944, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 0.4787558426271283, + "parameter.transformer_blocks.15.attn.to_q.weight": 0.6026810255486286, + "parameter.transformer_blocks.15.attn.to_v.weight": 0.5504704747572478, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 0.162270031903906, + "parameter.transformer_blocks.15.ff.net.2.weight": 0.1453468746058194, + "parameter.transformer_blocks.15.norm1.weight": 0.6601224321425502, + "parameter.transformer_blocks.15.norm2.weight": 0.20789091452427774, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 0.17192920023271227, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 0.15172464758144125, + "parameter.transformer_blocks.16.attn.norm_k.weight": 0.5411303219731743, + "parameter.transformer_blocks.16.attn.norm_q.weight": 0.5627860173783555, + "parameter.transformer_blocks.16.attn.to_k.weight": 0.5147295853848322, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.37626756060993605, + "parameter.transformer_blocks.16.attn.to_q.weight": 0.5211155316562834, + "parameter.transformer_blocks.16.attn.to_v.weight": 0.4601833083803108, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 0.1611149175734667, + "parameter.transformer_blocks.16.ff.net.2.weight": 0.13647756377949197, + "parameter.transformer_blocks.16.norm1.weight": 0.5305065646401627, + "parameter.transformer_blocks.16.norm2.weight": 0.16299974442833884, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 0.15750520234871465, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 0.1442080584059882, + "parameter.transformer_blocks.17.attn.norm_k.weight": 0.5064534732130231, + "parameter.transformer_blocks.17.attn.norm_q.weight": 0.5585420282834562, + "parameter.transformer_blocks.17.attn.to_k.weight": 0.5945243422467011, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 0.5920735852348725, + "parameter.transformer_blocks.17.attn.to_q.weight": 0.5748802118736813, + "parameter.transformer_blocks.17.attn.to_v.weight": 0.6715329239199194, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 0.15044139611347626, + "parameter.transformer_blocks.17.ff.net.2.weight": 0.13287263448993833, + "parameter.transformer_blocks.17.norm1.weight": 0.6910566752579711, + "parameter.transformer_blocks.17.norm2.weight": 0.15128633761209442, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 0.1434070672190233, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 0.13647082108089864, + "parameter.transformer_blocks.18.attn.norm_k.weight": 0.7149893895321376, + "parameter.transformer_blocks.18.attn.norm_q.weight": 0.6497792945045617, + "parameter.transformer_blocks.18.attn.to_k.weight": 0.6257317399231312, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 0.4735823683215769, + "parameter.transformer_blocks.18.attn.to_q.weight": 0.6227941989918219, + "parameter.transformer_blocks.18.attn.to_v.weight": 0.5951698295070317, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 0.141487463750178, + "parameter.transformer_blocks.18.ff.net.2.weight": 0.12124054956470803, + "parameter.transformer_blocks.18.norm1.weight": 0.6196138920784167, + "parameter.transformer_blocks.18.norm2.weight": 0.14495111923251633, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 0.1414025275323694, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 0.13174059786195905, + "parameter.transformer_blocks.19.attn.norm_k.weight": 0.6492313403519856, + "parameter.transformer_blocks.19.attn.norm_q.weight": 0.6377195588380845, + "parameter.transformer_blocks.19.attn.to_k.weight": 0.6360903900367645, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 0.4765003925776794, + "parameter.transformer_blocks.19.attn.to_q.weight": 0.629942788876498, + "parameter.transformer_blocks.19.attn.to_v.weight": 0.5691853563022023, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 0.12873415655769765, + "parameter.transformer_blocks.19.ff.net.2.weight": 0.10139029534587113, + "parameter.transformer_blocks.19.norm1.weight": 0.6353850059551215, + "parameter.transformer_blocks.19.norm2.weight": 0.1508616045251074, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 0.17733209220723198, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 0.1706727714082375, + "parameter.transformer_blocks.2.attn.norm_k.weight": 0.5551056072295623, + "parameter.transformer_blocks.2.attn.norm_q.weight": 0.5369920381607626, + "parameter.transformer_blocks.2.attn.to_k.weight": 0.5302745024408893, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 0.4763106748686571, + "parameter.transformer_blocks.2.attn.to_q.weight": 0.4913420936131699, + "parameter.transformer_blocks.2.attn.to_v.weight": 0.6139169897083216, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 0.17577450548046156, + "parameter.transformer_blocks.2.ff.net.2.weight": 0.16919691639079415, + "parameter.transformer_blocks.2.norm1.weight": 0.5438326294406551, + "parameter.transformer_blocks.2.norm2.weight": 0.17908548259702073, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 0.15516000757003306, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 0.13473269748319175, + "parameter.transformer_blocks.20.attn.norm_k.weight": 0.8717416419892866, + "parameter.transformer_blocks.20.attn.norm_q.weight": 0.9198161269361697, + "parameter.transformer_blocks.20.attn.to_k.weight": 0.8572434742105496, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 0.3892794559046116, + "parameter.transformer_blocks.20.attn.to_q.weight": 0.791981844639221, + "parameter.transformer_blocks.20.attn.to_v.weight": 0.478350420897483, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 0.14013424458483464, + "parameter.transformer_blocks.20.ff.net.2.weight": 0.11631561270226247, + "parameter.transformer_blocks.20.norm1.weight": 0.6788613775770197, + "parameter.transformer_blocks.20.norm2.weight": 0.13949718891009422, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 0.1482196449124827, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 0.1300058383040362, + "parameter.transformer_blocks.21.attn.norm_k.weight": 0.7481017080404762, + "parameter.transformer_blocks.21.attn.norm_q.weight": 0.7592458857742853, + "parameter.transformer_blocks.21.attn.to_k.weight": 0.7171547011740366, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 0.5013809935656205, + "parameter.transformer_blocks.21.attn.to_q.weight": 0.676030423741261, + "parameter.transformer_blocks.21.attn.to_v.weight": 0.6191617303664513, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 0.1420040837772908, + "parameter.transformer_blocks.21.ff.net.2.weight": 0.1180330600974897, + "parameter.transformer_blocks.21.norm1.weight": 0.640965339782901, + "parameter.transformer_blocks.21.norm2.weight": 0.1434696448275675, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 0.14936310971176525, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 0.132431273579337, + "parameter.transformer_blocks.22.attn.norm_k.weight": 0.6627742980495887, + "parameter.transformer_blocks.22.attn.norm_q.weight": 0.6658411201577026, + "parameter.transformer_blocks.22.attn.to_k.weight": 0.6216034582065753, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 0.6723760585981233, + "parameter.transformer_blocks.22.attn.to_q.weight": 0.6059599015397963, + "parameter.transformer_blocks.22.attn.to_v.weight": 0.772219315272921, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 0.1426257615377658, + "parameter.transformer_blocks.22.ff.net.2.weight": 0.11360215573011313, + "parameter.transformer_blocks.22.norm1.weight": 0.6697341123206519, + "parameter.transformer_blocks.22.norm2.weight": 0.14625720523831526, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 0.13884457252088103, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 0.12478032787063244, + "parameter.transformer_blocks.23.attn.norm_k.weight": 0.6814540343533161, + "parameter.transformer_blocks.23.attn.norm_q.weight": 0.684038690122064, + "parameter.transformer_blocks.23.attn.to_k.weight": 0.6908737360594973, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 0.49140031547262497, + "parameter.transformer_blocks.23.attn.to_q.weight": 0.6379720333772122, + "parameter.transformer_blocks.23.attn.to_v.weight": 0.6241045529727879, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 0.13312269435414717, + "parameter.transformer_blocks.23.ff.net.2.weight": 0.1143442047795657, + "parameter.transformer_blocks.23.norm1.weight": 0.7184202157238492, + "parameter.transformer_blocks.23.norm2.weight": 0.14232311638461428, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 0.14042778599981806, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 0.12344647130721911, + "parameter.transformer_blocks.24.attn.norm_k.weight": 0.5957750471379901, + "parameter.transformer_blocks.24.attn.norm_q.weight": 0.5858490156955045, + "parameter.transformer_blocks.24.attn.to_k.weight": 0.5651326224307139, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 0.3861928792070328, + "parameter.transformer_blocks.24.attn.to_q.weight": 0.5318470735889794, + "parameter.transformer_blocks.24.attn.to_v.weight": 0.5294161705020743, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 0.13253318691665164, + "parameter.transformer_blocks.24.ff.net.2.weight": 0.08400841861636524, + "parameter.transformer_blocks.24.norm1.weight": 0.5916204321295466, + "parameter.transformer_blocks.24.norm2.weight": 0.13317067128474694, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 0.10912880779097324, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 0.1021035135517177, + "parameter.transformer_blocks.25.attn.norm_k.weight": 0.6487957220130574, + "parameter.transformer_blocks.25.attn.norm_q.weight": 0.6523428868096084, + "parameter.transformer_blocks.25.attn.to_k.weight": 0.6660812607231732, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 0.5174420452991362, + "parameter.transformer_blocks.25.attn.to_q.weight": 0.6397013876867854, + "parameter.transformer_blocks.25.attn.to_v.weight": 0.660878895059845, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 0.10628172036996353, + "parameter.transformer_blocks.25.ff.net.2.weight": 0.07574625326211328, + "parameter.transformer_blocks.25.norm1.weight": 0.5410480441062008, + "parameter.transformer_blocks.25.norm2.weight": 0.10128470936739821, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 0.1280437196985238, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 0.10888463268784916, + "parameter.transformer_blocks.26.attn.norm_k.weight": 0.6441648947168217, + "parameter.transformer_blocks.26.attn.norm_q.weight": 0.6484796721945846, + "parameter.transformer_blocks.26.attn.to_k.weight": 0.6491471682152515, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.21242931299844153, + "parameter.transformer_blocks.26.attn.to_q.weight": 0.6467800475358854, + "parameter.transformer_blocks.26.attn.to_v.weight": 0.32522441779907785, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 0.11659499821369074, + "parameter.transformer_blocks.26.ff.net.2.weight": 0.08279338541473577, + "parameter.transformer_blocks.26.norm1.weight": 0.5092815994108456, + "parameter.transformer_blocks.26.norm2.weight": 0.12215671300627264, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 0.1390507418408379, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 0.11275655533933447, + "parameter.transformer_blocks.27.attn.norm_k.weight": 0.7734440691053224, + "parameter.transformer_blocks.27.attn.norm_q.weight": 0.7953313142746422, + "parameter.transformer_blocks.27.attn.to_k.weight": 0.8133527123298837, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 0.6348785766665711, + "parameter.transformer_blocks.27.attn.to_q.weight": 0.7918463015928954, + "parameter.transformer_blocks.27.attn.to_v.weight": 0.7301223125733811, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 0.1263956489072167, + "parameter.transformer_blocks.27.ff.net.2.weight": 0.09920134398712835, + "parameter.transformer_blocks.27.norm1.weight": 0.8697917621100352, + "parameter.transformer_blocks.27.norm2.weight": 0.11998925710626789, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 0.1332879678030465, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 0.10131284826629, + "parameter.transformer_blocks.28.attn.norm_k.weight": 0.9210698288645108, + "parameter.transformer_blocks.28.attn.norm_q.weight": 0.9192904857217888, + "parameter.transformer_blocks.28.attn.to_k.weight": 0.9299335039349791, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 0.8426652469633431, + "parameter.transformer_blocks.28.attn.to_q.weight": 0.9482478007716462, + "parameter.transformer_blocks.28.attn.to_v.weight": 0.8578499632178054, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 0.12237747003347534, + "parameter.transformer_blocks.28.ff.net.2.weight": 0.08989724385841778, + "parameter.transformer_blocks.28.norm1.weight": 0.9234679202995794, + "parameter.transformer_blocks.28.norm2.weight": 0.13418668037996598, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 0.12769710071712426, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 0.09735381316720622, + "parameter.transformer_blocks.29.attn.norm_k.weight": 1.086104090000951, + "parameter.transformer_blocks.29.attn.norm_q.weight": 1.0901827753128521, + "parameter.transformer_blocks.29.attn.to_k.weight": 0.9437641228092345, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 0.8368198991044917, + "parameter.transformer_blocks.29.attn.to_q.weight": 0.951135692714322, + "parameter.transformer_blocks.29.attn.to_v.weight": 0.8757337926823521, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 0.11756568410754052, + "parameter.transformer_blocks.29.ff.net.2.weight": 0.07980869084435108, + "parameter.transformer_blocks.29.norm1.weight": 0.855527361817073, + "parameter.transformer_blocks.29.norm2.weight": 0.11290706902206783, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 0.1736136707981538, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 0.16752519826321585, + "parameter.transformer_blocks.3.attn.norm_k.weight": 0.7638730850842748, + "parameter.transformer_blocks.3.attn.norm_q.weight": 0.7465127871209892, + "parameter.transformer_blocks.3.attn.to_k.weight": 0.8239210843040061, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 0.5116743271863021, + "parameter.transformer_blocks.3.attn.to_q.weight": 0.7078958561704025, + "parameter.transformer_blocks.3.attn.to_v.weight": 0.584106712366063, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 0.1734745564727508, + "parameter.transformer_blocks.3.ff.net.2.weight": 0.1647559335686392, + "parameter.transformer_blocks.3.norm1.weight": 0.7427319268410297, + "parameter.transformer_blocks.3.norm2.weight": 0.17633769659016266, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 0.1614848620480023, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 0.09714668064766162, + "parameter.transformer_blocks.30.attn.norm_k.weight": 0.9262721419142987, + "parameter.transformer_blocks.30.attn.norm_q.weight": 0.9425131749841288, + "parameter.transformer_blocks.30.attn.to_k.weight": 1.0460792057173103, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 0.6686605200736664, + "parameter.transformer_blocks.30.attn.to_q.weight": 0.9757513962929028, + "parameter.transformer_blocks.30.attn.to_v.weight": 0.7299885484468922, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 0.12726474297280962, + "parameter.transformer_blocks.30.ff.net.2.weight": 0.08961859178187792, + "parameter.transformer_blocks.30.norm1.weight": 0.9618845175797996, + "parameter.transformer_blocks.30.norm2.weight": 0.12280940291215998, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 0.1450064713873809, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 0.10459082237081521, + "parameter.transformer_blocks.31.attn.norm_k.weight": 0.5653750877917905, + "parameter.transformer_blocks.31.attn.norm_q.weight": 0.5560376981215728, + "parameter.transformer_blocks.31.attn.to_k.weight": 0.6695545906672978, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 0.6623203254172305, + "parameter.transformer_blocks.31.attn.to_q.weight": 0.6356325720384296, + "parameter.transformer_blocks.31.attn.to_v.weight": 0.7361700881562175, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 0.14358311859226364, + "parameter.transformer_blocks.31.ff.net.2.weight": 0.09965486066250556, + "parameter.transformer_blocks.31.norm1.weight": 0.6471378646598047, + "parameter.transformer_blocks.31.norm2.weight": 0.1389725319732575, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 0.152713623242756, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 0.0902834080472529, + "parameter.transformer_blocks.32.attn.norm_k.weight": 1.1410089440835625, + "parameter.transformer_blocks.32.attn.norm_q.weight": 1.1158782942368046, + "parameter.transformer_blocks.32.attn.to_k.weight": 1.0004746487190592, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 0.889554343452924, + "parameter.transformer_blocks.32.attn.to_q.weight": 0.9849043486731729, + "parameter.transformer_blocks.32.attn.to_v.weight": 0.9132571793910841, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 0.13161416846598228, + "parameter.transformer_blocks.32.ff.net.2.weight": 0.09298675626882538, + "parameter.transformer_blocks.32.norm1.weight": 0.9349206141687386, + "parameter.transformer_blocks.32.norm2.weight": 0.13212903942288529, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 0.13979614110149272, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 0.09553208066293542, + "parameter.transformer_blocks.33.attn.norm_k.weight": 0.9155705517529243, + "parameter.transformer_blocks.33.attn.norm_q.weight": 0.9154534423974686, + "parameter.transformer_blocks.33.attn.to_k.weight": 0.8314888895641861, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 0.8453351563331559, + "parameter.transformer_blocks.33.attn.to_q.weight": 0.8044708403916625, + "parameter.transformer_blocks.33.attn.to_v.weight": 0.8490922798253322, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 0.13045382202277842, + "parameter.transformer_blocks.33.ff.net.2.weight": 0.0840068196187977, + "parameter.transformer_blocks.33.norm1.weight": 0.7570142897318359, + "parameter.transformer_blocks.33.norm2.weight": 0.14674158754893832, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 0.13477742308433482, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 0.08323276800712222, + "parameter.transformer_blocks.34.attn.norm_k.weight": 0.9791420574656059, + "parameter.transformer_blocks.34.attn.norm_q.weight": 0.9652674685185889, + "parameter.transformer_blocks.34.attn.to_k.weight": 0.8292524820974579, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 0.7930937938479834, + "parameter.transformer_blocks.34.attn.to_q.weight": 0.7589795609253875, + "parameter.transformer_blocks.34.attn.to_v.weight": 0.8549580578578369, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 0.11995605351038738, + "parameter.transformer_blocks.34.ff.net.2.weight": 0.0838291183465568, + "parameter.transformer_blocks.34.norm1.weight": 0.9403059130060623, + "parameter.transformer_blocks.34.norm2.weight": 0.12189935037058316, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 0.10189507844898495, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 0.07609252844677573, + "parameter.transformer_blocks.35.attn.norm_k.weight": 0.8454882627709907, + "parameter.transformer_blocks.35.attn.norm_q.weight": 0.8218902481390318, + "parameter.transformer_blocks.35.attn.to_k.weight": 0.8230795944844709, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 0.7995565691361847, + "parameter.transformer_blocks.35.attn.to_q.weight": 0.8884296214201013, + "parameter.transformer_blocks.35.attn.to_v.weight": 0.8713868733028493, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 0.11097041565052795, + "parameter.transformer_blocks.35.ff.net.2.weight": 0.07438589686382978, + "parameter.transformer_blocks.35.norm1.weight": 0.7699673515094488, + "parameter.transformer_blocks.35.norm2.weight": 0.130018302136186, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 0.1425909917060937, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 0.10349439246805126, + "parameter.transformer_blocks.36.attn.norm_k.weight": 0.9145770498663861, + "parameter.transformer_blocks.36.attn.norm_q.weight": 0.9484815620003966, + "parameter.transformer_blocks.36.attn.to_k.weight": 1.1298876570819247, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 0.5106671007649235, + "parameter.transformer_blocks.36.attn.to_q.weight": 1.1055395676246789, + "parameter.transformer_blocks.36.attn.to_v.weight": 0.5525522622895881, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 0.13211208971682759, + "parameter.transformer_blocks.36.ff.net.2.weight": 0.07943197232542645, + "parameter.transformer_blocks.36.norm1.weight": 0.9174376768312953, + "parameter.transformer_blocks.36.norm2.weight": 0.3773931626682203, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 0.12218262090280432, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 0.09287378428596499, + "parameter.transformer_blocks.37.attn.norm_k.weight": 0.9123803015450559, + "parameter.transformer_blocks.37.attn.norm_q.weight": 0.9302709725598957, + "parameter.transformer_blocks.37.attn.to_k.weight": 1.000961027885407, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 0.7427503354176709, + "parameter.transformer_blocks.37.attn.to_q.weight": 0.9484356139329904, + "parameter.transformer_blocks.37.attn.to_v.weight": 0.7748231151004865, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 0.1467101168691995, + "parameter.transformer_blocks.37.ff.net.2.weight": 0.06782686777903887, + "parameter.transformer_blocks.37.norm1.weight": 0.8829208798747256, + "parameter.transformer_blocks.37.norm2.weight": 0.17577249861019198, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 0.1254128039706359, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 0.0690761047640983, + "parameter.transformer_blocks.38.attn.norm_k.weight": 0.8418897031323198, + "parameter.transformer_blocks.38.attn.norm_q.weight": 0.8231596183694648, + "parameter.transformer_blocks.38.attn.to_k.weight": 0.8308298748231153, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 0.826364197591171, + "parameter.transformer_blocks.38.attn.to_q.weight": 0.8747986913890288, + "parameter.transformer_blocks.38.attn.to_v.weight": 0.8249675556685392, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 0.11340375170577656, + "parameter.transformer_blocks.38.ff.net.2.weight": 0.06541392236676424, + "parameter.transformer_blocks.38.norm1.weight": 0.7646574747887432, + "parameter.transformer_blocks.38.norm2.weight": 0.10926704894255639, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 0.10287534823972645, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 0.06472699315158703, + "parameter.transformer_blocks.39.attn.norm_k.weight": 0.9668497714880203, + "parameter.transformer_blocks.39.attn.norm_q.weight": 0.9802806969412944, + "parameter.transformer_blocks.39.attn.to_k.weight": 0.9821426346333626, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 0.7643879929788783, + "parameter.transformer_blocks.39.attn.to_q.weight": 1.023908223876388, + "parameter.transformer_blocks.39.attn.to_v.weight": 0.775839462822879, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 0.09925381460058395, + "parameter.transformer_blocks.39.ff.net.2.weight": 0.06097360985984232, + "parameter.transformer_blocks.39.norm1.weight": 0.7778455401258393, + "parameter.transformer_blocks.39.norm2.weight": 0.16601035712608037, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 0.17159093153329363, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 0.1638525474502011, + "parameter.transformer_blocks.4.attn.norm_k.weight": 0.8085543754044254, + "parameter.transformer_blocks.4.attn.norm_q.weight": 0.7691899487095933, + "parameter.transformer_blocks.4.attn.to_k.weight": 0.6670061012138586, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 0.6291277678966494, + "parameter.transformer_blocks.4.attn.to_q.weight": 0.6630847996693151, + "parameter.transformer_blocks.4.attn.to_v.weight": 0.7004136365823667, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 0.17499912855333324, + "parameter.transformer_blocks.4.ff.net.2.weight": 0.15735854396481597, + "parameter.transformer_blocks.4.norm1.weight": 0.7134663928601898, + "parameter.transformer_blocks.4.norm2.weight": 0.16745667286470833, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 0.0995681843335852, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 0.059305223345764504, + "parameter.transformer_blocks.40.attn.norm_k.weight": 0.9435521245892633, + "parameter.transformer_blocks.40.attn.norm_q.weight": 0.9891083149347497, + "parameter.transformer_blocks.40.attn.to_k.weight": 0.9628497990950426, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 0.7036645759214321, + "parameter.transformer_blocks.40.attn.to_q.weight": 1.0394286869298817, + "parameter.transformer_blocks.40.attn.to_v.weight": 0.7389363281243745, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 0.10314438485333412, + "parameter.transformer_blocks.40.ff.net.2.weight": 0.05791788845953324, + "parameter.transformer_blocks.40.norm1.weight": 0.9204294948298577, + "parameter.transformer_blocks.40.norm2.weight": 0.08970946823680186, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 0.07167722151177615, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 0.05385167994476892, + "parameter.transformer_blocks.41.attn.norm_k.weight": 0.7415739676265102, + "parameter.transformer_blocks.41.attn.norm_q.weight": 0.7214113518977477, + "parameter.transformer_blocks.41.attn.to_k.weight": 0.8375918144818602, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 0.7181446985424611, + "parameter.transformer_blocks.41.attn.to_q.weight": 0.8633605626921834, + "parameter.transformer_blocks.41.attn.to_v.weight": 0.7340670197934589, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 0.08418384433323474, + "parameter.transformer_blocks.41.ff.net.2.weight": 0.055784632945377714, + "parameter.transformer_blocks.41.norm1.weight": 0.7416192510696097, + "parameter.transformer_blocks.41.norm2.weight": 0.07178941417708969, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 0.06136013817900501, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 0.04869513884491029, + "parameter.transformer_blocks.42.attn.norm_k.weight": 0.5782331610559092, + "parameter.transformer_blocks.42.attn.norm_q.weight": 0.5610783050313803, + "parameter.transformer_blocks.42.attn.to_k.weight": 0.5961723816303608, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 0.4970380312375704, + "parameter.transformer_blocks.42.attn.to_q.weight": 0.6264928024228196, + "parameter.transformer_blocks.42.attn.to_v.weight": 0.5601926606479264, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 0.08107068095989993, + "parameter.transformer_blocks.42.ff.net.2.weight": 0.052568672636339825, + "parameter.transformer_blocks.42.norm1.weight": 0.611284124711771, + "parameter.transformer_blocks.42.norm2.weight": 0.08594602357157691, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 0.08455479479594155, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 0.06298888911994822, + "parameter.transformer_blocks.43.attn.norm_k.weight": 0.8367560596212958, + "parameter.transformer_blocks.43.attn.norm_q.weight": 0.8293012138736153, + "parameter.transformer_blocks.43.attn.to_k.weight": 0.8355610570416689, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 0.5048074999452841, + "parameter.transformer_blocks.43.attn.to_q.weight": 0.8589702799195638, + "parameter.transformer_blocks.43.attn.to_v.weight": 0.5237629013042355, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 0.08180036987950778, + "parameter.transformer_blocks.43.ff.net.2.weight": 0.05492635350283063, + "parameter.transformer_blocks.43.norm1.weight": 0.625254111632119, + "parameter.transformer_blocks.43.norm2.weight": 0.08185618169640825, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 0.0747035755085786, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 0.06660390438644587, + "parameter.transformer_blocks.44.attn.norm_k.weight": 0.7054325330091511, + "parameter.transformer_blocks.44.attn.norm_q.weight": 0.7153772458747647, + "parameter.transformer_blocks.44.attn.to_k.weight": 0.8151165065352909, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 0.4896220425182629, + "parameter.transformer_blocks.44.attn.to_q.weight": 0.8899012780688226, + "parameter.transformer_blocks.44.attn.to_v.weight": 0.5503025423219586, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 0.08021322348273459, + "parameter.transformer_blocks.44.ff.net.2.weight": 0.05349598808605537, + "parameter.transformer_blocks.44.norm1.weight": 0.6243989599839498, + "parameter.transformer_blocks.44.norm2.weight": 0.11799748672207468, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 0.0497453175382903, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 0.04457923582338126, + "parameter.transformer_blocks.45.attn.norm_k.weight": 0.7810134190748194, + "parameter.transformer_blocks.45.attn.norm_q.weight": 0.7955276983624541, + "parameter.transformer_blocks.45.attn.to_k.weight": 0.7142341183040185, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.3006170712170958, + "parameter.transformer_blocks.45.attn.to_q.weight": 0.5358664533020208, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.3333600839476919, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 0.06397931304789603, + "parameter.transformer_blocks.45.ff.net.2.weight": 0.04759898484212081, + "parameter.transformer_blocks.45.norm1.weight": 0.5074765325765653, + "parameter.transformer_blocks.45.norm2.weight": 0.07059256845556657, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 0.03973402481862122, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 0.039342269700501455, + "parameter.transformer_blocks.46.attn.norm_k.weight": 0.5281689931014756, + "parameter.transformer_blocks.46.attn.norm_q.weight": 0.5023460667545797, + "parameter.transformer_blocks.46.attn.to_k.weight": 0.5710493648291428, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.20458083566418112, + "parameter.transformer_blocks.46.attn.to_q.weight": 0.6043269611697017, + "parameter.transformer_blocks.46.attn.to_v.weight": 0.21222939602038893, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 0.053632445429609044, + "parameter.transformer_blocks.46.ff.net.2.weight": 0.043544126943884376, + "parameter.transformer_blocks.46.norm1.weight": 0.20378245936065947, + "parameter.transformer_blocks.46.norm2.weight": 0.12094285973244451, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 0.05022102495701803, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 0.04967724417388883, + "parameter.transformer_blocks.47.attn.norm_k.weight": 0.40581739556096796, + "parameter.transformer_blocks.47.attn.norm_q.weight": 0.40086424817021526, + "parameter.transformer_blocks.47.attn.to_k.weight": 0.30190780592594835, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 0.1607475406502512, + "parameter.transformer_blocks.47.attn.to_q.weight": 0.32646706522384894, + "parameter.transformer_blocks.47.attn.to_v.weight": 0.17795703897959772, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 0.05300780342409638, + "parameter.transformer_blocks.47.ff.net.2.weight": 0.03804820836236104, + "parameter.transformer_blocks.47.norm1.weight": 0.2770958821336741, + "parameter.transformer_blocks.47.norm2.weight": 0.05101560480227945, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 0.03333339913016972, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 0.03236733810896702, + "parameter.transformer_blocks.48.attn.norm_k.weight": 0.19457433853619419, + "parameter.transformer_blocks.48.attn.norm_q.weight": 0.2039308465959488, + "parameter.transformer_blocks.48.attn.to_k.weight": 0.20109957601376408, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 0.06825592517444747, + "parameter.transformer_blocks.48.attn.to_q.weight": 0.23789554363993, + "parameter.transformer_blocks.48.attn.to_v.weight": 0.06767046317004026, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 0.04273345975348288, + "parameter.transformer_blocks.48.ff.net.2.weight": 0.030467613849704606, + "parameter.transformer_blocks.48.norm1.weight": 0.1027296515699527, + "parameter.transformer_blocks.48.norm2.weight": 0.036571067769884724, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 0.08135452791594873, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 0.026706518747267025, + "parameter.transformer_blocks.49.attn.norm_k.weight": 0.21441411702278998, + "parameter.transformer_blocks.49.attn.norm_q.weight": 0.18738898861526473, + "parameter.transformer_blocks.49.attn.to_k.weight": 0.2729530254337735, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 0.05813700787592316, + "parameter.transformer_blocks.49.attn.to_q.weight": 0.7771356836911435, + "parameter.transformer_blocks.49.attn.to_v.weight": 0.0805187373024994, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 0.030523150592913086, + "parameter.transformer_blocks.49.ff.net.2.weight": 0.01776529192249012, + "parameter.transformer_blocks.49.norm1.weight": 0.7904321091518686, + "parameter.transformer_blocks.49.norm2.weight": 0.03715688505236718, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 0.17641330345085307, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 0.1601695998366226, + "parameter.transformer_blocks.5.attn.norm_k.weight": 0.9930043052829239, + "parameter.transformer_blocks.5.attn.norm_q.weight": 0.9196082022003691, + "parameter.transformer_blocks.5.attn.to_k.weight": 0.9654403710082802, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 0.3941884282252417, + "parameter.transformer_blocks.5.attn.to_q.weight": 1.0187611646153703, + "parameter.transformer_blocks.5.attn.to_v.weight": 0.42468143004203623, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 0.16811237874287088, + "parameter.transformer_blocks.5.ff.net.2.weight": 0.1576003582577836, + "parameter.transformer_blocks.5.norm1.weight": 0.76708667521147, + "parameter.transformer_blocks.5.norm2.weight": 0.37211491427211013, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 0.19024775578310005, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 0.15817875829841233, + "parameter.transformer_blocks.6.attn.norm_k.weight": 0.4710895866015602, + "parameter.transformer_blocks.6.attn.norm_q.weight": 0.43767176224906923, + "parameter.transformer_blocks.6.attn.to_k.weight": 0.49508575795553283, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 0.5608999448263835, + "parameter.transformer_blocks.6.attn.to_q.weight": 0.49082194900600334, + "parameter.transformer_blocks.6.attn.to_v.weight": 0.5702360994896027, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 0.1721629763762543, + "parameter.transformer_blocks.6.ff.net.2.weight": 0.15903645990316964, + "parameter.transformer_blocks.6.norm1.weight": 0.5200554294016089, + "parameter.transformer_blocks.6.norm2.weight": 0.16227311821904442, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 0.20821190039464113, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 0.1750668113552889, + "parameter.transformer_blocks.7.attn.norm_k.weight": 0.767506548561085, + "parameter.transformer_blocks.7.attn.norm_q.weight": 0.724381651371519, + "parameter.transformer_blocks.7.attn.to_k.weight": 0.7161623747587325, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 0.6452027625697531, + "parameter.transformer_blocks.7.attn.to_q.weight": 0.7234735440770655, + "parameter.transformer_blocks.7.attn.to_v.weight": 0.6694764858944942, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 0.18568970462110013, + "parameter.transformer_blocks.7.ff.net.2.weight": 0.1577253307671931, + "parameter.transformer_blocks.7.norm1.weight": 0.7940419976295847, + "parameter.transformer_blocks.7.norm2.weight": 0.1862207062091121, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 0.177743452577233, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 0.16363789239079013, + "parameter.transformer_blocks.8.attn.norm_k.weight": 0.7775983592302588, + "parameter.transformer_blocks.8.attn.norm_q.weight": 0.8258994911211479, + "parameter.transformer_blocks.8.attn.to_k.weight": 0.796482110330491, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 0.49648199918847563, + "parameter.transformer_blocks.8.attn.to_q.weight": 0.7766663436028602, + "parameter.transformer_blocks.8.attn.to_v.weight": 0.5437818585307683, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 0.17455968335439612, + "parameter.transformer_blocks.8.ff.net.2.weight": 0.15671060212477078, + "parameter.transformer_blocks.8.norm1.weight": 0.7127991715379555, + "parameter.transformer_blocks.8.norm2.weight": 0.1771072697254998, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 0.18711840743085115, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 0.1618146031833099, + "parameter.transformer_blocks.9.attn.norm_k.weight": 0.8153320809159224, + "parameter.transformer_blocks.9.attn.norm_q.weight": 0.8138568088872633, + "parameter.transformer_blocks.9.attn.to_k.weight": 0.7583346203665146, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 0.4967006237582202, + "parameter.transformer_blocks.9.attn.to_q.weight": 0.7777404983883319, + "parameter.transformer_blocks.9.attn.to_v.weight": 0.5323776028893205, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 0.1668908492336362, + "parameter.transformer_blocks.9.ff.net.2.weight": 0.15348178168709498, + "parameter.transformer_blocks.9.norm1.weight": 0.704262611952027, + "parameter.transformer_blocks.9.norm2.weight": 0.16444916255759165, + "video": 0.006463553178251248 + }, + "ensemble_max_abs": { + "audio": 25.881103515625, + "input_gradients.audio_hidden_states": 1659.3463134765625, + "input_gradients.encoder_hidden_states": 84.62007141113281, + "input_gradients.hidden_states": 1499.3985595703125, + "nodes.audio_proj_in": 0.0, + "nodes.audio_proj_in.grad": 459.2862548828125, + "nodes.audio_proj_out": 25.881103515625, + "nodes.audio_proj_out.grad": 0.4043922424316406, + "nodes.block.0": 140320768.0, + "nodes.block.0.grad": 0.0034225070849061012, + "nodes.block.1": 143748096.0, + "nodes.block.1.grad": 0.0011117004323750734, + "nodes.block.10": 162256896.0, + "nodes.block.10.grad": 0.00033802539110183716, + "nodes.block.11": 164428800.0, + "nodes.block.11.grad": 0.000311677431454882, + "nodes.block.12": 200382464.0, + "nodes.block.12.grad": 0.00032012685551308095, + "nodes.block.13": 205912064.0, + "nodes.block.13.grad": 0.0003194731834810227, + "nodes.block.14": 181512192.0, + "nodes.block.14.grad": 0.000302475003991276, + "nodes.block.15": 229742592.0, + "nodes.block.15.grad": 0.0002795343170873821, + "nodes.block.16": 291952640.0, + "nodes.block.16.grad": 0.00028924597427248955, + "nodes.block.17": 237170688.0, + "nodes.block.17.grad": 0.0002339307829970494, + "nodes.block.18": 187920384.0, + "nodes.block.18.grad": 0.0002275904844282195, + "nodes.block.19": 216483840.0, + "nodes.block.19.grad": 0.00020624657918233424, + "nodes.block.2": 149746688.0, + "nodes.block.2.grad": 0.0010053855367004871, + "nodes.block.20": 209188864.0, + "nodes.block.20.grad": 0.00018948683282360435, + "nodes.block.21": 199888896.0, + "nodes.block.21.grad": 0.00018563814228400588, + "nodes.block.22": 248029184.0, + "nodes.block.22.grad": 0.00018029732746072114, + "nodes.block.23": 290775040.0, + "nodes.block.23.grad": 0.00017719886091072112, + "nodes.block.24": 463339520.0, + "nodes.block.24.grad": 0.0001746931520756334, + "nodes.block.25": 753213440.0, + "nodes.block.25.grad": 0.00016678865358699113, + "nodes.block.26": 1228570624.0, + "nodes.block.26.grad": 0.0001614233769942075, + "nodes.block.27": 1964654592.0, + "nodes.block.27.grad": 0.00018509774236008525, + "nodes.block.28": 928874496.0, + "nodes.block.28.grad": 0.00013470835983753204, + "nodes.block.29": 952860672.0, + "nodes.block.29.grad": 0.00011252994590904564, + "nodes.block.3": 175917056.0, + "nodes.block.3.grad": 0.0008070977637544274, + "nodes.block.30": 1398390784.0, + "nodes.block.30.grad": 8.35541941341944e-05, + "nodes.block.31": 1785257984.0, + "nodes.block.31.grad": 8.580036228522658e-05, + "nodes.block.32": 1542389760.0, + "nodes.block.32.grad": 6.19091370026581e-05, + "nodes.block.33": 1992736768.0, + "nodes.block.33.grad": 5.435024286271073e-05, + "nodes.block.34": 2048032768.0, + "nodes.block.34.grad": 5.132899968884885e-05, + "nodes.block.35": 2313650176.0, + "nodes.block.35.grad": 4.8147980123758316e-05, + "nodes.block.36": 2989260800.0, + "nodes.block.36.grad": 3.9090195059543476e-05, + "nodes.block.37": 3730833408.0, + "nodes.block.37.grad": 3.7519050238188356e-05, + "nodes.block.38": 4208328704.0, + "nodes.block.38.grad": 3.63895705959294e-05, + "nodes.block.39": 5305368576.0, + "nodes.block.39.grad": 3.337875386932865e-05, + "nodes.block.4": 149050368.0, + "nodes.block.4.grad": 0.0005667342338711023, + "nodes.block.40": 5325160448.0, + "nodes.block.40.grad": 3.524516796460375e-05, + "nodes.block.41": 5474811904.0, + "nodes.block.41.grad": 3.4466203942429274e-05, + "nodes.block.42": 5105451008.0, + "nodes.block.42.grad": 3.930190359824337e-05, + "nodes.block.43": 10013900800.0, + "nodes.block.43.grad": 3.9428363379556686e-05, + "nodes.block.44": 14327480320.0, + "nodes.block.44.grad": 4.173372872173786e-05, + "nodes.block.45": 15591800832.0, + "nodes.block.45.grad": 4.137665018788539e-05, + "nodes.block.46": 16392912896.0, + "nodes.block.46.grad": 4.1100756789091974e-05, + "nodes.block.47": 24343740416.0, + "nodes.block.47.grad": 4.002563582616858e-05, + "nodes.block.48": 19063373824.0, + "nodes.block.48.grad": 3.985674993600696e-05, + "nodes.block.49": 23810801664.0, + "nodes.block.49.grad": 4.638408768187219e-09, + "nodes.block.5": 139992064.0, + "nodes.block.5.grad": 0.0005438930820673704, + "nodes.block.6": 236392448.0, + "nodes.block.6.grad": 0.0005069520557299256, + "nodes.block.7": 190355456.0, + "nodes.block.7.grad": 0.00044696839177049696, + "nodes.block.8": 188913664.0, + "nodes.block.8.grad": 0.0004436581220943481, + "nodes.block.9": 189696000.0, + "nodes.block.9.grad": 0.0004370134265627712, + "nodes.context_embedder": 0.12600326538085938, + "nodes.context_embedder.grad": 3.906439781188965, + "nodes.norm_out": 111.63232421875, + "nodes.norm_out.grad": 1.3662357330322266, + "nodes.proj_in": 0.0, + "nodes.proj_in.grad": 177.72679138183594, + "nodes.proj_out": 9.693748474121094, + "nodes.proj_out.grad": 0.004043877124786377, + "nodes.refiner.0": 46.240234375, + "nodes.refiner.0.grad": 0.05538301542401314, + "nodes.refiner.1": 43.7119140625, + "nodes.refiner.1.grad": 0.056103259325027466, + "nodes.time_embedder": 0.0, + "nodes.time_embedder.grad": 1092.195556640625, + "nodes.time_proj": 0.0, + "nodes.token_refiner": 0.13956832885742188, + "nodes.token_refiner.grad": 29.5113582611084, + "parameter.audio_proj_in.bias": 959.9775390625, + "parameter.audio_proj_in.weight": 1504.7664794921875, + "parameter.audio_proj_out.bias": 0.6827297210693359, + "parameter.audio_proj_out.weight": 3522.9375, + "parameter.context_embedder.bias": 6.652046203613281, + "parameter.context_embedder.weight": 15.703447341918945, + "parameter.norm_out.linear.bias": 3.800384521484375, + "parameter.norm_out.linear.weight": 967.58203125, + "parameter.norm_out.norm.weight": 1044.5390625, + "parameter.proj_in.bias": 283.2672119140625, + "parameter.proj_in.weight": 538.5294189453125, + "parameter.proj_out.bias": 0.009851455688476562, + "parameter.proj_out.weight": 8.053604125976562, + "parameter.time_embedder.linear_1.bias": 588.97265625, + "parameter.time_embedder.linear_1.weight": 588.4208984375, + "parameter.time_embedder.linear_2.bias": 1201.2255859375, + "parameter.time_embedder.linear_2.weight": 7405.4541015625, + "parameter.token_refiner.final_norm.weight": 89.69309997558594, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 83.28971862792969, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 83.48487854003906, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 15.937400817871094, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 21.675384521484375, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 14.408926010131836, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 14.414284706115723, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 27.380104064941406, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 46.86830520629883, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 79.89678955078125, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 116.41938781738281, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 8.864188194274902, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 8.92668342590332, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 1.5547832250595093, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 3.674546718597412, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 2.011857509613037, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 2.641394853591919, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 31.77730941772461, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 54.34901809692383, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 45.30865478515625, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 29.382301330566406, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 1196.275390625, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 26491.515625, + "parameter.transformer_blocks.0.attn.norm_k.weight": 4221.015625, + "parameter.transformer_blocks.0.attn.norm_q.weight": 3929.76025390625, + "parameter.transformer_blocks.0.attn.to_k.weight": 2417.5927734375, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 1356.8145751953125, + "parameter.transformer_blocks.0.attn.to_q.weight": 8516.501953125, + "parameter.transformer_blocks.0.attn.to_v.weight": 747.5717163085938, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 3595.708740234375, + "parameter.transformer_blocks.0.ff.net.2.weight": 708.46337890625, + "parameter.transformer_blocks.0.norm1.weight": 2407.3369140625, + "parameter.transformer_blocks.0.norm2.weight": 7685.876953125, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 26.141883850097656, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 6241.2802734375, + "parameter.transformer_blocks.1.attn.norm_k.weight": 7.931436538696289, + "parameter.transformer_blocks.1.attn.norm_q.weight": 7.464145660400391, + "parameter.transformer_blocks.1.attn.to_k.weight": 7.599443435668945, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 3.394611358642578, + "parameter.transformer_blocks.1.attn.to_q.weight": 6.851078033447266, + "parameter.transformer_blocks.1.attn.to_v.weight": 4.701122760772705, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 846.48876953125, + "parameter.transformer_blocks.1.ff.net.2.weight": 481.535400390625, + "parameter.transformer_blocks.1.norm1.weight": 18.98525047302246, + "parameter.transformer_blocks.1.norm2.weight": 1647.931640625, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 5.168201446533203, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 1230.390625, + "parameter.transformer_blocks.10.attn.norm_k.weight": 3.2422327995300293, + "parameter.transformer_blocks.10.attn.norm_q.weight": 2.38580060005188, + "parameter.transformer_blocks.10.attn.to_k.weight": 2.1780014038085938, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 0.6007388830184937, + "parameter.transformer_blocks.10.attn.to_q.weight": 1.3468154668807983, + "parameter.transformer_blocks.10.attn.to_v.weight": 1.3952274322509766, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 643.40576171875, + "parameter.transformer_blocks.10.ff.net.2.weight": 187.44125366210938, + "parameter.transformer_blocks.10.norm1.weight": 5.480815410614014, + "parameter.transformer_blocks.10.norm2.weight": 2093.111328125, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 17.13475799560547, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 3995.078125, + "parameter.transformer_blocks.11.attn.norm_k.weight": 2.1234676837921143, + "parameter.transformer_blocks.11.attn.norm_q.weight": 2.16512393951416, + "parameter.transformer_blocks.11.attn.to_k.weight": 1.8319860696792603, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 0.587449848651886, + "parameter.transformer_blocks.11.attn.to_q.weight": 1.8382810354232788, + "parameter.transformer_blocks.11.attn.to_v.weight": 0.7373056411743164, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 535.1204833984375, + "parameter.transformer_blocks.11.ff.net.2.weight": 328.0040283203125, + "parameter.transformer_blocks.11.norm1.weight": 4.554431438446045, + "parameter.transformer_blocks.11.norm2.weight": 3220.455078125, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 5.411174774169922, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 1297.1611328125, + "parameter.transformer_blocks.12.attn.norm_k.weight": 3.307208776473999, + "parameter.transformer_blocks.12.attn.norm_q.weight": 3.3316569328308105, + "parameter.transformer_blocks.12.attn.to_k.weight": 1.5812640190124512, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 0.367980033159256, + "parameter.transformer_blocks.12.attn.to_q.weight": 2.3092703819274902, + "parameter.transformer_blocks.12.attn.to_v.weight": 0.45317959785461426, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 244.15122985839844, + "parameter.transformer_blocks.12.ff.net.2.weight": 103.07870483398438, + "parameter.transformer_blocks.12.norm1.weight": 2.0407357215881348, + "parameter.transformer_blocks.12.norm2.weight": 526.4371337890625, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 7.499958038330078, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 1779.9697265625, + "parameter.transformer_blocks.13.attn.norm_k.weight": 2.817554473876953, + "parameter.transformer_blocks.13.attn.norm_q.weight": 2.7233643531799316, + "parameter.transformer_blocks.13.attn.to_k.weight": 0.8722900152206421, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 0.5390692353248596, + "parameter.transformer_blocks.13.attn.to_q.weight": 1.7306736707687378, + "parameter.transformer_blocks.13.attn.to_v.weight": 0.88811194896698, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 151.831787109375, + "parameter.transformer_blocks.13.ff.net.2.weight": 205.3848876953125, + "parameter.transformer_blocks.13.norm1.weight": 5.726235866546631, + "parameter.transformer_blocks.13.norm2.weight": 912.68212890625, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 5.064319610595703, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 1209.17236328125, + "parameter.transformer_blocks.14.attn.norm_k.weight": 1.3951852321624756, + "parameter.transformer_blocks.14.attn.norm_q.weight": 1.3705897331237793, + "parameter.transformer_blocks.14.attn.to_k.weight": 1.9639244079589844, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 0.41162002086639404, + "parameter.transformer_blocks.14.attn.to_q.weight": 1.3147220611572266, + "parameter.transformer_blocks.14.attn.to_v.weight": 1.1094074249267578, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 241.68624877929688, + "parameter.transformer_blocks.14.ff.net.2.weight": 93.27825927734375, + "parameter.transformer_blocks.14.norm1.weight": 3.7792177200317383, + "parameter.transformer_blocks.14.norm2.weight": 551.478759765625, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 6.347064971923828, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 1483.4775390625, + "parameter.transformer_blocks.15.attn.norm_k.weight": 2.051187038421631, + "parameter.transformer_blocks.15.attn.norm_q.weight": 2.2277870178222656, + "parameter.transformer_blocks.15.attn.to_k.weight": 0.8541905879974365, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 0.7474973201751709, + "parameter.transformer_blocks.15.attn.to_q.weight": 1.0169093608856201, + "parameter.transformer_blocks.15.attn.to_v.weight": 0.6167980432510376, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 696.7670288085938, + "parameter.transformer_blocks.15.ff.net.2.weight": 140.95977783203125, + "parameter.transformer_blocks.15.norm1.weight": 3.8490424156188965, + "parameter.transformer_blocks.15.norm2.weight": 1405.282958984375, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 5.253047466278076, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 1252.7545166015625, + "parameter.transformer_blocks.16.attn.norm_k.weight": 1.5553314685821533, + "parameter.transformer_blocks.16.attn.norm_q.weight": 1.941348671913147, + "parameter.transformer_blocks.16.attn.to_k.weight": 1.1562013626098633, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.4928184151649475, + "parameter.transformer_blocks.16.attn.to_q.weight": 0.8413535356521606, + "parameter.transformer_blocks.16.attn.to_v.weight": 0.7481282949447632, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 268.7306823730469, + "parameter.transformer_blocks.16.ff.net.2.weight": 95.50224304199219, + "parameter.transformer_blocks.16.norm1.weight": 3.8275866508483887, + "parameter.transformer_blocks.16.norm2.weight": 1096.85888671875, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 9.43035888671875, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 2227.4814453125, + "parameter.transformer_blocks.17.attn.norm_k.weight": 1.7146778106689453, + "parameter.transformer_blocks.17.attn.norm_q.weight": 1.5827388763427734, + "parameter.transformer_blocks.17.attn.to_k.weight": 1.1171233654022217, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 0.5116302967071533, + "parameter.transformer_blocks.17.attn.to_q.weight": 1.061017394065857, + "parameter.transformer_blocks.17.attn.to_v.weight": 0.7643781900405884, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 312.38629150390625, + "parameter.transformer_blocks.17.ff.net.2.weight": 247.89013671875, + "parameter.transformer_blocks.17.norm1.weight": 4.868851661682129, + "parameter.transformer_blocks.17.norm2.weight": 607.411865234375, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 7.408510208129883, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 1765.047607421875, + "parameter.transformer_blocks.18.attn.norm_k.weight": 1.7119410037994385, + "parameter.transformer_blocks.18.attn.norm_q.weight": 1.1268473863601685, + "parameter.transformer_blocks.18.attn.to_k.weight": 0.8185585141181946, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 0.5142297744750977, + "parameter.transformer_blocks.18.attn.to_q.weight": 0.8801079988479614, + "parameter.transformer_blocks.18.attn.to_v.weight": 0.8254555463790894, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 580.1936645507812, + "parameter.transformer_blocks.18.ff.net.2.weight": 428.6776123046875, + "parameter.transformer_blocks.18.norm1.weight": 2.4846653938293457, + "parameter.transformer_blocks.18.norm2.weight": 1136.5263671875, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 6.499603271484375, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 1539.50634765625, + "parameter.transformer_blocks.19.attn.norm_k.weight": 2.237013578414917, + "parameter.transformer_blocks.19.attn.norm_q.weight": 2.2605738639831543, + "parameter.transformer_blocks.19.attn.to_k.weight": 1.7174041271209717, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 0.6810566186904907, + "parameter.transformer_blocks.19.attn.to_q.weight": 1.4289164543151855, + "parameter.transformer_blocks.19.attn.to_v.weight": 1.0487899780273438, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 294.0188293457031, + "parameter.transformer_blocks.19.ff.net.2.weight": 261.71728515625, + "parameter.transformer_blocks.19.norm1.weight": 6.750887870788574, + "parameter.transformer_blocks.19.norm2.weight": 1107.74365234375, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 33.29774475097656, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 8001.673828125, + "parameter.transformer_blocks.2.attn.norm_k.weight": 10.58828353881836, + "parameter.transformer_blocks.2.attn.norm_q.weight": 9.941483497619629, + "parameter.transformer_blocks.2.attn.to_k.weight": 6.533878803253174, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 2.6041641235351562, + "parameter.transformer_blocks.2.attn.to_q.weight": 5.902656078338623, + "parameter.transformer_blocks.2.attn.to_v.weight": 5.091973304748535, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 561.8043212890625, + "parameter.transformer_blocks.2.ff.net.2.weight": 381.8466796875, + "parameter.transformer_blocks.2.norm1.weight": 10.919822692871094, + "parameter.transformer_blocks.2.norm2.weight": 1716.290283203125, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 7.788921356201172, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 1854.8935546875, + "parameter.transformer_blocks.20.attn.norm_k.weight": 2.3954052925109863, + "parameter.transformer_blocks.20.attn.norm_q.weight": 2.0874135494232178, + "parameter.transformer_blocks.20.attn.to_k.weight": 1.894758939743042, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 0.5271312594413757, + "parameter.transformer_blocks.20.attn.to_q.weight": 1.0881155729293823, + "parameter.transformer_blocks.20.attn.to_v.weight": 0.8003813028335571, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 155.64996337890625, + "parameter.transformer_blocks.20.ff.net.2.weight": 190.9737548828125, + "parameter.transformer_blocks.20.norm1.weight": 5.4173431396484375, + "parameter.transformer_blocks.20.norm2.weight": 334.54669189453125, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 4.380756378173828, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 1040.6484375, + "parameter.transformer_blocks.21.attn.norm_k.weight": 2.262387990951538, + "parameter.transformer_blocks.21.attn.norm_q.weight": 2.1462185382843018, + "parameter.transformer_blocks.21.attn.to_k.weight": 1.493720531463623, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 0.5406997203826904, + "parameter.transformer_blocks.21.attn.to_q.weight": 0.84881591796875, + "parameter.transformer_blocks.21.attn.to_v.weight": 0.9486402869224548, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 208.00839233398438, + "parameter.transformer_blocks.21.ff.net.2.weight": 202.0222930908203, + "parameter.transformer_blocks.21.norm1.weight": 4.693154335021973, + "parameter.transformer_blocks.21.norm2.weight": 287.727294921875, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 5.165304183959961, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 1219.8623046875, + "parameter.transformer_blocks.22.attn.norm_k.weight": 1.7254773378372192, + "parameter.transformer_blocks.22.attn.norm_q.weight": 1.5732178688049316, + "parameter.transformer_blocks.22.attn.to_k.weight": 1.8890459537506104, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 0.524356484413147, + "parameter.transformer_blocks.22.attn.to_q.weight": 0.6948188543319702, + "parameter.transformer_blocks.22.attn.to_v.weight": 1.7596900463104248, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 253.32228088378906, + "parameter.transformer_blocks.22.ff.net.2.weight": 169.9898681640625, + "parameter.transformer_blocks.22.norm1.weight": 7.527772426605225, + "parameter.transformer_blocks.22.norm2.weight": 492.393798828125, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 5.38870906829834, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 1281.716796875, + "parameter.transformer_blocks.23.attn.norm_k.weight": 1.782730221748352, + "parameter.transformer_blocks.23.attn.norm_q.weight": 1.8004608154296875, + "parameter.transformer_blocks.23.attn.to_k.weight": 2.6039834022521973, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 0.3485468327999115, + "parameter.transformer_blocks.23.attn.to_q.weight": 0.8478580117225647, + "parameter.transformer_blocks.23.attn.to_v.weight": 1.2505159378051758, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 173.40414428710938, + "parameter.transformer_blocks.23.ff.net.2.weight": 195.6923828125, + "parameter.transformer_blocks.23.norm1.weight": 5.559892654418945, + "parameter.transformer_blocks.23.norm2.weight": 351.5206298828125, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 4.866369247436523, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 1158.40771484375, + "parameter.transformer_blocks.24.attn.norm_k.weight": 1.8183107376098633, + "parameter.transformer_blocks.24.attn.norm_q.weight": 1.4892802238464355, + "parameter.transformer_blocks.24.attn.to_k.weight": 1.4131574630737305, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 0.6507848501205444, + "parameter.transformer_blocks.24.attn.to_q.weight": 0.6354339122772217, + "parameter.transformer_blocks.24.attn.to_v.weight": 0.9285950660705566, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 348.76458740234375, + "parameter.transformer_blocks.24.ff.net.2.weight": 370.705078125, + "parameter.transformer_blocks.24.norm1.weight": 6.474706649780273, + "parameter.transformer_blocks.24.norm2.weight": 533.75048828125, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 5.577777862548828, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 1323.35205078125, + "parameter.transformer_blocks.25.attn.norm_k.weight": 2.0326733589172363, + "parameter.transformer_blocks.25.attn.norm_q.weight": 1.8265665769577026, + "parameter.transformer_blocks.25.attn.to_k.weight": 2.347825050354004, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 0.5998079776763916, + "parameter.transformer_blocks.25.attn.to_q.weight": 1.2458218336105347, + "parameter.transformer_blocks.25.attn.to_v.weight": 1.4591829776763916, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 409.814697265625, + "parameter.transformer_blocks.25.ff.net.2.weight": 540.3759765625, + "parameter.transformer_blocks.25.norm1.weight": 2.727937698364258, + "parameter.transformer_blocks.25.norm2.weight": 473.1929931640625, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 4.299077033996582, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 1016.748291015625, + "parameter.transformer_blocks.26.attn.norm_k.weight": 2.916896343231201, + "parameter.transformer_blocks.26.attn.norm_q.weight": 2.9030942916870117, + "parameter.transformer_blocks.26.attn.to_k.weight": 2.8463592529296875, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.7373838424682617, + "parameter.transformer_blocks.26.attn.to_q.weight": 1.6248605251312256, + "parameter.transformer_blocks.26.attn.to_v.weight": 1.5674855709075928, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 259.17681884765625, + "parameter.transformer_blocks.26.ff.net.2.weight": 319.775390625, + "parameter.transformer_blocks.26.norm1.weight": 5.834235668182373, + "parameter.transformer_blocks.26.norm2.weight": 694.09619140625, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 4.972381114959717, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 723.10498046875, + "parameter.transformer_blocks.27.attn.norm_k.weight": 2.396589994430542, + "parameter.transformer_blocks.27.attn.norm_q.weight": 2.8167920112609863, + "parameter.transformer_blocks.27.attn.to_k.weight": 4.00289249420166, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 0.48014622926712036, + "parameter.transformer_blocks.27.attn.to_q.weight": 2.6316230297088623, + "parameter.transformer_blocks.27.attn.to_v.weight": 3.5429189205169678, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 211.07672119140625, + "parameter.transformer_blocks.27.ff.net.2.weight": 83.690185546875, + "parameter.transformer_blocks.27.norm1.weight": 17.755285263061523, + "parameter.transformer_blocks.27.norm2.weight": 436.00537109375, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 4.6777191162109375, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 1108.4140625, + "parameter.transformer_blocks.28.attn.norm_k.weight": 2.4593915939331055, + "parameter.transformer_blocks.28.attn.norm_q.weight": 2.287010669708252, + "parameter.transformer_blocks.28.attn.to_k.weight": 2.3165788650512695, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 0.8154343366622925, + "parameter.transformer_blocks.28.attn.to_q.weight": 2.055948257446289, + "parameter.transformer_blocks.28.attn.to_v.weight": 3.6044247150421143, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 163.99996948242188, + "parameter.transformer_blocks.28.ff.net.2.weight": 396.37841796875, + "parameter.transformer_blocks.28.norm1.weight": 7.464785575866699, + "parameter.transformer_blocks.28.norm2.weight": 983.402099609375, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 4.8386688232421875, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 1127.87890625, + "parameter.transformer_blocks.29.attn.norm_k.weight": 6.578639507293701, + "parameter.transformer_blocks.29.attn.norm_q.weight": 5.110691547393799, + "parameter.transformer_blocks.29.attn.to_k.weight": 4.1849684715271, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 0.9435173869132996, + "parameter.transformer_blocks.29.attn.to_q.weight": 2.8432507514953613, + "parameter.transformer_blocks.29.attn.to_v.weight": 3.2650294303894043, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 344.19805908203125, + "parameter.transformer_blocks.29.ff.net.2.weight": 282.1435546875, + "parameter.transformer_blocks.29.norm1.weight": 8.832938194274902, + "parameter.transformer_blocks.29.norm2.weight": 376.4471435546875, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 31.928050994873047, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 7585.13330078125, + "parameter.transformer_blocks.3.attn.norm_k.weight": 7.919702053070068, + "parameter.transformer_blocks.3.attn.norm_q.weight": 7.012223720550537, + "parameter.transformer_blocks.3.attn.to_k.weight": 4.903183460235596, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 3.3478007316589355, + "parameter.transformer_blocks.3.attn.to_q.weight": 11.044124603271484, + "parameter.transformer_blocks.3.attn.to_v.weight": 1.4261471033096313, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 801.4122924804688, + "parameter.transformer_blocks.3.ff.net.2.weight": 260.0557861328125, + "parameter.transformer_blocks.3.norm1.weight": 9.486372947692871, + "parameter.transformer_blocks.3.norm2.weight": 2734.337890625, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 5.7413330078125, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 1350.337890625, + "parameter.transformer_blocks.30.attn.norm_k.weight": 5.465303897857666, + "parameter.transformer_blocks.30.attn.norm_q.weight": 5.46435546875, + "parameter.transformer_blocks.30.attn.to_k.weight": 5.119855880737305, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 0.9348153471946716, + "parameter.transformer_blocks.30.attn.to_q.weight": 3.1151809692382812, + "parameter.transformer_blocks.30.attn.to_v.weight": 2.243969440460205, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 159.632080078125, + "parameter.transformer_blocks.30.ff.net.2.weight": 159.099365234375, + "parameter.transformer_blocks.30.norm1.weight": 11.394719123840332, + "parameter.transformer_blocks.30.norm2.weight": 301.135498046875, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 5.110774993896484, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 739.34765625, + "parameter.transformer_blocks.31.attn.norm_k.weight": 3.8568966388702393, + "parameter.transformer_blocks.31.attn.norm_q.weight": 3.906276226043701, + "parameter.transformer_blocks.31.attn.to_k.weight": 1.2982009649276733, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 0.7056084275245667, + "parameter.transformer_blocks.31.attn.to_q.weight": 1.1121916770935059, + "parameter.transformer_blocks.31.attn.to_v.weight": 1.6242420673370361, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 139.44873046875, + "parameter.transformer_blocks.31.ff.net.2.weight": 104.65277099609375, + "parameter.transformer_blocks.31.norm1.weight": 4.170856475830078, + "parameter.transformer_blocks.31.norm2.weight": 417.41650390625, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 5.7159857749938965, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 770.952392578125, + "parameter.transformer_blocks.32.attn.norm_k.weight": 5.170241832733154, + "parameter.transformer_blocks.32.attn.norm_q.weight": 5.211503982543945, + "parameter.transformer_blocks.32.attn.to_k.weight": 3.110686779022217, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 1.5664355754852295, + "parameter.transformer_blocks.32.attn.to_q.weight": 2.448791027069092, + "parameter.transformer_blocks.32.attn.to_v.weight": 1.7731201648712158, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 140.49258422851562, + "parameter.transformer_blocks.32.ff.net.2.weight": 68.61331176757812, + "parameter.transformer_blocks.32.norm1.weight": 9.787178039550781, + "parameter.transformer_blocks.32.norm2.weight": 297.69940185546875, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 4.766090393066406, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 774.0791015625, + "parameter.transformer_blocks.33.attn.norm_k.weight": 4.196714878082275, + "parameter.transformer_blocks.33.attn.norm_q.weight": 4.320964813232422, + "parameter.transformer_blocks.33.attn.to_k.weight": 2.6738462448120117, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 2.137948513031006, + "parameter.transformer_blocks.33.attn.to_q.weight": 1.688585877418518, + "parameter.transformer_blocks.33.attn.to_v.weight": 1.9710156917572021, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 206.24911499023438, + "parameter.transformer_blocks.33.ff.net.2.weight": 71.080810546875, + "parameter.transformer_blocks.33.norm1.weight": 10.011297225952148, + "parameter.transformer_blocks.33.norm2.weight": 667.790771484375, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 7.7424468994140625, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 1036.16064453125, + "parameter.transformer_blocks.34.attn.norm_k.weight": 8.191692352294922, + "parameter.transformer_blocks.34.attn.norm_q.weight": 8.252612113952637, + "parameter.transformer_blocks.34.attn.to_k.weight": 4.370151519775391, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 1.28891122341156, + "parameter.transformer_blocks.34.attn.to_q.weight": 2.719578742980957, + "parameter.transformer_blocks.34.attn.to_v.weight": 2.4563608169555664, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 129.66920471191406, + "parameter.transformer_blocks.34.ff.net.2.weight": 223.3360595703125, + "parameter.transformer_blocks.34.norm1.weight": 25.90180778503418, + "parameter.transformer_blocks.34.norm2.weight": 483.9815673828125, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 4.1862640380859375, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 971.1640625, + "parameter.transformer_blocks.35.attn.norm_k.weight": 3.796952247619629, + "parameter.transformer_blocks.35.attn.norm_q.weight": 3.838136672973633, + "parameter.transformer_blocks.35.attn.to_k.weight": 2.2202110290527344, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 1.6180669069290161, + "parameter.transformer_blocks.35.attn.to_q.weight": 2.895798683166504, + "parameter.transformer_blocks.35.attn.to_v.weight": 1.8437247276306152, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 138.53948974609375, + "parameter.transformer_blocks.35.ff.net.2.weight": 124.335205078125, + "parameter.transformer_blocks.35.norm1.weight": 4.438076972961426, + "parameter.transformer_blocks.35.norm2.weight": 495.2100830078125, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 6.160249710083008, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 1463.8109130859375, + "parameter.transformer_blocks.36.attn.norm_k.weight": 3.277141809463501, + "parameter.transformer_blocks.36.attn.norm_q.weight": 3.317197799682617, + "parameter.transformer_blocks.36.attn.to_k.weight": 5.186769962310791, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 0.866863489151001, + "parameter.transformer_blocks.36.attn.to_q.weight": 3.173929214477539, + "parameter.transformer_blocks.36.attn.to_v.weight": 1.6129945516586304, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 206.4375, + "parameter.transformer_blocks.36.ff.net.2.weight": 71.85504150390625, + "parameter.transformer_blocks.36.norm1.weight": 15.505533218383789, + "parameter.transformer_blocks.36.norm2.weight": 1674.610595703125, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 5.213738441467285, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 1062.2938232421875, + "parameter.transformer_blocks.37.attn.norm_k.weight": 2.9631030559539795, + "parameter.transformer_blocks.37.attn.norm_q.weight": 3.029543876647949, + "parameter.transformer_blocks.37.attn.to_k.weight": 3.477022409439087, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 0.7805371284484863, + "parameter.transformer_blocks.37.attn.to_q.weight": 1.8869693279266357, + "parameter.transformer_blocks.37.attn.to_v.weight": 1.9407142400741577, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 447.45355224609375, + "parameter.transformer_blocks.37.ff.net.2.weight": 49.300811767578125, + "parameter.transformer_blocks.37.norm1.weight": 7.276589870452881, + "parameter.transformer_blocks.37.norm2.weight": 364.8778076171875, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 4.798489093780518, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 683.1319580078125, + "parameter.transformer_blocks.38.attn.norm_k.weight": 4.893096923828125, + "parameter.transformer_blocks.38.attn.norm_q.weight": 4.971759796142578, + "parameter.transformer_blocks.38.attn.to_k.weight": 2.918351173400879, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 0.5770251154899597, + "parameter.transformer_blocks.38.attn.to_q.weight": 1.906826376914978, + "parameter.transformer_blocks.38.attn.to_v.weight": 1.2777994871139526, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 97.0000228881836, + "parameter.transformer_blocks.38.ff.net.2.weight": 127.248291015625, + "parameter.transformer_blocks.38.norm1.weight": 9.897371292114258, + "parameter.transformer_blocks.38.norm2.weight": 468.759765625, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 7.547196388244629, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 1791.341552734375, + "parameter.transformer_blocks.39.attn.norm_k.weight": 3.645379066467285, + "parameter.transformer_blocks.39.attn.norm_q.weight": 3.6634902954101562, + "parameter.transformer_blocks.39.attn.to_k.weight": 2.178628444671631, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 0.7987926602363586, + "parameter.transformer_blocks.39.attn.to_q.weight": 1.5225403308868408, + "parameter.transformer_blocks.39.attn.to_v.weight": 1.4502114057540894, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 83.56243896484375, + "parameter.transformer_blocks.39.ff.net.2.weight": 121.65234375, + "parameter.transformer_blocks.39.norm1.weight": 5.666716575622559, + "parameter.transformer_blocks.39.norm2.weight": 720.703857421875, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 21.32648468017578, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 5028.822265625, + "parameter.transformer_blocks.4.attn.norm_k.weight": 4.957897186279297, + "parameter.transformer_blocks.4.attn.norm_q.weight": 4.433900833129883, + "parameter.transformer_blocks.4.attn.to_k.weight": 6.1132001876831055, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 1.9584068059921265, + "parameter.transformer_blocks.4.attn.to_q.weight": 5.3016510009765625, + "parameter.transformer_blocks.4.attn.to_v.weight": 1.5592217445373535, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 773.1954345703125, + "parameter.transformer_blocks.4.ff.net.2.weight": 268.1669921875, + "parameter.transformer_blocks.4.norm1.weight": 14.660944938659668, + "parameter.transformer_blocks.4.norm2.weight": 887.4386596679688, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 4.005348205566406, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 854.2412109375, + "parameter.transformer_blocks.40.attn.norm_k.weight": 2.949293613433838, + "parameter.transformer_blocks.40.attn.norm_q.weight": 4.571164608001709, + "parameter.transformer_blocks.40.attn.to_k.weight": 3.0293209552764893, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 0.6123585104942322, + "parameter.transformer_blocks.40.attn.to_q.weight": 3.568161964416504, + "parameter.transformer_blocks.40.attn.to_v.weight": 1.8074942827224731, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 140.12953186035156, + "parameter.transformer_blocks.40.ff.net.2.weight": 152.8375244140625, + "parameter.transformer_blocks.40.norm1.weight": 5.113841533660889, + "parameter.transformer_blocks.40.norm2.weight": 129.9615478515625, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 4.217306137084961, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 987.54638671875, + "parameter.transformer_blocks.41.attn.norm_k.weight": 2.3821609020233154, + "parameter.transformer_blocks.41.attn.norm_q.weight": 2.412914752960205, + "parameter.transformer_blocks.41.attn.to_k.weight": 1.237536072731018, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 0.7576071619987488, + "parameter.transformer_blocks.41.attn.to_q.weight": 1.2888585329055786, + "parameter.transformer_blocks.41.attn.to_v.weight": 1.4355392456054688, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 127.60647583007812, + "parameter.transformer_blocks.41.ff.net.2.weight": 152.33447265625, + "parameter.transformer_blocks.41.norm1.weight": 3.5250301361083984, + "parameter.transformer_blocks.41.norm2.weight": 445.77001953125, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 3.4761009216308594, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 846.796875, + "parameter.transformer_blocks.42.attn.norm_k.weight": 2.687297821044922, + "parameter.transformer_blocks.42.attn.norm_q.weight": 2.7009596824645996, + "parameter.transformer_blocks.42.attn.to_k.weight": 1.5777289867401123, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 0.6879916191101074, + "parameter.transformer_blocks.42.attn.to_q.weight": 0.9663207530975342, + "parameter.transformer_blocks.42.attn.to_v.weight": 1.3828082084655762, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 321.3841552734375, + "parameter.transformer_blocks.42.ff.net.2.weight": 144.8499755859375, + "parameter.transformer_blocks.42.norm1.weight": 4.310583114624023, + "parameter.transformer_blocks.42.norm2.weight": 276.15570068359375, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 5.307106018066406, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 1281.81640625, + "parameter.transformer_blocks.43.attn.norm_k.weight": 6.124272346496582, + "parameter.transformer_blocks.43.attn.norm_q.weight": 6.2348408699035645, + "parameter.transformer_blocks.43.attn.to_k.weight": 2.432459831237793, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 0.7861360311508179, + "parameter.transformer_blocks.43.attn.to_q.weight": 1.5843608379364014, + "parameter.transformer_blocks.43.attn.to_v.weight": 1.4354814291000366, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 281.88043212890625, + "parameter.transformer_blocks.43.ff.net.2.weight": 68.20425415039062, + "parameter.transformer_blocks.43.norm1.weight": 8.903609275817871, + "parameter.transformer_blocks.43.norm2.weight": 188.57647705078125, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 11.816520690917969, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 2805.9755859375, + "parameter.transformer_blocks.44.attn.norm_k.weight": 2.603389263153076, + "parameter.transformer_blocks.44.attn.norm_q.weight": 2.5335259437561035, + "parameter.transformer_blocks.44.attn.to_k.weight": 1.9004267454147339, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 0.9370132684707642, + "parameter.transformer_blocks.44.attn.to_q.weight": 1.4908075332641602, + "parameter.transformer_blocks.44.attn.to_v.weight": 1.3552401065826416, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 266.840576171875, + "parameter.transformer_blocks.44.ff.net.2.weight": 86.6273193359375, + "parameter.transformer_blocks.44.norm1.weight": 7.830251693725586, + "parameter.transformer_blocks.44.norm2.weight": 565.53564453125, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 5.412269592285156, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 1297.621337890625, + "parameter.transformer_blocks.45.attn.norm_k.weight": 2.0145363807678223, + "parameter.transformer_blocks.45.attn.norm_q.weight": 2.071850299835205, + "parameter.transformer_blocks.45.attn.to_k.weight": 2.3817930221557617, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.45184069871902466, + "parameter.transformer_blocks.45.attn.to_q.weight": 30.458595275878906, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.9223560690879822, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 193.89520263671875, + "parameter.transformer_blocks.45.ff.net.2.weight": 166.10183715820312, + "parameter.transformer_blocks.45.norm1.weight": 8.281604766845703, + "parameter.transformer_blocks.45.norm2.weight": 309.4776611328125, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 6.8365631103515625, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 1614.33984375, + "parameter.transformer_blocks.46.attn.norm_k.weight": 2.728790283203125, + "parameter.transformer_blocks.46.attn.norm_q.weight": 2.4782252311706543, + "parameter.transformer_blocks.46.attn.to_k.weight": 2.4199938774108887, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.3956280052661896, + "parameter.transformer_blocks.46.attn.to_q.weight": 2.4051148891448975, + "parameter.transformer_blocks.46.attn.to_v.weight": 1.3095606565475464, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 296.1636047363281, + "parameter.transformer_blocks.46.ff.net.2.weight": 853.3876953125, + "parameter.transformer_blocks.46.norm1.weight": 9.261405944824219, + "parameter.transformer_blocks.46.norm2.weight": 2051.6435546875, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 17.834823608398438, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 4247.4453125, + "parameter.transformer_blocks.47.attn.norm_k.weight": 5.551021575927734, + "parameter.transformer_blocks.47.attn.norm_q.weight": 5.684200286865234, + "parameter.transformer_blocks.47.attn.to_k.weight": 12.534589767456055, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 1.4007599353790283, + "parameter.transformer_blocks.47.attn.to_q.weight": 4.060039520263672, + "parameter.transformer_blocks.47.attn.to_v.weight": 10.387222290039062, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 795.94580078125, + "parameter.transformer_blocks.47.ff.net.2.weight": 1438.3203125, + "parameter.transformer_blocks.47.norm1.weight": 10.565017700195312, + "parameter.transformer_blocks.47.norm2.weight": 1186.939453125, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 21.9805908203125, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 4901.0, + "parameter.transformer_blocks.48.attn.norm_k.weight": 2.7005364894866943, + "parameter.transformer_blocks.48.attn.norm_q.weight": 2.77207088470459, + "parameter.transformer_blocks.48.attn.to_k.weight": 4.525178909301758, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 1.0088205337524414, + "parameter.transformer_blocks.48.attn.to_q.weight": 6.6492791175842285, + "parameter.transformer_blocks.48.attn.to_v.weight": 2.435667037963867, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 954.7305297851562, + "parameter.transformer_blocks.48.ff.net.2.weight": 712.043212890625, + "parameter.transformer_blocks.48.norm1.weight": 5.328737258911133, + "parameter.transformer_blocks.48.norm2.weight": 1848.45703125, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 6.877416610717773, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 960.73828125, + "parameter.transformer_blocks.49.attn.norm_k.weight": 5.781886100769043, + "parameter.transformer_blocks.49.attn.norm_q.weight": 5.775971412658691, + "parameter.transformer_blocks.49.attn.to_k.weight": 120.58187866210938, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 1.635695219039917, + "parameter.transformer_blocks.49.attn.to_q.weight": 5951.6728515625, + "parameter.transformer_blocks.49.attn.to_v.weight": 12.527933120727539, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 1403.05126953125, + "parameter.transformer_blocks.49.ff.net.2.weight": 1077.126953125, + "parameter.transformer_blocks.49.norm1.weight": 88.3868179321289, + "parameter.transformer_blocks.49.norm2.weight": 731.34423828125, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 22.666481018066406, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 5374.4609375, + "parameter.transformer_blocks.5.attn.norm_k.weight": 3.1488566398620605, + "parameter.transformer_blocks.5.attn.norm_q.weight": 3.0029892921447754, + "parameter.transformer_blocks.5.attn.to_k.weight": 2.35954213142395, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 2.4877138137817383, + "parameter.transformer_blocks.5.attn.to_q.weight": 3.6635231971740723, + "parameter.transformer_blocks.5.attn.to_v.weight": 0.8009382486343384, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 417.8459167480469, + "parameter.transformer_blocks.5.ff.net.2.weight": 180.81723022460938, + "parameter.transformer_blocks.5.norm1.weight": 8.364011764526367, + "parameter.transformer_blocks.5.norm2.weight": 7164.46484375, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 6.457424163818359, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 1539.517333984375, + "parameter.transformer_blocks.6.attn.norm_k.weight": 3.098749876022339, + "parameter.transformer_blocks.6.attn.norm_q.weight": 3.6933846473693848, + "parameter.transformer_blocks.6.attn.to_k.weight": 2.2403249740600586, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 0.9580168724060059, + "parameter.transformer_blocks.6.attn.to_q.weight": 4.970716953277588, + "parameter.transformer_blocks.6.attn.to_v.weight": 1.5979580879211426, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 856.0050659179688, + "parameter.transformer_blocks.6.ff.net.2.weight": 134.00628662109375, + "parameter.transformer_blocks.6.norm1.weight": 5.737216949462891, + "parameter.transformer_blocks.6.norm2.weight": 756.50048828125, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 6.233695030212402, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 1038.2841796875, + "parameter.transformer_blocks.7.attn.norm_k.weight": 3.7008655071258545, + "parameter.transformer_blocks.7.attn.norm_q.weight": 3.2430977821350098, + "parameter.transformer_blocks.7.attn.to_k.weight": 1.7681100368499756, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 0.9606469869613647, + "parameter.transformer_blocks.7.attn.to_q.weight": 1.8657087087631226, + "parameter.transformer_blocks.7.attn.to_v.weight": 1.2660351991653442, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 1445.8360595703125, + "parameter.transformer_blocks.7.ff.net.2.weight": 104.09483337402344, + "parameter.transformer_blocks.7.norm1.weight": 11.101558685302734, + "parameter.transformer_blocks.7.norm2.weight": 1281.765869140625, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 9.229780197143555, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 2193.487548828125, + "parameter.transformer_blocks.8.attn.norm_k.weight": 1.500385046005249, + "parameter.transformer_blocks.8.attn.norm_q.weight": 1.459054946899414, + "parameter.transformer_blocks.8.attn.to_k.weight": 1.144398808479309, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 0.9251453876495361, + "parameter.transformer_blocks.8.attn.to_q.weight": 1.5949926376342773, + "parameter.transformer_blocks.8.attn.to_v.weight": 0.6618691086769104, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 718.069580078125, + "parameter.transformer_blocks.8.ff.net.2.weight": 114.80775451660156, + "parameter.transformer_blocks.8.norm1.weight": 3.614859104156494, + "parameter.transformer_blocks.8.norm2.weight": 2606.68359375, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 5.06758975982666, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 1203.7950439453125, + "parameter.transformer_blocks.9.attn.norm_k.weight": 2.6946144104003906, + "parameter.transformer_blocks.9.attn.norm_q.weight": 2.358290195465088, + "parameter.transformer_blocks.9.attn.to_k.weight": 1.6546334028244019, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 0.5754650235176086, + "parameter.transformer_blocks.9.attn.to_q.weight": 3.5623562335968018, + "parameter.transformer_blocks.9.attn.to_v.weight": 1.1003837585449219, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 407.1868896484375, + "parameter.transformer_blocks.9.ff.net.2.weight": 90.68083190917969, + "parameter.transformer_blocks.9.norm1.weight": 6.272876739501953, + "parameter.transformer_blocks.9.norm2.weight": 448.25537109375, + "video": 0.77642822265625 + }, + "bounds_relative_l2": { + "audio": 0.055874562519869435, + "input_gradients.audio_hidden_states": 0.5322077855718023, + "input_gradients.encoder_hidden_states": 0.8201112660925163, + "input_gradients.hidden_states": 0.570961125155766, + "nodes.audio_proj_in": 0.001, + "nodes.audio_proj_in.grad": 0.5647759182449158, + "nodes.audio_proj_out": 0.05583448800632685, + "nodes.audio_proj_out.grad": 0.05987456251986943, + "nodes.block.0": 0.020856273681871214, + "nodes.block.0.grad": 2.2223851901162193, + "nodes.block.1": 0.01963606863771287, + "nodes.block.1.grad": 2.1711386259276138, + "nodes.block.10": 0.025624545745054676, + "nodes.block.10.grad": 2.3827599847554692, + "nodes.block.11": 0.025985391477088116, + "nodes.block.11.grad": 2.3844366994775155, + "nodes.block.12": 0.026510570243016597, + "nodes.block.12.grad": 2.464081837113086, + "nodes.block.13": 0.027265984942517207, + "nodes.block.13.grad": 2.495653322894215, + "nodes.block.14": 0.027774847708480276, + "nodes.block.14.grad": 2.5094531896273926, + "nodes.block.15": 0.028750013046869617, + "nodes.block.15.grad": 2.5264188071831617, + "nodes.block.16": 0.02945524012717312, + "nodes.block.16.grad": 2.5791103746642503, + "nodes.block.17": 0.03002846913698215, + "nodes.block.17.grad": 2.597714194651892, + "nodes.block.18": 0.03270233573383202, + "nodes.block.18.grad": 2.5959965556849545, + "nodes.block.19": 0.03500400443185651, + "nodes.block.19.grad": 2.647298785191703, + "nodes.block.2": 0.02025170917229934, + "nodes.block.2.grad": 2.561946800938579, + "nodes.block.20": 0.03603438523204576, + "nodes.block.20.grad": 2.654514680209731, + "nodes.block.21": 0.036850674168122315, + "nodes.block.21.grad": 2.6803915843962587, + "nodes.block.22": 0.03625456245606433, + "nodes.block.22.grad": 2.7575750756326936, + "nodes.block.23": 0.03428651375960175, + "nodes.block.23.grad": 2.8131499193681324, + "nodes.block.24": 0.029550852410779385, + "nodes.block.24.grad": 2.877172825167145, + "nodes.block.25": 0.026820069691058424, + "nodes.block.25.grad": 2.8977817427103716, + "nodes.block.26": 0.02514449582691894, + "nodes.block.26.grad": 2.901376893866698, + "nodes.block.27": 0.05127208996315919, + "nodes.block.27.grad": 2.9120452755602653, + "nodes.block.28": 0.04379588326751174, + "nodes.block.28.grad": 2.937083189568119, + "nodes.block.29": 0.0431860226385089, + "nodes.block.29.grad": 2.803516189221703, + "nodes.block.3": 0.021995502188648874, + "nodes.block.3.grad": 2.360842556421442, + "nodes.block.30": 0.04213505991316521, + "nodes.block.30.grad": 2.64568640018582, + "nodes.block.31": 0.04331209295755942, + "nodes.block.31.grad": 2.7451072939109578, + "nodes.block.32": 0.040014718762034836, + "nodes.block.32.grad": 2.615002999331556, + "nodes.block.33": 0.040392899720599576, + "nodes.block.33.grad": 2.609574000490925, + "nodes.block.34": 0.04067843997826269, + "nodes.block.34.grad": 2.834503210447525, + "nodes.block.35": 0.03641098525432868, + "nodes.block.35.grad": 2.9165366050627264, + "nodes.block.36": 0.03297402876678478, + "nodes.block.36.grad": 2.8276718551698936, + "nodes.block.37": 0.03160852819071327, + "nodes.block.37.grad": 2.77520294901385, + "nodes.block.38": 0.03233105914246465, + "nodes.block.38.grad": 2.8283303457677595, + "nodes.block.39": 0.03216289190451429, + "nodes.block.39.grad": 2.8540157299753397, + "nodes.block.4": 0.023591763377124136, + "nodes.block.4.grad": 2.4098772178669474, + "nodes.block.40": 0.030933532049720696, + "nodes.block.40.grad": 2.8697604169567468, + "nodes.block.41": 0.03119100676158183, + "nodes.block.41.grad": 2.8825146889968614, + "nodes.block.42": 0.026289492238989076, + "nodes.block.42.grad": 3.0928782062807088, + "nodes.block.43": 0.020140799108251242, + "nodes.block.43.grad": 3.0576597403511196, + "nodes.block.44": 0.019506680953228443, + "nodes.block.44.grad": 3.013652237870902, + "nodes.block.45": 0.02115346901345907, + "nodes.block.45.grad": 3.0030285081630783, + "nodes.block.46": 0.02105395340522597, + "nodes.block.46.grad": 2.9673756043551, + "nodes.block.47": 0.016002538734894753, + "nodes.block.47.grad": 2.973243601760091, + "nodes.block.48": 0.018486101654955133, + "nodes.block.48.grad": 3.050383267832265, + "nodes.block.49": 0.02039832231310697, + "nodes.block.49.grad": 0.05773916498393257, + "nodes.block.5": 0.024130894300617163, + "nodes.block.5.grad": 2.270758112999734, + "nodes.block.6": 0.024268191662500875, + "nodes.block.6.grad": 2.3733618325946417, + "nodes.block.7": 0.024324743369856273, + "nodes.block.7.grad": 2.3597214946400884, + "nodes.block.8": 0.025012925768023244, + "nodes.block.8.grad": 2.3545103657037747, + "nodes.block.9": 0.025473614701246088, + "nodes.block.9.grad": 2.3469190835276272, + "nodes.context_embedder": 0.007199885083778527, + "nodes.context_embedder.grad": 0.8232039140070144, + "nodes.norm_out": 0.04593889301942204, + "nodes.norm_out.grad": 0.058526997835089106, + "nodes.proj_in": 0.001, + "nodes.proj_in.grad": 0.6102260678480784, + "nodes.proj_out": 0.0496271698625265, + "nodes.proj_out.grad": 0.021805234505499992, + "nodes.refiner.0": 0.014741722004295996, + "nodes.refiner.0.grad": 0.833198641928383, + "nodes.refiner.1": 0.016572113205235068, + "nodes.refiner.1.grad": 0.832172955022045, + "nodes.time_embedder": 0.001, + "nodes.time_embedder.grad": 1.3471330701300817, + "nodes.time_proj": 0.001, + "nodes.token_refiner": 0.01681921465529777, + "nodes.token_refiner.grad": 0.8266076567858954, + "parameter.audio_proj_in.bias": 0.4924406210254883, + "parameter.audio_proj_in.weight": 0.5645161238198128, + "parameter.audio_proj_out.bias": 0.04115526495753882, + "parameter.audio_proj_out.weight": 0.04956870486156885, + "parameter.context_embedder.bias": 0.8921846392352075, + "parameter.context_embedder.weight": 0.8225462304542064, + "parameter.norm_out.linear.bias": 0.04634026248145849, + "parameter.norm_out.linear.weight": 0.04625262281004235, + "parameter.norm_out.norm.weight": 0.03589542548700971, + "parameter.proj_in.bias": 0.5710820501685063, + "parameter.proj_in.weight": 0.6181584003502457, + "parameter.proj_out.bias": 0.019366519519612328, + "parameter.proj_out.weight": 0.033329808130047274, + "parameter.time_embedder.linear_1.bias": 0.09295946902006715, + "parameter.time_embedder.linear_1.weight": 0.08581757990692288, + "parameter.time_embedder.linear_2.bias": 1.2949600448078376, + "parameter.time_embedder.linear_2.weight": 0.3858624288589434, + "parameter.token_refiner.final_norm.weight": 0.8545591164818493, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": 0.6977844496170106, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": 0.696436119960276, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": 0.8274289290532341, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": 0.8548751035821915, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": 0.7967628175599382, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": 0.8552145288519823, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": 0.8421136257842587, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": 0.8541830952817251, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": 0.8461923977581175, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": 0.867156538337202, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": 1.143592452947768, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": 1.144926517758613, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": 0.9203645497892541, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": 0.876771597022621, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": 0.9368330199913044, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": 0.8917728579157279, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": 0.9538617435520901, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": 0.8776271074767351, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": 0.6944313461838573, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": 0.9626344119635454, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": 0.8343857459191878, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": 0.4695623484222038, + "parameter.transformer_blocks.0.attn.norm_k.weight": 0.7579721892678787, + "parameter.transformer_blocks.0.attn.norm_q.weight": 0.8173884177839518, + "parameter.transformer_blocks.0.attn.to_k.weight": 0.598048869065734, + "parameter.transformer_blocks.0.attn.to_out.0.weight": 0.4607183400750576, + "parameter.transformer_blocks.0.attn.to_q.weight": 0.6095098053999932, + "parameter.transformer_blocks.0.attn.to_v.weight": 0.4565274641005145, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": 0.47368161351989685, + "parameter.transformer_blocks.0.ff.net.2.weight": 0.46706704734661963, + "parameter.transformer_blocks.0.norm1.weight": 0.5690199829802196, + "parameter.transformer_blocks.0.norm2.weight": 0.4713643572348861, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": 0.48625070765289713, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": 0.46631960433629827, + "parameter.transformer_blocks.1.attn.norm_k.weight": 1.5928846935443048, + "parameter.transformer_blocks.1.attn.norm_q.weight": 1.6802262756178572, + "parameter.transformer_blocks.1.attn.to_k.weight": 1.7463617619368006, + "parameter.transformer_blocks.1.attn.to_out.0.weight": 1.3992411530210536, + "parameter.transformer_blocks.1.attn.to_q.weight": 1.5308964954536752, + "parameter.transformer_blocks.1.attn.to_v.weight": 1.5078440376921296, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": 0.48165415622186136, + "parameter.transformer_blocks.1.ff.net.2.weight": 0.4602352329789112, + "parameter.transformer_blocks.1.norm1.weight": 1.6781901777194925, + "parameter.transformer_blocks.1.norm2.weight": 0.4499919762222125, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": 0.47148786502263684, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": 0.4272072578235439, + "parameter.transformer_blocks.10.attn.norm_k.weight": 1.6046638120314505, + "parameter.transformer_blocks.10.attn.norm_q.weight": 1.6108075051113542, + "parameter.transformer_blocks.10.attn.to_k.weight": 1.546673665186303, + "parameter.transformer_blocks.10.attn.to_out.0.weight": 1.3221676077517956, + "parameter.transformer_blocks.10.attn.to_q.weight": 1.5250613246280351, + "parameter.transformer_blocks.10.attn.to_v.weight": 1.439652952408004, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": 0.4726732753392043, + "parameter.transformer_blocks.10.ff.net.2.weight": 0.3963444861154397, + "parameter.transformer_blocks.10.norm1.weight": 1.6772450755020096, + "parameter.transformer_blocks.10.norm2.weight": 0.6002672020768325, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": 0.42361490222844905, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": 0.39662066333443935, + "parameter.transformer_blocks.11.attn.norm_k.weight": 1.840572307580958, + "parameter.transformer_blocks.11.attn.norm_q.weight": 1.9188744283003905, + "parameter.transformer_blocks.11.attn.to_k.weight": 1.7703977150032377, + "parameter.transformer_blocks.11.attn.to_out.0.weight": 1.2876843768443307, + "parameter.transformer_blocks.11.attn.to_q.weight": 1.737723211961002, + "parameter.transformer_blocks.11.attn.to_v.weight": 1.3995188087961634, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": 0.4197019364819283, + "parameter.transformer_blocks.11.ff.net.2.weight": 0.37612409124730006, + "parameter.transformer_blocks.11.norm1.weight": 1.6079513206454785, + "parameter.transformer_blocks.11.norm2.weight": 0.458253057250772, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": 0.43840053394416933, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": 0.40671239044508817, + "parameter.transformer_blocks.12.attn.norm_k.weight": 1.7488444150259954, + "parameter.transformer_blocks.12.attn.norm_q.weight": 1.7632868615827981, + "parameter.transformer_blocks.12.attn.to_k.weight": 1.3142226322965935, + "parameter.transformer_blocks.12.attn.to_out.0.weight": 1.1359072994114097, + "parameter.transformer_blocks.12.attn.to_q.weight": 1.4349954944475813, + "parameter.transformer_blocks.12.attn.to_v.weight": 1.2742728285071432, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": 0.45273844195521246, + "parameter.transformer_blocks.12.ff.net.2.weight": 0.3847020860814938, + "parameter.transformer_blocks.12.norm1.weight": 1.3172006511246992, + "parameter.transformer_blocks.12.norm2.weight": 0.46636478397989606, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": 0.44600884145665415, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": 0.40674649755173176, + "parameter.transformer_blocks.13.attn.norm_k.weight": 1.7799406417442298, + "parameter.transformer_blocks.13.attn.norm_q.weight": 1.7573791861082868, + "parameter.transformer_blocks.13.attn.to_k.weight": 1.6074768422290602, + "parameter.transformer_blocks.13.attn.to_out.0.weight": 1.6132410052966137, + "parameter.transformer_blocks.13.attn.to_q.weight": 1.5997128757317836, + "parameter.transformer_blocks.13.attn.to_v.weight": 1.6464668337321637, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": 0.4442169048972016, + "parameter.transformer_blocks.13.ff.net.2.weight": 0.380902675691008, + "parameter.transformer_blocks.13.norm1.weight": 1.7105338913592616, + "parameter.transformer_blocks.13.norm2.weight": 0.46961815484790537, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": 0.45725653215188217, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": 0.41858402192577837, + "parameter.transformer_blocks.14.attn.norm_k.weight": 1.469986824442482, + "parameter.transformer_blocks.14.attn.norm_q.weight": 1.5763017849092813, + "parameter.transformer_blocks.14.attn.to_k.weight": 1.4510501970223415, + "parameter.transformer_blocks.14.attn.to_out.0.weight": 1.2970826939393076, + "parameter.transformer_blocks.14.attn.to_q.weight": 1.5147021262359528, + "parameter.transformer_blocks.14.attn.to_v.weight": 1.4810157564197193, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": 0.43087334365484353, + "parameter.transformer_blocks.14.ff.net.2.weight": 0.37976161539859005, + "parameter.transformer_blocks.14.norm1.weight": 1.7076552397075138, + "parameter.transformer_blocks.14.norm2.weight": 0.5074845627613467, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": 0.4401152871230969, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": 0.39990535961965595, + "parameter.transformer_blocks.15.attn.norm_k.weight": 1.7406358338702776, + "parameter.transformer_blocks.15.attn.norm_q.weight": 1.815698683567046, + "parameter.transformer_blocks.15.attn.to_k.weight": 1.6943845045609855, + "parameter.transformer_blocks.15.attn.to_out.0.weight": 1.2497651908305334, + "parameter.transformer_blocks.15.attn.to_q.weight": 1.5719706664264341, + "parameter.transformer_blocks.15.attn.to_v.weight": 1.4362232343688441, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": 0.4269020829501556, + "parameter.transformer_blocks.15.ff.net.2.weight": 0.3829018739751305, + "parameter.transformer_blocks.15.norm1.weight": 1.7213183235706304, + "parameter.transformer_blocks.15.norm2.weight": 0.5455163777631221, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": 0.4520159206050519, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": 0.39948408371174726, + "parameter.transformer_blocks.16.attn.norm_k.weight": 1.4119388371302533, + "parameter.transformer_blocks.16.attn.norm_q.weight": 1.4682436451837244, + "parameter.transformer_blocks.16.attn.to_k.weight": 1.3432969220005635, + "parameter.transformer_blocks.16.attn.to_out.0.weight": 0.9832956575858338, + "parameter.transformer_blocks.16.attn.to_q.weight": 1.359900382306337, + "parameter.transformer_blocks.16.attn.to_v.weight": 1.201476601788808, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": 0.4238987856910134, + "parameter.transformer_blocks.16.ff.net.2.weight": 0.35984166582667915, + "parameter.transformer_blocks.16.norm1.weight": 1.3843170680644228, + "parameter.transformer_blocks.16.norm2.weight": 0.428799335513681, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": 0.4145135261066581, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": 0.3799409518555693, + "parameter.transformer_blocks.17.attn.norm_k.weight": 1.3217790303538601, + "parameter.transformer_blocks.17.attn.norm_q.weight": 1.4572092735369861, + "parameter.transformer_blocks.17.attn.to_k.weight": 1.550763289841423, + "parameter.transformer_blocks.17.attn.to_out.0.weight": 1.5443913216106686, + "parameter.transformer_blocks.17.attn.to_q.weight": 1.4996885508715714, + "parameter.transformer_blocks.17.attn.to_v.weight": 1.7509856021917902, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": 0.3961476298950383, + "parameter.transformer_blocks.17.ff.net.2.weight": 0.3504688496738397, + "parameter.transformer_blocks.17.norm1.weight": 1.8017473556707249, + "parameter.transformer_blocks.17.norm2.weight": 0.3983444777914455, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": 0.3778583747694606, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": 0.3598241348103365, + "parameter.transformer_blocks.18.attn.norm_k.weight": 1.8639724127835577, + "parameter.transformer_blocks.18.attn.norm_q.weight": 1.6944261657118604, + "parameter.transformer_blocks.18.attn.to_k.weight": 1.6319025238001412, + "parameter.transformer_blocks.18.attn.to_out.0.weight": 1.2363141576360999, + "parameter.transformer_blocks.18.attn.to_q.weight": 1.624264917378737, + "parameter.transformer_blocks.18.attn.to_v.weight": 1.5524415567182823, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": 0.37286740575046284, + "parameter.transformer_blocks.18.ff.net.2.weight": 0.3202254288682409, + "parameter.transformer_blocks.18.norm1.weight": 1.6159961194038832, + "parameter.transformer_blocks.18.norm2.weight": 0.3818729100045425, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": 0.3726465715841605, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": 0.34752555444109356, + "parameter.transformer_blocks.19.attn.norm_k.weight": 1.6930014849151624, + "parameter.transformer_blocks.19.attn.norm_q.weight": 1.6630708529790197, + "parameter.transformer_blocks.19.attn.to_k.weight": 1.6588350140955876, + "parameter.transformer_blocks.19.attn.to_out.0.weight": 1.2439010207019663, + "parameter.transformer_blocks.19.attn.to_q.weight": 1.6428512510788948, + "parameter.transformer_blocks.19.attn.to_v.weight": 1.484881926385726, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": 0.33970880705001394, + "parameter.transformer_blocks.19.ff.net.2.weight": 0.268614767899265, + "parameter.transformer_blocks.19.norm1.weight": 1.657001015483316, + "parameter.transformer_blocks.19.norm2.weight": 0.3972401717652793, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": 0.46606343973880315, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": 0.4487492056614175, + "parameter.transformer_blocks.2.attn.norm_k.weight": 1.4482745787968618, + "parameter.transformer_blocks.2.attn.norm_q.weight": 1.4011792992179826, + "parameter.transformer_blocks.2.attn.to_k.weight": 1.383713706346312, + "parameter.transformer_blocks.2.attn.to_out.0.weight": 1.2434077546585083, + "parameter.transformer_blocks.2.attn.to_q.weight": 1.2824894433942415, + "parameter.transformer_blocks.2.attn.to_v.weight": 1.6011841732416363, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": 0.46201371424920007, + "parameter.transformer_blocks.2.ff.net.2.weight": 0.4449119826160648, + "parameter.transformer_blocks.2.norm1.weight": 1.4189648365457033, + "parameter.transformer_blocks.2.norm2.weight": 0.4706222547522539, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": 0.40841601968208596, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": 0.35530501345629856, + "parameter.transformer_blocks.20.attn.norm_k.weight": 2.271528269172145, + "parameter.transformer_blocks.20.attn.norm_q.weight": 2.3965219300340412, + "parameter.transformer_blocks.20.attn.to_k.weight": 2.233833032947429, + "parameter.transformer_blocks.20.attn.to_out.0.weight": 1.0171265853519902, + "parameter.transformer_blocks.20.attn.to_q.weight": 2.0641527960619745, + "parameter.transformer_blocks.20.attn.to_v.weight": 1.2487110943334556, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": 0.3693490359205701, + "parameter.transformer_blocks.20.ff.net.2.weight": 0.3074205930258824, + "parameter.transformer_blocks.20.norm1.weight": 1.7700395817002512, + "parameter.transformer_blocks.20.norm2.weight": 0.367692691166245, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": 0.3903710767724551, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": 0.34301517959049416, + "parameter.transformer_blocks.21.attn.norm_k.weight": 1.950064440905238, + "parameter.transformer_blocks.21.attn.norm_q.weight": 1.9790393030131417, + "parameter.transformer_blocks.21.attn.to_k.weight": 1.8696022230524951, + "parameter.transformer_blocks.21.attn.to_out.0.weight": 1.3085905832706133, + "parameter.transformer_blocks.21.attn.to_q.weight": 1.7626791017272785, + "parameter.transformer_blocks.21.attn.to_v.weight": 1.6148204989527732, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": 0.37421061782095605, + "parameter.transformer_blocks.21.ff.net.2.weight": 0.31188595625347326, + "parameter.transformer_blocks.21.norm1.weight": 1.6715098834355426, + "parameter.transformer_blocks.21.norm2.weight": 0.3780210765516755, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": 0.39334408525058967, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": 0.34932131130627625, + "parameter.transformer_blocks.22.attn.norm_k.weight": 1.7282131749289305, + "parameter.transformer_blocks.22.attn.norm_q.weight": 1.7361869124100267, + "parameter.transformer_blocks.22.attn.to_k.weight": 1.6211689913370957, + "parameter.transformer_blocks.22.attn.to_out.0.weight": 1.7531777523551206, + "parameter.transformer_blocks.22.attn.to_q.weight": 1.5804957440034704, + "parameter.transformer_blocks.22.attn.to_v.weight": 2.0127702197095947, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": 0.3758269799981911, + "parameter.transformer_blocks.22.ff.net.2.weight": 0.30036560489829417, + "parameter.transformer_blocks.22.norm1.weight": 1.746308692033695, + "parameter.transformer_blocks.22.norm2.weight": 0.38526873361961966, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": 0.3659958885542907, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": 0.32942885246364434, + "parameter.transformer_blocks.23.attn.norm_k.weight": 1.7767804893186219, + "parameter.transformer_blocks.23.attn.norm_q.weight": 1.7835005943173663, + "parameter.transformer_blocks.23.attn.to_k.weight": 1.801271713754693, + "parameter.transformer_blocks.23.attn.to_out.0.weight": 1.2826408202288249, + "parameter.transformer_blocks.23.attn.to_q.weight": 1.6637272867807518, + "parameter.transformer_blocks.23.attn.to_v.weight": 1.6276718377292485, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": 0.35111900532078266, + "parameter.transformer_blocks.23.ff.net.2.weight": 0.30229493242687083, + "parameter.transformer_blocks.23.norm1.weight": 1.8728925608820077, + "parameter.transformer_blocks.23.norm2.weight": 0.37504010259999715, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": 0.370112243599527, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": 0.3259608253987697, + "parameter.transformer_blocks.24.attn.norm_k.weight": 1.5540151225587742, + "parameter.transformer_blocks.24.attn.norm_q.weight": 1.5282074408083117, + "parameter.transformer_blocks.24.attn.to_k.weight": 1.4743448183198562, + "parameter.transformer_blocks.24.attn.to_out.0.weight": 1.0091014859382852, + "parameter.transformer_blocks.24.attn.to_q.weight": 1.3878023913313464, + "parameter.transformer_blocks.24.attn.to_v.weight": 1.3814820433053931, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": 0.3495862859832943, + "parameter.transformer_blocks.24.ff.net.2.weight": 0.22342188840254962, + "parameter.transformer_blocks.24.norm1.weight": 1.5432131235368212, + "parameter.transformer_blocks.24.norm2.weight": 0.3512437453403421, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": 0.2887349002565304, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": 0.27046913523446603, + "parameter.transformer_blocks.25.attn.norm_k.weight": 1.691868877233949, + "parameter.transformer_blocks.25.attn.norm_q.weight": 1.7010915057049818, + "parameter.transformer_blocks.25.attn.to_k.weight": 1.7368112778802502, + "parameter.transformer_blocks.25.attn.to_out.0.weight": 1.350349317777754, + "parameter.transformer_blocks.25.attn.to_q.weight": 1.668223607985642, + "parameter.transformer_blocks.25.attn.to_v.weight": 1.723285127155597, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": 0.28133247296190517, + "parameter.transformer_blocks.25.ff.net.2.weight": 0.20194025848149452, + "parameter.transformer_blocks.25.norm1.weight": 1.411724914676122, + "parameter.transformer_blocks.25.norm2.weight": 0.26834024435523535, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": 0.3379136712161619, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": 0.28810004498840786, + "parameter.transformer_blocks.26.attn.norm_k.weight": 1.6798287262637364, + "parameter.transformer_blocks.26.attn.norm_q.weight": 1.69104714770592, + "parameter.transformer_blocks.26.attn.to_k.weight": 1.692782637359654, + "parameter.transformer_blocks.26.attn.to_out.0.weight": 0.557316213795948, + "parameter.transformer_blocks.26.attn.to_q.weight": 1.6866281235933018, + "parameter.transformer_blocks.26.attn.to_v.weight": 0.8505834862776024, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": 0.3081469953555959, + "parameter.transformer_blocks.26.ff.net.2.weight": 0.220262802078313, + "parameter.transformer_blocks.26.norm1.weight": 1.3291321584681985, + "parameter.transformer_blocks.26.norm2.weight": 0.32260745381630884, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": 0.3665319287861786, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": 0.29816704388226967, + "parameter.transformer_blocks.27.attn.norm_k.weight": 2.015954579673838, + "parameter.transformer_blocks.27.attn.norm_q.weight": 2.0728614171140696, + "parameter.transformer_blocks.27.attn.to_k.weight": 2.119717052057698, + "parameter.transformer_blocks.27.attn.to_out.0.weight": 1.655684299333085, + "parameter.transformer_blocks.27.attn.to_q.weight": 2.063800384141528, + "parameter.transformer_blocks.27.attn.to_v.weight": 1.903318012690791, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": 0.3336286871587634, + "parameter.transformer_blocks.27.ff.net.2.weight": 0.2629234943665337, + "parameter.transformer_blocks.27.norm1.weight": 2.2664585814860914, + "parameter.transformer_blocks.27.norm2.weight": 0.3169720684762965, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": 0.3515487162879209, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": 0.26841340549235404, + "parameter.transformer_blocks.28.attn.norm_k.weight": 2.399781555047728, + "parameter.transformer_blocks.28.attn.norm_q.weight": 2.395155262876651, + "parameter.transformer_blocks.28.attn.to_k.weight": 2.4228271102309455, + "parameter.transformer_blocks.28.attn.to_out.0.weight": 2.195929642104692, + "parameter.transformer_blocks.28.attn.to_q.weight": 2.47044428200628, + "parameter.transformer_blocks.28.attn.to_v.weight": 2.235409904366294, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": 0.3231814220870359, + "parameter.transformer_blocks.28.ff.net.2.weight": 0.23873283403188625, + "parameter.transformer_blocks.28.norm1.weight": 2.406016592778906, + "parameter.transformer_blocks.28.norm2.weight": 0.35388536898791156, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": 0.3370124618645231, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": 0.2581199142347362, + "parameter.transformer_blocks.29.attn.norm_k.weight": 2.828870634002473, + "parameter.transformer_blocks.29.attn.norm_q.weight": 2.8394752158134153, + "parameter.transformer_blocks.29.attn.to_k.weight": 2.4587867193040096, + "parameter.transformer_blocks.29.attn.to_out.0.weight": 2.1807317376716786, + "parameter.transformer_blocks.29.attn.to_q.weight": 2.477952801057237, + "parameter.transformer_blocks.29.attn.to_v.weight": 2.2819078609741155, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": 0.31067077867960535, + "parameter.transformer_blocks.29.ff.net.2.weight": 0.2125025961953128, + "parameter.transformer_blocks.29.norm1.weight": 2.22937114072439, + "parameter.transformer_blocks.29.norm2.weight": 0.2985583794573764, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": 0.4563955440751999, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": 0.44056551548436124, + "parameter.transformer_blocks.3.attn.norm_k.weight": 1.9910700212191144, + "parameter.transformer_blocks.3.attn.norm_q.weight": 1.945933246514572, + "parameter.transformer_blocks.3.attn.to_k.weight": 2.1471948191904158, + "parameter.transformer_blocks.3.attn.to_out.0.weight": 1.3353532506843855, + "parameter.transformer_blocks.3.attn.to_q.weight": 1.8455292260430465, + "parameter.transformer_blocks.3.attn.to_v.weight": 1.5236774521517638, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": 0.45603384682915216, + "parameter.transformer_blocks.3.ff.net.2.weight": 0.43336542727846195, + "parameter.transformer_blocks.3.norm1.weight": 1.9361030097866772, + "parameter.transformer_blocks.3.norm2.weight": 0.46347801113442294, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": 0.424860641324806, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": 0.2575813696839202, + "parameter.transformer_blocks.30.attn.norm_k.weight": 2.4133075689771766, + "parameter.transformer_blocks.30.attn.norm_q.weight": 2.4555342549587347, + "parameter.transformer_blocks.30.attn.to_k.weight": 2.7248059348650067, + "parameter.transformer_blocks.30.attn.to_out.0.weight": 1.7435173521915326, + "parameter.transformer_blocks.30.attn.to_q.weight": 2.5419536303615473, + "parameter.transformer_blocks.30.attn.to_v.weight": 1.9029702259619197, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": 0.33588833172930505, + "parameter.transformer_blocks.30.ff.net.2.weight": 0.2380083386328826, + "parameter.transformer_blocks.30.norm1.weight": 2.505899745707479, + "parameter.transformer_blocks.30.norm2.weight": 0.324304447571616, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": 0.38201682560719036, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": 0.27693613816411955, + "parameter.transformer_blocks.31.attn.norm_k.weight": 1.4749752282586552, + "parameter.transformer_blocks.31.attn.norm_q.weight": 1.450698015116089, + "parameter.transformer_blocks.31.attn.to_k.weight": 1.7458419357349744, + "parameter.transformer_blocks.31.attn.to_out.0.weight": 1.7270328460847992, + "parameter.transformer_blocks.31.attn.to_q.weight": 1.6576446872999169, + "parameter.transformer_blocks.31.attn.to_v.weight": 1.9190422292061655, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": 0.3783161083398855, + "parameter.transformer_blocks.31.ff.net.2.weight": 0.26410263772251447, + "parameter.transformer_blocks.31.norm1.weight": 1.6875584481154922, + "parameter.transformer_blocks.31.norm2.weight": 0.36632858313046956, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": 0.4020554204311656, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": 0.23973686092285756, + "parameter.transformer_blocks.32.attn.norm_k.weight": 2.9716232546172625, + "parameter.transformer_blocks.32.attn.norm_q.weight": 2.906283565015692, + "parameter.transformer_blocks.32.attn.to_k.weight": 2.606234086669554, + "parameter.transformer_blocks.32.attn.to_out.0.weight": 2.317841292977602, + "parameter.transformer_blocks.32.attn.to_q.weight": 2.5657513065502493, + "parameter.transformer_blocks.32.attn.to_v.weight": 2.3794686664168188, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": 0.34719683801155393, + "parameter.transformer_blocks.32.ff.net.2.weight": 0.246765566298946, + "parameter.transformer_blocks.32.norm1.weight": 2.4357935968387205, + "parameter.transformer_blocks.32.norm2.weight": 0.34853550249950177, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": 0.36846996686388106, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": 0.2533834097236321, + "parameter.transformer_blocks.33.attn.norm_k.weight": 2.3854834345576035, + "parameter.transformer_blocks.33.attn.norm_q.weight": 2.3851789502334184, + "parameter.transformer_blocks.33.attn.to_k.weight": 2.166871112866884, + "parameter.transformer_blocks.33.attn.to_out.0.weight": 2.202871406466205, + "parameter.transformer_blocks.33.attn.to_q.weight": 2.0966241850183223, + "parameter.transformer_blocks.33.attn.to_v.weight": 2.2126399275458635, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": 0.3441799372592239, + "parameter.transformer_blocks.33.ff.net.2.weight": 0.22341773100887405, + "parameter.transformer_blocks.33.norm1.weight": 1.9732371533027735, + "parameter.transformer_blocks.33.norm2.weight": 0.38652812762723965, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": 0.3554213000192706, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": 0.22140519681851778, + "parameter.transformer_blocks.34.attn.norm_k.weight": 2.5507693494105754, + "parameter.transformer_blocks.34.attn.norm_q.weight": 2.514695418148331, + "parameter.transformer_blocks.34.attn.to_k.weight": 2.1610564534533907, + "parameter.transformer_blocks.34.attn.to_out.0.weight": 2.0670438640047566, + "parameter.transformer_blocks.34.attn.to_q.weight": 1.9783468584060075, + "parameter.transformer_blocks.34.attn.to_v.weight": 2.2278909504303757, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": 0.3168857391270072, + "parameter.transformer_blocks.34.ff.net.2.weight": 0.2229557077010477, + "parameter.transformer_blocks.34.norm1.weight": 2.449795373815762, + "parameter.transformer_blocks.34.norm2.weight": 0.3219383109635162, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": 0.26992720396736086, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": 0.20284057396161692, + "parameter.transformer_blocks.35.attn.norm_k.weight": 2.2032694832045756, + "parameter.transformer_blocks.35.attn.norm_q.weight": 2.1419146451614828, + "parameter.transformer_blocks.35.attn.to_k.weight": 2.1450069456596244, + "parameter.transformer_blocks.35.attn.to_out.0.weight": 2.0838470797540802, + "parameter.transformer_blocks.35.attn.to_q.weight": 2.3149170156922634, + "parameter.transformer_blocks.35.attn.to_v.weight": 2.270605870587408, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": 0.2935230806913727, + "parameter.transformer_blocks.35.ff.net.2.weight": 0.19840333184595743, + "parameter.transformer_blocks.35.norm1.weight": 2.006915113924567, + "parameter.transformer_blocks.35.norm2.weight": 0.34304758555408366, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": 0.3757365784358437, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": 0.2740854204169333, + "parameter.transformer_blocks.36.attn.norm_k.weight": 2.3829003296526037, + "parameter.transformer_blocks.36.attn.norm_q.weight": 2.471052061201031, + "parameter.transformer_blocks.36.attn.to_k.weight": 2.9427079084130043, + "parameter.transformer_blocks.36.attn.to_out.0.weight": 1.332734461988801, + "parameter.transformer_blocks.36.attn.to_q.weight": 2.879402875824165, + "parameter.transformer_blocks.36.attn.to_v.weight": 1.441635881952929, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": 0.34849143326375176, + "parameter.transformer_blocks.36.ff.net.2.weight": 0.2115231280461088, + "parameter.transformer_blocks.36.norm1.weight": 2.390337959761368, + "parameter.transformer_blocks.36.norm2.weight": 0.9862222229373728, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": 0.32267481434729123, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": 0.24647183914350898, + "parameter.transformer_blocks.37.attn.norm_k.weight": 2.3771887840171453, + "parameter.transformer_blocks.37.attn.norm_q.weight": 2.423704528655729, + "parameter.transformer_blocks.37.attn.to_k.weight": 2.6074986725020586, + "parameter.transformer_blocks.37.attn.to_out.0.weight": 1.9361508720859442, + "parameter.transformer_blocks.37.attn.to_q.weight": 2.470932596225775, + "parameter.transformer_blocks.37.attn.to_v.weight": 2.019540099261265, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": 0.3864463038599187, + "parameter.transformer_blocks.37.ff.net.2.weight": 0.1813498562255011, + "parameter.transformer_blocks.37.norm1.weight": 2.3005942876742864, + "parameter.transformer_blocks.37.norm2.weight": 0.46200849638649916, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": 0.33107329032365335, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": 0.1845978723866556, + "parameter.transformer_blocks.38.attn.norm_k.weight": 2.1939132281440314, + "parameter.transformer_blocks.38.attn.norm_q.weight": 2.1452150077606085, + "parameter.transformer_blocks.38.attn.to_k.weight": 2.1651576745401, + "parameter.transformer_blocks.38.attn.to_out.0.weight": 2.1535469137370447, + "parameter.transformer_blocks.38.attn.to_q.weight": 2.279476597611475, + "parameter.transformer_blocks.38.attn.to_v.weight": 2.149915644738202, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": 0.29984975443501904, + "parameter.transformer_blocks.38.ff.net.2.weight": 0.17507619815358705, + "parameter.transformer_blocks.38.norm1.weight": 1.9931094344507323, + "parameter.transformer_blocks.38.norm2.weight": 0.2890943272506466, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": 0.2724759054232888, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": 0.17329018219412629, + "parameter.transformer_blocks.39.attn.norm_k.weight": 2.518809405868853, + "parameter.transformer_blocks.39.attn.norm_q.weight": 2.5537298120473655, + "parameter.transformer_blocks.39.attn.to_k.weight": 2.558570850046743, + "parameter.transformer_blocks.39.attn.to_out.0.weight": 1.9924087817450835, + "parameter.transformer_blocks.39.attn.to_q.weight": 2.6671613820786084, + "parameter.transformer_blocks.39.attn.to_v.weight": 2.0221826033394854, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": 0.2630599179615183, + "parameter.transformer_blocks.39.ff.net.2.weight": 0.16353138563559005, + "parameter.transformer_blocks.39.norm1.weight": 2.027398404327182, + "parameter.transformer_blocks.39.norm2.weight": 0.436626928527809, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": 0.45113642198656345, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": 0.43101662337052293, + "parameter.transformer_blocks.4.attn.norm_k.weight": 2.107241376051506, + "parameter.transformer_blocks.4.attn.norm_q.weight": 2.004893866644943, + "parameter.transformer_blocks.4.attn.to_k.weight": 1.7392158631560324, + "parameter.transformer_blocks.4.attn.to_out.0.weight": 1.6407321965312882, + "parameter.transformer_blocks.4.attn.to_q.weight": 1.7290204791402193, + "parameter.transformer_blocks.4.attn.to_v.weight": 1.8260754551141534, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": 0.45999773423866647, + "parameter.transformer_blocks.4.ff.net.2.weight": 0.41413221430852154, + "parameter.transformer_blocks.4.norm1.weight": 1.8600126214364934, + "parameter.transformer_blocks.4.norm2.weight": 0.4403873494482417, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": 0.2638772792673216, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": 0.15919358069898773, + "parameter.transformer_blocks.40.attn.norm_k.weight": 2.4582355239320846, + "parameter.transformer_blocks.40.attn.norm_q.weight": 2.5766816188303494, + "parameter.transformer_blocks.40.attn.to_k.weight": 2.5084094776471106, + "parameter.transformer_blocks.40.attn.to_out.0.weight": 1.8345278973957233, + "parameter.transformer_blocks.40.attn.to_q.weight": 2.7075145860176923, + "parameter.transformer_blocks.40.attn.to_v.weight": 1.9262344531233737, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": 0.27317540061866874, + "parameter.transformer_blocks.40.ff.net.2.weight": 0.15558650999478643, + "parameter.transformer_blocks.40.norm1.weight": 2.3981166865576298, + "parameter.transformer_blocks.40.norm2.weight": 0.23824461741568484, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": 0.191360775930618, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": 0.1450143678563992, + "parameter.transformer_blocks.41.attn.norm_k.weight": 1.9330923158289264, + "parameter.transformer_blocks.41.attn.norm_q.weight": 1.8806695149341441, + "parameter.transformer_blocks.41.attn.to_k.weight": 2.1827387176528363, + "parameter.transformer_blocks.41.attn.to_out.0.weight": 1.8721762162103988, + "parameter.transformer_blocks.41.attn.to_q.weight": 2.249737462999677, + "parameter.transformer_blocks.41.attn.to_v.weight": 1.913574251462993, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": 0.22387799526641033, + "parameter.transformer_blocks.41.ff.net.2.weight": 0.15004004565798207, + "parameter.transformer_blocks.41.norm1.weight": 1.9332100527809852, + "parameter.transformer_blocks.41.norm2.weight": 0.19165247686043319, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": 0.16453635926541302, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": 0.13160736099676676, + "parameter.transformer_blocks.42.attn.norm_k.weight": 1.5084062187453637, + "parameter.transformer_blocks.42.attn.norm_q.weight": 1.463803593081589, + "parameter.transformer_blocks.42.attn.to_k.weight": 1.555048192238938, + "parameter.transformer_blocks.42.attn.to_out.0.weight": 1.2972988812176829, + "parameter.transformer_blocks.42.attn.to_q.weight": 1.6338812862993308, + "parameter.transformer_blocks.42.attn.to_v.weight": 1.4615009176846085, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": 0.21578377049573982, + "parameter.transformer_blocks.42.ff.net.2.weight": 0.14167854885448355, + "parameter.transformer_blocks.42.norm1.weight": 1.5943387242506044, + "parameter.transformer_blocks.42.norm2.weight": 0.22845966128609996, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": 0.22484246646944805, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": 0.1687711117118654, + "parameter.transformer_blocks.43.attn.norm_k.weight": 2.180565755015369, + "parameter.transformer_blocks.43.attn.norm_q.weight": 2.1611831560713997, + "parameter.transformer_blocks.43.attn.to_k.weight": 2.1774587483083394, + "parameter.transformer_blocks.43.attn.to_out.0.weight": 1.3174994998577385, + "parameter.transformer_blocks.43.attn.to_q.weight": 2.238322727790866, + "parameter.transformer_blocks.43.attn.to_v.weight": 1.3667835433910123, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": 0.21768096168672024, + "parameter.transformer_blocks.43.ff.net.2.weight": 0.14780851910735965, + "parameter.transformer_blocks.43.norm1.weight": 1.6306606902435092, + "parameter.transformer_blocks.43.norm2.weight": 0.21782607241066146, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": 0.19922929632230438, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": 0.17817015140475925, + "parameter.transformer_blocks.44.attn.norm_k.weight": 1.8391245858237928, + "parameter.transformer_blocks.44.attn.norm_q.weight": 1.864980839274388, + "parameter.transformer_blocks.44.attn.to_k.weight": 2.1243029169917564, + "parameter.transformer_blocks.44.attn.to_out.0.weight": 1.2780173105474835, + "parameter.transformer_blocks.44.attn.to_q.weight": 2.3187433229789387, + "parameter.transformer_blocks.44.attn.to_v.weight": 1.4357866100370924, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": 0.21355438105510993, + "parameter.transformer_blocks.44.ff.net.2.weight": 0.14408956902374398, + "parameter.transformer_blocks.44.norm1.weight": 1.6284372959582694, + "parameter.transformer_blocks.44.norm2.weight": 0.31179346547739417, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": 0.1343378255995548, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": 0.12090601314079129, + "parameter.transformer_blocks.45.attn.norm_k.weight": 2.0356348895945304, + "parameter.transformer_blocks.45.attn.norm_q.weight": 2.0733720157423807, + "parameter.transformer_blocks.45.attn.to_k.weight": 1.8620087075904481, + "parameter.transformer_blocks.45.attn.to_out.0.weight": 0.7866043851644491, + "parameter.transformer_blocks.45.attn.to_q.weight": 1.3982527785852539, + "parameter.transformer_blocks.45.attn.to_v.weight": 0.8717362182639989, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": 0.1713462139245297, + "parameter.transformer_blocks.45.ff.net.2.weight": 0.1287573605895141, + "parameter.transformer_blocks.45.norm1.weight": 1.3244389846990698, + "parameter.transformer_blocks.45.norm2.weight": 0.18854067798447308, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": 0.10830846452841518, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": 0.10728990122130379, + "parameter.transformer_blocks.46.attn.norm_k.weight": 1.3782393820638366, + "parameter.transformer_blocks.46.attn.norm_q.weight": 1.3110997735619072, + "parameter.transformer_blocks.46.attn.to_k.weight": 1.4897283485557713, + "parameter.transformer_blocks.46.attn.to_out.0.weight": 0.5369101727268709, + "parameter.transformer_blocks.46.attn.to_q.weight": 1.5762500990412245, + "parameter.transformer_blocks.46.attn.to_v.weight": 0.5567964296530112, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": 0.14444435811698353, + "parameter.transformer_blocks.46.ff.net.2.weight": 0.11821473005409938, + "parameter.transformer_blocks.46.norm1.weight": 0.5348343943377146, + "parameter.transformer_blocks.46.norm2.weight": 0.3194514353043557, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": 0.13557466488824688, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": 0.13416083485211097, + "parameter.transformer_blocks.47.attn.norm_k.weight": 1.0601252284585165, + "parameter.transformer_blocks.47.attn.norm_q.weight": 1.0472470452425595, + "parameter.transformer_blocks.47.attn.to_k.weight": 0.7899602954074657, + "parameter.transformer_blocks.47.attn.to_out.0.weight": 0.42294360569065315, + "parameter.transformer_blocks.47.attn.to_q.weight": 0.8538143695820073, + "parameter.transformer_blocks.47.attn.to_v.weight": 0.4676883013469541, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": 0.1428202889026506, + "parameter.transformer_blocks.47.ff.net.2.weight": 0.1039253417421387, + "parameter.transformer_blocks.47.norm1.weight": 0.7254492935475526, + "parameter.transformer_blocks.47.norm2.weight": 0.1376405724859266, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": 0.09166683773844128, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": 0.08915507908331426, + "parameter.transformer_blocks.48.attn.norm_k.weight": 0.510893280194105, + "parameter.transformer_blocks.48.attn.norm_q.weight": 0.5352202011494669, + "parameter.transformer_blocks.48.attn.to_k.weight": 0.5278588976357866, + "parameter.transformer_blocks.48.attn.to_out.0.weight": 0.18246540545356343, + "parameter.transformer_blocks.48.attn.to_q.weight": 0.623528413463818, + "parameter.transformer_blocks.48.attn.to_v.weight": 0.1809432042421047, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": 0.1161069953590555, + "parameter.transformer_blocks.48.ff.net.2.weight": 0.08421579600923199, + "parameter.transformer_blocks.48.norm1.weight": 0.27209709408187704, + "parameter.transformer_blocks.48.norm2.weight": 0.1000847762017003, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": 0.2165217725814667, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": 0.07443694874289428, + "parameter.transformer_blocks.49.attn.norm_k.weight": 0.562476704259254, + "parameter.transformer_blocks.49.attn.norm_q.weight": 0.4922113703996883, + "parameter.transformer_blocks.49.attn.to_k.weight": 0.714677866127811, + "parameter.transformer_blocks.49.attn.to_out.0.weight": 0.15615622047740021, + "parameter.transformer_blocks.49.attn.to_q.weight": 2.0255527775969733, + "parameter.transformer_blocks.49.attn.to_v.weight": 0.21434871698649843, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": 0.08436019154157402, + "parameter.transformer_blocks.49.ff.net.2.weight": 0.05118975899847431, + "parameter.transformer_blocks.49.norm1.weight": 2.0601234837948583, + "parameter.transformer_blocks.49.norm2.weight": 0.10160790113615467, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": 0.463674588972218, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": 0.42144095957521877, + "parameter.transformer_blocks.5.attn.norm_k.weight": 2.586811193735602, + "parameter.transformer_blocks.5.attn.norm_q.weight": 2.3959813257209595, + "parameter.transformer_blocks.5.attn.to_k.weight": 2.5151449646215283, + "parameter.transformer_blocks.5.attn.to_out.0.weight": 1.0298899133856283, + "parameter.transformer_blocks.5.attn.to_q.weight": 2.6537790279999625, + "parameter.transformer_blocks.5.attn.to_v.weight": 1.1091717181092942, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": 0.44209218473146433, + "parameter.transformer_blocks.5.ff.net.2.weight": 0.41476093147023735, + "parameter.transformer_blocks.5.norm1.weight": 1.999425355549822, + "parameter.transformer_blocks.5.norm2.weight": 0.9724987771074863, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": 0.49964416503606013, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": 0.41626477157587205, + "parameter.transformer_blocks.6.attn.norm_k.weight": 1.2298329251640565, + "parameter.transformer_blocks.6.attn.norm_q.weight": 1.14294658184758, + "parameter.transformer_blocks.6.attn.to_k.weight": 1.2922229706843853, + "parameter.transformer_blocks.6.attn.to_out.0.weight": 1.463339856548597, + "parameter.transformer_blocks.6.attn.to_q.weight": 1.2811370674156086, + "parameter.transformer_blocks.6.attn.to_v.weight": 1.487613858672967, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": 0.4526237385782612, + "parameter.transformer_blocks.6.ff.net.2.weight": 0.4184947957482411, + "parameter.transformer_blocks.6.norm1.weight": 1.357144116444183, + "parameter.transformer_blocks.6.norm2.weight": 0.4269101073695155, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": 0.546350941026067, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": 0.46017370952375114, + "parameter.transformer_blocks.7.attn.norm_k.weight": 2.000517026258821, + "parameter.transformer_blocks.7.attn.norm_q.weight": 1.8883922935659494, + "parameter.transformer_blocks.7.attn.to_k.weight": 1.8670221743727042, + "parameter.transformer_blocks.7.attn.to_out.0.weight": 1.682527182681358, + "parameter.transformer_blocks.7.attn.to_q.weight": 1.8860312146003704, + "parameter.transformer_blocks.7.attn.to_v.weight": 1.7456388633256847, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": 0.48779323201486036, + "parameter.transformer_blocks.7.ff.net.2.weight": 0.4150858599947021, + "parameter.transformer_blocks.7.norm1.weight": 2.0695091938369203, + "parameter.transformer_blocks.7.norm2.weight": 0.4891738361436915, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": 0.46713297670080584, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": 0.43045852021605435, + "parameter.transformer_blocks.8.attn.norm_k.weight": 2.0267557339986726, + "parameter.transformer_blocks.8.attn.norm_q.weight": 2.1523386769149844, + "parameter.transformer_blocks.8.attn.to_k.weight": 2.0758534868592764, + "parameter.transformer_blocks.8.attn.to_out.0.weight": 1.2958531978900365, + "parameter.transformer_blocks.8.attn.to_q.weight": 2.0243324933674365, + "parameter.transformer_blocks.8.attn.to_v.weight": 1.4188328321799977, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": 0.45885517672142995, + "parameter.transformer_blocks.8.ff.net.2.weight": 0.41244756552440404, + "parameter.transformer_blocks.8.norm1.weight": 1.8582778459986842, + "parameter.transformer_blocks.8.norm2.weight": 0.46547890128629954, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": 0.491507859320213, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": 0.42571796827660574, + "parameter.transformer_blocks.9.attn.norm_k.weight": 2.124863410381398, + "parameter.transformer_blocks.9.attn.norm_q.weight": 2.1210277031068845, + "parameter.transformer_blocks.9.attn.to_k.weight": 1.9766700129529378, + "parameter.transformer_blocks.9.attn.to_out.0.weight": 1.2964216217713724, + "parameter.transformer_blocks.9.attn.to_q.weight": 2.027125295809663, + "parameter.transformer_blocks.9.attn.to_v.weight": 1.389181767512233, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": 0.4389162080074541, + "parameter.transformer_blocks.9.ff.net.2.weight": 0.40405263238644695, + "parameter.transformer_blocks.9.norm1.weight": 1.8360827910752702, + "parameter.transformer_blocks.9.norm2.weight": 0.4325678226497383, + "video": 0.017805238263453246 + } +} diff --git a/examples/models/minimax_h3/full-model-thresholds.json b/examples/models/minimax_h3/full-model-thresholds.json new file mode 100644 index 000000000..82b29d58c --- /dev/null +++ b/examples/models/minimax_h3/full-model-thresholds.json @@ -0,0 +1,3064 @@ +{ + "reference_report_sha256": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "other_report_sha256": "d49a80d6e438eb1f7d91193ad3f3ef5b5b3692b2c86a7532b86f23359ad33d79", + "status": "frozen_reference_bounds", + "rule": "1.5 * independent BF16 provider error + fixed floor", + "bounds": { + "video": { + "max_abs": 1.0716748962402343, + "relative_l2": 0.010294890439757714 + }, + "audio": { + "max_abs": 34.65894563293457, + "relative_l2": 0.030129260422720988 + }, + "nodes.audio_proj_in": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.audio_proj_in.grad": { + "max_abs": 656.3094281005859, + "relative_l2": 0.2961049755980539 + }, + "nodes.audio_proj_out": { + "max_abs": 34.65894563293457, + "relative_l2": 0.03022506434096451 + }, + "nodes.audio_proj_out.grad": { + "max_abs": 0.5465304005146027, + "relative_l2": 0.034129260422720985 + }, + "nodes.block.0": { + "max_abs": 86476800.001, + "relative_l2": 0.009254473682794001 + }, + "nodes.block.0.grad": { + "max_abs": 0.007892671618610621, + "relative_l2": 1.0282748009575065 + }, + "nodes.block.1": { + "max_abs": 135257088.001, + "relative_l2": 0.00974318315515435 + }, + "nodes.block.1.grad": { + "max_abs": 0.006624704105779529, + "relative_l2": 1.0305674729163048 + }, + "nodes.block.10": { + "max_abs": 243385344.001, + "relative_l2": 0.014183544640477025 + }, + "nodes.block.10.grad": { + "max_abs": 0.005376999407308176, + "relative_l2": 1.1788970789964635 + }, + "nodes.block.11": { + "max_abs": 246643200.001, + "relative_l2": 0.014389428832832524 + }, + "nodes.block.11.grad": { + "max_abs": 0.0053723131376318635, + "relative_l2": 1.1852776533292766 + }, + "nodes.block.12": { + "max_abs": 291520512.001, + "relative_l2": 0.014761840861898093 + }, + "nodes.block.12.grad": { + "max_abs": 0.005355812138877809, + "relative_l2": 1.2346247101218673 + }, + "nodes.block.13": { + "max_abs": 308868096.001, + "relative_l2": 0.015281297232727049 + }, + "nodes.block.13.grad": { + "max_abs": 0.005372488284483552, + "relative_l2": 1.2471660330417302 + }, + "nodes.block.14": { + "max_abs": 272268288.001, + "relative_l2": 0.015601342139096724 + }, + "nodes.block.14.grad": { + "max_abs": 0.005347652858181391, + "relative_l2": 1.2606747758455823 + }, + "nodes.block.15": { + "max_abs": 269202432.001, + "relative_l2": 0.016103171225580006 + }, + "nodes.block.15.grad": { + "max_abs": 0.005317293059197255, + "relative_l2": 1.2776078942880231 + }, + "nodes.block.16": { + "max_abs": 268137984.001, + "relative_l2": 0.016473838830060608 + }, + "nodes.block.16.grad": { + "max_abs": 0.005337768455210608, + "relative_l2": 1.2982750317435832 + }, + "nodes.block.17": { + "max_abs": 279490560.001, + "relative_l2": 0.01700976352461589 + }, + "nodes.block.17.grad": { + "max_abs": 0.005300611849525012, + "relative_l2": 1.304552845914669 + }, + "nodes.block.18": { + "max_abs": 281880576.001, + "relative_l2": 0.018818761997117264 + }, + "nodes.block.18.grad": { + "max_abs": 0.0053020529038622045, + "relative_l2": 1.3132094053007688 + }, + "nodes.block.19": { + "max_abs": 324725760.001, + "relative_l2": 0.020049400375340646 + }, + "nodes.block.19.grad": { + "max_abs": 0.005287827766442206, + "relative_l2": 1.331032060282588 + }, + "nodes.block.2": { + "max_abs": 123956736.001, + "relative_l2": 0.010425392138777783 + }, + "nodes.block.2.grad": { + "max_abs": 0.00622760480735451, + "relative_l2": 1.2071082994824671 + }, + "nodes.block.20": { + "max_abs": 313783296.001, + "relative_l2": 0.02069248717055676 + }, + "nodes.block.20.grad": { + "max_abs": 0.005284230249235407, + "relative_l2": 1.3333540488485576 + }, + "nodes.block.21": { + "max_abs": 299833344.001, + "relative_l2": 0.0211042005101494 + }, + "nodes.block.21.grad": { + "max_abs": 0.005257708161370829, + "relative_l2": 1.3444240875437083 + }, + "nodes.block.22": { + "max_abs": 372043776.001, + "relative_l2": 0.020841782421864578 + }, + "nodes.block.22.grad": { + "max_abs": 0.005248349679168314, + "relative_l2": 1.371844846626805 + }, + "nodes.block.23": { + "max_abs": 436162560.001, + "relative_l2": 0.019816861965668092 + }, + "nodes.block.23.grad": { + "max_abs": 0.005236955092987046, + "relative_l2": 1.3894633432736114 + }, + "nodes.block.24": { + "max_abs": 695009280.001, + "relative_l2": 0.016932971507712787 + }, + "nodes.block.24.grad": { + "max_abs": 0.005229643585043959, + "relative_l2": 1.418867550586571 + }, + "nodes.block.25": { + "max_abs": 1129820160.001, + "relative_l2": 0.015896194052533705 + }, + "nodes.block.25.grad": { + "max_abs": 0.005218857559957542, + "relative_l2": 1.4337955436014975 + }, + "nodes.block.26": { + "max_abs": 1842855936.001, + "relative_l2": 0.014708202945811977 + }, + "nodes.block.26.grad": { + "max_abs": 0.005217536799027585, + "relative_l2": 1.4267937146443628 + }, + "nodes.block.27": { + "max_abs": 2946981888.001, + "relative_l2": 0.029825333550548658 + }, + "nodes.block.27.grad": { + "max_abs": 0.005218868648516946, + "relative_l2": 1.42950586952009 + }, + "nodes.block.28": { + "max_abs": 1393311744.001, + "relative_l2": 0.02539093504291072 + }, + "nodes.block.28.grad": { + "max_abs": 0.005193499094166327, + "relative_l2": 1.4388679605626724 + }, + "nodes.block.29": { + "max_abs": 1097785344.001, + "relative_l2": 0.024924610281992207 + }, + "nodes.block.29.grad": { + "max_abs": 0.0051151480319094845, + "relative_l2": 1.2236330922921828 + }, + "nodes.block.3": { + "max_abs": 131919360.001, + "relative_l2": 0.011420138219926868 + }, + "nodes.block.3.grad": { + "max_abs": 0.005938520999625325, + "relative_l2": 1.136913714467664 + }, + "nodes.block.30": { + "max_abs": 1956667392.001, + "relative_l2": 0.023872450055785956 + }, + "nodes.block.30.grad": { + "max_abs": 0.005070673440859537, + "relative_l2": 1.16461148905023 + }, + "nodes.block.31": { + "max_abs": 2265833472.001, + "relative_l2": 0.024456562630709992 + }, + "nodes.block.31.grad": { + "max_abs": 0.005072068880399456, + "relative_l2": 1.194078650490841 + }, + "nodes.block.32": { + "max_abs": 2067637248.001, + "relative_l2": 0.023105217172228312 + }, + "nodes.block.32.grad": { + "max_abs": 0.005050663617003011, + "relative_l2": 1.1115885666094312 + }, + "nodes.block.33": { + "max_abs": 2405548032.001, + "relative_l2": 0.023271942704352414 + }, + "nodes.block.33.grad": { + "max_abs": 0.005049794474762166, + "relative_l2": 1.1311506019932958 + }, + "nodes.block.34": { + "max_abs": 2844155904.001, + "relative_l2": 0.02325442521193387 + }, + "nodes.block.34.grad": { + "max_abs": 0.005045931737986393, + "relative_l2": 1.2013872772100338 + }, + "nodes.block.35": { + "max_abs": 3470475264.001, + "relative_l2": 0.02134088292154418 + }, + "nodes.block.35.grad": { + "max_abs": 0.005041465970753052, + "relative_l2": 1.233510966039407 + }, + "nodes.block.36": { + "max_abs": 4340367360.001, + "relative_l2": 0.0180463405343924 + }, + "nodes.block.36.grad": { + "max_abs": 0.005039195159843075, + "relative_l2": 1.2096277002335567 + }, + "nodes.block.37": { + "max_abs": 5596250112.001, + "relative_l2": 0.017395995963329733 + }, + "nodes.block.37.grad": { + "max_abs": 0.005037690586395911, + "relative_l2": 1.1923039041300327 + }, + "nodes.block.38": { + "max_abs": 6312493056.001, + "relative_l2": 0.01715249891038743 + }, + "nodes.block.38.grad": { + "max_abs": 0.005036132002605882, + "relative_l2": 1.218786553128028 + }, + "nodes.block.39": { + "max_abs": 7958052864.001, + "relative_l2": 0.01719502785204246 + }, + "nodes.block.39.grad": { + "max_abs": 0.005035766980545304, + "relative_l2": 1.2318844221994263 + }, + "nodes.block.4": { + "max_abs": 169340928.001, + "relative_l2": 0.012555142596216557 + }, + "nodes.block.4.grad": { + "max_abs": 0.0056757682422176004, + "relative_l2": 1.17920320434519 + }, + "nodes.block.40": { + "max_abs": 7987740672.001, + "relative_l2": 0.017254775666558798 + }, + "nodes.block.40.grad": { + "max_abs": 0.0050362005121132826, + "relative_l2": 1.2320456657478307 + }, + "nodes.block.41": { + "max_abs": 8212217856.001, + "relative_l2": 0.017141120518213926 + }, + "nodes.block.41.grad": { + "max_abs": 0.005036102641388425, + "relative_l2": 1.2386607608755411 + }, + "nodes.block.42": { + "max_abs": 7658176512.001, + "relative_l2": 0.015161639630964174 + }, + "nodes.block.42.grad": { + "max_abs": 0.005034739564398478, + "relative_l2": 1.2944377183229072 + }, + "nodes.block.43": { + "max_abs": 15020851200.001, + "relative_l2": 0.010976331445410721 + }, + "nodes.block.43.grad": { + "max_abs": 0.005036297476981417, + "relative_l2": 1.2822040421243264 + }, + "nodes.block.44": { + "max_abs": 21491220480.001, + "relative_l2": 0.010692049224040517 + }, + "nodes.block.44.grad": { + "max_abs": 0.005036446064768825, + "relative_l2": 1.274834343770201 + }, + "nodes.block.45": { + "max_abs": 23387701248.001, + "relative_l2": 0.011531374588086778 + }, + "nodes.block.45.grad": { + "max_abs": 0.005037532945498242, + "relative_l2": 1.287014693377177 + }, + "nodes.block.46": { + "max_abs": 24589369344.001, + "relative_l2": 0.011653401897591616 + }, + "nodes.block.46.grad": { + "max_abs": 0.005035186672903365, + "relative_l2": 1.2755247924762254 + }, + "nodes.block.47": { + "max_abs": 36515610624.001, + "relative_l2": 0.009354598313891512 + }, + "nodes.block.47.grad": { + "max_abs": 0.005034289247196284, + "relative_l2": 1.2837355978859173 + }, + "nodes.block.48": { + "max_abs": 28595060736.001, + "relative_l2": 0.00991133297587967 + }, + "nodes.block.48.grad": { + "max_abs": 0.005034035918361042, + "relative_l2": 1.3231800171294437 + }, + "nodes.block.49": { + "max_abs": 35716202496.001, + "relative_l2": 0.01111432926598677 + }, + "nodes.block.49.grad": { + "max_abs": 0.005000006957613152, + "relative_l2": 0.03542644133688417 + }, + "nodes.block.5": { + "max_abs": 209988096.001, + "relative_l2": 0.01310967836419382 + }, + "nodes.block.5.grad": { + "max_abs": 0.005638774537947029, + "relative_l2": 1.1153482427434613 + }, + "nodes.block.6": { + "max_abs": 286685184.001, + "relative_l2": 0.013200938568209863 + }, + "nodes.block.6.grad": { + "max_abs": 0.005473614592920058, + "relative_l2": 1.1466965710007662 + }, + "nodes.block.7": { + "max_abs": 223693824.001, + "relative_l2": 0.013329474464451103 + }, + "nodes.block.7.grad": { + "max_abs": 0.005471000399556942, + "relative_l2": 1.1582630616146892 + }, + "nodes.block.8": { + "max_abs": 283370496.001, + "relative_l2": 0.013902013824406895 + }, + "nodes.block.8.grad": { + "max_abs": 0.005416834795032628, + "relative_l2": 1.1586090943783933 + }, + "nodes.block.9": { + "max_abs": 284544000.001, + "relative_l2": 0.014139204849605763 + }, + "nodes.block.9.grad": { + "max_abs": 0.0054162817203905434, + "relative_l2": 1.160017287530927 + }, + "nodes.context_embedder": { + "max_abs": 0.19000489807128906, + "relative_l2": 0.004576856779102996 + }, + "nodes.context_embedder.grad": { + "max_abs": 4.159599070549011, + "relative_l2": 0.4770407196194314 + }, + "nodes.norm_out": { + "max_abs": 79.44240625, + "relative_l2": 0.0241334814018383 + }, + "nodes.norm_out.grad": { + "max_abs": 1.6565254974365233, + "relative_l2": 0.034355922337465906 + }, + "nodes.proj_in": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.proj_in.grad": { + "max_abs": 192.61858642578125, + "relative_l2": 0.3541688852969683 + }, + "nodes.proj_out": { + "max_abs": 14.54162271118164, + "relative_l2": 0.0290541364591499 + }, + "nodes.proj_out.grad": { + "max_abs": 0.010576401948928834, + "relative_l2": 0.014294886490478394 + }, + "nodes.refiner.0": { + "max_abs": 69.3613515625, + "relative_l2": 0.008927916540939999 + }, + "nodes.refiner.0.grad": { + "max_abs": 0.08660967938601971, + "relative_l2": 0.48280690880483634 + }, + "nodes.refiner.1": { + "max_abs": 65.56887109375, + "relative_l2": 0.009971913328415988 + }, + "nodes.refiner.1.grad": { + "max_abs": 0.0854927796125412, + "relative_l2": 0.48221516635887207 + }, + "nodes.time_embedder": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.time_embedder.grad": { + "max_abs": 1082.1651867675782, + "relative_l2": 0.3660479013861043 + }, + "nodes.time_proj": { + "max_abs": 0.001, + "relative_l2": 0.001 + }, + "nodes.token_refiner": { + "max_abs": 0.2103524932861328, + "relative_l2": 0.010126469993441023 + }, + "nodes.token_refiner.grad": { + "max_abs": 42.15805328369141, + "relative_l2": 0.47900441737647814 + }, + "input_gradients.audio_hidden_states": { + "max_abs": 2186.4831494140626, + "relative_l2": 0.2894651121905773 + }, + "input_gradients.encoder_hidden_states": { + "max_abs": 126.93510711669921, + "relative_l2": 0.4752564996687594 + }, + "input_gradients.hidden_states": { + "max_abs": 1468.9695385742189, + "relative_l2": 0.28346112575648574 + }, + "parameter.audio_proj_in.bias": { + "max_abs": 1127.2670849609376, + "relative_l2": 0.27396454330724807 + }, + "parameter.audio_proj_in.weight": { + "max_abs": 2011.3793896484376, + "relative_l2": 0.30171097338711617 + }, + "parameter.audio_proj_out.bias": { + "max_abs": 1.0290945816040038, + "relative_l2": 0.025858806706272403 + }, + "parameter.audio_proj_out.weight": { + "max_abs": 5284.41125, + "relative_l2": 0.030712714343212797 + }, + "parameter.context_embedder.bias": { + "max_abs": 9.983069305419923, + "relative_l2": 0.5168372918664659 + }, + "parameter.context_embedder.weight": { + "max_abs": 23.560171012878417, + "relative_l2": 0.4766612868005037 + }, + "parameter.norm_out.linear.bias": { + "max_abs": 5.705576782226562, + "relative_l2": 0.028850151431610672 + }, + "parameter.norm_out.linear.weight": { + "max_abs": 1451.378046875, + "relative_l2": 0.028799590082716747 + }, + "parameter.norm_out.norm.weight": { + "max_abs": 1272.362421875, + "relative_l2": 0.02272375101128014 + }, + "parameter.proj_in.bias": { + "max_abs": 388.4746044921875, + "relative_l2": 0.33158579817413825 + }, + "parameter.proj_in.weight": { + "max_abs": 785.5183972167969, + "relative_l2": 0.3587452309712956 + }, + "parameter.proj_out.bias": { + "max_abs": 0.01868856430053711, + "relative_l2": 0.012918536528168385 + }, + "parameter.proj_out.weight": { + "max_abs": 12.085406188964845, + "relative_l2": 0.021091651937679996 + }, + "parameter.time_embedder.linear_1.bias": { + "max_abs": 783.52501953125, + "relative_l2": 0.04770286067239733 + }, + "parameter.time_embedder.linear_1.weight": { + "max_abs": 853.1036328125, + "relative_l2": 0.04727435712437821 + }, + "parameter.time_embedder.linear_2.bias": { + "max_abs": 1232.4355419921876, + "relative_l2": 0.33858927966178637 + }, + "parameter.time_embedder.linear_2.weight": { + "max_abs": 11108.18615234375, + "relative_l2": 0.2193347288362636 + }, + "parameter.token_refiner.final_norm.weight": { + "max_abs": 90.47127807617187, + "relative_l2": 0.4951302595087592 + }, + "parameter.token_refiner.refiner_blocks.0.attn.norm_k.weight": { + "max_abs": 124.09522140502929, + "relative_l2": 0.4046833363175061 + }, + "parameter.token_refiner.refiner_blocks.0.attn.norm_q.weight": { + "max_abs": 125.23231781005859, + "relative_l2": 0.40390545382323617 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_k.weight": { + "max_abs": 16.892093544006347, + "relative_l2": 0.47947822829994274 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_out.0.weight": { + "max_abs": 25.671869163513183, + "relative_l2": 0.49531255975895666 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_q.weight": { + "max_abs": 20.153124694824218, + "relative_l2": 0.46178624089996434 + }, + "parameter.token_refiner.refiner_blocks.0.attn.to_v.weight": { + "max_abs": 21.626427059173583, + "relative_l2": 0.4955083820299898 + }, + "parameter.token_refiner.refiner_blocks.0.ff.net.0.proj.weight": { + "max_abs": 33.706654434204104, + "relative_l2": 0.4879501687216877 + }, + "parameter.token_refiner.refiner_blocks.0.ff.net.2.weight": { + "max_abs": 70.30745780944824, + "relative_l2": 0.4949133242009952 + }, + "parameter.token_refiner.refiner_blocks.0.norm1.weight": { + "max_abs": 119.85018432617187, + "relative_l2": 0.49030330639891395 + }, + "parameter.token_refiner.refiner_blocks.0.norm2.weight": { + "max_abs": 174.6340817260742, + "relative_l2": 0.5023980028868472 + }, + "parameter.token_refiner.refiner_blocks.1.attn.norm_k.weight": { + "max_abs": 13.301282291412354, + "relative_l2": 0.6618802613160201 + }, + "parameter.token_refiner.refiner_blocks.1.attn.norm_q.weight": { + "max_abs": 13.395025138854981, + "relative_l2": 0.6626499140915075 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_k.weight": { + "max_abs": 2.337174837589264, + "relative_l2": 0.5330949325707235 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_out.0.weight": { + "max_abs": 5.516820077896118, + "relative_l2": 0.5079451521284352 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_q.weight": { + "max_abs": 3.0227862644195556, + "relative_l2": 0.5425959730719063 + }, + "parameter.token_refiner.refiner_blocks.1.attn.to_v.weight": { + "max_abs": 3.9670922803878783, + "relative_l2": 0.5165997257206122 + }, + "parameter.token_refiner.refiner_blocks.1.ff.net.0.proj.weight": { + "max_abs": 34.41474044799805, + "relative_l2": 0.5385256887991842 + }, + "parameter.token_refiner.refiner_blocks.1.ff.net.2.weight": { + "max_abs": 81.52852714538574, + "relative_l2": 0.5084387158519625 + }, + "parameter.token_refiner.refiner_blocks.1.norm1.weight": { + "max_abs": 43.96798217773438, + "relative_l2": 0.3892203088479711 + }, + "parameter.token_refiner.refiner_blocks.1.norm2.weight": { + "max_abs": 41.61128318786621, + "relative_l2": 0.55748139151743 + }, + "parameter.transformer_blocks.0.adaln_proj.linear.bias": { + "max_abs": 618.4180859375, + "relative_l2": 0.47075492675148795 + }, + "parameter.transformer_blocks.0.adaln_proj.linear.weight": { + "max_abs": 29961.772578125, + "relative_l2": 0.27301673947434835 + }, + "parameter.transformer_blocks.0.attn.norm_k.weight": { + "max_abs": 4459.5284375, + "relative_l2": 0.37966807565997 + }, + "parameter.transformer_blocks.0.attn.norm_q.weight": { + "max_abs": 4694.645380859375, + "relative_l2": 0.4193497891611958 + }, + "parameter.transformer_blocks.0.attn.to_k.weight": { + "max_abs": 2666.39416015625, + "relative_l2": 0.3421879560960825 + }, + "parameter.transformer_blocks.0.attn.to_out.0.weight": { + "max_abs": 1295.931025390625, + "relative_l2": 0.26791442696637935 + }, + "parameter.transformer_blocks.0.attn.to_q.weight": { + "max_abs": 3227.72814453125, + "relative_l2": 0.3076763356611271 + }, + "parameter.transformer_blocks.0.attn.to_v.weight": { + "max_abs": 967.9125012207031, + "relative_l2": 0.265496613904143 + }, + "parameter.transformer_blocks.0.ff.net.0.proj.weight": { + "max_abs": 5272.87560546875, + "relative_l2": 0.27539323856917125 + }, + "parameter.transformer_blocks.0.ff.net.2.weight": { + "max_abs": 931.70568359375, + "relative_l2": 0.27157714269997285 + }, + "parameter.transformer_blocks.0.norm1.weight": { + "max_abs": 2536.1047314453126, + "relative_l2": 0.30649315358287554 + }, + "parameter.transformer_blocks.0.norm2.weight": { + "max_abs": 8055.1895703125, + "relative_l2": 0.25356405562355117 + }, + "parameter.transformer_blocks.1.adaln_proj.linear.bias": { + "max_abs": 34.71782577514649, + "relative_l2": 0.28264463903051756 + }, + "parameter.transformer_blocks.1.adaln_proj.linear.weight": { + "max_abs": 8285.483515625, + "relative_l2": 0.2711459255786336 + }, + "parameter.transformer_blocks.1.attn.norm_k.weight": { + "max_abs": 9.839654808044434, + "relative_l2": 0.6866799937412663 + }, + "parameter.transformer_blocks.1.attn.norm_q.weight": { + "max_abs": 8.86622465133667, + "relative_l2": 0.6603229434845339 + }, + "parameter.transformer_blocks.1.attn.to_k.weight": { + "max_abs": 6.760801439285278, + "relative_l2": 0.8091230188837336 + }, + "parameter.transformer_blocks.1.attn.to_out.0.weight": { + "max_abs": 5.096917037963867, + "relative_l2": 0.716292400958965 + }, + "parameter.transformer_blocks.1.attn.to_q.weight": { + "max_abs": 6.934056406021118, + "relative_l2": 0.7213361749357194 + }, + "parameter.transformer_blocks.1.attn.to_v.weight": { + "max_abs": 5.85525155544281, + "relative_l2": 0.7914735134494816 + }, + "parameter.transformer_blocks.1.ff.net.0.proj.weight": { + "max_abs": 1124.0305126953126, + "relative_l2": 0.27999278243568926 + }, + "parameter.transformer_blocks.1.ff.net.2.weight": { + "max_abs": 722.3081005859375, + "relative_l2": 0.2676357113339872 + }, + "parameter.transformer_blocks.1.norm1.weight": { + "max_abs": 19.10787570953369, + "relative_l2": 0.7871692466073358 + }, + "parameter.transformer_blocks.1.norm2.weight": { + "max_abs": 1245.1071728515626, + "relative_l2": 0.2617261401281995 + }, + "parameter.transformer_blocks.10.adaln_proj.linear.bias": { + "max_abs": 7.491905097961426, + "relative_l2": 0.27412761443613665 + }, + "parameter.transformer_blocks.10.adaln_proj.linear.weight": { + "max_abs": 1787.8017529296876, + "relative_l2": 0.2485811102828138 + }, + "parameter.transformer_blocks.10.attn.norm_k.weight": { + "max_abs": 3.919461612701416, + "relative_l2": 0.8939340974928423 + }, + "parameter.transformer_blocks.10.attn.norm_q.weight": { + "max_abs": 3.357305293083191, + "relative_l2": 0.91772600843942 + }, + "parameter.transformer_blocks.10.attn.to_k.weight": { + "max_abs": 1.9735518741607665, + "relative_l2": 0.7451501197444428 + }, + "parameter.transformer_blocks.10.attn.to_out.0.weight": { + "max_abs": 0.7348604547977448, + "relative_l2": 0.6738693588687075 + }, + "parameter.transformer_blocks.10.attn.to_q.weight": { + "max_abs": 1.4673158574104308, + "relative_l2": 0.7542932763495229 + }, + "parameter.transformer_blocks.10.attn.to_v.weight": { + "max_abs": 1.669773106575012, + "relative_l2": 0.7198321350834086 + }, + "parameter.transformer_blocks.10.ff.net.0.proj.weight": { + "max_abs": 946.703486328125, + "relative_l2": 0.2748115050033871 + }, + "parameter.transformer_blocks.10.ff.net.2.weight": { + "max_abs": 281.16688049316406, + "relative_l2": 0.23077566506659983 + }, + "parameter.transformer_blocks.10.norm1.weight": { + "max_abs": 6.122771863937378, + "relative_l2": 0.8132083951553788 + }, + "parameter.transformer_blocks.10.norm2.weight": { + "max_abs": 2005.6719921875, + "relative_l2": 0.2858285655459195 + }, + "parameter.transformer_blocks.11.adaln_proj.linear.bias": { + "max_abs": 10.026488189697266, + "relative_l2": 0.24650859743948983 + }, + "parameter.transformer_blocks.11.adaln_proj.linear.weight": { + "max_abs": 2381.72375, + "relative_l2": 0.23093499807756118 + }, + "parameter.transformer_blocks.11.attn.norm_k.weight": { + "max_abs": 2.3551847982406615, + "relative_l2": 0.9131311070090182 + }, + "parameter.transformer_blocks.11.attn.norm_q.weight": { + "max_abs": 2.4871248054504393, + "relative_l2": 0.9289690749562688 + }, + "parameter.transformer_blocks.11.attn.to_k.weight": { + "max_abs": 1.9326666045188903, + "relative_l2": 0.831162001950036 + }, + "parameter.transformer_blocks.11.attn.to_out.0.weight": { + "max_abs": 0.6217981028556824, + "relative_l2": 0.6472011831701678 + }, + "parameter.transformer_blocks.11.attn.to_q.weight": { + "max_abs": 2.5200601863861083, + "relative_l2": 0.8275675088169366 + }, + "parameter.transformer_blocks.11.attn.to_v.weight": { + "max_abs": 0.9879325973987579, + "relative_l2": 0.7010794076362376 + }, + "parameter.transformer_blocks.11.ff.net.0.proj.weight": { + "max_abs": 681.3028210449219, + "relative_l2": 0.24425111720111248 + }, + "parameter.transformer_blocks.11.ff.net.2.weight": { + "max_abs": 300.01104248046875, + "relative_l2": 0.2191100526426731 + }, + "parameter.transformer_blocks.11.norm1.weight": { + "max_abs": 4.620777015686035, + "relative_l2": 0.7356413076782026 + }, + "parameter.transformer_blocks.11.norm2.weight": { + "max_abs": 596.8696850585937, + "relative_l2": 0.14912470039405729 + }, + "parameter.transformer_blocks.12.adaln_proj.linear.bias": { + "max_abs": 8.121762161254884, + "relative_l2": 0.2550387695831746 + }, + "parameter.transformer_blocks.12.adaln_proj.linear.weight": { + "max_abs": 1945.74669921875, + "relative_l2": 0.2367571483337047 + }, + "parameter.transformer_blocks.12.attn.norm_k.weight": { + "max_abs": 2.9853444147109984, + "relative_l2": 0.7001627545316077 + }, + "parameter.transformer_blocks.12.attn.norm_q.weight": { + "max_abs": 3.0688916492462157, + "relative_l2": 0.7122141032106739 + }, + "parameter.transformer_blocks.12.attn.to_k.weight": { + "max_abs": 1.6972085285186767, + "relative_l2": 0.5879689607748688 + }, + "parameter.transformer_blocks.12.attn.to_out.0.weight": { + "max_abs": 0.49494848132133484, + "relative_l2": 0.5773987091429704 + }, + "parameter.transformer_blocks.12.attn.to_q.weight": { + "max_abs": 2.5782805728912352, + "relative_l2": 0.6454071223125697 + }, + "parameter.transformer_blocks.12.attn.to_v.weight": { + "max_abs": 0.6551486897468567, + "relative_l2": 0.6474215209980447 + }, + "parameter.transformer_blocks.12.ff.net.0.proj.weight": { + "max_abs": 354.23184478759765, + "relative_l2": 0.26331063958954565 + }, + "parameter.transformer_blocks.12.ff.net.2.weight": { + "max_abs": 130.85853088378906, + "relative_l2": 0.2240588958162464 + }, + "parameter.transformer_blocks.12.norm1.weight": { + "max_abs": 2.7988663363456725, + "relative_l2": 0.6063861929230794 + }, + "parameter.transformer_blocks.12.norm2.weight": { + "max_abs": 789.6607006835937, + "relative_l2": 0.2711719907576323 + }, + "parameter.transformer_blocks.13.adaln_proj.linear.bias": { + "max_abs": 9.051812057495118, + "relative_l2": 0.2594281777634543 + }, + "parameter.transformer_blocks.13.adaln_proj.linear.weight": { + "max_abs": 2153.95958984375, + "relative_l2": 0.23677682551061446 + }, + "parameter.transformer_blocks.13.attn.norm_k.weight": { + "max_abs": 3.3524254608154296, + "relative_l2": 0.8954921306191933 + }, + "parameter.transformer_blocks.13.attn.norm_q.weight": { + "max_abs": 2.6193434047698974, + "relative_l2": 0.8582240172254604 + }, + "parameter.transformer_blocks.13.attn.to_k.weight": { + "max_abs": 1.079060022830963, + "relative_l2": 0.8219867581710362 + }, + "parameter.transformer_blocks.13.attn.to_out.0.weight": { + "max_abs": 0.7083306360244751, + "relative_l2": 0.8192653910512632 + }, + "parameter.transformer_blocks.13.attn.to_q.weight": { + "max_abs": 1.8803073811531066, + "relative_l2": 0.7905141412197559 + }, + "parameter.transformer_blocks.13.attn.to_v.weight": { + "max_abs": 1.1077276515960692, + "relative_l2": 0.8386404244712579 + }, + "parameter.transformer_blocks.13.ff.net.0.proj.weight": { + "max_abs": 227.7526806640625, + "relative_l2": 0.2558986494449096 + }, + "parameter.transformer_blocks.13.ff.net.2.weight": { + "max_abs": 218.08233154296875, + "relative_l2": 0.22186692828327384 + }, + "parameter.transformer_blocks.13.norm1.weight": { + "max_abs": 3.649041299819946, + "relative_l2": 0.8200788291571931 + }, + "parameter.transformer_blocks.13.norm2.weight": { + "max_abs": 574.981806640625, + "relative_l2": 0.24275651736493498 + }, + "parameter.transformer_blocks.14.adaln_proj.linear.bias": { + "max_abs": 6.945229415893555, + "relative_l2": 0.26591723008762436 + }, + "parameter.transformer_blocks.14.adaln_proj.linear.weight": { + "max_abs": 1645.763544921875, + "relative_l2": 0.24360616649564137 + }, + "parameter.transformer_blocks.14.attn.norm_k.weight": { + "max_abs": 2.0977778482437133, + "relative_l2": 0.6688106360712017 + }, + "parameter.transformer_blocks.14.attn.norm_q.weight": { + "max_abs": 1.3061669516563414, + "relative_l2": 0.6895363279211654 + }, + "parameter.transformer_blocks.14.attn.to_k.weight": { + "max_abs": 1.5680741119384765, + "relative_l2": 0.6876301501287961 + }, + "parameter.transformer_blocks.14.attn.to_out.0.weight": { + "max_abs": 0.5789981383085251, + "relative_l2": 0.6718812940143374 + }, + "parameter.transformer_blocks.14.attn.to_q.weight": { + "max_abs": 1.2793809819221496, + "relative_l2": 0.7055059914862861 + }, + "parameter.transformer_blocks.14.attn.to_v.weight": { + "max_abs": 1.4136269736289977, + "relative_l2": 0.758456532499533 + }, + "parameter.transformer_blocks.14.ff.net.0.proj.weight": { + "max_abs": 358.0343731689453, + "relative_l2": 0.2506961598008713 + }, + "parameter.transformer_blocks.14.ff.net.2.weight": { + "max_abs": 138.15642822265625, + "relative_l2": 0.22120862426841734 + }, + "parameter.transformer_blocks.14.norm1.weight": { + "max_abs": 3.1570816087722777, + "relative_l2": 0.7692865880722793 + }, + "parameter.transformer_blocks.14.norm2.weight": { + "max_abs": 827.2231396484375, + "relative_l2": 0.2948949400546231 + }, + "parameter.transformer_blocks.15.adaln_proj.linear.bias": { + "max_abs": 8.299082641601563, + "relative_l2": 0.25602805026332515 + }, + "parameter.transformer_blocks.15.adaln_proj.linear.weight": { + "max_abs": 2009.92736328125, + "relative_l2": 0.23283001516518612 + }, + "parameter.transformer_blocks.15.attn.norm_k.weight": { + "max_abs": 2.2380307364463805, + "relative_l2": 0.7711324359701877 + }, + "parameter.transformer_blocks.15.attn.norm_q.weight": { + "max_abs": 2.4794930267333983, + "relative_l2": 0.7798992025704999 + }, + "parameter.transformer_blocks.15.attn.to_k.weight": { + "max_abs": 1.2511296319961547, + "relative_l2": 0.7908391612072826 + }, + "parameter.transformer_blocks.15.attn.to_out.0.weight": { + "max_abs": 0.8765870380401611, + "relative_l2": 0.6449978432370598 + }, + "parameter.transformer_blocks.15.attn.to_q.weight": { + "max_abs": 1.296251540184021, + "relative_l2": 0.7386447155926827 + }, + "parameter.transformer_blocks.15.attn.to_v.weight": { + "max_abs": 0.8481322872638702, + "relative_l2": 0.7339279795477406 + }, + "parameter.transformer_blocks.15.ff.net.0.proj.weight": { + "max_abs": 995.5617932128906, + "relative_l2": 0.24840504785585898 + }, + "parameter.transformer_blocks.15.ff.net.2.weight": { + "max_abs": 154.35751998901367, + "relative_l2": 0.22302031190872912 + }, + "parameter.transformer_blocks.15.norm1.weight": { + "max_abs": 4.56264889717102, + "relative_l2": 0.8064254345561096 + }, + "parameter.transformer_blocks.15.norm2.weight": { + "max_abs": 2107.9294384765626, + "relative_l2": 0.3168363717864166 + }, + "parameter.transformer_blocks.16.adaln_proj.linear.bias": { + "max_abs": 7.884571199417114, + "relative_l2": 0.2628938003490684 + }, + "parameter.transformer_blocks.16.adaln_proj.linear.weight": { + "max_abs": 1879.1367749023439, + "relative_l2": 0.2325869713721619 + }, + "parameter.transformer_blocks.16.attn.norm_k.weight": { + "max_abs": 2.1271226453781127, + "relative_l2": 0.7322144551851629 + }, + "parameter.transformer_blocks.16.attn.norm_q.weight": { + "max_abs": 2.4217547225952147, + "relative_l2": 0.7826353295239833 + }, + "parameter.transformer_blocks.16.attn.to_k.weight": { + "max_abs": 1.7393020439147948, + "relative_l2": 0.7223069267374671 + }, + "parameter.transformer_blocks.16.attn.to_out.0.weight": { + "max_abs": 0.5886289674043655, + "relative_l2": 0.5151741110535427 + }, + "parameter.transformer_blocks.16.attn.to_q.weight": { + "max_abs": 1.2670303034782409, + "relative_l2": 0.7348671488129176 + }, + "parameter.transformer_blocks.16.attn.to_v.weight": { + "max_abs": 0.9935271489620209, + "relative_l2": 0.6170055938519416 + }, + "parameter.transformer_blocks.16.ff.net.0.proj.weight": { + "max_abs": 384.8454083251953, + "relative_l2": 0.24667237636020004 + }, + "parameter.transformer_blocks.16.ff.net.2.weight": { + "max_abs": 133.7273129272461, + "relative_l2": 0.20971634566923797 + }, + "parameter.transformer_blocks.16.norm1.weight": { + "max_abs": 4.996379976272583, + "relative_l2": 0.7024999699709787 + }, + "parameter.transformer_blocks.16.norm2.weight": { + "max_abs": 1165.293330078125, + "relative_l2": 0.24776021583071195 + }, + "parameter.transformer_blocks.17.adaln_proj.linear.bias": { + "max_abs": 13.588038330078126, + "relative_l2": 0.24125780352307197 + }, + "parameter.transformer_blocks.17.adaln_proj.linear.weight": { + "max_abs": 3197.22716796875, + "relative_l2": 0.2213120876089823 + }, + "parameter.transformer_blocks.17.attn.norm_k.weight": { + "max_abs": 2.577016716003418, + "relative_l2": 0.7274001496154391 + }, + "parameter.transformer_blocks.17.attn.norm_q.weight": { + "max_abs": 2.37910831451416, + "relative_l2": 0.7732131521152253 + }, + "parameter.transformer_blocks.17.attn.to_k.weight": { + "max_abs": 1.0713100481033324, + "relative_l2": 0.7684038527415521 + }, + "parameter.transformer_blocks.17.attn.to_out.0.weight": { + "max_abs": 0.77244544506073, + "relative_l2": 0.7915539672303255 + }, + "parameter.transformer_blocks.17.attn.to_q.weight": { + "max_abs": 1.3973073410987853, + "relative_l2": 0.7598057776059467 + }, + "parameter.transformer_blocks.17.attn.to_v.weight": { + "max_abs": 0.9992235350608826, + "relative_l2": 0.8915697210644591 + }, + "parameter.transformer_blocks.17.ff.net.0.proj.weight": { + "max_abs": 182.5413426208496, + "relative_l2": 0.23066209417021438 + }, + "parameter.transformer_blocks.17.ff.net.2.weight": { + "max_abs": 367.0468701171875, + "relative_l2": 0.20430895173490748 + }, + "parameter.transformer_blocks.17.norm1.weight": { + "max_abs": 3.7926521348953246, + "relative_l2": 0.8373703936130231 + }, + "parameter.transformer_blocks.17.norm2.weight": { + "max_abs": 863.1227978515625, + "relative_l2": 0.23192950641814164 + }, + "parameter.transformer_blocks.18.adaln_proj.linear.bias": { + "max_abs": 11.117765312194825, + "relative_l2": 0.22011060082853495 + }, + "parameter.transformer_blocks.18.adaln_proj.linear.weight": { + "max_abs": 2647.5764111328126, + "relative_l2": 0.20970623162134797 + }, + "parameter.transformer_blocks.18.attn.norm_k.weight": { + "max_abs": 1.8815052556991576, + "relative_l2": 0.7812328618762753 + }, + "parameter.transformer_blocks.18.attn.norm_q.weight": { + "max_abs": 0.9511458325386047, + "relative_l2": 0.7175421337378723 + }, + "parameter.transformer_blocks.18.attn.to_k.weight": { + "max_abs": 0.9281016337871552, + "relative_l2": 0.7242657368447106 + }, + "parameter.transformer_blocks.18.attn.to_out.0.weight": { + "max_abs": 0.6100253510475159, + "relative_l2": 0.6399585470232062 + }, + "parameter.transformer_blocks.18.attn.to_q.weight": { + "max_abs": 0.8241854357719421, + "relative_l2": 0.7026953919302968 + }, + "parameter.transformer_blocks.18.attn.to_v.weight": { + "max_abs": 1.0181871104240416, + "relative_l2": 0.7925326835764924 + }, + "parameter.transformer_blocks.18.ff.net.0.proj.weight": { + "max_abs": 870.2954968261719, + "relative_l2": 0.21723119562526702 + }, + "parameter.transformer_blocks.18.ff.net.2.weight": { + "max_abs": 535.0214184570312, + "relative_l2": 0.18686082434706205 + }, + "parameter.transformer_blocks.18.norm1.weight": { + "max_abs": 1.7391082096099852, + "relative_l2": 0.7427907656180031 + }, + "parameter.transformer_blocks.18.norm2.weight": { + "max_abs": 760.1420849609375, + "relative_l2": 0.20643528437354 + }, + "parameter.transformer_blocks.19.adaln_proj.linear.bias": { + "max_abs": 9.754404907226563, + "relative_l2": 0.21710379129855412 + }, + "parameter.transformer_blocks.19.adaln_proj.linear.weight": { + "max_abs": 2309.264521484375, + "relative_l2": 0.2026108967929386 + }, + "parameter.transformer_blocks.19.attn.norm_k.weight": { + "max_abs": 3.3605203676223754, + "relative_l2": 0.9235690859452362 + }, + "parameter.transformer_blocks.19.attn.norm_q.weight": { + "max_abs": 3.3958607959747313, + "relative_l2": 0.9027156051345979 + }, + "parameter.transformer_blocks.19.attn.to_k.weight": { + "max_abs": 2.5811061906814574, + "relative_l2": 0.9278696870441093 + }, + "parameter.transformer_blocks.19.attn.to_out.0.weight": { + "max_abs": 0.9226974296569824, + "relative_l2": 0.6308925415130978 + }, + "parameter.transformer_blocks.19.attn.to_q.weight": { + "max_abs": 1.8905621814727782, + "relative_l2": 0.9197084146310966 + }, + "parameter.transformer_blocks.19.attn.to_v.weight": { + "max_abs": 1.4341714429855346, + "relative_l2": 0.7554498717910688 + }, + "parameter.transformer_blocks.19.ff.net.0.proj.weight": { + "max_abs": 428.02108276367187, + "relative_l2": 0.19810123483654649 + }, + "parameter.transformer_blocks.19.ff.net.2.weight": { + "max_abs": 348.52154052734375, + "relative_l2": 0.1570854430188067 + }, + "parameter.transformer_blocks.19.norm1.weight": { + "max_abs": 10.131331806182862, + "relative_l2": 0.9411943421518384 + }, + "parameter.transformer_blocks.19.norm2.weight": { + "max_abs": 1661.620478515625, + "relative_l2": 0.2312924067876611 + }, + "parameter.transformer_blocks.2.adaln_proj.linear.bias": { + "max_abs": 38.34929931640625, + "relative_l2": 0.27099813831084796 + }, + "parameter.transformer_blocks.2.adaln_proj.linear.weight": { + "max_abs": 9139.78625, + "relative_l2": 0.2610091571123563 + }, + "parameter.transformer_blocks.2.attn.norm_k.weight": { + "max_abs": 11.38742530822754, + "relative_l2": 0.7203443894744488 + }, + "parameter.transformer_blocks.2.attn.norm_q.weight": { + "max_abs": 10.698475246429444, + "relative_l2": 0.7308512325109128 + }, + "parameter.transformer_blocks.2.attn.to_k.weight": { + "max_abs": 7.837068204879761, + "relative_l2": 0.7116781769614272 + }, + "parameter.transformer_blocks.2.attn.to_out.0.weight": { + "max_abs": 3.5596836853027343, + "relative_l2": 0.5998331316236671 + }, + "parameter.transformer_blocks.2.attn.to_q.weight": { + "max_abs": 7.5656751441955565, + "relative_l2": 0.6356281954464494 + }, + "parameter.transformer_blocks.2.attn.to_v.weight": { + "max_abs": 7.642959957122803, + "relative_l2": 0.7641174143769899 + }, + "parameter.transformer_blocks.2.ff.net.0.proj.weight": { + "max_abs": 835.4142712402344, + "relative_l2": 0.2686617582206923 + }, + "parameter.transformer_blocks.2.ff.net.2.weight": { + "max_abs": 552.6587780761719, + "relative_l2": 0.25879537458619123 + }, + "parameter.transformer_blocks.2.norm1.weight": { + "max_abs": 13.009734039306641, + "relative_l2": 0.6925190518757222 + }, + "parameter.transformer_blocks.2.norm2.weight": { + "max_abs": 1900.908564453125, + "relative_l2": 0.2736282238955311 + }, + "parameter.transformer_blocks.20.adaln_proj.linear.bias": { + "max_abs": 11.688382034301759, + "relative_l2": 0.2377400113550496 + }, + "parameter.transformer_blocks.20.adaln_proj.linear.weight": { + "max_abs": 2782.34533203125, + "relative_l2": 0.20709904622478764 + }, + "parameter.transformer_blocks.20.attn.norm_k.weight": { + "max_abs": 2.9652954387664794, + "relative_l2": 1.1392256242682777 + }, + "parameter.transformer_blocks.20.attn.norm_q.weight": { + "max_abs": 2.5787290382385253, + "relative_l2": 1.1756615296147424 + }, + "parameter.transformer_blocks.20.attn.to_k.weight": { + "max_abs": 2.683075909614563, + "relative_l2": 1.1153968264775866 + }, + "parameter.transformer_blocks.20.attn.to_out.0.weight": { + "max_abs": 0.7956968891620636, + "relative_l2": 0.528389423061085 + }, + "parameter.transformer_blocks.20.attn.to_q.weight": { + "max_abs": 1.4907193827629088, + "relative_l2": 1.0275720720510009 + }, + "parameter.transformer_blocks.20.attn.to_v.weight": { + "max_abs": 1.0080752718448638, + "relative_l2": 0.6380884034094756 + }, + "parameter.transformer_blocks.20.ff.net.0.proj.weight": { + "max_abs": 220.31978881835937, + "relative_l2": 0.21520136687725197 + }, + "parameter.transformer_blocks.20.ff.net.2.weight": { + "max_abs": 286.46563232421875, + "relative_l2": 0.1794734190533937 + }, + "parameter.transformer_blocks.20.norm1.weight": { + "max_abs": 6.677889709472656, + "relative_l2": 0.9076549223343746 + }, + "parameter.transformer_blocks.20.norm2.weight": { + "max_abs": 384.90670288085937, + "relative_l2": 0.21424578336514133 + }, + "parameter.transformer_blocks.21.adaln_proj.linear.bias": { + "max_abs": 6.576134567260742, + "relative_l2": 0.22732946736872406 + }, + "parameter.transformer_blocks.21.adaln_proj.linear.weight": { + "max_abs": 1560.97765625, + "relative_l2": 0.2000087574560543 + }, + "parameter.transformer_blocks.21.attn.norm_k.weight": { + "max_abs": 3.070456986427307, + "relative_l2": 0.9745551614882625 + }, + "parameter.transformer_blocks.21.attn.norm_q.weight": { + "max_abs": 2.8962028074264525, + "relative_l2": 0.9897367111919141 + }, + "parameter.transformer_blocks.21.attn.to_k.weight": { + "max_abs": 1.5696979212760924, + "relative_l2": 0.9636766966378697 + }, + "parameter.transformer_blocks.21.attn.to_out.0.weight": { + "max_abs": 0.8160495805740356, + "relative_l2": 0.6789384150673038 + }, + "parameter.transformer_blocks.21.attn.to_q.weight": { + "max_abs": 1.102442626953125, + "relative_l2": 0.9350146184006228 + }, + "parameter.transformer_blocks.21.attn.to_v.weight": { + "max_abs": 1.402775650024414, + "relative_l2": 0.8394066497398459 + }, + "parameter.transformer_blocks.21.ff.net.0.proj.weight": { + "max_abs": 172.8049038696289, + "relative_l2": 0.2180061256659362 + }, + "parameter.transformer_blocks.21.ff.net.2.weight": { + "max_abs": 207.50718963623046, + "relative_l2": 0.18204959014623456 + }, + "parameter.transformer_blocks.21.norm1.weight": { + "max_abs": 5.1228568601608275, + "relative_l2": 0.8691599430647023 + }, + "parameter.transformer_blocks.21.norm2.weight": { + "max_abs": 395.5959423828125, + "relative_l2": 0.22020446724135123 + }, + "parameter.transformer_blocks.22.adaln_proj.linear.bias": { + "max_abs": 7.377956275939941, + "relative_l2": 0.22904466456764788 + }, + "parameter.transformer_blocks.22.adaln_proj.linear.weight": { + "max_abs": 1733.79845703125, + "relative_l2": 0.2036469103690055 + }, + "parameter.transformer_blocks.22.attn.norm_k.weight": { + "max_abs": 2.5229035067558288, + "relative_l2": 0.9044992781713185 + }, + "parameter.transformer_blocks.22.attn.norm_q.weight": { + "max_abs": 2.3648268032073974, + "relative_l2": 0.8756975510736572 + }, + "parameter.transformer_blocks.22.attn.to_k.weight": { + "max_abs": 1.9956800985336303, + "relative_l2": 0.85522739174773 + }, + "parameter.transformer_blocks.22.attn.to_out.0.weight": { + "max_abs": 0.7181864130496979, + "relative_l2": 0.8913080694256287 + }, + "parameter.transformer_blocks.22.attn.to_q.weight": { + "max_abs": 0.8015691387653351, + "relative_l2": 0.8314708081082037 + }, + "parameter.transformer_blocks.22.attn.to_v.weight": { + "max_abs": 2.644535069465637, + "relative_l2": 1.0197672534558853 + }, + "parameter.transformer_blocks.22.ff.net.0.proj.weight": { + "max_abs": 357.3243588256836, + "relative_l2": 0.2189386423066487 + }, + "parameter.transformer_blocks.22.ff.net.2.weight": { + "max_abs": 207.0422848510742, + "relative_l2": 0.1754032335951697 + }, + "parameter.transformer_blocks.22.norm1.weight": { + "max_abs": 3.6996399211883544, + "relative_l2": 0.9034981506161427 + }, + "parameter.transformer_blocks.22.norm2.weight": { + "max_abs": 522.5956982421875, + "relative_l2": 0.21523293735280408 + }, + "parameter.transformer_blocks.23.adaln_proj.linear.bias": { + "max_abs": 8.08806360244751, + "relative_l2": 0.21326685878132157 + }, + "parameter.transformer_blocks.23.adaln_proj.linear.weight": { + "max_abs": 1922.5801953125, + "relative_l2": 0.19217049180594867 + }, + "parameter.transformer_blocks.23.attn.norm_k.weight": { + "max_abs": 2.515032832622528, + "relative_l2": 1.0121865478435368 + }, + "parameter.transformer_blocks.23.attn.norm_q.weight": { + "max_abs": 2.541628723144531, + "relative_l2": 1.0178988868215497 + }, + "parameter.transformer_blocks.23.attn.to_k.weight": { + "max_abs": 3.535975103378296, + "relative_l2": 0.9932593402245732 + }, + "parameter.transformer_blocks.23.attn.to_out.0.weight": { + "max_abs": 0.48513928532600403, + "relative_l2": 0.644548323931213 + }, + "parameter.transformer_blocks.23.attn.to_q.weight": { + "max_abs": 1.0787883150577544, + "relative_l2": 0.9093563448523482 + }, + "parameter.transformer_blocks.23.attn.to_v.weight": { + "max_abs": 1.7548005628585814, + "relative_l2": 0.8057142720129602 + }, + "parameter.transformer_blocks.23.ff.net.0.proj.weight": { + "max_abs": 239.49857604980468, + "relative_l2": 0.20468404153122077 + }, + "parameter.transformer_blocks.23.ff.net.2.weight": { + "max_abs": 293.54357421875, + "relative_l2": 0.17651630716934855 + }, + "parameter.transformer_blocks.23.norm1.weight": { + "max_abs": 6.985463981628418, + "relative_l2": 1.0053685364073084 + }, + "parameter.transformer_blocks.23.norm2.weight": { + "max_abs": 456.5930432128906, + "relative_l2": 0.21848467457692144 + }, + "parameter.transformer_blocks.24.adaln_proj.linear.bias": { + "max_abs": 7.304553871154785, + "relative_l2": 0.21564167899972708 + }, + "parameter.transformer_blocks.24.adaln_proj.linear.weight": { + "max_abs": 1737.616572265625, + "relative_l2": 0.1901697069608287 + }, + "parameter.transformer_blocks.24.attn.norm_k.weight": { + "max_abs": 2.216841106414795, + "relative_l2": 0.7926188934928867 + }, + "parameter.transformer_blocks.24.attn.norm_q.weight": { + "max_abs": 1.5476918268203734, + "relative_l2": 0.7777840222904054 + }, + "parameter.transformer_blocks.24.attn.to_k.weight": { + "max_abs": 1.8396434831619262, + "relative_l2": 0.7579427336702298 + }, + "parameter.transformer_blocks.24.attn.to_out.0.weight": { + "max_abs": 0.9694585251808167, + "relative_l2": 0.5167980950331458 + }, + "parameter.transformer_blocks.24.attn.to_q.weight": { + "max_abs": 0.9235194373130798, + "relative_l2": 0.7017912116896795 + }, + "parameter.transformer_blocks.24.attn.to_v.weight": { + "max_abs": 1.2338300991058349, + "relative_l2": 0.6999683243661816 + }, + "parameter.transformer_blocks.24.ff.net.0.proj.weight": { + "max_abs": 487.1518811035156, + "relative_l2": 0.20379978037497748 + }, + "parameter.transformer_blocks.24.ff.net.2.weight": { + "max_abs": 556.0626171875, + "relative_l2": 0.13101262792454788 + }, + "parameter.transformer_blocks.24.norm1.weight": { + "max_abs": 6.81080997467041, + "relative_l2": 0.7571361191504765 + }, + "parameter.transformer_blocks.24.norm2.weight": { + "max_abs": 800.630732421875, + "relative_l2": 0.20475600692712043 + }, + "parameter.transformer_blocks.25.adaln_proj.linear.bias": { + "max_abs": 8.371666793823243, + "relative_l2": 0.16869321168645987 + }, + "parameter.transformer_blocks.25.adaln_proj.linear.weight": { + "max_abs": 1985.033076171875, + "relative_l2": 0.15815527032757656 + }, + "parameter.transformer_blocks.25.attn.norm_k.weight": { + "max_abs": 2.4211975383758544, + "relative_l2": 0.8235665728543179 + }, + "parameter.transformer_blocks.25.attn.norm_q.weight": { + "max_abs": 2.205787365436554, + "relative_l2": 0.8263968452846001 + }, + "parameter.transformer_blocks.25.attn.to_k.weight": { + "max_abs": 2.5189250755310058, + "relative_l2": 0.8722942908154744 + }, + "parameter.transformer_blocks.25.attn.to_out.0.weight": { + "max_abs": 0.9047119665145874, + "relative_l2": 0.6978460086808806 + }, + "parameter.transformer_blocks.25.attn.to_q.weight": { + "max_abs": 1.404982750415802, + "relative_l2": 0.8480601970016735 + }, + "parameter.transformer_blocks.25.attn.to_v.weight": { + "max_abs": 2.1768286275863646, + "relative_l2": 0.8885601899867497 + }, + "parameter.transformer_blocks.25.ff.net.0.proj.weight": { + "max_abs": 408.89319885253906, + "relative_l2": 0.1644225805549453 + }, + "parameter.transformer_blocks.25.ff.net.2.weight": { + "max_abs": 753.0723828125, + "relative_l2": 0.11861937989316992 + }, + "parameter.transformer_blocks.25.norm1.weight": { + "max_abs": 3.4406565475463866, + "relative_l2": 0.7214860130902033 + }, + "parameter.transformer_blocks.25.norm2.weight": { + "max_abs": 537.69640625, + "relative_l2": 0.15529167748626888 + }, + "parameter.transformer_blocks.26.adaln_proj.linear.bias": { + "max_abs": 4.974030380249023, + "relative_l2": 0.1970655795477857 + }, + "parameter.transformer_blocks.26.adaln_proj.linear.weight": { + "max_abs": 1166.8911083984376, + "relative_l2": 0.16832694903177375 + }, + "parameter.transformer_blocks.26.attn.norm_k.weight": { + "max_abs": 3.5365945148468017, + "relative_l2": 0.8607390883814103 + }, + "parameter.transformer_blocks.26.attn.norm_q.weight": { + "max_abs": 3.5158914375305175, + "relative_l2": 0.863374440509692 + }, + "parameter.transformer_blocks.26.attn.to_k.weight": { + "max_abs": 3.6291722106933593, + "relative_l2": 0.8664953205185234 + }, + "parameter.transformer_blocks.26.attn.to_out.0.weight": { + "max_abs": 1.0129927444458007, + "relative_l2": 0.2929568074262903 + }, + "parameter.transformer_blocks.26.attn.to_q.weight": { + "max_abs": 1.9735407876968383, + "relative_l2": 0.8663906879530491 + }, + "parameter.transformer_blocks.26.attn.to_v.weight": { + "max_abs": 2.1177463579177855, + "relative_l2": 0.4406044410223554 + }, + "parameter.transformer_blocks.26.ff.net.0.proj.weight": { + "max_abs": 371.87398803710937, + "relative_l2": 0.1798924973205361 + }, + "parameter.transformer_blocks.26.ff.net.2.weight": { + "max_abs": 457.3574169921875, + "relative_l2": 0.12919007812210365 + }, + "parameter.transformer_blocks.26.norm1.weight": { + "max_abs": 7.477498416900635, + "relative_l2": 0.689315960809784 + }, + "parameter.transformer_blocks.26.norm2.weight": { + "max_abs": 594.1813305664062, + "relative_l2": 0.18823506950940896 + }, + "parameter.transformer_blocks.27.adaln_proj.linear.bias": { + "max_abs": 5.869821672439575, + "relative_l2": 0.21336485205109756 + }, + "parameter.transformer_blocks.27.adaln_proj.linear.weight": { + "max_abs": 1084.662470703125, + "relative_l2": 0.1741348330090017 + }, + "parameter.transformer_blocks.27.attn.norm_k.weight": { + "max_abs": 3.150988345146179, + "relative_l2": 1.062841235494389 + }, + "parameter.transformer_blocks.27.attn.norm_q.weight": { + "max_abs": 3.6676880168914794, + "relative_l2": 1.0967533088519172 + }, + "parameter.transformer_blocks.27.attn.to_k.weight": { + "max_abs": 4.32183874130249, + "relative_l2": 1.0950772065469172 + }, + "parameter.transformer_blocks.27.attn.to_out.0.weight": { + "max_abs": 0.6107107448577881, + "relative_l2": 0.8264677710738436 + }, + "parameter.transformer_blocks.27.attn.to_q.weight": { + "max_abs": 3.3824980306625365, + "relative_l2": 1.070568942653313 + }, + "parameter.transformer_blocks.27.attn.to_v.weight": { + "max_abs": 4.458062772750854, + "relative_l2": 0.9492867035892784 + }, + "parameter.transformer_blocks.27.ff.net.0.proj.weight": { + "max_abs": 259.40442932128906, + "relative_l2": 0.19459347336082505 + }, + "parameter.transformer_blocks.27.ff.net.2.weight": { + "max_abs": 125.5402783203125, + "relative_l2": 0.15380201598069254 + }, + "parameter.transformer_blocks.27.norm1.weight": { + "max_abs": 11.696521644592286, + "relative_l2": 0.962568903689285 + }, + "parameter.transformer_blocks.27.norm2.weight": { + "max_abs": 630.013056640625, + "relative_l2": 0.18498388565940183 + }, + "parameter.transformer_blocks.28.adaln_proj.linear.bias": { + "max_abs": 6.271578674316406, + "relative_l2": 0.20026362707009437 + }, + "parameter.transformer_blocks.28.adaln_proj.linear.weight": { + "max_abs": 1470.62609375, + "relative_l2": 0.156969272399435 + }, + "parameter.transformer_blocks.28.attn.norm_k.weight": { + "max_abs": 3.694087390899658, + "relative_l2": 1.1393344702953399 + }, + "parameter.transformer_blocks.28.attn.norm_q.weight": { + "max_abs": 3.2554990100860595, + "relative_l2": 1.132576919349636 + }, + "parameter.transformer_blocks.28.attn.to_k.weight": { + "max_abs": 2.87049311876297, + "relative_l2": 1.190224288047722 + }, + "parameter.transformer_blocks.28.attn.to_out.0.weight": { + "max_abs": 1.1344015049934386, + "relative_l2": 1.1170925582000946 + }, + "parameter.transformer_blocks.28.attn.to_q.weight": { + "max_abs": 2.7373598861694335, + "relative_l2": 1.2229476895309177 + }, + "parameter.transformer_blocks.28.attn.to_v.weight": { + "max_abs": 4.708512072563171, + "relative_l2": 1.1407878598746384 + }, + "parameter.transformer_blocks.28.ff.net.0.proj.weight": { + "max_abs": 246.0049542236328, + "relative_l2": 0.18856620505021301 + }, + "parameter.transformer_blocks.28.ff.net.2.weight": { + "max_abs": 281.19640625, + "relative_l2": 0.13984586578762667 + }, + "parameter.transformer_blocks.28.norm1.weight": { + "max_abs": 11.20217836380005, + "relative_l2": 1.2122775780156614 + }, + "parameter.transformer_blocks.28.norm2.weight": { + "max_abs": 359.1081494140625, + "relative_l2": 0.16797781699942133 + }, + "parameter.transformer_blocks.29.adaln_proj.linear.bias": { + "max_abs": 6.513003234863281, + "relative_l2": 0.19296473938499076 + }, + "parameter.transformer_blocks.29.adaln_proj.linear.weight": { + "max_abs": 1499.823359375, + "relative_l2": 0.15103071975080934 + }, + "parameter.transformer_blocks.29.attn.norm_k.weight": { + "max_abs": 9.076084260940553, + "relative_l2": 1.6341561350014264 + }, + "parameter.transformer_blocks.29.attn.norm_q.weight": { + "max_abs": 7.577287321090698, + "relative_l2": 1.640274162969278 + }, + "parameter.transformer_blocks.29.attn.to_k.weight": { + "max_abs": 6.094952707290649, + "relative_l2": 1.4206461842138516 + }, + "parameter.transformer_blocks.29.attn.to_out.0.weight": { + "max_abs": 1.0908178436756133, + "relative_l2": 0.9675961856620527 + }, + "parameter.transformer_blocks.29.attn.to_q.weight": { + "max_abs": 4.269876127243042, + "relative_l2": 1.4317035390714827 + }, + "parameter.transformer_blocks.29.attn.to_v.weight": { + "max_abs": 3.974481587409973, + "relative_l2": 1.0206224707618594 + }, + "parameter.transformer_blocks.29.ff.net.0.proj.weight": { + "max_abs": 516.3020886230469, + "relative_l2": 0.18134852616131078 + }, + "parameter.transformer_blocks.29.ff.net.2.weight": { + "max_abs": 341.305537109375, + "relative_l2": 0.12471303626652662 + }, + "parameter.transformer_blocks.29.norm1.weight": { + "max_abs": 13.254407291412354, + "relative_l2": 1.201718957896938 + }, + "parameter.transformer_blocks.29.norm2.weight": { + "max_abs": 412.28252685546875, + "relative_l2": 0.17436060353310176 + }, + "parameter.transformer_blocks.3.adaln_proj.linear.bias": { + "max_abs": 26.373812561035155, + "relative_l2": 0.2654205061972307 + }, + "parameter.transformer_blocks.3.adaln_proj.linear.weight": { + "max_abs": 6267.35216796875, + "relative_l2": 0.25628779739482377 + }, + "parameter.transformer_blocks.3.attn.norm_k.weight": { + "max_abs": 10.009553079605103, + "relative_l2": 1.0049191909026862 + }, + "parameter.transformer_blocks.3.attn.norm_q.weight": { + "max_abs": 8.742085580825806, + "relative_l2": 0.9629340394722943 + }, + "parameter.transformer_blocks.3.attn.to_k.weight": { + "max_abs": 7.2191501903533934, + "relative_l2": 1.0015157745442331 + }, + "parameter.transformer_blocks.3.attn.to_out.0.weight": { + "max_abs": 3.831388597488403, + "relative_l2": 0.6430799687149178 + }, + "parameter.transformer_blocks.3.attn.to_q.weight": { + "max_abs": 9.674795513153077, + "relative_l2": 0.8449424927333038 + }, + "parameter.transformer_blocks.3.attn.to_v.weight": { + "max_abs": 1.722345654964447, + "relative_l2": 0.7180312690570306 + }, + "parameter.transformer_blocks.3.ff.net.0.proj.weight": { + "max_abs": 1106.1234387207032, + "relative_l2": 0.26521183470912624 + }, + "parameter.transformer_blocks.3.ff.net.2.weight": { + "max_abs": 289.41797912597656, + "relative_l2": 0.2521339003529588 + }, + "parameter.transformer_blocks.3.norm1.weight": { + "max_abs": 11.262598876953126, + "relative_l2": 0.9202157109443916 + }, + "parameter.transformer_blocks.3.norm2.weight": { + "max_abs": 4101.5118359375, + "relative_l2": 0.269506544885244 + }, + "parameter.transformer_blocks.30.adaln_proj.linear.bias": { + "max_abs": 4.232561950683594, + "relative_l2": 0.21645874428585316 + }, + "parameter.transformer_blocks.30.adaln_proj.linear.weight": { + "max_abs": 1016.190546875, + "relative_l2": 0.15072002097149242 + }, + "parameter.transformer_blocks.30.attn.norm_k.weight": { + "max_abs": 6.234205846786499, + "relative_l2": 1.024429056236177 + }, + "parameter.transformer_blocks.30.attn.norm_q.weight": { + "max_abs": 6.326533203125, + "relative_l2": 1.0369492200530552 + }, + "parameter.transformer_blocks.30.attn.to_k.weight": { + "max_abs": 4.29625440120697, + "relative_l2": 1.1695632920193195 + }, + "parameter.transformer_blocks.30.attn.to_out.0.weight": { + "max_abs": 0.9467710602283478, + "relative_l2": 0.7817796444989354 + }, + "parameter.transformer_blocks.30.attn.to_q.weight": { + "max_abs": 3.7871464538574218, + "relative_l2": 1.0976680529287743 + }, + "parameter.transformer_blocks.30.attn.to_v.weight": { + "max_abs": 2.525110249519348, + "relative_l2": 0.8613141581962568 + }, + "parameter.transformer_blocks.30.ff.net.0.proj.weight": { + "max_abs": 239.4531201171875, + "relative_l2": 0.19589711445921443 + }, + "parameter.transformer_blocks.30.ff.net.2.weight": { + "max_abs": 238.6540478515625, + "relative_l2": 0.13942788767281689 + }, + "parameter.transformer_blocks.30.norm1.weight": { + "max_abs": 8.512052898406983, + "relative_l2": 1.0641770769293457 + }, + "parameter.transformer_blocks.30.norm2.weight": { + "max_abs": 451.7082470703125, + "relative_l2": 0.18921410436824 + }, + "parameter.transformer_blocks.31.adaln_proj.linear.bias": { + "max_abs": 7.6711624908447265, + "relative_l2": 0.2123018243287658 + }, + "parameter.transformer_blocks.31.adaln_proj.linear.weight": { + "max_abs": 799.74376953125, + "relative_l2": 0.16188623355622284 + }, + "parameter.transformer_blocks.31.attn.norm_k.weight": { + "max_abs": 4.266907458305359, + "relative_l2": 0.7253633062716792 + }, + "parameter.transformer_blocks.31.attn.norm_q.weight": { + "max_abs": 4.317539339065552, + "relative_l2": 0.7171574495366576 + }, + "parameter.transformer_blocks.31.attn.to_k.weight": { + "max_abs": 1.8372212696075438, + "relative_l2": 0.8527959018711727 + }, + "parameter.transformer_blocks.31.attn.to_out.0.weight": { + "max_abs": 0.7584056901931763, + "relative_l2": 0.7532493180744085 + }, + "parameter.transformer_blocks.31.attn.to_q.weight": { + "max_abs": 1.5440465259552, + "relative_l2": 0.8164946551799815 + }, + "parameter.transformer_blocks.31.attn.to_v.weight": { + "max_abs": 1.8497110652923583, + "relative_l2": 0.8529554209495858 + }, + "parameter.transformer_blocks.31.ff.net.0.proj.weight": { + "max_abs": 154.2845066833496, + "relative_l2": 0.22037467788839546 + }, + "parameter.transformer_blocks.31.ff.net.2.weight": { + "max_abs": 112.23086822509765, + "relative_l2": 0.15448229099375835 + }, + "parameter.transformer_blocks.31.norm1.weight": { + "max_abs": 5.710070734024048, + "relative_l2": 0.8192065322889842 + }, + "parameter.transformer_blocks.31.norm2.weight": { + "max_abs": 349.50445068359375, + "relative_l2": 0.21345879795988626 + }, + "parameter.transformer_blocks.32.adaln_proj.linear.bias": { + "max_abs": 4.077955131530762, + "relative_l2": 0.20541480550039418 + }, + "parameter.transformer_blocks.32.adaln_proj.linear.weight": { + "max_abs": 980.7512158203125, + "relative_l2": 0.14042511207087938 + }, + "parameter.transformer_blocks.32.attn.norm_k.weight": { + "max_abs": 6.775987749099731, + "relative_l2": 1.288531562899753 + }, + "parameter.transformer_blocks.32.attn.norm_q.weight": { + "max_abs": 6.837880973815918, + "relative_l2": 1.2434342562895009 + }, + "parameter.transformer_blocks.32.attn.to_k.weight": { + "max_abs": 4.671030168533325, + "relative_l2": 1.1509488580181235 + }, + "parameter.transformer_blocks.32.attn.to_out.0.weight": { + "max_abs": 1.6398096132278441, + "relative_l2": 0.992625717903801 + }, + "parameter.transformer_blocks.32.attn.to_q.weight": { + "max_abs": 2.183147554397583, + "relative_l2": 1.1746264337260195 + }, + "parameter.transformer_blocks.32.attn.to_v.weight": { + "max_abs": 1.9204409170150756, + "relative_l2": 1.0286533770205062 + }, + "parameter.transformer_blocks.32.ff.net.0.proj.weight": { + "max_abs": 209.76043975830078, + "relative_l2": 0.20242125269897343 + }, + "parameter.transformer_blocks.32.ff.net.2.weight": { + "max_abs": 89.70086181640625, + "relative_l2": 0.1444801344032381 + }, + "parameter.transformer_blocks.32.norm1.weight": { + "max_abs": 9.582114820480347, + "relative_l2": 1.1194710214665151 + }, + "parameter.transformer_blocks.32.norm2.weight": { + "max_abs": 336.048212890625, + "relative_l2": 0.20319355913432793 + }, + "parameter.transformer_blocks.33.adaln_proj.linear.bias": { + "max_abs": 7.154135589599609, + "relative_l2": 0.19469886900881594 + }, + "parameter.transformer_blocks.33.adaln_proj.linear.weight": { + "max_abs": 966.8830822753906, + "relative_l2": 0.14829812099440315 + }, + "parameter.transformer_blocks.33.attn.norm_k.weight": { + "max_abs": 5.550072317123413, + "relative_l2": 1.1274773312509463 + }, + "parameter.transformer_blocks.33.attn.norm_q.weight": { + "max_abs": 5.736447219848633, + "relative_l2": 1.1253863704143687 + }, + "parameter.transformer_blocks.33.attn.to_k.weight": { + "max_abs": 2.679831688404083, + "relative_l2": 0.9850919152514763 + }, + "parameter.transformer_blocks.33.attn.to_out.0.weight": { + "max_abs": 2.063485448360443, + "relative_l2": 0.9390572904823536 + }, + "parameter.transformer_blocks.33.attn.to_q.weight": { + "max_abs": 1.764441316127777, + "relative_l2": 0.9817237812251518 + }, + "parameter.transformer_blocks.33.attn.to_v.weight": { + "max_abs": 2.399023537635803, + "relative_l2": 0.9501950616159477 + }, + "parameter.transformer_blocks.33.ff.net.0.proj.weight": { + "max_abs": 309.37867248535156, + "relative_l2": 0.2006807330341676 + }, + "parameter.transformer_blocks.33.ff.net.2.weight": { + "max_abs": 80.20376480102539, + "relative_l2": 0.13101022942819657 + }, + "parameter.transformer_blocks.33.norm1.weight": { + "max_abs": 15.021945838928223, + "relative_l2": 0.9127326214831956 + }, + "parameter.transformer_blocks.33.norm2.weight": { + "max_abs": 1001.6911572265625, + "relative_l2": 0.2251123813234075 + }, + "parameter.transformer_blocks.34.adaln_proj.linear.bias": { + "max_abs": 8.243670349121095, + "relative_l2": 0.1806048426281073 + }, + "parameter.transformer_blocks.34.adaln_proj.linear.weight": { + "max_abs": 1148.6772412109376, + "relative_l2": 0.12984915201068334 + }, + "parameter.transformer_blocks.34.attn.norm_k.weight": { + "max_abs": 7.219762449264526, + "relative_l2": 1.032384395529538 + }, + "parameter.transformer_blocks.34.attn.norm_q.weight": { + "max_abs": 7.133918170928955, + "relative_l2": 1.0149967125277624 + }, + "parameter.transformer_blocks.34.attn.to_k.weight": { + "max_abs": 4.497727279663086, + "relative_l2": 0.9591314698602402 + }, + "parameter.transformer_blocks.34.attn.to_out.0.weight": { + "max_abs": 1.03602308511734, + "relative_l2": 0.8783822880153688 + }, + "parameter.transformer_blocks.34.attn.to_q.weight": { + "max_abs": 2.7484306144714354, + "relative_l2": 0.8864441081344899 + }, + "parameter.transformer_blocks.34.attn.to_v.weight": { + "max_abs": 3.29366970539093, + "relative_l2": 0.9451329266505385 + }, + "parameter.transformer_blocks.34.ff.net.0.proj.weight": { + "max_abs": 194.5088070678711, + "relative_l2": 0.18493408026558106 + }, + "parameter.transformer_blocks.34.ff.net.2.weight": { + "max_abs": 245.620478515625, + "relative_l2": 0.1307436775198352 + }, + "parameter.transformer_blocks.34.norm1.weight": { + "max_abs": 24.60771167755127, + "relative_l2": 1.0233308374050827 + }, + "parameter.transformer_blocks.34.norm2.weight": { + "max_abs": 319.7883786010742, + "relative_l2": 0.15874418206460217 + }, + "parameter.transformer_blocks.35.adaln_proj.linear.bias": { + "max_abs": 5.534396057128906, + "relative_l2": 0.1467721804613339 + }, + "parameter.transformer_blocks.35.adaln_proj.linear.weight": { + "max_abs": 1264.75109375, + "relative_l2": 0.11913879267016361 + }, + "parameter.transformer_blocks.35.attn.norm_k.weight": { + "max_abs": 4.387928371429443, + "relative_l2": 0.942640852962107 + }, + "parameter.transformer_blocks.35.attn.norm_q.weight": { + "max_abs": 4.449705009460449, + "relative_l2": 0.919191599683251 + }, + "parameter.transformer_blocks.35.attn.to_k.weight": { + "max_abs": 1.7174570608139037, + "relative_l2": 0.9179817867494432 + }, + "parameter.transformer_blocks.35.attn.to_out.0.weight": { + "max_abs": 1.459444110393524, + "relative_l2": 0.8814012609524116 + }, + "parameter.transformer_blocks.35.attn.to_q.weight": { + "max_abs": 2.8721355247497558, + "relative_l2": 1.0034159715694124 + }, + "parameter.transformer_blocks.35.attn.to_v.weight": { + "max_abs": 1.7862120914459227, + "relative_l2": 0.9437613436042724 + }, + "parameter.transformer_blocks.35.ff.net.0.proj.weight": { + "max_abs": 191.31423461914062, + "relative_l2": 0.17145562347579194 + }, + "parameter.transformer_blocks.35.ff.net.2.weight": { + "max_abs": 186.5078076171875, + "relative_l2": 0.11657884529574467 + }, + "parameter.transformer_blocks.35.norm1.weight": { + "max_abs": 6.005865459442139, + "relative_l2": 0.8703644599721103 + }, + "parameter.transformer_blocks.35.norm2.weight": { + "max_abs": 679.1898754882812, + "relative_l2": 0.200027453204279 + }, + "parameter.transformer_blocks.36.adaln_proj.linear.bias": { + "max_abs": 9.245374565124513, + "relative_l2": 0.2087435271818816 + }, + "parameter.transformer_blocks.36.adaln_proj.linear.weight": { + "max_abs": 2195.7213696289064, + "relative_l2": 0.1602415887020769 + }, + "parameter.transformer_blocks.36.attn.norm_k.weight": { + "max_abs": 4.618183856010437, + "relative_l2": 1.0325610439726254 + }, + "parameter.transformer_blocks.36.attn.norm_q.weight": { + "max_abs": 4.725456480979919, + "relative_l2": 1.0538600913481053 + }, + "parameter.transformer_blocks.36.attn.to_k.weight": { + "max_abs": 5.769529585838318, + "relative_l2": 1.221919776036673 + }, + "parameter.transformer_blocks.36.attn.to_out.0.weight": { + "max_abs": 1.0138367462158202, + "relative_l2": 0.5783897208725631 + }, + "parameter.transformer_blocks.36.attn.to_q.weight": { + "max_abs": 3.6643313217163085, + "relative_l2": 1.2180755207186575 + }, + "parameter.transformer_blocks.36.attn.to_v.weight": { + "max_abs": 1.8370685625076293, + "relative_l2": 0.6205732856999983 + }, + "parameter.transformer_blocks.36.ff.net.0.proj.weight": { + "max_abs": 309.66125, + "relative_l2": 0.20316813457524138 + }, + "parameter.transformer_blocks.36.ff.net.2.weight": { + "max_abs": 107.78756225585937, + "relative_l2": 0.12236753561577896 + }, + "parameter.transformer_blocks.36.norm1.weight": { + "max_abs": 22.513299827575683, + "relative_l2": 1.0778564356317475 + }, + "parameter.transformer_blocks.36.norm2.weight": { + "max_abs": 2511.9208935546876, + "relative_l2": 0.5710897440023305 + }, + "parameter.transformer_blocks.37.adaln_proj.linear.bias": { + "max_abs": 6.676037197113037, + "relative_l2": 0.1808701004652349 + }, + "parameter.transformer_blocks.37.adaln_proj.linear.weight": { + "max_abs": 1593.4457348632814, + "relative_l2": 0.1443106764289475 + }, + "parameter.transformer_blocks.37.attn.norm_k.weight": { + "max_abs": 2.832467083930969, + "relative_l2": 1.0101240952318908 + }, + "parameter.transformer_blocks.37.attn.norm_q.weight": { + "max_abs": 2.861815993785858, + "relative_l2": 1.0127782878012312 + }, + "parameter.transformer_blocks.37.attn.to_k.weight": { + "max_abs": 4.04865861415863, + "relative_l2": 1.1374777364018036 + }, + "parameter.transformer_blocks.37.attn.to_out.0.weight": { + "max_abs": 0.854398285150528, + "relative_l2": 0.8419456640459979 + }, + "parameter.transformer_blocks.37.attn.to_q.weight": { + "max_abs": 2.1792039918899535, + "relative_l2": 1.1004063862861377 + }, + "parameter.transformer_blocks.37.attn.to_v.weight": { + "max_abs": 2.1426338601112365, + "relative_l2": 0.8787250645274481 + }, + "parameter.transformer_blocks.37.ff.net.0.proj.weight": { + "max_abs": 656.2790783691406, + "relative_l2": 0.2233693593556396 + }, + "parameter.transformer_blocks.37.ff.net.2.weight": { + "max_abs": 50.611208801269534, + "relative_l2": 0.10374361375682521 + }, + "parameter.transformer_blocks.37.norm1.weight": { + "max_abs": 8.482384805679322, + "relative_l2": 1.0256669036930872 + }, + "parameter.transformer_blocks.37.norm2.weight": { + "max_abs": 547.3217114257812, + "relative_l2": 0.268658747915288 + }, + "parameter.transformer_blocks.38.adaln_proj.linear.bias": { + "max_abs": 5.421483640670776, + "relative_l2": 0.161460090239007 + }, + "parameter.transformer_blocks.38.adaln_proj.linear.weight": { + "max_abs": 771.6396435546875, + "relative_l2": 0.10861415714614744 + }, + "parameter.transformer_blocks.38.attn.norm_k.weight": { + "max_abs": 5.844645385742187, + "relative_l2": 0.9242552679150525 + }, + "parameter.transformer_blocks.38.attn.norm_q.weight": { + "max_abs": 5.868889694213867, + "relative_l2": 0.9179216618472709 + }, + "parameter.transformer_blocks.38.attn.to_k.weight": { + "max_abs": 2.9294017601013183, + "relative_l2": 0.9094237527007583 + }, + "parameter.transformer_blocks.38.attn.to_out.0.weight": { + "max_abs": 0.668610428571701, + "relative_l2": 0.9234331916237591 + }, + "parameter.transformer_blocks.38.attn.to_q.weight": { + "max_abs": 1.974614565372467, + "relative_l2": 0.9598699375006258 + }, + "parameter.transformer_blocks.38.attn.to_v.weight": { + "max_abs": 1.4428296136856078, + "relative_l2": 0.9253544945282148 + }, + "parameter.transformer_blocks.38.ff.net.0.proj.weight": { + "max_abs": 145.5050343322754, + "relative_l2": 0.17510562755866485 + }, + "parameter.transformer_blocks.38.ff.net.2.weight": { + "max_abs": 65.21968353271484, + "relative_l2": 0.09506300458020124 + }, + "parameter.transformer_blocks.38.norm1.weight": { + "max_abs": 13.351056938171388, + "relative_l2": 0.8662618342953593 + }, + "parameter.transformer_blocks.38.norm2.weight": { + "max_abs": 703.1446484375, + "relative_l2": 0.16890057341383458 + }, + "parameter.transformer_blocks.39.adaln_proj.linear.bias": { + "max_abs": 6.127394561767578, + "relative_l2": 0.13168527159369647 + }, + "parameter.transformer_blocks.39.adaln_proj.linear.weight": { + "max_abs": 1456.505, + "relative_l2": 0.08965412934922393 + }, + "parameter.transformer_blocks.39.attn.norm_k.weight": { + "max_abs": 4.7785280990600585, + "relative_l2": 1.0694236943687239 + }, + "parameter.transformer_blocks.39.attn.norm_q.weight": { + "max_abs": 4.7358480739593505, + "relative_l2": 1.0685993433522698 + }, + "parameter.transformer_blocks.39.attn.to_k.weight": { + "max_abs": 2.62364173412323, + "relative_l2": 1.124200114084775 + }, + "parameter.transformer_blocks.39.attn.to_out.0.weight": { + "max_abs": 0.757876490354538, + "relative_l2": 0.8481142640103078 + }, + "parameter.transformer_blocks.39.attn.to_q.weight": { + "max_abs": 1.915063087940216, + "relative_l2": 1.1942952625059793 + }, + "parameter.transformer_blocks.39.attn.to_v.weight": { + "max_abs": 1.7589626359939574, + "relative_l2": 0.8567941748841181 + }, + "parameter.transformer_blocks.39.ff.net.0.proj.weight": { + "max_abs": 117.97527587890625, + "relative_l2": 0.14981088979566995 + }, + "parameter.transformer_blocks.39.ff.net.2.weight": { + "max_abs": 182.483515625, + "relative_l2": 0.0897690232787083 + }, + "parameter.transformer_blocks.39.norm1.weight": { + "max_abs": 8.069322471618653, + "relative_l2": 0.9237021151344867 + }, + "parameter.transformer_blocks.39.norm2.weight": { + "max_abs": 1081.0607861328126, + "relative_l2": 0.25401553568912055 + }, + "parameter.transformer_blocks.4.adaln_proj.linear.bias": { + "max_abs": 27.046805267333983, + "relative_l2": 0.26238639729994045 + }, + "parameter.transformer_blocks.4.adaln_proj.linear.weight": { + "max_abs": 6488.14513671875, + "relative_l2": 0.25077882117530165 + }, + "parameter.transformer_blocks.4.attn.norm_k.weight": { + "max_abs": 4.5355954217910766, + "relative_l2": 0.9585719035587216 + }, + "parameter.transformer_blocks.4.attn.norm_q.weight": { + "max_abs": 3.9605391073226928, + "relative_l2": 0.9581588861708147 + }, + "parameter.transformer_blocks.4.attn.to_k.weight": { + "max_abs": 3.6364445734024047, + "relative_l2": 0.8205567058445037 + }, + "parameter.transformer_blocks.4.attn.to_out.0.weight": { + "max_abs": 2.0168289589881896, + "relative_l2": 0.789058322187198 + }, + "parameter.transformer_blocks.4.attn.to_q.weight": { + "max_abs": 5.848101501464844, + "relative_l2": 0.8094523151615062 + }, + "parameter.transformer_blocks.4.attn.to_v.weight": { + "max_abs": 1.6423358964920043, + "relative_l2": 0.8591711598148015 + }, + "parameter.transformer_blocks.4.ff.net.0.proj.weight": { + "max_abs": 1159.7981518554689, + "relative_l2": 0.2674986928299999 + }, + "parameter.transformer_blocks.4.ff.net.2.weight": { + "max_abs": 402.25548828125, + "relative_l2": 0.24103781594722395 + }, + "parameter.transformer_blocks.4.norm1.weight": { + "max_abs": 17.1214174079895, + "relative_l2": 0.8361116592439933 + }, + "parameter.transformer_blocks.4.norm2.weight": { + "max_abs": 801.8566235351562, + "relative_l2": 0.2561850092970625 + }, + "parameter.transformer_blocks.40.adaln_proj.linear.bias": { + "max_abs": 3.2005223083496093, + "relative_l2": 0.12521540524849173 + }, + "parameter.transformer_blocks.40.adaln_proj.linear.weight": { + "max_abs": 705.36681640625, + "relative_l2": 0.08374239247167095 + }, + "parameter.transformer_blocks.40.attn.norm_k.weight": { + "max_abs": 3.4289771366119384, + "relative_l2": 1.0478765270951127 + }, + "parameter.transformer_blocks.40.attn.norm_q.weight": { + "max_abs": 4.846121912002563, + "relative_l2": 1.0795186674130992 + }, + "parameter.transformer_blocks.40.attn.to_k.weight": { + "max_abs": 3.5802519369125365, + "relative_l2": 1.140589196399972 + }, + "parameter.transformer_blocks.40.attn.to_out.0.weight": { + "max_abs": 0.6302498775720596, + "relative_l2": 0.7935126763845027 + }, + "parameter.transformer_blocks.40.attn.to_q.weight": { + "max_abs": 5.073410515785217, + "relative_l2": 1.3077308914794323 + }, + "parameter.transformer_blocks.40.attn.to_v.weight": { + "max_abs": 2.0222979831695556, + "relative_l2": 0.8269094116055061 + }, + "parameter.transformer_blocks.40.ff.net.0.proj.weight": { + "max_abs": 210.19929779052734, + "relative_l2": 0.15512173606582044 + }, + "parameter.transformer_blocks.40.ff.net.2.weight": { + "max_abs": 229.26128662109375, + "relative_l2": 0.0881041806052548 + }, + "parameter.transformer_blocks.40.norm1.weight": { + "max_abs": 4.467885737419128, + "relative_l2": 0.9969333401675976 + }, + "parameter.transformer_blocks.40.norm2.weight": { + "max_abs": 194.94732177734375, + "relative_l2": 0.1290674790515727 + }, + "parameter.transformer_blocks.41.adaln_proj.linear.bias": { + "max_abs": 4.080959205627441, + "relative_l2": 0.09656603372527352 + }, + "parameter.transformer_blocks.41.adaln_proj.linear.weight": { + "max_abs": 953.324580078125, + "relative_l2": 0.08195072423360418 + }, + "parameter.transformer_blocks.41.attn.norm_k.weight": { + "max_abs": 2.757928853034973, + "relative_l2": 0.7563508971808889 + }, + "parameter.transformer_blocks.41.attn.norm_q.weight": { + "max_abs": 2.8040596294403075, + "relative_l2": 0.7409621277067486 + }, + "parameter.transformer_blocks.41.attn.to_k.weight": { + "max_abs": 1.3457395482063292, + "relative_l2": 0.8211947853401178 + }, + "parameter.transformer_blocks.41.attn.to_out.0.weight": { + "max_abs": 0.731994201540947, + "relative_l2": 0.8026639137656787 + }, + "parameter.transformer_blocks.41.attn.to_q.weight": { + "max_abs": 0.8410365927219391, + "relative_l2": 0.8438698268641297 + }, + "parameter.transformer_blocks.41.attn.to_v.weight": { + "max_abs": 1.642683868408203, + "relative_l2": 0.8220985051578446 + }, + "parameter.transformer_blocks.41.ff.net.0.proj.weight": { + "max_abs": 154.27525985717773, + "relative_l2": 0.12454783262895078 + }, + "parameter.transformer_blocks.41.ff.net.2.weight": { + "max_abs": 228.506708984375, + "relative_l2": 0.0842293967140266 + }, + "parameter.transformer_blocks.41.norm1.weight": { + "max_abs": 5.2925452041625975, + "relative_l2": 0.7890362572700109 + }, + "parameter.transformer_blocks.41.norm2.weight": { + "max_abs": 668.660029296875, + "relative_l2": 0.11268412126563454 + }, + "parameter.transformer_blocks.42.adaln_proj.linear.bias": { + "max_abs": 4.917359237670898, + "relative_l2": 0.09542605967411237 + }, + "parameter.transformer_blocks.42.adaln_proj.linear.weight": { + "max_abs": 1169.406611328125, + "relative_l2": 0.07712454346483907 + }, + "parameter.transformer_blocks.42.attn.norm_k.weight": { + "max_abs": 4.035946731567383, + "relative_l2": 0.8238593889325051 + }, + "parameter.transformer_blocks.42.attn.norm_q.weight": { + "max_abs": 4.056439523696899, + "relative_l2": 0.7996155756600228 + }, + "parameter.transformer_blocks.42.attn.to_k.weight": { + "max_abs": 1.6770926761627196, + "relative_l2": 0.7641442861987374 + }, + "parameter.transformer_blocks.42.attn.to_out.0.weight": { + "max_abs": 0.5470972406864166, + "relative_l2": 0.5442167760593228 + }, + "parameter.transformer_blocks.42.attn.to_q.weight": { + "max_abs": 1.1088188636302947, + "relative_l2": 0.8081196804260867 + }, + "parameter.transformer_blocks.42.attn.to_v.weight": { + "max_abs": 1.2766135680675506, + "relative_l2": 0.6157095221288996 + }, + "parameter.transformer_blocks.42.ff.net.0.proj.weight": { + "max_abs": 155.3827389526367, + "relative_l2": 0.11766137098501636 + }, + "parameter.transformer_blocks.42.ff.net.2.weight": { + "max_abs": 129.46139038085937, + "relative_l2": 0.07740887728938908 + }, + "parameter.transformer_blocks.42.norm1.weight": { + "max_abs": 4.736500029563904, + "relative_l2": 0.7671526716518661 + }, + "parameter.transformer_blocks.42.norm2.weight": { + "max_abs": 270.3913067626953, + "relative_l2": 0.10944300490817432 + }, + "parameter.transformer_blocks.43.adaln_proj.linear.bias": { + "max_abs": 5.132880096435547, + "relative_l2": 0.11273468825966387 + }, + "parameter.transformer_blocks.43.adaln_proj.linear.weight": { + "max_abs": 1240.3321484375, + "relative_l2": 0.08187268981591711 + }, + "parameter.transformer_blocks.43.attn.norm_k.weight": { + "max_abs": 6.753908519744873, + "relative_l2": 0.9205244132931526 + }, + "parameter.transformer_blocks.43.attn.norm_q.weight": { + "max_abs": 6.826011304855347, + "relative_l2": 0.9200605195684711 + }, + "parameter.transformer_blocks.43.attn.to_k.weight": { + "max_abs": 2.645546441078186, + "relative_l2": 0.9111346814717839 + }, + "parameter.transformer_blocks.43.attn.to_out.0.weight": { + "max_abs": 0.7447509664297104, + "relative_l2": 0.571198163788551 + }, + "parameter.transformer_blocks.43.attn.to_q.weight": { + "max_abs": 1.889353756904602, + "relative_l2": 0.9313925335595282 + }, + "parameter.transformer_blocks.43.attn.to_v.weight": { + "max_abs": 1.4668158936500548, + "relative_l2": 0.599385356863508 + }, + "parameter.transformer_blocks.43.ff.net.0.proj.weight": { + "max_abs": 300.6612957763672, + "relative_l2": 0.1136452440596051 + }, + "parameter.transformer_blocks.43.ff.net.2.weight": { + "max_abs": 66.31138122558593, + "relative_l2": 0.08056046834619426 + }, + "parameter.transformer_blocks.43.norm1.weight": { + "max_abs": 13.360413913726807, + "relative_l2": 0.7114983351779064 + }, + "parameter.transformer_blocks.43.norm2.weight": { + "max_abs": 156.96438110351562, + "relative_l2": 0.09389628263383776 + }, + "parameter.transformer_blocks.44.adaln_proj.linear.bias": { + "max_abs": 17.729781036376952, + "relative_l2": 0.11705536326286792 + }, + "parameter.transformer_blocks.44.adaln_proj.linear.weight": { + "max_abs": 4208.96837890625, + "relative_l2": 0.1049058565796688 + }, + "parameter.transformer_blocks.44.attn.norm_k.weight": { + "max_abs": 2.269458894729614, + "relative_l2": 0.7091757295194187 + }, + "parameter.transformer_blocks.44.attn.norm_q.weight": { + "max_abs": 2.164663915634155, + "relative_l2": 0.7054317049993656 + }, + "parameter.transformer_blocks.44.attn.to_k.weight": { + "max_abs": 1.7072026181221007, + "relative_l2": 0.8324279544785277 + }, + "parameter.transformer_blocks.44.attn.to_out.0.weight": { + "max_abs": 0.8187230277061462, + "relative_l2": 0.5460598839539066 + }, + "parameter.transformer_blocks.44.attn.to_q.weight": { + "max_abs": 1.233398710489273, + "relative_l2": 0.8872870872880635 + }, + "parameter.transformer_blocks.44.attn.to_v.weight": { + "max_abs": 1.7586665201187133, + "relative_l2": 0.6282548434192357 + }, + "parameter.transformer_blocks.44.ff.net.0.proj.weight": { + "max_abs": 149.37920272827148, + "relative_l2": 0.10937401822764103 + }, + "parameter.transformer_blocks.44.ff.net.2.weight": { + "max_abs": 92.94004333496093, + "relative_l2": 0.07597955498059568 + }, + "parameter.transformer_blocks.44.norm1.weight": { + "max_abs": 8.42225254058838, + "relative_l2": 0.6669851115364491 + }, + "parameter.transformer_blocks.44.norm2.weight": { + "max_abs": 848.308466796875, + "relative_l2": 0.18199623008311203 + }, + "parameter.transformer_blocks.45.adaln_proj.linear.bias": { + "max_abs": 6.1834706115722655, + "relative_l2": 0.07386063824844913 + }, + "parameter.transformer_blocks.45.adaln_proj.linear.weight": { + "max_abs": 1467.5729931640626, + "relative_l2": 0.06511389998485793 + }, + "parameter.transformer_blocks.45.attn.norm_k.weight": { + "max_abs": 2.2533670711517333, + "relative_l2": 0.771607820457134 + }, + "parameter.transformer_blocks.45.attn.norm_q.weight": { + "max_abs": 2.2690254497528075, + "relative_l2": 0.7876834677150717 + }, + "parameter.transformer_blocks.45.attn.to_k.weight": { + "max_abs": 2.8169800090789794, + "relative_l2": 0.768766108578039 + }, + "parameter.transformer_blocks.45.attn.to_out.0.weight": { + "max_abs": 0.510720004439354, + "relative_l2": 0.34819230093551706 + }, + "parameter.transformer_blocks.45.attn.to_q.weight": { + "max_abs": 25.44289291381836, + "relative_l2": 0.4792674547547203 + }, + "parameter.transformer_blocks.45.attn.to_v.weight": { + "max_abs": 1.0604091036319732, + "relative_l2": 0.3754161340434127 + }, + "parameter.transformer_blocks.45.ff.net.0.proj.weight": { + "max_abs": 201.54845703125, + "relative_l2": 0.09227611490400561 + }, + "parameter.transformer_blocks.45.ff.net.2.weight": { + "max_abs": 119.0656689453125, + "relative_l2": 0.07023417264246236 + }, + "parameter.transformer_blocks.45.norm1.weight": { + "max_abs": 9.380532150268555, + "relative_l2": 0.569506798252519 + }, + "parameter.transformer_blocks.45.norm2.weight": { + "max_abs": 368.22149169921875, + "relative_l2": 0.1011082779913207 + }, + "parameter.transformer_blocks.46.adaln_proj.linear.bias": { + "max_abs": 8.651469116210938, + "relative_l2": 0.06329771186354746 + }, + "parameter.transformer_blocks.46.adaln_proj.linear.weight": { + "max_abs": 1905.06359375, + "relative_l2": 0.05911674769537031 + }, + "parameter.transformer_blocks.46.attn.norm_k.weight": { + "max_abs": 2.558462266921997, + "relative_l2": 0.585320050556141 + }, + "parameter.transformer_blocks.46.attn.norm_q.weight": { + "max_abs": 2.6533019590377807, + "relative_l2": 0.5598057701695375 + }, + "parameter.transformer_blocks.46.attn.to_k.weight": { + "max_abs": 2.455328290462494, + "relative_l2": 0.6239648127845449 + }, + "parameter.transformer_blocks.46.attn.to_out.0.weight": { + "max_abs": 0.549632688164711, + "relative_l2": 0.24239122366787058 + }, + "parameter.transformer_blocks.46.attn.to_q.weight": { + "max_abs": 2.604859833717346, + "relative_l2": 0.6585910041072844 + }, + "parameter.transformer_blocks.46.attn.to_v.weight": { + "max_abs": 1.3638730692863463, + "relative_l2": 0.24932692934413414 + }, + "parameter.transformer_blocks.46.ff.net.0.proj.weight": { + "max_abs": 410.6479443359375, + "relative_l2": 0.07948161125168707 + }, + "parameter.transformer_blocks.46.ff.net.2.weight": { + "max_abs": 971.273310546875, + "relative_l2": 0.06522679281801508 + }, + "parameter.transformer_blocks.46.norm1.weight": { + "max_abs": 10.147108917236329, + "relative_l2": 0.2227172172388825 + }, + "parameter.transformer_blocks.46.norm2.weight": { + "max_abs": 2693.47033203125, + "relative_l2": 0.16338661750553896 + }, + "parameter.transformer_blocks.47.adaln_proj.linear.bias": { + "max_abs": 16.257235412597655, + "relative_l2": 0.07627798102860456 + }, + "parameter.transformer_blocks.47.adaln_proj.linear.weight": { + "max_abs": 3683.17296875, + "relative_l2": 0.06780493854220315 + }, + "parameter.transformer_blocks.47.attn.norm_k.weight": { + "max_abs": 5.5190323638916015, + "relative_l2": 0.4231322946674436 + }, + "parameter.transformer_blocks.47.attn.norm_q.weight": { + "max_abs": 5.7188004302978515, + "relative_l2": 0.4193198161245315 + }, + "parameter.transformer_blocks.47.attn.to_k.weight": { + "max_abs": 11.963134651184083, + "relative_l2": 0.330995632161804 + }, + "parameter.transformer_blocks.47.attn.to_out.0.weight": { + "max_abs": 1.202230279445648, + "relative_l2": 0.18708588305373947 + }, + "parameter.transformer_blocks.47.attn.to_q.weight": { + "max_abs": 4.0040234375, + "relative_l2": 0.35734895930583077 + }, + "parameter.transformer_blocks.47.attn.to_v.weight": { + "max_abs": 11.460832719802857, + "relative_l2": 0.20559483258091166 + }, + "parameter.transformer_blocks.47.ff.net.0.proj.weight": { + "max_abs": 1193.923701171875, + "relative_l2": 0.07835797547764027 + }, + "parameter.transformer_blocks.47.ff.net.2.weight": { + "max_abs": 1576.885859375, + "relative_l2": 0.051879736665692025 + }, + "parameter.transformer_blocks.47.norm1.weight": { + "max_abs": 15.85252655029297, + "relative_l2": 0.34502870028341526 + }, + "parameter.transformer_blocks.47.norm2.weight": { + "max_abs": 1780.4141796875, + "relative_l2": 0.08152340720341918 + }, + "parameter.transformer_blocks.48.adaln_proj.linear.bias": { + "max_abs": 26.97588623046875, + "relative_l2": 0.04928975797603363 + }, + "parameter.transformer_blocks.48.adaln_proj.linear.weight": { + "max_abs": 7351.505, + "relative_l2": 0.049891999286638504 + }, + "parameter.transformer_blocks.48.attn.norm_k.weight": { + "max_abs": 2.1515078592300414, + "relative_l2": 0.19186781900909602 + }, + "parameter.transformer_blocks.48.attn.norm_q.weight": { + "max_abs": 2.2295125770568847, + "relative_l2": 0.1988307599059829 + }, + "parameter.transformer_blocks.48.attn.to_k.weight": { + "max_abs": 4.589643363952637, + "relative_l2": 0.2042588278502384 + }, + "parameter.transformer_blocks.48.attn.to_out.0.weight": { + "max_abs": 1.014324550628662, + "relative_l2": 0.096045697959378 + }, + "parameter.transformer_blocks.48.attn.to_q.weight": { + "max_abs": 7.781969909667969, + "relative_l2": 0.2424748187686756 + }, + "parameter.transformer_blocks.48.attn.to_v.weight": { + "max_abs": 2.8530453491210936, + "relative_l2": 0.0977908986332009 + }, + "parameter.transformer_blocks.48.ff.net.0.proj.weight": { + "max_abs": 1301.2202099609376, + "relative_l2": 0.06038765587406777 + }, + "parameter.transformer_blocks.48.ff.net.2.weight": { + "max_abs": 1068.0698193359376, + "relative_l2": 0.045891303085466316 + }, + "parameter.transformer_blocks.48.norm1.weight": { + "max_abs": 6.685605888366699, + "relative_l2": 0.12466980346726744 + }, + "parameter.transformer_blocks.48.norm2.weight": { + "max_abs": 2388.690546875, + "relative_l2": 0.05236776060520865 + }, + "parameter.transformer_blocks.49.adaln_proj.linear.bias": { + "max_abs": 8.821124916076661, + "relative_l2": 0.10311497979458151 + }, + "parameter.transformer_blocks.49.adaln_proj.linear.weight": { + "max_abs": 1441.112421875, + "relative_l2": 0.043180249258132096 + }, + "parameter.transformer_blocks.49.attn.norm_k.weight": { + "max_abs": 4.704728012084961, + "relative_l2": 0.22726517993617218 + }, + "parameter.transformer_blocks.49.attn.norm_q.weight": { + "max_abs": 3.2715772438049315, + "relative_l2": 0.22292634803167102 + }, + "parameter.transformer_blocks.49.attn.to_k.weight": { + "max_abs": 180.87781799316406, + "relative_l2": 0.3738389842182632 + }, + "parameter.transformer_blocks.49.attn.to_out.0.weight": { + "max_abs": 2.1655740785598754, + "relative_l2": 0.07821954800563145 + }, + "parameter.transformer_blocks.49.attn.to_q.weight": { + "max_abs": 7475.51427734375, + "relative_l2": 0.9976771172686144 + }, + "parameter.transformer_blocks.49.attn.to_v.weight": { + "max_abs": 3.403834466934204, + "relative_l2": 0.07990269277842786 + }, + "parameter.transformer_blocks.49.ff.net.0.proj.weight": { + "max_abs": 1048.5288037109376, + "relative_l2": 0.04472210669022003 + }, + "parameter.transformer_blocks.49.ff.net.2.weight": { + "max_abs": 1158.47814453125, + "relative_l2": 0.028978917566814043 + }, + "parameter.transformer_blocks.49.norm1.weight": { + "max_abs": 79.21369064331054, + "relative_l2": 0.8639755692388044 + }, + "parameter.transformer_blocks.49.norm2.weight": { + "max_abs": 558.84435546875, + "relative_l2": 0.04921952358353799 + }, + "parameter.transformer_blocks.5.adaln_proj.linear.bias": { + "max_abs": 34.00472152709961, + "relative_l2": 0.26746757834431667 + }, + "parameter.transformer_blocks.5.adaln_proj.linear.weight": { + "max_abs": 8061.69640625, + "relative_l2": 0.2452543997549339 + }, + "parameter.transformer_blocks.5.attn.norm_k.weight": { + "max_abs": 3.9417589950561522, + "relative_l2": 1.2664483396718582 + }, + "parameter.transformer_blocks.5.attn.norm_q.weight": { + "max_abs": 3.946983938217163, + "relative_l2": 1.2224533507336455 + }, + "parameter.transformer_blocks.5.attn.to_k.weight": { + "max_abs": 2.9919197607040404, + "relative_l2": 1.1727527541387843 + }, + "parameter.transformer_blocks.5.attn.to_out.0.weight": { + "max_abs": 2.8459457206726073, + "relative_l2": 0.510954297011832 + }, + "parameter.transformer_blocks.5.attn.to_q.weight": { + "max_abs": 5.031534795761108, + "relative_l2": 1.2450305679568294 + }, + "parameter.transformer_blocks.5.attn.to_v.weight": { + "max_abs": 1.0378571796417235, + "relative_l2": 0.539951216146229 + }, + "parameter.transformer_blocks.5.ff.net.0.proj.weight": { + "max_abs": 493.2738751220703, + "relative_l2": 0.25716856811430633 + }, + "parameter.transformer_blocks.5.ff.net.2.weight": { + "max_abs": 234.38091552734375, + "relative_l2": 0.24140053738667538 + }, + "parameter.transformer_blocks.5.norm1.weight": { + "max_abs": 10.863516931533814, + "relative_l2": 0.946200087408264 + }, + "parameter.transformer_blocks.5.norm2.weight": { + "max_abs": 10746.702265625, + "relative_l2": 0.5631723714081652 + }, + "parameter.transformer_blocks.6.adaln_proj.linear.bias": { + "max_abs": 8.736480836868287, + "relative_l2": 0.28376338375463944 + }, + "parameter.transformer_blocks.6.adaln_proj.linear.weight": { + "max_abs": 2082.391901855469, + "relative_l2": 0.2422681374476185 + }, + "parameter.transformer_blocks.6.attn.norm_k.weight": { + "max_abs": 3.949999814033508, + "relative_l2": 0.6616773576672742 + }, + "parameter.transformer_blocks.6.attn.norm_q.weight": { + "max_abs": 4.279451971054077, + "relative_l2": 0.6088499948792981 + }, + "parameter.transformer_blocks.6.attn.to_k.weight": { + "max_abs": 3.365487461090088, + "relative_l2": 0.671770525013 + }, + "parameter.transformer_blocks.6.attn.to_out.0.weight": { + "max_abs": 1.4420253086090087, + "relative_l2": 0.7280617994287809 + }, + "parameter.transformer_blocks.6.attn.to_q.weight": { + "max_abs": 7.461075429916382, + "relative_l2": 0.671384648994924 + }, + "parameter.transformer_blocks.6.attn.to_v.weight": { + "max_abs": 1.8169648694992064, + "relative_l2": 0.7366241094623122 + }, + "parameter.transformer_blocks.6.ff.net.0.proj.weight": { + "max_abs": 1212.0285290527345, + "relative_l2": 0.26324446456438144 + }, + "parameter.transformer_blocks.6.ff.net.2.weight": { + "max_abs": 201.01442993164062, + "relative_l2": 0.24355468985475448 + }, + "parameter.transformer_blocks.6.norm1.weight": { + "max_abs": 8.35041940689087, + "relative_l2": 0.7145662697849823 + }, + "parameter.transformer_blocks.6.norm2.weight": { + "max_abs": 942.755732421875, + "relative_l2": 0.23473965726291585 + }, + "parameter.transformer_blocks.7.adaln_proj.linear.bias": { + "max_abs": 7.761792545318603, + "relative_l2": 0.3173178505919617 + }, + "parameter.transformer_blocks.7.adaln_proj.linear.weight": { + "max_abs": 1557.43126953125, + "relative_l2": 0.26760021703293335 + }, + "parameter.transformer_blocks.7.attn.norm_k.weight": { + "max_abs": 4.056298260688782, + "relative_l2": 0.8996356894661431 + }, + "parameter.transformer_blocks.7.attn.norm_q.weight": { + "max_abs": 3.4165216732025145, + "relative_l2": 0.8404564580857118 + }, + "parameter.transformer_blocks.7.attn.to_k.weight": { + "max_abs": 1.8837275552749633, + "relative_l2": 0.8492421136620586 + }, + "parameter.transformer_blocks.7.attn.to_out.0.weight": { + "max_abs": 1.1330161440372466, + "relative_l2": 0.83062026638771 + }, + "parameter.transformer_blocks.7.attn.to_q.weight": { + "max_abs": 2.040867094993591, + "relative_l2": 0.8312702507157053 + }, + "parameter.transformer_blocks.7.attn.to_v.weight": { + "max_abs": 1.3269965398311614, + "relative_l2": 0.8473707200751844 + }, + "parameter.transformer_blocks.7.ff.net.0.proj.weight": { + "max_abs": 2168.759089355469, + "relative_l2": 0.2835345569316502 + }, + "parameter.transformer_blocks.7.ff.net.2.weight": { + "max_abs": 156.14725006103515, + "relative_l2": 0.24158799615078969 + }, + "parameter.transformer_blocks.7.norm1.weight": { + "max_abs": 16.6573380279541, + "relative_l2": 0.9899422775170169 + }, + "parameter.transformer_blocks.7.norm2.weight": { + "max_abs": 1497.3561962890626, + "relative_l2": 0.2729074793630908 + }, + "parameter.transformer_blocks.8.adaln_proj.linear.bias": { + "max_abs": 13.849670295715333, + "relative_l2": 0.2716151788658495 + }, + "parameter.transformer_blocks.8.adaln_proj.linear.weight": { + "max_abs": 3290.2363232421876, + "relative_l2": 0.24767352098462966 + }, + "parameter.transformer_blocks.8.attn.norm_k.weight": { + "max_abs": 1.5559553551673888, + "relative_l2": 0.8947608029700306 + }, + "parameter.transformer_blocks.8.attn.norm_q.weight": { + "max_abs": 1.506643121242523, + "relative_l2": 0.9384623469999616 + }, + "parameter.transformer_blocks.8.attn.to_k.weight": { + "max_abs": 1.2076490569114684, + "relative_l2": 0.9773064344087591 + }, + "parameter.transformer_blocks.8.attn.to_out.0.weight": { + "max_abs": 1.310686831474304, + "relative_l2": 0.633518587303856 + }, + "parameter.transformer_blocks.8.attn.to_q.weight": { + "max_abs": 2.280301456451416, + "relative_l2": 0.9573002642277444 + }, + "parameter.transformer_blocks.8.attn.to_v.weight": { + "max_abs": 0.8934161114692688, + "relative_l2": 0.6871977728427118 + }, + "parameter.transformer_blocks.8.ff.net.0.proj.weight": { + "max_abs": 832.8469494628906, + "relative_l2": 0.2663712599558485 + }, + "parameter.transformer_blocks.8.ff.net.2.weight": { + "max_abs": 145.30633056640625, + "relative_l2": 0.24006590318715618 + }, + "parameter.transformer_blocks.8.norm1.weight": { + "max_abs": 4.724163656234741, + "relative_l2": 0.8735804939347517 + }, + "parameter.transformer_blocks.8.norm2.weight": { + "max_abs": 3142.030390625, + "relative_l2": 0.23299972377987604 + }, + "parameter.transformer_blocks.9.adaln_proj.linear.bias": { + "max_abs": 6.839249973297119, + "relative_l2": 0.2846913416140406 + }, + "parameter.transformer_blocks.9.adaln_proj.linear.weight": { + "max_abs": 1627.18126953125, + "relative_l2": 0.24772190477496486 + }, + "parameter.transformer_blocks.9.attn.norm_k.weight": { + "max_abs": 3.0111261653900145, + "relative_l2": 0.9238515507972674 + }, + "parameter.transformer_blocks.9.attn.norm_q.weight": { + "max_abs": 3.0268102931976317, + "relative_l2": 0.932745375224013 + }, + "parameter.transformer_blocks.9.attn.to_k.weight": { + "max_abs": 2.0182001042366027, + "relative_l2": 0.9300877917722118 + }, + "parameter.transformer_blocks.9.attn.to_out.0.weight": { + "max_abs": 0.7163272845745087, + "relative_l2": 0.6560424752082701 + }, + "parameter.transformer_blocks.9.attn.to_q.weight": { + "max_abs": 4.336956386566162, + "relative_l2": 0.9655093562800879 + }, + "parameter.transformer_blocks.9.attn.to_v.weight": { + "max_abs": 1.6555756378173827, + "relative_l2": 0.6974471762705327 + }, + "parameter.transformer_blocks.9.ff.net.0.proj.weight": { + "max_abs": 450.18311584472656, + "relative_l2": 0.2535606699307462 + }, + "parameter.transformer_blocks.9.ff.net.2.weight": { + "max_abs": 119.15561950683593, + "relative_l2": 0.23522267253064247 + }, + "parameter.transformer_blocks.9.norm1.weight": { + "max_abs": 6.78931510925293, + "relative_l2": 0.8447450481755553 + }, + "parameter.transformer_blocks.9.norm2.weight": { + "max_abs": 476.0451611328125, + "relative_l2": 0.23699570586893662 + } + } +} diff --git a/examples/models/minimax_h3/reference-manifest.json b/examples/models/minimax_h3/reference-manifest.json new file mode 100644 index 000000000..acfe57486 --- /dev/null +++ b/examples/models/minimax_h3/reference-manifest.json @@ -0,0 +1,1641 @@ +{ + "source_commit": "2681cadb4dd14c8980f6d760ba11751be657f261", + "note": "Regenerated on 2026-10-16 with the original commands at source_commit after the first run's artifacts were deleted. 1614 of the first run's 1635 files are byte-identical; the others are metadata (report.json, the thresholds file that embeds their hashes) and the video.mp4 / audio.wav containers, which differ even between two runs of one job.", + "files": { + "fp32-reference/report.json": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "fp32-reference/gradients.json": "1e933436ca157876aa7dd5b6f1f4b17d1a51329ad52ac9c20107033be16292e2", + "fp32-reference/tensors.pt": "3944d33e8331efcd6b6ef3a0aa7a7d5528bde3cb7a51cf3bbe27d47bbf5264c6", + "fp32-reference/gradients/0637.pt": "44f5bbaed42fff5acf99e8c204882a0c574de0b29d65fef631383e2dd5098d74", + "fp32-reference/gradients/0636.pt": "254731e038755b142a49f7751856177c062266a3a6c3b7bd59efd4187f0f1158", + "fp32-reference/gradients/0635.pt": "a48608cbefc9a5aa90a1981238b8aacef860f7225f0ed6b6e0b217e2fa89effa", + "fp32-reference/gradients/0634.pt": "a7909e90edad95fc8c1721f7e8dd5dff6b326bfdd4c525eb93e373b4dc1efbc7", + "fp32-reference/gradients/0631.pt": "0d75a0b1311b45f7d2d74622a61bfcbcdbd59b32f455fb16a45cb5d9a444f2aa", + "fp32-reference/gradients/0633.pt": "665c07fb50f760eaf86cb0de99a9b0550fff06a4f0133c34991ad08d99d7fb4d", + "fp32-reference/gradients/0632.pt": "54ca30b8ff8c43b8b41e49e145d526c8b9a678ac2d80b4cb645293fd4f4e8c2b", + "fp32-reference/gradients/0628.pt": "e853596f30bf518fccd1508de3682c15b89fe823f7a2916ede958b54efd6c2e8", + "fp32-reference/gradients/0627.pt": "90f3f9cb0982f472ccc91fd6e696c42bb204c1faef5de058bdf01ebdfdff2976", + "fp32-reference/gradients/0626.pt": "bc42369ec3de7ea1166281989736b960ee1bd92dfc34d28b3e5b8bc1e9c22570", + "fp32-reference/gradients/0625.pt": "10beec67d5f8a74e03451f12e1ce7867c3b9f8740287d2ee6a64962bb290570d", + "fp32-reference/gradients/0624.pt": "decd792208d9969a670b04bf357c0bf7a08061a8d44fbc401a378d0f2665be49", + "fp32-reference/gradients/0623.pt": "b70e894832b85566ff63c954bc817f2d22ec8edb0a1591f89fbd9986c194ee7b", + "fp32-reference/gradients/0622.pt": "b74e7fe0e27483e4bc02030c72f9bd8d9cf13fe3e9b36981d5e7c766520a2f72", + "fp32-reference/gradients/0621.pt": "3a21a0af227127a43604e0dc12ec4019d34ce597481ea019e85f5d2382c1bb2c", + "fp32-reference/gradients/0620.pt": "31679f9ddaebeecc04ed14f0566dabb117ee185e8665500232a691e95decb1e9", + "fp32-reference/gradients/0619.pt": "e90a3f84a10a51357abf17e880cf5f26ec4875423cfab369acc9837d5950a24c", + "fp32-reference/gradients/0630.pt": "00c5f09323896833fdb4e8f7e472ccc92a87e2c8d77ba635f5e3e9655fe1f67f", + "fp32-reference/gradients/0629.pt": "05658b8175118bdd51cdf1703f585e62381593314a06a9bcf296ce698dbf593a", + "fp32-reference/gradients/0616.pt": "dc2d855b1c2a8c0689a31f26dbcc49b52f59342d8abbfd02c53555653c5a7a99", + "fp32-reference/gradients/0615.pt": "273f5f9e8f5f285402b951ceaf434d44ba8fa594b32e57b6af586880c33b929d", + "fp32-reference/gradients/0614.pt": "8efca892a52bdfffa7abc41f2bdee17397e73e0c4bcdc28be888dd0959579c28", + "fp32-reference/gradients/0613.pt": "37e41946f8b7dc1de33eed0275f37562249fe0d55ce037e13f1c35a5683d211c", + "fp32-reference/gradients/0612.pt": "3dc925c6fdba3cd5033a637ff8ca873f9ae655eb0e58ab1360a864c22c122995", + "fp32-reference/gradients/0611.pt": "fb09ec9e84b92033dc0d4644f6ec72e0868830ecc362670782f198f3795dfc55", + "fp32-reference/gradients/0610.pt": "7529a771ef1c37b3415ca8de67e454c88c83952047457f6334b0d8ea0124d3ed", + "fp32-reference/gradients/0609.pt": "22e05312b01f606c9481481686665ec98bd68dd23bb522098769726c0081c0a3", + "fp32-reference/gradients/0608.pt": "180abc034fbed4b13c7558c20e5ba3d2a5ccca4cf1647e3e0d2addc25a01accb", + "fp32-reference/gradients/0607.pt": "1a34a0d1bb63396ee46373f80034e76cd06dcb7c56582f3eea15715eeaa681ca", + "fp32-reference/gradients/0618.pt": "bf8e30cfd617e225b5aa4beaa44a44be00e89e3510a0f46d3ed946b0817fb9ae", + "fp32-reference/gradients/0617.pt": "b2814b4164c79a314e50b652d97c7dde7bd18033a9e7af86ad3a6ce1be6bc839", + "fp32-reference/gradients/0604.pt": "33b5102eceb55ce3ef632e08e5f0d258d060833e9d34ab7aced4e1a4513f804f", + "fp32-reference/gradients/0603.pt": "bf302df2dc6a5fcdd4c94587eafdf84bec36dd716d902a0c35d01449a474ca45", + "fp32-reference/gradients/0602.pt": "3fe3279733a24f54602fc0b76e0413dbc967cbcf4a6890f67b3df16bc27ef631", + "fp32-reference/gradients/0601.pt": "a720456c6e189915151fe83a96c991c4dfec8327c83ff97c9f52627e54077273", + "fp32-reference/gradients/0600.pt": "baf9bd4336b597dcf9775959c63be60eb094687945e88a5702ede50f9c226a42", + "fp32-reference/gradients/0599.pt": "f9536b7ade5e21e340c122636b7761964d863e5abb25c97bf26ca26b37221087", + "fp32-reference/gradients/0598.pt": "d14af50b4470ca9245f1a2998949e09531af4d06f0d56b71b65fcff925ea3a93", + "fp32-reference/gradients/0597.pt": "bb9016a2949975d81fe4cbe92f4b3600078c360ad170715054596fc9876fbf06", + "fp32-reference/gradients/0596.pt": "33e3f6636ce02508301b8891be3551908e82da23db45b48b2ac461deb51b0abc", + "fp32-reference/gradients/0595.pt": "197e55d14ca2f0b0a59c16a8532add4336471be3efdeb16d67608cf512fed9c3", + "fp32-reference/gradients/0606.pt": "ca5e63333d5e4c48af48fb4c2aefa3420701d67b9d7c5cfde98aa306fc90faf7", + "fp32-reference/gradients/0605.pt": "646f515364bfc1ab4372208d51d405066c62e541a86f8e6021d8d07f30a2fe56", + "fp32-reference/gradients/0592.pt": "41f2515a72db0c0cbfc31691bdcfc2f683d67ab5877618c83480afa36a40fcfe", + "fp32-reference/gradients/0591.pt": "f453e824a82d73817f087d9a18dc6702d18075212b2280ede5268c49fe91c0df", + "fp32-reference/gradients/0590.pt": "2535ff23d17149d142ad967251ec4a659a281f0624488cf514f9ed0e658c98ce", + "fp32-reference/gradients/0589.pt": "8a3d3f193b56a3ede7ea543ea59a890f923a4fd91927667a02672744c249ebdf", + "fp32-reference/gradients/0588.pt": "c780e2182cc311a799ce3dd6399fff3a03fc1a4e3efa807ce22c25a5ba1508d6", + "fp32-reference/gradients/0587.pt": "562fc28ed2dce596e281aec9ed03f5a0757cbd8ddba265940bfb723866113bed", + "fp32-reference/gradients/0586.pt": "921043d6bf593d325aa59acf2d54e6e1a68c3effdf9b129293945f3477366794", + "fp32-reference/gradients/0585.pt": "c972070823d14a2f3a93b508204d9b69dd34b974a4bc408b5bc704de971b0553", + "fp32-reference/gradients/0584.pt": "d583d795b1478b779ade820ad6d5d578afd3c57d225237a7022926ae90ee0b53", + "fp32-reference/gradients/0583.pt": "6193692f06308374c4b957e22eac7903d1b3ed2ef8ea3bc9ede70db7440b23d5", + "fp32-reference/gradients/0594.pt": "74e34ec8619044dda13074a19b9142e24154a8172b8f0b942f005f29f2c7be0d", + "fp32-reference/gradients/0593.pt": "bae9313909674635cf9f33352f58f1b67bd568aca424ae9b25e6c5db2aecf2a9", + "fp32-reference/gradients/0580.pt": "c04fe8669c0ecac74a1b733fe1fd7debb0ee0d046528ad533913b5e3089963f0", + "fp32-reference/gradients/0579.pt": "fa351a8a3fe71204357f9e802709beba5fd9efffe0d5e2b849faa8da300e99b1", + "fp32-reference/gradients/0578.pt": "1f40e8332afa6a18a9ca57add6205fbf08e8d1bb7c6eeedcf62b3f89e77a4dcb", + "fp32-reference/gradients/0577.pt": "aea66f63e383a5d071ce9a20061b17fbeb9807604529431f6522f9dedbf07d65", + "fp32-reference/gradients/0576.pt": "3075eb7cd56feab500821f11a8097809949095c2a31e3941d189a0375507731e", + "fp32-reference/gradients/0575.pt": "e47d9e80455259ebe88c8991d38b07b0a7d8617625a62a536874e42dd9edb5e2", + "fp32-reference/gradients/0574.pt": "b8e196fff622d71b28ff84c19e70c78bea5b803c4eb43c7f5e2f1dce6a6278e7", + "fp32-reference/gradients/0573.pt": "dac0a85d7bc8ae9b0c08ad76374ed7881285b1d552552ddf92bc426127a02d3b", + "fp32-reference/gradients/0572.pt": "6aa08bc98cbcca883cb53fa12f80f5f0bdc5be204d33a46734697709c40d288c", + "fp32-reference/gradients/0571.pt": "88eef77acf516a8cc66a709b738746bb732f33ee9d20e5a98035ce1b21b8ee62", + "fp32-reference/gradients/0582.pt": "f0a6f8cf798c92ce501a2ecb43727bbc0cb656dcf67af11fc0bc17f96518e576", + "fp32-reference/gradients/0581.pt": "c9f1cc11bbd50411948b1785c09c7a1b0b2bc0f43dcc39ae0f6a968f80d1a0b2", + "fp32-reference/gradients/0568.pt": "25eba65a0b9ab0c565008aafebc65f3e0ad5ca8e7d73a39c504f88ef92d30bee", + "fp32-reference/gradients/0567.pt": "c0a0a485d1e25c2c275bc4614b5b1251ad04bc07a0e1fc682fcce8a0a33ec056", + "fp32-reference/gradients/0566.pt": "12c60adc2b6d12959d1d7d6bc6e959ccdf47a31d5fb657203736b2721474b382", + "fp32-reference/gradients/0565.pt": "abe6e3af6bccd98f5418ff0994a99f38f55b71b849a6a6bd48300162a092cc16", + "fp32-reference/gradients/0564.pt": "f6c11d204ae2f6a1f8cf53b3d59f0cf8618ec0f518650784ac12367e9aaf785b", + "fp32-reference/gradients/0563.pt": "13a275ed3a508c43e5e8116a75c6db44ea27c26fa3f68bc8b65bf1e43f389659", + "fp32-reference/gradients/0562.pt": "9b3839bbe302539a908f053a0b96662f657313a74bea71f56460411f32a1c1bb", + "fp32-reference/gradients/0561.pt": "c1614ca2877742c4bc056e33e1b7e47e7f525ece5ec2f9f045875ef2688ef2a2", + "fp32-reference/gradients/0560.pt": "bd1f4e9f86449803a63fde68633487f426d2503208fa6a2a74354440138254c8", + "fp32-reference/gradients/0559.pt": "069e800a99f054f11ffb3709c942e76049e7b05314cf2a1cf370ac62c17ae99f", + "fp32-reference/gradients/0570.pt": "ae867dd8f4e994bc4275508ae2c1caf8f1f8a35c6485a12b02121257471d3586", + "fp32-reference/gradients/0569.pt": "efd950800dd04a89d271131c81a0715c53b2fc98240e84948eb7c996ac14a952", + "fp32-reference/gradients/0556.pt": "f5096f499f47380174c2b551409907df82f929c2225cb37c1466344b9ecb7e3a", + "fp32-reference/gradients/0555.pt": "54d1f6946665e11ae7c8a9f6a8f24c34f53da005f58bbe1a49e83325b2b473de", + "fp32-reference/gradients/0554.pt": "e266c4857acb712b3f34394d24982838318358afcf3422e6638270677fa7b3c7", + "fp32-reference/gradients/0553.pt": "de2138d1665e1e5fb59a830d82602010e22f9ecc70958bc51028e330d3dfa92f", + "fp32-reference/gradients/0552.pt": "0d478c3f601c1d04769663777ec496a96d095e936006d5995e348b1fd74993fc", + "fp32-reference/gradients/0551.pt": "cf98d824f0f80bf32329cb91290685c5246a590a477477157942a3ba9f8444dc", + "fp32-reference/gradients/0550.pt": "ca7a07a9ad353f2b5f334333c661501c782f35c24a7bb2ee1bbf24b2849740f2", + "fp32-reference/gradients/0549.pt": "942b6b6d0fd0e9c81d84705bf7aa9c9da6a63f269f2a0fdeb466a77742b7b6df", + "fp32-reference/gradients/0548.pt": "96f2fcddc5d2ce4bf5199432f70f08dbeab0fc5f8910bd0667fa7a6e1cf1d841", + "fp32-reference/gradients/0547.pt": "1250c5a774a0837114d2927fe25f9abcbd29f2541a06cb6fc711862a50f86c06", + "fp32-reference/gradients/0558.pt": "98d6c6c7b9cf59fc506ee12a825da8d9818091c3e2ac3dbc2571372b77223ada", + "fp32-reference/gradients/0557.pt": "b10177d042eded9e6ab9905cdb666be6741a462bbc68983f98bc9482d2a580f7", + "fp32-reference/gradients/0544.pt": "50edc8355df398821f71866d565d5315410bb7e315c38793a2bba56a4777362d", + "fp32-reference/gradients/0543.pt": "8fa3752eff89183e99bb8e2cad6580752007ec066a50c0befdae0cf56edd80a5", + "fp32-reference/gradients/0542.pt": "2f3305e9de3502d8c4eb39f6ec7edd88d095bc89a3c7564e653e11aac36478bb", + "fp32-reference/gradients/0541.pt": "fbd3c1d13b2b3c77c669423deb4b6bab16f8533b83ab21157fa8e7faf4cbf23c", + "fp32-reference/gradients/0540.pt": "c3345fc3b40045674a6056adb5a8b950c3d68809169f4b9f869b107c2be20d11", + "fp32-reference/gradients/0539.pt": "6a2034e9f9fc7f97545623164c6294dc5d8ad144dd47496757e9fa2c9a4b1754", + "fp32-reference/gradients/0538.pt": "9b69cd5698548300b19f15a6f47bf2da4c67cfe7fdcb1c8cd98ffff7237e4311", + "fp32-reference/gradients/0537.pt": "c501a28230fff172649e2582d00bd0045245be2579cba72a10c1bafc8ee71a60", + "fp32-reference/gradients/0536.pt": "e608e623f2f61c554f61f4ea3dba5d3a19652190609623150e30a0555d6b6540", + "fp32-reference/gradients/0535.pt": "f46efdb2f5aefd044af14259bc8aabb2401f805a7ee0040880133c128d260770", + "fp32-reference/gradients/0546.pt": "85c068aacfed9e8192ff677e7737d061da7d856223270bcc1c03f289a1304585", + "fp32-reference/gradients/0545.pt": "85f48b3f6304874c95a9809f6985ba2945a5a6030347c6dae41372001a1c1bad", + "fp32-reference/gradients/0532.pt": "943695429503883526618ca2bf73208139d79471fa5bb741f91ee921723af1a5", + "fp32-reference/gradients/0531.pt": "77a97794ff01fa1888a3d06512e4103da118d821fb255d81876455a0ea058fb7", + "fp32-reference/gradients/0530.pt": "28a5f17c058e33274d209ca9ad4ea9b0d9156804a99dc3efe19641800217f54c", + "fp32-reference/gradients/0529.pt": "90a288c366c9cdf0aa20c4ac71820ecf9c97ddc5ea2a3a73cba8bf61786eb85b", + "fp32-reference/gradients/0528.pt": "e90a01fb080473bebb60405512d74810ea034960d9e48c14a1b2375ccbeb5011", + "fp32-reference/gradients/0527.pt": "70d05a36b918f5413b63b8a09034a64d7b7c6fb223f6aaa6abe171724ae5ee96", + "fp32-reference/gradients/0526.pt": "ce091cef2919bcc951a8ef1cfd8128fdd2c46f628dd4dd1aff13bb2141e5bf2f", + "fp32-reference/gradients/0525.pt": "1207b3d5a9a08f6fd525365a405529c67a9f86f959e0c361881310e38037be2a", + "fp32-reference/gradients/0524.pt": "1c5058f32db55c0615d936c8ff4ee8b2da5e745fac42e88dcd01e8f77a0358b2", + "fp32-reference/gradients/0523.pt": "1796a50450a530a31f47ed133de263736b0034c0c8ced0981ed572fd882f6f58", + "fp32-reference/gradients/0534.pt": "4d7bf480f2b928a86e909594ae851461d42bdef139e1cf9a62cbdad51dad8ca4", + "fp32-reference/gradients/0533.pt": "670ea78d38faca647db06d2c2c6d1e53a19280e4d55bc1f5f851ed38d26a5537", + "fp32-reference/gradients/0520.pt": "43cc5f2f0e18e1f9ec61a750c87bd9db3cf097f453bfe6f45fa9c67a65edb647", + "fp32-reference/gradients/0519.pt": "2e72ba1a0ea33725378e093d98bf853cdc68886ce67cfe0d03169709b83e7197", + "fp32-reference/gradients/0518.pt": "da36f02ee00005f0702d5a0a01687bc4bfe87ee617ed4e7ede9ebfe1370cfa34", + "fp32-reference/gradients/0517.pt": "7ec01591947add93c475af44b6e1d8728e0d0fea292f998b2e92001da5a8ff22", + "fp32-reference/gradients/0516.pt": "8b36ed32494a658f7d6037e57bb90f5b0f2c44d10c86c63d3188cbe2701970e4", + "fp32-reference/gradients/0515.pt": "73794c97a169cec101c036e22b867b1f80c665eaae74454119b08a6cabf228cb", + "fp32-reference/gradients/0514.pt": "21ac37ba854658a299d18635b581b1a731fe7a873090b6bf2527609744c2f706", + "fp32-reference/gradients/0513.pt": "c45611708b5dcfd3f37fc561e29706b20552566dfe07ec802639a9fb53dc76cc", + "fp32-reference/gradients/0512.pt": "f17e67901af8cee77e500cd7ccc816658254d0e9f0e7f97ac966e37eec414066", + "fp32-reference/gradients/0511.pt": "72ec8e7ff122c89b9ad5f081bb38780a05798baf5db199ec65b0bb9531861a03", + "fp32-reference/gradients/0522.pt": "006b34aa59e38dd230e0deb54437a4d865f21dd5973904e65180fce39e10d9ac", + "fp32-reference/gradients/0521.pt": "6568995b15407cbd88fb3acf9d8d53b66a73ece10977a45a130a9e63f61a8fe1", + "fp32-reference/gradients/0508.pt": "b06e0d4f4ec0c738ffa9aec689ab50ccfb9f1f5b21347ed722779a9a9270c6c8", + "fp32-reference/gradients/0507.pt": "543df7123d3a2a58123ab577238b08f4ae9a996550ac427f9b4e67b0c410beae", + "fp32-reference/gradients/0506.pt": "c1281ee12c79c65278789b3f3088a3878cd180bee4b356eee859f1e15f521662", + "fp32-reference/gradients/0505.pt": "28d4541906dca31a95663cd8871ad0d1bd8f2f9a9b0ec4fec2fbd131e929678c", + "fp32-reference/gradients/0504.pt": "3e0a93c47d2e1aeecf7b88523df493ed221452dca89ad7a7303350ac7e97b49b", + "fp32-reference/gradients/0503.pt": "6e5574a0ac9f3ebf84943be5e0ea709dfff1d14736162babd9dc6822a64cf236", + "fp32-reference/gradients/0502.pt": "2a58c4eccb5676c4a62438925c86a5b4cddb1fcac574428979dbbe0ccc7e6e49", + "fp32-reference/gradients/0501.pt": "13d5eac783b4d9fd547ac31888e1a81187e810d13de9d4edb7f3a74a9579c490", + "fp32-reference/gradients/0500.pt": "9db2150d5ffa5de04cec3cea3c918955b56ebcd2abdccf012553b47b1e09c52a", + "fp32-reference/gradients/0499.pt": "a6ce0ef449017393dd6d0f193f43a390c1694d50cfa234c5ba87575db6c255d0", + "fp32-reference/gradients/0510.pt": "dc6b7b5902b4fdd6cc8f95a4664ef737148c4f658a37e672f599aa49aa878800", + "fp32-reference/gradients/0509.pt": "f06a43d2844738a53cc4465b536eda78a51363c42c4eb301191451ddb72649bd", + "fp32-reference/gradients/0496.pt": "40e8e219f53a6fd32cd13f25f016b5fee910167ea3693d277088df2e30aae283", + "fp32-reference/gradients/0495.pt": "fa09566b9cc7aa8282c5c633a6e020818489d151dc860771a7aeba5130569182", + "fp32-reference/gradients/0494.pt": "a8e1bacaf41297dd7a02e4c0694683410db0cc83f74289f48dc4ef54d3881c20", + "fp32-reference/gradients/0493.pt": "38826546a90bfe5288484fbbc7e16fbfcdb6f99b73afda65f0ae2319b1cf62e8", + "fp32-reference/gradients/0492.pt": "206d571b87c88493d5e302d424c1e21c9b4f8fefb1c5d5b676deafa9dab5c943", + "fp32-reference/gradients/0491.pt": "2ff197a7f8d707fece3e2471469cbbe1e82936f3ee36822c6f95a87ab2cdd726", + "fp32-reference/gradients/0490.pt": "0390cd25a8adf59249cda2d30c3d06aad0d7bade0eb0f76ae049481fd8d42532", + "fp32-reference/gradients/0489.pt": "f4ea4a3182afcff0da048f60ecd1bca38def8ed141ae3e0f38ce36c63c95aa91", + "fp32-reference/gradients/0488.pt": "400662f9763b4ae3850e2d38ab235b2ed26041215462c2c2d8908312465e7aa0", + "fp32-reference/gradients/0487.pt": "30a3ed098eb7eb5f14732fd40665cf8688b246cd8c0f473b6e714525b2c15fdc", + "fp32-reference/gradients/0498.pt": "d825eee41db8392a5a5732534a5ad28a235d0de217b3f3e53b5a752d2cfd2a92", + "fp32-reference/gradients/0497.pt": "704279a36d508ded75e63d94e52be769c79853cc9dd890a72b644962413ccede", + "fp32-reference/gradients/0484.pt": "9e0cb0aa6de8230b78bf2e6c3ce1e545b697ae5953c5fd40497f5cd291f10673", + "fp32-reference/gradients/0483.pt": "e661fe7eb319d398ccf230f331f638c6ec766e1a9da286450ab52cdc0a320658", + "fp32-reference/gradients/0482.pt": "dcb47d66beb78c80a357abdbc490ab890fdf10a1c3ba5585a8f07637c42ab6d3", + "fp32-reference/gradients/0481.pt": "326ecce695eb2154fc11dcb8183c6d6d3f2c539492c7dda9cface28f4a830f10", + "fp32-reference/gradients/0480.pt": "6a1bc745f3205e625de2982a281dbfa8fe9e64a0720dda995b39a4f79159ea05", + "fp32-reference/gradients/0479.pt": "c261a24a15163c2363e7a3a22b8a25dbe3612ac588c1ad8881a9a6f80365a323", + "fp32-reference/gradients/0478.pt": "14b50f0bb5050805c84ad92dc420254d1ec031c4e8574f37f4a40e4f80f76603", + "fp32-reference/gradients/0477.pt": "4c3a42f151a26de43655da95b7987019eff5a4f041e2da6dcc67fac8d3fae685", + "fp32-reference/gradients/0476.pt": "713271643a14adf133f189ebbb69ba97aea5761a975572a8be93d1dd8dc4d9ba", + "fp32-reference/gradients/0475.pt": "5a399f516b0e67d766153961a6f4031f82d1ad223bc169a1d851c7e5321c5f6b", + "fp32-reference/gradients/0486.pt": "46fdef3e713423cb93df30ee7285674bc26d14e5e7060fb3407987a763be2583", + "fp32-reference/gradients/0485.pt": "d8975b3bd49a354ed8a856863bb38cce8090082bca012f7bd6238958567f2b1d", + "fp32-reference/gradients/0472.pt": "cbbb64c0c541c893a28a09427a0d1ebfc3892f4d48dd648fa2ed48d16b39dbe2", + "fp32-reference/gradients/0471.pt": "322be960b040310dc6037828a1ab58439059daf66fe1a36a1b5436e710624cd0", + "fp32-reference/gradients/0470.pt": "ef81d5a07a2ba847257b5cbef97233b346cdce42417a4df78748802ad0251fa8", + "fp32-reference/gradients/0469.pt": "59a1240c5f65c6f87a658123d7c25db803a29a30839503649fa5f75bbb7382c8", + "fp32-reference/gradients/0468.pt": "957a78f6d9061c3f7f4a0e2c8dce6cf958425c63178dec09a9789fd20944211f", + "fp32-reference/gradients/0467.pt": "4aee1927a6078a296a97226231b13dabc98d6e79ed5584f462dbe60e1187e4db", + "fp32-reference/gradients/0466.pt": "8dd58690d2b120c1f9a6519ec7f030351206f11abc01b49e1c513faf3c9d7a22", + "fp32-reference/gradients/0465.pt": "f21a5c5a9cc17a30f1d095e64d9e2d0e6442c80ab3f4c6fe185fe4b129ac8019", + "fp32-reference/gradients/0464.pt": "f425133d2df900b0fdb4c0308fcb0bb624fd8b4e5d82c8bfd0e6e2a85cf53b6a", + "fp32-reference/gradients/0463.pt": "3e2b380eeeaec92c66c2042d35c5ac491dc0870025eb0f578d7070db0324a80d", + "fp32-reference/gradients/0474.pt": "930be3a0d6fa35fe355628d8a3e203604c5174baf0f56e91de547cd36ebb9595", + "fp32-reference/gradients/0473.pt": "5e9f2affab05ae19018a73d6317635b76bca52196226d4b36f50081a3acbe1e4", + "fp32-reference/gradients/0460.pt": "c6a71ba438f79c7e7cc15f581aceebab346bf2b3ea3aa97bbd28cff471c28205", + "fp32-reference/gradients/0459.pt": "35267c9a1f22e7ea83e08a7e1b47760de3a119f173a941ea58b5e03d5fa2564e", + "fp32-reference/gradients/0458.pt": "7b3cc3bfc0c6c03c9694830532949aba765d6c8eac45843532328973f7a9e878", + "fp32-reference/gradients/0457.pt": "c6ff9b6a3ad24ba2ed0dcbf80a8c42bbe3d23e96d2f64daef69f8cfc07b6ffe9", + "fp32-reference/gradients/0456.pt": "603609358dbbb1edf4d134e4edf8f8e2eafb7980a6d9ec4a1847e3c92c1abb6c", + "fp32-reference/gradients/0455.pt": "8c8805ecad1a9281fb9aff6f25c8239ac38bd35fc10e7afa5055d607dd18bcf2", + "fp32-reference/gradients/0454.pt": "984b2d3a175b67d0b2b3ee91dc6a72c0e2074faa0958386c2f9dfab2ed65bf7f", + "fp32-reference/gradients/0453.pt": "deab52b241bfbe5a72a47bfd8115fa91df25f0d3732cb8d8de8db17afdd42345", + "fp32-reference/gradients/0452.pt": "f2c89000d816ae526c707d1806862160a80f69cc9932a131f28b07453c10a302", + "fp32-reference/gradients/0451.pt": "dd33693e1fc442b2779b7192aa95e6607b396fd8eb34133fb1259dc0adc4bdfa", + "fp32-reference/gradients/0462.pt": "98cfb654e544a673ca011163aed6f1e9d9f9c0e551ba96147ddd8729470f3d06", + "fp32-reference/gradients/0461.pt": "2c03f7c8e551207b401d849529c1861edad03f1c18a7087f9b0dd80116348054", + "fp32-reference/gradients/0448.pt": "81cd65ae67a59af052232c49299b08cd30d669c216602207bab48f49182d8828", + "fp32-reference/gradients/0447.pt": "af8e51ba5a0d5e4bbecb1dc7b107a77a14a1ac487da2ee16185573f7be5e3455", + "fp32-reference/gradients/0446.pt": "677a8910044036fce0c62f5bcb84b7a372da4c3d7181250c0ddb0b8a63b13e09", + "fp32-reference/gradients/0445.pt": "487db7dcc022a52162cd98cf26dd8147355c608a21d13c6af16155bb82318070", + "fp32-reference/gradients/0444.pt": "df3201fb9a082cdb9d93c8066e74a5318681a60f5f924ab7e064e61030ee8974", + "fp32-reference/gradients/0443.pt": "97af11e81de3f569e077cb031cd98f894e6c5d3d4588867302e776def0f5d1c2", + "fp32-reference/gradients/0442.pt": "127a82048329e8f6784e0774a03c1f19397b9a69feaeab3e6b7b3fc244c84435", + "fp32-reference/gradients/0441.pt": "6a76887e80fc3135b73a551ce908e64e23e78b67dfaa5f4f4988b7a8a4d3ba83", + "fp32-reference/gradients/0440.pt": "7a6edef7f1ee390b8c52a06161c7fe053acc7e2f38c178f28203643715627021", + "fp32-reference/gradients/0439.pt": "993197a391f39470dd4d0f5868ac165767b54b7abec0465e77ef4a2e62a7f6f8", + "fp32-reference/gradients/0450.pt": "973638a3a1b1673a22a62cbdabec4617f0aee24a103978c2cd6aba5ee672047d", + "fp32-reference/gradients/0449.pt": "2d54af03139f375dede9804624777a461dcbb4e65be05c0e4e34fa53c0c475ad", + "fp32-reference/gradients/0436.pt": "8ccb65336ec154b5ceec97e7c7e77771959094c970304cbc9e669eedafbc984b", + "fp32-reference/gradients/0435.pt": "a06840af610d2d695f155cb579496efc71a48eec470900a9367e4b4d2d6d0589", + "fp32-reference/gradients/0434.pt": "498768b33a2ddfe59e2d05643197b65d88532805e5ff2424bb13ea34183b335c", + "fp32-reference/gradients/0433.pt": "402b8c894be2381a1b3effb145f8d7c2b315fe7d861a2b554171f1dd5c5932f5", + "fp32-reference/gradients/0432.pt": "281948db17d5cba6b3147be99d40d83a61a0eb584afb64a5a3559931cbb1f1d0", + "fp32-reference/gradients/0431.pt": "45f30f12ed0ed990450967f11f708e01131c7010ea8a911cc30dfceac68c54ed", + "fp32-reference/gradients/0430.pt": "aba14a507adb1457d5bf2f1bd78252a402195b9b79eca8daa88b5ed5b8af64bf", + "fp32-reference/gradients/0429.pt": "c9da0de7771e8ad19613ce08bb8de95adc5b4c9577ba57946f13d7c1ff3e8832", + "fp32-reference/gradients/0428.pt": "4979d33eb8e70d2c8951a180cab912154ef015b4d7808aeddf061043ee5c4331", + "fp32-reference/gradients/0427.pt": "1e4030f3ca19efea712ac4f1939f3522fb3dc630cc79766b86f0fbd9ef543b23", + "fp32-reference/gradients/0438.pt": "81fd57d56eb2af15ce2d8e45c56250cbc5e9b2ece6227bf71ecb12c73277b7be", + "fp32-reference/gradients/0437.pt": "eb3313558184ac8a8dc339736ff5c5658a2ed3d865d43cd585072c66fa8ff2fa", + "fp32-reference/gradients/0424.pt": "0cc8f805a9744b814159f4e6b272e70a1b92662f85d604db8d46acca1128c38a", + "fp32-reference/gradients/0423.pt": "5530b03cf876b7fb8a5450f1f0709eb6e91ab4b5e23ac76d5fed093714375eaa", + "fp32-reference/gradients/0422.pt": "690554a5d80a0b98ba6378f3c6ed6d41357e4899bf9bc88cebb32b8bae9abfa7", + "fp32-reference/gradients/0421.pt": "9879ec70748ad0f7c274d709b1d8a24335398aaef2db6e09a2cd7d8f912ae3d1", + "fp32-reference/gradients/0420.pt": "c2e67c2a62f16d8f19068d3b9bad00bf1dd9ca62cd2ae6e54232f6dfb71e3b66", + "fp32-reference/gradients/0419.pt": "31be6eff12f5ad14cdcbb43dabe3c493911a953f32dc1c60f06306fa6310e2c7", + "fp32-reference/gradients/0418.pt": "ed3d0feb5ed1eb58bdb4e516aee3ee1caa866d0bf9d9c04b37f25e464317e108", + "fp32-reference/gradients/0417.pt": "0740fdd7326f0ea2c1bb4ea7140359b1f329341d3f20db1bc1f8c3d0b5d927d8", + "fp32-reference/gradients/0416.pt": "556ac4f1e52dcdb880c4288242a7d1761021b7cfe7d3f9c3c7671c49d0ab1a65", + "fp32-reference/gradients/0415.pt": "46a96495be56bb52abdfd7267938fd55044658cbb2894767f4f9ca014f698fb3", + "fp32-reference/gradients/0426.pt": "e070b09cf801ed3a8a65e5b2ec63ecc1642b2389bbf8f7486a9e792af4562c3e", + "fp32-reference/gradients/0425.pt": "55c1d8439a2cd53dc49bc35c43f76ca4a3f76a2d2054cd85e025bea41cd3991b", + "fp32-reference/gradients/0412.pt": "61dcac766920bc765df7fc69fd30e77fa8a07f0874bfa9ed4889164a60e37cd3", + "fp32-reference/gradients/0411.pt": "99c45965ccd1495cc6e7821ff2030354791c435be19706609dffc55f52197429", + "fp32-reference/gradients/0410.pt": "81b4d92fbcef872dc767575bd04fc04d0a99e7d73c9bd7b534753277bb4acf20", + "fp32-reference/gradients/0409.pt": "8304957c813471f601c71ea518fe5a49249638d0413a3f14a20e20df00a17ae1", + "fp32-reference/gradients/0408.pt": "ca69b37d38be66fafa21ef99922f3e6ce737701ff17c862db68779f15d9eda15", + "fp32-reference/gradients/0407.pt": "b927fc05062ecbae856da3b7bb845d5f46988a2b11763d02720a60abffa2de35", + "fp32-reference/gradients/0406.pt": "b2188a96cd251ab54b15c7f05be6168823cdec771bc8a5563cc2ed0ee713c57c", + "fp32-reference/gradients/0405.pt": "cd009b42b8316e84b06670df68c3f1056058ce1b3c6ea5eebcf2e39ab7fdc7a4", + "fp32-reference/gradients/0404.pt": "0c092e2448b5f7c815bbd67c06052019551ce387d77000e192841c8aab18878f", + "fp32-reference/gradients/0403.pt": "1c10e4c4aee8e5232206d8b8268d2e98f768ea5075122e50d99f7273d13c282c", + "fp32-reference/gradients/0414.pt": "a1a5c98c122831855aaadffe0f55a746acb3562b5cbacbe983640a880edf74ab", + "fp32-reference/gradients/0413.pt": "d226b6558e5ef7bb333da5318dfc2c998674634c47c59333fb4812bff8df1565", + "fp32-reference/gradients/0400.pt": "babdf8d72aa4ebadf25177e8755a63e216a639ca599c868583d71158aa4b2ad1", + "fp32-reference/gradients/0399.pt": "f4c1f48734bfc6a92d095bc358b0bca49e29058ec8bb0bb312d4662e8fe4a6e0", + "fp32-reference/gradients/0398.pt": "993c0beb924ca34bd796c74ff2ace0e3b6b7b63189e6822fed76ae022e40c151", + "fp32-reference/gradients/0397.pt": "cb1b248337dd8869f80f077de96241ee6a3e0f1ad89ee70ca32797bcebd7d2b1", + "fp32-reference/gradients/0396.pt": "479869b785c2e46f81e047701031624a34ddd0daa1fce350a4c4c13f5708f046", + "fp32-reference/gradients/0395.pt": "02bf140d6ca2c7b05b627c671771d58a12525b556fea24d0e6b3cb1da8a19894", + "fp32-reference/gradients/0394.pt": "6ea8d33e05f192a3904d0df8238b4d3dc06d60885170d677504109b0b2f42f7b", + "fp32-reference/gradients/0393.pt": "497dd32db340a8602a5667f0967854fd47efd8b0eedb0d78a1b1892dae69b0c2", + "fp32-reference/gradients/0392.pt": "e241fdb37146d5df78d6a5cd2a386af505a314af8cce700d4b31b858412aa07e", + "fp32-reference/gradients/0391.pt": "4cd2bc752abfee65b34d656a719e3d41ccfaaca69b83252dbdf441b0ced575cc", + "fp32-reference/gradients/0402.pt": "b269ddd15d6b3d1c7eb93abaa14eefb096f11a07ae2e1a7b40b00f8479969364", + "fp32-reference/gradients/0401.pt": "c3f791af8aa2f3dd31b82faf4e24be613ee1af339fb3dd0b72307e21d905b3f6", + "fp32-reference/gradients/0388.pt": "055d0f42b1d00d33caafb3b36e0a191d77e664ac9fa6d1622fb73de6ed29cf8f", + "fp32-reference/gradients/0387.pt": "09eadfc79f8d76fc9fc1527980e87fb32cdabbf47820fba6a2bc9a96e417e112", + "fp32-reference/gradients/0386.pt": "bcbac108b3e79a5920c78755848acc74d3ea88ab83c7d4cfb66a3ddae5e49701", + "fp32-reference/gradients/0385.pt": "dab9514a1d728bed815212d68f49ccd086f65685636de4f6428d3a4b62aa25ca", + "fp32-reference/gradients/0384.pt": "61c28bf78e73158fa599ae5e4fca1fff10f631fcb242fbb19603dc063f42e5b6", + "fp32-reference/gradients/0383.pt": "f98993334332d6b3fb603bff01157ef4121f68e83ac4ffe866d64ce7e7f79e74", + "fp32-reference/gradients/0382.pt": "da0dfb211b031fac433ddc0f22bc5837e944ea9724e8912866b8085f4337c6f1", + "fp32-reference/gradients/0381.pt": "5c4b234b4780b7aaaa601c4cfe152bba2369655900fdb39f4989c8e8f3899c2e", + "fp32-reference/gradients/0380.pt": "0bcafb5cf2c1c7521b15dfaadd32d254675adb96e9b0bdf2c2350f60ebf1fda2", + "fp32-reference/gradients/0379.pt": "fc1adaad6c208cf4ed78501d6ca918624f9b5abe395267703a9b449c3a391d89", + "fp32-reference/gradients/0390.pt": "926fc65ea61633bc44e4d42e1e32eec564dce7b3043d617512e0fabb5f90b087", + "fp32-reference/gradients/0389.pt": "d7f9ef3c7b1aab9cf97eba892e19793aaca7b4270d2a04411eed323c608b6fd4", + "fp32-reference/gradients/0376.pt": "fca479e9c367975ec12aef6bcdcbbff639faca4441503a7c4a7388ddbd1e81cf", + "fp32-reference/gradients/0375.pt": "8fedbec498e53baadf9b7921de73e53fc3b4ef06ac7a39a7717b09511909eacf", + "fp32-reference/gradients/0374.pt": "cb14cd46f5564f1f558ca0ed65cdb0782157ba931887de4efdc973d0c4d79f94", + "fp32-reference/gradients/0373.pt": "6deadc0364719fe6f8fc6c5dcf3c308bc540ed97a46da3b17e2304eb9a197dfc", + "fp32-reference/gradients/0372.pt": "6d87e6164df3fb9e066120461671f47be986b399d2a97890fb86e1ccc250dcc4", + "fp32-reference/gradients/0371.pt": "26acf2f7ba06424123a879fc2c08251fc1afcc72fe71ce13764fb14f21a9e247", + "fp32-reference/gradients/0370.pt": "21a95af2c977ede2b79e6a14e83e70f44b5eea58fae021bb3ba189505823f6d5", + "fp32-reference/gradients/0369.pt": "5040b554bd3a9e027cdb0c3a53b609b6eb2b90636932c575a00b5f9867244a29", + "fp32-reference/gradients/0368.pt": "c5f6a9a097c9fd6aa9c978ffc1d9b5d223370be963f293e2d9ebdff9acffe1ef", + "fp32-reference/gradients/0367.pt": "141fc2a9740cca0e795932932f082b5f0a6fd2f4bd235f5ba88b3d3550703767", + "fp32-reference/gradients/0378.pt": "2a5160f53f07d9545f97166fbc156e0a998557d251201100b5c8ad04e72c3216", + "fp32-reference/gradients/0377.pt": "4b891cea1186ea29963536e51bff3f61f4b131e9902396104eab8fe0883a6301", + "fp32-reference/gradients/0364.pt": "297e433afd8df2f4e363d2deb099d38e41c13ddc4acc0c5ac4d5beae92353886", + "fp32-reference/gradients/0363.pt": "830c8fa1f764f5d3c1a89351070f1ad76e4340273e46b38fdcc068b2f1a4fe04", + "fp32-reference/gradients/0362.pt": "b71b38fbcc78f6d31641241633e001f20db6f57529c570c8071e1147169fa42f", + "fp32-reference/gradients/0361.pt": "cfc0ea845a88b63b1e22c922564d39fe64a6d7e9eb0f049905bf513245625f5d", + "fp32-reference/gradients/0360.pt": "c1b8a40a0859cdf37efd0133c492c043a5231f12f616cb7acf46601ac3af6899", + "fp32-reference/gradients/0359.pt": "6da7d1c92ea12e9a8562f985157967e36f840e07f5868d41f1d9873941d289ca", + "fp32-reference/gradients/0358.pt": "98381516e776803417d7441f825aeef0cce21a357310943c3851a2bb914790c1", + "fp32-reference/gradients/0357.pt": "d230e6f3b0f30ad25069564c7f6fe583acf3ca88840c86bac1b404c30f590caa", + "fp32-reference/gradients/0356.pt": "f142bbc7005fd0daad1cf4f2b9e5e698752eda421b75b18bee14a1f3fb0f0ae3", + "fp32-reference/gradients/0355.pt": "5622fee6dcaeae422325f614b9feb9d809d1cb0d86dc35f90735bb015baf51ef", + "fp32-reference/gradients/0366.pt": "be69db3dcc33b8749c8b21bcf9da4afcccfb68fb671884772afb623c57b2fd99", + "fp32-reference/gradients/0365.pt": "4c254dd353ebc1b09b7eab21170765e489b8d9ddd798dd99addf772285fe0cdc", + "fp32-reference/gradients/0352.pt": "d0ea48a2b67f353d9f3f8148e93069b85182cb3f306c5c7749e989f9776be585", + "fp32-reference/gradients/0351.pt": "664941612ea2d708d020617616bdecc811efc94c6baf4b3d8e8fa060a3f322f2", + "fp32-reference/gradients/0350.pt": "a46c006861b58b0959738f1a7e4af5385b9e9ff72bf8f6e1085856c1cef6adcf", + "fp32-reference/gradients/0349.pt": "01646aee7e5ec1f0c3e0b270a89dd93955c977981b0a835d88018e771f191078", + "fp32-reference/gradients/0348.pt": "dad5413edd983c8a8bfde79146d99dfe938564fa8d0029cf1279e28b48690186", + "fp32-reference/gradients/0347.pt": "8e67e8e8331059794a4da7d91441617ebaeadde4879a1acfd23d61c1fc2703f7", + "fp32-reference/gradients/0346.pt": "7c44095544782f10c97dbee924dd647e6297553676ab13543a3d832f2c32aadf", + "fp32-reference/gradients/0345.pt": "3b91220e343c52b71eac538752c6e2661b4756ecdb3a5d314b8546e5d0c7a77f", + "fp32-reference/gradients/0344.pt": "18529e38c196dac0c5fe86c3f7bbfbfb02db2875d9fd89d0a722e6c699e12573", + "fp32-reference/gradients/0343.pt": "0ebb81d928088e9a98ae13c93db5964545e382beee8cbf9c58d366c59e19f97e", + "fp32-reference/gradients/0354.pt": "7f28f9f54e3efc7e43c30f4beb199b0784eaf365e86cb343f8579f2cb31cceef", + "fp32-reference/gradients/0353.pt": "e9bf6f6ec4046508264dc20f914f629f77e9d6e3010a26c47334cf22cf22a0a7", + "fp32-reference/gradients/0340.pt": "30c1ca0449204c49165ee0ae647516732b457697e4a0c46dbbeff6bb9bfe4103", + "fp32-reference/gradients/0339.pt": "bbb21d4990a263e4bd00bc6a685d963beb4e8ead1397e6710678804a78a697a3", + "fp32-reference/gradients/0338.pt": "1ec28f042a16d01964213fe873cd8f4358fa8b84341763307d0393fd7ab7f20c", + "fp32-reference/gradients/0337.pt": "31cbfb86edcd646f8e67050b218451fa2696bb22ba4599a6a3b4e32849c57236", + "fp32-reference/gradients/0336.pt": "26e3d7b7d6ea02900d252c202288afdb36b34b580393ddfd90f877ac1d158792", + "fp32-reference/gradients/0335.pt": "ad309b703b42366a9bd5df22967d2663240da28c9056a32f7a039f164b439c32", + "fp32-reference/gradients/0334.pt": "e6140fd83b826824ebe2dce3131947772fb0d95210a110b101b54ad8d99d32d5", + "fp32-reference/gradients/0333.pt": "91061cdce387adfbadecd2b52e31be797d093ccff436cf7c337483d26df30f9b", + "fp32-reference/gradients/0332.pt": "6784885842c58905df6382ba505a76d4271d5faae11d8486203ad93db3b8e086", + "fp32-reference/gradients/0331.pt": "f87214d58ec9e88064e24e472391264823ae025a939c9c0cff9408b63d72798b", + "fp32-reference/gradients/0342.pt": "c80472d05577d563baeeb3da54f51511a4d588b29e95aac30872ee23d617265d", + "fp32-reference/gradients/0341.pt": "1b97fe915bab23a8f2acdcb9a4e14e12c1bde2b495c19cfad61a5d19dd52a0f0", + "fp32-reference/gradients/0328.pt": "a46f533054a001ff28f32b915cf85a05d8b95bd4cd13ea4205842b611b1e6f9b", + "fp32-reference/gradients/0327.pt": "9e85dcbfb51f3d492c3c93b87d5cae0bc961110b3be6c94a6281fbb19608d151", + "fp32-reference/gradients/0326.pt": "61446f3b0ab7612c75cd61a325ec49977d827985ef78c0c9d4407073599eb586", + "fp32-reference/gradients/0325.pt": "f30573b2637ee6c35c97042a913a2d48112169bdcd6b31ab5155aabe4febdd0b", + "fp32-reference/gradients/0324.pt": "9ab6259674209f759dc294e731a10a80205973b98c201e8a2cec04a63a51bfbf", + "fp32-reference/gradients/0323.pt": "e7bc8e825b3b558a0740fd26425d103647650801e3cc1103fd0d39c4cf3fbfda", + "fp32-reference/gradients/0322.pt": "cac33dfe30e43034a3d814b379cd9ec6390e3307941ad5bed24a1c0a79d0f2f7", + "fp32-reference/gradients/0321.pt": "932fa66b71a1bdf1e8ababcfd9ecae9ce9ed1bd8589791b7342140ecb4cc5bc0", + "fp32-reference/gradients/0320.pt": "94d781055abb8fa924731009a5ac4c6f23b27d69364b92a367eecc5917ac9536", + "fp32-reference/gradients/0319.pt": "6047d1b068cf0252ef47033e1111f070a19919ae4deba4e95e6486b2fdd1367d", + "fp32-reference/gradients/0330.pt": "43b519b8c6d8a59c735918b465fb461b8026356941cfa3fc0fb777f4f8439f4d", + "fp32-reference/gradients/0329.pt": "c31311bb28b3f451096ee4017f227c457f0f9f4d9af8044ee906c20af5e37516", + "fp32-reference/gradients/0316.pt": "aa535cfd1f5e50eb830a4451df49b5b4a017b4a5ffa0a9cc3bd0ce7b39c7af06", + "fp32-reference/gradients/0315.pt": "0e3a47131d9328e1f4d9eb6d202ef20decab26ff352d84b3bd8c16c232cb998e", + "fp32-reference/gradients/0314.pt": "c900945c0b6af837c551a71df7edbafed618c66d15f8ac8be9340608c3bcfa49", + "fp32-reference/gradients/0313.pt": "c98d8cc2c560a21ca866f8d4a0b3ad3b52ec9997e8329e61f18db9aa8af1de09", + "fp32-reference/gradients/0312.pt": "529b655af8096c58d58fb1b3c3aff067b97a68729ea84e9604748c4ac54a7844", + "fp32-reference/gradients/0311.pt": "07369812478dbcca4562706fa2523f4b0262f8c2b8ca5e1ea5f6b10649ca599e", + "fp32-reference/gradients/0310.pt": "7d810f9a97fb1a1a5f3f0ad9c4919b6cc05ce2439b1989a2cd5ff28a90aad19f", + "fp32-reference/gradients/0309.pt": "8aae11e69cea48219c2994364b21b5939b6d72a082393881cb50ddbb4658fc4d", + "fp32-reference/gradients/0308.pt": "89bec3cbe3662565b5a735087e2468b2284b3be9bcd6efa646dfc783803f0747", + "fp32-reference/gradients/0307.pt": "3e5769ec13b31cc8dfb80caa425e24defb8c245815f3b04bcb2a021bba2356b9", + "fp32-reference/gradients/0318.pt": "6a58c2cd5662cf5f600b467b827730026bb8e7fefaf1ed0dd59dc31ac1b50b03", + "fp32-reference/gradients/0317.pt": "8349ec6d6d454cbee6161a21fe93381a31d99f041f545c0420df84511e3ed59e", + "fp32-reference/gradients/0304.pt": "56927f97e433dc8c4a4ac52ff00747a03e7c8f4c9e3014540ff256fdaaccfc2c", + "fp32-reference/gradients/0303.pt": "7eca12e84511596c21ffb80a2b25d80cded57e5c8ad6b0de1e785a6b9501f590", + "fp32-reference/gradients/0302.pt": "86bc7606eca062317ec2bcfefa4eda0cd3e53c5f2b5d4175281b4ad55a53672e", + "fp32-reference/gradients/0301.pt": "e826563d9b100cd3e3e22b7538d4961642fbfd1381b3973664dd96e5d7e50755", + "fp32-reference/gradients/0300.pt": "f1a98109766ef78c2eecb238dec287d4ff647fb8dce0c2caad229f853aaff82d", + "fp32-reference/gradients/0299.pt": "f948e050467b4d540dc4a9339f551c317e8ac12f0cdcc7b1ddac35bb50778d71", + "fp32-reference/gradients/0298.pt": "36d0016c7ca4e22235e2e676b5ad3b82678375744e2a3da7f86a63dc39506099", + "fp32-reference/gradients/0297.pt": "2e7e293b196e88cfcdda3313ee7cb9f2586af250b2dfd64894c52a71380636e6", + "fp32-reference/gradients/0296.pt": "44fc8f38aad56d152da5931db0c48a1d1adb09cf93da010feed7e718a6542d53", + "fp32-reference/gradients/0295.pt": "d428d1ae88d44f78d842e2f86c83a55c8614dd98d9146615c5bb86c76c5839da", + "fp32-reference/gradients/0306.pt": "47977dee08ee2be3dc9b3457e67fbc668b0748879f61f320e37e3f4802a0a60f", + "fp32-reference/gradients/0305.pt": "83a717de444b9e8c67f092a1486764954c44df1c5a3169994f02491ae3e2febb", + "fp32-reference/gradients/0292.pt": "e6f26acfa525135c68d7aa2a621385489211f29234ffaa1391f71d9136ba1abd", + "fp32-reference/gradients/0291.pt": "6568c1c9332be4f04d93f7c9df5403705244be1680065951f8d3111c10a76311", + "fp32-reference/gradients/0290.pt": "ac9d5e5ee0f40b4da80cc7b250c479e98d16fd24080344313298ff263978c342", + "fp32-reference/gradients/0289.pt": "9ccf2de8ca240f9ae55c7d54b517c71c39c9eec1e6e328cb0309c72255878cd5", + "fp32-reference/gradients/0288.pt": "b44208f51754dcf8dc83e82a48b3bc9390c5fbbe8ddb8d20c05c556ddcc0fd1c", + "fp32-reference/gradients/0287.pt": "4e9ee7340584f74e2769eaf5397214455a126302c5c9c769dda55c5429a9fdf1", + "fp32-reference/gradients/0286.pt": "e958e029137c0eaf4321f4e2becd71fcc591d7e41fc56f003019ab6af5cbea44", + "fp32-reference/gradients/0285.pt": "2eb3c1cf558bd3efbf1ba135be3309ea7f4625a69ba53d8ef7b7f7d540e5ab77", + "fp32-reference/gradients/0284.pt": "c9f18d0ba332d302fe1db2ae201c3159759650f697db7583df26fde6f1450460", + "fp32-reference/gradients/0283.pt": "57d52095ba5b541d7bf1e96ab47b16d12071ac63ad48386c120ff458492f9ca4", + "fp32-reference/gradients/0294.pt": "7f54a9b76f96e3c5a15d276f057d7215db701fbaff25104b7461a6a4e534d992", + "fp32-reference/gradients/0293.pt": "445dab90ea8620c59f2ede7de5ee4a79a9da227aab31767d45ab433116e2ceaa", + "fp32-reference/gradients/0280.pt": "bb173c55f854dc97199e915da49e6c61f763decf02e3af4fd2a9229932f600ca", + "fp32-reference/gradients/0279.pt": "09918cb3e216463cfe2229411b587edf64e184db29ac4449d1db6924241a93c8", + "fp32-reference/gradients/0278.pt": "24ead138531983bdc29613177e2a12d1e7c5a99521100e444930ffd84779ee23", + "fp32-reference/gradients/0277.pt": "0791e037b3685357f850875d4eaa239127eeb3170ff72985fd509950c7b1926e", + "fp32-reference/gradients/0276.pt": "bb7689d72d84a0d2ad58677ca28a9747c33fcb1f5e50a94843aa0425f5380270", + "fp32-reference/gradients/0275.pt": "f8b425462d1239aa42a7d33410233d39ac02a9df052d3d0cc1d4aca8d4594ab1", + "fp32-reference/gradients/0274.pt": "9e3a878665c09d337a3a349a2536f4e2dfceb3dc68f44dc9be4e2863a925d55f", + "fp32-reference/gradients/0273.pt": "c3f20c91bc0b816d38ec0c46219e0776fdb1041b59b5c3df7c3708052974ed7c", + "fp32-reference/gradients/0272.pt": "f9f149bc940f7cd388b1d2303857732ca510494c392a303d2404ef598e61a408", + "fp32-reference/gradients/0271.pt": "89c8029ad82790e44abbdbe41c813e520967b96cdcebb45e98362cb419c03d4d", + "fp32-reference/gradients/0282.pt": "01461d30708557a310882c612cdaea64dcf2f4db11372f6c7c932ff110bb8236", + "fp32-reference/gradients/0281.pt": "83a5000db45c95e38d4be269495bc82fc456ad2ce7294fe6728ce65e49bcabc9", + "fp32-reference/gradients/0268.pt": "74b1618f5397c50676856eb6ec11be0dfd0934599535f89fc2a2709791247620", + "fp32-reference/gradients/0267.pt": "48aa3a0fb63758a5ab02cceb6a0b2a3d6d339b8420085091b5a05e572819ae72", + "fp32-reference/gradients/0266.pt": "a06bf21815dbd6ef3b6b7f08022ad523ea583af27ccdd82b6a58221274a57f16", + "fp32-reference/gradients/0265.pt": "87d805f43b18dac430d0d63ca05dadd55a5ca32e4852856f0310381a281dcab8", + "fp32-reference/gradients/0264.pt": "a469ad817abaeec10adeebaffb9aea07c8befb7279634d2d788b29c40821a6a3", + "fp32-reference/gradients/0263.pt": "08a25e62a6f8680731ae589381bae0429081c03cfa255cafbe06d04b4ea08675", + "fp32-reference/gradients/0262.pt": "3904581f9ef042d39c459e0c62070846962d3d5b7200920080d815af59307c94", + "fp32-reference/gradients/0261.pt": "015443a3746e4b298cea82ff383735370802948a008fa478d831ae47fde942de", + "fp32-reference/gradients/0260.pt": "62672784add9ad400444f1ce37f8cdc5c05a2c04275920a2408fd4818a7aaa99", + "fp32-reference/gradients/0259.pt": "16cfb35a77918631f4affa6bf494d7b7ef28de3e3cb3ba078112bb380117f008", + "fp32-reference/gradients/0270.pt": "a798d90c7b1e37acfc75eb3c8fb9609ec51a607cf02350ddce86fdd1f96c1c65", + "fp32-reference/gradients/0269.pt": "4042dde7e9a5f8a2137dd9f3b21adf3444cbbe0eee4094dde3532a31e2616bbf", + "fp32-reference/gradients/0256.pt": "5d9ba1422fc50d1ced3a8e7845cf09dfae1e2034f62cb1646887000a2aa2278c", + "fp32-reference/gradients/0255.pt": "0f6ed31fe4b00d96ce35c3754d1439f08f4201a23869164840ad921fc7aa523a", + "fp32-reference/gradients/0254.pt": "13bfc45f4ff6d72e810ace53e61656914ea406da30913629b7dce9e6b7fcdd22", + "fp32-reference/gradients/0253.pt": "957610f17e94a964b3ce98682830f9d5f033338888d431e0fcdf4636192478ae", + "fp32-reference/gradients/0252.pt": "cdefa2367ac2f76cb315427623e7751731247cfc463fb6961e830209a8f36565", + "fp32-reference/gradients/0251.pt": "ff6c7251f77ffa0380a5b216f474e5cb89b2d0c65e78df0fadb9b782a2a8a7b6", + "fp32-reference/gradients/0250.pt": "331a77e539eb03f25188dbbba071e4890eeb5ee0353efa2d6e04fc3a0e4791fc", + "fp32-reference/gradients/0249.pt": "58cb63a6e35cb65fbb897bd7fb45cfd59a5a1fde8be6ea089e76f7428f4937c8", + "fp32-reference/gradients/0248.pt": "6573b84ca394825d338b9b3d7b3f6e3a329ca505ebd81480ada17009071343da", + "fp32-reference/gradients/0247.pt": "53f565fc6d5264df52f5b0947e017d072cbd57f578932041e2506e6290a15820", + "fp32-reference/gradients/0258.pt": "9710d5632b54a5317ff9563a004544de71fd583e2b328293e3977d42d46a333e", + "fp32-reference/gradients/0257.pt": "805d4b08afd97ace82c902e0aaee865d417b1999cf52f445219deeafd8b70d7d", + "fp32-reference/gradients/0244.pt": "15da11791244a0b7a82e9c4b8fd4a7ca450d999b8296ff93767bb52f13c09321", + "fp32-reference/gradients/0243.pt": "f52c438fc875bf0fa379b8f89a49bf45a3871a967290445066ff196a9005eba4", + "fp32-reference/gradients/0242.pt": "68fbe79a87f5151b3d6aac46b7915788eb0f06e8df15ee7d4d98ab7a97c473fa", + "fp32-reference/gradients/0241.pt": "32018a1c61616fe55d4eb061f4fe97691dd20384c47cc91ae78c9874e9c8c540", + "fp32-reference/gradients/0240.pt": "fe17142d8523a3e82c108c252aee0ef93968eeebf2b39109651cd22b11a2d1c7", + "fp32-reference/gradients/0239.pt": "52eeb688c6bde3a35c304e8b3723e80ca13ec8bdf0fe9a650c564ddc23016b5a", + "fp32-reference/gradients/0238.pt": "90523ef79ff4d4822ec45016d7d033d848220f9a215b54a2ac376aaad2ee73d2", + "fp32-reference/gradients/0237.pt": "4c844c2ddf60ebc6711e4f4df72bcfcdb2985016ff30affa67f1226c92f028b1", + "fp32-reference/gradients/0236.pt": "3cce77166f9b8a42cbcb57bca033688f6b6bff865884449454437acb361b1859", + "fp32-reference/gradients/0235.pt": "63dd5be6524df313c429b7c64049a8a710b5e0b8e3e7c67e456e99ea110e96b9", + "fp32-reference/gradients/0246.pt": "5c74a456b051370097290c6aae8c654804ef4f5d819ff5f36f21a0758855b6ef", + "fp32-reference/gradients/0245.pt": "f34ddd4fdc299207f52de538234b5071ff28a6f0f7ffe817916af075dd1f0f12", + "fp32-reference/gradients/0232.pt": "63a832e049d3a280a339d9667adc479fa409493e08f7911244114a128cd3e854", + "fp32-reference/gradients/0231.pt": "53a28778723343eeec78750c57c7ff4d69d610ff906a4d2915a50cb3b16d9af2", + "fp32-reference/gradients/0230.pt": "2a01daadbf34e70ecc1b5cc1fa8d44f6add72ad9a7f80c0e7f6310810afba0b3", + "fp32-reference/gradients/0229.pt": "ec95345a5bb65c8b09f7bfa56a6f18f10b693b81878dae1b13bf157b27f5f9f9", + "fp32-reference/gradients/0228.pt": "88893e25ac13fbb7f61549868f5efd41e0a310fad2ef92efe3a67f6a9bbd0f50", + "fp32-reference/gradients/0227.pt": "ece7264aeb17f5cbccfebe102f64ca822a9d36d786958126da75454f61d3a8ca", + "fp32-reference/gradients/0226.pt": "737db907ea2d2bc477b0f713340f6b6a3968ee2c4c29e199bd6343d96cf74ada", + "fp32-reference/gradients/0225.pt": "780d7021878199e7385d612bc2736ef6afa30bab555c9a5bd1d3551751a0094b", + "fp32-reference/gradients/0224.pt": "f31570780d9507fe089fc08ea55913db84390b42eb433ab42e0b95f771d2fe0b", + "fp32-reference/gradients/0223.pt": "2ffdf58e282f2776a4e13b1fea9aa737dc7dc7ea0486bc14f0489a5f6c34fa81", + "fp32-reference/gradients/0234.pt": "e16320be47fef92d457a1282ee5bacbd32d8406e99840016852c942dcb9191e7", + "fp32-reference/gradients/0233.pt": "ab8ffb2e37b4fd738831b1e0d14532175235cc9563ca59e6127fbf8816e4a011", + "fp32-reference/gradients/0220.pt": "b79aaeaef0a4df18ab2e8363ae21d175185f78b088b79232a8b1ee4a11b6f566", + "fp32-reference/gradients/0219.pt": "f7f8ed60e67a53edbaf38770c80f4dac3db2683381c7df9d2ce346283645cf67", + "fp32-reference/gradients/0218.pt": "d1919d831be54c10c1fff2a78f604774d280f7bc6e33d42a4e38a3a18ffd975d", + "fp32-reference/gradients/0217.pt": "a62dc9c713e64abf9433f7d2baa7bd0209649f44009d6a3a3d114808bbe1beef", + "fp32-reference/gradients/0216.pt": "1b0d2e81da8cc84a5297d622b0ce1760f83a125319b32ca04461d092608946e6", + "fp32-reference/gradients/0215.pt": "4181323351f413d3b8d96352c742893d072f220d5dec5c6cc2894fa049dc6e76", + "fp32-reference/gradients/0214.pt": "52aec816316351f1cf09e3d2f8a1fac7ddf60c2be166f442a344958ec465e915", + "fp32-reference/gradients/0213.pt": "cc75dcf14ed94d45a8f993f1e36e73124ec946134ff88e28d02ee11ee08087c9", + "fp32-reference/gradients/0212.pt": "7a0ab54bfd179305ef69f6555cdd724c2c0e017be27385fa6966088717c32757", + "fp32-reference/gradients/0211.pt": "1a4525da9aa3cc3ebbb07f10edb4c03a22eecb299868e253a4df142c64f49a87", + "fp32-reference/gradients/0222.pt": "7034b0dccf92c1fea7caba281b139c41527a853db8a60e560616b149a92a686e", + "fp32-reference/gradients/0221.pt": "e422fc8c1978ab439c4920978ec8148784135544af2d9abda46c2e317e8a4da9", + "fp32-reference/gradients/0208.pt": "17afd90ee3fd66673308cbca0d91138d6f63637f73bceb1e0e734b74bfb61961", + "fp32-reference/gradients/0207.pt": "73284b75bfa35fb6f4453632ad0021175c0ffb2a8cfb0192f9598d67d4b7abcc", + "fp32-reference/gradients/0206.pt": "be3853417df32c614a41562ba6ce11e2a3d12ba341b0113acae8fb27a7dd25e3", + "fp32-reference/gradients/0205.pt": "392c3545647a01f50144be8c0908bb5974cc995d8f416521c269fce3b24fe426", + "fp32-reference/gradients/0204.pt": "5c76c7f71014a0abeaf47753763483186bb4ee1d9b83db15a886fe9b56a4d128", + "fp32-reference/gradients/0203.pt": "ed5588d737e6ca4f5d5c71708516fe43ab90afd49f966308201cef6db624a4d6", + "fp32-reference/gradients/0202.pt": "0b1b89065e42478eb33ce7c71f580f528dda416b65c997b75202e4fe27e913b9", + "fp32-reference/gradients/0201.pt": "e95e304dd8bbf323ba3a644314d8901dfa3f5cf6f65ac14ee1b6446463029a85", + "fp32-reference/gradients/0200.pt": "21fe5ea00c1eb8648dcefc6a352d04c56d47b3d3975ad4f8339dd1091b570c54", + "fp32-reference/gradients/0199.pt": "1f22dac9532faa8d265e9678f8e63f759b4162ebe3fd236aa72a01e78889535a", + "fp32-reference/gradients/0210.pt": "3962095170b0967eeee1c86e6fc27256225160abf987133fcba67ee5e0bb88ba", + "fp32-reference/gradients/0209.pt": "77231e7a1d4b5b3a9c99ed04d022864ccd5a321e290515c3111764ec9408fdde", + "fp32-reference/gradients/0196.pt": "9973b017c9345214248b04363544ef1d3ece268ffd383d505ff510a6c8dc824b", + "fp32-reference/gradients/0195.pt": "9a960ae7c115ffc2bdc2458017ed9c33d2e1746f0a615d725b08f5f72e920388", + "fp32-reference/gradients/0194.pt": "e14d1ba9ab68ede0a055ac665761f62bfaa5acfb12fba96fec91fcfaf165c3b2", + "fp32-reference/gradients/0193.pt": "c3394f2d71d8bf6572ae31239fdfc7a7d925a93894911efbd778fbdee4b5889f", + "fp32-reference/gradients/0192.pt": "0191597afc0d3ad11d02adc401028a9967b5d3cf1984f2137161b8c1219b4b90", + "fp32-reference/gradients/0191.pt": "50a83e66da7b52aa46618325411680efdf7a008df939e82ece17aa7d4eff36a0", + "fp32-reference/gradients/0190.pt": "05e90464879175abba5316f3e852091d11748e6c1bedc53c102ab7e4b9721b16", + "fp32-reference/gradients/0189.pt": "7ac70df33e3b5a715eeebed46db148f3c12060ba3453dc97426fba2b8afd3fc0", + "fp32-reference/gradients/0188.pt": "5158f69cd770cf5db251a96b2b1634fd36dc781ed1249e59880caad14f1a107f", + "fp32-reference/gradients/0187.pt": "2bf8c64a232aa5c0259287f76214eac70cba6282dfb06db4b361cd02e1064adc", + "fp32-reference/gradients/0198.pt": "78ef8746e59eab1471ef933617d8aec8c4dea3803a132e047b1d2fc44509753b", + "fp32-reference/gradients/0197.pt": "408b3bdbb4f384d2459866122b893d3810f0b07191b790a6d1cef3041b76f372", + "fp32-reference/gradients/0184.pt": "b3e02e5674aa74e6371679edc8485ee573ee3d861b205ab8c6db88858ba23c68", + "fp32-reference/gradients/0183.pt": "dde7a8bc4595275db1c6fb26bfb62a802c9dddde2807a3bf0981d08de2febf2b", + "fp32-reference/gradients/0182.pt": "be56c87888d39287928f45af5d828976a34ec73a4a9437eef6d4a881b61b51cd", + "fp32-reference/gradients/0181.pt": "7777e4eb88e1e2312c8dbe1fdd7210d2fc9086b687300526ef0fe1f5418c5681", + "fp32-reference/gradients/0180.pt": "f4532dbd854c7b95ad09dc80841dcf11410b6e495f6710f5126cf04bb7a91434", + "fp32-reference/gradients/0179.pt": "54be67e63f3159371053dc664a051b60fb52bc42ed8a883087cd3b60cace98e3", + "fp32-reference/gradients/0178.pt": "09360ac4602a57af14b53ed431f50f7dac1d41bd65c2d396a4c991fdd2c8b6c0", + "fp32-reference/gradients/0177.pt": "330cd7a890d573598367fa63e37f56ac86c1b416eb26cbf7125137ef6b4518e0", + "fp32-reference/gradients/0176.pt": "9b71236e3685061d9a4bf032ba06ec5faef50b82ece1b7255208a37578c5e5b1", + "fp32-reference/gradients/0175.pt": "520dfc2f5c73201a764a36b927e88cc259708d7c791a84974b94f26875f9bfc4", + "fp32-reference/gradients/0186.pt": "4437758b852edc34f8cfad3abfd5a9a0fcca6092a3d6cee9bec375c189d969c0", + "fp32-reference/gradients/0185.pt": "020c5c3262063b7de91284e2c0cea160ac1ea2aa2f8cd49ddec81fe720d5713f", + "fp32-reference/gradients/0172.pt": "882862c9b0142712b1c4df5601ea8c6f169cc68bc9c4500901d84715ceb3fd19", + "fp32-reference/gradients/0171.pt": "096ceea720f53c78afff13b465a9640bfac01d1c6a3b5b361ccc07b033c299d0", + "fp32-reference/gradients/0170.pt": "d2e0784fc06524ffecdb4a3482e986fd07621ccebe1d6cbf56a25c98351d9255", + "fp32-reference/gradients/0169.pt": "f3dddfc6ae5b77c9a1d39cb902cc856f3ce8ab85f512061046c0e8f146f69806", + "fp32-reference/gradients/0168.pt": "7e6fbf962e9166dbd5c271eafdddc8071a40a446cac930b77c676afb4731fbe3", + "fp32-reference/gradients/0167.pt": "ee003a6246d66822acca80d7e4a6135cac6e6e61215738a978ac6fa23a48474b", + "fp32-reference/gradients/0166.pt": "dad720f418df89b6a9076ccfda4c4454f0b921a48589043a4932318cfab7ae6c", + "fp32-reference/gradients/0165.pt": "d3c36343db437f2ec2d69fc6182256271d153a68f98a25ce01a7d09dd82555bd", + "fp32-reference/gradients/0164.pt": "c4a8a78c37e9f71bc4bf273d8af0b428ad3cc400142dec2cefa8c1f064a4ebc1", + "fp32-reference/gradients/0163.pt": "4a99b33fb4182893bea399da2f290b312b3cb252683292aae11d2e96ebddb581", + "fp32-reference/gradients/0174.pt": "cc2ff6280cae269b3d0f817c88f8a0a715b4a50d5fb17d764c79b3382325a30d", + "fp32-reference/gradients/0173.pt": "4591792c167fccfa8630c57547f1c0f01ba2182e4edbe93f488a6018dc537289", + "fp32-reference/gradients/0160.pt": "7f1f4106d8097ee5c5531c8e38e2b7bd16c81d81aebf60a9ff056889fc11bf31", + "fp32-reference/gradients/0159.pt": "befc5346742535b4b009d4800bc69403bbd9e14b21a640b7f58c9267e23f5c14", + "fp32-reference/gradients/0158.pt": "19df487803b9e7534d28c91fcd2aae4abd47eacc59c3bf7398645d273cba5c53", + "fp32-reference/gradients/0157.pt": "c5d25f9369eb7b457e4b708dd8703adacf76a15b206004d1b34dd7345421f1a5", + "fp32-reference/gradients/0156.pt": "3166f23d79b4ed7aadd48d6532ca02e8090521d14e35159f11f60ccca493db3e", + "fp32-reference/gradients/0155.pt": "077f1bd6d24a041db7f4ae9032d897821608eff5a61db7a53f468bf2eac1eedc", + "fp32-reference/gradients/0154.pt": "4c4575ac1020a8efb8ea4925507ac2876423c053ab8823271fac7bd4eb2a10d7", + "fp32-reference/gradients/0153.pt": "0c9efab1fe6e2eb94d425d1c575803bc7225dac499b9d1210b8e7abae56913bf", + "fp32-reference/gradients/0152.pt": "f8be0b6b44bfe1b67bebeeca2bbc8f1ab8145b857136115982572b680d77eed0", + "fp32-reference/gradients/0151.pt": "1c8be24f986eefed4dbe1959588edf1422ad1756ffb2283125c5cb6548b0f766", + "fp32-reference/gradients/0162.pt": "d4ac31763f94cefac0e7ffdf139a9b1b0a1f834ba8041a2b2111ae9ae8c15afd", + "fp32-reference/gradients/0161.pt": "d2ff019717053fc34e5f0206174a770866f699ba21962a03cbb924a7ba1cbb6d", + "fp32-reference/gradients/0148.pt": "f597968fbe2d77bee43d5559d1431d5b7da81b18874f932559393406eaaf7177", + "fp32-reference/gradients/0147.pt": "fe3fc60bc9005807ff5c21cdc3ceea0973edc448ac5b44dce4604e9d18b4043c", + "fp32-reference/gradients/0146.pt": "20e4bfe11ba93df308bc986e6c4a45c36b7ae8813fd5b2e1006e7c4820af1dfe", + "fp32-reference/gradients/0145.pt": "4ab42890c9082364010b4df131170f989a2538c486d89907de1104bba2fe4ef6", + "fp32-reference/gradients/0144.pt": "cd60b2a8446b40d6a7a6f0a889ca5c421bbd73bf539cf73fc3ce2041699b63f9", + "fp32-reference/gradients/0143.pt": "4f495c888e6cf02d397f37339c50aa0ab1d098d57f2068aa100c5ecd3541e30c", + "fp32-reference/gradients/0142.pt": "577149715278de3eb1e73c6875e6840ef7d8329baf33ffc424c6006a703221f5", + "fp32-reference/gradients/0141.pt": "c90e0e8ede5f868e097ac2fb04e8215cb4b34659c3114205f2a96db9ee400e34", + "fp32-reference/gradients/0140.pt": "505f0ce6e6462c38703163cf6733dbf28990f0985ebf6e9c1dfac35aeefc069c", + "fp32-reference/gradients/0139.pt": "dcb2c21276b9bfaa597605352caa86b66019441b19785058d0d8b9f3ae6f49a6", + "fp32-reference/gradients/0150.pt": "91ecd6deb6402a6b9654562418a8dd8f231d19735576982188ce886a0173bfc2", + "fp32-reference/gradients/0149.pt": "f4b515d08ef82923144e9ea06e82ddfbae76d7e017277746f7307912af6f90f3", + "fp32-reference/gradients/0136.pt": "4db5a0dc8f6127148d007e28d163968f820508625e2d8cba0c30adb245b367df", + "fp32-reference/gradients/0135.pt": "ab09687d9b1df75194c431a43f041784e1da0dc3e0279b9001f208a0269e90a0", + "fp32-reference/gradients/0134.pt": "e82046c52150ba940b6fb1a17bef72518e96923a5e358920ad0d3d5b99bccc19", + "fp32-reference/gradients/0133.pt": "b96f4bfcf59e9660377e7f4a09c3d716c6a310bac91c3108e439c7a716b374a4", + "fp32-reference/gradients/0132.pt": "9b8a47fa6fda0b78fc3230473f5e3397c66d286508faf91121e33aa1b8eb0a0d", + "fp32-reference/gradients/0131.pt": "197e9a7ab040df42415317b624ae6a8179d2cb897a048e8ec8b1585253d987d1", + "fp32-reference/gradients/0130.pt": "a82998d9cb3f69b29f02e973d4f6b05d7fd5b71a8b3274c6e01d22d7d083283d", + "fp32-reference/gradients/0129.pt": "c603922d13f5d498e416addfd24cdc692f0c88826463857985e91f51b6242d61", + "fp32-reference/gradients/0128.pt": "eac06846dde403b093c414e46211ae5420c95e6676028cc0c19e0fb820221a2f", + "fp32-reference/gradients/0127.pt": "584dd5c4dafe5526ada168b79564d2d5a2d6ac9885475235c79b7c6db5f7dc17", + "fp32-reference/gradients/0138.pt": "31b9f1e938dd2b239e872c6d83599349f210bb52c9236e84765c04ea4ceda70a", + "fp32-reference/gradients/0137.pt": "6af6a3eaaa857bd1b34a075c77bdecd0bc232448bb762c49c95202a15b771219", + "fp32-reference/gradients/0124.pt": "b5846f779c74725d02a6cedcd676eb43734d4d20bcde782f3a3012fd25daf113", + "fp32-reference/gradients/0123.pt": "1e130cfd1e73fdf877088f9c998f8f5754ecdda5ecdf6ed42ba96faee3bf2dac", + "fp32-reference/gradients/0122.pt": "78ca0aa4d199db71f297f54c2b52a0aac683760322cfdc61476857a5661e85fd", + "fp32-reference/gradients/0121.pt": "20231c6a4f9db313eb50f34a38583e245fd0c32f52257bbce19b3483fde2569b", + "fp32-reference/gradients/0120.pt": "cf26d994ef0dd0b6077a9f79699287d032d04e196d7ac942433d78af6ee1b7b6", + "fp32-reference/gradients/0119.pt": "3ff84f127c5bf5e89ff6a817d4b3be0b4f457ca915ac666f38d1f9aeba2ecb76", + "fp32-reference/gradients/0118.pt": "d744872d12bbba24546dd2ba89c4109e921367a974937d6c8628e4adfc58065d", + "fp32-reference/gradients/0117.pt": "e2f4c477bd9e8d41f05b433baa278277b1888797e4d40ab7eb6cce8814649e29", + "fp32-reference/gradients/0116.pt": "4e0a4f1866f621f004de2b20a08eb631f572833b5232bc0acfad693cda1a67f4", + "fp32-reference/gradients/0115.pt": "44fe298e7d2d8f0a9ac3013172486e1c7db7862e5510c4e278f4e7c8f0c3c42a", + "fp32-reference/gradients/0126.pt": "bb5b91b0fc8f5505dd13c59279fb36545020ac46543afba7d1029eabff7d5a18", + "fp32-reference/gradients/0125.pt": "4e605ef9b9b5a6dd2cf3c0d097f07264e8c86fb5315173386a52907d1444a283", + "fp32-reference/gradients/0112.pt": "1d81d90143fdee2c5bbceba8b10e18f5460c91c79892f70f7bd452ebe853b4d0", + "fp32-reference/gradients/0111.pt": "1f4cf67047321d713c447c5182377ac7e70cd792e2185b2ba8f498affa30caa1", + "fp32-reference/gradients/0110.pt": "af8b9e1ecbc48d297314e2886a80be8196c3ad50b62780891f47ba833c9f0d2f", + "fp32-reference/gradients/0109.pt": "f28d836783666399312c65d077a44478f890459bdb8b7f62bba025b05a0f791c", + "fp32-reference/gradients/0108.pt": "ad1506ed4e9d06d43381af842d2e603771e39d8ea35cbffcc169f7fff6584226", + "fp32-reference/gradients/0107.pt": "ee3522f5f906c621463224feb56f7f75021586c4c78d02aee28defc866954c19", + "fp32-reference/gradients/0106.pt": "8ed67af3aba1a42c0173f5fee9da3d65fd63683568ec679bf7ff0c556b7d6da1", + "fp32-reference/gradients/0105.pt": "fe1228cda80642b2d3e541909b26e3e53e99bb3ad5254aa9b80d79c8f32af793", + "fp32-reference/gradients/0104.pt": "c60e6980543b98ebd0dd271687cbd4bebd75f5dfb2244267b4c5c7ed7bbe0aea", + "fp32-reference/gradients/0103.pt": "11dac88c8b802a30dcf3e77a08400e7fbe3eecf3564538966df03c6ede4f9c34", + "fp32-reference/gradients/0114.pt": "493169ce6544ee958d39d54f7ae2c8b68aea1b998987794c6a2fd9d6d5231535", + "fp32-reference/gradients/0113.pt": "98ad3557a69c4218a09241a3476ff8b5049d6056ee29ad0fb574d504facaf968", + "fp32-reference/gradients/0100.pt": "70f66af21bb9f5d5f868eb3c064c7b810377123da2d57c253e70f9dd2505c44d", + "fp32-reference/gradients/0099.pt": "7a4e1dd9733c55b88d5275e40325b8ffefbf9c2c70ecd19d0d8d9b0ecd3f02c4", + "fp32-reference/gradients/0098.pt": "4dbc44dc7e03ad611ca70fb20179749afc75dc296809b9ece2a762729411032b", + "fp32-reference/gradients/0097.pt": "b769dd57b82e763c726e908f275998bb3854901ae6278695ef797668c2ebc8b4", + "fp32-reference/gradients/0096.pt": "a7f51e770c17c1a83fb39ca6561268efe315421e7aa3a1da7be8ca1740857487", + "fp32-reference/gradients/0095.pt": "edb1849e30227917f2c234c0e808499493d25866b14cb523379f62285458c349", + "fp32-reference/gradients/0094.pt": "eee1eb0387492d2eeef00eb8224c9a524e55b3fd555264e3d1034291988194d2", + "fp32-reference/gradients/0093.pt": "f21d6b20698c4daf711e87917064e7ce2291d13cb775f90b4bdf0a852d4a5fdb", + "fp32-reference/gradients/0092.pt": "82000f2b130e0585f1609129b09d8ecb5cef98a3412ce47b7107c2b48bb6cfd5", + "fp32-reference/gradients/0091.pt": "75b0f1d8c206663516eb65d47f94b1fbc5d20c91b4853ea3544e02449077d70a", + "fp32-reference/gradients/0102.pt": "b625032731794390d28521187feac7c92745d2ed56789d821540b73b64f854e2", + "fp32-reference/gradients/0101.pt": "fbaa0c0ff79a307daa72da00b9f3d0d16f60e17a872621d2c0931f4813085816", + "fp32-reference/gradients/0088.pt": "6630823b8faa6cda492942db91e73524e2dc9d5760175e9858b7bca512b77690", + "fp32-reference/gradients/0087.pt": "7001edfd5dc9854b2c6c55031590d9cacb1db146b3d148b6b62f334ab2c5b81c", + "fp32-reference/gradients/0086.pt": "a5c226ce3763a07fb29e3b2656bbb69f6ee4505b232830b0ccfc14772e2a6be2", + "fp32-reference/gradients/0085.pt": "2b66738ad8814a1ba47dcd9776029b3c27b19cfdec8d018720622e764e223708", + "fp32-reference/gradients/0084.pt": "83bb85db5b6579a0bb7a7b4af6f276b8130991a32805051755c91a2e69f8e308", + "fp32-reference/gradients/0083.pt": "9722fb5d0e3560b5cd1dee15e382ae36c85077219aefe1d6d9cb86faad7a6079", + "fp32-reference/gradients/0082.pt": "40d597f6fe90c7f9c0482008f02b3d3674ff62e24fc1a09d825d8737cdaabbea", + "fp32-reference/gradients/0081.pt": "5152cc9686464f57f751f8f50c9ec3a28dab0869ba3450e33c3ab3188319587e", + "fp32-reference/gradients/0080.pt": "1c410fb3553d74495ae5621ad82060807a463ce82732701c5e7919345a4a1cc5", + "fp32-reference/gradients/0079.pt": "2a5b84b3868ae17c3fa593118ebb50f3010aa0c80366e76e8dc638f503eecc50", + "fp32-reference/gradients/0090.pt": "dfa763e291bd9fc454ab24d0afb6bb554d485cd309eb691459484e3ae5d0d71f", + "fp32-reference/gradients/0089.pt": "cdddaa734b91a56d3982037ddc06b31b3de1b127da0c18d24d418c4a22e5aba5", + "fp32-reference/gradients/0076.pt": "728eaaf85f2fe00fca2709fc058ec4251489905d0b80e0c125d7b5f9aa91cfe5", + "fp32-reference/gradients/0075.pt": "20013f7b225182f41902427568763ed16c57a37ee8c6f3fccbcb6dc2db0a0240", + "fp32-reference/gradients/0074.pt": "4bd48f9b26a284f88094ee2d1b5e25d7a8551dd0adcd2a1022e61e7033cc1e2e", + "fp32-reference/gradients/0073.pt": "2d8c58d640fff1c73e8ff384be388e516ff7c141707ca5fabb63058783d51eaa", + "fp32-reference/gradients/0072.pt": "8f465620b1bc98917d2afa8afa9df1d7b598b620b9b43ab3bab980b672ba59d5", + "fp32-reference/gradients/0071.pt": "7c3b0806b92db3c8b232d47bf69b780c79d818d7a18a34d59547cc933ed52369", + "fp32-reference/gradients/0070.pt": "61ef4c0077c8eeb8425fc99afc26fbb80bd305e508528770434aec19477a4424", + "fp32-reference/gradients/0069.pt": "f3cccc0dd4318165e0a57c5ccc6e350396f175e88fb225535041cc57b6766814", + "fp32-reference/gradients/0068.pt": "6d2ad0be502f304db305602b70e10ff016717a25cff20bee9c21cf7d6870d11a", + "fp32-reference/gradients/0067.pt": "97522b486c15336bbcad70a081771d91bd6b0c496b9e7b68580b180338f3c32d", + "fp32-reference/gradients/0078.pt": "1aba7ec1b18b858cab2a82553b1b5b22ad8a1ec8d5ac619c46fc2e652330e4b8", + "fp32-reference/gradients/0077.pt": "7a8a4c1eb4380d436f070aeb4b062c56e3e389ec20b8e0ec75c777252140e729", + "fp32-reference/gradients/0064.pt": "9d487ed2abd411d54da33a569dd4ab64f37f126f76c5327b72794c0d0b77ad2c", + "fp32-reference/gradients/0063.pt": "53f2f6cd0c8985d68e2ba59c4374ee61ee308047cc41e6c70d9e515eb6a5fc9b", + "fp32-reference/gradients/0062.pt": "06c1a9ad1b28e22cee769d513b771784c6ace0130ac97497c1c2a8070b36ba67", + "fp32-reference/gradients/0061.pt": "7071b184165b47a92d5f415a47db3ee8705b9187af41bc5b57a49a188961f1ce", + "fp32-reference/gradients/0060.pt": "4ef21ba0a0a59600d8a33c4b8bfad87463d0988343a30add0827b7715945e7c5", + "fp32-reference/gradients/0059.pt": "fe2eca1fe0316fd0917bd5abf263e19d1324513f53b11ee94f1372c089fecf94", + "fp32-reference/gradients/0058.pt": "c4332f8a32817da839859cfdba35217a3d48ee77ef99bf776a807f6d1c92479e", + "fp32-reference/gradients/0057.pt": "d788d2c6d6d2136cde78c9b2a87328f7f1f473d6b8170414c579c40cc7a44bdc", + "fp32-reference/gradients/0056.pt": "b18fe46471fa08b76922c5c8469b45e32f4cce5b667fb140b5c9ee13850e2a62", + "fp32-reference/gradients/0055.pt": "7f5f5d80d2d0413c175d5d7f1fbb4e979a4a5e75d66eddaf269e7a5a15406c18", + "fp32-reference/gradients/0066.pt": "674e9a06e7d0a44fe1c680cbadb8b694714cf8988c68e625ce82a2cf1078e3e5", + "fp32-reference/gradients/0065.pt": "2170eb1d748979d92954df19908340928198ad5637faadc399ffe4356b0789c3", + "fp32-reference/gradients/0052.pt": "65eb7545e27ff9452acde808b36b588a2f9ebdda83c9cb34fa29ef2c55a72438", + "fp32-reference/gradients/0051.pt": "1b1345f6ac31172a92562a5f0ba8b82d94ceb7f476bfc8c679e203799b1a660a", + "fp32-reference/gradients/0050.pt": "55e84dd9a515859978ae9a5a96ce178baf057da1cc14cc49864dc56897539727", + "fp32-reference/gradients/0049.pt": "10a7c34b3b1997aee50246b7fb5a51b87866153f57e54d3faa75cac826cca248", + "fp32-reference/gradients/0048.pt": "d4c84c0b1e38e077e69dfbb2eaef2dbfffb022d38bc6284514aeecfd36444b83", + "fp32-reference/gradients/0047.pt": "ec1a717e2a51887118ab8f4c62af44ce5176a16cffff35e25af6a75b955b8845", + "fp32-reference/gradients/0046.pt": "d6670b2182ffd8f77f953328dcfea550cf00ee03122ab2e804b629b516302814", + "fp32-reference/gradients/0045.pt": "a42872fcc9b0203c79329e0f731c62e2cffaddfffc3dc105ee136157d0f777d1", + "fp32-reference/gradients/0044.pt": "e68139314f366b7c18709f828e37f7c063c5d44d198808f59c1dcd042c711f00", + "fp32-reference/gradients/0043.pt": "8ac7a0cab906637797d6919c6a2e26a536babee044e564deeddb16e18cee45cc", + "fp32-reference/gradients/0054.pt": "4dc8058a898db51a5310c51838d3500feb29e00a833ce8c244b714c92691bdab", + "fp32-reference/gradients/0053.pt": "5148c5396672518eec5c4fc641fa5301231971fa70d3c15db300ac230bfa7158", + "fp32-reference/gradients/0040.pt": "50f4495cde40013a908e8110c47fc15e6de4c11254649b154573187b3131d8da", + "fp32-reference/gradients/0039.pt": "8a0b2d7962117ea6f70ca44abd2d23748c71bad596388e397a780780c318b0e6", + "fp32-reference/gradients/0038.pt": "e816f0674c99299f65154c3a43fc21ff40b0fb05b7a3a5f49462526b2fed2fab", + "fp32-reference/gradients/0037.pt": "40605f89ea7f0dd0fdae830b978d06dd216240beb9d93bdb0306fe9f7e3271b3", + "fp32-reference/gradients/0036.pt": "99ccc3a8e9739279978cf612bf1703f9ea64874fd0a83f9a39fb8a289d43361c", + "fp32-reference/gradients/0035.pt": "733a0981c95af7b1e5a4f14f05bb576f905796e9d09ba065193fcafc4b91c7aa", + "fp32-reference/gradients/0034.pt": "2ad4201fe5b25d1ee3b13a32c794f94a85e73962d614ec4f9b91a6e2cf5cdfb5", + "fp32-reference/gradients/0033.pt": "3906ebb90fd6a62739df341571e2c64375bdfeedb47536513c4e52433c9c6aa6", + "fp32-reference/gradients/0032.pt": "06fc4c4d064a07037b709e41059bac406daa7cf468f371f1ce5c9b635482e400", + "fp32-reference/gradients/0031.pt": "c0aae5701c42c167c4b65ac48f0f55086bd3afe390a2fcce405655ae2be7ed39", + "fp32-reference/gradients/0042.pt": "3984142de083c8d2c12c2ff72e672a3ef2ab16ac182c80b2e4765c444f29c1fb", + "fp32-reference/gradients/0041.pt": "f0e5722180894e058da84b2bec560b937e35e59378ca095d53a4d6745935649d", + "fp32-reference/gradients/0009.pt": "97d4a4c63a4d19a156195d96cd0489d2c94b8fc05b9ba5b07f5cecd083263924", + "fp32-reference/gradients/0008.pt": "ff74c7140174f036a208d51310f75fcd7c7aa8d10b03d0c172d58eb245b94157", + "fp32-reference/gradients/0007.pt": "a9ed52ea232dc4559e4fbcf024bad48a5aaafb79ee6e399d68b5844931c3b24e", + "fp32-reference/gradients/0006.pt": "ea30e4874686fd68cdeb741a2f0fc7dab4704a250b3b46b808da46fdaf1ca392", + "fp32-reference/gradients/0030.pt": "2559c36a9f9532b818bfa734baefdaa5bd6ab465a7456f3306153377715011b5", + "fp32-reference/gradients/0029.pt": "77d6954cfd6fb5b4dce2fdb75f75bc24b6e20bbe61f61eba6d381929450438e1", + "fp32-reference/gradients/0028.pt": "3a4bc8d1864ce3048424bb7760800882aa41aff44acfa94b405d1d053fac9ba5", + "fp32-reference/gradients/0027.pt": "0831a649430c7fccdfe0d09f9c2836f151ed1ad9a2d812440bea3dbc4877e4a4", + "fp32-reference/gradients/0026.pt": "539679efb96862b5792975b15de270da58ac14a45d6dac71c824dda9fa1d1723", + "fp32-reference/gradients/0025.pt": "e69675232c8ac02fb7874276df98c44fe1b570ff3dd217ad92c39c4a0cde15de", + "fp32-reference/gradients/0024.pt": "27a8fc9823a9d73b3f2fda8a39fa87af6d7368b9ca4f8a036a23967257f3b035", + "fp32-reference/gradients/0023.pt": "cdbee4034ed3595a82c19c2e128019a9b54624e1b904e9f8b98cecd8f1f422eb", + "fp32-reference/gradients/0022.pt": "ba98182fec10ffb5faf61fe825fe24348d47524a6afb9e6b73861164dda17e44", + "fp32-reference/gradients/0021.pt": "3ac846501b778b7ac829a350173d68e24c4655901027d9d99858c00126e5543d", + "fp32-reference/gradients/0020.pt": "6373d0f0a778e4c4e76ae7910ac98ec5cad79716d870a593f8956fca3d22f690", + "fp32-reference/gradients/0019.pt": "8d5a5f8d5d4f6721c6a4447ea9d91b57c5a16e57d0f8c333a9cfe359202cdb68", + "fp32-reference/gradients/0018.pt": "260dba92837138bfd2c58a4dd4c2e23f224a32c1b607146d88d0c9770dedbc03", + "fp32-reference/gradients/0017.pt": "f4e2784dc14cc048acc47cd0e1448c87af3a2b94656e555b68124ff6edc59c60", + "fp32-reference/gradients/0016.pt": "ee3e74c0568d935973f196a7db97c83437d9c923713f84665ad2187017dbba8f", + "fp32-reference/gradients/0015.pt": "11b82267677ec4b6ec327c5c34b578fec10e55e73d50b0d16e51bf7f9ad43422", + "fp32-reference/gradients/0014.pt": "a06fe168ab2e23c279b8aa1d3facde62a7f764423d3195e0366f71884f33b4df", + "fp32-reference/gradients/0013.pt": "bb5e1e290216b6ff0747caf7e5f9dc44f7ab386d71d7a5cbd439d13a5358c11a", + "fp32-reference/gradients/0012.pt": "630b5c9bf756301a9a4a14ee045202b463fdbc9bc8c297d90a61dabdbae536fd", + "fp32-reference/gradients/0011.pt": "60fcc18f64a16e2f0b29fefd7d2c248abe2ef3587a0b3ff6b82a8323f856d66c", + "fp32-reference/gradients/0010.pt": "319cc63c3a9858022c03b06f513070c6ace8615a5d144a58bd0be3c9cc386a9a", + "fp32-reference/gradients/0005.pt": "d454aff83dcd6e14cbf9685128797a913c73539ddb470fd286c47e7ad5e14c3f", + "fp32-reference/gradients/0004.pt": "2d20102884ed602eb8ec075554844fc5840bf9464a4bc353be5f993ad741074f", + "fp32-reference/gradients/0003.pt": "53c91edfe9637cc99b14ff5e1503e0e6a21996c73fdb94f3149099a52a384137", + "fp32-reference/gradients/0002.pt": "0419ee6d6ad48307830cc52cdeb3421584ce0e43175cdf5ef832b13f10ac02ab", + "fp32-reference/gradients/0001.pt": "a13fa20784d73f4073e29ccc80a63d6cd8d1d3e99127c939d1f3e0532d74aadc", + "fp32-reference/gradients/0000.pt": "82f4f19f10c9540f0f9e7daf898b793723970846f0e2af0f19e5cb6137ddc335", + "bf16-provider/report.json": "d49a80d6e438eb1f7d91193ad3f3ef5b5b3692b2c86a7532b86f23359ad33d79", + "bf16-provider/gradients.json": "4aa8c3e141a8649ba4384ef1d7852cef2d1ea1aabcc3a06f64acd17efed61e9f", + "bf16-provider/tensors.pt": "0946fbd60bef1996693d2f97d69f13e01165c9094fca9fe988cfc95ec64773db", + "bf16-provider/gradients/0637.pt": "e3844a24c49f10a55cb8cd4fcd97f71d814bccebad3edd81988a49502bdf3b86", + "bf16-provider/gradients/0636.pt": "4f0171f5de3f271bdbab29b9250b2e7b9415209be2c43fde44314b75be9082ff", + "bf16-provider/gradients/0635.pt": "0a93c6cebf6871f3a8872fca367c5b26a23800fc924d9ed5543df2f2a335f0bf", + "bf16-provider/gradients/0634.pt": "bd513cac9f098d35c3b83e59d5dba5400e32865958783fd2e21e151ecad331c1", + "bf16-provider/gradients/0631.pt": "51598eb239d21a6e2cd3079821ea65660515f1468b8df864c3a5e566b93c6214", + "bf16-provider/gradients/0633.pt": "33806dce327b8335d65d3b898fa90056564828a353b823b625f11593e73d5514", + "bf16-provider/gradients/0632.pt": "0e60adf736145369074cf76f434d7cf9eb221d6f49a67da2f4978e24485a92bc", + "bf16-provider/gradients/0628.pt": "4ef0834886c1f5916a6f6305766e1024b1203483797a21460c28b06bb5beba6a", + "bf16-provider/gradients/0627.pt": "913665b9659bdb29a70bc3d02eb2d445c85a5009c83169a4147330676b21e264", + "bf16-provider/gradients/0626.pt": "28c56ce68138d83acaa74e7b2a26e4f7bbcf5f76551df90ef6c01ccd1bb6f113", + "bf16-provider/gradients/0625.pt": "ff8c9ec07cec56bb2d12787fffbea708f8b734291f62730e0019dd0baf15b1fa", + "bf16-provider/gradients/0624.pt": "3fc9857d0ddac770e071e5d998908b0988d5ef7405eec8138f8f45c462082ef2", + "bf16-provider/gradients/0623.pt": "a0c593df280bde9d01eeef67decdd64a51828b36d51046c114cf9bd5bb590a1d", + "bf16-provider/gradients/0622.pt": "5cb69d6bc4671abe1a18e59031c26a746e6dd585c6bbe3cc2be15aa813cd54d8", + "bf16-provider/gradients/0621.pt": "851a411f41d4b9e9a278c97a2a92043aab6366a7a4bc4a9cbcb364979dcb18fb", + "bf16-provider/gradients/0620.pt": "9db3e1ffb4b9ae818bdf8bfe0f06ecf09cca5dcaea67e7d901d5db79e53c0b10", + "bf16-provider/gradients/0619.pt": "d1babab5f9512b93122dd571e60da19cb76da15dfadf504509d67d03a94f68d8", + "bf16-provider/gradients/0630.pt": "ca7f24a16c16afeda5d876e29ef3b7e3293ec77645bb33cc7b1b8fda6925cc64", + "bf16-provider/gradients/0629.pt": "38ec65bd40e3f255dbbf49da0cadc63becbe53a5d5c6c89f2770fca50bfbbc8f", + "bf16-provider/gradients/0616.pt": "fa914d3a6b6cbecc3a38cccbcddbff1fa34aca0cbc2775cb91055e43a76d1b65", + "bf16-provider/gradients/0615.pt": "2b3075e68c3d6f92b541d642810ad7b11f505bb5d33b67637a44968688093ac0", + "bf16-provider/gradients/0614.pt": "ea0d5c550b383dab6225d6a7790239a77c0a90b6e5fbe4a5f3df8f680bccdf9a", + "bf16-provider/gradients/0613.pt": "62d8bca5994bd2ca25a853951ff7449554c7c1389644deef17da1b72ae5d17c1", + "bf16-provider/gradients/0612.pt": "7a7229e523859a3d0582b2ca62d54dcfaea25d6c4757be3bd8895515b724261c", + "bf16-provider/gradients/0611.pt": "bdf44f033071dd300e9e55b20e5db52b7838b67e2db3561532b642b3c3765f8a", + "bf16-provider/gradients/0610.pt": "d3fdd2b2236b3bc9f013f9efa2fa36226b6aefe2c9f2c6701ed99070b7d7605a", + "bf16-provider/gradients/0609.pt": "e8d0bfbd34d5ebe63d08c85c2eb61f040425e8885dea6dbe89282861358fac32", + "bf16-provider/gradients/0608.pt": "2b2fc55e2d8f80a9a81354b259bc49cbaaa109b1e509eb75521d7ac68fe00760", + "bf16-provider/gradients/0607.pt": "7fb3904db359e9aade77abd89a6adceccc3bbe3c781bc1c21b99b0efe5ae6245", + "bf16-provider/gradients/0618.pt": "018aa39a33eb4c1b40782f074212f339aae5b9fa0b3e892838d90abdef0a1bc5", + "bf16-provider/gradients/0617.pt": "016dc3451f14408615f3ea626beb19d7572df1c340c76fd12e839a9a5f69bb6e", + "bf16-provider/gradients/0604.pt": "bcefb03c9741e6625375bbe3bf23a4e183e15b94e04ec9f17937b490b6913643", + "bf16-provider/gradients/0603.pt": "0ed996bb8bf4df3dd76ec08056f76f03aab6eedd219f13a2f160cc782222a0f6", + "bf16-provider/gradients/0602.pt": "54e929bb4912a7ea5c0df1d971c6d7f2592608e03882d34337e2d440627b1434", + "bf16-provider/gradients/0601.pt": "6a9e96d568953e43e0cfdad1bf21f0c7af94ab58ec20d9f7455e4e0897b7f39f", + "bf16-provider/gradients/0600.pt": "fbdbc4a352ba654f2c3a9fcf102e2a4b297876e9e5bdacc0ebd2cf9b4535e83d", + "bf16-provider/gradients/0599.pt": "88dd012e28d652906053a614309260f8dff041b49f5bd25acc735694016ee9a2", + "bf16-provider/gradients/0598.pt": "221d3d71a7f2ae8df65643da20c1598506dd329486e48bfcefd2ae90fa15a1d2", + "bf16-provider/gradients/0597.pt": "8ae935a939083e830c030362f857d70dc3b403f31aa60b069c8502fe4c78d92c", + "bf16-provider/gradients/0596.pt": "9888c82184122e0b7f1b24eb98d886545a11871e6f765be63c4250d58d25b6a9", + "bf16-provider/gradients/0595.pt": "890446a2799cd772077e1b3a96b9c28b44bd67a238835915352c72d69bbc9b62", + "bf16-provider/gradients/0606.pt": "d878219a9c322bdc610b840e657905ebece4bcf489d4f60e20ecb8a634cc9b27", + "bf16-provider/gradients/0605.pt": "ad169fbd40159df0b2a04665eccf1ede575c60d63c1c7837fd484f0aa38de5a3", + "bf16-provider/gradients/0592.pt": "8bbdd22256b837d712b1ee07cd7638feacebdb19cb05d123a0b14aef46f7dda0", + "bf16-provider/gradients/0591.pt": "fd3c0839a61f3c8e8b3387b0d64aa57c0cdc8314c1cab47b2af6fcf0de18e6ea", + "bf16-provider/gradients/0590.pt": "689f0c9b67e7e3fcd0fa372b130626ec3b3912940caf34b39c2f35a79ae54d1a", + "bf16-provider/gradients/0589.pt": "0a8199f7a5c05ac343dec7fff537a821690c2eab374e9bb4b72e985920637188", + "bf16-provider/gradients/0588.pt": "bd960d7750933208a3a0deaf760806e8f41d1b809fc952ca4721efb30112e196", + "bf16-provider/gradients/0587.pt": "0ecedec069314be2c5c804ac9c080aed436a26bb617e00e16143482e17abb9ca", + "bf16-provider/gradients/0586.pt": "c68f6b73c56a0364d8726b63b5014cb85ab0d44fe4a6b97d28cb881acb77c11e", + "bf16-provider/gradients/0585.pt": "8cb2b4527d51586788fce5328f2046fdd02fb3dff61b3d4db1bca938797cbb82", + "bf16-provider/gradients/0584.pt": "fd9e2170bfdaee78cbad2af809c172ed1cc011c39fe5ce5e81126b5e21e2d485", + "bf16-provider/gradients/0583.pt": "945e8f576364faffaaab6ddde5e24d26b159369d124c04510da98221f40f82af", + "bf16-provider/gradients/0594.pt": "8602f38b3b8d337dff24593156a4adecf274da9a80cfeb66c477c131547f240e", + "bf16-provider/gradients/0593.pt": "46dbd687badab03b9018c139d00adc6bb4b4e8dbb577bf1eebeec2458ab871fe", + "bf16-provider/gradients/0580.pt": "1491b47c7978d9ea8359a56292fc44c30b89a89ea59f208715c78daf9d6ac2d6", + "bf16-provider/gradients/0579.pt": "c63604ee2ffc3ac441857498a5cc8e563ae0a1a68d79d378ebdbb654726c9985", + "bf16-provider/gradients/0578.pt": "8573eed9b76c6d67e917a9568684fbf073e4967f1b47aa2f1b3ab6f330fe08cc", + "bf16-provider/gradients/0577.pt": "c5222d4c103c9860ea6fddcfb2bb08cf1dd0986784e3abeab51e0fedba6e54ec", + "bf16-provider/gradients/0576.pt": "23aef65cf4f34badb6f54f79937f785907a60361f8ff8f376d38ec4f9325ff3c", + "bf16-provider/gradients/0575.pt": "a5a3e293f026a3ccfc1f3a9f69eaf9c2b48fb36b307e4e5c63d4bcb273150402", + "bf16-provider/gradients/0574.pt": "ad03e712eae857ae84a4390c896367ade6323952d7ced6eed7e4a8d7eb40da31", + "bf16-provider/gradients/0573.pt": "57ee0afb6f94214756c9c54a490823b27abdb9005a849b0fdc5b640747a32a89", + "bf16-provider/gradients/0572.pt": "8768cade03c13d12225068b4809d7cb13d89a387055fc3ece71bd1cbd01a04d1", + "bf16-provider/gradients/0571.pt": "7be2bc4784f7e9950a73628a98feb51c287665f55c8d471a6e63f27c73df57db", + "bf16-provider/gradients/0582.pt": "54baeabca7ef89c73e80f7380f99f0de89601beafe54afde037852fb0eb6f7ab", + "bf16-provider/gradients/0581.pt": "9359cbd8a00cb77ccd6dcd98303949a3d661bdd1f5194be1c3302565d0a67ab7", + "bf16-provider/gradients/0568.pt": "be46d5c4351a616510ab17d76b5e806bb29dbcd48b00554ca91d91effa60ecb4", + "bf16-provider/gradients/0567.pt": "5770498c485f2adba56c536b11ec8e6b310b125f35e3d161c9114b50e48c08fd", + "bf16-provider/gradients/0566.pt": "984de7b29f1ec4a997c47ac9cde37713c5318e38a03974f225e435980322e5c5", + "bf16-provider/gradients/0565.pt": "0edd6067054df8df147fab513af9d769bd5ffcbcae67417b1fd14a4cf7d3af2f", + "bf16-provider/gradients/0564.pt": "a7a9edb3ab3d9d41d05df54033c91f211eb4ff312ef936f7cb55c83f679cd673", + "bf16-provider/gradients/0563.pt": "4a025606f6cfe00c2de7eaf2b80cd472133eebf6108e7d64f71588453b4339fd", + "bf16-provider/gradients/0562.pt": "9fecc5de4630537cf580421569821b71bafcfe65dcc38ad81633c08134045ed2", + "bf16-provider/gradients/0561.pt": "533684836475cff6d2bc5d95a923ebc2ba5e0a313f91b98d597f5f6a7c80c47b", + "bf16-provider/gradients/0560.pt": "27ed55d0e7e8f8ccb0d3e8d5cd88f31df7cba976ee71537919c6106cf3fdf908", + "bf16-provider/gradients/0559.pt": "53a2f2a9f5ff38cd615ea1d0b716554685ccc76b1c1c0ed574a419e3a61ab48b", + "bf16-provider/gradients/0570.pt": "17ac192b08997134047773026351be1d1db636cc7f6f8b7932db24e40031ccfb", + "bf16-provider/gradients/0569.pt": "86c813be3a9da399fc2b9ef799f5362f2113a34cb8dfa09759004d819adc70de", + "bf16-provider/gradients/0556.pt": "5d38387ad6a82948ce726466261d5080e0be20ba316cdd2a733063555bc6dbad", + "bf16-provider/gradients/0555.pt": "22cc0d5842e8410830f2c4bfa8aece9afd10cf946cb8bec33c76fdf8fac329a7", + "bf16-provider/gradients/0554.pt": "aefb9d245e82e0489b09a907b22cbce33049b23a6d2eae2288a4ce8d6d791414", + "bf16-provider/gradients/0553.pt": "a7a46058da0b6bbba09abc1eb9ba6476c8e6b79b1a1b9c44a46640337ac7a0d8", + "bf16-provider/gradients/0552.pt": "4f5217e95931f275f7f8f123e11bc785e993ae3e1666097fcea3e3d9a3d8b04e", + "bf16-provider/gradients/0551.pt": "a4316761352fa318b00679b1ac45ac02e367d12e1e3cb684d76fd9f06fccede3", + "bf16-provider/gradients/0550.pt": "5146dc1f6469a40d648a442cc26722c6a936ec0d61ab9f68d5a4582ae5c0dfc8", + "bf16-provider/gradients/0549.pt": "fac434ba5e865bd6692d8a6547e2884ef72aa51784ccfa5b2534cc30599f8a1f", + "bf16-provider/gradients/0548.pt": "5c102395455ce8168e4be4c20b397924957ed1d5eed6fb190ed706188f57c7cc", + "bf16-provider/gradients/0547.pt": "ffc8141ac3d86477b39d0866968032db4e834753e6280a38439383705ba296e2", + "bf16-provider/gradients/0558.pt": "ef70685243de893d4f8a4ca105ee66d81ff4f65b791a4ad971cca9b1069689f7", + "bf16-provider/gradients/0557.pt": "eb10aeab6ce010ee9836d3fa25e99d0847483c30be9b5215e01712306a1081cf", + "bf16-provider/gradients/0544.pt": "f5812b13181808ded4de9c07802f9b30b5181d1a1d3724177a4e90803bff9c60", + "bf16-provider/gradients/0543.pt": "ccc434434e34eccaf5744e962d07bc25e461f57f8e1e862f2da4bd5b9eca6c16", + "bf16-provider/gradients/0542.pt": "f065482e4865b0847c0a4c677177a9db80c52fbde3315460a52b82d2b01cae32", + "bf16-provider/gradients/0541.pt": "c87bd58dc250a8d960ab75703713325c8d64d944ac80cf59cb6185625bf54e7d", + "bf16-provider/gradients/0540.pt": "9d8a6ab3d9f4f746760ea4e590fb645062a13d0da36bf49ebb90ee254aa33328", + "bf16-provider/gradients/0539.pt": "1db2f753557e53896f7a3b55594405472f91dc51389e3fc77f1f74dadd9a2ffe", + "bf16-provider/gradients/0538.pt": "a37e2b4389b788567f0748fa937358f1d92decdcbbee2aa30f5be376081f227d", + "bf16-provider/gradients/0537.pt": "63e82dcd38eb8f7e9bfd752299bd89642ee4650111b083b048127d8c276dbd2f", + "bf16-provider/gradients/0536.pt": "178f0d0b39cecb680d755ebba612815be2bca21b2b31adf43ab52a3ba8e9e382", + "bf16-provider/gradients/0535.pt": "b1cbc58016470fb40996db3f6ee8e4d49c6872edf3cf1c0ec748d1d9986385dc", + "bf16-provider/gradients/0546.pt": "8aba3a6f188e499801e293ea99e6a24872a0b3da322a0f2686472fbadb97bcd0", + "bf16-provider/gradients/0545.pt": "13cbdc90c566f913377702a0579a5ba75c5517ce23a68e3120fb9d5a4ccacb3e", + "bf16-provider/gradients/0532.pt": "2c615efd2923d609bf1cf6f77d44dcb3e351cc7dbf0283fe843be0198c04c77c", + "bf16-provider/gradients/0531.pt": "6e045c469beefbc41e97702d42c908ccb777c3682e278ade0ba47bef890580c6", + "bf16-provider/gradients/0530.pt": "b703f285f960fcf2f4e0984ee64f3119ce26820ce9204db657d0c21043ec4fd0", + "bf16-provider/gradients/0529.pt": "eaef32f9eeb96ba92f46432c8114feb21dd6d1d4fc3bde96d6300d84bd8665ed", + "bf16-provider/gradients/0528.pt": "e0592d062f6a01396ccb01cc53deed7a84edafa914af63241a8d06c16cc3b60d", + "bf16-provider/gradients/0527.pt": "8e0ff55072cf1798c0963f4a00b0b9ec25e20052cc8cd79b3ed633c826e24d96", + "bf16-provider/gradients/0526.pt": "18d9af8ebcd68e06a3dc9a0cea36ca6d6efd2231d7289a3e385e0fbe88d4cb22", + "bf16-provider/gradients/0525.pt": "e441d5ec143c790b19762e2f217028ab7266b8d8b92864e61c07749499a19158", + "bf16-provider/gradients/0524.pt": "3ae5eaca1a440c8373b5b61584644db4d650c30a024d446c06d631f310ef19f1", + "bf16-provider/gradients/0523.pt": "92f7e448dd792f3cf7ae0b006ccb190edc0310fd08028cd388f194ae2f98f9a2", + "bf16-provider/gradients/0534.pt": "b0b7832cfecec6c81a0f1910d5141a8a8b5b13bef75599abdf588ebd80bc39df", + "bf16-provider/gradients/0533.pt": "5cb82cd6cd1f7709dd1402d3029f28703e7eb599dec844e8e3b422bbbfbc3c04", + "bf16-provider/gradients/0520.pt": "dc6193ebf4b7e1f9acf742f6c049b44914dee467f09b4ec704de28547f773962", + "bf16-provider/gradients/0519.pt": "fde3ad4471b0504fe2fbf6795e385394831e9eb98822ecc390922e8f3cc43b1b", + "bf16-provider/gradients/0518.pt": "306f556e6b7df14938884265422e149b638a6ed52888432e9c7a598ef35f7a97", + "bf16-provider/gradients/0517.pt": "0e6ffec908b0b43cada76470fea711bddac5c6104a3db49412e357eec175769c", + "bf16-provider/gradients/0516.pt": "835066eb37267c91e1fa65874e5700976430603a6e7a8fbe99699c402642196a", + "bf16-provider/gradients/0515.pt": "90043d814566b6a61e2ea7888f15d681101165a8833a51f503d58a2a6d90ca1b", + "bf16-provider/gradients/0514.pt": "f3f633bdeabcf26df4bb70dcde9f5e10fc50c4479f190cf09ceb4395e569e67e", + "bf16-provider/gradients/0513.pt": "211944e3666590a18abef957222a5999fe7c9895cde6f34d02ba5d1a20f0c8e9", + "bf16-provider/gradients/0512.pt": "f7d95606e6c293746562cda0a7697ccc9a73b20cb2d2062681ae73802730513d", + "bf16-provider/gradients/0511.pt": "7dcbb1587b8a88bdf4df3016df63b7db60ec82e350c01d8d1b60324b76f879f2", + "bf16-provider/gradients/0522.pt": "3c623cfae38c555d279e4ebc55dce1580241c46d9fd80e7d79b9a7f3bb3aceba", + "bf16-provider/gradients/0521.pt": "c7fef0cd677c2c2a5c0906590d3fe37f28459f1b1a9124ae8813a60260523b39", + "bf16-provider/gradients/0508.pt": "b9962e684b42a46ea85472f2677b753424c202c4fc47dfc402fb7bedb36db7c1", + "bf16-provider/gradients/0507.pt": "a4ab90249ece77aaf07ba8ada94c057904c1892d54873aa3461aed852724c6f7", + "bf16-provider/gradients/0506.pt": "cc059e7e82af46b0fb5fca3377ac0ec774411bf8fe79224d86350023b49fcf8d", + "bf16-provider/gradients/0505.pt": "35e5ad1f88d24fbda6cb315a64c991646c22a9101cd84bee0866b730e25261cf", + "bf16-provider/gradients/0504.pt": "01f2cf395c8e0a322301ab737620f025ff4c51241497bfb5b156f48b39fd0f0b", + "bf16-provider/gradients/0503.pt": "77cff9445ba8f1019aa7b9828f7b8d38c3bcdb1225fa28c06b296f822c29e64f", + "bf16-provider/gradients/0502.pt": "7130564a5ed1b02ea68e72aee594c2bff4302db8ab5e0f45fbaad2b75c19c9bf", + "bf16-provider/gradients/0501.pt": "8ff81f867922902edb9453e61fd21d7896b12c63b1b226c691bdfd7171b0f9b3", + "bf16-provider/gradients/0500.pt": "94263e7b0e65c6b51ee84fe45c61873723e89d320b08d5a4e1ca78af04e04756", + "bf16-provider/gradients/0499.pt": "663aaff4d99ec8b2f58085cc3bd5cf6f4062467b980aabdb0099d0c04bad9b07", + "bf16-provider/gradients/0510.pt": "5658660b4191fe7a0b7b9ae7c3ecc5ceef8b0576be6a83c270785dd129094a46", + "bf16-provider/gradients/0509.pt": "817adb51b2e64d38622d83fd59e360243862be28863311a019792583f3f7f5c2", + "bf16-provider/gradients/0496.pt": "1ae60be582155b48fe58a92a2a9a1e25c1c97b0d58b07be1334b70882795f0af", + "bf16-provider/gradients/0495.pt": "b7f691c0beae28cc99ecb050dae1154ca27d5a83ebf759bdf27b785e69593514", + "bf16-provider/gradients/0494.pt": "2f441517cc9479f5cf538ac4e0dcc86f37df201b80e9bdeab69b1daa7a64b615", + "bf16-provider/gradients/0493.pt": "f20240e0b40e95b910eb4a262948fe8ca21d4d2dab6f0236d50662619d94dcd0", + "bf16-provider/gradients/0492.pt": "f56d7184fa1fc74eab5e8ffe910151b7df772170d0beba7ff099ee9e4e6cd9d0", + "bf16-provider/gradients/0491.pt": "03ebd4d2a58ebdbc22ebf2e6bdcdee6533c256c7c3c1184cdb0e8a2fef1feeb6", + "bf16-provider/gradients/0490.pt": "4d796d39f7b8bb74e311c777d5ba00ba259c4556fb8fb0d54e917e6901cdcd2a", + "bf16-provider/gradients/0489.pt": "84d87af311c8b91664283ab5b2aae4299313b999d188df6798b309a9cd1c1fae", + "bf16-provider/gradients/0488.pt": "78b4e4da3bb7a0b2b0c68d203f70631143beb55af9e38e55eb73731426dc1915", + "bf16-provider/gradients/0487.pt": "6c52d0b39d7d7d8521743459744aac609ad2310464333b7635e2307bf96bbce0", + "bf16-provider/gradients/0498.pt": "ac4a9f164643ddb499717218c034112daf3419c7d5355fd160b8522843c9eb6b", + "bf16-provider/gradients/0497.pt": "ed7fd638cc5f89b9d794f9851ac983683ba48f54d1bd4a6fca55c09e93d10fbe", + "bf16-provider/gradients/0484.pt": "d677b348266831ebbc4c523e3d84ec83885bb851a00e266b0a51fa2ab5493621", + "bf16-provider/gradients/0483.pt": "22a12b0ef8de6043639f920c524e8c930d2a02b1d20451109a8b6fca0a919529", + "bf16-provider/gradients/0482.pt": "c5bced8939f0d1005c917c2d4d3f1e42eb115b4cf95f6e9ec086e9103f5d3a37", + "bf16-provider/gradients/0481.pt": "9c6cb5e1b11bf8f01a34aa7955cbd31b0b02d4eabc29250f11edb94eff21f0fc", + "bf16-provider/gradients/0480.pt": "5da555815a3fe105b85057c507eb8ff2498f3102d13028f65969daf0ffa90195", + "bf16-provider/gradients/0479.pt": "c59623111a24369ef542337593b4244453ec11121f932831d845a09fa10767eb", + "bf16-provider/gradients/0478.pt": "51b3a8f563657f5fb55bec5774d9e3b431a6fca3be71a95c789e0838ad4bd431", + "bf16-provider/gradients/0477.pt": "891d737107ca018bb54c89fb5c919af92709baf5e55281108cb6d1b62f5b577e", + "bf16-provider/gradients/0476.pt": "03c78db510bb19dae85acb75d11e5f5ad65d6e54d7aa6ede0c5f5ae1e598ff29", + "bf16-provider/gradients/0475.pt": "bf623529e4399a3b5f45d58fbaa3e994e101e35d4e7ef4bf174d11038216eb03", + "bf16-provider/gradients/0486.pt": "844838a0d4a67972838a0fa63c0d19c52924a10bcf7e64a6a16631b501325c84", + "bf16-provider/gradients/0485.pt": "d34cf4a12d9f02908d36a7f7775411de7e3a1e3e6539ee98c6cf88e6b812078d", + "bf16-provider/gradients/0472.pt": "cffca0566f6916795e9a852648b53c5a12ecc297141bb2d28dbcb1cea1f94140", + "bf16-provider/gradients/0471.pt": "47789d466380fa6c27a12fbed584875b4adfbecdb92be3ab8bf86996097d0fe8", + "bf16-provider/gradients/0470.pt": "335519ee69f6d2b2732de1f686500bb6cf1c0e1989f0b13e6f05aee7107f84dc", + "bf16-provider/gradients/0469.pt": "852f5894d5bd0e2e876b2dbcbc818348aa4b8828cbecf8992bedc826460d56ba", + "bf16-provider/gradients/0468.pt": "3b8e415537faefd522f822b0def895db00fda3fa56eb1a9e099b084ceb8a76cd", + "bf16-provider/gradients/0467.pt": "00de339693e9e42a50d2d68053f506fb74dd2edd0cbcb5e36c600c05386516f7", + "bf16-provider/gradients/0466.pt": "7c4d30defcd08033f135e5b18b8755561b05cf5d483b96459a5cbe571245a065", + "bf16-provider/gradients/0465.pt": "3b78037be58521e544113886f171d23c1373f2737a67045e19718f3d2419199e", + "bf16-provider/gradients/0464.pt": "b051ca08bef62117ec3460c5401114293d18e85e7425bcf4865998c3846105aa", + "bf16-provider/gradients/0463.pt": "05e46077a1beaca7d3424fdc71b7caa74f7819351b07c567a994374a4b01bc13", + "bf16-provider/gradients/0474.pt": "71be156216ba221a97e1423a03217a7a7167c88f150f603efd2d8c3a11f38503", + "bf16-provider/gradients/0473.pt": "bf55e7d09d5c5b856d6feca13155c7c4e82ecca6a4520856365e706ab00f0d5b", + "bf16-provider/gradients/0460.pt": "7fabc4c35ccc0907d2ecd5b49a57c1efb64cfef04c8828dc96de9c22529cbde9", + "bf16-provider/gradients/0459.pt": "459398bfa18a8c7021ea269175b9d639fdf138ee51dd84cc9d1101f52bbb6298", + "bf16-provider/gradients/0458.pt": "c1c19be6ddf5f7d84700f5bb13f9154b39cfd949f51a730d946bb66a8677c2bd", + "bf16-provider/gradients/0457.pt": "30dd1380984026f13f99f8d9a3df21ceeba1b3dc526684db5fd5acee94c64715", + "bf16-provider/gradients/0456.pt": "830b6ce8c09e1e0d1d51acd72d2ad4ad8595d23072ceae4efc3970bed5570cb3", + "bf16-provider/gradients/0455.pt": "d243fbca25f5f6aa29ad48552f38fec86f8a01c8534eae18f5fefa3187210d99", + "bf16-provider/gradients/0454.pt": "8d6507b3747f7224ea8d168b034911ba1a7f1837a52fc070faa696186b8c9c02", + "bf16-provider/gradients/0453.pt": "4f1524673d3334f5b0884719e71620c95df2719d8d68e990bb8fbe89c99827fe", + "bf16-provider/gradients/0452.pt": "81a00bb1487bdc2dd595e0e3f7ee75e7d9148eb943f566d7bf3777c97c5e915b", + "bf16-provider/gradients/0451.pt": "d31dd660f1c1e8ae3fe121978d12fe12d006a3fe4ea06b43b0ca948c43c56261", + "bf16-provider/gradients/0462.pt": "99cee9ccb61645d1dbdfb9abdc91a677fda2c0c9d67c654f4ee265e6d3df614d", + "bf16-provider/gradients/0461.pt": "0c0bcf12ef370674906bdf31844ee8cbc460516270da6b924faa8eb22354a7b4", + "bf16-provider/gradients/0448.pt": "d669605d10e1954f66a509a33d4cbe8f01820a0aae89accb562a946699bd0a43", + "bf16-provider/gradients/0447.pt": "f0991a36dc1b3f2de055077396835da8cd918b592291f189b11b21968fedfe6f", + "bf16-provider/gradients/0446.pt": "8fee8797705a82f026c1d4369af14836a07113a3000396c9afdec448c149e539", + "bf16-provider/gradients/0445.pt": "fead9a7e8c04fc2624058097546e874eb151025d806206c7b8c0b78583111fa0", + "bf16-provider/gradients/0444.pt": "f94353bd4cdfac7d0d14be7193cf401a5bf920d0790deaa7841459878614bfff", + "bf16-provider/gradients/0443.pt": "a9e90619d6e5521aeb2334f3f6be6c803274580eec082d91c209219b5d5d51b3", + "bf16-provider/gradients/0442.pt": "8e21e37aaf689d5ac7a935c6a00b271ae54ac323d891970c8f0c7e2f3503b6b8", + "bf16-provider/gradients/0441.pt": "5a0ab2c9b9163f3a568e2bfdffa89227dcd6986f41f295cc03a570f24e748d41", + "bf16-provider/gradients/0440.pt": "36a3c16a220274d2bc149cf95b75770ffb0660b21f893dae6c75e20d2ab5c9d5", + "bf16-provider/gradients/0439.pt": "b171415dc97e7661edf1ad12a2674373d5dd9f12f1a7ae45c6fb2c7143464d1e", + "bf16-provider/gradients/0450.pt": "8e6f07c28eae7a00bc37f8d6f9d92a8c0b06d6de8231b0361bf4bb7b3d9bdee5", + "bf16-provider/gradients/0449.pt": "7b2a0f5393f7c4b6342495cae101f5d0fa69194bf5c190bed72bf28fc75003e6", + "bf16-provider/gradients/0436.pt": "834ec333f78bd28821895823aca1b30ed8fe08da07b718b05dce0ac5e6dc8bf0", + "bf16-provider/gradients/0435.pt": "34cb3acae0d595bca5bdafa674db93e07af5f2b622d1d522c51891886cc6baca", + "bf16-provider/gradients/0434.pt": "73068876b071591f76196e1c2c8ec8056a40b8c76e47a1791cfbb6a164ddecf8", + "bf16-provider/gradients/0433.pt": "dac1afd8bcd9a2aff4726ab8ae02fc31a549f990e957584b2f1f3f943f355c67", + "bf16-provider/gradients/0432.pt": "eda8de51540f5547b3bf15019803b443f0c53a34bfae2c69c7895bf0b603fb72", + "bf16-provider/gradients/0431.pt": "9595b4baac45a409602faa397f33dcc17748636ce674617fb5fc4e7334c1bce9", + "bf16-provider/gradients/0430.pt": "940334deb170ddfed4b1a31e8c229628aa009f745d708ab6104bea92c8fc1e0f", + "bf16-provider/gradients/0429.pt": "e4d4a1c4cc701cd55c69c00c508ac3bf76925cf2b31d66d4e07a5ee929485a97", + "bf16-provider/gradients/0428.pt": "14d3e7302550ae94a58643aed0a2511a3bfa63daecb3cede0d62eaa682a8377c", + "bf16-provider/gradients/0427.pt": "3d104208990f1ea302885f316d2b08c2774a3d5e028ddae08c94eda048d88f32", + "bf16-provider/gradients/0438.pt": "1baed369a2ce360a0de3eb5863270a5f257117221db024f0f98932a86fc26405", + "bf16-provider/gradients/0437.pt": "d6dc7e416eb2fdb75eaac86bd4f1c8b78efd500cfafae9a8f065429d56ad530b", + "bf16-provider/gradients/0424.pt": "1b898b8bc3b14d42334488f9df4165690876762702441029edf32b6e753d9035", + "bf16-provider/gradients/0423.pt": "f87f47ba3ad6649007e9020f8706aa331b98513294ea61f62b3420f382d7aff1", + "bf16-provider/gradients/0422.pt": "10288c7ee2c446c385058a64c8c095fdf56bda47f389999a513c382a63748e12", + "bf16-provider/gradients/0421.pt": "50b3c863396334d9f8ea937c0c716052a4edefaa2bfd31f70e314d61b5d9e04a", + "bf16-provider/gradients/0420.pt": "e00669d8f3e71b772e5a0f7fa8911c7cda6e59195f36df868659c5e2bf1bfd8f", + "bf16-provider/gradients/0419.pt": "2c2f4bead308123ce434bdb8ba3ad24a06e4b45770b50066885d02484b470006", + "bf16-provider/gradients/0418.pt": "dc69f98edaa1b6b26ea2c0df259f60f097129f3cba6ff943ea8a8277ff3d05b7", + "bf16-provider/gradients/0417.pt": "e723311caaf808dc9c0459b41f94618b7c2e2bcf806cd1691a14875eb239f41b", + "bf16-provider/gradients/0416.pt": "68d79df4f2a59dbb98ea8758ea1fe1a5b9a649909531df08e3ef5b19bee21790", + "bf16-provider/gradients/0415.pt": "c657bfb1ec98bd4b6e56e8a8e6dae63e716cf4b91688bde1f39066ca70222c64", + "bf16-provider/gradients/0426.pt": "7cf0e8f7f7b12fbc639cef97ace25ca674d93f1d88d4bf2597bf64653679f483", + "bf16-provider/gradients/0425.pt": "0500fb0d766b281637ba24e2050de2e3fc10601d564c674a8640dfe3c6f373fc", + "bf16-provider/gradients/0412.pt": "4bfc5c032d26964efa3fd2a84e1d8968b769d40eec8bbe9718d392855c626f71", + "bf16-provider/gradients/0411.pt": "2f4ae0f3c0e4096397e7786b3ea0259099187206e5c1b7b1b6b3c5fa86b25d69", + "bf16-provider/gradients/0410.pt": "42f419cbd8544eda7be7fa2f667097d8bd88e006538e9d67d1bce3961b11225a", + "bf16-provider/gradients/0409.pt": "eef39b3b6b1d834cf0e506ca09c0de4f3843fc891e2995ebf09174704b4371dd", + "bf16-provider/gradients/0408.pt": "b3adae34f2641f63b9d9429a286631d22ce93a8b119572b47594f3f04439e6ab", + "bf16-provider/gradients/0407.pt": "ce7cf574a586123e9cc0f8d0117e409b1411663269383a346af027ce19db29cf", + "bf16-provider/gradients/0406.pt": "a499999820e53a0f239336712c66ee76ead36582c149c9d3d6f905f65f6653ec", + "bf16-provider/gradients/0405.pt": "3111bb92fd28d64f382b4aa0767397cd2514d5602b0a2de635f47009b7133861", + "bf16-provider/gradients/0404.pt": "e751f5cc45343da810a307d273962c735c8edbe8d28ec19b1957e80c9c392c96", + "bf16-provider/gradients/0403.pt": "6336ce3fbc6f62af1a59f40650229c6a72791e15c6a78c730a5445029e9edc23", + "bf16-provider/gradients/0414.pt": "983932ac4f61e0a7c2a592a58ddeea75dea31f7aba0ced0b922e213064e72ed3", + "bf16-provider/gradients/0413.pt": "3c3e33a6f32818988c23262eed7ef1a226d69f5a5f68ef5c43582bae43f99420", + "bf16-provider/gradients/0400.pt": "6e20fa0e1e3756137131e6744ffd2d8ee67d285accc58638ed14b56f641daa2b", + "bf16-provider/gradients/0399.pt": "199826583b91b8a1e4ab99510e120815e327d477876a48cd06436f80568061ab", + "bf16-provider/gradients/0398.pt": "ad1c0c75231bb2ccf4dddbf7252ba64c7f85011718730af27f15ab91866f42ae", + "bf16-provider/gradients/0397.pt": "6126e25b6ba7c4f39aa4169f82efb2edd8866a0f7ebab22901fe4806280bc99d", + "bf16-provider/gradients/0396.pt": "b9401ae1183b64c57979e10df3fd738439960d17687bd10fcfbe23fd74c72d6d", + "bf16-provider/gradients/0395.pt": "00cb1f8b781f7df14ccbeaccdfe13d0f184258680d09e5b121ed4f62d4b50b43", + "bf16-provider/gradients/0394.pt": "7c41ba4078643c2c007c57a88417a1d82ba8159083c979b518f4ff440c48ce05", + "bf16-provider/gradients/0393.pt": "dee0364f82594a2d9d3880ea80a46c85322f38155f6b45b6fa46b4a8728da63c", + "bf16-provider/gradients/0392.pt": "7d2210cf08a9090efe94c510c1ea16cb1f547511f5dad337d55cbdfa0c72d10d", + "bf16-provider/gradients/0391.pt": "a8ad0998896ee8ffe5972b3f214746319e4e52733984e8e1df0424bf5ee799a4", + "bf16-provider/gradients/0402.pt": "3f8dbe7fadc5ea145906611363bf22173a47368175f76ab266384018ea60017a", + "bf16-provider/gradients/0401.pt": "de62139743fb85546f4a2c5a71e9c1ecdeb53c48a49b7d97bc257748623203e4", + "bf16-provider/gradients/0388.pt": "c01a779947fb30a631e357fd880decdd520c4caa90f31473d4b58b546b596eef", + "bf16-provider/gradients/0387.pt": "1af9a5d52305c25c7ef28239e6b73cc6b69aed344f9c2c677d8580c189405401", + "bf16-provider/gradients/0386.pt": "83f94e71e1ac221cb1e56132ade7704b2856cf3d126227a8b217c0300e5144ca", + "bf16-provider/gradients/0385.pt": "93daf48ecb6e04cf0dd96764571546a3d8b0cbc08eebc39978a87dbbc360eb3f", + "bf16-provider/gradients/0384.pt": "4c889512e4ba10b671cfd286a629b3d19f965584a50b7f5916593109eb05dd2c", + "bf16-provider/gradients/0383.pt": "83a4d89b8d6fbdd88a4f3224565faf09180735e7da060fd6a44395b78b4131f4", + "bf16-provider/gradients/0382.pt": "736fdc79e54f720023033cf411b5a56eed176d259cda852d6fac8e4fe0070962", + "bf16-provider/gradients/0381.pt": "e1650693f27a890330a91b847ff669f2845d962df460ff9f18e7706b0b8f0ced", + "bf16-provider/gradients/0380.pt": "6af3e26fcaf59b4ca8d600dffdded8191cfa7d31773b4bad11dd689d8fe1891c", + "bf16-provider/gradients/0379.pt": "7a9026ef060cb6d7ef16718bac6f3aab381f38bb4afb1bed63cf75213cd45ae5", + "bf16-provider/gradients/0390.pt": "9d01f008020aa5f7bab6bd45d2321186ae97ecb3f747f21c38d56d13cd08563a", + "bf16-provider/gradients/0389.pt": "fed763e3838a95b12a4a4683b1a863eaf4f8a17e1ca7ff596b8f1f1824e80fb8", + "bf16-provider/gradients/0376.pt": "80d991324b2a369e92703b4d330cf9f95a2625c05a6db6e2a1b33cad19ef6cc8", + "bf16-provider/gradients/0375.pt": "9a2548d35c9168c4b76fa30fae9a755e8b2401130261a645fab732f9ba58b631", + "bf16-provider/gradients/0374.pt": "366e7a501ca2023befd3a1f1a6599a72f59f49f9af9cd52895a9e1d245b67d47", + "bf16-provider/gradients/0373.pt": "2f7113d23b7e40fd5f5c3940bb0b54b8caa829e8b278ea720ba99ccbd4533640", + "bf16-provider/gradients/0372.pt": "fb22006855623f04e5823a46ce0ba4e0ac7ae939bfe5e81efbb25a812ec196c1", + "bf16-provider/gradients/0371.pt": "c04413582df24961b0f65b52f7a0480f922105e877c465f574dd9ec4bdb66df8", + "bf16-provider/gradients/0370.pt": "df53702eb87cd51d31af9d714012b7fdfc031b0b6f334903c9d02507324e5db5", + "bf16-provider/gradients/0369.pt": "5cd145903aeaddfd0523565eabc48ed2ec3a9f9e2de099ca02d3085d76387672", + "bf16-provider/gradients/0368.pt": "442b10d48479cd3bb6637e77028c7f37275c0498e0b12c657461af6f9e479dfd", + "bf16-provider/gradients/0367.pt": "2820a3ddd747ee37c3f92ef04921cfb014f1bbcafbd47304f48aa4cd64557401", + "bf16-provider/gradients/0378.pt": "7bc6ac4c284eef24a85062aeab27d2db52d9f52294f8249473186a2b9619b406", + "bf16-provider/gradients/0377.pt": "0fb26d3d09283c9c93ee5ff99e8d015058afa2759f59678297a05465d586d941", + "bf16-provider/gradients/0364.pt": "7086c5889af19c994a10c9fee802da9a0d3505be3ca936a39bc819379635b3c1", + "bf16-provider/gradients/0363.pt": "b10767c3b65bd0440f5374265e67bdf2998868413214e1d5348b7c7a6496af7b", + "bf16-provider/gradients/0362.pt": "7540995f2696f4fd1fd660595c81d78d2ca64ae36553ac4eb0cac167c3f3fd36", + "bf16-provider/gradients/0361.pt": "cfd05b67644a4c309d6a72c22f399f36f977156a150d66ffa6e2586d9a49cbd9", + "bf16-provider/gradients/0360.pt": "c587df6077f08da4b7e0619960cca589f5248db0c32bb82cdb781602cebc59c5", + "bf16-provider/gradients/0359.pt": "92b58b50957eaa4c9bde642ae32635a31f1fab21f5979b1a4b98ba09edbde855", + "bf16-provider/gradients/0358.pt": "10c4b56975dd96a39c64a0fc1296f89e0c2ba0b63d363bdb65d333eed53648c7", + "bf16-provider/gradients/0357.pt": "16f230ee1e7ea6f01259b29c1747d9126043a8bf5939d8a44dacd673cd2ef072", + "bf16-provider/gradients/0356.pt": "e8ec4e58e9d1df18453191536147588efc7a9d7734c81cb748997e6234ee26a5", + "bf16-provider/gradients/0355.pt": "215feb49ffa57c65ead59544bf40b64cddbc2b408f5650684218810e9f9428ab", + "bf16-provider/gradients/0366.pt": "8dc4a8cb1ef2ef03b7261b9057a2d1b3415925b8e102dd713d7952324bac9906", + "bf16-provider/gradients/0365.pt": "048dfe529664b2f752583875838bf01c6382bd1356ba5c6ad4181b533aba765c", + "bf16-provider/gradients/0352.pt": "73b885e92b77e1a87a2ef06a0b9ba5b6106294192b8bfa925fd1a7ceb0b0bd46", + "bf16-provider/gradients/0351.pt": "8ada9649ccd6e091746424be7c0e9f3c4d7f4b5de9d0bf1624d6c94616638c61", + "bf16-provider/gradients/0350.pt": "b897d88524957c46b4b5c829385b1ded1c97a5c60995b1de494e047fa542c019", + "bf16-provider/gradients/0349.pt": "d4f9dda2f1686a3daf3cb6303d07bf76b14f2d162a6ec68edd2826c26bfffab3", + "bf16-provider/gradients/0348.pt": "114a6555be1fea1a2411da97bd39ee547da87ff8401f7c884b38111d285769af", + "bf16-provider/gradients/0347.pt": "d99e95c384ac1e1a79269ab365e7a641fb51cec7e53f48f19023b591162ed321", + "bf16-provider/gradients/0346.pt": "c8196dfd759a551f9073919e1184d9ae2c1e7df7f679b53ff72e5fb323828c90", + "bf16-provider/gradients/0345.pt": "061fbb5691890cb52bb3d4c4d523c30f19441631617f27f52a2156ab1e713d41", + "bf16-provider/gradients/0344.pt": "bc7afb0a9f9b5bf678de90d4e32760c9dd2133b2ff7665542f68bd10c73ca3ac", + "bf16-provider/gradients/0343.pt": "c89979544fadc5881e4db713ef85ba71fa97f3875b750ab6cfa1f61d6c8cbda5", + "bf16-provider/gradients/0354.pt": "9d28f231db5552c04e92baeeed5c1039326d284d896354b39151fb13059d9acd", + "bf16-provider/gradients/0353.pt": "2cdfa89c6fa01d30190c3561c3dcfa1a99cf345d252352b3d36a1d85d3b40f7c", + "bf16-provider/gradients/0340.pt": "16e736c422313236fa0687931f559902adbbc030a71a6a9265f0e0d48a744ca7", + "bf16-provider/gradients/0339.pt": "c44c4f8a6e6b690c9a676688a102a57a557d6b0fadaa2e5ab3914c37831d514e", + "bf16-provider/gradients/0338.pt": "432eb47149d2699aa0938b44e8311cf02bd578b2e4e3d9026e1b1ccecd078af7", + "bf16-provider/gradients/0337.pt": "9d2996d3e10e01646574bb35f9e4ea39d32041dda99842ba88443f6b1e90a187", + "bf16-provider/gradients/0336.pt": "5ac82f7000cebb3952880b1624f7c822661dd1e80e36272205c130eaf2c7de9a", + "bf16-provider/gradients/0335.pt": "5f96493dd608cb99d6d6b55638f872230e4074eefcaeb2ad19e1f74ccd52f775", + "bf16-provider/gradients/0334.pt": "ab438df3ff6b41104fbb062d14be1d32db9acd9bcce90586eff6d28af0952db7", + "bf16-provider/gradients/0333.pt": "c0b604625ceea9e2be57f6712581231b1c8a7fb19a7466a57203e1ef9b55fc95", + "bf16-provider/gradients/0332.pt": "89fa2ab19f9530652a5589e5dec63dbc5658d6292728abfceb6e971a7361cd68", + "bf16-provider/gradients/0331.pt": "a28d779aa3c7ea69749cfe92f72150ab4253b3620a25fd26157bc50eba51249f", + "bf16-provider/gradients/0342.pt": "dbe29d903f6f33fb03d6e57fc57066a77783484ac332a7001aec997aac07f2d3", + "bf16-provider/gradients/0341.pt": "15db439a8cc038776a700c98ff1af3dd5535d5b9a81f1f14d33102c3697266d3", + "bf16-provider/gradients/0328.pt": "9c2096499920a1f2ab43be7737404f8c7ca77df77ee8da8c4d175ffe79b1b8e7", + "bf16-provider/gradients/0327.pt": "6c209b6e47c1dfec4b212471dcc803ae83376d08853117a860d4d4800481ee4d", + "bf16-provider/gradients/0326.pt": "9c65c872e4526ca10382a1a58cf2f917aaf21e12ab9e23ac78a6f711b1c83c79", + "bf16-provider/gradients/0325.pt": "7f2fca93366b019fa9cf2399d3614722bf4026d2c62be5998cfa4f4ae8436e71", + "bf16-provider/gradients/0324.pt": "a383beb4c6dc0a9b81668c4ebd111b5a6e527c576253ccbbcc574967277fb386", + "bf16-provider/gradients/0323.pt": "bd611186afd49b0f5ba6c530c7c2307f30bd03e0a1fbfeeb4d15ad01818a6c27", + "bf16-provider/gradients/0322.pt": "9f7a464f463566fc87aa8800106925523a6e651f42595abc6969936574bb9845", + "bf16-provider/gradients/0321.pt": "07f5921da45d7d5a3d10203bb49465d56b92106aa9799baa0bf5f30947a287cf", + "bf16-provider/gradients/0320.pt": "dab1df649ff666c8e37eee09779a8b39eefd171a8f99da0102af1590bf4cb0d8", + "bf16-provider/gradients/0319.pt": "a706aca7f87492191c7c7434ee5d185764eb1f59d94ae7c4299d316aaeef51d3", + "bf16-provider/gradients/0330.pt": "134317c0c9e4b21682d825d0c35c9436e269c9517e42dff6440f73aba83c7867", + "bf16-provider/gradients/0329.pt": "46b37f67a291d05e35ededf73764f5763d68dab2ffb115ffe5d85afb8b8f7fe4", + "bf16-provider/gradients/0316.pt": "9390fd6efd83be43ee17922a778e5be65806fc0a88fd127c250d3139bf9a4eb6", + "bf16-provider/gradients/0315.pt": "d56e52357f154ae94ce733262f7284c38bbcb2cb4ac06079bc8c7008a485055a", + "bf16-provider/gradients/0314.pt": "bf43b0a5eab07b921a7c7d587f90fae7abca00158cce1e089ed0acc182ccec13", + "bf16-provider/gradients/0313.pt": "6136f702dce8910065120cd350ebf5cf3cef47ed4fe158f87d8421cfb3da669d", + "bf16-provider/gradients/0312.pt": "d9a8b18e267e976a7346bf311cfd79fa4625f091d03c104b74f0724a011231d1", + "bf16-provider/gradients/0311.pt": "5b295970ff06425a7b72ba8021c74ce779e14c49b4d3fe2be355e05d5655ff6b", + "bf16-provider/gradients/0310.pt": "0d653d6145f8ced6c4fa61a187bb1773503d188b51729343434c153d0c954877", + "bf16-provider/gradients/0309.pt": "e010c4c69dadd88e4b434a3c39f4fd87b84ae264659ce400b6343e935aadf4a0", + "bf16-provider/gradients/0308.pt": "da3c2b653c9a47bd73805db2bd693b463e8ef20beb16273b11484df3c76b19b4", + "bf16-provider/gradients/0307.pt": "5023bfd386b924ee3b89c869ef73022bc9a77fc7bdbe3d1c90a636874ddf071f", + "bf16-provider/gradients/0318.pt": "cfccb4ffc0d753db41273a3fc30c11773a5b36a3c9ad58e575a2f9f3e2d6328c", + "bf16-provider/gradients/0317.pt": "47aed8bf8daa5dcd663f8e5bdfa0ccf9942fdc886baf5b374b383adb86500462", + "bf16-provider/gradients/0304.pt": "97739c64dced790b37ae7c068e9c6d53f3e20d30e6e54c6a179a1b9f84453d71", + "bf16-provider/gradients/0303.pt": "3304aa31f19d4af1973f91e58484b35ebeed42f1ccbfa5979be754e9f2e11149", + "bf16-provider/gradients/0302.pt": "ca8c7d81f786b9ce66cb5889306f513d6e03024376bdb3d44acdacae6dcbbfeb", + "bf16-provider/gradients/0301.pt": "787675e431dc8cbea94c09decc3149fca55cb1e70293365770e2214615b47451", + "bf16-provider/gradients/0300.pt": "668ef60e15794443875978942371c0e340578c7991eceb3ee832830ec0c82130", + "bf16-provider/gradients/0299.pt": "6eef743f79d8f93aa8e010987b292b0f4a72c7d4a5203cb0c74f37da09c88ef4", + "bf16-provider/gradients/0298.pt": "233c25284295be84a203fe10833fa9801ab197a468a9b7bc1153e7aa4bc3d8ab", + "bf16-provider/gradients/0297.pt": "7f2bb4b77e5380e5caa61cb8781b66a7b68e20679497e7cc334dbd4e48af6ce5", + "bf16-provider/gradients/0296.pt": "49b573ab96c800002052bf1d470d39ab94810c9ee7a7135d7b03e9f1ce62f4bf", + "bf16-provider/gradients/0295.pt": "f88f6163f48a80473a21d129f619ab42a4d88334f1ecfa74005d1627e4057333", + "bf16-provider/gradients/0306.pt": "796f1ecc30c0895b2645a17d2820d9c0c3c437f0c1629f086869d764236758d3", + "bf16-provider/gradients/0305.pt": "fb7e07b6d951996018c55a750b4de1a7587e50a9b1e8b1ace03d381423b5166d", + "bf16-provider/gradients/0292.pt": "d774e15a385a6e2affe56fd4bf7565b8a62f90a7984dc1121cace9520140d5c2", + "bf16-provider/gradients/0291.pt": "0f295a53ea57aa75fa7607fb78bec5b95247494f29651269b54385f9f3d63f31", + "bf16-provider/gradients/0290.pt": "5e4ae29bfeb14e558df98e786d1a10bc280afa49d5f2ccdb89ded768f3f2bec8", + "bf16-provider/gradients/0289.pt": "0e50704ebdab957ad87686cf78306b329410a668bd177568538388e5cf8c9e38", + "bf16-provider/gradients/0288.pt": "3f5034461985f9f4a3d3c03d67dca680b3693b7bedc4e98abcb661baf72c9762", + "bf16-provider/gradients/0287.pt": "86c0f75adb9d77407f34f834b3123a66113a0e4433efad21dc6d644919897d02", + "bf16-provider/gradients/0286.pt": "35913b01eb81cf78be88754a54e9d3001561f33363145b6df93492f39bef5486", + "bf16-provider/gradients/0285.pt": "afc40a230f35770796db2500ae780c5f3be23bb0ae7f25f77fc0bc64cc807665", + "bf16-provider/gradients/0284.pt": "26e2bec40c56e38cc2ff769a5668677cfc914242de5d996fb799d0bf525dbbb7", + "bf16-provider/gradients/0283.pt": "4417139e2f7fa76dbc6bd9234fab69923b8ec6dfb2ab8006e2d4e97d5c0225eb", + "bf16-provider/gradients/0294.pt": "808a1412bd7e2daef88ea1c267c0135b823be7e1038336a884721c56cabc4bb6", + "bf16-provider/gradients/0293.pt": "76ba8b80c0d1b68afd1d6b058511a7b61d2672d1fb1f265d53d1d750f43d64d4", + "bf16-provider/gradients/0280.pt": "dc6b8ea09a1e08f2ced13ff7112a2ed17c3c137b50a6cf8dbbd962eb0650c421", + "bf16-provider/gradients/0279.pt": "3f29a1ebd7a98e252bad3ff33801d197d9e955445217e878cd5dbe032336cb90", + "bf16-provider/gradients/0278.pt": "9f0d2d5267f3011ac21138c731cbc0099ba70eccb81d11a5c8e2c45e73021c4c", + "bf16-provider/gradients/0277.pt": "8d9315411992cf348e28707747169acf8084122ef7e78839a51aa615ab33dcde", + "bf16-provider/gradients/0276.pt": "2bc86872c0427eef27c8d41828ed65185cd20d43c97faa02faa276b435112a6d", + "bf16-provider/gradients/0275.pt": "d18373120bfc7561ba43364cc7ef2952c1412a934b11bf19aaa43446b6319b6e", + "bf16-provider/gradients/0274.pt": "d03638a9a7cdc8097cdb71710c5f6fde6a4cb0c72e9f904ed29bcc3be379b993", + "bf16-provider/gradients/0273.pt": "d2a9a7ea8eaffffe8c698ca8eec6c4e3ca404eb00ec77d808fc3834c62c3b331", + "bf16-provider/gradients/0272.pt": "6db399f8af1c9925068b79ddd174c9a3a8707db1f6d33b6c89e53ff2d43d673f", + "bf16-provider/gradients/0271.pt": "a186a972fad6db0a157ccb39a88f0abc5926f4129b582090077aafd84f22d74f", + "bf16-provider/gradients/0282.pt": "9ee552757a1496ec847ba7c77445854cf51d5b857f5a869e7c779883e0154605", + "bf16-provider/gradients/0281.pt": "0d0252a452b5a900e0e1cb3211fc88b44f2b9f6aa38b183d55e19a37f3a2b9ad", + "bf16-provider/gradients/0268.pt": "eda4537d480a47f9b8c854ad5da8c51c14e11e9ef79445db161bb87ead56034c", + "bf16-provider/gradients/0267.pt": "a543f6e2eef7495a8ed588339eabc7df314634ab816e08e8b1e430a9dcba66e5", + "bf16-provider/gradients/0266.pt": "e759cff02bd6196fe5de0c1028b9ad2668d467354bdb16cd416708cdb15ae185", + "bf16-provider/gradients/0265.pt": "042bb32b5a2a98e22fb7f4e7c9b27867c70d2eedb1198a637dc703c5cb3703f0", + "bf16-provider/gradients/0264.pt": "35b9ad420c48646419a247e113f0319c3b5b28dd85ce64c009b83943e254db70", + "bf16-provider/gradients/0263.pt": "c7964735b2b3bc3b030a62306d2f99c3ffa91487c92a8fee0dcc4832a7173faa", + "bf16-provider/gradients/0262.pt": "f9ff815ad5e7325b80892fb584703cd6e8735e17efce777cff2e1144fc65bbdf", + "bf16-provider/gradients/0261.pt": "4ad2f2b3328ddfa1828d3ebd55a5f3605231fe5b919e0c5d17594aa0151e5210", + "bf16-provider/gradients/0260.pt": "53f7d3922bbea80523b20f065de6c4d58161a69404bda4010828e119bc14c5a3", + "bf16-provider/gradients/0259.pt": "deb98987e5cdc89e76beb9fc92c52409bf8beb9e3fd66f29b5c557ff944a898c", + "bf16-provider/gradients/0270.pt": "331af297b425fca0d9f996e2afe83661f98533818922f3bc7ef35ecd90c7458f", + "bf16-provider/gradients/0269.pt": "21e1849f0dc4290caa005a924f36dee7aa4a5ca43718f7b2aafa8fd1652f976e", + "bf16-provider/gradients/0256.pt": "d9a3f3ad49d2ec14f3e1c98939daac81d748bd5ce3223209fd0f79c50e4ccb40", + "bf16-provider/gradients/0255.pt": "8827e510654d2727d8d9a24316e2332c30c43d31db69dd0b988df016fcfc5824", + "bf16-provider/gradients/0254.pt": "0e3bfcea4d4857a0affab4319659a6de26be790c102e76a1c142de78a8e877b3", + "bf16-provider/gradients/0253.pt": "e182a9865b8f41da2e50b5c625a502aa4b69afd3a7f3e027f90bc7ac22b6b95d", + "bf16-provider/gradients/0252.pt": "64044c1e4986c3f95008f294d1dd807b27e1aece6ed48745422c13cbd3de32fc", + "bf16-provider/gradients/0251.pt": "995d73480e764bfeca642a3e98aad010e02aeb1f382bb6a43b7968f23becd61d", + "bf16-provider/gradients/0250.pt": "e4d24a090de0e3b44a7dc0a0e7b37eb776cf600ad4f7fbf47f0e56a31ca10e92", + "bf16-provider/gradients/0249.pt": "2b47350687d676b10fe5b477f8448b396276b40618645b866c5292ea907d28d9", + "bf16-provider/gradients/0248.pt": "3eee607af2409ff8291636a53d8653616c2155633d69f07f35a72192721bfb88", + "bf16-provider/gradients/0247.pt": "7b4921e404de60678f39763f3f01dc9aa0525b3700f6a666868175d4ed6d2a70", + "bf16-provider/gradients/0258.pt": "7fa9e2361f768b3c11a05912ffad5db282a8c92627cb5c1328015bfa1ed66b16", + "bf16-provider/gradients/0257.pt": "35a0be6b96d695c4bace48f372f6d6da79d529917241a0fc408f542a92da8f79", + "bf16-provider/gradients/0244.pt": "2edda58756a5909e00df680ba20935aaf4b2560219d6795271c7f07cf63fff32", + "bf16-provider/gradients/0243.pt": "0d5dd5c1bc41d326425e58aa87814b63c524da5202202f65ed38e5f955a6e236", + "bf16-provider/gradients/0242.pt": "5e0d121d7d055e4c666de5ea54194efff27fb38124dc46efdf88379fbc952282", + "bf16-provider/gradients/0241.pt": "97ff757508a1f23ea0a5037c65db9afc86d6438abe1b613ebd036f9d0e674918", + "bf16-provider/gradients/0240.pt": "4c994d82c228bf5a5c39d3eaba7fb8290a4593176396979cd05732a34b02ef2b", + "bf16-provider/gradients/0239.pt": "81390d41da237daaf46382ae3d9281a258228ef0d7fa88b2f1bb90abe1f6a63b", + "bf16-provider/gradients/0238.pt": "cce6e8691c4e078d1bc2f6a660fb3b68cdbbad9cd605a062c43f1f4735021dba", + "bf16-provider/gradients/0237.pt": "ababa08f47c4071fa17302e9029a220d69c1d076dd62680ea60aa78ed560724a", + "bf16-provider/gradients/0236.pt": "5c2b4c74861056484b085664f3be32a3dfeb088f4a792876bdc6474c9faa35fa", + "bf16-provider/gradients/0235.pt": "e0818c2664fe961bda4170d740112c554245340796bfc88af79a2465e0aaa1ab", + "bf16-provider/gradients/0246.pt": "96769ebad77b2012bcf50570d35720984269ad4886376203150a0315b2e77b03", + "bf16-provider/gradients/0245.pt": "856296f7d8d4c371d7e6496c781b73b753d90fd1bb93933a308d9e0c3c426969", + "bf16-provider/gradients/0232.pt": "138e34674452dacd239e811f4d0fd95a03ee33c4e2e4ffc82ddc3ab81fb0747a", + "bf16-provider/gradients/0231.pt": "ef62f36d3acd46a543240d3e4eb2457c078041e0d6d7c247266495c32688423d", + "bf16-provider/gradients/0230.pt": "028f6ec9ac2acb4596391ab04bed596121ba3590e8ea542025b3a7b36ba76233", + "bf16-provider/gradients/0229.pt": "f6a26bdf2356aecce69b55a10ed5f757376a572b4c1bf9adf1ea5b48c60889e5", + "bf16-provider/gradients/0228.pt": "3337413bb7b2476214add605a710ddad5b52b0b884c0909625e6aed1fd2fa903", + "bf16-provider/gradients/0227.pt": "f8e336ec4317ab0959462b1514b00488be782431a251f8226d5aeb327bb95ddd", + "bf16-provider/gradients/0226.pt": "e5acb44c93e7eba97d8766da911f7f2cec1d9a81e8494bc0fcee36ac39628392", + "bf16-provider/gradients/0225.pt": "fd46ba7fe63dd918a34eb099b720e3a2cbea2fe8d2a32ce710487dd7e8825249", + "bf16-provider/gradients/0224.pt": "66feef305f8e73b239261d70d32282cc6a4862fc6389284e667b2f9cb1bcdc5c", + "bf16-provider/gradients/0223.pt": "4c00f3c929248a54acb02ebcb2691f3c98eea5dba89b836f641bfe87deeddd4e", + "bf16-provider/gradients/0234.pt": "3a00acaf15200bda3868ebf8657e48a66d743fe15b4a84ad46b8bbbbe6c3ce1e", + "bf16-provider/gradients/0233.pt": "f3713c4b0601d73011409a8a7786a3a2c5005575040d7b7b67e64c3b9db2d951", + "bf16-provider/gradients/0220.pt": "9b9a49fe27a7d4885a41c31084cda29cbfe820861ceec6931c9fb9f272cc2900", + "bf16-provider/gradients/0219.pt": "43a94aa6cda2a92ac3d0b458f1c6b6aab9f8f94149f08fce0d6290179160492c", + "bf16-provider/gradients/0218.pt": "231c88e71b80ba4b5ee9658309e540e75dbae979fd203d5f91f6570bb8c01545", + "bf16-provider/gradients/0217.pt": "02ea797c01ff789685a3daaaccd55b4e4b3da1bc69afc9d015732890d27d2d51", + "bf16-provider/gradients/0216.pt": "a68db9b9da5c22cc5015d4db7f5eff44ea8048ce7aa74a3d5c1948a3d23c4de8", + "bf16-provider/gradients/0215.pt": "e6a8c4f881f31c148bf2424cdcd3f62a89492385659db4cf638d1768f6c3e7d4", + "bf16-provider/gradients/0214.pt": "6268b4b6c92da5faadfeaf8a5268f20f0783fb127d64f25a10916fab0d296e93", + "bf16-provider/gradients/0213.pt": "d1ab32df4a95afc81a053f3087a05f0c6d8dd88eca14ebac174d415bbbb1aaec", + "bf16-provider/gradients/0212.pt": "accce9d5f93aa2cea3c290073756588ae67a72e875031660a5d711ce8744669d", + "bf16-provider/gradients/0211.pt": "c7bb5e8a8094d7addf1260486a3768ce296f18fa7a3790b44afb734d9dc4c6eb", + "bf16-provider/gradients/0222.pt": "bb262babfe633986374c150d32a6c2c4918008707e38d79118cfb3138d0a4a8a", + "bf16-provider/gradients/0221.pt": "367d15082c9c3443e1d302044f0ba89d48a6c4607c1e75a7e00f2bc66af6f081", + "bf16-provider/gradients/0208.pt": "ec69a37acf4fd5ec8b48896c0984b2f36ff580d3c847f0d3f821bc4ddd399d5f", + "bf16-provider/gradients/0207.pt": "ef7b8ad9b92f94a04799bf5f0e470499ad35e382baa15f4c8fa7e9ef402ac66d", + "bf16-provider/gradients/0206.pt": "7ea4af6d3be7edf6d195cacaf5360ab95b4256f4f41fb4cc10a8ed4bb1771d6f", + "bf16-provider/gradients/0205.pt": "3da5315785f276d61cef3108d748074a5a16cc9c6a94e625a059755113e523a4", + "bf16-provider/gradients/0204.pt": "c860184bf7f61c1b423b7a76c496923b20eaa48104265b04aa3da36f4bcd4662", + "bf16-provider/gradients/0203.pt": "142f97eeb0d5376046d69e5e325a0b8a4797adca669687e6f393a64e6bd168fb", + "bf16-provider/gradients/0202.pt": "1f070a3c00ea16a81f7bb03fdc123837f9a6f72eca06c4bed0e7d74d05fe5ea5", + "bf16-provider/gradients/0201.pt": "583c07b39f496470dabd088eee80f0b2e4b0dcb0f1c5da9a796025bd80303955", + "bf16-provider/gradients/0200.pt": "8d2b63e23ee19dd7f9cbfeafcd61d4d5a4052dc7dcbe1344762c7e4473d7b05e", + "bf16-provider/gradients/0199.pt": "a68321b3801cf41e9fa06723164feb62e8d285986cbb7bb3422c85bb4b4cce1f", + "bf16-provider/gradients/0210.pt": "3d1b6ef3a7257378e4c0745b0e8140ab68b1a83b513acf1507417c74429aaf08", + "bf16-provider/gradients/0209.pt": "8da497487256c7cebbeb3acc10d59bf7c438108fa2d0b75fede9c3c6d482e922", + "bf16-provider/gradients/0196.pt": "d56d202cdd76c0302322836571ab2463c93479fe288bfed1857826694fa42742", + "bf16-provider/gradients/0195.pt": "57b41605efeadb40d6aa9108db0e7a5aa7e1f46e5b0f41551fdff305bf2a7f07", + "bf16-provider/gradients/0194.pt": "5d23277607ee99ce66886822f8914ce125a695c1c4a507d9b94d91ce529f9aa0", + "bf16-provider/gradients/0193.pt": "8b7104e81471686c7793303c31404cc359890fb701afedfdc1f0825e8ff03b4d", + "bf16-provider/gradients/0192.pt": "a972ea28abae2e9606fa3ec4e30bf39052e18aff339b66c2de785f6104e8f4ef", + "bf16-provider/gradients/0191.pt": "cad419f09a97deb1912989597ca98c29148705ebb4768f6ede89a5de7d0a4f20", + "bf16-provider/gradients/0190.pt": "dc61442a3ae246cf85e03efed12725f4ae033237e3b92760b81bff1a61b40b05", + "bf16-provider/gradients/0189.pt": "22bc489135f9d4c1ff0f56d8c52977857b32922900555ec1f5bd097ebd9c1a3e", + "bf16-provider/gradients/0188.pt": "b2614dc807f92b54e0ea8589b3fc30241d572e90b8351d5d1db9e467bfc9cab3", + "bf16-provider/gradients/0187.pt": "daba5072a2576fce5fb57cf32e3ee225e4c323ecb5ca67890b7f0e59351658d2", + "bf16-provider/gradients/0198.pt": "45f6987de952c6c13e8379d76188cd83bc648c7a2a05f4292f2f206338ba2211", + "bf16-provider/gradients/0197.pt": "dd4d032f2330d6d263bb9e8ec402ccf5b8103e5ca453286637f67cf373009625", + "bf16-provider/gradients/0184.pt": "7ffd4d1b9fe71b0f514abe5b564a7bc60bdc08551c5fd76ca643589f40279e6e", + "bf16-provider/gradients/0183.pt": "9975bac895e20c1ddc9937b17c851960c2b35a9e5975e9529a36fd246dfe366d", + "bf16-provider/gradients/0182.pt": "a9c98f1544a7c4d2ec3c412b01250ebb273779d6e7d1cfcd76c51533a2bf1b1f", + "bf16-provider/gradients/0181.pt": "6bc9d0c23aae1a4377728a5f80fe0159283d0cc5ed475e8524e2c6427ab4f1c6", + "bf16-provider/gradients/0180.pt": "c5739199f54a40ce008e8500c1dab2b3bd12b7e4c467796152957c652eb3b800", + "bf16-provider/gradients/0179.pt": "e2c76c1bbe74302a29de725df3c8b2e0f14611df60c94a0a2dd6d5c7d8089309", + "bf16-provider/gradients/0178.pt": "cb668f84ff41f5caab5b717d5ca608228d770e208f842de229d71d4c8b550e69", + "bf16-provider/gradients/0177.pt": "7b3dc7301357f3702cd3fd5694d2e5cd456f07b43c439c9e22a2548d30a2de8c", + "bf16-provider/gradients/0176.pt": "6d7ebff37525d991c1879f82dee3e5fa2dbe4dc6373a0c709108c661b2768701", + "bf16-provider/gradients/0175.pt": "a4eeab0695a22cf78fd777301257a2589261390e1cc6f9a14c4d87f6a883b559", + "bf16-provider/gradients/0186.pt": "8434eeded31e0c7ea8b047c94ff5f164c7b50062b574dc0939dd18cbd248d8a8", + "bf16-provider/gradients/0185.pt": "1c1cc9022678617e00a2e806b8e18dafa8b85e6165fcbf78e0590be822441d43", + "bf16-provider/gradients/0172.pt": "016a36a9d3567779fe23d91f99ff855e1a9af8066e373eef9f23299bfffc5837", + "bf16-provider/gradients/0171.pt": "b0c8b6b6aa50924a8139da67dee1efce799e4a8cba77c994ee4087a15ceeca04", + "bf16-provider/gradients/0170.pt": "dc0221a72663a02aa2c447e624e9f5a81d3c8324d1c0918329b72fa5b9ea5682", + "bf16-provider/gradients/0169.pt": "ccc96b95774d18294008e74868b40e27006ccc8c2bd914b5386346524080a97d", + "bf16-provider/gradients/0168.pt": "505b0a5d3af2787f104dda0b8ef704250f737816488c52fbea7d5d7d18a53330", + "bf16-provider/gradients/0167.pt": "711d2fd41c74af125cec317c7cd3b2fda5722f7572124486e2b4913b94591225", + "bf16-provider/gradients/0166.pt": "9c74031094116729122f0f311dbc62f65f56a36e4c2df75ab6a96afff189d3a6", + "bf16-provider/gradients/0165.pt": "58804c807ce3e933f23acc19a9f7298fe75a921a9ea41bd148db05e1bdcf94bc", + "bf16-provider/gradients/0164.pt": "69db65b8d18eeed754cc5896693ad5f42ef0a1b8187416cb761de299cce77ee0", + "bf16-provider/gradients/0163.pt": "ef62de4a59ed2b7e58a66dd933affab1263387e3798e07cfa3490b84f6e11b83", + "bf16-provider/gradients/0174.pt": "09ede8866a958b2ee9e8126369fef6ce7e4323409fe3c00563d8360d9d42101f", + "bf16-provider/gradients/0173.pt": "6a1096389ef27a462216c6b69f687496093ac17f42981176afb6dc45fb6413a0", + "bf16-provider/gradients/0160.pt": "9bde85c82bffb4798e008201297dfd29dbfa0724eb98d499ee1ade1e9f43a549", + "bf16-provider/gradients/0159.pt": "1c559ae77b1bd082da5e9c5e54f2b50c4e58f9e1b49fc37f2ecc303642d7ddba", + "bf16-provider/gradients/0158.pt": "018d969699505bdc6bf01a63d9d49a24004ca43d4e00558fbfede39ad7f7731a", + "bf16-provider/gradients/0157.pt": "9f18ab851c0c1d9c73813dc72daec46bc826635fc0818bdf48df0295b8386e81", + "bf16-provider/gradients/0156.pt": "cd2a439b79d958254592d7f5960001b3f3e6ac4bd9930f7d7774a57f06b957f8", + "bf16-provider/gradients/0155.pt": "c4e09428b8557906d650939be1935fa41c04757640720cd1e8d82351d458821d", + "bf16-provider/gradients/0154.pt": "13b9ecfc3c2af43fde4b48e02654633c6fe0ddd7423190ee23fe785d46c90e56", + "bf16-provider/gradients/0153.pt": "29c671db274b7200b9a57226cc85ab5241cd7c1a30a2ce5c9461470c27484e17", + "bf16-provider/gradients/0152.pt": "3a7377c5348ea6165e9b7b82683e7a84d73e3047bfa84275598c7c2135abd8a4", + "bf16-provider/gradients/0151.pt": "593708635544f6403cf94e36c59324a516cc8cfed11530c16da1dc54ab90e5d4", + "bf16-provider/gradients/0162.pt": "afae5321af111190bbe3a27f8633ca05f7a95b7c195178682c5b3d18062eda0f", + "bf16-provider/gradients/0161.pt": "15bb8d7e4ef8e2b4bff6ab0ac4641aec9faf6f9e6bcac66c62cba59b851315b0", + "bf16-provider/gradients/0148.pt": "e23cd4f3ae6d0f9b19affbd1c48dc80a47d16ef05c624ca126577e0d0495cb7b", + "bf16-provider/gradients/0147.pt": "85f26343f8242e0ff98a6ea9faf60cffaafe9076d4e415476121ae9c28547a68", + "bf16-provider/gradients/0146.pt": "815508d3a5721fb5850dc275553142796cc47f9fcbcc3c6d97a003c66d95df28", + "bf16-provider/gradients/0145.pt": "a759e15ec38a04639058443c47a2a1fa5660eb0d439524f6b986ee652f6936fc", + "bf16-provider/gradients/0144.pt": "acb1925dc915c8cc874d66abd51c5e944b98c629dd334639ff34cc0a97ec2245", + "bf16-provider/gradients/0143.pt": "10da5b0274b175a6223cf48b88ddf37ebf16e0083abfcf7c1a7e96b948f26283", + "bf16-provider/gradients/0142.pt": "162739213b2cd276a285823e42b40e49b737108a5802cce1132cb3f5b812222d", + "bf16-provider/gradients/0141.pt": "cd4c8c4f39356449b76bfa1ca2e7126f6b5e843345ce221cc34cc590524dea3b", + "bf16-provider/gradients/0140.pt": "ec04f8fe998018da5fbd1820a5856d53c690fa17206212057a59e33f321d935e", + "bf16-provider/gradients/0139.pt": "53a5cd74f51b7995c2c0418ae70d94c27d24a05ef83e5dcd1a27bebe5471770a", + "bf16-provider/gradients/0150.pt": "07932e7ee2dbc633b8ab0f8e213bc5f793058f7b59a4b8ebd8c6d5f98485e2d3", + "bf16-provider/gradients/0149.pt": "4f66852763363e3cc77cdc62cfe13ded71b03f9b607849f32443b50b5b520931", + "bf16-provider/gradients/0136.pt": "1d712a653fcf0981214b1924d30d3e8c8c04c50f57e075d70c39aaf65df3eaba", + "bf16-provider/gradients/0135.pt": "c75dc1ed5222c4fa46a564d6032b151bdd36c46cd789f9d75593bbcd14ffb6a2", + "bf16-provider/gradients/0134.pt": "826cc623be29f797ecf4546c34870dbe86bcad0a20b6a61043c7771851153ede", + "bf16-provider/gradients/0133.pt": "1eae2273a5afb2b3acdae724666600f103af57b3d072cd4967bfd0107442d840", + "bf16-provider/gradients/0132.pt": "534b45c8e194ec24845c3ab685390f5d6a8194b2058c85f1e2f91f4bd5210312", + "bf16-provider/gradients/0131.pt": "087678ebe2ed7fb82735eda24a4631aa309a2f40546753fda03d6d691e1b6d00", + "bf16-provider/gradients/0130.pt": "b33178c0f5a3921fba3d9adf920f222a4d20c33f1aeeb206975bdc258c92dc98", + "bf16-provider/gradients/0129.pt": "b701fbc13ec95e85222495fec6c6028b1a1e93aff213eca68c3e7105e450b6c4", + "bf16-provider/gradients/0128.pt": "4c7621976840a1f8459727e71edd47c42ca19f20475c49597355c25344e4f842", + "bf16-provider/gradients/0127.pt": "274dc09a775da8f4a61b32173d2178e8df21ea1b83bbe765c99baa2714f64870", + "bf16-provider/gradients/0138.pt": "877c7a18047c4e3ceadaf73b508d90877081685c75ba1102249a6939ec1d4c25", + "bf16-provider/gradients/0137.pt": "18b570db82f4deeecbf75fb8dc029cca837e6a8d0da23a5fe45f73b663e41892", + "bf16-provider/gradients/0124.pt": "cf19e165eaf709bfe0359d8e486e78497e2b0e00db7c024df03fadf42ebea5f5", + "bf16-provider/gradients/0123.pt": "ac8722337c50c2fb9798afb3afb098e71fc5bf036dc081b6479525d12481959d", + "bf16-provider/gradients/0122.pt": "774e8af6593a0dea7dd67593f6485c293625ff5645fcc251f25d841ea572d7a7", + "bf16-provider/gradients/0121.pt": "9161e0f787bfcc9cbf4e1108de217bc77d25d37811eb4b641c95fe4443e047de", + "bf16-provider/gradients/0120.pt": "89e5bbe7aa9f2dd6dbc0bf76c5e405fb784ced9c68b1de06133a43cc59a65fd3", + "bf16-provider/gradients/0119.pt": "55b11262ad988c34f26974d3cd40a86ae59543f911c8c591517ee27f85efd640", + "bf16-provider/gradients/0118.pt": "0a098c417551c7e109535459167c92d26294628cf3fcbca6f3f17cc22442c2e3", + "bf16-provider/gradients/0117.pt": "fe8c8212c7b9673111ec777f46013df5fd477d8bf00b86c7e4bdab7bfc5e2c93", + "bf16-provider/gradients/0116.pt": "8c3033828275a2771b6bf640047642bebc6c05c0f7369723e8e9e2195dca89bb", + "bf16-provider/gradients/0115.pt": "31cc3f5422a20c671f2cce549e93cc8c9ca1da510f16e46a17a3b4e5aa32bc63", + "bf16-provider/gradients/0126.pt": "3d5e9096ac850b1f195a13943969b7f9ee28a241943c1f70850473ca4672d4d0", + "bf16-provider/gradients/0125.pt": "144130e8520ffbcca2a2ab0f41f6f84807f181c1d1077c3435206965b75b53ab", + "bf16-provider/gradients/0112.pt": "5c7af8d04a634bededab66333074d7626ad411ac306d693dc95ed6def0459cd8", + "bf16-provider/gradients/0111.pt": "f59db7383925bd7ea207b917235a06ede25592c747ee02a38c721677a6003314", + "bf16-provider/gradients/0110.pt": "ac80f3091f2725f92ba02eeb582c806eb863da016ff92af423bf777814181f0d", + "bf16-provider/gradients/0109.pt": "f0a2afc6c0fc3844b3529e2be8ffb7255a2ca864a8ae2d0a9e3f669a20c14ca8", + "bf16-provider/gradients/0108.pt": "9bb5c3eb397301d96fd1490756031a1f8d689f9b6242e5f1d6de836a34292d1c", + "bf16-provider/gradients/0107.pt": "7c613e45624c62ae2e08c0c477923fd46fdb9092cb733edd403976c4a04b4b6b", + "bf16-provider/gradients/0106.pt": "63c7b0a3b11ccf9e9d6011884d9cb36b6284bbe67d160f14028fbfb1de61bbac", + "bf16-provider/gradients/0105.pt": "aecae0cef955e9e1e1fb66d901716569c8e5d74dd3db01eee052e81ea3c6dc42", + "bf16-provider/gradients/0104.pt": "d6d8d40daeac243f6fd9975d4ccac10bcf30354f35173a562ba6e05f6ad6a7a1", + "bf16-provider/gradients/0103.pt": "0db7dbf32f62f6659ee1205f91b81ca1805b529ae75e725eb5ca7b1909ff44cd", + "bf16-provider/gradients/0114.pt": "6da9b3527255cd82ac91ea421bb0802a9474902783191eb0f5d65b0ac062dde4", + "bf16-provider/gradients/0113.pt": "10630f0a39d2d0abfb49aa3ac08d173f34ab6eb11ce15bf3672f207a887ead14", + "bf16-provider/gradients/0100.pt": "2832a9cc77db09443da100feb1553606e3ea625218f5ef4574724c909cc54a9e", + "bf16-provider/gradients/0099.pt": "5e37cd9e586a650413176939381063b26530930fa26daabb00e60d50b1353d2c", + "bf16-provider/gradients/0098.pt": "8bbf8b390b4441065205c70a7e5e8b011bd1b8f35f5f2da2a827562190c3d9c7", + "bf16-provider/gradients/0097.pt": "6ebc32a32feb456765d1e3ad7bd96987e6ff8cad906c1b7d3e352e08e410f711", + "bf16-provider/gradients/0096.pt": "4bee3136156a423a43b4ff236a72d3a2a56932233ab25e89a984bbaaf2c8e4f4", + "bf16-provider/gradients/0095.pt": "bb34ac908f20a239d98c1f5c0416d91cefa95f5c105ecd241f24de8bebacfa3a", + "bf16-provider/gradients/0094.pt": "17d4322dd958c3e096b66230312b58cf14646e24b500234e3c6f302b70f2568d", + "bf16-provider/gradients/0093.pt": "dd70dcf21fb8ccf93a51991e89415380bf421b2f550245076f7fd4c877d786b0", + "bf16-provider/gradients/0092.pt": "904a0e7949f730a8257cbb4fb7afe68e05bf491ad40bf36e51e55aacf2438eee", + "bf16-provider/gradients/0091.pt": "15ce62a31c2871e75845a26cb32f8668c13c6d544e0ee941e090eddf48af1f41", + "bf16-provider/gradients/0102.pt": "2105ddf33f8b5d864df4e9c8f005099fda2d241cea290e0c8c0edc7f7f9d1f44", + "bf16-provider/gradients/0101.pt": "3f6b9693e70bc787745771556135dfed5ccf4265de5eef6da488622de0517577", + "bf16-provider/gradients/0088.pt": "2cb9cf8ce9dd39c7f7f57f4ad225f9a5bc99e9113b243095905f5edcedde9cbd", + "bf16-provider/gradients/0087.pt": "435fa661238b6c1b4551c5491d18099d7468149a4d48a1da7a45b3cdd8a116e8", + "bf16-provider/gradients/0086.pt": "b9d949c5cadedcfe6d46f8192de6f307fff9f7b05c9faa0d73b165f349f8c05f", + "bf16-provider/gradients/0085.pt": "4c2a4273bd9defa879591794a590eab1242047245433c618cf23e95f41cb87e8", + "bf16-provider/gradients/0084.pt": "f35f309d595786a6e929928b5b50c9a93bf0a91d8b9a2e744836ab53978eea6c", + "bf16-provider/gradients/0083.pt": "166cfdaf62b8893a83d92dcf8bbdd162edfa88bc48f0fdfd4d2e5aa8fab1d8d0", + "bf16-provider/gradients/0082.pt": "143aae86d17c6a1935980f0b2d1eff88d1aab8c942d28fe982a8091f3f9a0031", + "bf16-provider/gradients/0081.pt": "5f473c39fc3374181c611d1eb265679ddcdda92ca1fab207f75a4cf5762bb459", + "bf16-provider/gradients/0080.pt": "4d42a16f74698f5bcee7ae76b980878506cb23b4b8f3801bd1f4bd0cbc014ba4", + "bf16-provider/gradients/0079.pt": "c61dd2aa8d7b9c462fbb35803ca2232ebafe460a83147eb4aa8db48922035b2b", + "bf16-provider/gradients/0090.pt": "0141bc5cf23698c2f891fe0824f2f7d1bea7400b01f1ffb38c908574260fe5b9", + "bf16-provider/gradients/0089.pt": "938b32ab210aa1ce7cff95866361c0230371d44cf54c67dd945324613926972a", + "bf16-provider/gradients/0076.pt": "293d60dffa0f358385c13b7e3d456b5598dd9966e3fd121c6ff82455c318f374", + "bf16-provider/gradients/0075.pt": "f36571750844e76325093b7c37ac3bc916ce5374afec6f260b7d046c03c3a70a", + "bf16-provider/gradients/0074.pt": "9038d73341d28c607d92439ff89ca0bda9c372953ad4c6ea5251e7e748c9cb09", + "bf16-provider/gradients/0073.pt": "339cd2fcee85b12299e2cb9548f33ceaff90765e4eedfb086edb6c13a9ecd65f", + "bf16-provider/gradients/0072.pt": "1ab591630cbed9bdc50b6efc206608f81ff91db1398c80174d6b941f181ecdea", + "bf16-provider/gradients/0071.pt": "d09d2cdc8fa64016333be8b54e51b169944e4a4a408e216152760d79f2ada1dd", + "bf16-provider/gradients/0070.pt": "5d91967d801a1adc4cfcc5a317fad0c0fce6aed02f7b406ba30883f493b9d33f", + "bf16-provider/gradients/0069.pt": "d793217a0b139aa0130f525a81bf5ca4d70b91ccf920782d2f6d0ed49e199640", + "bf16-provider/gradients/0068.pt": "f4d4e7601a3d80129447decf647ead76d36587b33ffa68cf135b54c68395a202", + "bf16-provider/gradients/0067.pt": "2d646468bfd531007d63d06c684387e9a86b9c634746ba8c20790a9e348f58c5", + "bf16-provider/gradients/0078.pt": "1b8c659f0574cf4febdf7f8fa8df027d81d411244705639760530716780647de", + "bf16-provider/gradients/0077.pt": "e48b2388258704879784d4b0bd4e589e85f84285c333fb04d0d23e3c1af332b6", + "bf16-provider/gradients/0064.pt": "b7d99f27a3c4b724fc6b7cc74903657b04465c0432e52c92121ff65ffef82eb7", + "bf16-provider/gradients/0063.pt": "4133fb7012619a578275b9011dfc5eb37c3c1619d22a18379d69d7808f7d3e08", + "bf16-provider/gradients/0062.pt": "0fa4df1da173e0fbc15ede70c252eed208b1b3e6c1b4a6ef77f064bc160be685", + "bf16-provider/gradients/0061.pt": "39b847e7d283acdd0acada980458ea5272dacf5d9a89902228ae58106f596de2", + "bf16-provider/gradients/0060.pt": "97872df7510332413337eca50c394132423e2c26c231339d1328261332892216", + "bf16-provider/gradients/0059.pt": "152c079726dd0ae673badbc4c878db1ba144bbdffa781947a9018dc3a901530d", + "bf16-provider/gradients/0058.pt": "17996bdb44c46cdba536edcb3271ac25a6e914c11f6f8b9392702475012e19da", + "bf16-provider/gradients/0057.pt": "e6a2d5f495f9dfa98402fcfe5292a85b31927351a00fb0aa0c01e139f21ed7c8", + "bf16-provider/gradients/0056.pt": "6c7eb067b1ef06be780b031196b052f2065f7646798962af6ee5910251ab814a", + "bf16-provider/gradients/0055.pt": "ee511ef6481d96e4178843075ea5d3c65d5fdb5a0640b3fb80b83215a01d79b0", + "bf16-provider/gradients/0066.pt": "71fa3630d746d844ef8332249afde4763182fb108b603d90d32807e60ac085ea", + "bf16-provider/gradients/0065.pt": "3a5690dfac45adec64f184ee74a2e514bd435c2ab6bb7eea078e924e7b3b9975", + "bf16-provider/gradients/0052.pt": "55f1a640d777a587dae8ecba98331694099dcee1ec05f2801e3e34a890c675ec", + "bf16-provider/gradients/0051.pt": "10d6da1d756568758347c08247e50184ada2aafaf5c3ddbdf73d661d672f71e6", + "bf16-provider/gradients/0050.pt": "56bed646e4f450fa1d6f8a8bd5ec45484a69c874efb7ba9b2fd7d6149e1b1318", + "bf16-provider/gradients/0049.pt": "dbbd40845e50516541f796388f91efaaf805d1594880bdb63292f97bd6a9f025", + "bf16-provider/gradients/0048.pt": "2bc9cb41dd36ffe8d4681a45e69969a29b97a30ecd717cf1fe6f776dc239a76e", + "bf16-provider/gradients/0047.pt": "4c14231efb114fea51cd04115a861cb9d1e3acc56257c84cac2c3a8d3ff63a95", + "bf16-provider/gradients/0046.pt": "1f6dfc8143744f1bdfc3b6dd2a42e041ec1dcf863369a02b4431da4c2b0e0b46", + "bf16-provider/gradients/0045.pt": "e42525b6185b5070da925ba8ecea153090ba1d08edf3f9d89f0bd7dddd105bf7", + "bf16-provider/gradients/0044.pt": "911f5953d0c2a5f4746549c1e9703f1f596c02496198d2831797a3e37c4f5d71", + "bf16-provider/gradients/0043.pt": "44d0fda8dc1ce5e177bbd028228a60f3c71fe96689c0811ad700476ebec370b4", + "bf16-provider/gradients/0054.pt": "1ffec37aa9359996c99b32f7d9c785d89bb107806e8fa1ba3bbb1f5d6bf60402", + "bf16-provider/gradients/0053.pt": "d874a0151cfe590905adcd4b089cb0cf8d0193ec4073da7dcac82b3f3999eb6d", + "bf16-provider/gradients/0040.pt": "4895de0e001cb9602018260946f4b3c5f508d526aaac4659b8b18b30e677c441", + "bf16-provider/gradients/0039.pt": "12eebd64cd38985f12a7f1d39ca002b35206c2d37a00c94846209fe96b2b9811", + "bf16-provider/gradients/0038.pt": "44d110c789407ada35b4b984fed865ade732f26d79c78762afcba528897ed0e1", + "bf16-provider/gradients/0037.pt": "24fb83adcaa5f42f51ea86c546221a85ab7937605150460652459d44c6dd29e4", + "bf16-provider/gradients/0036.pt": "ba28120bda4a87b324e20c4e03e9497e9df18e489e1b89df82eb5b9233ec5537", + "bf16-provider/gradients/0035.pt": "c5f967b3ef117e01c047e6b578deaf86154b44619055e82f10378833356c571e", + "bf16-provider/gradients/0034.pt": "fcf2313216dd2b8fdc8ed31a389b459d8704e281685f83b8be8893cbbdf8236a", + "bf16-provider/gradients/0033.pt": "48ed901ebc30a8eb15c8aa00bec6bee7fd589c70031fa84227a34bd830dc51c9", + "bf16-provider/gradients/0032.pt": "9f23a16e0d67d0d5781e55fe6d90230f2c7a69f668867d084dc064843f861e6a", + "bf16-provider/gradients/0031.pt": "89b10eef54ff52932a865f92174f3d97306ced96759aed3a069e8c11a6eee8ac", + "bf16-provider/gradients/0042.pt": "080dafd3d8a7c0c8077931a5a7f3a160e4b1add36f84cd70ce7d359b0313853b", + "bf16-provider/gradients/0041.pt": "07d1ba8a4a7d26df9da7ca67280d374596a320739659b21500814befcfce439d", + "bf16-provider/gradients/0009.pt": "de47e82f2f3500d93f56e4c028f0a4fa8d2790f4e62479754dd2412a99fa7b7d", + "bf16-provider/gradients/0008.pt": "1dfc2144fa6d768523cb6898d6a20b99cca90b46e3515ea4fb1ec738310a5b74", + "bf16-provider/gradients/0007.pt": "17f995c829b58ad1bb6a584222727be244251458f564b7fef5c7dd88d4c198dd", + "bf16-provider/gradients/0006.pt": "62bcb914fc74eff342f30a65eb2cf3d1509d550912f1bc58bf684250b7017534", + "bf16-provider/gradients/0030.pt": "e88add9ff07e2e340f831ddeb0d6cdb7234da718dbeacf483e96fcddb340b3dc", + "bf16-provider/gradients/0029.pt": "565007293ed7d6ccd577c0879757e489f63aa81ee5846f268a55831610b14adf", + "bf16-provider/gradients/0028.pt": "ae8bf5675c319d904ad3a47466ac56616eea4cd76dcad9f21a860eaea6702cbe", + "bf16-provider/gradients/0027.pt": "5b3b2dba09b7e4811283d4797c97eef57378dfee7c4f0bb0bc89b3732f3c3d79", + "bf16-provider/gradients/0026.pt": "ed74b7bca7380650b552056563196958929d6d130ebff126550df717eef15123", + "bf16-provider/gradients/0025.pt": "d15c5a45bb7858ce7a2033573d3bd8d15b3d502db00ddc9335ac2880d5f20ef6", + "bf16-provider/gradients/0024.pt": "6970eb0e2285784f4c5495009a87dfbb36ac612c5570e04d220c7484669e727d", + "bf16-provider/gradients/0023.pt": "de6c1a0c38b17d5054a85664dd4587a942e036305754bf3796bead97ce2b9957", + "bf16-provider/gradients/0022.pt": "55761c3189e84cc1bc2de51c4f270e679ffbd532145994deaf253930117d6e7d", + "bf16-provider/gradients/0021.pt": "b232701911cda69913bff55f8c305b3d6da1b08f3ddd1eef63c1b3883a4e2313", + "bf16-provider/gradients/0020.pt": "b113fe0092d38ff7c209ca0cae8bb438e437e1d9e80cc59b7f1933a84fdf6be1", + "bf16-provider/gradients/0019.pt": "73c9019f4869b3233e55b8d3ddcf641f27a2c398e6d843cf50040b6c0215ffa8", + "bf16-provider/gradients/0018.pt": "94a2d5447adf7806a9b1c012b4ba95e180dec535f9b658b8fea5c160dc108e0f", + "bf16-provider/gradients/0017.pt": "d721ce7e8891d57d9a5e34c06211d2e3e5aafeb610fb4b00c23ed21d61a17591", + "bf16-provider/gradients/0016.pt": "6ffaf1248e2c88315943bdfbece1abc8e24dc655634d801a6203ea6016572e85", + "bf16-provider/gradients/0015.pt": "d1d638a00f6a41657f8073129fd3f7c3d57cf58915aab47a24d79f66f69a6aec", + "bf16-provider/gradients/0014.pt": "b12ab91a2aa7392166ad5d2090a26a3ca1f62b98ccdabab9594d685806ecfa42", + "bf16-provider/gradients/0013.pt": "ba45fd0ff6f5600ecb2a9ec6babda8836c863225430c5430d5b78c16b0ccb75a", + "bf16-provider/gradients/0012.pt": "816a2593b6597196ee735b7909ba63ff33f8003e98e0411a834daad804333047", + "bf16-provider/gradients/0011.pt": "a0b24512b80d5a7de9c1c5d636fc977c263d9507d30352e537f4c4af04d5b8ac", + "bf16-provider/gradients/0010.pt": "aa8e3825d6987a4ae2ebe5af3e3d7326d2102f1a98ac4fc4516327897e5b6a7b", + "bf16-provider/gradients/0005.pt": "5a661a8f6aa437196011babdafe14d8c754ad2e4be75900786e6507cb134b700", + "bf16-provider/gradients/0004.pt": "4beae74ac56b4eaed53c2b17a7d31f98d045e64e01d067d5067851304a0efcb3", + "bf16-provider/gradients/0003.pt": "b3427ad172c897d9ca3f684d70286d98f99f9074b37a369e7aad5e2fb69eb3df", + "bf16-provider/gradients/0002.pt": "1c56aa9e0a604d21e52062270e0c6ce5e319b79694986338da13dbaa8dd81f08", + "bf16-provider/gradients/0001.pt": "f575932d4d780dcf5da597987286d0c9cccb1bc589113a38d0a0b010b56a5e87", + "bf16-provider/gradients/0000.pt": "369cf6e3c0a16b175a5eb8d95e43b339b96a1f7f01b1b765f05eb8a38048897a", + "generation-reference/42/report.json": "e55e877570fa3f089b158be76a4423f8b25443af5c85fca45ba3dd378974c123", + "generation-reference/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference/42/audio.wav": "373efac29bd3073f1dfebf0cba30c51848866351fb6d2c8e55921047c82cd296", + "generation-reference/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference/42/video.mp4": "560d134d198e725880f308be6b8ebafd71718d01f323ba7b5f3fbee379302d01", + "generation-reference/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference/43/report.json": "7e181902ac5984cd04c8773a2c0bae77ba10618b63b0e5b2a3b56599c4ac5525", + "generation-reference/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference/43/audio.wav": "eb404d2f89fbf3eb9795ef2704c245618e47e552a3b4d782b52602762569de6a", + "generation-reference/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference/43/video.mp4": "14dc8307aab50e41849f02e3726612a5e2b5d197dc04987ae9196d2b63c2e5c7", + "generation-reference/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference/44/report.json": "50f571441bad710594cb53fb861cc73fd6b8b9e81a2fbe82de1f9dea4feaddfc", + "generation-reference/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference/44/audio.wav": "0cab0190275bb143bbd997f50928cea3e2c0208f90cd71b9633046ec6e4ddf0c", + "generation-reference/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference/44/video.mp4": "e68e18a0fd4a1ea9c5da84969c70299d9d21db6ce9aa8592bd649f84c6b9b595", + "generation-reference/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "generation-reference-repeat/42/report.json": "e1f254991fa8c8e72f2dffc4e30a056d38e9e302fd2f684d8c593895afc020e4", + "generation-reference-repeat/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference-repeat/42/audio.wav": "b9891b3a34e097361f06e63a9e890dc4551600a9ec575a08dc07ef9493c19fe8", + "generation-reference-repeat/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference-repeat/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference-repeat/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference-repeat/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference-repeat/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference-repeat/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference-repeat/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference-repeat/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference-repeat/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference-repeat/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference-repeat/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference-repeat/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference-repeat/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference-repeat/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference-repeat/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference-repeat/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference-repeat/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference-repeat/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference-repeat/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference-repeat/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference-repeat/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference-repeat/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference-repeat/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference-repeat/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference-repeat/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference-repeat/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference-repeat/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference-repeat/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference-repeat/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference-repeat/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference-repeat/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference-repeat/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference-repeat/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference-repeat/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference-repeat/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference-repeat/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference-repeat/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference-repeat/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference-repeat/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference-repeat/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference-repeat/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference-repeat/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference-repeat/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference-repeat/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference-repeat/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference-repeat/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference-repeat/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference-repeat/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference-repeat/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference-repeat/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference-repeat/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference-repeat/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference-repeat/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference-repeat/42/video.mp4": "f97c122eada187d8805bb088a2326774af2524e598091446c1302924ee043d9e", + "generation-reference-repeat/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference-repeat/43/report.json": "914a6e564097b23305d6bce5ba4b06a30603d5dd3c0866e3688f3a8a8bd76987", + "generation-reference-repeat/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference-repeat/43/audio.wav": "1cc055d49792fa9760701ad092f1f9e9dba691503ca7a4205702fc3e0c6d63e1", + "generation-reference-repeat/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference-repeat/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference-repeat/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference-repeat/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference-repeat/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference-repeat/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference-repeat/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference-repeat/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference-repeat/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference-repeat/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference-repeat/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference-repeat/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference-repeat/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference-repeat/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference-repeat/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference-repeat/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference-repeat/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference-repeat/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference-repeat/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference-repeat/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference-repeat/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference-repeat/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference-repeat/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference-repeat/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference-repeat/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference-repeat/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference-repeat/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference-repeat/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference-repeat/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference-repeat/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference-repeat/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference-repeat/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference-repeat/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference-repeat/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference-repeat/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference-repeat/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference-repeat/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference-repeat/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference-repeat/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference-repeat/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference-repeat/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference-repeat/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference-repeat/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference-repeat/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference-repeat/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference-repeat/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference-repeat/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference-repeat/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference-repeat/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference-repeat/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference-repeat/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference-repeat/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference-repeat/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference-repeat/43/video.mp4": "51af10f6386bbe0b12db43d7c5d4727fe5e6c1aa91e5ee7502ff4b9e71c96f3a", + "generation-reference-repeat/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference-repeat/44/report.json": "d72c8d8cdc680c7af21481cd1fd35e821fc231020caa8c310497e25833433100", + "generation-reference-repeat/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference-repeat/44/audio.wav": "edda8f495c4288cd872d2a585bfa17d9a4f7e6d0df79893a0cf1b08f3adaff74", + "generation-reference-repeat/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference-repeat/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference-repeat/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference-repeat/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference-repeat/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference-repeat/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference-repeat/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference-repeat/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference-repeat/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference-repeat/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference-repeat/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference-repeat/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference-repeat/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference-repeat/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference-repeat/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference-repeat/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference-repeat/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference-repeat/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference-repeat/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference-repeat/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference-repeat/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference-repeat/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference-repeat/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference-repeat/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference-repeat/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference-repeat/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference-repeat/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference-repeat/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference-repeat/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference-repeat/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference-repeat/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference-repeat/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference-repeat/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference-repeat/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference-repeat/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference-repeat/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference-repeat/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference-repeat/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference-repeat/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference-repeat/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference-repeat/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference-repeat/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference-repeat/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference-repeat/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference-repeat/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference-repeat/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference-repeat/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference-repeat/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference-repeat/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference-repeat/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference-repeat/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference-repeat/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference-repeat/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference-repeat/44/video.mp4": "929285c63024073b07eab1f6b7c3fbf5d2e49cc0f07b2e426dff5fb006d09196", + "generation-reference-repeat/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "full-model-thresholds.json": "9542bce25d394e72f24c2cc02012fc47dc16100a37d5f86feda59c132132eb30", + "reference-repeat-quality.json": "3817569746ac3a1c99cffe32bdf183fbe12de2817c3e60fd6c21c17711bdd691", + "commit.txt": "7c9a380b0ba89a2288ca661bd2fda1e102a02606521681ab0bcf3077e5d2d600", + "dependencies.txt": "6139443eb38b47a805548b0d892808f5b6a26c7bc3b3796ff96ef6b9b0185cbf", + "extra-dependencies.json": "d2655f41325039e28fd4aae62d5637de0df49d08ea8dc9545295cdbbf1f47ce6" + } +} diff --git a/examples/models/minimax_h3/reference.sbatch b/examples/models/minimax_h3/reference.sbatch new file mode 100644 index 000000000..ce7429710 --- /dev/null +++ b/examples/models/minimax_h3/reference.sbatch @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +result_root=$4 +reference_image=$5 +provider_source=$6 +extra_dependencies=$7 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +"$python_bin" -m pip list --path "$extra_dependencies" --format=json > "$result_dir/extra-dependencies.json" +for mode in fp32-reference bf16-provider; do + "$python_bin" -m tools.validation.models.h3_full_model \ + --checkpoint "$checkpoint" --manifest rl_engine/validation/models/h3_full_manifest.json \ + --mode "$mode" --output "$result_dir/$mode" > "$result_dir/$mode.log" 2>&1 +done +"$python_bin" -m tools.validation.models.h3_full_compare \ + --reference "$result_dir/fp32-reference" --other "$result_dir/bf16-provider" \ + --output "$result_dir/full-model-thresholds.json" > "$result_dir/calibration.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --preflight --output "$result_dir/generation-preflight" \ + > "$result_dir/generation-preflight.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference" \ + > "$result_dir/generation-reference.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference-repeat" \ + > "$result_dir/generation-reference-repeat.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_generation_compare --reference "$result_dir/generation-reference" \ + --candidate "$result_dir/generation-reference-repeat" --provider-repeat \ + --contract examples/models/minimax_h3/generation-contract.json --output "$result_dir/reference-repeat-quality.json" \ + > "$result_dir/reference-repeat-quality.log" 2>&1 +# Candidate jobs are submitted only after these immutable reference artifacts +# are reviewed and the numerical thresholds have been committed. diff --git a/rl_engine/validation/models/h3_full_manifest.json b/rl_engine/validation/models/h3_full_manifest.json new file mode 100644 index 000000000..6cc2e5b5a --- /dev/null +++ b/rl_engine/validation/models/h3_full_manifest.json @@ -0,0 +1,311 @@ +{ + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "files": [ + { + "path": "audio_scheduler/scheduler_config.json", + "size": 96, + "sha256": "804780f7133477067bd6bbfbc02dc8b3cf9feeb400f97c08f5b1d5f6cbab3840" + }, + { + "path": "audio_vae/config.json", + "size": 2271, + "sha256": "9a3c645ff892b376c6f5f4c8685964cd75474731af594ff058492a0000caabb6" + }, + { + "path": "audio_vae/diffusion_pytorch_model.safetensors", + "size": 605429340, + "sha256": "52c59e67ba8de5477c81bfbced0327aabf500f1bfdeefd5ee754529241cb26cb" + }, + { + "path": "model_index.json", + "size": 2936, + "sha256": "5a587fe13b2371427415ac892463142683aefcd8d322e274a3a095eac37ac7d2" + }, + { + "path": "modular_model_index.json", + "size": 2935, + "sha256": "a2b6a210e482ffb78e613b553f570c44e101afce6741bd4ed91429d0559af031" + }, + { + "path": "processor/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "processor/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "processor/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "processor/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "processor/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "processor/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "processor/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "scheduler/scheduler_config.json", + "size": 97, + "sha256": "8fa6c3aa70dc9e691e1a6df899fd1b6f75f70481a27cee6e18a303817075c304" + }, + { + "path": "text_encoder/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "text_encoder/config.json", + "size": 1474, + "sha256": "d2dd0c60d01b9e195d9447c52da61c7302d28828524914c044d9c6e1b81d0427" + }, + { + "path": "text_encoder/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "text_encoder/model-00001-of-00014.safetensors", + "size": 4932328944, + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb" + }, + { + "path": "text_encoder/model-00002-of-00014.safetensors", + "size": 4875990528, + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa" + }, + { + "path": "text_encoder/model-00003-of-00014.safetensors", + "size": 4875990552, + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940" + }, + { + "path": "text_encoder/model-00004-of-00014.safetensors", + "size": 4875990584, + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f" + }, + { + "path": "text_encoder/model-00005-of-00014.safetensors", + "size": 4875990584, + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2" + }, + { + "path": "text_encoder/model-00006-of-00014.safetensors", + "size": 4875990584, + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974" + }, + { + "path": "text_encoder/model-00007-of-00014.safetensors", + "size": 4875990584, + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934" + }, + { + "path": "text_encoder/model-00008-of-00014.safetensors", + "size": 4875990584, + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f" + }, + { + "path": "text_encoder/model-00009-of-00014.safetensors", + "size": 4875990584, + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47" + }, + { + "path": "text_encoder/model-00010-of-00014.safetensors", + "size": 4875990584, + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f" + }, + { + "path": "text_encoder/model-00011-of-00014.safetensors", + "size": 4875990584, + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f" + }, + { + "path": "text_encoder/model-00012-of-00014.safetensors", + "size": 4875990584, + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e" + }, + { + "path": "text_encoder/model-00013-of-00014.safetensors", + "size": 4875990584, + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c" + }, + { + "path": "text_encoder/model-00014-of-00014.safetensors", + "size": 3270697008, + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b" + }, + { + "path": "text_encoder/model.safetensors.index.json", + "size": 97831, + "sha256": "06c952c569285870b811989b794b9766493e280fb77fbcb957fc4e5fcf25403a" + }, + { + "path": "text_encoder/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "text_encoder/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "text_encoder/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "text_encoder/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "text_encoder/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "tokenizer/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "tokenizer/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "tokenizer/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "tokenizer/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "transformer/config.json", + "size": 546, + "sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00014.safetensors", + "size": 4825958704, + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00014.safetensors", + "size": 4702158032, + "sha256": "2c4d362eddd2802180ac9c744849eb9ba8d9c8b984bdf9822cb02ed004b29184" + }, + { + "path": "transformer/diffusion_pytorch_model-00003-of-00014.safetensors", + "size": 4933368192, + "sha256": "949c5aafbbfa5654da730a6a7fafd75adb164d0857b095a30e8bb6d390887d69" + }, + { + "path": "transformer/diffusion_pytorch_model-00004-of-00014.safetensors", + "size": 4567069608, + "sha256": "eef7616790105ee839766bb2027203bf2c0d87c6aa038dca84145a8675f5ce28" + }, + { + "path": "transformer/diffusion_pytorch_model-00005-of-00014.safetensors", + "size": 4702158080, + "sha256": "43fdf42d638e8bc6745f713fae80c93bb301807a1a5ae7249344ce28e202a494" + }, + { + "path": "transformer/diffusion_pytorch_model-00006-of-00014.safetensors", + "size": 4933368232, + "sha256": "6442510b34d173653f0cce5c964b935395a8f7accf0b9cc0aa31aec59805239d" + }, + { + "path": "transformer/diffusion_pytorch_model-00007-of-00014.safetensors", + "size": 4567069608, + "sha256": "29f48f535c91dac76496ca821eeb16ca24bc4caf3f0cae8b920a89b1f966da6d" + }, + { + "path": "transformer/diffusion_pytorch_model-00008-of-00014.safetensors", + "size": 4702158080, + "sha256": "c711b096c764bd60f0b8b6ad49518bfab6d614fb788c725add8741c0674a4cd8" + }, + { + "path": "transformer/diffusion_pytorch_model-00009-of-00014.safetensors", + "size": 4933368232, + "sha256": "44428defe3976cbb87635ad200b958199e739986697cd29fdf27aeb7294b5944" + }, + { + "path": "transformer/diffusion_pytorch_model-00010-of-00014.safetensors", + "size": 4567069608, + "sha256": "3d44939c374c9da382e9c6877e1946adf7b84e08c7a881c068f228d6849411c9" + }, + { + "path": "transformer/diffusion_pytorch_model-00011-of-00014.safetensors", + "size": 4702158080, + "sha256": "224d24430b58127a5577721084e0e704a0e74ec96dd7c35bc6fc0994ebd87c33" + }, + { + "path": "transformer/diffusion_pytorch_model-00012-of-00014.safetensors", + "size": 4933368232, + "sha256": "48fa2bd8fe134eef565ab2464f1c2589a6657cba0d14283dfc06b532f8961f3c" + }, + { + "path": "transformer/diffusion_pytorch_model-00013-of-00014.safetensors", + "size": 4567069608, + "sha256": "be5b4b1809f9d546ffd4b3fcf41e5c1e02b819125caa6bc105c109b04c051bd3" + }, + { + "path": "transformer/diffusion_pytorch_model-00014-of-00014.safetensors", + "size": 4644161920, + "sha256": "8fbd5e6c1fb1df7ce988ca90f3d59e7610e465c7517e4b344eda4a214ba4b97d" + }, + { + "path": "transformer/diffusion_pytorch_model.safetensors.index.json", + "size": 64488, + "sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e" + }, + { + "path": "vae/config.json", + "size": 2011, + "sha256": "78f67deec3d63aae807f2bfe7154bc1e26f6372cb20b63265fcbae1b62bb5745" + }, + { + "path": "vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "size": 5061033024, + "sha256": "72f4c6be84ac0674f27398cde991dd9d719762f3952c4921aa66b2ce542f6374" + }, + { + "path": "vae/diffusion_pytorch_model-00002-of-00003.safetensors", + "size": 4955986528, + "sha256": "2e05e8bc23fa4071043e17fd242be8acd0685e781a43987432b2eae925be4198" + }, + { + "path": "vae/diffusion_pytorch_model-00003-of-00003.safetensors", + "size": 398539336, + "sha256": "c05d6ac4b1a33de372799d708531da6320f6a3ce6d1ce6d895e770988e004a39" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors.index.json", + "size": 74228, + "sha256": "15f6d44553c3c616b0dc999920aa784f92ecee7e4201f1f99ac405cfbf3061ca" + } + ] +} diff --git a/rl_engine/validation/models/h3_full_model.py b/rl_engine/validation/models/h3_full_model.py new file mode 100644 index 000000000..1219ef89f --- /dev/null +++ b/rl_engine/validation/models/h3_full_model.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Instrument the full pinned H3 provider with the actual RL-Kernel block graph. + +Input/output projections and both text-refiner layers retain provider arithmetic. +Each of the 50 distinct blocks executes the existing candidate graph with its own +parameters. No weights are copied from block zero. This is integration evidence, +not qualification of the reference/interim bindings as final RFC kernels. +""" + +from __future__ import annotations + +from collections import Counter +from contextlib import contextmanager +from types import MethodType + +import torch + +from rl_engine.validation.models import h3_block + +MODEL_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe" +PROVIDER_REVISION = "f53d552036a0d1bd5570782a39cd40cfabf112bc" + + +def validate_full_structure(model): + expected = dict( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + for key, value in expected.items(): + if getattr(model.config, key) != value: + raise ValueError(f"full H3 requires {key}={value}") + blocks = model.transformer_blocks + if len(blocks) != 50 or len({id(b) for b in blocks}) != 50: + raise ValueError("50 distinct transformer blocks are required") + owned = set() + for i, block in enumerate(blocks): + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError(f"block {i} has missing or shared parameters") + owned.update(id(p) for p in params) + required = ( + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ) + if any(not hasattr(model, name) for name in required): + raise ValueError("full-model boundary components are missing") + refiners = model.token_refiner.refiner_blocks + if len(refiners) != 2 or len({id(b) for b in refiners}) != 2: + raise ValueError("two distinct real refiner blocks are required") + for block in refiners: + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError("refiner parameters are missing or shared") + owned.update(id(p) for p in params) + return required + + +class FullModelTrace: + def __init__(self): + self.calls = Counter() + self.backward_calls = Counter() + self.snapshots = {} + self.capture = False + self.actual_backends = {} + self.candidate = False + self.candidate_calls = Counter() + self.forward_order = [] + self.backward_order = [] + self.execution_order = [] + self._events = set() + + def record_event(self, name): + if name not in self._events: + self._events.add(name) + self.execution_order.append(name) + + def observe(self, name, value): + self.calls[name] += 1 + self.record_event("nodes." + name) + if name not in self.forward_order: + self.forward_order.append(name) + if not isinstance(value, torch.Tensor): + return + if self.capture: + self.snapshots[name] = value.detach().cpu() + if value.requires_grad: + + def record(grad): + self.backward_calls[name] += 1 + self.record_event("nodes." + name + ".grad") + if name not in self.backward_order: + self.backward_order.append(name) + if self.capture: + self.snapshots[name + ".grad"] = grad.detach().cpu() + + value.register_hook(record) + + def validate(self, *, backward=False): + required = ( + [f"block.{i}" for i in range(50)] + + ["refiner.0", "refiner.1"] + + [ + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ] + ) + missing = [name for name in required if not self.calls[name]] + if backward: + missing += [name + ".backward" for name in required if not self.backward_calls[name]] + if missing: + raise ValueError(f"incomplete full-model execution: {missing}") + if self.candidate: + missing_nodes = [ + f"block.{i}.{node.name}" + for i in range(50) + for node in h3_block.NODES + if not self.candidate_calls[f"block.{i}.{node.name}"] + ] + if missing_nodes: + raise ValueError(f"candidate nodes did not execute: {missing_nodes}") + return dict( + candidate_calls=dict(self.candidate_calls), + execution_order=self.execution_order, + forward_order=self.forward_order, + backward_order=self.backward_order, + actual_backends=self.actual_backends, + calls=dict(self.calls), + backward_calls=dict(self.backward_calls), + bindings={ + n.name: dict( + rfc_row=n.binding.rfc_row, status=n.binding.status, backend=n.binding.backend + ) + for n in h3_block.NODES + }, + boundary_binding="provider/reference", + full_transformer_blocks=50, + refiner_layers=2, + ) + + +def require_candidate_cuda(ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + if ".cuda." not in backends[name]: + raise ValueError(f"candidate {name} fell back from its CUDA implementation") + return backends + + +@contextmanager +def full_model_execution( + model, *, candidate: bool, trace: FullModelTrace, ops=None, provider_nodes=() +): + """Temporarily bind each real block, preserving module parameters and restoring on error. + + ``provider_nodes`` (diagnostic hybrids only) names candidate graph nodes that run + their provider replay in every block; such a run localises drift and is never a + candidate result. + """ + unknown = set(provider_nodes) - {n.name for n in h3_block.NODES} + if unknown: + raise ValueError(f"unknown block nodes: {sorted(unknown)}") + boundaries = validate_full_structure(model) + originals, handles = [], [] + trace.candidate = candidate + if candidate: + ops = ops or h3_block.CandidateOps() + if hasattr(ops, "backends"): + trace.actual_backends = require_candidate_cuda(ops) + try: + for i, block in enumerate(model.transformer_blocks): + if candidate: + original = block.forward + originals.append((block, original)) + + def forward( + self, + hidden_states, + temb, + adaln_indices, + rotary_emb, + attention_mask=None, + _block_index=i, + ): + if attention_mask is not None: + raise ValueError( + "the qualified H3 candidate requires an unmasked packed sequence" + ) + weights = {h3_block.BLOCK + "." + n: p for n, p in self.named_parameters()} + ctx = dict( + weights=weights, + hidden=hidden_states, + temb=temb, + index=adaln_indices, + rope=rotary_emb, + provider_nodes=frozenset(provider_nodes), + node_observer=lambda name: trace.candidate_calls.update( + [f"block.{_block_index}.{name}"] + ), + ) + return h3_block.run_graph("candidate", ctx, ops)[h3_block.OUTPUT] + + block.forward = MethodType(forward, block) + + def hook(module, args, output, name=f"block.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for i, block in enumerate(model.token_refiner.refiner_blocks): + + def hook(module, args, output, name=f"refiner.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for name in boundaries: + + def hook(module, args, output, name=name): + trace.observe(name, output) + + handles.append(getattr(model, name).register_forward_hook(hook)) + yield trace + finally: + for handle in handles: + handle.remove() + for block, original in originals: + block.forward = original diff --git a/tests/models/minimax_h3/test_h3_full_model.py b/tests/models/minimax_h3/test_h3_full_model.py new file mode 100644 index 000000000..b10898105 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_full_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_block +from rl_engine.validation.models.h3_full_model import ( + FullModelTrace, + full_model_execution, + validate_full_structure, +) + + +class Block(torch.nn.Module): + def __init__(self, value): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(float(value))) + + def forward(self, x, *args): + return x + self.weight + + +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + self.transformer_blocks = torch.nn.ModuleList([Block(i) for i in range(50)]) + for name in [ + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + setattr(self, name, torch.nn.Identity()) + self.token_refiner.refiner_blocks = torch.nn.ModuleList([Block(50), Block(51)]) + + +def test_full_stack_uses_each_layers_own_parameters_and_keeps_gradients(monkeypatch): + model = Model() + seen = [] + + def run(mode, ctx, ops): + p = ctx["weights"][h3_block.BLOCK + ".weight"] + seen.append(p.item()) + return {h3_block.OUTPUT: ctx["hidden"] + p} + + monkeypatch.setattr(h3_block, "run_graph", run) + original = model.transformer_blocks[0].forward + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + x = torch.ones(1, requires_grad=True) + for block in model.transformer_blocks: + x = block(x, None, None, None) + x.sum().backward() + assert x.item() == 1 + sum(range(50)) + assert seen == list(range(50)) + assert all(b.weight.grad.item() == 1 for b in model.transformer_blocks) + assert model.transformer_blocks[0].forward == original + + +def test_repeated_block_zero_is_not_full_model(): + model = Model() + model.transformer_blocks[49] = model.transformer_blocks[0] + with pytest.raises(ValueError, match="distinct"): + validate_full_structure(model) + + +def test_partial_execution_cannot_publish_full_model_pass(): + trace = FullModelTrace() + trace.observe("block.0", torch.ones(1)) + with pytest.raises(ValueError, match="incomplete"): + trace.validate() + + +def test_binding_restored_after_failure(): + model = Model() + original = model.transformer_blocks[12].forward + with pytest.raises(RuntimeError, match="injected"): + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + raise RuntimeError("injected") + assert model.transformer_blocks[12].forward == original + + +def test_candidate_execution_requires_every_node(): + trace = FullModelTrace() + trace.candidate = True + for name in [f"block.{i}" for i in range(50)] + [ + "refiner.0", + "refiner.1", + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + trace.calls[name] = 1 + with pytest.raises(ValueError, match="candidate nodes did not execute"): + trace.validate() + + +def test_drift_order_follows_execution_instead_of_lexical_names(): + from tools.validation.models.h3_full_compare import drift_order + + report = { + "trace": { + "calls": {"block.2": 1, "block.10": 1}, + "backward_calls": {"block.10": 1, "block.2": 1}, + } + } + names = [ + "video", + "nodes.block.10", + "nodes.block.2", + "nodes.block.2.grad", + "nodes.block.10.grad", + ] + assert drift_order(report, names) == [ + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "nodes.block.2.grad", + "video", + ] + + +def test_reference_inventory_rejects_tampered_tensor(tmp_path): + import hashlib + import json + + from tools.validation.models.h3_evidence_integrity import verify_inventory + + tensor = tmp_path / "tensor.pt" + tensor.write_bytes(b"original") + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps({"files": {"tensor.pt": hashlib.sha256(b"original").hexdigest()}}) + ) + verify_inventory(tmp_path, manifest) + tensor.write_bytes(b"changed") + with pytest.raises(ValueError, match="changed"): + verify_inventory(tmp_path, manifest) + + +def test_cuda_guard_rejects_one_fallback_among_native_bindings(): + from rl_engine.validation.models.h3_full_model import require_candidate_cuda + + native = "rl_engine.backends.cuda.model_specific.minimax_h3.NativeOp" + ops = SimpleNamespace( + backends=lambda: {"projection": native, "norm": "torch.nn.RMSNorm", "gate": native} + ) + with pytest.raises(ValueError, match="candidate norm fell back"): + require_candidate_cuda(ops) + + +def test_drift_order_interleaves_parameter_and_activation_backward_hooks(): + from tools.validation.models.h3_full_compare import drift_order + + trace = FullModelTrace() + for name in ( + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "parameter.block.10.weight", + "nodes.block.2.grad", + ): + trace.record_event(name) + trace.record_event("nodes.block.2") # checkpoint recomputation is not a new first event + names = list(reversed(trace.execution_order)) + assert drift_order({"trace": {"execution_order": trace.execution_order}}, names) == ( + trace.execution_order + ) diff --git a/tools/validation/models/h3_candidate_prepare.py b/tools/validation/models/h3_candidate_prepare.py new file mode 100644 index 000000000..ae17b6600 --- /dev/null +++ b/tools/validation/models/h3_candidate_prepare.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fail closed on missing native H3 kernels and pin the candidate/reference identities.""" +import argparse +import json +import os +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from rl_engine.backends.extension import _C +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import require_candidate_cuda +from rl_engine.validation.models.h3_weights import load_h3_manifest, sha256_tensor +from tools.validation.models.h3_evidence_integrity import verify_inventory +from tools.validation.models.h3_full_model import digest, verify_checkpoint + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + a = p.parse_args() + a.output.mkdir() + manifest = Path("examples/models/minimax_h3/reference-manifest.json") + verify_inventory(a.reference_root, manifest) + verify_checkpoint( + a.checkpoint, Path("rl_engine/validation/models/h3_full_manifest.json"), all_files=True + ) + if _C is None or not hasattr(_C, "h3_timestep_sinusoid_forward"): + raise ValueError("candidate extension lacks H3 kernels") + ops = CandidateOps() + backends = require_candidate_cuda(ops) + # Inject a real eager CUDA fallback and require the same guard used by the + # full-model path to reject it before it can count as candidate execution. + original = ops.norm + try: + ops.norm = torch.nn.RMSNorm(128).cuda() + try: + require_candidate_cuda(ops) + except ValueError as exc: + if "candidate norm fell back" not in str(exc): + raise + else: + raise AssertionError("eager CUDA norm fallback was accepted") + finally: + ops.norm = original + require_candidate_cuda(ops) + extension = Path(_C.__file__).resolve() + if not extension.is_relative_to(Path.cwd().resolve()): + raise ValueError("candidate loaded an extension outside its clean source snapshot") + spec = load_h3_manifest() + index = json.loads( + (a.checkpoint / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + )["weight_map"] + weights = a.output / "weights" + weights.mkdir() + for section, filename in [ + ("tensors", "h3_conditioning.safetensors"), + ("block_tensors", "h3_block0.safetensors"), + ]: + values = {} + for shard in sorted({index[name] for name in spec[section]}): + with safe_open(str(a.checkpoint / "transformer" / shard), framework="pt") as handle: + for name, row in spec[section].items(): + if index[name] == shard: + value = handle.get_tensor(name).contiguous() + if sha256_tensor(value) != row["sha256"]: + raise ValueError(f"block weight identity changed: {name}") + values[name] = value + save_file(values, str(weights / filename)) + report = dict( + extension=str(extension), + extension_sha256=digest(extension), + torch=torch.__version__, + backends=backends, + negative_controls={"eager_cuda_norm_fallback_rejected": True}, + reference_manifest_sha256=digest(manifest), + build={ + key: os.environ.get(key) + for key in ("CUDA_HOME", "MAX_JOBS", "KERNEL_ALIGN_USE_FAST_MATH") + }, + ) + (a.output / "identity.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_compare.py b/tools/validation/models/h3_full_compare.py new file mode 100644 index 000000000..6f1645e67 --- /dev/null +++ b/tools/validation/models/h3_full_compare.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Freeze full-H3 reference error bounds, or evaluate a candidate without changing them.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from tools.validation.models.h3_full_model import digest + + +def tensors(root): + report = json.loads((root / "report.json").read_text()) + if digest(root / "tensors.pt") != report["tensors_sha256"]: + raise ValueError("full-model tensor evidence changed") + payload = torch.load(root / "tensors.pt", map_location="cpu", weights_only=True) + yield "video", payload["video"] + yield "audio", payload["audio"] + for section in ("nodes", "input_gradients"): + for name, value in sorted(payload[section].items()): + yield section + "." + name, value + for name, row in sorted(json.loads((root / "gradients.json").read_text()).items()): + path = root / row["path"] + if digest(path) != row["sha256"]: + raise ValueError(f"gradient artifact changed: {path}") + yield "parameter." + name, torch.load(path, map_location="cpu", weights_only=True) + + +def compare(left, right): + a, b = tensors(left), tensors(right) + from itertools import zip_longest + + for x, y in zip_longest(a, b): + if x is None or y is None or x[0] != y[0]: + raise ValueError("different full-model tensor inventories") + name, value = x + other = y[1] + if ( + value.shape != other.shape + or not torch.isfinite(value).all() + or not torch.isfinite(other).all() + ): + raise ValueError(f"nonfinite/shape mismatch: {name}") + delta = value.float() - other.float() + yield name, dict( + max_abs=delta.abs().max().item(), + relative_l2=(delta.double().norm() / value.double().norm().clamp_min(1e-12)).item(), + ) + + +def drift_order(report, names): + trace = report["trace"] + if "execution_order" in trace: + observed = trace["execution_order"] + return [name for name in observed if name in names] + [ + name for name in names if name not in observed + ] + forward = trace.get("forward_order", list(trace["calls"])) + backward = trace.get("backward_order", list(trace["backward_calls"])) + ordered = ["nodes." + name for name in forward] + ordered += ["nodes." + name + ".grad" for name in backward] + ordered += [name for name in names if name not in ordered] + return [name for name in ordered if name in names] + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--reference", type=Path, required=True) + p.add_argument("--other", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--thresholds", type=Path) + p.add_argument("--reference-manifest", type=Path) + p.add_argument("--reference-root", type=Path) + args = p.parse_args() + if args.output.exists(): + raise ValueError("never overwrite calibration or failed evidence") + from tools.validation.models.h3_evidence_integrity import verify_inventory + + if args.thresholds and not (args.reference_root and args.reference_manifest): + raise ValueError("candidate comparison requires pinned reference inventory") + if args.reference_manifest: + verify_inventory(args.reference_root, args.reference_manifest) + ref = json.loads((args.reference / "report.json").read_text()) + other = json.loads((args.other / "report.json").read_text()) + if ref["mode"] != "fp32-reference" or ref["status"] != "execution_complete_acceptance_pending": + raise ValueError("independent completed FP32 reference required") + if other["status"] != "execution_complete_acceptance_pending": + raise ValueError("other execution is incomplete") + reference_sha = digest(args.reference / "report.json") + if args.thresholds: + contract = json.loads(args.thresholds.read_text()) + if other["mode"] != "candidate" or contract["reference_report_sha256"] != reference_sha: + raise ValueError("wrong candidate/reference contract") + if other.get("thresholds_sha256") != digest(args.thresholds): + raise ValueError("thresholds were not frozen before candidate execution") + elif other["mode"] != "bf16-provider": + raise ValueError("only independent provider can establish thresholds") + rows = dict(compare(args.reference, args.other)) + report = dict( + reference_report_sha256=reference_sha, + other_report_sha256=digest(args.other / "report.json"), + ) + if args.thresholds: + if rows.keys() != contract["bounds"].keys(): + raise ValueError("threshold tensor inventory differs") + failures = [ + name + for name in drift_order(other if "execution_order" in other["trace"] else ref, rows) + if any(rows[name][k] > contract["bounds"][name][k] for k in rows[name]) + ] + hybrid = other.get("diagnostic_provider_nodes") or [] + report.update( + status="diagnostic_only" if hybrid else ("failed" if failures else "passed"), + diagnostic_provider_nodes=hybrid, + errors=rows, + failures=failures, + first_drift=failures[0] if failures else None, + thresholds_sha256=digest(args.thresholds), + ) + else: + # Committed before candidate execution. Forward and gradient tensors use + # separate floors; the existing block-0 golden contract is unchanged. + bounds = { + name: { + key: 1.5 * value + + ( + 5e-3 + if ".grad" in name + or name.startswith("parameter.") + or name.startswith("input_gradients.") + else 1e-3 + ) + for key, value in row.items() + } + for name, row in rows.items() + } + report.update( + status="frozen_reference_bounds", + rule="1.5 * independent BF16 provider error + fixed floor", + bounds=bounds, + errors=rows, + ) + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.thresholds and failures and not report.get("diagnostic_provider_nodes"): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_ensemble_rule.py b/tools/validation/models/h3_full_ensemble_rule.py new file mode 100644 index 000000000..72c397533 --- /dev/null +++ b/tools/validation/models/h3_full_ensemble_rule.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Full-model acceptance from an ensemble of independent BF16 provider configurations. + +The single-configuration rule (``1.5 * one provider's error + floor``) accepts only the +configuration it was frozen from: providers that differ only in their SDPA backend or +BLAS library fail it on tens of tensors. This rule is calibrated on several provider +configurations, compared with the same frozen FP32 reference, and contains no candidate +data: + +* E[t] is the largest ``relative_l2`` error any configuration has on tensor ``t``; +* the bound is ``k * E[t] + floor`` (floors as in the single-configuration rule); +* ``k`` is the smallest factor, rounded up to 0.1, for which every configuration held + out passes against the maximum of the others (leave-one-out); +* ``max_abs``, set by one element, is reported but not gated: held-out configurations + exceed the others' maximum by several times on it. + + h3_full_ensemble_rule.py calibrate --provider NAME=comparison.json ... \\ + --output rule.json + h3_full_ensemble_rule.py check --rule rule.json --candidate comparison.json \\ + --output out.json +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path + + +def _floor(name: str) -> float: + gradient = ( + ".grad" in name or name.startswith("parameter.") or name.startswith("input_gradients.") + ) + return 5e-3 if gradient else 1e-3 + + +def _digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def calibrate(args) -> None: + """Build the ensemble bounds and the leave-one-out factor from provider comparisons.""" + + providers, sources = {}, {} + for item in args.provider: + name, _, path = item.partition("=") + report = json.loads(Path(path).read_text()) + if report.get("diagnostic_provider_nodes"): + raise ValueError(f"{name} is a hybrid diagnostic, not a provider configuration") + providers[name] = report["errors"] + sources[name] = { + "sha256": _digest(Path(path)), + "reference": report["reference_report_sha256"], + } + if len({s["reference"] for s in sources.values()}) != 1: + raise ValueError("provider comparisons use different FP32 references") + # Configurations whose errors are identical to an earlier one are one configuration. + distinct = {} + for name in sorted(providers): + if not any(providers[name] == providers[kept] for kept in distinct): + distinct[name] = providers[name] + if len(distinct) < 3: + raise ValueError("need at least three distinct provider configurations") + names = sorted(next(iter(distinct.values()))) + loo = {} + for held in distinct: + others = [distinct[o] for o in distinct if o != held] + worst = 0.0 + for t in names: + top = max(o[t]["relative_l2"] for o in others) + if top > 0: + worst = max(worst, (distinct[held][t]["relative_l2"] - _floor(t)) / top) + loo[held] = worst + factor = math.ceil(max(loo.values()) * 10) / 10 + ensemble = {t: max(d[t]["relative_l2"] for d in distinct.values()) for t in names} + rule = { + "status": "frozen_ensemble_rule", + "rule": "relative_l2 <= k * max over provider configurations + floor; max_abs reported", + "factor_k": factor, + "leave_one_out_worst_ratio": loo, + "providers": sorted(distinct), + "duplicates_dropped": sorted(set(providers) - set(distinct)), + "sources": sources, + "reference_report_sha256": next(iter(sources.values()))["reference"], + "ensemble_relative_l2": ensemble, + "ensemble_max_abs": {t: max(d[t]["max_abs"] for d in distinct.values()) for t in names}, + "bounds_relative_l2": {t: factor * ensemble[t] + _floor(t) for t in names}, + } + args.output.write_text(json.dumps(rule, indent=2) + "\n") + print( + f"k = {factor} from leave-one-out {loo}; " + f"{len(distinct)} configurations, {len(names)} tensors" + ) + + +def check(args) -> None: + """Apply a frozen ensemble rule to one candidate comparison; exit 1 on any failure.""" + + rule = json.loads(args.rule.read_text()) + report = json.loads(args.candidate.read_text()) + if report["reference_report_sha256"] != rule["reference_report_sha256"]: + raise ValueError("candidate was compared with a different FP32 reference") + if report.get("diagnostic_provider_nodes"): + raise ValueError("a hybrid diagnostic is not a candidate result") + errors = report["errors"] + if errors.keys() != rule["bounds_relative_l2"].keys(): + raise ValueError("tensor inventory differs from the rule") + failures = [ + t for t in sorted(errors) if errors[t]["relative_l2"] > rule["bounds_relative_l2"][t] + ] + ratios = {t: errors[t]["relative_l2"] / rule["bounds_relative_l2"][t] for t in errors} + max_abs_ratio = { + t: errors[t]["max_abs"] / rule["ensemble_max_abs"][t] + for t in errors + if rule["ensemble_max_abs"][t] + } + out = { + "status": "failed" if failures else "passed", + "rule_sha256": _digest(args.rule), + "candidate_comparison_sha256": _digest(args.candidate), + "failures": failures, + "worst_relative_l2_over_bound": sorted(ratios.items(), key=lambda x: -x[1])[:10], + "outputs": {t: errors[t] for t in ("video", "audio") if t in errors}, + "max_abs_over_ensemble_max_reported": sorted(max_abs_ratio.items(), key=lambda x: -x[1])[ + :10 + ], + } + args.output.write_text(json.dumps(out, indent=2) + "\n") + print( + f"{out['status']}: {len(failures)} failures; worst {out['worst_relative_l2_over_bound'][0]}" + ) + if failures: + raise SystemExit(1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="cmd", required=True) + cal = sub.add_parser("calibrate") + cal.add_argument("--provider", action="append", required=True, help="name=comparison.json") + cal.add_argument("--output", type=Path, required=True) + chk = sub.add_parser("check") + chk.add_argument("--rule", type=Path, required=True) + chk.add_argument("--candidate", type=Path, required=True) + chk.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + calibrate(args) if args.cmd == "calibrate" else check(args) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_model.py b/tools/validation/models/h3_full_model.py new file mode 100644 index 000000000..09a8b07dc --- /dev/null +++ b/tools/validation/models/h3_full_model.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Full pinned H3 forward/backward evidence; reference calibration precedes candidates. + +The small packed input checks the entire training graph, not generation quality. +FP32 reference and BF16 provider runs never instantiate candidate operators. +""" +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import subprocess +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_cases import h3_fl2va_layout +from rl_engine.validation.models.h3_full_model import ( + MODEL_REVISION, + PROVIDER_REVISION, + FullModelTrace, + full_model_execution, +) + + +def digest(path): + h = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024**2), b""): + h.update(block) + return h.hexdigest() + + +def verify_checkpoint(root, manifest, *, all_files=False): + spec = json.loads(manifest.read_text()) + if spec["revision"] != MODEL_REVISION: + raise ValueError("wrong H3 revision") + needed = [r for r in spec["files"] if (all_files or r["path"].startswith("transformer/"))] + if not needed: + raise ValueError("no pinned transformer shards") + for record in needed: + path = root / record["path"] + if path.stat().st_size != record["size"] or digest(path) != record["sha256"]: + raise ValueError(f"checkpoint identity mismatch: {path}") + return spec + + +def packed_inputs(device): + layout = h3_fl2va_layout( + num_text_tokens=4, + num_latent_frames=1, + latent_height=4, + latent_width=4, + num_audio_latents=2, + device=device, + ) + generator = torch.Generator().manual_seed(420539) + values = dict( + hidden_states=(1, 4, 96), audio_hidden_states=(1, 4, 32), encoder_hidden_states=(1, 4, 5120) + ) + result = { + name: torch.randn(shape, generator=generator).to(device).requires_grad_() + for name, shape in values.items() + } + tags = layout["token_tags"] + return dict( + **result, + **layout, + timestep=torch.tensor([0.0, 500.0], device=device), + video_indices=(tags == 0).nonzero().flatten(), + text_indices=(tags == 1).nonzero().flatten(), + audio_indices=(tags == 2).nonzero().flatten(), + return_dict=False, + ) + + +# Provider passes that change only which kernels compute the same diffusers graph, to +# measure how much an independent BF16 provider's errors move on their own. Each forces +# one choice; an unavailable kernel raises instead of falling back. +PROVIDER_VARIANTS = ("default", "sdpa-flash", "sdpa-efficient", "sdpa-cudnn", "blas-cublaslt") + + +def provider_variant_context(variant: str): + """Context manager forcing one SDPA backend, or a no-op (BLAS choice is set globally).""" + + from contextlib import nullcontext + + from torch.nn.attention import SDPBackend, sdpa_kernel + + backends = { + "sdpa-flash": SDPBackend.FLASH_ATTENTION, + "sdpa-efficient": SDPBackend.EFFICIENT_ATTENTION, + "sdpa-cudnn": SDPBackend.CUDNN_ATTENTION, + } + if variant == "blas-cublaslt": + torch.backends.cuda.preferred_blas_library("cublaslt") + return sdpa_kernel([backends[variant]]) if variant in backends else nullcontext() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--manifest", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--mode", choices=["fp32-reference", "bf16-provider", "candidate"], required=True + ) + parser.add_argument("--thresholds", type=Path) + parser.add_argument( + "--provider-variant", + choices=PROVIDER_VARIANTS, + default="default", + help="bf16-provider only: force one SDPA backend or cuBLASLt (variability study)", + ) + parser.add_argument( + "--provider-nodes", + default="", + help="candidate only, diagnostic: comma list of block nodes that run the provider " + "replay in every block (the result is labelled diagnostic, never acceptance)", + ) + args = parser.parse_args() + provider_nodes = tuple(n for n in args.provider_nodes.split(",") if n) + if provider_nodes and args.mode != "candidate": + parser.error("--provider-nodes applies to --mode candidate only") + if args.mode == "candidate" and not args.thresholds: + parser.error("candidate execution requires previously frozen calibration thresholds") + if args.provider_variant != "default" and args.mode != "bf16-provider": + parser.error("--provider-variant applies to --mode bf16-provider only") + if args.thresholds: + thresholds = json.loads(args.thresholds.read_text()) + if thresholds.get("status") != "frozen_reference_bounds" or not thresholds.get( + "reference_report_sha256" + ): + raise ValueError("candidate requires immutable independent FP32 calibration") + args.output.mkdir(parents=True, exist_ok=False) + report = dict( + status="running", + mode=args.mode, + provider_variant=args.provider_variant, + diagnostic_provider_nodes=list(provider_nodes), + job_id=os.environ.get("SLURM_JOB_ID"), + model_revision=MODEL_REVISION, + command=list(os.sys.argv), + ) + report_path = args.output / "report.json" + started = time.perf_counter() + try: + commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("evidence requires a clean source commit") + report["commit"] = commit + verify_checkpoint(args.checkpoint, args.manifest) + import diffusers + from diffusers import MiniMaxH3Transformer3DModel + + provider_root = Path(diffusers.__file__).resolve().parents[2] + provider_commit = subprocess.check_output( + ["git", "-C", str(provider_root), "rev-parse", "HEAD"], text=True + ).strip() + if provider_commit != PROVIDER_REVISION: + raise ValueError("diffusers provider revision changed") + if subprocess.check_output( + ["git", "-C", str(provider_root), "status", "--porcelain"], text=True + ).strip(): + raise ValueError("diffusers provider source is dirty") + report.update( + provider_revision=provider_commit, + torch=torch.__version__, + diffusers=diffusers.__version__, + device=torch.cuda.get_device_name(), + checkpoint_manifest_sha256=digest(args.manifest), + ) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.manual_seed(420539) + dtype = torch.float32 if args.mode == "fp32-reference" else torch.bfloat16 + model, loading = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=dtype, + local_files_only=True, + low_cpu_mem_usage=True, + output_loading_info=True, + ) + report["checkpoint_loading"] = loading + if any( + loading.get(key) + for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") + ): + raise ValueError("full H3 checkpoint did not load every tensor exactly") + model = model.to("cuda").eval() + model.enable_gradient_checkpointing() + inputs = packed_inputs("cuda") + trace = FullModelTrace() + trace.capture = True + gradients = {} + handles = [] + # Persist and release each completed parameter gradient. This preserves + # the complete backward chain while bounding resident FP32 gradient memory. + directory = args.output / "gradients" + directory.mkdir() + for i, (name, parameter) in enumerate(model.named_parameters()): + + def save_grad(p, name=name, i=i): + trace.record_event("parameter." + name) + grad = p.grad.detach().cpu() + if not torch.isfinite(grad).all(): + raise ValueError(f"non-finite parameter gradient: {name}") + path = directory / f"{i:04}.pt" + torch.save(grad, path) + gradients[name] = dict( + path=str(path.relative_to(args.output)), + shape=list(grad.shape), + dtype=str(grad.dtype), + sha256=digest(path), + ) + p.grad = None + + handles.append(parameter.register_post_accumulate_grad_hook(save_grad)) + for name in ("hidden_states", "audio_hidden_states", "encoder_hidden_states"): + + def record_input(grad, name=name): + trace.record_event("input_gradients." + name) + return grad + + handles.append(inputs[name].register_hook(record_input)) + torch.cuda.reset_peak_memory_stats() + begin = time.perf_counter() + report["preferred_blas_library"] = None + with ( + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + provider_nodes=provider_nodes, + ), + provider_variant_context(args.provider_variant), + ): + report["preferred_blas_library"] = str(torch.backends.cuda.preferred_blas_library()) + with torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = model(**inputs) + if not torch.isfinite(video).all() or not torch.isfinite(audio).all(): + raise ValueError("non-finite full-model prediction") + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - begin + begin = time.perf_counter() + loss = video.float().square().mean() + audio.float().square().mean() + loss.backward() + torch.cuda.synchronize() + backward_seconds = time.perf_counter() - begin + for h in handles: + h.remove() + if set(gradients) != set(dict(model.named_parameters())): + raise ValueError("missing full-model parameter gradients") + input_grads = {} + for name in ("hidden_states", "audio_hidden_states", "encoder_hidden_states"): + grad = inputs[name].grad + if grad is None or not torch.isfinite(grad).all(): + raise ValueError(f"missing/non-finite input gradient: {name}") + input_grads[name] = grad.cpu() + torch.save( + dict( + video=video.detach().cpu(), + audio=audio.detach().cpu(), + input_gradients=input_grads, + nodes=trace.snapshots, + ), + args.output / "tensors.pt", + ) + (args.output / "gradients.json").write_text(json.dumps(gradients, indent=2) + "\n") + report.update( + trace=trace.validate(backward=True), + forward_seconds=forward_seconds, + backward_seconds=backward_seconds, + peak_allocated_bytes=torch.cuda.max_memory_allocated(), + peak_reserved_bytes=torch.cuda.max_memory_reserved(), + loss=loss.item(), + parameter_gradients=len(gradients), + tensors_sha256=digest(args.output / "tensors.pt"), + ) + # Numerical acceptance is a separate comparison against immutable reference + # artifacts; merely finishing all layers is never a numerical PASS. + report["status"] = "execution_complete_acceptance_pending" + if args.thresholds: + report["thresholds_sha256"] = digest(args.thresholds) + except BaseException as exc: + report.update(status="failed", error=f"{type(exc).__name__}: {exc}") + raise + finally: + report["total_seconds"] = time.perf_counter() - started + report_path.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_provider_spread.py b/tools/validation/models/h3_full_provider_spread.py new file mode 100644 index 000000000..733fc2d2b --- /dev/null +++ b/tools/validation/models/h3_full_provider_spread.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""How far independent BF16 provider passes spread, and where the candidate falls. + +Inputs are ``h3_full_compare`` outputs against the same frozen FP32 reference: one +for the candidate and one per provider variant (``h3_full_model --mode bf16-provider +--provider-variant ...``). For every tensor and metric (``max_abs``, ``relative_l2``) +it reports the providers' minimum and maximum error, the candidate's error, and +whether the candidate lies above the providers' maximum. It also applies the frozen +bounds to each provider variant, which shows whether the bounds hold for providers +themselves. Diagnostic only: it changes no bound and accepts nothing. + + python tools/validation/models/h3_full_provider_spread.py \\ + --thresholds examples/models/minimax_h3/full-model-thresholds.json \\ + --candidate candidate-comparison.json \\ + --provider default=default.json --provider sdpa-flash=flash.json ... \\ + --output spread.json +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +METRICS = ("max_abs", "relative_l2") + + +def main() -> None: + """Load the comparisons, compute per-tensor spreads, and write a summary.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--thresholds", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--provider", action="append", required=True, help="name=comparison.json") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + bounds = json.loads(args.thresholds.read_text())["bounds"] + candidate = json.loads(args.candidate.read_text())["errors"] + providers = {} + for item in args.provider: + name, _, path = item.partition("=") + providers[name] = json.loads(Path(path).read_text())["errors"] + names = sorted(bounds) + for label, errors in [("candidate", candidate), *providers.items()]: + missing = set(names) - set(errors) + if missing: + raise ValueError(f"{label} lacks {len(missing)} tensors, e.g. {sorted(missing)[:3]}") + + tensors = {} + summary = { + "tensors": len(names), + "providers": sorted(providers), + "candidate_failures": {m: 0 for m in METRICS}, + "candidate_above_provider_max": {m: 0 for m in METRICS}, + "provider_failures": {p: {m: 0 for m in METRICS} for p in providers}, + } + for name in names: + entry = {} + for m in METRICS: + values = [providers[p][name][m] for p in providers] + cand = candidate[name][m] + bound = bounds[name][m] + entry[m] = { + "bound": bound, + "candidate": cand, + "provider_min": min(values), + "provider_max": max(values), + "candidate_over_provider_max": cand / max(values) if max(values) else None, + } + summary["candidate_failures"][m] += cand > bound + summary["candidate_above_provider_max"][m] += cand > max(values) + for p in providers: + summary["provider_failures"][p][m] += providers[p][name][m] > bound + tensors[name] = entry + worst = sorted( + names, + key=lambda n: tensors[n]["max_abs"]["candidate_over_provider_max"] or 0.0, + reverse=True, + )[:20] + summary["worst_candidate_over_provider_max"] = [ + {"tensor": n, **tensors[n]["max_abs"]} for n in worst + ] + args.output.write_text(json.dumps({"summary": summary, "tensors": tensors}, indent=2) + "\n") + print(json.dumps(summary, indent=2)[:4000]) + + +if __name__ == "__main__": + main() From 6fe346df65ab48b047d83d09fb3efd8f632a9b9b Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Fri, 23 Oct 2026 14:16:28 +0000 Subject: [PATCH 43/62] [WS1][CUDA][MiniMax-H3] generation gate: 50-block generation compare and resume check (RFC #420, #539 split 3/3) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/generation-compare.sbatch | 26 ++ .../minimax_h3/generation-contract.json | 18 + .../models/minimax_h3/generation-inputs.json | 8 + .../models/minimax_h3/generation-prompt.txt | 7 + .../minimax_h3/generation-resume-check.sbatch | 77 ++++ examples/models/minimax_h3/generation.sbatch | 56 +++ .../validation/models/h3_generation_resume.py | 320 +++++++++++++++ .../minimax_h3/test_h3_generation_resume.py | 384 ++++++++++++++++++ tools/validation/models/h3_full_generation.py | 365 +++++++++++++++++ .../validation/models/h3_generation_adopt.py | 63 +++ .../models/h3_generation_compare.py | 150 +++++++ 11 files changed, 1474 insertions(+) create mode 100644 examples/models/minimax_h3/generation-compare.sbatch create mode 100644 examples/models/minimax_h3/generation-contract.json create mode 100644 examples/models/minimax_h3/generation-inputs.json create mode 100644 examples/models/minimax_h3/generation-prompt.txt create mode 100644 examples/models/minimax_h3/generation-resume-check.sbatch create mode 100644 examples/models/minimax_h3/generation.sbatch create mode 100644 rl_engine/validation/models/h3_generation_resume.py create mode 100644 tests/models/minimax_h3/test_h3_generation_resume.py create mode 100644 tools/validation/models/h3_full_generation.py create mode 100644 tools/validation/models/h3_generation_adopt.py create mode 100644 tools/validation/models/h3_generation_compare.py diff --git a/examples/models/minimax_h3/generation-compare.sbatch b/examples/models/minimax_h3/generation-compare.sbatch new file mode 100644 index 000000000..c188e3c4a --- /dev/null +++ b/examples/models/minimax_h3/generation-compare.sbatch @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Quality comparison of the three candidate seeds with the provider reference (CPU only). +# Submit after the generation array finished successfully (--dependency=afterok:). +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=8 +#SBATCH --mem=128G +#SBATCH --time=06:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +reference_root=$3 +result_root=$4 +extra_dependencies=$5 # overlay with scikit-image / scipy (the reference job used it too) +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +commit=$(git rev-parse HEAD) +export PYTHONPATH="$source_dir:$extra_dependencies" +"$python_bin" -m tools.validation.models.h3_generation_compare \ + --reference "$reference_root/generation-reference" \ + --candidate "$result_root/generation-${commit:0:8}" \ + --contract examples/models/minimax_h3/generation-contract.json \ + --reference-root "$reference_root" \ + --reference-manifest examples/models/minimax_h3/reference-manifest.json \ + --output "$result_root/generation-comparison-${commit:0:8}.json" diff --git a/examples/models/minimax_h3/generation-contract.json b/examples/models/minimax_h3/generation-contract.json new file mode 100644 index 000000000..7c29221c6 --- /dev/null +++ b/examples/models/minimax_h3/generation-contract.json @@ -0,0 +1,18 @@ +{ + "schema": 1, + "model_revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "provider_revision": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "workflow": "fl2va", + "seeds": [42, 43, 44], + "num_frames": 124, + "height": 768, + "width": 1344, + "fps": 24, + "sampling_rate": 32000, + "num_inference_steps": 50, + "preflight": {"num_frames": 124, "height": 128, "width": 128, "num_inference_steps": 2}, + "quality_bounds": {"minimum_psnr_db": 30, "minimum_ssim": 0.99, "maximum_temporal_rmse": 0.01, "maximum_audio_relative_l2": 0.05}, + "quality_bounds_policy": "Conservative absolute gates fixed before candidate generation; a failure remains a failure. Reference repeats must first satisfy these same bounds.", + "attention_query_chunk": 256, + "status": "unqualified_until_full_reference_and_candidate_runs_pass" +} diff --git a/examples/models/minimax_h3/generation-inputs.json b/examples/models/minimax_h3/generation-inputs.json new file mode 100644 index 000000000..ea62adc23 --- /dev/null +++ b/examples/models/minimax_h3/generation-inputs.json @@ -0,0 +1,8 @@ +{ + "image_sha256": "352e959e08b886406d3766cee25ebe7bbf679e6d53eeffbd1c607b8dc93e2068", + "image_url": "https://cdn.hailuoai.com/prod/hailuo_demo/testsets/H3_AA_I2VA/gallery/sr_v17_variants_seed42_43_20260724/inputs/4a3a90bf9100_KDmcbkhzYo5sjjxr9FqcVmWVnzb.png", + "prompt_reused_verbatim": true, + "prompt_sha256": "fba3868d31e9911a81d60c191261ec2192c827bcaa25c937660bf35aeb0b1673", + "source_revision": "bbee488da3baf65817f4f8826b1f6e89cc28fda7", + "source_suite": "tests/e2e/accuracy/minimax_h3/test_minimax_h3_i2va_ref2va_similarity.py" +} diff --git a/examples/models/minimax_h3/generation-prompt.txt b/examples/models/minimax_h3/generation-prompt.txt new file mode 100644 index 000000000..ba5be1ceb --- /dev/null +++ b/examples/models/minimax_h3/generation-prompt.txt @@ -0,0 +1,7 @@ +For the target video, at 0.00 seconds into the target video, (from [Shot 1]) is fully referenced. + +integrated_multimodal_description: [Shot 1] This is a live-action, cinematic shot with a shallow depth of field. The camera holds a perfectly static shot throughout the entire eight-second duration, capturing a cozy family gathering in a traditional Japanese dining room. The scene opens with a large, intricately patterned blue and white ceramic bowl of ramen in the immediate foreground, rendered in crisp, sharp focus. The bowl sits on a smooth, polished long wooden table. Inside the bowl, a rich, oily golden-brown broth surrounds yellow wavy noodles, topped with two thick, round slices of chashu pork featuring visible fat marbling and a distinct spiral meat pattern. A generous mound of freshly chopped, bright green scallions rests in the center, and a crisp, dark green rectangular sheet of nori seaweed is tucked into the right edge. To the left of the bowl, a pair of light brown wooden chopsticks rests horizontally on a small, dark rectangular chopstick rest, near a small cylindrical ceramic teacup with blue painted patterns. On the right side of the table, a spherical paper lantern with a ribbed bamboo frame sits on a black wooden base. In the background, a large family of seven is gathered around the table, initially appearing as a soft, blurred presence. Behind them, traditional Japanese sliding shoji screens with wooden lattice frames are open, revealing a bright outdoor scene with lush green trees. Early in the clip, the thick, white steam rising from the hot ramen broth immediately intensifies, billowing upwards in thick, swirling clouds that dance continuously above the bowl. As the clip progresses into the middle seconds, the camera maintains its static position while the focus begins a deliberate, smooth shift deeper into the room. The foreground ramen bowl, its vibrant ingredients, and the rising steam gradually soften into a hazy, out-of-focus blur. Simultaneously, the family members in the background come into sharp, detailed clarity. The heavy steam continues to rise from the foreground, creating a dynamic, translucent veil between the camera and the family. With the focus now firmly locked on the background, the vibrant family dinner comes alive. The man in the dark navy blue long-sleeved shirt on the left leans forward, his mouth moving animatedly in a silent exchange. The young girl in the crisp white short-sleeved t-shirt beside him smiles brightly, looking toward the center of the table. The woman on the far left, wearing a soft light blue long-sleeved blouse, turns her head slightly, smiling gently. Across the table, the woman in the light grey button-down shirt smiles broadly, her eyes crinkling, as she rests her hands near her plate. The woman in the dark grey top further back uses her wooden chopsticks to pick up a small piece of food from a central ceramic dish filled with bright red pickled vegetables. The woman in the center back in the light grey sweater smiles gently, her hands clasped softly in front of her, observing the interaction. Throughout the remainder of the clip, the family continues their lively physical interaction, their mouths moving in continuous, silent cadences of conversation, while the thick, white steam from the blurred ramen bowl in the foreground never stops rising, adding a comforting atmosphere to the warm gathering. + +overall_soundscape: The soundscape begins with a quiet room tone mixed with the faint, airy rustle of the thick steam billowing from the hot ramen bowl in the foreground, accompanied by the subtle, continuous hissing and bubbling of the rich broth. As the visual focus shifts deeper into the room, the physical sounds of the bustling family dinner become dominant in the foreground. The clear, sharp clinking of ceramic bowls and wooden chopsticks touching plates is clearly heard as the family members reach for food. This is followed by the faint, muffled thud of a cup being set down on the smooth wooden table, and the subtle, rhythmic rustle of cotton and wool clothing as the family members lean forward and gesture, perfectly capturing the lively, physical atmosphere of the shared meal. + +non_diegetic_music: A gentle, heartwarming acoustic guitar melody plays softly in the background, accompanied by the subtle, resonant notes of a traditional Japanese koto. The music maintains a slow, comforting tempo that enhances the cozy, nostalgic, and joyful atmosphere of the family gathering. diff --git a/examples/models/minimax_h3/generation-resume-check.sbatch b/examples/models/minimax_h3/generation-resume-check.sbatch new file mode 100644 index 000000000..be544b38f --- /dev/null +++ b/examples/models/minimax_h3/generation-resume-check.sbatch @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Resume check on the real pipeline at preflight size (128x128, 6 sampling steps, 5 transformer +# calls): one run without interruption, and two preempted runs that are resumed -- "cut" +# after 2, 4 and 5 calls (the last one after the final step, before decoding, so the final +# attempt has nothing left to compute) and "cutverify" after 2 calls with --verify-last-step. +# Every step file, the initial inputs and the decoded media must be byte-equal to "whole". +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=04:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +reference_image=$4 +provider_source=$5 +extra_dependencies=$6 +result_root=$7 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 +export PATH="$CUDA_HOME/bin:$PATH" +out="$result_root/resume-check-$(git rev-parse --short=8 HEAD)" +mkdir "$out" +gen() { + "$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --preflight --preflight-steps 6 "$@" +} +gen --output "$out/whole" > "$out/whole.log" 2>&1 +gen --output "$out/cut" --stop-after-step 2 > "$out/cut-1.log" 2>&1 || true +gen --output "$out/cut" --resume --stop-after-step 4 > "$out/cut-2.log" 2>&1 || true +gen --output "$out/cut" --resume --stop-after-step 5 > "$out/cut-3.log" 2>&1 || true +gen --output "$out/cut" --resume > "$out/cut-4.log" 2>&1 +gen --output "$out/cutverify" --stop-after-step 2 > "$out/cutverify-1.log" 2>&1 || true +gen --output "$out/cutverify" --resume --verify-last-step > "$out/cutverify-2.log" 2>&1 +"$python_bin" - "$out" <<'PY' +import json +import sys +from pathlib import Path + +root = Path(sys.argv[1]) +whole = root / "whole" / "42" +whole_report = json.loads((whole / "report.json").read_text()) +names = sorted( + p.name + for p in whole.iterdir() + if p.name.startswith("step-") and p.suffix == ".pt" + or p.name in ("initial.pt", "video.npy", "audio.npy") +) +result = {"compared_files": len(names), "steps": whole_report["steps"], "runs": {}} +for name in ("cut", "cutverify"): + directory = root / name / "42" + report = json.loads((directory / "report.json").read_text()) + result["runs"][name] = { + "differing": [n for n in names if (whole / n).read_bytes() != (directory / n).read_bytes()], + "resume": {k: v for k, v in report["resume"].items() if k != "attempts"}, + "attempts": len(report["resume"]["attempts"]), + "trace_calls_match": report["trace"]["calls"].keys() == whole_report["trace"]["calls"].keys() + and report["trace"]["candidate_calls"].keys() + == whole_report["trace"]["candidate_calls"].keys(), + "failure_markers": sorted(p.name for p in (root / name).glob("failure*.json")), + } +(root / "resume-check.json").write_text(json.dumps(result, indent=2) + "\n") +print(json.dumps(result, indent=2)) +cut, verify = result["runs"]["cut"], result["runs"]["cutverify"] +assert not cut["differing"] and not verify["differing"], result +assert cut["resume"]["real_forward_calls"] == 0, "the final cut attempt should replay everything" +assert cut["resume"]["recomputed_and_verified_step"] is None +assert verify["resume"]["recomputed_and_verified_step"] is not None +assert cut["trace_calls_match"] and verify["trace_calls_match"] +PY diff --git a/examples/models/minimax_h3/generation.sbatch b/examples/models/minimax_h3/generation.sbatch new file mode 100644 index 000000000..ca1a97e0e --- /dev/null +++ b/examples/models/minimax_h3/generation.sbatch @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +# Full-resolution candidate generation, one seed per array task (--array=0-2). Preemption +# requeues the task; the generation continues from its saved steps (--resume). +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +#SBATCH --requeue +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +gate_job=$4 # job id of the candidate.sbatch run whose gates passed for this commit +result_root=$5 +reference_image=$6 +provider_source=$7 +extra_dependencies=$8 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +commit=$(git rev-parse HEAD) +# Only a passing numerical comparison of the same model code permits full generation: the gates +# passed at a commit whose model paths (h3_generation_resume.MODEL_PATHS) are byte-identical to +# this checkout's. The same commit qualifies trivially. +gate_dir= +for attempt in "$result_root"/job-"$gate_job"-attempt-*; do + if [[ -f $attempt/commit.txt ]] && "$python_bin" - "$attempt/commit.txt" "$attempt/full-model-acceptance.json" \ + 2> /dev/null <<'PY' +import json +import sys + +from rl_engine.validation.models.h3_generation_resume import model_trees + +gate_commit = open(sys.argv[1]).read().strip() +assert json.load(open(sys.argv[2]))["status"] == "passed" +assert model_trees(gate_commit) == model_trees("HEAD") +PY + then + gate_dir=$attempt + fi +done +[[ -n $gate_dir ]] || { echo "no passing numerical gate for these model paths in job $gate_job" >&2; exit 1; } +seeds=($("$python_bin" -c 'import json; print(*json.load(open("examples/models/minimax_h3/generation-contract.json"))["seeds"])')) +seed=${seeds[$SLURM_ARRAY_TASK_ID]} +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false KERNEL_ALIGN_USE_FAST_MATH=0 +export CUDA_HOME=/usr/local/cuda-13.0 +export PATH="$CUDA_HOME/bin:$PATH" +out="$result_root/generation-${commit:0:8}" +mkdir -p "$out" +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --candidate --seeds "$seed" --resume --output "$out" \ + >> "$out/seed-$seed.log" 2>&1 diff --git a/rl_engine/validation/models/h3_generation_resume.py b/rl_engine/validation/models/h3_generation_resume.py new file mode 100644 index 000000000..4d1cfc428 --- /dev/null +++ b/rl_engine/validation/models/h3_generation_resume.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable denoising for the full-H3 candidate generation (RFC #420). + +A candidate generation takes about 20 hours per seed and its cluster partition +preempts jobs every few hours, so a restart must not redo finished steps. The +generation already keeps every step's predicted velocities (``step-NNN.pt``); +``StepStore`` adds, next to each, a sidecar with the digest of the transformer's +inputs for that step and the digest of the saved file, both written atomically +(the sidecar last, so a step counts only when both exist). + +On resume every saved transformer call is *replayed*: the saved velocities are +returned instead of recomputed, and the pipeline's own scheduler update turns them +into the next latents exactly as before. The input digest of every replayed call is +checked against the one saved, so the replay is proven to have rebuilt the trajectory +the earlier attempt had. Recomputing the last saved step (``verify_last_step``) +additionally checks that the forward pass reproduces its stored output bitwise; it +costs one full step per restart, which on a heavily preempted partition is most of an +attempt, so it is opt-in. The execution trace is saved after the first real step of +an attempt, so a run whose remaining work is only decoding keeps a complete trace. + +A run is tied to the numerics it started with (``model_trees``: the git tree of every +path in ``MODEL_PATHS``), not to a commit, so a commit that changes only orchestration +code can continue it. ``adopt_identity`` moves a run started under a commit-based +identity onto this one and refuses unless those paths are byte-identical. +""" + +from __future__ import annotations + +import datetime +import functools +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Callable + +import torch + +RUN_FILE = "run.json" +TRACE_FILE = "trace.json" + +# Everything that determines the generated numbers. The resume and orchestration code is +# deliberately absent: changing it must not invalidate saved steps. +MODEL_PATHS = ( + "rl_engine/backends", + "rl_engine/contracts", + "rl_engine/ops", + "rl_engine/reference", + "rl_engine/runtime", + "rl_engine/validation/models/h3_block.py", + "rl_engine/validation/models/h3_cases.py", + "rl_engine/validation/models/h3_full_manifest.json", + "rl_engine/validation/models/h3_full_model.py", + "rl_engine/validation/models/h3_provider.py", + "rl_engine/validation/models/h3_weights.py", + "csrc", + "build_tools", + "examples/models/minimax_h3/generation-contract.json", + "examples/models/minimax_h3/generation-inputs.json", + "examples/models/minimax_h3/generation-prompt.txt", +) + + +class ResumeError(ValueError): + """A saved run cannot be continued, or a replay disagrees with what was saved.""" + + +def tensor_digest(named: dict[str, torch.Tensor]) -> str: + """sha256 over tensor names, dtypes, shapes and raw bytes, in name order.""" + + digest = hashlib.sha256() + for name in sorted(named): + tensor = named[name].detach().cpu().contiguous() + digest.update(f"{name}|{tensor.dtype}|{tuple(tensor.shape)}|".encode()) + if tensor.numel(): + digest.update(memoryview(tensor.reshape(-1).view(torch.uint8).numpy())) + return digest.hexdigest() + + +def inputs_digest(args: tuple, kwargs: dict[str, Any]) -> str: + """Digest of every tensor passed to one transformer call (non-tensors are ignored).""" + + named = {f"arg{i}": v for i, v in enumerate(args) if isinstance(v, torch.Tensor)} + named.update({k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)}) + return tensor_digest(named) + + +def model_trees(revision: str = "HEAD", cwd: Path | None = None) -> dict[str, str]: + """The git tree (or blob) id of every model path at ``revision``.""" + + import subprocess + + trees = {} + for path in MODEL_PATHS: + trees[path] = subprocess.check_output( + ["git", "rev-parse", f"{revision}:{path}"], text=True, cwd=cwd + ).strip() + return trees + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1 << 24): + digest.update(block) + return digest.hexdigest() + + +def _atomic_write(path: Path, write: Callable[[Path], object]) -> None: + temporary = path.with_name(path.name + ".tmp") + write(temporary) + os.replace(temporary, path) + + +class StepStore: + """The per-step velocities and input digests of one seed, in one directory.""" + + def __init__(self, directory: Path): + self.directory = directory + + def _step(self, step: int) -> Path: + return self.directory / f"step-{step:03}.pt" + + def _sidecar(self, step: int) -> Path: + return self.directory / f"step-{step:03}.in.json" + + def open(self, identity: dict[str, Any], *, resume: bool) -> int: + """Create the directory, or on ``resume`` validate it; return the saved step count. + + ``identity`` (commit, contract, seed, flags, ...) must equal the one stored, + so steps from a different source or setting can never be mixed in. + """ + + run_file = self.directory / RUN_FILE + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + if not self.directory.exists(): + self.directory.mkdir(parents=True) + elif not resume: + raise ResumeError(f"{self.directory} exists; pass resume to continue it") + if run_file.exists(): + stored = json.loads(run_file.read_text()) + if stored["identity"] != identity: + raise ResumeError("saved run has a different identity; refusing to continue it") + else: + stored = {"identity": identity, "attempts": []} + saved = self._complete_steps() + stored["attempts"].append({"started": now, "saved_steps": saved}) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) + return saved + + def attempts(self) -> list[dict[str, Any]]: + return json.loads((self.directory / RUN_FILE).read_text())["attempts"] + + def _complete_steps(self) -> int: + count = 0 + while self._step(count).exists() and self._sidecar(count).exists(): + sidecar = json.loads(self._sidecar(count).read_text()) + if _file_sha256(self._step(count)) != sidecar["output_sha256"]: + break + count += 1 + if count and not (self.directory / "initial.pt").exists(): + return 0 + return count + + def save(self, step: int, values: tuple[torch.Tensor, ...], digest: str) -> None: + """Persist one step's outputs, then its sidecar; both atomically.""" + + _atomic_write(self._step(step), lambda p: torch.save(tuple(values), p)) + sidecar = {"inputs_sha256": digest, "output_sha256": _file_sha256(self._step(step))} + _atomic_write( + self._sidecar(step), lambda p: p.write_text(json.dumps(sidecar, indent=2) + "\n") + ) + + def save_trace(self, trace: dict[str, Any]) -> None: + """Keep the validated execution trace of this seed's real forward passes.""" + + _atomic_write( + self.directory / TRACE_FILE, lambda p: p.write_text(json.dumps(trace, indent=2) + "\n") + ) + + def load_trace(self) -> dict[str, Any] | None: + path = self.directory / TRACE_FILE + return json.loads(path.read_text()) if path.exists() else None + + def saved_digest(self, step: int) -> str: + return json.loads(self._sidecar(step).read_text())["inputs_sha256"] + + def load(self, step: int) -> tuple[torch.Tensor, ...]: + sidecar = json.loads(self._sidecar(step).read_text()) + if _file_sha256(self._step(step)) != sidecar["output_sha256"]: + raise ResumeError(f"saved step {step} changed on disk") + return tuple(torch.load(self._step(step), map_location="cpu", weights_only=True)) + + +class ResumableRun: + """Replays saved steps and records new ones for one seed's transformer calls.""" + + def __init__( + self, + store: StepStore, + identity: dict[str, Any], + *, + resume: bool, + verify_last_step: bool = False, + ): + self.store = store + self.saved = store.open(identity, resume=resume) + self.verify_last_step = verify_last_step + self.replay = self.saved - 1 if verify_last_step and self.saved else self.saved + self.verified_step: int | None = None + self._forward_calls = 0 + self.recorded = 0 + self.real_calls = 0 + + def install(self, module: torch.nn.Module, make_output: Callable[[tuple, dict], Any]): + """Serve the first ``replay`` calls of ``module`` from disk; return an uninstaller.""" + + original = module.forward + + # The pipeline picks the transformer's layout inputs from inspect.signature(forward), + # so the wrapper must keep the original signature (functools.wraps sets __wrapped__). + @functools.wraps(original) + def forward(*args, **kwargs): + index = self._forward_calls + self._forward_calls += 1 + if index >= self.replay: + return original(*args, **kwargs) + if inputs_digest(args, kwargs) != self.store.saved_digest(index): + raise ResumeError(f"replayed step {index} has different inputs than when saved") + device = next( + v.device for v in (*args, *kwargs.values()) if isinstance(v, torch.Tensor) + ) + values = tuple(v.to(device) for v in self.store.load(index)) + return make_output(values, kwargs) + + module.forward = forward + + def uninstall(): + module.forward = original + + return uninstall + + def record(self, args: tuple, kwargs: dict[str, Any], values: tuple[torch.Tensor, ...]): + """Handle one finished call (forward hook): save it, or verify a recomputed step.""" + + index = self.recorded + self.recorded += 1 + if index < self.replay: + return + self.real_calls += 1 + digest = inputs_digest(args, kwargs) + values = tuple(v.detach().cpu() for v in values) + if index < self.saved: + if digest != self.store.saved_digest(index): + raise ResumeError(f"recomputed step {index} has different inputs than saved") + stored = self.store.load(index) + if len(stored) != len(values) or not all( + torch.equal(a, b) for a, b in zip(stored, values, strict=True) + ): + raise ResumeError(f"recomputed step {index} differs from the saved output") + self.verified_step = index + return + self.store.save(index, values, digest) + + def summary(self) -> dict[str, Any]: + return dict( + saved_steps_at_start=self.saved, + replayed_steps=self.replay, + verify_last_step=self.verify_last_step, + recomputed_and_verified_step=self.verified_step, + real_forward_calls=self.real_calls, + attempts=self.store.attempts(), + ) + + +def adopt_identity( + directory: Path, + new_identity: dict[str, Any], + from_commit: str, + cwd: Path | None = None, + *, + dry_run: bool = False, +) -> None: + """Move a run started under a commit-based identity onto a ``model_trees`` identity. + + Refuses unless everything but the commit matches and the model paths at + ``from_commit`` are byte-identical to the ones in ``new_identity``; records the + adoption in the run's attempt list. ``dry_run`` performs every check and writes nothing. + """ + + run_file = directory / RUN_FILE + stored = json.loads(run_file.read_text()) + old = stored["identity"] + if "model_trees" in old: + raise ResumeError("run already has a model-tree identity") + if old.get("commit", "").startswith(from_commit) is False: + raise ResumeError(f"run was started at {old.get('commit')}, not {from_commit}") + rest_old = {k: v for k, v in old.items() if k != "commit"} + rest_new = {k: v for k, v in new_identity.items() if k != "model_trees"} + if rest_old != rest_new: + raise ResumeError("run differs from the new identity in more than the commit") + then = model_trees(old["commit"], cwd=cwd) + if then != new_identity["model_trees"]: + changed = sorted(k for k in then if then[k] != new_identity["model_trees"][k]) + raise ResumeError(f"model paths changed since {from_commit}: {changed}") + if dry_run: + return + stored["identity"] = new_identity + stored["attempts"].append( + { + "adopted": { + "from_commit": old["commit"], + "model_paths_identical": True, + "when": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + } + ) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) diff --git a/tests/models/minimax_h3/test_h3_generation_resume.py b/tests/models/minimax_h3/test_h3_generation_resume.py new file mode 100644 index 000000000..5c0d3d56b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_generation_resume.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable candidate generation: a preempted run continues to the same result. + +A fake transformer and denoising loop stand in for the pipeline, so these run on CPU. +The loop feeds each call's output back into the next call's inputs (as the scheduler +does), so any wrong replay would change every later step. +""" + +from __future__ import annotations + +import inspect +import subprocess +import types + +import pytest +import torch + +from rl_engine.validation.models import h3_generation_resume as resume_module +from rl_engine.validation.models.h3_generation_resume import ( + ResumableRun, + ResumeError, + StepStore, + adopt_identity, + inputs_digest, + model_trees, +) + +STEPS = 6 +IDENTITY = {"commit": "abc", "seed": 42, "candidate": True, "preflight": False} + + +class FakeTransformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.computed = 0 + + def forward(self, hidden_states, audio_hidden_states, *, return_dict=True): + self.computed += 1 + video = torch.tanh(hidden_states) * 0.5 + audio = torch.sin(audio_hidden_states) * 0.25 + return _output((video, audio), {"return_dict": return_dict}) + + +def _output(values, kwargs): + if kwargs.get("return_dict", True): + return types.SimpleNamespace(sample=values[0], audio_sample=values[1]) + return values + + +def _initial(seed=0): + generator = torch.Generator().manual_seed(seed) + return torch.randn(8, 4, generator=generator), torch.randn(3, 5, generator=generator) + + +def denoise( + directory, *, resume, stop_after=None, initial=None, identity=IDENTITY, verify_last_step=False +): + """Run the fake loop through a ResumableRun; return (latents, transformer, run).""" + + video, audio = initial if initial is not None else _initial() + store = StepStore(directory) + run = ResumableRun(store, identity, resume=resume, verify_last_step=verify_last_step) + if not (directory / "initial.pt").exists(): + torch.save({"hidden_states": video, "audio_hidden_states": audio}, directory / "initial.pt") + module = FakeTransformer() + run.install(module, _output) + + def hook(mod, args, kwargs, output): + values = (output.sample, output.audio_sample) + run.record(args, kwargs, values) + if stop_after is not None and run.recorded >= stop_after: + raise KeyboardInterrupt("simulated preemption") + + module.register_forward_hook(hook, with_kwargs=True) + for _ in range(STEPS): + out = module(hidden_states=video, audio_hidden_states=audio) + # The scheduler's update: the next inputs depend on this call's outputs. + video = video + out.sample + audio = audio - 0.5 * out.audio_sample + return (video, audio), module, run + + +def _step_bytes(directory): + return { + p.name: p.read_bytes() + for p in sorted(directory.iterdir()) + if p.name.startswith("step-") and p.suffix in (".pt", ".json") + } + + +def test_uninterrupted_run_saves_every_step(tmp_path): + (video, _), module, run = denoise(tmp_path / "a", resume=False) + assert module.computed == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.pt")]) == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.in.json")]) == STEPS + assert run.summary()["replayed_steps"] == 0 + + +def test_resumed_run_with_verification_matches_an_uninterrupted_one(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True, verify_last_step=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert _step_bytes(tmp_path / "whole").keys() == _step_bytes(tmp_path / "cut").keys() + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + # Four steps were saved: three replayed, the fourth recomputed and verified, two new. + summary = run.summary() + assert summary["saved_steps_at_start"] == 4 + assert summary["replayed_steps"] == 3 + assert summary["recomputed_and_verified_step"] == 3 + assert module.computed == STEPS - 3 + assert [a["saved_steps"] for a in summary["attempts"]] == [0, 4] + + +def test_repeated_preemptions_still_converge(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + for stop in (2, 3, 5): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=stop > 2, stop_after=stop) + resumed, _, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert len(run.summary()["attempts"]) == 4 + + +def test_with_verification_a_complete_run_recomputes_only_its_last_step(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True, verify_last_step=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 1 + assert run.summary()["recomputed_and_verified_step"] == STEPS - 1 + + +def test_different_inputs_are_rejected_during_replay(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + # Rejected at the first replayed call, not only at the recomputed last saved step. + with pytest.raises(ResumeError, match="replayed step 0 has different inputs"): + denoise(tmp_path / "a", resume=True, initial=_initial(seed=1)) + + +def test_a_nondeterministic_recompute_is_rejected_when_verifying(tmp_path, monkeypatch): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + + original = FakeTransformer.forward + + def noisy(self, hidden_states, audio_hidden_states, *, return_dict=True): + out = original(self, hidden_states, audio_hidden_states, return_dict=return_dict) + return _output((out.sample + 1e-6, out.audio_sample), {"return_dict": return_dict}) + + monkeypatch.setattr(FakeTransformer, "forward", noisy) + with pytest.raises(ResumeError, match="differs from the saved output"): + denoise(tmp_path / "a", resume=True, verify_last_step=True) + # Without verification nothing recomputes the saved steps, so the drift goes unnoticed. + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "b", resume=False, stop_after=3) + denoise(tmp_path / "b", resume=True) + + +def test_a_different_identity_cannot_be_continued(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + with pytest.raises(ResumeError, match="different identity"): + denoise(tmp_path / "a", resume=True, identity={**IDENTITY, "commit": "other"}) + + +def test_an_existing_directory_needs_resume(tmp_path): + denoise(tmp_path / "a", resume=False, stop_after=None) + with pytest.raises(ResumeError, match="resume"): + denoise(tmp_path / "a", resume=False) + + +def test_a_corrupted_step_ends_the_valid_prefix(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=4) + (tmp_path / "a" / "step-002.pt").write_bytes(b"damaged") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 2 + + +def test_leftover_temporary_files_are_ignored(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "step-003.pt.tmp").write_bytes(b"half written") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 3 + + +def test_a_run_without_initial_inputs_is_not_resumable(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "initial.pt").unlink() + assert ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True).saved == 0 + + +def test_inputs_digest_tracks_value_dtype_and_shape_and_ignores_non_tensors(): + x = torch.arange(6.0).reshape(2, 3) + base = inputs_digest((), {"a": x, "flag": True}) + assert base == inputs_digest((), {"a": x.clone(), "flag": False}) + assert base != inputs_digest((), {"a": x + 1}) + assert base != inputs_digest((), {"a": x.double()}) + assert base != inputs_digest((), {"a": x.reshape(3, 2)}) + assert base != inputs_digest((), {"b": x}) + assert inputs_digest((x,), {}) != inputs_digest((), {"a": x}) + assert inputs_digest((), {"s": torch.tensor(3.0)}) # zero-dim tensors hash + + +def test_replay_wrapper_keeps_the_forward_signature(tmp_path): + """The pipeline selects the transformer's layout inputs from inspect.signature(forward).""" + + module = FakeTransformer() + before = list(inspect.signature(module.forward).parameters) + (tmp_path / "a").mkdir() + run = ResumableRun(StepStore(tmp_path / "a" / "x"), IDENTITY, resume=False) + uninstall = run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + uninstall() + assert list(inspect.signature(module.forward).parameters) == before + + +def test_replay_wrapper_keeps_a_wrapped_forward_signature(tmp_path): + """Same through an accelerate-style partial that itself carries __wrapped__.""" + + import functools + + module = FakeTransformer() + plain = module.forward + + def new_forward(m, *args, **kwargs): + return plain(*args, **kwargs) + + hooked = functools.update_wrapper(functools.partial(new_forward, module), plain) + module.forward = hooked + before = list(inspect.signature(module.forward).parameters) + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + run = ResumableRun(StepStore(tmp_path / "b"), IDENTITY, resume=False) + run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + + +# --- default resume: replay everything saved, recompute nothing -------------------------- + + +def test_default_resume_replays_every_saved_step_and_computes_only_the_rest(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + summary = run.summary() + assert (summary["saved_steps_at_start"], summary["replayed_steps"]) == (4, 4) + assert summary["recomputed_and_verified_step"] is None + assert summary["real_forward_calls"] == STEPS - 4 == module.computed + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + + +def test_a_complete_run_resumes_with_no_forward_pass_at_all(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 0 + assert run.real_calls == 0 + + +def test_each_attempt_only_needs_a_few_steps_to_make_progress(tmp_path): + """Attempts that each finish one new step still add up to the whole run.""" + + whole, _, _ = denoise(tmp_path / "whole", resume=False) + saved = 0 + result = run = None + for attempt in range(STEPS): + stop = saved + 1 if saved + 1 < STEPS else None # the last attempt is left to finish + try: + result, _, run = denoise(tmp_path / "cut", resume=attempt > 0, stop_after=stop) + except KeyboardInterrupt: + saved += 1 + continue + break + assert result is not None and saved == STEPS - 1 + assert all(torch.equal(a, b) for a, b in zip(whole, result, strict=True)) + assert run.real_calls == 1 and len(run.summary()["attempts"]) == STEPS + + +def test_the_validated_trace_survives_for_a_later_replay_only_attempt(tmp_path): + store = StepStore(tmp_path / "a") + assert store.load_trace() is None + (tmp_path / "a").mkdir() + trace = {"calls": {"block.0": 3}, "execution_order": ["block.0", "block.1"]} + store.save_trace(trace) + assert StepStore(tmp_path / "a").load_trace() == trace + assert not list((tmp_path / "a").glob("*.tmp")) + + +# --- identity tied to the model paths, and adopting a commit-based run ------------------- + + +def _git(repo, *args): + return subprocess.check_output(["git", *args], cwd=repo, text=True).strip() + + +@pytest.fixture +def repo(tmp_path, monkeypatch): + root = tmp_path / "repo" + (root / "model").mkdir(parents=True) + (root / "orchestration").mkdir() + (root / "model" / "kernel.py").write_text("v1\n") + (root / "orchestration" / "resume.py").write_text("v1\n") + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@example.com") + _git(root, "config", "user.name", "t") + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "a") + monkeypatch.setattr(resume_module, "MODEL_PATHS", ("model",)) + return root + + +def _commit(repo, path, text): + (repo / path).write_text(text) + _git(repo, "commit", "-q", "-am", f"change {path}") + return _git(repo, "rev-parse", "HEAD") + + +def _old_run(directory, commit, trees): + base = {k: v for k, v in IDENTITY.items() if k != "commit"} + store = StepStore(directory) + store.open({"commit": commit, **base}, resume=False) + return {**base, "model_trees": trees} + + +def test_model_trees_ignore_orchestration_changes_but_not_model_changes(repo): + first = model_trees("HEAD", cwd=repo) + _commit(repo, "orchestration/resume.py", "v2\n") + assert model_trees("HEAD", cwd=repo) == first + _commit(repo, "model/kernel.py", "v2\n") + assert model_trees("HEAD", cwd=repo) != first + + +def test_adopting_a_commit_based_run_lets_a_later_commit_continue_it(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + new_identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "orchestration/resume.py", "v2\n") + new_identity["model_trees"] = model_trees("HEAD", cwd=repo) + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert StepStore(tmp_path / "run").open(new_identity, resume=True) == 0 + attempts = StepStore(tmp_path / "run").attempts() + assert attempts[-2]["adopted"]["from_commit"] == old + assert attempts[-2]["adopted"]["model_paths_identical"] is True + + +def test_adoption_is_refused_when_a_model_path_changed(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "model/kernel.py", "v2\n") + new_identity = { + **{k: v for k, v in IDENTITY.items() if k != "commit"}, + "model_trees": model_trees("HEAD", cwd=repo), + } + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="model paths changed"): + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert (tmp_path / "run" / "run.json").read_text() == before + + +def test_adoption_refuses_other_differences_and_dry_run_writes_nothing(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="more than the commit"): + adopt_identity(tmp_path / "run", {**identity, "seed": 7}, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="not abc"): + adopt_identity(tmp_path / "run", identity, "abc", cwd=repo) + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo, dry_run=True) + assert (tmp_path / "run" / "run.json").read_text() == before + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="already"): + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) diff --git a/tools/validation/models/h3_full_generation.py b/tools/validation/models/h3_full_generation.py new file mode 100644 index 000000000..64b215e9f --- /dev/null +++ b/tools/validation/models/h3_full_generation.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Pinned three-seed FL2VA generation evidence for the full H3 candidate stack. + +A candidate seed takes about 20 hours and its partition preempts jobs, so a seed can be +run on its own (``--seeds``) and continued after a preemption (``--resume``): saved steps +are replayed from disk and verified (see ``h3_generation_resume``), never recomputed. +""" +from __future__ import annotations + +import argparse +import datetime +import functools +import json +import os +import subprocess +import time +from pathlib import Path + +import numpy as np +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from rl_engine.validation.models.h3_generation_resume import ResumableRun, StepStore, model_trees +from tools.validation.models.h3_full_model import ( + PROVIDER_VARIANTS, + digest, + provider_variant_context, + verify_checkpoint, +) + + +class QueryChunkAttention: + """Bound attention workspace; every query still attends to every key.""" + + def __init__(self, op, chunk): + self.op, self.chunk = op, chunk + + def __call__(self, q, k, v, **kwargs): + if torch.is_grad_enabled(): + raise ValueError("generation query chunking is inference only") + return torch.cat( + [self.op(part.contiguous(), k, v, **kwargs) for part in q.split(self.chunk, dim=2)], + dim=2, + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-image", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--candidate", action="store_true") + p.add_argument("--preflight", action="store_true") + p.add_argument( + "--seeds", help="comma list of contract seeds to run (default: all; preflight: first)" + ) + p.add_argument( + "--resume", + action="store_true", + help="continue a preempted run in --output: skip finished seeds, replay saved steps", + ) + p.add_argument( + "--verify-last-step", + action="store_true", + help="on resume, recompute the last saved step and require it to match bitwise " + "(costs one step per restart)", + ) + p.add_argument( + "--provider-variant", + choices=PROVIDER_VARIANTS, + default="default", + help="provider only, a variability study: force one SDPA backend or cuBLASLt", + ) + p.add_argument("--preflight-steps", type=int, help="test only: preflight num_inference_steps") + p.add_argument( + "--stop-after-step", type=int, help="test only: raise after N calls (preemption)" + ) + args = p.parse_args() + if args.provider_variant != "default" and args.candidate: + p.error("--provider-variant applies to provider generation only") + if args.preflight_steps is not None and not args.preflight: + p.error("--preflight-steps applies to --preflight only") + args.output.mkdir(parents=True, exist_ok=args.resume) + stale = args.output / failure_name(args.seeds) + if args.resume and stale.exists(): + # A failure of an earlier attempt is kept, not overwritten or hidden. + stale.rename(args.output / f"{stale.stem}-{int(time.time())}.json") + source = Path(__file__).resolve().parents[3] + config_root = source / "examples/models/minimax_h3" + contract = json.loads((config_root / "generation-contract.json").read_text()) + identity_inputs = json.loads((config_root / "generation-inputs.json").read_text()) + prompt = config_root / "generation-prompt.txt" + if ( + digest(prompt) != identity_inputs["prompt_sha256"] + or digest(args.reference_image) != identity_inputs["image_sha256"] + ): + raise ValueError("fixed prompt/reference input identity changed") + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("generation requires clean source") + verify_checkpoint( + args.checkpoint, + source / "rl_engine/validation/models/h3_full_manifest.json", + all_files=True, + ) + import diffusers + + provider_root = Path(diffusers.__file__).resolve().parents[2] + if ( + subprocess.check_output( + ["git", "-C", str(provider_root), "rev-parse", "HEAD"], text=True + ).strip() + != contract["provider_revision"] + ): + raise ValueError("generation provider revision differs") + if subprocess.check_output( + ["git", "-C", str(provider_root), "status", "--porcelain"], text=True + ).strip(): + raise ValueError("generation provider is dirty") + import soundfile as sf + from diffusers import ComponentsManager, ModularPipeline + from diffusers.utils.export_utils import encode_video + from PIL import Image + from transformers import Qwen3VLProcessor + + # Load before the large models: ModularPipeline otherwise catches processor + # dependency failures and leaves None, failing much later in conditioning. + processor = Qwen3VLProcessor.from_pretrained( + str(args.checkpoint / "processor"), local_files_only=True + ) + manager = ComponentsManager() + manager.enable_auto_cpu_offload(device="cuda", memory_reserve_margin="20GB") + pipe = ModularPipeline.from_pretrained( + str(args.checkpoint), workflow="fl2va", components_manager=manager + ) + pipe.update_components(processor=processor) + pipe.load_components( + dtype=torch.bfloat16, + pretrained_model_name_or_path=str(args.checkpoint), + local_files_only=True, + ) + for name in ( + "processor", + "tokenizer", + "text_encoder", + "transformer", + "vae", + "audio_vae", + "scheduler", + "audio_scheduler", + ): + if getattr(pipe, name, None) is None: + raise RuntimeError(f"Required FL2VA component failed to load: {name}") + ops = None + if args.candidate: + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + settings = contract["preflight"] if args.preflight else contract + if args.preflight_steps is not None: + settings = dict(settings, num_inference_steps=args.preflight_steps) + seeds = contract["seeds"][:1] if args.preflight else contract["seeds"] + if args.seeds: + seeds = [int(s) for s in args.seeds.split(",")] + if not set(seeds) <= set(contract["seeds"]): + raise ValueError(f"seeds {seeds} are not all in the contract {contract['seeds']}") + from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput + + def make_output(values, kwargs): + if kwargs.get("return_dict", True): + return MiniMaxH3TransformerOutput(sample=values[0], audio_sample=values[1]) + return values + + for seed in seeds: + directory = args.output / str(seed) + if args.resume and (directory / "report.json").exists(): + print(f"seed {seed} already complete; skipping", flush=True) + continue + identity = run_identity( + config_root, + contract, + identity_inputs, + candidate=args.candidate, + preflight=args.preflight, + seed=seed, + settings=settings, + provider_variant=args.provider_variant, + ) + run = ResumableRun( + StepStore(directory), + identity, + resume=args.resume, + verify_last_step=args.verify_last_step, + ) + trace = FullModelTrace() + initial = {} + + def capture(module, positional, kwargs, output, run=run, directory=directory): + # Preserve the complete first packed denoiser input to prove identical + # noise and conditioning, then retain every step's predicted velocities. + values = (output.sample, output.audio_sample) if hasattr(output, "sample") else output + if run.saved == 0 and run.recorded == 0: + initial.update( + {k: v.detach().cpu() for k, v in kwargs.items() if isinstance(v, torch.Tensor)} + ) + tmp = directory / "initial.pt.tmp" + torch.save(initial, tmp) + os.replace(tmp, directory / "initial.pt") + run.record(positional, kwargs, tuple(values)) + if run.real_calls == 1: + # One real forward pass has run every block and boundary: keep that trace, so a + # later attempt that only replays and decodes still reports a complete one. + run.store.save_trace(trace.validate()) + if args.stop_after_step is not None and run.recorded >= args.stop_after_step: + raise KeyboardInterrupt(f"simulated preemption after {run.recorded} calls") + + if args.provider_variant != "default": + # Only the transformer runs under the forced kernel: the text encoder and VAE + # (which have their own attention masks) and the conditioning stay the provider's. + plain_forward = pipe.transformer.forward + + @functools.wraps(plain_forward) + def variant_forward(*a, _plain=plain_forward, **kw): + with provider_variant_context(args.provider_variant): + return _plain(*a, **kw) + + pipe.transformer.forward = variant_forward + uninstall = run.install(pipe.transformer, make_output) + handle = pipe.transformer.register_forward_hook(capture, with_kwargs=True) + started = time.perf_counter() + try: + with ( + torch.inference_mode(), + full_model_execution( + pipe.transformer, candidate=args.candidate, trace=trace, ops=ops + ), + ): + result = pipe( + prompt=prompt.read_text(), + image=Image.open(args.reference_image).convert("RGB"), + num_frames=settings["num_frames"], + height=settings["height"], + width=settings["width"], + num_inference_steps=settings["num_inference_steps"], + generator=torch.Generator().manual_seed(seed), + output_type="np", + output=["videos", "audio", "sampling_rate"], + ) + finally: + handle.remove() + uninstall() + if args.provider_variant != "default": + pipe.transformer.forward = plain_forward + video = np.asarray(result["videos"][0]) + audio = result["audio"][0].detach().cpu().numpy() + if video.shape != ( + settings["num_frames"], + settings["height"], + settings["width"], + 3, + ): + raise ValueError(f"wrong video dimensions {video.shape}") + if not np.isfinite(video).all() or not np.isfinite(audio).all(): + raise ValueError("non-finite generated media") + if ( + result["sampling_rate"] != contract["sampling_rate"] + or abs( + audio.shape[-1] / result["sampling_rate"] - settings["num_frames"] / contract["fps"] + ) + > 0.05 + ): + raise ValueError("incorrect audio sample rate/duration") + np.save(directory / "video.npy", video) + np.save(directory / "audio.npy", audio) + sf.write(directory / "audio.wav", audio.T, result["sampling_rate"], subtype="FLOAT") + encode_video( + video, + fps=contract["fps"], + output_path=str(directory / "video.mp4"), + audio=result["audio"][0], + audio_sample_rate=result["sampling_rate"], + ) + for index in (0, len(video) // 2, len(video) - 1): + Image.fromarray(np.round(video[index] * 255).clip(0, 255).astype("uint8")).save( + directory / f"frame-{index:03}.png" + ) + final_trace = trace.validate() if run.real_calls else run.store.load_trace() + if final_trace is None: + raise ValueError("no execution trace: this seed never ran a real forward pass") + report = dict( + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + model_revision=contract["model_revision"], + provider_revision=contract["provider_revision"], + status="generation_complete_comparison_pending", + candidate=args.candidate, + provider_variant=args.provider_variant, + seed=seed, + preflight=args.preflight, + trace=final_trace, + steps=run.recorded, + resume=run.summary(), + seconds=time.perf_counter() - started, + peak_allocated_bytes=torch.cuda.max_memory_allocated(), + contract_sha256=digest(config_root / "generation-contract.json"), + ) + report["artifacts"] = { + path.name: digest(path) for path in sorted(directory.iterdir()) if path.is_file() + } + (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +def run_identity( + config_root, + contract, + inputs, + *, + candidate, + preflight, + seed, + settings, + provider_variant="default", +): + """What a saved run is tied to: the numerics (model-path git trees), inputs and settings.""" + + return dict( + model_trees=model_trees(), + contract_sha256=digest(config_root / "generation-contract.json"), + image_sha256=inputs["image_sha256"], + prompt_sha256=inputs["prompt_sha256"], + candidate=candidate, + preflight=preflight, + seed=seed, + **({"provider_variant": provider_variant} if provider_variant != "default" else {}), + settings={k: settings[k] for k in ("num_frames", "height", "width")} + | {"num_inference_steps": settings["num_inference_steps"]}, + ) + + +def failure_name(seeds): + """Failure marker name: per seed selection, so parallel seed tasks do not collide.""" + + return "failure.json" if not seeds else f"failure-seeds-{seeds.replace(',', '-')}.json" + + +if __name__ == "__main__": + try: + main() + except BaseException as exc: + import sys + + if "--output" in sys.argv: + directory = Path(sys.argv[sys.argv.index("--output") + 1]) + seeds = sys.argv[sys.argv.index("--seeds") + 1] if "--seeds" in sys.argv else None + marker = directory / failure_name(seeds) + if directory.exists() and not marker.exists(): + marker.write_text( + json.dumps( + { + "status": "failed", + "error": f"{type(exc).__name__}: {exc}", + "when": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + ) + ) + raise diff --git a/tools/validation/models/h3_generation_adopt.py b/tools/validation/models/h3_generation_adopt.py new file mode 100644 index 000000000..b87876878 --- /dev/null +++ b/tools/validation/models/h3_generation_adopt.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Move saved candidate-generation runs onto the model-tree identity (RFC #420, #539). + +Runs started before the identity moved from a commit to the git trees of the model paths +(``h3_generation_resume.MODEL_PATHS``) carry the commit they started at. This re-labels +them for the current checkout, once, and only if every model path is byte-identical +between that commit and ``HEAD``; anything else is refused and nothing is written. + + python tools/validation/models/h3_generation_adopt.py --output \\ + --from-commit 4792c8a9 +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +from rl_engine.validation.models.h3_generation_resume import adopt_identity # noqa: E402 +from tools.validation.models.h3_full_generation import run_identity # noqa: E402 + + +def main() -> None: + """Adopt every seed directory of ``--output`` (all contract seeds by default).""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--from-commit", required=True) + parser.add_argument("--seeds", help="comma list; default: the contract's seeds") + args = parser.parse_args() + source = Path(__file__).resolve().parents[3] + config_root = source / "examples/models/minimax_h3" + contract = json.loads((config_root / "generation-contract.json").read_text()) + inputs = json.loads((config_root / "generation-inputs.json").read_text()) + seeds = [int(s) for s in args.seeds.split(",")] if args.seeds else contract["seeds"] + # Check every seed before changing any, so a refusal leaves all runs untouched. + plans = [] + for seed in seeds: + identity = run_identity( + config_root, + contract, + inputs, + candidate=True, + preflight=False, + seed=seed, + settings=contract, + ) + plans.append((args.output / str(seed), identity)) + for directory, identity in plans: + adopt_identity(directory, identity, args.from_commit, cwd=source, dry_run=True) + for directory, identity in plans: + adopt_identity(directory, identity, args.from_commit, cwd=source) + print(f"adopted {directory.name} (from {args.from_commit})") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_generation_compare.py b/tools/validation/models/h3_generation_compare.py new file mode 100644 index 000000000..27510767a --- /dev/null +++ b/tools/validation/models/h3_generation_compare.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Compare fixed full-generation runs, preserving failed quality gates.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from skimage.metrics import structural_similarity + +from tools.validation.models.h3_full_model import digest + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--reference", type=Path, required=True) + p.add_argument("--candidate", type=Path, required=True) + p.add_argument("--contract", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--provider-repeat", action="store_true") + p.add_argument("--reference-manifest", type=Path) + p.add_argument("--reference-root", type=Path) + a = p.parse_args() + contract = json.loads(a.contract.read_text()) + from tools.validation.models.h3_evidence_integrity import verify_generation, verify_inventory + + if not a.provider_repeat and not (a.reference_root and a.reference_manifest): + raise ValueError("candidate comparison requires pinned reference inventory") + if a.reference_manifest: + verify_inventory(a.reference_root, a.reference_manifest) + rows = [] + if a.output.exists(): + raise ValueError("never overwrite a prior comparison") + for seed in contract["seeds"]: + left, right = a.reference / str(seed), a.candidate / str(seed) + reports = [verify_generation(root) for root in (left, right)] + if any( + r["preflight"] or r["seed"] != seed or r["contract_sha256"] != digest(a.contract) + for r in reports + ): + raise ValueError("wrong generation evidence/contract") + if reports[0]["candidate"] or reports[1]["candidate"] == a.provider_repeat: + raise ValueError("reference/candidate path mismatch") + initial = [ + torch.load(root / "initial.pt", map_location="cpu", weights_only=True) + for root in (left, right) + ] + if initial[0].keys() != initial[1].keys() or not initial[0]: + raise ValueError("missing packed initial noise evidence") + if any(not torch.equal(value, initial[1][name]) for name, value in initial[0].items()): + raise ValueError("different initial noise/conditioning") + x = np.load(left / "video.npy", mmap_mode="r") + y = np.load(right / "video.npy", mmap_mode="r") + if x.shape != y.shape or x.shape != ( + contract["num_frames"], + contract["height"], + contract["width"], + 3, + ): + raise ValueError("video geometry differs") + mse = 0.0 + ssim = 0.0 + temporal = 0.0 + previous = None + for frame in range(len(x)): + if not np.isfinite(x[frame]).all() or not np.isfinite(y[frame]).all(): + raise ValueError("nonfinite video frame") + delta = x[frame].astype("float64") - y[frame] + mse += float(np.mean(delta**2)) + ssim += float( + structural_similarity(x[frame], y[frame], channel_axis=-1, data_range=1.0) + ) + if previous is not None: + temporal += float(np.mean((delta - previous) ** 2)) + previous = delta + audio = np.load(left / "audio.npy") + other = np.load(right / "audio.npy") + if not np.isfinite(audio).all() or not np.isfinite(other).all(): + raise ValueError("nonfinite audio") + if audio.shape != other.shape: + raise ValueError("audio dimensions differ") + audio_error = float( + np.linalg.norm((audio - other).astype("float64")) + / max(np.linalg.norm(audio.astype("float64")), 1e-12) + ) + row = dict( + seed=seed, + psnr_db=float(-10 * np.log10(max(mse / len(x), 1e-30))), + ssim=ssim / len(x), + temporal_rmse=float(np.sqrt(temporal / (len(x) - 1))), + audio_relative_l2=audio_error, + audio_max_abs=float(np.max(np.abs(audio - other))), + ) + predictions = [] + if reports[0]["steps"] != reports[1]["steps"]: + raise ValueError("different scheduler step counts") + for step in range(reports[0]["steps"]): + v, w = [ + torch.load(root / f"step-{step:03}.pt", map_location="cpu", weights_only=True) + for root in (left, right) + ] + if len(v) != 2 or len(w) != 2: + raise ValueError("missing video/audio prediction") + if any( + b.shape != c.shape or not torch.isfinite(b).all() or not torch.isfinite(c).all() + for b, c in zip(v, w, strict=True) + ): + raise ValueError("nonfinite or malformed step prediction") + predictions.append( + [ + dict( + max_abs=(b.float() - c.float()).abs().max().item(), + relative_l2=( + (b.float() - c.float()).double().norm() + / b.double().norm().clamp_min(1e-12) + ).item(), + ) + for b, c in zip(v, w, strict=True) + ] + ) + bound = contract["quality_bounds"] + row["passed"] = ( + row["psnr_db"] >= bound["minimum_psnr_db"] + and row["ssim"] >= bound["minimum_ssim"] + and row["temporal_rmse"] <= bound["maximum_temporal_rmse"] + and audio_error <= bound["maximum_audio_relative_l2"] + ) + row["step_errors"] = predictions + rows.append(row) + passed = all(r["passed"] for r in rows) + a.output.write_text( + json.dumps( + dict( + status="passed" if passed else "failed", + contract_sha256=digest(a.contract), + seeds=rows, + ), + indent=2, + ) + + "\n" + ) + if not passed: + raise SystemExit(1) + + +if __name__ == "__main__": + main() From 669513d223d33fbc8d0af14b1b00733a50c858fd Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:07:41 +0000 Subject: [PATCH 44/62] fix(h3): keep the gate-residual forward out of the timed backward _gate_perf had no candidate_backward_setup/provider_backward_setup, so each timed backward call also created the leaves, ran the forward and built the graph while measure() reported backward_timing_scope = backward_only. The forward moves into the setup like in _gather_perf and _norm_perf. A test checks that every perf case with a timed backward has its setup. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/validation/models/h3_report.py | 13 +++++++++---- .../models/minimax_h3/test_h3_report_timing.py | 18 ++++++++++++++++++ 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 492bfc1b6..aa1b53498 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -606,9 +606,12 @@ def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: gate = _gate_view(table) grad = torch.randn_like(residual) - def backward(fn, residual=residual, y=y, table=table, grad=grad): + def prepare_backward(fn, residual=residual, y=y, table=table): leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] - fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward(grad) + return fn(leaves[0], leaves[1], _gate_view(leaves[2])) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -622,10 +625,12 @@ def backward(fn, residual=residual, y=y, table=table, grad=grad): "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( r, g, i, y ), - "candidate_backward": lambda b=backward, i=index: b( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) ), - "provider_backward": lambda b=backward, i=index: b( + "provider_backward_setup": lambda b=prepare_backward, i=index: b( lambda r, y_, g: provider_gate_residual(r, g, i, y_) ), } diff --git a/tests/models/minimax_h3/test_h3_report_timing.py b/tests/models/minimax_h3/test_h3_report_timing.py index 853680a67..9b77b16cf 100644 --- a/tests/models/minimax_h3/test_h3_report_timing.py +++ b/tests/models/minimax_h3/test_h3_report_timing.py @@ -52,3 +52,21 @@ def test_the_reported_medians_come_from_the_stored_samples(): assert row[f"{key}_peak_mib"] >= 0 # the larger matmul is slower than the smaller one in the same run assert row["provider_us"] > row["candidate_us"] + + +def test_backward_timings_of_every_perf_case_exclude_the_forward_pass(monkeypatch): + """A case with a timed backward must prepare the graph outside the CUDA-event window.""" + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.models import h3_report + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (257,)) + registry = KernelRegistry() + checked = 0 + for op, build in h3_report.PERF_CASES.items(): + for case in build(registry): + for side in ("candidate", "provider"): + if f"{side}_backward" in case: + checked += 1 + assert f"{side}_backward_setup" in case, f"{op} {case['case']} {side}" + assert checked From e5318ad2ab1a75100336762b08d057b65d5a853d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 16:11:52 +0000 Subject: [PATCH 45/62] fix(h3): do not ask for the checkpoint when the gate-residual evidence never reads it h3_evidence.py treated every op but timestep_sinusoid_h3 as weighted. adaln_gate_residual runs on seeded inputs only, so the CLI refused to run without the 4.8 GB shard and a run with it would have been labelled pinned_checkpoint. The weight-free ops are now an explicit set. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- tests/models/minimax_h3/test_h3_cli.py | 1 + tools/validation/models/h3_evidence.py | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 8e97fc6a2..3f9b726ed 100644 --- a/tests/models/minimax_h3/test_h3_cli.py +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -145,6 +145,7 @@ def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, we "op, weight_source", [ ("timestep_sinusoid_h3", "not_applicable"), + ("adaln_gate_residual", "not_applicable"), ("timestep_mlp_fp32", "pinned_checkpoint"), ("adaln_projection_3mod", "pinned_checkpoint"), ], diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py index ed8267a94..892366786 100644 --- a/tools/validation/models/h3_evidence.py +++ b/tools/validation/models/h3_evidence.py @@ -37,6 +37,9 @@ load_h3_manifest, ) +# Operators whose evidence runs on seeded inputs only and never reads the checkpoint. +WEIGHT_FREE_OPS = frozenset({"timestep_sinusoid_h3", "adaln_gate_residual"}) + _NEEDS_PINNED_MLP = ( "timestep_mlp_fp32 evidence needs the pinned time_embedder weights; set " "RL_KERNEL_H3_WEIGHTS (tools/weights/prepare_h3_weights.py) " @@ -85,7 +88,7 @@ def main() -> None: ) args = parser.parse_args() require_pinned_weights(args.op, args.allow_random_weights) - uses_weights = args.op != "timestep_sinusoid_h3" + uses_weights = args.op not in WEIGHT_FREE_OPS smoke_run = args.allow_random_weights and args.op == "timestep_mlp_fp32" if uses_weights and h3_weights_dir() is None and not smoke_run: raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") From 239ef808e279b6f62dec4592c4311c7cc1b0c32f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:27:36 +0000 Subject: [PATCH 46/62] fix(h3): assert the gate row in the native kernels and check the backward position repeat - gate_residual.cu: Gate carries the table row count and the row kernel asserts 0 <= index < rows before using it as an offset, like the RMSNorm and row-gather kernels (check_range=False callers and direct _C calls get a device-side assert instead of an out-of-range read). - The backward uses _repeat_positions() instead of index.repeat(rows // len(index)). - Test: rows -1 and R in forward and backward expect the device-side assert (subprocess). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 7 ++- .../minimax_h3/gate_residual.py | 9 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 52 +++++++++++++++++++ 3 files changed, 63 insertions(+), 5 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index c2bd7921a..6e0b5a85b 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -50,6 +50,7 @@ struct Gate { int64_t stride; const int64_t* index; // (S,) int64_t seq; + int64_t num_rows; }; // One block per row (grid-strided); 16-byte vectors when the row allows it. @@ -60,7 +61,9 @@ __global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __re Gate gate, bool vectorized) { constexpr int kVecElems = 16 / sizeof(T); for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { - const T* g = static_cast(gate.rows) + gate.index[r % gate.seq] * gate.stride; + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; const T* ar = a + r * n; const T* br = kDy ? nullptr : b + r * n; T* orow = out + r * n; @@ -148,7 +151,7 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); - return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0)}; + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; } void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 2c6c6be55..ad9bb6d8f 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -7,9 +7,12 @@ import torch -from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) from rl_engine.backends.extension import _C, _EXT_AVAILABLE -from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" @@ -40,7 +43,7 @@ def forward(ctx, residual, y, gate, index): def backward(ctx, grad): y2, gate, index = ctx.saved_tensors g2 = grad.contiguous().view_as(y2) - positions = index.repeat(y2.shape[0] // index.shape[0]) + positions = _repeat_positions(index, y2.shape[0]) tiles = _segment_tiles(positions, gate.shape[0]) dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) record_backward( diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index e4c303e10..4dd9521f8 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -11,6 +11,12 @@ from __future__ import annotations +import os +import subprocess +import sys +import textwrap +from pathlib import Path + import pytest import torch @@ -113,3 +119,49 @@ def test_registry_dispatches_cuda(self): _cuda_op() op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") assert type(op).__name__ == "H3GateResidualCudaOp" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the gate table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_forward"): + sys.exit(77) + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + gate = torch.ones(2, 8, device="cuda") # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_gate_residual_forward(x, x, gate, index) + else: + i64 = dict(dtype=torch.int64, device="cuda") + _C.h3_gate_residual_backward( + torch.ones_like(x), x, gate, index, + torch.tensor([0, 2, 1, 3], **i64), torch.tensor([0, 2], **i64), + torch.tensor([2, 4], **i64), torch.tensor([0, 1, 2], **i64), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + if completed.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert completed.returncode != 0, "the kernel accepted a gate row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" From 1d26287f4550ce64a2f52ecc615fbdbde3d4ab3a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 17:36:20 +0000 Subject: [PATCH 47/62] fix(h3): check the position repeat in final_adaln_out backward and keep its forward out of the timed backward - final_adaln_out backward uses _repeat_positions() instead of floor division. - _final_perf had no backward setup, so the timed backward also ran the forward and built the graph (the guard test added with the gate-residual fix finds exactly this); the forward moves into the setup like in the other perf cases. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../model_specific/minimax_h3/final_adaln_out.py | 11 +++++++---- rl_engine/validation/models/h3_report.py | 15 +++++++++++---- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py index 2ff1fd16a..55993caad 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -8,9 +8,10 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward -from rl_engine.backends.extension 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.adaln_row_gather import ( + _repeat_positions, + _segment_tiles, +) from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( det_linear_available, det_linear_backward_input, @@ -19,6 +20,8 @@ silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS @@ -54,7 +57,7 @@ def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): def backward(ctx, grad): x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors shift, scale = table.chunk(2, dim=-1) - positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + positions = _repeat_positions(timestep_indices, x2.shape[0]) tiles = _segment_tiles(positions, table.shape[0]) dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( grad.contiguous().view_as(x2), diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 4e1033945..36eb1c713 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -717,9 +717,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 backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + def prepare_backward(fn, tensors=(x, nw, temb, w, b)): leaves = [t.detach().requires_grad_(True) for t in tensors] - fn(*leaves).backward(grad) + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) cases.append( { @@ -730,8 +733,12 @@ def backward(fn, tensors=(x, nw, temb, w, b), grad=grad): "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out(*t, i), - "candidate_backward": lambda bw=backward, i=ti: bw(lambda *t: op(*t, i)), - "provider_backward": lambda bw=backward, i=ti: bw( + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=ti: b( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda b=prepare_backward, i=ti: b( lambda *t: provider_final_adaln_out(*t, i) ), } From a2d0674bd8b8e975c3493abff6351f6429b7b8e2 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:16:56 +0000 Subject: [PATCH 48/62] docs(h3): gate-residual latencies from the backward-only rerun The table and the report links now come from the B200 report rerun at 239ef808, where the timed backward no longer includes the forward pass: backward 1.02 / 1.08 / 2.21 ms for S = 4097 / 32768 / 131072 against 0.55 / 2.24 / 8.92 ms for diffusers (the speedup at S = 131072 grows from 3.2x to 4.0x). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-adaln-gate-residual.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md index b35bc41c1..e869f4363 100644 --- a/docs/operators/h3-adaln-gate-residual.md +++ b/docs/operators/h3-adaln-gate-residual.md @@ -65,9 +65,9 @@ B200, B = 1, H = 5376, BF16: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.05 ms | 0.06 ms | 1.19 ms | 0.68 ms | -| 32768 | 0.24 ms | 0.41 ms | 1.37 ms | 2.71 ms | -| 131072 | 0.92 ms | 1.59 ms | 3.29 ms | 10.6 ms | +| 4097 | 0.06 ms | 0.07 ms | 1.02 ms | 0.55 ms | +| 32768 | 0.24 ms | 0.43 ms | 1.08 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.21 ms | 8.92 ms | The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and @@ -75,10 +75,11 @@ tile setup. ## Evidence -![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/figure.png) -The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), -written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `e7bc53a`. It also records +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr488/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `239ef808` on an otherwise idle B200 (backward timings exclude the leaf +creation, the forward pass and the graph construction; the execution order alternates and is recorded). It also records forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. ## Existing implementations (RFC #420 reuse rule) From 337ee6041a96d5f8691a58d2b4eeea27420c8fc0 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 20:17:08 +0000 Subject: [PATCH 49/62] docs(h3): final_adaln_out latencies from the backward-only rerun The table and the report links now come from the B200 report rerun at 1d26287f (pinned norm_out weights), where the timed backward no longer includes the forward pass: backward 1.38 / 2.56 / 7.04 ms for S = 4097 / 32768 / 131072 against 0.86 / 4.47 / 17.7 ms for diffusers. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 73eac42d6..90fbbfd17 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -62,22 +62,23 @@ B200, B = 1, T = 3, pinned `norm_out` weights: | S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | | --- | --- | --- | --- | --- | -| 4097 | 0.18 ms | 0.19 ms | 1.50 ms | 1.07 ms | -| 32768 | 0.44 ms | 0.84 ms | 2.75 ms | 5.34 ms | -| 131072 | 1.29 ms | 2.99 ms | 7.76 ms | 20.8 ms | +| 4097 | 0.18 ms | 0.18 ms | 1.38 ms | 0.86 ms | +| 32768 | 0.42 ms | 0.82 ms | 2.56 ms | 4.47 ms | +| 131072 | 1.28 ms | 2.96 ms | 7.04 ms | 17.7 ms | At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the projection backward. ## Evidence -![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/figure.png) +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr489/reports/experiments/h3-final-adaln-out-b200/figure.png) There are two data files: -- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/3848e6fcb227415ce8ad4d21ca120701f58aacf6/pr489/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean - tree at `a4bc99f9` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + tree at `1d26287f` on an otherwise idle B200 (backward timings exclude the leaf creation, the + forward pass and the graph construction; the previous revision came from `a4bc99f9`), with FP64 leaves and upstream gradients, after the golden's declared cast was changed to round the FP32 SiLU as the model does. The backward reference keeps only the SiLU and table roundings, so the errors include the BF16 rounding of `norm(x)` and `1 + scale`. From 1857fa69aed91b8333b9ba7cb16e985d0bfd1eba Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:58:21 +0000 Subject: [PATCH 50/62] fix(h3): check the gate backward's tile metadata and lengths, and scan the reference's rows once gate_residual.cu: the backward runs the shared tile-metadata pre-pass (tile_metadata.cuh, #486) before the segment kernels, so sorted positions, tiles and segment boundaries that do not describe the rows end in a device-side assert instead of an out-of-bounds read, including the case of no tiles with a non-trivial seg_first_tile (the partial buffer then holds one uninitialized row). Host checks: sorted_pos has M entries, tile_begin and tile_end have the same length, seg_first_tile has at least one entry, and the column blocks fit grid.y (65535); the fold launch is skipped for zero segments. Native gate backward, idle B200, min of 400: 86.6 -> 90.7 us at S = 4097, 362.7 -> 367.5 us at S = 32768, 1315.0 -> 1319.7 us at S = 131072 (one extra launch of about 4 us). Tests run corrupted metadata in a child process and the host checks. gate_residual.py: dy is viewed as the shape of y (validate_h3_gate_residual makes it the residual's shape, so nothing changes for valid calls) and the saved shapes are named for the tensor they belong to. Reference: the gate gradient takes the positions of all rows from positions_by_row (one sort and one host read instead of one torch.nonzero scan and sync per row); a test compares the gradient bitwise with the loop it replaces. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 22 ++++- .../minimax_h3/gate_residual.py | 6 +- .../reference/minimax_h3/gate_residual.py | 5 +- .../minimax_h3/test_h3_adaln_gate_residual.py | 87 +++++++++++++++++++ 4 files changed, 112 insertions(+), 8 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 6e0b5a85b..123b434fa 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -27,6 +27,8 @@ #include #include +#include "tile_metadata.cuh" + namespace { __device__ __forceinline__ float to_f(float v) { return v; } @@ -216,8 +218,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: } const int64_t rows = grad.size(0); const int64_t n = grad.size(1); + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); const int64_t segments = seg_first_tile.numel() - 1; TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + TORCH_CHECK(sorted_pos.numel() == rows, "sorted_pos must have M entries"); + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); const int64_t tiles = tile_begin.numel(); const c10::cuda::CUDAGuard guard(grad.device()); auto stream = at::cuda::getCurrentCUDAStream(); @@ -225,7 +230,14 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); auto dgate = torch::empty({segments, n}, grad.options()); const int threads = 256; - const unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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 tile list is caller-built metadata: check it before the kernels read through it. + h3_tiles::check_tile_metadata(stream, sorted_pos.data_ptr(), rows, + tile_begin.data_ptr(), tile_end.data_ptr(), tiles, + seg_first_tile.data_ptr(), segments); H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { const T* gp = reinterpret_cast(grad.data_ptr()); const T* yp = reinterpret_cast(y.data_ptr()); @@ -237,9 +249,11 @@ std::vector h3_gate_residual_backward(torch::Tensor grad, torch:: gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), partial.data_ptr(), n); } - gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( - partial.data_ptr(), seg_first_tile.data_ptr(), - reinterpret_cast(dgate.data_ptr()), n); + if (segments > 0) { // a zero-sized grid is not a valid launch + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + } }); C10_CUDA_KERNEL_LAUNCH_CHECK(); return {dy, dgate}; diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index ad9bb6d8f..04f6bdf7d 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -36,7 +36,9 @@ def forward(ctx, residual, y, gate, index): y2 = y.contiguous().view(-1, hidden) out = _C.h3_gate_residual_forward(res2, y2, gate, index) ctx.save_for_backward(y2, gate, index) - ctx.shape = residual.shape + # validate_h3_gate_residual makes y and residual one shape; each gradient is still viewed + # as the shape of its own input. + ctx.residual_shape, ctx.y_shape = residual.shape, y.shape return out.view(residual.shape) @staticmethod @@ -49,7 +51,7 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.shape), dgate, None + return grad, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py index f68a0fdf4..0ad2c1b0e 100644 --- a/rl_engine/reference/minimax_h3/gate_residual.py +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -18,6 +18,8 @@ import torch +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + _DTYPES = (torch.float32, torch.bfloat16, torch.float16) @@ -69,8 +71,7 @@ def backward(ctx, grad): contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) - for row in range(ctx.gate_rows): - positions = torch.nonzero(index == row).flatten() + for row, positions in enumerate(positions_by_row(index, ctx.gate_rows)): if positions.numel(): d_gate[row] = per_position.index_select(0, positions).sum(dim=0) return grad, grad * gathered, d_gate.to(gathered.dtype), None diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 4dd9521f8..b26ab8a54 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -52,6 +52,25 @@ def _gate(table, chunk=2): class TestReference: + def test_gate_backward_is_bitwise_the_per_row_scan_it_replaced(self): + from rl_engine.reference.minimax_h3.gate_residual import _DeterministicGateResidual + + g = torch.Generator(device="cpu").manual_seed(5) + residual, y = torch.randn(2, 300, 16, generator=g), torch.randn(2, 300, 16, generator=g) + gate = torch.randn(7, 16, generator=g) + index = torch.randint(0, 5, (300,), generator=g) # rows 5 and 6 are never referenced + grad = torch.randn(2, 300, 16, generator=g) + leaf = gate.clone().requires_grad_(True) + _DeterministicGateResidual.apply(residual, y, leaf, index).backward(grad) + + per_position = (grad * y).reshape(-1, 300, 16).sum(dim=0) + expected = torch.zeros(7, 16) + for row in range(7): # the loop of the earlier implementation + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + expected[row] = per_position.index_select(0, positions).sum(dim=0) + assert torch.equal(leaf.grad, expected) + def test_rejects_bad_inputs(self): op = NativeH3GateResidualOp() res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) @@ -165,3 +184,71 @@ def test_gate_row_outside_the_table_asserts_in_the_kernel(entrypoint, bad_row): assert ( "device-side assert triggered" in completed.stderr ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + if not hasattr(_C, "h3_gate_residual_backward"): + sys.exit(77) + mutation = sys.argv[1] + x = torch.randn(4, 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 == "no_tiles": + begin, end = begin[:0], end[:0] + elif mutation == "short_sorted_pos": + sorted_pos = sorted_pos[:3] + elif mutation == "no_boundary": + first = first[:0] + grad = torch.ones_like(x) + _C.h3_gate_residual_backward(grad, x, gate, index, sorted_pos, begin, end, first) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(): + control = _gate_backward_in_child("none") + if control.returncode == 77: + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + assert control.returncode == 0, control.stderr + for mutation in ( + "tile_end_past_rows", + "position_past_rows", + "segments_past_tiles", + "no_tiles", + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) + # Lengths are checked on the host, before any launch. + for mutation, message in ( + ("short_sorted_pos", "sorted_pos must have M entries"), + ("no_boundary", "seg_first_tile must hold R + 1 entries"), + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0 and message in done.stderr, (mutation, done.stderr) From fcf0a588b191d7b321757911e9e9331de26e4c7a Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:04:07 +0000 Subject: [PATCH 51/62] fix(h3): run the final AdaLN backward only for the inputs that need a gradient, refuse pre-SM80 at construction, fix three doc statements final_adaln_out.py: backward reads ctx.needs_input_grad. The projection's weight and bias folds run only when one of them needs a gradient (the bias fold only when the bias does), d_temb (a SiLU VJP over a full input-gradient pass) only when temb does, and every input that does not need a gradient gets None instead of a tensor; a test freezes each input in turn and compares the others bitwise with the all-trainable gradients. H3FinalAdaLNOutCudaOp.__init__ calls det_linear_device_supported() like the projection and modulation ops, so the registry falls back to the PyTorch reference on a pre-SM80 device instead of every forward raising. docs/operators/h3-final-adaln-out.md: the projection is split with chunk(2, dim=-1); the gtest sentence says the gtest failed on d_norm_w from S = 257 and on d_scale/dW/d_temb from S = 1024 because the golden did not round where the model does; the gradient statement of the bitwise comparison names the exception recorded in the report (all five gradients differ at S = 257, 4097 and 32768; at S = 1 the bias gradient is equal). Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- docs/operators/h3-final-adaln-out.md | 9 ++--- .../minimax_h3/final_adaln_out.py | 29 +++++++++++++--- .../minimax_h3/test_h3_final_adaln_out.py | 33 +++++++++++++++++++ 3 files changed, 63 insertions(+), 8 deletions(-) diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md index 90fbbfd17..067e8ec1e 100644 --- a/docs/operators/h3-final-adaln-out.md +++ b/docs/operators/h3-final-adaln-out.md @@ -6,7 +6,7 @@ blocks (RFC #420, WS1 step 7): ```text -shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2, dim=-1) # (T, 5376) each, shift first out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] ``` @@ -49,8 +49,9 @@ out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timeste cast, `norm_out.linear`'s BF16 output table, `norm_out.norm`'s BF16 output and the BF16 `1 + scale`. All four are rounded straight-through for the gradient. The table and `1 + scale` are shared by every position of a timestep. Without them the golden's `d_norm_w` drifts - systematically with S, and the gtest missed it from S = 257. `norm_out.norm`'s rounding - enters `d_scale`, `dW` and `d_temb` summed over S, and the gtest missed those from S = 1024. + systematically with S, and the gtest failed on `d_norm_w` from S = 257 because the golden did not + round where the model does. `norm_out.norm`'s rounding enters `d_scale`, `dW` and `d_temb` + summed over S, and the gtest failed on those from S = 1024 for the same reason. ## Performance Notes @@ -121,7 +122,7 @@ gradient with diffusers' path and with this kernel. ![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) -There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and the gradients differ as expected for the deterministic FP32 reductions: all five at S = 257, 4097 and 32768, and all but the bias gradient at S = 1, where one position leaves a single term to sum. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `8cb4abed` by ```bash diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py index 55993caad..042e50f7d 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -16,6 +16,7 @@ det_linear_available, det_linear_backward_input, det_linear_backward_weight, + det_linear_device_supported, det_linear_forward, silu_backward_fp32, ) @@ -56,6 +57,7 @@ def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): @staticmethod def backward(ctx, grad): x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + need_x, need_norm_weight, need_temb, need_weight, need_bias = ctx.needs_input_grad[:5] shift, scale = table.chunk(2, dim=-1) positions = _repeat_positions(timestep_indices, x2.shape[0]) tiles = _segment_tiles(positions, table.shape[0]) @@ -69,14 +71,29 @@ def backward(ctx, grad): timestep_indices, *tiles, ) - # The table gradient stays FP32 into the projection backward. + # The table gradient stays FP32 into the projection backward; the projection's weight, + # bias and input folds run only for the inputs that need a gradient. d_table = torch.cat([d_shift, d_scale], dim=1) - dw, db = det_linear_backward_weight(d_table, act, weight.dtype) - d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + d_temb = dw = db = None + if need_weight or need_bias: + dw, *rest = det_linear_backward_weight(d_table, act, weight.dtype, with_bias=need_bias) + db = rest[0] if need_bias else None + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) record_backward( "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + return ( + dx.view(ctx.x_shape) if need_x else None, + d_norm_weight if need_norm_weight else None, + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + ) class H3FinalAdaLNOutCudaOp: @@ -95,6 +112,10 @@ class H3FinalAdaLNOutCudaOp: def __init__(self) -> None: if not (det_linear_available() and h3_rmsnorm_available()): raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + if not det_linear_device_supported(): + # The BF16 projection is an mma.sync kernel; refuse here so the registry falls back to + # the PyTorch reference instead of failing in every forward. + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py index 5fb964677..dc2bb8153 100644 --- a/tests/models/minimax_h3/test_h3_final_adaln_out.py +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -130,6 +130,39 @@ def grads(fn, dtype=None): rel = ((g.double() - r).abs().max() / r.abs().max()).item() assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + @pytest.mark.parametrize( + "frozen", ["x", "norm_weight", "temb", "weight", "bias", "weight_and_bias", "temb_and_bias"] + ) + def test_frozen_inputs_get_no_gradient_and_the_others_are_unchanged(self, norm_out, frozen): + x, temb, ti = _inputs(1500, seed=6) + grad = torch.randn_like(x) + names = ("x", "norm_weight", "temb", "weight", "bias") + tensors = (x, norm_out[0], temb, *norm_out[1:]) + + def run(frozen_names): + leaves = [ + t.detach().clone().requires_grad_(name not in frozen_names) + for name, t in zip(names, tensors) + ] + _cuda_op()(*leaves, ti).backward(grad) + return [leaf.grad for leaf in leaves] + + full = run(()) + partial = run(frozen.split("_and_") if "_and_" in frozen else [frozen]) + frozen_set = frozen.split("_and_") if "_and_" in frozen else [frozen] + for name, got, expected in zip(names, partial, full): + if name in frozen_set: + assert got is None, name + else: + assert torch.equal(got, expected), name + + def test_unsupported_device_is_refused_at_construction(self, monkeypatch): + from rl_engine.backends.cuda.model_specific.minimax_h3 import final_adaln_out as module + + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3FinalAdaLNOutCudaOp() + def test_report_backward_uses_fp64_golden(self, monkeypatch): from types import SimpleNamespace From b3747d90ec5db2e4bd4a20c23cd4cbc9fd9d62b7 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:18:32 +0000 Subject: [PATCH 52/62] fix(h3): refuse pre-SM80 before the first collective, check the TP table gradient, handle a TP1-only figure, say which TP sizes run tp_adaln_projection.py: the op refuses a device without SM80 when it is built (so the registry can fall back) and again, for BF16 weights, in the input validation that runs before the first collective, so a rank that cannot run the BF16 forward fails before the other ranks start waiting in gather_rows (a test counts the collective calls). The backward checks that the table gradient is (T, N) before it slices this rank's columns; autograd already checks the shape of a gradient that reaches the node through the graph, so the check guards a direct call of the node's backward, which the test makes. det_linear.cu: h3_det_linear_fold_chunks checks the grid.x block count like h3_det_linear_backward_input (a bound of 5e11 elements, so it cannot be exercised with real memory). plot_h3_evidence.py: a TP evidence report without a run of TP > 1 (--worlds 1) no longer fails in the bar chart on an empty sequence; the right panel says so. Tested for TP1-only, TP1/2/4 and TP2/4 reports. docs/operators/h3-tp-adaln-3mod.md: TP 3, 6 and 7 pass the shard-alignment check but only 1, 2, 4 and 8 can run, because DeterministicCollective rejects the other world sizes when it is constructed. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/det_linear.cu | 4 ++ docs/operators/h3-tp-adaln-3mod.md | 6 ++- .../minimax_h3/tp_adaln_projection.py | 18 +++++++- .../minimax_h3/test_h3_bitwise_reuse_plot.py | 41 ++++++++++++++++++- tests/models/minimax_h3/test_h3_tp_adaln.py | 40 ++++++++++++++++++ tools/validation/models/plot_h3_evidence.py | 13 ++++++ 6 files changed, 118 insertions(+), 4 deletions(-) diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index cc84bfa5a..56e6be84f 100644 --- a/csrc/cuda/h3/det_linear.cu +++ b/csrc/cuda/h3/det_linear.cu @@ -561,6 +561,10 @@ torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType o if (elems == 0) return out; auto stream = at::cuda::getCurrentCUDAStream(); const int threads = 256; + // Same bound as h3_det_linear_backward_input: grid.x holds at most INT32_MAX blocks. (That is 5e11 + // elements, beyond any tensor that fits in memory; the check keeps the two entry points alike.) + TORCH_CHECK((elems + threads - 1) / threads <= std::numeric_limits::max(), + "rows * K = ", elems, " needs more blocks than one grid.x launch holds"); const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); if (out_dtype == at::kFloat) { det_fold_chunks_kernel<<>>( diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md index f1b1a2409..6487c6a8d 100644 --- a/docs/operators/h3-tp-adaln-3mod.md +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -28,8 +28,10 @@ op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the `N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. `N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard -never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 qualify. Anything else fails closed when -the op is built. +never splits a WS1 partial sum. TP 1, 2, 3, 4, 6, 7 and 8 pass this alignment check and anything +else fails closed when the op is built. Only 1, 2, 4 and 8 can run: `DeterministicCollective` +supports exactly those world sizes and rejects a 3-, 6- or 7-way group when it is constructed, so +those sizes are covered by the ownership tests only. | TP | Columns per rank | Rank 0 owns | | --- | --- | --- | 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 d262ac030..371f0ae74 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 @@ -31,18 +31,19 @@ import torch import torch.nn.functional as F -from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( CONTRACT, DINPUT_CHUNK, det_linear_available, det_linear_backward_input_partials, det_linear_backward_weight, + det_linear_device_supported, det_linear_fold_chunks, det_linear_forward, silu_backward_fp32, ) from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table @@ -122,6 +123,10 @@ def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: raise TypeError( f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" ) + if weight.dtype == torch.bfloat16 and not det_linear_device_supported(weight.device): + # Checked here, before any collective: a rank that fails inside the forward would leave the + # other ranks waiting in gather_rows. + raise RuntimeError(f"the BF16 h3 det_linear forward needs SM80 or newer ({weight.device})") width = shard.end - shard.begin if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): raise ValueError( @@ -148,6 +153,13 @@ def forward(ctx, temb, weight_shard, bias_shard, collective, shard): def backward(ctx, grad_table): temb, act, weight = ctx.saved_tensors shard = ctx.shard + if grad_table.dim() != 2 or grad_table.shape != (temb.shape[0], shard.n_total): + # The backward reads this rank's columns of a replicated (T, N) table gradient; any + # other width would slice the wrong columns without an error. + raise ValueError( + f"table gradient must be ({temb.shape[0]}, {shard.n_total}), " + f"got {tuple(grad_table.shape)}" + ) grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) partial = det_linear_backward_input_partials(grad_local, weight) @@ -176,6 +188,10 @@ def __init__(self, collective, n_total: int) -> None: raise RuntimeError( "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" ) + if not det_linear_device_supported(): + # Best effort at construction so the registry can fall back; forward checks the + # weight's own device before the first collective. + raise RuntimeError("the BF16 h3 det_linear forward needs SM80 or newer") self.collective = collective self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) 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 ffb87bd65..e386f4ec4 100644 --- a/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py +++ b/tests/models/minimax_h3/test_h3_bitwise_reuse_plot.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""A heatmap cell counts every forward output of the operator, not only the first.""" +"""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. +""" from __future__ import annotations @@ -71,3 +74,39 @@ def result(frac_differ): texts = [t.get_text() for t in ax.texts] assert "n/a" in texts and "0" in texts plt.close(fig) + + +def _tp_run(tp: int) -> dict: + perf = { + "num_timesteps": 3, + "tp_forward_us": 110.0, + "tp_forward_backward_us": 2000.0, + "ws1_forward_us": 100.0, + "ws1_forward_backward_us": 1900.0, + } + equal = dict.fromkeys(("table", "d_temb", "d_weight_shard", "d_bias_shard"), True) + return { + "tp": tp, + "ranks": [{"rank": 0, "equality": [{"num_timesteps": 3, **equal}], "perf": [perf]}], + } + + +def _tp_report(worlds: list[int]) -> dict: + return { + "op": "tp_adaln_3mod", + "environment": {"gpu": "gpu", "torch": "t", "cuda": "c"}, + "rl_kernel_commit": "0123456789", + "model_revision": "0123456789", + "runs": [_tp_run(tp) for tp in worlds], + } + + +@pytest.mark.parametrize("worlds", [[1], [1, 2, 4], [2, 4]]) +def test_tp_figure_for_reports_with_and_without_tp_above_one(worlds): + import matplotlib.pyplot as plt + from plot_h3_evidence import plot_tp_adaln + + fig = plot_tp_adaln(_tp_report(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) diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py index 2a62aa502..6d3737824 100644 --- a/tests/models/minimax_h3/test_h3_tp_adaln.py +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -128,6 +128,46 @@ def test_fails_closed(self): H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) assert op.readback()["columns"] == [2304, 4608] + def test_pre_sm80_is_refused_at_construction_and_before_the_first_collective(self, monkeypatch): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3 import tp_adaln_projection as module + + calls = [] + collective = SimpleNamespace( + rank=0, world_size=2, backend_id="unused", all_gather=lambda t: calls.append(t) + ) + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + op = module.H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w, b = module.shard_adaln_projection(weight, bias, 2, 0) + monkeypatch.setattr(module, "det_linear_device_supported", lambda *args: False) + with pytest.raises(RuntimeError, match="SM80"): + module.H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + with pytest.raises(RuntimeError, match="SM80"): + op(temb, w, b) + assert not calls # no rank is left waiting in a collective + + def test_backward_rejects_a_table_gradient_of_the_wrong_width(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + _TPAdaLNProjection, + shard_adaln_projection, + ) + + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(2, 3 * 6 * 256, 136) + op = H3TPAdaLNProjectionCudaOp(one, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 1, 0) + leaf = temb.clone().requires_grad_(True) + table = op.forward_table(leaf, w, b) + node = table.grad_fn # the autograd node is the ``ctx`` its backward receives + # Autograd checks the shape of a gradient that reaches this node through the graph, so the + # check guards a direct call of the node's backward, which this makes. + for bad in (grad[:, :-64], grad[:1]): + with pytest.raises(ValueError, match="table gradient must be"): + _TPAdaLNProjection.backward(node, bad.to(table.dtype)) + assert _TPAdaLNProjection.backward(node, grad.to(table.dtype))[0].shape == temb.shape + @pytest.mark.parametrize("world", [2, 4, 8]) @pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py index 564cc3ce6..35688f318 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/plot_h3_evidence.py @@ -599,6 +599,19 @@ def worst(run, key): tp_runs = [run for run in runs if run["tp"] > 1] _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + if not tp_runs: + # A TP1-only report (--worlds 1) has nothing to compare with the single-GPU WS1 call. + right.text( + 0.5, + 0.5, + "no run with TP > 1 in this report", + transform=right.transAxes, + ha="center", + va="center", + fontsize=10, + color=MUTED, + ) + return fig _bars( right, [f"TP{run['tp']}" for run in tp_runs], From 5d50d34388e3e79dfd76101d741513d85f5c375d Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Mon, 26 Oct 2026 23:27:58 +0000 Subject: [PATCH 53/62] fix(h3): check the lists of the sequence-parallel native entry points, and pin that the exchange padding is never read The SP entry points take caller-built row lists and tile lists, which the pre-pass of the other backward kernels did not cover. h3_gate_grad_partials, h3_gate_grad_fold, h3_rmsnorm_backward_partials and h3_rmsnorm_fold_partials now run a pre-pass (h3_tiles::check_lists in tile_metadata.cuh, one small launch) that asserts on the device that every listed row is a row of the activation, every tile lies inside its list, and the segment boundaries are non-decreasing and inside the partial buffer; they check that the column blocks fit grid.y (65535) and N fits the int indices of the segment kernels, require at least one segment boundary, and skip a fold launch with no segments. One test per entry point runs the corrupted lists in a child process and expects the device assert; the rmsnorm table-layout tests of #487 (overlapping rows rejected, a view of a wider table read by its row pitch) are ported into test_h3_native_validation.py. sp_norm_adaln: no code change. The padding slots of the send blocks and of the partial gathers hold a valid table row but no tile list refers to them; a test now pads both with NaN (sp = 3, 4 and 8, B = 1 and 2, shards smaller than a tile) and the gradients still equal WS1 bitwise, so a read of the padding would show. The plain-norm backward passes None segment tiles because the binding takes optionals; it is covered by test_plain_norm_ranks_reproduce_ws1, and plan.families[1] is only read when the plan was built with an index. The local-slice validator receives x with S_local positions and the matching slice of the index, which is what validate_h3_modulation checks (S = len(index)), against the full tables for the range check. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/adaln_row_gather.cu | 5 +- csrc/cuda/h3/gate_residual.cu | 19 ++- csrc/cuda/h3/rmsnorm_modulate.cu | 26 +++- csrc/cuda/h3/tile_metadata.cuh | 41 ++++++ .../minimax_h3/test_h3_adaln_row_gather.py | 6 +- .../minimax_h3/test_h3_native_validation.py | 121 +++++++++++++++++- .../minimax_h3/test_h3_sp_norm_adaln.py | 75 +++++++++++ 7 files changed, 278 insertions(+), 15 deletions(-) diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index dbff808e4..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"); diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 8439c037b..f9cd6db6f 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -280,7 +280,15 @@ torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch:: auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); if (tiles == 0) return partial; const int threads = 256; - const unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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()>>>( @@ -302,12 +310,19 @@ torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_t 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 unsigned col_blocks = static_cast((n + threads - 1) / threads); + 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()>>>( diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 6dbdb31a4..25e907b70 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -505,7 +505,14 @@ std::vector h3_rmsnorm_backward_partials( 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); if (tiles > 0) { H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, @@ -523,6 +530,9 @@ std::vector h3_rmsnorm_backward_partials( 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) { @@ -572,14 +582,22 @@ std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, to 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); - 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(); + 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)}; } diff --git a/csrc/cuda/h3/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index a959054e2..7cda74cc4 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -58,4 +58,45 @@ 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) { + const int64_t items = std::max({n_rows, tiles, segments + 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(); +} + } // namespace h3_tiles 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 26a9c7921..65a927577 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -393,7 +393,7 @@ def test_the_backward_stops_on_metadata_that_does_not_describe_the_tiles(self): assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) - 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() @@ -401,8 +401,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_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py index bcb06f1f2..1c8040e5d 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -286,7 +286,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -318,7 +319,8 @@ def test_native_table_bounds_fail_in_isolated_process(native_extension, operatio sys.exit(0) raise raise AssertionError("invalid index reached native table lookup without an assertion") - """) + """ + ) result = subprocess.run( [sys.executable, "-c", script, operation, str(bad_index)], capture_output=True, @@ -352,7 +354,8 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada import sys import textwrap - script = textwrap.dedent(""" + script = textwrap.dedent( + """ import sys import torch from rl_engine import _C @@ -372,9 +375,119 @@ def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metada 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) +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) +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_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/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py index 4866e68f3..b3213fd1e 100644 --- a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -165,6 +165,79 @@ def rank(coll): ) 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 @@ -284,7 +357,9 @@ def _worker_never_reports(rank, world, init, queue, path, target, kwargs): @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.models import h3_ws2 workers = { From ab407b39ccc7f09066ce7300a00c8ff17c9c4401 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:38:27 +0000 Subject: [PATCH 54/62] fix(h3): reject gate tables whose rows overlap, and state the residual gradient shape * `make_gate` refuses a table whose row stride is below N (an expanded table would make every row alias the same memory), as the modulation tables already did; `H3GateResidualCudaOp` copies such a table to dense rows first, so an expanded gate still works and gives the bytes of its dense copy. Views of a wider AdaLN table (unit column stride, larger row stride) go to the kernel unchanged. * The residual's VJP is the incoming cotangent, which already has the residual's shape because forward returns `out.view(residual.shape)`; `ctx.residual_shape` was saved and unused, which read like a missing reshape. The return is now `grad.reshape(ctx.residual_shape)` (a no-op for equal shapes, and it accepts a non-contiguous gradient where `view` would not), and a test feeds a (B, S, H) residual with a transposed cotangent and compares every gradient with the eager VJP bitwise. * `_C.pyi`: contracts for `h3_gate_residual_forward` / `_backward`. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- csrc/cuda/h3/gate_residual.cu | 4 ++ .../minimax_h3/gate_residual.py | 10 ++++- .../minimax_h3/test_h3_adaln_gate_residual.py | 37 +++++++++++++++++++ 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu index 123b434fa..81051fcc6 100644 --- a/csrc/cuda/h3/gate_residual.cu +++ b/csrc/cuda/h3/gate_residual.cu @@ -149,6 +149,10 @@ Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torc "gate must be a CUDA tensor with the activations' dtype"); TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, "gate must be (R, N) with unit column stride"); + // Rows are read as dense N-element spans at `stride(0)`: a view of a wider table is fine, an + // expanded or overlapping one (stride(0) < N) is not a table of rows. + TORCH_CHECK(gate.size(0) <= 1 || gate.stride(0) >= gate.size(1), "gate row stride ", + gate.stride(0), " is smaller than N=", gate.size(1)); TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && index.dim() == 1 && index.is_contiguous() && index.numel() > 0, "index must be a non-empty contiguous int64 tensor on ", x.device()); diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index 04f6bdf7d..da7319bfb 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -51,7 +51,9 @@ def backward(ctx, grad): record_backward( "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return grad, dy.view(ctx.y_shape), dgate, None + # The cotangent has the output's shape, which is the residual's; the residual's VJP is that + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. + return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: @@ -80,7 +82,11 @@ def forward(self, residual, y, gate, index, *, check_range: bool = True): validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) if not residual.is_cuda: raise ValueError("H3GateResidualCudaOp needs CUDA tensors") - if gate.stride(1) != 1: + # A view of a wider table (``table.view(R, 6, H)[:, 2]``: unit column stride, row stride 6H) + # goes to the kernel as it is, which indexes rows by the stride. Anything that is not a + # table of dense rows (non-unit columns, or rows that overlap because the stride is below + # H, as for an expanded table) is copied first. + if gate.stride(1) != 1 or (gate.shape[0] > 1 and gate.stride(0) < gate.shape[1]): gate = gate.contiguous() return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index b26ab8a54..193abce60 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -125,6 +125,43 @@ def grads(fn, dtype=None): assert rel < 5e-3 # only the BF16 rounding of the table gradient assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + def test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_gradients(self): + """A (B, S, H) residual gets a (B, S, H) gradient even when the incoming one is a view.""" + + table, residual, y, index = _case(300, hidden=64, seed=5) + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # a non-contiguous cotangent: the transpose of a (B, H, S) tensor + grad = torch.randn(2, 64, 300, device="cuda").bfloat16().transpose(1, 2) + assert not grad.is_contiguous() and grad.shape == out.shape + out.backward(grad) + ref = [t.detach().clone().requires_grad_(True) for t in (residual, y, table)] + provider_gate_residual(ref[0], _gate(ref[2]), index, ref[1]).backward(grad) + assert leaves[0].grad.shape == residual.shape and leaves[1].grad.shape == y.shape + assert torch.equal(leaves[0].grad, ref[0].grad) + assert torch.equal(leaves[1].grad, ref[1].grad) + + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): + """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" + + from rl_engine.backends.extension import _C + + hidden = 64 + row = (torch.randn(1, hidden, device="cuda") * 0.5).bfloat16() + expanded = row.expand(9, hidden) # stride (0, 1): all nine rows alias one + assert expanded.stride(0) == 0 + _, residual, y, index = _case(100, hidden=hidden, seed=6) + index = index % 9 + dense = expanded.contiguous() + got = _cuda_op()(residual, y, expanded, index) + assert torch.equal(got, _cuda_op()(residual, y, dense, index)) + assert torch.equal(got, provider_gate_residual(residual, dense, index, y)) + # the native entry point refuses the layout instead of reading overlapping rows + res2, y2 = residual.reshape(-1, hidden), y.reshape(-1, hidden) + with pytest.raises(RuntimeError, match="row stride"): + _C.h3_gate_residual_forward(res2, y2, expanded, index) + _C.h3_gate_residual_forward(res2, y2, dense, index) # the dense copy is accepted + def test_rejects_cpu_and_bad_index(self): table, residual, y, index = _case(10, hidden=16) with pytest.raises(IndexError): From 9aa200d236bd7410342c021067575db81b1f1a44 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:46:07 +0000 Subject: [PATCH 55/62] docs(h3): document the det_linear partial and fold stubs used by the tensor-parallel path h3_det_linear_backward_input_partials returns the (C, T, K) FP32 per-chunk partials (64 output columns per chunk, ascending), h3_det_linear_fold_chunks left-folds them over C and casts once; both were bare signatures in _C.pyi. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/_C.pyi | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 5408f1557..c910add64 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -374,10 +374,23 @@ def h3_det_linear_backward_input( """ ... -def h3_det_linear_backward_input_partials( - grad: torch.Tensor, weight: torch.Tensor -) -> torch.Tensor: ... -def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... +def h3_det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Return the ``(C, T, K)`` FP32 per-chunk partials of ``grad @ weight`` before the fold. + + Same operands as ``h3_det_linear_backward_input``. Chunk ``c`` of ``C = ceil(N / 64)`` sums the + output columns ``64c .. 64c + 63`` in ascending order; the tensor-parallel path reduces its + shard's chunks with ``h3_det_linear_fold_chunks`` after the shards are gathered. + """ + ... + +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Left-fold ``partial (C, T, K)`` FP32 over ``C`` in ascending order, cast once to ``out_dtype``. + + Require a non-empty contiguous FP32 CUDA tensor and an FP32 or BF16 ``out_dtype``; return + ``(T, K)``. + """ + ... + def h3_det_linear_backward_weight( grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True ) -> list[torch.Tensor]: @@ -425,6 +438,7 @@ def h3_adaln_row_gather_backward( inside ``sorted_pos``, segments that cut the tile list exactly and in order. """ ... + def h3_rmsnorm_forward( x: torch.Tensor, weight: torch.Tensor, @@ -469,6 +483,7 @@ def h3_rmsnorm_backward( modulation they must be omitted. """ ... + def h3_gate_residual_forward( residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor ) -> torch.Tensor: @@ -501,4 +516,3 @@ def h3_gate_residual_backward( checked on the device before use. """ ... - From 002fb393da77917cccb92cf70a2390a38019cea8 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 02:52:23 +0000 Subject: [PATCH 56/62] fix(h3): check the modulation index in the sequence-parallel partials, floor the list check's launch, pin the no-H3-symbols build * `h3_rmsnorm_backward_partials` runs `h3_tiles::check_index_range` over the S index values before its dweight kernel, which reads them inside its row loop without an assert (an assert there costs about 11% of the backward). This entry point can be called without the dx entry point having run, so it cannot lean on dx's per-row assert. The native-validation child-process test calls it with an index past the table and a negative one (and a valid control) and requires the device-side assert. * `check_lists` launches at least one block, like `check_tile_metadata`. * `test_h3_extension_without_h3_symbols.py`: with an extension that lacks every H3 binding (what a ROCm, fast-math or relaxed-flag build produces), every CUDA op, including the sequence-parallel one, refuses to construct, so the registry falls back instead of failing with AttributeError; the SP availability probe is False when any single symbol it calls is missing. * `_C.pyi`: contracts for the sequence-parallel entry points. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../test_h3_extension_without_h3_symbols.py | 73 +++++++++++++++++++ .../minimax_h3/test_h3_native_validation.py | 11 +++ 2 files changed, 84 insertions(+) create mode 100644 tests/models/minimax_h3/test_h3_extension_without_h3_symbols.py 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 1c8040e5d..657f0b4cb 100644 --- a/tests/models/minimax_h3/test_h3_native_validation.py +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -432,6 +432,10 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension 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) @@ -450,6 +454,12 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension 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() """ @@ -459,6 +469,7 @@ def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension "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",), } From 60e7785c754a2f08b4f2e207977430063a181dd1 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Tue, 27 Oct 2026 11:05:10 +0000 Subject: [PATCH 57/62] fix(h3): return no residual gradient for a residual that needs none The gate residual backward handed back the cotangent as the residual's gradient whether or not the residual required one. It now follows needs_input_grad[0], like the other ports of this group. The tile-metadata mutation test gains a gap, an overlap and a short last tile, and expects the empty tile list to be refused on the host. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/gate_residual.py | 6 ++-- .../minimax_h3/test_h3_adaln_gate_residual.py | 31 ++++++++++++++++++- 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py index da7319bfb..f39cf351b 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -52,8 +52,10 @@ def backward(ctx, grad): "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) # The cotangent has the output's shape, which is the residual's; the residual's VJP is that - # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. - return grad.reshape(ctx.residual_shape), dy.view(ctx.y_shape), dgate, None + # cotangent itself. reshape (not view) so a non-contiguous gradient is accepted. A residual + # that needs no gradient (a frozen skip path) gets none. + d_residual = grad.reshape(ctx.residual_shape) if ctx.needs_input_grad[0] else None + return d_residual, dy.view(ctx.y_shape), dgate, None class H3GateResidualCudaOp: diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py index 193abce60..f93561263 100644 --- a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -141,6 +141,26 @@ def test_residual_gradient_keeps_the_residual_shape_for_batched_and_strided_grad assert torch.equal(leaves[0].grad, ref[0].grad) assert torch.equal(leaves[1].grad, ref[1].grad) + def test_no_residual_gradient_is_returned_for_a_residual_that_needs_none(self): + """The backward node hands back ``None`` for a frozen residual, the cotangent otherwise.""" + + table, residual, y, index = _case(100, hidden=64, seed=7) + grad = torch.randn_like(residual) + for needs_residual in (False, True): + leaves = [ + residual.detach().clone().requires_grad_(needs_residual), + y.detach().clone().requires_grad_(True), + table.detach().clone().requires_grad_(True), + ] + out = _cuda_op()(leaves[0], leaves[1], _gate(leaves[2]), index) + # grad_fn.apply returns what backward returned, before autograd drops the unused ones + d_residual, d_y, *_ = out.grad_fn.apply(grad) + if needs_residual: + assert torch.equal(d_residual, grad) + else: + assert d_residual is None + assert d_y.shape == y.shape + def test_a_gate_table_whose_rows_overlap_is_copied_not_read_in_place(self): """An expanded table (row stride 0) is a legal gate: it is copied to dense rows first.""" @@ -248,6 +268,12 @@ def _gate_backward_in_child(mutation: str) -> subprocess.CompletedProcess: first = torch.tensor([0, 1, 5], **i64) elif mutation == "no_tiles": begin, end = begin[:0], end[:0] + elif mutation == "gap_between_tiles": + begin = torch.tensor([0, 3], **i64) + elif mutation == "overlapping_tiles": + begin = torch.tensor([0, 1], **i64) + elif mutation == "last_tile_stops_short": + end = torch.tensor([2, 3], **i64) elif mutation == "short_sorted_pos": sorted_pos = sorted_pos[:3] elif mutation == "no_boundary": @@ -277,13 +303,16 @@ def test_the_gate_backward_stops_on_tile_metadata_that_does_not_describe_the_row "tile_end_past_rows", "position_past_rows", "segments_past_tiles", - "no_tiles", + "gap_between_tiles", + "overlapping_tiles", + "last_tile_stops_short", ): done = _gate_backward_in_child(mutation) assert done.returncode != 0, mutation assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) # Lengths are checked on the host, before any launch. for mutation, message in ( + ("no_tiles", "the tile list is empty but there are 4 positions to sum"), ("short_sorted_pos", "sorted_pos must have M entries"), ("no_boundary", "seg_first_tile must hold R + 1 entries"), ): From 2768c0f17e378f3b3616cfc438cdfb79686b10e5 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:19:01 +0000 Subject: [PATCH 58/62] test(h3): restore the rmsnorm native validation tests dropped from this branch The file is on test-h3 and on every branch below this one, but earlier merges of the tensor-parallel branch into this one left it out. Three of its tests (modulation row outside the table asserts in the kernel, a view of a wider table is read by its row pitch, the backward stops on tile metadata that does not describe the rows) exist nowhere else. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../test_h3_rmsnorm_native_validation.py | 282 ++++++++++++++++++ 1 file changed, 282 insertions(+) create mode 100644 tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py new file mode 100644 index 000000000..e97d7f6c8 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct native calls reject malformed pointers and layouts before launching CUDA work.""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +pytestmark = [ + pytest.mark.cuda_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), +] + +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}") + + +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.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.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.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.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.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.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.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("entrypoint", ["forward", "backward"]) +@pytest.mark.parametrize("bad_row", [-1, 2], ids=["negative", "past-table"]) +def test_modulation_row_outside_the_table_asserts_in_the_kernel(native, entrypoint, bad_row): + """check_range=False callers get a device-side assert, not a read past the shift/scale table.""" + + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + entrypoint, bad_row = sys.argv[1], int(sys.argv[2]) + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) # a table of two rows + index = torch.tensor([0, bad_row], dtype=torch.int64, device="cuda") + if entrypoint == "forward": + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + else: + _C.h3_rmsnorm_backward( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=index, + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(bad_row)], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "the kernel accepted a modulation row outside the table" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + +def test_rejects_modulation_rows_that_overlap(native, inputs): + """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") + inputs["shift"] = inputs["scale"] = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + _call(native, inputs, "forward") + + +def test_a_view_of_a_wider_table_is_read_by_its_row_pitch(native, inputs): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + table = torch.randn(2, 16, generator=g).cuda() + views = table.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 = [] + for shift, scale in (views, copies): + inputs["shift"], inputs["scale"] = shift, scale + out.append(_call(native, inputs, "forward")[0]) + assert torch.equal(out[0], out[1]) + + +def _backward_in_child(mutation: str) -> subprocess.CompletedProcess: + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + + mutation = sys.argv[1] + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device="cuda") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + args = dict( + grad=torch.ones_like(x), x=x, weight=weight, rstd=rstd, shift=shift, scale=scale, + index=torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda"), + sorted_pos=torch.tensor([0, 2, 1, 3], dtype=torch.int64, device="cuda"), + tile_begin=torch.tensor([0, 2], dtype=torch.int64, device="cuda"), + tile_end=torch.tensor([2, 4], dtype=torch.int64, device="cuda"), + seg_first_tile=torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"), + ) + if mutation == "tile_end_past_rows": + args["tile_end"] = torch.tensor([2, 9], dtype=torch.int64, device="cuda") + elif mutation == "position_past_rows": + args["sorted_pos"] = torch.tensor([0, 2, 1, 4], dtype=torch.int64, device="cuda") + elif mutation == "segments_past_tiles": + args["seg_first_tile"] = torch.tensor([0, 1, 5], dtype=torch.int64, device="cuda") + _C.h3_rmsnorm_backward(**args) + torch.cuda.synchronize() + """ + ) + return subprocess.run( + [sys.executable, "-c", probe, mutation], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + +def test_the_backward_stops_on_tile_metadata_that_does_not_describe_the_rows(native): + control = _backward_in_child("none") + assert control.returncode == 0, control.stderr + for mutation in ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"): + done = _backward_in_child(mutation) + assert done.returncode != 0, mutation + assert "device-side assert triggered" in done.stderr, (mutation, done.stderr) From 7b8594136f50e9a86050d420a01f43f82f662b9f Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 14:33:07 +0000 Subject: [PATCH 59/62] fix(h3): skip the weight and bias gradients of a frozen tensor-parallel shard The projection backward folded dW and db for every call. A frozen shard now returns None for them and skips the folds, which are local to the rank. The input-gradient gather is a collective, so it stays unconditional. Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .../minimax_h3/tp_adaln_projection.py | 13 ++++++-- tests/models/minimax_h3/test_h3_tp_adaln.py | 32 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) 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 371f0ae74..28768e881 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 @@ -161,13 +161,22 @@ def backward(ctx, grad_table): f"got {tuple(grad_table.shape)}" ) grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() - dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + # The weight and bias folds are local to this rank: skip them for a frozen shard. The input + # gradient below stays unconditional, because its gather is a collective every rank must + # reach, and ``temb`` is replicated. + _, need_weight, need_bias = ctx.needs_input_grad[:3] + dw = db = None + if need_weight or need_bias: + dw, *rest = det_linear_backward_weight( + grad_local, act, weight.dtype, with_bias=need_bias + ) + db = rest[0] if need_bias else None partial = det_linear_backward_input_partials(grad_local, weight) d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) record_backward( "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" ) - return silu_backward_fp32(d_act, temb), dw, db, None, None + return silu_backward_fp32(d_act, temb), dw if need_weight else None, db, None, None class H3TPAdaLNProjectionCudaOp: diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py index 6d3737824..74281f750 100644 --- a/tests/models/minimax_h3/test_h3_tp_adaln.py +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -168,6 +168,38 @@ def test_backward_rejects_a_table_gradient_of_the_wrong_width(self): _TPAdaLNProjection.backward(node, bad.to(table.dtype)) assert _TPAdaLNProjection.backward(node, grad.to(table.dtype))[0].shape == temb.shape + @pytest.mark.parametrize("frozen", [("weight",), ("bias",), ("weight", "bias")]) + def test_backward_skips_the_gradients_of_a_frozen_shard(self, frozen): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(2, 3 * 6 * 256, 136) + op = H3TPAdaLNProjectionCudaOp(one, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 1, 0) + leaves = { + "temb": temb.clone().requires_grad_(True), + "weight": w.clone().requires_grad_("weight" not in frozen), + "bias": b.clone().requires_grad_("bias" not in frozen), + } + table = op.forward_table(*leaves.values()) + got = table.grad_fn.apply(grad.to(table.dtype)) + reference = { + name: leaf.detach().clone().requires_grad_(True) for name, leaf in leaves.items() + } + for name in leaves: + reference[name].requires_grad_(name not in frozen) + reference["temb"].requires_grad_(True) + op.forward_table(*reference.values()).backward(grad.to(table.dtype)) + for name, value in zip(("temb", "weight", "bias"), got[:3], strict=True): + if name in frozen: + assert value is None + else: + assert torch.equal(value, reference[name].grad) + @pytest.mark.parametrize("world", [2, 4, 8]) @pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") From 74cc5531c8eda8549687c0c101ff5721a2796e6e Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:10:49 +0000 Subject: [PATCH 60/62] ci: run the pipeline for pull requests and pushes to test-h3 (lost in the #494 sync) Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 328dd8530..d25d9db74 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-h3 ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-h3 ] permissions: contents: read From 692f132808445cff0ed150f022c475adda7955ee Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 15:18:54 +0000 Subject: [PATCH 61/62] Restore the projection-row registry, block-validation and one-block test hunks dropped by the #588 sync Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- rl_engine/runtime/registry.py | 50 ++++++++++++++++++++ rl_engine/validation/models/h3_block.py | 40 ++++++++++++++-- tests/models/minimax_h3/test_h3_one_block.py | 7 ++- 3 files changed, 93 insertions(+), 4 deletions(-) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 887799d57..7c2042ae7 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -232,6 +232,28 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_H3_TIMESTEP_MLP = ( "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" ) + PYTORCH_H3_AUDIO_INPUT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3AudioInputProjectionOp" + ) + CUDA_H3_AUDIO_INPUT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioInputProjectionCudaOp" + ) + PYTORCH_H3_AUDIO_OUTPUT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3AudioOutputProjectionOp" + ) + CUDA_H3_AUDIO_OUTPUT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3AudioOutputProjectionCudaOp" + ) + PYTORCH_H3_TEXT_CONTEXT_PROJECTION = ( + "rl_engine.reference.minimax_h3.projection.NativeH3TextContextProjectionOp" + ) + CUDA_H3_TEXT_CONTEXT_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3TextContextProjectionCudaOp" + ) + PYTORCH_H3_QKV_GEMM = "rl_engine.reference.minimax_h3.projection.NativeH3QKVGemmOp" + CUDA_H3_QKV_GEMM = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.projection.H3QKVGemmCudaOp" + ) PYTORCH_H3_ADALN_PROJECTION = ( "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" ) @@ -832,6 +854,22 @@ def __init__(self): OpBackend.CUDA_H3_TIMESTEP_MLP, OpBackend.PYTORCH_H3_TIMESTEP_MLP, ], + "audio_input_projection_fp32": [ + OpBackend.CUDA_H3_AUDIO_INPUT_PROJECTION, + OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION, + ], + "audio_output_projection_fp32": [ + OpBackend.CUDA_H3_AUDIO_OUTPUT_PROJECTION, + OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION, + ], + "text_context_projection": [ + OpBackend.CUDA_H3_TEXT_CONTEXT_PROJECTION, + OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION, + ], + "h3_qkv_gemm": [ + OpBackend.CUDA_H3_QKV_GEMM, + OpBackend.PYTORCH_H3_QKV_GEMM, + ], "adaln_projection_3mod": [ OpBackend.CUDA_H3_ADALN_PROJECTION, OpBackend.PYTORCH_H3_ADALN_PROJECTION, @@ -893,6 +931,10 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], @@ -950,6 +992,10 @@ def __init__(self): "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], @@ -984,6 +1030,10 @@ def __init__(self): "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "audio_input_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_INPUT_PROJECTION], + "audio_output_projection_fp32": [OpBackend.PYTORCH_H3_AUDIO_OUTPUT_PROJECTION], + "text_context_projection": [OpBackend.PYTORCH_H3_TEXT_CONTEXT_PROJECTION], + "h3_qkv_gemm": [OpBackend.PYTORCH_H3_QKV_GEMM], "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py index b20e885a1..0f895613e 100644 --- a/rl_engine/validation/models/h3_block.py +++ b/rl_engine/validation/models/h3_block.py @@ -93,6 +93,7 @@ def __init__(self, registry: KernelRegistry | None = None) -> None: self.norm = registry.get_op("h3_rmsnorm", device="cuda") self.gate = registry.get_op("adaln_gate_residual", device="cuda") self.gemm = TritonDetGemmOp() + self.qkv = registry.get_op("h3_qkv_gemm", device="cuda") self.attention = DeterministicAttentionOp() self.swiglu = SwiGLUCudaOp() @@ -169,6 +170,39 @@ def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> ) +def _row_keys(x: torch.Tensor) -> torch.Tensor: + """Logical identity ``(batch row, token position)`` of every flattened row of ``[B, S, C]``. + + The parameter-gradient reduction of a projection sorts its rows by this key, so dW does not + depend on the physical order the rows reach the GEMM in. + """ + + batch, seq = x.shape[0], x.shape[1] + rows = torch.arange(batch, device=x.device).repeat_interleave(seq) + positions = torch.arange(seq, device=x.device).repeat(batch) + return torch.stack([rows, positions], dim=1) + + +def _qkv_node(name: str, weight: str) -> Node: + """Q, K or V projection: the ``h3_qkv_gemm`` row (fixed-order Triton GEMM, BF16).""" + + def candidate(ops, ctx, x): + flat = x.reshape(-1, x.shape[-1]) + return ops.qkv.project(flat, _w(ctx, weight), logical_keys=_row_keys(x)).view( + *x.shape[:-1], -1 + ) + + return Node( + name=name, + inputs=("norm1",), + binding=Binding("h3_qkv_gemm", "row", "H3QKVGemmCudaOp.project"), + candidate=candidate, + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # fixed-order K loop vs cuBLAS + ) + + def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: return Node( name=name, @@ -240,9 +274,9 @@ def _adaln_params(ctx): provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS ), _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), - _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), - _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), - _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _qkv_node("q_proj", "attn.to_q.weight"), + _qkv_node("k_proj", "attn.to_k.weight"), + _qkv_node("v_proj", "attn.to_v.weight"), _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), _rope_node("rope_q", "q_norm"), diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py index 8f64a4508..b55e4ffcd 100644 --- a/tests/models/minimax_h3/test_h3_one_block.py +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -72,7 +72,12 @@ def test_graph_is_well_formed(): seen.add(node.name) assert h3_block.OUTPUT == "residual_mlp" rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} - assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + assert rows == { + "adaln_projection_3mod", + "h3_rmsnorm", + "adaln_gate_residual", + "h3_qkv_gemm", + } def test_partial_rope_rotates_96_of_128_channels(): From 07a77a01692d25b32c586177ebd7b6df45a4eb98 Mon Sep 17 00:00:00 2001 From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> Date: Thu, 29 Oct 2026 21:25:10 +0000 Subject: [PATCH 62/62] style: end the sigma-schedule benchmark JSON with a newline and format the h3 stubs in _C.pyi with the pinned black Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com> --- reports/experiments/h3-sigma-schedule/benchmark-4090d.json | 2 +- rl_engine/_C.pyi | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/reports/experiments/h3-sigma-schedule/benchmark-4090d.json b/reports/experiments/h3-sigma-schedule/benchmark-4090d.json index 4aee5d792..65c5c0fec 100644 --- a/reports/experiments/h3-sigma-schedule/benchmark-4090d.json +++ b/reports/experiments/h3-sigma-schedule/benchmark-4090d.json @@ -2705,4 +2705,4 @@ } } } -} \ No newline at end of file +} diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index d2a740225..85fda464c 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -516,6 +516,7 @@ def h3_gate_residual_backward( checked on the device before use. """ ... + def h3_rmsnorm_backward_partials( grad: torch.Tensor, x: torch.Tensor,