diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 328dd8530..deb74d419 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 @@ -75,6 +75,9 @@ jobs: run: | python -m pytest tests/reference/minimax_h3 tests/runtime/test_h3_sigma_schedule_dispatch.py -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/build_tools/extensions.py b/build_tools/extensions.py index b3e8257c6..47ffd9c9e 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -26,6 +26,7 @@ "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" # nvcc options that relax FP32 math. They can also arrive through the environment, where PyTorch's diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 3d5500406..0046fce21 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -510,6 +510,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); @@ -530,6 +532,33 @@ 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); +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 @@ -803,6 +832,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); @@ -825,5 +860,35 @@ 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")); + m.def("h3_rmsnorm_backward_partials", &h3_rmsnorm_backward_partials, + "WS1 RMSNorm dweight / table-gradient tile partials over explicit row lists (SP)", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), py::arg("shift"), + py::arg("scale"), py::arg("index"), py::arg("dw_rows"), py::arg("dw_begin"), + py::arg("dw_end"), py::arg("seg_rows") = py::none(), py::arg("seg_begin") = py::none(), + py::arg("seg_end") = py::none()); + m.def("h3_rmsnorm_fold_partials", &h3_rmsnorm_fold_partials, + "WS1 ascending folds of RMSNorm dweight and per-segment table-gradient partials", + py::arg("dw_partial"), py::arg("weight"), py::arg("seg_partial") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_grad_partials", &h3_gate_grad_partials, + "WS1 d_gate tile partials over explicit row lists (SP)", py::arg("grad"), py::arg("y"), + py::arg("rows"), py::arg("tile_begin"), py::arg("tile_end")); + m.def("h3_gate_grad_fold", &h3_gate_grad_fold, + "WS1 per-segment ascending fold of d_gate partials, cast once", py::arg("partial"), + py::arg("seg_first_tile"), py::arg("dtype")); + m.def("h3_rmsnorm_backward_dx", &h3_rmsnorm_backward_dx, + "Row-local dx of the WS1 RMSNorm(+modulation) backward", py::arg("grad"), py::arg("x"), + py::arg("weight"), py::arg("rstd"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_gate_residual_backward_dy", &h3_gate_residual_backward_dy, + "Row-local d_y = grad * gate[index] of the WS1 gated-residual backward", + py::arg("grad"), py::arg("gate"), py::arg("index")); #endif } diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu index 4a557e7b6..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() && t->scalar_type() == at::kLong, - "tile metadata must be contiguous 1-D int64 tensors on grad's device"); + 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(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); @@ -225,9 +226,10 @@ torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sor const int64_t num_rows = seg_first_tile.numel() - 1; const int64_t tiles = tile_begin.numel(); TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); const c10::cuda::CUDAGuard device_guard(grad.device()); auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); - if (num_rows == 0) return out; // a zero-sized grid is not a valid launch auto partial = torch::empty({std::max(tiles, 1), width}, grad.options().dtype(at::kFloat)); auto stream = at::cuda::getCurrentCUDAStream(); diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu index c8625d0bf..bfe045a98 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. db is the same @@ -509,25 +512,22 @@ 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); 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()); - auto out = torch::empty({rows, k_in}, grad.options().dtype(out_dtype)); - if (k_in == 0) return out; - if (n_out == 0) return out.zero_(); const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; // chunks is blockIdx.y and rows * k_in / threads is blockIdx.x of the fold: fail with a clear message // instead of "invalid configuration argument" when a launch dimension would overflow. @@ -536,6 +536,7 @@ torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor wei TORCH_CHECK((rows * k_in + 255) / 256 <= std::numeric_limits::max(), "rows * K = ", rows * k_in, " is too large for one input-gradient fold launch"); 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)); @@ -549,7 +550,27 @@ 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; + // 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<<>>( @@ -562,6 +583,18 @@ 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"); + 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, diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..79c5e64e4 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,359 @@ +// 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 + +#include "tile_metadata.cuh" + +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"); + // 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()); + 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); + 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(); + 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 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()); + 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); + } + 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}; +} + +// Sequence parallel: the WS1 d_gate tile partials over explicit row lists of +// (grad, y) rows -> (tiles, N) FP32. +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + for (const auto* t : {&rows, &tile_begin, &tile_end}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on grad's device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin and tile_end differ in length"); + const int64_t n = grad.size(1); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); + if (tiles == 0) return partial; + const int threads = 256; + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + // The row and tile lists are caller-built: every row must be a row of grad and every tile must + // lie inside the list before the kernel reads through them. + h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), rows.data_ptr(), rows.numel(), + grad.size(0), tile_begin.data_ptr(), tile_end.data_ptr(), + tiles, nullptr, 0, 0); + H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(y.data_ptr()), + rows.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), + partial.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// The WS1 d_gate fold: each segment's tiles in ascending order, cast once to dtype. +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat, + "partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(seg_first_tile.is_cuda() && seg_first_tile.device() == partial.device() && + seg_first_tile.scalar_type() == at::kLong && seg_first_tile.dim() == 1 && + seg_first_tile.is_contiguous(), + "seg_first_tile must be a contiguous int64 CUDA tensor on partial's device"); + TORCH_CHECK(seg_first_tile.numel() >= 1, "seg_first_tile must hold R + 1 entries"); + const int64_t n = partial.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + const c10::cuda::CUDAGuard guard(partial.device()); + auto out = torch::empty({segments, n}, partial.options().dtype(dtype)); + if (segments == 0) return out; // a zero-sized grid is not a valid launch + const int threads = 256; + const int64_t col_blocks64 = (n + threads - 1) / threads; + TORCH_CHECK(col_blocks64 <= 65535, "N=", n, " needs ", col_blocks64, + " column blocks, more than the grid.y limit of 65535"); + const unsigned col_blocks = static_cast(col_blocks64); + h3_tiles::check_lists(at::cuda::getCurrentCUDAStream(), nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile.data_ptr(), segments, partial.size(0)); + H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Row-local part of the WS1 backward only: d_y = cast(grad * gate[index]). +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + const Gate g = make_gate(gate, index, grad); + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const c10::cuda::CUDAGuard guard(grad.device()); + auto dy = torch::empty_like(grad); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward_dy", [&] { + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + reinterpret_cast(grad.data_ptr()), nullptr, reinterpret_cast(dy.data_ptr()), + rows, n, g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dy; +} diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu index 6d0146408..52b1de85d 100644 --- a/csrc/cuda/h3/rmsnorm_modulate.cu +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -26,6 +26,10 @@ // autograd's RMSNorm backward, which forms rstd^3 differently. // dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order // dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) +// +// The tile partials and the folds are also bound separately (`*_partials`, +// `*_fold_*`) so a sequence-parallel caller can compute each WS1 tile on the +// rank that holds its rows and fold the gathered partials in WS1 order. #include #include @@ -36,6 +40,7 @@ #include #include +#include #include #include "tile_metadata.cuh" @@ -112,27 +117,11 @@ struct Modulation { const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr const void* scale; int64_t row_stride; + int64_t num_rows; const int64_t* index; // (S,) table row per sequence position int64_t seq; // S: row r of x uses index[r % S] - int64_t rows; // R: table rows; an index outside [0, R) is a caller error }; -// Table row of x row `r`. check_range=False callers skip the Python range check, so the native side -// asserts the index before it is used as an offset (as the row-gather kernel does). Kernels that read -// it once per row use this; the dweight partials read it inside their row loop, where an assert -// costs about 11% of the backward, so they use the unchecked read and the entry points run -// h3_tiles::check_index_range over the S index values before launching them. -__device__ __forceinline__ int64_t modulation_row(const Modulation& mod, int64_t r) { - const int64_t i = mod.index[r % mod.seq]; - CUDA_KERNEL_ASSERT(i >= 0 && i < mod.rows); - return i; -} - -__device__ __forceinline__ int64_t modulation_row_checked_by_entry_point(const Modulation& mod, - int64_t r) { - return mod.index[r % mod.seq]; -} - template __global__ void __launch_bounds__(32 * kWarps) rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, @@ -144,7 +133,8 @@ __global__ void __launch_bounds__(32 * kWarps) const T* shift = nullptr; const T* scale = nullptr; if (mod.index != nullptr) { - const int64_t i = modulation_row(mod, r); + 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; } @@ -181,9 +171,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) + modulation_row(mod, r) * 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; @@ -206,22 +199,31 @@ __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) { - const T* scale = mod.index != nullptr - ? static_cast(mod.scale) + - modulation_row_checked_by_entry_point(mod, r) * mod.row_stride - : nullptr; + for (int64_t p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; + // Read unchecked: this loop is the kernel's whole cost, and an assert in it slows the backward by + // about 11%. The entry points run h3_tiles::check_index_range over the S index values first. + const T* scale = nullptr; + if (mod.index != nullptr) { + scale = static_cast(mod.scale) + mod.index[r % mod.seq] * 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; @@ -282,19 +284,13 @@ void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& l TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); } -void check_index_tensor(const torch::Tensor& t, const char* name, const torch::Tensor& x) { - TORCH_CHECK(t.is_cuda() && t.device() == x.device() && t.scalar_type() == at::kLong && - t.dim() == 1 && t.is_contiguous(), - name, " must be a contiguous 1-D int64 CUDA tensor on x's device"); -} - Modulation make_modulation(const c10::optional& shift, const c10::optional& scale, const c10::optional& index, const torch::Tensor& x, int64_t n) { // No modulation: kernels test `index == nullptr`. seq = 1 (not 0) keeps `r % seq` defined for // any code that computes it before looking at the index. - Modulation mod{nullptr, nullptr, 0, nullptr, 1, 0}; + Modulation mod{nullptr, nullptr, 0, 0, nullptr, 1}; if (!index.has_value()) { TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); return mod; @@ -315,16 +311,17 @@ Modulation make_modulation(const c10::optional& shift, // expanded or overlapping one (row_stride < N) is not a table of rows. TORCH_CHECK(sh.size(0) <= 1 || sh.stride(0) >= n, "shift/scale row stride ", sh.stride(0), " is smaller than N=", n); - check_index_tensor(ix, "row index", x); - TORCH_CHECK(ix.numel() > 0, "row index must be nonempty"); + 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); - mod.rows = sh.size(0); return mod; } @@ -402,10 +399,15 @@ std::vector h3_rmsnorm_backward( 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"); - check_index_tensor(*sorted_pos, "sorted_pos", x); - check_index_tensor(*tile_begin, "tile_begin", x); - check_index_tensor(*tile_end, "tile_end", x); - check_index_tensor(*seg_first_tile, "seg_first_tile", x); + 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"); @@ -419,7 +421,7 @@ std::vector h3_rmsnorm_backward( const c10::cuda::CUDAGuard guard(x.device()); auto stream = at::cuda::getCurrentCUDAStream(); // The dweight partials read the index inside their row loop: check its values first. - if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.rows); + if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.num_rows); 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)); @@ -481,3 +483,168 @@ std::vector h3_rmsnorm_backward( } return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; } + +namespace { + +void check_tiles(const torch::Tensor& rows, const torch::Tensor& begin, const torch::Tensor& end, + const torch::Tensor& like) { + for (const auto* t : {&rows, &begin, &end}) { + TORCH_CHECK(t->is_cuda() && t->device() == like.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on x's device"); + } + TORCH_CHECK(begin.numel() == end.numel(), "tile_begin and tile_end differ in length"); +} + +} // namespace + +// Sequence parallel: the dweight and (with modulation) table-gradient tile +// partials of the WS1 backward, for explicit tiles over the rows of x. `index` +// gives each row's table row (one entry per row of x). Returns +// {dweight_partial (tiles, N)} or {dweight_partial, table_partial (seg_tiles, 2N)}. +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.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(); + TORCH_CHECK(2 * n <= std::numeric_limits::max() && (2 * n + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + // The row and tile lists are caller-built: every row must be a row of x and every tile must lie + // inside its list before the kernels read through them. + h3_tiles::check_lists(stream, dw_rows.data_ptr(), dw_rows.numel(), rows, + dw_begin.data_ptr(), dw_end.data_ptr(), tiles, nullptr, 0, + 0); + // The dweight partials read the modulation index inside their row loop: check its values first + // (this entry point can be called without the dx entry point having run). + if (mod.index != nullptr) h3_tiles::check_index_range(stream, mod.index, mod.seq, mod.num_rows); + if (tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, + 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + rstd.data_ptr(), dw_partial.data_ptr(), rows, static_cast(n), mod, + dw_rows.data_ptr(), dw_begin.data_ptr(), dw_end.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + if (mod.index == nullptr) return {dw_partial}; + + TORCH_CHECK(seg_rows.has_value() && seg_begin.has_value() && seg_end.has_value(), + "modulated partials need the segment tiles"); + check_tiles(*seg_rows, *seg_begin, *seg_end, x); + const int64_t seg_tiles = seg_begin->numel(); + const int64_t width = 2 * n; + h3_tiles::check_lists(stream, seg_rows->data_ptr(), seg_rows->numel(), rows, + seg_begin->data_ptr(), seg_end->data_ptr(), seg_tiles, + nullptr, 0, 0); + auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_partials", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + seg_rows->data_ptr(), seg_begin->data_ptr(), + seg_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dw_partial, seg_partial}; +} + +// The WS1 folds: dweight = ascending fold of all tiles, cast to weight's dtype; +// table rows = ascending fold of each segment's tiles, FP32 (R, 2N). +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile) { + TORCH_CHECK(dw_partial.is_cuda() && dw_partial.dim() == 2 && dw_partial.is_contiguous() && + dw_partial.scalar_type() == at::kFloat, + "dw_partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(weight.is_cuda() && weight.device() == dw_partial.device() && weight.dim() == 1 && + weight.size(0) == dw_partial.size(1), + "weight must be (N,) on dw_partial's device"); + const int64_t n = dw_partial.size(1); + const c10::cuda::CUDAGuard guard(dw_partial.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + auto dweight = torch::empty_like(weight); + H3_DISPATCH(weight.scalar_type(), "h3_rmsnorm_fold_partials", [&] { + fold_tiles_kernel<<>>( + dw_partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), + dw_partial.size(0), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!seg_partial.has_value()) return {dweight}; + + TORCH_CHECK(seg_first_tile.has_value() && seg_first_tile->is_cuda() && + seg_first_tile->device() == dw_partial.device() && + seg_first_tile->scalar_type() == at::kLong && seg_first_tile->is_contiguous(), + "segment fold needs seg_first_tile (int64, on dw_partial's device)"); + TORCH_CHECK(seg_partial->is_cuda() && seg_partial->device() == dw_partial.device() && + seg_partial->dim() == 2 && seg_partial->is_contiguous() && + seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, + "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + TORCH_CHECK(seg_first_tile->numel() >= 1, "seg_first_tile must hold R + 1 entries"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + TORCH_CHECK(width <= std::numeric_limits::max() && (width + threads - 1) / threads <= 65535, + "N=", n, " is too large for the kernels (int indices, grid.y limit of 65535)"); + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (segments > 0) { // a zero-sized grid is not a valid launch + h3_tiles::check_lists(stream, nullptr, 0, 0, nullptr, nullptr, 0, + seg_first_tile->data_ptr(), segments, seg_partial->size(0)); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, + stream>>>(seg_partial->data_ptr(), + seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +// Row-local part of the WS1 backward only: dx (sequence-parallel callers +// compute the cross-row reductions with h3_rmsnorm_backward_partials). +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.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/tile_metadata.cuh b/csrc/cuda/h3/tile_metadata.cuh index 24c473e5f..17901f91e 100644 --- a/csrc/cuda/h3/tile_metadata.cuh +++ b/csrc/cuda/h3/tile_metadata.cuh @@ -79,6 +79,48 @@ inline void check_tile_metadata(cudaStream_t stream, const int64_t* sorted_pos, C10_CUDA_KERNEL_LAUNCH_CHECK(); } +// Sequence-parallel entry points take explicit lists. `rows` (n_rows entries, each a row of an +// M-row activation, so in [0, row_limit)) and the tiles [tile_begin, tile_end) over it, and/or the +// segment boundaries seg_first_tile (segments + 1 entries, in [0, seg_limit], non-decreasing). A +// null pointer skips its part. Same device-side asserts, same single small launch. +static __global__ void check_lists_kernel(const int64_t* __restrict__ rows, int64_t n_rows, + int64_t row_limit, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, int64_t tiles, + const int64_t* __restrict__ seg_first_tile, + int64_t segments, int64_t seg_limit) { + const int64_t first_thread = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + if (rows != nullptr) { + for (int64_t i = first_thread; i < n_rows; i += stride) { + CUDA_KERNEL_ASSERT(rows[i] >= 0 && rows[i] < row_limit); + } + } + if (tile_begin != nullptr) { + for (int64_t i = first_thread; i < tiles; i += stride) { + CUDA_KERNEL_ASSERT(tile_begin[i] >= 0 && tile_begin[i] <= tile_end[i] && + tile_end[i] <= n_rows); + } + } + if (seg_first_tile != nullptr) { + for (int64_t i = first_thread; i <= segments; i += stride) { + CUDA_KERNEL_ASSERT(seg_first_tile[i] >= 0 && seg_first_tile[i] <= seg_limit); + CUDA_KERNEL_ASSERT(i == segments || seg_first_tile[i] <= seg_first_tile[i + 1]); + } + } +} + +inline void check_lists(cudaStream_t stream, const int64_t* rows, int64_t n_rows, int64_t row_limit, + const int64_t* tile_begin, const int64_t* tile_end, int64_t tiles, + const int64_t* seg_first_tile, int64_t segments, int64_t seg_limit) { + // At least one block, so an empty list is a no-op and not an invalid launch. + const int64_t items = std::max({n_rows, tiles, segments + 1, int64_t{1}}); + const unsigned blocks = static_cast(std::min((items + 255) / 256, 1024)); + check_lists_kernel<<>>(rows, n_rows, row_limit, tile_begin, tile_end, + tiles, seg_first_tile, segments, seg_limit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + // Table-row indices read inside a kernel's loop: `n` entries, each a row of a table with `limit` // rows. The assert would cost the loop its scheduling (the same 6-12% as above, measured on the // rmsnorm dweight partials), so the entry point checks the S index values once, here, before it diff --git a/docs/.nav.yml b/docs/.nav.yml index 6e1c7357c..fcf820f03 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -38,6 +38,10 @@ nav: - 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 3fd506034..7f39ee295 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -37,5 +37,15 @@ 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) +- [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) +- [MiniMax-H3 Architecture Fingerprint](h3-arch-fingerprint.md) +- [MiniMax-H3 Operator Trace](h3-operator-trace.md) - [H3 Flash Attention (optional fast backend)](h3-flash-attention.md) - [Operator Doc Template](../contributing/operator-doc-template.md) +- [MiniMax-H3 Projection Rows](h3-projection-rows.md) +- [MiniMax-H3 Token Refiner Block](h3-token-refiner.md) +- [MiniMax-H3 Trajectory Artifact](h3-trajectory-artifact.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..ffd200b0a --- /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 `a746bff0` 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 index 5de612b3f..a33419272 100644 --- a/docs/operators/h3-adaln-projection.md +++ b/docs/operators/h3-adaln-projection.md @@ -137,6 +137,11 @@ or in a batch of 9. ![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 | diff --git a/docs/operators/h3-arch-fingerprint.md b/docs/operators/h3-arch-fingerprint.md new file mode 100644 index 000000000..cb0551ac0 --- /dev/null +++ b/docs/operators/h3-arch-fingerprint.md @@ -0,0 +1,80 @@ +# MiniMax-H3 Architecture Fingerprint + +## Summary + +`h3_arch_fingerprint` (RFC #420, foundation) is the single, machine-checked description of +the model every other H3 row refers to: the pinned checkpoint, its config, the exact +tensor table, the mixed-precision map, the scheduler contract and the strict operator +graph with the RFC row that owns each node. It is pure Python (stdlib only) and has no +kernel; its job is to make "which model, which graph, which row" a testable fact. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227ee7df40d41602854ae760620d6eb651fe` +(`FL2VA` = `transformer/`, `Ref2VA` = `transformer_ref/`; identical config and tensor table). + +## What Is Pinned + +| Fact | Value | +| --- | --- | +| Architecture | `MiniMaxH3Transformer3DModel`, 33,122,992,896 parameters, 638 tensors | +| Blocks | 50 transformer blocks (17 graph nodes each) + 2 token-refiner blocks (14 nodes each) | +| Widths | hidden 5376; 56 heads x 128 = 7168 attention inner dim (differs from hidden); FFN 14336, gate/up fused to 28672 | +| AdaLN | 3 modalities (video 0, text 1, audio 2), row = `timestep_index * 3 + tag`, projection 2688 -> 96768 | +| RoPE | MM-RoPE rotates 96 of 128 channels; 32 pass through | +| FP32 modules | `proj_in`, `audio_proj_in`, `time_embedder`, `proj_out`, `audio_proj_out`; FP32 compute without parameters: `time_proj`, `rope` | +| BF16 | every other module and the packed residual stream | +| Scheduler | rectified-flow Euler, `eta = 0`, video shift 12.0, audio shift 3.0, `t = 1 - sigma` | + +The public sampler is an ODE (`eta = 0`), so the fingerprint carries no log-probability. +That is the subject of the `stochastic_transition_contract` row. + +## How It Is Kept Honest + +1. `tools/weights/snapshot_h3_checkpoint.py` reads the Hub tree, configs and safetensors + headers at the pinned revision (no weights downloaded) and writes + `rl_engine/models/minimax_h3/pinned_checkpoint.json`. +2. `expected_tensor_table(config)` derives every tensor name, dtype and shape from + `config.json` alone. `build_fingerprint()` raises unless its digest and parameter + count equal the checkpoint headers for both variants. +3. `validate_diffusers_model(model)` compares a live (meta-device is enough) Diffusers + model against the table; `validate_checkpoint_headers(dir)` compares local shards. +4. `arch_fingerprint.json` is the committed golden. `tools/validation/models/h3_fingerprint.py check` + fails when it differs from the recomputed fingerprint. + +## Entry Points + +```python +from rl_engine.models.minimax_h3 import fingerprint as fp + +fingerprint = fp.build_fingerprint() # raises on any inconsistency +graph = fp.build_graph() # 17-node blocks, refiner, stem, tail +fp.validate_diffusers_model(model) # class, config, state-dict names/shapes +fp.validate_checkpoint_headers(path) # local safetensors headers, dtype and shape +``` + +```bash +python tools/validation/models/h3_fingerprint.py print|write|check +python tools/validation/models/h3_fingerprint.py checkpoint $RL_KERNEL_H3_WEIGHTS/hf/transformer +``` + +## Graph And Row Ownership + +`BLOCK_NODE_NAMES` equals `rl_engine.validation.models.h3_block.NODE_NAMES`, and the owning +row of every block node equals the `Binding.rfc_row` there (both are tested). The graph is +the vocabulary used by `h3_operator_trace` (what must be traced) and by +`h3_trajectory_artifact` (what a sealed run covers). Refiner residual adds are owned by +`token_refiner_block`; the stem rows `audio_input_projection_fp32`, `text_context_projection` +and the tail row `audio_output_projection_fp32` are visible in the graph even before their +kernels land. + +## Tests + +`tests/models/minimax_h3/test_h3_arch_fingerprint.py` (CPU, seconds): analytic table equals the +checkpoint digest for both variants; mixed precision per tensor; graph counts, uniqueness, +order and ownership against `h3_block`; scheduler shifts against the released configs; +corrupted shard dtype/shape/unknown names, an edited fixture and a changed config all fail +closed; local shard headers and a meta-device Diffusers model are checked when available. + +## Out Of Scope + +No numerical operator, no weights, no runtime dispatch. Changing the pinned revision means +re-running the snapshot tool, refreshing the golden file and updating dependent rows. diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md new file mode 100644 index 000000000..029211e2c --- /dev/null +++ b/docs/operators/h3-final-adaln-out.md @@ -0,0 +1,176 @@ +# 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, dim=-1) # (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 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 + +```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 `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 `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). + +## 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 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 +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-operator-trace.md b/docs/operators/h3-operator-trace.md new file mode 100644 index 000000000..9f55a8828 --- /dev/null +++ b/docs/operators/h3-operator-trace.md @@ -0,0 +1,79 @@ +# MiniMax-H3 Operator Trace + +## Summary + +`h3_operator_trace` (RFC #420, foundation) records, for every node of the fingerprinted H3 +graph, what actually ran: owning RFC row, binding status, requested and actual backend, +implementation identity, dtype, accumulator, split policy, topology and fallback state. +A strict run is only evidence when this ledger is complete and clean; it is the gate the +later rows (`ws1_full_h3_transformer`, `ws2_*`, `final_*`) report through. + +It is pure Python (no torch import in the core) and builds on `h3_arch_fingerprint` +(`fingerprint.build_graph()` is the vocabulary: 2 refiner blocks x 14 nodes, 50 blocks x 17 +nodes, stem and tail) and on the semantic operator registry +(`rl_engine/runtime/semantic_registry.py`). + +## Entry Points + +```python +from rl_engine.models.minimax_h3 import operator_trace as ot + +trace = ot.OperatorTrace(mode="candidate", forwards=1) +trace.record_node("transformer_blocks.0.q_proj", status="row", backend="...", dtype="bfloat16") +call = ot.trace_call_from_resolution(path, resolution=resolution, provenance=provenance) +trace.record(call) # registry resolution + instance provenance -> one call + +trace.assert_strict() # raises ValueError listing every violation +trace.violations(paths=..., allowed_statuses=("row",)) +trace.digest; trace.dumps(); ot.OperatorTrace.loads(text) +ot.first_difference(a, b) # first call where two traces disagree +``` + +`block_observer(trace, block_index=0)` (`stage="token_refiner"` for the 14-node refiner block) plugs into `h3_block.run_graph` through the existing +`ctx["node_observer"]` hook. It records each node with the binding `h3_block.NODES` declares, +so the one-block replay and the full-model replay emit a trace without any change in arithmetic. + +## What `violations` Reports + +| Check | Reported when | +| --- | --- | +| Binding | status is not in `allowed_statuses` (default only `row`: `interim` stand-ins and `reference` provider replays are not strict evidence) | +| Fallback | any fallback path was taken (`OperatorResolutionTrace.fallback_attempts` non-empty) | +| Dispatch | `requested_backend` is set and differs from the backend that ran | +| Split policy | anything other than `none` (split-K, stream-K, atomics) | +| Accumulator | a reduction row (GEMMs, RMSNorm, attention, AdaLN projection, projections) without an FP32 accumulator; element-wise rows must declare `none` | +| Backend | no backend recorded | +| Coverage | a node ran a different number of times than `forwards`, or a node outside the expected set ran | + +`record` also fails immediately for paths outside the graph, a row that disagrees with the +graph's owner, or an unknown status. A trace serialises with a content digest +(`trace_sha256`) and the graph digest; loading rejects tampered content or a trace recorded +against a different graph. + +## Today's One-Block Result + +On block 0 (tiny layout, H100) the trace reports 17 calls and 9 non-`row` nodes: 7 +`interim` (O, gate-up and down GEMMs, QK norm, attention, SwiGLU, ...) and 2 `reference` +(RoPE of Q and K). That list is exactly the work the unclaimed and claimed rows still owe; +`h3_operator_trace.py verify --block 0` exits non-zero until it is empty, and +`--allow row interim reference` records the current honest state. + +```bash +python tools/validation/models/h3_operator_trace.py block --layout tiny --out trace.json +python tools/validation/models/h3_operator_trace.py verify trace.json --block 0 +python tools/validation/models/h3_operator_trace.py diff a.json b.json +``` + +## Tests + +`tests/models/minimax_h3/test_h3_operator_trace.py`: strict full-graph trace; each violation +class individually; coverage (missing, repeated, foreign, subset, multiple forwards); +digest stability and tamper detection; `first_difference`; the registry bridge with real +`OperatorResolution`/`OperatorInstanceProvenance`; and (CUDA) a block-0 replay whose recorded +statuses and backends equal `h3_block`'s declared bindings and are repeatable. + +## Out Of Scope + +The trace does not measure numerics and does not change dispatch. Accumulator and split +policy are the contract values declared per row, not hardware counters; kernels that can +change them must record the real value through `record`/`trace_call_from_resolution`. diff --git a/docs/operators/h3-projection-rows.md b/docs/operators/h3-projection-rows.md new file mode 100644 index 000000000..241244744 --- /dev/null +++ b/docs/operators/h3-projection-rows.md @@ -0,0 +1,119 @@ +# 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`, #567)): + +| 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 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. + +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"`). + +## 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/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md index 50fd4bf17..70e7b9648 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.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) diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..eeb90f72f --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,117 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `ecde5bb`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. + +## Tests + +```bash +python -m pytest tests/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-token-refiner.md b/docs/operators/h3-token-refiner.md new file mode 100644 index 000000000..61799fb56 --- /dev/null +++ b/docs/operators/h3-token-refiner.md @@ -0,0 +1,67 @@ +# MiniMax-H3 Token Refiner Block + +## Summary + +RFC #420 row `token_refiner_block`: the two text-refiner blocks (`token_refiner.refiner_blocks.{0,1}`) +and the final RMSNorm that closes the refiner. Weights are the released checkpoint's +(`MiniMaxAI/MiniMax-H3@42ed227`, shard 14 of 14). A refiner block is a plain pre-norm +transformer block with **no AdaLN, no RoPE and no timestep input**, 14 nodes: + +```text +norm1 -> q_proj, k_proj, v_proj -> q_norm, k_norm -> attention -> o_proj -> residual_attn +residual_attn -> norm2 -> gate_up -> swiglu -> down -> residual_mlp +``` + +The row owns the two residual adds (`residual_attn`, `residual_mlp`, FP32 add of the +BF16 operands rounded once to BF16) and the graph itself; every other node is bound to +the operator that already exists for it: `h3_rmsnorm`, `h3_qkv_gemm` (Q/K/V, from the +[projection rows](h3-projection-rows.md)) and the still-`interim` stand-ins for attention, +SwiGLU and the O/gate_up/down GEMMs (claimed by other contributors; they are replaced when +those rows land, without changing the graph). + +## Entry Points + +```python +from rl_engine.validation.models import h3_refiner as hr, h3_block as hb + +ops = hb.CandidateOps() +text = hr.run_refiner("candidate", weights, hidden, ops) # 2 blocks + final_norm +report = hr.run_refiner_case(weights, block=0, tokens=257) # per-node evidence +``` + +`operator_trace.block_observer(trace, block_index=b, stage="token_refiner")` records the +refiner nodes under their fingerprint paths (`token_refiner.refiner_blocks.b.`), so +`trace.coverage(paths=...)` can be checked against `fingerprint.build_graph()`. + +## Contract + +- Node names and owner rows equal `fingerprint.REFINER_NODES` (a test fails otherwise). +- Bitwise repeatable; each batch row independent of its neighbours. +- Nodes that promise it (norms) equal diffusers' result on diffusers' inputs. +- Every node, the block output and the 2-block stack have FP32-golden error + <= 1.25 x diffusers' error + 1e-3. +- Backward: the gradient reaching the input and each of the 10 parameter leaves is + repeatable and within 1.5 x diffusers' error + 5e-3 of the FP32 golden's. +- A fused `[Q;K;V]` rollout layout gives bitwise the same Q/K/V as the strict one + (inherited from `h3_qkv_gemm`). + +## Evidence + +`tests/models/minimax_h3/test_h3_token_refiner.py` (needs `RL_KERNEL_H3_WEIGHTS` pointing at a +directory with the shard index and shard 14; CPU graph tests run anywhere, CUDA tests skip +without weights). On H100 (Hopper) 15 passed. Tokens 1, 64 and 257 per block, both blocks. + +## Limits + +Attention, the O/gate_up/down GEMMs and SwiGLU are `interim` until their rows land, so a +refiner trace is not yet strict-clean (`assert_strict` reports exactly those nodes). +Fail-closed strict mode is therefore enforced per node by the trace, not by the block. + +## 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-token-refiner-h100/report.json` (from `tools/validation/models/h3_refiner_report.py`, real block-0 weights) records per-node agreement, +batch/repeat invariance, relative error vs FP64 and block time at 12, 257 and 925 tokens. +Time: 3.7 ms vs 0.54 ms (12 tokens), 8.5 vs 0.62 ms (257), 32.3 vs 1.42 ms (925) for RL-Kernel vs provider. Almost all of it is the nodes +that still run interim bindings owned by other rows (full attention, O/gate_up/down GEMMs); the rows this PR adds (QKV, norms, residual) are +within about 2x of the provider. Replacing the interim bindings is tracked by those rows, not by this PR. diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..868df445a --- /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 `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. + +| 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/docs/operators/h3-trajectory-artifact.md b/docs/operators/h3-trajectory-artifact.md new file mode 100644 index 000000000..2239e2180 --- /dev/null +++ b/docs/operators/h3-trajectory-artifact.md @@ -0,0 +1,78 @@ +# MiniMax-H3 Trajectory Artifact + +## Summary + +`h3_trajectory_artifact` (RFC #420, foundation) is the sealed record that end-to-end claims +are made about: "native and strict first differ at step 3", "strict is bitwise identical +across runs". It stores digests, not tensors, so it is small and diffable; tensors can be kept +beside it when you need to know *how far* apart two runs are. + +It builds on `h3_arch_fingerprint` (the artifact is bound to the fingerprint digest) and can +reference an operator trace digest (`h3_operator_trace`), so what ran and what it produced +travel together. The core is stdlib; torch is imported only to hash or load tensors. + +## Contents + +Per artifact: variant (`FL2VA` / `Ref2VA`), label (`native`, `strict`, ...), seed, steps, +modalities (`video`, `audio`), per-modality shift, exact FP32 sigma grids (stored as +`float.hex` of the FP32 value), per-modality digest of the initial noise (dtype, shape, +sha256), optional operator-trace digest and free metadata. + +Per step: the sigma and next sigma used, the digest of each modality's velocity and of the +latent after the step, and `chain_sha256 = H(previous chain, step body)`. The record is sealed +by `artifact_sha256`. + +Tensor digests are sha256 over dtype, shape and raw bytes (any device, any stride), so one +flipped bit, a reshape or a dtype change is a different digest. + +## Entry Points + +```python +from rl_engine.models.minimax_h3 import trajectory as tj + +rec = tj.TrajectoryRecorder(variant="FL2VA", label="strict", seed=7, + sigma_grids={"video": grid_v, "audio": grid_a}, + initial_noise={"video": z_v, "audio": z_a}, + shifts={"video": 12.0, "audio": 3.0}, + operator_trace_sha256=trace.digest, keep_tensors=True) +for i in range(steps): + rec.record_step(i, velocity={"video": v_v, "audio": v_a}, latent={"video": x_v, "audio": x_a}) +artifact = rec.seal(); rec.save_tensors("run.safetensors") +open("run.json", "w").write(tj.dumps(artifact)) + +tj.verify(artifact) # fail closed +tj.first_drift(a, b, tensors_lhs=pa, tensors_rhs=pb) # {"step", "kind", "modality", "mismatch_count", "max_abs_diff", ...} +tj.drift_report(a, b) # every differing (step, kind, modality) +``` + +```bash +python tools/validation/models/h3_trajectory.py verify run.json +python tools/validation/models/h3_trajectory.py compare native.json strict.json --require-equal +``` + +## Guarantees + +* `verify` rejects: a changed seal, an edited/dropped/reordered step (hash chain), a sigma + that disagrees with the grid, a grid that is not non-increasing or leaves `(0, 1]`, a final + latent that differs from the last step, a schema mismatch, and an artifact recorded + against a different architecture fingerprint. +* A comparison first checks that the two runs were *given* the same variant, seed, step + count, shifts, sigma grids and initial noise; otherwise it raises, because a drift + between different inputs says nothing about the operators (`require_same_inputs=False` + overrides, deliberately). +* Within a step the velocity is compared before the latent it feeds, and video before audio, + so the first drift names the earliest observable cause. + +## Tests + +`tests/models/minimax_h3/test_h3_trajectory_artifact.py` (CPU): deterministic artifacts; the +seal/chain on every kind of edit including a re-sealed one; recorder contract (order, +completeness, modalities, grids); exact FP32 grids; digest sensitivity to bytes, dtype and +shape; first drift on a synthetic Euler trajectory (step, kind, modality, propagation to +later latents) and with kept tensors (mismatch count, max abs diff); incomparable inputs. + +## Out Of Scope + +It does not define the sigma schedule (`h3_sigma_schedule`), the Euler step (`h3_ode_step`) +or the stochastic transition (`stochastic_transition_contract`); it records what those +produced. 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/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 286fbb02c..85fda464c 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -374,6 +374,23 @@ def h3_det_linear_backward_input( """ ... +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]: @@ -421,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, @@ -466,3 +484,114 @@ def h3_rmsnorm_backward( """ ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: + """Return ``residual + gate[index[r % S]] * y`` for contiguous CUDA ``(M, N)`` tensors. + + ``residual`` and ``y`` share a dtype (FP32, FP16 or BF16); ``gate`` is an ``(R, N)`` view of an + AdaLN table with unit column stride and a row stride of at least ``N``; ``index`` is a + contiguous int64 ``(S,)`` tensor with ``M`` a multiple of ``S``. The gate row is gathered in the + kernel and the two roundings (gate * y, then the add) are the eager expression's, so the result + is bitwise the eager one. An index outside ``[0, R)`` stops with a device-side assert. + """ + ... + +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]: + """Return ``[d_y, d_gate]`` of ``h3_gate_residual_forward`` (``d_residual`` is ``grad`` itself). + + ``d_y = grad * gate[index]`` is bitwise the eager VJP. ``d_gate`` is ``(R, N)`` in the gate + dtype: a deterministic FP32 sum over stable-sorted tiles of rows per table row, from + ``sorted_pos``, ``tile_begin``, ``tile_end`` and ``seg_first_tile`` as built by + ``adaln_row_gather._segment_tiles`` (``R`` is ``seg_first_tile.numel() - 1``); the metadata is + checked on the device before use. + """ + ... + +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Sequence parallel: FP32 tile partials of the rmsnorm weight (and table) gradients. + + Same operands as ``h3_rmsnorm_backward`` plus explicit int64 row lists: ``dw_rows`` with tiles + ``[dw_begin, dw_end)`` over it, and with modulation ``seg_rows`` with ``[seg_begin, seg_end)`` + tiles; ``index`` then has one entry per row of ``x``. Returns ``[dw_partial (tiles, N)]`` or + ``[dw_partial, seg_partial (seg_tiles, 2N)]``. Lists and index values are checked on the device + before use; ``h3_rmsnorm_fold_partials`` folds the gathered partials. + """ + ... + +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: + """Fold gathered ``h3_rmsnorm_backward_partials`` in tile order. + + Return ``[dweight]`` in ``weight``'s dtype, or ``[dweight, d_shift, d_scale]`` (FP32 ``(R, N)``) + when ``seg_partial`` and ``seg_first_tile`` (``R + 1`` entries) are given. + """ + ... + +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: + """Sequence parallel: FP32 ``(tiles, N)`` partials of ``sum(grad * y)`` over tiles of ``rows``. + + ``rows`` is an int64 list of rows of ``grad``/``y`` and ``[tile_begin, tile_end)`` cut it into + tiles; the lists are checked on the device before use. + """ + ... + +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + """Fold gathered ``h3_gate_grad_partials`` per table row; return ``(R, N)`` in ``dtype``. + + ``seg_first_tile`` has ``R + 1`` entries into the tile axis of ``partial``, checked on the device. + """ + ... + +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: + """The row-local ``dx`` of ``h3_rmsnorm_backward`` alone, in ``x``'s dtype (sequence parallel).""" + ... + +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: + """``d_y = grad * gate[index[r % S]]``, the exact part of ``h3_gate_residual_backward``.""" + ... diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py index ff35ee019..4be94e91f 100644 --- a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py +++ b/rl_engine/backends/cuda/model_specific/minimax_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", ) @@ -105,6 +108,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/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..042e50f7d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -0,0 +1,134 @@ +# 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 ( + _repeat_positions, + _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_device_supported, + 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 + 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]) + 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; 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) + 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) 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: + """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") + 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) + + 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..f39cf351b --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -0,0 +1,96 @@ +# 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 ( + _repeat_positions, + _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) + # 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 + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + 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( + "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. 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: + """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") + # 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()) + + 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/projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py new file mode 100644 index 000000000..a7d63e33c --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/projection.py @@ -0,0 +1,269 @@ +# 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 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 + +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_IMPL_KEYED = "h3_projection_gemm_v1_ascending_n_dinput_logical_key_order_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)" + ) + + +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, 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. + """ + + 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)] + 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): + """``(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, 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:4] + dx = dw = db = None + if need_x: + dx = gemm.linear_backward_input(grad_out, weight) + if need_w or 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_KEYED if ctx.keyed else BACKWARD_IMPL, + family="cuda", + ) + return None, dx, dw, db, None, None + + +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, 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)) + + @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 = [ + _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_KEYED if ctx.keyed else BACKWARD_IMPL, + family="cuda", + ) + return (dx, *dws, None, None) + + +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, *, logical_keys=None): + return self.forward(x, weight, bias, logical_keys=logical_keys) + + 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, *_keyed(logical_keys, x, weight, bias) + ) + + 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, logical_keys=logical_keys) + + +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, *, 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, *, logical_keys=None): + 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, *_keyed(logical_keys, x, weight_q, weight_k, weight_v) + ) + + 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, *_keyed(logical_keys, x, weight)) + + def forward_fused(self, x, weight_qkv, *, logical_keys=None) -> 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, *_keyed(logical_keys, x, weight_qkv) + ) + + 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/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/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py new file mode 100644 index 000000000..22ed0a114 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -0,0 +1,238 @@ +# 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_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 + +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}" + ) + 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( + 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 + 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() + # 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 if need_weight else None, 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" + ) + 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) + + 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/backends/shared/triton/gemm/h3_projection.py b/rl_engine/backends/shared/triton/gemm/h3_projection.py new file mode 100644 index 000000000..4b1f548b5 --- /dev/null +++ b/rl_engine/backends/shared/triton/gemm/h3_projection.py @@ -0,0 +1,194 @@ +# 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/models/minimax_h3/arch_fingerprint.json b/rl_engine/models/minimax_h3/arch_fingerprint.json new file mode 100644 index 000000000..92861f48f --- /dev/null +++ b/rl_engine/models/minimax_h3/arch_fingerprint.json @@ -0,0 +1,133 @@ +{ + "checkpoint": { + "architecture": "MiniMaxH3Transformer3DModel", + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe" + }, + "components": { + "FL2VA": "35b2e8271752518b683abd327d5a7ce5a164b90faa038bc807efb44f666f5490", + "Ref2VA": "9681cadaf0c954a6b28f0962b54a8d19287a6c21924dde127ef38fc8e4328719", + "audio_scheduler": "dc8edef0db5c133b968057cc6e0f2ef744444a9ece7ae5814cf77b877c40c356", + "scheduler": "e60054558dcfb74246af86d5da946fdb79e5445c0c710bfcfe165e3c5da1f3a9", + "transformer": "0785226a50a6480f99dbbe06f81cca5f974bc69e59b54479638125c194875b7a", + "transformer_ref": "d41d0949f6bf83c8cce1f2492fd5bc2811f1a6603e7b84f02bb3aacdc7e9bb71" + }, + "derived": { + "adaln_modalities": 3, + "adaln_projection_out": 96768, + "attention_inner_dim": 7168, + "attention_inner_dim_differs_from_hidden": true, + "audio_width": 32, + "ffn_gate_up_out": 28672, + "final_adaln_out": 10752, + "rotary_channels": 96, + "rotary_pass_through_channels": 32, + "video_patch_width": 96 + }, + "fingerprint_sha256": "805d0ebf601a09d65ede184441fda2694aa374d80f3dc274ba4670cfcfdf0839", + "graph": { + "block_nodes_per_block": 17, + "blocks": 50, + "digest": "67c07e832b7af5e3aa060e23db39afbb59310a692aa5d9f7e4f5836aff703041", + "nodes": 891, + "refiner_blocks": 2, + "refiner_nodes_per_block": 14, + "rows": [ + "adaln_gate_residual", + "adaln_projection_3mod", + "audio_input_projection_fp32", + "audio_output_projection_fp32", + "final_adaln_out", + "h3_attention_o_gemm", + "h3_ffn_down_gemm", + "h3_ffn_gate_up_gemm", + "h3_full_attention", + "h3_mm_rope_3axis_partial", + "h3_qk_rmsnorm_d128", + "h3_qkv_gemm", + "h3_rmsnorm", + "h3_swiglu", + "packed_index_copy", + "packed_index_select", + "text_context_projection", + "timestep_mlp_fp32", + "timestep_sinusoid_h3", + "token_refiner_block", + "video_input_projection_fp32", + "video_output_projection_fp32", + "video_patch_pack_unpack" + ], + "stem_nodes": 8 + }, + "mixed_precision": { + "bf16": "every other module and the packed residual stream", + "fp32_compute_without_parameters": [ + "time_proj", + "rope" + ], + "fp32_modules": [ + "proj_in", + "audio_proj_in", + "time_embedder", + "proj_out", + "audio_proj_out" + ] + }, + "modality_tags": { + "audio": 2, + "text": 1, + "video": 0 + }, + "parameters": 33122992896, + "scheduler": { + "audio_shift": 3.0, + "eta": 0.0, + "euler_update": "x_next = r * xt + (1 - r) * x0, r = sigma_next / sigma", + "family": "rectified_flow_euler", + "grid": "shifted linspace(1, 0, steps), consecutive FP32 duplicates removed", + "sigma_transform": "s * sigma / (1 + (s - 1) * sigma)", + "timestep": "t = 1 - sigma", + "velocity": "data-ward, x0 = xt + sigma * v", + "video_shift": 12.0 + }, + "schema_version": "rlkernel.h3.arch_fingerprint.v1", + "tensors": 638, + "transformer_config": { + "attention_head_dim": 128, + "audio_in_channels": 32, + "ffn_dim": 14336, + "final_norm_eps": 1e-05, + "freq_dim": 256, + "hidden_size": 5376, + "in_channels": 24, + "norm_eps": 1e-05, + "num_attention_heads": 56, + "num_layers": 50, + "num_refiner_layers": 2, + "patch_size": [ + 1, + 2, + 2 + ], + "qk_norm_eps": 1e-05, + "rope_freq_dim": 16, + "rope_theta": 10000.0, + "text_dim": 5120, + "time_embed_dim": 2688, + "time_embed_hidden_dim": 5376 + }, + "variants": { + "FL2VA": { + "bytes": 66280569250, + "directory": "transformer", + "files": 16, + "tensor_table_digest": "9b5e42469bcc46d65399f5625cbb1e766742faf6c460d294d499523990290e0d" + }, + "Ref2VA": { + "bytes": 66280569250, + "directory": "transformer_ref", + "files": 16, + "tensor_table_digest": "9b5e42469bcc46d65399f5625cbb1e766742faf6c460d294d499523990290e0d" + } + } +} diff --git a/rl_engine/models/minimax_h3/fingerprint.py b/rl_engine/models/minimax_h3/fingerprint.py new file mode 100644 index 000000000..1c4485da3 --- /dev/null +++ b/rl_engine/models/minimax_h3/fingerprint.py @@ -0,0 +1,533 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 architecture fingerprint (RFC #420 ``h3_arch_fingerprint``). + +One machine-readable description of what "the H3 graph" means to RL-Kernel: + +* the pinned checkpoint (repo, revision, both variants, file hashes); +* the transformer configuration and the quantities derived from it; +* the declared mixed-precision boundaries; +* the deterministic rectified-flow scheduler settings; +* the executable graph: stem, two token-refiner blocks, 50 transformer blocks + and the output path, each node tagged with the RFC row that owns its + arithmetic. + +The module needs neither torch nor a GPU. Everything is derived from +``pinned_checkpoint.json`` (written by ``tools/weights/snapshot_h3_checkpoint.py`` +from the real Hub files) plus the analytic constants below, and the two are +cross-checked: the analytic tensor table must hash to the same digest as the +real safetensors headers. + +``validate_diffusers_model`` and ``validate_checkpoint_headers`` compare a live +Diffusers module, or a local checkpoint directory, against the fingerprint. +Any mismatch raises; nothing is normalised away. +""" + +from __future__ import annotations + +import hashlib +import json +import struct +from dataclasses import dataclass +from functools import lru_cache +from pathlib import Path +from typing import Any, Iterable, Mapping + +SCHEMA_VERSION = "rlkernel.h3.arch_fingerprint.v1" +PINNED_CHECKPOINT_PATH = Path(__file__).with_name("pinned_checkpoint.json") + +REPO = "MiniMaxAI/MiniMax-H3" +REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe" +ARCHITECTURE = "MiniMaxH3Transformer3DModel" + +# Variants differ only in weights and inputs: identical config and tensor table. +VARIANTS: Mapping[str, str] = {"FL2VA": "transformer", "Ref2VA": "transformer_ref"} + +# AdaLN modality tags (row = timestep_index * 3 + tag). +MODALITY_TAGS: Mapping[str, int] = {"video": 0, "text": 1, "audio": 2} + +# Declared mixed precision. Everything not listed as FP32 is BF16. +FP32_MODULES = ("proj_in", "audio_proj_in", "time_embedder", "proj_out", "audio_proj_out") +FP32_COMPUTE = ("time_proj", "rope") # no parameters; frequencies/features built in FP32 +_SAFETENSORS_DTYPE = {"fp32": "F32", "bf16": "BF16"} + +# Deterministic rectified-flow Euler scheduler (eta = 0). Shifts come from the +# released scheduler configs and are re-checked against pinned_checkpoint.json. +SCHEDULER: Mapping[str, Any] = { + "family": "rectified_flow_euler", + "eta": 0.0, + "video_shift": 12.0, + "audio_shift": 3.0, + "sigma_transform": "s * sigma / (1 + (s - 1) * sigma)", + "timestep": "t = 1 - sigma", + "velocity": "data-ward, x0 = xt + sigma * v", + "euler_update": "x_next = r * xt + (1 - r) * x0, r = sigma_next / sigma", + "grid": "shifted linspace(1, 0, steps), consecutive FP32 duplicates removed", +} + +# --- Graph vocabulary ------------------------------------------------------------------------- + +# The 17 nodes of one transformer block, in graph order, with the RFC row that +# owns the arithmetic. ``h3_block.NODE_NAMES`` must equal these names. +BLOCK_NODES: tuple[tuple[str, str], ...] = ( + ("adaln_projection", "adaln_projection_3mod"), + ("norm1", "h3_rmsnorm"), + ("q_proj", "h3_qkv_gemm"), + ("k_proj", "h3_qkv_gemm"), + ("v_proj", "h3_qkv_gemm"), + ("q_norm", "h3_qk_rmsnorm_d128"), + ("k_norm", "h3_qk_rmsnorm_d128"), + ("rope_q", "h3_mm_rope_3axis_partial"), + ("rope_k", "h3_mm_rope_3axis_partial"), + ("attention", "h3_full_attention"), + ("o_proj", "h3_attention_o_gemm"), + ("residual_attn", "adaln_gate_residual"), + ("norm2", "h3_rmsnorm"), + ("ffn_gate_up", "h3_ffn_gate_up_gemm"), + ("swiglu", "h3_swiglu"), + ("ffn_down", "h3_ffn_down_gemm"), + ("residual_mlp", "adaln_gate_residual"), +) +BLOCK_NODE_NAMES = tuple(name for name, _ in BLOCK_NODES) + +# The 14 nodes of one token-refiner block: plain pre-norm, no AdaLN, no RoPE, +# plain (ungated) residual adds. +REFINER_NODES: tuple[tuple[str, str], ...] = ( + ("norm1", "h3_rmsnorm"), + ("q_proj", "h3_qkv_gemm"), + ("k_proj", "h3_qkv_gemm"), + ("v_proj", "h3_qkv_gemm"), + ("q_norm", "h3_qk_rmsnorm_d128"), + ("k_norm", "h3_qk_rmsnorm_d128"), + ("attention", "h3_full_attention"), + ("o_proj", "h3_attention_o_gemm"), + ("residual_attn", "token_refiner_block"), + ("norm2", "h3_rmsnorm"), + ("ffn_gate_up", "h3_ffn_gate_up_gemm"), + ("swiglu", "h3_swiglu"), + ("ffn_down", "h3_ffn_down_gemm"), + ("residual_mlp", "token_refiner_block"), +) + +# Stem nodes in execution order: (path, owning row, precision). +STEM_NODES: tuple[tuple[str, str, str], ...] = ( + ("video_patch_pack", "video_patch_pack_unpack", "fp32"), + ("proj_in", "video_input_projection_fp32", "fp32"), + ("audio_proj_in", "audio_input_projection_fp32", "fp32"), + ("context_embedder", "text_context_projection", "bf16"), + ("time_proj", "timestep_sinusoid_h3", "fp32"), + ("time_embedder", "timestep_mlp_fp32", "fp32"), + ("rope", "h3_mm_rope_3axis_partial", "fp32"), +) +PACK_NODE = ("packed_index_copy", "packed_index_copy", "bf16") +TAIL_NODES: tuple[tuple[str, str, str], ...] = ( + ("norm_out", "final_adaln_out", "bf16"), + ("proj_out", "video_output_projection_fp32", "fp32"), + ("audio_proj_out", "audio_output_projection_fp32", "fp32"), + ("packed_index_select", "packed_index_select", "fp32"), +) +REFINER_FINAL_NODE = ("token_refiner.final_norm", "h3_rmsnorm", "bf16") + + +@dataclass(frozen=True) +class GraphNode: + """One node of the strict H3 graph.""" + + path: str # unique, e.g. ``transformer_blocks.7.attention`` + stage: str # stem | token_refiner | transformer_blocks | output + layer: int | None # block index within its stack, or None + node: str # node name inside the block (``attention``) or the module name + row: str # RFC #420 row owning the arithmetic + precision: str # "fp32" or "bf16": the declared compute precision at the node + + def to_dict(self) -> dict[str, Any]: + return { + "path": self.path, + "stage": self.stage, + "layer": self.layer, + "node": self.node, + "row": self.row, + "precision": self.precision, + } + + +# Checkpoint-compatible FP32 boundary for block nodes: everything in the blocks is BF16. +_BLOCK_PRECISION = "bf16" + + +# --- Pinned checkpoint ----------------------------------------------------------------------- + + +@lru_cache(maxsize=1) +def pinned_checkpoint() -> Mapping[str, Any]: + """Load the pinned Hub snapshot (file hashes, configs, tensor-table digests).""" + + data = json.loads(PINNED_CHECKPOINT_PATH.read_text()) + if data["repo"] != REPO or data["revision"] != REVISION: + raise ValueError("pinned_checkpoint.json does not match the pinned repo/revision") + return data + + +def transformer_config(variant: str = "FL2VA") -> dict[str, Any]: + """Return the public (underscore-free) transformer config of ``variant``.""" + + directory = VARIANTS[variant] + raw = pinned_checkpoint()["configs"][directory] + return {key: value for key, value in raw.items() if not key.startswith("_")} + + +def derived_quantities(config: Mapping[str, Any]) -> dict[str, Any]: + """Quantities the strict operators depend on, derived from the config.""" + + heads, head_dim, hidden = ( + config["num_attention_heads"], + config["attention_head_dim"], + config["hidden_size"], + ) + patch_t, patch_h, patch_w = config["patch_size"] + modalities = len(MODALITY_TAGS) + return { + "attention_inner_dim": heads * head_dim, + "attention_inner_dim_differs_from_hidden": heads * head_dim != hidden, + "rotary_channels": 6 * config["rope_freq_dim"], + "rotary_pass_through_channels": head_dim - 6 * config["rope_freq_dim"], + "video_patch_width": config["in_channels"] * patch_t * patch_h * patch_w, + "audio_width": config["audio_in_channels"], + "adaln_modalities": modalities, + "adaln_projection_out": 6 * modalities * hidden, + "final_adaln_out": 2 * hidden, + "ffn_gate_up_out": 2 * config["ffn_dim"], + } + + +# --- Tensor table --------------------------------------------------------------------------- + + +def _precision_of(name: str) -> str: + top = name.split(".", 1)[0] + return "fp32" if top in FP32_MODULES else "bf16" + + +def _attention_tensors( + prefix: str, config: Mapping[str, Any] +) -> Iterable[tuple[str, tuple[int, ...]]]: + hidden = config["hidden_size"] + inner = config["num_attention_heads"] * config["attention_head_dim"] + head_dim = config["attention_head_dim"] + ffn = config["ffn_dim"] + yield f"{prefix}.attn.to_q.weight", (inner, hidden) + yield f"{prefix}.attn.to_k.weight", (inner, hidden) + yield f"{prefix}.attn.to_v.weight", (inner, hidden) + yield f"{prefix}.attn.norm_q.weight", (head_dim,) + yield f"{prefix}.attn.norm_k.weight", (head_dim,) + yield f"{prefix}.attn.to_out.0.weight", (hidden, inner) + yield f"{prefix}.norm2.weight", (hidden,) + yield f"{prefix}.ff.net.0.proj.weight", (2 * ffn, hidden) + yield f"{prefix}.ff.net.2.weight", (hidden, ffn) + + +def expected_tensor_table(config: Mapping[str, Any]) -> dict[str, tuple[str, tuple[int, ...]]]: + """Analytic ``name -> (safetensors dtype, shape)`` of the released checkpoint.""" + + hidden = config["hidden_size"] + patch_t, patch_h, patch_w = config["patch_size"] + video_width = config["in_channels"] * patch_t * patch_h * patch_w + shapes: dict[str, tuple[int, ...]] = { + "proj_in.weight": (hidden, video_width), + "proj_in.bias": (hidden,), + "audio_proj_in.weight": (hidden, config["audio_in_channels"]), + "audio_proj_in.bias": (hidden,), + "context_embedder.weight": (hidden, config["text_dim"]), + "context_embedder.bias": (hidden,), + "time_embedder.linear_1.weight": (config["time_embed_hidden_dim"], config["freq_dim"]), + "time_embedder.linear_1.bias": (config["time_embed_hidden_dim"],), + "time_embedder.linear_2.weight": ( + config["time_embed_dim"], + config["time_embed_hidden_dim"], + ), + "time_embedder.linear_2.bias": (config["time_embed_dim"],), + "norm_out.norm.weight": (hidden,), + "norm_out.linear.weight": (2 * hidden, config["time_embed_dim"]), + "norm_out.linear.bias": (2 * hidden,), + "proj_out.weight": (video_width, hidden), + "proj_out.bias": (video_width,), + "audio_proj_out.weight": (config["audio_in_channels"], hidden), + "audio_proj_out.bias": (config["audio_in_channels"],), + "token_refiner.final_norm.weight": (hidden,), + } + for index in range(config["num_refiner_layers"]): + prefix = f"token_refiner.refiner_blocks.{index}" + shapes[f"{prefix}.norm1.weight"] = (hidden,) + shapes.update(_attention_tensors(prefix, config)) + for index in range(config["num_layers"]): + prefix = f"transformer_blocks.{index}" + shapes[f"{prefix}.norm1.weight"] = (hidden,) + shapes.update(_attention_tensors(prefix, config)) + shapes[f"{prefix}.adaln_proj.linear.weight"] = ( + 6 * len(MODALITY_TAGS) * hidden, + config["time_embed_dim"], + ) + shapes[f"{prefix}.adaln_proj.linear.bias"] = (6 * len(MODALITY_TAGS) * hidden,) + return { + name: (_SAFETENSORS_DTYPE[_precision_of(name)], shape) for name, shape in shapes.items() + } + + +def tensor_table_digest(table: Mapping[str, tuple[str, Iterable[int]]]) -> str: + """SHA-256 over the sorted ``name|dtype|shape`` lines of a tensor table.""" + + lines = [ + f"{name}|{dtype}|{','.join(str(int(dim)) for dim in shape)}" + for name, (dtype, shape) in sorted(table.items()) + ] + return hashlib.sha256("\n".join(lines).encode()).hexdigest() + + +def parameter_count(table: Mapping[str, tuple[str, Iterable[int]]]) -> int: + total = 0 + for _, shape in table.values(): + count = 1 + for dim in shape: + count *= int(dim) + total += count + return total + + +# --- Graph ----------------------------------------------------------------------------------- + + +def build_graph(config: Mapping[str, Any] | None = None) -> tuple[GraphNode, ...]: + """Return the full strict graph in execution order.""" + + config = dict(config or transformer_config()) + nodes: list[GraphNode] = [] + for path, row, precision in STEM_NODES: + nodes.append(GraphNode(path, "stem", None, path, row, precision)) + for index in range(config["num_refiner_layers"]): + for name, row in REFINER_NODES: + nodes.append( + GraphNode( + f"token_refiner.refiner_blocks.{index}.{name}", + "token_refiner", + index, + name, + row, + _BLOCK_PRECISION, + ) + ) + path, row, precision = REFINER_FINAL_NODE + nodes.append(GraphNode(path, "token_refiner", None, "final_norm", row, precision)) + path, row, precision = PACK_NODE + nodes.append(GraphNode(path, "stem", None, path, row, precision)) + for index in range(config["num_layers"]): + for name, row in BLOCK_NODES: + nodes.append( + GraphNode( + f"transformer_blocks.{index}.{name}", + "transformer_blocks", + index, + name, + row, + _BLOCK_PRECISION, + ) + ) + for path, row, precision in TAIL_NODES: + nodes.append(GraphNode(path, "output", None, path, row, precision)) + paths = [node.path for node in nodes] + if len(set(paths)) != len(paths): + raise ValueError("graph paths must be unique") + return tuple(nodes) + + +def graph_digest(graph: Iterable[GraphNode]) -> str: + payload = json.dumps([node.to_dict() for node in graph], sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode()).hexdigest() + + +def rows_in_graph(graph: Iterable[GraphNode] | None = None) -> tuple[str, ...]: + """Distinct owning rows of the graph, sorted.""" + + return tuple(sorted({node.row for node in (graph or build_graph())})) + + +# --- Components ------------------------------------------------------------------------------ + + +def component_hash(files: Iterable[Mapping[str, Any]]) -> str: + """SHA-256 over sorted ``path|size|sha256`` lines of a directory's files.""" + + lines = sorted(f"{f['path']}|{f['size']}|{f['sha256'] or f['git_oid']}" for f in files) + return hashlib.sha256("\n".join(lines).encode()).hexdigest() + + +def component_hashes() -> dict[str, str]: + """Hash of every pinned component directory (variant transformers, schedulers, wrappers).""" + + return { + directory: component_hash(files) + for directory, files in sorted(pinned_checkpoint()["dirs"].items()) + } + + +# --- Fingerprint --------------------------------------------------------------------------- + + +def build_fingerprint() -> dict[str, Any]: + """Assemble the fingerprint; every cross-check runs and raises on disagreement.""" + + snapshot = pinned_checkpoint() + config = transformer_config("FL2VA") + if transformer_config("Ref2VA") != config: + raise ValueError( + "FL2VA and Ref2VA transformer configs differ; the fingerprint is per-variant" + ) + table = expected_tensor_table(config) + digest = tensor_table_digest(table) + for variant, directory in VARIANTS.items(): + pinned = snapshot["tensor_tables"][directory] + if pinned["digest"] != digest: + raise ValueError( + f"{variant}: analytic tensor table does not match the checkpoint headers" + ) + if pinned["parameters"] != parameter_count(table): + raise ValueError(f"{variant}: parameter count mismatch") + for key, shift in (("scheduler", "video_shift"), ("audio_scheduler", "audio_shift")): + if float(snapshot["configs"][key]["shift"]) != float(SCHEDULER[shift]): + raise ValueError(f"{key} shift does not match the pinned scheduler settings") + graph = build_graph(config) + rows = rows_in_graph(graph) + fingerprint = { + "schema_version": SCHEMA_VERSION, + "checkpoint": {"repo": REPO, "revision": REVISION, "architecture": ARCHITECTURE}, + "variants": { + variant: { + "directory": directory, + "tensor_table_digest": snapshot["tensor_tables"][directory]["digest"], + "files": len(snapshot["dirs"][directory]), + "bytes": sum(f["size"] for f in snapshot["dirs"][directory]), + } + for variant, directory in VARIANTS.items() + }, + "transformer_config": config, + "derived": derived_quantities(config), + "parameters": parameter_count(table), + "tensors": len(table), + "mixed_precision": { + "fp32_modules": list(FP32_MODULES), + "fp32_compute_without_parameters": list(FP32_COMPUTE), + "bf16": "every other module and the packed residual stream", + }, + "modality_tags": dict(MODALITY_TAGS), + "scheduler": dict(SCHEDULER), + "graph": { + "nodes": len(graph), + "stem_nodes": len(STEM_NODES) + 1, + "refiner_blocks": config["num_refiner_layers"], + "refiner_nodes_per_block": len(REFINER_NODES), + "blocks": config["num_layers"], + "block_nodes_per_block": len(BLOCK_NODES), + "digest": graph_digest(graph), + "rows": list(rows), + }, + "components": component_hashes(), + } + fingerprint["fingerprint_sha256"] = fingerprint_digest(fingerprint) + return fingerprint + + +def fingerprint_digest(fingerprint: Mapping[str, Any]) -> str: + """Digest of a fingerprint without its own ``fingerprint_sha256`` key.""" + + body = {key: value for key, value in fingerprint.items() if key != "fingerprint_sha256"} + return hashlib.sha256( + json.dumps(body, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +# --- Validation against a live model / checkpoint --------------------------------------------- + + +def _shape_tuple(shape: Iterable[int]) -> tuple[int, ...]: + return tuple(int(dim) for dim in shape) + + +def _diff_tables( + actual: Mapping[str, tuple[str | None, tuple[int, ...]]], + expected: Mapping[str, tuple[str, tuple[int, ...]]], + *, + check_dtype: bool, +) -> list[str]: + problems: list[str] = [] + for name in sorted(set(expected) - set(actual)): + problems.append(f"missing tensor {name}") + for name in sorted(set(actual) - set(expected)): + problems.append(f"unexpected tensor {name}") + for name in sorted(set(actual) & set(expected)): + dtype, shape = actual[name] + exp_dtype, exp_shape = expected[name] + if shape != exp_shape: + problems.append(f"{name}: shape {shape} != {exp_shape}") + if check_dtype and dtype != exp_dtype: + problems.append(f"{name}: dtype {dtype} != {exp_dtype}") + return problems + + +def validate_diffusers_model(model: Any, variant: str = "FL2VA") -> None: + """Check a Diffusers ``MiniMaxH3Transformer3DModel`` (meta device is fine). + + Compares the class name, config and the state-dict names/shapes to the + fingerprint. Dtypes are not compared because Diffusers builds modules in + one default dtype; use ``validate_checkpoint_headers`` for stored dtypes. + """ + + problems: list[str] = [] + if type(model).__name__ != ARCHITECTURE: + problems.append(f"class {type(model).__name__} != {ARCHITECTURE}") + config = transformer_config(variant) + live = {k: v for k, v in dict(model.config).items() if not k.startswith("_")} + for key, value in config.items(): + if live.get(key) != value: + problems.append(f"config[{key}] = {live.get(key)!r} != {value!r}") + actual = { + name: (None, _shape_tuple(tensor.shape)) for name, tensor in model.state_dict().items() + } + problems += _diff_tables(actual, expected_tensor_table(config), check_dtype=False) + if problems: + raise ValueError("H3 model does not match the fingerprint:\n " + "\n ".join(problems)) + + +def read_safetensors_header(path: Path) -> dict[str, tuple[str, tuple[int, ...]]]: + """Read ``name -> (dtype, shape)`` from a safetensors file without loading tensors.""" + + with Path(path).open("rb") as handle: + (length,) = struct.unpack(" int: + """Check every local shard header in ``directory`` against the fingerprint. + + Shards may be partially present: every tensor found must exist in the + expected table with the declared dtype and shape. Returns the number of + tensors checked. A directory with no shards is an error. + """ + + directory = Path(directory) + expected = expected_tensor_table(transformer_config(variant)) + shards = sorted(directory.glob("*.safetensors")) + if not shards: + raise ValueError(f"no safetensors shards in {directory}") + found: dict[str, tuple[str, tuple[int, ...]]] = {} + for shard in shards: + found.update(read_safetensors_header(shard)) + problems = [ + problem + for problem in _diff_tables(found, expected, check_dtype=True) + if not problem.startswith("missing tensor") + ] + if problems: + raise ValueError(f"{directory} does not match the fingerprint:\n " + "\n ".join(problems)) + return len(found) diff --git a/rl_engine/models/minimax_h3/operator_trace.py b/rl_engine/models/minimax_h3/operator_trace.py new file mode 100644 index 000000000..7a3f9648f --- /dev/null +++ b/rl_engine/models/minimax_h3/operator_trace.py @@ -0,0 +1,425 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Strict operator trace for MiniMax-H3 (RFC #420 ``h3_operator_trace``). + +A strict run is only evidence if one can say, per graph node, *what actually ran*. +``OperatorTrace`` is that ledger: one ``OperatorCall`` per executed node of the +fingerprinted graph (``fingerprint.build_graph``), carrying the owning RFC row, +how the node is bound (``row`` / ``interim`` / ``reference``), the requested and +the actual backend, the implementation identity, dtype, accumulator, split policy, +topology and fallback state. + +The trace is checked, not just stored: + +* ``coverage`` -- every graph node ran exactly ``forwards`` times, nothing outside the graph ran; +* ``violations`` / ``assert_strict`` -- fail-closed gate: non-row bindings, fallbacks, + requested != actual backend, a reduction row without an FP32 accumulator, any split + policy other than ``none``, and missing/extra nodes are all reported; +* ``digest`` -- stable hash of the semantic content, so two runs (or two ranks) can be + compared; ``first_difference`` names the first call at which two traces disagree. + +It has no torch dependency. ``trace_call_from_resolution`` bridges the semantic operator +registry (``OperatorResolution`` / ``OperatorInstanceProvenance``), and +``block_observer`` plugs into ``h3_block.run_graph``'s ``node_observer`` hook so the +existing one-block replay and full-model replay can emit a trace without changing their +arithmetic. +""" + +from __future__ import annotations + +import hashlib +import json +from collections import Counter +from dataclasses import dataclass, field, fields +from typing import Any, Callable, Iterable, Mapping, Sequence + +from rl_engine.models.minimax_h3 import fingerprint as fp + +SCHEMA_VERSION = "rlkernel.h3.operator_trace.v1" + +# How a node is bound today (``h3_block.Binding.status``). +STATUS_ROW = "row" # the owning RFC row's operator +STATUS_INTERIM = "interim" # another deterministic RL-Kernel operator stands in +STATUS_REFERENCE = "reference" # provider replay; contributes no RL-Kernel evidence +STATUSES = (STATUS_ROW, STATUS_INTERIM, STATUS_REFERENCE) + +SPLIT_NONE = "none" # no split-K / stream-K / atomics: one fixed reduction order + +# Rows whose arithmetic reduces over a dimension and so must accumulate in FP32. +# Everything else is element-wise or pure data movement ("none"). +_REDUCING_ROWS = frozenset( + { + "adaln_projection_3mod", + "audio_input_projection_fp32", + "audio_output_projection_fp32", + "final_adaln_out", + "h3_attention_o_gemm", + "h3_ffn_down_gemm", + "h3_ffn_gate_up_gemm", + "h3_full_attention", + "h3_qk_rmsnorm_d128", + "h3_qkv_gemm", + "h3_rmsnorm", + "text_context_projection", + "timestep_mlp_fp32", + "video_input_projection_fp32", + "video_output_projection_fp32", + } +) +ACCUMULATOR_FP32 = "fp32" +ACCUMULATOR_NONE = "none" + + +def expected_accumulator(row: str) -> str: + """Accumulator the strict contract requires for a row's node.""" + + return ACCUMULATOR_FP32 if row in _REDUCING_ROWS else ACCUMULATOR_NONE + + +@dataclass(frozen=True) +class OperatorCall: + """One executed graph node.""" + + path: str # graph path, e.g. ``transformer_blocks.7.attention`` + row: str # owning RFC row (checked against the graph) + status: str # row | interim | reference + backend: str # backend/operator that actually ran + requested_backend: str = "" # what dispatch was asked for; "" when not resolved by a registry + implementation: str = "" # concrete class/function reference + instance_fingerprint: str = "" # OperatorInstanceProvenance.instance_fingerprint, if known + dtype: str = "" # output dtype + accumulator: str = "" # fp32 | none | other + split_policy: str = SPLIT_NONE + topology: Mapping[str, Any] = field(default_factory=dict) # e.g. {"tp": 4, "cp": 2, "sp": True} + shapes: tuple[tuple[int, ...], ...] = () # input shapes, when known + fallback: bool = False # True if any fallback path was taken + + def to_dict(self) -> dict[str, Any]: + out: dict[str, Any] = {} + for item in fields(self): + value = getattr(self, item.name) + out[item.name] = ( + [list(shape) for shape in value] + if item.name == "shapes" + else dict(value) if item.name == "topology" else value + ) + return out + + @staticmethod + def from_dict(data: Mapping[str, Any]) -> OperatorCall: + known = {item.name for item in fields(OperatorCall)} + extra = set(data) - known + if extra: + raise ValueError(f"unknown OperatorCall fields: {sorted(extra)}") + values = dict(data) + values["shapes"] = tuple(tuple(shape) for shape in values.get("shapes", ())) + values["topology"] = dict(values.get("topology", {})) + return OperatorCall(**values) + + +class OperatorTrace: + """Ordered ledger of the nodes one run executed.""" + + def __init__( + self, + *, + mode: str = "candidate", + forwards: int = 1, + graph: Sequence[fp.GraphNode] | None = None, + metadata: Mapping[str, Any] | None = None, + ) -> None: + if forwards < 1: + raise ValueError("forwards must be >= 1") + self.mode = mode + self.forwards = forwards + self.graph = tuple(graph) if graph is not None else fp.build_graph() + self.metadata = dict(metadata or {}) + self._by_path = {node.path: node for node in self.graph} + self.calls: list[OperatorCall] = [] + + # -- recording --------------------------------------------------------------------------- + + def record(self, call: OperatorCall) -> OperatorCall: + """Append a call. Unknown paths and a row that disagrees with the graph raise at once.""" + + node = self._by_path.get(call.path) + if node is None: + raise ValueError(f"{call.path!r} is not a node of the H3 graph") + if call.row != node.row: + raise ValueError(f"{call.path}: row {call.row!r} but the graph assigns {node.row!r}") + if call.status not in STATUSES: + raise ValueError(f"{call.path}: status {call.status!r} not in {STATUSES}") + self.calls.append(call) + return call + + def record_node(self, path: str, **fields_: Any) -> OperatorCall: + """Record by path; the row is taken from the graph.""" + + node = self._by_path.get(path) + if node is None: + raise ValueError(f"{path!r} is not a node of the H3 graph") + fields_.setdefault("accumulator", expected_accumulator(node.row)) + return self.record(OperatorCall(path=path, row=node.row, **fields_)) + + def merge(self, other: OperatorTrace) -> None: + """Append another trace of the same graph (e.g. another rank's or stage's calls).""" + + if other.graph != self.graph: + raise ValueError("cannot merge traces of different graphs") + self.calls.extend(other.calls) + + # -- coverage ---------------------------------------------------------------------------- + + def counts(self) -> Counter[str]: + return Counter(call.path for call in self.calls) + + def coverage(self, paths: Iterable[str] | None = None) -> dict[str, list[str]]: + """Nodes missing, repeated the wrong number of times, or outside the expectation. + + ``paths`` narrows the expectation to a subset of the graph (a one-block replay + expects only that block); the default is the whole graph. + """ + + expected = list(paths) if paths is not None else [node.path for node in self.graph] + counts = self.counts() + wrong = [ + f"{path}: ran {counts.get(path, 0)}x, expected {self.forwards}x" + for path in expected + if counts.get(path, 0) != self.forwards + ] + unexpected = sorted(set(counts) - set(expected)) + return {"wrong_count": wrong, "unexpected": unexpected} + + # -- strictness -------------------------------------------------------------------------- + + def violations( + self, + *, + paths: Iterable[str] | None = None, + allowed_statuses: Iterable[str] = (STATUS_ROW,), + require_coverage: bool = True, + ) -> list[str]: + """Every reason this trace is not strict evidence; empty means it is.""" + + allowed = set(allowed_statuses) + problems: list[str] = [] + for call in self.calls: + where = call.path + if call.status not in allowed: + problems.append( + f"{where}: bound as {call.status!r} ({call.backend}); allowed {sorted(allowed)}" + ) + if call.fallback: + problems.append(f"{where}: a fallback path was taken") + if call.requested_backend and call.requested_backend != call.backend: + problems.append( + f"{where}: requested {call.requested_backend!r} but {call.backend!r} ran" + ) + if call.split_policy != SPLIT_NONE: + problems.append(f"{where}: split policy {call.split_policy!r} is not 'none'") + want = expected_accumulator(call.row) + if call.accumulator != want: + problems.append( + f"{where}: accumulator {call.accumulator!r}, strict contract needs {want!r}" + ) + if not call.backend: + problems.append(f"{where}: no backend recorded") + if require_coverage: + cov = self.coverage(paths) + problems += cov["wrong_count"] + [ + f"{p}: not in the expected set" for p in cov["unexpected"] + ] + return problems + + def assert_strict(self, **kwargs: Any) -> None: + """Raise ``ValueError`` listing every violation (see ``violations``).""" + + problems = self.violations(**kwargs) + if problems: + shown = problems[:40] + more = f"\n ... and {len(problems) - 40} more" if len(problems) > 40 else "" + raise ValueError("operator trace is not strict:\n " + "\n ".join(shown) + more) + + def status_summary(self) -> dict[str, dict[str, int]]: + """Per row, how many calls ran under each binding status.""" + + summary: dict[str, Counter[str]] = {} + for call in self.calls: + summary.setdefault(call.row, Counter())[call.status] += 1 + return {row: dict(counter) for row, counter in sorted(summary.items())} + + # -- serialisation and comparison ----------------------------------------------------------- + + def to_dict(self) -> dict[str, Any]: + graph = {"digest": fp.graph_digest(self.graph), "nodes": len(self.graph)} + body = { + "schema_version": SCHEMA_VERSION, + "mode": self.mode, + "forwards": self.forwards, + "metadata": self.metadata, + "graph": graph, + "calls": [call.to_dict() for call in self.calls], + } + body["trace_sha256"] = _digest(body) + return body + + @property + def digest(self) -> str: + return self.to_dict()["trace_sha256"] + + @staticmethod + def from_dict(data: Mapping[str, Any]) -> OperatorTrace: + if data.get("schema_version") != SCHEMA_VERSION: + raise ValueError(f"unsupported trace schema {data.get('schema_version')!r}") + trace = OperatorTrace( + mode=data["mode"], forwards=data["forwards"], metadata=data.get("metadata", {}) + ) + if data["graph"]["digest"] != fp.graph_digest(trace.graph): + raise ValueError("trace was recorded against a different H3 graph") + for item in data["calls"]: + trace.record(OperatorCall.from_dict(item)) + if data.get("trace_sha256") != trace.digest: + raise ValueError("trace_sha256 does not match its content") + return trace + + def dumps(self) -> str: + return json.dumps(self.to_dict(), indent=1, sort_keys=True) + "\n" + + @staticmethod + def loads(text: str) -> OperatorTrace: + return OperatorTrace.from_dict(json.loads(text)) + + +def first_difference(lhs: OperatorTrace, rhs: OperatorTrace) -> dict[str, Any] | None: + """First call at which two traces disagree, or ``None`` if they are identical.""" + + for index, (a, b) in enumerate(zip(lhs.calls, rhs.calls)): + da, db = a.to_dict(), b.to_dict() + if da != db: + changed = sorted(key for key in da if da[key] != db[key]) + return {"index": index, "path": a.path, "fields": changed, "lhs": da, "rhs": db} + if len(lhs.calls) != len(rhs.calls): + index = min(len(lhs.calls), len(rhs.calls)) + longer = lhs if len(lhs.calls) > len(rhs.calls) else rhs + return {"index": index, "path": longer.calls[index].path, "fields": ["length"]} + return None + + +def _digest(body: Mapping[str, Any]) -> str: + payload = {key: value for key, value in body.items() if key != "trace_sha256"} + return hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +# --- Bridges -------------------------------------------------------------------------------- + + +def trace_call_from_resolution( + path: str, + *, + resolution: Any, + provenance: Any | None = None, + status: str = STATUS_ROW, + dtype: str | None = None, + accumulator: str | None = None, + split_policy: str = SPLIT_NONE, + shapes: Iterable[Iterable[int]] = (), + row: str | None = None, + graph: Sequence[fp.GraphNode] | None = None, +) -> OperatorCall: + """Build a call from a semantic-registry ``OperatorResolution`` (+ instance provenance). + + ``backend`` is the resolved concrete backend, ``requested_backend`` what the caller + asked for, the fingerprint is the session's ``instance_fingerprint``, and a + non-empty ``fallback_attempts`` marks the call as a fallback. Unresolved + resolutions raise: there is nothing that ran. + """ + + trace = resolution.trace + if trace.status != "resolved" or trace.concrete_backend is None: + raise ValueError(f"{path}: resolution status {trace.status!r}; nothing ran") + graph_row = row or {n.path: n.row for n in (graph or fp.build_graph())}.get(path) + if graph_row is None: + raise ValueError(f"{path!r} is not a node of the H3 graph") + return OperatorCall( + path=path, + row=graph_row, + status=status, + backend=trace.concrete_backend, + requested_backend=trace.requested_backend, + implementation=( + provenance.concrete_implementation + if provenance is not None + else (trace.implementation_reference or "") + ), + instance_fingerprint=provenance.instance_fingerprint if provenance is not None else "", + dtype=dtype or resolution.requirements.dtype, + accumulator=accumulator if accumulator is not None else expected_accumulator(graph_row), + split_policy=split_policy, + topology=dict(resolution.requirements.topology), + shapes=tuple(tuple(shape) for shape in shapes), + fallback=bool(trace.fallback_attempts), + ) + + +# The one-block replay (h3_block.run_graph) names nodes by their block-local name; +# the full-model replay reports ``block..``. Both map onto graph paths. +def graph_path(name: str, *, block_index: int | None = None) -> str: + """Graph path of a ``block..`` / ``refiner..`` / bare block-node name.""" + + parts = name.split(".") + if len(parts) == 3 and parts[0] == "block": + return f"transformer_blocks.{int(parts[1])}.{parts[2]}" + if len(parts) == 3 and parts[0] == "refiner": + return f"token_refiner.refiner_blocks.{int(parts[1])}.{parts[2]}" + if block_index is not None and name in fp.BLOCK_NODE_NAMES: + return f"transformer_blocks.{block_index}.{name}" + return name + + +def block_observer( + trace: OperatorTrace, + *, + block_index: int = 0, + bindings: Mapping[str, Any] | None = None, + topology: Mapping[str, Any] | None = None, + dtype: str = "bfloat16", + stage: str = "transformer_blocks", +) -> Callable[[str], None]: + """``node_observer`` for ``h3_block.run_graph`` / ``h3_refiner.run_graph``. + + Records each node as the block runs it. ``stage`` is ``"transformer_blocks"`` (the 17-node + block) or ``"token_refiner"`` (the 14-node refiner block). ``bindings`` maps a node name + to a ``Binding`` (status, backend); by default it is taken from the stage's ``NODES`` so + the trace states exactly how every node is bound today. Node order and row ownership are + verified by ``record``. + """ + + if stage not in ("transformer_blocks", "token_refiner"): + raise ValueError(f"unknown stage {stage!r}") + if bindings is None: + import importlib # lazy: the graph modules need torch + + module = "h3_refiner" if stage == "token_refiner" else "h3_block" + source = importlib.import_module(f"rl_engine.validation.models.{module}") + + bindings = {node.name: node.binding for node in source.NODES} + + def observe(name: str) -> None: + binding = bindings[name] + path = ( + f"token_refiner.refiner_blocks.{block_index}.{name}" + if stage == "token_refiner" + else graph_path(name, block_index=block_index) + ) + trace.record_node( + path, + status=binding.status, + backend=binding.backend, + dtype=dtype, + topology=dict(topology or {}), + ) + + return observe diff --git a/rl_engine/models/minimax_h3/pinned_checkpoint.json b/rl_engine/models/minimax_h3/pinned_checkpoint.json new file mode 100644 index 000000000..03a5b4955 --- /dev/null +++ b/rl_engine/models/minimax_h3/pinned_checkpoint.json @@ -0,0 +1,1270 @@ +{ + "configs": { + "audio_scheduler": { + "_class_name": "MiniMaxH3Scheduler", + "_diffusers_version": "0.36.0.dev0", + "shift": 3.0 + }, + "scheduler": { + "_class_name": "MiniMaxH3Scheduler", + "_diffusers_version": "0.36.0.dev0", + "shift": 12.0 + }, + "transformer": { + "_class_name": "MiniMaxH3Transformer3DModel", + "_diffusers_version": "0.36.0.dev0", + "attention_head_dim": 128, + "audio_in_channels": 32, + "ffn_dim": 14336, + "final_norm_eps": 1e-05, + "freq_dim": 256, + "hidden_size": 5376, + "in_channels": 24, + "norm_eps": 1e-05, + "num_attention_heads": 56, + "num_layers": 50, + "num_refiner_layers": 2, + "patch_size": [ + 1, + 2, + 2 + ], + "qk_norm_eps": 1e-05, + "rope_freq_dim": 16, + "rope_theta": 10000.0, + "text_dim": 5120, + "time_embed_dim": 2688, + "time_embed_hidden_dim": 5376 + }, + "transformer_ref": { + "_class_name": "MiniMaxH3Transformer3DModel", + "_diffusers_version": "0.36.0.dev0", + "attention_head_dim": 128, + "audio_in_channels": 32, + "ffn_dim": 14336, + "final_norm_eps": 1e-05, + "freq_dim": 256, + "hidden_size": 5376, + "in_channels": 24, + "norm_eps": 1e-05, + "num_attention_heads": 56, + "num_layers": 50, + "num_refiner_layers": 2, + "patch_size": [ + 1, + 2, + 2 + ], + "qk_norm_eps": 1e-05, + "rope_freq_dim": 16, + "rope_theta": 10000.0, + "text_dim": 5120, + "time_embed_dim": 2688, + "time_embed_hidden_dim": 5376 + } + }, + "dirs": { + "FL2VA": [ + { + "git_oid": "bc2a8dc3bfe9037e1f0e82e97f309019b6a398e6", + "path": "FL2VA/audio_vae/config.json", + "sha256": null, + "size": 1973 + }, + { + "git_oid": "4ad8704ca5bf410fb2f471978fe9537c98ca4aa0", + "path": "FL2VA/audio_vae/config.yaml", + "sha256": null, + "size": 91 + }, + { + "git_oid": "234a4223d1fc7340f0dc450b7c7e286555c5065d", + "path": "FL2VA/audio_vae/dac_activations.py", + "sha256": null, + "size": 2189 + }, + { + "git_oid": "d39ca2802eec1bd98e31cb2065c088ac6b57c2b0", + "path": "FL2VA/audio_vae/dac_alias_free_act.py", + "sha256": null, + "size": 835 + }, + { + "git_oid": "ad025952f5ada1c57be051f8e6dc999267b58d36", + "path": "FL2VA/audio_vae/dac_alias_free_filter.py", + "sha256": null, + "size": 3300 + }, + { + "git_oid": "af4b2b44891ac10d5f4a544a73a0d1a5d3342637", + "path": "FL2VA/audio_vae/dac_alias_free_resample.py", + "sha256": null, + "size": 1740 + }, + { + "git_oid": "d3b9b1fcba4e8e2ab4ce323577a1f7321a9f7ab1", + "path": "FL2VA/audio_vae/dac_attn_proj.py", + "sha256": null, + "size": 3317 + }, + { + "git_oid": "b8b8c8b3ce0f2144b46663ea560bd1904b654925", + "path": "FL2VA/audio_vae/dac_audio_vae.py", + "sha256": null, + "size": 7266 + }, + { + "git_oid": "00d639800dbe91e91f1cef39f2642848622e47ff", + "path": "FL2VA/audio_vae/dac_bigvgan.py", + "sha256": null, + "size": 7102 + }, + { + "git_oid": "b1d30c34fa506c62c338a9bb7e10816e3bbc32cb", + "path": "FL2VA/audio_vae/dac_utils.py", + "sha256": null, + "size": 362 + }, + { + "git_oid": "db6aa0b9d2022da4b04f1f80fa2a5260231c5c83", + "path": "FL2VA/audio_vae/metadata.json", + "sha256": null, + "size": 440 + }, + { + "git_oid": "b03f1e5027888da226701e20c097053068e3e310", + "path": "FL2VA/audio_vae/minimax_h3_audio_vae.py", + "sha256": null, + "size": 3532 + }, + { + "git_oid": "122a6864b663b50ddc7e48c06cf35fbff1233012", + "path": "FL2VA/audio_vae/model.safetensors", + "sha256": "37dddc2f3e6d5d5139d823d5ea283bbf304dadcb885b1ccda818aa13dade5ea2", + "size": 605429308 + }, + { + "git_oid": "ed4dfb751d8417e5ab8cf58a6915214bda6ba676", + "path": "FL2VA/model_index.json", + "sha256": null, + "size": 719 + }, + { + "git_oid": "1081bacf1af7c7c6de4a585ce02cd0fd34e382da", + "path": "FL2VA/processor/chat_template.json", + "sha256": null, + "size": 5499 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "FL2VA/processor/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "2ea84a437d448ff71b08df68fdd949d5cc4ebb64", + "path": "FL2VA/processor/preprocessor_config.json", + "sha256": null, + "size": 390 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "FL2VA/processor/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "FL2VA/processor/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "3ba673a5ad7d4d13f54155ecd38b2a94a6dac8fe", + "path": "FL2VA/processor/video_preprocessor_config.json", + "sha256": null, + "size": 385 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "FL2VA/processor/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "1081bacf1af7c7c6de4a585ce02cd0fd34e382da", + "path": "FL2VA/text_encoder/chat_template.json", + "sha256": null, + "size": 5499 + }, + { + "git_oid": "29e925092dee9f14278c53e7e7d876cea8a997bc", + "path": "FL2VA/text_encoder/config.json", + "sha256": null, + "size": 1474 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "FL2VA/text_encoder/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "614247ef75d01ea40af9ef31e07103221340214d", + "path": "FL2VA/text_encoder/model-00001-of-00014.safetensors", + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb", + "size": 4932328944 + }, + { + "git_oid": "0ee7c24d1c85e72bc46fd07821a1cfb4af8efd76", + "path": "FL2VA/text_encoder/model-00002-of-00014.safetensors", + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa", + "size": 4875990528 + }, + { + "git_oid": "cabcc450fe74db21934f8d99d1bc103942f721ba", + "path": "FL2VA/text_encoder/model-00003-of-00014.safetensors", + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940", + "size": 4875990552 + }, + { + "git_oid": "da3ca77f0ee7d9bdbe0696750c094f0dc6ad3c68", + "path": "FL2VA/text_encoder/model-00004-of-00014.safetensors", + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f", + "size": 4875990584 + }, + { + "git_oid": "abb15edfb6d8c352e87914ce6a0d96a98e927e04", + "path": "FL2VA/text_encoder/model-00005-of-00014.safetensors", + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2", + "size": 4875990584 + }, + { + "git_oid": "66ace514a9c58011ac2bd4c796c21654b49443b2", + "path": "FL2VA/text_encoder/model-00006-of-00014.safetensors", + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974", + "size": 4875990584 + }, + { + "git_oid": "71e871eb7cda8e711788e283e0fb512998a952be", + "path": "FL2VA/text_encoder/model-00007-of-00014.safetensors", + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934", + "size": 4875990584 + }, + { + "git_oid": "bf1d7c575e6d7cb1117cb8fe06c102193ced09a8", + "path": "FL2VA/text_encoder/model-00008-of-00014.safetensors", + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f", + "size": 4875990584 + }, + { + "git_oid": "9dfb35899c209de09e18aa4bf8346efed8269d76", + "path": "FL2VA/text_encoder/model-00009-of-00014.safetensors", + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47", + "size": 4875990584 + }, + { + "git_oid": "4d94038272e1f5c0368639babad1f7b595a3690b", + "path": "FL2VA/text_encoder/model-00010-of-00014.safetensors", + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f", + "size": 4875990584 + }, + { + "git_oid": "826ba879b38cc6566b5aea1c3c04741a3087ba20", + "path": "FL2VA/text_encoder/model-00011-of-00014.safetensors", + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f", + "size": 4875990584 + }, + { + "git_oid": "c19df94cbd35798336f17365b1115200f1eec91f", + "path": "FL2VA/text_encoder/model-00012-of-00014.safetensors", + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e", + "size": 4875990584 + }, + { + "git_oid": "39fab9be2b33a5527cbb5131e13c387544a695a9", + "path": "FL2VA/text_encoder/model-00013-of-00014.safetensors", + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c", + "size": 4875990584 + }, + { + "git_oid": "055162e6edd9914d2652ccf99f2371c000b613fe", + "path": "FL2VA/text_encoder/model-00014-of-00014.safetensors", + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b", + "size": 3270697008 + }, + { + "git_oid": "b7f68f12031a25c81d72d46a284fe9e1daa448b8", + "path": "FL2VA/text_encoder/model.safetensors.index.json", + "sha256": null, + "size": 97831 + }, + { + "git_oid": "2ea84a437d448ff71b08df68fdd949d5cc4ebb64", + "path": "FL2VA/text_encoder/preprocessor_config.json", + "sha256": null, + "size": 390 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "FL2VA/text_encoder/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "FL2VA/text_encoder/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "3ba673a5ad7d4d13f54155ecd38b2a94a6dac8fe", + "path": "FL2VA/text_encoder/video_preprocessor_config.json", + "sha256": null, + "size": 385 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "FL2VA/text_encoder/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "FL2VA/tokenizer/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "FL2VA/tokenizer/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "FL2VA/tokenizer/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "FL2VA/tokenizer/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "94cc6ca1254d9d04b3f7cc06a8c9fc05d0ded065", + "path": "FL2VA/transformer/config.json", + "sha256": null, + "size": 604 + }, + { + "git_oid": "513241f2b5fc05ae4c1ba19d4fd84abe4c6013cc", + "path": "FL2VA/transformer/model-00001-of-00013.safetensors", + "sha256": "0b3386565e476bfdea287e9ea9f269d036e5c649ed14bb9b4afac1dc4661bd2a", + "size": 5227812968 + }, + { + "git_oid": "7793de07ceb51b49bdb3884967310d3fe1678ee3", + "path": "FL2VA/transformer/model-00002-of-00013.safetensors", + "sha256": "9c98fd4579c9bc96d1b1bbf65c1243f9256df60bccd229570cecc61897773003", + "size": 5164578856 + }, + { + "git_oid": "62aecfbe9a858299e61ce00c11f9c708fff708a6", + "path": "FL2VA/transformer/model-00003-of-00013.safetensors", + "sha256": "fa484c940d4170199fb69f5016739df2daf6ecd1150631b12a297e06e3964730", + "size": 5164578872 + }, + { + "git_oid": "f59e9950d10f7d4b8ba1efa8d98272071604461d", + "path": "FL2VA/transformer/model-00004-of-00013.safetensors", + "sha256": "4a851df37cce2d14b39cb58df41def7b591b1f11b8520a32b421a5fa3e997c9f", + "size": 5164578896 + }, + { + "git_oid": "4f6952c0c7c973eabc4bdf20ae18fc5f6db68351", + "path": "FL2VA/transformer/model-00005-of-00013.safetensors", + "sha256": "3fe6ff94dc9d3107c6776a25277d2b72b5443f88cf0f4ed835fd5b65a94ea5d4", + "size": 5164578896 + }, + { + "git_oid": "d71684218adddaa38121a7f18675d084bc7c0ac0", + "path": "FL2VA/transformer/model-00006-of-00013.safetensors", + "sha256": "79b47e1b9ff03a2c6f06dd15a28972acf32c0867172ee34e5aa6e96b99494127", + "size": 5164578896 + }, + { + "git_oid": "003c8d5d3c99fb9ee5adfe3397518515debf9adf", + "path": "FL2VA/transformer/model-00007-of-00013.safetensors", + "sha256": "6ac4d6ce639786b722a8dd99843ed56dcab48a0d174a4b0f9b8cdcbe5160fb0f", + "size": 5164578896 + }, + { + "git_oid": "8475056b4301fefe0e2ccb3161cb576d94883965", + "path": "FL2VA/transformer/model-00008-of-00013.safetensors", + "sha256": "03531812243fb0a333e01efa7e3d39750580bca8cbf2ba0529da419278eab736", + "size": 5164578896 + }, + { + "git_oid": "645a8b2038cfc60d43b7ba5417463ca3812072bc", + "path": "FL2VA/transformer/model-00009-of-00013.safetensors", + "sha256": "7e82a80b0d0d267026eb841b81523069cc7b956c1ebd8a18c07e462ffb262487", + "size": 5164578896 + }, + { + "git_oid": "01919aedf83ef64923df3fc7aff5f1c8079fd3af", + "path": "FL2VA/transformer/model-00010-of-00013.safetensors", + "sha256": "ef0a1f6b65145232543dd2232d58d91a4d31f44b0da5084acf22cea123255481", + "size": 5164578896 + }, + { + "git_oid": "cff9ae71bdbe60801777f8e5d7d9c0c4e08cf65e", + "path": "FL2VA/transformer/model-00011-of-00013.safetensors", + "sha256": "bf885b8f2f078c5499b75fded21094a26bfe0e6e00bee276e499443ecd2f298d", + "size": 5164578896 + }, + { + "git_oid": "57843db09194e8b1465ba54cd3d33f2c8927edb4", + "path": "FL2VA/transformer/model-00012-of-00013.safetensors", + "sha256": "dc717572d014ebf8527e35af815d10eb2f5a6bc87253b8a7d2f456965192db69", + "size": 5164578896 + }, + { + "git_oid": "ad8450684951494e4f6b1f8e7392bcb8d95bdfc0", + "path": "FL2VA/transformer/model-00013-of-00013.safetensors", + "sha256": "8bfd852d5817e9836de1d3ec8dbac1c5446b167568b717371f08282b22291aa2", + "size": 4242305176 + }, + { + "git_oid": "89d6b0862f66080854b239ad7a0b50c31a6de986", + "path": "FL2VA/transformer/model.safetensors.index.json", + "sha256": null, + "size": 38323 + }, + { + "git_oid": "0443bd31cefe7bd5022aa7f36fbc6aced9fdd84f", + "path": "FL2VA/video_vae/attention.py", + "sha256": null, + "size": 5785 + }, + { + "git_oid": "b13137b8d163afea25c235beef1476039d7f6c0e", + "path": "FL2VA/video_vae/base_module.py", + "sha256": null, + "size": 9520 + }, + { + "git_oid": "20f91e40ea9bc24c21a0e757715f93acf49eb841", + "path": "FL2VA/video_vae/config.json", + "sha256": null, + "size": 1807 + }, + { + "git_oid": "f10f5002b13bfa46c8eb9a9b4ba8d88bc891a517", + "path": "FL2VA/video_vae/conv.py", + "sha256": null, + "size": 4487 + }, + { + "git_oid": "40b7cac47cabc3755c8c06dc620c689418cb769d", + "path": "FL2VA/video_vae/flash.py", + "sha256": null, + "size": 5763 + }, + { + "git_oid": "f142c48a932430b67d77164a6bfba3f61102ac3d", + "path": "FL2VA/video_vae/func.py", + "sha256": null, + "size": 5840 + }, + { + "git_oid": "3e66f174c2c82600b6bc924ede9fb3fb068d1f3e", + "path": "FL2VA/video_vae/klvae.py", + "sha256": null, + "size": 48594 + }, + { + "git_oid": "231d9c77a7a3dc74de007c29332cfe92e3782f7d", + "path": "FL2VA/video_vae/minimax_h3_video_vae.py", + "sha256": null, + "size": 5123 + }, + { + "git_oid": "00a96b83c6ee3f78cf75873e46d96be5d50879a6", + "path": "FL2VA/video_vae/norm.py", + "sha256": null, + "size": 10686 + }, + { + "git_oid": "fd4038037c22d9c7b2e375e455e975ef4c56db7b", + "path": "FL2VA/video_vae/normalize.py", + "sha256": null, + "size": 1200 + }, + { + "git_oid": "e63b13c9b20738d4a2835e8fde4504a8145e121f", + "path": "FL2VA/video_vae/parallel.py", + "sha256": null, + "size": 12997 + }, + { + "git_oid": "fb5105116e094aface836f5182b126bb4fbdc699", + "path": "FL2VA/video_vae/source/config.json", + "sha256": null, + "size": 1164 + }, + { + "git_oid": "cc1fe048902a2d790834962551620c2c6dbb2b31", + "path": "FL2VA/video_vae/source/model.safetensors", + "sha256": "5f0c2e161d895a9fee7645ca32d4a7e3a22b90cacfcbeba62ec999cdbbefe0d3", + "size": 10415548320 + }, + { + "git_oid": "89c13d28758834b1fa8a1cd69018ec00807402f6", + "path": "FL2VA/video_vae/utils.py", + "sha256": null, + "size": 764 + }, + { + "git_oid": "40a7917a1144162372f53328753f00e2f4c40a20", + "path": "FL2VA/video_vae/vae_cnn.py", + "sha256": null, + "size": 8836 + }, + { + "git_oid": "6f2bf8d091c35eba2c4b6e92ee0038c0bfc9f137", + "path": "FL2VA/video_vae/vae_module.py", + "sha256": null, + "size": 1884 + }, + { + "git_oid": "0f2611987b6b0a0bfde2a3192121eede32dc9d8d", + "path": "FL2VA/video_vae/vae_processor.py", + "sha256": null, + "size": 8369 + }, + { + "git_oid": "9021bcdd1dd2120ef5d0f7e0d0521199dbc1d3f9", + "path": "FL2VA/video_vae/vae_vit.py", + "sha256": null, + "size": 13490 + } + ], + "Ref2VA": [ + { + "git_oid": "bc2a8dc3bfe9037e1f0e82e97f309019b6a398e6", + "path": "Ref2VA/audio_vae/config.json", + "sha256": null, + "size": 1973 + }, + { + "git_oid": "4ad8704ca5bf410fb2f471978fe9537c98ca4aa0", + "path": "Ref2VA/audio_vae/config.yaml", + "sha256": null, + "size": 91 + }, + { + "git_oid": "234a4223d1fc7340f0dc450b7c7e286555c5065d", + "path": "Ref2VA/audio_vae/dac_activations.py", + "sha256": null, + "size": 2189 + }, + { + "git_oid": "d39ca2802eec1bd98e31cb2065c088ac6b57c2b0", + "path": "Ref2VA/audio_vae/dac_alias_free_act.py", + "sha256": null, + "size": 835 + }, + { + "git_oid": "ad025952f5ada1c57be051f8e6dc999267b58d36", + "path": "Ref2VA/audio_vae/dac_alias_free_filter.py", + "sha256": null, + "size": 3300 + }, + { + "git_oid": "af4b2b44891ac10d5f4a544a73a0d1a5d3342637", + "path": "Ref2VA/audio_vae/dac_alias_free_resample.py", + "sha256": null, + "size": 1740 + }, + { + "git_oid": "d3b9b1fcba4e8e2ab4ce323577a1f7321a9f7ab1", + "path": "Ref2VA/audio_vae/dac_attn_proj.py", + "sha256": null, + "size": 3317 + }, + { + "git_oid": "b8b8c8b3ce0f2144b46663ea560bd1904b654925", + "path": "Ref2VA/audio_vae/dac_audio_vae.py", + "sha256": null, + "size": 7266 + }, + { + "git_oid": "00d639800dbe91e91f1cef39f2642848622e47ff", + "path": "Ref2VA/audio_vae/dac_bigvgan.py", + "sha256": null, + "size": 7102 + }, + { + "git_oid": "b1d30c34fa506c62c338a9bb7e10816e3bbc32cb", + "path": "Ref2VA/audio_vae/dac_utils.py", + "sha256": null, + "size": 362 + }, + { + "git_oid": "db6aa0b9d2022da4b04f1f80fa2a5260231c5c83", + "path": "Ref2VA/audio_vae/metadata.json", + "sha256": null, + "size": 440 + }, + { + "git_oid": "b03f1e5027888da226701e20c097053068e3e310", + "path": "Ref2VA/audio_vae/minimax_h3_audio_vae.py", + "sha256": null, + "size": 3532 + }, + { + "git_oid": "122a6864b663b50ddc7e48c06cf35fbff1233012", + "path": "Ref2VA/audio_vae/model.safetensors", + "sha256": "37dddc2f3e6d5d5139d823d5ea283bbf304dadcb885b1ccda818aa13dade5ea2", + "size": 605429308 + }, + { + "git_oid": "b6dd080b974d63e301a049cd09ffe45cbec5b58e", + "path": "Ref2VA/model_index.json", + "sha256": null, + "size": 707 + }, + { + "git_oid": "1081bacf1af7c7c6de4a585ce02cd0fd34e382da", + "path": "Ref2VA/processor/chat_template.json", + "sha256": null, + "size": 5499 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "Ref2VA/processor/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "2ea84a437d448ff71b08df68fdd949d5cc4ebb64", + "path": "Ref2VA/processor/preprocessor_config.json", + "sha256": null, + "size": 390 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "Ref2VA/processor/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "Ref2VA/processor/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "3ba673a5ad7d4d13f54155ecd38b2a94a6dac8fe", + "path": "Ref2VA/processor/video_preprocessor_config.json", + "sha256": null, + "size": 385 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "Ref2VA/processor/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "1081bacf1af7c7c6de4a585ce02cd0fd34e382da", + "path": "Ref2VA/text_encoder/chat_template.json", + "sha256": null, + "size": 5499 + }, + { + "git_oid": "29e925092dee9f14278c53e7e7d876cea8a997bc", + "path": "Ref2VA/text_encoder/config.json", + "sha256": null, + "size": 1474 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "Ref2VA/text_encoder/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "614247ef75d01ea40af9ef31e07103221340214d", + "path": "Ref2VA/text_encoder/model-00001-of-00014.safetensors", + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb", + "size": 4932328944 + }, + { + "git_oid": "0ee7c24d1c85e72bc46fd07821a1cfb4af8efd76", + "path": "Ref2VA/text_encoder/model-00002-of-00014.safetensors", + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa", + "size": 4875990528 + }, + { + "git_oid": "cabcc450fe74db21934f8d99d1bc103942f721ba", + "path": "Ref2VA/text_encoder/model-00003-of-00014.safetensors", + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940", + "size": 4875990552 + }, + { + "git_oid": "da3ca77f0ee7d9bdbe0696750c094f0dc6ad3c68", + "path": "Ref2VA/text_encoder/model-00004-of-00014.safetensors", + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f", + "size": 4875990584 + }, + { + "git_oid": "abb15edfb6d8c352e87914ce6a0d96a98e927e04", + "path": "Ref2VA/text_encoder/model-00005-of-00014.safetensors", + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2", + "size": 4875990584 + }, + { + "git_oid": "66ace514a9c58011ac2bd4c796c21654b49443b2", + "path": "Ref2VA/text_encoder/model-00006-of-00014.safetensors", + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974", + "size": 4875990584 + }, + { + "git_oid": "71e871eb7cda8e711788e283e0fb512998a952be", + "path": "Ref2VA/text_encoder/model-00007-of-00014.safetensors", + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934", + "size": 4875990584 + }, + { + "git_oid": "bf1d7c575e6d7cb1117cb8fe06c102193ced09a8", + "path": "Ref2VA/text_encoder/model-00008-of-00014.safetensors", + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f", + "size": 4875990584 + }, + { + "git_oid": "9dfb35899c209de09e18aa4bf8346efed8269d76", + "path": "Ref2VA/text_encoder/model-00009-of-00014.safetensors", + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47", + "size": 4875990584 + }, + { + "git_oid": "4d94038272e1f5c0368639babad1f7b595a3690b", + "path": "Ref2VA/text_encoder/model-00010-of-00014.safetensors", + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f", + "size": 4875990584 + }, + { + "git_oid": "826ba879b38cc6566b5aea1c3c04741a3087ba20", + "path": "Ref2VA/text_encoder/model-00011-of-00014.safetensors", + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f", + "size": 4875990584 + }, + { + "git_oid": "c19df94cbd35798336f17365b1115200f1eec91f", + "path": "Ref2VA/text_encoder/model-00012-of-00014.safetensors", + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e", + "size": 4875990584 + }, + { + "git_oid": "39fab9be2b33a5527cbb5131e13c387544a695a9", + "path": "Ref2VA/text_encoder/model-00013-of-00014.safetensors", + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c", + "size": 4875990584 + }, + { + "git_oid": "055162e6edd9914d2652ccf99f2371c000b613fe", + "path": "Ref2VA/text_encoder/model-00014-of-00014.safetensors", + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b", + "size": 3270697008 + }, + { + "git_oid": "b7f68f12031a25c81d72d46a284fe9e1daa448b8", + "path": "Ref2VA/text_encoder/model.safetensors.index.json", + "sha256": null, + "size": 97831 + }, + { + "git_oid": "2ea84a437d448ff71b08df68fdd949d5cc4ebb64", + "path": "Ref2VA/text_encoder/preprocessor_config.json", + "sha256": null, + "size": 390 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "Ref2VA/text_encoder/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "Ref2VA/text_encoder/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "3ba673a5ad7d4d13f54155ecd38b2a94a6dac8fe", + "path": "Ref2VA/text_encoder/video_preprocessor_config.json", + "sha256": null, + "size": 385 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "Ref2VA/text_encoder/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "20024bfe7c83998e9aeaf98a0cd6a2ce6306c2f0", + "path": "Ref2VA/tokenizer/merges.txt", + "sha256": null, + "size": 1671839 + }, + { + "git_oid": "c6cc1014128b19d1fc46b1d30a23e3b1d35db421", + "path": "Ref2VA/tokenizer/tokenizer.json", + "sha256": null, + "size": 7032403 + }, + { + "git_oid": "204d76f78dac6dedc820418c30bf01145de78a21", + "path": "Ref2VA/tokenizer/tokenizer_config.json", + "sha256": null, + "size": 11003 + }, + { + "git_oid": "4783fe10ac3adce15ac8f358ef5462739852c569", + "path": "Ref2VA/tokenizer/vocab.json", + "sha256": null, + "size": 2776833 + }, + { + "git_oid": "94cc6ca1254d9d04b3f7cc06a8c9fc05d0ded065", + "path": "Ref2VA/transformer/config.json", + "sha256": null, + "size": 604 + }, + { + "git_oid": "f48c7fcbbe943ff177671ceb3c202f19e3d8449b", + "path": "Ref2VA/transformer/model-00001-of-00013.safetensors", + "sha256": "902d6ae787f6d394d0104c1ef3ce19528209715fd5b0292f09d4ef7ca19c8504", + "size": 5227812968 + }, + { + "git_oid": "542e6fae21dedab37c5cd37b8e58b9fab72951f0", + "path": "Ref2VA/transformer/model-00002-of-00013.safetensors", + "sha256": "97b907e4d9b3ba533c820827ab4a03ec0d6be2a5cb8925a7e570de849ea8b183", + "size": 5164578856 + }, + { + "git_oid": "d491a1813aa3deb2a4d0cec04197d7a1efd87bbc", + "path": "Ref2VA/transformer/model-00003-of-00013.safetensors", + "sha256": "fbdf425e351622103419d8f1ddf50f2515917235588849afb365967615141f55", + "size": 5164578872 + }, + { + "git_oid": "29b59b83328586a890808084181f2c56ae3ba5b2", + "path": "Ref2VA/transformer/model-00004-of-00013.safetensors", + "sha256": "3831e9ff2640a93e3601d957c81640cc6ba7aa317dc7eae1d6ac3e2428767db3", + "size": 5164578896 + }, + { + "git_oid": "99bbc0f6b4ce6f0b74f7cdd0c8b23c130d064e04", + "path": "Ref2VA/transformer/model-00005-of-00013.safetensors", + "sha256": "ec7c6bdd26985a7ac6c83a563d5d4ad6fd961c51a2e5eeabeaaa3b45b6053e1a", + "size": 5164578896 + }, + { + "git_oid": "973f1d5e9da5fc2d3f63ed3945198a2ad12aeffe", + "path": "Ref2VA/transformer/model-00006-of-00013.safetensors", + "sha256": "536c4fe01ffdeaca8f28d8df8c53cd7deb62942ec630aa1f058d0423db4438a7", + "size": 5164578896 + }, + { + "git_oid": "855dcf9066129b72aa9e8f5d6649c5132a9208d3", + "path": "Ref2VA/transformer/model-00007-of-00013.safetensors", + "sha256": "c7b80189722b211af43bf04066ea98ef5ca5bc4f1f13d3b06444d4c30aafead7", + "size": 5164578896 + }, + { + "git_oid": "4bd259c78558aef907ec0547f435c358ba065be3", + "path": "Ref2VA/transformer/model-00008-of-00013.safetensors", + "sha256": "1a7da8bb7477d27cfcc03bd7ff81b3be015411e74006da6e542ce4ce32f7ecb4", + "size": 5164578896 + }, + { + "git_oid": "6848e69e79fdbab390e97bf665fa8163b7c8da5f", + "path": "Ref2VA/transformer/model-00009-of-00013.safetensors", + "sha256": "33f44f03128b6d47e3c17a12d3d6f9c67a06b2facdb18a7f2b2121312ba43909", + "size": 5164578896 + }, + { + "git_oid": "c440c39f7697f3675f39bc39b413cd6bccc42c94", + "path": "Ref2VA/transformer/model-00010-of-00013.safetensors", + "sha256": "adb8df74dcffebed919e15ca1e99f6a21f738ac3dbcb111f27168d9b1f03a304", + "size": 5164578896 + }, + { + "git_oid": "1f02d0bc15715b2f6af6bd3e1d56891e3ac39bec", + "path": "Ref2VA/transformer/model-00011-of-00013.safetensors", + "sha256": "7175cdde6953c234ba0f518e81628ae4ed97790e1eaadabaf1783bddba030c90", + "size": 5164578896 + }, + { + "git_oid": "56803a7f185f1a41f30adf649de5ae383b71cd0a", + "path": "Ref2VA/transformer/model-00012-of-00013.safetensors", + "sha256": "1c5c0bf33ab63e018e5923afaa8ce0dbb0feb07dbda1671605b11b99192450b6", + "size": 5164578896 + }, + { + "git_oid": "7d48095e28a37271e8df4b761caae961fb973e15", + "path": "Ref2VA/transformer/model-00013-of-00013.safetensors", + "sha256": "572e060b9b727bba6ee6c639e988d74978adfbf94f8bac9aa8e7f5bde88c1ec7", + "size": 4242305176 + }, + { + "git_oid": "89d6b0862f66080854b239ad7a0b50c31a6de986", + "path": "Ref2VA/transformer/model.safetensors.index.json", + "sha256": null, + "size": 38323 + }, + { + "git_oid": "0443bd31cefe7bd5022aa7f36fbc6aced9fdd84f", + "path": "Ref2VA/video_vae/attention.py", + "sha256": null, + "size": 5785 + }, + { + "git_oid": "b13137b8d163afea25c235beef1476039d7f6c0e", + "path": "Ref2VA/video_vae/base_module.py", + "sha256": null, + "size": 9520 + }, + { + "git_oid": "20f91e40ea9bc24c21a0e757715f93acf49eb841", + "path": "Ref2VA/video_vae/config.json", + "sha256": null, + "size": 1807 + }, + { + "git_oid": "f10f5002b13bfa46c8eb9a9b4ba8d88bc891a517", + "path": "Ref2VA/video_vae/conv.py", + "sha256": null, + "size": 4487 + }, + { + "git_oid": "40b7cac47cabc3755c8c06dc620c689418cb769d", + "path": "Ref2VA/video_vae/flash.py", + "sha256": null, + "size": 5763 + }, + { + "git_oid": "f142c48a932430b67d77164a6bfba3f61102ac3d", + "path": "Ref2VA/video_vae/func.py", + "sha256": null, + "size": 5840 + }, + { + "git_oid": "3e66f174c2c82600b6bc924ede9fb3fb068d1f3e", + "path": "Ref2VA/video_vae/klvae.py", + "sha256": null, + "size": 48594 + }, + { + "git_oid": "231d9c77a7a3dc74de007c29332cfe92e3782f7d", + "path": "Ref2VA/video_vae/minimax_h3_video_vae.py", + "sha256": null, + "size": 5123 + }, + { + "git_oid": "00a96b83c6ee3f78cf75873e46d96be5d50879a6", + "path": "Ref2VA/video_vae/norm.py", + "sha256": null, + "size": 10686 + }, + { + "git_oid": "fd4038037c22d9c7b2e375e455e975ef4c56db7b", + "path": "Ref2VA/video_vae/normalize.py", + "sha256": null, + "size": 1200 + }, + { + "git_oid": "e63b13c9b20738d4a2835e8fde4504a8145e121f", + "path": "Ref2VA/video_vae/parallel.py", + "sha256": null, + "size": 12997 + }, + { + "git_oid": "fb5105116e094aface836f5182b126bb4fbdc699", + "path": "Ref2VA/video_vae/source/config.json", + "sha256": null, + "size": 1164 + }, + { + "git_oid": "cc1fe048902a2d790834962551620c2c6dbb2b31", + "path": "Ref2VA/video_vae/source/model.safetensors", + "sha256": "5f0c2e161d895a9fee7645ca32d4a7e3a22b90cacfcbeba62ec999cdbbefe0d3", + "size": 10415548320 + }, + { + "git_oid": "89c13d28758834b1fa8a1cd69018ec00807402f6", + "path": "Ref2VA/video_vae/utils.py", + "sha256": null, + "size": 764 + }, + { + "git_oid": "40a7917a1144162372f53328753f00e2f4c40a20", + "path": "Ref2VA/video_vae/vae_cnn.py", + "sha256": null, + "size": 8836 + }, + { + "git_oid": "6f2bf8d091c35eba2c4b6e92ee0038c0bfc9f137", + "path": "Ref2VA/video_vae/vae_module.py", + "sha256": null, + "size": 1884 + }, + { + "git_oid": "0f2611987b6b0a0bfde2a3192121eede32dc9d8d", + "path": "Ref2VA/video_vae/vae_processor.py", + "sha256": null, + "size": 8369 + }, + { + "git_oid": "9021bcdd1dd2120ef5d0f7e0d0521199dbc1d3f9", + "path": "Ref2VA/video_vae/vae_vit.py", + "sha256": null, + "size": 13490 + } + ], + "audio_scheduler": [ + { + "git_oid": "28a3d4ed602d925031c7e5d6eeb07df8958ff0f1", + "path": "audio_scheduler/scheduler_config.json", + "sha256": null, + "size": 96 + } + ], + "scheduler": [ + { + "git_oid": "1d44f95b955a648ad4fab455e4fe11053888c30b", + "path": "scheduler/scheduler_config.json", + "sha256": null, + "size": 97 + } + ], + "transformer": [ + { + "git_oid": "c251a4133c81852a150c6668a570e142dc776b6e", + "path": "transformer/config.json", + "sha256": null, + "size": 546 + }, + { + "git_oid": "78e7c04f1246f7fd29127ba802f4c1fd65776404", + "path": "transformer/diffusion_pytorch_model-00001-of-00014.safetensors", + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size": 4825958704 + }, + { + "git_oid": "4a79d68cef731db07fafbe78d9d5c6d5ddf9f395", + "path": "transformer/diffusion_pytorch_model-00002-of-00014.safetensors", + "sha256": "2c4d362eddd2802180ac9c744849eb9ba8d9c8b984bdf9822cb02ed004b29184", + "size": 4702158032 + }, + { + "git_oid": "f69102656be6332afea3376b76a81a2190b2aa89", + "path": "transformer/diffusion_pytorch_model-00003-of-00014.safetensors", + "sha256": "949c5aafbbfa5654da730a6a7fafd75adb164d0857b095a30e8bb6d390887d69", + "size": 4933368192 + }, + { + "git_oid": "663bd8d4ca0dfa722c99ca01d54049dd7efa0515", + "path": "transformer/diffusion_pytorch_model-00004-of-00014.safetensors", + "sha256": "eef7616790105ee839766bb2027203bf2c0d87c6aa038dca84145a8675f5ce28", + "size": 4567069608 + }, + { + "git_oid": "338b0e9ce0fadd8aa27a253e50a3dbd565ded8ec", + "path": "transformer/diffusion_pytorch_model-00005-of-00014.safetensors", + "sha256": "43fdf42d638e8bc6745f713fae80c93bb301807a1a5ae7249344ce28e202a494", + "size": 4702158080 + }, + { + "git_oid": "aff5d426da7d9972f3adb29f02cffaf62f94f9c1", + "path": "transformer/diffusion_pytorch_model-00006-of-00014.safetensors", + "sha256": "6442510b34d173653f0cce5c964b935395a8f7accf0b9cc0aa31aec59805239d", + "size": 4933368232 + }, + { + "git_oid": "e23672b7490bcd26fcb50976b89903acd85fc6ff", + "path": "transformer/diffusion_pytorch_model-00007-of-00014.safetensors", + "sha256": "29f48f535c91dac76496ca821eeb16ca24bc4caf3f0cae8b920a89b1f966da6d", + "size": 4567069608 + }, + { + "git_oid": "a7c89c8303ae2cace7518f9d995522e00872788e", + "path": "transformer/diffusion_pytorch_model-00008-of-00014.safetensors", + "sha256": "c711b096c764bd60f0b8b6ad49518bfab6d614fb788c725add8741c0674a4cd8", + "size": 4702158080 + }, + { + "git_oid": "a63bcdf9be4d78b32124fcc032fbc1245f39ec05", + "path": "transformer/diffusion_pytorch_model-00009-of-00014.safetensors", + "sha256": "44428defe3976cbb87635ad200b958199e739986697cd29fdf27aeb7294b5944", + "size": 4933368232 + }, + { + "git_oid": "e3b94e045a76f7cc87b17add5983165a800214a5", + "path": "transformer/diffusion_pytorch_model-00010-of-00014.safetensors", + "sha256": "3d44939c374c9da382e9c6877e1946adf7b84e08c7a881c068f228d6849411c9", + "size": 4567069608 + }, + { + "git_oid": "72a25189806a418d770b37ae8efc9b0d53fc312b", + "path": "transformer/diffusion_pytorch_model-00011-of-00014.safetensors", + "sha256": "224d24430b58127a5577721084e0e704a0e74ec96dd7c35bc6fc0994ebd87c33", + "size": 4702158080 + }, + { + "git_oid": "9545750c42b012ae13089df388d1569d10ad2a72", + "path": "transformer/diffusion_pytorch_model-00012-of-00014.safetensors", + "sha256": "48fa2bd8fe134eef565ab2464f1c2589a6657cba0d14283dfc06b532f8961f3c", + "size": 4933368232 + }, + { + "git_oid": "5827d7644a1ebc1ce9fcf40607857257f14f5659", + "path": "transformer/diffusion_pytorch_model-00013-of-00014.safetensors", + "sha256": "be5b4b1809f9d546ffd4b3fcf41e5c1e02b819125caa6bc105c109b04c051bd3", + "size": 4567069608 + }, + { + "git_oid": "5c97657ddfa2a1cf31bb6460a25c1caeccfe5a0c", + "path": "transformer/diffusion_pytorch_model-00014-of-00014.safetensors", + "sha256": "8fbd5e6c1fb1df7ce988ca90f3d59e7610e465c7517e4b344eda4a214ba4b97d", + "size": 4644161920 + }, + { + "git_oid": "da038def25abb5a5e4be509816ee55a0b7105e09", + "path": "transformer/diffusion_pytorch_model.safetensors.index.json", + "sha256": null, + "size": 64488 + } + ], + "transformer_ref": [ + { + "git_oid": "c251a4133c81852a150c6668a570e142dc776b6e", + "path": "transformer_ref/config.json", + "sha256": null, + "size": 546 + }, + { + "git_oid": "9b0c05d22cf4683ecb0d959abf91b0ecd8cd1e4c", + "path": "transformer_ref/diffusion_pytorch_model-00001-of-00014.safetensors", + "sha256": "7a3fcad885f51560e550b2e84c9a8d8b35e62996cfd9076937e992bd23478df9", + "size": 4825958704 + }, + { + "git_oid": "51e0598afcf143178f49a45b326de0947588d51d", + "path": "transformer_ref/diffusion_pytorch_model-00002-of-00014.safetensors", + "sha256": "1638ae1dc8ae26c4ba43ad28a6d851ad8983847324bb2b468719c7c81f219706", + "size": 4702158032 + }, + { + "git_oid": "b30b53bc2e884978e8aa857dc3f941c3ea7ac816", + "path": "transformer_ref/diffusion_pytorch_model-00003-of-00014.safetensors", + "sha256": "1ef3c4954ffe5a664c2e3028e2a3241190d9c159dce6ba1136002c6af1db5353", + "size": 4933368192 + }, + { + "git_oid": "0d13668925c5f8d8441b6d10db7cf868dcb4598d", + "path": "transformer_ref/diffusion_pytorch_model-00004-of-00014.safetensors", + "sha256": "12d92f2975cfd5c5b786126385c52e5bf64884d4b4d6e60c3ef5d857c3f7469f", + "size": 4567069608 + }, + { + "git_oid": "aa394e3cfcf04c26e949cfb3db37de989ea71a79", + "path": "transformer_ref/diffusion_pytorch_model-00005-of-00014.safetensors", + "sha256": "304d41ce03d59ac94bceb055935bf4e034df0badf8b0df4ded327c08a288a4cc", + "size": 4702158080 + }, + { + "git_oid": "7f8fcc49f1b389332e3171e3f6aad4a6feb2d768", + "path": "transformer_ref/diffusion_pytorch_model-00006-of-00014.safetensors", + "sha256": "12a134b7c76d86edbe8fa2dc315f6cdaf4e1aca1b6ea4dfe4cad92df03d42eeb", + "size": 4933368232 + }, + { + "git_oid": "94bb031935b29cb39e5831635b46e1c1191e7668", + "path": "transformer_ref/diffusion_pytorch_model-00007-of-00014.safetensors", + "sha256": "b96395261359937c00fb42f4eb29306dc59b1a3368eeba52af4fb66e3e142c69", + "size": 4567069608 + }, + { + "git_oid": "0bb15268c4dee8a62b7c15ece7f0ded6aaff7654", + "path": "transformer_ref/diffusion_pytorch_model-00008-of-00014.safetensors", + "sha256": "1897a6bf3b4fc834bb82d73ca02a7afc7d38c07f50ec5382cd54cd2f91b604d1", + "size": 4702158080 + }, + { + "git_oid": "2f6e04bfeba92b1916c5b960c125cd0d55c85801", + "path": "transformer_ref/diffusion_pytorch_model-00009-of-00014.safetensors", + "sha256": "edfb38235adc96b99f55a401849befce59075a745e99c2d8c63ff358dd36443d", + "size": 4933368232 + }, + { + "git_oid": "23aa8735c6a03426bd681a44ae0704f5c8a81d08", + "path": "transformer_ref/diffusion_pytorch_model-00010-of-00014.safetensors", + "sha256": "f8710775cf3413670edd7e23861b650a3431a71a6cc14cb1080623ab6b052385", + "size": 4567069608 + }, + { + "git_oid": "0082cfd7c050bf55c405d67d026514012dd259a6", + "path": "transformer_ref/diffusion_pytorch_model-00011-of-00014.safetensors", + "sha256": "9e18acc09f84edb5b34df9628efa15cfcab8bb76e8e20c1c2e979a107a0f7215", + "size": 4702158080 + }, + { + "git_oid": "f60fcc3c5e79e132ffadf06e5e98dcbf9c7b6401", + "path": "transformer_ref/diffusion_pytorch_model-00012-of-00014.safetensors", + "sha256": "ea2e18228f8bdba1a4e0f32b155e4586df055997c45356213d05b971ba13e2f4", + "size": 4933368232 + }, + { + "git_oid": "d9ac20a69b58fcd4dabe8bfef78d29640bc291a2", + "path": "transformer_ref/diffusion_pytorch_model-00013-of-00014.safetensors", + "sha256": "1e12083b1875678f7414ff55b09cd8bb1c30b861243f9bb7ff1e75b6ad3f1bdc", + "size": 4567069608 + }, + { + "git_oid": "b1dfd11195ff1b5cd520cb7e83730f092d4ea0ff", + "path": "transformer_ref/diffusion_pytorch_model-00014-of-00014.safetensors", + "sha256": "b340f44b5690cc745d48ae399381ec15b26a4fe25d483f677ccb4960dadb50d4", + "size": 4644161920 + }, + { + "git_oid": "da038def25abb5a5e4be509816ee55a0b7105e09", + "path": "transformer_ref/diffusion_pytorch_model.safetensors.index.json", + "sha256": null, + "size": 64488 + } + ] + }, + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "tensor_tables": { + "transformer": { + "digest": "9b5e42469bcc46d65399f5625cbb1e766742faf6c460d294d499523990290e0d", + "parameters": 33122992896, + "tensors": 638 + }, + "transformer_ref": { + "digest": "9b5e42469bcc46d65399f5625cbb1e766742faf6c460d294d499523990290e0d", + "parameters": 33122992896, + "tensors": 638 + } + } +} diff --git a/rl_engine/models/minimax_h3/trajectory.py b/rl_engine/models/minimax_h3/trajectory.py new file mode 100644 index 000000000..339c2e93d --- /dev/null +++ b/rl_engine/models/minimax_h3/trajectory.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sealed trajectory artifact for MiniMax-H3 (RFC #420 ``h3_trajectory_artifact``). + +A sampling (or SDE rollout) trajectory is the unit that end-to-end claims are made about: +"native and strict differ from step 3", "strict is bitwise identical across runs". This +module defines the record that carries such a claim and the comparison that locates it. + +An artifact holds, per modality (video, audio): + +* the initial noise digest and the exact FP32 sigma grid; + +and per step: the sigmas used, and the digest of each modality's velocity and of the +latent after the step. Steps are hash-chained (``chain_sha256``) and the whole record is +sealed (``artifact_sha256``), so an edited, truncated or reordered artifact fails to load. +It is bound to the architecture fingerprint (``h3_arch_fingerprint``) and optionally to an +operator trace (``h3_operator_trace``), so "what ran" and "what it produced" travel together. + +Digests cover dtype, shape and raw bytes: one flipped bit anywhere is a different digest, +and a reshape or dtype change is too. Tensors can be kept next to the artifact +(``TrajectoryRecorder(keep_tensors=True)`` + ``save_tensors``) so ``first_drift`` can also +report how far apart two diverging tensors are (mismatch count, max abs difference). + +The core is stdlib; torch is imported only to hash or load tensors. +""" + +from __future__ import annotations + +import hashlib +import json +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +from rl_engine.models.minimax_h3 import fingerprint as fp + +SCHEMA_VERSION = "rlkernel.h3.trajectory.v1" +MODALITIES = ("video", "audio") +# Order in which a step's tensors are compared: a velocity diverges before the latent it feeds. +KINDS = ("velocity", "latent") + + +# --- hashing ------------------------------------------------------------------------------ + + +def tensor_digest(tensor: Any) -> str: + """sha256 over dtype, shape and raw bytes of a tensor (any device, any stride).""" + + import torch + + cpu = tensor.detach().contiguous().cpu() + header = f"{cpu.dtype}|{tuple(cpu.shape)}|".encode() + raw = cpu.reshape(-1).view(torch.uint8).numpy().tobytes() if cpu.numel() else b"" + return hashlib.sha256(header + raw).hexdigest() + + +def f32_hex(value: float) -> str: + """Exact text form of a FP32 value (round to FP32, then ``float.hex``).""" + + return struct.unpack(" float: + return float.fromhex(text) + + +def _sha(payload: Any) -> str: + return hashlib.sha256( + json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() + ).hexdigest() + + +# --- recording ---------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class StepRecord: + index: int + sigma: Mapping[str, str] # modality -> FP32 hex, the sigma this step starts from + sigma_next: Mapping[str, str] + velocity_sha256: Mapping[str, str] + latent_sha256: Mapping[str, str] # after the step + chain_sha256: str + + def to_dict(self) -> dict[str, Any]: + return { + "index": self.index, + "sigma": dict(self.sigma), + "sigma_next": dict(self.sigma_next), + "velocity_sha256": dict(self.velocity_sha256), + "latent_sha256": dict(self.latent_sha256), + "chain_sha256": self.chain_sha256, + } + + +def _chain(previous: str, body: Mapping[str, Any]) -> str: + return _sha({"previous": previous, "step": body}) + + +GENESIS = "0" * 64 + + +class TrajectoryRecorder: + """Builds an artifact step by step; ``seal()`` returns the immutable dict.""" + + def __init__( + self, + *, + variant: str, + label: str, + seed: int, + sigma_grids: Mapping[str, Sequence[float]], + initial_noise: Mapping[str, Any], + shifts: Mapping[str, float] | None = None, + operator_trace_sha256: str = "", + metadata: Mapping[str, Any] | None = None, + keep_tensors: bool = False, + ) -> None: + if variant not in fp.VARIANTS: + raise ValueError(f"unknown variant {variant!r}") + if set(sigma_grids) != set(initial_noise) or not set(sigma_grids) <= set(MODALITIES): + raise ValueError( + f"modalities must be a subset of {MODALITIES}, the same for grids and noise" + ) + steps = {len(grid) - 1 for grid in sigma_grids.values()} + if len(steps) != 1 or steps.pop() < 1: + raise ValueError("every sigma grid needs the same length of at least 2 (steps + 1)") + self.variant, self.label, self.seed = variant, label, int(seed) + self.modalities = tuple(m for m in MODALITIES if m in sigma_grids) + self.grids = {m: [f32_hex(s) for s in sigma_grids[m]] for m in self.modalities} + self.steps = len(self.grids[self.modalities[0]]) - 1 + self.shifts = {m: float((shifts or {}).get(m, 0.0)) for m in self.modalities} + self.operator_trace_sha256 = operator_trace_sha256 + self.metadata = dict(metadata or {}) + self.initial = { + m: { + "sha256": tensor_digest(initial_noise[m]), + "shape": list(initial_noise[m].shape), + "dtype": str(initial_noise[m].dtype), + } + for m in self.modalities + } + self.records: list[StepRecord] = [] + self.keep_tensors = keep_tensors + self.tensors: dict[str, Any] = {} + if keep_tensors: + for m in self.modalities: + self.tensors[tensor_key(-1, "latent", m)] = initial_noise[m].detach().cpu().clone() + + def record_step( + self, + index: int, + *, + velocity: Mapping[str, Any], + latent: Mapping[str, Any], + ) -> StepRecord: + """Record step ``index`` (consecutive from 0): its velocities and resulting latents.""" + + if index != len(self.records): + raise ValueError(f"expected step {len(self.records)}, got {index}") + if index >= self.steps: + raise ValueError(f"trajectory has only {self.steps} steps") + for name, group in (("velocity", velocity), ("latent", latent)): + if set(group) != set(self.modalities): + raise ValueError(f"{name} must cover exactly {self.modalities}") + body: dict[str, Any] = { + "index": index, + "sigma": {m: self.grids[m][index] for m in self.modalities}, + "sigma_next": {m: self.grids[m][index + 1] for m in self.modalities}, + "velocity_sha256": {m: tensor_digest(velocity[m]) for m in self.modalities}, + "latent_sha256": {m: tensor_digest(latent[m]) for m in self.modalities}, + } + previous = self.records[-1].chain_sha256 if self.records else GENESIS + record = StepRecord(chain_sha256=_chain(previous, body), **body) + self.records.append(record) + if self.keep_tensors: + for kind, group in (("velocity", velocity), ("latent", latent)): + for m in self.modalities: + self.tensors[tensor_key(index, kind, m)] = group[m].detach().cpu().clone() + return record + + def seal(self) -> dict[str, Any]: + if len(self.records) != self.steps: + raise ValueError(f"recorded {len(self.records)} of {self.steps} steps") + artifact = { + "schema_version": SCHEMA_VERSION, + "architecture_fingerprint_sha256": fp.build_fingerprint()["fingerprint_sha256"], + "variant": self.variant, + "label": self.label, + "seed": self.seed, + "steps": self.steps, + "modalities": list(self.modalities), + "shifts": self.shifts, + "sigma_grids": self.grids, + "initial_noise": self.initial, + "operator_trace_sha256": self.operator_trace_sha256, + "metadata": self.metadata, + "step_records": [record.to_dict() for record in self.records], + "final_latent_sha256": dict(self.records[-1].latent_sha256), + } + artifact["artifact_sha256"] = _sha(artifact) + return artifact + + def save_tensors(self, path: Path) -> None: + """Write kept tensors as safetensors (keys ``step.kind.modality``; -1 = initial noise).""" + + if not self.keep_tensors: + raise ValueError("recorder was created with keep_tensors=False") + from safetensors.torch import save_file + + save_file({key: value.contiguous() for key, value in self.tensors.items()}, str(path)) + + +def tensor_key(step: int, kind: str, modality: str) -> str: + return f"{step}.{kind}.{modality}" + + +# --- validation --------------------------------------------------------------------------- + + +def verify(artifact: Mapping[str, Any], *, check_fingerprint: bool = True) -> None: + """Fail closed unless the artifact is internally consistent and sealed.""" + + if artifact.get("schema_version") != SCHEMA_VERSION: + raise ValueError(f"unsupported trajectory schema {artifact.get('schema_version')!r}") + body = {k: v for k, v in artifact.items() if k != "artifact_sha256"} + if artifact.get("artifact_sha256") != _sha(body): + raise ValueError("artifact_sha256 does not match the content") + modalities = tuple(artifact["modalities"]) + if not modalities or not set(modalities) <= set(MODALITIES): + raise ValueError("invalid modalities") + steps = artifact["steps"] + if len(artifact["step_records"]) != steps: + raise ValueError("step_records length differs from steps") + for modality in modalities: + grid = [f32_from_hex(s) for s in artifact["sigma_grids"][modality]] + if len(grid) != steps + 1: + raise ValueError(f"{modality}: sigma grid needs steps + 1 entries") + if any(b > a for a, b in zip(grid, grid[1:])): + raise ValueError(f"{modality}: sigma grid is not non-increasing") + if grid[0] <= 0.0 or grid[-1] < 0.0 or grid[0] > 1.0: + raise ValueError(f"{modality}: sigma grid must start in (0, 1] and end >= 0") + previous = GENESIS + for index, record in enumerate(artifact["step_records"]): + body = {k: v for k, v in record.items() if k != "chain_sha256"} + if record["index"] != index: + raise ValueError(f"step {index}: out-of-order record {record['index']}") + for modality in modalities: + if record["sigma"][modality] != artifact["sigma_grids"][modality][index]: + raise ValueError(f"step {index}: sigma disagrees with the grid") + if record["sigma_next"][modality] != artifact["sigma_grids"][modality][index + 1]: + raise ValueError(f"step {index}: sigma_next disagrees with the grid") + if record["chain_sha256"] != _chain(previous, body): + raise ValueError(f"step {index}: hash chain broken") + previous = record["chain_sha256"] + if artifact["final_latent_sha256"] != artifact["step_records"][-1]["latent_sha256"]: + raise ValueError("final_latent_sha256 differs from the last step") + if check_fingerprint and ( + artifact["architecture_fingerprint_sha256"] != fp.build_fingerprint()["fingerprint_sha256"] + ): + raise ValueError("artifact was recorded against a different H3 architecture fingerprint") + + +def dumps(artifact: Mapping[str, Any]) -> str: + verify(artifact) + return json.dumps(artifact, indent=1, sort_keys=True) + "\n" + + +def loads(text: str, *, check_fingerprint: bool = True) -> dict[str, Any]: + artifact = json.loads(text) + verify(artifact, check_fingerprint=check_fingerprint) + return artifact + + +# --- comparison --------------------------------------------------------------------------- + + +def input_mismatches(lhs: Mapping[str, Any], rhs: Mapping[str, Any]) -> list[str]: + """Differences in what the two runs were *given*: a drift comparison needs none.""" + + problems: list[str] = [] + for key in ("variant", "seed", "steps", "modalities", "shifts", "sigma_grids"): + if lhs[key] != rhs[key]: + problems.append(f"{key} differs") + for modality in lhs["modalities"]: + if lhs["initial_noise"].get(modality) != rhs["initial_noise"].get(modality): + problems.append(f"initial noise differs for {modality}") + return problems + + +def first_drift( + lhs: Mapping[str, Any], + rhs: Mapping[str, Any], + *, + tensors_lhs: Path | None = None, + tensors_rhs: Path | None = None, + require_same_inputs: bool = True, +) -> dict[str, Any] | None: + """First (step, kind, modality) at which two trajectories differ, or ``None`` if bitwise equal. + + With both tensor files the result also carries ``mismatch_count``, ``numel`` and + ``max_abs_diff``. Differing inputs (seed, noise, grids) make a drift meaningless, so + they raise unless ``require_same_inputs=False``. + """ + + verify(lhs) + verify(rhs) + if require_same_inputs: + problems = input_mismatches(lhs, rhs) + if problems: + raise ValueError("trajectories are not comparable: " + "; ".join(problems)) + for a, b in zip(lhs["step_records"], rhs["step_records"]): + for kind in KINDS: + for modality in lhs["modalities"]: + if a[f"{kind}_sha256"][modality] != b[f"{kind}_sha256"][modality]: + drift = {"step": a["index"], "kind": kind, "modality": modality} + if tensors_lhs is not None and tensors_rhs is not None: + drift.update( + _tensor_difference( + tensors_lhs, tensors_rhs, tensor_key(a["index"], kind, modality) + ) + ) + return drift + return None + + +def _tensor_difference(path_a: Path, path_b: Path, key: str) -> dict[str, Any]: + import torch + from safetensors import safe_open + + with safe_open(str(path_a), framework="pt") as fa, safe_open(str(path_b), framework="pt") as fb: + a, b = fa.get_tensor(key), fb.get_tensor(key) + if a.shape != b.shape or a.dtype != b.dtype: + return {"shape_or_dtype_differs": True} + mismatch = a != b + nan_mismatch = torch.isnan(a) != torch.isnan(b) + differing = mismatch & ~(torch.isnan(a) & torch.isnan(b)) | nan_mismatch + diff = (a.double() - b.double()).abs() + diff = torch.where(torch.isfinite(diff), diff, torch.zeros_like(diff)) + return { + "mismatch_count": int(differing.sum()), + "numel": int(a.numel()), + "max_abs_diff": float(diff.max()) if diff.numel() else 0.0, + } + + +def drift_report( + lhs: Mapping[str, Any], + rhs: Mapping[str, Any], + *, + tensors_lhs: Path | None = None, + tensors_rhs: Path | None = None, +) -> dict[str, Any]: + """Every (step, kind, modality) that differs, in order, plus the headline numbers. + + ``bitwise_equal`` is the strict claim; ``differing`` is how many of the compared + tensors differ; with tensor files each entry carries mismatch count and max abs diff. + """ + + verify(lhs) + verify(rhs) + problems = input_mismatches(lhs, rhs) + if problems: + raise ValueError("trajectories are not comparable: " + "; ".join(problems)) + entries: list[dict[str, Any]] = [] + total = 0 + for a, b in zip(lhs["step_records"], rhs["step_records"]): + for kind in KINDS: + for modality in lhs["modalities"]: + total += 1 + if a[f"{kind}_sha256"][modality] != b[f"{kind}_sha256"][modality]: + entry = {"step": a["index"], "kind": kind, "modality": modality} + if tensors_lhs is not None and tensors_rhs is not None: + entry.update( + _tensor_difference( + tensors_lhs, tensors_rhs, tensor_key(a["index"], kind, modality) + ) + ) + entries.append(entry) + return { + "labels": [lhs["label"], rhs["label"]], + "steps": lhs["steps"], + "compared": total, + "differing": len(entries), + "bitwise_equal": not entries, + "first_drift": entries[0] if entries else None, + "final_latent_equal": lhs["final_latent_sha256"] == rhs["final_latent_sha256"], + "entries": entries, + } + + +def sigmas_of(artifact: Mapping[str, Any], modality: str) -> list[float]: + return [f32_from_hex(text) for text in artifact["sigma_grids"][modality]] + + +def modalities_of(artifact: Mapping[str, Any]) -> Iterable[str]: + return tuple(artifact["modalities"]) 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/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py new file mode 100644 index 000000000..0ad2c1b0e --- /dev/null +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -0,0 +1,98 @@ +# 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 + +from rl_engine.reference.minimax_h3.adaln_row_gather import positions_by_row + +_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, 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 + + +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/projection.py b/rl_engine/reference/minimax_h3/projection.py new file mode 100644 index 000000000..1ab933156 --- /dev/null +++ b/rl_engine/reference/minimax_h3/projection.py @@ -0,0 +1,186 @@ +# 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`` is 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 3b589a13d..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" ) @@ -246,6 +268,16 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): ) 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, ...]: @@ -822,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, @@ -831,6 +879,14 @@ 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, + ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -875,9 +931,15 @@ 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], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "musa": { "logp": [ @@ -930,9 +992,15 @@ 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], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -962,9 +1030,15 @@ 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], + "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..0f895613e --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,752 @@ +# 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.qkv = registry.get_op("h3_qkv_gemm", device="cuda") + 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 _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, + 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), + _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"), + _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 8eb509d33..92455bae0 100644 --- a/rl_engine/validation/models/h3_cases.py +++ b/rl_engine/validation/models/h3_cases.py @@ -38,3 +38,151 @@ 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) + + +# 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 317598fa9..3aa1aa191 100644 --- a/rl_engine/validation/models/h3_chain.py +++ b/rl_engine/validation/models/h3_chain.py @@ -33,6 +33,8 @@ 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, @@ -150,6 +152,62 @@ 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, + ) +) + + +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", @@ -209,6 +267,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(), } @@ -241,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"]} @@ -313,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_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 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 83a7b6b82..3eec1818f 100644 --- a/rl_engine/validation/models/h3_provider.py +++ b/rl_engine/validation/models/h3_provider.py @@ -92,3 +92,101 @@ 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 + + +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_refiner.py b/rl_engine/validation/models/h3_refiner.py new file mode 100644 index 000000000..35f1e3b35 --- /dev/null +++ b/rl_engine/validation/models/h3_refiner.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 token-refiner blocks, node by node (RFC #420 ``token_refiner_block``). + +The text stream is refined by two plain pre-norm blocks and a final RMSNorm before it is +scattered into the packed sequence. A refiner block is the transformer block minus AdaLN, +minus RoPE and with plain (ungated) residual adds:: + + h = h + attn(norm1(h)) attn: q/k/v -> per-head RMSNorm -> full attention -> o_proj + h = h + ff(norm2(h)) ff: gate_up -> SwiGLU -> down + +The block is a fixed graph of 14 nodes (``NODES``, the vocabulary pinned by +``fingerprint.REFINER_NODES``) over the same operators the transformer block binds, so +every qualified row is reused and the only arithmetic owned by this row is the two plain +BF16 residual adds. Each node runs three ways, as in ``h3_block``: provider (diffusers op +sequence), candidate (RL-Kernel bindings, chained and isolated) and an FP32 golden, and +names the RFC row that owns it with its binding status (``row`` / ``interim`` / +``reference``; see ``h3_block``). The report carries repeat and batch-row bitwise checks, +agreement with the provider, error against the golden, and the first drifting node. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable + +import torch + +from rl_engine.models.minimax_h3 import fingerprint as fp +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.validation.models import h3_block as hb +from rl_engine.validation.models import h3_provider as prov + +HIDDEN = hb.HIDDEN +PREFIX = "token_refiner.refiner_blocks" +FINAL_NORM = "token_refiner.final_norm.weight" +Binding = hb.Binding + + +@dataclass(frozen=True) +class Node: + name: str + inputs: tuple[str, ...] + binding: Binding + candidate: Callable[..., Any] # (ops, ctx, *inputs) + provider: Callable[..., Any] # (ctx, *inputs) + golden: Callable[..., Any] # (ctx, *inputs) + provider_bitwise_isolated: bool + + +def _w(ctx: dict[str, Any], name: str) -> torch.Tensor: + return ctx["weights"][f"{ctx['prefix']}.{name}"] + + +def _norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_rmsnorm", "row", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(x, _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(x, _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _qkv_node(name: str, weight: str) -> Node: + 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, + ) + + +def _head_norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_qk_rmsnorm_d128", "interim", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(hb._heads(x), _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(hb._heads(x), _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(hb._heads(x), _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _linear_node(name: str, source: str, weight: str, row: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding( + row, "interim", "triton _det_gemm_fp32_kernel (FP32 K loop, one BF16 rounding)" + ), + candidate=lambda ops, ctx, x: hb._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, + ) + + +def _residual_node(name: str, residual: str, sublayer: str) -> Node: + """Plain residual add: one BF16 add per element, owned by this row.""" + + return Node( + name=name, + inputs=(residual, sublayer), + binding=Binding("token_refiner_block", "row", "torch.add (BF16 elementwise)"), + candidate=lambda ops, ctx, r, y: r + y, + provider=lambda ctx, r, y: r + y, + golden=lambda ctx, r, y: r.float() + y.float(), + provider_bitwise_isolated=True, + ) + + +NODES: list[Node] = [ + _norm_node("norm1", "hidden", "norm1.weight"), + _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"), + Node( + name="attention", + inputs=("q_norm", "k_norm", "v_proj"), + binding=Binding("h3_full_attention", "interim", "DeterministicAttentionOp(causal=False)"), + candidate=lambda ops, ctx, q, k, v: hb._cand_attention(ops, q, k, v), + provider=lambda ctx, q, k, v: prov.provider_attention(q, k, hb._heads(v)), + golden=lambda ctx, q, k, v: gold.golden_attention(q, k, hb._heads(v)), + provider_bitwise_isolated=False, + ), + _linear_node("o_proj", "attention", "attn.to_out.0.weight", "h3_attention_o_gemm"), + _residual_node("residual_attn", "hidden", "o_proj"), + _norm_node("norm2", "residual_attn", "norm2.weight"), + _linear_node("ffn_gate_up", "norm2", "ff.net.0.proj.weight", "h3_ffn_gate_up_gemm"), + Node( + name="swiglu", + inputs=("ffn_gate_up",), + binding=Binding("h3_swiglu", "interim", "SwiGLUCudaOp"), + candidate=lambda ops, ctx, x: hb._cand_swiglu(ops, x), + provider=lambda ctx, x: prov.provider_swiglu(x), + golden=lambda ctx, x: gold.golden_swiglu(x), + provider_bitwise_isolated=False, + ), + _linear_node("ffn_down", "swiglu", "ff.net.2.weight", "h3_ffn_down_gemm"), + _residual_node("residual_mlp", "residual_attn", "ffn_down"), +] +NODE_NAMES = tuple(node.name for node in NODES) +OUTPUT = NODES[-1].name +PARAM_LEAVES = ( + "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", +) + + +def make_refiner_context( + weights: dict[str, torch.Tensor], *, block: int, tokens: int, seed: int = 0, batch: int = 1 +) -> dict[str, Any]: + """Weights of refiner ``block`` and a seeded BF16 text-stream input ``(batch, tokens, 5376)``. + + Batch row 0 is the same tensor at every ``batch``; further rows get their own seeds. + """ + + rows = [] + for b in range(batch): + gen = torch.Generator(device="cpu").manual_seed(seed if b == 0 else seed * 1009 + b) + rows.append(torch.randn(1, tokens, HIDDEN, generator=gen)) + hidden = torch.cat(rows).to(torch.bfloat16).cuda() + return { + "weights": weights, + "prefix": f"{PREFIX}.{block}", + "block": block, + "hidden": hidden, + "tokens": tokens, + } + + +def run_graph(mode: str, ctx: dict[str, Any], ops: hb.CandidateOps | None = None) -> dict[str, Any]: + """Every node's output for one refiner block ("candidate", "provider" or "golden").""" + + values: dict[str, Any] = {"hidden": ctx["hidden"]} + observer = ctx.get("node_observer") + with hb.strict_fp32_matmul(): + for node in NODES: + args = [values[name] for name in node.inputs] + if 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) + if observer is not None: + observer(node.name) + return values + + +def run_isolated( + ctx: dict[str, Any], ops: hb.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 run_refiner( + mode: str, + weights: dict[str, torch.Tensor], + hidden: torch.Tensor, + ops: hb.CandidateOps | None = None, + *, + num_blocks: int = 2, + observer_factory: Callable[[int], Callable[[str], None]] | None = None, +) -> torch.Tensor: + """The whole text refiner: ``num_blocks`` blocks then the final RMSNorm.""" + + h = hidden + for block in range(num_blocks): + ctx = {"weights": weights, "prefix": f"{PREFIX}.{block}", "hidden": h} + if observer_factory is not None: + ctx["node_observer"] = observer_factory(block) + h = run_graph(mode, ctx, ops)[OUTPUT] + weight = weights[FINAL_NORM] + if mode == "candidate": + assert ops is not None + return ops.norm(h, weight) + if mode == "provider": + return prov.provider_head_rmsnorm(h, weight) + return gold.golden_rmsnorm(h, weight) + + +def run_refiner_case( + weights: dict[str, torch.Tensor], + *, + block: int = 0, + tokens: int = 257, + seed: int = 0, + ops: hb.CandidateOps | None = None, +) -> dict[str, Any]: + """Forward replay of one refiner block: per-node agreement, invariance, accuracy, drift.""" + + ops = ops or hb.CandidateOps() + ctx = make_refiner_context(weights, block=block, tokens=tokens, seed=seed, batch=1) + ctx2 = make_refiner_context(weights, block=block, tokens=tokens, 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] = { + "block": block, + "tokens": tokens, + "seed": seed, + "backends": ops.backends(), + "nodes": [], + } + firsts: dict[str, str | None] = dict.fromkeys( + ("first_drift", "first_isolated_drift", "first_repeat_drift", "first_batch_drift") + ) + for node in NODES: + name = node.name + row0 = hb._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": hb.bitwise(repeat[name], chained[name]), + "batch_row_bitwise_equal": hb.bitwise(row0, chained[name]), + "chained_vs_provider": { + "bitwise_equal": hb.bitwise(chained[name], provider[name]), + "max_abs": hb.max_abs(chained[name], provider[name]), + }, + "isolated_vs_provider": { + "bitwise_equal": hb.bitwise(isolated[name], provider[name]), + "max_abs": hb.max_abs(isolated[name], provider[name]), + }, + "candidate_rel_err_vs_golden": hb.rel_err(chained[name], golden[name]), + "provider_rel_err_vs_golden": hb.rel_err(provider[name], golden[name]), + "candidate_sha256": hb.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": hb.rel_err(chained[OUTPUT], golden[OUTPUT]), + "provider_rel_err_vs_golden": hb.rel_err(provider[OUTPUT], golden[OUTPUT]), + "candidate_vs_provider_max_abs": hb.max_abs(chained[OUTPUT], provider[OUTPUT]), + } + return report + + +def graph_matches_fingerprint() -> bool: + """The node vocabulary and owners equal ``fingerprint.REFINER_NODES``.""" + + return [(n.name, n.binding.rfc_row) for n in NODES] == [ + (name, row) for name, row in fp.REFINER_NODES + ] diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py index 90aa4463d..6f2852dc8 100644 --- a/rl_engine/validation/models/h3_report.py +++ b/rl_engine/validation/models/h3_report.py @@ -20,11 +20,17 @@ 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 h3_weights_dir, load_h3_conditioning_weights +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 = ( @@ -94,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 @@ -110,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"))) @@ -157,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) @@ -228,6 +238,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 = [], [] @@ -303,6 +315,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): @@ -388,10 +402,11 @@ def backward(outputs, grads=grads): def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: from rl_engine.validation.models.h3_chain import run_backward_case - from rl_engine.validation.models.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]) + 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() @@ -400,7 +415,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) @@ -425,16 +442,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 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 # --------------------------------------------------------------------------- # @@ -497,15 +520,20 @@ def backward(output, grad=grad): def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: op = registry.get_op("h3_rmsnorm", device="cuda") - weights = load_h3_conditioning_weights("cuda") - gen = torch.Generator(device="cuda").manual_seed(0) - 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) + } + 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)) @@ -542,17 +570,20 @@ 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] = { - "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, @@ -560,12 +591,242 @@ 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 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, @@ -573,4 +834,6 @@ 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, + "final_adaln_out": _final_accuracy, } diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py index c6e5942c2..d1c051415 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,76 @@ 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), + } + + +# --- Tensors straight from the released shards -------------------------------------------------- + +SHARD_ROOTS = ("hf_full/transformer", "hf/transformer") # under $RL_KERNEL_H3_WEIGHTS + + +def h3_shard_root() -> Path | None: + """The first shard directory under the weights dir that has an index, or None.""" + + base = h3_weights_dir() + if base is None: + return None + for sub in SHARD_ROOTS: + root = base / sub + if (root / "diffusion_pytorch_model.safetensors.index.json").is_file(): + return root + return None + + +def load_h3_tensors( + names: list[str], device: torch.device | str = "cpu" +) -> dict[str, torch.Tensor]: + """Load named tensors from whichever released shards are present. + + Resolves each name through the shard index and reads only that tensor + (safetensors memory-maps). Raises ``FileNotFoundError`` when the weights + directory, the index or a needed shard is missing, so tests can skip. + """ + + from safetensors import safe_open + + root = h3_shard_root() + if root is None: + raise FileNotFoundError(f"no shard index under ${WEIGHTS_ENV}/{{{','.join(SHARD_ROOTS)}}}") + weight_map = json.loads((root / "diffusion_pytorch_model.safetensors.index.json").read_text())[ + "weight_map" + ] + by_shard: dict[str, list[str]] = {} + for name in names: + if name not in weight_map: + raise KeyError(f"{name} is not in the H3 shard index") + by_shard.setdefault(weight_map[name], []).append(name) + out: dict[str, torch.Tensor] = {} + for shard, wanted in by_shard.items(): + path = root / shard + if not path.is_file(): + raise FileNotFoundError(f"shard {shard} is not downloaded under {root}") + with safe_open(str(path), framework="pt", device="cpu") as handle: + for name in wanted: + out[name] = handle.get_tensor(name).to(device) + return out 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 35835fc8d..bac4b4158 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -55,6 +55,8 @@ 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, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -99,6 +101,8 @@ 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)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -523,6 +527,47 @@ 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 _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 157d16ba6..ee855c22c 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -368,6 +368,34 @@ 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.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/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_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..f93561263 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,320 @@ +# 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 os +import subprocess +import sys +import textwrap +from pathlib import Path + +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_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) + 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_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_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.""" + + 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): + _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" + + +@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}" + + +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 == "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": + 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", + "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"), + ): + done = _gate_backward_in_child(mutation) + assert done.returncode != 0 and message in done.stderr, (mutation, done.stderr) 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 b2d0b5560..8d3d985c1 100644 --- a/tests/models/minimax_h3/test_h3_adaln_row_gather.py +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -445,7 +445,7 @@ def test_backward_writes_the_requested_dtype_whatever_the_gradient_dtype( ).to(out_dtype) assert torch.equal(got, expected) - def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_boundaries(self): + def test_backward_with_no_table_rows_or_no_boundaries_is_rejected(self): from rl_engine.backends.extension import _C _cuda_op() @@ -453,8 +453,8 @@ def test_backward_with_no_table_rows_returns_an_empty_table_and_rejects_no_bound zero = torch.zeros(0, dtype=torch.int64, device="cuda") pos = torch.arange(4, device="cuda") one = torch.zeros(1, dtype=torch.int64, device="cuda") # R = 0 - empty = _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) - assert empty.shape == (0, 48) + with pytest.raises(RuntimeError, match="must be positive"): + _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, one, torch.float32) with pytest.raises(RuntimeError, match="R \\+ 1 entries"): _C.h3_adaln_row_gather_backward(grad, pos, zero, zero, zero, torch.float32) @@ -512,5 +512,5 @@ def test_metadata_on_another_device_than_grad_is_rejected(self, name): args = _backward_args() args[name] = args[name].to("cuda:1") - with pytest.raises(RuntimeError, match="grad's device"): + with pytest.raises(RuntimeError, match="tile metadata"): _C.h3_adaln_row_gather_backward(**args) diff --git a/tests/models/minimax_h3/test_h3_arch_fingerprint.py b/tests/models/minimax_h3/test_h3_arch_fingerprint.py new file mode 100644 index 000000000..59fb1f98a --- /dev/null +++ b/tests/models/minimax_h3/test_h3_arch_fingerprint.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""``h3_arch_fingerprint`` (RFC #420): the pinned MiniMax-H3 architecture facts. + +Everything here is stdlib/CPU. The tests pin three things: + +* the analytic tensor table (derived from ``config.json``) equals the digest of the + real checkpoint headers, so the fingerprint cannot drift from the release; +* the graph vocabulary equals the operator-validation graph in ``h3_block`` and + every node is owned by exactly one RFC row; +* corrupted inputs (wrong shape, wrong dtype, wrong config, edited fixture) fail closed. +""" + +from __future__ import annotations + +import copy +import json +import os +import struct +from pathlib import Path + +import pytest + +from rl_engine.models.minimax_h3 import fingerprint as fp +from rl_engine.validation.models import h3_block + +SNAPSHOT = fp.pinned_checkpoint() + + +def test_fixture_is_pinned_to_the_documented_revision() -> None: + assert SNAPSHOT["repo"] == fp.REPO + assert SNAPSHOT["revision"] == fp.REVISION + assert set(SNAPSHOT["tensor_tables"]) == set(fp.VARIANTS.values()) + + +@pytest.mark.parametrize("variant", sorted(fp.VARIANTS)) +def test_analytic_table_equals_checkpoint_headers(variant: str) -> None: + table = fp.expected_tensor_table(fp.transformer_config(variant)) + pinned = SNAPSHOT["tensor_tables"][fp.VARIANTS[variant]] + assert fp.tensor_table_digest(table) == pinned["digest"] + assert fp.parameter_count(table) == pinned["parameters"] == 33_122_992_896 + assert len(table) == pinned["tensors"] == 638 + + +def test_variants_share_one_architecture() -> None: + assert fp.transformer_config("FL2VA") == fp.transformer_config("Ref2VA") + digests = {info["digest"] for info in SNAPSHOT["tensor_tables"].values()} + assert len(digests) == 1 + + +def test_derived_quantities() -> None: + derived = fp.derived_quantities(fp.transformer_config()) + assert derived["attention_inner_dim"] == 7168 + assert derived["attention_inner_dim_differs_from_hidden"] is True + assert derived["rotary_channels"] == 96 + assert derived["rotary_pass_through_channels"] == 32 + assert derived["adaln_projection_out"] == 96768 + assert derived["ffn_gate_up_out"] == 28672 + + +def test_mixed_precision_matches_the_declared_modules() -> None: + table = fp.expected_tensor_table(fp.transformer_config()) + for name, (dtype, _shape) in table.items(): + expect_fp32 = any(name.startswith(module + ".") for module in fp.FP32_MODULES) + assert dtype == ("F32" if expect_fp32 else "BF16"), name + + +def test_fingerprint_is_stable_and_self_consistent() -> None: + first = fp.build_fingerprint() + second = fp.build_fingerprint() + assert first == second + assert first["fingerprint_sha256"] == fp.fingerprint_digest(first) + assert first["graph"]["blocks"] == 50 + assert first["graph"]["refiner_blocks"] == 2 + assert first["graph"]["block_nodes_per_block"] == 17 + assert first["graph"]["refiner_nodes_per_block"] == len(fp.REFINER_NODES) + expected_nodes = ( + len(fp.STEM_NODES) + + 1 # packed_index_copy + + 2 * len(fp.REFINER_NODES) + + 1 # refiner final norm + + 50 * len(fp.BLOCK_NODES) + + len(fp.TAIL_NODES) + ) + assert first["graph"]["nodes"] == expected_nodes + + +def test_committed_golden_matches() -> None: + golden = ( + Path(__file__).resolve().parents[3] / "rl_engine/models/minimax_h3/arch_fingerprint.json" + ) + if not golden.exists(): + pytest.skip("no committed golden fingerprint") + assert json.loads(golden.read_text()) == fp.build_fingerprint() + + +# --- graph ------------------------------------------------------------------------------------ + + +def test_block_vocabulary_equals_the_operator_validation_graph() -> None: + assert fp.BLOCK_NODE_NAMES == h3_block.NODE_NAMES + + +def test_block_node_owners_equal_h3_block_bindings() -> None: + owners = dict(fp.BLOCK_NODES) + bindings = {node.name: node.binding.rfc_row for node in h3_block.NODES} + assert owners == bindings + + +def test_every_graph_node_has_exactly_one_owner_and_a_precision() -> None: + graph = fp.build_graph() + paths = [node.path for node in graph] + assert len(paths) == len(set(paths)) + assert all(node.row for node in graph) + assert {node.precision for node in graph} <= {"fp32", "bf16"} + + +def test_graph_covers_both_precisions_in_the_right_places() -> None: + by_path = {node.path: node for node in fp.build_graph()} + for path in ("proj_in", "audio_proj_in", "time_embedder", "proj_out", "audio_proj_out"): + assert by_path[path].precision == "fp32", path + assert by_path["transformer_blocks.0.q_proj"].precision == "bf16" + assert by_path["transformer_blocks.49.ffn_down"].precision == "bf16" + + +def test_graph_order_is_execution_order() -> None: + paths = [node.path for node in fp.build_graph()] + assert paths.index("transformer_blocks.0.adaln_projection") < paths.index( + "transformer_blocks.1.adaln_projection" + ) + assert paths.index("token_refiner.refiner_blocks.1.ffn_down") < paths.index( + "transformer_blocks.0.adaln_projection" + ) + assert paths[-1] == "packed_index_select" + + +def test_unclaimed_rows_are_visible_in_the_graph() -> None: + rows = set(fp.rows_in_graph()) + for row in ( + "audio_input_projection_fp32", + "audio_output_projection_fp32", + "text_context_projection", + "h3_qkv_gemm", + "token_refiner_block", + ): + assert row in rows + + +# --- scheduler -------------------------------------------------------------------------------- + + +def test_scheduler_shifts_match_the_released_configs() -> None: + assert float(SNAPSHOT["configs"]["scheduler"]["shift"]) == fp.SCHEDULER["video_shift"] == 12.0 + assert ( + float(SNAPSHOT["configs"]["audio_scheduler"]["shift"]) == fp.SCHEDULER["audio_shift"] == 3.0 + ) + assert fp.SCHEDULER["eta"] == 0.0 + + +def test_component_hashes_cover_every_pinned_directory() -> None: + hashes = fp.component_hashes() + assert set(hashes) == set(SNAPSHOT["dirs"]) + assert all(len(value) == 64 for value in hashes.values()) + + +# --- fail-closed ------------------------------------------------------------------------------ + + +def _write_shard(path: Path, tensors: dict[str, tuple[str, tuple[int, ...]]]) -> None: + header = { + name: {"dtype": dtype, "shape": list(shape), "data_offsets": [0, 0]} + for name, (dtype, shape) in tensors.items() + } + raw = json.dumps(header).encode() + path.write_bytes(struct.pack(" None: + table = fp.expected_tensor_table(fp.transformer_config()) + subset = dict(list(table.items())[:40]) + _write_shard(tmp_path / "a.safetensors", subset) + assert fp.validate_checkpoint_headers(tmp_path) == 40 + + +def test_checkpoint_headers_reject_wrong_dtype_shape_and_unknown_names(tmp_path: Path) -> None: + table = fp.expected_tensor_table(fp.transformer_config()) + name = "transformer_blocks.0.attn.to_q.weight" + dtype, shape = table[name] + bad_dtype = "F32" if dtype == "BF16" else "BF16" + for tensors, needle in ( + ({name: (bad_dtype, shape)}, "dtype"), + ({name: (dtype, shape[::-1])}, "shape"), + ({"transformer_blocks.0.surprise.weight": (dtype, shape)}, "surprise"), + ): + directory = tmp_path / needle + directory.mkdir() + _write_shard(directory / "a.safetensors", tensors) + with pytest.raises(ValueError, match=needle): + fp.validate_checkpoint_headers(directory) + + +def test_checkpoint_headers_require_shards(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="no safetensors"): + fp.validate_checkpoint_headers(tmp_path) + + +def test_edited_fixture_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + broken = copy.deepcopy(dict(SNAPSHOT)) + broken["tensor_tables"]["transformer"] = {"digest": "0" * 64, "parameters": 1, "tensors": 1} + monkeypatch.setattr(fp, "pinned_checkpoint", lambda: broken) + with pytest.raises(ValueError, match="does not match"): + fp.build_fingerprint() + + +def test_config_change_changes_the_analytic_digest() -> None: + config = dict(fp.transformer_config()) + base = fp.tensor_table_digest(fp.expected_tensor_table(config)) + config["num_layers"] = config["num_layers"] - 1 + assert fp.tensor_table_digest(fp.expected_tensor_table(config)) != base + + +# --- live artifacts (skip when unavailable) --------------------------------------------------- + + +def test_local_shard_headers_match_the_fingerprint() -> None: + root = os.environ.get("RL_KERNEL_H3_WEIGHTS") + directory = Path(root) / "hf" / "transformer" if root else None + if directory is None or not list(directory.glob("*.safetensors")): + pytest.skip("RL_KERNEL_H3_WEIGHTS/hf/transformer has no shards") + assert fp.validate_checkpoint_headers(directory) > 0 + + +def test_diffusers_model_on_meta_device_matches_the_fingerprint() -> None: + diffusers = pytest.importorskip("diffusers") + torch = pytest.importorskip("torch") + cls = getattr(diffusers, fp.ARCHITECTURE, None) + if cls is None: + pytest.skip(f"diffusers has no {fp.ARCHITECTURE}") + with torch.device("meta"): + model = cls.from_config(fp.transformer_config()) + fp.validate_diffusers_model(model) + # A model with one fewer block must be rejected. + smaller = dict(fp.transformer_config()) + smaller["num_layers"] -= 1 + with torch.device("meta"): + bad = cls.from_config(smaller) + with pytest.raises(ValueError): + fp.validate_diffusers_model(bad) diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py index 95af1ff18..bd3a44363 100644 --- a/tests/models/minimax_h3/test_h3_benchmark_timing.py +++ b/tests/models/minimax_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 cpu_generators = {} 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 ae039698b..091001e36 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 @@ -87,3 +90,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_cli.py b/tests/models/minimax_h3/test_h3_cli.py index 0732760b6..d8c8db7c7 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/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_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..dc2bb8153 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -0,0 +1,236 @@ +# 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 + + @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 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..657f0b4cb --- /dev/null +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -0,0 +1,504 @@ +# 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 + + +def test_rmsnorm_rejects_modulation_rows_that_overlap(native_extension, native_case): + """A table of rows needs a row stride of at least N; an overlapping view is not one.""" + + table = torch.arange(32, dtype=torch.float32, device="cuda") + overlapping = table.as_strided((2, 8), (4, 1)) + with pytest.raises(RuntimeError, match="row stride 4 is smaller than N=8"): + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, overlapping, overlapping, + native_case["index"], + ) # fmt: skip + + +def test_rmsnorm_reads_a_view_of_a_wider_table_by_its_row_pitch(native_extension, native_case): + """``table.chunk(2, dim=-1)`` has row stride 2N; the kernel must index rows by it.""" + + g = torch.Generator(device="cpu").manual_seed(0) + views = torch.randn(2, 16, generator=g).cuda().chunk(2, dim=-1) + copies = tuple(v.contiguous() for v in views) + assert views[0].stride(0) == 16 and copies[0].stride(0) == 8 + out = [ + native_extension.h3_rmsnorm_forward( + native_case["x"], native_case["weight"], 1e-5, shift, scale, native_case["index"] + )[0] + for shift, scale in (views, copies) + ] + assert torch.equal(out[0], out[1]) + + +_BACKWARD_CHILD = """ +import sys +import torch +from rl_engine.backends.extension import _C + +kind, mutation = sys.argv[1:] +x = torch.randn(4, 8, device="cuda") +weight = torch.ones(8, device="cuda") +gate = torch.ones(2, 8, device="cuda") +index = torch.tensor([0, 1, 0, 1], dtype=torch.int64, device="cuda") +i64 = dict(dtype=torch.int64, device="cuda") +sorted_pos = torch.tensor([0, 2, 1, 3], **i64) +begin, end = torch.tensor([0, 2], **i64), torch.tensor([2, 4], **i64) +first = torch.tensor([0, 1, 2], **i64) +if mutation == "tile_end_past_rows": + end = torch.tensor([2, 9], **i64) +elif mutation == "position_past_rows": + sorted_pos = torch.tensor([0, 2, 1, 4], **i64) +elif mutation == "segments_past_tiles": + first = torch.tensor([0, 1, 5], **i64) +elif mutation == "index_past_table": + index = torch.tensor([0, 1, 0, 7], **i64) +elif mutation == "index_negative": + index = torch.tensor([0, 1, 0, -1], **i64) +if kind == "rmsnorm_backward": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index, sorted_pos, + begin, end, first) +elif kind == "gate_backward": + _C.h3_gate_residual_backward(torch.ones_like(x), x, gate, index, sorted_pos, begin, end, first) +elif kind == "gate_partials": # sequence parallel: explicit row list and tiles + _C.h3_gate_grad_partials(torch.ones_like(x), x, sorted_pos, begin, end) +elif kind == "gate_fold": + _C.h3_gate_grad_fold(torch.ones(2, 8, device="cuda"), first, torch.float32) +elif kind == "rmsnorm_partials": + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, None, None, None, + sorted_pos, begin, end, None, None, None) +elif kind == "rmsnorm_fold": + _C.h3_rmsnorm_fold_partials(torch.ones(2, 8, device="cuda"), weight, + torch.ones(2, 16, device="cuda"), first) +elif kind == "rmsnorm_partials_modulated": + # called without the dx entry point having run: the partials check the index values themselves + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device="cuda").chunk(2, dim=1) + _C.h3_rmsnorm_backward_partials(torch.ones_like(x), x, weight, rstd, shift, scale, index, + sorted_pos, begin, end, sorted_pos, begin, end) +torch.cuda.synchronize() +""" + +_METADATA_CASES = { + "rmsnorm_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_backward": ("tile_end_past_rows", "position_past_rows", "segments_past_tiles"), + "gate_partials": ("tile_end_past_rows", "position_past_rows"), + "gate_fold": ("segments_past_tiles",), + "rmsnorm_partials": ("tile_end_past_rows", "position_past_rows"), + "rmsnorm_partials_modulated": ("index_past_table", "index_negative"), + "rmsnorm_fold": ("segments_past_tiles",), +} + + +@pytest.mark.parametrize("kind", sorted(_METADATA_CASES)) +def test_segment_metadata_that_does_not_describe_the_rows_ends_in_a_device_assert( + native_extension, kind +): + """Every entry point that reads caller-built tile lists checks them before it reads.""" + + import os + import subprocess + import sys + + if not hasattr(native_extension, "h3_gate_grad_partials"): + pytest.skip("rl_engine._C lacks the H3 SP entry points") + + def run(mutation): + return subprocess.run( + [sys.executable, "-c", _BACKWARD_CHILD, kind, mutation], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + + control = run("none") + assert control.returncode == 0, control.stderr + for mutation in _METADATA_CASES[kind]: + done = run(mutation) + assert done.returncode != 0, (kind, mutation) + assert "device-side assert triggered" in done.stderr, (kind, mutation, done.stderr) diff --git a/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..b55e4ffcd --- /dev/null +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -0,0 +1,144 @@ +# 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", + "h3_qkv_gemm", + } + + +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_operator_trace.py b/tests/models/minimax_h3/test_h3_operator_trace.py new file mode 100644 index 000000000..042503138 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_operator_trace.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""``h3_operator_trace`` (RFC #420): the per-node ledger of what actually ran. + +CPU tests drive the ledger with synthetic runs over the full fingerprinted graph and +check that every way a run can fall short of "strict" is reported. A CUDA test runs +the one-block replay of block 0 with the trace attached and checks it reports the +bindings ``h3_block`` declares. +""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest + +from rl_engine.models.minimax_h3 import fingerprint as fp +from rl_engine.models.minimax_h3 import operator_trace as ot +from rl_engine.runtime.semantic_registry import ( + OperatorBackendDescriptor, + OperatorRequirements, + SemanticOperatorCatalog, +) + +GRAPH = fp.build_graph() +BLOCK0 = [n.path for n in GRAPH if n.path.startswith("transformer_blocks.0.")] + + +def _run(trace: ot.OperatorTrace, nodes=GRAPH, **overrides) -> ot.OperatorTrace: + for node in nodes: + values = dict(status="row", backend=f"cuda.{node.row}", dtype="bfloat16") + values.update(overrides) + trace.record_node(node.path, **values) + return trace + + +def test_full_graph_trace_with_all_rows_is_strict() -> None: + trace = _run(ot.OperatorTrace()) + assert len(trace.calls) == len(GRAPH) + trace.assert_strict() + assert trace.violations() == [] + assert trace.coverage() == {"wrong_count": [], "unexpected": []} + + +def test_accumulator_follows_the_row_contract() -> None: + trace = _run(ot.OperatorTrace()) + by_path = {call.path: call for call in trace.calls} + assert by_path["transformer_blocks.3.q_proj"].accumulator == "fp32" + assert by_path["transformer_blocks.3.attention"].accumulator == "fp32" + assert by_path["transformer_blocks.3.swiglu"].accumulator == "none" + assert by_path["transformer_blocks.3.rope_q"].accumulator == "none" + assert ot.expected_accumulator("video_input_projection_fp32") == "fp32" + + +def test_record_rejects_unknown_paths_wrong_rows_and_statuses() -> None: + trace = ot.OperatorTrace() + with pytest.raises(ValueError, match="not a node"): + trace.record_node("transformer_blocks.50.attention", status="row", backend="x") + with pytest.raises(ValueError, match="graph assigns"): + trace.record(ot.OperatorCall("transformer_blocks.0.attention", "h3_swiglu", "row", "x")) + with pytest.raises(ValueError, match="status"): + trace.record_node("transformer_blocks.0.attention", status="maybe", backend="x") + + +@pytest.mark.parametrize( + "overrides, needle", + [ + ({"status": "interim"}, "bound as 'interim'"), + ({"status": "reference"}, "bound as 'reference'"), + ({"fallback": True}, "fallback"), + ({"requested_backend": "rlkernel.other"}, "requested"), + ({"split_policy": "split_k"}, "split policy"), + ({"accumulator": "bf16"}, "accumulator"), + ({"backend": ""}, "no backend"), + ], +) +def test_every_strictness_violation_is_reported(overrides, needle) -> None: + trace = ot.OperatorTrace() + _run(trace, nodes=GRAPH) + bad = ot.OperatorTrace() + for call in trace.calls: + bad.record( + replace(call, **overrides) if call.path == "transformer_blocks.7.q_proj" else call + ) + problems = bad.violations() + assert any("transformer_blocks.7.q_proj" in p and needle in p for p in problems), problems + with pytest.raises(ValueError, match="not strict"): + bad.assert_strict() + + +def test_interim_and_reference_can_be_allowed_explicitly() -> None: + trace = ot.OperatorTrace() + for node in GRAPH: + status = "interim" if node.node in ("q_proj", "k_proj") else "row" + trace.record_node(node.path, status=status, backend="b") + assert trace.violations() + assert trace.violations(allowed_statuses=("row", "interim")) == [] + + +def test_coverage_reports_missing_repeated_and_foreign_nodes() -> None: + trace = _run(ot.OperatorTrace(), nodes=[n for n in GRAPH if n.path != "norm_out"]) + trace.record_node("proj_in", status="row", backend="b") # twice now + cov = trace.coverage() + assert any(line.startswith("norm_out: ran 0x") for line in cov["wrong_count"]) + assert any(line.startswith("proj_in: ran 2x") for line in cov["wrong_count"]) + assert trace.coverage(paths=["proj_in"])["unexpected"] != [] # everything else is foreign + + +def test_forwards_expects_each_node_that_many_times() -> None: + trace = ot.OperatorTrace(forwards=2) + _run(trace) + assert trace.violations() != [] + _run(trace) + assert trace.violations() == [] + + +def test_subset_expectation_for_a_one_block_replay() -> None: + trace = _run(ot.OperatorTrace(), nodes=[n for n in GRAPH if n.path in BLOCK0]) + assert trace.violations() != [] # the full graph did not run + assert trace.violations(paths=BLOCK0) == [] + + +def test_digest_is_stable_and_content_sensitive() -> None: + a, b = _run(ot.OperatorTrace()), _run(ot.OperatorTrace()) + assert a.digest == b.digest + c = _run(ot.OperatorTrace(), backend="cuda.other") + assert c.digest != a.digest + d = ot.OperatorTrace(metadata={"seed": 1}) + assert _run(d).digest != a.digest + + +def test_roundtrip_and_tamper_detection() -> None: + trace = _run(ot.OperatorTrace(metadata={"run": "x"})) + again = ot.OperatorTrace.loads(trace.dumps()) + assert again.digest == trace.digest and again.calls == trace.calls + raw = trace.to_dict() + raw["calls"][5]["backend"] = "edited" + with pytest.raises(ValueError, match="trace_sha256"): + ot.OperatorTrace.from_dict(raw) + raw = trace.to_dict() + raw["graph"]["digest"] = "0" * 64 + with pytest.raises(ValueError, match="different H3 graph"): + ot.OperatorTrace.from_dict(raw) + raw = trace.to_dict() + raw["schema_version"] = "v0" + with pytest.raises(ValueError, match="schema"): + ot.OperatorTrace.from_dict(raw) + + +def test_first_difference_names_the_first_diverging_call() -> None: + a, b = _run(ot.OperatorTrace()), ot.OperatorTrace() + for call in a.calls: + b.record(replace(call, backend="other") if call.path.endswith("4.attention") else call) + diff = ot.first_difference(a, b) + assert diff["path"] == "transformer_blocks.4.attention" + assert diff["fields"] == ["backend"] + assert ot.first_difference(a, a) is None + shorter = ot.OperatorTrace() + shorter.calls = a.calls[:-3] + assert ot.first_difference(a, shorter)["fields"] == ["length"] + + +def test_status_summary_counts_per_row() -> None: + trace = ot.OperatorTrace() + for node in GRAPH: + trace.record_node( + node.path, status="row" if node.node != "o_proj" else "interim", backend="b" + ) + summary = trace.status_summary() + assert summary["h3_attention_o_gemm"] == {"interim": 50 + 2} + assert summary["h3_full_attention"] == {"row": 50 + 2} + + +def test_merge_requires_the_same_graph() -> None: + a = ot.OperatorTrace(graph=GRAPH[:20]) + b = ot.OperatorTrace() + with pytest.raises(ValueError, match="different graphs"): + a.merge(b) + + +@pytest.mark.parametrize( + "name, block, expected", + [ + ("block.7.attention", None, "transformer_blocks.7.attention"), + ("refiner.1.q_proj", None, "token_refiner.refiner_blocks.1.q_proj"), + ("attention", 3, "transformer_blocks.3.attention"), + ("norm_out", None, "norm_out"), + ], +) +def test_graph_path_normalisation(name, block, expected) -> None: + assert ot.graph_path(name, block_index=block) == expected + + +# --- registry bridge ---------------------------------------------------------------------------- + + +class _FakeNorm: + def __call__(self, x): + return x + + +def _descriptor(**overrides) -> OperatorBackendDescriptor: + values = dict( + semantic_op="h3_rmsnorm", + backend_id="rlkernel.test_h3_rmsnorm", + supported_targets=frozenset({"rollout", "training"}), + supported_devices=frozenset({"cpu"}), + supported_dtypes=frozenset({"bfloat16"}), + supported_topologies={"*": "*"}, + determinism_or_alignment_properties={"strict_observable": True}, + lifecycle="request", + implementation_class_or_factory=_FakeNorm, + fallback_policy="error", + version_or_build_fingerprint="test-1", + is_smoke_only=True, + ) + values.update(overrides) + return OperatorBackendDescriptor(**values) + + +def _resolve(descriptor: OperatorBackendDescriptor, **topology): + from rl_engine.runtime.semantic_registry import OperatorResolutionPolicy + + catalog = SemanticOperatorCatalog([descriptor]) + session = catalog.session(OperatorResolutionPolicy(strict=True, allow_test_backends=True)) + resolution = session.resolve( + semantic_op=descriptor.semantic_op, + requested_backend=descriptor.backend_id, + target="training", + requirements=OperatorRequirements(device="cpu", dtype="bfloat16", topology=topology), + ) + instance = session.instantiate(resolution) + return resolution, session.instance_provenance(resolution, instance) + + +def test_call_from_a_resolution_carries_registry_provenance() -> None: + resolution, provenance = _resolve(_descriptor()) + call = ot.trace_call_from_resolution( + "transformer_blocks.0.norm1", resolution=resolution, provenance=provenance, shapes=[(1, 8)] + ) + assert call.row == "h3_rmsnorm" and call.backend == "rlkernel.test_h3_rmsnorm" + assert call.requested_backend == call.backend + assert call.instance_fingerprint == provenance.instance_fingerprint != "" + assert call.implementation.endswith("_FakeNorm") + assert call.fallback is False and call.shapes == ((1, 8),) + assert call.accumulator == "fp32" + trace = ot.OperatorTrace() + trace.record(call) + assert trace.violations(paths=["transformer_blocks.0.norm1"]) == [] + + +def test_call_from_a_resolution_keeps_topology_and_rejects_unknown_paths() -> None: + resolution, provenance = _resolve(_descriptor(), tp=4, cp=2) + call = ot.trace_call_from_resolution( + "transformer_blocks.0.norm1", resolution=resolution, provenance=provenance + ) + assert dict(call.topology) == {"tp": 4, "cp": 2} + with pytest.raises(ValueError, match="not a node"): + ot.trace_call_from_resolution("nope", resolution=resolution, provenance=provenance) + + +def test_unresolved_resolution_cannot_become_a_call() -> None: + resolution, _ = _resolve(_descriptor()) + broken = replace(resolution, trace=replace(resolution.trace, status="unsupported")) + with pytest.raises(ValueError, match="nothing ran"): + ot.trace_call_from_resolution("transformer_blocks.0.norm1", resolution=broken) + + +# --- the real one-block replay (CUDA) ---------------------------------------------------------- + + +def test_block_replay_trace_matches_h3_block_bindings(h3_block_weights_cuda) -> None: + torch = pytest.importorskip("torch") + from rl_engine.validation.models import h3_block + + ops = h3_block.CandidateOps() + ctx = h3_block.make_block_context(h3_block_weights_cuda, layout="tiny", seed=0, batch=1) + trace = ot.OperatorTrace(mode="candidate", metadata={"layout": "tiny", "block": 0}) + ctx["node_observer"] = ot.block_observer(trace, block_index=0) + with torch.no_grad(): + h3_block.run_graph("candidate", ctx, ops) + + assert [c.path for c in trace.calls] == BLOCK0 + assert trace.coverage(paths=BLOCK0) == {"wrong_count": [], "unexpected": []} + bindings = {n.name: n.binding for n in h3_block.NODES} + for call, node in zip(trace.calls, h3_block.NODES): + assert (call.status, call.backend, call.row) == ( + bindings[node.name].status, + bindings[node.name].backend, + bindings[node.name].rfc_row, + ) + # Today's stack still has interim and reference nodes: the strict gate must say so. + assert trace.violations(paths=BLOCK0) + assert trace.violations(paths=BLOCK0, allowed_statuses=("row", "interim", "reference")) == [] + # Replaying twice gives the same ledger: the trace is part of the repeatable surface. + again = ot.OperatorTrace(mode="candidate", metadata={"layout": "tiny", "block": 0}) + ctx["node_observer"] = ot.block_observer(again, block_index=0) + with torch.no_grad(): + h3_block.run_graph("candidate", ctx, ops) + assert again.digest == trace.digest 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..120729257 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_projection_rows.py @@ -0,0 +1,552 @@ +# 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) + + +def _splittable(row): + spec, _ = ROWS[row] + return spec.out_features % (2 * spec.shard_unit) == 0 + + +@requires_cuda +@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] + 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", [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): + 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 + + +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)) diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py index b79973eb1..ed74d1901 100644 --- a/tests/models/minimax_h3/test_h3_report.py +++ b/tests/models/minimax_h3/test_h3_report.py @@ -1,13 +1,190 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors -"""Comparative timing balances backend order and preserves its raw evidence.""" +"""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 -from rl_engine.validation.models import h3_report + +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): diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py index 309775562..c0de05aa1 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -229,6 +229,56 @@ def test_rejects_a_row_wider_than_the_int_the_kernels_count_with(self): _C.h3_rmsnorm_backward(x, x, w, torch.empty(1, device="cuda")) +_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): diff --git a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py index dda99680f..e97d7f6c8 100644 --- a/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py +++ b/tests/models/minimax_h3/test_h3_rmsnorm_native_validation.py @@ -96,7 +96,7 @@ def test_valid_plain_and_modulated_backward(native, inputs): @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 nonempty"): + with pytest.raises(RuntimeError, match="row index must be a non-empty"): _call(native, inputs, entrypoint) 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..b3213fd1e --- /dev/null +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,400 @@ +# 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]) + + @pytest.mark.parametrize("sp, batch, seq", [(4, 2, 777), (3, 2, 1000), (8, 1, 40)]) + def test_padding_of_the_exchange_and_partial_gathers_is_never_read( + self, monkeypatch, sp, batch, seq + ): + """Pad every send block and every partial gather with NaN: WS1 equality still holds. + + ``recv_global`` is zero in the padding slots, which map to a real table row, so the + padding is only harmless because no tile list refers to it; a NaN would reach the sums + of any tile that did. + """ + + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3 import sp_norm_adaln as module + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + + hidden = 256 + c = sp_case(batch, seq, hidden, layout="interleaved", dtype=torch.float32, seed=sp + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + nan = float("nan") + gather_rows = module.gather_rows + + def exchange(self, collective, *tensors): + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_full((self.max_send, *t.shape[1:]), nan) # padding is NaN + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family, partial): + fam = self.families[family] + padded = partial.new_full((fam["width"], *partial.shape[1:]), nan) # padding is NaN + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + monkeypatch.setattr(module._Plan, "exchange", exchange) + monkeypatch.setattr(module._Plan, "gather_partials", gather_partials) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = module.sp_plan(lay, idx, rows) + norm = module.sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + return norm, module.sp_gate_residual_backward( + local["grad"], local["y"], gate, plan, coll + ) + + for (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + assert torch.isfinite(dx).all() and torch.isfinite(dy).all() + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dsh.dtype)) + assert torch.equal(dsc, ws1_norm[3].to(dsc.dtype)) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + H3SPNormAdaLNCudaOp, + ) + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@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_token_refiner.py b/tests/models/minimax_h3/test_h3_token_refiner.py new file mode 100644 index 000000000..5e523de1c --- /dev/null +++ b/tests/models/minimax_h3/test_h3_token_refiner.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``token_refiner_block``: the two text-refiner blocks and their final norm. + +CPU tests pin the graph to the architecture fingerprint. CUDA tests replay both refiner +blocks of the pinned checkpoint (shard 14) node by node and check: + +* every node is bitwise repeatable and batch-row invariant; +* nodes that promise it equal diffusers on diffusers' inputs; +* every node, the block output and the whole refiner stack are about as close to the FP32 + golden as diffusers itself; gradients reaching the input and every parameter likewise; +* the operator trace of a replay equals the declared bindings and covers the fingerprint's + refiner paths. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.models.minimax_h3 import fingerprint as fp +from rl_engine.models.minimax_h3 import operator_trace as ot +from rl_engine.validation.models import h3_block as hb +from rl_engine.validation.models import h3_refiner as hr +from rl_engine.validation.models.h3_weights import load_h3_tensors + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) +FORWARD_FACTOR, FORWARD_FLOOR = 1.25, 1e-3 +BACKWARD_FACTOR, BACKWARD_FLOOR = 1.5, 5e-3 + + +def test_graph_equals_the_fingerprint_vocabulary(): + assert hr.graph_matches_fingerprint() + assert hr.NODE_NAMES == tuple(name for name, _ in fp.REFINER_NODES) + assert len(hr.NODES) == 14 + + +def test_graph_is_well_formed_and_owners_are_declared(): + seen = {"hidden"} + for node in hr.NODES: + assert set(node.inputs) <= seen, node.name + assert node.name not in seen, node.name + assert node.binding.status in ("row", "interim", "reference") + seen.add(node.name) + assert hr.OUTPUT == "residual_mlp" + own = {n.name for n in hr.NODES if n.binding.rfc_row == "token_refiner_block"} + assert own == {"residual_attn", "residual_mlp"} # the row owns only the plain adds + rows_with_kernels = {n.binding.rfc_row for n in hr.NODES if n.binding.status == "row"} + assert rows_with_kernels == {"h3_rmsnorm", "h3_qkv_gemm", "token_refiner_block"} + + +def test_refiner_has_no_adaln_and_no_rope_nodes(): + assert not any("adaln" in n.binding.rfc_row or "rope" in n.binding.rfc_row for n in hr.NODES) + + +@pytest.fixture(scope="module") +def refiner_weights(): + names = [f"{hr.PREFIX}.{b}.{leaf}" for b in (0, 1) for leaf in hr.PARAM_LEAVES] + names.append(hr.FINAL_NORM) + try: + weights = load_h3_tensors(names, "cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + table = fp.expected_tensor_table(fp.transformer_config()) + for name, tensor in weights.items(): + dtype, shape = table[name] + assert tuple(tensor.shape) == shape and dtype == "BF16", name + return weights + + +@needs_gpu +@pytest.mark.parametrize("tokens", [1, 64, 257]) +@pytest.mark.parametrize("block", [0, 1]) +def test_nodes_accuracy_repeat_batch_and_isolated(refiner_weights, block, tokens): + report = hr.run_refiner_case(refiner_weights, block=block, tokens=tokens, seed=3) + assert report["first_repeat_drift"] is None + assert report["first_batch_drift"] is None + for entry in report["nodes"]: + name = entry["node"] + assert entry["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * entry["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ), name + if entry["provider_bitwise_isolated_promised"]: + assert entry["isolated_vs_provider"]["bitwise_equal"], name + out = report["output"] + assert out["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * out["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ) + + +@needs_gpu +def test_report_names_rows_status_and_backend_per_node(refiner_weights): + report = hr.run_refiner_case(refiner_weights, block=0, tokens=16) + assert [e["node"] for e in report["nodes"]] == list(hr.NODE_NAMES) + by = {e["node"]: e for e in report["nodes"]} + assert by["q_proj"]["status"] == "row" and by["q_proj"]["rfc_row"] == "h3_qkv_gemm" + assert by["attention"]["status"] == "interim" + assert by["residual_attn"]["rfc_row"] == "token_refiner_block" + assert all(len(e["candidate_sha256"]) == 64 for e in report["nodes"]) + + +@needs_gpu +def test_full_refiner_stack_against_golden_and_provider(refiner_weights): + ops = hb.CandidateOps() + gen = torch.Generator("cpu").manual_seed(11) + hidden = torch.randn(1, 130, hr.HIDDEN, generator=gen).bfloat16().cuda() + with torch.no_grad(): + cand = hr.run_refiner("candidate", refiner_weights, hidden, ops) + again = hr.run_refiner("candidate", refiner_weights, hidden, ops) + provider = hr.run_refiner("provider", refiner_weights, hidden) + golden = hr.run_refiner("golden", refiner_weights, hidden) + assert torch.equal(cand, again) + assert cand.dtype == torch.bfloat16 and cand.shape == hidden.shape + assert hb.rel_err(cand, golden) <= FORWARD_FACTOR * hb.rel_err(provider, golden) + FORWARD_FLOOR + + +@needs_gpu +def test_refiner_rows_do_not_depend_on_batch_neighbours(refiner_weights): + ops = hb.CandidateOps() + gen = torch.Generator("cpu").manual_seed(5) + hidden = torch.randn(2, 65, hr.HIDDEN, generator=gen).bfloat16().cuda() + with torch.no_grad(): + both = hr.run_refiner("candidate", refiner_weights, hidden, ops) + first = hr.run_refiner("candidate", refiner_weights, hidden[:1].contiguous(), ops) + assert torch.equal(both[:1], first) + + +@needs_gpu +def test_backward_against_fp32_golden(refiner_weights): + ops = hb.CandidateOps() + block = 0 + leaves = [f"{hr.PREFIX}.{block}.{leaf}" for leaf in hr.PARAM_LEAVES] + gen = torch.Generator("cpu").manual_seed(2) + hidden0 = torch.randn(1, 33, hr.HIDDEN, generator=gen).bfloat16().cuda() + grad_out = torch.randn(1, 33, hr.HIDDEN, generator=gen).bfloat16().cuda() + + def grads(mode, dtype): + weights = { + n: refiner_weights[n].detach().to(dtype).clone().requires_grad_() for n in leaves + } + hidden = hidden0.detach().to(dtype).clone().requires_grad_() + ctx = {"weights": weights, "prefix": f"{hr.PREFIX}.{block}", "hidden": hidden} + out = hr.run_graph(mode, ctx, ops)[hr.OUTPUT] + out.backward(grad_out.to(out.dtype)) + return {"hidden": hidden.grad, **{n: weights[n].grad for n in leaves}} + + cand = grads("candidate", torch.bfloat16) + again = grads("candidate", torch.bfloat16) + provider = grads("provider", torch.bfloat16) + golden = grads("golden", torch.float32) + for name in cand: + assert torch.equal(cand[name], again[name]), f"{name}: backward is not repeatable" + c_err = hb.rel_err(cand[name], golden[name]) + p_err = hb.rel_err(provider[name], golden[name]) + assert c_err <= BACKWARD_FACTOR * p_err + BACKWARD_FLOOR, (name, c_err, p_err) + + +@needs_gpu +def test_operator_trace_covers_the_refiner_paths(refiner_weights): + ops = hb.CandidateOps() + gen = torch.Generator("cpu").manual_seed(1) + hidden = torch.randn(1, 17, hr.HIDDEN, generator=gen).bfloat16().cuda() + trace = ot.OperatorTrace(metadata={"stage": "token_refiner"}) + with torch.no_grad(): + hr.run_refiner( + "candidate", + refiner_weights, + hidden, + ops, + observer_factory=lambda b: ot.block_observer( + trace, block_index=b, stage="token_refiner" + ), + ) + paths = [ + n.path + for n in fp.build_graph() + if n.path.startswith("token_refiner.refiner_blocks.") and n.stage == "token_refiner" + ] + assert [c.path for c in trace.calls] == paths + assert trace.coverage(paths=paths) == {"wrong_count": [], "unexpected": []} + bound = {n.name: n.binding for n in hr.NODES} + for call in trace.calls: + node = call.path.rsplit(".", 1)[1] + assert (call.status, call.row) == (bound[node].status, bound[node].rfc_row) + assert trace.violations(paths=paths) # interim nodes still present + assert trace.violations(paths=paths, allowed_statuses=("row", "interim")) == [] + + +@needs_gpu +def test_text_path_context_projection_then_refiner(refiner_weights): + """context_embedder (row) feeding the refiner, against the provider chain.""" + + names = ["context_embedder.weight", "context_embedder.bias"] + ctxw = load_h3_tensors(names, "cuda") + from rl_engine.runtime.registry import KernelRegistry + + proj = KernelRegistry().get_op("text_context_projection", device="cuda") + gen = torch.Generator("cpu").manual_seed(9) + text = torch.randn(70, 5120, generator=gen).bfloat16().cuda() + ops = hb.CandidateOps() + with torch.no_grad(): + h_cand = proj.forward(text, ctxw[names[0]], ctxw[names[1]]).unsqueeze(0) + out = hr.run_refiner("candidate", refiner_weights, h_cand, ops) + h_prov = torch.nn.functional.linear(text, ctxw[names[0]], ctxw[names[1]]).unsqueeze(0) + provider = hr.run_refiner("provider", refiner_weights, h_prov) + golden = hr.run_refiner("golden", refiner_weights, h_prov) + assert hb.rel_err(out, golden) <= 2 * hb.rel_err(provider, golden) + 2e-3 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..0fe5b26d8 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -0,0 +1,231 @@ +# 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] + + 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("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") +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_trajectory_artifact.py b/tests/models/minimax_h3/test_h3_trajectory_artifact.py new file mode 100644 index 000000000..0d3de9b41 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_trajectory_artifact.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""``h3_trajectory_artifact`` (RFC #420): sealed trajectories and where two of them drift. + +CPU only. Synthetic Euler-like trajectories stand in for a model: what is under test +is the record (digests, hash chain, seal), its fail-closed validation and the drift locator. +""" + +from __future__ import annotations + +import copy +import json +from pathlib import Path + +import pytest +import torch + +from rl_engine.models.minimax_h3 import trajectory as tj + +STEPS = 6 +SHAPES = {"video": (1, 4, 3, 8), "audio": (1, 5, 2)} + + +def _grid(steps: int, shift: float) -> list[float]: + sigmas = [1.0 - i / steps for i in range(steps + 1)] + return [shift * s / (1 + (shift - 1) * s) for s in sigmas] + + +def _noise(seed: int) -> dict[str, torch.Tensor]: + gen = torch.Generator().manual_seed(seed) + return {m: torch.randn(shape, generator=gen).bfloat16() for m, shape in SHAPES.items()} + + +def _run( + label: str = "strict", + *, + seed: int = 7, + perturb: dict | None = None, + keep_tensors: bool = False, + operator_trace: str = "", +) -> tj.TrajectoryRecorder: + """A deterministic fake trajectory; ``perturb`` = {(step, modality): delta} on velocity.""" + + noise = _noise(seed) + recorder = tj.TrajectoryRecorder( + variant="FL2VA", + label=label, + seed=seed, + sigma_grids={"video": _grid(STEPS, 12.0), "audio": _grid(STEPS, 3.0)}, + initial_noise=noise, + shifts={"video": 12.0, "audio": 3.0}, + operator_trace_sha256=operator_trace, + metadata={"note": "synthetic"}, + keep_tensors=keep_tensors, + ) + latent = {m: t.clone() for m, t in noise.items()} + gen = torch.Generator().manual_seed(1234) + for step in range(STEPS): + velocity, new_latent = {}, {} + for m in recorder.modalities: + v = torch.randn(SHAPES[m], generator=gen).bfloat16() + delta = (perturb or {}).get((step, m)) + if delta is not None: + v = v.clone() + v.view(-1)[0] += delta + velocity[m] = v + grid = _grid(STEPS, 12.0 if m == "video" else 3.0) + sigma, sigma_next = grid[step], grid[step + 1] + r = sigma_next / sigma + x0 = latent[m].float() + sigma * v.float() # data-ward velocity + new_latent[m] = (r * latent[m].float() + (1 - r) * x0).bfloat16() + latent = new_latent + recorder.record_step(step, velocity=velocity, latent=latent) + return recorder + + +def test_artifact_is_deterministic_and_verifies() -> None: + a, b = _run().seal(), _run().seal() + assert a == b + tj.verify(a) + assert a["steps"] == STEPS and a["modalities"] == ["video", "audio"] + assert a["final_latent_sha256"] == a["step_records"][-1]["latent_sha256"] + + +def test_label_or_trace_changes_the_seal_but_not_the_trajectory() -> None: + a = _run("native").seal() + b = _run("strict", operator_trace="ab" * 32).seal() + assert a["artifact_sha256"] != b["artifact_sha256"] + assert a["step_records"] == b["step_records"] + + +def test_json_roundtrip() -> None: + artifact = _run().seal() + assert tj.loads(tj.dumps(artifact)) == artifact + + +def test_sigma_grids_are_exact_fp32() -> None: + artifact = _run().seal() + for modality in tj.modalities_of(artifact): + decoded = tj.sigmas_of(artifact, modality) + assert decoded == [tj.f32_from_hex(tj.f32_hex(s)) for s in decoded] + assert all(torch.tensor(s, dtype=torch.float32).item() == s for s in decoded) + assert tj.f32_hex(0.1) != (0.1).hex() # FP32 rounding is applied, not FP64 text + + +def test_tensor_digest_covers_bytes_dtype_and_shape() -> None: + x = torch.arange(12, dtype=torch.float32) + assert tj.tensor_digest(x) == tj.tensor_digest(x.clone()) + assert tj.tensor_digest(x) != tj.tensor_digest(x.reshape(3, 4)) + assert tj.tensor_digest(x) != tj.tensor_digest(x.to(torch.float64)) + y = x.clone() + y[5] = torch.nextafter(y[5], torch.tensor(1e9)) + assert tj.tensor_digest(x) != tj.tensor_digest(y) + assert tj.tensor_digest(x.reshape(3, 4).t()) == tj.tensor_digest( + x.reshape(3, 4).t().contiguous() + ) + assert tj.tensor_digest(torch.empty(0)) != tj.tensor_digest( + torch.empty(0, dtype=torch.bfloat16) + ) + + +# --- recorder contract ------------------------------------------------------------------------ + + +def test_recorder_enforces_step_order_and_completeness() -> None: + recorder = _run() + with pytest.raises(ValueError, match="only"): + recorder.record_step(STEPS, velocity=_noise(0), latent=_noise(0)) + partial = tj.TrajectoryRecorder( + variant="FL2VA", + label="x", + seed=0, + sigma_grids={"video": _grid(2, 12.0)}, + initial_noise={"video": _noise(0)["video"]}, + ) + with pytest.raises(ValueError, match="expected step 0"): + partial.record_step( + 1, velocity={"video": _noise(0)["video"]}, latent={"video": _noise(0)["video"]} + ) + with pytest.raises(ValueError, match="recorded 0 of 2"): + partial.seal() + with pytest.raises(ValueError, match="exactly"): + partial.record_step(0, velocity=_noise(0), latent={"video": _noise(0)["video"]}) + + +@pytest.mark.parametrize( + "kwargs, match", + [ + (dict(variant="Nope"), "unknown variant"), + (dict(sigma_grids={"video": [1.0]}), "length"), + ( + dict( + sigma_grids={"video": [1.0, 0.5], "audio": [1.0, 0.7, 0.0]}, + noise=("video", "audio"), + ), + "length", + ), + (dict(noise=("video", "audio")), "modalities"), + ], +) +def test_recorder_rejects_bad_construction(kwargs, match) -> None: + noise = _noise(0) + args = dict( + variant="FL2VA", + label="x", + seed=0, + sigma_grids={"video": _grid(2, 12.0)}, + initial_noise={"video": noise["video"]}, + ) + names = kwargs.pop("noise", None) + if names: + args["initial_noise"] = {m: noise[m] for m in names} + args.update(kwargs) + with pytest.raises(ValueError, match=match): + tj.TrajectoryRecorder(**args) + + +# --- fail closed --------------------------------------------------------------------------------- + + +def _mutations(): + def edit_digest(a): + a["step_records"][2]["velocity_sha256"]["video"] = "0" * 64 + + def drop_step(a): + a["step_records"].pop() + + def swap_steps(a): + a["step_records"][1], a["step_records"][2] = a["step_records"][2], a["step_records"][1] + + def edit_grid(a): + a["sigma_grids"]["video"][3] = tj.f32_hex(0.123) + + def edit_seed(a): + a["seed"] = 8 + + def edit_noise(a): + a["initial_noise"]["audio"]["sha256"] = "1" * 64 + + def edit_label(a): + a["label"] = "native" + + return [edit_digest, drop_step, swap_steps, edit_grid, edit_seed, edit_noise, edit_label] + + +@pytest.mark.parametrize("mutate", _mutations(), ids=lambda f: f.__name__) +def test_any_edit_fails_verification(mutate) -> None: + artifact = copy.deepcopy(_run().seal()) + mutate(artifact) + with pytest.raises(ValueError): + tj.verify(artifact) + + +def test_resealed_edit_is_still_caught_by_the_chain_or_grid() -> None: + artifact = copy.deepcopy(_run().seal()) + artifact["step_records"][2]["velocity_sha256"]["video"] = "0" * 64 + artifact["artifact_sha256"] = tj._sha( + {k: v for k, v in artifact.items() if k != "artifact_sha256"} + ) + with pytest.raises(ValueError, match="hash chain"): + tj.verify(artifact) + + +def test_non_monotone_grid_and_wrong_fingerprint_are_rejected() -> None: + artifact = copy.deepcopy(_run().seal()) + artifact["architecture_fingerprint_sha256"] = "f" * 64 + artifact["artifact_sha256"] = tj._sha( + {k: v for k, v in artifact.items() if k != "artifact_sha256"} + ) + with pytest.raises(ValueError, match="fingerprint"): + tj.verify(artifact) + tj.verify(artifact, check_fingerprint=False) + recorder = tj.TrajectoryRecorder( + variant="FL2VA", + label="x", + seed=0, + sigma_grids={"video": [1.0, 0.2, 0.5]}, + initial_noise={"video": _noise(0)["video"]}, + ) + for step in range(2): + recorder.record_step( + step, velocity={"video": _noise(0)["video"]}, latent={"video": _noise(0)["video"]} + ) + with pytest.raises(ValueError, match="non-increasing"): + tj.verify(recorder.seal()) + + +def test_unsupported_schema() -> None: + artifact = _run().seal() + artifact["schema_version"] = "v0" + with pytest.raises(ValueError, match="schema"): + tj.verify(artifact) + + +# --- drift -------------------------------------------------------------------------------------- + + +def test_identical_trajectories_have_no_drift() -> None: + a, b = _run("strict").seal(), _run("strict").seal() + assert tj.first_drift(a, b) is None + report = tj.drift_report(a, b) + assert report["bitwise_equal"] and report["differing"] == 0 + assert report["compared"] == STEPS * 2 * 2 and report["final_latent_equal"] + + +def test_first_drift_names_step_kind_and_modality() -> None: + a = _run("strict").seal() + b = _run("native", perturb={(2, "audio"): 0.25}).seal() + drift = tj.first_drift(a, b) + assert drift == {"step": 2, "kind": "velocity", "modality": "audio"} + report = tj.drift_report(a, b) + assert not report["bitwise_equal"] and not report["final_latent_equal"] + # The perturbed velocity changes its own latent, which every later latent inherits; + # the other modality and the other velocities are untouched. + assert [(e["step"], e["kind"], e["modality"]) for e in report["entries"]] == [ + (2, "velocity", "audio"), + (2, "latent", "audio"), + (3, "latent", "audio"), + (4, "latent", "audio"), + (5, "latent", "audio"), + ] + assert report["first_drift"] == drift + + +def test_velocity_is_reported_before_the_latent_of_the_same_step() -> None: + a = _run().seal() + b = _run(perturb={(0, "video"): 1.0}).seal() + assert tj.first_drift(a, b)["kind"] == "velocity" + + +def test_drift_with_tensors_reports_count_and_magnitude(tmp_path: Path) -> None: + ra, rb = _run(keep_tensors=True), _run("native", perturb={(3, "video"): 0.5}, keep_tensors=True) + ra.save_tensors(tmp_path / "a.safetensors") + rb.save_tensors(tmp_path / "b.safetensors") + drift = tj.first_drift( + ra.seal(), + rb.seal(), + tensors_lhs=tmp_path / "a.safetensors", + tensors_rhs=tmp_path / "b.safetensors", + ) + assert (drift["step"], drift["kind"], drift["modality"]) == (3, "velocity", "video") + assert drift["mismatch_count"] == 1 + assert drift["numel"] == 1 * 4 * 3 * 8 + assert drift["max_abs_diff"] > 0.0 + + +def test_different_inputs_are_not_comparable() -> None: + a, b = _run(seed=7).seal(), _run(seed=8).seal() + with pytest.raises(ValueError, match="not comparable"): + tj.first_drift(a, b) + assert tj.first_drift(a, b, require_same_inputs=False) is not None + + +def test_save_tensors_needs_keep_tensors(tmp_path: Path) -> None: + with pytest.raises(ValueError, match="keep_tensors"): + _run().save_tensors(tmp_path / "x.safetensors") + + +def test_artifact_json_is_plain_text() -> None: + text = tj.dumps(_run().seal()) + parsed = json.loads(text) + assert parsed["schema_version"] == tj.SCHEMA_VERSION + assert all(isinstance(v, str) for v in parsed["sigma_grids"]["video"]) 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/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 47c423aa6..c08e5fb9b 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -238,6 +238,8 @@ def fake_load_backend(backend): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", ) diff --git a/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index 6c04efadb..e7d1e5275 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/test_ws1_gtest_gpu.py @@ -56,6 +56,8 @@ def test_all_ws1_single_ops_are_registered(): "adaln_projection_3mod", "adaln_row_gather", "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", } <= names diff --git a/tools/migration/layout.json b/tools/migration/layout.json index 5b1fd4caa..e763520f9 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -467,15 +467,26 @@ "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", @@ -484,10 +495,14 @@ "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", @@ -702,9 +717,14 @@ "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", @@ -738,7 +758,9 @@ "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", @@ -819,6 +841,7 @@ "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", @@ -837,6 +860,7 @@ "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_bitwise_reuse.py": "tools/validation/models/plot_h3_bitwise_reuse.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", @@ -872,6 +896,7 @@ "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", @@ -880,12 +905,16 @@ "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_rmsnorm_native_validation.py": "tests/models/minimax_h3/test_h3_rmsnorm_native_validation.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", 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_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") 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_fingerprint.py b/tools/validation/models/h3_fingerprint.py new file mode 100644 index 000000000..5715697e8 --- /dev/null +++ b/tools/validation/models/h3_fingerprint.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Print, write or check the MiniMax-H3 architecture fingerprint. + + python tools/validation/models/h3_fingerprint.py print + python tools/validation/models/h3_fingerprint.py write # refresh the golden file + python tools/validation/models/h3_fingerprint.py check # golden == recomputed + python tools/validation/models/h3_fingerprint.py checkpoint DIR # local shard headers + +``check`` exits non-zero on any difference, so it is usable as a CI gate. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.models.minimax_h3 import fingerprint as fp # noqa: E402 + +GOLDEN = REPO_ROOT / "rl_engine/models/minimax_h3/arch_fingerprint.json" + + +def _dump(fingerprint: dict) -> str: + return json.dumps(fingerprint, indent=1, sort_keys=True) + "\n" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + sub.add_parser("print") + sub.add_parser("write") + sub.add_parser("check") + checkpoint = sub.add_parser("checkpoint") + checkpoint.add_argument("directory", type=Path) + checkpoint.add_argument("--variant", default="FL2VA", choices=sorted(fp.VARIANTS)) + args = parser.parse_args() + + if args.command == "checkpoint": + count = fp.validate_checkpoint_headers(args.directory, args.variant) + print(f"ok: {count} tensors in {args.directory} match the {args.variant} fingerprint") + return 0 + fingerprint = fp.build_fingerprint() + if args.command == "print": + sys.stdout.write(_dump(fingerprint)) + elif args.command == "write": + GOLDEN.write_text(_dump(fingerprint)) + print(f"wrote {GOLDEN} ({fingerprint['fingerprint_sha256']})") + else: + if not GOLDEN.exists() or json.loads(GOLDEN.read_text()) != fingerprint: + print("golden fingerprint differs from the recomputed one", file=sys.stderr) + return 1 + print(f"ok: {fingerprint['fingerprint_sha256']}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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_operator_trace.py b/tools/validation/models/h3_operator_trace.py new file mode 100644 index 000000000..34fbfa1b7 --- /dev/null +++ b/tools/validation/models/h3_operator_trace.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Record, verify and compare MiniMax-H3 operator traces. + + # one-block replay of block 0 with the trace attached (needs CUDA + RL_KERNEL_H3_WEIGHTS) + python tools/validation/models/h3_operator_trace.py block --layout tiny --out trace.json + + # fail-closed gate: exit 1 and list every violation + python tools/validation/models/h3_operator_trace.py verify trace.json [--allow row interim] + python tools/validation/models/h3_operator_trace.py verify trace.json --block 0 + + # first call at which two traces disagree + python tools/validation/models/h3_operator_trace.py diff a.json b.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.models.minimax_h3 import fingerprint as fp # noqa: E402 +from rl_engine.models.minimax_h3 import operator_trace as ot # noqa: E402 + + +def _record_block(layout: str, seed: int) -> ot.OperatorTrace: + import torch + + from rl_engine.validation.models import h3_block, h3_weights + + weights = h3_weights.load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + ctx = h3_block.make_block_context(weights, layout=layout, seed=seed, batch=1) + trace = ot.OperatorTrace( + mode="candidate", metadata={"layout": layout, "seed": seed, "block": 0} + ) + ctx["node_observer"] = ot.block_observer(trace, block_index=0) + with torch.no_grad(): + h3_block.run_graph("candidate", ctx, ops) + return trace + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + block = sub.add_parser("block") + block.add_argument("--layout", default="tiny") + block.add_argument("--seed", type=int, default=0) + block.add_argument("--out", type=Path, required=True) + verify = sub.add_parser("verify") + verify.add_argument("trace", type=Path) + verify.add_argument("--allow", nargs="*", default=["row"], choices=ot.STATUSES) + verify.add_argument( + "--block", type=int, default=None, help="expect only this transformer block" + ) + diff = sub.add_parser("diff") + diff.add_argument("lhs", type=Path) + diff.add_argument("rhs", type=Path) + args = parser.parse_args() + + if args.command == "block": + trace = _record_block(args.layout, args.seed) + args.out.write_text(trace.dumps()) + print(f"wrote {args.out} ({trace.digest})") + print(json.dumps(trace.status_summary(), indent=1)) + return 0 + if args.command == "verify": + trace = ot.OperatorTrace.loads(args.trace.read_text()) + paths = None + if args.block is not None: + prefix = f"transformer_blocks.{args.block}." + paths = [node.path for node in fp.build_graph() if node.path.startswith(prefix)] + problems = trace.violations(paths=paths, allowed_statuses=args.allow) + if problems: + print(f"{len(problems)} violation(s):", file=sys.stderr) + for problem in problems[:60]: + print(" " + problem, file=sys.stderr) + return 1 + print(f"ok: {len(trace.calls)} calls, digest {trace.digest}") + return 0 + lhs = ot.OperatorTrace.loads(args.lhs.read_text()) + rhs = ot.OperatorTrace.loads(args.rhs.read_text()) + difference = ot.first_difference(lhs, rhs) + if difference is None: + print("identical") + return 0 + print(json.dumps(difference, indent=1)) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/validation/models/h3_projection_bench.py b/tools/validation/models/h3_projection_bench.py new file mode 100644 index 000000000..a0cbef075 --- /dev/null +++ b/tools/validation/models/h3_projection_bench.py @@ -0,0 +1,108 @@ +#!/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) + 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, + "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, + "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}" + ) + 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() diff --git a/tools/validation/models/h3_projection_charts.py b/tools/validation/models/h3_projection_charts.py new file mode 100644 index 000000000..a47a0d8d0 --- /dev/null +++ b/tools/validation/models/h3_projection_charts.py @@ -0,0 +1,58 @@ +# 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])) diff --git a/tools/validation/models/h3_refiner_report.py b/tools/validation/models/h3_refiner_report.py new file mode 100644 index 000000000..8f10bbaed --- /dev/null +++ b/tools/validation/models/h3_refiner_report.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Accuracy and timing evidence (JSON + PNG) of the H3 token refiner on real weights (H100).""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time +from pathlib import Path + +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.validation.models import h3_block as hb # noqa: E402 +from rl_engine.validation.models import h3_refiner as hr # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_tensors # noqa: E402 + + +def _median_ms(fn, repeats=5): + fn() + torch.cuda.synchronize() + out = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + out.append((time.perf_counter() - start) * 1e3) + return statistics.median(out) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tokens", type=int, nargs="*", default=[12, 257, 925]) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--png", type=Path) + args = parser.parse_args() + names = [f"{hr.PREFIX}.0.{leaf}" for leaf in hr.PARAM_LEAVES] + [hr.FINAL_NORM] + weights = load_h3_tensors(names, "cuda") + ops = hb.CandidateOps() + report = {"device": torch.cuda.get_device_name(0), "cases": []} + for tokens in args.tokens: + case = hr.run_refiner_case(weights, block=0, tokens=tokens, seed=3, ops=ops) + ctx = hr.make_refiner_context(weights, block=0, tokens=tokens, seed=3) + with torch.no_grad(): + cand = _median_ms(lambda: hr.run_graph("candidate", ctx, ops)) + prov = _median_ms(lambda: hr.run_graph("provider", ctx)) + case["candidate_ms"], case["provider_ms"] = cand, prov + case["nodes"] = [ + {k: v for k, v in n.items() if k != "candidate_sha256"} for n in case["nodes"] + ] + report["cases"].append(case) + print(f"T={tokens}: candidate {cand:.2f} ms provider {prov:.2f} ms") + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1, default=str) + "\n") + if args.png: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(1, 2, figsize=(12, 4)) + big = report["cases"][-1] + xs = range(len(big["nodes"])) + ax[0].semilogy( + xs, + [max(n["candidate_rel_err_vs_golden"], 1e-9) for n in big["nodes"]], + "o-", + label="RL-Kernel", + ) + ax[0].semilogy( + xs, + [max(n["provider_rel_err_vs_golden"], 1e-9) for n in big["nodes"]], + "x--", + label="provider", + ) + ax[0].set_xticks(list(xs)) + ax[0].set_xticklabels( + [n["node"] for n in big["nodes"]], rotation=70, ha="right", fontsize=6 + ) + ax[0].set( + ylabel=f"relative error vs FP64 ({big['tokens']} tokens)", title="Accuracy per node" + ) + ax[0].legend() + toks = [c["tokens"] for c in report["cases"]] + ax[1].plot(toks, [c["candidate_ms"] for c in report["cases"]], "o-", label="RL-Kernel") + ax[1].plot(toks, [c["provider_ms"] for c in report["cases"]], "x--", label="provider") + ax[1].set( + xlabel="tokens", ylabel="forward ms, one refiner block (H100)", title="Performance" + ) + ax[1].legend() + args.png.parent.mkdir(parents=True, exist_ok=True) + fig.tight_layout() + fig.savefig(args.png, dpi=120) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_trajectory.py b/tools/validation/models/h3_trajectory.py new file mode 100644 index 000000000..7dc4b5692 --- /dev/null +++ b/tools/validation/models/h3_trajectory.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Verify and compare sealed MiniMax-H3 trajectory artifacts. + + python tools/validation/models/h3_trajectory.py verify run.json + python tools/validation/models/h3_trajectory.py compare a.json b.json \\ + [--tensors-a a.safetensors --tensors-b b.safetensors] [--out report.json] [--require-equal] + +``compare`` prints the first drift (step, kind, modality, and with tensor files the +mismatch count and max abs difference) and exits 1 with ``--require-equal`` when the two +runs are not bitwise identical, so it can gate a strict claim. +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.models.minimax_h3 import trajectory as tj # noqa: E402 + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="command", required=True) + verify = sub.add_parser("verify") + verify.add_argument("artifact", type=Path) + compare = sub.add_parser("compare") + compare.add_argument("lhs", type=Path) + compare.add_argument("rhs", type=Path) + compare.add_argument("--tensors-a", type=Path) + compare.add_argument("--tensors-b", type=Path) + compare.add_argument("--out", type=Path) + compare.add_argument("--require-equal", action="store_true") + args = parser.parse_args() + + if args.command == "verify": + artifact = tj.loads(args.artifact.read_text()) + print( + f"ok: {artifact['label']} {artifact['variant']} seed={artifact['seed']} " + f"steps={artifact['steps']} sha256={artifact['artifact_sha256']}" + ) + return 0 + lhs = tj.loads(args.lhs.read_text()) + rhs = tj.loads(args.rhs.read_text()) + report = tj.drift_report(lhs, rhs, tensors_lhs=args.tensors_a, tensors_rhs=args.tensors_b) + text = json.dumps(report, indent=1, sort_keys=True) + "\n" + if args.out: + args.out.write_text(text) + summary = {k: report[k] for k in ("labels", "steps", "compared", "differing", "bitwise_equal")} + summary["first_drift"] = report["first_drift"] + print(json.dumps(summary, indent=1)) + return 1 if (args.require_equal and not report["bitwise_equal"]) else 0 + + +if __name__ == "__main__": + raise SystemExit(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_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 index 4c289e628..c53a7846e 100644 --- a/tools/validation/models/plot_h3_evidence.py +++ b/tools/validation/models/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 (``tools/validation/models/h3_evidence.py``) as a PNG. +"""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: @@ -23,6 +24,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). @@ -330,6 +332,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"), @@ -437,12 +453,325 @@ def repeat(mode: str) -> str: 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") + 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], + [ + ("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, } 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}") diff --git a/tools/weights/snapshot_h3_checkpoint.py b/tools/weights/snapshot_h3_checkpoint.py new file mode 100644 index 000000000..1d0687ae6 --- /dev/null +++ b/tools/weights/snapshot_h3_checkpoint.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write ``pinned_checkpoint.json``: the Hub facts the H3 fingerprint is built on. + +Reads, at the pinned revision of ``MiniMaxAI/MiniMax-H3`` and without downloading +any weights: + +* the file tree (path, size, LFS sha256) of the variant transformers, the two + schedulers and the FL2VA / Ref2VA wrapper directories; +* ``config.json`` / ``scheduler_config.json`` of each; +* the safetensors headers (name, dtype, shape) of every transformer shard, reduced to a + digest, a parameter count and a tensor count per variant. + + python tools/weights/snapshot_h3_checkpoint.py [--out PATH] [--token-file PATH] + +The token comes from ``HF_TOKEN`` or the Hugging Face token file. The output is +committed; ``tests/models/minimax_h3/test_h3_arch_fingerprint.py`` checks it +against the analytic tensor table so a stale or edited file fails closed. +""" + +from __future__ import annotations + +import argparse +import json +import os +import struct +import sys +import urllib.request +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.models.minimax_h3 import fingerprint as fp # noqa: E402 + +DIRECTORIES = ("transformer", "transformer_ref", "scheduler", "audio_scheduler", "FL2VA", "Ref2VA") +CONFIG_FILES = { + "transformer": "config.json", + "transformer_ref": "config.json", + "scheduler": "scheduler_config.json", + "audio_scheduler": "scheduler_config.json", +} + + +def _token(token_file: Path | None) -> str: + token = os.environ.get("HF_TOKEN", "").strip() + if not token: + path = token_file or Path.home() / ".cache" / "huggingface" / "token" + token = path.read_text().strip() if path.is_file() else "" + if not token: + raise SystemExit("set HF_TOKEN or pass --token-file (the repo requires access)") + return token + + +class Hub: + """Minimal authenticated reader for one pinned revision.""" + + def __init__(self, token: str): + self.token = token + + def _open(self, url: str, headers: dict[str, str] | None = None): + request = urllib.request.Request( + url, headers={"Authorization": f"Bearer {self.token}", **(headers or {})} + ) + return urllib.request.urlopen(request, timeout=60) + + def tree(self, directory: str) -> list[dict]: + url = ( + f"https://huggingface.co/api/models/{fp.REPO}/tree/{fp.REVISION}/{directory}" + "?recursive=true" + ) + entries: list[dict] = [] + while url: + response = self._open(url) + entries += json.load(response) + link = response.headers.get("Link") or "" + url = link.split("<")[1].split(">")[0] if 'rel="next"' in link else None + return entries + + def resolve(self, path: str, headers: dict[str, str] | None = None) -> bytes: + url = f"https://huggingface.co/{fp.REPO}/resolve/{fp.REVISION}/{path}" + return self._open(url, headers).read() + + def safetensors_header(self, path: str) -> dict: + (length,) = struct.unpack(" dict: + snapshot: dict = {"repo": fp.REPO, "revision": fp.REVISION, "dirs": {}, "configs": {}} + for directory in DIRECTORIES: + files = [ + { + "path": entry["path"], + "size": entry["size"], + "sha256": (entry.get("lfs") or {}).get("oid"), + "git_oid": entry["oid"], + } + for entry in hub.tree(directory) + if entry["type"] == "file" + ] + snapshot["dirs"][directory] = sorted(files, key=lambda f: f["path"]) + if directory in CONFIG_FILES: + raw = hub.resolve(f"{directory}/{CONFIG_FILES[directory]}") + snapshot["configs"][directory] = json.loads(raw) + tables: dict = {} + for directory in ("transformer", "transformer_ref"): + table: dict[str, tuple[str, tuple[int, ...]]] = {} + for entry in snapshot["dirs"][directory]: + if entry["path"].endswith(".safetensors"): + for name, item in hub.safetensors_header(entry["path"]).items(): + table[name] = (item["dtype"], tuple(item["shape"])) + tables[directory] = { + "digest": fp.tensor_table_digest(table), + "parameters": fp.parameter_count(table), + "tensors": len(table), + } + snapshot["tensor_tables"] = tables + return snapshot + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, default=fp.PINNED_CHECKPOINT_PATH) + parser.add_argument("--token-file", type=Path, default=None) + args = parser.parse_args() + snapshot = build_snapshot(Hub(_token(args.token_file))) + args.out.write_text(json.dumps(snapshot, indent=1, sort_keys=True) + "\n") + print(f"wrote {args.out}") + for directory, info in snapshot["tensor_tables"].items(): + print(f"{directory}: {info['tensors']} tensors, {info['parameters']:,} parameters") + + +if __name__ == "__main__": + main()