diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4634569cd..58d2dd175 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -5,9 +5,9 @@ name: CI-Pipeline on: push: - branches: [ main, test ] + branches: [ main, test, test-h3 ] pull_request: - branches: [ main, test ] + branches: [ main, test, test-h3 ] permissions: contents: read @@ -71,6 +71,9 @@ jobs: - name: Check refactor packaging and compatibility run: python ci/run.py layout -q + - name: Run full H3 binding contracts + run: python -m pytest -q tests/models/minimax_h3/test_h3_full_model.py + - name: Run Mocked Hardware Discovery Tests run: | python -m pytest tests/runtime/test_dispatch.py -v diff --git a/benchmarks/models/benchmark_h3_conditioning.py b/benchmarks/models/benchmark_h3_conditioning.py new file mode 100644 index 000000000..261556809 --- /dev/null +++ b/benchmarks/models/benchmark_h3_conditioning.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark the MiniMax-H3 conditioning ops (RFC #420) against the provider path. + +CUDA-event medians, bandwidth and peak memory per case, plus the backend +the registry dispatched. Timings alternate candidate/provider execution order; +backward timings exclude forward setup. Cases live in ``rl_engine/validation/models/h3_report.py``. + + python benchmarks/models/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 + python benchmarks/models/benchmark_h3_conditioning.py --op all --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment # noqa: E402 +from rl_engine.validation.models.h3_report import PERF_CASES, TIMED_KEYS, measure # noqa: E402 + + +def main() -> None: + """Benchmark selected CUDA operators, print summaries, and optionally save JSON.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", default="all", choices=["all", *PERF_CASES]) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iters", type=int, default=200) + parser.add_argument("--json", type=Path, default=None) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + + registry = KernelRegistry() + env = environment() + print(f"device={env['gpu']} torch={env['torch']} cuda={env['cuda']}") + print("timing order alternates each iteration; backward timings exclude forward setup") + results = [] + for name in list(PERF_CASES) if args.op == "all" else [args.op]: + for case in PERF_CASES[name](registry): + row = measure(case, args.warmup, args.iters) + results.append(row) + parts = [ + f"{key}={row[f'{key}_us']:.2f}us ({row[f'{key}_gbps']:.1f} GB/s, " + f"peak {row[f'{key}_peak_mib']:.2f} MiB)" + for key in TIMED_KEYS + if f"{key}_us" in row + ] + order = row["execution_order"] + phases = [" -> ".join(order[f"iteration_{i}"]) for i in (0, 1)] + print( + f"{name} {row['case']} [{row['backend']}]: " + + "; ".join(parts) + + "; alternating order: " + + " / ".join(phases) + ) + if args.json is not None: + args.json.write_text(json.dumps({"environment": env, "results": results}, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/build_tools/extensions.py b/build_tools/extensions.py index 670fdb072..dbddf0c0a 100644 --- a/build_tools/extensions.py +++ b/build_tools/extensions.py @@ -14,6 +14,20 @@ _PROJECT_ROOT = Path(__file__).resolve().parents[1] +# MiniMax-H3 (RFC #420) kernels promise bitwise parity with diffusers' eager +# FP32 math. --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for +# approximations and flushes denormals, so these sources are left out of +# KERNEL_ALIGN_USE_FAST_MATH=1 builds; their bindings are gated on +# KERNEL_ALIGN_WITH_H3_KERNELS and dispatch falls back to the PyTorch reference. +H3_BITWISE_SOURCES = ( + "csrc/cuda/h3/timestep_sinusoid.cu", + "csrc/cuda/h3/det_linear.cu", + "csrc/cuda/h3/adaln_row_gather.cu", + "csrc/cuda/h3/rmsnorm_modulate.cu", + "csrc/cuda/h3/gate_residual.cu", +) +H3_KERNELS_DEFINE = "-DKERNEL_ALIGN_WITH_H3_KERNELS" + def _resolve_package_data_source( filename: str | os.PathLike[str], @@ -219,6 +233,7 @@ def get_extensions(): ) if is_rocm or torch.cuda.is_available(): + use_fast_math = envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH) cuda_sources = [ "csrc/bindings/ops.cpp", "csrc/cuda/logprob/fused_logp_kernel.cu", @@ -241,12 +256,14 @@ def get_extensions(): # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/collectives/deterministic_collective.cu") + if not use_fast_math: + cuda_sources.extend(H3_BITWISE_SOURCES) # This source contains NVIDIA PTX (cp.async, ldmatrix, and mma.sync). # The ROCm dispatcher falls back to PyTorch SDPA for this operator. cuda_sources.append("csrc/cuda/attention/prefix_shared_attention.cu") nvcc_flags = ["-O3", "-Xfatbin", "-compress-all"] - if envs.env_flag(envs.KERNEL_ALIGN_USE_FAST_MATH): + if use_fast_math: nvcc_flags.append("--use_fast_math") if not is_rocm: cc_major, cc_minor = torch.cuda.get_device_capability() @@ -333,6 +350,9 @@ def get_extensions(): if os.name != "nt" and not is_rocm: # CUDA IPC metadata queries use the driver API (cuPointerGetAttribute). extra_link_args.append("-lcuda") + if not is_rocm and not use_fast_math: + cxx_flags.append(H3_KERNELS_DEFINE) + nvcc_flags.append(H3_KERNELS_DEFINE) if not is_rocm: sm90_srcs = [ diff --git a/csrc/bindings/ops.cpp b/csrc/bindings/ops.cpp index 0c7c3e7f4..0046fce21 100644 --- a/csrc/bindings/ops.cpp +++ b/csrc/bindings/ops.cpp @@ -500,6 +500,67 @@ at::Tensor prefix_shared_attention( #endif #endif +// MiniMax-H3 (RFC #420) conditioning-path declarations. CUDA only. +// Fast-math builds omit the H3 sources (bitwise contract), so gate on the build define. +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period, bool check_range); +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation); +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype); +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight); +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype); +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias); +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num); +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype); +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index); +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile); +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index); +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile); +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end); +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile); +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end); +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype); +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index); +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index); +#endif + // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "RL-Kernel High-Performance Operator Extension Library"; @@ -758,4 +819,76 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Deterministic GPT-NeoX token-major RoPE apply for ROCm"); #endif #endif + +#if defined(KERNEL_ALIGN_WITH_H3_KERNELS) + m.def("h3_timestep_sinusoid_forward", torch::wrap_pybind_function(h3_timestep_sinusoid_forward), + "MiniMax-H3 FP32 [cos | sin] timestep features, bitwise to the diffusers CUDA path", + py::arg("timestep"), py::arg("num_channels") = 256, py::arg("max_period") = 10000.0, + py::arg("check_range") = true); + m.def("h3_det_linear_forward", &h3_det_linear_forward, + "Batch-invariant warp-per-column linear (contract h3-det-linear-v1), optional SiLU", + py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), + py::arg("activation") = 0, py::arg("save_pre_activation") = false); + m.def("h3_det_linear_backward_input", &h3_det_linear_backward_input, + "Deterministic grad @ weight with fixed 64-row N chunks folded in order", + py::arg("grad"), py::arg("weight"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_input_partials", &h3_det_linear_backward_input_partials, + "Per-64-row-chunk FP32 partials of grad @ weight, before the ascending fold", + py::arg("grad"), py::arg("weight")); + m.def("h3_det_linear_fold_chunks", &h3_det_linear_fold_chunks, + "Ascending left fold of h3_det_linear_backward_input_partials, cast once", + py::arg("partial"), py::arg("out_dtype")); + m.def("h3_det_linear_backward_weight", &h3_det_linear_backward_weight, + "Deterministic dW/dbias as ascending-row FP32 folds", + py::arg("grad"), py::arg("x"), py::arg("w_dtype"), py::arg("with_bias") = true); + m.def("h3_adaln_row_gather_forward", &h3_adaln_row_gather_forward, + "Fused six-way AdaLN row gather by timestep_index * 3 + token_tag (pure copy)", + py::arg("rows"), py::arg("timestep_indices"), py::arg("token_tags"), + py::arg("chunks") = 6, py::arg("modality_num") = 3); + m.def("h3_adaln_row_gather_backward", &h3_adaln_row_gather_backward, + "Deterministic segmented sum (sorted tiles folded in order) for the row gather", + py::arg("grad"), py::arg("sorted_pos"), py::arg("tile_begin"), py::arg("tile_end"), + py::arg("seg_first_tile"), py::arg("out_dtype")); + m.def("h3_rmsnorm_forward", &h3_rmsnorm_forward, + "RMSNorm replaying PyTorch's reduction order, with optional fused AdaLN modulation", + py::arg("x"), py::arg("weight"), py::arg("eps"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_rmsnorm_backward", &h3_rmsnorm_backward, + "Deterministic RMSNorm(+modulation) backward: row-local dx, tiled dweight, sorted table grads", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), + py::arg("shift") = py::none(), py::arg("scale") = py::none(), + py::arg("index") = py::none(), py::arg("sorted_pos") = py::none(), + py::arg("tile_begin") = py::none(), py::arg("tile_end") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_residual_forward", &h3_gate_residual_forward, + "residual + gate[index] * y with the gate row gathered in-kernel (eager rounding order)", + py::arg("residual"), py::arg("y"), py::arg("gate"), py::arg("index")); + m.def("h3_gate_residual_backward", &h3_gate_residual_backward, + "Gated-residual backward: exact dy, deterministic sorted segment sum for dgate", + py::arg("grad"), py::arg("y"), py::arg("gate"), py::arg("index"), py::arg("sorted_pos"), + py::arg("tile_begin"), py::arg("tile_end"), py::arg("seg_first_tile")); + m.def("h3_rmsnorm_backward_partials", &h3_rmsnorm_backward_partials, + "WS1 RMSNorm dweight / table-gradient tile partials over explicit row lists (SP)", + py::arg("grad"), py::arg("x"), py::arg("weight"), py::arg("rstd"), py::arg("shift"), + py::arg("scale"), py::arg("index"), py::arg("dw_rows"), py::arg("dw_begin"), + py::arg("dw_end"), py::arg("seg_rows") = py::none(), py::arg("seg_begin") = py::none(), + py::arg("seg_end") = py::none()); + m.def("h3_rmsnorm_fold_partials", &h3_rmsnorm_fold_partials, + "WS1 ascending folds of RMSNorm dweight and per-segment table-gradient partials", + py::arg("dw_partial"), py::arg("weight"), py::arg("seg_partial") = py::none(), + py::arg("seg_first_tile") = py::none()); + m.def("h3_gate_grad_partials", &h3_gate_grad_partials, + "WS1 d_gate tile partials over explicit row lists (SP)", py::arg("grad"), py::arg("y"), + py::arg("rows"), py::arg("tile_begin"), py::arg("tile_end")); + m.def("h3_gate_grad_fold", &h3_gate_grad_fold, + "WS1 per-segment ascending fold of d_gate partials, cast once", py::arg("partial"), + py::arg("seg_first_tile"), py::arg("dtype")); + m.def("h3_rmsnorm_backward_dx", &h3_rmsnorm_backward_dx, + "Row-local dx of the WS1 RMSNorm(+modulation) backward", py::arg("grad"), py::arg("x"), + py::arg("weight"), py::arg("rstd"), py::arg("shift") = py::none(), + py::arg("scale") = py::none(), py::arg("index") = py::none()); + m.def("h3_gate_residual_backward_dy", &h3_gate_residual_backward_dy, + "Row-local d_y = grad * gate[index] of the WS1 gated-residual backward", + py::arg("grad"), py::arg("gate"), py::arg("index")); +#endif } diff --git a/csrc/cuda/h3/adaln_row_gather.cu b/csrc/cuda/h3/adaln_row_gather.cu new file mode 100644 index 000000000..0de2d98d5 --- /dev/null +++ b/csrc/cuda/h3/adaln_row_gather.cu @@ -0,0 +1,259 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 AdaLN row gather (RFC #420 `adaln_row_gather`). +// +// rows: (3T, C * H) modulation rows, row r = timestep * 3 + modality and +// column block c = shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, +// gate_mlp. For every packed-sequence position s: +// +// r[s] = timestep_indices[s] * modality_num + token_tags[s] +// out[c, s, :] = rows[r[s], c * H : (c + 1) * H] +// +// Forward is a pure copy: bitwise equal to C index_select calls, one launch, +// independent of S, packing order and repeats. +// +// Backward d_rows[r, j] = sum_{s : r[s] = r} grad[c(j), s, h(j)] is the one +// cross-row reduction. Positions are visited in a fixed order: the caller +// sorts them by (r, s) with a stable sort and cuts every segment into tiles of +// at most kTile positions. Each tile is an ascending FP32 chain from 0; the +// tiles of a segment are then left-folded in ascending order and cast once. +// No atomics: the result is a function of (r[s], grad) only. + +#include +#include +#include +#include +#include + +#include + +namespace { + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_float(__half v) { return __half2float(v); } +__device__ __forceinline__ float to_float(double v) { return static_cast(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); +} +template <> +__device__ __forceinline__ __half from_float<__half>(float v) { return __float2half(v); } +template <> +__device__ __forceinline__ double from_float(float v) { return v; } + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; +template <> +struct CudaType { + using type = __half; +}; + +template +__device__ __forceinline__ int64_t adaln_row(const idx_t* ti, const idx_t* tags, int64_t s, + int64_t modality_num, int64_t num_timesteps) { + const int64_t timestep = static_cast(ti[s]); + const int64_t tag = static_cast(tags[s]); + // Check each semantic index before multiplication, including values that overflow int64. + CUDA_KERNEL_ASSERT(timestep >= 0 && timestep < num_timesteps); + CUDA_KERNEL_ASSERT(tag >= 0 && tag < modality_num); + return timestep * modality_num + tag; +} + +// One block per (s, c); 16-byte copies when every row start is 16-byte aligned. +template +__global__ void adaln_row_gather_kernel(const T* __restrict__ rows, int64_t row_stride, + const idx_t* __restrict__ ti, + const idx_t* __restrict__ tags, T* __restrict__ out, + int64_t seq, int64_t hidden, int64_t chunks, + int64_t modality_num, int64_t num_timesteps) { + for (int64_t s = blockIdx.x; s < seq; s += gridDim.x) { + const int64_t c = blockIdx.y; + const int64_t r = adaln_row(ti, tags, s, modality_num, num_timesteps); + const T* src = rows + r * row_stride + c * hidden; + T* dst = out + (c * seq + s) * hidden; + if constexpr (kVector) { + constexpr int kVec = 16 / sizeof(T); + const uint4* src4 = reinterpret_cast(src); + uint4* dst4 = reinterpret_cast(dst); + for (int64_t i = threadIdx.x; i < hidden / kVec; i += blockDim.x) dst4[i] = src4[i]; + } else { + for (int64_t i = threadIdx.x; i < hidden; i += blockDim.x) dst[i] = src[i]; + } + } +} + +// partial[tile, j] = sum over sorted positions p in tile (ascending) of grad[c, s_p, h] +template +__global__ void adaln_row_gather_partial_kernel(const g_t* __restrict__ grad, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t seq, + int64_t hidden, int64_t width) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + const int64_t c = j / hidden; + const int64_t h = j - c * hidden; + const g_t* g = grad + c * seq * hidden + h; + float acc = 0.0f; + const int64_t begin = tile_begin[tile], end = tile_end[tile]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= seq); + for (int64_t p = begin; p < end; ++p) { + CUDA_KERNEL_ASSERT(sorted_pos[p] >= 0 && sorted_pos[p] < seq); + acc += to_float(g[sorted_pos[p] * hidden]); + } + partial[tile * width + j] = acc; +} + +template +__global__ void adaln_row_gather_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + out_t* __restrict__ out, int64_t width, int64_t tiles) { + const int64_t r = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + const int64_t begin = seg_first_tile[r], end = seg_first_tile[r + 1]; + CUDA_KERNEL_ASSERT(begin >= 0 && begin <= end && end <= tiles); + for (int64_t tile = begin; tile < end; ++tile) { + acc += partial[tile * width + j]; + } + out[r * width + j] = from_float(acc); +} + +void check_index(const torch::Tensor& t, const char* name, int64_t seq, + const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 1 && t.size(0) == seq, name, " must be (", seq, ",)"); + TORCH_CHECK(t.scalar_type() == at::kLong || t.scalar_type() == at::kInt, name, + " must be int64 or int32"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +// rows (R, C * H) with unit column stride -> out (C, S, H) contiguous. +torch::Tensor h3_adaln_row_gather_forward(torch::Tensor rows, torch::Tensor timestep_indices, + torch::Tensor token_tags, int64_t chunks, + int64_t modality_num) { + TORCH_CHECK(rows.is_cuda() && rows.dim() == 2 && rows.stride(1) == 1, + "rows must be a 2-D CUDA tensor with unit column stride"); + TORCH_CHECK(chunks > 0 && rows.size(1) % chunks == 0, "rows width ", rows.size(1), + " is not a multiple of ", chunks); + TORCH_CHECK(modality_num > 0 && rows.size(0) % modality_num == 0, "rows count ", + rows.size(0), " is not a multiple of ", modality_num); + const int64_t num_timesteps = rows.size(0) / modality_num; + const int64_t seq = timestep_indices.numel(); + TORCH_CHECK(seq > 0, "the packed sequence must not be empty"); + check_index(timestep_indices, "timestep_indices", seq, rows); + check_index(token_tags, "token_tags", seq, rows); + TORCH_CHECK(timestep_indices.scalar_type() == token_tags.scalar_type(), + "timestep_indices and token_tags must share an integer dtype"); + const int64_t hidden = rows.size(1) / chunks; + const c10::cuda::CUDAGuard device_guard(rows.device()); + auto out = torch::empty({chunks, seq, hidden}, rows.options()); + auto stream = at::cuda::getCurrentCUDAStream(); + const dim3 grid(static_cast(std::min(seq, 1 << 20)), + static_cast(chunks)); + const int threads = 128; + const int64_t elem = rows.element_size(); + const bool vector = (hidden * elem) % 16 == 0 && (rows.stride(0) * elem) % 16 == 0 && + reinterpret_cast(rows.data_ptr()) % 16 == 0; + + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, rows.scalar_type(), "h3_adaln_row_gather_forward", [&] { + using T = typename CudaType::type; + const T* src = reinterpret_cast(rows.data_ptr()); + T* dst = reinterpret_cast(out.data_ptr()); + auto launch = [&](auto idx_tag) { + using idx_t = decltype(idx_tag); + const idx_t* ti = timestep_indices.data_ptr(); + const idx_t* tags = token_tags.data_ptr(); + if (vector) { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); + } else { + adaln_row_gather_kernel<<>>( + src, rows.stride(0), ti, tags, dst, seq, hidden, chunks, modality_num, + num_timesteps); + } + }; + if (timestep_indices.scalar_type() == at::kLong) { + launch(int64_t{0}); + } else { + launch(int32_t{0}); + } + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad (C, S, H) contiguous; sorted_pos (S,) positions sorted stably by row; +// tiles [tile_begin, tile_end) in sorted order; seg_first_tile (R + 1,). +torch::Tensor h3_adaln_row_gather_backward(torch::Tensor grad, torch::Tensor sorted_pos, + torch::Tensor tile_begin, torch::Tensor tile_end, + torch::Tensor seg_first_tile, + c10::ScalarType out_dtype) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 3 && grad.is_contiguous(), + "grad must be a contiguous (C, S, H) CUDA tensor"); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && + t->dim() == 1 && t->is_contiguous() && + t->scalar_type() == at::kLong, + "tile metadata must be contiguous 1-D int64 tensors on grad device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin/tile_end length mismatch"); + const int64_t chunks = grad.size(0); + const int64_t seq = grad.size(1); + const int64_t hidden = grad.size(2); + const int64_t width = chunks * hidden; + const int64_t num_rows = seg_first_tile.numel() - 1; + const int64_t tiles = tile_begin.numel(); + TORCH_CHECK(chunks > 0 && seq > 0 && hidden > 0 && num_rows > 0, + "gradient dimensions and table row count must be positive"); + TORCH_CHECK(sorted_pos.numel() == seq, "sorted_pos must have S entries"); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto out = torch::empty({num_rows, width}, grad.options().dtype(out_dtype)); + auto partial = torch::empty({std::max(tiles, 1), width}, + grad.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((width + threads - 1) / threads); + if (tiles > 0) { + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, grad.scalar_type(), "h3_adaln_row_gather_partial", [&] { + using G = typename CudaType::type; + adaln_row_gather_partial_kernel + <<(tiles), col_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), + sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), seq, hidden, width); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + AT_DISPATCH_FLOATING_TYPES_AND2( + at::kHalf, at::kBFloat16, out_dtype, "h3_adaln_row_gather_fold", [&] { + using O = typename CudaType::type; + adaln_row_gather_fold_kernel + <<(num_rows), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), width, tiles); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/csrc/cuda/h3/det_linear.cu b/csrc/cuda/h3/det_linear.cu new file mode 100644 index 000000000..96b954b0e --- /dev/null +++ b/csrc/cuda/h3/det_linear.cu @@ -0,0 +1,620 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Deterministic, batch-invariant row-wise linear layers for the MiniMax-H3 +// conditioning path (RFC #420: `timestep_mlp_fp32`, `adaln_projection_3mod`). +// +// These layers run on a handful of rows (one per distinct timestep) against +// large weights, so they are GEMVs bound by weight bandwidth. The reduction +// order is fixed and depends only on K (and N for the input gradient), never +// on the number of rows or their position. +// +// Reduction contract h3-det-linear-v1 +// ----------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * one warp per output column n; lane l owns the 16-byte K chunks +// c = l, l + 32, l + 64, ... (VEC = 16 / sizeof(w) elements each); +// * each lane accumulates its chunks in ascending order and the elements of +// a chunk in ascending order with fmaf into an FP32 accumulator from 0; +// * the 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1), +// whose result is identical on every lane (FP add is commutative); +// * the bias is added once after the butterfly; the activation runs in FP32 +// (SiLU as v / (1 + expf(-v)), PyTorch's formula); one cast at the store. +// * this FP32 forward is used for FP32 weights; BF16 weights take the +// tensor-core forward of contract h3-det-linear-bf16-mma-v1 below. +// d_input dx[t, k] = sum_n g[t, n] * W[n, k] +// * N is split into fixed chunks of 64 rows; each chunk is an ascending +// fmaf chain into FP32 from 0; the chunk partials are then left-folded in +// ascending chunk order and cast once. +// The partials and the fold are also bound separately: a contiguous N +// slice starting on a chunk boundary yields exactly the full call's +// partials for those chunks, which is what tensor-parallel shards gather. +// d_weight dW[n, k] = sum_t g[t, n] * x[t, k]; d_bias db[n] = sum_t g[t, n] +// * an ascending-t fmaf chain into FP32 from 0, cast once. This is the one +// cross-row reduction; its order is the logical row order. +// +// Reduction contract h3-det-linear-bf16-mma-v1 (BF16 x and W, FP32 accumulate) +// ----------------------------------------------------------------------------- +// forward y[t, n] = act(bias[n] + dot(x[t, :], W[n, :])) +// * a warp owns 16 output columns (MMA rows) and 8 input rows (MMA columns; +// rows past T are zero), so every launch runs the same instruction +// sequence for every T <= 8, and an MMA column never sees another's data; +// * K is visited in groups of 16 in ascending order. Within a 32-wide step, +// quad q supplies k = 32s + 8q + {0..3} to the first mma.sync m16n8k16 and +// k = 32s + 8q + {4..7} to the second (the same k map for W and x); +// * every mma.sync starts from a zero accumulator, so the tensor core only +// sums 16 products; its FP32 result is added to the running FP32 sum with +// an IEEE add, group by group in ascending k; +// * then bias, FP32 activation and one cast, as above. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +#ifndef H3_FWD_F32_COLS +#define H3_FWD_F32_COLS 2 +#define H3_FWD_F32_AHEAD 6 +#endif +#ifndef H3_FWD_CONFIGS +#define H3_FWD_CONFIGS H3_FWD_CASE(H3_FWD_F32_COLS, H3_FWD_F32_AHEAD) +#endif + +constexpr int kWarp = 32; +constexpr int kWarpsPerBlock = 8; +constexpr int kRowTile = 4; +constexpr int kDInputChunk = 64; + +enum Activation : int64_t { kActNone = 0, kActSilu = 1 }; + +__device__ __forceinline__ float to_float(float v) { return v; } +__device__ __forceinline__ float to_float(__nv_bfloat16 v) { return __bfloat162float(v); } + +template +__device__ __forceinline__ T from_float(float v); +template <> +__device__ __forceinline__ float from_float(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float v) { + return __float2bfloat16(v); // round to nearest even, like at::BFloat16 +} + +// 16-byte vector load converted to FP32. +template +struct Vec16 { + static constexpr int kN = 16 / sizeof(T); + __device__ __forceinline__ static void load(const T* ptr, float (&out)[kN]) { + unpack(load_raw(ptr), out); + } + __device__ __forceinline__ static uint4 load_raw(const T* ptr) { + return *reinterpret_cast(ptr); + } + // Streaming weights are read once: bypass L1 allocation. + __device__ __forceinline__ static uint4 load_stream(const T* ptr) { + return __ldcs(reinterpret_cast(ptr)); + } + __device__ __forceinline__ static void unpack(const uint4& raw, float (&out)[kN]) { + const T* vals = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < kN; ++i) out[i] = to_float(vals[i]); + } +}; + +__device__ __forceinline__ float warp_butterfly_sum(float v) { +#pragma unroll + for (int offset = kWarp / 2; offset > 0; offset >>= 1) { + v += __shfl_xor_sync(0xffffffffu, v, offset); + } + return v; +} + +// Each warp owns kCols adjacent output columns so that kLoadAhead x kCols +// 16-byte weight loads are in flight per lane. Grouping columns and chunks +// changes only when loads are issued: every column still accumulates its own +// chunks in ascending order, exactly as a one-column-per-warp kernel would. +template +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_kernel(const x_t* __restrict__ x, const w_t* __restrict__ w, + const w_t* __restrict__ bias, out_t* __restrict__ out, + float* __restrict__ pre_act, int64_t rows, int64_t n_out, + int64_t k_in, int64_t activation) { + using WV = Vec16; + using XV = Vec16; + constexpr int kVec = WV::kN; + static_assert(kCols * kRows <= kWarp, "one warp cannot store more than kWarp outputs per tile"); + static_assert(XV::kN == kVec, "x and weight share a dtype"); + constexpr int64_t kStride = static_cast(kWarp) * kVec; + + const int lane = threadIdx.x % kWarp; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + + for (int64_t n0 = warp_global * kCols; n0 < n_out; n0 += warp_stride * kCols) { + for (int64_t t0 = 0; t0 < rows; t0 += kRows) { + float acc[kCols][kRows]; +#pragma unroll + for (int c = 0; c < kCols; ++c) { +#pragma unroll + for (int r = 0; r < kRows; ++r) acc[c][r] = 0.0f; + } + for (int64_t g0 = static_cast(lane) * kVec; g0 < k_in; + g0 += kStride * kLoadAhead) { + uint4 raw[kLoadAhead][kCols]; +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; +#pragma unroll + for (int c = 0; c < kCols; ++c) { + if (k0 < k_in && n0 + c < n_out) raw[u][c] = WV::load_stream(w + (n0 + c) * k_in + k0); + } + } +#pragma unroll + for (int u = 0; u < kLoadAhead; ++u) { + const int64_t k0 = g0 + u * kStride; + if (k0 >= k_in) break; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + if (t0 + r < rows) { + float xv[kVec]; + XV::load(x + (t0 + r) * k_in + k0, xv); +#pragma unroll + for (int c = 0; c < kCols; ++c) { + float wv[kVec]; + WV::unpack(raw[u][c], wv); +#pragma unroll + for (int j = 0; j < kVec; ++j) acc[c][r] = fmaf(xv[j], wv[j], acc[c][r]); + } + } + } + } + } +#pragma unroll + for (int c = 0; c < kCols; ++c) { + const int64_t n = n0 + c; +#pragma unroll + for (int r = 0; r < kRows; ++r) { + const int64_t t = t0 + r; + const float sum = warp_butterfly_sum(acc[c][r]); + if (lane == c * kRows + r && t < rows && n < n_out) { + float v = sum + (bias != nullptr ? to_float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = from_float(v); + } + } + } + } + } +} + +__device__ __forceinline__ void mma_m16n8k16_bf16(float (&c)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 + // BF16 mma.sync needs SM80+; the host refuses to launch below that. + __trap(); +#else + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, " + "{%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(c[0]), "+f"(c[1]), "+f"(c[2]), "+f"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b0), "r"(b1)); +#endif +} + +constexpr int kMmaCols = 16; // output columns per warp (MMA M) +constexpr int kMmaRows = 8; // input rows per pass (MMA N) +constexpr int kMmaAhead = 4; // 32-wide K steps loaded before use + +__device__ __forceinline__ void mma_group_accumulate(float (&acc)[4], const uint32_t (&a)[4], + uint32_t b0, uint32_t b1) { + float part[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + mma_m16n8k16_bf16(part, a, b0, b1); +#pragma unroll + for (int i = 0; i < 4; ++i) acc[i] += part[i]; +} + +// Contract h3-det-linear-bf16-mma-v1. Lane = 4 * g + q: g picks the MMA row +// (output columns n0 + g and n0 + g + 8) and the MMA column (input row t0 + g). +__global__ void __launch_bounds__(kWarpsPerBlock * kWarp) + det_linear_forward_bf16_mma_kernel(const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ w, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ out, float* __restrict__ pre_act, + int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation) { + const int lane = threadIdx.x % kWarp; + const int g = lane >> 2; + const int q = lane & 3; + const int64_t warp_global = + static_cast(blockIdx.x) * kWarpsPerBlock + threadIdx.x / kWarp; + const int64_t warp_stride = static_cast(gridDim.x) * kWarpsPerBlock; + const uint4 zero = make_uint4(0u, 0u, 0u, 0u); + + for (int64_t n0 = warp_global * kMmaCols; n0 < n_out; n0 += warp_stride * kMmaCols) { + const bool lo_ok = n0 + g < n_out; + const bool hi_ok = n0 + g + 8 < n_out; + const __nv_bfloat16* w_lo = w + (n0 + g) * k_in; + const __nv_bfloat16* w_hi = w + (n0 + g + 8) * k_in; + for (int64_t t0 = 0; t0 < rows; t0 += kMmaRows) { + const bool x_ok = t0 + g < rows; + const __nv_bfloat16* x_row = x + (t0 + g) * k_in; + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + for (int64_t kb = 0; kb < k_in; kb += 32 * kMmaAhead) { + uint4 a_lo[kMmaAhead], a_hi[kMmaAhead], bx[kMmaAhead]; +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + const int64_t k = kb + 32 * u + 8 * q; + const bool k_ok = k < k_in; + a_lo[u] = (k_ok && lo_ok) ? __ldcs(reinterpret_cast(w_lo + k)) : zero; + a_hi[u] = (k_ok && hi_ok) ? __ldcs(reinterpret_cast(w_hi + k)) : zero; + bx[u] = (k_ok && x_ok) ? *reinterpret_cast(x_row + k) : zero; + } +#pragma unroll + for (int u = 0; u < kMmaAhead; ++u) { + if (kb + 32 * u >= k_in) break; + const uint32_t first[4] = {a_lo[u].x, a_hi[u].x, a_lo[u].y, a_hi[u].y}; + const uint32_t second[4] = {a_lo[u].z, a_hi[u].z, a_lo[u].w, a_hi[u].w}; + mma_group_accumulate(acc, first, bx[u].x, bx[u].y); + mma_group_accumulate(acc, second, bx[u].z, bx[u].w); + } + } + // acc = {(n0+g, t0+2q), (n0+g, t0+2q+1), (n0+g+8, t0+2q), (n0+g+8, t0+2q+1)} +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int64_t n = n0 + g + (i >= 2 ? 8 : 0); + const int64_t t = t0 + 2 * q + (i & 1); + if (n < n_out && t < rows) { + float v = acc[i] + (bias != nullptr ? __bfloat162float(bias[n]) : 0.0f); + if (pre_act != nullptr) pre_act[t * n_out + n] = v; + if (activation == kActSilu) v = v / (1.0f + expf(-v)); + out[t * n_out + n] = __float2bfloat16(v); + } + } + } + } +} + +// partial[c, t, k] = sum_{n in chunk c} g[t, n] * w[n, k] +template +__global__ void det_linear_dinput_partial_kernel(const g_t* __restrict__ grad, + const w_t* __restrict__ w, + float* __restrict__ partial, int64_t rows, + int64_t n_out, int64_t k_in) { + const int64_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t chunk = blockIdx.y; + if (k >= k_in) return; + const int64_t n_begin = chunk * kDInputChunk; + const int64_t n_end = min(n_begin + kDInputChunk, n_out); + for (int64_t t0 = 0; t0 < rows; t0 += kRowTile) { + float acc[kRowTile]; +#pragma unroll + for (int r = 0; r < kRowTile; ++r) acc[r] = 0.0f; + for (int64_t n = n_begin; n < n_end; ++n) { + const float wv = to_float(w[n * k_in + k]); +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) acc[r] = fmaf(to_float(grad[t * n_out + n]), wv, acc[r]); + } + } +#pragma unroll + for (int r = 0; r < kRowTile; ++r) { + const int64_t t = t0 + r; + if (t < rows) partial[(chunk * rows + t) * k_in + k] = acc[r]; + } + } +} + +template +__global__ void det_fold_chunks_kernel(const float* __restrict__ partial, out_t* __restrict__ out, + int64_t chunks, int64_t elems) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i >= elems) return; + float acc = partial[i]; + for (int64_t c = 1; c < chunks; ++c) acc += partial[c * elems + i]; + out[i] = from_float(acc); +} + +// dW[n, k] = sum_t g[t, n] * x[t, k] ; ascending t, FP32 from 0, one cast. +template +__global__ void det_linear_dweight_kernel(const g_t* __restrict__ grad, const x_t* __restrict__ x, + w_t* __restrict__ dw, int64_t rows, int64_t n_out, + int64_t k_in) { + const int64_t total = n_out * k_in; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < total; + i += static_cast(gridDim.x) * blockDim.x) { + const int64_t n = i / k_in; + const int64_t k = i - n * k_in; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) { + acc = fmaf(to_float(grad[t * n_out + n]), to_float(x[t * k_in + k]), acc); + } + dw[i] = from_float(acc); + } +} + +template +__global__ void det_linear_dbias_kernel(const g_t* __restrict__ grad, w_t* __restrict__ db, + int64_t rows, int64_t n_out) { + const int64_t n = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (n >= n_out) return; + float acc = 0.0f; + for (int64_t t = 0; t < rows; ++t) acc += to_float(grad[t * n_out + n]); + db[n] = from_float(acc); +} + +template +struct CudaType { + using type = T; +}; +template <> +struct CudaType { + using type = __nv_bfloat16; +}; + +template +const typename CudaType::type* cptr(const torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} +template +typename CudaType::type* mptr(torch::Tensor& t) { + return reinterpret_cast::type*>(t.data_ptr()); +} + +void check_matrix(const torch::Tensor& t, const char* name) { + TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(t.dim() == 2, name, " must be 2-D"); + TORCH_CHECK(t.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(t.scalar_type() == at::kFloat || t.scalar_type() == at::kBFloat16, name, + " must be float32 or bfloat16, got ", t.scalar_type()); + TORCH_CHECK(reinterpret_cast(t.data_ptr()) % 16 == 0, name, + " must be 16-byte aligned"); +} + +int64_t grid_for(int64_t work, int64_t per_block) { + const int64_t blocks = (work + per_block - 1) / per_block; + return std::max(1, std::min(blocks, 1 << 20)); +} + +// Launch shape only: the row tile, the columns per warp and the load depth +// change which loads are in flight, never a column's accumulation order. +struct ForwardConfig { + int rows; + int cols; + int ahead; +}; + +ForwardConfig pick_forward_config(int64_t rows) { + const int row_tile = rows <= 1 ? 1 : (rows <= 2 ? 2 : 4); + return {row_tile, H3_FWD_F32_COLS, H3_FWD_F32_AHEAD}; +} + +template +void launch_forward_impl(const T* x, const T* w, const T* bias, T* out, float* pre, + int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + const int64_t blocks = grid_for((n_out + kCols - 1) / kCols, kWarpsPerBlock); + det_linear_forward_kernel + <<(blocks), kWarpsPerBlock * kWarp, 0, stream>>>( + x, w, bias, out, pre, rows, n_out, k_in, activation); +} + +template +void launch_forward_rows(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, + int64_t activation, cudaStream_t stream) { +#define H3_FWD_CASE(C, A) \ + if (cfg.cols == C && cfg.ahead == A) { \ + launch_forward_impl(x, w, bias, out, pre, rows, n_out, k_in, activation, \ + stream); \ + return; \ + } + H3_FWD_CONFIGS +#undef H3_FWD_CASE + TORCH_CHECK(false, "unsupported det_linear forward config cols=", cfg.cols, + " ahead=", cfg.ahead); +} + +template +void launch_forward(const ForwardConfig& cfg, const T* x, const T* w, const T* bias, T* out, + float* pre, int64_t rows, int64_t n_out, int64_t k_in, int64_t activation, + cudaStream_t stream) { + if (cfg.rows == 1) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else if (cfg.rows == 2) { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } else { + launch_forward_rows(cfg, x, w, bias, out, pre, rows, n_out, k_in, activation, stream); + } +} + +} // namespace + +std::vector h3_det_linear_forward(torch::Tensor x, torch::Tensor weight, + c10::optional bias, + int64_t activation, bool save_pre_activation) { + check_matrix(x, "x"); + check_matrix(weight, "weight"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must be on the same device"); + TORCH_CHECK(x.scalar_type() == weight.scalar_type(), + "x and weight must share a dtype (cast at the declared boundary first), got ", + x.scalar_type(), " and ", weight.scalar_type()); + TORCH_CHECK(activation == kActNone || activation == kActSilu, "unknown activation ", + activation); + const int64_t rows = x.size(0); + const int64_t k_in = x.size(1); + const int64_t n_out = weight.size(0); + TORCH_CHECK(rows > 0, "x must have at least one row"); + TORCH_CHECK(weight.size(1) == k_in, "weight is [", n_out, ", ", weight.size(1), + "] but x has K=", k_in); + const int64_t vec = 16 / x.element_size(); + TORCH_CHECK(k_in % vec == 0, "K=", k_in, " must be a multiple of ", vec, + " for 16-byte vector loads"); + if (bias.has_value()) { + TORCH_CHECK(bias->is_cuda() && bias->dim() == 1 && bias->size(0) == n_out && + bias->is_contiguous() && bias->scalar_type() == weight.scalar_type(), + "bias must be a contiguous [N] CUDA tensor with the weight dtype"); + TORCH_CHECK(bias->device() == weight.device(), + "bias and weight must be on the same device"); + } + + const c10::cuda::CUDAGuard device_guard(x.device()); + auto out = torch::empty({rows, n_out}, x.options()); + torch::Tensor pre; + if (save_pre_activation) pre = torch::empty({rows, n_out}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + if (x.scalar_type() == at::kFloat) { + const ForwardConfig cfg = pick_forward_config(rows); + launch_forward(cfg, cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, + k_in, activation, stream); + } else { + TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8, + "the BF16 h3 det_linear forward uses mma.sync and needs SM80 or newer"); + const int64_t blocks = grid_for((n_out + kMmaCols - 1) / kMmaCols, kWarpsPerBlock); + det_linear_forward_bf16_mma_kernel<<(blocks), kWarpsPerBlock * kWarp, + 0, stream>>>( + cptr(x), cptr(weight), + bias.has_value() ? cptr(*bias) : nullptr, mptr(out), + save_pre_activation ? pre.data_ptr() : nullptr, rows, n_out, k_in, activation); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (save_pre_activation) return {out, pre}; + return {out}; +} + +// grad [T, N] (float32), weight [N, K] -> partial [ceil(N / 64), T, K] float32: +// the per-chunk sums of contract h3-det-linear-v1, before the ascending fold. +// A contiguous slice of N that starts on a chunk boundary yields exactly the +// matching rows of the full call (tensor-parallel shards rely on this). +torch::Tensor h3_det_linear_backward_input_partials(torch::Tensor grad, torch::Tensor weight) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(weight, "weight"); + TORCH_CHECK(grad.device() == weight.device(), "grad and weight must be on the same device"); + TORCH_CHECK(grad.size(1) == weight.size(0), "grad N != weight N"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = weight.size(0); + const int64_t k_in = weight.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + const int64_t chunks = (n_out + kDInputChunk - 1) / kDInputChunk; + auto partial = torch::empty({chunks, rows, k_in}, grad.options()); + if (partial.numel() == 0) return partial; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + dim3 grid(static_cast((k_in + threads - 1) / threads), static_cast(chunks)); + if (weight.scalar_type() == at::kFloat) { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, n_out, + k_in); + } else { + det_linear_dinput_partial_kernel<<>>( + grad.data_ptr(), cptr(weight), partial.data_ptr(), rows, + n_out, k_in); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// partial [C, T, K] float32 -> left fold over C in ascending order, cast once. +torch::Tensor h3_det_linear_fold_chunks(torch::Tensor partial, c10::ScalarType out_dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 3 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat && partial.size(0) > 0, + "partial must be a non-empty contiguous 3-D float32 CUDA tensor"); + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + const c10::cuda::CUDAGuard device_guard(partial.device()); + const int64_t chunks = partial.size(0); + auto out = torch::empty({partial.size(1), partial.size(2)}, partial.options().dtype(out_dtype)); + const int64_t elems = out.numel(); + if (elems == 0) return out; + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned fold_blocks = static_cast((elems + threads - 1) / threads); + if (out_dtype == at::kFloat) { + det_fold_chunks_kernel<<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } else { + det_fold_chunks_kernel<__nv_bfloat16><<>>( + partial.data_ptr(), mptr(out), chunks, elems); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// grad [T, N] (float32), weight [N, K] -> grad_input [T, K] in out_dtype. +torch::Tensor h3_det_linear_backward_input(torch::Tensor grad, torch::Tensor weight, + c10::ScalarType out_dtype) { + TORCH_CHECK(out_dtype == at::kFloat || out_dtype == at::kBFloat16, + "out_dtype must be float32 or bfloat16"); + auto partial = h3_det_linear_backward_input_partials(grad, weight); + if (partial.size(0) == 0) { // N == 0: no output column contributes + return torch::zeros({partial.size(1), partial.size(2)}, grad.options().dtype(out_dtype)); + } + return h3_det_linear_fold_chunks(partial, out_dtype); +} + +// grad [T, N] float32, x [T, K] (float32 or bfloat16) -> (dW [N, K], dbias [N]) in w_dtype. +std::vector h3_det_linear_backward_weight(torch::Tensor grad, torch::Tensor x, + c10::ScalarType w_dtype, + bool with_bias) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous() && + grad.scalar_type() == at::kFloat, + "grad must be a contiguous 2-D float32 CUDA tensor"); + check_matrix(x, "x"); + TORCH_CHECK(grad.device() == x.device(), "grad and x must be on the same device"); + TORCH_CHECK(grad.size(0) == x.size(0), "grad rows != x rows"); + TORCH_CHECK(w_dtype == at::kFloat || w_dtype == at::kBFloat16, + "w_dtype must be float32 or bfloat16"); + const int64_t rows = grad.size(0); + TORCH_CHECK(rows > 0, "grad must have at least one row"); + const int64_t n_out = grad.size(1); + const int64_t k_in = x.size(1); + const c10::cuda::CUDAGuard device_guard(grad.device()); + auto dw = torch::empty({n_out, k_in}, grad.options().dtype(w_dtype)); + if (n_out == 0) { + if (!with_bias) return {dw}; + return {dw, torch::empty({n_out}, grad.options().dtype(w_dtype))}; + } + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const int64_t total = n_out * k_in; + const unsigned blocks = static_cast(grid_for(total, threads)); + const float* g = grad.data_ptr(); +#define H3_DW_LAUNCH(XT, XC, WT, WC) \ + det_linear_dweight_kernel<<>>( \ + g, cptr(x), mptr(dw), rows, n_out, k_in) + if (x.scalar_type() == at::kFloat && w_dtype == at::kFloat) { + H3_DW_LAUNCH(float, float, float, float); + } else if (x.scalar_type() == at::kFloat) { + H3_DW_LAUNCH(float, float, at::BFloat16, __nv_bfloat16); + } else if (w_dtype == at::kFloat) { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, float, float); + } else { + H3_DW_LAUNCH(at::BFloat16, __nv_bfloat16, at::BFloat16, __nv_bfloat16); + } +#undef H3_DW_LAUNCH + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!with_bias) return {dw}; + auto db = torch::empty({n_out}, grad.options().dtype(w_dtype)); + const unsigned bias_blocks = static_cast((n_out + threads - 1) / threads); + if (w_dtype == at::kFloat) { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } else { + det_linear_dbias_kernel<<>>( + g, mptr(db), rows, n_out); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dw, db}; +} diff --git a/csrc/cuda/h3/gate_residual.cu b/csrc/cuda/h3/gate_residual.cu new file mode 100644 index 000000000..eb0b03f50 --- /dev/null +++ b/csrc/cuda/h3/gate_residual.cu @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 gated residual (RFC #420 `adaln_gate_residual`). +// +// out[r, j] = cast(residual[r, j] + cast(gate[i, j] * y[r, j])) i = index[r % S] +// +// The gate row is read from a strided view of the AdaLN table (gate_msa or +// gate_mlp) inside the kernel, and the two roundings sit exactly where the +// eager `residual + gate.index_select(0, i) * y` rounds (__fmul_rn/__fadd_rn +// keep FP32 from contracting them into one FMA), so the result is +// bitwise equal to diffusers. Elementwise: every output depends only on its +// own inputs. +// +// Backward: d_residual = grad; d_y = cast(grad * gate) (the product of two +// 16-bit values is exact in FP32, so this equals the eager VJP bitwise); +// d_gate[i] = sum over positions mapped to i of grad * y, an FP32 segmented +// sum over positions sorted stably by row, in fixed tiles folded in order. +// The tile partials and the fold are also bound separately for sequence +// parallel callers (each tile computed where its rows live, folded in order). + +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +struct Gate { + const void* rows; // (R, N) view, row stride `stride` + int64_t stride; + const int64_t* index; // (S,) + int64_t seq; + int64_t num_rows; +}; + +// One block per row (grid-strided); 16-byte vectors when the row allows it. +// `kVecElems` elements are loaded together; each is computed independently. +template +__global__ void gate_residual_rows_kernel(const T* __restrict__ a, const T* __restrict__ b, + T* __restrict__ out, int64_t rows, int64_t n, + Gate gate, bool vectorized) { + constexpr int kVecElems = 16 / sizeof(T); + for (int64_t r = blockIdx.x; r < rows; r += gridDim.x) { + const int64_t i = gate.index[r % gate.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < gate.num_rows); + const T* g = static_cast(gate.rows) + i * gate.stride; + const T* ar = a + r * n; + const T* br = kDy ? nullptr : b + r * n; + T* orow = out + r * n; + if (vectorized) { + for (int64_t v = threadIdx.x; v < n / kVecElems; v += blockDim.x) { + const uint4 av = reinterpret_cast(ar)[v]; + const uint4 gv = reinterpret_cast(g)[v]; + uint4 bv = av; + if (!kDy) bv = reinterpret_cast(br)[v]; + const T* ae = reinterpret_cast(&av); + const T* ge = reinterpret_cast(&gv); + const T* be = reinterpret_cast(&bv); + uint4 ov; + T* oe = reinterpret_cast(&ov); +#pragma unroll + for (int k = 0; k < kVecElems; ++k) { + if (kDy) { + oe[k] = from_f(to_f(ae[k]) * to_f(ge[k])); // grad * gate + } else { + const float p = round_t(__fmul_rn(to_f(ge[k]), to_f(be[k]))); // gate * y + oe[k] = from_f(__fadd_rn(to_f(ae[k]), p)); // residual + p + } + } + reinterpret_cast(orow)[v] = ov; + } + } else { + for (int64_t j = threadIdx.x; j < n; j += blockDim.x) { + if (kDy) { + orow[j] = from_f(to_f(ar[j]) * to_f(g[j])); + } else { + const float p = round_t(__fmul_rn(to_f(g[j]), to_f(br[j]))); + orow[j] = from_f(__fadd_rn(to_f(ar[j]), p)); + } + } + } + } +} + +// partial[tile, j] = sum over sorted positions of the tile (ascending) of grad * y +template +__global__ void gate_grad_partial_kernel(const T* __restrict__ grad, const T* __restrict__ y, + const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, + const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int64_t n) { + const int64_t tile = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t e = sorted_pos[p] * n + j; + acc = fmaf(to_f(grad[e]), to_f(y[e]), acc); + } + partial[tile * n + j] = acc; +} + +template +__global__ void gate_grad_fold_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + T* __restrict__ out, int64_t n) { + const int64_t seg = blockIdx.x; + const int64_t j = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * n + j]; + out[seg * n + j] = from_f(acc); +} + +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +Gate make_gate(const torch::Tensor& gate, const torch::Tensor& index, const torch::Tensor& x) { + TORCH_CHECK(gate.is_cuda() && gate.device() == x.device() && gate.scalar_type() == x.scalar_type(), + "gate must be a CUDA tensor with the activations' dtype"); + TORCH_CHECK(gate.dim() == 2 && gate.size(1) == x.size(1) && gate.stride(1) == 1, + "gate must be (R, N) with unit column stride"); + TORCH_CHECK(index.is_cuda() && index.device() == x.device() && index.scalar_type() == at::kLong && + index.dim() == 1 && index.is_contiguous() && index.numel() > 0, + "index must be a non-empty contiguous int64 tensor on ", x.device()); + TORCH_CHECK(x.size(0) % index.size(0) == 0, "rows must be a multiple of S"); + return Gate{gate.data_ptr(), gate.stride(0), index.data_ptr(), index.size(0), gate.size(0)}; +} + +void check_act(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device() && t.scalar_type() == like.scalar_type(), + name, " must match the residual's device and dtype"); + TORCH_CHECK(t.sizes() == like.sizes() && t.is_contiguous(), name, + " must be contiguous with the residual's shape"); +} + +unsigned row_blocks(int64_t rows) { + return static_cast(std::min(rows, 1 << 30)); +} + +bool can_vectorize(const torch::Tensor& gate, int64_t n, std::initializer_list acts) { + const int64_t vec = 16 / gate.element_size(); + if (n % vec != 0 || (gate.stride(0) * gate.element_size()) % 16 != 0) return false; + if (reinterpret_cast(gate.data_ptr()) % 16 != 0) return false; + for (const auto* t : acts) { + if (reinterpret_cast(t->data_ptr()) % 16 != 0) return false; + } + return true; +} + +} // namespace + +// residual, y: (M, N); gate: (R, N) view; index: (S,), M % S == 0. +torch::Tensor h3_gate_residual_forward(torch::Tensor residual, torch::Tensor y, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(residual.is_cuda() && residual.dim() == 2 && residual.is_contiguous() && + residual.numel() > 0, + "residual must be a non-empty contiguous (M, N) CUDA tensor"); + check_act(y, "y", residual); + const Gate g = make_gate(gate, index, residual); + const c10::cuda::CUDAGuard guard(residual.device()); + auto out = torch::empty_like(residual); + auto stream = at::cuda::getCurrentCUDAStream(); + const bool vec = can_vectorize(gate, residual.size(1), {&residual, &y, &out}); + H3_DISPATCH(residual.scalar_type(), "h3_gate_residual_forward", [&] { + gate_residual_rows_kernel<<>>( + reinterpret_cast(residual.data_ptr()), reinterpret_cast(y.data_ptr()), + reinterpret_cast(out.data_ptr()), residual.size(0), residual.size(1), g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Returns {d_y, d_gate (R, N) in the gate dtype}. d_residual is the incoming grad. +std::vector h3_gate_residual_backward(torch::Tensor grad, torch::Tensor y, + torch::Tensor gate, torch::Tensor index, + torch::Tensor sorted_pos, + torch::Tensor tile_begin, + torch::Tensor tile_end, + torch::Tensor seg_first_tile) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + const Gate g = make_gate(gate, index, grad); + for (const auto* t : {&sorted_pos, &tile_begin, &tile_end, &seg_first_tile}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 tensors on ", grad.device()); + } + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + TORCH_CHECK(segments == gate.size(0), "segments must equal the gate rows"); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dy = torch::empty_like(y); + auto partial = torch::empty({std::max(tiles, 1), n}, grad.options().dtype(at::kFloat)); + auto dgate = torch::empty({segments, n}, grad.options()); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward", [&] { + const T* gp = reinterpret_cast(grad.data_ptr()); + const T* yp = reinterpret_cast(y.data_ptr()); + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + gp, nullptr, reinterpret_cast(dy.data_ptr()), rows, n, g, vec); + if (tiles > 0) { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, stream>>>( + gp, yp, sorted_pos.data_ptr(), tile_begin.data_ptr(), + tile_end.data_ptr(), partial.data_ptr(), n); + } + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, stream>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(dgate.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dy, dgate}; +} + +// Sequence parallel: the WS1 d_gate tile partials over explicit row lists of +// (grad, y) rows -> (tiles, N) FP32. +torch::Tensor h3_gate_grad_partials(torch::Tensor grad, torch::Tensor y, torch::Tensor rows, + torch::Tensor tile_begin, torch::Tensor tile_end) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + check_act(y, "y", grad); + for (const auto* t : {&rows, &tile_begin, &tile_end}) { + TORCH_CHECK(t->is_cuda() && t->device() == grad.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on grad's device"); + } + TORCH_CHECK(tile_begin.numel() == tile_end.numel(), "tile_begin and tile_end differ in length"); + const int64_t n = grad.size(1); + const int64_t tiles = tile_begin.numel(); + const c10::cuda::CUDAGuard guard(grad.device()); + auto partial = torch::empty({tiles, n}, grad.options().dtype(at::kFloat)); + if (tiles == 0) return partial; + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(grad.scalar_type(), "h3_gate_grad_partials", [&] { + gate_grad_partial_kernel<<(tiles), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(y.data_ptr()), + rows.data_ptr(), tile_begin.data_ptr(), tile_end.data_ptr(), + partial.data_ptr(), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return partial; +} + +// The WS1 d_gate fold: each segment's tiles in ascending order, cast once to dtype. +torch::Tensor h3_gate_grad_fold(torch::Tensor partial, torch::Tensor seg_first_tile, + c10::ScalarType dtype) { + TORCH_CHECK(partial.is_cuda() && partial.dim() == 2 && partial.is_contiguous() && + partial.scalar_type() == at::kFloat, + "partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(seg_first_tile.is_cuda() && seg_first_tile.device() == partial.device() && + seg_first_tile.scalar_type() == at::kLong && seg_first_tile.dim() == 1 && + seg_first_tile.is_contiguous(), + "seg_first_tile must be a contiguous int64 CUDA tensor on partial's device"); + const int64_t n = partial.size(1); + const int64_t segments = seg_first_tile.numel() - 1; + const c10::cuda::CUDAGuard guard(partial.device()); + auto out = torch::empty({segments, n}, partial.options().dtype(dtype)); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(dtype, "h3_gate_grad_fold", [&] { + gate_grad_fold_kernel<<(segments), col_blocks), threads, 0, + at::cuda::getCurrentCUDAStream()>>>( + partial.data_ptr(), seg_first_tile.data_ptr(), + reinterpret_cast(out.data_ptr()), n); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} + +// Row-local part of the WS1 backward only: d_y = cast(grad * gate[index]). +torch::Tensor h3_gate_residual_backward_dy(torch::Tensor grad, torch::Tensor gate, + torch::Tensor index) { + TORCH_CHECK(grad.is_cuda() && grad.dim() == 2 && grad.is_contiguous(), + "grad must be a contiguous (M, N) CUDA tensor"); + const Gate g = make_gate(gate, index, grad); + const int64_t rows = grad.size(0); + const int64_t n = grad.size(1); + const c10::cuda::CUDAGuard guard(grad.device()); + auto dy = torch::empty_like(grad); + H3_DISPATCH(grad.scalar_type(), "h3_gate_residual_backward_dy", [&] { + const bool vec = can_vectorize(gate, n, {&grad, &dy}); + gate_residual_rows_kernel<<>>( + reinterpret_cast(grad.data_ptr()), nullptr, reinterpret_cast(dy.data_ptr()), + rows, n, g, vec); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dy; +} diff --git a/csrc/cuda/h3/rmsnorm_modulate.cu b/csrc/cuda/h3/rmsnorm_modulate.cu new file mode 100644 index 000000000..f9304ab13 --- /dev/null +++ b/csrc/cuda/h3/rmsnorm_modulate.cu @@ -0,0 +1,593 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 RMSNorm and fused AdaLN modulation (RFC #420 `h3_rmsnorm`). +// +// n[r] = cast(w * (rstd[r] * x[r])) rstd = rsqrtf(sum(x^2) / N + eps) +// out[r] = cast(cast(n[r] * cast(1 + scale[i])) + shift[i]) i = row_index[r % S] +// +// Forward contract h3-rmsnorm-v1: the statistics replay PyTorch's +// vectorized_layer_norm_kernel (torch cf30153): one +// (32, 4) block per row, 4-element vectors, thread t sums vectors t, t+128, ... +// in order, a shuffle-down tree (16..1), a cross-warp tree, sum / N, rsqrtf, +// then w * (rstd * x) and one cast. nn.RMSNorm therefore matches bitwise, and +// each modulation step rounds to the tensor dtype exactly where the eager +// `n * (1.0 + scale) + shift` expression does. Rows are independent: the result +// does not depend on batch size, position or the other rows. +// +// Backward (FP32, one cast per output, no atomics): +// d_n = g * cast(1 + scale) (g without modulation) +// dx = rstd * w * d_n - x * rstd^3 * sum(w * d_n * x) / N (row-local) +// dweight = sum_r d_n * x * rstd over fixed 256-row tiles, tiles folded in order +// dshift = sum_{r -> i} g, dscale = sum_{r -> i} g * n (sorted tiles) +// +// The tile partials and the folds are also bound separately (`*_partials`, +// `*_fold_*`) so a sequence-parallel caller can compute each WS1 tile on the +// rank that holds its rows and fold the gathered partials in WS1 order. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +constexpr int kVec = 4; // PyTorch's layer-norm vec_size +constexpr int kWarps = 4; // num_threads() / warp_size = 128 / 32 +constexpr int kRowTile = 256; // dweight row tile + +__device__ __forceinline__ float to_f(float v) { return v; } +__device__ __forceinline__ float to_f(__nv_bfloat16 v) { return __bfloat162float(v); } +__device__ __forceinline__ float to_f(__half v) { return __half2float(v); } + +template +__device__ __forceinline__ T from_f(float v); +template <> +__device__ __forceinline__ float from_f(float v) { return v; } +template <> +__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); } +template <> +__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); } + +// Round to T and back: the value an eager T-dtype op would produce. +template +__device__ __forceinline__ float round_t(float v) { return to_f(from_f(v)); } + +// PyTorch's rms statistics: sequential per thread, shuffle-down, cross-warp tree. +// Returns sum(v^2) / N on every thread. +template +__device__ __forceinline__ float row_mean_square(const T* __restrict__ row, int n, float* buf) { + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const float v = to_f(row[i * kVec + k]); + s = s + v * v; + } + } + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + if (thrx == 0) buf[0] = s / float(n); + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +// Same tree for an arbitrary per-element FP32 term; returns the row sum. +__device__ __forceinline__ float row_sum_tree(float s, float* buf) { + for (int off = 16; off > 0; off >>= 1) s = s + __shfl_down_sync(0xffffffffu, s, off); + for (int off = blockDim.y / 2; off > 0; off /= 2) { + if (threadIdx.x == 0 && threadIdx.y >= off && threadIdx.y < 2 * off) buf[threadIdx.y - off] = s; + __syncthreads(); + if (threadIdx.x == 0 && threadIdx.y < off) s = s + buf[threadIdx.y]; + __syncthreads(); + } + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + if (thrx == 0) buf[0] = s; + __syncthreads(); + const float out = buf[0]; + __syncthreads(); + return out; +} + +struct Modulation { + const void* shift; // (R, N) rows with stride `row_stride` (elements), or nullptr + const void* scale; + int64_t row_stride; + int64_t num_rows; + const int64_t* index; // (S,) table row per sequence position + int64_t seq; // S: row r of x uses index[r % S] +}; + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_fwd_kernel(const T* __restrict__ x, const T* __restrict__ w, T* __restrict__ y, + float* __restrict__ rstd_out, int n, float eps, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const float rstd = rsqrtf(row_mean_square(xr, n, buf) + eps); + const T* shift = nullptr; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + shift = static_cast(mod.shift) + i * mod.row_stride; + scale = static_cast(mod.scale) + i * mod.row_stride; + } + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + float v = round_t(__fmul_rn(to_f(w[j]), __fmul_rn(rstd, to_f(xr[j])))); + if (shift != nullptr) { + // Separately rounded ops (no FMA contraction), as the eager expression. + const float t1 = round_t(__fadd_rn(1.0f, to_f(scale[j]))); + v = round_t(__fadd_rn(round_t(__fmul_rn(v, t1)), to_f(shift[j]))); + } + y[r * n + j] = from_f(v); + } + } + if (thrx == 0) rstd_out[r] = rstd; +} + +template +__device__ __forceinline__ float d_norm_out(const T* g_row, const T* scale, int j) { + const float g = to_f(g_row[j]); + return scale == nullptr ? g : g * round_t(1.0f + to_f(scale[j])); +} + +template +__global__ void __launch_bounds__(32 * kWarps) + rmsnorm_modulate_dx_kernel(const T* __restrict__ g, const T* __restrict__ x, + const T* __restrict__ w, const float* __restrict__ rstd, + T* __restrict__ dx, int n, Modulation mod) { + __shared__ float buf[kWarps]; + const int64_t r = blockIdx.x; + const T* xr = x + r * n; + const T* gr = g + r * n; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } + const int numx = blockDim.x * blockDim.y; + const int thrx = threadIdx.x + threadIdx.y * blockDim.x; + float s = 0.0f; + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + s = fmaf(to_f(w[j]) * d_norm_out(gr, scale, j), to_f(xr[j]), s); + } + } + const float dot = row_sum_tree(s, buf); + const float rs = rstd[r]; + const float coef = rs * rs * rs * dot / float(n); + for (int i = thrx; i < n / kVec; i += numx) { +#pragma unroll + for (int k = 0; k < kVec; ++k) { + const int j = i * kVec + k; + dx[r * n + j] = from_f(rs * to_f(w[j]) * d_norm_out(gr, scale, j) - to_f(xr[j]) * coef); + } + } +} + +// partial[tile, j] = sum over rows of the tile (ascending) of d_n * x * rstd. +// Tiles are the fixed 256-row blocks, or (sequence parallel) explicit row lists +// rows[tile_begin[tile] .. tile_end[tile]) that hold the same rows in the same order. +template +__global__ void rmsnorm_dweight_partial_kernel(const T* __restrict__ g, const T* __restrict__ x, + const float* __restrict__ rstd, + float* __restrict__ partial, int64_t rows, int n, + Modulation mod, + const int64_t* __restrict__ row_list = nullptr, + const int64_t* __restrict__ tile_begin = nullptr, + const int64_t* __restrict__ tile_end = nullptr) { + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= n) return; + const int64_t tile = blockIdx.x; + const int64_t p0 = row_list != nullptr ? tile_begin[tile] : tile * kRowTile; + const int64_t p1 = row_list != nullptr ? tile_end[tile] : min(p0 + kRowTile, rows); + float acc = 0.0f; + for (int64_t p = p0; p < p1; ++p) { + const int64_t r = row_list != nullptr ? row_list[p] : p; + const T* scale = nullptr; + if (mod.index != nullptr) { + const int64_t i = mod.index[r % mod.seq]; + CUDA_KERNEL_ASSERT(i >= 0 && i < mod.num_rows); + scale = static_cast(mod.scale) + i * mod.row_stride; + } + acc = fmaf(d_norm_out(g + r * n, scale, j), to_f(x[r * n + j]) * rstd[r], acc); + } + partial[static_cast(blockIdx.x) * n + j] = acc; +} + +template +__global__ void fold_tiles_kernel(const float* __restrict__ partial, T* __restrict__ out, + int64_t tiles, int n) { + const int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n) return; + float acc = 0.0f; + for (int64_t t = 0; t < tiles; ++t) acc += partial[t * n + j]; + out[j] = from_f(acc); +} + +// Table gradient tiles: positions sorted by (index, position); columns j < n +// accumulate g (shift), columns j >= n accumulate g * n_out (scale). +template +__global__ void modulation_grad_partial_kernel( + const T* __restrict__ g, const T* __restrict__ x, const T* __restrict__ w, + const float* __restrict__ rstd, const int64_t* __restrict__ sorted_pos, + const int64_t* __restrict__ tile_begin, const int64_t* __restrict__ tile_end, + float* __restrict__ partial, int n) { + const int64_t tile = blockIdx.x; + const int j2 = blockIdx.y * blockDim.x + threadIdx.x; + if (j2 >= 2 * n) return; + const bool is_scale = j2 >= n; + const int j = is_scale ? j2 - n : j2; + const float wj = to_f(w[j]); + float acc = 0.0f; + for (int64_t p = tile_begin[tile]; p < tile_end[tile]; ++p) { + const int64_t r = sorted_pos[p]; + const float gv = to_f(g[r * n + j]); + if (is_scale) { + const float nv = round_t(wj * (rstd[r] * to_f(x[r * n + j]))); + acc = fmaf(gv, nv, acc); + } else { + acc += gv; + } + } + partial[tile * 2 * n + j2] = acc; +} + +__global__ void fold_segments_kernel(const float* __restrict__ partial, + const int64_t* __restrict__ seg_first_tile, + float* __restrict__ out, int width) { + const int64_t seg = blockIdx.x; + const int j = blockIdx.y * blockDim.x + threadIdx.x; + if (j >= width) return; + float acc = 0.0f; + for (int64_t t = seg_first_tile[seg]; t < seg_first_tile[seg + 1]; ++t) acc += partial[t * width + j]; + out[seg * width + j] = acc; +} + +void check_rows(const torch::Tensor& t, const char* name, const torch::Tensor& like) { + TORCH_CHECK(t.is_cuda() && t.device() == like.device(), name, " must be on ", like.device()); + TORCH_CHECK(t.dim() == 2 && t.is_contiguous(), name, " must be a contiguous 2-D tensor"); + TORCH_CHECK(t.scalar_type() == like.scalar_type(), name, " must have dtype ", like.scalar_type()); +} + +Modulation make_modulation(const c10::optional& shift, + const c10::optional& scale, + const c10::optional& index, const torch::Tensor& x, + int64_t n) { + Modulation mod{nullptr, nullptr, 0, 0, nullptr, 1}; + if (!index.has_value()) { + TORCH_CHECK(!shift.has_value() && !scale.has_value(), "shift/scale need a row index"); + return mod; + } + TORCH_CHECK(shift.has_value() && scale.has_value(), "modulation needs both shift and scale"); + const auto& sh = *shift; + const auto& sc = *scale; + const auto& ix = *index; + for (const auto* t : {&sh, &sc}) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == x.scalar_type(), + "shift/scale must be CUDA tensors with x's dtype"); + TORCH_CHECK(t->dim() == 2 && t->size(1) == n && t->stride(1) == 1, + "shift/scale must be (R, N) with unit column stride"); + } + TORCH_CHECK(sh.size(0) == sc.size(0) && sh.stride(0) == sc.stride(0), + "shift and scale must be views of the same table layout"); + TORCH_CHECK(ix.is_cuda() && ix.device() == x.device() && ix.scalar_type() == at::kLong && + ix.dim() == 1 && ix.is_contiguous() && ix.numel() > 0, + "row index must be a non-empty contiguous int64 tensor on ", x.device()); + TORCH_CHECK(x.size(0) % ix.size(0) == 0, "x rows (", x.size(0), ") must be a multiple of S (", + ix.size(0), ")"); + mod.shift = sh.data_ptr(); + mod.scale = sc.data_ptr(); + mod.row_stride = sh.stride(0); + mod.num_rows = sh.size(0); + mod.index = ix.data_ptr(); + mod.seq = ix.size(0); + return mod; +} + +void check_xw(const torch::Tensor& x, const torch::Tensor& w) { + TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.is_contiguous(), "x must be a contiguous (M, N) CUDA tensor"); + TORCH_CHECK(x.size(0) > 0, "x must have at least one row"); + TORCH_CHECK(w.is_cuda() && w.device() == x.device() && w.dim() == 1 && + w.size(0) == x.size(1) && w.is_contiguous() && + w.scalar_type() == x.scalar_type(), + "weight must be a contiguous (N,) CUDA tensor with x's dtype and device"); + TORCH_CHECK(x.size(1) > 0, "x must have at least one column"); + TORCH_CHECK(x.size(1) % kVec == 0, "N=", x.size(1), " must be a multiple of ", kVec, + " (PyTorch's vectorized RMSNorm path)"); +} + +// float32 / float16 / bfloat16 only (float64 has no PyTorch-matching vector path here). +#define H3_DISPATCH(TYPE, NAME, ...) \ + [&] { \ + switch (TYPE) { \ + case at::kFloat: { using T = float; __VA_ARGS__(); break; } \ + case at::kHalf: { using T = __half; __VA_ARGS__(); break; } \ + case at::kBFloat16: { using T = __nv_bfloat16; __VA_ARGS__(); break; } \ + default: TORCH_CHECK(false, NAME, ": unsupported dtype ", TYPE); \ + } \ + }() + +} // namespace + +// x (M, N); optional shift/scale (R, N) strided views and index (S,), M % S == 0. +std::vector h3_rmsnorm_forward(torch::Tensor x, torch::Tensor weight, double eps, + c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + TORCH_CHECK(x.scalar_type() != at::kDouble, "float64 is not supported"); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + const c10::cuda::CUDAGuard guard(x.device()); + auto y = torch::empty_like(x); + auto rstd = torch::empty({x.size(0)}, x.options().dtype(at::kFloat)); + auto stream = at::cuda::getCurrentCUDAStream(); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_forward", [&] { + rmsnorm_modulate_fwd_kernel<<(x.size(0)), dim3(32, kWarps), 0, stream>>>( + reinterpret_cast(x.data_ptr()), reinterpret_cast(weight.data_ptr()), + reinterpret_cast(y.data_ptr()), rstd.data_ptr(), static_cast(n), + static_cast(eps), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {y, rstd}; +} + +// Returns {dx, dweight} and, with modulation, {d_shift, d_scale} as FP32 (R, N). +std::vector h3_rmsnorm_backward( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, c10::optional sorted_pos, + c10::optional tile_begin, c10::optional tile_end, + c10::optional seg_first_tile) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && + rstd.scalar_type() == at::kFloat && rstd.dim() == 1 && + rstd.is_contiguous() && rstd.size(0) == x.size(0), + "rstd must be contiguous float32 (M,) statistics on x's device"); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + if (mod.index != nullptr) { + TORCH_CHECK(sorted_pos.has_value() && tile_begin.has_value() && tile_end.has_value() && + seg_first_tile.has_value(), + "modulated backward needs the sorted segment tiles"); + const std::pair tiles[] = {{&*sorted_pos, "sorted_pos"}, + {&*tile_begin, "tile_begin"}, + {&*tile_end, "tile_end"}, + {&*seg_first_tile, "seg_first_tile"}}; + for (const auto& [t, name] : tiles) { + TORCH_CHECK(t->is_cuda() && t->device() == x.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + name, ": tile metadata must be contiguous int64 tensors on ", x.device()); + } + TORCH_CHECK(sorted_pos->size(0) == rows, "sorted_pos must have M entries"); + TORCH_CHECK(seg_first_tile->size(0) == shift->size(0) + 1, + "seg_first_tile must have R + 1 entries"); + TORCH_CHECK(tile_begin->size(0) == tile_end->size(0), + "tile_begin and tile_end must have the same number of entries"); + } else { + TORCH_CHECK(!sorted_pos.has_value() && !tile_begin.has_value() && !tile_end.has_value() && + !seg_first_tile.has_value(), + "sorted segment tiles require modulation"); + } + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + auto dx = torch::empty_like(x); + const int64_t tiles = (rows + kRowTile - 1) / kRowTile; + auto partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + auto dweight = torch::empty_like(weight); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward", [&] { + const T* g = reinterpret_cast(grad.data_ptr()); + const T* xp = reinterpret_cast(x.data_ptr()); + const T* wp = reinterpret_cast(weight.data_ptr()); + rmsnorm_modulate_dx_kernel<<(rows), dim3(32, kWarps), 0, stream>>>( + g, xp, wp, rstd.data_ptr(), reinterpret_cast(dx.data_ptr()), + static_cast(n), mod); + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, 0, + stream>>>(g, xp, rstd.data_ptr(), + partial.data_ptr(), rows, + static_cast(n), mod); + fold_tiles_kernel<<>>( + partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), tiles, + static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (mod.index == nullptr) return {dx, dweight}; + + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t seg_tiles = tile_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({std::max(seg_tiles, 1), width}, x.options().dtype(at::kFloat)); + auto table_grad = torch::empty({segments, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_grad", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + sorted_pos->data_ptr(), tile_begin->data_ptr(), + tile_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( + seg_partial.data_ptr(), seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dx, dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +namespace { + +void check_tiles(const torch::Tensor& rows, const torch::Tensor& begin, const torch::Tensor& end, + const torch::Tensor& like) { + for (const auto* t : {&rows, &begin, &end}) { + TORCH_CHECK(t->is_cuda() && t->device() == like.device() && t->scalar_type() == at::kLong && + t->dim() == 1 && t->is_contiguous(), + "tile metadata must be contiguous int64 CUDA tensors on x's device"); + } + TORCH_CHECK(begin.numel() == end.numel(), "tile_begin and tile_end differ in length"); +} + +} // namespace + +// Sequence parallel: the dweight and (with modulation) table-gradient tile +// partials of the WS1 backward, for explicit tiles over the rows of x. `index` +// gives each row's table row (one entry per row of x). Returns +// {dweight_partial (tiles, N)} or {dweight_partial, table_partial (seg_tiles, 2N)}. +std::vector h3_rmsnorm_backward_partials( + torch::Tensor grad, torch::Tensor x, torch::Tensor weight, torch::Tensor rstd, + c10::optional shift, c10::optional scale, + c10::optional index, torch::Tensor dw_rows, torch::Tensor dw_begin, + torch::Tensor dw_end, c10::optional seg_rows, + c10::optional seg_begin, c10::optional seg_end) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + check_tiles(dw_rows, dw_begin, dw_end, x); + const int64_t rows = x.size(0); + const int64_t n = x.size(1); + const Modulation mod = make_modulation(shift, scale, index, x, n); + TORCH_CHECK(mod.index == nullptr || mod.seq == rows, "index needs one entry per row of x"); + const c10::cuda::CUDAGuard guard(x.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + const int64_t tiles = dw_begin.numel(); + auto dw_partial = torch::empty({tiles, n}, x.options().dtype(at::kFloat)); + if (tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_partials", [&] { + rmsnorm_dweight_partial_kernel<<(tiles), col_blocks), threads, + 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + rstd.data_ptr(), dw_partial.data_ptr(), rows, static_cast(n), mod, + dw_rows.data_ptr(), dw_begin.data_ptr(), dw_end.data_ptr()); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + if (mod.index == nullptr) return {dw_partial}; + + TORCH_CHECK(seg_rows.has_value() && seg_begin.has_value() && seg_end.has_value(), + "modulated partials need the segment tiles"); + check_tiles(*seg_rows, *seg_begin, *seg_end, x); + const int64_t seg_tiles = seg_begin->numel(); + const int64_t width = 2 * n; + auto seg_partial = torch::empty({seg_tiles, width}, x.options().dtype(at::kFloat)); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + if (seg_tiles > 0) { + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_modulation_partials", [&] { + modulation_grad_partial_kernel + <<(seg_tiles), wide_blocks), threads, 0, stream>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + seg_rows->data_ptr(), seg_begin->data_ptr(), + seg_end->data_ptr(), seg_partial.data_ptr(), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + } + return {dw_partial, seg_partial}; +} + +// The WS1 folds: dweight = ascending fold of all tiles, cast to weight's dtype; +// table rows = ascending fold of each segment's tiles, FP32 (R, 2N). +std::vector h3_rmsnorm_fold_partials(torch::Tensor dw_partial, torch::Tensor weight, + c10::optional seg_partial, + c10::optional seg_first_tile) { + TORCH_CHECK(dw_partial.is_cuda() && dw_partial.dim() == 2 && dw_partial.is_contiguous() && + dw_partial.scalar_type() == at::kFloat, + "dw_partial must be a contiguous float32 (tiles, N) CUDA tensor"); + TORCH_CHECK(weight.is_cuda() && weight.device() == dw_partial.device() && weight.dim() == 1 && + weight.size(0) == dw_partial.size(1), + "weight must be (N,) on dw_partial's device"); + const int64_t n = dw_partial.size(1); + const c10::cuda::CUDAGuard guard(dw_partial.device()); + auto stream = at::cuda::getCurrentCUDAStream(); + const int threads = 256; + const unsigned col_blocks = static_cast((n + threads - 1) / threads); + auto dweight = torch::empty_like(weight); + H3_DISPATCH(weight.scalar_type(), "h3_rmsnorm_fold_partials", [&] { + fold_tiles_kernel<<>>( + dw_partial.data_ptr(), reinterpret_cast(dweight.data_ptr()), + dw_partial.size(0), static_cast(n)); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (!seg_partial.has_value()) return {dweight}; + + TORCH_CHECK(seg_first_tile.has_value() && seg_first_tile->is_cuda() && + seg_first_tile->device() == dw_partial.device() && + seg_first_tile->scalar_type() == at::kLong && seg_first_tile->is_contiguous(), + "segment fold needs seg_first_tile (int64, on dw_partial's device)"); + TORCH_CHECK(seg_partial->is_cuda() && seg_partial->device() == dw_partial.device() && + seg_partial->dim() == 2 && seg_partial->is_contiguous() && + seg_partial->scalar_type() == at::kFloat && seg_partial->size(1) == 2 * n, + "seg_partial must be a contiguous float32 (tiles, 2N) CUDA tensor"); + const int64_t segments = seg_first_tile->numel() - 1; + const int64_t width = 2 * n; + auto table_grad = torch::empty({segments, width}, dw_partial.options()); + const unsigned wide_blocks = static_cast((width + threads - 1) / threads); + fold_segments_kernel<<(segments), wide_blocks), threads, 0, stream>>>( + seg_partial->data_ptr(), seg_first_tile->data_ptr(), + table_grad.data_ptr(), static_cast(width)); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {dweight, table_grad.narrow(1, 0, n), table_grad.narrow(1, n, n)}; +} + +// Row-local part of the WS1 backward only: dx (sequence-parallel callers +// compute the cross-row reductions with h3_rmsnorm_backward_partials). +torch::Tensor h3_rmsnorm_backward_dx(torch::Tensor grad, torch::Tensor x, torch::Tensor weight, + torch::Tensor rstd, c10::optional shift, + c10::optional scale, + c10::optional index) { + check_xw(x, weight); + check_rows(grad, "grad", x); + TORCH_CHECK(grad.sizes() == x.sizes(), "grad must match x"); + TORCH_CHECK(rstd.is_cuda() && rstd.device() == x.device() && rstd.scalar_type() == at::kFloat && + rstd.numel() == x.size(0), + "rstd must be the forward's float32 (M,) statistics on x's device"); + const Modulation mod = make_modulation(shift, scale, index, x, x.size(1)); + const c10::cuda::CUDAGuard guard(x.device()); + auto dx = torch::empty_like(x); + H3_DISPATCH(x.scalar_type(), "h3_rmsnorm_backward_dx", [&] { + rmsnorm_modulate_dx_kernel<<(x.size(0)), dim3(32, kWarps), 0, + at::cuda::getCurrentCUDAStream()>>>( + reinterpret_cast(grad.data_ptr()), reinterpret_cast(x.data_ptr()), + reinterpret_cast(weight.data_ptr()), rstd.data_ptr(), + reinterpret_cast(dx.data_ptr()), static_cast(x.size(1)), mod); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return dx; +} diff --git a/csrc/cuda/h3/timestep_sinusoid.cu b/csrc/cuda/h3/timestep_sinusoid.cu new file mode 100644 index 000000000..2acfd516c --- /dev/null +++ b/csrc/cuda/h3/timestep_sinusoid.cu @@ -0,0 +1,93 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// MiniMax-H3 sinusoidal timestep features (RFC #420 `timestep_sinusoid_h3`). +// +// freq[k] = expf((c * k) * (1 / half)) c = (float)(-ln(max_period)) +// arg[t, k] = t[t] * freq[k] +// out[t] = [cosf(arg) | sinf(arg)] (T, 2 * half) FP32 +// +// Every value is produced by the same FP32 operation sequence as diffusers' +// get_timestep_embedding on CUDA (scalar * arange, multiply by the reciprocal +// of the CPU-scalar divisor, exp, multiply, cos/sin), so the output is meant +// to be bitwise equal to that path. Each element depends only on (t, k): +// batch-, position- and repeat-invariant by construction. Build without +// --use_fast_math: expf/sinf/cosf must stay the precise libdevice functions. +// nvcc defines no macro for fast math, so build_tools/extensions.py lists this +// file in H3_BITWISE_SOURCES, which KERNEL_ALIGN_USE_FAST_MATH=1 builds leave +// out; the op then falls back to the PyTorch reference. + +#include +#include +#include +#include + +#include +#include + +namespace { + +__global__ void h3_timestep_sinusoid_kernel( + const float* __restrict__ timestep, + float* __restrict__ out, + int64_t num_timesteps, + int64_t half, + float neg_log_max_period, + float inv_half) { + const int64_t total = num_timesteps * half; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < total; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t row = idx / half; + const int64_t k = idx - row * half; + // __fmul_rn keeps each product a separately rounded FP32 multiply (no FMA + // contraction), matching the two elementwise torch kernels it replays. + const float exponent = __fmul_rn(__fmul_rn(neg_log_max_period, static_cast(k)), inv_half); + const float freq = expf(exponent); + const float arg = __fmul_rn(timestep[row], freq); + float* out_row = out + row * 2 * half; + out_row[k] = cosf(arg); + out_row[half + k] = sinf(arg); + } +} + +} // namespace + +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor timestep, int64_t num_channels, + double max_period, bool check_range) { + TORCH_CHECK(timestep.is_cuda(), "timestep must be a CUDA tensor"); + TORCH_CHECK(timestep.scalar_type() == at::kFloat, "timestep must be float32, got ", + timestep.scalar_type()); + TORCH_CHECK(timestep.dim() == 1, "timestep must be 1-D (num_timesteps,)"); + TORCH_CHECK(timestep.numel() > 0, "timestep must hold at least one timestep"); + TORCH_CHECK(num_channels > 0 && num_channels % 2 == 0, + "num_channels must be a positive even number, got ", num_channels); + TORCH_CHECK(max_period > 0.0, "max_period must be positive"); + + const c10::cuda::CUDAGuard device_guard(timestep.device()); + auto t = timestep.contiguous(); + // Check at the native boundary, including direct extension calls. The explicit + // opt-out is for already-validated inputs and kernel-only profiling. + if (check_range) { + TORCH_CHECK_VALUE(((t >= 0) & (t <= 1)).all().item(), + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + "unscaled"); + } + const int64_t num_timesteps = t.size(0); + const int64_t half = num_channels / 2; + auto out = torch::empty({num_timesteps, num_channels}, t.options()); + + // torch evaluates `-math.log(max_period) * arange` with the Python double + // rounded to the FP32 opmath type, and `x / half` as x * (1 / half) in FP32. + const float neg_log_max_period = static_cast(-std::log(max_period)); + const float inv_half = 1.0f / static_cast(half); + + const int threads = 256; + const int64_t total = num_timesteps * half; + const int64_t blocks = std::min((total + threads - 1) / threads, 65535); + auto stream = at::cuda::getCurrentCUDAStream(); + h3_timestep_sinusoid_kernel<<(blocks), threads, 0, stream>>>( + t.data_ptr(), out.data_ptr(), num_timesteps, half, neg_log_max_period, + inv_half); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return out; +} diff --git a/docs/.nav.yml b/docs/.nav.yml index 1037edce5..05e5aba93 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -30,6 +30,16 @@ nav: - operators/sampling.md - operators/det-gemm.md - operators/embedding.md + - MiniMax-H3: + - operators/h3-timestep-sinusoid.md + - operators/h3-timestep-mlp.md + - operators/h3-adaln-projection.md + - operators/h3-adaln-row-gather.md + - operators/h3-rmsnorm.md + - operators/h3-adaln-gate-residual.md + - operators/h3-final-adaln-out.md + - operators/h3-tp-adaln-3mod.md + - operators/h3-sp-norm-adaln.md - Developer Guide: - contributing/README.md - Contributor Guide: contributing/contributor-guide.md diff --git a/docs/operators/README.md b/docs/operators/README.md index 184d1d0e4..e7c97d1eb 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,5 +32,23 @@ Every operator page should include: - [Matmul](matmul.md) - [Sampling](sampling.md) - [Token Embedding](embedding.md) +- [MiniMax-H3 Timestep Sinusoid](h3-timestep-sinusoid.md) +- [MiniMax-H3 FP32 Timestep MLP](h3-timestep-mlp.md) +- [MiniMax-H3 AdaLN Projection](h3-adaln-projection.md) +- [MiniMax-H3 AdaLN Row Gather](h3-adaln-row-gather.md) +- [MiniMax-H3 RMSNorm and AdaLN Modulation](h3-rmsnorm.md) +- [MiniMax-H3 AdaLN Gated Residual](h3-adaln-gate-residual.md) +- [MiniMax-H3 Final AdaLN Output](h3-final-adaln-out.md) +- [MiniMax-H3 Tensor-Parallel AdaLN Projection](h3-tp-adaln-3mod.md) +- [MiniMax-H3 Sequence-Parallel Norm and AdaLN](h3-sp-norm-adaln.md) +- [MiniMax-H3 One Block, Forward and Backward](h3-one-block.md) +- [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 Full Transformer and Dual Trajectory](h3-full-transformer.md) +- [MiniMax-H3 Integration Rows](h3-integration.md) +- [MiniMax-H3 Trajectory Artifact](h3-trajectory-artifact.md) +- [MiniMax-H3 Tensor/Sequence/Context Parallel Rows](h3-parallel.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..9a5d815a5 --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,143 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python tools/validation/models/h3_prior_art.py --op gate_residual --out reports/experiments/h3-prior-art-b200/gate_residual.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `910d614a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python tools/validation/operators/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_gate_residual \ + --out reports/experiments/h3-adaln-gate-residual-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md new file mode 100644 index 000000000..e8fc09537 --- /dev/null +++ b/docs/operators/h3-adaln-projection.md @@ -0,0 +1,207 @@ +# MiniMax-H3 Three-Modality AdaLN Projection + +## Summary + +`adaln_projection_3mod` is the per-block `MiniMaxH3AdaLayerNormModulation` +(RFC #420, WS1 step 3). It turns the shared FP32 timestep embedding into the six +modulation tensors of one transformer block, for all three modalities: + +```text +act = silu(temb).to(bf16) SiLU in FP32, one declared cast +table = act @ W.T + b (T, 96768) BF16, W: (96768, 2688) +rows = table.view(3T, 32256) row t * 3 + m, m: 0 video, 1 text, 2 audio +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) # (3T, 5376) each +``` + +Projection output channel `o = m * 32256 + c * 5376 + h` is chunk `c`, hidden index `h` +of modality `m`. The table rows are laid out `[t0m0, t0m1, t0m2, t1m0, ...]`. That is the +layout `adaln_row_gather` addresses with `timestep_indices * 3 + token_tags`. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`, `transformer_blocks.0.adaln_proj.linear` +(BF16 weight and bias). All 50 blocks have the same shape. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_projection_3mod", device="cuda") +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, weight, bias) +table = op.forward_table(temb, weight, bias) # raw (T, 96768) for a fused gather +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp` | `rl_engine._C.h3_det_linear_*` | BF16 weights: contract `h3-det-linear-bf16-mma-v1`; FP32 weights: `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp` | n/a | `forward`: provider path; `forward_fp32`: declared-cast FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `temb` | `(T, D)`, `T >= 1` | float32 | H3: `D = 2688`. BF16 is rejected (RFC probe H7) | +| `weight` | `(6 * H * 3, D)` | bfloat16 (checkpoint) or float32 | First dim must be a multiple of 18. H3: `H = 5376` | +| `bias` | `(6 * H * 3,)` | same as `weight` | | +| outputs | six `(3T, H)` views | weight dtype | Views of one `(T, 18H)` table, as in diffusers | + +The CUDA kernel needs `D` to be a multiple of 8 for BF16 or 4 for FP32, and BF16 needs SM80+. These inputs fail +closed: a BF16 `temb`, mismatched dtypes or shapes, an empty `temb`, and mixed devices. + +## Numerics + +- **Mixed-precision boundary.** SiLU runs in FP32 at `temb`'s precision, and the result is + rounded to BF16 exactly once. These are the provider's own elementwise ops, so the + activation is bitwise equal to diffusers. Casting before the SiLU (probe H7) is + rejected at the API. It is also numerically visible: on block 0 it changes about 52% of + the BF16 outputs. +- **Projection (BF16 weights).** It follows contract `h3-det-linear-bf16-mma-v1` in + `csrc/cuda/h3/det_linear.cu`: + - Each warp owns 16 output columns and 8 input rows. Rows past T are zero, so every + T <= 8 runs the same instruction sequence. A tensor-core column never sees another + column's data. + - K is visited in groups of 16 in ascending order. Each `mma.sync m16n8k16` (BF16 in, + FP32 out) starts from zero, so the tensor core only sums 16 products. Its result is + added to an FP32 running sum with an IEEE add. + - Then the bias is added and the result is rounded to BF16 once. + + A timestep's 18 modulation rows are bitwise independent of the other timesteps in the + call and of their order. Restarting the accumulator for every 16-wide group is what keeps + accuracy at the level of an FP32 FMA chain. Letting the tensor core accumulate across + the whole of K, as cuBLAS does, loses about 0.15% of correctly rounded outputs. FP32 + weights use the FMA path of `h3-det-linear-v1` instead, because FP32 on tensor cores + would mean TF32, which the contract forbids. +- **Golden.** `forward_fp32` computes the SiLU in FP64, rounds it to BF16 at the declared + boundary (the cast is model semantics), and runs the projection in FP64. It returns + FP32 without the final rounding. Its gradient stays in FP64 end to end, because the + cast is applied straight-through (identity VJP). Writing the cast as + `.to(bf16).double()` would make autograd round the golden's own `d_temb` to BF16. +- **Backward.** No cuBLAS and no atomics. `d_act` is computed with the chunked + deterministic `grad @ W`, kept in FP32 through the cast (whose VJP is the identity), + and passed through the FP32 SiLU VJP; it is row-local and batch-invariant. + `dW` and `db` are ascending-row FP32 folds, rounded to BF16 once. Unlike the provider, + `d_temb` is not rounded to BF16 on its way through the cast. + +Measured on a B200 with the pinned block-0 weights and torch 2.13.0+cu130. Fractions are +medians over 20 random draws of `temb`, with T = 4: + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| outputs equal to the correctly rounded golden | 99.982% (worst draw 99.976%) | 99.842% (worst draw 99.787%) | +| CUDA equal to provider | 99.84% | | +| early-cast golden (probe H7) equal to CUDA | 43% | | +| `d_temb` vs FP64 golden (gtest, T = 3) | max abs 2.3e-5 | 6.1e-2 (`d_act` rounded to BF16) | +| `dW`, `db` vs FP64 golden rounded to BF16 | > 99.9% bitwise equal | | + +The contract tolerance is `reduction` / `bfloat16`: atol 5e-2, rtol 2e-2, and atol 1e-1 +for gradients. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_projection_3mod +``` + +B200, pinned BF16 weights (520 MB per block, streamed once per call): + +| T | CUDA op | provider (SiLU + cast + cuBLAS) | CUDA GEMV kernel | cuBLAS kernel | +| --- | --- | --- | --- | --- | +| 1 | 100 µs | 112 µs | 78.3 µs (6.6 TB/s) | 93 µs | +| 2 | 100 µs | 100 µs | 78.4 µs | 79 µs | +| 4 | 101 µs | 101 µs | 79.0 µs | 79 µs | + +The kernel takes the same time for every T <= 8, because it always computes 8 rows. It +matches cuBLAS for T >= 2 and is faster at T = 1. The remaining 22 µs of the op time is +the SiLU and cast kernels plus the Python wrapper. With those included, the op is within +about 1% of the provider for T >= 2. + +An earlier FMA-only BF16 kernel ran at 81, 98 and 166 µs for T = 1, 2 and 4. Its time +grew with T because every extra row adds FMAs per weight element. Packed FP32 FMA +(`__ffma2_rn`) gave identical bits but no speedup. Moving to tensor cores removed the +dependence on T. + +## Evidence + +![adaln_projection_3mod on B200: latency and correctly rounded outputs](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-adaln-projection-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `a0834552` on an +otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 +SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also +records that a timestep's 18 modulation rows are bitwise identical whether it runs alone +or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![adaln_projection vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.png) + +The sizes below are timestep (input row) counts T. `H3AdaLNProjectionCudaOp` +streams its 520 MB weight once per 8 rows; T = 256 and 2048 require additional +weight passes and are outside its intended T ≤ 8 range. This accounts for the +steep increase in the reported CUDA forward + backward times (77.4 and 627 ms). + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers AdaLN projection, BF16 F.linear [plain] | **no** (14214 rows, 3080 sub-batches, 2232 sweep cases) | 3.6e-03 / 3.9e-03 / 716 µs | 3.4e-03 / 3.2e-03 / 728 µs | 3.2e-03 / 3.5e-03 / 2346 µs | +| rl-kernel H3AdaLNProjectionCudaOp | yes | 3.6e-03 / 3.9e-03 / 2276 µs | 3.4e-03 / 3.2e-03 / 77446 µs | 3.2e-03 / 3.5e-03 / 626992 µs | +| diffusers AdaLN projection, BF16 F.linear [vllm] | yes | 3.6e-03 / 3.9e-03 / 681 µs | 3.4e-03 / 3.2e-03 / 722 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang] | yes | 3.6e-03 / 3.9e-03 / 3096 µs | 3.4e-03 / 3.2e-03 / 3690 µs | 3.2e-03 / 3.5e-03 / 14296 µs | +| diffusers AdaLN projection, BF16 F.linear [sglang_ieee] | yes | 3.6e-03 / 3.9e-03 / 3138 µs | 3.4e-03 / 3.2e-03 / 3740 µs | 3.2e-03 / 3.5e-03 / 14325 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_te_native] | **no** (7107 rows, 2056 sub-batches, 1104 sweep cases) | 3.6e-03 / 3.9e-03 / 702 µs | 3.4e-03 / 3.2e-03 / 706 µs | 3.2e-03 / 3.5e-03 / 2330 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton] | yes | 3.6e-03 / 3.9e-03 / 2532 µs | 3.4e-03 / 3.2e-03 / 3204 µs | 3.2e-03 / 3.5e-03 / 13771 µs | +| diffusers AdaLN projection, BF16 F.linear [megatron_triton_ieee] | yes | 3.6e-03 / 3.9e-03 / 2533 µs | 3.4e-03 / 3.2e-03 / 3206 µs | 3.2e-03 / 3.5e-03 / 13773 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-prior-art-b200/adaln_projection.json) +was written from a clean tree at `c88691c` by + +```bash +python tools/validation/models/h3_prior_art.py --op adaln_projection --out reports/experiments/h3-prior-art-b200/adaln_projection.json --megatron-src +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/adaln_projection.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_projection bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png) + +This kernel's BF16 output differs from diffusers in 0.08-0.2% of elements (final-rounding flips of a different, fixed reduction order). Every mode (vLLM's settings, SGLang, Megatron) is bitwise equal to diffusers from T = 4 but not at T = 1 (0.1% differ), and differs from this kernel in 0.1-0.2%. +[`adaln_projection.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr485/reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json) was written from a clean tree at `0e476483` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_projection --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_projection.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_projection.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_projection.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP -> projection +python tools/validation/operators/check_operator.py --op adaln_projection_3mod --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_projection_3mod \ + --out reports/experiments/h3-adaln-projection-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-projection-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- The BF16 path needs SM80 or newer (`mma.sync`). It is rejected at launch on older GPUs. +- T > 8 runs one more weight pass per 8 rows. +- Sharding the 96768-wide projection (`tp_adaln_3mod`) is a separate WS2 row. diff --git a/docs/operators/h3-adaln-row-gather.md b/docs/operators/h3-adaln-row-gather.md new file mode 100644 index 000000000..94e32b711 --- /dev/null +++ b/docs/operators/h3-adaln-row-gather.md @@ -0,0 +1,204 @@ +# MiniMax-H3 AdaLN Row Gather + +## Summary + +`adaln_row_gather` selects every packed-sequence position's six modulation vectors from +one block's AdaLN table (RFC #420, WS1 step 3): + +```text +adaln_indices = timestep_indices * 3 + token_tags (S,) 0 video, 1 text, 2 audio +shift_msa[adaln_indices], scale_msa[...], gate_msa[...], +shift_mlp[...], scale_mlp[...], gate_mlp[...] six (S, 5376) +``` + +`rows` is the `(3T, 6H)` view of the [AdaLN projection](h3-adaln-projection.md) table, and +its six column blocks are the six modulation tensors. `timestep_indices` and `token_tags` +are semantic inputs (RFC #420 §4): they are validated, never clamped. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_row_gather", device="cuda") +outs = op(table.view(-1, 6 * 5376), timestep_indices, token_tags) # six (S, 5376) +outs = op.gather_chunks((shift_msa, ..., gate_mlp), timestep_indices, token_tags) # drop-in +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather.H3AdaLNRowGatherCudaOp` | `rl_engine._C.h3_adaln_row_gather_{forward,backward}` | Forward bitwise equal to `index_select`; deterministic backward | +| PyTorch reference | `rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp` | n/a | `index_select` forward; deterministic per-row FP32 backward | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +The raw diffusers path, including its atomic backward, is kept as +`rl_engine.validation.models.h3_provider.provider_adaln_row_gather` for evidence. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `rows` | `(3T, 6H)` | bf16 / fp16 / fp32 | Unit column stride. Any row stride (16-byte copies when aligned) | +| `timestep_indices` | `(S,)`, `S >= 1` | int64 or int32 | In `[0, T)` | +| `token_tags` | `(S,)` | same as `timestep_indices` | In `{0, 1, 2}` | +| outputs | six `(S, H)` | `rows.dtype` | Contiguous; slices of one `(6, S, H)` buffer | + +These inputs fail closed: out-of-range tags (a fourth modality, probe H2) or timesteps +(an offset past `T`, probe H3), length or dtype mismatches, empty sequences, and tables +that are not 3 rows per timestep or 6 blocks wide. The range check is one host +read-back; pass `check_range=False` when the indices are already validated. + +## Numerics + +- **Forward** is a byte copy: one launch, one block per `(position, chunk)`, using + 16-byte vector copies. It is bitwise equal to diffusers' six `index_select` calls for + every packing, length and index dtype tested (S up to 131072). +- **Backward**: `d_rows[r] = sum over {s : r[s] = r} of grad[s]` is the op's only + reduction. Positions are sorted stably by `(row, position)` and cut into tiles of 256. + Each tile is an ascending FP32 chain from 0, and a row's tiles are left-folded in order + and rounded once. There are no atomics, so the result depends only on the indices and + the gradient. `index_select`'s own backward is an atomic scatter-add in the gradient + dtype (BF16), which is neither deterministic nor accurate. +- **Tolerance class.** The VJP is a segmented reduction, so gtest judges the op as + `reduction`. The forward is asserted bitwise in `tests/models/minimax_h3`. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison (T = 3, S = 4097, H = 5376, BF16) | CUDA | provider (`index_select`) | +| --- | --- | --- | +| Forward vs `index_select` | bitwise equal | | +| `d_rows` repeat-bitwise | yes | no | +| `d_rows` correctly rounded from the FP64 sum | 99.99931% | 7.9% | +| `d_rows` max abs error vs FP64 | 0.25 (1 ULP) | 4.99 | + +### Fused modulation: `H3AdaLNModulationCudaOp` + +`rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation.H3AdaLNModulationCudaOp(temb, weight, bias, +timestep_indices, token_tags)` runs [`adaln_projection_3mod`](h3-adaln-projection.md) and +this gather as one autograd node. Its forward uses the same kernels and is bitwise equal +to calling the two ops in sequence. + +The difference is in the backward. When the ops are called separately, autograd returns +the gather's FP32 segment sum to the projection in the table's BF16 dtype, which rounds it +once more. The fused node feeds the FP32 table gradient straight into the projection +backward. Use it wherever both ops run back to back, as they do in an H3 block. + +### Whole-chain backward + +`tools/validation/models/h3_chain_replay.py --backward` runs on the pinned weights for +T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. It checks six parameter gradients in each +of the 16 cases (the time embedder's `linear_{1,2}.{weight,bias}` and block 0's AdaLN +weight and bias) against an FP64 golden in which every declared cast is straight-through: + +| S >= 257 (72 gradients) | RL-Kernel, separate ops | RL-Kernel, fused | diffusers | +| --- | --- | --- | --- | +| repeat-bitwise | 72 / 72 | 72 / 72 | 0 / 72 | +| time embedder: max abs error / golden max | 5.3e-4 to 2.4e-3 | 6.5e-7 to 4.9e-6 | 3.8e-3 to 1.3e-1, varying run to run | +| AdaLN weight/bias: correctly rounded BF16, worst case | 59.8% | 99.57% | 1.6% | + +At S = 3 all three chains are deterministic. Both RL-Kernel variants reach 1.5e-6 on the +time embedder, against 3.2e-3 for diffusers. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_row_gather +``` + +B200, T = 3, H = 5376, BF16. "write BW" counts the 6·S·H outputs; the table stays in L2. + +| S | CUDA forward | provider forward | CUDA backward | provider backward | +| --- | --- | --- | --- | --- | +| 4097 | 0.07 ms | 0.11 ms | 0.88 ms | 1.28 ms | +| 32768 | 0.33 ms | 0.62 ms | 1.75 ms | 8.04 ms | +| 131072 | 1.17 ms | 2.42 ms | 5.62 ms | 32.45 ms | + +Backward timings and peak memory exclude leaf creation and the forward pass. The CUDA +backward stacks the six output gradients and materialises per-tile FP32 partials, adding +8130 MiB at S = 131072; the provider adds under 2 MiB. Candidate/provider execution order +alternates each iteration and is recorded in the report. + +## Evidence + +![adaln_row_gather on B200: latency and whole-chain gradient accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/figure.png) + +There are two data files: + +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-adaln-row-gather-b200/report.json): op timings, forward + bitwise checks, op-level backward, and the chain-backward cases plotted above. +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-adaln-row-gather-b200/chain_replay.json): the full + stage-wise forward replay and backward replay over T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}. + +`report.json` was regenerated from a clean tree at commit `f020f7d5` with backward-only +timings and alternating execution order, and with the op-backward check's rows and gradients +drawn from a seeded generator. Before, they came from the process's default CUDA generator, +so the fraction depended on how many draws the runner made first (99.99965% at `80e4609`). +`chain_replay.json` was written from a clean tree at commit `d3b37fe1`. Both were regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the +model does; every CUDA pass/fail flag is unchanged. + +## Existing implementations (RFC #420 reuse rule) + +![adaln_row_gather vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers six index_select calls | **no** (param/table grads not repeatable) | 0.0e+00 / 6.0e-02 / 1381 µs | 0.0e+00 / 2.1e-01 / 9753 µs | +| rl-kernel H3AdaLNRowGatherCudaOp | yes | 0.0e+00 / 2.5e-03 / 1106 µs | 0.0e+00 / 3.0e-03 / 2776 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-prior-art-b200/adaln_row_gather.json) +was written from a clean tree at `6c900ae` by + +```bash +python tools/validation/models/h3_prior_art.py --op adaln_row_gather --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/adaln_row_gather.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![adaln_row_gather bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png) + +There is no reuse candidate beyond diffusers' six index_select calls. The forward is bitwise identical to them at every S (1 to 32768). The table gradient matches at S = 1 and differs from S = 257 on: index_select accumulates it with BF16 atomics, this kernel in a fixed FP32 order. +[`adaln_row_gather.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr486/reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json) was written from a clean tree at `46347118` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_row_gather.py tests/models/minimax_h3/test_h3_adaln_modulation.py -v +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # whole chain, forward + backward +python tools/validation/operators/check_operator.py --op adaln_row_gather --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_row_gather \ + --out reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/h3_chain_replay.py --timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 \ + --backward --out reports/experiments/h3-adaln-row-gather-b200/chain_replay.json +``` + +## Known Limitations + +- Like diffusers, the op materialises six `(S, H)` tensors, which is 8 GB at + S = 131072. Fusing the gather into the norm/modulate and gate-residual consumers + (`h3_rmsnorm`, `adaln_gate_residual`) avoids this. Those are separate rows. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-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-checkpoint-mapper-vllm.md b/docs/operators/h3-checkpoint-mapper-vllm.md new file mode 100644 index 000000000..050653650 --- /dev/null +++ b/docs/operators/h3-checkpoint-mapper-vllm.md @@ -0,0 +1,32 @@ +# MiniMax-H3 Checkpoint Mapper and vLLM-Omni Rollout Adapter + +## Summary + +RFC #420 integration rows, validated on H100 against the installed vLLM-Omni name tables: + +| Row | Where | What it fixes | +| --- | --- | --- | +| `h3_checkpoint_mapper` | `models/minimax_h3/checkpoint_mapper.py` | Diffusers <-> vLLM-Omni native names and layouts | +| `vllm_runtime_adapter_h3` | `integrations/engines/rollout/vllm/h3_omni.py` | strict fixed-order kernels behind vLLM-Omni's `qkv_proj` / `out_proj` / `fc1` / `fc2` | + +## Checkpoint mapper + +`transformer_blocks.N` -> `blocks.N`, `token_refiner.refiner_blocks.N` -> `token_refiner.blocks.N`; +`to_q/k/v` -> one `attn.qkv_proj` (`[Q;K;V]` rows); `to_out.0` -> `out_proj`; `ff.net.0.proj` +(value-first) -> `mlp.fc1` (gate-first: halves swapped); `ff.net.2` -> `mlp.fc2`; model-level names +from vLLM-Omni's `fasth3` table. All 638 checkpoint tensors map; the round trip is bitwise; +unmapped names and incomplete q/k/v sets fail closed. A test compares the tables with +`vllm_omni.diffusion.models.minimax_h3.fasth3` when importable (passes with vLLM-Omni installed). + +## vLLM-Omni adapter + +`install_h3_rollout(model, collective=None, tp=1)` replaces the four fused linears' forwards +(returning vLLM's `(output, None)`): `qkv_proj` / `fc1` use `h3-projection-gemm-v1` (a fused column +equals the same column alone, so fused == separate bitwise, tested); `out_proj` / `fc2` use the +canonical row-parallel reduction of [`h3-parallel-v1`](h3-parallel.md). Wrong dtype, non-CUDA +weights or a model without those modules raise in strict mode. Attention, rotary, QK norm, +SwiGLU and AdaLN stay vLLM-Omni's and are reported `external`. `handle.remove()` restores. + +## Evidence + +`tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py`. The name-table test compares with `vllm_omni.diffusion.models.minimax_h3.fasth3` when it is importable and is skipped otherwise. diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md new file mode 100644 index 000000000..d40a49f5f --- /dev/null +++ b/docs/operators/h3-final-adaln-out.md @@ -0,0 +1,175 @@ +# MiniMax-H3 Final AdaLN Output + +## Summary + +`final_adaln_out` is H3's `norm_out` (`MiniMaxH3AdaLayerNormOut`), which runs once after the 50 +blocks (RFC #420, WS1 step 7): + +```text +shift, scale = norm_out.linear(silu(temb).to(bf16)).chunk(2) # (T, 5376) each, shift first +out = norm_out.norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] +``` + +The table has one row per distinct timestep and is indexed by `timestep_indices`, not by +`adaln_indices`. Diffusers then upcasts `out` to the FP32 output heads' dtype. That cast is exact +and is left to the caller. + +## Entry Point + +```python +op = kernel_registry.get_op("final_adaln_out", device="cuda") +out = op(x, norm_out_norm_w, temb, norm_out_linear_w, norm_out_linear_b, timestep_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+) | `rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out.H3FinalAdaLNOutCudaOp` | `rl_engine._C.h3_det_linear_*`, `rl_engine._C.h3_rmsnorm_*` | One autograd node | +| PyTorch reference | `rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp` | n/a | `forward`: diffusers replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Numerics + +- **Projection.** The [AdaLN projection](h3-adaln-projection.md)'s deterministic tensor-core + GEMV (`h3-det-linear-bf16-mma-v1`), applied to `norm_out.linear`. The FP32 SiLU is rounded + once at the declared cast. BF16 `temb` is rejected (probe H7). +- **Norm and modulation.** The [`h3_rmsnorm`](h3-rmsnorm.md) kernel indexed by + `timestep_indices`. Given the same table, it is bitwise equal to diffusers. Overall, about + 0.04% of output elements differ from diffusers by 1 ULP, all traced to the projection's + summation tree versus cuBLAS. Rows are batch- and position-invariant. +- **Backward.** One autograd node, so the table gradient (FP32 segment sums) goes straight into + the projection backward without an extra BF16 rounding: + - `d_temb`, `dW` and `db` are 21–37× closer to FP64 in the recorded run than diffusers, + whose error varies from run to run because it rounds the table gradient to BF16 and accumulates it with + `index_select` atomics; + - `dx` and `d_norm_w` are at the BF16 rounding level for both; + - every gradient is repeat-bitwise. +- **Golden.** FP64, rounding only where the model stores a value in its own dtype: the SiLU + cast, `norm_out.linear`'s BF16 output table, `norm_out.norm`'s BF16 output and the BF16 + `1 + scale`. All four are rounded straight-through for the gradient. The table and `1 + scale` + are shared by every position of a timestep. Without them the golden's `d_norm_w` drifts + systematically with S, and the gtest missed it from S = 257. `norm_out.norm`'s rounding + enters `d_scale`, `dW` and `d_temb` summed over S, and the gtest missed those from S = 1024. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op final_adaln_out +``` + +B200, B = 1, T = 3, pinned `norm_out` weights. Backward timings exclude input +and forward-graph setup: + +| S | CUDA fwd | diffusers fwd | CUDA fwd+bwd | diffusers fwd+bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.19 ms | 0.21 ms | 1.37 ms | 0.83 ms | +| 32768 | 0.42 ms | 0.82 ms | 2.53 ms | 4.45 ms | +| 131072 | 1.28 ms | 2.97 ms | 7.02 ms | 17.61 ms | + +At small S the backward is dominated by fixed setup: the stable sort for the segment sums and the +projection backward. + +The stored timing evidence includes forward graph setup in the columns labeled +`fwd+bwd`; those values are not backward-only measurements. The current benchmark +prepares the graph outside the timed backward region. Historical timings are retained +with their original commit and corrected scope. + +## Evidence + +![final_adaln_out on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/figure.png) + +There are two data files: + +- [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-final-adaln-out-b200/report.json): op timings, the + forward-equality fraction, row invariance and backward accuracy. Regenerated from a clean + tree at `0a2d5589` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The + backward reference keeps only the SiLU and table roundings, so the errors include the BF16 + rounding of `norm(x)` and `1 + scale`. +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the + whole conditioning chain (timestep → … → norm_out) replayed stage by stage over + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean + tree at `50e913fa` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). + +## Existing implementations (RFC #420 reuse rule) + +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay) | **no** (param/table grads not repeatable) | 8.3e-03 / 7.3e-02 / 1021 µs | 8.9e-03 / 2.4e-01 / 4525 µs | +| rl-kernel H3FinalAdaLNOutCudaOp | yes | 8.3e-03 / 5.2e-03 / 1642 µs | 8.9e-03 / 5.7e-03 / 2959 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) +was written from a clean tree at `03da729` by + +```bash +python tools/validation/models/h3_prior_art.py --op final_adaln_out --out reports/experiments/h3-prior-art-b200/final_adaln_out.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/final_adaln_out.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) + +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and every gradient differs as expected for the deterministic FP32 reductions. +[`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json) was written from a clean tree at `4430ef16` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_final_adaln_out.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out +python tools/validation/operators/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 4097 --normalized-dim 5376 --check-grad # also 257, 1024 +python tools/validation/models/h3_evidence.py --op final_adaln_out \ + --out reports/experiments/h3-final-adaln-out-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-final-adaln-out-b200/report.json +``` + +## Known Limitations + +- **FP32 gtest at larger S.** H3 runs this op with BF16 activations and weights, and the BF16 + gtest passes at S = 257, 1024 and 4097. With `--dtype fp32` (every input FP32), the FP32 + reduction tolerance (atol = rtol = 1e-4) misses `d_temb` from S = 257 and `dW` from + S = 1024. Both gradients are FP32 sums over 3·S rows. Where large terms cancel, the + summation error exceeds the 1e-4 absolute floor. The PyTorch reference misses the same + gradients at the same S (`--batch 3`, seed 123, max abs error): + + | S | `d_temb` CUDA | `d_temb` PyTorch | `dW` CUDA | `dW` PyTorch | + | --- | --- | --- | --- | --- | + | 64 | 9.2e-5 | 1.5e-4 | 1.5e-4 | 9.2e-5 | + | 257 | 3.1e-4 ✗ | 5.9e-4 ✗ | 8.5e-4 | 2.4e-4 | + | 1024 | 1.5e-3 ✗ | 2.4e-3 ✗ | 2.0e-3 ✗ | 1.5e-3 ✗ | + | 4097 | 5.9e-3 ✗ | 9.3e-3 ✗ | 5.9e-3 ✗ | 3.9e-3 ✗ | + + ✗ fails the gtest. Unmarked values above 1e-4 pass through the rtol term. In a PyTorch + emulation, accumulating only the projection backward in FP64 clears S = 257 but not + S ≥ 1024, because the segment sums contribute as well. Making every reduction FP64 would + change the `h3_det_linear` kernels shared with `timestep_mlp_fp32` and + `adaln_projection_3mod`. The alternative is an FP32 gradient + tolerance that scales with reduction length. Which of the two to take is an open question + for the maintainers. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-full-transformer.md b/docs/operators/h3-full-transformer.md new file mode 100644 index 000000000..30dbf9ffe --- /dev/null +++ b/docs/operators/h3-full-transformer.md @@ -0,0 +1,55 @@ +# MiniMax-H3 Full Transformer and Dual Trajectory + +## Summary + +RFC #420 rows `ws1_full_h3_transformer` and `ws1_dual_trajectory`, validated on H100 (Hopper) +with the pinned checkpoint (`MiniMaxAI/MiniMax-H3@42ed227`, all 14 shards, sha256-verified, +638 tensors, 33,122,992,896 parameters, one GPU, BF16). + +- `tools/validation/models/h3_full_transformer.py`: the full 33B forward and backward on a + packed input through the RL-Kernel block graph (50 real blocks) **and** the boundary rows + (`audio_input_projection_fp32`, `audio_output_projection_fp32`, `text_context_projection`, + `token_refiner_block`); `proj_in`, `proj_out`, `time_embedder`, `norm_out` stay on the + provider and the report says so (`boundary_bindings`). +- `tools/validation/models/h3_dual_trajectory.py`: a 4-step flow-matching Euler sampler + (video + audio, FP32 latents) run in two layouts - the request alone, and as row 1 of a + batch of 2 - sealed with the trajectory artifact and compared with `drift_report`. + +## Result (H100, `reports/experiments/h3-full-transformer-h100/`) + +| Check | Provider | RL-Kernel | +| --- | --- | --- | +| Full model bitwise repeatable (outputs, input grads, all 638 parameter grads) | yes | **yes** | +| Video / audio output vs provider | - | rel 2.1e-3 / 1.8e-2 | +| Same, provider with another SDPA kernel vs default | rel 1.9e-3 / 1.6e-2 | - | +| Dual trajectory single vs batched (4 steps, 16 tensors) | **16 / 16 differ**, first at step 0 (max abs 0.36) | 8 / 16 differ (max abs 2.9e-6) | +| Same, with `proj_in`/`proj_out` on the fixed-order GEMM (diagnostic) | - | **0 / 16 differ, final latents equal** | + +Reading: the candidate stays inside the spread of independent provider kernel choices +(`provider_variants_vs_provider`). The residual strict drift is entirely the provider's +video `proj_in`/`proj_out` (cuBLAS FP32, row-count dependent); binding them to the same +fixed-order GEMM makes the full 33B sampler bitwise layout-invariant. Those two rows belong +to another contributor, so the diagnostic binding is labelled in the artifact metadata +(`diagnostic_video_projections`) and is not claimed as a row. + +## Limits + +- No FP32 full-model golden (132 GB of weights); agreement is against the provider and its + kernel-choice spread. Per-block FP32 golden evidence is in the one-block and refiner rows. +- Attention, O/gate_up/down GEMMs, SwiGLU and QK norm are still `interim` bindings (see the + operator trace); `time_embedder` and `norm_out` use the provider in this harness. +- Parameter-gradient fingerprints are GPU-side position-weighted checksums, not SHA-256. +- Sampler is the deterministic Euler ODE; the stochastic `cps` transition is a separate row. + +## Hopper evidence chart + +These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description). +Three seeds of the single-vs-batched Euler control (`reports/experiments/h3-full-transformer-h100/dual_trajectory_seed{11,22,33}[_diag].json`, plus the original +`dual_trajectory.json`) and the strict-vs-provider comparison of the 33B forward/backward +(`full_transformer.json`). + +* provider: the two paths differ at every step (max abs 0.25-0.54 on step 0). +* RL-Kernel strict: the first difference is 4.3e-6 to 4.8e-6 and comes from the video `proj_in`/`proj_out`, which are another contributor's rows and + are not bound here; the difference then grows through the sampler on some seeds. +* strict with the video projections bound only as a diagnostic (`--diag-video-projections`): bitwise equal at every step on all three seeds. +* the wall-time panel includes warm-up; steady-state step times are a follow-up. diff --git a/docs/operators/h3-integration.md b/docs/operators/h3-integration.md new file mode 100644 index 000000000..c699255e2 --- /dev/null +++ b/docs/operators/h3-integration.md @@ -0,0 +1,39 @@ +# MiniMax-H3 Integration Rows (runtime modes, overlay, Diffusers adapter, training provider) + +## Summary + +RFC #420 integration rows, validated on H100 against the real verl-omni scheduler: + +| Row | Where | What it fixes | +| --- | --- | --- | +| `h3_runtime_modes` | `models/minimax_h3/runtime.py` | `RL_KERNEL_MODE` strict / audit / auto / off -> bind the RL-Kernel graph or not; strict raises instead of falling back, audit reports, case id `R/R` or `P/P` | +| `h3_ablation_overlay` | same | `AblationOverlay.parse("node:attention,boundary:token_refiner")` keeps parts on provider arithmetic; the run is labelled `ablation(...)`, never `R/R`; unknown names are rejected | +| `diffusers_runtime_adapter_h3` | same (`H3DiffusersAdapter`) | binds `MiniMaxH3Transformer3DModel` inside `with adapter.active():` (forward **and** backward, so checkpoint recomputation sees the same operators) and reports bindings | +| `h3_training_provider` | same (`H3TrainingProvider`) | actor-side `step()` scope + per-sample CPS transition | +| `h3_lora_gradient_gate` | same | LoRA over a strict projection: A/B/dX bitwise repeatable, forward rows independent of batching, base frozen, FP64-close | +| `stochastic_transition_contract` | `models/minimax_h3/transition.py` | the `cps`, `noise_level=0.8` step, layout invariant | + +## Stochastic transition contract + +```text +std = sigma_prev * sin(0.8 * pi / 2) +mean = (x - sigma v) * (1 - sigma_prev) + (x + v (1 - sigma)) * sqrt(sigma_prev^2 - std^2) +x_prev = mean + std * eps log_prob = mean over non-batch dims of -(x_prev - mean)^2 +``` + +- elementwise FP32 in this order; **bitwise equal** to + `FlowMatchSDEDiscreteScheduler.sample_previous_step(sde_type="cps")` (`prev_sample_mean`, tested + against the verl-omni checkout); +- `eps = step_noise(seed, step, sample_id, modality, shape)`: one generator per sample (SHA-256 of + the key), so noise never depends on batch order or composition; +- the log-prob mean is a fixed-order GEMM reduction per sample (the scheduler's `torch.mean` order is + shape dependent; ours agrees to FP32 rounding and is batch invariant); +- `noise_level = 0` (non-selected steps) yields `x_prev == mean`, `log_prob == 0`; scoring a stored + rollout (`prev_sample=`) reproduces the sampling log-prob bitwise. + +Open item for the maintainer: the formulas above are the scheduler's current CPS law; if the +recipe changes the noise law, only `transition.py` and its test change. + +## Evidence + +`tests/models/minimax_h3/test_h3_runtime_integration.py` (runtime modes, overlay, adapter, transition, LoRA gate). The transition tests compare with the verl-omni scheduler when it is importable and are skipped otherwise. 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-parallel.md b/docs/operators/h3-parallel.md new file mode 100644 index 000000000..3ffcacd42 --- /dev/null +++ b/docs/operators/h3-parallel.md @@ -0,0 +1,93 @@ +# MiniMax-H3 Tensor / Sequence / Context Parallel Rows (WS2) + +## Summary + +RFC #420 rows `tp_attention_7168`, `tp_ffn_14336`, `cp_full_attention` (Ulysses), +`cp_metadata_coshard`, `cp_arbitrary_length`, `distributed_packed_io`, and closeouts +`ws2_tp2_tp4_tp8`, `ws2_tp4_cp2`, `ws2_tp4_cp2_sp`. Contract `h3-parallel-v1` +(`rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py`); the block that composes them is +`rl_engine/validation/models/h3_tp.py`. Cross-node reductions use +`rl_engine/distributed/algorithms/hierarchical.py` (`hierarchical_ipc_nccl_tree`). + +Result (H100, 2 nodes x 4 GPUs, block 0 of the pinned checkpoint, 264 packed rows): every rank's +rows of the forward, the input gradient and every weight-gradient shard are **bitwise equal** to the +single-GPU canonical block for TP2, TP4, TP8 (two nodes), TP4 x CP2 (two nodes) and TP4 x CP2 + SP. +Evidence: `reports/experiments/h3-ws2-parallel-h100/*.json`. + +## The canonical reduction + +All cross-rank sums are defined on **8 canonical units** (the divisors of 8 are the TP sizes): + +| Layer | Sharding | Reduction | Rule | +| --- | --- | --- | --- | +| Q/K/V, gate/up forward | output rows (whole heads / halves) | none | the unsharded GEMM's columns (`h3-projection-gemm-v1`) | +| O (7168 -> 5376), down (14336 -> 5376) forward | reduction axis, `8/tp` K-units per rank | K | each K-unit is one fixed-order FP32 GEMM; a rank adds its units as a balanced tree; the collective adds ranks as the same tree; one BF16 rounding | +| O / down backward (`dX`, `dW`) | - | N / tokens | plain fixed-order GEMM, independent of the shard | +| Q/K/V, gate/up backward `dX` | output rows | N | unit `j` = block `j` of every fused weight, summed in the order q,k,v (hidden,gate); tree over units; collective; one rounding | + +The unit of the O projection is 896 K (14 K-tiles) and of down 1792 (28 K-tiles), so a unit never +splits a tile. FP32 partials cross ranks (`DeterministicCollective` fixed tree on a node; +`HierarchicalCollective` across nodes). The row-parallel contract differs from the single-sweep +K reduction of the unsharded O/down GEMM rows (claimed by another contributor): it is the +TP-invariant definition, and TP=1 is that definition, not the sweep. + +## Cross-node collective + +`DeterministicCollective` is CUDA IPC and therefore single node. `HierarchicalCollective(local_size)` +reduces each group of `local_size` ranks with that tree, moves the group partials with NCCL +`all_gather` (exact copies, no NCCL arithmetic) and adds them in one fixed balanced order. For +power-of-two sizes this is the 8-way tree; a unit test checks it equals the single-node tree of the +same world on 4 GPUs for FP32 and BF16 (`local_size` 1 and 2), and the TP8 run above uses it across +two nodes. + +## Context parallelism + +- `SeqShard`: contiguous row slices differing by at most one row, any length >= 0 (including fewer + rows than ranks): `cp_arbitrary_length`. +- `coshard`: hidden, AdaLN row index, token tags, rotary tables are cut with one `SeqShard` so the + metadata follows its rows: `cp_metadata_coshard`. +- `ulysses_attention`: all-to-all sequence <-> head (ragged split sizes), attention on `H/cp` heads + of the full sequence, all-to-all back; both moves are exact, so rows equal the unsharded + attention's if the attention is head independent (the repo's deterministic attention is). Tested at + lengths 1, 2, 7, 130 and in the block at 261 / 263 rows. +- `scatter_packed` / `gather_packed`: exact packed IO, ragged gather by padding: `distributed_packed_io`. +- `seq_sharded_linear`: weight gradient from gathered rows so it equals the unsharded token sweep. + +## Sequence parallelism + +With `sp=True` the residual stream is sharded by rows across the TP group: norm and gate-residual are +row-local, the column-parallel input is all-gathered (exact) and the row-parallel output is +reduce-scattered with the same tree (`rows % tp == 0` required, checked). + +## Tests and evidence + +`tests/models/minimax_h3/test_h3_parallel.py`: 70 tests on 4 H100 (sharding, canonical folding for +TP 2/4/8, accuracy against FP64, row independence, Ulysses ragged, ten block grids with backward, +hierarchical collective). `tools/validation/models/h3_parallel_run.py` runs a grid under `srun` on +any node count and compares against the single-GPU block. + +| Grid | GPUs / nodes | Collective | Forward | dX | dW shards | +| --- | --- | --- | --- | --- | --- | +| TP2 | 2 / 1 | IPC tree | bitwise | bitwise | bitwise | +| TP4 | 4 / 1 | IPC tree | bitwise | bitwise | bitwise | +| TP8 | 8 / 2 | hierarchical | bitwise | bitwise | bitwise | +| TP4 x CP2 | 8 / 2 | IPC tree (TP in node), NCCL all-to-all (CP across) | bitwise | bitwise | not compared (token axis sharded) | +| TP4 x CP2 + SP | 8 / 2 | same | bitwise | bitwise | not compared | + +Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4xCP2+SP 104 +(tiny 264-row case, launch-bound; not a scaling claim). + +## Limits + +- Attention, SwiGLU and QK norm inside the block are the repo's interim operators; the parallel + rows only require them to be row/head independent, which the bitwise results confirm. +- Pipeline parallelism is out of scope (`ws2_pipeline_parallel` is not part of this change). +- dW under CP is exact through `seq_sharded_linear` for non-TP linears; the TP x CP block compares + forward and `dX` (CP shards the token axis of the TP weight gradients). +- `HierarchicalCollective` assumes ranks are packed node-major and node counts of 1/2/4/8. + +## Hopper evidence chart + +These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description). +Per-layout results of the 2 nodes x 4 H100 runs (`reports/experiments/h3-ws2-parallel-h100/*.json`): every layout is bitwise equal to the single-GPU reference +(forward and input gradient); mean per-rank forward time of block 0 is in the figure. 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 new file mode 100644 index 000000000..75dd3d477 --- /dev/null +++ b/docs/operators/h3-rmsnorm.md @@ -0,0 +1,169 @@ +# MiniMax-H3 RMSNorm and AdaLN Modulation + +## Summary + +`h3_rmsnorm` covers every RMSNorm in MiniMax-H3 (RFC #420, WS1 step 4). That is the block +`norm1`/`norm2`, the token-refiner norms, the refiner `final_norm` and `norm_out.norm`. All of +them are `nn.RMSNorm(5376, eps=1e-5)` with a BF16 affine weight. In a transformer block and in +`norm_out`, the normalised rows are immediately modulated by per-row AdaLN parameters, and the +op fuses that step: + +```text +n = rms_norm(x, weight, eps) +out = n * (1.0 + scale[index]) + shift[index] +``` + +`index` is `adaln_indices` in a block (rows of the [projection](h3-adaln-projection.md) table) +and `timestep_indices` in `norm_out`. `shift`/`scale` are `(R, H)` row views of the AdaLN table, +gathered inside the kernel, so the `(S, H)` tensors that +[`adaln_row_gather`](h3-adaln-row-gather.md) materialises are never needed. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("h3_rmsnorm", device="cuda") +n = op(x, weight) # plain RMSNorm +out = op.forward_modulated(x, weight, shift_msa, scale_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM80+, validated on SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp` | `rl_engine._C.h3_rmsnorm_{forward,backward}` | Bitwise equal to `nn.RMSNorm` and to the diffusers modulation | +| PyTorch reference | `rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp` | n/a | `forward*`: provider path; `forward*_fp32`: FP64 golden (the modulated one stores `norm(x)` and `1 + scale` in BF16, straight-through, as the model does) | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(..., S, N)` | bf16 / fp16 / fp32 | `N % 4 == 0`; H3: `N = 5376` | +| `weight` | `(N,)` | `x.dtype` | | +| `shift`, `scale` | `(R, N)` row views, same stride | `x.dtype` | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; one entry per position, shared across the batch | + +Out-of-range indices, mismatched dtypes or shapes, a non-positive eps, and `N % 4 != 0` all fail +closed. + +## Numerics + +- **Statistics (contract `h3-rmsnorm-v1`).** The kernel replays PyTorch's own + `vectorized_layer_norm_kernel` (torch `cf30153`): + - one `(32, 4)` block per row, reading 4-element vectors; + - thread `t` sums vectors `t, t+128, …` in order; + - a shuffle-down tree (16 → 1), then a cross-warp tree; + - `rsqrtf(Σx²/N + eps)`, then `w * (rstd * x)` with one cast. + + The plain norm is therefore **bitwise equal to `nn.RMSNorm`**. Rows are independent of batch + size and position. +- **Modulation.** `1 + scale`, `n * (…)` and `+ shift` are each rounded to the tensor dtype, + exactly where the eager expression rounds. The fused output is **bitwise equal to diffusers' + `norm(x) * (1.0 + scale.index_select(0, i)) + shift.index_select(0, i)`**. +- **Backward.** Everything is in FP32 with one cast per output and no atomics: + - `dx` is row-local, using the same reduction tree; + - `dweight` is folded over fixed 256-row tiles in ascending order; + - `dshift`/`dscale` are segmented sums over positions sorted stably by table row, in the + same scheme as [`adaln_row_gather`](h3-adaln-row-gather.md). + + Diffusers' `index_select` backward accumulates `dshift`/`dscale` with BF16 atomics, so it is + non-deterministic and about 15–20× less accurate (its error varies from run to run). + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op h3_rmsnorm +``` + +B200, block `norm1` + MSA modulation, B = 1, H = 5376, BF16: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.08 ms | 0.12 ms | 1.23 ms | 0.85 ms | +| 32768 | 0.35 ms | 0.76 ms | 2.18 ms | 4.40 ms | +| 131072 | 1.21 ms | 2.90 ms | 6.17 ms | 17.55 ms | + +The forward is a single pass that never materialises the gathered rows. At small S the backward +is dominated by the fixed cost of the stable sort and tile setup. Backward timings and peak +memory exclude leaf creation and the forward pass. Candidate/provider execution order alternates +each iteration and is recorded in the report. + +## Evidence + +![h3_rmsnorm on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-rmsnorm-b200/report.json), written by +`tools/validation/models/h3_evidence.py` from a clean tree at commit `80e4609`. It also records: + +- bitwise equality with `nn.RMSNorm` for all four pinned norm weights; +- bitwise equality with diffusers for the modulation; +- row invariance. + +## Existing implementations (RFC #420 reuse rule) + +![norm_modulate vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers composition (F.rms_norm + index_select modulation) | **no** (param/table grads not repeatable) | 8.6e-03 / 4.5e-02 / 872 µs | 7.5e-03 / 1.9e-01 / 4434 µs | +| torch F.rms_norm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 331 µs | 2.0e-03 / 2.8e-03 / 850 µs | +| TE 2.20.2 RMSNorm (no modulation) | yes | 2.1e-03 / 2.1e-03 / 471 µs | 2.0e-03 / 2.8e-03 / 1597 µs | +| Liger 0.8.4 modulated RMSNorm + index_select | **no** (param/table grads not repeatable) | 6.6e-03 / 4.7e-02 / 810 µs | 6.9e-03 / 2.0e-01 / 3911 µs | +| Liger 0.8.4 modulated RMSNorm + rl-kernel row gather | yes | 6.6e-03 / 6.5e-03 / 1824 µs | 6.9e-03 / 8.0e-03 / 4763 µs | +| SGLang 0.5.21 fused_norm_scale_shift (forward only) | yes | 4.4e-03 / — / 85 µs (fwd only) | 4.1e-03 / — / 436 µs (fwd only) | +| rl-kernel H3RMSNormCudaOp.forward_modulated | yes | 8.6e-03 / 4.2e-03 / 1093 µs | 7.5e-03 / 5.0e-03 / 2653 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-prior-art-b200/norm_modulate.json) +was written from a clean tree at `ee83dec` by + +```bash +python tools/validation/models/h3_prior_art.py --op norm_modulate --out reports/experiments/h3-prior-art-b200/norm_modulate.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/norm_modulate.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![norm_modulate bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png) + +This kernel's modulated forward is bitwise identical to the diffusers composition at every S (1 to 32768). Liger 0.8.4 differs in 41.8-46.5% of the BF16 outputs, with either gather, and SGLang 0.5.21's fused_norm_scale_shift in 36.0-36.7%. torch F.rms_norm and Transformer Engine compute only the norm, not the modulation, so their rows are listed for completeness, not as parity candidates. +[`norm_modulate.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr487/reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json) was written from a clean tree at `228e50b9` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_rmsnorm.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, incl. block norm1 + modulation +python tools/validation/operators/check_operator.py --op h3_rmsnorm --candidate cuda --device cuda \ + --dtype bf16 --batch 2 --seq 257 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op h3_rmsnorm --out reports/experiments/h3-rmsnorm-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-rmsnorm-b200/report.json +``` + +## Known Limitations + +- Bitwise parity with `nn.RMSNorm` is tied to PyTorch's vectorized layer-norm path (`N % 4 == 0`, + aligned tensors). Other shapes are rejected rather than approximated. +- The token-refiner blocks' own norm weights live in shard 14, which is not pinned. They are the + same module (`nn.RMSNorm(5376, eps=1e-5)`) as the four pinned norms tested here. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-sp-norm-adaln.md b/docs/operators/h3-sp-norm-adaln.md new file mode 100644 index 000000000..8eb91cc38 --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,117 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `6c0d380`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. + +## Tests + +```bash +python -m pytest tests/models/minimax_h3/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-sp-norm-adaln-b200/report.json +``` + +The rank-backward tests run SP 2–8 on one GPU, with ranks as threads around the plain backward +functions. They cover odd `S`, `B = 2`, shards smaller than a tile, block and interleaved packing, +and bf16 and fp32, and compare every rank against WS1. The NCCL tests run the autograd region end +to end and need as many GPUs as ranks. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- With interleaved packing, up to a full shard of rows can move in the backward. The result is + still exact. +- Attention and FFN are other rows. Until they land, the region uses a seeded stand-in for the + sublayer output. diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md new file mode 100644 index 000000000..cc5809cee --- /dev/null +++ b/docs/operators/h3-timestep-mlp.md @@ -0,0 +1,184 @@ +# MiniMax-H3 FP32 Timestep MLP + +## Summary + +`timestep_mlp_fp32` is H3's `TimestepEmbedding(in_channels=256, time_embed_dim=5376, +out_dim=2688)`. It runs on the sinusoidal features of the distinct timesteps +(RFC #420, WS1 step 3): + +```text +temb = linear_2(silu(linear_1(features))) (T, 256) -> (T, 5376) -> (T, 2688) +``` + +`time_embedder` is a `_keep_in_fp32_modules` module. Its weights, biases, +activations and output are all FP32. `temb` stays FP32 because every AdaLN +projection applies its own SiLU before casting to BF16. + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227`. The tensors are +`time_embedder.linear_{1,2}.{weight,bias}` (see `rl_engine/validation/models/h3_manifest.json`). + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("timestep_mlp_fp32", device="cuda") +temb = op(features, w1, b1, w2, b2) # all float32 +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp` | `rl_engine._C.h3_det_linear_{forward,backward_input,backward_weight}` | Contract `h3-det-linear-v1` | +| PyTorch reference | `rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp` | n/a | `forward`: provider path (`F.linear`/`F.silu`); `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `x` | `(T, K)`, `T >= 1` | float32 | H3: `K = 256` | +| `w1`, `b1` | `(H, K)`, `(H,)` | float32 | H3: `H = 5376` | +| `w2`, `b2` | `(D, H)`, `(D,)` | float32 | H3: `D = 2688` | +| output | `(T, D)` | float32 | | + +The CUDA kernel needs `K` and `H` to be multiples of 4, for 16-byte vector loads. +These inputs fail closed: BF16 in any argument (RFC probe H7), mismatched shapes, +an empty `x`, mixed devices, and unaligned `K`. + +## Numerics: contract `h3-det-linear-v1` + +Forward, for each output `y[t, n]`: + +- One warp owns a group of output columns. Lane `l` reads the 16-byte K chunks + `l, l + 32, l + 64, ...`. +- Each lane accumulates its chunks in ascending order, and the elements inside a chunk + in ascending order, using `fmaf` into an FP32 accumulator that starts at 0. +- The 32 lane sums are combined with an xor butterfly (16, 8, 4, 2, 1). +- The bias is added once, then SiLU runs in FP32 as `v / (1 + expf(-v))` + (PyTorch's formula). The result is stored as FP32. + +How many columns a warp owns, the row tile (1, 2 or 4), and how many chunks are +loaded ahead change only when loads are issued, never the order in which a column +is summed. A timestep's `temb` is therefore bitwise identical however many other +timesteps share the call, and wherever it sits among them. + +Backward is deterministic. There is no cuBLAS and there are no atomics: + +- `d_hidden = g @ W2` and `dx = d_pre @ W1`: N is split into fixed 64-row chunks. + Each chunk is an ascending `fmaf` chain, and the chunks are then left-folded in + ascending order. These gradients are row-local and batch-invariant. +- `d_pre = d_hidden * s * (1 + z * (1 - s))`, elementwise in FP32, using the saved + pre-activation `z`. +- `dW[n, k] = sum_t g[t, n] * x[t, k]` and `db[n] = sum_t g[t, n]`: an ascending-`t` + FP32 fold that starts from 0, in logical row order. + +Accuracy, measured on a B200 with the pinned checkpoint weights, TF32 disabled and 200 +draws of 4 timesteps (each draw includes t = 0 and t = 1): + +| Comparison | CUDA | provider (cuBLAS) | +| --- | --- | --- | +| max abs error vs FP64 golden, median over draws | 1.1e-6 | 5.4e-7 | +| gradients vs FP64 autograd (`x`, `w1`, `b1`, `w2`, `b2`) | <= 1.2e-5 | | + +Both outputs are about 100× inside the contract (`reduction` / `float32`, atol and rtol +1e-4). Neither is more accurate in general: cuBLAS's tree happens to do slightly better +on these weights. What the CUDA kernel adds is a fixed summation order, which makes each +row batch- and position-invariant and every run repeat-bitwise, and it is 2× faster +for T >= 2. The two outputs are not bitwise equal, so `tools/validation/models/h3_chain_replay.py` +reports this stage as the chain's `first_drift`, which is expected for a reduction. + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op timestep_mlp_fp32 +``` + +B200, pinned weights (63.3 MB FP32). End-to-end op times include the Python wrapper: + +| T | CUDA op | provider | kernels only (CUDA) | +| --- | --- | --- | --- | +| 1 | 35.5 µs | 40.8 µs | 3.8 + 11.6 µs | +| 2 | 40.4 µs | 84.9 µs | 3.9 + 13.7 µs | +| 4 | 44.9 µs | 88.6 µs | 5.0 + 15.8 µs | + +The 5376→2688 layer reads its weights at about 5.0 TB/s at T = 1. For T >= 2 the +provider switches from cuBLAS GEMV to an SGEMM path, which takes about 52 µs of +kernel time. + +## Evidence + +![timestep_mlp_fp32 on B200: latency and per-draw error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-timestep-mlp-b200/report.json), written +by `tools/validation/models/h3_evidence.py` from a clean tree at commit `65ef7f6`. It also records that a +timestep's row is bitwise identical whether it runs alone or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![timestep_mlp vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers TimestepEmbedding, FP32 F.linear [plain] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 5.7e-07 / 4.1e-07 / 777 µs | 5.7e-07 / 1.0e-06 / 832 µs | 3.7e-06 / 2.0e-06 / 3609 µs | +| rl-kernel H3TimestepMLPCudaOp | yes | 2.6e-07 / 4.2e-07 / 748 µs | 2.3e-07 / 7.6e-07 / 4595 µs | 2.6e-07 / 2.3e-06 / 42993 µs | +| diffusers TimestepEmbedding, FP32 F.linear [vllm] | **no** (14214 rows, 4096 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.5e-07 / 752 µs | 3.7e-06 / 2.8e-06 / 921 µs | 3.7e-06 / 3.2e-06 / 3826 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang] | yes | 1.6e-03 / 1.6e-03 / 1036 µs | 1.5e-03 / 1.5e-03 / 1140 µs | 1.5e-03 / 1.7e-03 / 2636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [sglang_ieee] | yes | 3.3e-06 / 2.3e-06 / 2234 µs | 3.7e-06 / 2.8e-06 / 2428 µs | 3.7e-06 / 3.2e-06 / 7329 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_te_native] | **no** (14214 rows, 4108 sub-batches, 2232 sweep cases) | 1.8e-07 / 4.1e-07 / 729 µs | 3.7e-06 / 1.0e-06 / 814 µs | 3.7e-06 / 2.0e-06 / 3636 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton] | yes | 1.6e-03 / 1.6e-03 / 1027 µs | 1.5e-03 / 1.5e-03 / 1137 µs | 1.5e-03 / 1.7e-03 / 2634 µs | +| diffusers TimestepEmbedding, FP32 F.linear [megatron_triton_ieee] | yes | 3.3e-06 / 2.3e-06 / 2309 µs | 3.7e-06 / 2.8e-06 / 2435 µs | 3.7e-06 / 3.2e-06 / 7349 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +4096-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-prior-art-b200/timestep_mlp.json) +was written from a clean tree at `ffd958a` by + +```bash +python tools/validation/models/h3_prior_art.py --op timestep_mlp --out reports/experiments/h3-prior-art-b200/timestep_mlp.json --megatron-src +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/timestep_mlp.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![timestep_mlp bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png) + +No candidate reproduces this kernel's bits at any T (90.6-100% of the FP32 outputs differ). None is a bitwise drop-in for diffusers at H3's T = 1-4 either: vLLM's settings and Megatron te_native match diffusers at T = 1, 64 and 2048 but not at T = 4; SGLang/Megatron Triton with ieee match only from T = 64; the TF32 Triton modes never match. +[`timestep_mlp.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr484/reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json) was written from a clean tree at `0df83a5a` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_mlp --megatron-src \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_mlp.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_mlp.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: sinusoid -> MLP +python tools/validation/operators/check_operator.py --op timestep_mlp_fp32 --candidate cuda --device cuda \ + --dtype fp32 --batch 3 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_mlp_fp32 \ + --out reports/experiments/h3-timestep-mlp-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timestep-mlp-b200/report.json +``` + +## Known Limitations + +- FP32 only, by contract. The BF16 and FP16 gtest rows do not apply to this op. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- Weight gradients are reductions across timesteps, so they follow logical row + order. They are deterministic, but not batch-invariant (no cross-row reduction can be). diff --git a/docs/operators/h3-timestep-sinusoid.md b/docs/operators/h3-timestep-sinusoid.md new file mode 100644 index 000000000..75edf172b --- /dev/null +++ b/docs/operators/h3-timestep-sinusoid.md @@ -0,0 +1,163 @@ +# MiniMax-H3 Timestep Sinusoid + +## Summary + +`timestep_sinusoid_h3` produces the FP32 sinusoidal timestep features that feed the +MiniMax-H3 timestep MLP (RFC #420, WS1 step 3). H3 builds +`Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)` and calls it +on the *distinct* timesteps of the packed sequence, unscaled in `[0, 1]` +(`t = 1 - sigma`): + +```text +freq[k] = exp(-ln(10000) * k / 128) k = 0..127 +arg[t, k] = t * freq[k] +out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256), float32 +``` + +Pinned model: `MiniMaxAI/MiniMax-H3@42ed227` (`rl_engine/validation/models/h3_manifest.json`). +Provider reference: `huggingface/diffusers@f53d552`, `get_timestep_embedding`. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("timestep_sinusoid_h3", device="cuda") +features = op(timestep) # timestep: (T,) in [0, 1] -> (T, 256) float32 +``` + +## Backends + +| Backend | Wrapper | Native symbol | Status | +| --- | --- | --- | --- | +| CUDA (SM90, SM100) | `rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid.H3TimestepSinusoidCudaOp` | `rl_engine._C.h3_timestep_sinusoid_forward` | Bitwise equal to the provider path | +| PyTorch reference | `rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp` | n/a | `forward`: provider replay; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `timestep` | `(T,)`, `T >= 1` | fp32 / bf16 / fp16 | Finite, in `[0, 1]`; any stride; CUDA for the CUDA backend | +| `num_channels` | scalar | int | Positive and even; H3 uses 256 | +| output | `(T, num_channels)` | float32 | `[cos \| sin]` order | + +Inputs that break the contract fail closed: non-1-D, empty, integer or out-of-range +timesteps raise. `t * 1000` callers (RFC probe H10) therefore fail instead of +silently producing a different embedding. The range check costs one host read-back. +The native CUDA entrypoint performs the value check, so the wrapper does not repeat +the host synchronization. Pass `check_range=False` to `forward` or +`rl_engine._C.h3_timestep_sinusoid_forward` only when the caller has already validated +`t`; this trusted path skips value validation and is used for kernel-only profiling. + +## Numerics + +- **Operation order.** The kernel evaluates exactly the FP32 operations the provider + runs on CUDA: `(float)(-ln 10000) * k`, multiplied by the FP32 reciprocal of `128` + (PyTorch divides by a CPU scalar this way), `expf`, `t * freq`, `cosf` and `sinf`. + The multiplies use `__fmul_rn`, so they are never fused into an FMA. The extension + is built without `--use_fast_math`, so `expf`, `sinf` and `cosf` are the precise + libdevice functions. +- **Cast points.** bf16 and fp16 timesteps are upcast to FP32 before anything else, + matching `timesteps[:, None].float()`. The output is never rounded below FP32. +- **Invariance.** Each output element depends only on `(t, k)`. Results are + bitwise identical across batch size, position, permutation and repeated runs. +- **Backward.** `d/dt` is the analytic row-local VJP + (`-f * sin(t f)` for the cos half, `f * cos(t f)` for the sin half). It is summed over + the 256 channels with a fixed pairwise tree (`tree_sum_lastdim_fp32`), so it does not + depend on how many timesteps share the launch. + +Measured on a B200 (torch 2.13.0+cu130): + +| Comparison | Result | +| --- | --- | +| CUDA vs provider path, T in {1, 2, 3, 7, 64, 1000, 4097} | bitwise equal | +| CUDA vs FP64 golden | max abs 1.19e-7 (contract: FP32 elementwise atol 1e-5) | +| gradient vs FP64 golden (gtest, T = 7) | max abs 2.4e-7 | + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op timestep_sinusoid_h3 +``` + +On a B200 this is a single launch of about 15 µs, independent of `T` for `T <= 64`. With +the range check it takes about 49 µs. The provider path takes about 62–70 µs, because it is +six eager kernels plus two concatenations. Either way the op is launch-bound: it moves +about 1 KB per timestep. + +## Evidence + +![timestep_sinusoid_h3 on B200: latency and error vs FP64](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-timestep-sinusoid-b200/report.json). +`tools/validation/models/h3_evidence.py` wrote it from a clean tree at commit `0522865`, and the report +records that commit and the environment. + +## Existing implementations (RFC #420 reuse rule) + +![timestep_sinusoid vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.png) + +| Implementation | Batch-invariant | size 3: fwd err / worst grad err / fwd+bwd | size 256: fwd err / worst grad err / fwd+bwd | size 2048: fwd err / worst grad err / fwd+bwd | +|---|---|---|---|---| +| diffusers get_timestep_embedding (op-for-op replay) | yes | 6.7e-08 / — / 64 µs (fwd only) | 8.6e-08 / — / 64 µs (fwd only) | 9.9e-08 / — / 65 µs (fwd only) | +| SGLang 0.5.21 timestep_embedding | yes | 6.7e-08 / — / 13 µs (fwd only) | 8.6e-08 / — / 12 µs (fwd only) | 9.9e-08 / — / 13 µs (fwd only) | +| rl-kernel H3TimestepSinusoidCudaOp (check_range=False) | yes | 6.7e-08 / — / 16 µs (fwd only) | 8.6e-08 / — / 16 µs (fwd only) | 9.9e-08 / — / 16 µs (fwd only) | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +8192-timestep batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-prior-art-b200/timestep_sinusoid.json) +was written from a clean tree at `7a82917` by + +```bash +python tools/validation/models/h3_prior_art.py --op timestep_sinusoid --out reports/experiments/h3-prior-art-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/timestep_sinusoid.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs: +2^20 uniform timesteps, every 997th FP32 bit pattern in [0, 1], and a 100001-point grid. It +compares each output with diffusers' path and with this kernel. `--fast-math-variant` adds this +kernel compiled with `--use_fast_math`, which the build excludes (see Known Limitations). + +![timestep_sinusoid bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png) + +SGLang's `timestep_embedding`, diffusers' path and this kernel are bitwise identical on all +2.22M timesteps (567.6M outputs). The fast-math build differs in 53.7–82.7% of outputs. +[`timestep_sinusoid.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr483/reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json) +was written from a clean tree at `0eb5ddd5` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_bitwise_reuse.py \ + reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python tools/validation/operators/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timestep-sinusoid-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- `max_period` is fixed at 10000, the only value H3 uses. +- `KERNEL_ALIGN_USE_FAST_MATH=1` builds leave out every H3 kernel (`H3_BITWISE_SOURCES` in + `build_tools/extensions.py`): fast math replaces `expf`/`sinf`/`cosf` with approximations and + breaks bitwise parity. Such builds dispatch the PyTorch reference. diff --git a/docs/operators/h3-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..318dbd2d0 --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,113 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `4947996`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +python tools/validation/models/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \ + --out reports/experiments/h3-tp-adaln-3mod-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/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/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..40fb56213 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,153 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... +def h3_timestep_sinusoid_forward( + timestep: torch.Tensor, + num_channels: int = 256, + max_period: float = 10000.0, + check_range: bool = True, +) -> torch.Tensor: + """Return FP32 CUDA cosine-then-sine features of shape ``(T, num_channels)``. + + Require a nonempty 1-D FP32 CUDA timestep tensor, positive even channels + and positive ``max_period``. Check finite values in ``[0, 1]`` when enabled; + the Python operator supplies the analytic backward for this native forward. + """ + ... + +def h3_det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: int = 0, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Return a deterministic CUDA projection and optional FP32 pre-activation. + + Require same-device/dtype FP32 or BF16 contiguous, 16-byte-aligned matrices + ``x (T, K)`` and ``weight (N, K)``, ``T > 0`` and ``K`` divisible by the + 16-byte vector width. Optional contiguous bias is ``(N,)`` in that dtype. + Activation 0 is identity and 1 is SiLU; BF16 forward requires SM80 or newer. + Return ``[out]`` in the input dtype or ``[out, pre]`` when saving activation. + """ + ... + +def h3_det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Return a deterministic ``(T, K)`` CUDA input gradient in ``out_dtype``. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``weight (N, K)``. Output dtype + must be FP32 or BF16; output-column chunks accumulate in fixed order. + """ + ... + +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... +def h3_det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True +) -> list[torch.Tensor]: + """Return CUDA weight and optional bias gradients using an ascending row fold. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``x (T, K)``. Return ``[dweight]`` + with shape ``(N, K)``, or ``[dweight, dbias]`` with ``dbias (N,)`` when + requested, in FP32 or BF16 as specified by ``w_dtype``. + """ + ... + +def h3_adaln_row_gather_forward( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + chunks: int = 6, + modality_num: int = 3, +) -> torch.Tensor: ... +def h3_adaln_row_gather_backward( + grad: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... +def h3_rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, + sorted_pos: torch.Tensor | None = None, + tile_begin: torch.Tensor | None = None, + tile_end: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/__init__.py b/rl_engine/backends/cuda/model_specific/__init__.py new file mode 100644 index 000000000..3b09d3c11 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA implementations with model-specific arithmetic contracts.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py new file mode 100644 index 000000000..17a3a3884 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused H3 AdaLN modulation: ``adaln_projection_3mod`` followed by ``adaln_row_gather``. + +The forward is exactly the two RFC #420 ops back to back (same kernels, same +bits). The difference is the backward. With two separate ops, autograd hands +the gather's table gradient to the projection in the table's dtype (BF16), +which rounds a long FP32 segment sum once more before it reaches the +projection weights and ``temb``. Here the segment sum stays FP32 into the +projection backward, so the parameter gradients carry no extra BF16 rounding. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + adaln_row_gather_available, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import validate_h3_adaln_projection +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_adaln_row_gather_forward" +) +BACKWARD_IMPL = "fp32_segment_sum_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_adaln_row_gather_backward[fp32]" + "+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNModulationCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias) + rows = table.view(-1, H3_ADALN_CHUNKS * hidden) + packed = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(temb, act, weight, timestep_indices, token_tags) + ctx.table_shape = table.shape + return packed + + @staticmethod + def backward(ctx, grad_packed): + temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] + num_timesteps, width = ctx.table_shape + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) + d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) + d_table = d_rows.view(num_timesteps, width) + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) + record_backward( + "adaln_modulation_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return ( + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + None, + ) + + +class H3AdaLNModulationCudaOp: + """``adaln_projection_3mod`` + ``adaln_row_gather`` with an FP32 table gradient. + + Returns the six ``(S, H)`` modulation tensors for every packed position, + bitwise equal to running the two ops separately. + """ + + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and adaln_row_gather_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + + def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") + _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + packed = _H3AdaLNModulationCuda.apply( + temb, + weight, + bias, + timestep_indices.contiguous(), + token_tags.contiguous(), + hidden, + ) + return tuple(packed.unbind(0)) + + +def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): + # Reuse the gather's checks against a stand-in of the right row count. + stand_in = torch.empty( + (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device + ) + validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..f02431903 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 three-modality AdaLN projection (RFC #420 ``adaln_projection_3mod``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.adaln_projection import ( + split_adaln_table, + validate_h3_adaln_projection, +) + +KERNEL_ID = "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNProjectionCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + """Project FP32 ``(T, D)`` embeddings to a weight-dtype ``(T, 18H)`` table. + + Save the embedding, rounded SiLU activation and weight for the VJP. + """ + + temb = temb.contiguous() + # Declared mixed-precision boundary: SiLU at temb's FP32 precision, + # then exactly one rounding to the projection dtype. These are the + # provider's own elementwise ops, so the activation is bitwise equal. + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + def backward(ctx, grad_table): + """Return requested embedding, weight and bias gradients in their input dtypes. + + Accumulate projection gradients deterministically and evaluate the + straight-through cast and SiLU VJP in FP32 for the embedding gradient. + """ + + temb, act, weight = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad + grad_table = grad_table.float().contiguous() + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + if need_temb: + # The cast's VJP is the identity; the SiLU VJP runs in FP32. + d_act = det_linear_backward_input(grad_table, weight, torch.float32) + d_temb = silu_backward_fp32(d_act, temb) + record_backward( + "adaln_projection_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return d_temb, dw if need_weight else None, db if need_bias else None + + +class H3AdaLNProjectionCudaOp: + """CUDA candidate: FP32 SiLU, one BF16 cast, deterministic GEMV, six views. + + The projection weight (96768 x 2688 BF16, 520 MB per block) is streamed + once per call; every output element follows contract ``h3-det-linear-v1``, + so a timestep's 3 x 6 modulation rows do not depend on the other + timesteps in the call. The six outputs are views of one ``(T, 6H*3)`` + table, laid out exactly like diffusers' ``view(-1, 6H).chunk(6)``. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six ``(3T, H)`` views after FP32 SiLU and a deterministic projection. + + Require CUDA FP32 ``temb`` of shape ``(T, D)`` with ``T > 0`` and + same-device FP32 or BF16 weight/bias of shapes ``(18H, D)``/``(18H,)``. + Outputs share one table in the weight dtype; autograd covers all inputs. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + table = _H3AdaLNProjectionCuda.apply(temb, weight, bias) + return split_adaln_table(table, hidden) + + def forward_table(self, temb, weight, bias) -> torch.Tensor: + """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" + + validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + return _H3AdaLNProjectionCuda.apply(temb, weight, bias) + + def forward_fp32(self, temb, weight, bias): + """Run the CUDA projection with its declared weight-dtype output boundary.""" + + return self.forward(temb, weight, bias) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..ee681454d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 AdaLN row gather (RFC #420 ``adaln_row_gather``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = "rl_engine._C.h3_adaln_row_gather_forward" +BACKWARD_IMPL = "stable_sorted_segment_tiles_fp32_ordered_fold" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_adaln_row_gather_backward" +# Positions per backward tile; part of the fixed summation order. +BACKWARD_TILE = 256 + + +def adaln_row_gather_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_adaln_row_gather_forward") + and hasattr(_C, "h3_adaln_row_gather_backward") + ) + + +def _segment_tiles(index: torch.Tensor, num_rows: int): + """Sorted positions and fixed-size tiles per row segment (all integer ops).""" + + sorted_pos = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows) + seg_start = torch.cumsum(counts, 0) - counts + tiles_per_seg = (counts + BACKWARD_TILE - 1) // BACKWARD_TILE + seg_first_tile = torch.zeros(num_rows + 1, dtype=torch.int64, device=index.device) + seg_first_tile[1:] = torch.cumsum(tiles_per_seg, 0) + tile_seg = torch.repeat_interleave(torch.arange(num_rows, device=index.device), tiles_per_seg) + tile_local = torch.arange(tile_seg.numel(), device=index.device) - seg_first_tile[tile_seg] + tile_begin = seg_start[tile_seg] + tile_local * BACKWARD_TILE + tile_end = torch.minimum(tile_begin + BACKWARD_TILE, seg_start[tile_seg] + counts[tile_seg]) + return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile + + +class _H3AdaLNRowGatherCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, rows, timestep_indices, token_tags): + out = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(timestep_indices, token_tags) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return out + + @staticmethod + def backward(ctx, grad_out): + timestep_indices, token_tags = ctx.saved_tensors + num_rows, rows_dtype = ctx.rows_meta + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_rows) + grad_rows = _C.h3_adaln_row_gather_backward(grad_out.contiguous(), *tiles, rows_dtype) + record_backward( + "adaln_row_gather", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad_rows, None, None + + +class H3AdaLNRowGatherCudaOp: + """CUDA candidate: one launch gathers all six modulation tensors. + + The forward is a byte copy (bitwise equal to six ``index_select`` calls). + The backward replaces ``index_select``'s atomic BF16 scatter-add with a + deterministic FP32 segmented sum in stable sorted order, cast once. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not adaln_row_gather_available(): + raise RuntimeError( + "rl_engine._C lacks h3_adaln_row_gather_*; rebuild with " + "csrc/cuda/h3/adaln_row_gather.cu" + ) + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + if not rows.is_cuda: + raise ValueError("H3AdaLNRowGatherCudaOp needs CUDA tensors") + if rows.stride(1) != 1: + rows = rows.contiguous() + packed = _H3AdaLNRowGatherCuda.apply( + rows, timestep_indices.contiguous(), token_tags.contiguous() + ) + return tuple(packed.unbind(0)) + + def gather_chunks(self, chunks, timestep_indices, token_tags, **kwargs): + """Drop-in for diffusers' six ``(3T, H)`` tensors: concatenate, then gather.""" + + if len(chunks) != H3_ADALN_CHUNKS: + raise ValueError(f"expected {H3_ADALN_CHUNKS} modulation tensors, got {len(chunks)}") + return self.forward(torch.cat(list(chunks), dim=1), timestep_indices, token_tags, **kwargs) + + def forward_fp32(self, rows, timestep_indices, token_tags): + return tuple(out.float() for out in self.forward(rows, timestep_indices, token_tags)) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py new file mode 100644 index 000000000..1274364f2 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python surface of the H3 deterministic row-wise linear (contract h3-det-linear-v1). + +See ``csrc/cuda/h3/det_linear.cu`` for the reduction order. Every helper +raises instead of falling back to cuBLAS when the extension is missing. +""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE + +ACT_NONE = 0 +ACT_SILU = 1 +CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) +_SYMBOLS = ( + "h3_det_linear_forward", + "h3_det_linear_backward_input", + "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", +) + + +def det_linear_available() -> bool: + """Return whether the extension exposes all three H3 deterministic linear APIs.""" + + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +def _require() -> None: + """Raise when a deterministic linear entry point is unavailable.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " + "with csrc/cuda/h3/det_linear.cu (no cuBLAS fallback is allowed)" + ) + + +def det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + activation: int = ACT_NONE, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Project CUDA ``x`` of shape ``(T, K)`` with ``(N, K)`` weights and optional bias. + + Inputs share an FP32 or BF16 dtype/device, ``T > 0``, and ``K`` supports + 16-byte vector loads; BF16 forward needs SM80 or newer. Return ``[out]`` + in the input dtype, plus an FP32 pre-activation when requested. Activation + is identity or SiLU; contiguous copies preserve the reduction contract. + """ + + _require() + return _C.h3_det_linear_forward( + x.contiguous(), + weight.contiguous(), + None if bias is None else bias.contiguous(), + int(activation), + bool(save_pre_activation), + ) + + +def det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Compute a deterministic ``(T, K)`` input VJP on the inputs' CUDA device. + + Convert nonempty ``(T, N)`` gradients to FP32 and combine them with FP32 + or BF16 ``(N, K)`` weights; return FP32 or BF16 as specified by ``out_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_input( + grad.float().contiguous(), weight.contiguous(), out_dtype + ) + + +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + +def det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True +) -> list[torch.Tensor]: + """Fold nonempty CUDA ``(T, N)`` gradients and ``(T, K)`` inputs in row order. + + Convert gradients to FP32; ``x`` is FP32 or BF16 on the same device. + Return ``[dweight]`` of shape ``(N, K)``, optionally followed by ``dbias`` + of shape ``(N,)``, both in the requested FP32 or BF16 ``w_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_weight( + grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) + ) + + +def silu_backward_fp32(grad: torch.Tensor, pre_activation: torch.Tensor) -> torch.Tensor: + """Elementwise SiLU VJP in FP32: g * s * (1 + z * (1 - s)), s = sigmoid(z).""" + + z = pre_activation.float() + s = torch.sigmoid(z) + return grad.float() * (s * (1.0 + z * (1.0 - s))) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..741a0a85e --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 final AdaLN output (RFC #420 ``final_adaln_out``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out +from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_rmsnorm_forward[modulated]" +) +BACKWARD_IMPL = "fp32_table_grad_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_rmsnorm_backward+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3FinalAdaLNOutCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): + hidden = x.shape[-1] + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) # declared boundary: one cast after the FP32 SiLU + (table,) = det_linear_forward(act, weight, bias) + shift, scale = table.chunk(2, dim=-1) + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward( + x2, norm_weight.contiguous(), float(eps), shift, scale, timestep_indices + ) + ctx.save_for_backward(x2, norm_weight, rstd, temb, act, weight, table, timestep_indices) + ctx.x_shape = x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + shift, scale = table.chunk(2, dim=-1) + positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + tiles = _segment_tiles(positions, table.shape[0]) + dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( + grad.contiguous().view_as(x2), + x2, + norm_weight.contiguous(), + rstd, + shift, + scale, + timestep_indices, + *tiles, + ) + # The table gradient stays FP32 into the projection backward. + d_table = torch.cat([d_shift, d_scale], dim=1) + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + record_backward( + "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + + +class H3FinalAdaLNOutCudaOp: + """CUDA candidate: tensor-core shift/scale projection + fused norm/modulation. + + The projection is ``adaln_projection_3mod``'s deterministic GEMV + (``h3-det-linear-bf16-mma-v1``) on ``norm_out.linear``; the norm and + modulation are ``h3_rmsnorm``'s kernel indexed by ``timestep_indices``. + One autograd node, so the table gradient reaches the projection in FP32. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and h3_rmsnorm_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + if weight.dtype != x.dtype: + raise TypeError("norm_out.linear must share the activations' dtype") + if not x.is_cuda: + raise ValueError("H3FinalAdaLNOutCudaOp needs CUDA tensors") + return _H3FinalAdaLNOutCuda.apply( + x, norm_weight, temb, weight, bias, timestep_indices.contiguous(), eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py new file mode 100644 index 000000000..3d040e9e8 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py new file mode 100644 index 000000000..f61a7dbcd --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 tensor-, sequence- and context-parallel building blocks (contract ``h3-parallel-v1``). + +Rows: ``tp_attention_7168``, ``tp_ffn_14336``, ``cp_full_attention`` (Ulysses), +``cp_metadata_coshard``, ``cp_arbitrary_length``, ``distributed_packed_io``. + +Every reduction that crosses ranks is defined on **8 canonical units** so TP 1/2/4/8 add the +same numbers in the same order: + +* *row-parallel* (O and FFN-down, ``K`` = 7168 / 14336): the reduction axis is cut into 8 + equal K-units; a unit is one fixed-order FP32 GEMM (``h3-projection-gemm-v1`` tile, ``K`` + swept ascending). A rank holding ``8/tp`` consecutive units adds them in balanced-tree + order and the collective adds the ranks' partials in the same tree: the result is the + balanced tree over the 8 units for every TP size, rounded to BF16 once. +* *column-parallel backward* (``dX`` of Q/K/V and gate/up, reduction over the output axis): + the same 8 units, where unit ``j`` is output block ``j`` of every weight of the fused + projection, added in the order given (q, k, v / hidden, gate). + +The forward of a column-parallel projection and the backward of a row-parallel one reduce +over axes that do not depend on the sharding (they are plain ``h3-projection-gemm-v1``), so +they are bitwise the unsharded result without any collective. + +Context parallelism (Ulysses) moves data only: sequence <-> head all-to-alls are exact, so +attention (per head, full sequence) is the unsharded attention. Weight gradients under +sequence sharding are computed from the gathered rows, so they equal the unsharded GEMM. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from rl_engine.backends.shared.triton.gemm import h3_projection as g +from rl_engine.distributed.algorithms.hierarchical import tree_add + +CONTRACT = "h3-parallel-v1" +UNITS = 8 # canonical reduction units; TP sizes are the divisors 1, 2, 4, 8 +TP_SIZES = (1, 2, 4, 8) +HEADS, HEAD_DIM = 56, 128 +INNER = HEADS * HEAD_DIM # 7168 +FFN = 14336 + + +# --- Sharding helpers --------------------------------------------------------------------------- + + +def check_tp(tp: int, rank: int) -> None: + if tp not in TP_SIZES or not 0 <= rank < tp: + raise ValueError( + f"tp must be one of {TP_SIZES} and rank in [0, tp); got tp={tp} rank={rank}" + ) + + +def shard_rows(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Output-axis shard of a column-parallel weight ``(N, K)``: rows ``[r*N/tp, (r+1)*N/tp)``.""" + + check_tp(tp, rank) + n = weight.shape[0] + if n % UNITS: + raise ValueError(f"{n} output rows do not split into {UNITS} units") + return weight.narrow(0, rank * n // tp, n // tp).contiguous() + + +def shard_cols(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Reduction-axis shard of a row-parallel weight ``(N, K)``: columns of ``8/tp`` K-units.""" + + check_tp(tp, rank) + k = weight.shape[1] + unit = k // UNITS + if unit * UNITS != k or unit % g.CONFIG_BF16["BLOCK_K"]: + raise ValueError(f"K={k} must split into {UNITS} units of a multiple of the GEMM K tile") + return weight.narrow(1, rank * k // tp, k // tp).contiguous() + + +def shard_heads(x: torch.Tensor, tp: int, rank: int, dim: int = -1) -> torch.Tensor: + """Head shard of a ``(..., 7168)`` activation or weight axis (whole 128-wide heads).""" + + check_tp(tp, rank) + return x.narrow(dim, rank * x.shape[dim] // tp, x.shape[dim] // tp).contiguous() + + +def shard_gate_up(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor: + """Shard of the fused ``(2*FFN, K)`` gate/up weight keeping ``[hidden; gate]`` halves. + + Each half is sharded on its own, so local rows are ``[hidden shard; gate shard]`` and + ``chunk(2)`` still splits them. + """ + + half = weight.shape[0] // 2 + return torch.cat( + [shard_rows(weight[:half], tp, rank), shard_rows(weight[half:], tp, rank)] + ).contiguous() + + +# --- Row-parallel linear (reduction over the sharded axis) --------------------------------------- + + +def _units_partial(x: torch.Tensor, w: torch.Tensor, tp: int) -> torch.Tensor: + """FP32 balanced-tree sum of this rank's ``8/tp`` K-units of ``x @ w.T``.""" + + local = UNITS // tp + unit = x.shape[1] // local + parts = [ + g.gemm( + x[:, i * unit : (i + 1) * unit], + w[:, i * unit : (i + 1) * unit], + b_t=True, + out_dtype=torch.float32, + ) + for i in range(local) + ] + return tree_add(parts) + + +class _RowParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w, collective, tp, scatter): + ctx.save_for_backward(x, w) + ctx.scatter, ctx.collective = scatter, collective + partial = _units_partial(x, w, tp) + if tp > 1: + partial = ( + collective.reduce_scatter(partial) if scatter else collective.all_reduce(partial) + ) + return partial.to(x.dtype) + + @staticmethod + def backward(ctx, grad): + x, w = ctx.saved_tensors + if ctx.scatter: # sequence-parallel: the rows of dY are sharded; gather them (exact) + grad = ctx.collective.all_gather(grad.contiguous()) + grad = grad.contiguous() + dx = g.gemm(grad, w) # reduction over N: independent of the shard + dw = g.gemm(grad, x, a_t=True) # reduction over tokens: independent of the shard + return dx, dw, None, None, None + + +def row_parallel_linear(x, w_shard, collective=None, tp: int = 1, scatter: bool = False): + """``y = x @ W.T`` with ``x (T, K/tp)`` and ``w_shard (N, K/tp)``; ``y`` bitwise TP-invariant. + + ``scatter`` gives this rank's ``T/tp`` rows of ``y`` (sequence parallel) instead of all of it. + """ + + if x.dim() != 2 or w_shard.dim() != 2 or x.shape[1] != w_shard.shape[1]: + raise ValueError("row_parallel_linear needs x (T, K/tp) and w (N, K/tp)") + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a collective") + return _RowParallel.apply(x.contiguous(), w_shard, collective, tp, scatter) + + +# --- Column-parallel linear (backward reduces over the sharded axis) ------------------------------ + + +class _ColumnParallel(torch.autograd.Function): + @staticmethod + def forward(ctx, x, collective, tp, gather, *weights): + ctx.save_for_backward(x, *weights) + ctx.collective, ctx.tp, ctx.gather = collective, tp, gather + full_x = collective.all_gather(x) if gather else x + ctx.full_x = full_x if gather else None + return tuple(g.gemm(full_x, w, b_t=True) for w in weights) + + @staticmethod + def backward(ctx, *grads): + x, *weights = ctx.saved_tensors + full_x = ctx.full_x if ctx.gather else x + local = UNITS // ctx.tp + grads = [t.contiguous() for t in grads] + units = [] + for j in range(local): + acc = None + for grad, w in zip(grads, weights): + width = grad.shape[1] // local + piece = g.gemm( + grad[:, j * width : (j + 1) * width], + w[j * width : (j + 1) * width], + out_dtype=torch.float32, + ) + acc = piece if acc is None else acc + piece + units.append(acc) + partial = tree_add(units) + if ctx.tp > 1: + partial = ( + ctx.collective.reduce_scatter(partial) + if ctx.gather + else ctx.collective.all_reduce(partial) + ) + dx = partial.to(x.dtype) + dws = tuple(g.gemm(grad, full_x, a_t=True) for grad in grads) + return (dx, None, None, None, *dws) + + +def column_parallel_linear( + x, weights: Sequence[torch.Tensor], collective=None, tp: int = 1, gather: bool = False +): + """``[x @ w.T for w in weights]`` with each ``w`` a row shard ``(N/tp, K)``. + + The forward is the unsharded projection's columns (no collective). With ``gather``, ``x`` is + this rank's rows of a sequence-parallel activation, gathered first (exact), and ``dX`` is + reduce-scattered back to rows. ``dX`` is the canonical 8-unit tree for every ``tp``. + """ + + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a collective") + return _ColumnParallel.apply(x.contiguous(), collective, tp, gather, *weights) + + +# --- Context parallelism (Ulysses) --------------------------------------------------------------- + + +@dataclass(frozen=True) +class SeqShard: + """This rank's contiguous slice of a packed sequence of ``length`` rows over ``cp`` ranks. + + Slices differ by at most one row (first ``length % cp`` ranks get the extra row), so any + length is legal, including ``length < cp`` (some ranks hold nothing). + """ + + length: int + cp: int + rank: int + + @property + def sizes(self) -> list[int]: + base, extra = divmod(self.length, self.cp) + return [base + (r < extra) for r in range(self.cp)] + + @property + def lo(self) -> int: + return sum(self.sizes[: self.rank]) + + @property + def hi(self) -> int: + return self.lo + self.sizes[self.rank] + + +def coshard(shard: SeqShard, *tensors: torch.Tensor, dim: int = 0) -> list[torch.Tensor]: + """Slice every per-row tensor (hidden, timestep_indices, token_tags, position_ids, ...) alike. + + The metadata that selects AdaLN rows and rotary angles must follow the rows it describes; + slicing them in one call with one ``SeqShard`` is what keeps them aligned. + """ + + out = [] + for t in tensors: + if t.shape[dim] != shard.length: + raise ValueError( + f"tensor has {t.shape[dim]} rows along dim {dim}, expected {shard.length}" + ) + out.append(t.narrow(dim, shard.lo, shard.hi - shard.lo).contiguous()) + return out + + +def scatter_packed(packed: torch.Tensor, shard: SeqShard, dim: int = 0) -> torch.Tensor: + """This rank's rows of a packed tensor that every rank holds in full.""" + + return coshard(shard, packed, dim=dim)[0] + + +def gather_packed(local: torch.Tensor, shard: SeqShard, group=None, dim: int = 0) -> torch.Tensor: + """Reassemble the packed tensor from every rank's rows (exact, any length).""" + + if shard.cp == 1: + return local + sizes = shard.sizes + widest = max(sizes) + moved = local.movedim(dim, 0).contiguous() + pad = moved.new_zeros((widest - moved.shape[0], *moved.shape[1:])) + chunks = [torch.empty_like(torch.cat([moved, pad])) for _ in range(shard.cp)] + dist.all_gather(chunks, torch.cat([moved, pad]), group=group) + return torch.cat([c[:n] for c, n in zip(chunks, sizes)]).movedim(0, dim) + + +class _AllToAll(torch.autograd.Function): + @staticmethod + def forward(ctx, x, in_sizes, out_sizes, group): + ctx.sizes, ctx.group = (in_sizes, out_sizes), group + out = x.new_empty((sum(out_sizes), *x.shape[1:])) + dist.all_to_all_single(out, x.contiguous(), out_sizes, in_sizes, group=group) + return out + + @staticmethod + def backward(ctx, grad): + in_sizes, out_sizes = ctx.sizes + back = grad.new_empty((sum(in_sizes), *grad.shape[1:])) + dist.all_to_all_single(back, grad.contiguous(), in_sizes, out_sizes, group=ctx.group) + return back, None, None, None + + +def ulysses_attention(q, k, v, attention, shard: SeqShard, group=None): + """Full attention of a sequence-sharded ``q, k, v (B, S_local, H, D)`` by head scatter. + + All-to-all makes each rank hold ``H/cp`` heads for the whole sequence, ``attention(q, k, v)`` + (``(B, H', S, D)`` layout, per-head independent) runs, and the result goes back. Both moves + are exact, so the output rows equal the unsharded attention's rows bitwise as long as + ``attention`` is batch/head invariant. + """ + + cp = shard.cp + if cp == 1: + return attention(q, k, v) + batch, local_rows, heads, dim = q.shape + if heads % cp: + raise ValueError(f"{heads} heads do not divide over cp={cp}") + sizes, hp = shard.sizes, heads // cp + if local_rows != sizes[shard.rank]: + raise ValueError(f"expected {sizes[shard.rank]} local rows, got {local_rows}") + + def to_heads(t): + # peer r receives head block r of my rows; I receive every peer's rows of my head block + send = t.reshape(batch, local_rows, cp, hp, dim).permute(2, 1, 0, 3, 4) + send = send.reshape(cp * local_rows, batch, hp, dim) + got = _AllToAll.apply(send, [local_rows] * cp, sizes, group) + return got.permute(1, 0, 2, 3) # (B, S, hp, D), sequence in rank order + + out = attention(to_heads(q), to_heads(k), to_heads(v)) # (B, S, hp, D) + send = out.permute(1, 0, 2, 3).contiguous() + got = _AllToAll.apply(send, sizes, [local_rows] * cp, group) + got = got.reshape(cp, local_rows, batch, hp, dim) + return got.permute(2, 1, 0, 3, 4).reshape(batch, local_rows, heads, dim) + + +# --- Weight gradients under sequence sharding ---------------------------------------------------- + + +class _SeqLinear(torch.autograd.Function): + """Row-local linear on a sequence shard whose ``dW`` is computed from gathered rows.""" + + @staticmethod + def forward(ctx, x, w, shard, group): + ctx.save_for_backward(x, w) + ctx.shard, ctx.group = shard, group + return g.gemm(x, w, b_t=True) + + @staticmethod + def backward(ctx, grad): + x, w = ctx.saved_tensors + grad = grad.contiguous() + dx = g.gemm(grad, w) + full_x = gather_packed(x, ctx.shard, ctx.group) + full_g = gather_packed(grad, ctx.shard, ctx.group) + return dx, g.gemm(full_g, full_x, a_t=True), None, None + + +def seq_sharded_linear(x, w, shard: SeqShard, group=None): + """``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM (token sweep).""" + + return ( + _SeqLinear.apply(x.contiguous(), w, shard, group) + if shard.cp > 1 + else g.linear_forward(x, w, None) + ) 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/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..d9266e759 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 RMSNorm with fused AdaLN modulation (RFC #420 ``h3_rmsnorm``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +KERNEL_ID = "rl_engine._C.h3_rmsnorm_forward" +BACKWARD_IMPL = "row_local_dx_tiled_dweight_sorted_table_grads" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_rmsnorm_backward" + + +def h3_rmsnorm_available() -> bool: + return bool( + _EXT_AVAILABLE and hasattr(_C, "h3_rmsnorm_forward") and hasattr(_C, "h3_rmsnorm_backward") + ) + + +class _H3RMSNormCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, index, eps): + hidden = x.shape[-1] + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward(x2, weight.contiguous(), float(eps), shift, scale, index) + ctx.save_for_backward(x2, weight, rstd, shift, scale, index) + ctx.x_shape = x.shape + ctx.modulated = index is not None + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + if not ctx.modulated: + dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) + d_shift = d_scale = None + else: + positions = index.repeat(x2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, shift.shape[0]) + dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( + g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles + ) + d_shift, d_scale = d_shift.to(shift.dtype), d_scale.to(scale.dtype) + record_backward( + "h3_rmsnorm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None + + +class H3RMSNormCudaOp: + """CUDA candidate. + + The statistics replay PyTorch's own RMSNorm reduction order, so the plain + norm is bitwise equal to ``nn.RMSNorm``; the modulated form gathers its + ``shift``/``scale`` rows in the kernel (no ``(S, H)`` intermediates) and + rounds each step where the eager expression does, so it is bitwise equal to + diffusers' ``n * (1.0 + scale[i]) + shift[i]`` as well. Rows are + independent of batch size and position. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_rmsnorm_available(): + raise RuntimeError( + "rl_engine._C lacks h3_rmsnorm_*; rebuild with csrc/cuda/h3/rmsnorm_modulate.cu" + ) + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + return _H3RMSNormCuda.apply(x, weight, None, None, None, eps) + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + return self.forward(x, weight, eps).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py new file mode 100644 index 000000000..37cea620b --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sequence-parallel H3 RMSNorm / AdaLN modulation / gated residual (RFC #420 ``sp_norm_adaln``). + +Ownership: rank ``r`` of ``sp`` holds packed positions ``[r * S // sp, (r + 1) * S // sp)`` +of every batch item. The full ``(S,)`` row index is replicated; it is small, and every +rank derives the same reduction plan from it. + +Forward and the row-local gradients (``dx``, ``d_sublayer``, ``d_residual``) are the WS1 +kernels on the local rows. Rows are independent, so they are byte-equal by construction. + +The cross-row reductions (``d_norm_weight``, ``d_shift``/``d_scale``, ``d_gate``) are WS1 +two-level folds: FP32 sums over fixed 256-element tiles, then an ascending fold of the +tiles. Every rank builds the global WS1 tile list. A tile is computed by the rank that +holds its first row; the tile's rows from other ranks are all-gathered beforehand (only +rows of tiles that straddle a shard boundary move). The tile partials are all-gathered, +put back in WS1 tile order, and every rank runs the WS1 fold. The partial and fold kernels +are the WS1 kernels, so every gradient equals WS1 on every rank. The only collectives are +rank-ordered all-gathers (copies). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + BACKWARD_TILE, + _segment_tiles, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +_SYMBOLS = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward_dx", + "h3_rmsnorm_backward_partials", + "h3_rmsnorm_fold_partials", + "h3_gate_residual_forward", + "h3_gate_residual_backward_dy", + "h3_gate_grad_partials", + "h3_gate_grad_fold", +) +BACKWARD_IMPL = "row_local_ws1+gather_straddling_rows+ws1_tile_partials+gather+ws1_fold" + + +def sp_norm_adaln_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +@dataclass(frozen=True) +class SPRowLayout: + """Which packed positions each sequence-parallel rank holds.""" + + seq_len: int + batch: int + sp: int + rank: int + + def bounds(self, rank: int) -> tuple[int, int]: + return rank * self.seq_len // self.sp, (rank + 1) * self.seq_len // self.sp + + @property + def lo(self) -> int: + return self.bounds(self.rank)[0] + + @property + def hi(self) -> int: + return self.bounds(self.rank)[1] + + @property + def local_len(self) -> int: + return self.hi - self.lo + + +def sp_row_layout(seq_len: int, batch: int, sp: int, rank: int) -> SPRowLayout: + if sp < 1 or not 0 <= rank < sp: + raise ValueError(f"need 0 <= rank < sp, got rank={rank}, sp={sp}") + if batch < 1 or seq_len < sp: + raise ValueError(f"every rank needs a row: seq_len={seq_len}, sp={sp}, batch={batch}") + return SPRowLayout(seq_len, batch, sp, rank) + + +def _ordinal_by_group(group: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per-group counts, and each element's position within its group (in input order).""" + + counts = torch.bincount(group, minlength=groups) + order = torch.argsort(group, stable=True) + ordinal = torch.empty_like(group) + first = torch.cumsum(counts, 0) - counts + ordinal[order] = torch.arange(group.numel(), device=group.device) - first[group[order]] + return counts, ordinal + + +class _Plan: + """Global WS1 tiles of one backward, and this rank's part in computing them. + + ``families`` are ``(rows, tile_begin, tile_end)`` over global flattened rows + ``b * S + s``; consecutive tiles cover ``rows`` in order. Every rank builds the same + plan from replicated metadata, so collective calls and shapes agree. + """ + + def __init__(self, layout: SPRowLayout, families, device) -> None: + self.layout, sp = layout, layout.sp + self._his = torch.tensor([layout.bounds(r)[1] for r in range(sp)], device=device) + + # Rows of a tile computed by another rank: the only activation rows that move. + tile_owner, shared = [], [] + for rows, begin, end in families: + owner = self.owner(rows[begin]) + tile_owner.append(owner) + shared.append(rows[torch.repeat_interleave(owner, end - begin) != self.owner(rows)]) + shared = torch.unique(torch.cat(shared)) + shared_owner = self.owner(shared) + counts, slot = _ordinal_by_group(shared_owner, sp) + self.max_send = int(counts.max()) if shared.numel() else 0 + + # Buffer row of every global row this rank reads: its own rows first, then + # every rank's send block (padded to max_send) in rank order. + own = self.global_rows(layout.rank) + buffer_index = torch.full((layout.batch * layout.seq_len,), -1, device=device) + buffer_index[own] = torch.arange(own.numel(), device=device) + foreign = shared_owner != layout.rank + buffer_index[shared[foreign]] = ( + own.numel() + shared_owner[foreign] * self.max_send + slot[foreign] + ) + mine = ~foreign + self.send_local = self.local_index(shared[mine][torch.argsort(slot[mine])]) + self.recv_global = torch.zeros(sp * self.max_send, dtype=torch.long, device=device) + self.recv_global[shared_owner * self.max_send + slot] = shared # padding: never read + + self.families = [ + self._my_tiles(rows, begin, end, owner, buffer_index) + for (rows, begin, end), owner in zip(families, tile_owner) + ] + + def owner(self, rows: torch.Tensor) -> torch.Tensor: + return torch.searchsorted(self._his, rows % self.layout.seq_len, right=True) + + def global_rows(self, rank: int) -> torch.Tensor: + lo, hi = self.layout.bounds(rank) + s = torch.arange(lo, hi, device=self._his.device) + b = torch.arange(self.layout.batch, device=self._his.device) + return (b[:, None] * self.layout.seq_len + s[None, :]).reshape(-1) + + def local_index(self, rows: torch.Tensor) -> torch.Tensor: + b, s = rows // self.layout.seq_len, rows % self.layout.seq_len + return b * self.layout.local_len + s - self.layout.lo + + def _my_tiles(self, rows, begin, end, owner, buffer_index): + mine = torch.nonzero(owner == self.layout.rank).flatten() + lengths = (end - begin)[mine] + my_end = torch.cumsum(lengths, 0) + my_begin = my_end - lengths + elems = torch.repeat_interleave(begin[mine] - my_begin, lengths) + elems = elems + torch.arange(elems.numel(), device=rows.device) + counts, ordinal = _ordinal_by_group(owner, self.layout.sp) + width = int(counts.max()) if owner.numel() else 0 + return { + "tiles": (buffer_index[rows[elems]].contiguous(), my_begin, my_end), + "width": width, # tiles per rank in the padded partial gather + "gather_index": owner * width + ordinal, # WS1 tile order within that gather + } + + def exchange(self, collective, *tensors: torch.Tensor) -> list[torch.Tensor]: + """For each ``(M_local, ...)`` tensor: local rows, then the gathered send blocks.""" + + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_zeros((self.max_send, *t.shape[1:])) + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family: int, partial: torch.Tensor) -> torch.Tensor: + """This rank's tile partials -> every tile's partial, in WS1 tile order.""" + + fam = self.families[family] + padded = partial.new_zeros((fam["width"], *partial.shape[1:])) + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + +def _dweight_tiles(total: int, device): + begin = torch.arange(0, total, BACKWARD_TILE, device=device) + return torch.arange(total, device=device), begin, torch.clamp(begin + BACKWARD_TILE, max=total) + + +@dataclass(frozen=True) +class SPPlan: + """Everything an SP backward needs that depends only on the layout and the row index. + + Built once per ``(row index, table rows)`` and shared by every norm and gated + residual that uses that index (every block of a forward pass). + """ + + layout: SPRowLayout + plan: _Plan + local_index: torch.Tensor | None # (S_local,) table row of each local position + index_buf: torch.Tensor | None # table row of each exchange-buffer row + seg_first_tile: torch.Tensor | None + num_rows: int | None + + +def sp_plan( + layout: SPRowLayout, index_full=None, num_rows: int | None = None, *, device=None +) -> SPPlan: + """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" + + device = index_full.device if index_full is not None else torch.device(device or "cuda") + families = [_dweight_tiles(layout.batch * layout.seq_len, device)] + if index_full is None: + return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + if num_rows is None: + raise ValueError("indexed SP plans require num_rows") + *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) + plan = _Plan(layout, [*families, tuple(seg_tiles)], device) + local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) + received = index_full[plan.recv_global % layout.seq_len] + index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() + return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) + + +def sp_rmsnorm_backward(grad, x, weight, rstd, shift, scale, sp: SPPlan, collective): + """One rank's backward on its ``(B * S_local, H)`` rows: ``(dx, dweight, d_shift, d_scale)``. + + ``dx`` is row-local; the other gradients are the full WS1 reductions, identical on + every rank. ``shift``/``scale`` are ``None`` for the plain norm. + """ + + weight, plan, modulated = weight.contiguous(), sp.plan, sp.local_index is not None + dx = _C.h3_rmsnorm_backward_dx(grad, x, weight, rstd, shift, scale, sp.local_index) + g_buf, x_buf, rstd_buf = plan.exchange(collective, grad, x, rstd) + seg = plan.families[1]["tiles"] if modulated else (None, None, None) + partials = _C.h3_rmsnorm_backward_partials( + g_buf, x_buf, weight, rstd_buf, shift, scale, sp.index_buf, + *plan.families[0]["tiles"], *seg, + ) # fmt: skip + gathered = [plan.gather_partials(collective, i, p) for i, p in enumerate(partials)] + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_rmsnorm_backward_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + if not modulated: + (dweight,) = _C.h3_rmsnorm_fold_partials(gathered[0], weight) + return dx, dweight, None, None + dweight, d_shift, d_scale = _C.h3_rmsnorm_fold_partials( + gathered[0], weight, gathered[1], sp.seg_first_tile + ) + return dx, dweight, d_shift.to(shift.dtype), d_scale.to(scale.dtype) + + +def sp_gate_residual_backward(grad, y, gate, sp: SPPlan, collective): + """One rank's backward on its rows: ``(d_sublayer, d_gate)``; ``d_residual`` is ``grad``.""" + + plan = sp.plan + dy = _C.h3_gate_residual_backward_dy(grad, gate, sp.local_index) + g_buf, y_buf = plan.exchange(collective, grad, y) + partial = _C.h3_gate_grad_partials(g_buf, y_buf, *plan.families[1]["tiles"]) + dgate = _C.h3_gate_grad_fold( + plan.gather_partials(collective, 1, partial), sp.seg_first_tile, gate.dtype + ) + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_gate_grad_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + return dy, dgate + + +class _SPRMSNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, sp, collective, eps): + x2 = x.contiguous().view(-1, x.shape[-1]) + y, rstd = _C.h3_rmsnorm_forward( + x2, weight.contiguous(), float(eps), shift, scale, sp.local_index + ) + ctx.save_for_backward(x2, weight, rstd, shift, scale) + ctx.sp, ctx.collective, ctx.x_shape = sp, collective, x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + dx, dweight, d_shift, d_scale = sp_rmsnorm_backward( + g2, x2, weight, rstd, shift, scale, ctx.sp, ctx.collective + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None, None + + +class _SPGateResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, sp, collective): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, sp.local_index) + ctx.save_for_backward(y2, gate) + ctx.sp, ctx.collective, ctx.shape = sp, collective, residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + dy, dgate = sp_gate_residual_backward(g2, y2, gate, ctx.sp, ctx.collective) + return grad, dy.view(ctx.shape), dgate, None, None + + +class H3SPNormAdaLNCudaOp: + """Sequence-parallel ``h3_rmsnorm`` (+ modulation) and ``adaln_gate_residual``. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` and + ``backend_id`` (e.g. ``DeterministicCollective``). Activations are this rank's + ``(B, S_local, H)`` rows; the norm weight, the table views and the ``(S,)`` row + index are the full, replicated tensors. Gradients of the replicated tensors come + back identical on every rank, equal to WS1. + """ + + op_class = "reduction" + backward_impl = BACKWARD_IMPL + + def __init__(self, collective, seq_len: int, batch: int) -> None: + if not sp_norm_adaln_available(): + raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") + self.collective = collective + self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None, *, device=None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + device = index.device if index is not None else torch.device(device or "cuda") + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + version = -1 if index is None else index._version + for obj, ver, rows, cached_device, plan in self._plans: + if obj is index and ver == version and rows == num_rows and cached_device == device: + return plan + plan = sp_plan(self.layout, index, num_rows, device=device) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, device, plan), *self._plans[:3]] + return plan + + def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: + expected = (self.layout.batch, self.layout.local_len) + if not isinstance(x, torch.Tensor) or x.dim() != 3 or tuple(x.shape[:2]) != expected: + shape = tuple(x.shape) if isinstance(x, torch.Tensor) else type(x).__name__ + raise ValueError(f"activations must be this rank's {expected} + (H,) rows, got {shape}") + if not x.is_cuda: + raise ValueError("H3SPNormAdaLNCudaOp needs CUDA tensors") + if index is None: + return + if index.dim() != 1 or index.shape[0] != self.layout.seq_len: + raise ValueError( + f"index must be the full ({self.layout.seq_len},) row index, " + f"got {tuple(index.shape)}" + ) + if num_rows is None: + raise ValueError("indexed SP validation requires num_rows") + if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): + raise IndexError(f"index must be in [0, {num_rows})") + + def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + self._check(x, None, None) + validate_h3_rmsnorm(x, weight, eps) + return _SPRMSNorm.apply( + x, weight, None, None, self.plan(None, None, device=x.device), self.collective, eps + ) + + def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + self._check(x, index, shift.shape[0]) + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index[self.layout.lo : self.layout.hi]) + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + sp = self.plan(index, shift.shape[0]) + return _SPRMSNorm.apply(x, weight, shift, scale, sp, self.collective, eps) + + def gate_residual(self, residual, y, gate, index) -> torch.Tensor: + self._check(residual, index, gate.shape[0]) + validate_h3_gate_residual(residual, y, gate, index[self.layout.lo : self.layout.hi]) + if gate.stride(1) != 1: + gate = gate.contiguous() + sp = self.plan(index, gate.shape[0]) + return _SPGateResidual.apply(residual, y, gate, sp, self.collective) + + def readback(self) -> dict: + return { + "op": "sp_norm_adaln", + "backward_impl": BACKWARD_IMPL, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_256_element_tiles_then_ascending_fold", + "sp": self.layout.sp, + "rank": self.layout.rank, + "positions": [self.layout.lo, self.layout.hi], + "batch": self.layout.batch, + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..4d58f6d9f --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 FP32 timestep MLP (RFC #420 ``timestep_mlp_fp32``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + ACT_SILU, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3TimestepMLPCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2): + """Evaluate the FP32 CUDA linear-SiLU-linear graph and save VJP inputs.""" + + x = x.contiguous() + hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) + (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) + ctx.save_for_backward(x, w1, w2, pre, hidden) + return out + + @staticmethod + def backward(ctx, grad_out): + """Return requested FP32 input and parameter VJPs with fixed reduction order.""" + + x, w1, w2, pre, hidden = ctx.saved_tensors + need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad + grad_out = grad_out.float().contiguous() + dw2 = db2 = dw1 = db1 = dx = None + if need_w2 or need_b2: + dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + d_pre = None + if need_x or need_w1 or need_b1: + d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) + d_pre = silu_backward_fp32(d_hidden, pre).contiguous() + if need_w1 or need_b1: + dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + if need_x: + dx = det_linear_backward_input(d_pre, w1, torch.float32) + record_backward( + "timestep_mlp_fp32", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return ( + dx, + dw1 if need_w1 else None, + db1 if need_b1 else None, + dw2 if need_w2 else None, + db2 if need_b2 else None, + ) + + +class H3TimestepMLPCudaOp: + """CUDA candidate: two warp-per-column GEMVs, SiLU fused into the first. + + The weights dominate the traffic (5.5 MB + 57.8 MB FP32), the rows are the + handful of distinct timesteps, and every output element follows contract + ``h3-det-linear-v1``, so a timestep's ``temb`` bytes do not depend on which + other timesteps share the call. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` timestep embeddings for same-device CUDA inputs.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Apply two deterministic FP32 CUDA projections with an intervening SiLU. + + Require same-device FP32 tensors: ``x`` is nonempty ``(T, K)``, + weights are ``(H, K)`` and ``(D, H)``, and biases are ``(H,)`` and + ``(D,)``. Return ``(T, D)`` with autograd support for all five inputs. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + if not x.is_cuda: + raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") + return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Use the standard CUDA MLP, whose inputs and output are already FP32.""" + + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..65934e1e9 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 sinusoidal timestep features (RFC #420 ``timestep_sinusoid_h3``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD +from rl_engine.reference.minimax_h3.fixed_order import tree_sum_lastdim_fp32 +from rl_engine.reference.minimax_h3.timestep_sinusoid import ( + NativeH3TimestepSinusoidOp, + validate_h3_timesteps, +) + +KERNEL_ID = "rl_engine._C.h3_timestep_sinusoid_forward" +BACKWARD_IMPL = "row_local_fp32_analytic_tree_sum" + + +def h3_sinusoid_cuda_available() -> bool: + """Return whether the extension provides the CUDA H3 sinusoid entry point.""" + + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +class _H3TimestepSinusoidCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): + """Compute FP32 CUDA features and save them for the analytic timestep VJP.""" + + t32 = timestep.detach().float().contiguous() + out = _C.h3_timestep_sinusoid_forward( + t32, int(num_channels), float(H3_MAX_PERIOD), check_range + ) + ctx.save_for_backward(out) + ctx.num_channels = int(num_channels) + ctx.timestep_dtype = timestep.dtype + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + """Reduce channel derivatives with a fixed FP32 tree into the timestep dtype.""" + + (out,) = ctx.saved_tensors + half = ctx.num_channels // 2 + freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) + cos_part, sin_part = out[:, :half], out[:, half:] + g = grad_out.float() + # d cos(t f)/dt = -f sin(t f); d sin(t f)/dt = f cos(t f). Row-local, + # summed over channels with a fixed pairwise tree. + contrib = torch.cat( + [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 + ) + grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) + record_backward( + "timestep_sinusoid_h3", + kernel_id="rl_engine.reference.minimax_h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None, None + + +class H3TimestepSinusoidCudaOp: + """CUDA candidate: one thread per (timestep, frequency), FP32 output. + + Construction fails when the extension lacks the symbol, so the registry + falls back to the PyTorch reference instead of silently mis-dispatching. + """ + + op_class = "elementwise" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + """Require the native H3 CUDA sinusoid symbol before creating this operator.""" + + if not h3_sinusoid_cuda_available(): + raise RuntimeError( + "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu (it is omitted " + "from KERNEL_ALIGN_USE_FAST_MATH=1 builds to keep bitwise parity)" + ) + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 CUDA features for timesteps in ``[0, 1]`` with gradients.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on CUDA. + + Require nonempty 1-D FP32, FP16 or BF16 timesteps and a positive even + channel count. Native validation rejects nonfinite/out-of-range values + when ``check_range`` is true; backward returns the original input dtype. + """ + + # Native code owns value validation, avoiding a second host sync here. + validate_h3_timesteps(timestep, num_channels, check_range=False) + if not timestep.is_cuda: + raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") + return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) + + def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Use the standard CUDA sinusoid path, which already returns FP32 features.""" + + # The op is FP32 end to end; the FP32 path is the op itself. + return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py new file mode 100644 index 000000000..c4ceed334 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: + if tp < 1 or not 0 <= rank < tp: + raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total <= 0: + raise ValueError(f"N must be positive, got {n_total}") + if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): + raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") + if n_total % (tp * DINPUT_CHUNK): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/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/distributed/algorithms/hierarchical.py b/rl_engine/distributed/algorithms/hierarchical.py new file mode 100644 index 000000000..253b34399 --- /dev/null +++ b/rl_engine/distributed/algorithms/hierarchical.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Cross-node deterministic collective: node-local IPC fixed tree, then one exact gather. + +``DeterministicCollective`` (``cuda_ipc_fixed_tree``) reads peers' memory through CUDA IPC, +so it is single node. This wrapper (``hierarchical_ipc_nccl_tree``) reaches the same bits on +several nodes without relying on any NCCL reduction order: + +1. every group of ``local_size`` consecutive ranks reduces with the node-local fixed tree; +2. the group partials move with NCCL ``all_gather`` (bit-exact data movement, no arithmetic); +3. each rank adds the group partials in one fixed, balanced order. + +With ``world_size = nodes * local_size`` and ``local_size``, ``nodes`` powers of two, the +tree ``((0+1)+(2+3)) + ((4+5)+(6+7))`` of an 8-way node-local reduction is exactly "the 4-way +tree of each half, then the two halves added". The input dtype is the accumulation dtype, as in +``DeterministicCollective`` (FP32 inputs keep every partial in FP32, which is what the H3 +parallel operators use), so the result is bitwise the single-node tree of the same world size. +""" + +from __future__ import annotations + +import torch +import torch.distributed as dist + +from rl_engine.distributed.algorithms.collectives import DeterministicCollective + +_GROUP_COUNTS = (1, 2, 4, 8) + + +def tree_add(parts: list[torch.Tensor]) -> torch.Tensor: + """Balanced pairwise sum ``((p0+p1)+(p2+p3))+...`` of a power-of-two number of tensors.""" + + if len(parts) & (len(parts) - 1) or not parts: + raise ValueError(f"tree reduction needs a power-of-two part count, got {len(parts)}") + while len(parts) > 1: + parts = [parts[i] + parts[i + 1] for i in range(0, len(parts), 2)] + return parts[0] + + +class HierarchicalCollective: + """TP-invariant all-reduce / reduce-scatter over ``world_size`` ranks on several nodes.""" + + backend_id = "hierarchical_ipc_nccl_tree" + + def __init__( + self, + local_size: int, + group: dist.ProcessGroup | None = None, + device: torch.device | str | int | None = None, + *, + max_size_bytes: int = 64 << 20, + ) -> None: + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before collectives") + self.group = group if group is not None else dist.group.WORLD + self.world_size = dist.get_world_size(group=self.group) + self.rank = dist.get_rank(group=self.group) + if self.world_size % local_size or local_size not in _GROUP_COUNTS: + raise ValueError( + f"local_size {local_size} must be 1/2/4/8 and divide world {self.world_size}" + ) + self.local_size = local_size + self.num_groups = self.world_size // local_size + if self.num_groups not in _GROUP_COUNTS: + raise ValueError(f"{self.num_groups} node groups; need 1, 2, 4 or 8") + self.group_index, self.local_rank = divmod(self.rank, local_size) + # Every rank creates every group in the same order (a torch.distributed requirement). + ranks = dist.get_process_group_ranks(self.group) # global rank of each group member + local_groups = [ + dist.new_group(ranks[g * local_size : (g + 1) * local_size]) + for g in range(self.num_groups) + ] + cross_groups = [ + dist.new_group([ranks[g * local_size + r] for g in range(self.num_groups)]) + for r in range(local_size) + ] + self._cross = cross_groups[self.local_rank] + self._local = DeterministicCollective( + local_groups[self.group_index], device, max_size_bytes=max_size_bytes + ) + + def all_reduce(self, input: torch.Tensor) -> torch.Tensor: + """The fixed-tree sum of every rank's ``input`` on every rank (new tensor).""" + + if input.dtype not in (torch.float32, torch.bfloat16, torch.float16): + raise TypeError(f"unsupported dtype {input.dtype}") + wide = ( + input.contiguous() + ) # the input dtype is the accumulation dtype, as in the node-local tree + partial = self._local.all_reduce(wide) if self.local_size > 1 else wide.clone() + if self.num_groups > 1: + gathered = [torch.empty_like(partial) for _ in range(self.num_groups)] + dist.all_gather(gathered, partial, group=self._cross) + partial = tree_add(gathered) + return partial + + def reduce_scatter(self, input: torch.Tensor) -> torch.Tensor: + """``all_reduce`` then this rank's equal share of dimension 0.""" + + if input.shape[0] % self.world_size: + raise ValueError("dimension 0 must divide by the world size") + rows = input.shape[0] // self.world_size + return self.all_reduce(input)[self.rank * rows : (self.rank + 1) * rows].contiguous() + + def all_gather(self, input: torch.Tensor) -> torch.Tensor: + """Rank-ordered concatenation along dimension 0 (exact data movement).""" + + out = [torch.empty_like(input) for _ in range(self.world_size)] + dist.all_gather(out, input.contiguous(), group=self.group) + return torch.cat(out) + + def close(self) -> None: + self._local.close() diff --git a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py new file mode 100644 index 000000000..3d48b8caf --- /dev/null +++ b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``vllm_runtime_adapter_h3``: strict H3 projections inside a vLLM-Omni DiT. + +vLLM-Omni's ``MiniMaxH3Attention`` / MLP keep fused linears: ``attn.qkv_proj`` (``[Q;K;V]`` rows), +``attn.out_proj`` (row parallel), ``mlp.fc1`` (gate-first), ``mlp.fc2`` (row parallel); each +returns ``(output, bias)``. ``install_h3_rollout(model)`` replaces those four forwards with the +fixed-order kernels so rollout reads the same bits as the strict training graph: + +* ``qkv_proj`` / ``fc1``: ``h3-projection-gemm-v1`` over the fused rows (a column of a fused + projection equals the same column computed alone, so fused == separate bitwise); +* ``out_proj`` / ``fc2``: the canonical 8-unit row-parallel reduction of ``h3-parallel-v1`` + (``tp`` ranks with a deterministic collective, ``tp=1`` alone), one BF16 rounding. + +Attention, rotary and norm stay vLLM-Omni's (their rows are separate); the report lists them as +``external``. Everything is restored by ``handle.remove()``. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.backends.shared.triton.gemm import h3_projection as g + +LINEAR_NAMES = {"qkv_proj": "column", "fc1": "column", "out_proj": "row", "fc2": "row"} + + +@dataclass +class H3RolloutHandle: + patched: list[str] = field(default_factory=list) + _restore: list = field(default_factory=list) + + def remove(self) -> None: + for module in self._restore: + module.__dict__.pop("forward", None) + self._restore.clear() + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.remove() + + def report(self) -> dict: + return { + "contract": par.CONTRACT, + "patched": list(self.patched), + "external": ["attention", "rotary", "qk_norm", "swiglu", "adaln"], + } + + +def _wrap(module, kind, collective, tp): + def forward(x): + shape = x.shape + flat = x.reshape(-1, shape[-1]).contiguous() + if kind == "column": + out = g.gemm(flat, module.weight, b_t=True) + else: + out = par.row_parallel_linear(flat, module.weight, collective, tp) + return out.view(*shape[:-1], out.shape[-1]), None + + return forward + + +def install_h3_rollout( + model: torch.nn.Module, *, collective=None, tp: int = 1, strict: bool = True +) -> H3RolloutHandle: + """Patch every ``qkv_proj`` / ``out_proj`` / ``fc1`` / ``fc2`` of ``blocks`` and the refiner.""" + + if tp > 1 and collective is None: + raise ValueError("tp > 1 needs a deterministic collective") + handle = H3RolloutHandle() + for name, module in model.named_modules(): + leaf = name.rsplit(".", 1)[-1] + kind = LINEAR_NAMES.get(leaf) + if kind is None or not hasattr(module, "weight"): + continue + weight = module.weight + if not weight.is_cuda or weight.dtype != torch.bfloat16: + if strict: + raise RuntimeError( + f"{name}: strict H3 rollout needs CUDA BF16 weights, got " + f"{weight.dtype} on {weight.device}" + ) + continue + module.forward = _wrap(module, kind, collective, tp) + handle._restore.append(module) + handle.patched.append(name) + if strict and not handle.patched: + raise RuntimeError("no H3 projection modules found (qkv_proj/out_proj/fc1/fc2)") + return handle 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/checkpoint_mapper.py b/rl_engine/models/minimax_h3/checkpoint_mapper.py new file mode 100644 index 000000000..945900e56 --- /dev/null +++ b/rl_engine/models/minimax_h3/checkpoint_mapper.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_checkpoint_mapper``: Diffusers <-> vLLM-Omni native H3 names and layouts. + +The training side (diffusers, FSDP2) and the rollout side (vLLM-Omni) hold the same 638 +tensors under different names and fused layouts. This module is the single, tested mapping, so +the strict operators on both sides read bit-identical weights: + +* ``transformer_blocks.N`` -> ``blocks.N``; + ``token_refiner.refiner_blocks.N`` -> ``token_refiner.blocks.N``; +* ``attn.to_q/to_k/to_v`` -> one fused ``attn.qkv_proj`` laid out ``[Q; K; V]`` (rows), the + layout ``h3_qkv_gemm.forward_fused`` is bitwise-equal to; +* ``attn.to_out.0`` -> ``attn.out_proj``, ``attn.norm_q/k`` -> ``attn.q_norm/k_norm``; +* ``ff.net.0.proj`` (value-first ``[hidden; gate]``) -> ``mlp.fc1`` (gate-first, halves swapped), + ``ff.net.2`` -> ``mlp.fc2``; +* model-level names per vLLM-Omni's ``fasth3`` table (``proj_in`` -> ``video_patch_proj`` ...). + +The tables mirror ``vllm_omni.diffusion.models.minimax_h3.fasth3`` (checked by a test when that +package is importable). Mapping is pure data movement: +``native_to_diffusers(diffusers_to_native(x))`` is bitwise ``x``. +""" + +from __future__ import annotations + +from collections.abc import Mapping + +import torch + +PLAIN, SWAP_HALVES = "plain", "swap_halves" +QKV_SLOTS = ("q", "k", "v") + +MODEL_LEVEL = { + "proj_in": "video_patch_proj", + "proj_out": "final_layer.video_out", + "audio_proj_in": "audio_patch_proj", + "audio_proj_out": "final_layer.audio_out", + "context_embedder": "condition_proj", + "time_embedder.linear_1": "time_embedder.proj_in", + "time_embedder.linear_2": "time_embedder.proj_out", + "norm_out.linear": "final_layer.adaln_proj.linear", + "norm_out.norm": "final_layer.norm", + "token_refiner.final_norm": "token_refiner.final_norm", +} +BLOCK_PREFIXES = ( + ("token_refiner.refiner_blocks.", "token_refiner.blocks."), + ("transformer_blocks.", "blocks."), +) +BLOCK_TARGETS = { + "attn.to_q": ("attn.qkv_proj", "q"), + "attn.to_k": ("attn.qkv_proj", "k"), + "attn.to_v": ("attn.qkv_proj", "v"), + "attn.to_out.0": ("attn.out_proj", PLAIN), + "attn.norm_q": ("attn.q_norm", PLAIN), + "attn.norm_k": ("attn.k_norm", PLAIN), + "ff.net.0.proj": ("mlp.fc1", SWAP_HALVES), + "ff.net.2": ("mlp.fc2", PLAIN), + "adaln_proj.linear": ("adaln_proj.linear", PLAIN), + "norm1": ("norm1", PLAIN), + "norm2": ("norm2", PLAIN), +} + + +def resolve(module: str) -> tuple[str, str] | None: + """Diffusers module path -> ``(native module path, layout)`` or ``None``.""" + + if module in MODEL_LEVEL: + return MODEL_LEVEL[module], PLAIN + for source, target in BLOCK_PREFIXES: + if module.startswith(source): + index, _, suffix = module[len(source) :].partition(".") + mapped = BLOCK_TARGETS.get(suffix) + if index.isdigit() and mapped is not None: + return f"{target}{index}.{mapped[0]}", mapped[1] + return None + + +def swap_halves(tensor: torch.Tensor) -> torch.Tensor: + first, second = tensor.chunk(2, dim=0) + return torch.cat((second, first), dim=0) + + +def fuse_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + """``[Q; K; V]`` rows (the serving layout).""" + + return torch.cat([q, k, v], dim=0) + + +def split_qkv(fused: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if fused.shape[0] % 3: + raise ValueError(f"fused QKV rows {fused.shape[0]} are not a multiple of 3") + q, k, v = fused.chunk(3, dim=0) + return q, k, v + + +def grouped_qkv_to_qkv(weight: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor: + """Native checkpoint's per-head ``[q,k,v]`` groups -> ``[Q; K; V]`` (vLLM-Omni's loader).""" + + per_group = 3 * head_dim + if weight.shape[0] != heads * per_group: + raise ValueError(f"grouped QKV rows {weight.shape[0]} != {heads * per_group}") + grouped = weight.reshape(heads, per_group, *weight.shape[1:]) + q, k, v = grouped.split(head_dim, dim=1) + return torch.cat([t.reshape(heads * head_dim, *weight.shape[1:]) for t in (q, k, v)]) + + +def qkv_to_grouped(fused: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor: + q, k, v = split_qkv(fused) + parts = [t.reshape(heads, head_dim, *fused.shape[1:]) for t in (q, k, v)] + return torch.cat(parts, dim=1).reshape(heads * 3 * head_dim, *fused.shape[1:]) + + +def diffusers_to_native(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Diffusers state dict -> native names, fused QKV, gate-first ``fc1``; unmapped keys fail.""" + + out: dict[str, torch.Tensor] = {} + qkv: dict[str, dict[str, torch.Tensor]] = {} + for name, tensor in state.items(): + module, _, kind = name.rpartition(".") + target = resolve(module) + if target is None or kind not in ("weight", "bias"): + raise ValueError(f"unmapped Diffusers H3 tensor: {name}") + native, layout = target + key = f"{native}.{kind}" + if layout in QKV_SLOTS: + qkv.setdefault(key, {})[layout] = tensor + elif layout == SWAP_HALVES: + out[key] = swap_halves(tensor) + else: + out[key] = tensor + for key, parts in qkv.items(): + if set(parts) != set(QKV_SLOTS): + raise ValueError(f"{key} needs q, k and v, got {sorted(parts)}") + out[key] = fuse_qkv(parts["q"], parts["k"], parts["v"]) + return out + + +def _native_block(module: str) -> tuple[str, str, str] | None: + for source, target in BLOCK_PREFIXES: + if module.startswith(target): + index, _, suffix = module[len(target) :].partition(".") + if index.isdigit(): + return source, index, suffix + return None + + +def native_to_diffusers(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """Inverse of :func:`diffusers_to_native` (splits fused QKV, swaps ``fc1`` halves back).""" + + model = {native: diff for diff, native in MODEL_LEVEL.items()} + block = { + native: (diff, layout) + for diff, (native, layout) in BLOCK_TARGETS.items() + if layout not in QKV_SLOTS + } + out: dict[str, torch.Tensor] = {} + for name, tensor in state.items(): + module, _, kind = name.rpartition(".") + if module in model: + out[f"{model[module]}.{kind}"] = tensor + continue + parsed = _native_block(module) + if parsed is None: + raise ValueError(f"unmapped native H3 tensor: {name}") + source, index, suffix = parsed + if suffix == "attn.qkv_proj": + for slot, part in zip(QKV_SLOTS, split_qkv(tensor)): + out[f"{source}{index}.attn.to_{slot}.{kind}"] = part + elif suffix in block: + diff, layout = block[suffix] + out[f"{source}{index}.{diff}.{kind}"] = ( + swap_halves(tensor) if layout == SWAP_HALVES else tensor + ) + else: + raise ValueError(f"unmapped native H3 tensor: {name}") + return out 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/runtime.py b/rl_engine/models/minimax_h3/runtime.py new file mode 100644 index 000000000..85d31e5ba --- /dev/null +++ b/rl_engine/models/minimax_h3/runtime.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 runtime modes, ablation overlay, Diffusers adapter, training provider and LoRA gradient gate. + +Rows: ``h3_runtime_modes``, ``h3_ablation_overlay``, ``diffusers_runtime_adapter_h3``, +``h3_training_provider``, ``h3_lora_gradient_gate``. + +* ``h3_runtime_policy(mode)`` maps ``RL_KERNEL_MODE`` (strict / audit / auto / off) to what the H3 + stack does: bind the RL-Kernel graph or not, fail on a fallback or report it, and the case id + (``R/R`` strict and audit, ``P/P`` auto and off) the experiment records. +* ``AblationOverlay`` keeps selected block nodes or boundary modules on provider arithmetic while + everything else is strict, labels the run ``ablation(...)`` (never ``R/R``) and validates names. +* ``H3DiffusersAdapter`` binds a ``MiniMaxH3Transformer3DModel`` for the duration of a + ``with`` block (forward *and* backward, so checkpoint recomputation sees the same operators) + and reports what was bound. +* ``H3TrainingProvider`` is the actor-side entry point: ``step()`` scope, strict forward / + backward, and the layout-invariant CPS transition for log-probs. +* ``lora_gradient_gate`` checks that a LoRA over a strict projection trains bitwise + reproducibly and layout invariantly with a frozen base. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from contextlib import contextmanager +from dataclasses import dataclass + +import torch + +from rl_engine.runtime.policy import RLKernelMode, rl_kernel_mode_policy + +CONTRACT = "h3-runtime-v1" + + +@dataclass(frozen=True) +class H3RuntimePolicy: + mode: RLKernelMode + bind: bool # run the RL-Kernel graph instead of provider arithmetic + fail_on_fallback: bool + report_only: bool + case_id: str + + +def h3_runtime_policy(mode: str | RLKernelMode | None = None) -> H3RuntimePolicy: + """Policy for ``mode`` (default ``RL_KERNEL_MODE``, else strict).""" + + base = rl_kernel_mode_policy(mode) + # strict and audit bind the RL-Kernel graph (R/R); auto and off leave the provider (P/P) + bind = base.enabled and base.provider_mode == "strict" + return H3RuntimePolicy(base.mode, bind, base.fail_on_fallback, base.report_only, base.case_id) + + +@dataclass(frozen=True) +class AblationOverlay: + """Parts that stay on provider arithmetic inside an otherwise strict run.""" + + native_nodes: frozenset[str] = frozenset() + native_boundaries: frozenset[str] = frozenset() + + def __post_init__(self) -> None: + from rl_engine.validation.models import h3_block + from rl_engine.validation.models.h3_full_model import STRICT_BOUNDARY_ROWS + + bad_nodes = set(self.native_nodes) - set(h3_block.NODE_NAMES) + bad_boundaries = set(self.native_boundaries) - set(STRICT_BOUNDARY_ROWS) + if bad_nodes or bad_boundaries: + raise ValueError( + f"unknown ablation targets: nodes {sorted(bad_nodes)}, " + f"boundaries {sorted(bad_boundaries)}" + ) + + @classmethod + def parse(cls, spec: str) -> AblationOverlay: + """``"node:attention,boundary:token_refiner"`` -> overlay; empty -> no ablation.""" + + nodes, boundaries = set(), set() + for item in filter(None, (part.strip() for part in spec.split(","))): + kind, _, name = item.partition(":") + if kind == "node": + nodes.add(name) + elif kind == "boundary": + boundaries.add(name) + else: + raise ValueError( + f"ablation item must be node: or boundary:, got {item!r}" + ) + return cls(frozenset(nodes), frozenset(boundaries)) + + @property + def empty(self) -> bool: + return not (self.native_nodes or self.native_boundaries) + + @property + def label(self) -> str: + if self.empty: + return "R/R" + parts = [f"node:{n}" for n in sorted(self.native_nodes)] + parts += [f"boundary:{n}" for n in sorted(self.native_boundaries)] + return "ablation(native: " + ",".join(parts) + ")" + + +class H3DiffusersAdapter: + """Bind a Diffusers H3 transformer according to the runtime policy and an optional overlay.""" + + def __init__( + self, + model, + mode: str | RLKernelMode | None = None, + overlay: AblationOverlay | None = None, + ops=None, + ) -> None: + from rl_engine.validation.models.h3_full_model import FullModelTrace + + self.model, self.policy = model, h3_runtime_policy(mode) + self.overlay = overlay or AblationOverlay() + self.ops, self.trace = ops, FullModelTrace() + self.fell_back: str | None = None + self._scope = None + + def _bind(self) -> bool: + if not self.policy.bind: + return False + if not (torch.cuda.is_available() and next(self.model.parameters()).is_cuda): + if self.policy.fail_on_fallback: + raise RuntimeError("strict H3 needs CUDA parameters; refusing to fall back") + self.fell_back = "no CUDA device" + return False + return True + + @contextmanager + def active(self): + """Operators bound until the block exits; keep ``backward()`` inside it.""" + + from rl_engine.validation.models.h3_full_model import full_model_execution + + bind = self._bind() + try: + with full_model_execution( + self.model, + candidate=bind, + trace=self.trace, + ops=self.ops, + provider_nodes=tuple(self.overlay.native_nodes), + strict_boundaries=bind, + native_boundaries=tuple(self.overlay.native_boundaries), + ): + yield self + except RuntimeError as exc: + if self.policy.fail_on_fallback or not bind: + raise + self.fell_back = f"{type(exc).__name__}: {exc}" + raise + + def report(self) -> dict: + bound = self.policy.bind and self.fell_back is None + return { + "contract": CONTRACT, + "mode": self.policy.mode.value, + "case_id": ("R/R" if self.overlay.empty else "ablation") if bound else "P/P", + "bound": bound, + "fell_back": self.fell_back, + "label": self.overlay.label if bound else "P/P", + "native_nodes": sorted(self.overlay.native_nodes), + "native_boundaries": sorted(self.overlay.native_boundaries), + "boundary_bindings": self.trace.boundary_bindings, + } + + +class H3TrainingProvider: + """Actor-side strict training scope: ``with provider.step(): forward(...); loss.backward()``.""" + + def __init__( + self, model, mode: str | RLKernelMode | None = None, overlay: AblationOverlay | None = None + ) -> None: + self.adapter = H3DiffusersAdapter(model, mode, overlay) + self.model = model + + @contextmanager + def step(self): + with self.adapter.active(): + yield self + + def forward(self, **inputs): + return self.model(**inputs) + + def transition( + self, + sample, + velocity, + sigma, + sigma_prev, + *, + seed, + step, + sample_ids, + modality, + noise_level=None, + prev_sample=None, + ): + """CPS transition for a batch whose per-sample noise depends on ``sample_ids`` only.""" + + from rl_engine.models.minimax_h3 import transition as tr + + level = tr.NOISE_LEVEL if noise_level is None else noise_level + noise = None + if prev_sample is None: + noise = torch.stack( + [ + tr.step_noise( + seed, step, int(i), modality, tuple(sample.shape[1:]), sample.device + ) + for i in sample_ids + ] + ) + return tr.cps_transition( + sample.float(), + velocity.float(), + sigma, + sigma_prev, + noise_level=level, + noise=noise, + prev_sample=prev_sample, + ) + + +# --- LoRA over a strict projection ---------------------------------------------------------- + + +class _StrictLoRA(torch.autograd.Function): + """``y = round_bf16(x W^T + scale * (x A^T) B^T)``; W frozen, A and B trainable (FP32).""" + + @staticmethod + def forward(ctx, x, weight, a, b, scale): + from rl_engine.backends.shared.triton.gemm import h3_projection as g + + ctx.save_for_backward(x, weight, a, b) + ctx.scale = scale + base = g.gemm(x, weight, b_t=True, out_dtype=torch.float32) + x32 = x.float() + low = g.gemm(g.gemm(x32, a, b_t=True), b, b_t=True) + return (base + scale * low).to(x.dtype) + + @staticmethod + def backward(ctx, grad): + from rl_engine.backends.shared.triton.gemm import h3_projection as g + + x, weight, a, b = ctx.saved_tensors + g32 = grad.float().contiguous() + x32 = x.float().contiguous() + mid = g.gemm(x32, a, b_t=True) # (T, r) + d_mid = g.gemm(g32, b) * ctx.scale # (T, r) + d_b = g.gemm(g32, mid, a_t=True) * ctx.scale # (N, r) + d_a = g.gemm(d_mid, x32, a_t=True) # (r, K) + dx = g.gemm(grad.contiguous(), weight, out_dtype=torch.float32) + g.gemm(d_mid, a) + return dx.to(x.dtype), None, d_a, d_b, None + + +def strict_lora_linear(x, weight, a, b, scale: float = 1.0): + """LoRA projection of ``x (T, K)``: frozen BF16 ``weight``, FP32 ``a (r, K)``, ``b (N, r)``.""" + + return _StrictLoRA.apply(x.contiguous(), weight, a, b, scale) + + +def lora_gradient_gate( + weight: torch.Tensor, + x: torch.Tensor, + grad: torch.Tensor, + *, + rank: int = 8, + scale: float = 2.0, + seed: int = 0, +) -> Mapping[str, object]: + """Gate: A/B gradients are bitwise repeatable and independent of how rows are batched, the + base gets no gradient, and the result is close to an FP64 golden. + + ``x (T, K)`` rows are also run split in two halves (gradients of the halves summed in the fixed + order of the full GEMM are not claimed equal; the gate checks the *forward* of a half and the + full-batch gradients, which is what a data-parallel micro-batch split changes). + """ + + gen = torch.Generator().manual_seed(seed) + n, k = weight.shape + a0 = (torch.randn(rank, k, generator=gen) / k**0.5).to(x.device) + b0 = (torch.randn(n, rank, generator=gen) * 0.1).to(x.device) + + frozen: list[bool] = [] + + def run(rows): + a, b = a0.clone().requires_grad_(True), b0.clone().requires_grad_(True) + w = weight.detach().clone() + xi = x[rows].detach().clone().requires_grad_(True) + y = strict_lora_linear(xi, w, a, b, scale) + y.backward(grad[rows]) + frozen.append(w.grad is None and not w.requires_grad) + return y.detach(), xi.grad, a.grad, b.grad + + full = run(slice(None)) + again = run(slice(None)) + half = run(slice(0, x.shape[0] // 2)) + ref_y = ( + x.double() @ weight.double().t() + scale * (x.double() @ a0.double().t()) @ b0.double().t() + ) + err = ((full[0].double() - ref_y).norm() / ref_y.norm()).item() + return { + "repeatable": all(torch.equal(p, q) for p, q in zip(full, again)), + "forward_rows_independent": torch.equal(full[0][: half[0].shape[0]], half[0]), + "base_frozen": all(frozen), + "forward_rel_err_vs_fp64": err, + } 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/models/minimax_h3/transition.py b/rl_engine/models/minimax_h3/transition.py new file mode 100644 index 000000000..0b9d4683a --- /dev/null +++ b/rl_engine/models/minimax_h3/transition.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``stochastic_transition_contract``: the H3 flow-matching CPS step, layout invariant. + +Targets the verl-omni H3 recipe: ``sde_type="cps"``, ``noise_level=0.8`` +(``FlowMatchSDEDiscreteScheduler.sample_previous_step``). With ``s = sin(noise_level * pi / 2)``:: + + std = sigma_prev * s + pred_x0 = x - sigma * v + noise_est = x + v * (1 - sigma) + mean = pred_x0 * (1 - sigma_prev) + noise_est * sqrt(sigma_prev**2 - std**2) + x_prev = mean + std * eps + log_prob = mean over non-batch dims of -(x_prev - mean)**2 (CPS has no normaliser) + +The contract makes every number a function of the *sample*, never of how samples are batched: + +* FP32 elementwise ops in exactly the order above (separate roundings, no fused contraction); +* ``eps`` comes from ``step_noise(seed, step, sample_id, modality, shape)``, a per-sample + generator, not one stream consumed in batch order; +* the log-prob mean uses a fixed-order GEMM reduction (``h3-projection-gemm-v1``, FP32 ``ieee``), + each sample's row independent of the batch; +* ``noise_level == 0`` (a non-selected step) gives ``std == 0``, ``x_prev == mean`` and + ``log_prob == 0``. +""" + +from __future__ import annotations + +import hashlib +import math +from dataclasses import dataclass + +import torch + +SDE_TYPE = "cps" +NOISE_LEVEL = 0.8 +CONTRACT = "h3-cps-transition-v1" + + +def step_noise( + seed: int, + step: int, + sample_id: int, + modality: str, + shape: tuple[int, ...], + device: torch.device | str = "cpu", +) -> torch.Tensor: + """FP32 standard normal noise for one sample at one step, independent of the batch layout.""" + + digest = hashlib.sha256(f"{seed}/{step}/{sample_id}/{modality}".encode()).digest() + gen = torch.Generator(device="cpu").manual_seed(int.from_bytes(digest[:8], "little") >> 1) + return torch.randn(shape, generator=gen, dtype=torch.float32).to(device) + + +@dataclass +class Transition: + prev_sample: torch.Tensor + prev_sample_mean: torch.Tensor + std_dev_t: torch.Tensor + log_prob: torch.Tensor | None + + +def fixed_order_mean(x: torch.Tensor) -> torch.Tensor: + """Per-sample mean over all non-batch elements with a fixed reduction order (``(B,)`` FP32).""" + + flat = x.reshape(x.shape[0], -1).contiguous().float() + ones = torch.ones(flat.shape[1], 1, device=flat.device, dtype=torch.float32) + from rl_engine.backends.shared.triton.gemm import h3_projection as g + + return g.gemm(flat, ones)[:, 0] / flat.shape[1] + + +def cps_transition( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: torch.Tensor | float, + sigma_prev: torch.Tensor | float, + *, + noise_level: float = NOISE_LEVEL, + noise: torch.Tensor | None = None, + prev_sample: torch.Tensor | None = None, + return_logprob: bool = True, +) -> Transition: + """One CPS step. ``noise`` (or ``prev_sample`` when scoring a stored rollout) is required.""" + + if sample.dtype != torch.float32 or velocity.dtype != torch.float32: + raise TypeError("the CPS step runs in FP32 (the scheduler asserts the same)") + if (noise is None) == (prev_sample is None): + raise ValueError("pass exactly one of noise (sampling) or prev_sample (scoring)") + sigma = torch.as_tensor(sigma, dtype=torch.float32, device=sample.device) + sigma_prev = torch.as_tensor(sigma_prev, dtype=torch.float32, device=sample.device) + sin_noise = math.sin(noise_level * math.pi / 2) + std_dev_t = sigma_prev * sin_noise + pred_original = sample - sigma * velocity + noise_estimate = sample + velocity * (1 - sigma) + mean = pred_original * (1 - sigma_prev) + noise_estimate * torch.sqrt( + sigma_prev**2 - std_dev_t**2 + ) + if prev_sample is None: + prev_sample = mean + std_dev_t * noise + log_prob = None + if return_logprob: + log_prob = fixed_order_mean(-((prev_sample.detach() - mean) ** 2)) + return Transition(prev_sample, mean, std_dev_t, log_prob) diff --git a/rl_engine/reference/minimax_h3/__init__.py b/rl_engine/reference/minimax_h3/__init__.py new file mode 100644 index 000000000..a01f9218b --- /dev/null +++ b/rl_engine/reference/minimax_h3/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 (RFC #420) PyTorch goldens. + +Shape constants are pinned to ``MiniMaxAI/MiniMax-H3@42ed227`` (see +``rl_engine/validation/models/h3_manifest.json``). The ops accept other sizes so that +tiny synthetic shapes can be tested, but the H3 layout rules (three modality +rows per timestep, six modulation chunks) are fixed. +""" + +from __future__ import annotations + +H3_FREQ_DIM = 256 +H3_MAX_PERIOD = 10000 +H3_TIME_EMBED_HIDDEN_DIM = 5376 +H3_TIME_EMBED_DIM = 2688 +H3_HIDDEN_SIZE = 5376 +# Modality tags of the packed sequence: 0 video, 1 text, 2 audio. +H3_MODALITY_NUM = 3 +# shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp +H3_ADALN_CHUNKS = 6 +H3_ADALN_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..be83239c5 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 three-modality AdaLN projection (RFC #420 row ``adaln_projection_3mod``). + +One ``MiniMaxH3AdaLayerNormModulation`` per transformer block:: + + act = silu(temb).to(weight.dtype) SiLU in FP32, one cast to BF16 + table = act @ W.T + b (T, 6 * H * 3), BF16 + rows = table.view(3 * T, 6 * H) row t * 3 + m, m in {video, text, audio} + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) + +Output channel ``o = m * 6H + c * H + h`` of the projection is chunk ``c``, hidden +index ``h`` of modality ``m``. Each of the six outputs has shape ``(3T, H)``, +and row ``t * 3 + m`` is what ``timestep_indices * 3 + token_tags`` addresses. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +def adaln_hidden_size(weight: torch.Tensor) -> int: + """Infer ``H`` from a 2-D projection weight whose row count is divisible by 18.""" + + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM + if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: + raise ValueError( + f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " + f"dim is not a multiple of {rows_per_hidden}" + ) + return weight.shape[0] // rows_per_hidden + + +def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """Check the AdaLN input shapes, shared device and precision boundary; return ``H``. + + Require nonempty FP32 ``temb`` of shape ``(T, D)`` and FP32 or BF16 weight + of shape ``(18H, D)`` with matching-dtype bias of shape ``(18H,)``. + """ + + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != temb.device: + raise ValueError(f"{name} is on {tensor.device}, temb is on {temb.device}") + if temb.dtype != torch.float32: + # RFC probe H7: the activation must run at the timestep-embedding + # precision; a BF16 temb means the cast happened before the SiLU. + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in _WEIGHT_DTYPES or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" + ) + hidden = adaln_hidden_size(weight) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + if bias.shape != (weight.shape[0],): + raise ValueError(f"bias must be ({weight.shape[0]},), got {tuple(bias.shape)}") + return hidden + + +def split_adaln_table(table: torch.Tensor, hidden: int) -> tuple[torch.Tensor, ...]: + """(T, 6H*3) projection output -> six (3T, H) views, diffusers order.""" + + return table.view(-1, H3_ADALN_CHUNKS * hidden).chunk(H3_ADALN_CHUNKS, dim=-1) + + +class NativeH3AdaLNProjectionOp: + """PyTorch reference for the H3 AdaLN projection. + + ``forward`` is the provider path (``F.silu`` in FP32, cast, ``F.linear``). + ``forward_fp32`` is the golden: the SiLU in FP64, rounded once to the + weight dtype at the declared boundary (the cast is model semantics, not + an implementation detail), then the projection in FP64. The six outputs are + returned in FP32 without the final BF16 rounding. Its gradient is FP64 + end to end: the cast is applied straight-through (identity VJP). + """ + + op_class = "reduction" + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` views using the PyTorch provider path.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Apply FP32 SiLU, cast to the weight dtype and project on the input device. + + Validate ``(T, D)`` embeddings and ``(18H, D)``/``(18H,)`` parameters; + return six weight-dtype ``(3T, H)`` views with PyTorch autograd support. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + return split_adaln_table(table, hidden) + + def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six FP32 golden outputs from FP64 SiLU and projection arithmetic. + + Preserve the declared activation rounding to the weight dtype while + using an identity VJP at that cast boundary for the FP64 golden graph. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # The declared cast rounds the FP32 SiLU, as the model (and the kernel) + # computes it: an FP64 SiLU can round to the other weight-dtype neighbour + # near a tie and shift a whole output column. Its VJP is the identity, + # written straight-through so autograd does not round the FP64 gradient + # to the weight dtype on its way back (``.to(bf16)`` would). + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/rl_engine/reference/minimax_h3/adaln_row_gather.py b/rl_engine/reference/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..2cd713d45 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_row_gather.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 AdaLN row gather (RFC #420 row ``adaln_row_gather``). + +H3 computes one row index per packed-sequence position and uses it to pick +every position's six modulation vectors:: + + adaln_indices = timestep_indices * 3 + token_tags (S,) + shift_msa[adaln_indices], scale_msa[adaln_indices], ... six (S, H) + +``rows`` is the ``(3T, 6H)`` view of one block's AdaLN projection +(``table.view(-1, 6H)``); its six column blocks are the six modulation +tensors. ``timestep_indices`` and ``token_tags`` are semantic inputs +(RFC #420 §4): out-of-range values are an error, never clamped. +""" + +from __future__ import annotations + +import torch + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_INDEX_DTYPES = (torch.int64, torch.int32) + + +def validate_h3_adaln_row_gather( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + *, + check_range: bool = True, +) -> int: + """Return the hidden size; raise on anything H3 does not define.""" + + if not isinstance(rows, torch.Tensor) or rows.dim() != 2: + raise ValueError("rows must be a 2-D (3T, 6H) tensor") + if not rows.is_floating_point(): + raise TypeError(f"rows must be floating point, got {rows.dtype}") + num_rows, width = rows.shape + if num_rows == 0 or num_rows % H3_MODALITY_NUM != 0: + raise ValueError(f"rows must hold 3 rows per timestep, got {num_rows}") + if width == 0 or width % H3_ADALN_CHUNKS != 0: + raise ValueError(f"rows width must be 6 * H, got {width}") + for name, index in (("timestep_indices", timestep_indices), ("token_tags", token_tags)): + if not isinstance(index, torch.Tensor) or index.dim() != 1: + raise ValueError(f"{name} must be a 1-D (S,) tensor") + if index.dtype not in _INDEX_DTYPES: + raise TypeError(f"{name} must be int64 or int32, got {index.dtype}") + if index.device != rows.device: + raise ValueError(f"{name} is on {index.device}, rows are on {rows.device}") + if timestep_indices.shape != token_tags.shape: + raise ValueError( + f"timestep_indices {tuple(timestep_indices.shape)} and token_tags " + f"{tuple(token_tags.shape)} must have the same length" + ) + if timestep_indices.numel() == 0: + raise ValueError("the packed sequence must not be empty") + if timestep_indices.dtype != token_tags.dtype: + raise TypeError("timestep_indices and token_tags must share an integer dtype") + if check_range: + num_timesteps = num_rows // H3_MODALITY_NUM + bad = ( + (timestep_indices < 0) + | (timestep_indices >= num_timesteps) + | (token_tags < 0) + | (token_tags >= H3_MODALITY_NUM) + ) + if bool(bad.any()): # one host sync + raise IndexError( + f"timestep_indices must lie in [0, {num_timesteps}) and token_tags in " + f"[0, {H3_MODALITY_NUM}) (0 video, 1 text, 2 audio)" + ) + return width // H3_ADALN_CHUNKS + + +def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) -> torch.Tensor: + return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() + + +class _DeterministicRowGather(torch.autograd.Function): + """``index_select`` forward; FP32 per-row segment sums for the backward. + + ``index_select``'s own backward is an atomic scatter-add in the grad dtype + (BF16 for H3), which is neither deterministic nor accurate. There are + only 3T table rows, so the reference sums each row's positions with one + FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. + """ + + @staticmethod + def forward(ctx, rows, index): + ctx.save_for_backward(index) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return torch.stack( + [chunk.index_select(0, index) for chunk in rows.chunk(H3_ADALN_CHUNKS, dim=-1)] + ) + + @staticmethod + def backward(ctx, grad): + (index,) = ctx.saved_tensors + num_rows, dtype = ctx.rows_meta + grad32 = grad.float() # (6, S, H) + out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) + for row in range(num_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + out[row] = grad32.index_select(1, positions).sum(dim=1) + return out.reshape(num_rows, -1).to(dtype), None + + +class NativeH3AdaLNRowGatherOp: + """PyTorch reference: ``index_select`` on each of the six column blocks. + + The forward is byte-identical to diffusers; the backward is deterministic + (see ``_DeterministicRowGather``). The raw diffusers path, including its + atomic backward, is ``rl_engine.validation.models.h3_provider.provider_adaln_row_gather``. + ``forward_fp32`` returns the same rows in FP32 with an FP64 gradient. + + The forward is a copy, but the VJP is a segmented reduction over the + packed sequence, so the op is judged with the ``reduction`` tolerance + class; forward bitwise equality is asserted separately. + """ + + op_class = "reduction" + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple(_DeterministicRowGather.apply(rows, index).unbind(0)) + + def forward_fp32(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple( + chunk.index_select(0, index).float() + for chunk in rows.double().chunk(H3_ADALN_CHUNKS, dim=-1) + ) diff --git a/rl_engine/reference/minimax_h3/block.py b/rl_engine/reference/minimax_h3/block.py new file mode 100644 index 000000000..427d7de2e --- /dev/null +++ b/rl_engine/reference/minimax_h3/block.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 golden of one MiniMax-H3 transformer block (RFC #420 ``ws1_one_h3_block``). + +Every node takes FP32 inputs (BF16 weights and activations upcast exactly), +computes in FP32 with TF32 off, and returns FP32 without rounding back to +BF16, so the golden is the block's arithmetic without the BF16 storage the +provider and the candidate both have. It is the yardstick their errors are +measured against, not a bitwise target. The one declared cast kept from the +checkpoint contract is the AdaLN input: ``silu(temb)`` is rounded to BF16 +before the projection, as every implementation must do it. + +The partial MM-RoPE uses exact FP32 ``cos``/``sin`` tables on the leading 96 +of 128 head channels (the provider rounds the tables to BF16 first). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +H3_HEADS = 56 +H3_HEAD_DIM = 128 +H3_ROTARY_DIM = 96 +H3_NORM_EPS = 1e-5 + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +def golden_adaln_tables(temb, weight, bias, hidden_size: int = 5376) -> tuple[torch.Tensor, ...]: + """Six ``(3T, H)`` FP32 tables; ``silu(temb)`` rounded to the weight dtype first.""" + + act = F.silu(_f32(temb)).to(weight.dtype) + table = F.linear(_f32(act), _f32(weight), _f32(bias)) + return table.view(-1, 6 * hidden_size).chunk(6, dim=-1) + + +def golden_rmsnorm(x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + x = _f32(x) + rstd = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + return x * rstd * _f32(weight) + + +def golden_norm_modulate(x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + n = golden_rmsnorm(x, weight, eps) + return n * (1.0 + _f32(scale).index_select(0, index)) + _f32(shift).index_select(0, index) + + +def golden_linear(x, weight) -> torch.Tensor: + return F.linear(_f32(x), _f32(weight)) + + +def golden_rope_tables(position_ids, rope_freq_dim: int = 16, rope_theta: float = 10000.0): + """FP32 ``(S, 96)`` cos/sin, the same angles as the provider's rotary module.""" + + exponent = torch.arange( + 0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device + ) + inv_freq = 1.0 / (rope_theta ** (exponent / (2 * rope_freq_dim))) + freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs = torch.cat(freqs.unbind(dim=1), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def golden_apply_rotary(x, cos, sin) -> torch.Tensor: + """Rotate-half on the leading ``cos.shape[-1]`` channels of ``(B, S, H, D)``; FP32 tables.""" + + x = _f32(x) + rot = cos.shape[-1] + head, tail = x[..., :rot], x[..., rot:] + x1, x2 = head.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + head = head * cos[None, :, None, :] + rotated * sin[None, :, None, :] + return torch.cat((head, tail), dim=-1) + + +def golden_attention(q, k, v) -> torch.Tensor: + """Non-causal full attention over one packed document, ``(B, S, H, D)`` -> ``(B, S, H*D)``.""" + + q, k, v = (_f32(t).permute(0, 2, 1, 3) for t in (q, k, v)) + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.shape[-1]) + out = torch.matmul(torch.softmax(scores, dim=-1), v) + return out.permute(0, 2, 1, 3).flatten(2, 3) + + +def golden_swiglu(projected) -> torch.Tensor: + hidden, gate = _f32(projected).chunk(2, dim=-1) + return hidden * F.silu(gate) + + +def golden_gate_residual(residual, gate, index, y) -> torch.Tensor: + return _f32(residual) + _f32(gate).index_select(0, index) * _f32(y) diff --git a/rl_engine/reference/minimax_h3/final_adaln_out.py b/rl_engine/reference/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..a21a5cbdb --- /dev/null +++ b/rl_engine/reference/minimax_h3/final_adaln_out.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 final AdaLN output (RFC #420 row ``final_adaln_out``). + +``MiniMaxH3AdaLayerNormOut`` runs after the 50 blocks:: + + shift, scale = linear(silu(temb).to(bf16)).chunk(2) # (T, H) each, shift first + out = norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] + +``linear`` is ``norm_out.linear`` (``2688 -> 2 * 5376``, BF16), ``norm`` is +``norm_out.norm`` (``nn.RMSNorm(5376, eps=1e-5)``), and the table is indexed +by ``timestep_indices`` (one row per distinct timestep), not by +``adaln_indices``. Diffusers then casts the result to the FP32 output heads' +dtype, an exact upcast left to the caller. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + NativeH3RMSNormOp, + validate_h3_rmsnorm, +) + + +def validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, *, same_dtype: bool = True +) -> int: + hidden = x.shape[-1] + if same_dtype: + validate_h3_rmsnorm(x, norm_weight, eps) + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if temb.dim() != 2 or temb.shape[0] == 0: + raise ValueError(f"temb must be a non-empty (T, D) matrix, got {tuple(temb.shape)}") + if weight.shape != (2 * hidden, temb.shape[1]) or bias.shape != (2 * hidden,): + raise ValueError( + f"norm_out.linear must be ({2 * hidden}, {temb.shape[1]}) with a ({2 * hidden},) " + f"bias, got {tuple(weight.shape)} / {tuple(bias.shape)}" + ) + if bias.dtype != weight.dtype: + raise TypeError("norm_out.linear weight and bias must share a dtype") + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if timestep_indices.dtype != torch.int64 or timestep_indices.dim() != 1: + raise TypeError("timestep_indices must be a 1-D int64 tensor") + if x.dim() < 2 or x.shape[-2] != timestep_indices.shape[0]: + raise ValueError("x must be (..., S, H) with S = len(timestep_indices)") + if bool(((timestep_indices < 0) | (timestep_indices >= temb.shape[0])).any()): + raise IndexError(f"timestep_indices must lie in [0, {temb.shape[0]})") + return hidden + + +class NativeH3FinalAdaLNOutOp: + """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only where + the model itself stores a value in its own dtype (the SiLU cast, the BF16 + table, the ``norm_out.norm`` output and ``1 + scale``), straight-through for + the gradient, and computes everything else in FP64.""" + + op_class = "reduction" + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + return NativeH3RMSNormOp().forward_modulated( + x, norm_weight, shift, scale, timestep_indices, eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, same_dtype=False + ) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # Declared cast of the FP32 SiLU, as the model computes it (an FP64 SiLU can + # round to the other neighbour near a tie); identity VJP. + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the + # weight dtype, and every position of a timestep shares that row, so the + # rounding is model semantics (straight-through for the gradient). + table = table + (table.to(weight.dtype).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() + # norm_out.norm is an x.dtype module and diffusers forms 1 + scale in + # x.dtype, so both are stored rounded. 1 + scale is shared by every + # position of a timestep: left unrounded, its error is systematic and + # grows with S in d_norm_weight; n's rounding enters d_scale (and so + # dW, d_temb) summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.index_select(0, timestep_indices) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.index_select(0, timestep_indices) + return out.float() diff --git a/rl_engine/reference/minimax_h3/fixed_order.py b/rl_engine/reference/minimax_h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/reference/minimax_h3/fixed_order.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shape-independent FP32 reductions shared by the H3 goldens and backwards. + +``torch.sum`` picks its reduction tree from the tensor shape, so a row's +result can change with the number of rows next to it. These helpers spell +the order out with elementwise ops only, which keeps every row's bytes a +function of that row alone. +""" + +from __future__ import annotations + +import torch + + +def tree_sum_lastdim_fp32(x: torch.Tensor) -> torch.Tensor: + """Pairwise tree over the last dim: pad to a power of two, halve until 1. + + Level ``l`` adds element ``i`` and ``i + width/2`` for every ``i``. The + tree depends only on the last-dim size. + """ + + x = x.float() + width = x.shape[-1] + if width == 0: + return x.new_zeros(x.shape[:-1]) + padded = 1 << (width - 1).bit_length() + if padded != width: + x = torch.nn.functional.pad(x, (0, padded - width)) + while x.shape[-1] > 1: + half = x.shape[-1] // 2 + x = x[..., :half] + x[..., half:] + return x[..., 0] + + +def fold_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left fold over dim 0 in ascending row order: ((r0 + r1) + r2) + ...""" + + rows = rows.float() + if rows.shape[0] == 0: + return rows.new_zeros(rows.shape[1:]) + acc = rows[0].clone() + for index in range(1, rows.shape[0]): + acc = acc + rows[index] + return acc diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py new file mode 100644 index 000000000..f68a0fdf4 --- /dev/null +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.validation.models.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/reference/minimax_h3/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/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..1727fd1d0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 RMSNorm and its fused AdaLN modulation (RFC #420 row ``h3_rmsnorm``). + +Every H3 RMSNorm is ``nn.RMSNorm(5376, eps=1e-5)`` with an affine weight in the +block dtype: the block ``norm1``/``norm2``, the token-refiner norms and the +final ``norm_out.norm``. Inside a transformer block (and in ``norm_out``) the +normalised rows are immediately modulated by per-row AdaLN parameters:: + + n = rms_norm(x, weight, eps) + out = n * (1.0 + scale[index]) + shift[index] # eager, rounded after every op + +``index`` is ``adaln_indices`` in a block and ``timestep_indices`` in +``norm_out``; ``shift``/``scale`` are ``(R, H)`` rows (views of the AdaLN +table). ``x`` is ``(..., S, H)`` and ``index`` has one entry per position ``S``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +H3_NORM_EPS = 1e-5 +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> int: + for name, tensor in (("x", x), ("weight", weight)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype not in _DTYPES: + raise TypeError(f"{name} must be fp32, bf16 or fp16, got {tensor.dtype}") + if weight.dtype != x.dtype: + raise TypeError(f"weight dtype {weight.dtype} must match x dtype {x.dtype}") + if weight.device != x.device: + raise ValueError(f"weight is on {weight.device}, x is on {x.device}") + if x.dim() < 2 or x.numel() == 0: + raise ValueError(f"x must be a non-empty (..., N) tensor, got {tuple(x.shape)}") + if weight.shape != (x.shape[-1],): + raise ValueError(f"weight must be ({x.shape[-1]},), got {tuple(weight.shape)}") + if not eps > 0: + raise ValueError(f"eps must be positive, got {eps}") + return x.shape[-1] + + +def validate_h3_modulation( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + hidden = x.shape[-1] + for name, tensor in (("shift", shift), ("scale", scale)): + if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): + raise TypeError(f"{name} must match x's dtype and device") + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") + if shift.shape != scale.shape: + raise ValueError("shift and scale must have the same shape") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != x.device: + raise TypeError("index must be a 1-D int64 tensor on x's device") + if x.dim() < 2 or x.shape[-2] != index.shape[0]: + raise ValueError( + f"x must be (..., S, {hidden}) with S = len(index) = {index.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if check_range and bool(((index < 0) | (index >= shift.shape[0])).any()): + raise IndexError(f"index must lie in [0, {shift.shape[0]})") + + +class _DeterministicModulation(torch.autograd.Function): + """Eager ``n * (1.0 + scale[i]) + shift[i]``; the ``shift``/``scale`` + gradients as FP32 per-row sums instead of ``index_select``'s atomic + scatter-add (``d_n`` is the eager VJP).""" + + @staticmethod + def forward(ctx, n, shift, scale, index): + t1 = 1.0 + scale.index_select(0, index) + ctx.save_for_backward(n, t1, index) + ctx.rows = shift.shape[0] + return n * t1 + shift.index_select(0, index) + + @staticmethod + def backward(ctx, grad): + n, t1, index = ctx.saved_tensors + hidden = n.shape[-1] + g = grad.float().reshape(-1, index.shape[0], hidden) + per_shift = g.sum(dim=0) + per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) + d_shift = per_shift.new_zeros((ctx.rows, hidden)) + d_scale = per_scale.new_zeros((ctx.rows, hidden)) + for row in range(ctx.rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + return grad * t1, d_shift.to(t1.dtype), d_scale.to(t1.dtype), None + + +class NativeH3RMSNormOp: + """PyTorch reference. + + ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and + the eager modulation expression, with a deterministic shift/scale gradient; + the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / + ``forward_modulated_fp32`` are the golden: the same math in FP64, returned in + FP32. The modulated golden rounds, straight-through, only the two values the + model stores in ``x``'s dtype: ``norm(x)`` and ``1 + scale``. + """ + + op_class = "reduction" + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + hidden = validate_h3_rmsnorm(x, weight, eps) + return F.rms_norm(x, (hidden,), weight, eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) + return (x64 * rstd * weight.double()).float() + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + n = self.forward(x, weight, eps) + return _DeterministicModulation.apply(n, shift, scale, index) + + def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + # The golden may be fed FP32 golden modulation rows for a BF16 x. + validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() + # The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype. + # 1 + scale is shared by every position of a row, so leaving it unrounded + # skews d_norm_weight systematically with S; n's rounding enters d_scale + # summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.double().index_select(0, index) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.double().index_select(0, index) + return out.float() diff --git a/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..51ef48db7 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_mlp.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 timestep MLP (RFC #420 row ``timestep_mlp_fp32``). + +``TimestepEmbedding(in_channels=256, time_embed_dim=5376, out_dim=2688)``:: + + temb = linear_2(silu(linear_1(features))) 256 -> 5376 -> 2688 + +``time_embedder`` is in ``_keep_in_fp32_modules``: weights, biases, +activations and the output are FP32, and ``temb`` stays FP32 because every +AdaLN projection applies its own SiLU to it before casting (RFC #420 §4). +A BF16 argument is a contract violation (RFC probe H7), not a lower-precision +variant, so it is rejected. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def validate_h3_timestep_mlp( + x: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> None: + """Check a same-device FP32 ``(T, K) -> (T, H) -> (T, D)`` MLP contract. + + Require ``T > 0``, weights of shapes ``(H, K)``/``(D, H)`` and biases + of shapes ``(H,)``/``(D,)``; reject lower-precision parameters or inputs. + """ + + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} + for name, tensor in named.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError( + f"{name} must be float32: the H3 time_embedder is an FP32 module " + f"(got {tensor.dtype})" + ) + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") + hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if k_in != x.shape[1]: + raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") + if b1.shape != (hidden,): + raise ValueError(f"b1 must be ({hidden},), got {tuple(b1.shape)}") + if w2.dim() != 2 or w2.shape[1] != hidden: + raise ValueError(f"w2 must be (D, {hidden}), got {tuple(w2.shape)}") + if b2.shape != (w2.shape[0],): + raise ValueError(f"b2 must be ({w2.shape[0]},), got {tuple(b2.shape)}") + + +class NativeH3TimestepMLPOp: + """PyTorch reference for the H3 timestep MLP. + + ``forward`` is the provider path (``nn.Linear`` / ``F.silu`` in FP32; run + it with TF32 disabled, as the RL-Kernel contract requires). + ``forward_fp32`` is the independent golden: the same graph in FP64, + rounded once to FP32, so its own rounding error (~1e-16 relative) is far + below the FP32 contract and its reduction order does not matter. + """ + + op_class = "reduction" + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` embeddings using the PyTorch provider graph.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Validate and evaluate linear-SiLU-linear in FP32 on the input device. + + Inputs follow the ``(T, K) -> (T, H) -> (T, D)`` contract and retain + PyTorch autograd support; CUDA callers must disable TF32 for this path. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Evaluate the validated MLP in FP64 and round its ``(T, D)`` output to FP32.""" + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + z = F.linear(x.double(), w1.double(), b1.double()) + h = z * torch.sigmoid(z) + return F.linear(h, w2.double(), b2.double()).float() diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..0d95e58f0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 sinusoidal timestep features (RFC #420 row ``timestep_sinusoid_h3``). + +H3 builds ``Timesteps(num_channels=256, flip_sin_to_cos=True, +downscale_freq_shift=0)`` and feeds it the *distinct* timestep values of the +packed sequence, unscaled in ``[0, 1]`` (``t = 1 - sigma``):: + + freq[k] = exp(-ln(10000) * k / 128) k = 0..127, FP32 + arg[t, k] = t * freq[k] FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + +The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` +module in the checkpoint, so the features never pass through BF16. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD + +_TIMESTEP_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_timesteps( + timestep: torch.Tensor, num_channels: int, *, check_range: bool = True +) -> None: + """Fail closed on inputs the H3 convention does not define. + + ``check_range`` rejects values outside ``[0, 1]``: H3 consumes + ``t = 1 - sigma`` unscaled, so a ``t * 1000`` caller (RFC probe H10) is a + convention error, not a different valid input. It reads the tensor back + to the host, which is cheap for the handful of distinct timesteps. + """ + + if not isinstance(timestep, torch.Tensor): + raise TypeError("timestep must be a torch.Tensor") + if timestep.dim() != 1: + raise ValueError(f"timestep must be 1-D (num_timesteps,), got {tuple(timestep.shape)}") + if timestep.numel() == 0: + raise ValueError("timestep must hold at least one timestep") + if timestep.dtype not in _TIMESTEP_DTYPES: + raise TypeError(f"timestep must be fp32, bf16 or fp16, got {timestep.dtype}") + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") + if check_range: + t32 = timestep.detach().float() + # One host sync: NaN fails both comparisons, so this also rejects it. + if not bool(((t32 >= 0) & (t32 <= 1)).all()): + raise ValueError( + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + f"unscaled (got {t32.cpu().tolist()[:8]})" + ) + + +class NativeH3TimestepSinusoidOp: + """PyTorch reference for the H3 sinusoidal timestep features. + + ``forward`` replays diffusers' ``get_timestep_embedding`` op for op on the + input's device (the provider path). ``forward_fp32`` is the independent + golden: the same formula evaluated in FP64 and rounded once to FP32. + """ + + op_class = "elementwise" + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 sinusoidal features for nonempty timesteps in ``[0, 1]``.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on the input device. + + Require nonempty 1-D FP32, BF16 or FP16 timesteps and positive even + channels. Validate finite values in ``[0, 1]`` when ``check_range`` is + true and preserve PyTorch autograd through the provider formula. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + # Same op sequence and dtypes as diffusers get_timestep_embedding with + # flip_sin_to_cos=True, downscale_freq_shift=0, scale=1. + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / half + freq = torch.exp(exponent) + arg = timestep[:, None].float() * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1) + + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Evaluate the validated sinusoid formula in FP64 and round once to FP32. + + Return ``(T, num_channels)`` on the timestep device with cosine channels + followed by sine channels and gradients through the FP64 golden graph. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + k = torch.arange(half, dtype=torch.float64, device=timestep.device) + freq = torch.exp(-math.log(H3_MAX_PERIOD) * k / half) + arg = timestep.double()[:, None] * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1).float() + + @staticmethod + def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: + """The FP32 frequency table the provider path multiplies by.""" + + half = num_channels // 2 + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=device + ) + return torch.exp(exponent / half) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 528c77b98..e84fc4ddb 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -204,6 +204,65 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + # MiniMax-H3 (RFC #420) conditioning path + PYTORCH_H3_TIMESTEP_SINUSOID = ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + PYTORCH_H3_TIMESTEP_MLP = "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp" + CUDA_H3_TIMESTEP_MLP = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" + ) + PYTORCH_H3_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" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) + PYTORCH_H3_ADALN_ROW_GATHER = ( + "rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ) + CUDA_H3_ADALN_ROW_GATHER = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" + ) + PYTORCH_H3_RMSNORM = "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp" + CUDA_H3_RMSNORM = "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp" + CUDA_H3_GATE_RESIDUAL = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp" + ) + PYTORCH_H3_FINAL_ADALN_OUT = ( + "rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp" + ) + CUDA_H3_FINAL_ADALN_OUT = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out.H3FinalAdaLNOutCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -502,6 +561,7 @@ class KernelRegistry: """ def __init__(self): + """Initialize backend caches, capability contracts, and device-specific priorities.""" self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) @@ -673,6 +733,47 @@ def __init__(self): OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], + "timestep_sinusoid_h3": [ + OpBackend.CUDA_H3_TIMESTEP_SINUSOID, + OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, + ], + "timestep_mlp_fp32": [ + OpBackend.CUDA_H3_TIMESTEP_MLP, + OpBackend.PYTORCH_H3_TIMESTEP_MLP, + ], + "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, + ], + "adaln_row_gather": [ + OpBackend.CUDA_H3_ADALN_ROW_GATHER, + OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, + ], + "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -714,6 +815,17 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "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": [ @@ -763,6 +875,17 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "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], @@ -789,6 +912,17 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "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..1d5c3df2b --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,771 @@ +# 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, attention: str = "auto") -> 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.attention_h3 = None + if attention != "interim": + # h3_full_attention (#553): query-chunked, bounded workspace, same #240 arithmetic + try: + from rl_engine.backends.cuda.attention.h3_full_attention import ( + H3FullAttentionOp, + _extension, + ) + + _extension() + self.attention_h3 = H3FullAttentionOp() + except (ImportError, RuntimeError): + if attention == "h3_full": + raise + 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() + if op is not None + } + + +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: + if ops.attention_h3 is not None: # BSHD, bounded workspace + q, k, v = (t.contiguous() for t in (q, k, _heads(v))) + return ops.attention_h3(q, k, v).flatten(2, 3) + 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 new file mode 100644 index 000000000..92455bae0 --- /dev/null +++ b/rl_engine/validation/models/h3_cases.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic MiniMax-H3 (RFC #420) test inputs shared by tests and scripts.""" + +from __future__ import annotations + +import torch + + +def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: + """Return seeded FP32 timesteps of shape ``(num,)`` on ``device``. + + Requires positive ``num``; the first value is 0 and the last is 1 when + ``num > 1``, with all remaining values sampled uniformly in [0, 1]. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + t = torch.rand(num, generator=generator) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(device) + + +def h3_packed_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """Return seeded int64 timestep indices and modality tags, each shape ``(seq_len,)``. + + Each row independently samples a timestep in ``[0, num_timesteps)`` and + a video/text/audio tag in ``[0, 3)``. The first three rows cover all tags + when available; both tensors are returned on ``device``. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + token_tags = torch.randint(0, 3, (seq_len,), generator=generator) + timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) + token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) + return timestep_indices.to(device), token_tags.to(device) + + +def h3_block_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for block-structured packing, as H3 packs requests. + + Each timestep owns one contiguous run of text, video and audio blocks (about + 5% / 80% / 15% of its tokens), so table rows form contiguous position runs. + ``h3_packed_layout`` is the interleaved stress case. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + cuts = torch.sort(torch.randperm(seq_len - 1, generator=generator)[: num_timesteps - 1] + 1) + edges = [0, *cuts.values.tolist(), seq_len] + timestep_indices = torch.empty(seq_len, dtype=torch.long) + token_tags = torch.empty(seq_len, dtype=torch.long) + for t, (lo, hi) in enumerate(zip(edges, edges[1:])): + text, audio = (hi - lo) // 20, (hi - lo) * 3 // 20 + timestep_indices[lo:hi] = t + token_tags[lo:hi] = 0 # video + token_tags[lo : lo + text] = 1 + token_tags[hi - audio : hi] = 2 + return timestep_indices.to(device), token_tags.to(device) + + +# diffusers@f53d552 modular_pipelines/minimax_h3/before_denoise.py constants. +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +def _spatial_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + """One aspect-normalised spatial axis, as numpy's ``linspace(endpoint=False)`` builds it.""" + + import numpy as np + + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def h3_fl2va_layout( + *, + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + audio_channels: int = 2, + patch: tuple[int, int] = (2, 2), + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Packed text/audio/video rows of a text-only FL2VA request, laid out as diffusers packs them. + + Replays ``position_ids``/``token_tags`` of the pinned pipeline's FL2VA + layout step (no keyframe anchors): text rows on the time axis at their row + index, channel-major audio rows pinned to the width extremes, then the video + frames on the non-uniform ``5/3 * (1, 4, 4, 4, 4)`` clock. ``position_ids`` + is ``(S, 3)`` float64 as the pipeline makes it; the rotary module casts it. + ``timestep_indices`` puts text on table row 0 and the generated media on + row 1, a two-timestep stand-in for the conditioning rows. + """ + + import numpy as np + + patch_h, patch_w = patch + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_audio_rows = num_audio_latents * audio_channels + num_video_rows = num_latent_frames * rows_per_frame + seq_len = num_text_tokens + num_audio_rows + num_video_rows + audio_start = num_text_tokens + video_start = audio_start + num_audio_rows + + position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + frame_grid = torch.stack([grid.reshape(-1) for grid in grids], dim=-1) + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(audio_channels) + position_ids[audio_start:video_start, 2] = torch.cat( + [ + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full( + (num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 + ), + ] + ) + spans = torch.tensor( + [ + _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] + for i in range(num_latent_frames) + ], + dtype=torch.float64, + ) + frame_time = float(num_text_tokens) + torch.cat( + [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] + ) + video = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video[:, :, 0] = frame_time[:, None] + video[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video.reshape(-1, 3) + + token_tags = torch.full((seq_len,), VIDEO_TAG, dtype=torch.long) + token_tags[:num_text_tokens] = TEXT_TAG + token_tags[audio_start:video_start] = AUDIO_TAG + timestep_indices = torch.ones(seq_len, dtype=torch.long) + timestep_indices[:num_text_tokens] = 0 + return { + "position_ids": position_ids.to(device), + "token_tags": token_tags.to(device), + "timestep_indices": timestep_indices.to(device), + } + + +# Named FL2VA layouts for the one-block replay: (text, latent frames, latent H, latent W, audio). +H3_BLOCK_LAYOUTS = { + # 32 + 2*20 + 2*(8*12) = 264 rows: the CI-sized case. + "tiny": dict( + num_text_tokens=32, + num_latent_frames=2, + latent_height=16, + latent_width=24, + num_audio_latents=20, + ), + # 77 + 2*40 + 3*(16*16) = 925 rows: a non-power-of-two, three-frame case. + "small": dict( + num_text_tokens=77, + num_latent_frames=3, + latent_height=32, + latent_width=32, + num_audio_latents=40, + ), + # 120 + 2*80 + 5*(18*32) = 3160 rows: a 16:9 five-frame case. + "medium": dict( + num_text_tokens=120, + num_latent_frames=5, + latent_height=36, + latent_width=64, + num_audio_latents=80, + ), +} diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py new file mode 100644 index 000000000..3aa1aa191 --- /dev/null +++ b/rl_engine/validation/models/h3_chain.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stage-by-stage replay of the MiniMax-H3 conditioning chain (RFC #420). + +Each stage runs three ways on the same device: + +* provider: the diffusers op sequence (``rl_engine.validation.models.h3_provider``); +* candidate: the dispatched RL-Kernel op, fed the previous candidate output + (chained) and, separately, the previous provider output (isolated); +* golden: the FP32/FP64 reference, fed the previous golden output. + +A stage also declares what it promises, so the end-to-end test and the +evidence JSON check the same thing: whether its isolated output is bitwise +equal to the provider, and its tolerance against the golden (rows of +``tolerance_contract.json``). ``first_drift`` is the first stage whose chained +output differs from the provider; ``first_isolated_drift`` the first stage +that differs even on provider inputs. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +HIDDEN = 5376 + + +@dataclass(frozen=True) +class Stage: + name: str + op_type: str + # (op, ctx, upstream) -> outputs; ``upstream`` is the previous stage's output. + candidate: Callable[[Any, dict, Any], Any] + provider: Callable[[dict, Any], Any] + golden: Callable[[Any, dict, Any], Any] + # Promises checked by tests/models/minimax_h3/test_h3_conditioning_e2e.py. + provider_bitwise_isolated: bool + golden_atol: float + golden_rtol: float + + +def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + """Return the context's MLP weights and biases in linear_1, linear_2 call order.""" + + weights = ctx["weights"] + return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] + + +def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the context's block-0 AdaLN projection weight and bias.""" + + prefix = "transformer_blocks.0.adaln_proj.linear" + return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] + + +STAGES: list[Stage] = [ + Stage( + name="timestep_sinusoid_h3", + op_type="timestep_sinusoid_h3", + candidate=lambda op, ctx, _up: op(ctx["timestep"]), + provider=lambda ctx, _up: provider_time_proj(ctx["timestep"]), + golden=lambda op, ctx, _up: op.forward_fp32(ctx["timestep"]), + provider_bitwise_isolated=True, + golden_atol=1e-5, # elementwise / float32 + golden_rtol=1e-5, + ), + Stage( + name="timestep_mlp_fp32", + op_type="timestep_mlp_fp32", + candidate=lambda op, ctx, up: op(up, *mlp_params(ctx)), + provider=lambda ctx, up: provider_time_embedder(up, *mlp_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *mlp_params(ctx)), + provider_bitwise_isolated=False, # reduction: different tree from cuBLAS + golden_atol=1e-4, # reduction / float32 + golden_rtol=1e-4, + ), + Stage( + name="adaln_projection_3mod", + op_type="adaln_projection_3mod", + candidate=lambda op, ctx, up: op(up, *adaln_params(ctx)), + provider=lambda ctx, up: provider_adaln_modulation(up, *adaln_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *adaln_params(ctx)), + provider_bitwise_isolated=False, # reduction: tensor-core tree differs from cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ), + Stage( + name="adaln_row_gather", + op_type="adaln_row_gather", + candidate=lambda op, ctx, up: op.gather_chunks( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + provider=lambda ctx, up: provider_adaln_row_gather( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + golden=lambda op, ctx, up: op.forward_fp32( + torch.cat(list(up), dim=1), ctx["timestep_indices"], ctx["token_tags"] + ), + provider_bitwise_isolated=True, # a copy + golden_atol=5e-2, # carries the projection's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ), +] + + +def adaln_indices(ctx: dict[str, Any]) -> torch.Tensor: + return ctx["timestep_indices"] * 3 + ctx["token_tags"] + + +def _norm1_modulated(ctx: dict[str, Any], modulation, call): + shift_msa, scale_msa = modulation[0], modulation[1] + weight = ctx["weights"]["transformer_blocks.0.norm1.weight"] + return call(ctx["hidden"], weight, shift_msa, scale_msa, adaln_indices(ctx)) + + +STAGES.append( + Stage( + name="h3_rmsnorm", + op_type="h3_rmsnorm", + # Block norm1 followed by the MSA shift/scale of the projection stage. + candidate=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated + ), + provider=lambda ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], provider_norm_modulate + ), + golden=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated_fp32 + ), + provider_bitwise_isolated=True, # replays nn.RMSNorm's reduction order + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + + +def _final(ctx: dict[str, Any], call): + weights = ctx["weights"] + history = ctx["history"] + return call( + history["adaln_gate_residual"], # the residual stream after the gated sublayer + weights["norm_out.norm.weight"], + history["timestep_mlp_fp32"], # temb + weights["norm_out.linear.weight"], + weights["norm_out.linear.bias"], + ctx["timestep_indices"], + ) + + +STAGES.append( + Stage( + name="final_adaln_out", + op_type="final_adaln_out", + candidate=lambda op, ctx, _up: _final(ctx, op), + provider=lambda ctx, _up: _final(ctx, provider_final_adaln_out), + golden=lambda op, ctx, _up: _final(ctx, op.forward_fp32), + provider_bitwise_isolated=False, # shift/scale projection: tensor-core tree vs cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +# The backward replay covers the conditioning chain up to the row gather. +CONDITIONING_STAGES = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) + +# Parameters whose gradients the backward replay reports, in chain order. +GRAD_LEAVES = ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +) +BACKWARD_MODES = ("candidate", "candidate_fused", "provider") + + +def _flat(value: Any) -> list[torch.Tensor]: + """Collect tensors from nested sequences in their original order.""" + + if isinstance(value, torch.Tensor): + return [value] + return [tensor for item in value for tensor in _flat(item)] + + +def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + """Report tensor equality, maximum FP32 absolute error, and tolerance checks.""" + + lhs_t, rhs_t = _flat(lhs), _flat(rhs) + bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) + max_abs = 0.0 + within = True + for a, b in zip(lhs_t, rhs_t, strict=True): + diff = (a.float() - b.float()).abs() + if diff.numel(): + max_abs = max(max_abs, float(diff.max())) + within = within and bool((diff <= atol + rtol * b.float().abs()).all()) + return {"bitwise_equal": bitwise, "max_abs": max_abs, "within_tolerance": within} + + +def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + """Build seeded CUDA timesteps and packed row labels with the supplied weights.""" + + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) + return { + "timestep": h3_timesteps(num_timesteps, seed=seed), + "timestep_indices": timestep_indices, + "token_tags": token_tags, + "weights": weights, + # Stand-in for the packed hidden states (the patch/text projections are + # other RFC rows): BF16 (1, S, H), seeded. + "hidden": torch.randn( + (1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2) + ) + .to(torch.bfloat16) + .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), + } + + +def golden_op(registry: KernelRegistry, op_type: str): + """Return the PyTorch reference at the end of the operator's CPU priority list.""" + + return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) + + +def run_case( + registry: KernelRegistry, + weights, + *, + num_timesteps: int, + seq_len: int, + seed: int = 0, + stages: list[Stage] | None = None, +) -> dict[str, Any]: + """Replay CUDA candidate, provider, and golden paths and report each stage's drift. + + Supplied stages must include their prerequisites in execution order. + Omitting ``stages`` runs the full registered chain without gradients. + """ + + stages = STAGES if stages is None else stages + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + report: dict[str, Any] = { + "num_timesteps": num_timesteps, + "seq_len": seq_len, + "seed": seed, + "stages": [], + } + chained: Any = None + provider: Any = None + golden: Any = None + first_drift = first_isolated = None + # Each mode sees its own earlier outputs (a stage may read any earlier stage). + history: dict[str, dict[str, Any]] = {"candidate": {}, "provider": {}, "golden": {}} + ctx_c = {**ctx, "history": history["candidate"]} + ctx_p = {**ctx, "history": history["provider"]} + ctx_g = {**ctx, "history": history["golden"]} + for stage in stages: + op = registry.get_op(stage.op_type, device="cuda") + gold = golden_op(registry, stage.op_type) + with torch.no_grad(): + isolated = stage.candidate(op, ctx_p, provider) + chained_input = chained + chained = stage.candidate(op, ctx_c, chained_input) + repeat = stage.candidate(op, ctx_c, chained_input) + golden = stage.golden(gold, ctx_g, golden) + provider = stage.provider(ctx_p, provider) + history["candidate"][stage.name] = chained + history["provider"][stage.name] = provider + history["golden"][stage.name] = golden + entry = { + "stage": stage.name, + "backend": type(op).__name__, + "kernel_id": getattr(op, "kernel_id", type(op).__name__), + "repeat_bitwise_equal": compare(repeat, chained)["bitwise_equal"], + "chained_vs_provider": compare(chained, provider), + "isolated_vs_provider": compare(isolated, provider), + "chained_vs_golden": compare(chained, golden, stage.golden_atol, stage.golden_rtol), + "provider_vs_golden": compare(provider, golden, stage.golden_atol, stage.golden_rtol), + } + if first_drift is None and not entry["chained_vs_provider"]["bitwise_equal"]: + first_drift = stage.name + if first_isolated is None and not entry["isolated_vs_provider"]["bitwise_equal"]: + first_isolated = stage.name + report["stages"].append(entry) + report["first_drift"] = first_drift + report["first_isolated_drift"] = first_isolated + return report + + +def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torch.Tensor]: + """Parameter gradients of the full chain for one execution mode. + + ``candidate`` runs the four RL-Kernel ops separately, ``candidate_fused`` + runs projection + gather as ``H3AdaLNModulationCudaOp``, ``provider`` the + diffusers path, ``golden`` the FP64 references. + """ + + leaves = { + name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES + } + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}, "history": {}} + value = None + if mode == "candidate_fused": + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + + for stage in STAGES[:2]: + value = stage.candidate(registry.get_op(stage.op_type, device="cuda"), run_ctx, value) + value = H3AdaLNModulationCudaOp()( + value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] + ) + else: + for stage in (s for s in STAGES if s.name in CONDITIONING_STAGES): + if mode == "provider": + value = stage.provider(run_ctx, value) + elif mode == "candidate": + op = registry.get_op(stage.op_type, device="cuda") + value = stage.candidate(op, run_ctx, value) + else: + value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + if value is None: + raise ValueError("conditioning chain produced no output") + outputs = list(value) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) + return [leaves[name].grad for name in GRAD_LEAVES] + + +def run_backward_case( + registry: KernelRegistry, weights, *, num_timesteps: int, seq_len: int, seed: int = 0 +) -> dict[str, Any]: + """Determinism and accuracy of the chain's parameter gradients, per mode.""" + + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + generator = torch.Generator(device="cuda").manual_seed(seed + 1) + hidden = weights[GRAD_LEAVES[4]].shape[0] // 18 + upstream = [ + torch.randn(seq_len, hidden, device="cuda", generator=generator).bfloat16() + for _ in range(6) + ] + golden = chain_grads("golden", registry, ctx, upstream) + runs = { + mode: [chain_grads(mode, registry, ctx, upstream) for _ in range(2)] + for mode in BACKWARD_MODES + } + report: dict[str, Any] = {"num_timesteps": num_timesteps, "seq_len": seq_len, "leaves": {}} + for index, name in enumerate(GRAD_LEAVES): + gold = golden[index].float() + entry: dict[str, Any] = {"golden_absmax": float(gold.abs().max())} + for mode, (first, second) in runs.items(): + grad = first[index] + err = (grad.float() - gold).abs().max() + entry[mode] = { + "repeat_bitwise_equal": bool(torch.equal(first[index], second[index])), + "max_abs_vs_golden": float(err), + "max_abs_vs_golden_over_absmax": float(err / gold.abs().max().clamp_min(1e-30)), + "correctly_rounded_fraction": float( + (grad == golden[index].to(grad.dtype)).float().mean() + ), + } + report["leaves"][name] = entry + return report + + +def git_state() -> dict[str, Any]: + """Return HEAD and tracked-tree dirtiness, defaulting to unknown/dirty if Git fails.""" + + def run(*args: str) -> str: + """Run Git in the repository root and return stripped stdout.""" + + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() + + try: + sha = run("rev-parse", "HEAD") + dirty = bool(run("status", "--porcelain", "--untracked-files=no")) + except Exception: # noqa: BLE001 + sha, dirty = "unknown", True + return {"rl_kernel_commit": sha, "tracked_tree_dirty": dirty} + + +def environment() -> dict[str, Any]: + """Record the active CUDA device, software versions, and matmul TF32 setting.""" + + return { + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": bool(torch.backends.cuda.matmul.allow_tf32), + } diff --git a/rl_engine/validation/models/h3_full_manifest.json b/rl_engine/validation/models/h3_full_manifest.json new file mode 100644 index 000000000..6cc2e5b5a --- /dev/null +++ b/rl_engine/validation/models/h3_full_manifest.json @@ -0,0 +1,311 @@ +{ + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "files": [ + { + "path": "audio_scheduler/scheduler_config.json", + "size": 96, + "sha256": "804780f7133477067bd6bbfbc02dc8b3cf9feeb400f97c08f5b1d5f6cbab3840" + }, + { + "path": "audio_vae/config.json", + "size": 2271, + "sha256": "9a3c645ff892b376c6f5f4c8685964cd75474731af594ff058492a0000caabb6" + }, + { + "path": "audio_vae/diffusion_pytorch_model.safetensors", + "size": 605429340, + "sha256": "52c59e67ba8de5477c81bfbced0327aabf500f1bfdeefd5ee754529241cb26cb" + }, + { + "path": "model_index.json", + "size": 2936, + "sha256": "5a587fe13b2371427415ac892463142683aefcd8d322e274a3a095eac37ac7d2" + }, + { + "path": "modular_model_index.json", + "size": 2935, + "sha256": "a2b6a210e482ffb78e613b553f570c44e101afce6741bd4ed91429d0559af031" + }, + { + "path": "processor/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "processor/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "processor/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "processor/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "processor/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "processor/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "processor/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "scheduler/scheduler_config.json", + "size": 97, + "sha256": "8fa6c3aa70dc9e691e1a6df899fd1b6f75f70481a27cee6e18a303817075c304" + }, + { + "path": "text_encoder/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "text_encoder/config.json", + "size": 1474, + "sha256": "d2dd0c60d01b9e195d9447c52da61c7302d28828524914c044d9c6e1b81d0427" + }, + { + "path": "text_encoder/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "text_encoder/model-00001-of-00014.safetensors", + "size": 4932328944, + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb" + }, + { + "path": "text_encoder/model-00002-of-00014.safetensors", + "size": 4875990528, + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa" + }, + { + "path": "text_encoder/model-00003-of-00014.safetensors", + "size": 4875990552, + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940" + }, + { + "path": "text_encoder/model-00004-of-00014.safetensors", + "size": 4875990584, + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f" + }, + { + "path": "text_encoder/model-00005-of-00014.safetensors", + "size": 4875990584, + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2" + }, + { + "path": "text_encoder/model-00006-of-00014.safetensors", + "size": 4875990584, + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974" + }, + { + "path": "text_encoder/model-00007-of-00014.safetensors", + "size": 4875990584, + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934" + }, + { + "path": "text_encoder/model-00008-of-00014.safetensors", + "size": 4875990584, + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f" + }, + { + "path": "text_encoder/model-00009-of-00014.safetensors", + "size": 4875990584, + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47" + }, + { + "path": "text_encoder/model-00010-of-00014.safetensors", + "size": 4875990584, + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f" + }, + { + "path": "text_encoder/model-00011-of-00014.safetensors", + "size": 4875990584, + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f" + }, + { + "path": "text_encoder/model-00012-of-00014.safetensors", + "size": 4875990584, + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e" + }, + { + "path": "text_encoder/model-00013-of-00014.safetensors", + "size": 4875990584, + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c" + }, + { + "path": "text_encoder/model-00014-of-00014.safetensors", + "size": 3270697008, + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b" + }, + { + "path": "text_encoder/model.safetensors.index.json", + "size": 97831, + "sha256": "06c952c569285870b811989b794b9766493e280fb77fbcb957fc4e5fcf25403a" + }, + { + "path": "text_encoder/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "text_encoder/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "text_encoder/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "text_encoder/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "text_encoder/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "tokenizer/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "tokenizer/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "tokenizer/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "tokenizer/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "transformer/config.json", + "size": 546, + "sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00014.safetensors", + "size": 4825958704, + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00014.safetensors", + "size": 4702158032, + "sha256": "2c4d362eddd2802180ac9c744849eb9ba8d9c8b984bdf9822cb02ed004b29184" + }, + { + "path": "transformer/diffusion_pytorch_model-00003-of-00014.safetensors", + "size": 4933368192, + "sha256": "949c5aafbbfa5654da730a6a7fafd75adb164d0857b095a30e8bb6d390887d69" + }, + { + "path": "transformer/diffusion_pytorch_model-00004-of-00014.safetensors", + "size": 4567069608, + "sha256": "eef7616790105ee839766bb2027203bf2c0d87c6aa038dca84145a8675f5ce28" + }, + { + "path": "transformer/diffusion_pytorch_model-00005-of-00014.safetensors", + "size": 4702158080, + "sha256": "43fdf42d638e8bc6745f713fae80c93bb301807a1a5ae7249344ce28e202a494" + }, + { + "path": "transformer/diffusion_pytorch_model-00006-of-00014.safetensors", + "size": 4933368232, + "sha256": "6442510b34d173653f0cce5c964b935395a8f7accf0b9cc0aa31aec59805239d" + }, + { + "path": "transformer/diffusion_pytorch_model-00007-of-00014.safetensors", + "size": 4567069608, + "sha256": "29f48f535c91dac76496ca821eeb16ca24bc4caf3f0cae8b920a89b1f966da6d" + }, + { + "path": "transformer/diffusion_pytorch_model-00008-of-00014.safetensors", + "size": 4702158080, + "sha256": "c711b096c764bd60f0b8b6ad49518bfab6d614fb788c725add8741c0674a4cd8" + }, + { + "path": "transformer/diffusion_pytorch_model-00009-of-00014.safetensors", + "size": 4933368232, + "sha256": "44428defe3976cbb87635ad200b958199e739986697cd29fdf27aeb7294b5944" + }, + { + "path": "transformer/diffusion_pytorch_model-00010-of-00014.safetensors", + "size": 4567069608, + "sha256": "3d44939c374c9da382e9c6877e1946adf7b84e08c7a881c068f228d6849411c9" + }, + { + "path": "transformer/diffusion_pytorch_model-00011-of-00014.safetensors", + "size": 4702158080, + "sha256": "224d24430b58127a5577721084e0e704a0e74ec96dd7c35bc6fc0994ebd87c33" + }, + { + "path": "transformer/diffusion_pytorch_model-00012-of-00014.safetensors", + "size": 4933368232, + "sha256": "48fa2bd8fe134eef565ab2464f1c2589a6657cba0d14283dfc06b532f8961f3c" + }, + { + "path": "transformer/diffusion_pytorch_model-00013-of-00014.safetensors", + "size": 4567069608, + "sha256": "be5b4b1809f9d546ffd4b3fcf41e5c1e02b819125caa6bc105c109b04c051bd3" + }, + { + "path": "transformer/diffusion_pytorch_model-00014-of-00014.safetensors", + "size": 4644161920, + "sha256": "8fbd5e6c1fb1df7ce988ca90f3d59e7610e465c7517e4b344eda4a214ba4b97d" + }, + { + "path": "transformer/diffusion_pytorch_model.safetensors.index.json", + "size": 64488, + "sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e" + }, + { + "path": "vae/config.json", + "size": 2011, + "sha256": "78f67deec3d63aae807f2bfe7154bc1e26f6372cb20b63265fcbae1b62bb5745" + }, + { + "path": "vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "size": 5061033024, + "sha256": "72f4c6be84ac0674f27398cde991dd9d719762f3952c4921aa66b2ce542f6374" + }, + { + "path": "vae/diffusion_pytorch_model-00002-of-00003.safetensors", + "size": 4955986528, + "sha256": "2e05e8bc23fa4071043e17fd242be8acd0685e781a43987432b2eae925be4198" + }, + { + "path": "vae/diffusion_pytorch_model-00003-of-00003.safetensors", + "size": 398539336, + "sha256": "c05d6ac4b1a33de372799d708531da6320f6a3ce6d1ce6d895e770988e004a39" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors.index.json", + "size": 74228, + "sha256": "15f6d44553c3c616b0dc999920aa784f92ecee7e4201f1f99ac405cfbf3061ca" + } + ] +} diff --git a/rl_engine/validation/models/h3_full_model.py b/rl_engine/validation/models/h3_full_model.py new file mode 100644 index 000000000..f8f27ddc3 --- /dev/null +++ b/rl_engine/validation/models/h3_full_model.py @@ -0,0 +1,344 @@ +# 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() + self.boundary_bindings = {} + + 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=( + "strict rows + provider" if self.boundary_bindings else "provider/reference" + ), + boundary_bindings=self.boundary_bindings, + 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 + + +# Boundary modules bound to RL-Kernel rows by ``strict_boundaries``; the others keep provider +# arithmetic and are reported as such (``proj_in``/``proj_out`` belong to the video-projection +# rows claimed by another contributor; ``time_embedder``/``norm_out`` have their own rows +# but are not part of this binding yet). +STRICT_BOUNDARY_ROWS = { + "audio_proj_in": "audio_input_projection_fp32", + "audio_proj_out": "audio_output_projection_fp32", + "context_embedder": "text_context_projection", + "token_refiner": "token_refiner_block", +} +PROVIDER_BOUNDARIES = ("proj_in", "proj_out", "time_embedder", "norm_out") + + +@contextmanager +def _strict_boundaries(model, trace, ops, native=frozenset()): + """Rebind the audio projections, the text projection and the whole token refiner to rows.""" + + from rl_engine.runtime.registry import KernelRegistry + from rl_engine.validation.models import h3_refiner + + registry = KernelRegistry() + row_ops = { + name: registry.get_op(row, device="cuda") + for name, row in STRICT_BOUNDARY_ROWS.items() + if name != "token_refiner" and name not in native + } + originals = [] + + def project(name): + module, op = getattr(model, name), row_ops[name] + + def forward(x): + flat = x.reshape(-1, x.shape[-1]) + out = op.forward(flat, module.weight, module.bias) + return out.view(*x.shape[:-1], out.shape[-1]) + + originals.append((module, module.forward)) + module.forward = forward + + for name in row_ops: + project(name) + + refiner = model.token_refiner + params = dict(model.named_parameters()) + + def refine(hidden): + if hidden.dim() != 3: + raise ValueError("token refiner expects (batch, tokens, hidden)") + return h3_refiner.run_refiner( + "candidate", + params, + hidden, + ops, + num_blocks=len(refiner.refiner_blocks), + on_block=lambda b, h: trace.observe(f"refiner.{b}", h), + observer_factory=lambda b: ( + lambda node, _b=b: trace.candidate_calls.update([f"refiner.{_b}.{node}"]) + ), + ) + + if "token_refiner" not in native: + originals.append((refiner, refiner.forward)) + refiner.forward = refine + trace.boundary_bindings = { + **{ + n: dict(row=r, status="provider" if n in native else "row") + for n, r in STRICT_BOUNDARY_ROWS.items() + }, + **{n: dict(row=None, status="provider") for n in PROVIDER_BOUNDARIES}, + } + try: + yield + finally: + for module, _original in originals: + module.__dict__.pop("forward", None) # back to the class method + + +@contextmanager +def full_model_execution( + model, + *, + candidate: bool, + trace: FullModelTrace, + ops=None, + provider_nodes=(), + strict_boundaries: bool = False, + native_boundaries=(), +): + """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 strict_boundaries and not candidate: + raise ValueError("strict_boundaries applies to candidate execution only") + unknown_boundaries = set(native_boundaries) - set(STRICT_BOUNDARY_ROWS) + if unknown_boundaries: + raise ValueError(f"unknown boundary modules: {sorted(unknown_boundaries)}") + if candidate: + ops = ops or h3_block.CandidateOps() + if hasattr(ops, "backends"): + trace.actual_backends = require_candidate_cuda(ops) + boundary_context = ( + _strict_boundaries(model, trace, ops, frozenset(native_boundaries)) + if strict_boundaries + else None + ) + try: + if boundary_context is not None: + boundary_context.__enter__() + 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 + if boundary_context is not None: + boundary_context.__exit__(None, None, None) diff --git a/rl_engine/validation/models/h3_generation_resume.py b/rl_engine/validation/models/h3_generation_resume.py new file mode 100644 index 000000000..4d1cfc428 --- /dev/null +++ b/rl_engine/validation/models/h3_generation_resume.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable denoising for the full-H3 candidate generation (RFC #420). + +A candidate generation takes about 20 hours per seed and its cluster partition +preempts jobs every few hours, so a restart must not redo finished steps. The +generation already keeps every step's predicted velocities (``step-NNN.pt``); +``StepStore`` adds, next to each, a sidecar with the digest of the transformer's +inputs for that step and the digest of the saved file, both written atomically +(the sidecar last, so a step counts only when both exist). + +On resume every saved transformer call is *replayed*: the saved velocities are +returned instead of recomputed, and the pipeline's own scheduler update turns them +into the next latents exactly as before. The input digest of every replayed call is +checked against the one saved, so the replay is proven to have rebuilt the trajectory +the earlier attempt had. Recomputing the last saved step (``verify_last_step``) +additionally checks that the forward pass reproduces its stored output bitwise; it +costs one full step per restart, which on a heavily preempted partition is most of an +attempt, so it is opt-in. The execution trace is saved after the first real step of +an attempt, so a run whose remaining work is only decoding keeps a complete trace. + +A run is tied to the numerics it started with (``model_trees``: the git tree of every +path in ``MODEL_PATHS``), not to a commit, so a commit that changes only orchestration +code can continue it. ``adopt_identity`` moves a run started under a commit-based +identity onto this one and refuses unless those paths are byte-identical. +""" + +from __future__ import annotations + +import datetime +import functools +import hashlib +import json +import os +from pathlib import Path +from typing import Any, Callable + +import torch + +RUN_FILE = "run.json" +TRACE_FILE = "trace.json" + +# Everything that determines the generated numbers. The resume and orchestration code is +# deliberately absent: changing it must not invalidate saved steps. +MODEL_PATHS = ( + "rl_engine/backends", + "rl_engine/contracts", + "rl_engine/ops", + "rl_engine/reference", + "rl_engine/runtime", + "rl_engine/validation/models/h3_block.py", + "rl_engine/validation/models/h3_cases.py", + "rl_engine/validation/models/h3_full_manifest.json", + "rl_engine/validation/models/h3_full_model.py", + "rl_engine/validation/models/h3_provider.py", + "rl_engine/validation/models/h3_weights.py", + "csrc", + "build_tools", + "examples/models/minimax_h3/generation-contract.json", + "examples/models/minimax_h3/generation-inputs.json", + "examples/models/minimax_h3/generation-prompt.txt", +) + + +class ResumeError(ValueError): + """A saved run cannot be continued, or a replay disagrees with what was saved.""" + + +def tensor_digest(named: dict[str, torch.Tensor]) -> str: + """sha256 over tensor names, dtypes, shapes and raw bytes, in name order.""" + + digest = hashlib.sha256() + for name in sorted(named): + tensor = named[name].detach().cpu().contiguous() + digest.update(f"{name}|{tensor.dtype}|{tuple(tensor.shape)}|".encode()) + if tensor.numel(): + digest.update(memoryview(tensor.reshape(-1).view(torch.uint8).numpy())) + return digest.hexdigest() + + +def inputs_digest(args: tuple, kwargs: dict[str, Any]) -> str: + """Digest of every tensor passed to one transformer call (non-tensors are ignored).""" + + named = {f"arg{i}": v for i, v in enumerate(args) if isinstance(v, torch.Tensor)} + named.update({k: v for k, v in kwargs.items() if isinstance(v, torch.Tensor)}) + return tensor_digest(named) + + +def model_trees(revision: str = "HEAD", cwd: Path | None = None) -> dict[str, str]: + """The git tree (or blob) id of every model path at ``revision``.""" + + import subprocess + + trees = {} + for path in MODEL_PATHS: + trees[path] = subprocess.check_output( + ["git", "rev-parse", f"{revision}:{path}"], text=True, cwd=cwd + ).strip() + return trees + + +def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(1 << 24): + digest.update(block) + return digest.hexdigest() + + +def _atomic_write(path: Path, write: Callable[[Path], object]) -> None: + temporary = path.with_name(path.name + ".tmp") + write(temporary) + os.replace(temporary, path) + + +class StepStore: + """The per-step velocities and input digests of one seed, in one directory.""" + + def __init__(self, directory: Path): + self.directory = directory + + def _step(self, step: int) -> Path: + return self.directory / f"step-{step:03}.pt" + + def _sidecar(self, step: int) -> Path: + return self.directory / f"step-{step:03}.in.json" + + def open(self, identity: dict[str, Any], *, resume: bool) -> int: + """Create the directory, or on ``resume`` validate it; return the saved step count. + + ``identity`` (commit, contract, seed, flags, ...) must equal the one stored, + so steps from a different source or setting can never be mixed in. + """ + + run_file = self.directory / RUN_FILE + now = datetime.datetime.now(datetime.timezone.utc).isoformat() + if not self.directory.exists(): + self.directory.mkdir(parents=True) + elif not resume: + raise ResumeError(f"{self.directory} exists; pass resume to continue it") + if run_file.exists(): + stored = json.loads(run_file.read_text()) + if stored["identity"] != identity: + raise ResumeError("saved run has a different identity; refusing to continue it") + else: + stored = {"identity": identity, "attempts": []} + saved = self._complete_steps() + stored["attempts"].append({"started": now, "saved_steps": saved}) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) + return saved + + def attempts(self) -> list[dict[str, Any]]: + return json.loads((self.directory / RUN_FILE).read_text())["attempts"] + + def _complete_steps(self) -> int: + count = 0 + while self._step(count).exists() and self._sidecar(count).exists(): + sidecar = json.loads(self._sidecar(count).read_text()) + if _file_sha256(self._step(count)) != sidecar["output_sha256"]: + break + count += 1 + if count and not (self.directory / "initial.pt").exists(): + return 0 + return count + + def save(self, step: int, values: tuple[torch.Tensor, ...], digest: str) -> None: + """Persist one step's outputs, then its sidecar; both atomically.""" + + _atomic_write(self._step(step), lambda p: torch.save(tuple(values), p)) + sidecar = {"inputs_sha256": digest, "output_sha256": _file_sha256(self._step(step))} + _atomic_write( + self._sidecar(step), lambda p: p.write_text(json.dumps(sidecar, indent=2) + "\n") + ) + + def save_trace(self, trace: dict[str, Any]) -> None: + """Keep the validated execution trace of this seed's real forward passes.""" + + _atomic_write( + self.directory / TRACE_FILE, lambda p: p.write_text(json.dumps(trace, indent=2) + "\n") + ) + + def load_trace(self) -> dict[str, Any] | None: + path = self.directory / TRACE_FILE + return json.loads(path.read_text()) if path.exists() else None + + def saved_digest(self, step: int) -> str: + return json.loads(self._sidecar(step).read_text())["inputs_sha256"] + + def load(self, step: int) -> tuple[torch.Tensor, ...]: + sidecar = json.loads(self._sidecar(step).read_text()) + if _file_sha256(self._step(step)) != sidecar["output_sha256"]: + raise ResumeError(f"saved step {step} changed on disk") + return tuple(torch.load(self._step(step), map_location="cpu", weights_only=True)) + + +class ResumableRun: + """Replays saved steps and records new ones for one seed's transformer calls.""" + + def __init__( + self, + store: StepStore, + identity: dict[str, Any], + *, + resume: bool, + verify_last_step: bool = False, + ): + self.store = store + self.saved = store.open(identity, resume=resume) + self.verify_last_step = verify_last_step + self.replay = self.saved - 1 if verify_last_step and self.saved else self.saved + self.verified_step: int | None = None + self._forward_calls = 0 + self.recorded = 0 + self.real_calls = 0 + + def install(self, module: torch.nn.Module, make_output: Callable[[tuple, dict], Any]): + """Serve the first ``replay`` calls of ``module`` from disk; return an uninstaller.""" + + original = module.forward + + # The pipeline picks the transformer's layout inputs from inspect.signature(forward), + # so the wrapper must keep the original signature (functools.wraps sets __wrapped__). + @functools.wraps(original) + def forward(*args, **kwargs): + index = self._forward_calls + self._forward_calls += 1 + if index >= self.replay: + return original(*args, **kwargs) + if inputs_digest(args, kwargs) != self.store.saved_digest(index): + raise ResumeError(f"replayed step {index} has different inputs than when saved") + device = next( + v.device for v in (*args, *kwargs.values()) if isinstance(v, torch.Tensor) + ) + values = tuple(v.to(device) for v in self.store.load(index)) + return make_output(values, kwargs) + + module.forward = forward + + def uninstall(): + module.forward = original + + return uninstall + + def record(self, args: tuple, kwargs: dict[str, Any], values: tuple[torch.Tensor, ...]): + """Handle one finished call (forward hook): save it, or verify a recomputed step.""" + + index = self.recorded + self.recorded += 1 + if index < self.replay: + return + self.real_calls += 1 + digest = inputs_digest(args, kwargs) + values = tuple(v.detach().cpu() for v in values) + if index < self.saved: + if digest != self.store.saved_digest(index): + raise ResumeError(f"recomputed step {index} has different inputs than saved") + stored = self.store.load(index) + if len(stored) != len(values) or not all( + torch.equal(a, b) for a, b in zip(stored, values, strict=True) + ): + raise ResumeError(f"recomputed step {index} differs from the saved output") + self.verified_step = index + return + self.store.save(index, values, digest) + + def summary(self) -> dict[str, Any]: + return dict( + saved_steps_at_start=self.saved, + replayed_steps=self.replay, + verify_last_step=self.verify_last_step, + recomputed_and_verified_step=self.verified_step, + real_forward_calls=self.real_calls, + attempts=self.store.attempts(), + ) + + +def adopt_identity( + directory: Path, + new_identity: dict[str, Any], + from_commit: str, + cwd: Path | None = None, + *, + dry_run: bool = False, +) -> None: + """Move a run started under a commit-based identity onto a ``model_trees`` identity. + + Refuses unless everything but the commit matches and the model paths at + ``from_commit`` are byte-identical to the ones in ``new_identity``; records the + adoption in the run's attempt list. ``dry_run`` performs every check and writes nothing. + """ + + run_file = directory / RUN_FILE + stored = json.loads(run_file.read_text()) + old = stored["identity"] + if "model_trees" in old: + raise ResumeError("run already has a model-tree identity") + if old.get("commit", "").startswith(from_commit) is False: + raise ResumeError(f"run was started at {old.get('commit')}, not {from_commit}") + rest_old = {k: v for k, v in old.items() if k != "commit"} + rest_new = {k: v for k, v in new_identity.items() if k != "model_trees"} + if rest_old != rest_new: + raise ResumeError("run differs from the new identity in more than the commit") + then = model_trees(old["commit"], cwd=cwd) + if then != new_identity["model_trees"]: + changed = sorted(k for k in then if then[k] != new_identity["model_trees"][k]) + raise ResumeError(f"model paths changed since {from_commit}: {changed}") + if dry_run: + return + stored["identity"] = new_identity + stored["attempts"].append( + { + "adopted": { + "from_commit": old["commit"], + "model_paths_identical": True, + "when": datetime.datetime.now(datetime.timezone.utc).isoformat(), + } + } + ) + _atomic_write(run_file, lambda p: p.write_text(json.dumps(stored, indent=2) + "\n")) diff --git a/rl_engine/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json new file mode 100644 index 000000000..e590a8473 --- /dev/null +++ b/rl_engine/validation/models/h3_manifest.json @@ -0,0 +1,137 @@ +{ + "version": "h3-conditioning-v1", + "rfc": "RL-Align/RL-Kernel#420", + "model_identity": { + "hf_repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "subfolder": "transformer", + "architecture": "MiniMaxH3Transformer3DModel", + "config_sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae", + "index_file": "diffusion_pytorch_model.safetensors.index.json", + "index_sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e", + "config_fingerprint": { + "hidden_size": 5376, + "num_attention_heads": 56, + "attention_head_dim": 128, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "norm_eps": 1e-05, + "final_norm_eps": 1e-05 + } + }, + "weight_shards": { + "diffusion_pytorch_model-00001-of-00014.safetensors": { + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size_bytes": 4825958704 + } + }, + "tensors": { + "time_embedder.linear_1.weight": { + "dtype": "float32", "shape": [5376, 256], + "sha256": "4c50eabe3c434dc05422e6501cdf7baa5cecbb1a7706ae9599c7a5b0e34e2913" + }, + "time_embedder.linear_1.bias": { + "dtype": "float32", "shape": [5376], + "sha256": "64763a5e39cb078eb0d8c5be5d56ad36cac3dec37be93a87e106d82ea6c071b7" + }, + "time_embedder.linear_2.weight": { + "dtype": "float32", "shape": [2688, 5376], + "sha256": "697819dd9413e0fec3a50987a964f1fd5db568ca545267414238f452294a8154" + }, + "time_embedder.linear_2.bias": { + "dtype": "float32", "shape": [2688], + "sha256": "096a3c7539f0248698cc65376c8da260e7365bf4ccef847c09bb9a36782a743c" + }, + "transformer_blocks.0.adaln_proj.linear.weight": { + "dtype": "bfloat16", "shape": [96768, 2688], + "sha256": "a0000193b30ea36a1996f53fc0d6c4c37ce5fd4585284f0f9bae570d25c6b8cf" + }, + "transformer_blocks.0.adaln_proj.linear.bias": { + "dtype": "bfloat16", "shape": [96768], + "sha256": "d59f9de9e07e6b4e3b57009627e276a26971ce379555fc99dbb06d3a237c8fc0" + }, + "transformer_blocks.0.norm1.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "5e32e4d7379a9f9302755af395ed50b5b998352a221dc2c99204a85dd9f9b49f" + }, + "transformer_blocks.0.norm2.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "a2c513ec459414e211a3d37f77c1f556f11117f56516b22eb5912f8ff5f98872" + }, + "token_refiner.final_norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "cf10a56a11370a216a4ac3afe9e8b8a3133206bccd6011a92b9aa791a9196b93" + }, + "norm_out.norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "91ac17792929e0a84533b28caf71ac831ec3d666c07a473393639be20392f5ce" + }, + "norm_out.linear.weight": { + "dtype": "bfloat16", "shape": [10752, 2688], + "sha256": "38063ae4b492aacd320e71d4a38a32325c0d47b29f59978f0f4979ca080fe673" + }, + "norm_out.linear.bias": { + "dtype": "bfloat16", "shape": [10752], + "sha256": "3b12cdc7bbda772e7e31539ffa9a32f529dba7f4148273d2128bdc621fd27438" + } + }, + "block_tensors": { + "transformer_blocks.0.attn.to_q.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "196f9120ad018784bcefd4224f0e08fae88ba1adf079ecdca18410a414a6a2a4" + }, + "transformer_blocks.0.attn.to_k.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "8e9171f3b8cf894904d1ddbd96510739d7291d2992310de6a52e1d733ec895cd" + }, + "transformer_blocks.0.attn.to_v.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "275bbc993b1ba738831e04d86f290c9988108d96fe8c97c3b576a043c7ff3fe1" + }, + "transformer_blocks.0.attn.norm_q.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "34a3f13557a2c0a9af260e48d431041e381779a0f61616cfe9800ec839693519" + }, + "transformer_blocks.0.attn.norm_k.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "8991949ddea3bc2136a0fcf0b980fa63bd0df760207e3116f7ce96df4e6ca058" + }, + "transformer_blocks.0.attn.to_out.0.weight": { + "dtype": "bfloat16", "shape": [5376, 7168], + "sha256": "822cb426aaebe57d3d219b3c4f6855ca5c13821cb84ebe0a7474fe3a87869fd6" + }, + "transformer_blocks.0.ff.net.0.proj.weight": { + "dtype": "bfloat16", "shape": [28672, 5376], + "sha256": "49587044f325aae73f046c4cda9637052670227ec3c2c289359d9d1ebf29f95d" + }, + "transformer_blocks.0.ff.net.2.weight": { + "dtype": "bfloat16", "shape": [5376, 14336], + "sha256": "e19d4e1649e4b55b1931842e5cc4bb46892c2e635b9315b2996630a86284e88f" + } + }, + "reference_implementation": { + "repo": "huggingface/diffusers", + "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "file": "src/diffusers/models/transformers/transformer_minimax_h3.py", + "file_sha256": "435873d9bc26b8770342961e18cb20324aeeaa4d045e85e0941638d4857d7db3", + "timestep_embedding_file": "src/diffusers/models/embeddings.py" + }, + "conventions": { + "timestep": "distinct timesteps t = 1 - sigma, unscaled in [0, 1], shape (T,)", + "modality_tags": {"video": 0, "text": 1, "audio": 2}, + "adaln_row": "timestep_index * 3 + token_tag", + "adaln_table_rows": "[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]", + "adaln_chunk_order": ["shift_msa", "scale_msa", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"], + "precision": { + "timestep_sinusoid": "float32", + "timestep_mlp": "float32 weights, float32 activations", + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output", + "rmsnorm": "nn.RMSNorm(5376, eps=1e-5), bfloat16 affine weight, statistics in float32", + "final_modulation": "norm_out.linear(silu(temb).bfloat16) -> (shift, scale), indexed by timestep_indices" + } + } +} diff --git a/rl_engine/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py new file mode 100644 index 000000000..3eec1818f --- /dev/null +++ b/rl_engine/validation/models/h3_provider.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The diffusers MiniMax-H3 conditioning path, op for op, without diffusers. + +Each function replays the tensor ops of the pinned reference +(``huggingface/diffusers@f53d552``, see ``h3_manifest.json``) in the same +order and dtypes, so it runs the same kernels the provider does on a given +device. This is the *provider* side of the RFC #420 comparisons; the FP32/FP64 +goldens live in ``rl_engine.reference.minimax_h3``. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + + +def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: + """Replay provider FP32 cosine/sine features for ``(T,)`` timesteps. + + Matches ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)`` + with the default channel count, returning shape ``(T, 256)``. + """ + + half_dim = num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / (half_dim - 0) + emb = torch.exp(exponent) + emb = timestep[:, None].float() * emb[None, :] + emb = 1 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb + + +def provider_time_embedder( + features: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> torch.Tensor: + """Replay linear_1, SiLU, and linear_2 using the checkpoint weight dtype. + + H3's FP32 weights map features of shape ``(T, 256)`` to ``(T, 2688)``. + """ + + sample = F.linear(features.to(w1.dtype), w1, b1) + sample = F.silu(sample) + return F.linear(sample, w2, b2) + + +def provider_adaln_modulation( + temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 +) -> tuple[torch.Tensor, ...]: + """Replay provider SiLU, the weight-dtype cast, and the AdaLN projection. + + Return six ``(3T, hidden_size)`` views in timestep-major modality order + for checkpoint weights shaped ``(18 * hidden_size, 2688)``. + """ + + temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + temb = temb.view(-1, 6 * hidden_size) + return temb.chunk(6, dim=-1) + + +def provider_adaln_row_gather( + modulation: tuple[torch.Tensor, ...], timestep_indices: torch.Tensor, token_tags: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """``adaln_indices = timestep_indices * 3 + token_tags`` then six ``index_select``.""" + + adaln_indices = timestep_indices * 3 + token_tags + return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) + + +def provider_norm_modulate( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``norm(x) * (1.0 + scale[i]) + shift[i]`` with ``index_select``; block and norm_out.""" + + norm_hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( + 0, indices + ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output + + +def provider_final_adaln_out( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + temb: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + timestep_indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``MiniMaxH3AdaLayerNormOut.forward``: shift first, indexed by timestep.""" + + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +# --- One transformer block (RFC #420 ``ws1_one_h3_block``) ------------------------------------ + + +def provider_rope( + position_ids: torch.Tensor, rope_freq_dim: int = 16, rope_theta: float = 10000.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """``MiniMaxH3RotaryPosEmbed``: FP32 ``(S, 96)`` cos/sin over the ``(t, h, w)`` axes.""" + + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device) + / (2 * rope_freq_dim) + ) + ) + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def provider_apply_rotary( + hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """``_apply_rotary_emb``: rotate-half on the leading 96 of 128 channels of ``(B, S, H, D)``. + + The FP32 tables are cast to the activation dtype (BF16) before the rotation. + """ + + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + x1, x2 = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-x2, x1), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + +def provider_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A bias-free ``nn.Linear`` of the block (Q/K/V, O and both FFN projections).""" + + return F.linear(x, weight) + + +def provider_head_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + """``norm_q``/``norm_k``: ``nn.RMSNorm(128)`` on ``(B, S, H, D)``.""" + + return F.rms_norm(x, (x.shape[-1],), weight, eps) + + +def provider_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """The native diffusers backend: SDPA on ``(B, H, S, D)`` views, non-causal, no mask. + + Returns ``(B, S, H*D)`` after ``flatten(2, 3).type_as(query)``. + """ + + q, k, v = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + return out.permute(0, 2, 1, 3).flatten(2, 3).type_as(query) + + +def provider_swiglu(projected: torch.Tensor) -> torch.Tensor: + """diffusers ``SwiGLU`` after its projection: ``hidden * silu(gate)`` on the two halves.""" + + hidden_states, gate = projected.chunk(2, dim=-1) + return hidden_states * F.silu(gate) diff --git a/rl_engine/validation/models/h3_refiner.py b/rl_engine/validation/models/h3_refiner.py new file mode 100644 index 000000000..60fcd9383 --- /dev/null +++ b/rl_engine/validation/models/h3_refiner.py @@ -0,0 +1,330 @@ +# 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, + on_block: Callable[[int, torch.Tensor], 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] + if on_block is not None: + on_block(block, h) + 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 new file mode 100644 index 000000000..16b212e6d --- /dev/null +++ b/rl_engine/validation/models/h3_report.py @@ -0,0 +1,831 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Per-operator performance and accuracy measurements for the H3 evidence. + +``PERF_CASES[op]`` builds the timed cases that ``benchmarks/models/benchmark_h3_conditioning.py`` +prints and ``tools/validation/models/h3_evidence.py`` stores. ``ACCURACY[op]`` measures the +op against the provider path and its golden. Each RFC #420 row registers both. +""" + +from __future__ import annotations + +import statistics +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) + +# Keys of a perf case that hold a timed callable, in display order. +TIMED_KEYS = ( + "candidate", + "candidate_checked", + "provider", + "candidate_backward", + "provider_backward", +) + + +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ + + for _ in range(warmup): + _prepared_call(fn, setup)() + torch.cuda.synchronize() + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + run() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + row["timing_samples_us"] = samples + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) + return row + + +# --------------------------------------------------------------------------- # +# timestep_sinusoid_h3 +# --------------------------------------------------------------------------- # + + +def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA sinusoid timing cases with checked, unchecked, and provider calls.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + cases = [] + for num in (1, 2, 4, 64): + t = h3_timesteps(num) + cases.append( + { + "op": "timestep_sinusoid_h3", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": num * 4 + num * 256 * 4, + "candidate": lambda t=t: op.forward(t, check_range=False), + "candidate_checked": lambda t=t: op.forward(t), + "provider": lambda t=t: provider_time_proj(t), + } + ) + return cases + + +def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + """Compare sinusoid outputs to provider and FP64 golden across timestep counts.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + rows = [] + for num in (1, 2, 3, 7, 64, 1000, 4097): + t = h3_timesteps(num, seed=num) + out, ref, gold = op(t), provider_time_proj(t), golden.forward_fp32(t) + rows.append( + { + "num_timesteps": num, + "bitwise_equal_to_provider": bool(torch.equal(out, ref)), + "max_abs_vs_fp64": float((out - gold).abs().max()), + "provider_max_abs_vs_fp64": float((ref - gold).abs().max()), + } + ) + return {"contract_atol": 1e-5, "cases": rows} + + +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Return CUDA checkpoint parameters, or seeded random tensors when weights are unset.""" + + if h3_weights_dir() is not None: + weights = load_h3_conditioning_weights("cuda", names) + return [weights[name] for name in names] + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 + for shape in shapes + ] + + +# --------------------------------------------------------------------------- # +# timestep_mlp_fp32 +# --------------------------------------------------------------------------- # + +MLP_NAMES = [f"time_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")] +MLP_SHAPES = [(5376, 256), (5376,), (2688, 5376), (2688,)] + + +def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA MLP candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + weight_bytes = sum(p.numel() * p.element_size() for p in params) + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + x = torch.rand(num, 256, device="cuda", generator=gen) * 2 - 1 + cases.append( + { + "op": "timestep_mlp_fp32", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight_bytes, + "candidate": lambda x=x: op(x, *params), + "provider": lambda x=x: provider_time_embedder(x, *params), + } + ) + return cases + + +def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + """Measure per-draw MLP error against FP64 and test single-row/batched equality.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + cuda_err, provider_err = [], [] + for seed in range(draws): + x = sinusoid(h3_timesteps(4, seed=seed)) + gold = golden.forward_fp32(x, *params) + cuda_err.append(float((op(x, *params) - gold).abs().max())) + provider_err.append(float((provider_time_embedder(x, *params) - gold).abs().max())) + # Row invariance: each timestep alone equals its row in the batch of 9. + x = sinusoid(h3_timesteps(9, seed=1)) + full = op(x, *params) + invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) + return { + "contract_atol": 1e-4, + "draws": draws, + "num_timesteps": 4, + "cuda_max_abs_vs_fp64": cuda_err, + "provider_max_abs_vs_fp64": provider_err, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_projection_3mod +# --------------------------------------------------------------------------- # + +ADALN_NAMES = [ + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +] + + +def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + """Return block-0 AdaLN projection weights and bias as BF16 CUDA tensors.""" + + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) + return weight.bfloat16(), bias.bfloat16() + + +def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA AdaLN candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + weight, bias = _adaln_params() + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) + cases.append( + { + "op": "adaln_projection_3mod", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight.numel() * weight.element_size(), + "candidate": lambda temb=temb: op(temb, weight, bias), + "provider": lambda temb=temb: provider_adaln_modulation(temb, weight, bias), + } + ) + return cases + + +def _flat_cat(outputs) -> torch.Tensor: + """Concatenate flattened modulation tensors in their output order.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: + """Report per-draw BF16 equality to FP64 goldens and single-row/batch invariance.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["adaln_projection_3mod"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + weight, bias = _adaln_params() + cuda_frac, provider_frac, early_frac = [], [], [] + for seed in range(draws): + g = torch.Generator(device="cuda").manual_seed(seed) + temb = torch.randn(4, 2688, device="cuda", generator=g) * 2 + gold = _flat_cat(golden.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat_cat(golden.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + ours = _flat_cat(op(temb, weight, bias)) + theirs = _flat_cat(provider_adaln_modulation(temb, weight, bias)) + cuda_frac.append(float((ours == gold).float().mean())) + provider_frac.append(float((theirs == gold).float().mean())) + early_frac.append(float((ours == early).float().mean())) + gen = torch.Generator(device="cuda").manual_seed(draws) + temb = torch.randn(9, 2688, device="cuda", generator=gen) + full = op(temb, weight, bias) + invariant = all( + all( + torch.equal(s, f[3 * i : 3 * i + 3]) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) + ) + for i in range(9) + ) + return { + "draws": draws, + "num_timesteps": 4, + "cuda_correctly_rounded": cuda_frac, + "provider_correctly_rounded": provider_frac, + "early_cast_golden_match": early_frac, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_row_gather +# --------------------------------------------------------------------------- # + +GATHER_SEQ_LENS = (4097, 32768, 131072) + + +def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_row_gather", device="cuda") + num_timesteps, hidden = 3, 5376 + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda", generator=gen).bfloat16() + chunks = rows.chunk(6, dim=-1) + cases = [] + for seq in GATHER_SEQ_LENS: + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + grads = [ + torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) + ] + + def prepare_backward(fn, ti=ti, tags=tags): + leaf = rows.detach().requires_grad_(True) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) + + cases.append( + { + "op": "adaln_row_gather", + "case": f"S={seq}", + "backend": type(op).__name__, + # bytes written (six (S, H) BF16 outputs); the 3T-row table stays in L2 + "bytes": 6 * seq * hidden * 2, + "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), + "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) + ), + } + ) + return cases + + +def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: + from rl_engine.validation.models.h3_chain import run_backward_case + + op = registry.get_op("adaln_row_gather", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + # Seeded so the op-backward fractions below are reproducible from this commit. + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() + forward_bitwise = {} + for seq in (1, 257, 4097, 32768): + ti, tags = h3_packed_layout(seq, 3, seed=seq) + ours = op(rows, ti, tags) + theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) + + # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = [torch.randn(4097, 5376, device="cuda", generator=gen).bfloat16() for _ in range(6)] + + def grad_of(fn): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(fn(leaf)), grads) + return leaf.grad + + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward(list(golden.forward_fp32(ref, ti, tags)), [g.float() for g in grads]) + gold = ref.grad.bfloat16() + backward = {} + for name, fn in ( + ("cuda", lambda leaf: op(leaf, ti, tags)), + ("provider", lambda leaf: provider_adaln_row_gather(leaf.chunk(6, dim=-1), ti, tags)), + ): + first, second = grad_of(fn), grad_of(fn) + backward[name] = { + "repeat_bitwise_equal": bool(torch.equal(first, second)), + "correctly_rounded_fraction": float((first == gold).float().mean()), + } + + result: dict[str, Any] = { + "forward_bitwise_vs_index_select": forward_bitwise, + "op_backward": backward, + "chain_backward": [], + } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return result + + +# --------------------------------------------------------------------------- # +# h3_rmsnorm (block norm1 + MSA modulation) +# --------------------------------------------------------------------------- # + +NORM_SEQ_LENS = (4097, 32768, 131072) + + +def _norm_inputs(seq: int, seed: int = 0): + weight = h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0].bfloat16() + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).bfloat16() + shift, scale = table.view(9, 6, 5376)[:, 0], table.view(9, 6, 5376)[:, 1] + ti, tags = h3_packed_layout(seq, 3, seed=seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + return x, weight, shift, scale, ti * 3 + tags + + +def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("h3_rmsnorm", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, weight, shift, scale, index = _norm_inputs(seq) + grad = torch.randn_like(x) + + def prepare_backward(fn, x=x, weight=weight, shift=shift, scale=scale): + leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) + + cases.append( + { + "op": "h3_rmsnorm", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 2 * seq * 5376 * 2, # read x, write y + "candidate": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + op.forward_modulated(x, w, sh, sc, i, check_range=False) + ), + "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + provider_norm_modulate(x, w, sh, sc, i) + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: op.forward_modulated( + x_, w_, sh_, sc_, i, check_range=False + ) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) + ), + } + ) + return cases + + +def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("h3_rmsnorm", device="cuda") + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + gen = torch.Generator(device="cuda").manual_seed(0) + plain = {} + for name in names: + x = torch.randn(2, 777, 5376, device="cuda", generator=gen).bfloat16() + ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) + plain[name] = bool(torch.equal(op(x, weights[name]), ref)) + x, weight, shift, scale, index = _norm_inputs(4097, seed=1) + modulated = bool( + torch.equal( + op.forward_modulated(x, weight, shift, scale, index), + provider_norm_modulate(x, weight, shift, scale, index), + ) + ) + invariant = bool( + torch.equal( + op.forward_modulated(x[:, 100:140], weight, shift, scale, index[100:140])[0], + op.forward_modulated(x, weight, shift, scale, index)[0, 100:140], + ) + ) + + # Backward against FP64, for the CUDA op and the diffusers expression. + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, weight, shift, scale)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, w_, sh_, sc_): + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * w_ + return n * (1 + sc_.index_select(0, index)) + sh_.index_select(0, index) + + ref = grads(golden, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op.forward_modulated(*t, index)), + ("provider", lambda *t: provider_norm_modulate(t[0], t[1], t[2], t[3], index)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) + }, + } + return { + "weight_source": weight_source, + "plain_bitwise_vs_nn_rmsnorm": plain, + "modulated_bitwise_vs_diffusers": modulated, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# final_adaln_out (norm_out) +# --------------------------------------------------------------------------- # + +FINAL_NAMES = ["norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias"] + + +def _final_inputs(seq: int, seed: int = 0): + norm_weight, weight, bias = h3_params(FINAL_NAMES, [(5376,), (10752, 2688), (10752,)]) + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, _ = h3_packed_layout(seq, 3, seed=seed) + return x, norm_weight.bfloat16(), temb, weight.bfloat16(), bias.bfloat16(), ti + + +def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("final_adaln_out", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, nw, temb, w, b, ti = _final_inputs(seq) + grad = torch.randn_like(x) + + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + leaves = [t.detach().requires_grad_(True) for t in tensors] + return fn(*leaves), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "final_adaln_out", + "case": f"S={seq}", + "backend": type(op).__name__, + # read x + norm_out.linear, write out + "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, + "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out( + t[0], t[1], t[2], t[3], t[4], i + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], i) + ), + } + ) + return cases + + +def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("final_adaln_out", device="cuda") + golden_op = registry._get_or_create_backend( + registry._priority_map["cpu"]["final_adaln_out"][-1] + ) + if golden_op is None: + raise RuntimeError("required H3 golden backend is unavailable") + x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) + ours = op(x, nw, temb, w, b, ti) + theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) + golden = golden_op.forward_fp32(x, nw, temb, w, b, ti) + invariant = bool( + torch.equal(op(x[:, 100:160], nw, temb, w, b, ti[100:160])[0], ours[0, 100:160]) + ) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, nw, temb, w, b)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden_backward(x_, nw_, t_, w_, b_): + # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. + act = t_ * torch.sigmoid(t_) + # The declared cast rounds the FP32 SiLU, as the model computes it. + declared = torch.nn.functional.silu(t_.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = torch.nn.functional.linear(act, w_, b_) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw_ + return n * (1.0 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ref = grads(golden_backward, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], ti)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) + }, + } + return { + "equal_to_diffusers_fraction": float((ours == theirs).float().mean()), + "max_abs_vs_golden": float((ours.float() - golden).abs().max()), + "provider_max_abs_vs_golden": float((theirs.float() - golden).abs().max()), + "rows_batch_invariant": invariant, + "backward": backward, + } + + +PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { + "timestep_sinusoid_h3": _sinusoid_perf, + "timestep_mlp_fp32": _mlp_perf, + "adaln_projection_3mod": _projection_perf, + "adaln_row_gather": _gather_perf, + "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, + "final_adaln_out": _final_perf, +} +ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { + "timestep_sinusoid_h3": _sinusoid_accuracy, + "timestep_mlp_fp32": _mlp_accuracy, + "adaln_projection_3mod": _projection_accuracy, + "adaln_row_gather": _gather_accuracy, + "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, + "final_adaln_out": _final_accuracy, +} diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py new file mode 100644 index 000000000..2234b2400 --- /dev/null +++ b/rl_engine/validation/models/h3_tp.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One H3 block under tensor / sequence / context parallelism (RFC #420 WS2 closeouts). + +``run_tp_block`` is the block graph with the attention and FFN projections replaced by the +``h3-parallel-v1`` operators: column-parallel Q/K/V and gate/up, per-head norm, rotary and +attention on the rank's heads, Ulysses head<->sequence all-to-all across the CP group, +row-parallel O and down with the deterministic collective, optional sequence parallelism +(rows of the residual stream sharded across the TP group). With ``tp=cp=1`` it is the +canonical single-GPU block every other layout must equal bitwise. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.validation.models import h3_block as hb + +HIDDEN, HEAD_DIM = hb.HIDDEN, hb.HEAD_DIM + + +@dataclass +class Grid: + tp: int + cp: int + tp_rank: int + cp_rank: int + tp_collective: Any = None + cp_group: Any = None + + def seq_shard(self, length: int) -> par.SeqShard: + return par.SeqShard(length, self.cp, self.cp_rank) + + +def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid: + """Rank ``r`` is ``(cp_rank, tp_rank) = divmod(r, tp)``; TP groups are consecutive ranks. + + ``local_size`` (ranks per node) selects the cross-node collective when the TP group spans + nodes; ``None`` means the TP group is on one node (CUDA IPC fixed tree). + """ + + from rl_engine.distributed.algorithms.collectives import DeterministicCollective + from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective + + world, rank = dist.get_world_size(), dist.get_rank() + if world != tp * cp: + raise ValueError(f"world {world} != tp {tp} x cp {cp}") + cp_rank, tp_rank = divmod(rank, tp) + tp_groups = [dist.new_group(list(range(c * tp, (c + 1) * tp))) for c in range(cp)] + cp_groups = [dist.new_group(list(range(t, world, tp))) for t in range(tp)] + collective: DeterministicCollective | HierarchicalCollective | None = None + if tp > 1: + mine = tp_groups[cp_rank] + if local_size is not None and local_size < tp: + collective = HierarchicalCollective( + local_size, mine, torch.cuda.current_device(), max_size_bytes=256 << 20 + ) + else: + collective = DeterministicCollective( + mine, torch.cuda.current_device(), max_size_bytes=256 << 20 + ) + return Grid(tp, cp, tp_rank, cp_rank, collective, cp_groups[tp_rank] if cp > 1 else None) + + +def _attention_fn(ops): + def attend(q, k, v): # (B, S, H', D) -> (B, S, H', D) + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, v)) + return ops.attention(q, k, v, causal=False).transpose(1, 2).contiguous() + + return attend + + +def run_tp_block( + weights: dict[str, torch.Tensor], + hidden: torch.Tensor, + temb: torch.Tensor, + index: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor], + ops: hb.CandidateOps, + grid: Grid, + *, + sp: bool = False, + seq_len: int | None = None, + leaves: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + """Block 0 for this rank. ``hidden/index/rope`` are the full packed sequence on every rank. + + The rank slices its CP rows (``cp_metadata_coshard``: hidden, AdaLN row index and both + rotary tables with the same ``SeqShard``) and, with ``sp``, its TP row share of those. + Returns this rank's rows of the block output (CP rows, or the SP share of them). With + ``leaves`` the weight shards are tracked and stored there by name so the caller can read + their gradients. + """ + + if hidden.shape[0] != 1: + raise ValueError("the parallel block takes one packed sequence (batch 1)") + tp, w = grid.tp, (lambda name: weights[f"{hb.BLOCK}.{name}"]) + seq = seq_len or hidden.shape[1] + shard = grid.seq_shard(seq) + h = par.coshard(shard, hidden, dim=1)[0] + idx = par.coshard(shard, index, dim=0)[0] + cos, sin = par.coshard(shard, rope[0], rope[1], dim=0) + rows = h.shape[1] + if sp: + if rows % tp: + raise ValueError(f"sequence parallelism needs {rows} local rows divisible by tp={tp}") + share = rows // tp + lo, hi = grid.tp_rank * share, (grid.tp_rank + 1) * share + h_res, idx_res = h[:, lo:hi].contiguous(), idx[lo:hi].contiguous() + else: + h_res, idx_res = h, idx + coll = grid.tp_collective + + def leaf(name, tensor): + tensor = tensor.detach().clone() + if leaves is not None: + tensor.requires_grad_(True) + leaves[name] = tensor + return tensor + + table = ops.projection(temb, w("adaln_proj.linear.weight"), w("adaln_proj.linear.bias")) + + def proj_in(x, name_weights, gather): + flat = x.reshape(-1, HIDDEN) + return par.column_parallel_linear(flat, name_weights, coll, tp, gather=gather) + + x = ops.norm.forward_modulated(h_res, w("norm1.weight"), table[0], table[1], idx_res) + q, k, v = proj_in( + x, + [ + leaf(f"attn.to_{n}.weight", par.shard_rows(w(f"attn.to_{n}.weight"), tp, grid.tp_rank)) + for n in "qkv" + ], + sp, + ) + local_heads = par.HEADS // tp + + def heads(t): + return t.view(1, rows, local_heads, HEAD_DIM) + + q = ops.norm(heads(q), w("attn.norm_q.weight")) + k = ops.norm(heads(k), w("attn.norm_k.weight")) + from rl_engine.validation.models import h3_provider as prov + + q, k = prov.provider_apply_rotary(q, cos, sin), prov.provider_apply_rotary(k, cos, sin) + attn = par.ulysses_attention(q, k, heads(v), _attention_fn(ops), shard, grid.cp_group) + o = par.row_parallel_linear( + attn.reshape(rows, local_heads * HEAD_DIM), + leaf("attn.to_out.0.weight", par.shard_cols(w("attn.to_out.0.weight"), tp, grid.tp_rank)), + coll, + tp, + scatter=sp, + ) + h1 = ops.gate(h_res, o.view(1, -1, HIDDEN), table[2], idx_res) + x2 = ops.norm.forward_modulated(h1, w("norm2.weight"), table[3], table[4], idx_res) + fused = w("ff.net.0.proj.weight") + half = fused.shape[0] // 2 + hid, gate = proj_in( + x2, + [ + leaf("ff.hidden", par.shard_rows(fused[:half], tp, grid.tp_rank)), + leaf("ff.gate", par.shard_rows(fused[half:], tp, grid.tp_rank)), + ], + sp, + ) + act = ops.swiglu(gate.contiguous(), hid.contiguous()) + down = par.row_parallel_linear( + act, + leaf("ff.net.2.weight", par.shard_cols(w("ff.net.2.weight"), tp, grid.tp_rank)), + coll, + tp, + scatter=sp, + ) + return ops.gate(h1, down.view(1, -1, HIDDEN), table[5], idx_res) + + +def block_target( + collective, + dummy, + *, + tp: int, + cp: int, + sp: bool = False, + seq_len: int = 0, + backward: bool = False, + local_size: int = -1, + layout: str = "tiny", + seed: int = 0, +): + """``run_world`` target: this rank's rows of the parallel block (and gradients) for one case.""" + + del collective # the grid builds its own TP/CP groups + from rl_engine.validation.models.h3_weights import load_h3_block_weights + + grid = make_grid(tp, cp, local_size=local_size if local_size > 0 else None) + return block_case( + load_h3_block_weights("cuda"), + grid, + sp=sp, + seq_len=seq_len, + backward=backward, + layout=layout, + seed=seed, + ) + + +def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layout="tiny", seed=0): + """Run block 0 on ``grid``; the global rows this rank returns are in ``rows``.""" + + ctx = hb.make_block_context(weights, layout=layout, seed=seed) + n = seq_len or ctx["seq_len"] + hidden = ctx["hidden"][:, :n].clone().requires_grad_(backward) + index = ctx["index"][:n] + rope = (ctx["rope"][0][:n], ctx["rope"][1][:n]) + ops = hb.CandidateOps() + leaves: dict[str, torch.Tensor] | None = {} if backward else None + with torch.set_grad_enabled(backward): + out = run_tp_block( + weights, hidden, ctx["temb"], index, rope, ops, grid, sp=sp, seq_len=n, leaves=leaves + ) + shard = grid.seq_shard(n) + lo, hi = shard.lo, shard.hi + if sp: + share = (hi - lo) // grid.tp + lo, hi = lo + grid.tp_rank * share, lo + (grid.tp_rank + 1) * share + result = {"out": out.detach(), "rows": torch.tensor([lo, hi])} + if backward: + gen = torch.Generator().manual_seed(seed + 99) + full_grad = torch.randn(1, n, HIDDEN, generator=gen).to(out.dtype).cuda() + out.backward(full_grad[:, lo:hi]) + result["dx"] = hidden.grad[:, lo:hi].detach() + assert leaves is not None + for name, tensor in leaves.items(): + result["dw." + name] = tensor.grad.detach() + return result diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py new file mode 100644 index 000000000..d1c051415 --- /dev/null +++ b/rl_engine/validation/models/h3_weights.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pinned MiniMax-H3 conditioning weights for the RFC #420 tests. + +Only the tensors listed in ``h3_manifest.json`` are used. They all live in +shard 1 of 14, so ``tools/weights/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. With ``--block`` it also +extracts block 0's attention and FFN weights (manifest ``block_tensors``, +about 1.1 GB) into ``h3_block0.safetensors`` for ``ws1_one_h3_block``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch + +MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") +WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" +EXTRACTED_FILE = "h3_conditioning.safetensors" +BLOCK_FILE = "h3_block0.safetensors" + + +def load_h3_manifest() -> dict[str, Any]: + """Read the bundled checkpoint identity and pinned tensor contracts.""" + + return json.loads(MANIFEST_PATH.read_text()) + + +def h3_weights_dir() -> Path | None: + """Expand the configured extraction directory, or return None when unset or blank.""" + + value = os.environ.get(WEIGHTS_ENV, "").strip() + return Path(value).expanduser() if value else None + + +def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + """Return a file SHA256 while reading bounded chunks.""" + + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def sha256_tensor(tensor: torch.Tensor) -> str: + """Hash a CPU tensor's contiguous bytes without safetensors metadata.""" + + tensor_bytes = tensor.contiguous().view(torch.uint8).numpy() + return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() + + +def _load_pinned( + filename: str, section: str, device: torch.device | str, names: list[str] | None +) -> dict[str, torch.Tensor]: + """Load one extracted file and check every tensor of a manifest section against it.""" + + from safetensors.torch import load_file + + root = h3_weights_dir() + if root is None: + raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") + path = root / filename + if not path.is_file(): + flag = " --block" if section == "block_tensors" else "" + raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py{flag}") + specs = load_h3_manifest()[section] + tensors = load_file(str(path), device="cpu") + wanted = names if names is not None else list(specs) + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + out: dict[str, torch.Tensor] = {} + for name in wanted: + if name not in specs: + raise KeyError(name) + out[name] = tensors[name].to(device=device) + return out + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + return _load_pinned(EXTRACTED_FILE, "tensors", device, names) + + +def load_h3_block_weights(device: torch.device | str = "cpu") -> dict[str, torch.Tensor]: + """Conditioning tensors plus block 0's attention/FFN weights, every one sha-checked.""" + + return { + **load_h3_conditioning_weights(device), + **_load_pinned(BLOCK_FILE, "block_tensors", device, None), + } + + +# --- 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 aaebc279f..15d7adc91 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -16,6 +16,10 @@ DEFAULT_VOCAB = 151936 DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 +H3_FREQ_DIM = 256 +H3_TIME_HIDDEN = 5376 +H3_TIME_EMBED = 2688 +H3_HIDDEN = 5376 def make_operator_inputs( @@ -24,6 +28,7 @@ def make_operator_inputs( dtype: torch.dtype, device: torch.device, ) -> dict[str, Any]: + """Build operator keyword inputs on the requested device from CLI shape and seed options.""" builders = { "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, @@ -44,6 +49,13 @@ def make_operator_inputs( "embedding": _make_embedding_inputs, "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, + "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, + "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, + "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, + "adaln_row_gather": _make_adaln_row_gather_inputs, + "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -52,6 +64,7 @@ def make_operator_inputs( def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: + """Return the operator's dimension label for benchmark and evidence reports.""" if op_name == "video_patch_pack_unpack": shape = "x".join(str(dim) for dim in _video_patch_shape(args)) return f"{_arg_str(args, 'patch_direction', 'pack')}-{shape}" @@ -78,6 +91,16 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", + "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", + "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" + f"x{H3_TIME_EMBED}", + "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" + f"x{6 * 3 * _h3_hidden(args)}", + "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" + f"->{batch * seq}", + "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -370,6 +393,151 @@ def _make_kv_cache_attention_inputs( } +def _h3_num_timesteps(args: argparse.Namespace) -> int: + """Read the packed H3 timestep count, falling back to the CLI batch size.""" + # H3 packs a handful of distinct timesteps; reuse --batch as their count. + return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) + + +def _h3_timesteps( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """Build seeded timesteps in [0, 1] with endpoint cases in the requested storage dtype.""" + num = _h3_num_timesteps(args) + mode = _arg_str(args, "input_mode", "random") + if mode == "constant": + t = torch.full((num,), 0.5, device=device) + else: + t = torch.rand((num,), generator=_generator(args, device, offset=7), device=device) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(dtype) + + +def _make_timestep_sinusoid_h3_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Supply a one-dimensional packed timestep tensor to the H3 sinusoid operator.""" + return {"timestep": _h3_timesteps(args, dtype, device)} + + +def _make_timestep_mlp_fp32_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build [T, 256] sinusoid features and seeded FP32 256→5376→2688 MLP parameters.""" + # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. + del dtype + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + features = NativeH3TimestepSinusoidOp().forward(_h3_timesteps(args, torch.float32, device)) + scale1, scale2 = H3_FREQ_DIM**-0.5, H3_TIME_HIDDEN**-0.5 + return { + "x": features, + "w1": _floating_tensor((H3_TIME_HIDDEN, H3_FREQ_DIM), args, torch.float32, device, 1) + * scale1, + "b1": _floating_tensor((H3_TIME_HIDDEN,), args, torch.float32, device, 2) * 0.1, + "w2": _floating_tensor((H3_TIME_EMBED, H3_TIME_HIDDEN), args, torch.float32, device, 3) + * scale2, + "b2": _floating_tensor((H3_TIME_EMBED,), args, torch.float32, device, 4) * 0.1, + } + + +def _h3_hidden(args: argparse.Namespace) -> int: + """Read the AdaLN channel width, defaulting to the checkpoint's 5376 channels.""" + return _arg_int(args, "normalized_dim", H3_HIDDEN) + + +def _make_adaln_projection_3mod_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build FP32 [T, 2688] embeddings and 18H projection parameters in the chosen dtype.""" + # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the + # projection's weight dtype, BF16 in the checkpoint. + n_out = 6 * 3 * _h3_hidden(args) + temb = _floating_tensor( + (_h3_num_timesteps(args), H3_TIME_EMBED), args, torch.float32, device, 0 + ) + weight = _floating_tensor((n_out, H3_TIME_EMBED), args, torch.float32, device, 1) + bias = _floating_tensor((n_out,), args, torch.float32, device, 2) + return { + "temb": temb * 2.0, + "weight": (weight * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (bias * 0.1).to(dtype), + } + + +def _make_adaln_row_gather_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # T distinct timesteps (--batch) and a packed sequence of batch * seq rows + # mixing all three modalities. + batch, seq = _batch_seq(args) + num_timesteps = _h3_num_timesteps(args) + generator = _generator(args, device, offset=29) + packed = batch * seq + token_tags = torch.randint(0, 3, (packed,), generator=generator, device=device) + timestep_indices = torch.randint( + 0, num_timesteps, (packed,), generator=generator, device=device + ) + rows = _floating_tensor((3 * num_timesteps, 6 * _h3_hidden(args)), args, dtype, device, 0) + return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} + + +def _make_h3_rmsnorm_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + weight = _floating_tensor((hidden,), args, torch.float32, device, 1).abs() + 0.5 + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "weight": weight.to(dtype), + "eps": 1e-5, + } + + +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + +def _make_final_adaln_out_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract; norm_out.linear/norm/x use ``dtype`` (BF16 in H3). + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + num_timesteps = _h3_num_timesteps(args) + # Realistic scales (norm weight in [0.5, 1], unit activations): with unscaled + # inputs the parameter gradients sum hundreds of rows into values ~1e3, and + # their cancelling elements cannot meet a per-element FP32 atol. + linear = _floating_tensor((2 * hidden, H3_TIME_EMBED), args, torch.float32, device, 3) + norm = 0.5 + 0.5 * torch.sigmoid(_floating_tensor((hidden,), args, torch.float32, device, 4)) + generator = _generator(args, device, offset=37) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "norm_weight": norm.to(dtype), + "temb": _floating_tensor((num_timesteps, H3_TIME_EMBED), args, torch.float32, device, 1), + "weight": (linear * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (_floating_tensor((2 * hidden,), args, torch.float32, device, 2) * 0.1).to(dtype), + "timestep_indices": torch.randint( + 0, num_timesteps, (seq,), generator=generator, device=device + ), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 92995b938..d5e6a7c85 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -272,6 +272,110 @@ def _load_object(path: str) -> Any: }, grad_input_names=("logits",), ), + # MiniMax-H3 (RFC #420) conditioning path. The op is FP32 end to end, so + # its FP32 rows are the declared contract; other dtypes only change the + # timestep input's storage dtype. + "timestep_sinusoid_h3": OperatorSpec( + name="timestep_sinusoid_h3", + op_class="elementwise", + gold_path="rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ), + }, + grad_input_names=("timestep",), + ), + "timestep_mlp_fp32": OperatorSpec( + name="timestep_mlp_fp32", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_mlp.H3TimestepMLPCudaOp" + ), + }, + grad_input_names=("x", "w1", "b1", "w2", "b2"), + ), + "adaln_projection_3mod": OperatorSpec( + name="adaln_projection_3mod", + op_class="reduction", + gold_path=("rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_projection.H3AdaLNProjectionCudaOp" + ), + }, + grad_input_names=("temb", "weight", "bias"), + ), + # The gather's forward is a copy (asserted bitwise in tests/models/minimax_h3), but its VJP + # sums every packed position of a table row, so it is judged as a reduction. + "adaln_row_gather": OperatorSpec( + name="adaln_row_gather", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_row_gather.H3AdaLNRowGatherCudaOp" + ), + }, + grad_input_names=("rows",), + ), + "h3_rmsnorm": OperatorSpec( + name="h3_rmsnorm", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + "cuda": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "gate_residual.H3GateResidualCudaOp" + ), + }, + grad_input_names=("residual", "y", "gate"), + ), + "final_adaln_out": OperatorSpec( + name="final_adaln_out", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "final_adaln_out.H3FinalAdaLNOutCudaOp" + ), + }, + grad_input_names=("x", "norm_weight", "temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..820347673 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -10,7 +10,12 @@ import torch from torch.utils import cpp_extension -from build_tools.extensions import _resolve_package_data_source, get_cmdclass +from build_tools.extensions import ( + H3_BITWISE_SOURCES, + H3_KERNELS_DEFINE, + _resolve_package_data_source, + get_cmdclass, +) def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: @@ -83,3 +88,28 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert _resolve_package_data_source(source, project_root=tmp_path) == source assert "build_py" in get_cmdclass() + + +def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + extension = _load_extension_config(monkeypatch, hip=None) + + assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["nvcc"] + assert "--use_fast_math" not in extension["extra_compile_args"]["nvcc"] + + +def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: + # --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for approximations + # and enables FTZ, which breaks the H3 kernels' bitwise parity promise; nvcc + # exposes no macro to reject it in-source, so the build drops the kernels + # and their bindings. + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + extension = _load_extension_config(monkeypatch, hip=None) + + assert "--use_fast_math" in extension["extra_compile_args"]["nvcc"] + assert not set(H3_BITWISE_SOURCES) & set(extension["sources"]) + assert not [s for s in extension["sources"] if s.startswith("csrc/cuda/h3/")] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py new file mode 100644 index 000000000..f9c2b000f --- /dev/null +++ b/tests/models/minimax_h3/conftest.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared fixtures for the MiniMax-H3 (RFC #420) operator tests. + +Real-weight tests read the pinned tensors from ``$RL_KERNEL_H3_WEIGHTS`` +(written by ``tools/weights/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_block_weights, + load_h3_conditioning_weights, +) + + +@pytest.fixture(scope="session") +def h3_weights_cpu() -> dict[str, torch.Tensor]: + """Load pinned CPU weights once per session, skipping unavailable artifacts.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") + try: + return load_h3_conditioning_weights("cpu") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(scope="session") +def h3_block_weights_cuda() -> dict[str, torch.Tensor]: + """Conditioning plus block-0 tensors on CUDA, skipping when they were not extracted.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py --block") + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + try: + return load_h3_block_weights("cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +def h3_kernels_built() -> bool: + """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +@pytest.fixture +def require_h3_kernels() -> None: + if not h3_kernels_built(): + pytest.skip("rl_engine._C has no H3 kernels (not built, or a fast-math build)") diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..e4c303e10 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.gate_residual import NativeH3GateResidualOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_gate_residual + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py new file mode 100644 index 000000000..066445fc0 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused AdaLN modulation (projection + row gather) with an FP32 table gradient. + +The forward must be bitwise equal to the two RFC #420 ops run separately; the +backward must be deterministic and must not re-round the table gradient to +BF16 (so ``d_temb`` reaches FP64-golden accuracy). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models.h3_cases import h3_packed_layout + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _ops(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks the H3 kernels") + return H3AdaLNModulationCudaOp(), H3AdaLNProjectionCudaOp(), H3AdaLNRowGatherCudaOp() + + +def _params(hidden, dim, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = 18 * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).bfloat16().cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).bfloat16().cuda() + temb = (torch.randn(3, dim, generator=g) * 2).cuda() + return temb, weight, bias + + +def _grads(seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + +def _leaf_grads(fn, temb, weight, bias, grads): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(fn(*leaves)), grads) + return [leaf.grad for leaf in leaves] + + +@requires_cuda +class TestFusedModulation: + def test_forward_bitwise_equal_to_separate_ops(self): + fused, proj, gather = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + ours = fused(temb, weight, bias, ti, tags) + separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) + assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + + def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): + fused, proj, gather = _ops() + weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() + bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() + g = torch.Generator(device="cuda").manual_seed(0) + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = _grads(4097, 5376, seed=1) + ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + separate = _leaf_grads( + lambda t, w, b: gather.gather_chunks(proj(t, w, b), ti, tags), + temb, + weight, + bias, + grads, + ) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column (about 1 temb draw in 30 dropped below 0.9999 that way). + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # identity-VJP cast + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) + index = ti * 3 + tags + torch.autograd.backward( + [chunk.index_select(0, index) for chunk in table.chunk(6, dim=-1)], + [g.double() for g in grads], + ) + + def rel(grad, golden): + return ((grad.double() - golden).abs().max() / golden.abs().max()).item() + + # d_temb: no BF16 rounding anywhere on its path, so it reaches FP32 accuracy. + assert rel(ours[0], ref[0].grad) < 1e-5 + assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 + # dW, db: only their own final BF16 rounding remains. + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + + def test_backward_repeat_bitwise(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=3) + ti, tags = h3_packed_layout(3000, 3, seed=3) + grads = _grads(3000, 64, seed=3) + first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + assert all(torch.equal(a, b) for a, b in zip(first, again)) + + def test_rejects_bad_indices(self): + fused, _, _ = _ops() + temb, weight, bias = _params(8, 16) + ti, tags = h3_packed_layout(10, 3) + with pytest.raises(IndexError): + fused(temb, weight, bias, ti + 3, tags) + with pytest.raises(TypeError): + fused(temb.bfloat16(), weight, bias, ti, tags) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..5fb10a73f --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_projection_3mod``: SiLU + 2688 -> 96768 projection, 6 x 3 table. + +* layout: the six outputs and the three modality rows per timestep land + exactly where diffusers puts them (RFC probe H4); +* mixed precision: SiLU in FP32, one cast to BF16; an early cast (probe H7) + is rejected at the API and is measurably different numerically; +* accuracy: forward and gradients against the declared-cast golden on the + pinned block-0 weights; +* invariance and determinism: a timestep's 18 modulation rows do not depend + on the other timesteps, and repeated runs are bitwise equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.validation.models.h3_provider import provider_adaln_modulation + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3AdaLNProjectionCudaOp() + + +def _temb(num, dim=2688, seed=0, device="cuda"): + """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(num, dim, generator=g) * 2).to(device) + + +def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + """Create seeded projection weights for six chunks and three modalities per channel.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +def _flat(outputs): + """Flatten and concatenate all modulation chunks for elementwise comparisons.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +class TestReference: + def test_provider_layout_matches_diffusers(self): + """Match the provider values and timestep-major modality layout on CPU.""" + + temb = _temb(2, dim=16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + for a, b in zip(ours, theirs, strict=True): + assert a.shape == (6, 8) and torch.equal(a, b) + + def test_rejects_bf16_temb(self): + # RFC probe H7: casting temb before the SiLU. + """Reject an early BF16 cast so SiLU receives the declared FP32 embeddings.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(TypeError, match="float32"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) + + def test_rejects_bad_shapes(self): + """Reject invalid projection dimensions, empty batches, and mismatched bias dtype.""" + + weight, bias = _synthetic(8, 16, device="cpu") + op = NativeH3AdaLNProjectionOp() + with pytest.raises(ValueError): # not a multiple of 6 * 3 rows + op.forward(_temb(2, 16, device="cpu"), weight[:-1], bias[:-1]) + with pytest.raises(ValueError): + op.forward(_temb(2, 15, device="cpu"), weight, bias) + with pytest.raises(ValueError): + op.forward(_temb(0, 16, device="cpu"), weight, bias) + with pytest.raises(TypeError): + op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + + +@requires_cuda +class TestLayout: + def test_channel_identity_lands_in_its_chunk_and_modality_row(self): + """Probe H4: bias o = m*6H + c*H + h tagged (m, c); zero weight passes it through.""" + + hidden, dim, num = 16, 32, 3 + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + o = torch.arange(n_out) + modality, chunk = o // (H3_ADALN_CHUNKS * hidden), (o // hidden) % H3_ADALN_CHUNKS + bias = (modality * 10 + chunk).to(torch.bfloat16).cuda() + weight = torch.zeros(n_out, dim, dtype=torch.bfloat16, device="cuda") + outputs = _cuda_op()(_temb(num, dim), weight, bias) + assert len(outputs) == 6 + for c, out in enumerate(outputs): + assert out.shape == (num * 3, hidden) + for t in range(num): + for m in range(3): + expected = torch.full((hidden,), float(m * 10 + c), device="cuda") + assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) + + def test_outputs_are_views_of_one_table_like_diffusers(self): + """Require all six modulation chunks to share storage with the expected offsets.""" + + weight, bias = _synthetic(16, 32) + outputs = _cuda_op()(_temb(2, 32), weight, bias) + base = outputs[0].untyped_storage().data_ptr() + assert all(out.untyped_storage().data_ptr() == base for out in outputs) + assert outputs[1].data_ptr() - outputs[0].data_ptr() == 16 * outputs[0].element_size() + + +@pytest.fixture(scope="module") +def block0(h3_weights_cpu): + """Move the pinned block-zero projection weight and bias to CUDA for this module.""" + + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4]) + def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + """Check checkpoint accuracy and near-exact BF16 rounding of the declared golden.""" + + weight, bias = block0 + temb = _temb(num, seed=num) + outputs = _cuda_op()(temb, weight, bias) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + for out, gold in zip(outputs, golden, strict=True): + assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) + # tolerance_contract.json forward_accuracy / reduction / bfloat16. + torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) + # Beyond the contract: almost every element is the correctly rounded + # golden value. The rest come from the FP32 tensor-core accumulation over + # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact + # equality is not reachable: over 300 temb draws per T the worst match was + # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). + declared = _flat(golden).bfloat16() + assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 + + def test_early_cast_is_numerically_detectable(self, block0): + """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" + + weight, bias = block0 + temb = _temb(3, seed=7) + op = NativeH3AdaLNProjectionOp() + ours = _flat(_cuda_op()(temb, weight, bias)) + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + assert (ours == declared).float().mean() >= 0.9995 + assert (ours == early).float().mean() < 0.8 + + def test_rows_are_batch_and_position_invariant(self, block0): + """Keep each timestep modulation row identical alone and under batch reordering.""" + + weight, bias = block0 + op = _cuda_op() + temb = _temb(5, seed=3) + full = op(temb, weight, bias) + for i in range(5): + single = op(temb[i : i + 1], weight, bias) + for s, f in zip(single, full, strict=True): + assert torch.equal(s, f[3 * i : 3 * i + 3]) + swapped = op(temb.flip(0), weight, bias) + for s, f in zip(swapped, full, strict=True): + assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) + + def test_backward_against_golden(self, block0): + """Check embedding and parameter gradients against a high-precision declared-cast VJP.""" + + weight, bias = block0 + temb = _temb(2, seed=5) + g = torch.Generator(device="cuda").manual_seed(0) + grads = [torch.randn(6, 5376, device="cuda", generator=g).bfloat16() for _ in range(6)] + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(_cuda_op()(*leaves)), grads) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column. + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # declared cast, identity VJP + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) + torch.autograd.backward(list(table), [gr.double() for gr in grads]) + # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad + # and BF16 weight/bias grads (rounded once from an FP32 accumulation). + torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) + for leaf, r in zip(leaves[1:], ref[1:], strict=True): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + # With T = 2 each dW/db element sums two products, and the one BF16 cast + # of the FP32 sum is the correctly rounded value of every element + # (checked on 500 of 500 temb draws). + assert torch.equal(leaf.grad, r.grad.to(torch.bfloat16)) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) + def test_other_sizes(self, dtype, hidden, dim): + """Match the projection golden at non-checkpoint dimensions in FP32 and BF16.""" + + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) + temb = _temb(3, dim, seed=hidden) + tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) + for out, gold in zip( + _cuda_op()(temb, weight, bias), + NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + strict=True, + ): + torch.testing.assert_close(out.float(), gold, **tol) + + def test_repeat_forward_backward_bitwise(self): + """Require repeated projection outputs and all input gradients to match bitwise.""" + + weight, bias = _synthetic(64, 2688, seed=9) + temb = _temb(3, seed=9) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + outs = _cuda_op()(*leaves) + torch.autograd.backward(list(outs), grads) + runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_temb_grad_rows_are_batch_invariant(self): + """Keep embedding gradients bitwise equal in full batches and single-row calls.""" + + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(4, seed=11) + grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] + + def d_temb(rows): + """Compute selected embedding-row gradients with their matching modality gradients.""" + + leaf = temb[rows].detach().clone().requires_grad_(True) + outs = _cuda_op()(leaf, weight, bias) + torch.autograd.backward( + list(outs), [g.view(4, 3, -1)[rows].reshape(-1, 64) for g in grads] + ) + return leaf.grad + + full = d_temb(slice(0, 4)) + for i in range(4): + assert torch.equal(d_temb(slice(i, i + 1))[0], full[i]) + + @pytest.mark.parametrize("num", [8, 9, 17]) + def test_rows_invariant_across_tensor_core_row_tiles(self, num): + """The BF16 path computes rows in tiles of 8; a row's bits must not depend on its tile.""" + + weight, bias = _synthetic(64, 2688, seed=13) + temb = _temb(num, seed=num) + full = _cuda_op()(temb, weight, bias) + for i in (0, 7, num - 1): + single = _cuda_op()(temb[i : i + 1], weight, bias) + for s_out, f_out in zip(single, full, strict=True): + assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + + def test_rejects_cpu(self): + """Reject CPU tensors at the CUDA projection boundary.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError): + _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) + + def test_registry_dispatches_cuda(self): + """Resolve the projection registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py new file mode 100644 index 000000000..0f9fec6c3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_row_gather``: ``timestep_index * 3 + token_tag`` lookup of six tensors. + +* forward is a copy: bitwise equal to diffusers' six ``index_select`` calls, + for every packing, length and index dtype; +* semantic indices fail closed (out-of-range tags or timesteps, RFC probes + H2/H3), and flipping a tag or offsetting a timestep picks a different row; +* backward is a deterministic FP32 segmented sum: repeat-bitwise and + correctly rounded, unlike the atomic BF16 ``index_select`` backward; +* batch invariance: a position's output does not depend on which other + positions share the call, and a row's gradient does not depend on the + positions that reference other rows. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +from rl_engine.reference.minimax_h3.adaln_row_gather import NativeH3AdaLNRowGatherOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_adaln_row_gather + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks h3_adaln_row_gather_*") + return H3AdaLNRowGatherCudaOp() + + +def _rows(num_timesteps, hidden=5376, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(3 * num_timesteps, 6 * hidden, generator=g).to(dtype).to(device) + + +class TestReference: + def test_matches_diffusers_index_select(self): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = h3_packed_layout(11, 2, device="cpu") + index = ti * 3 + tags + chunks = rows.chunk(6, dim=-1) + for ours, chunk in zip(NativeH3AdaLNRowGatherOp().forward(rows, ti, tags), chunks): + assert torch.equal(ours, chunk.index_select(0, index)) + + @pytest.mark.parametrize( + "mutate, error", + [ + (lambda ti, tags: (ti, tags.clone().fill_(3)), IndexError), # H2: no 4th modality + (lambda ti, tags: (ti, tags - 1), IndexError), + (lambda ti, tags: (ti + 2, tags), IndexError), # H3: timestep offset past T + (lambda ti, tags: (ti, tags[:-1]), ValueError), + (lambda ti, tags: (ti[:0], tags[:0]), ValueError), + (lambda ti, tags: (ti.float(), tags), TypeError), + (lambda ti, tags: (ti.int(), tags), TypeError), # mixed index dtypes + ], + ) + def test_rejects_invalid_indices(self, mutate, error): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = mutate(*h3_packed_layout(9, 2, device="cpu")) + with pytest.raises(error): + NativeH3AdaLNRowGatherOp().forward(rows, ti, tags) + + def test_rejects_bad_rows(self): + ti, tags = h3_packed_layout(4, 1, device="cpu") + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(4, 24), ti, tags) # not 3 per t + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) + @pytest.mark.parametrize("num_timesteps", [1, 3]) + def test_bitwise_equal_to_index_select(self, seq, num_timesteps): + rows = _rows(num_timesteps, seed=seq) + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + ours = _cuda_op()(rows, ti, tags) + ref = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + assert len(ours) == 6 + for a, b in zip(ours, ref): + assert a.shape == (seq, 5376) and a.is_contiguous() and torch.equal(a, b) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [24, 5, 5376]) + def test_dtypes_and_unaligned_hidden(self, dtype, hidden): + rows = _rows(2, hidden=hidden, dtype=dtype) + ti, tags = h3_packed_layout(33, 2) + for a, b in zip( + _cuda_op()(rows, ti.int(), tags.int()), NativeH3AdaLNRowGatherOp()(rows, ti, tags) + ): + assert torch.equal(a, b) + + def test_strided_rows_view(self): + table = _rows(4, hidden=16) + rows = table[::2] # row stride 2 * 6H, still 3 rows per timestep after slicing 12 -> 6 + ti, tags = h3_packed_layout(20, 2) + for a, b in zip(_cuda_op()(rows, ti, tags), NativeH3AdaLNRowGatherOp()(rows, ti, tags)): + assert torch.equal(a, b) + + def test_tag_flip_and_timestep_offset_select_other_rows(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(30, 2) + base = _cuda_op()(rows, ti, tags) + flipped = _cuda_op()(rows, ti, (tags + 1) % 3) # H2 + shifted = _cuda_op()(rows, 1 - ti, tags) # H3 within range + assert not torch.equal(base[0], flipped[0]) + assert not torch.equal(base[0], shifted[0]) + # Every output row is exactly the addressed table row, nothing else. + index = ti * 3 + tags + for c, out in enumerate(base): + assert torch.equal(out, rows[index, c * 16 : (c + 1) * 16]) + + def test_rejects_out_of_range_on_device(self): + rows = _rows(1, hidden=16) + ti, tags = h3_packed_layout(8, 1) + with pytest.raises(IndexError): + _cuda_op()(rows, ti, tags + 3) + + @pytest.mark.parametrize("entrypoint", ["native", "unchecked_wrapper"]) + @pytest.mark.parametrize( + "timestep, tag, index_dtype", + [ + pytest.param(-1, 0, "int32", id="negative-timestep"), + pytest.param(2, 0, "int64", id="timestep-past-table"), + pytest.param(0, 3, "int32", id="tag-past-modality-valid-row"), + pytest.param(1, -1, "int64", id="negative-tag-valid-row"), + # Multiplication by 3 wraps this value to row 2 in signed int64. + pytest.param((2**64 + 2) // 3, 0, "int64", id="timestep-overflow-valid-row"), + pytest.param(0, 2**63 - 1, "int64", id="int64-max-tag"), + ], + ) + def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): + """Reject semantic index violations without poisoning pytest's CUDA context.""" + + _cuda_op() + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + entrypoint, timestep, tag, index_dtype = sys.argv[1:] + rows = torch.zeros((6, 6), dtype=torch.float32, device="cuda") + dtype = getattr(torch, index_dtype) + ti = torch.tensor([int(timestep)], dtype=dtype, device="cuda") + tags = torch.tensor([int(tag)], dtype=dtype, device="cuda") + if entrypoint == "native": + _C.h3_adaln_row_gather_forward(rows, ti, tags, 6, 3) + else: + H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "native gather accepted invalid semantic indices" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + def test_gather_chunks_drop_in(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(10, 2) + chunks = rows.chunk(6, dim=-1) + for a, b in zip(_cuda_op().gather_chunks(chunks, ti, tags), _cuda_op()(rows, ti, tags)): + assert torch.equal(a, b) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + def _cuda_grad(self, rows, ti, tags, grads): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(_cuda_op()(leaf, ti, tags)), grads) + return leaf.grad + + def test_correctly_rounded_fp32_segment_sum(self): + rows = _rows(3, seed=1) + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = self._grads(4097, 5376) + ours = self._cuda_grad(rows, ti, tags, grads) + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward( + list(NativeH3AdaLNRowGatherOp().forward_fp32(ref, ti, tags)), [g.float() for g in grads] + ) + assert ours.dtype == torch.bfloat16 + # tolerance_contract.json gradient_accuracy / elementwise / bfloat16. + torch.testing.assert_close(ours.double(), ref.grad, atol=2e-2, rtol=1.6e-2) + assert (ours == ref.grad.bfloat16()).float().mean() > 0.9999 + + def test_repeat_bitwise(self): + rows = _rows(2, hidden=512, seed=2) + ti, tags = h3_packed_layout(5000, 2, seed=2) + grads = self._grads(5000, 512, seed=2) + first = self._cuda_grad(rows, ti, tags, grads) + for _ in range(3): + assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + + def test_forward_rows_invariant_to_batch_size_and_position(self): + rows = _rows(3, hidden=512, seed=5) + ti, tags = h3_packed_layout(4097, 3, seed=5) + full = _cuda_op()(rows, ti, tags) + perm = torch.randperm(4097, generator=torch.Generator().manual_seed(5)).cuda() + for pick in ( + torch.tensor([0]), + torch.tensor([2048]), + torch.arange(100, 357), + torch.arange(4000, 4097), + perm, + ): + pick = pick.cuda() + part = _cuda_op()(rows, ti[pick], tags[pick]) + for a, b in zip(part, full): + assert torch.equal(a, b[pick]) + + def test_backward_row_gradient_independent_of_other_rows_tokens(self): + # Each row's segment is tiled on its own, so a row's gradient from the + # full packing equals the one from a packing of only its own positions + # (same order), even though ~455 positions per row cross tile boundaries. + rows = _rows(3, hidden=512, seed=6) + ti, tags = h3_packed_layout(4097, 3, seed=6) + grads = self._grads(4097, 512, seed=6) + full = self._cuda_grad(rows, ti, tags, grads) + flat = ti * 3 + tags + for r in range(rows.shape[0]): + own = (flat == r).nonzero().squeeze(1) + assert own.numel() > 256 + alone = self._cuda_grad(rows, ti[own], tags[own], [g[own] for g in grads]) + assert torch.equal(alone[r], full[r]) + assert torch.count_nonzero(alone[torch.arange(rows.shape[0], device="cuda") != r]) == 0 + + def test_unreferenced_rows_get_zero_grad(self): + rows = _rows(3, hidden=16) + ti = torch.zeros(7, dtype=torch.long, device="cuda") + tags = torch.zeros(7, dtype=torch.long, device="cuda") # only row 0 is used + grad = self._cuda_grad(rows, ti, tags, self._grads(7, 16)) + assert torch.count_nonzero(grad[1:]) == 0 + assert torch.count_nonzero(grad[0]) > 0 + + def test_segments_longer_than_one_tile(self): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import BACKWARD_TILE + + seq = 3 * BACKWARD_TILE + 17 + rows = _rows(1, hidden=8, dtype=torch.float32) + ti = torch.zeros(seq, dtype=torch.long, device="cuda") + tags = torch.zeros(seq, dtype=torch.long, device="cuda") + grads = [torch.ones(seq, 8, device="cuda") for _ in range(6)] + grad = self._cuda_grad(rows, ti, tags, grads) + assert torch.equal(grad[0], torch.full((48,), float(seq), device="cuda")) + + def test_reference_backward_is_deterministic_and_accurate(self): + rows = _rows(2, hidden=256, seed=4) + ti, tags = h3_packed_layout(3000, 2, seed=4) + grads = self._grads(3000, 256, seed=4) + + def ref_grad(): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(NativeH3AdaLNRowGatherOp()(leaf, ti, tags)), grads) + return leaf.grad + + first = ref_grad() + assert torch.equal(ref_grad(), first) + torch.testing.assert_close( + first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 + ) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_row_gather", device="cuda") + assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" diff --git a/tests/models/minimax_h3/test_h3_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 new file mode 100644 index 000000000..bd3a44363 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU regressions for benchmark ordering and backward timing boundaries.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_report + + +@pytest.fixture +def cuda_clock(monkeypatch): + clock = SimpleNamespace(time=0.0, active=False, log=[], allocated=0, peak=0, events=0) + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + self.start = clock.events % 2 == 0 + clock.events += 1 + + def record(self): + clock.active = self.start + clock.log.append("start" if self.start else "end") + self.time = clock.time + + def synchronize(self): + assert not clock.active + + def elapsed_time(self, end): + return end.time - self.time + + def reset_peak(): + clock.log.append("reset_peak") + clock.peak = clock.allocated + + def allocated(): + clock.log.append("memory_baseline") + return clock.allocated + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: clock.log.append("synchronize")) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset_peak) + monkeypatch.setattr(torch.cuda, "memory_allocated", allocated) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: clock.peak) + return clock + + +def test_setup_is_outside_backward_events_and_memory_baseline(cuda_clock): + states = [] + + def setup(): + assert not cuda_clock.active + cuda_clock.log.append("setup") + cuda_clock.time += 50.0 + cuda_clock.allocated += 64 * 2**20 + return object() + + def backward(state): + states.append(state) + cuda_clock.log.append("backward") + cuda_clock.time += 2.0 + cuda_clock.allocated += 3 * 2**20 + cuda_clock.peak = cuda_clock.allocated + + assert h3_report.time_us(backward, warmup=1, iters=2, setup=setup) == 2000.0 + assert len({id(state) for state in states}) == 3 + assert cuda_clock.log.count("start") == 2 + for index, entry in enumerate(cuda_clock.log): + if entry == "start": + assert cuda_clock.log[index - 1] == "setup" + assert cuda_clock.log[index + 1 : index + 3] == ["backward", "end"] + + cuda_clock.log.clear() + assert h3_report.peak_mib(backward, setup=setup) == 3.0 + assert cuda_clock.log == [ + "setup", + "synchronize", + "reset_peak", + "memory_baseline", + "backward", + "synchronize", + ] + + +@pytest.mark.parametrize("keys", [("candidate", "provider"), h3_report.TIMED_KEYS]) +def test_measure_interleaves_and_records_both_orders(cuda_clock, keys): + calls = [] + case = {"op": "test", "case": "tiny", "backend": "cpu", "bytes": 4096} + for position, key in enumerate(keys, start=1): + + def run(key=key, duration=position): + calls.append(key) + cuda_clock.time += duration + + case[key] = run + + row = h3_report.measure(case, warmup=2, iters=4) + first, second = list(keys), list(reversed(keys)) + assert calls == (first + second) * 3 + first # warmup, samples, then memory calls + assert row["execution_order"] == { + "policy": "alternating", + "iteration_0": first, + "iteration_1": second, + } + for position, key in enumerate(keys, start=1): + assert row[f"{key}_us"] == position * 1000.0 + assert row[f"{key}_gbps"] == pytest.approx(4096 / (position * 1e-3) / 1e9) + assert row[f"{key}_peak_mib"] == 0.0 + assert not any(key.endswith("_setup") for key in row) + + +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) +def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( + monkeypatch, cuda_clock, operator +): + prepared_leaves = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def track_backward(outputs): + tensors = outputs if isinstance(outputs, tuple) else (outputs,) + + def backward_hook(grad): + cuda_clock.time += 2.0 / len(tensors) + return grad + + for tensor in tensors: + tensor.register_hook(backward_hook) + return outputs + + def track(fn, leaves, *args): + if leaves[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([leaf if leaf.is_leaf else leaf._base for leaf in leaves]) + cuda_clock.time += 1.0 + outputs = fn(*leaves, *args) + return track_backward(outputs) if leaves[0].requires_grad else outputs + + if operator == "norm": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 4), + torch.randn(3, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_norm_modulate + + def forward(x, weight, shift, scale, index, **kwargs): + return track(provider, [x, weight, shift, scale], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_norm_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) + op = SimpleNamespace(forward_modulated=forward) + build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g: provider(r, g, index, y_), [residual, y, gate]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(x, nw, temb, w, b, ti): + def compute(x, nw, temb, w, b): + rows = torch.nn.functional.linear(torch.nn.functional.silu(temb), w, b) + shift, scale = rows.chunk(2, dim=-1) + norm = torch.nn.functional.rms_norm(x, (4,), nw) + return norm * (1 + scale[ti]) + shift[ti] + + return track(compute, [x, nw, temb, w, b]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf + else: + randn = torch.randn + cpu_generators = {} + + def cpu_randn(*args, **kwargs): + kwargs["device"] = "cpu" + # A seeded CUDA generator cannot draw a CPU tensor: continue the same seed on + # one CPU generator per CUDA generator. + gen = kwargs.get("generator") + if gen is not None and gen.device.type != "cpu": + if id(gen) not in cpu_generators: + cpu_generators[id(gen)] = torch.Generator().manual_seed(gen.initial_seed()) + kwargs["generator"] = cpu_generators[id(gen)] + return randn(*args, **kwargs) + + provider = h3_report.provider_adaln_row_gather + + def forward(rows, ti, tags, **kwargs): + def gather(leaf, ti, tags): + return provider(leaf.chunk(6, dim=-1), ti, tags) + + return track(gather, [rows], ti, tags) + + def provider_forward(chunks, ti, tags): + if chunks[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([chunks[0]._base]) + cuda_clock.time += 1.0 + outputs = provider(chunks, ti, tags) + return track_backward(outputs) if chunks[0].requires_grad else outputs + + monkeypatch.setattr(torch, "randn", cpu_randn) + monkeypatch.setattr(h3_report, "GATHER_SEQ_LENS", (2, 3)) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: (torch.arange(seq) % num, torch.arange(seq) % 3), + ) + monkeypatch.setattr(h3_report, "provider_adaln_row_gather", provider_forward) + op = SimpleNamespace(forward=forward) + build = h3_report._gather_perf + + try: + registry = SimpleNamespace(get_op=lambda *args, **kwargs: op) + for case in build(registry): + row = h3_report.measure(case, warmup=1, iters=2) + assert row["backward_timing_scope"] == "backward_only" + assert row["candidate_backward_us"] == pytest.approx(2000.0) + assert row["provider_backward_us"] == pytest.approx(2000.0) + finally: + torch.set_num_threads(previous_threads) + + assert prepared_leaves + assert len({id(leaves[0]) for leaves in prepared_leaves}) == len(prepared_leaves) + assert all(leaf.grad is not None for leaves in prepared_leaves for leaf in leaves) diff --git a/tests/models/minimax_h3/test_h3_chain_replay_cli.py b/tests/models/minimax_h3/test_h3_chain_replay_cli.py new file mode 100644 index 000000000..b9baea84d --- /dev/null +++ b/tests/models/minimax_h3/test_h3_chain_replay_cli.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Reject incomplete chain dependencies before checking CUDA or loading weights.""" + +from __future__ import annotations + +import sys + +import pytest + +from tools.validation.models import h3_chain_replay + +STAGE_NAMES = [stage.name for stage in h3_chain_replay.h3_chain.STAGES] + + +@pytest.fixture +def no_cuda_check(monkeypatch): + """Fail if invalid CLI arguments reach a CUDA availability check.""" + + def unexpected_check(): + """Report an environment check performed before argument rejection.""" + + pytest.fail("invalid arguments must be rejected before checking CUDA") + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", unexpected_check) + + +@pytest.mark.parametrize("stages, message", [("", "unknown stages"), ("unknown", "unknown stages")]) +def test_reject_invalid_stages(monkeypatch, capsys, no_cuda_check, stages, message): + """Reject unknown stages before any device check.""" + + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES))) +def test_backward_requires_all_stages(monkeypatch, capsys, no_cuda_check, count): + """Require the full chain for parameter-gradient replay.""" + + monkeypatch.setattr( + sys, + "argv", + ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count]), "--backward"], + ) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert "--backward requires every stage" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "stages", + [ + *(",".join(STAGE_NAMES[:count]) for count in range(1, len(STAGE_NAMES) + 1)), + STAGE_NAMES[1], + STAGE_NAMES[-1], + ",".join(STAGE_NAMES[:1] + STAGE_NAMES[2:]), + ",".join(reversed(STAGE_NAMES)), + ",".join([STAGE_NAMES[0], STAGE_NAMES[0]]), + ], +) +def test_accept_any_known_selection(monkeypatch, stages): + """Any selection of known stages reaches the device requirement; the replay + runs its prerequisites (tests/models/minimax_h3/test_h3_cli.py).""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() + + +@pytest.mark.parametrize("options", [[], ["--backward"]]) +def test_accept_all_stages_by_default(monkeypatch, options): + """Keep the complete chain as the default for forward and backward replay.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", *options]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() diff --git a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py new file mode 100644 index 000000000..5686f1b10 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 integration rows: checkpoint mapper and the vLLM-Omni rollout adapter.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.models.minimax_h3 import checkpoint_mapper as cm +from rl_engine.models.minimax_h3 import fingerprint as fp + +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +# --- checkpoint mapper ---------------------------------------------------------------------- + + +def _names(): + return list(fp.expected_tensor_table(fp.transformer_config())) + + +def test_every_checkpoint_tensor_maps_and_native_names_are_unique(): + names = _names() + assert len(names) == 638 + state = {n: torch.zeros(6, 2) for n in names} + native = cm.diffusers_to_native(state) + # q/k/v of each of the 52 blocks (50 + 2 refiner) fuse per kind (weight): 3 -> 1 + assert len(native) == 638 - 2 * 52 + assert len(set(native)) == len(native) + + +def test_diffusers_native_roundtrip_is_bitwise(): + names = _names() + gen = torch.Generator().manual_seed(0) + state = {n: torch.randn(6, 2, generator=gen) for n in names} + back = cm.native_to_diffusers(cm.diffusers_to_native(state)) + assert set(back) == set(state) + assert all(torch.equal(back[n], state[n]) for n in state) + + +def test_fc1_halves_are_swapped_and_qkv_is_q_then_k_then_v(): + w = torch.arange(8.0).view(8, 1) + state = { + "transformer_blocks.3.ff.net.0.proj.weight": w, + "transformer_blocks.3.attn.to_q.weight": torch.full((2, 1), 1.0), + "transformer_blocks.3.attn.to_k.weight": torch.full((2, 1), 2.0), + "transformer_blocks.3.attn.to_v.weight": torch.full((2, 1), 3.0), + } + native = cm.diffusers_to_native(state) + assert torch.equal(native["blocks.3.mlp.fc1.weight"], torch.cat([w[4:], w[:4]])) + assert native["blocks.3.attn.qkv_proj.weight"].flatten().tolist() == [1, 1, 2, 2, 3, 3] + + +def test_unmapped_names_fail_closed(): + with pytest.raises(ValueError, match="unmapped"): + cm.diffusers_to_native({"transformer_blocks.0.surprise.weight": torch.zeros(1)}) + with pytest.raises(ValueError, match="unmapped"): + cm.native_to_diffusers({"blocks.0.surprise.weight": torch.zeros(1)}) + with pytest.raises(ValueError, match="q, k and v"): + cm.diffusers_to_native({"transformer_blocks.0.attn.to_q.weight": torch.zeros(1, 1)}) + + +def test_grouped_checkpoint_layout_roundtrip(): + w = torch.randn(4 * 3 * 8, 5) + fused = cm.grouped_qkv_to_qkv(w, heads=4, head_dim=8) + assert torch.equal(cm.qkv_to_grouped(fused, heads=4, head_dim=8), w) + + +def test_tables_match_vllm_omni_when_importable(): + fasth3 = pytest.importorskip("vllm_omni.diffusion.models.minimax_h3.fasth3") + assert dict(fasth3._MODEL_LEVEL_TARGETS) == cm.MODEL_LEVEL + assert {k: v[0] for k, v in fasth3._BLOCK_TARGETS.items() if k in cm.BLOCK_TARGETS} == { + k: v[0] for k, v in cm.BLOCK_TARGETS.items() + } + + +class _Lin(torch.nn.Module): + def __init__(self, n, k): + super().__init__() + self.weight = torch.nn.Parameter( + (torch.randn(n, k) / 30).bfloat16().cuda(), requires_grad=False + ) + + def forward(self, x): + raise AssertionError("must be replaced") + + +@needs_gpu +def test_vllm_adapter_patches_projections_and_restores(): + pytest.importorskip("triton") + from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + torch.manual_seed(0) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList([torch.nn.Module()]) + blk = model.blocks[0] + blk.attn = torch.nn.Module() + blk.attn.qkv_proj, blk.attn.out_proj = _Lin(3 * 1024, 640), _Lin(640, 1024) + blk.mlp = torch.nn.Module() + blk.mlp.fc1, blk.mlp.fc2 = _Lin(2 * 2048, 640), _Lin(640, 2048) + x = torch.randn(2, 17, 640).bfloat16().cuda() + with install_h3_rollout(model) as handle: + qkv, bias = blk.attn.qkv_proj(x) + assert bias is None and qkv.shape == (2, 17, 3 * 1024) + q, k, v = (blk.attn.qkv_proj.weight[i * 1024 : (i + 1) * 1024] for i in range(3)) + sep = [par.g.gemm(x.reshape(-1, 640), w, b_t=True) for w in (q, k, v)] + assert torch.equal(qkv.reshape(-1, 3072), torch.cat(sep, dim=1)) # fused == separate + y, _ = blk.attn.out_proj(torch.randn(2, 17, 1024).bfloat16().cuda()) + assert y.shape == (2, 17, 640) + assert sorted(handle.patched) == [ + "blocks.0.attn.out_proj", + "blocks.0.attn.qkv_proj", + "blocks.0.mlp.fc1", + "blocks.0.mlp.fc2", + ] + with pytest.raises(AssertionError, match="must be replaced"): + blk.attn.qkv_proj(x) + + +@needs_gpu +def test_vllm_adapter_strict_rejects_wrong_dtype_and_empty_models(): + from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout + + with pytest.raises(RuntimeError, match="no H3 projection"): + install_h3_rollout(torch.nn.Linear(2, 2).cuda()) + holder = torch.nn.Module() + holder.qkv_proj = torch.nn.Linear(8, 8).cuda() # float32 + with pytest.raises(RuntimeError, match="BF16"): + install_h3_rollout(holder) diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py new file mode 100644 index 000000000..8e97fc6a2 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 CLIs preserve chain dependencies and require pinned report weights.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_chain +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV +from tools.validation.models import h3_chain_replay, h3_evidence + + +@pytest.fixture +def cpu_chain_cli(monkeypatch): + """Replace CUDA chain stages with CPU stubs that track dependencies and provider drift.""" + + executed = [] + state = {"provider_shift": 0} + + class CpuOp: + def __call__(self, value): + """Advance the synthetic chain value by one for candidate and golden replay.""" + + return value + 1 + + forward_fp32 = __call__ + + stages = [] + for index, stage in enumerate(h3_chain.STAGES): + + def stage_input(ctx, upstream, index=index): + """Use raw timesteps for stage zero and require predecessor outputs thereafter.""" + + if index == 0: + return ctx["timestep"] + assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" + return upstream + + stages.append( + replace( + stage, + candidate=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + golden=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + provider=lambda ctx, up, stage_input=stage_input, index=index: ( + stage_input(ctx, up) + 1 + (state["provider_shift"] if index == 0 else 0) + ), + ) + ) + + def get_op(op_type, *, device): + """Record each dispatched stage and return its deterministic CPU stub.""" + + executed.append(op_type) + return CpuOp() + + monkeypatch.setattr(h3_chain, "STAGES", stages) + monkeypatch.setattr(h3_chain, "make_context", lambda *a, **kw: {"timestep": torch.zeros(1)}) + monkeypatch.setattr(h3_chain, "golden_op", lambda *a: CpuOp()) + monkeypatch.setattr(h3_chain_replay, "KernelRegistry", lambda: SimpleNamespace(get_op=get_op)) + monkeypatch.setattr(h3_chain_replay, "load_h3_conditioning_weights", lambda *a: {}) + monkeypatch.setattr(h3_chain, "environment", lambda: {}) + monkeypatch.setattr(h3_chain, "git_state", lambda: {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + return executed, state + + +@pytest.mark.parametrize("provider_shift", [0, 1]) +@pytest.mark.parametrize( + "requested, execution_count", + [ + ("timestep_mlp_fp32", 2), + ("adaln_projection_3mod", 3), + ("timestep_sinusoid_h3,adaln_projection_3mod", 3), + ("timestep_mlp_fp32,timestep_sinusoid_h3", 2), + ], +) +def test_chain_selection_runs_dependencies( + cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift +): + """Execute prerequisite stages but report requested stages and preserve first drift.""" + + executed, state = cpu_chain_cli + state["provider_shift"] = provider_shift + out = tmp_path / "chain.json" + monkeypatch.setattr( + sys, + "argv", + [ + "h3_chain_replay.py", + "--stages", + requested, + "--timesteps", + "1", + "--seq-lens", + "3", + "--out", + str(out), + ], + ) + h3_chain_replay.main() + + report = json.loads(out.read_text()) + stage_names = [stage.name for stage in h3_chain.STAGES] + expected_report = [name for name in stage_names if name in requested.split(",")] + assert executed == stage_names[:execution_count] + assert report["executed_stages"] == executed + assert report["stages"] == expected_report + case = report["cases"][0] + assert [entry["stage"] for entry in case["stages"]] == expected_report + assert all(entry["chained_vs_golden"]["bitwise_equal"] for entry in case["stages"]) + assert case["first_drift"] == (stage_names[0] if provider_shift else None) + assert case["first_isolated_drift"] == (stage_names[0] if provider_shift else None) + summary = capsys.readouterr().out.split("; first_drift=", 1)[0] + for name in stage_names: + assert (f"{name}=" in summary) == (name in expected_report) + + +@pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) +@pytest.mark.parametrize("weights_env", [None, " "]) +def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + + if weights_env is None: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, weights_env) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + with pytest.raises(SystemExit, match=WEIGHTS_ENV): + h3_evidence.main() + assert not out.exists() + + +@pytest.mark.parametrize( + "op, weight_source", + [ + ("timestep_sinusoid_h3", "not_applicable"), + ("timestep_mlp_fp32", "pinned_checkpoint"), + ("adaln_projection_3mod", "pinned_checkpoint"), + ], +) +def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + + if weight_source == "not_applicable": + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + monkeypatch.setattr(h3_evidence, "KernelRegistry", lambda: object()) + monkeypatch.setattr( + h3_evidence, + "git_state", + lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, + ) + monkeypatch.setattr(h3_evidence, "environment", lambda: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py new file mode 100644 index 000000000..07cd49776 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End to end: the H3 conditioning chain on the pinned checkpoint, through the registry. + +Runs every stage registered in ``rl_engine.validation.models.h3_chain.STAGES`` (each +RFC #420 row adds its own), chained exactly as the model calls them, and +checks every stage's promises: + +* the registry dispatched the CUDA backend; +* a repeated run is bitwise equal; +* the chained output is within the stage's contract tolerance of the golden; +* stages that promise it are bitwise equal to diffusers on diffusers' inputs, + so the first isolated drift can only be a stage that declares a reduction. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_chain + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. +CASES = [(1, 3), (3, 257), (4, 4097)] + + +@pytest.fixture(scope="module") +def chain_weights(h3_weights_cpu): + """Move all pinned conditioning tensors to CUDA once for the end-to-end module.""" + + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} + + +@pytest.fixture(scope="module") +def registry(): + """Provide the production registry used to dispatch every conditioning stage.""" + + return KernelRegistry() + + +@pytest.mark.parametrize("num_timesteps, seq_len", CASES) +def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Check CUDA dispatch, golden tolerance, repeatability, and declared provider parity.""" + + report = h3_chain.run_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + stages = {stage.name: stage for stage in h3_chain.STAGES} + for entry in report["stages"]: + stage = stages[entry["stage"]] + assert entry["backend"].endswith("CudaOp"), entry + assert entry["repeat_bitwise_equal"], entry["stage"] + assert entry["chained_vs_golden"]["within_tolerance"], entry + if stage.provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] + drift = report["first_isolated_drift"] + assert drift is None or not stages[drift].provider_bitwise_isolated + + +@pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) +def test_chain_backward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" + + report = h3_chain.run_backward_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + for name, entry in report["leaves"].items(): + for mode in ("candidate", "candidate_fused"): + assert entry[mode]["repeat_bitwise_equal"], (mode, name) + fused = entry["candidate_fused"] + if name.startswith("time_embedder"): + # FP32 parameters: no BF16 rounding anywhere on the fused path. + assert fused["max_abs_vs_golden_over_absmax"] < 1e-5, (name, fused) + else: + # BF16 parameters: only their own final rounding remains. + assert fused["correctly_rounded_fraction"] > 0.99, (name, fused) diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py new file mode 100644 index 000000000..ca11c4e53 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Native input validation for the H3 deterministic linear kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available +from rl_engine.backends.extension import _C + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), + pytest.mark.skipif(not det_linear_available(), reason="rl_engine._C lacks h3_det_linear_*"), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_input_rejects_zero_rows(dtype): + """Reject empty gradient batches in native input-gradient kernels for both dtypes.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + weight = torch.empty(16, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_input(grad, weight, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_rejects_zero_rows(dtype, with_bias): + """Reject empty gradient batches before parameter-gradient or bias-gradient reduction.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + x = torch.empty(0, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +def test_backward_input_handles_empty_columns(dtype, out_dtype, k_in, n_out): + """Match the input VJP of empty linear dimensions in the requested output dtype.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + weight = torch.empty(n_out, k_in, device="cuda", dtype=dtype) + expected = (grad @ weight.float()).to(out_dtype) + actual = _C.h3_det_linear_backward_input(grad, weight, out_dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with_bias): + """Return empty weight gradients while preserving any nonempty bias reduction.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + x = torch.arange(3 * k_in, device="cuda", dtype=dtype).reshape(3, k_in) + expected_weight = (grad.T @ x.float()).to(w_dtype) + actual = _C.h3_det_linear_backward_weight(grad, x, w_dtype, with_bias) + assert len(actual) == (2 if with_bias else 1) + torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) + if with_bias: + torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..5fb964677 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``final_adaln_out``: norm_out shift/scale projection + final norm/modulation. + +* layout: ``shift`` is the first half of ``norm_out.linear``'s output, rows + are indexed by ``timestep_indices`` (not ``adaln_indices``); +* forward matches diffusers' ``norm_out`` to the projection's 1-ULP ties + (the norm and modulation are bitwise; the GEMV tree differs from cuBLAS); +* backward is deterministic, with the table gradient kept in FP32; +* rows are batch/position invariant; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.final_adaln_out import NativeH3FinalAdaLNOutOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_final_adaln_out + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3FinalAdaLNOutCudaOp() + + +@pytest.fixture(scope="module") +def norm_out(h3_weights_cpu): + return [ + h3_weights_cpu[name].cuda() + for name in ("norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias") + ] + + +def _inputs(seq, num_timesteps=3, seed=0, batch=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(batch, seq, 5376, generator=g) * 2).bfloat16().cuda() + temb = (torch.randn(num_timesteps, 2688, generator=g) * 2).cuda() + ti, _ = h3_packed_layout(seq, num_timesteps, seed=seed) + return x, temb, ti + + +def _golden(x, norm_weight, temb, weight, bias, timestep_indices): + act = temb * torch.sigmoid(temb) + # Declared cast of the FP32 SiLU, as the model computes it; identity VJP. + declared = F.silu(temb.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = F.linear(act, weight, bias) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) * norm_weight + return n * (1 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3FinalAdaLNOutOp() + x, nw = torch.zeros(1, 4, 8), torch.ones(8) + temb, w, b = torch.zeros(2, 16), torch.zeros(16, 16), torch.zeros(16) + ti = torch.tensor([0, 1, 1, 0]) + with pytest.raises(TypeError): # BF16 temb (probe H7) + op(x, nw, temb.bfloat16(), w, b, ti) + with pytest.raises(ValueError): # not 2H rows + op(x, nw, temb, w[:10], b[:10], ti) + with pytest.raises(IndexError): + op(x, nw, temb, w, b, ti + 2) + + +@requires_cuda +class TestCuda: + def test_shift_first_layout_and_timestep_indexing(self): + """Zero weights: the bias halves are (shift, scale); rows follow timestep_indices.""" + + hidden, seq = 16, 6 + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + nw = torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(2 * hidden, 32, device="cuda", dtype=torch.bfloat16) + b = torch.cat([torch.full((hidden,), 1.0), torch.full((hidden,), -1.0)]).bfloat16().cuda() + ti = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda") + out = _cuda_op()(x, nw, torch.zeros(2, 32, device="cuda"), w, b, ti) + # scale = -1 -> norm(x) * 0 + shift = 1 everywhere + assert torch.equal(out, torch.ones_like(out)) + + @pytest.mark.parametrize("seq, num_timesteps", [(1, 1), (257, 2), (4097, 3)]) + def test_forward_against_diffusers_and_golden(self, norm_out, seq, num_timesteps): + x, temb, ti = _inputs(seq, num_timesteps, seed=seq) + ours = _cuda_op()(x, *norm_out[:1], temb, *norm_out[1:], ti) + theirs = provider_final_adaln_out(x, norm_out[0], temb, *norm_out[1:], ti) + golden = NativeH3FinalAdaLNOutOp().forward_fp32(x, norm_out[0], temb, *norm_out[1:], ti) + assert (ours == theirs).float().mean() > 0.999 + torch.testing.assert_close(ours.float(), golden, atol=5e-2, rtol=2e-2) + + def test_rows_are_batch_and_position_invariant(self, norm_out): + x, temb, ti = _inputs(600, seed=1, batch=2) + full = _cuda_op()(x, norm_out[0], temb, *norm_out[1:], ti) + part = _cuda_op()(x[1:2, 200:260], norm_out[0], temb, *norm_out[1:], ti[200:260]) + assert torch.equal(part[0], full[1, 200:260]) + + def test_backward_deterministic_and_fp32_table_gradient(self, norm_out): + x, temb, ti = _inputs(4097, seed=2) + grad = torch.randn_like(x) + + def grads(fn, dtype=None): + tensors = [x, norm_out[0], temb, *norm_out[1:]] + leaves = [ + (t if dtype is None else t.to(dtype)).detach().clone().requires_grad_(True) + for t in tensors + ] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda *t: _cuda_op()(*t, ti)) + again = grads(lambda *t: _cuda_op()(*t, ti)) + ref = grads(lambda *t: _golden(*t, ti), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): + rel = ((g.double() - r).abs().max() / r.abs().max()).item() + assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + + def test_report_backward_uses_fp64_golden(self, monkeypatch): + from types import SimpleNamespace + + from rl_engine.validation.models import h3_report + + rng = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn(1, 160, 8, device="cuda", generator=rng).bfloat16() + nw = torch.randn(8, device="cuda", generator=rng).bfloat16() + temb = torch.randn(3, 4, device="cuda", generator=rng) + w = torch.randn(16, 4, device="cuda", generator=rng).bfloat16() + b = torch.randn(16, device="cuda", generator=rng).bfloat16() + ti = torch.arange(160, device="cuda") % 3 + inputs = (x, nw, temb, w, b) + native = NativeH3FinalAdaLNOutOp() + registry = SimpleNamespace( + get_op=lambda *args, **kwargs: native, + _get_or_create_backend=lambda _: native, + _priority_map={"cpu": {"final_adaln_out": [None]}}, + ) + monkeypatch.setattr(h3_report, "_final_inputs", lambda *args, **kwargs: (*inputs, ti)) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", native) + report = h3_report._final_accuracy(registry) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + leaves = [t.detach().clone().requires_grad_(True) for t in inputs] + native(*leaves, ti).backward(grad) + ref = [t.double().detach().requires_grad_(True) for t in inputs] + _golden(*ref, ti).backward(grad.double()) + for name, leaf, reference in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), leaves, ref): + assert reference.grad.dtype == torch.float64 + expected = float( + (leaf.grad.double() - reference.grad).abs().max() / reference.grad.abs().max() + ) + for backend in ("cuda", "provider"): + assert report["backward"][backend]["rel_error"][name] == pytest.approx(expected) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("final_adaln_out", device="cuda") + assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_bf16_gradients_meet_gtest_contract(self, seq): + # The check_operator path; S=257 used to miss d_norm_weight before the + # golden stored norm(x) and 1 + scale in BF16 like the model does. + import argparse + + from rl_engine.validation.operators import run_operator_suite + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + _cuda_op() + args = argparse.Namespace( + op="final_adaln_out", + candidate="cuda", + arch_key=None, + batch=3, + seq=seq, + normalized_dim=5376, + seed=123, + input_mode="random", + ) + report = run_operator_suite( + "final_adaln_out", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.bfloat16, torch.device("cuda"))], + check_grad=True, + ) + assert report.passed, report.candidates[0].cases[0].outputs diff --git a/tests/models/minimax_h3/test_h3_full_model.py b/tests/models/minimax_h3/test_h3_full_model.py new file mode 100644 index 000000000..ccab371a7 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_full_model.py @@ -0,0 +1,208 @@ +# 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 + ) + + +def test_strict_boundaries_apply_to_candidate_execution_only(): + from rl_engine.validation.models.h3_full_model import PROVIDER_BOUNDARIES, STRICT_BOUNDARY_ROWS + + with pytest.raises(ValueError, match="candidate execution only"): + with full_model_execution( + Model(), candidate=False, trace=FullModelTrace(), strict_boundaries=True + ): + pass + # every boundary module is either bound to a row or explicitly left on the provider + assert set(STRICT_BOUNDARY_ROWS) | set(PROVIDER_BOUNDARIES) == { + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + } diff --git a/tests/models/minimax_h3/test_h3_generation_resume.py b/tests/models/minimax_h3/test_h3_generation_resume.py new file mode 100644 index 000000000..5c0d3d56b --- /dev/null +++ b/tests/models/minimax_h3/test_h3_generation_resume.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Resumable candidate generation: a preempted run continues to the same result. + +A fake transformer and denoising loop stand in for the pipeline, so these run on CPU. +The loop feeds each call's output back into the next call's inputs (as the scheduler +does), so any wrong replay would change every later step. +""" + +from __future__ import annotations + +import inspect +import subprocess +import types + +import pytest +import torch + +from rl_engine.validation.models import h3_generation_resume as resume_module +from rl_engine.validation.models.h3_generation_resume import ( + ResumableRun, + ResumeError, + StepStore, + adopt_identity, + inputs_digest, + model_trees, +) + +STEPS = 6 +IDENTITY = {"commit": "abc", "seed": 42, "candidate": True, "preflight": False} + + +class FakeTransformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.computed = 0 + + def forward(self, hidden_states, audio_hidden_states, *, return_dict=True): + self.computed += 1 + video = torch.tanh(hidden_states) * 0.5 + audio = torch.sin(audio_hidden_states) * 0.25 + return _output((video, audio), {"return_dict": return_dict}) + + +def _output(values, kwargs): + if kwargs.get("return_dict", True): + return types.SimpleNamespace(sample=values[0], audio_sample=values[1]) + return values + + +def _initial(seed=0): + generator = torch.Generator().manual_seed(seed) + return torch.randn(8, 4, generator=generator), torch.randn(3, 5, generator=generator) + + +def denoise( + directory, *, resume, stop_after=None, initial=None, identity=IDENTITY, verify_last_step=False +): + """Run the fake loop through a ResumableRun; return (latents, transformer, run).""" + + video, audio = initial if initial is not None else _initial() + store = StepStore(directory) + run = ResumableRun(store, identity, resume=resume, verify_last_step=verify_last_step) + if not (directory / "initial.pt").exists(): + torch.save({"hidden_states": video, "audio_hidden_states": audio}, directory / "initial.pt") + module = FakeTransformer() + run.install(module, _output) + + def hook(mod, args, kwargs, output): + values = (output.sample, output.audio_sample) + run.record(args, kwargs, values) + if stop_after is not None and run.recorded >= stop_after: + raise KeyboardInterrupt("simulated preemption") + + module.register_forward_hook(hook, with_kwargs=True) + for _ in range(STEPS): + out = module(hidden_states=video, audio_hidden_states=audio) + # The scheduler's update: the next inputs depend on this call's outputs. + video = video + out.sample + audio = audio - 0.5 * out.audio_sample + return (video, audio), module, run + + +def _step_bytes(directory): + return { + p.name: p.read_bytes() + for p in sorted(directory.iterdir()) + if p.name.startswith("step-") and p.suffix in (".pt", ".json") + } + + +def test_uninterrupted_run_saves_every_step(tmp_path): + (video, _), module, run = denoise(tmp_path / "a", resume=False) + assert module.computed == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.pt")]) == STEPS + assert len([p for p in (tmp_path / "a").glob("step-*.in.json")]) == STEPS + assert run.summary()["replayed_steps"] == 0 + + +def test_resumed_run_with_verification_matches_an_uninterrupted_one(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True, verify_last_step=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert _step_bytes(tmp_path / "whole").keys() == _step_bytes(tmp_path / "cut").keys() + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + # Four steps were saved: three replayed, the fourth recomputed and verified, two new. + summary = run.summary() + assert summary["saved_steps_at_start"] == 4 + assert summary["replayed_steps"] == 3 + assert summary["recomputed_and_verified_step"] == 3 + assert module.computed == STEPS - 3 + assert [a["saved_steps"] for a in summary["attempts"]] == [0, 4] + + +def test_repeated_preemptions_still_converge(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + for stop in (2, 3, 5): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=stop > 2, stop_after=stop) + resumed, _, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + assert len(run.summary()["attempts"]) == 4 + + +def test_with_verification_a_complete_run_recomputes_only_its_last_step(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True, verify_last_step=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 1 + assert run.summary()["recomputed_and_verified_step"] == STEPS - 1 + + +def test_different_inputs_are_rejected_during_replay(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + # Rejected at the first replayed call, not only at the recomputed last saved step. + with pytest.raises(ResumeError, match="replayed step 0 has different inputs"): + denoise(tmp_path / "a", resume=True, initial=_initial(seed=1)) + + +def test_a_nondeterministic_recompute_is_rejected_when_verifying(tmp_path, monkeypatch): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + + original = FakeTransformer.forward + + def noisy(self, hidden_states, audio_hidden_states, *, return_dict=True): + out = original(self, hidden_states, audio_hidden_states, return_dict=return_dict) + return _output((out.sample + 1e-6, out.audio_sample), {"return_dict": return_dict}) + + monkeypatch.setattr(FakeTransformer, "forward", noisy) + with pytest.raises(ResumeError, match="differs from the saved output"): + denoise(tmp_path / "a", resume=True, verify_last_step=True) + # Without verification nothing recomputes the saved steps, so the drift goes unnoticed. + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "b", resume=False, stop_after=3) + denoise(tmp_path / "b", resume=True) + + +def test_a_different_identity_cannot_be_continued(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + with pytest.raises(ResumeError, match="different identity"): + denoise(tmp_path / "a", resume=True, identity={**IDENTITY, "commit": "other"}) + + +def test_an_existing_directory_needs_resume(tmp_path): + denoise(tmp_path / "a", resume=False, stop_after=None) + with pytest.raises(ResumeError, match="resume"): + denoise(tmp_path / "a", resume=False) + + +def test_a_corrupted_step_ends_the_valid_prefix(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=4) + (tmp_path / "a" / "step-002.pt").write_bytes(b"damaged") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 2 + + +def test_leftover_temporary_files_are_ignored(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "step-003.pt.tmp").write_bytes(b"half written") + run = ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True) + assert run.saved == 3 + + +def test_a_run_without_initial_inputs_is_not_resumable(tmp_path): + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "a", resume=False, stop_after=3) + (tmp_path / "a" / "initial.pt").unlink() + assert ResumableRun(StepStore(tmp_path / "a"), IDENTITY, resume=True).saved == 0 + + +def test_inputs_digest_tracks_value_dtype_and_shape_and_ignores_non_tensors(): + x = torch.arange(6.0).reshape(2, 3) + base = inputs_digest((), {"a": x, "flag": True}) + assert base == inputs_digest((), {"a": x.clone(), "flag": False}) + assert base != inputs_digest((), {"a": x + 1}) + assert base != inputs_digest((), {"a": x.double()}) + assert base != inputs_digest((), {"a": x.reshape(3, 2)}) + assert base != inputs_digest((), {"b": x}) + assert inputs_digest((x,), {}) != inputs_digest((), {"a": x}) + assert inputs_digest((), {"s": torch.tensor(3.0)}) # zero-dim tensors hash + + +def test_replay_wrapper_keeps_the_forward_signature(tmp_path): + """The pipeline selects the transformer's layout inputs from inspect.signature(forward).""" + + module = FakeTransformer() + before = list(inspect.signature(module.forward).parameters) + (tmp_path / "a").mkdir() + run = ResumableRun(StepStore(tmp_path / "a" / "x"), IDENTITY, resume=False) + uninstall = run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + uninstall() + assert list(inspect.signature(module.forward).parameters) == before + + +def test_replay_wrapper_keeps_a_wrapped_forward_signature(tmp_path): + """Same through an accelerate-style partial that itself carries __wrapped__.""" + + import functools + + module = FakeTransformer() + plain = module.forward + + def new_forward(m, *args, **kwargs): + return plain(*args, **kwargs) + + hooked = functools.update_wrapper(functools.partial(new_forward, module), plain) + module.forward = hooked + before = list(inspect.signature(module.forward).parameters) + assert before == ["hidden_states", "audio_hidden_states", "return_dict"] + run = ResumableRun(StepStore(tmp_path / "b"), IDENTITY, resume=False) + run.install(module, _output) + assert list(inspect.signature(module.forward).parameters) == before + + +# --- default resume: replay everything saved, recompute nothing -------------------------- + + +def test_default_resume_replays_every_saved_step_and_computes_only_the_rest(tmp_path): + whole, _, _ = denoise(tmp_path / "whole", resume=False) + with pytest.raises(KeyboardInterrupt): + denoise(tmp_path / "cut", resume=False, stop_after=4) + resumed, module, run = denoise(tmp_path / "cut", resume=True) + assert all(torch.equal(a, b) for a, b in zip(whole, resumed, strict=True)) + summary = run.summary() + assert (summary["saved_steps_at_start"], summary["replayed_steps"]) == (4, 4) + assert summary["recomputed_and_verified_step"] is None + assert summary["real_forward_calls"] == STEPS - 4 == module.computed + for name, data in _step_bytes(tmp_path / "whole").items(): + if name.endswith(".pt"): + assert data == _step_bytes(tmp_path / "cut")[name], name + + +def test_a_complete_run_resumes_with_no_forward_pass_at_all(tmp_path): + whole, _, _ = denoise(tmp_path / "a", resume=False) + again, module, run = denoise(tmp_path / "a", resume=True) + assert all(torch.equal(x, y) for x, y in zip(whole, again, strict=True)) + assert module.computed == 0 + assert run.real_calls == 0 + + +def test_each_attempt_only_needs_a_few_steps_to_make_progress(tmp_path): + """Attempts that each finish one new step still add up to the whole run.""" + + whole, _, _ = denoise(tmp_path / "whole", resume=False) + saved = 0 + result = run = None + for attempt in range(STEPS): + stop = saved + 1 if saved + 1 < STEPS else None # the last attempt is left to finish + try: + result, _, run = denoise(tmp_path / "cut", resume=attempt > 0, stop_after=stop) + except KeyboardInterrupt: + saved += 1 + continue + break + assert result is not None and saved == STEPS - 1 + assert all(torch.equal(a, b) for a, b in zip(whole, result, strict=True)) + assert run.real_calls == 1 and len(run.summary()["attempts"]) == STEPS + + +def test_the_validated_trace_survives_for_a_later_replay_only_attempt(tmp_path): + store = StepStore(tmp_path / "a") + assert store.load_trace() is None + (tmp_path / "a").mkdir() + trace = {"calls": {"block.0": 3}, "execution_order": ["block.0", "block.1"]} + store.save_trace(trace) + assert StepStore(tmp_path / "a").load_trace() == trace + assert not list((tmp_path / "a").glob("*.tmp")) + + +# --- identity tied to the model paths, and adopting a commit-based run ------------------- + + +def _git(repo, *args): + return subprocess.check_output(["git", *args], cwd=repo, text=True).strip() + + +@pytest.fixture +def repo(tmp_path, monkeypatch): + root = tmp_path / "repo" + (root / "model").mkdir(parents=True) + (root / "orchestration").mkdir() + (root / "model" / "kernel.py").write_text("v1\n") + (root / "orchestration" / "resume.py").write_text("v1\n") + _git(root, "init", "-q") + _git(root, "config", "user.email", "t@example.com") + _git(root, "config", "user.name", "t") + _git(root, "add", "-A") + _git(root, "commit", "-q", "-m", "a") + monkeypatch.setattr(resume_module, "MODEL_PATHS", ("model",)) + return root + + +def _commit(repo, path, text): + (repo / path).write_text(text) + _git(repo, "commit", "-q", "-am", f"change {path}") + return _git(repo, "rev-parse", "HEAD") + + +def _old_run(directory, commit, trees): + base = {k: v for k, v in IDENTITY.items() if k != "commit"} + store = StepStore(directory) + store.open({"commit": commit, **base}, resume=False) + return {**base, "model_trees": trees} + + +def test_model_trees_ignore_orchestration_changes_but_not_model_changes(repo): + first = model_trees("HEAD", cwd=repo) + _commit(repo, "orchestration/resume.py", "v2\n") + assert model_trees("HEAD", cwd=repo) == first + _commit(repo, "model/kernel.py", "v2\n") + assert model_trees("HEAD", cwd=repo) != first + + +def test_adopting_a_commit_based_run_lets_a_later_commit_continue_it(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + new_identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "orchestration/resume.py", "v2\n") + new_identity["model_trees"] = model_trees("HEAD", cwd=repo) + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert StepStore(tmp_path / "run").open(new_identity, resume=True) == 0 + attempts = StepStore(tmp_path / "run").attempts() + assert attempts[-2]["adopted"]["from_commit"] == old + assert attempts[-2]["adopted"]["model_paths_identical"] is True + + +def test_adoption_is_refused_when_a_model_path_changed(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + _commit(repo, "model/kernel.py", "v2\n") + new_identity = { + **{k: v for k, v in IDENTITY.items() if k != "commit"}, + "model_trees": model_trees("HEAD", cwd=repo), + } + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="model paths changed"): + adopt_identity(tmp_path / "run", new_identity, old[:8], cwd=repo) + assert (tmp_path / "run" / "run.json").read_text() == before + + +def test_adoption_refuses_other_differences_and_dry_run_writes_nothing(repo, tmp_path): + old = _git(repo, "rev-parse", "HEAD") + identity = _old_run(tmp_path / "run", old, model_trees("HEAD", cwd=repo)) + before = (tmp_path / "run" / "run.json").read_text() + with pytest.raises(ResumeError, match="more than the commit"): + adopt_identity(tmp_path / "run", {**identity, "seed": 7}, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="not abc"): + adopt_identity(tmp_path / "run", identity, "abc", cwd=repo) + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo, dry_run=True) + assert (tmp_path / "run" / "run.json").read_text() == before + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) + with pytest.raises(ResumeError, match="already"): + adopt_identity(tmp_path / "run", identity, old[:8], cwd=repo) diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py new file mode 100644 index 000000000..03a70f1a3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +def _native_call(extension, operation, boundary, case): + x, index = case["x"], case["index"] + tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("device", ["cpu", "other_cuda"]) +def test_native_rejects_mixed_device_tiles( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) + + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +@pytest.mark.cuda_only +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.cuda_only +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize( + "operation", ["gate_forward", "gate_dy", "norm_forward", "norm_dx", "norm_partials"] +) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + script = textwrap.dedent( + """ + import sys + import torch + from rl_engine import _C + op, bad = sys.argv[1], int(sys.argv[2]) + x = torch.ones(2, 8, device="cuda") + w = torch.ones(8, device="cuda") + table = torch.ones(2, 8, device="cuda") + index = torch.tensor([bad, 0], device="cuda") + rstd = torch.ones(2, device="cuda") + rows = torch.arange(2, device="cuda") + begin = torch.tensor([0], device="cuda") + end = torch.tensor([2], device="cuda") + try: + if op == "gate_forward": + _C.h3_gate_residual_forward(x, x, table, index) + elif op == "gate_dy": + _C.h3_gate_residual_backward_dy(x, table, index) + elif op == "norm_forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, table, table, index) + elif op == "norm_dx": + _C.h3_rmsnorm_backward_dx(x, x, w, rstd, table, table, index) + else: + _C.h3_rmsnorm_backward_partials( + x, x, w, rstd, table, table, index, rows, begin, end, rows, begin, end) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid index reached native table lookup without an assertion") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +def test_row_gather_rejects_metadata_on_other_device(native_extension, metadata): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + grad = torch.ones(6, 2, 8, device="cuda:0") + tensors = dict( + sorted_pos=torch.tensor([0, 1], device="cuda:0"), + tile_begin=torch.tensor([0], device="cuda:0"), + tile_end=torch.tensor([2], device="cuda:0"), + seg_first_tile=torch.tensor([0, 1], device="cuda:0"), + ) + tensors[metadata] = tensors[metadata].to("cuda:1") + with pytest.raises(RuntimeError, match="grad device"): + native_extension.h3_adaln_row_gather_backward(grad, **tensors, out_dtype=torch.float32) + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("bound", ["negative", "past_end"]) +def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metadata, bound): + import subprocess + import sys + import textwrap + + script = textwrap.dedent( + """ + import sys + import torch + from rl_engine import _C + name, bound = sys.argv[1:] + grad = torch.ones(6, 2, 8, device="cuda") + metadata = dict(sorted_pos=torch.tensor([0, 1], device="cuda"), + tile_begin=torch.tensor([0], device="cuda"), + tile_end=torch.tensor([2], device="cuda"), + seg_first_tile=torch.tensor([0, 1], device="cuda")) + metadata[name][0] = -1 if bound == "negative" else 3 + try: + _C.h3_adaln_row_gather_backward(grad, **metadata, out_dtype=torch.float32) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid metadata was not rejected") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py new file mode 100644 index 000000000..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_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py new file mode 100644 index 000000000..1d1e9d65e --- /dev/null +++ b/tests/models/minimax_h3/test_h3_parallel.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 WS2 rows: TP attention/FFN, Ulysses CP, metadata co-shard, arbitrary length, packed IO. + +Single-process tests pin the canonical 8-unit reduction (summing per-rank partials in tree +order reproduces the TP=1 result for TP 2/4/8) and the sharding helpers. Multi-process tests +(2-4 GPUs) run the whole block 0 under TP / CP / SP grids and require every rank's rows to be +bitwise equal to the single-GPU canonical block, forward and backward. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par +from rl_engine.distributed.algorithms.hierarchical import tree_add +from rl_engine.validation.models import h3_tp +from rl_engine.validation.models.h3_ws2 import run_world + +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def gpus(n): + return pytest.mark.skipif(torch.cuda.device_count() < n, reason=f"needs {n} GPUs") + + +def _triton(): + pytest.importorskip("triton") + + +class TestSharding: + @pytest.mark.parametrize("tp", [1, 2, 4, 8]) + def test_row_and_column_shards_tile_once(self, tp): + w = torch.arange(56 * 128 * 16, dtype=torch.float32).view(56 * 128, 16) + assert torch.equal(torch.cat([par.shard_rows(w, tp, r) for r in range(tp)]), w) + w2 = torch.arange(8 * 64 * 4, dtype=torch.float32).view(4, 8 * 64) + assert torch.equal(torch.cat([par.shard_cols(w2, tp, r) for r in range(tp)], 1), w2) + + @pytest.mark.parametrize("tp", [3, 5, 6, 16]) + def test_illegal_tp_is_rejected(self, tp): + with pytest.raises(ValueError): + par.shard_rows(torch.zeros(56 * 128, 4), tp, 0) + + def test_shards_are_whole_heads(self): + for tp in (1, 2, 4, 8): + assert (par.HEADS // tp) * par.HEAD_DIM == par.INNER // tp + + def test_gate_up_shards_keep_both_halves(self): + w = torch.arange(2 * 64 * 3, dtype=torch.float32).view(128, 3) + s = par.shard_gate_up(w, 4, 1) + hidden, gate = s.chunk(2) + assert torch.equal(hidden, w[:64][16:32]) and torch.equal(gate, w[64:][16:32]) + + @pytest.mark.parametrize("length", [1, 2, 3, 7, 261, 264, 925]) + @pytest.mark.parametrize("cp", [1, 2, 3, 4]) + def test_seq_shards_tile_any_length(self, length, cp): + spans = [ + (par.SeqShard(length, cp, r).lo, par.SeqShard(length, cp, r).hi) for r in range(cp) + ] + assert spans[0][0] == 0 and spans[-1][1] == length + assert all(a[1] == b[0] for a, b in zip(spans, spans[1:])) + sizes = [hi - lo for lo, hi in spans] + assert max(sizes) - min(sizes) <= 1 + + def test_coshard_slices_every_per_row_tensor_alike(self): + n = 11 + idx, tags, rows = torch.arange(n), torch.arange(n) * 3, torch.arange(n * 2).view(n, 2) + for r in range(3): + s = par.SeqShard(n, 3, r) + a, b, c = par.coshard(s, idx, tags, rows) + assert torch.equal(a * 3, b) and torch.equal(c[:, 0] // 2, a) + + def test_coshard_rejects_wrong_length(self): + with pytest.raises(ValueError, match="rows"): + par.coshard(par.SeqShard(5, 2, 0), torch.zeros(4)) + + +@needs_gpu +class TestCanonicalReduction: + """Per-rank partials added in tree order are bitwise the TP=1 result (no collective needed).""" + + @pytest.mark.parametrize("k", [7168, 14336]) + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_row_parallel_partials_fold_to_tp1(self, k, tp): + _triton() + gen = torch.Generator().manual_seed(1) + x = torch.randn(33, k, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, k, generator=gen) / k**0.5).bfloat16().cuda() + full = par._units_partial(x, w, 1) + parts = [ + par._units_partial(par.shard_heads(x, tp, r, 1), par.shard_cols(w, tp, r), tp) + for r in range(tp) + ] + assert torch.equal(tree_add(parts), full) + assert torch.equal(par.row_parallel_linear(x, w), full.bfloat16()) + + def test_row_parallel_close_to_fp64(self): + _triton() + gen = torch.Generator().manual_seed(2) + x = torch.randn(64, 7168, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda() + got = par.row_parallel_linear(x, w).double() + ref = x.double() @ w.double().t() + provider = torch.nn.functional.linear(x, w).double() + err = (got - ref).norm() / ref.norm() + assert err <= 1.25 * (provider - ref).norm() / ref.norm() + 1e-3 + + def test_each_row_is_independent_of_the_others(self): + _triton() + gen = torch.Generator().manual_seed(3) + x = torch.randn(70, 7168, generator=gen).bfloat16().cuda() + w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda() + whole = par.row_parallel_linear(x, w) + for rows in (slice(0, 1), slice(5, 6), slice(30, 31), slice(0, 33)): + assert torch.equal(par.row_parallel_linear(x[rows].contiguous(), w), whole[rows]) + + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_column_parallel_backward_folds_to_tp1(self, tp): + _triton() + gen = torch.Generator().manual_seed(4) + k, n = 5376, 7168 + ws = [(torch.randn(n, k, generator=gen) / k**0.5).bfloat16().cuda() for _ in range(3)] + grads = [torch.randn(40, n, generator=gen).bfloat16().cuda() for _ in range(3)] + + def unit_partials(weights, gs, local): + units = [] + for j in range(local): + acc = None + for grad, w in zip(gs, weights): + width = grad.shape[1] // local + piece = par.g.gemm( + grad[:, j * width : (j + 1) * width], + w[j * width : (j + 1) * width], + out_dtype=torch.float32, + ) + acc = piece if acc is None else acc + piece + units.append(acc) + return tree_add(units) + + full = unit_partials(ws, grads, 8) + ranks = [ + unit_partials( + [par.shard_rows(w, tp, r) for w in ws], + [par.shard_heads(g_, tp, r, 1) for g_ in grads], + 8 // tp, + ) + for r in range(tp) + ] + assert torch.equal(tree_add(ranks), full) + + def test_column_parallel_forward_is_the_unsharded_columns(self): + _triton() + gen = torch.Generator().manual_seed(5) + x = torch.randn(40, 5376, generator=gen).bfloat16().cuda() + w = (torch.randn(7168, 5376, generator=gen) / 5376**0.5).bfloat16().cuda() + (full,) = par.column_parallel_linear(x, [w]) + for tp in (2, 4, 8): + for r in range(tp): + (part,) = par.column_parallel_linear(x, [par.shard_rows(w, tp, r)]) + assert torch.equal(part, par.shard_heads(full, tp, r, 1)) + + +@needs_gpu +class TestPackedIO: + @pytest.mark.parametrize("length", [1, 5, 261]) + def test_scatter_then_gather_is_identity_single_rank(self, length): + t = torch.randn(length, 7).cuda() + shard = par.SeqShard(length, 1, 0) + assert torch.equal(par.gather_packed(par.scatter_packed(t, shard), shard), t) + + +# --- Multi-process grids -------------------------------------------------------------------- + + +def _reference(seq_len=0, backward=False): + from rl_engine.validation.models.h3_weights import load_h3_block_weights + + return h3_tp.block_case( + load_h3_block_weights("cuda"), h3_tp.Grid(1, 1, 0, 0), seq_len=seq_len, backward=backward + ) + + +def _assemble(ranks, tp, cp, key="out"): + """Global rows from the ranks; replicated ranks must agree bitwise.""" + + pieces: dict[tuple[int, int], torch.Tensor] = {} + for r in ranks: + span = tuple(int(v) for v in r["rows"]) + if span in pieces: + assert torch.equal(pieces[span], r[key]), f"ranks disagree on rows {span}" + pieces[span] = r[key] + return torch.cat([pieces[s] for s in sorted(pieces)], dim=1), sorted(pieces) + + +GRIDS = [ + pytest.param(2, 1, False, 0, id="tp2", marks=gpus(2)), + pytest.param(4, 1, False, 0, id="tp4", marks=gpus(4)), + pytest.param(1, 2, False, 0, id="cp2", marks=gpus(2)), + pytest.param(2, 2, False, 0, id="tp2_cp2", marks=gpus(4)), + pytest.param(2, 1, True, 0, id="tp2_sp", marks=gpus(2)), + pytest.param(4, 1, True, 0, id="tp4_sp", marks=gpus(4)), + pytest.param(2, 2, True, 0, id="tp2_cp2_sp", marks=gpus(4)), + pytest.param(1, 2, False, 261, id="cp2_len261", marks=gpus(2)), + pytest.param(2, 2, False, 261, id="tp2_cp2_len261", marks=gpus(4)), + pytest.param(4, 1, False, 263, id="tp4_len263", marks=gpus(4)), +] + + +@needs_gpu +@pytest.mark.parametrize(("tp", "cp", "sp", "length"), GRIDS) +def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length): + _triton() + ranks = run_world( + tp * cp, + h3_tp.block_target, + {"dummy": torch.zeros(1)}, + tp=tp, + cp=cp, + sp=sp, + seq_len=length, + backward=True, + ) + ref = _reference(length, backward=True) + out, spans = _assemble(ranks, tp, cp) + assert out.shape == ref["out"].shape, (spans, out.shape) + assert torch.equal(out.cpu(), ref["out"].cpu()), "forward differs from the single-GPU block" + dx, _ = _assemble(ranks, tp, cp, "dx") + assert torch.equal(dx.cpu(), ref["dx"].cpu()), "input gradient differs" + if cp == 1: # weight gradients: shards of the unsharded gradient + for name, axis in ( + ("attn.to_out.0.weight", 1), + ("ff.net.2.weight", 1), + ("attn.to_q.weight", 0), + ("attn.to_k.weight", 0), + ("attn.to_v.weight", 0), + ("ff.hidden", 0), + ("ff.gate", 0), + ): + if sp: + continue # SP shards sequence rows; dW is reduced over rows below + full = ref["dw." + name].cpu() + shards = [r["dw." + name].cpu() for r in sorted(ranks, key=lambda r: r["rank"])] + assert torch.equal(torch.cat(shards, dim=axis), full), f"dW {name} differs" + + +@needs_gpu +@gpus(2) +def test_ulysses_matches_unsharded_attention_for_ragged_lengths(): + _triton() + for length in (1, 2, 7, 130): + ranks = run_world(2, _ulysses_target, {"dummy": torch.zeros(1)}, length=length) + gathered = torch.cat([r["out"] for r in ranks], dim=1) + assert torch.equal(gathered.cpu(), ranks[0]["ref"].cpu()), f"length {length}" + + +def _ulysses_target(collective, dummy, *, length): + import torch.distributed as dist + + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + + del collective, dummy + rank = dist.get_rank() + gen = torch.Generator().manual_seed(7) + q, k, v = (torch.randn(1, length, 8, 128, generator=gen).bfloat16().cuda() for _ in range(3)) + op = DeterministicAttentionOp() + + def attend(q_, k_, v_): + t = lambda x: x.transpose(1, 2).contiguous() # noqa: E731 + return op(t(q_), t(k_), t(v_), causal=False).transpose(1, 2).contiguous() + + shard = par.SeqShard(length, 2, rank) + local = [par.scatter_packed(t, shard, dim=1) for t in (q, k, v)] + out = par.ulysses_attention(*local, attend, shard, dist.group.WORLD) + return {"out": out, "ref": attend(q, k, v)} + + +# --- Hierarchical (cross-node design) collective, exercised inside one node ---------------- + + +def _hier_target(collective, dummy, *, local_size, dtype): + import torch.distributed as dist + + from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective + + del dummy + h = HierarchicalCollective( + local_size, device=torch.cuda.current_device(), max_size_bytes=64 << 20 + ) + gen = torch.Generator().manual_seed(100 + dist.get_rank()) + x = ( + (torch.randn(257, 129, generator=gen) * 10 ** torch.randint(-3, 4, (257, 1), generator=gen)) + .to(dtype) + .cuda() + ) + flat = collective.all_reduce(x.clone()) + mine = h.all_reduce(x) + rs = h.reduce_scatter(x[:256].contiguous()) + ag = h.all_gather(x[:2].contiguous()) + return {"flat": flat, "hier": mine, "rs": rs, "ag": ag, "x": x} + + +@needs_gpu +@gpus(4) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("local_size", [1, 2]) +def test_hierarchical_equals_single_node_fixed_tree(local_size, dtype): + ranks = run_world( + 4, _hier_target, {"dummy": torch.zeros(1)}, local_size=local_size, dtype=dtype + ) + for r in ranks: + assert torch.equal(r["flat"], r["hier"]), f"rank {r['rank']}: hierarchical != fixed tree" + assert torch.equal(ranks[0]["rs"], ranks[0]["hier"][:64]) + assert torch.equal(ranks[3]["rs"], ranks[3]["hier"][192:256]) + assert torch.equal(ranks[1]["ag"], torch.cat([r["x"][:2] for r in ranks])) 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 new file mode 100644 index 000000000..ed74d1901 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements, and +comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models import h3_chain, h3_report +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import ( + provider_adaln_row_gather, + provider_norm_modulate, +) +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run tools/weights/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = ( + Path(__file__).resolve().parents[3] + / "tools" + / "validation" + / "models" + / "plot_h3_evidence.py" + ) + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + + assert enable_timing + + def record(self): + """Capture the synthetic elapsed clock at this event.""" + + self.timestamp = elapsed_us + + def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + + pass + + def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py new file mode 100644 index 000000000..ced9a9377 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_rmsnorm``: RMSNorm (eps 1e-5) and its fused AdaLN modulation. + +* forward is bitwise equal to ``nn.RMSNorm`` and to diffusers' + ``norm(x) * (1.0 + scale[i]) + shift[i]`` (block: ``adaln_indices``; + ``norm_out``: ``timestep_indices``) on the pinned norm weights; +* rows are independent of batch size and position; +* backward is deterministic and checked against an FP64 golden; +* malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.rmsnorm import NativeH3RMSNormOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_norm_modulate + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NORMS = ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", +) + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import ( + H3RMSNormCudaOp, + h3_rmsnorm_available, + ) + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3RMSNormCudaOp() + + +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(shape, generator=g) * scale).to(device=device, dtype=dtype) + + +def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(rows, chunks * hidden, generator=g) * 0.5).to(dtype).cuda() + return table.chunk(chunks, dim=-1) # strided (R, H) views, like the AdaLN outputs + + +def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * weight.double() + return n * (1 + scale.double().index_select(0, index)) + shift.double().index_select(0, index) + + +class TestReference: + def test_golden_close_to_provider(self): + x, w = _x((3, 64), torch.float32, device="cpu"), torch.rand(64) + 0.5 + op = NativeH3RMSNormOp() + torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + + def test_modulated_with_prevalidated_indices(self): + x, w = _x((2, 8), torch.float32, device="cpu"), torch.ones(8) + shift, scale = torch.randn(3, 8), torch.randn(3, 8) + index = torch.tensor([2, 0]) + op = NativeH3RMSNormOp() + expected = op.forward(x, w) * (1 + scale[index]) + shift[index] + torch.testing.assert_close( + op.forward_modulated(x, w, shift, scale, index, check_range=False), + expected, + ) + + def test_rejects_bad_inputs(self): + op = NativeH3RMSNormOp() + x, w = torch.randn(2, 8), torch.ones(8) + with pytest.raises(TypeError): + op.forward(x, w.double()) + with pytest.raises(TypeError): + op.forward(x, w.bfloat16()) + with pytest.raises(ValueError): + op.forward(x, torch.ones(7)) + with pytest.raises(ValueError): + op.forward(x, w, eps=0.0) + shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(ValueError): # one index per position S + op.forward_modulated(x, w, shift, scale, torch.tensor([0])) + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("name", NORMS) + def test_bitwise_equal_to_nn_rmsnorm_on_pinned_weights(self, h3_weights_cpu, name): + weight = h3_weights_cpu[name].cuda() + module = torch.nn.RMSNorm(5376, eps=1e-5).cuda().bfloat16() + module.weight.data.copy_(weight) + x = _x((2, 777, 5376), seed=1) + assert torch.equal(_cuda_op()(x, weight), module(x)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [12, 256, 4096, 5376]) + def test_bitwise_equal_to_f_rms_norm(self, dtype, hidden): + x = _x((5, 33, hidden), dtype, seed=hidden, scale=30.0) + w = (torch.rand(hidden, device="cuda") * 2).to(dtype) + ref = torch.nn.functional.rms_norm(x, (hidden,), w, 1e-5) + assert torch.equal(_cuda_op()(x, w), ref) + + def test_block_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["transformer_blocks.0.norm1.weight"].cuda() + shift, scale, *_ = _table(9, 5376, 6, seed=2) # (3T, H) rows, T = 3 + ti, tags = h3_packed_layout(4097, 3, seed=2) + index = ti * 3 + tags + x = _x((2, 4097, 5376), seed=2) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, index)) + + def test_norm_out_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["norm_out.norm.weight"].cuda() + shift, scale = _table(3, 5376, 2, seed=3) # (T, H) rows indexed by timestep + ti, _ = h3_packed_layout(1000, 3, seed=3) + x = _x((1, 1000, 5376), seed=3) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, ti) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, ti)) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_modulation_bitwise_in_every_dtype(self, dtype): + """FP32 must not contract the separately rounded eager ops into an FMA.""" + + w = (torch.rand(5376, device="cuda") + 0.5).to(dtype) + shift, scale, *_ = _table(9, 5376, 6, dtype=dtype, seed=9) + ti, tags = h3_packed_layout(333, 3, seed=9) + index = ti * 3 + tags + x = _x((2, 333, 5376), dtype, seed=9) + ours = _cuda_op().forward_modulated(x, w, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + + def test_rows_are_batch_and_position_invariant(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + shift, scale, *_ = _table(9, 5376, 6, seed=4) + ti, tags = h3_packed_layout(300, 3, seed=4) + index = ti * 3 + tags + x = _x((3, 300, 5376), seed=4) + full = _cuda_op().forward_modulated(x, w, shift, scale, index) + part = _cuda_op().forward_modulated(x[2:3, 50:90], w, shift, scale, index[50:90]) + assert torch.equal(part[0], full[2, 50:90]) + + def test_rejects_unsupported(self): + w = torch.ones(6, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(_x((2, 6)), w) + with pytest.raises(ValueError): + _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + + +_DEVICE_MESSAGES = { + "weight": "with x's dtype and device", + "index": r"row index must be a non-empty contiguous int64 tensor on cuda:0", + "rstd": "on x's device", +} + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +class TestCudaNativeDeviceValidation: + @pytest.mark.parametrize("argument", ["weight", "index"]) + def test_forward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + if argument == "weight": + weight = weight.to("cuda:1") + else: + index = index.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @pytest.mark.parametrize("argument", ["weight", "index", "rstd"]) + def test_backward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + if argument == "weight": + weight = weight.to("cuda:1") + elif argument == "index": + index = index.to("cuda:1") + else: + rstd = rstd.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, fn, tensors, grad): + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + def test_modulated_backward_against_fp64(self): + g = torch.Generator(device="cpu").manual_seed(5) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.5).bfloat16().cuda() + ti, tags = h3_packed_layout(2049, 3, seed=5) + index = ti * 3 + tags + x = _x((2, 2049, 5376), seed=5) + grad = _x((2, 2049, 5376), seed=6, scale=1.0) + op = _cuda_op() + + def ours(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return op.forward_modulated(x_, w_, shift, scale, index) + + def golden(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return _fp64_modulated(x_, w_, shift, scale, index) + + got = self._grads(ours, (x, w, table), grad) + ref = self._grads(golden, (x.double(), w.double(), table.double()), grad.double()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + rel = ((g_.double() - r).abs().max() / r.abs().max()).item() + # Two BF16 roundings remain: the eager VJP uses round(1 + scale), and + # every gradient is rounded once to BF16 (2^-9 relative each). + assert rel < 1e-2, (name, rel) + again = self._grads(ours, (x, w, table), grad) + assert all(torch.equal(a, b) for a, b in zip(got, again)) + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_modulated_gradients_meet_contract_against_golden(self, seq): + # Against forward_modulated_fp32 before it stored norm(x) and 1 + scale + # in BF16, d_weight missed from S=1024 and d_scale at S=4097. + from rl_engine.contracts.numerical import load_contract, resolve_tolerance + + spec = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=torch.bfloat16, + ) + g = torch.Generator(device="cpu").manual_seed(seq) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.1).bfloat16().cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seq) + index = ti * 3 + tags + x = _x((1, seq, 5376), seed=seq) + grad = _x((1, seq, 5376), seed=seq + 1, scale=1.0) + op, native = _cuda_op(), NativeH3RMSNormOp() + + def run(fn, upstream): + def call(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return fn(x_, w_, shift, scale, index) + + return self._grads(call, (x, w, table), upstream) + + got = run(op.forward_modulated, grad) + ref = run(native.forward_modulated_fp32, grad.float()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + torch.testing.assert_close( + g_.float(), r.float(), atol=spec.atol, rtol=spec.rtol, msg=lambda m: f"{name}: {m}" + ) + + def test_plain_backward_against_fp64_and_row_local_dx(self): + w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() + x = _x((4, 64, 5376), seed=7) + grad = _x((4, 64, 5376), seed=8, scale=1.0) + op = _cuda_op() + got = self._grads(lambda a, b: op(a, b), (x, w), grad) + + def golden(a, b): + return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + 1e-5) * b + + ref = self._grads(golden, (x.double(), w.double()), grad.double()) + for g, r in zip(got, ref): + assert ((g.double() - r).abs().max() / r.abs().max()).item() < 5e-3 + part = self._grads(lambda a, b: op(a, b), (x[1:2, 10:20], w), grad[1:2, 10:20]) + assert torch.equal(part[0][0], got[0][1, 10:20]) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + assert ( + type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" + ) diff --git a/tests/models/minimax_h3/test_h3_runtime_integration.py b/tests/models/minimax_h3/test_h3_runtime_integration.py new file mode 100644 index 000000000..273a99afc --- /dev/null +++ b/tests/models/minimax_h3/test_h3_runtime_integration.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 integration rows: runtime modes, ablation overlay, Diffusers adapter, training +provider, LoRA gradient gate and the CPS transition contract.""" + +from __future__ import annotations + +import importlib.util +import math +import os +from pathlib import Path + +import pytest +import torch + +from rl_engine.models.minimax_h3 import runtime as rt +from rl_engine.runtime.policy import RLKernelMode + +needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +# --- runtime modes and overlay -------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("mode", "bind", "fail", "report", "case"), + [ + ("strict", True, True, False, "R/R"), + ("audit", True, False, True, "R/R"), + ("auto", False, False, False, "P/P"), + ("off", False, False, False, "P/P"), + ], +) +def test_runtime_modes(mode, bind, fail, report, case): + p = rt.h3_runtime_policy(mode) + assert (p.bind, p.fail_on_fallback, p.report_only, p.case_id) == (bind, fail, report, case) + + +def test_runtime_mode_reads_environment(monkeypatch): + monkeypatch.setenv("RL_KERNEL_MODE", "off") + assert rt.h3_runtime_policy().mode is RLKernelMode.OFF + monkeypatch.setenv("RL_KERNEL_MODE", "bogus") + with pytest.raises(RuntimeError, match="RL_KERNEL_MODE"): + rt.h3_runtime_policy() + + +def test_overlay_parse_label_and_validation(): + assert rt.AblationOverlay.parse("").label == "R/R" + o = rt.AblationOverlay.parse("node:attention, boundary:token_refiner") + assert o.native_nodes == {"attention"} and o.native_boundaries == {"token_refiner"} + assert o.label == "ablation(native: node:attention,boundary:token_refiner)" + for bad in ("node:nope", "boundary:proj_in", "layer:x"): + with pytest.raises(ValueError): + rt.AblationOverlay.parse(bad) + + +class _FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.p = torch.nn.Parameter(torch.zeros(1)) + + +def test_adapter_off_binds_nothing_and_reports_provider(): + adapter = rt.H3DiffusersAdapter(_FakeModel(), "off") + assert adapter._bind() is False + rep = adapter.report() + assert (rep["bound"], rep["label"], rep["case_id"]) == (False, "P/P", "P/P") + + +def test_adapter_strict_refuses_cpu_instead_of_falling_back(): + with pytest.raises(RuntimeError, match="refusing to fall back"): + rt.H3DiffusersAdapter(_FakeModel(), "strict")._bind() + + +def test_adapter_audit_reports_the_fallback(): + adapter = rt.H3DiffusersAdapter(_FakeModel(), "audit") + assert adapter._bind() is False and adapter.fell_back == "no CUDA device" + assert adapter.report()["label"] == "P/P" + + +# --- CPS transition ----------------------------------------------------------------------- + + +def _scheduler_cps(x, v, sigma, sigma_prev, noise, level): + sin_noise = math.sin(level * math.pi / 2) + std = sigma_prev * sin_noise + pred = x - sigma * v + est = x + v * (1 - sigma) + mean = pred * (1 - sigma_prev) + est * torch.sqrt(sigma_prev**2 - std**2) + prev = mean + std * noise + return prev, mean, std, (-((prev - mean) ** 2)).mean(dim=tuple(range(1, prev.ndim))) + + +@needs_gpu +class TestTransition: + def _case(self, batch=3, shape=(257, 96)): + from rl_engine.models.minimax_h3 import transition as tr + + g = torch.Generator().manual_seed(1) + x = torch.randn(batch, *shape, generator=g).cuda() + v = torch.randn(batch, *shape, generator=g).cuda() + noise = torch.stack([tr.step_noise(7, 2, i, "video", shape, "cuda") for i in range(batch)]) + return tr, x, v, noise + + def test_matches_the_scheduler_formula(self): + pytest.importorskip("triton") + tr, x, v, noise = self._case() + sigma = torch.tensor(0.75, device="cuda") + sigma_prev = torch.tensor(0.5, device="cuda") + out = tr.cps_transition(x, v, sigma, sigma_prev, noise=noise) + prev, mean, std, logp = _scheduler_cps(x, v, sigma, sigma_prev, noise, tr.NOISE_LEVEL) + assert torch.equal(out.prev_sample, prev) and torch.equal(out.prev_sample_mean, mean) + assert torch.equal(out.std_dev_t, std) + assert torch.allclose(out.log_prob, logp, rtol=1e-4, atol=1e-12) + + def test_targets_cps_at_noise_level_08(self): + from rl_engine.models.minimax_h3 import transition as tr + + assert (tr.SDE_TYPE, tr.NOISE_LEVEL) == ("cps", 0.8) + + def test_each_sample_is_independent_of_the_batch(self): + pytest.importorskip("triton") + tr, x, v, noise = self._case() + whole = tr.cps_transition(x, v, 0.75, 0.5, noise=noise) + for i in range(x.shape[0]): + one = tr.cps_transition(x[i : i + 1], v[i : i + 1], 0.75, 0.5, noise=noise[i : i + 1]) + assert torch.equal(one.prev_sample, whole.prev_sample[i : i + 1]) + assert torch.equal(one.log_prob, whole.log_prob[i : i + 1]) + + def test_noise_depends_on_sample_id_only(self): + from rl_engine.models.minimax_h3 import transition as tr + + a = tr.step_noise(1, 3, 5, "video", (4, 8), "cuda") + assert torch.equal(a, tr.step_noise(1, 3, 5, "video", (4, 8), "cpu").cuda()) + for other in ( + (2, 3, 5, "video"), + (1, 4, 5, "video"), + (1, 3, 6, "video"), + (1, 3, 5, "audio"), + ): + assert not torch.equal(a, tr.step_noise(*other, (4, 8), "cuda")) + + def test_scoring_a_stored_rollout_reproduces_sampling_logprob(self): + pytest.importorskip("triton") + tr, x, v, noise = self._case() + sampled = tr.cps_transition(x, v, 0.75, 0.5, noise=noise) + scored = tr.cps_transition(x, v, 0.75, 0.5, prev_sample=sampled.prev_sample) + assert torch.equal(scored.log_prob, sampled.log_prob) + + def test_zero_noise_level_is_the_mean_with_zero_logprob(self): + pytest.importorskip("triton") + tr, x, v, noise = self._case() + out = tr.cps_transition(x, v, 0.75, 0.5, noise_level=0.0, noise=noise) + assert torch.equal(out.prev_sample, out.prev_sample_mean) + assert torch.count_nonzero(out.log_prob) == 0 + + def test_gradient_reaches_the_velocity(self): + pytest.importorskip("triton") + tr, x, v, noise = self._case() + v.requires_grad_(True) + out = tr.cps_transition(x, v, 0.75, 0.5, prev_sample=(x + 0.1).detach()) + out.prev_sample_mean.sum().backward() + assert v.grad is not None and torch.isfinite(v.grad).all() + + def test_matches_verl_omni_scheduler_when_available(self): + root = os.environ.get("VERL_OMNI_PATH") + if not root: + pytest.skip("set VERL_OMNI_PATH to a verl-omni checkout") + # load by path: putting the checkout on sys.path would shadow our own `tests` package + path = Path(root) / "verl_omni/pipelines/schedulers/flow_match_sde.py" + if not path.exists(): + pytest.skip(f"no scheduler at {path}") + spec = importlib.util.spec_from_file_location("flow_match_sde_under_test", path) + sched_mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(sched_mod) + pytest.importorskip("triton") + tr, x, v, noise = self._case(batch=2, shape=(64, 96)) + scheduler = sched_mod.FlowMatchSDEDiscreteScheduler(num_train_timesteps=1000, shift=1.0) + scheduler.sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + scheduler.timesteps = scheduler.sigmas[:-1] * 1000 + scheduler._step_index = 1 + # the scheduler draws its own noise: score the sampled rollout instead + ours = tr.cps_transition(x, v, 0.75, 0.5, noise=noise) + ref = scheduler.sample_previous_step( + sample=x.cpu(), + model_output=v.cpu(), + prev_sample=ours.prev_sample.cpu(), + noise_level=tr.NOISE_LEVEL, + sde_type=tr.SDE_TYPE, + ) + assert torch.equal(ref[2], ours.prev_sample_mean.cpu()) + assert torch.allclose(ref[1], ours.log_prob.cpu(), rtol=1e-4, atol=1e-12) + + +# --- LoRA gradient gate and vLLM-Omni adapter ------------------------------------------------- + + +@needs_gpu +def test_lora_gradient_gate(): + pytest.importorskip("triton") + g = torch.Generator().manual_seed(3) + w = (torch.randn(512, 384, generator=g) / 20).bfloat16().cuda() + x = torch.randn(70, 384, generator=g).bfloat16().cuda() + grad = torch.randn(70, 512, generator=g).bfloat16().cuda() + report = rt.lora_gradient_gate(w, x, grad) + assert report["repeatable"] and report["forward_rows_independent"] and report["base_frozen"] + assert report["forward_rel_err_vs_fp64"] < 5e-3 + + +@needs_gpu +def test_lora_gradients_match_fp64_autograd(): + pytest.importorskip("triton") + g = torch.Generator().manual_seed(4) + n, k, r, t = 256, 192, 4, 33 + w = (torch.randn(n, k, generator=g) / 14).bfloat16().cuda() + x = torch.randn(t, k, generator=g).bfloat16().cuda().requires_grad_() + a = (torch.randn(r, k, generator=g) / 14).cuda().requires_grad_() + b = (torch.randn(n, r, generator=g) * 0.1).cuda().requires_grad_() + grad = torch.randn(t, n, generator=g).bfloat16().cuda() + rt.strict_lora_linear(x, w, a, b, 2.0).backward(grad) + xd, ad, bd = (t_.detach().double().requires_grad_() for t_ in (x, a, b)) + ((xd @ w.double().t() + 2.0 * (xd @ ad.t()) @ bd.t()) * grad.double()).sum().backward() + for got, ref in ((a.grad, ad.grad), (b.grad, bd.grad), (x.grad, xd.grad)): + assert (got.double() - ref).norm() / ref.norm() < 5e-3 diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..a60489a04 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``sp_norm_adaln``: sequence-parallel norm / modulation / gated residual. + +* ownership: contiguous position shards per batch item; bad layouts fail closed; +* backward (one GPU, ranks as threads around the plain backward functions): every + rank's row gradients equal WS1's rows, and its ``dweight``/``d_shift``/``d_scale``/ + ``d_gate`` equal WS1's, for SP 2..8, odd S, B > 1, interleaved modality rows and + bf16/fp32; +* end to end (real NCCL, autograd): the MLP-side block region + ``norm2(residual + gate_msa * y)`` with modulation is byte-equal to WS1 on every rank. +""" + +from __future__ import annotations + +import os +import threading +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + SPRowLayout, + sp_row_layout, +) +from rl_engine.validation.models.h3_ws2 import ( + run_world, + sp_case, + sp_matches_ws1, + sp_region_rank, + ws1_sp_region, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +ALL_TRUE = dict.fromkeys(("out", "d_residual", "d_y", "d_weight", "d_table"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_norm_adaln_available, + ) + + if not sp_norm_adaln_available(): + pytest.skip("rl_engine._C lacks the H3 SP symbols") + + +class _ThreadGroup: + """Ranks as threads on one GPU; ``all_gather`` is the rank-ordered concatenation. + + Only for the plain backward functions: autograd runs every backward of a device + on one worker thread, so autograd-driven ranks need real processes. + """ + + def __init__(self, world: int) -> None: + self.barrier, self.slots = threading.Barrier(world), [None] * world + + def handle(self, rank: int): + def all_gather(t): + self.slots[rank] = t + self.barrier.wait() + out = torch.cat(self.slots) + self.barrier.wait() + return out + + return SimpleNamespace(rank=rank, world_size=len(self.slots), all_gather=all_gather) + + def run(self, fn): + results, errors = [None] * len(self.slots), [] + + def body(rank): + try: + results[rank] = fn(self.handle(rank)) + except BaseException as exc: # re-raised below + errors.append(exc) + self.barrier.abort() + + threads = [threading.Thread(target=body, args=(r,)) for r in range(len(self.slots))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise next( + (e for e in errors if not isinstance(e, threading.BrokenBarrierError)), errors[0] + ) + return results + + +class TestLayout: + @pytest.mark.parametrize("seq, sp", [(8, 8), (257, 2), (1000, 3), (4097, 8)]) + def test_shards_cover_every_position_once(self, seq, sp): + bounds = [SPRowLayout(seq, 1, sp, r).bounds(r) for r in range(sp)] + assert bounds[0][0] == 0 and bounds[-1][1] == seq + assert all(a[1] == b[0] and a[0] < a[1] for a, b in zip(bounds, bounds[1:])) + + @pytest.mark.parametrize("seq, batch, sp, rank", [(3, 1, 4, 0), (8, 0, 2, 0), (8, 1, 2, 2)]) + def test_rejects(self, seq, batch, sp, rank): + with pytest.raises(ValueError): + sp_row_layout(seq, batch, sp, rank) + + +@requires_cuda +class TestBackwardRanks: + @pytest.mark.parametrize( + "sp, batch, seq, hidden, layout, dtype", + [ + (2, 1, 4097, 5376, "block", torch.bfloat16), + (2, 1, 4097, 5376, "interleaved", torch.bfloat16), + (3, 2, 1000, 256, "interleaved", torch.bfloat16), + (4, 2, 777, 512, "block", torch.float32), + (5, 1, 1301, 128, "interleaved", torch.bfloat16), + (8, 2, 2049, 256, "block", torch.bfloat16), + (8, 1, 40, 64, "interleaved", torch.float32), # shards far smaller than a tile + ], + ) + def test_ranks_reproduce_ws1(self, sp, batch, seq, hidden, layout, dtype): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_gate_residual_backward, + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(batch, seq, hidden, layout=layout, dtype=dtype, seed=sp * 100 + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] # shift/scale_mlp, gate_msa + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = sp_plan(lay, idx, rows) + norm = sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + gate_grads = sp_gate_residual_backward(local["grad"], local["y"], gate, plan, coll) + return lay, norm, gate_grads + + for lay, (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + part = slice(lay.lo, lay.hi) + assert torch.equal( + dx.view(batch, -1, hidden), ws1_norm[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dtype)) and torch.equal( + dsc, ws1_norm[3].to(dtype) + ) + assert torch.equal( + dy.view(batch, -1, hidden), ws1_gate[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + H3SPNormAdaLNCudaOp, + ) + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_plain_norm_uses_activation_device_and_separate_cache(): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + op = H3SPNormAdaLNCudaOp(collective, 5, 1) + with torch.cuda.device(0): + for device in ("cuda:1", "cuda:0"): + x = torch.randn(1, 5, 8, device=device, requires_grad=True) + weight = torch.randn(8, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) + result = op.norm(x, weight) + expected = torch.nn.functional.rms_norm(ref_x, (8,), ref_w, 1e-5) + result.sum().backward() + expected.sum().backward() + torch.testing.assert_close(result, expected) + torch.testing.assert_close(x.grad, ref_x.grad) + torch.testing.assert_close(weight.grad, ref_w.grad) + assert len(op._plans) == 2 + assert op.plan(None, None, device="cuda:1").plan.send_local.device.index == 1 + assert op.plan(None, None, device="cuda:0").plan.send_local.device.index == 0 + + +def _worker_exit_without_result(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + os._exit(kwargs["exit_code"]) + time.sleep(60) + + +def _worker_reports_error(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + queue.put({"rank": rank, "error": "intentional worker failure"}) + time.sleep(60) + + +def _worker_never_reports(rank, world, init, queue, path, target, kwargs): + import time + + time.sleep(60) + + +@pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) +def test_run_world_cleans_up_failed_workers(monkeypatch, mode): + import time + + import torch.multiprocessing as mp + + from rl_engine.validation.models import h3_ws2 + + workers = { + "clean_exit": _worker_exit_without_result, + "crash": _worker_exit_without_result, + "reported_error": _worker_reports_error, + "timeout": _worker_never_reports, + } + monkeypatch.setattr(h3_ws2, "_bootstrap", workers[mode]) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + expected = TimeoutError if mode == "timeout" else RuntimeError + with pytest.raises(expected): + run_world( + 2, + None, + {}, + timeout=1 if mode == "timeout" else 30, + exit_code=7 if mode == "crash" else 0, + ) + assert time.monotonic() - start < 20 + assert {p.pid for p in mp.active_children()} <= before + + +def _worker_reports_result(rank, world, init, queue, path, target, kwargs): + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + + +def test_run_world_returns_rank_ordered_results(monkeypatch): + from rl_engine.validation.models import h3_ws2 + + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker_reports_result) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + results = run_world(2, None, {}, timeout=30) + assert [result["rank"] for result in results] == [0, 1] + assert [result["value"].item() for result in results] == [0, 1] diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py new file mode 100644 index 000000000..555a28ebe --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_mlp_fp32``: FP32 256 -> 5376 -> 2688 timestep MLP. + +* accuracy: forward and every gradient against an FP64 golden, on the pinned + checkpoint weights and on synthetic non-H3 sizes; +* invariance: a timestep's ``temb`` bytes (and its row-local ``dx``) do not + depend on how many timesteps share the call or where it sits; +* determinism: repeated forward/backward runs are bitwise equal; +* contract: BF16 inputs (RFC probe H7) and malformed shapes fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.timestep_mlp import NativeH3TimestepMLPOp +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NAMES = ("x", "w1", "b1", "w2", "b2") + + +def _cuda_op(): + """Construct the CUDA MLP operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import H3TimestepMLPCudaOp + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3TimestepMLPCudaOp() + + +def _features(num, device="cuda", seed=0): + """Generate seeded H3 sinusoidal features on the requested test device.""" + + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) + + +def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + """Create seeded FP32 parameters for a two-layer SiLU MLP of the requested sizes.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 + b1 = torch.randn(hidden, generator=g) * 0.1 + w2 = torch.randn(out, hidden, generator=g) / hidden**0.5 + b2 = torch.randn(out, generator=g) * 0.1 + return [t.to(device) for t in (w1, b1, w2, b2)] + + +def _real_params(weights): + """Return pinned timestep MLP weights and biases on CUDA in layer order.""" + + return [ + weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") + ] + + +def _fp64_grads(x, params, grad_out): + """Compute gradients of both linear layers and SiLU using FP64 autograd.""" + + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] + z = F.linear(leaves[0], leaves[1], leaves[2]) + F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) + return [leaf.grad for leaf in leaves] + + +class TestReference: + def test_provider_and_golden_agree(self): + """Keep the FP32 provider close to the high-precision MLP replay on CPU.""" + + x = _features(5, device="cpu") + params = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + torch.testing.assert_close( + op.forward(x, *params), op.forward_fp32(x, *params), atol=1e-5, rtol=1e-5 + ) + + @pytest.mark.parametrize("which", range(5)) + def test_rejects_bf16_anywhere(self, which): + # RFC probe H7: casting a declared FP32 path to BF16 early. + """Reject BF16 in every argument of the declared FP32 MLP path.""" + + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] + args[which] = args[which].bfloat16() + with pytest.raises(TypeError, match="float32"): + NativeH3TimestepMLPOp().forward(*args) + + def test_rejects_shape_mismatches(self): + """Reject incompatible layer dimensions, malformed biases, and empty feature batches.""" + + x = _features(2, device="cpu") + w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + with pytest.raises(ValueError): + op.forward(x[:, :128], w1, b1, w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1[:10], w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1, w2[:, :10], b2) + with pytest.raises(ValueError): + op.forward(x[:0], w1, b1, w2, b2) + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) + def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + """Check checkpoint forward accuracy, FP32 dtype, and the timestep output shape.""" + + params = _real_params(h3_weights_cpu) + x = _features(num, seed=num) + out = _cuda_op()(x, *params) + gold = NativeH3TimestepMLPOp().forward_fp32(x, *params) + assert out.dtype == torch.float32 and out.shape == (num, 2688) + # tolerance_contract.json forward_accuracy / reduction / float32. + torch.testing.assert_close(out, gold, atol=1e-4, rtol=1e-4) + assert (out - gold).abs().max().item() < 1e-5 + + def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + """Keep each checkpoint embedding row bitwise unchanged across batch sizes and positions.""" + + op = _cuda_op() + params = _real_params(h3_weights_cpu) + x_all = _features(9, seed=1) + full = op(x_all, *params) + for i in range(9): + assert torch.equal(op(x_all[i : i + 1], *params)[0], full[i]) + for num in (2, 3, 4, 5, 8): + for pos in (0, num - 1): + x = _features(num, seed=10 + num) + x[pos] = x_all[4] + assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) + + def test_backward_against_fp64_golden(self, h3_weights_cpu): + """Match all five checkpoint input gradients to the FP64 autograd golden.""" + + params = _real_params(h3_weights_cpu) + x = _features(3, seed=2) + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + grad_out = torch.randn( + 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) + ) + _cuda_op()(*leaves).backward(grad_out) + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out), strict=True): + # tolerance_contract.json gradient_accuracy / reduction / float32. + torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) + def test_other_sizes(self, k_in, hidden, out): + """Match the high-precision forward golden for small and checkpoint-sized MLPs.""" + + params = _synthetic_params(k_in, hidden, out, seed=k_in) + x = torch.randn(5, k_in, device="cuda") + torch.testing.assert_close( + _cuda_op()(x, *params), + NativeH3TimestepMLPOp().forward_fp32(x, *params), + atol=1e-4, + rtol=1e-4, + ) + + def test_zero_hidden_size_matches_cpu_forward_backward(self): + """Match CPU forward values and all input gradients with a zero-width hidden layer.""" + + generator = torch.Generator(device="cpu").manual_seed(0) + cpu_inputs = [ + torch.randn(3, 8, generator=generator), + torch.empty(0, 8), + torch.empty(0), + torch.empty(16, 0), + torch.linspace(-1.0, 1.0, 16), + ] + cpu_inputs = [tensor.requires_grad_(True) for tensor in cpu_inputs] + cuda_inputs = [tensor.detach().cuda().requires_grad_(True) for tensor in cpu_inputs] + cpu_out = NativeH3TimestepMLPOp()(*cpu_inputs) + cuda_out = _cuda_op()(*cuda_inputs) + torch.testing.assert_close(cuda_out.cpu(), cpu_out, atol=0, rtol=0) + grad = torch.randn(3, 16, generator=generator) + cpu_out.backward(grad) + cuda_out.backward(grad.cuda()) + for name, cpu_input, cuda_input in zip(NAMES, cpu_inputs, cuda_inputs, strict=True): + torch.testing.assert_close( + cuda_input.grad.cpu(), cpu_input.grad, atol=0, rtol=0, msg=name + ) + + def test_rejects_unaligned_k(self): + """Reject inner dimensions that violate the native FP32 four-element alignment.""" + + params = _synthetic_params(6, 8, 8) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(torch.randn(2, 6, device="cuda"), *params) + + def test_rejects_cpu_and_mixed_devices(self): + """Reject CPU-only and mixed-device inputs at the CUDA MLP boundary.""" + + params = _synthetic_params(8, 24, 16) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *params) + + def test_repeat_and_backward_are_deterministic(self): + """Require repeated MLP outputs and all input gradients to be bitwise equal.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=3) + x = _features(4, seed=3) + grad_out = torch.randn(4, 2688, device="cuda") + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + out = op(*leaves) + out.backward(grad_out) + runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_dx_rows_are_batch_invariant(self): + """Keep feature gradients bitwise unchanged when rows are evaluated independently.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=4) + x = _features(6, seed=4) + grad_out = torch.randn(6, 2688, device="cuda") + + def dx(rows): + """Compute selected feature-row gradients using matching upstream rows.""" + + leaf = x[rows].detach().clone().requires_grad_(True) + op(leaf, *params).backward(grad_out[rows]) + return leaf.grad + + full = dx(slice(0, 6)) + for i in range(6): + assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + + def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") + assert type(op).__name__ == "H3TimestepMLPCudaOp" + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +class TestCudaDetLinearDeviceContract: + def test_forward_rejects_bias_on_another_device(self, dtype): + """Reject a projection bias on a different CUDA device from its weight.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_forward + + _cuda_op() + x = torch.randn(2, 8, device="cuda:0", dtype=dtype) + weight = torch.randn(16, 8, device="cuda:0", dtype=dtype) + bias = torch.randn(16, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="bias and weight must be on the same device"): + det_linear_forward(x, weight, bias) + + def test_backward_input_rejects_weight_on_another_device(self, dtype): + """Reject native input-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_input, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + weight = torch.randn(16, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and weight must be on the same device"): + det_linear_backward_input(grad, weight, dtype) + + def test_backward_weight_rejects_input_on_another_device(self, dtype): + """Reject native parameter-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_weight, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + x = torch.randn(2, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and x must be on the same device"): + det_linear_backward_weight(grad, x, dtype) diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..277beadea --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_sinusoid_h3``: FP32 [cos | sin] timestep features. + +Three properties, each checked separately: + +* provider parity: the CUDA kernel is bitwise equal to the diffusers + ``get_timestep_embedding`` replay on the same device (an elementwise op + with the same FP32 operation sequence); +* accuracy: within the FP32 elementwise contract of an FP64 golden; +* invariance: a row's bytes do not depend on how many timesteps share the + launch, their order, or repeated runs. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps +from rl_engine.validation.models.h3_provider import provider_time_proj + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + """Construct the CUDA sinusoid operator, skipping a missing native forward entrypoint.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + h3_sinusoid_cuda_available, + ) + + if not h3_sinusoid_cuda_available(): + pytest.skip("rl_engine._C lacks h3_timestep_sinusoid_forward") + return H3TimestepSinusoidCudaOp() + + +def _dense_unit_sweep() -> torch.Tensor: + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + edges = torch.tensor([0x00000001, 0x007FFFFF, 0x00800000, 0x3F800000], dtype=torch.int32) + return torch.cat([bits, edges.cuda()]).view(torch.float32) + + +class TestReference: + def test_layout_is_cos_then_sin(self): + """Check FP32 cosine-then-sine layout and the unit-frequency boundary values.""" + + t = torch.tensor([0.0, 1.0]) + out = NativeH3TimestepSinusoidOp().forward(t) + assert out.shape == (2, 256) and out.dtype == torch.float32 + assert torch.equal(out[0, :128], torch.ones(128)) + assert torch.equal(out[0, 128:], torch.zeros(128)) + # k = 0 has frequency 1: cos(1), sin(1) at t = 1. + assert out[1, 0].item() == pytest.approx(math.cos(1.0), abs=1e-7) + assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) + + def test_provider_path_matches_diffusers_cpu(self): + """Require CPU sinusoid values to match the diffusers provider bitwise.""" + + t = h3_timesteps(9, device="cpu") + ours = NativeH3TimestepSinusoidOp().forward(t) + assert torch.equal(ours, provider_time_proj(t, 256)) + + def test_golden_close_to_provider_path(self): + """Keep the CPU provider within FP32 error of the high-precision sinusoid golden.""" + + t = h3_timesteps(33, device="cpu") + op = NativeH3TimestepSinusoidOp() + torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) + + @pytest.mark.parametrize( + "bad, error", + [ + (torch.tensor([[0.5]]), ValueError), # not 1-D + (torch.empty(0), ValueError), # empty + (torch.tensor([1, 0]), TypeError), # integer + (torch.tensor([0.5, 1.5]), ValueError), # H10: t outside [0, 1] + (torch.tensor([500.0]), ValueError), # H10: t * 1000 convention + (torch.tensor([float("nan")]), ValueError), + ], + ) + def test_rejects_invalid_timesteps(self, bad, error): + """Reject empty, malformed, non-floating, non-finite, or out-of-range timesteps.""" + + with pytest.raises(error): + NativeH3TimestepSinusoidOp().forward(bad) + + def test_rejects_odd_channels(self): + """Reject odd feature counts that cannot split into equal cosine and sine halves.""" + + with pytest.raises(ValueError): + NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) + def test_bitwise_equal_to_diffusers_cuda_path(self, num): + """Require CUDA sinusoid values to match both provider replays for varied batch sizes.""" + + t = h3_timesteps(num, seed=num) + out = _cuda_op()(t) + assert out.dtype == torch.float32 and out.shape == (num, 256) + assert torch.equal(out, provider_time_proj(t, 256)) + assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + + def test_bitwise_over_dense_fp32_sweep_of_unit_interval(self): + # Every 997th FP32 bit pattern in [0, 1] plus the subnormals' edge: a + # fast-math build (approximate __expf/__sinf/__cosf, FTZ) fails here. + t = _dense_unit_sweep() + assert torch.equal(_cuda_op().forward(t, check_range=False), provider_time_proj(t, 256)) + + def test_sglang_timestep_embedding_is_bitwise_equal(self): + # RFC #420 reuse rule: parity with SGLang is checked bit for bit, not + # through matching FP64 error summaries. + jit = pytest.importorskip("sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit") + t = torch.cat([_dense_unit_sweep(), h3_timesteps(4097, seed=3)]) + theirs = jit.timestep_embedding(t, 256, flip_sin_to_cos=True, downscale_freq_shift=0.0) + assert torch.equal(theirs, provider_time_proj(t, 256)) + assert torch.equal(theirs, _cuda_op().forward(t, check_range=False)) + + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) + def test_other_channel_counts_match_reference(self, num_channels): + """Preserve bitwise reference parity across supported even channel counts.""" + + t = h3_timesteps(5) + out = _cuda_op()(t, num_channels=num_channels) + ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) + assert torch.equal(out, ref) + + def test_fp32_contract_against_fp64_golden(self): + """Bound CUDA sinusoid error against the high-precision golden by FP32 precision.""" + + t = h3_timesteps(257) + out = _cuda_op()(t) + gold = NativeH3TimestepSinusoidOp().forward_fp32(t) + # tolerance_contract.json forward_accuracy / elementwise / float32. + torch.testing.assert_close(out, gold, atol=1e-5, rtol=1e-5) + assert (out - gold).abs().max().item() <= 2.0**-23 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_low_precision_timesteps_are_upcast_first(self, dtype): + """Match provider FP32 upcasting for BF16 and FP16 timesteps.""" + + t = h3_timesteps(6).to(dtype) + assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) + + def test_batch_position_and_repeat_invariance(self): + """Keep row bits unchanged across batch positions, repeated calls, and permutations.""" + + op = _cuda_op() + base = torch.tensor(0.3141592, device="cuda") + single = op(base.view(1)) + for num in (2, 5, 64, 1023): + for pos in (0, num // 2, num - 1): + t = h3_timesteps(num, seed=num) + t[pos] = base + assert torch.equal(op(t)[pos], single[0]), (num, pos) + t = h3_timesteps(129, seed=11) + first = op(t) + for _ in range(3): + assert torch.equal(op(t), first) + perm = torch.randperm(129, generator=torch.Generator().manual_seed(0)).cuda() + assert torch.equal(op(t[perm]), first[perm]) + + def test_non_contiguous_input(self): + """Treat strided timesteps identically to their contiguous copies.""" + + t = h3_timesteps(20) + strided = t[::2] + assert not strided.is_contiguous() + assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) + + def test_rejects_cpu_input(self): + """Reject CPU timesteps at the CUDA sinusoid boundary.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.5])) + + def test_rejects_out_of_range(self): + """Reject scaled timesteps outside the declared unit interval.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + + @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) + def test_native_entrypoint_rejects_invalid_values(self, bad): + """Reject invalid values in native and wrapper calls even when neighbors are valid.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + # Valid neighbors must not hide an invalid timestep in the native call. + t = torch.tensor([0.0, bad, 1.0], device="cuda") + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _C.h3_timestep_sinusoid_forward(t) + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _cuda_op()(t) + + def test_native_entrypoint_accepts_boundaries(self): + """Accept both interval endpoints and match provider values in the native call.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + t = torch.tensor([0.0, 0.5, 1.0], device="cuda") + assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) + + def test_trusted_path_supports_cuda_graph(self): + """Capture and replay the prevalidated sinusoid path without changing its output.""" + + op = _cuda_op() + t = h3_timesteps(7) + expected = op(t) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op.forward(t, check_range=False) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op.forward(t, check_range=False) + graph.replay() + assert torch.equal(out, expected) + + def test_backward_matches_autograd_reference(self): + """Match timestep gradients to the high-precision autograd sinusoid replay.""" + + t = h3_timesteps(17).requires_grad_(True) + t_ref = t.detach().clone().requires_grad_(True) + grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) + torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + + def test_backward_is_row_invariant(self): + """Keep a timestep gradient bitwise unchanged when unrelated rows join the batch.""" + + op = _cuda_op() + grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) + + def grad_of(t, pos): + """Differentiate one selected timestep using a fixed upstream feature-gradient row.""" + + t = t.clone().requires_grad_(True) + grad = torch.zeros(t.shape[0], 256, device="cuda") + grad[pos] = grad_row + op(t).backward(grad) + return t.grad[pos] + + single = grad_of(torch.tensor([0.7], device="cuda"), 0) + t = h3_timesteps(40) + t[13] = 0.7 + assert torch.equal(grad_of(t, 13), single) + + def test_registry_dispatches_cuda(self): + """Resolve the sinusoid registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") + assert type(op).__name__ == "H3TimestepSinusoidCudaOp" diff --git a/tests/models/minimax_h3/test_h3_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..c9b511153 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import adaln_column_shard +from rl_engine.validation.models.h3_ws2 import ( + run_world, + tp_matches_ws1, + tp_projection_rank, + ws1_projection, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8, 9]) + def test_shards_tile_the_table_once(self, tp): + shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] + assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + assert tp_matches_ws1(ws1, ranks) == ALL_TRUE + assert [r["readback"]["rank"] for r in ranks] == list(range(world)) + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@pytest.mark.parametrize("n_total", [0, -1152]) +def test_rejects_nonpositive_projection(n_total): + with pytest.raises(ValueError, match="N must be positive"): + adaln_column_shard(n_total, 1, 0) diff --git a/tests/models/minimax_h3/test_h3_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_weights.py b/tests/models/minimax_h3/test_h3_weights.py new file mode 100644 index 000000000..13b9b2489 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weights.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Extracted H3 weights must retain the pinned tensor contents.""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.validation.models import h3_weights + + +@pytest.mark.parametrize( + "dtype,raw_bytes", + [ + (torch.float32, bytes.fromhex("000000000000803f")), + (torch.bfloat16, bytes.fromhex("0000803f")), + ], +) +def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + """Hash the exact FP32 and BF16 bit patterns without numerical conversion.""" + + tensor = torch.tensor([0.0, 1.0], dtype=dtype) + assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +@pytest.mark.parametrize("names", [None, ["bias"]]) +def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + + from safetensors.torch import save_file + + tensors = { + "weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=dtype), + "bias": torch.tensor([0.0, 1.0], dtype=dtype), + } + manifest = { + "tensors": { + name: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "sha256": h3_weights.sha256_tensor(tensor), + } + for name, tensor in tensors.items() + } + } + if mismatch == "dtype": + tensors["weight"] = tensors["weight"].to(torch.float64) + elif mismatch == "shape": + tensors["weight"] = tensors["weight"].reshape(4) + elif mismatch == "sha256": + tensors["weight"][0, 0] = 1.0 + # Tensor checksums do not depend on serializer metadata or tensor key order. + save_file( + dict(reversed(list(tensors.items()))), + str(tmp_path / h3_weights.EXTRACTED_FILE), + metadata={"revision": "a different serializer metadata value"}, + ) + monkeypatch.setenv(h3_weights.WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(h3_weights, "load_h3_manifest", lambda: manifest) + + if mismatch: + with pytest.raises(ValueError, match=f"weight: {mismatch} .* != manifest"): + h3_weights.load_h3_conditioning_weights("cpu", names=names) + else: + actual = h3_weights.load_h3_conditioning_weights("cpu", names=names) + assert list(actual) == (names if names is not None else list(tensors)) + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/tests/models/minimax_h3/test_h3_ws2_processes.py b/tests/models/minimax_h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..f448ba580 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.validation.models import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert [r["rank"] for r in ranks] == [0, 1] + assert [r["value"].item() for r in ranks] == [0, 1] + elif mode == "timeout": + with pytest.raises(TimeoutError, match="timed out"): + h3_ws2.run_world(2, None, {}, timeout=0.2, mode=mode) + else: + message = "original rank traceback" if mode == "error" else "rank 0 exited with code 7" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before diff --git a/tests/models/minimax_h3/test_prepare_h3_weights.py b/tests/models/minimax_h3/test_prepare_h3_weights.py new file mode 100644 index 000000000..5c9a966b9 --- /dev/null +++ b/tests/models/minimax_h3/test_prepare_h3_weights.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Weight extraction rejects off-manifest tensors even under Python optimization.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.validation.models.h3_weights import EXTRACTED_FILE, sha256_file, sha256_tensor + + +@pytest.mark.parametrize("optimize", [0, 1]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + + import huggingface_hub + from safetensors.torch import load_file, save_file + + tensor = torch.tensor([0.0, 1.0]) + config = tmp_path / "config.json" + config.write_text("{}") + index = tmp_path / "index.json" + index.write_text(json.dumps({"weight_map": {"weight": "shard.safetensors"}})) + shard = tmp_path / "shard.safetensors" + save_file({"weight": tensor}, str(shard)) + manifest = { + "model_identity": { + "hf_repo": "test/h3", + "subfolder": "dit", + "revision": "pinned", + "config_sha256": sha256_file(config), + "index_file": index.name, + "index_sha256": sha256_file(index), + }, + "weight_shards": {shard.name: {"sha256": sha256_file(shard)}}, + "tensors": { + "weight": { + "dtype": "bfloat16" if mismatch == "dtype" else "float32", + "shape": [3] if mismatch == "shape" else [2], + "sha256": "0" * 64 if mismatch == "sha256" else sha256_tensor(tensor), + } + }, + } + monkeypatch.setattr( + huggingface_hub, + "hf_hub_download", + lambda repo, filename, **kwargs: str(tmp_path / Path(filename).name), + ) + out_dir = tmp_path / "extracted" + monkeypatch.setattr(sys, "argv", ["prepare_h3_weights.py", "--out", str(out_dir)]) + script = Path(__file__).resolve().parents[3] / "tools" / "weights" / "prepare_h3_weights.py" + namespace = {"__name__": "prepare_h3_weights_test", "__file__": str(script)} + # Compiling the actual script at optimize=1 reproduces python -O's assert removal. + exec(compile(script.read_text(), str(script), "exec", optimize=optimize), namespace) + namespace["load_h3_manifest"] = lambda: manifest + target = out_dir / EXTRACTED_FILE + if mismatch: + with pytest.raises(SystemExit, match=f"weight: {mismatch} .* != manifest"): + namespace["main"]() + assert not target.exists() + else: + namespace["main"]() + assert torch.equal(load_file(str(target))["weight"], tensor) diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 3c47bbbcf..647053ce1 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -4,6 +4,7 @@ import sys from types import ModuleType +import pytest import torch import rl_engine.platforms.device as device_module @@ -231,6 +232,44 @@ def fake_load_backend(backend): ] +H3_OPS = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", +) + + +@pytest.mark.parametrize("op_name", H3_OPS) +def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + """Prefer H3 CUDA backends on CUDA and use their PyTorch fallbacks elsewhere.""" + + registry = KernelRegistry() + cuda_chain = registry._priority_map["cuda"][op_name] + assert cuda_chain[0].name.startswith("CUDA_H3_") + assert cuda_chain[-1].name.startswith("PYTORCH_H3_") + for platform in ("rocm", "musa", "cpu", "npu"): + chain = registry._priority_map[platform][op_name] + assert [backend.name for backend in chain] == [cuda_chain[-1].name] + + +def test_h3_timestep_falls_back_when_fast_math_build_omits_kernel(monkeypatch): + # KERNEL_ALIGN_USE_FAST_MATH=1 builds leave out the bitwise timestep kernel + # and its binding; dispatch must then pick the PyTorch reference. + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_sinusoid + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + monkeypatch.setattr(timestep_sinusoid, "h3_sinusoid_cuda_available", lambda: False) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda:0") + + assert isinstance(op, NativeH3TimestepSinusoidOp) + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod diff --git a/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index bcaebc2f5..e7d1e5275 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/test_ws1_gtest_gpu.py @@ -33,6 +33,8 @@ def _run(script: str, *args: str, timeout: int = 300) -> None: def test_all_ws1_single_ops_are_registered(): + """Require gtest specs for all WS1 single operators and the H3 conditioning stages.""" + names = set(operator_names()) assert { "rms_norm", @@ -49,6 +51,13 @@ def test_all_ws1_single_ops_are_registered(): "pack", "linear_logp", "prefix_shared_attention", + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", } <= names diff --git a/tools/migration/README.md b/tools/migration/README.md index 4a6cee300..48491e383 100644 --- a/tools/migration/README.md +++ b/tools/migration/README.md @@ -8,3 +8,6 @@ Historical workload/operator identifiers and archived evidence intentionally retain their original values. New source imports should use the mapped canonical module. Remove compatibility paths only in a separately announced migration once downstream launchers, patches, imports and saved objects have migrated. + +The MiniMax-H3 entries also record the open RFC #420 PR stack migration to +`test-h3`; those H3 files were not part of the original main baseline. diff --git a/tools/migration/layout.json b/tools/migration/layout.json index cb8fe469a..6c7fc9898 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -347,6 +347,7 @@ "benchmarks/benchmark_deterministic_attention.py": "benchmarks/operators/attention/benchmark_deterministic_attention.py", "benchmarks/benchmark_grpo_loss.py": "benchmarks/operators/loss/benchmark_grpo_loss.py", "benchmarks/benchmark_grpo_op.py": "benchmarks/operators/loss/benchmark_grpo_op.py", + "benchmarks/benchmark_h3_conditioning.py": "benchmarks/models/benchmark_h3_conditioning.py", "benchmarks/benchmark_linear_logp.py": "benchmarks/operators/logprob/benchmark_linear_logp.py", "benchmarks/benchmark_pack.py": "benchmarks/operators/packing/benchmark_pack.py", "benchmarks/benchmark_qwen_ffn_layout.py": "benchmarks/layers/benchmark_qwen_ffn_layout.py", @@ -466,6 +467,42 @@ "docs/design/ws2-attention-single-gpu-harness.md": "docs/archive/design/ws2-attention-single-gpu-harness.md", "docs/design/ws2-cp-attention-contract.md": "docs/contracts/ws2-cp-attention-contract.md", "docs/design/ws2-module-debug-matrix.md": "docs/archive/design/ws2-module-debug-matrix.md", + "docs/usage/evidence/h3-adaln-gate-residual-b200/figure.png": "reports/experiments/h3-adaln-gate-residual-b200/figure.png", + "docs/usage/evidence/h3-adaln-gate-residual-b200/report.json": "reports/experiments/h3-adaln-gate-residual-b200/report.json", + "docs/usage/evidence/h3-adaln-projection-b200/figure.png": "reports/experiments/h3-adaln-projection-b200/figure.png", + "docs/usage/evidence/h3-adaln-projection-b200/report.json": "reports/experiments/h3-adaln-projection-b200/report.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/chain_replay.json": "reports/experiments/h3-adaln-row-gather-b200/chain_replay.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/figure.png": "reports/experiments/h3-adaln-row-gather-b200/figure.png", + "docs/usage/evidence/h3-adaln-row-gather-b200/report.json": "reports/experiments/h3-adaln-row-gather-b200/report.json", + "docs/usage/evidence/h3-final-adaln-out-b200/chain_replay.json": "reports/experiments/h3-final-adaln-out-b200/chain_replay.json", + "docs/usage/evidence/h3-final-adaln-out-b200/figure.png": "reports/experiments/h3-final-adaln-out-b200/figure.png", + "docs/usage/evidence/h3-final-adaln-out-b200/report.json": "reports/experiments/h3-final-adaln-out-b200/report.json", + "docs/usage/evidence/h3-one-block-b200/figure.png": "reports/experiments/h3-one-block-b200/figure.png", + "docs/usage/evidence/h3-one-block-b200/report.json": "reports/experiments/h3-one-block-b200/report.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json": "reports/experiments/h3-prior-art-b200/adaln_row_gather.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.png": "reports/experiments/h3-prior-art-b200/adaln_row_gather.png", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json": "reports/experiments/h3-prior-art-b200/final_adaln_out.json", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.png": "reports/experiments/h3-prior-art-b200/final_adaln_out.png", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.json": "reports/experiments/h3-prior-art-b200/gate_residual.json", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.png": "reports/experiments/h3-prior-art-b200/gate_residual.png", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.json": "reports/experiments/h3-prior-art-b200/norm_modulate.json", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.png": "reports/experiments/h3-prior-art-b200/norm_modulate.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json": "reports/experiments/h3-prior-art-b200/timestep_mlp.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.png": "reports/experiments/h3-prior-art-b200/timestep_mlp.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.png": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.png", + "docs/usage/evidence/h3-rmsnorm-b200/figure.png": "reports/experiments/h3-rmsnorm-b200/figure.png", + "docs/usage/evidence/h3-rmsnorm-b200/report.json": "reports/experiments/h3-rmsnorm-b200/report.json", + "docs/usage/evidence/h3-sp-norm-adaln-b200/figure.png": "reports/experiments/h3-sp-norm-adaln-b200/figure.png", + "docs/usage/evidence/h3-sp-norm-adaln-b200/report.json": "reports/experiments/h3-sp-norm-adaln-b200/report.json", + "docs/usage/evidence/h3-timestep-mlp-b200/figure.png": "reports/experiments/h3-timestep-mlp-b200/figure.png", + "docs/usage/evidence/h3-timestep-mlp-b200/report.json": "reports/experiments/h3-timestep-mlp-b200/report.json", + "docs/usage/evidence/h3-timestep-sinusoid-b200/figure.png": "reports/experiments/h3-timestep-sinusoid-b200/figure.png", + "docs/usage/evidence/h3-timestep-sinusoid-b200/report.json": "reports/experiments/h3-timestep-sinusoid-b200/report.json", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/figure.png": "reports/experiments/h3-tp-adaln-3mod-b200/figure.png", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json": "reports/experiments/h3-tp-adaln-3mod-b200/report.json", "envs.py": "build_tools/envs.py", "examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json": "configs/experiments/cross_config/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json", "examples/cross_config_s0_cpu_smoke.json": "configs/experiments/cross_config/cross_config_s0_cpu_smoke.json", @@ -675,6 +712,19 @@ "rl_engine/kernels/ops/cuda/attention/prefix_shared_attn.py": "rl_engine/backends/cuda/attention/prefix_shared_attn.py", "rl_engine/kernels/ops/cuda/attention/strict_runtime.py": "rl_engine/backends/cuda/attention/strict_runtime.py", "rl_engine/kernels/ops/cuda/ffn.py": "rl_engine/backends/cuda/ffn/ffn.py", + "rl_engine/kernels/ops/cuda/h3/__init__.py": "rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py", + "rl_engine/kernels/ops/cuda/h3/adaln_modulation.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py", + "rl_engine/kernels/ops/cuda/h3/adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/cuda/h3/det_linear.py": "rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py", + "rl_engine/kernels/ops/cuda/h3/final_adaln_out.py": "rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py", + "rl_engine/kernels/ops/cuda/h3/gate_residual.py": "rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py", + "rl_engine/kernels/ops/cuda/h3/rmsnorm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py": "rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py", + "rl_engine/kernels/ops/cuda/h3/timestep_mlp.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py", + "rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py", + "rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py", + "rl_engine/kernels/ops/cuda/h3/ws2_comm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py", "rl_engine/kernels/ops/cuda/linear/__init__.py": "rl_engine/backends/cuda/embedding/__init__.py", "rl_engine/kernels/ops/cuda/linear/embedding.py": "rl_engine/backends/cuda/embedding/embedding.py", "rl_engine/kernels/ops/cuda/linear/lm_head.py": "rl_engine/backends/cuda/embedding/lm_head.py", @@ -705,6 +755,15 @@ "rl_engine/kernels/ops/pytorch/attention/stateful_kv.py": "rl_engine/reference/attention/stateful_kv.py", "rl_engine/kernels/ops/pytorch/ffn/__init__.py": "rl_engine/ops/ffn/__init__.py", "rl_engine/kernels/ops/pytorch/ffn/ffn.py": "rl_engine/ops/ffn/qwen3.py", + "rl_engine/kernels/ops/pytorch/h3/__init__.py": "rl_engine/reference/minimax_h3/__init__.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_projection.py": "rl_engine/reference/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py": "rl_engine/reference/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py": "rl_engine/reference/minimax_h3/final_adaln_out.py", + "rl_engine/kernels/ops/pytorch/h3/fixed_order.py": "rl_engine/reference/minimax_h3/fixed_order.py", + "rl_engine/kernels/ops/pytorch/h3/gate_residual.py": "rl_engine/reference/minimax_h3/gate_residual.py", + "rl_engine/kernels/ops/pytorch/h3/rmsnorm.py": "rl_engine/reference/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py": "rl_engine/reference/minimax_h3/timestep_mlp.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py": "rl_engine/reference/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/pytorch/linear/__init__.py": "rl_engine/reference/embedding/__init__.py", "rl_engine/kernels/ops/pytorch/linear/embedding.py": "rl_engine/reference/embedding/embedding.py", "rl_engine/kernels/ops/pytorch/linear/lm_head.py": "rl_engine/reference/embedding/lm_head.py", @@ -776,6 +835,13 @@ "rl_engine/testing/__init__.py": "rl_engine/validation/common/__init__.py", "rl_engine/testing/attention_comparison.py": "rl_engine/validation/common/attention_comparison.py", "rl_engine/testing/distributed_logprob_comparison.py": "rl_engine/validation/common/distributed_logprob_comparison.py", + "rl_engine/testing/h3_cases.py": "rl_engine/validation/models/h3_cases.py", + "rl_engine/testing/h3_chain.py": "rl_engine/validation/models/h3_chain.py", + "rl_engine/testing/h3_manifest.json": "rl_engine/validation/models/h3_manifest.json", + "rl_engine/testing/h3_provider.py": "rl_engine/validation/models/h3_provider.py", + "rl_engine/testing/h3_report.py": "rl_engine/validation/models/h3_report.py", + "rl_engine/testing/h3_weights.py": "rl_engine/validation/models/h3_weights.py", + "rl_engine/testing/h3_ws2.py": "rl_engine/validation/models/h3_ws2.py", "rl_engine/testing/logprob_comparison.py": "rl_engine/validation/common/logprob_comparison.py", "rl_engine/testing/logprob_drift.py": "rl_engine/validation/common/logprob_drift.py", "rl_engine/testing/reference_ops.py": "rl_engine/reference/definitions.py", @@ -790,6 +856,13 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", + "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", + "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/h3_ws2_evidence.py": "tools/validation/models/h3_ws2_evidence.py", + "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", + "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", + "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", "scripts/prepare_ws1_weights.py": "tools/weights/prepare_ws1_weights.py", "scripts/run_perf.py": "tools/benchmarking/run_perf.py", "scripts/run_profile_suite.py": "tools/benchmarking/run_profile_suite.py", @@ -820,6 +893,27 @@ "tests/distributed/test_rocm_attention_transport.py": "tests/distributed/cp/test_rocm_attention_transport.py", "tests/distributed/test_rocm_strict_attention_cp.py": "tests/distributed/cp/test_rocm_strict_attention_cp.py", "tests/distributed/test_transport_deterministic_collective.py": "tests/distributed/collectives/test_transport_deterministic_collective.py", + "tests/h3/conftest.py": "tests/models/minimax_h3/conftest.py", + "tests/h3/test_h3_adaln_gate_residual.py": "tests/models/minimax_h3/test_h3_adaln_gate_residual.py", + "tests/h3/test_h3_adaln_modulation.py": "tests/models/minimax_h3/test_h3_adaln_modulation.py", + "tests/h3/test_h3_adaln_projection.py": "tests/models/minimax_h3/test_h3_adaln_projection.py", + "tests/h3/test_h3_adaln_row_gather.py": "tests/models/minimax_h3/test_h3_adaln_row_gather.py", + "tests/h3/test_h3_benchmark_timing.py": "tests/models/minimax_h3/test_h3_benchmark_timing.py", + "tests/h3/test_h3_chain_replay_cli.py": "tests/models/minimax_h3/test_h3_chain_replay_cli.py", + "tests/h3/test_h3_cli.py": "tests/models/minimax_h3/test_h3_cli.py", + "tests/h3/test_h3_conditioning_e2e.py": "tests/models/minimax_h3/test_h3_conditioning_e2e.py", + "tests/h3/test_h3_det_linear.py": "tests/models/minimax_h3/test_h3_det_linear.py", + "tests/h3/test_h3_final_adaln_out.py": "tests/models/minimax_h3/test_h3_final_adaln_out.py", + "tests/h3/test_h3_native_validation.py": "tests/models/minimax_h3/test_h3_native_validation.py", + "tests/h3/test_h3_report.py": "tests/models/minimax_h3/test_h3_report.py", + "tests/h3/test_h3_rmsnorm.py": "tests/models/minimax_h3/test_h3_rmsnorm.py", + "tests/h3/test_h3_sp_norm_adaln.py": "tests/models/minimax_h3/test_h3_sp_norm_adaln.py", + "tests/h3/test_h3_timestep_mlp.py": "tests/models/minimax_h3/test_h3_timestep_mlp.py", + "tests/h3/test_h3_timestep_sinusoid.py": "tests/models/minimax_h3/test_h3_timestep_sinusoid.py", + "tests/h3/test_h3_tp_adaln.py": "tests/models/minimax_h3/test_h3_tp_adaln.py", + "tests/h3/test_h3_weights.py": "tests/models/minimax_h3/test_h3_weights.py", + "tests/h3/test_h3_ws2_processes.py": "tests/models/minimax_h3/test_h3_ws2_processes.py", + "tests/h3/test_prepare_h3_weights.py": "tests/models/minimax_h3/test_prepare_h3_weights.py", "tests/linear_logp_tp.py": "tests/distributed/tp/linear_logp_tp.py", "tests/test_alignment_model_wrappers.py": "tests/models/qwen3/test_alignment_model_wrappers.py", "tests/test_alignment_wrapper_interfaces.py": "tests/models/qwen3/test_alignment_wrapper_interfaces.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py new file mode 100644 index 000000000..d43e5c70e --- /dev/null +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct bitwise comparison of one RFC #420 operator with its reuse candidates. + +``h3_prior_art.py`` reports each implementation's batch invariance and its error against FP64. +Matching error summaries do not show that two implementations produce the same bits, so before +a candidate is adopted this runner compares outputs directly. Every candidate that +``h3_prior_art.candidates()`` defines runs on the same inputs (each global batch-invariant mode +in its own process, as in ``h3_prior_art.py``), and every forward output and leaf gradient is +compared byte for byte with diffusers' path and with the rl-kernel op: + + python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \\ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \\ + [--megatron-src /path/to/Megatron-LM] [--fast-math-variant] + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png + +For ``timestep_sinusoid`` the inputs are three timestep sets: 2^20 uniform draws, every 997th +FP32 bit pattern in [0, 1] (a dense sweep, mostly tiny values), and a 100001-point grid. +``--fast-math-variant`` also compiles ``csrc/cuda/h3/timestep_sinusoid.cu`` with +``--use_fast_math`` and compares it, to show what the fast-math build guard prevents. Other +operators use each candidate's own input generator at four sizes. + +Intermediate tensors are written next to the report (not to /tmp, which cluster job epilogs +can clear) and removed afterwards. Run on a clean tree; the report records the commit, the tree +state and every library version. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import h3_prior_art as PA +import torch + +SEED = 7 +SIZES = { + "timestep_mlp": (1, 4, 64, 2048), + "adaln_projection": (1, 4, 64, 2048), + "adaln_row_gather": (1, 257, 4097, 32768), + "norm_modulate": (1, 257, 4097, 32768), + "gate_residual": (1, 257, 4097, 32768), + "final_adaln_out": (1, 257, 4097, 32768), +} +SINUSOID_SETS = ("random_2^20", "dense_fp32_sweep", "grid_100001") +_FAST_MATH_BINDING = """ +#include +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor, int64_t, double, bool); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &h3_timestep_sinusoid_forward); } +""" + + +def _sinusoid_inputs(name: str) -> dict[str, torch.Tensor]: + if name == "random_2^20": + t = torch.rand(1 << 20, device="cuda", generator=torch.Generator("cuda").manual_seed(0)) + elif name == "dense_fp32_sweep": + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + t = bits.view(torch.float32) + else: + t = torch.linspace(0, 1, 100001, device="cuda") + return {"t": t} + + +def _fast_math_sinusoid(build_dir: Path): + from torch.utils.cpp_extension import load + + build_dir.mkdir(parents=True, exist_ok=True) + binding = build_dir / "binding.cpp" + binding.write_text(_FAST_MATH_BINDING) + root = Path(__file__).resolve().parents[3] + mod = load( + name="h3_timestep_sinusoid_fast_math", + sources=[str(root / "csrc/cuda/h3/timestep_sinusoid.cu"), str(binding)], + build_directory=str(build_dir), + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + return { + "source": "rl-kernel timestep_sinusoid.cu compiled with --use_fast_math", + "grads": [], + "fn": lambda t: mod.forward(t, PA.F1, 10000.0, False), + } + + +def _run(cand: dict[str, Any], inputs: dict[str, torch.Tensor]): + args = {} + for k, v in inputs.items(): + args[k] = v.detach().clone().requires_grad_(True) if k in cand["grads"] else v + out = cand["fn"](**args) + outs = list(out) if isinstance(out, (tuple, list)) else [out] + live = [o for o in outs if o.requires_grad] + if live: + gen = torch.Generator("cuda").manual_seed(SEED) + ups = [torch.randn(o.shape, device="cuda", generator=gen).to(o.dtype) for o in live] + torch.autograd.backward(live, ups) + grads = {k: args[k].grad.cpu() if args[k].grad is not None else None for k in cand["grads"]} + return [o.detach().cpu() for o in outs], grads + + +def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: bool) -> None: + """Child process: run every candidate of one mode and save its outputs.""" + + try: + PA.enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + (parts / f"{mode}.unavailable").write_text(f"{type(exc).__name__}: {exc}"[:300]) + return + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cands = PA.candidates(op, mode) + if fast_math and op == "timestep_sinusoid" and mode == "plain": + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + for cand in cands: + name = cand.pop("name") + if "unavailable" in cand: + (parts / f"{name}.unavailable").write_text(cand["unavailable"][:300]) + continue + cases = SINUSOID_SETS if op == "timestep_sinusoid" else SIZES[op] + for case in cases: + inputs = ( + _sinusoid_inputs(case) if op == "timestep_sinusoid" else cand["make"](SEED, case) + ) + outs, grads = _run(cand, inputs) + torch.save( + {"source": cand["source"], "out": outs, "grads": grads}, + parts / f"{name}__{case}.pt", + ) + print(f" [{mode}] {name}: done", flush=True) + torch.cuda.empty_cache() + + +def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | None: + if a is None or b is None: + return None + if a.shape != b.shape or a.dtype != b.dtype: + return {"equal": False, "note": f"{a.dtype}{tuple(a.shape)} vs {b.dtype}{tuple(b.shape)}"} + width = a.element_size() + differ = ( + a.contiguous().view(torch.uint8).view(-1, width) + != b.contiguous().view(torch.uint8).view(-1, width) + ).any(1) + return { + "equal": not bool(differ.any()), + "differ": int(differ.sum()), + "numel": a.numel(), + "max_abs": float((a.double() - b.double()).abs().max()), + } + + +def compare(parts: Path) -> dict[str, Any]: + runs: dict[str, dict[str, Any]] = {} + for f in sorted(parts.glob("*.pt")): + name, case = f.stem.rsplit("__", 1) + runs.setdefault(name, {})[case] = torch.load(f) + refs = [r for r in ("diffusers[plain]", "diffusers", "rl_kernel") if r in runs] + out: dict[str, Any] = {} + for name, by_case in runs.items(): + entry: dict[str, Any] = {"source": next(iter(by_case.values()))["source"]} + for ref in refs: + if ref == name: + continue + entry[f"vs {ref}"] = { + case: { + "out": [_same(a, b) for a, b in zip(run["out"], runs[ref][case]["out"])], + "grads": { + g: _same(run["grads"].get(g), runs[ref][case]["grads"].get(g)) + for g in run["grads"] + }, + } + for case, run in by_case.items() + } + out[name] = entry + for f in parts.glob("*.unavailable"): + out[f.stem] = {"unavailable": f.read_text()} + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(PA.OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument( + "--fast-math-variant", + action="store_true", + help="timestep_sinusoid only: also compare a --use_fast_math build of the kernel", + ) + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(PA.OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {PA.OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + parts = args.out.with_name(f".{args.out.stem}.parts") + + if args.mode is not None: # child process: one global mode + dump(args.op, args.mode, parts, args.megatron_src, args.fast_math_variant) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + shutil.rmtree(parts, ignore_errors=True) + parts.mkdir(parents=True) + modes = list(PA.MODES) if args.op in PA.GEMM_OPS else ["plain"] + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(args.out), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.fast_math_variant: + cmd.append("--fast-math-variant") + proc = subprocess.run(cmd, env={**os.environ, **PA.MODES[mode]}) + if proc.returncode != 0: + (parts / f"{mode}.unavailable").write_text(f"subprocess exited with {proc.returncode}") + results = compare(parts) + shutil.rmtree(parts, ignore_errors=True) + report = { + "kind": "h3_bitwise_reuse", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: PA._version(lib) for lib in PA.LIBS}, + "mode_environment": {m: PA.MODES[m] for m in modes}, + }, + "cases": list(SINUSOID_SETS if args.op == "timestep_sinusoid" else SIZES[args.op]), + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_diagnostic.py b/tools/validation/models/h3_block_diagnostic.py new file mode 100644 index 000000000..19e5b111a --- /dev/null +++ b/tools/validation/models/h3_block_diagnostic.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Localize a failed block-0 gradient gate without changing its numerical contract.""" +import argparse +import json +import os +import subprocess +from dataclasses import replace +from pathlib import Path + +import torch + +from rl_engine.backends.extension import _C +from rl_engine.validation.models import h3_block as block +from rl_engine.validation.models.h3_weights import load_h3_block_weights +from tools.validation.models.h3_full_model import digest + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--failed-job", required=True) + parser.add_argument("--causal-sweep", action="store_true") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=False) + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("diagnostic requires clean source") + ops = block.CandidateOps() + weights = load_h3_block_weights("cuda") + ctx = block.make_block_context(weights, layout="tiny", seed=0, batch=1) + upstream = ( + torch.randn((1, ctx["seq_len"], block.HIDDEN), generator=torch.Generator().manual_seed(11)) + .bfloat16() + .cuda() + ) + report = dict( + status="diagnostic_only_not_acceptance", + failed_job=args.failed_job, + job_id=os.environ.get("SLURM_JOB_ID"), + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + extension_sha256=digest(Path(_C.__file__)), + torch=torch.__version__, + device=torch.cuda.get_device_name(), + backends=ops.backends(), + default_matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision_reduction=( + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + ), + leaves={}, + nodes=[], + ) + # Keep FP32 arithmetic strict through backward as well as forward. + with block.strict_fp32_matmul(): + golden, golden_nodes = block._backward("golden", ctx, None, upstream) + provider, provider_nodes = block._backward("provider", ctx, None, upstream) + candidate, candidate_nodes = block._backward("candidate", ctx, ops, upstream) + for name in golden: + report["leaves"][name] = dict( + candidate=block.rel_err(candidate[name], golden[name]), + provider=block.rel_err(provider[name], golden[name]), + ) + for node in reversed(block.NODES): + report["nodes"].append( + dict( + node=node.name, + errors=[ + dict(candidate=block.rel_err(c, g), provider=block.rel_err(p, g)) + for g, p, c in zip( + golden_nodes[node.name], + provider_nodes[node.name], + candidate_nodes[node.name], + strict=True, + ) + ], + ) + ) + # Same input and upstream gradient isolate the RMSNorm implementation from + # error accumulated through other candidate nodes. + with torch.no_grad(): + values = block.run_graph("provider", ctx) + isolated = {} + for name, source, weight_name in ( + ("q_norm", "q_proj", "attn.norm_q.weight"), + ("k_norm", "k_proj", "attn.norm_k.weight"), + ): + grads = {} + node = next(n for n in block.NODES if n.name == name) + for mode in ("golden", "provider", "candidate"): + x = values[source].detach().clone() + w = weights[block.BLOCK + "." + weight_name].detach().clone() + if mode == "golden": + x, w = x.float(), w.float() + x.requires_grad_() + w.requires_grad_() + local = {**ctx, "weights": {**weights, block.BLOCK + "." + weight_name: w}} + y = ( + node.candidate(ops, local, x) + if mode == "candidate" + else getattr(node, mode)(local, x) + ) + torch.autograd.backward(y, provider_nodes[name][0].to(y.dtype)) + grads[mode] = w.grad + isolated[name] = dict( + candidate=block.rel_err(grads["candidate"], grads["golden"]), + provider=block.rel_err(grads["provider"], grads["golden"]), + ) + report["isolated_norm_weight_gradients"] = isolated + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + if args.causal_sweep: + # These hybrids intentionally use provider nodes for localization only. + # They cannot be submitted to the candidate acceptance comparator. + report["causal_sweep"] = [] + original_nodes = block.NODES + try: + for node in original_nodes: + for baseline in ("candidate", "provider"): + mixed = [] + for item in original_nodes: + use_provider = (item.name == node.name) == (baseline == "candidate") + mixed.append( + replace( + item, + candidate=lambda ops, ctx, *xs, item=item: item.provider(ctx, *xs), + ) + if use_provider + else item + ) + block.NODES = tuple(mixed) + with block.strict_fp32_matmul(): + leaves, _ = block._backward("candidate", ctx, ops, upstream) + report["causal_sweep"].append( + dict( + baseline=baseline, + replaced_node=node.name, + errors={ + name: block.rel_err(leaves[name], golden[name]) for name in golden + }, + ) + ) + del leaves + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + finally: + block.NODES = original_nodes + report["causal_sweep_complete"] = True + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_replay.py b/tools/validation/models/h3_block_replay.py new file mode 100644 index 000000000..7d8230294 --- /dev/null +++ b/tools/validation/models/h3_block_replay.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay one real MiniMax-H3 block forward and backward and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_block.py`` for what is compared. + + python tools/weights/prepare_h3_weights.py --out --block + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_replay.py --out block_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_block_weights, + load_h3_manifest, +) + + +def main() -> None: + """Run the forward (and optionally backward) block replay per layout and save the report.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--layouts", + default="tiny,small,medium", + help=f"comma list of packed FL2VA layouts: {', '.join(H3_BLOCK_LAYOUTS)}", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-backward", action="store_true", help="forward replay only") + parser.add_argument("--timing-repeats", type=int, default=5) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + layouts = args.layouts.split(",") + unknown = sorted(set(layouts) - set(H3_BLOCK_LAYOUTS)) + if unknown: + parser.error(f"unknown layouts: {unknown}") + if not torch.cuda.is_available(): + raise SystemExit("the block replay needs a CUDA device") + + weights = load_h3_block_weights("cuda") + manifest = load_h3_manifest() + ops = h3_block.CandidateOps() + + forward, backward = [], [] + for layout in layouts: + case = h3_block.run_block_case( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + forward.append(case) + print( + f"{layout} S={case['seq_len']}: first_drift={case['first_drift']} " + f"first_isolated_drift={case['first_isolated_drift']} " + f"repeat_drift={case['first_repeat_drift']} batch_drift={case['first_batch_drift']} " + f"out err cand {case['output']['candidate_rel_err_vs_golden']:.2e} " + f"prov {case['output']['provider_rel_err_vs_golden']:.2e}" + ) + if not args.no_backward: + case = h3_block.run_block_backward( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + backward.append(case) + det = all(v["repeat_bitwise_equal"] for v in case["leaves"].values()) + print( + f"{layout} backward: leaves {'det' if det else 'NONDET'}, " + f"grad_repeat_drift={case['first_grad_repeat_drift']} " + f"grad_batch_drift={case['first_grad_batch_drift']} " + f"dX batch-row {case['leaves']['hidden']['batch_row_bitwise_equal']}" + ) + torch.cuda.empty_cache() + + evidence = { + "kind": "h3_one_block_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "nodes": [ + { + "node": n.name, + "inputs": list(n.inputs), + "rfc_row": n.binding.rfc_row, + "status": n.binding.status, + "backend": n.binding.backend, + } + for n in h3_block.NODES + ], + "forward_cases": forward, + "backward_cases": backward, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_ulp_sweep.py b/tools/validation/models/h3_block_ulp_sweep.py new file mode 100644 index 000000000..7fe66646b --- /dev/null +++ b/tools/validation/models/h3_block_ulp_sweep.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sweep the one-block backward over layouts and seeds and record leaf errors in ulps. + +Backs the ulp allowance of ``tests/models/minimax_h3/test_h3_one_block.py``: +for every leaf gradient it records the candidate's and diffusers' error against +the FP32 golden both as ``max|err| / max|golden|`` and in ulps of the leaf's +storage dtype at ``max|golden|``, and the per-leaf excess +``candidate_ulps - provider_ulps``. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_ulp_sweep.py --layouts tiny,small --seeds 0,1,2,3,4 \\ + --out ulp_sweep.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_block_weights # noqa: E402 + + +def main() -> None: + """Run the backward replay for every (layout, seed) and write per-leaf ulp errors.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layouts", default="tiny,small") + parser.add_argument("--seeds", default="0,1,2,3,4") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("the sweep needs a CUDA device") + + weights = load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + cases = [] + for layout in args.layouts.split(","): + for seed in (int(s) for s in args.seeds.split(",")): + report = h3_block.run_block_backward(weights, layout=layout, seed=seed, ops=ops) + leaves = { + name: { + key: entry[key] + for key in ( + "storage_dtype", + "candidate_rel_err_vs_golden", + "provider_rel_err_vs_golden", + "candidate_ulps_vs_golden", + "provider_ulps_vs_golden", + ) + } + | { + "excess_ulps": entry["candidate_ulps_vs_golden"] + - entry["provider_ulps_vs_golden"] + } + for name, entry in report["leaves"].items() + } + worst = max(leaves, key=lambda n: leaves[n]["excess_ulps"]) + print( + f"{layout} S={report['seq_len']} seed={seed}: worst excess " + f"{leaves[worst]['excess_ulps']:+.2f} ulp ({worst})" + ) + cases.append( + {"layout": layout, "seq_len": report["seq_len"], "seed": seed, "leaves": leaves} + ) + torch.cuda.empty_cache() + + bf16 = [ + leaf["excess_ulps"] + for case in cases + for leaf in case["leaves"].values() + if leaf["storage_dtype"] == "bfloat16" + ] + evidence = { + "kind": "h3_one_block_backward_ulp_sweep", + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "cases": cases, + "max_bf16_excess_ulps": max(bf16), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"max BF16 excess: {max(bf16):+.2f} ulp; wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_candidate_prepare.py b/tools/validation/models/h3_candidate_prepare.py new file mode 100644 index 000000000..ae17b6600 --- /dev/null +++ b/tools/validation/models/h3_candidate_prepare.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fail closed on missing native H3 kernels and pin the candidate/reference identities.""" +import argparse +import json +import os +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from rl_engine.backends.extension import _C +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import require_candidate_cuda +from rl_engine.validation.models.h3_weights import load_h3_manifest, sha256_tensor +from tools.validation.models.h3_evidence_integrity import verify_inventory +from tools.validation.models.h3_full_model import digest, verify_checkpoint + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + a = p.parse_args() + a.output.mkdir() + manifest = Path("examples/models/minimax_h3/reference-manifest.json") + verify_inventory(a.reference_root, manifest) + verify_checkpoint( + a.checkpoint, Path("rl_engine/validation/models/h3_full_manifest.json"), all_files=True + ) + if _C is None or not hasattr(_C, "h3_timestep_sinusoid_forward"): + raise ValueError("candidate extension lacks H3 kernels") + ops = CandidateOps() + backends = require_candidate_cuda(ops) + # Inject a real eager CUDA fallback and require the same guard used by the + # full-model path to reject it before it can count as candidate execution. + original = ops.norm + try: + ops.norm = torch.nn.RMSNorm(128).cuda() + try: + require_candidate_cuda(ops) + except ValueError as exc: + if "candidate norm fell back" not in str(exc): + raise + else: + raise AssertionError("eager CUDA norm fallback was accepted") + finally: + ops.norm = original + require_candidate_cuda(ops) + extension = Path(_C.__file__).resolve() + if not extension.is_relative_to(Path.cwd().resolve()): + raise ValueError("candidate loaded an extension outside its clean source snapshot") + spec = load_h3_manifest() + index = json.loads( + (a.checkpoint / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + )["weight_map"] + weights = a.output / "weights" + weights.mkdir() + for section, filename in [ + ("tensors", "h3_conditioning.safetensors"), + ("block_tensors", "h3_block0.safetensors"), + ]: + values = {} + for shard in sorted({index[name] for name in spec[section]}): + with safe_open(str(a.checkpoint / "transformer" / shard), framework="pt") as handle: + for name, row in spec[section].items(): + if index[name] == shard: + value = handle.get_tensor(name).contiguous() + if sha256_tensor(value) != row["sha256"]: + raise ValueError(f"block weight identity changed: {name}") + values[name] = value + save_file(values, str(weights / filename)) + report = dict( + extension=str(extension), + extension_sha256=digest(extension), + torch=torch.__version__, + backends=backends, + negative_controls={"eager_cuda_norm_fallback_rejected": True}, + reference_manifest_sha256=digest(manifest), + build={ + key: os.environ.get(key) + for key in ("CUDA_HOME", "MAX_JOBS", "KERNEL_ALIGN_USE_FAST_MATH") + }, + ) + (a.output / "identity.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_chain_replay.py b/tools/validation/models/h3_chain_replay.py new file mode 100644 index 000000000..11068259c --- /dev/null +++ b/tools/validation/models/h3_chain_replay.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the MiniMax-H3 conditioning chain and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_chain_replay.py --out chain_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models import h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) + + +def _ints(text: str) -> list[int]: + """Parse a comma-separated list of integer case sizes.""" + + return [int(v) for v in text.split(",")] + + +def main() -> None: + """Replay requested CUDA stages with prerequisites and optionally save JSON evidence.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages to report; preceding stages run automatically", + ) + parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") + parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None) + parser.add_argument( + "--backward", + action="store_true", + help="also replay the chain backward (needs every stage): parameter-gradient " + "determinism and accuracy for the separate-op, fused and diffusers chains", + ) + args = parser.parse_args() + + wanted = args.stages.split(",") + stage_names = [s.name for s in h3_chain.STAGES] + unknown = sorted(set(wanted) - set(stage_names)) + if unknown: + parser.error(f"unknown stages: {unknown}") + stages = [s for s in h3_chain.STAGES if s.name in wanted] + if args.backward and len(stages) != len(h3_chain.STAGES): + parser.error("--backward requires every stage of the chain") + last_stage = max(i for i, stage in enumerate(h3_chain.STAGES) if stage.name in wanted) + execution_stages = h3_chain.STAGES[: last_stage + 1] + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + weights = load_h3_conditioning_weights("cuda") + registry = KernelRegistry() + manifest = load_h3_manifest() + + cases = [] + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_case( + registry, + weights, + num_timesteps=num_timesteps, + seq_len=seq_len, + seed=args.seed, + stages=execution_stages, + ) + case["stages"] = [entry for entry in case["stages"] if entry["stage"] in wanted] + cases.append(case) + summary = ", ".join( + f"{e['stage']}=" + f"{'eq' if e['chained_vs_provider']['bitwise_equal'] else 'drift'}" + f"(gold {e['chained_vs_golden']['max_abs']:.3e})" + for e in case["stages"] + ) + print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + + backward_cases = [] + if args.backward: + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_backward_case( + registry, weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed + ) + backward_cases.append(case) + parts = [] + for mode in h3_chain.BACKWARD_MODES: + entries = case["leaves"].values() + det = all(e[mode]["repeat_bitwise_equal"] for e in entries) + worst = max(e[mode]["max_abs_vs_golden_over_absmax"] for e in entries) + parts.append(f"{mode}: {'det' if det else 'NONDET'} {worst:.1e}") + summary = " | ".join(parts) + print(f"backward T={num_timesteps} S={seq_len}: {summary}") + + evidence = { + "kind": "h3_conditioning_chain_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "stages": [s.name for s in stages], + "executed_stages": [s.name for s in execution_stages], + "cases": cases, + "backward_cases": backward_cases, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_dual_trajectory.py b/tools/validation/models/h3_dual_trajectory.py new file mode 100644 index 000000000..7d880c5d3 --- /dev/null +++ b/tools/validation/models/h3_dual_trajectory.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""RFC #420 ``ws1_dual_trajectory``: two layouts of the same sampling run must agree bitwise. + +Both trajectories run the same multi-step flow-matching Euler sampler (``x += (s' - s) v``, +FP32 latents, video + audio) on the full pinned 33B H3 and record it with the sealed +trajectory artifact. They differ only in how the request is laid out for the model: + +* ``single`` - the request alone in the batch; +* ``batched`` - the request as row 1 of a batch of 2 beside an unrelated request. + +Under RL-Kernel strict operators the two are bitwise identical at every step; under the +provider's own kernels the sampler drifts (and the artifact comparison names the first step). +That is the train/rollout-layout control the end-to-end experiment needs. +""" + +from __future__ import annotations + +import argparse +import json +import time +from contextlib import contextmanager +from pathlib import Path + +import torch + +from rl_engine.models.minimax_h3 import trajectory as tj +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_model import packed_inputs + +SIGMAS = (1.0, 0.75, 0.5, 0.25, 0.0) + + +@contextmanager +def diagnostic_video_projections(model, enabled): + """DIAGNOSTIC ONLY: run ``proj_in``/``proj_out`` through the fixed-order Triton GEMM. + + The video projection rows belong to another contributor; this binding exists only to + attribute the remaining single-vs-batched drift and is labelled in the artifact metadata. + """ + + if not enabled: + yield + return + from rl_engine.backends.shared.triton.gemm import h3_projection as g + + saved = [] + for name in ("proj_in", "proj_out"): + module = getattr(model, name) + + def forward(x, module=module): + flat = x.reshape(-1, x.shape[-1]) + out = g.linear_forward(flat.contiguous(), module.weight, module.bias) + return out.view(*x.shape[:-1], out.shape[-1]) + + saved.append(module) + module.forward = forward + try: + yield + finally: + for module in saved: + module.__dict__.pop("forward", None) + + +def sample( + model, + base, + noise, + *, + candidate, + strict_boundaries, + layout, + label, + keep_tensors, + seed, + diag_video=False, +): + video, audio = noise["video"].clone(), noise["audio"].clone() + other = { + k: torch.randn(v.shape, generator=torch.Generator().manual_seed(77 + i)).to(v.device) + for i, (k, v) in enumerate(noise.items()) + } + recorder = tj.TrajectoryRecorder( + variant="FL2VA", + label=label, + seed=seed, + sigma_grids={"video": SIGMAS, "audio": SIGMAS}, + initial_noise={"video": video, "audio": audio}, + metadata={ + "layout": layout, + "candidate": candidate, + "strict_boundaries": strict_boundaries, + "sampler": "euler-ode", + "diagnostic_video_projections": diag_video, + }, + keep_tensors=keep_tensors, + ) + trace = FullModelTrace() + with ( + torch.no_grad(), + diagnostic_video_projections(model, diag_video), + full_model_execution( + model, candidate=candidate, trace=trace, strict_boundaries=strict_boundaries + ), + ): + for step in range(len(SIGMAS) - 1): + sigma, sigma_next = SIGMAS[step], SIGMAS[step + 1] + enc = base["encoder_hidden_states"].detach() + if layout == "single": + hs, ahs, es = video.bfloat16(), audio.bfloat16(), enc + else: # the request is row 1, an unrelated one is row 0 + hs = torch.cat([other["video"].bfloat16(), video.bfloat16()]) + ahs = torch.cat([other["audio"].bfloat16(), audio.bfloat16()]) + es = torch.cat([enc * 0.5, enc]) + inputs = { + **base, + "hidden_states": hs, + "audio_hidden_states": ahs, + "encoder_hidden_states": es, + "timestep": torch.tensor([0.0, sigma], device=hs.device), + } + v_video, v_audio = model(**inputs) + row = 0 if layout == "single" else 1 + v_video, v_audio = v_video[row : row + 1].float(), v_audio[row : row + 1].float() + video = video + (sigma_next - sigma) * v_video + audio = audio + (sigma_next - sigma) * v_audio + recorder.record_step( + step, + velocity={"video": v_video, "audio": v_audio}, + latent={"video": video, "audio": audio}, + ) + return recorder + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--seed", type=int, default=420539) + parser.add_argument( + "--diag-video-projections", + action="store_true", + help="diagnostic: bind proj_in/proj_out to the fixed-order GEMM (strict arm only)", + ) + parser.add_argument("--strict-only", action="store_true") + args = parser.parse_args() + from diffusers import MiniMaxH3Transformer3DModel + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + started = time.perf_counter() + model = ( + MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + ) + .to("cuda") + .eval() + ) + base = packed_inputs("cuda") + gen = torch.Generator().manual_seed(args.seed) + noise = { + "video": torch.randn(base["hidden_states"].shape, generator=gen).cuda(), + "audio": torch.randn(base["audio_hidden_states"].shape, generator=gen).cuda(), + } + args.output.mkdir(parents=True, exist_ok=True) + result = {"load_seconds": time.perf_counter() - started, "pairs": {}} + arms = (("provider", False), ("rlkernel_strict", True)) + if args.strict_only: + arms = arms[1:] + for name, candidate in arms: + recs = {} + for layout in ("single", "batched"): + t0 = time.perf_counter() + rec = sample( + model, + base, + noise, + candidate=candidate, + strict_boundaries=candidate, + layout=layout, + label=f"{name}-{layout}", + keep_tensors=True, + seed=args.seed, + diag_video=args.diag_video_projections and candidate, + ) + recs[layout] = (rec, time.perf_counter() - t0) + paths = {} + for layout, (rec, _) in recs.items(): + art = rec.seal() + (args.output / f"{name}-{layout}.json").write_text(tj.dumps(art)) + rec.save_tensors(args.output / f"{name}-{layout}.safetensors") + paths[layout] = art + report = tj.drift_report( + paths["single"], + paths["batched"], + tensors_lhs=args.output / f"{name}-single.safetensors", + tensors_rhs=args.output / f"{name}-batched.safetensors", + ) + report["seconds"] = {k: v[1] for k, v in recs.items()} + result["pairs"][name] = report + print( + name, + { + k: report[k] + for k in ("steps", "compared", "differing", "bitwise_equal", "first_drift") + }, + flush=True, + ) + (args.output / "report.json").write_text(json.dumps(result, indent=1) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py new file mode 100644 index 000000000..213d5beb3 --- /dev/null +++ b/tools/validation/models/h3_evidence.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write one RFC #420 operator's evidence report (performance + accuracy) as JSON. + +Run it on a clean tree; the report records the commit and whether the tree +was dirty. ``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) + + +def main() -> None: + """Collect CUDA accuracy and performance evidence, requiring weights for weighted ops.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + uses_weights = args.op != "timestep_sinusoid_h3" + if uses_weights and h3_weights_dir() is None: + raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + registry = KernelRegistry() + manifest = load_h3_manifest() + report = { + "kind": "h3_operator_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weight_source": "pinned_checkpoint" if uses_weights else "not_applicable", + "reference_commit": manifest["reference_implementation"]["commit"], + **git_state(), + "environment": environment(), + "accuracy": ACCURACY[args.op](registry), + "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence_integrity.py b/tools/validation/models/h3_evidence_integrity.py new file mode 100644 index 000000000..42d9b383d --- /dev/null +++ b/tools/validation/models/h3_evidence_integrity.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate immutable reference inventories before consuming any tensor evidence.""" +from __future__ import annotations + +import json +from pathlib import Path + +from tools.validation.models.h3_full_model import digest + + +def verify_inventory(root: Path, manifest: Path): + records = json.loads(manifest.read_text())["files"] + for name, expected in records.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError("evidence path escapes root") + if digest(path) != expected: + raise ValueError(f"reference evidence changed: {name}") + + +def verify_generation(root): + report = json.loads((root / "report.json").read_text()) + if report["status"] != "generation_complete_comparison_pending": + raise ValueError("generation did not complete") + required = {"initial.pt", "video.npy", "audio.npy"} | { + f"step-{i:03}.pt" for i in range(report["steps"]) + } + if not required.issubset(report["artifacts"]): + raise ValueError("incomplete generation artifact inventory") + for name, expected in report["artifacts"].items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()) or digest(path) != expected: + raise ValueError(f"generation artifact changed: {name}") + return report diff --git a/tools/validation/models/h3_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_charts.py b/tools/validation/models/h3_full_charts.py new file mode 100644 index 000000000..5b8a4548e --- /dev/null +++ b/tools/validation/models/h3_full_charts.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Charts of the full-transformer and dual-trajectory evidence (drift, accuracy, time).""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 + +FLOOR = 1e-12 # plotted value for a bitwise-equal step + + +def _per_step(pair: dict) -> dict[int, float]: + steps = {} + for entry in pair.get("entries", []): + steps[entry["step"]] = max(steps.get(entry["step"], 0.0), entry["max_abs_diff"]) + return steps + + +def main(evidence: Path, out: Path) -> None: + fig, ax = plt.subplots(1, 3, figsize=(16, 4)) + series = ( + ("provider", "", "provider", "tab:red"), + ("RL-Kernel strict", "", "rlkernel_strict", "tab:blue"), + ("strict + video proj. diagnostic", "_diag", "rlkernel_strict", "tab:green"), + ) + for label, suffix, key, color in series: + for i, seed in enumerate((11, 22, 33)): + path = evidence / f"h3_dual_trajectory_seed{seed}{suffix}.json" + pair = json.loads(path.read_text())["pairs"].get(key) + if pair is None: + continue + steps = _per_step(pair) + xs = range(4) + ax[0].semilogy( + xs, + [max(steps.get(s, 0.0), FLOOR) for s in xs], + "o-", + color=color, + alpha=0.8, + label=label if i == 0 else None, + ) + ax[0].set( + xlabel="denoising step", + ylabel=f"max |single - batched| (floor {FLOOR:g} = bitwise equal)", + title="Rollout vs train path drift, 3 seeds", + ) + ax[0].legend(fontsize=7) + full = json.loads((evidence / "h3_full_transformer_h100.json").read_text()) + cmp_ = full["candidate_vs_provider"] + names = ["video_rel", "audio_rel"] + [f"d{k}" for k in cmp_["input_grad_rel"]] + values = [cmp_["video_rel"], cmp_["audio_rel"], *cmp_["input_grad_rel"].values()] + ax[1].bar(names, values) + ax[1].set( + yscale="log", ylabel="relative difference to provider", title="Strict vs provider (33B)" + ) + ax[1].tick_params(axis="x", rotation=60, labelsize=7) + sec = full["seconds"] + ax[2].bar( + ["provider", "strict", "provider fwd", "strict fwd"], + [sec["provider"], sec["candidate"], sec["provider_forward"], sec["candidate_forward"]], + ) + ax[2].set(ylabel="seconds (H100)", title="Wall time incl. warm-up, tiny request") + 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_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_full_speed.py b/tools/validation/models/h3_full_speed.py new file mode 100644 index 000000000..13644c6dd --- /dev/null +++ b/tools/validation/models/h3_full_speed.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Time one full-model forward + backward step, provider against the strict RL-Kernel graph. + + python tools/validation/models/h3_full_speed.py --checkpoint \ + --sizes 12,925,4253,8349 --out speed.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path + +import torch + +SIZES = { + "12": None, + "925": dict(text=77, frames=3, side=32, audio=40), + "4253": dict(text=77, frames=1, side=128, audio=40), + "8349": dict(text=77, frames=2, side=128, audio=40), +} + + +def sized_inputs(device, *, text, frames, side, audio): + from rl_engine.validation.models.h3_cases import h3_fl2va_layout + + layout = h3_fl2va_layout( + num_text_tokens=text, + num_latent_frames=frames, + latent_height=side, + latent_width=side, + num_audio_latents=audio, + device=device, + ) + tags = layout["token_tags"] + video, text_idx, audio_idx = ((tags == t).nonzero().flatten() for t in (0, 1, 2)) + gen = torch.Generator().manual_seed(420539) + + def rnd(*shape): + return torch.randn(shape, generator=gen).to(device).requires_grad_() + + return dict( + **layout, + hidden_states=rnd(1, len(video), 96), + audio_hidden_states=rnd(1, len(audio_idx), 32), + encoder_hidden_states=rnd(1, len(text_idx), 5120), + timestep=torch.tensor([0.0, 500.0], device=device), + video_indices=video, + text_indices=text_idx, + audio_indices=audio_idx, + return_dict=False, + ) + + +def timed(model, inputs, candidate, repeats, attention="auto"): + from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution + + grad_inputs = ("hidden_states", "audio_hidden_states", "encoder_hidden_states") + backends = {} + + def once(): + fresh = { + k: (v.detach().clone().requires_grad_() if k in grad_inputs else v) + for k, v in inputs.items() + } + trace = FullModelTrace() + torch.cuda.synchronize() + t0 = time.perf_counter() + ops = None + if candidate: + from rl_engine.validation.models.h3_block import CandidateOps + + ops = CandidateOps(attention=attention) + with full_model_execution( + model, candidate=candidate, trace=trace, ops=ops, strict_boundaries=candidate + ): + with torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = model(**fresh) + torch.cuda.synchronize() + t1 = time.perf_counter() + (video.float().square().mean() + audio.float().square().mean()).backward() + torch.cuda.synchronize() + backends.update(getattr(trace, "actual_backends", None) or {}) + return t1 - t0, time.perf_counter() - t0 + + torch.cuda.reset_peak_memory_stats() + once() + runs = [once() for _ in range(repeats)] + return dict( + forward_s=statistics.median(r[0] for r in runs), + total_s=statistics.median(r[1] for r in runs), + peak_gib=torch.cuda.max_memory_allocated() / 2**30, + attention_backend=backends.get("attention_h3") or backends.get("attention"), + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--sizes", default="12,925,4253,8349") + p.add_argument("--attention", default="auto", choices=("auto", "h3_full", "interim")) + p.add_argument("--repeats", type=int, default=3) + p.add_argument("--out", type=Path, required=True) + args = p.parse_args() + from diffusers import MiniMaxH3Transformer3DModel + + from tools.validation.models.h3_full_model import packed_inputs + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + ).to("cuda") + model.eval().enable_gradient_checkpointing() + for q in model.parameters(): + q.register_post_accumulate_grad_hook(lambda t: setattr(t, "grad", None)) + result = dict(device=torch.cuda.get_device_name(0), torch=torch.__version__, sizes={}) + for size in args.sizes.split(","): + inputs = ( + packed_inputs("cuda") if SIZES[size] is None else sized_inputs("cuda", **SIZES[size]) + ) + entry = dict(tokens=int(inputs["position_ids"].shape[0])) + for name, candidate in (("provider", False), ("strict", True)): + try: + entry[name] = timed(model, inputs, candidate, args.repeats, args.attention) + except torch.OutOfMemoryError: + torch.cuda.empty_cache() + entry[name] = dict(error="CUDA out of memory") + if "total_s" in entry["provider"] and "total_s" in entry["strict"]: + entry["ratio"] = entry["strict"]["total_s"] / entry["provider"]["total_s"] + entry["fwd_ratio"] = entry["strict"]["forward_s"] / entry["provider"]["forward_s"] + result["sizes"][size] = entry + print(size, json.dumps(entry), flush=True) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(result, indent=1)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_transformer.py b/tools/validation/models/h3_full_transformer.py new file mode 100644 index 000000000..09451d801 --- /dev/null +++ b/tools/validation/models/h3_full_transformer.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""RFC #420 ``ws1_full_h3_transformer``: the whole pinned 33B H3 DiT, forward and backward. + +Runs the released BF16 checkpoint on one GPU three ways on a small packed input: + +* ``provider`` - diffusers' own arithmetic (cuBLAS / SDPA); +* ``candidate`` - the 50 real blocks through the RL-Kernel block graph, and (default) the + audio projections, the text projection and the token refiner through their rows + (``--provider-boundaries`` keeps those on the provider); +* ``candidate`` again - to prove the strict run is bitwise repeatable. + +Reports forward/backward agreement with the provider, bitwise repeatability of the outputs, +of every parameter gradient and of the input gradients, and which boundary modules are bound +to rows. Gradients are hashed or compared on the fly and released, so memory stays bounded. +This is integration evidence for a pinned checkpoint; interim block bindings are listed in +the report and are not final-kernel qualification. +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_model import packed_inputs, provider_variant_context + + +def tensor_sha(t: torch.Tensor) -> str: + """Bit-level GPU fingerprint: two position-weighted sums of the raw 16-bit words. + + Not cryptographic; equality of two fingerprints means the bytes agree for any realistic + difference (any single changed word changes the plain sum or the weighted sum). + """ + + t = t.detach().contiguous() + words = t.view(torch.uint8).view(-1) + words = words[: words.numel() // 2 * 2].view(torch.int16).to(torch.int64) + weights = torch.arange(words.numel(), device=words.device, dtype=torch.int64) % 65521 + 1 + plain, weighted = int(words.sum()), int((words * weights).sum()) + return f"{t.dtype}:{tuple(t.shape)}:{plain:x}:{weighted:x}" + + +def rel(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.double(), b.double() + return float((a - b).norm() / b.norm().clamp_min(1e-30)) + + +def run_once( + model, + inputs, + *, + candidate, + strict_boundaries, + reference_grads=None, + keep_grads=False, + variant="default", +): + """One forward+backward. Returns outputs, input grads, per-parameter (sha, rel-to-reference).""" + + trace = FullModelTrace() + stats, kept, handles = {}, {}, [] + for name, parameter in model.named_parameters(): + + def hook(p, name=name): + g = p.grad.detach() + if not torch.isfinite(g).all(): + raise ValueError(f"non-finite gradient: {name}") + entry = dict(sha=tensor_sha(g)) + if reference_grads is not None: + entry["rel_vs_reference"] = rel(g, reference_grads[name].to(g.device)) + if keep_grads: + kept[name] = g.cpu() + stats[name] = entry + p.grad = None + + handles.append(parameter.register_post_accumulate_grad_hook(hook)) + grad_inputs = ("hidden_states", "audio_hidden_states", "encoder_hidden_states") + fresh = { + k: (v.detach().clone().requires_grad_() if k in grad_inputs else v) + for k, v in inputs.items() + } + torch.cuda.synchronize() + begin = time.perf_counter() + with ( + full_model_execution( + model, candidate=candidate, trace=trace, strict_boundaries=strict_boundaries + ), + provider_variant_context(variant), + ): + with torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = model(**fresh) + torch.cuda.synchronize() + forward_s = time.perf_counter() - begin + loss = video.float().square().mean() + audio.float().square().mean() + loss.backward() + torch.cuda.synchronize() + total_s = time.perf_counter() - begin + for h in handles: + h.remove() + in_grads = { + k: fresh[k].grad.detach().cpu() + for k in ("hidden_states", "audio_hidden_states", "encoder_hidden_states") + } + return dict( + video=video.detach().cpu(), + audio=audio.detach().cpu(), + in_grads=in_grads, + stats=stats, + kept=kept, + trace=trace.validate(backward=True), + forward_s=forward_s, + total_s=total_s, + loss=float(loss), + ) + + +def summarise(cand, ref): + rels = [v["rel_vs_reference"] for v in cand["stats"].values() if "rel_vs_reference" in v] + return dict( + video_rel=rel(cand["video"], ref["video"]), + audio_rel=rel(cand["audio"], ref["audio"]), + input_grad_rel={k: rel(cand["in_grads"][k], ref["in_grads"][k]) for k in cand["in_grads"]}, + param_grad_rel_max=max(rels) if rels else None, + param_grad_rel_median=sorted(rels)[len(rels) // 2] if rels else None, + param_grad_count=len(cand["stats"]), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True, help="dir holding transformer/") + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--provider-boundaries", action="store_true") + args = parser.parse_args() + from diffusers import MiniMaxH3Transformer3DModel + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + started = time.perf_counter() + model, loading = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + output_loading_info=True, + ) + if any( + loading.get(k) for k in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs") + ): + raise ValueError("checkpoint did not load every tensor exactly") + model = model.to("cuda").eval() + model.enable_gradient_checkpointing() + load_s = time.perf_counter() - started + inputs = packed_inputs("cuda") + strict = not args.provider_boundaries + + provider = run_once(model, inputs, candidate=False, strict_boundaries=False, keep_grads=True) + provider_again = run_once(model, inputs, candidate=False, strict_boundaries=False) + cand = run_once( + model, inputs, candidate=True, strict_boundaries=strict, reference_grads=provider["kept"] + ) + cand2 = run_once(model, inputs, candidate=True, strict_boundaries=strict) + # How far an independent provider kernel choice moves the same numbers (calibration). + variants = {} + for name in ("sdpa-efficient", "blas-cublaslt"): + try: + other = run_once( + model, + inputs, + candidate=False, + strict_boundaries=False, + reference_grads=provider["kept"], + variant=name, + ) + variants[name] = summarise(other, provider) + except Exception as exc: # a forced kernel may be unavailable; record, do not fail + variants[name] = {"error": f"{type(exc).__name__}: {exc}"} + finally: + torch.backends.cuda.preferred_blas_library("default") + + def same(a, b): + return ( + torch.equal(a["video"], b["video"]) + and torch.equal(a["audio"], b["audio"]) + and all(torch.equal(a["in_grads"][k], b["in_grads"][k]) for k in a["in_grads"]) + and all(a["stats"][k]["sha"] == b["stats"][k]["sha"] for k in a["stats"]) + ) + + report = dict( + status="execution_complete", + device=torch.cuda.get_device_name(), + torch=torch.__version__, + parameters=sum(p.numel() for p in model.parameters()), + load_seconds=load_s, + strict_boundaries=strict, + candidate_vs_provider=summarise(cand, provider), + provider_repeat_vs_provider=summarise(provider_again, provider), + provider_variants_vs_provider=variants, + provider_bitwise_repeatable=same(provider, provider_again), + candidate_bitwise_repeatable=same(cand, cand2), + seconds=dict( + provider=provider["total_s"], + candidate=cand["total_s"], + provider_forward=provider["forward_s"], + candidate_forward=cand["forward_s"], + ), + peak_allocated_bytes=torch.cuda.max_memory_allocated(), + trace=dict( + bindings=cand["trace"]["bindings"], + boundary_bindings=cand["trace"]["boundary_bindings"], + candidate_calls=len(cand["trace"]["candidate_calls"]), + ), + loss=dict(provider=provider["loss"], candidate=cand["loss"]), + ) + args.output.mkdir(parents=True, exist_ok=True) + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps(report, indent=2)) + + +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_mapper_vllm_report.py b/tools/validation/models/h3_mapper_vllm_report.py new file mode 100644 index 000000000..1b3e50fbe --- /dev/null +++ b/tools/validation/models/h3_mapper_vllm_report.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Measurements of the RFC #420 ``h3_checkpoint_mapper`` and ``vllm_runtime_adapter_h3`` rows on +the real checkpoint: JSON report plus a PNG. + + python tools/validation/models/h3_mapper_vllm_report.py --checkpoint \ + --out report.json --png figure.png +""" + +from __future__ import annotations + +import argparse +import json +import re +import time +from collections import defaultdict +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from rl_engine.backends.shared.triton.gemm import h3_projection as g +from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout +from rl_engine.models.minimax_h3 import checkpoint_mapper as cm + +GROUP = re.compile(r"^((?:token_refiner\.refiner_blocks|transformer_blocks)\.\d+)\.") + + +def rel(a, b): + a, b = a.double(), b.double() + return ((a - b).norm() / b.norm()).item() + + +def mapping_report(root): + index = {} + for f in sorted(root.glob("*.safetensors")): + with safe_open(f, "pt") as h: + for k in h.keys(): + index[k] = f + groups = defaultdict(list) + for name in index: + m = GROUP.match(name) + groups[m.group(1) if m else "model-level"].append(name) + tensors = native_count = 0 + bitwise, t_fwd, t_back = True, 0.0, 0.0 + native_names = set() + for names in groups.values(): + state = {} + for n in names: + with safe_open(index[n], "pt") as h: + state[n] = h.get_tensor(n) + t0 = time.perf_counter() + native = cm.diffusers_to_native(state) + t1 = time.perf_counter() + back = cm.native_to_diffusers(native) + t_back += time.perf_counter() - t1 + t_fwd += t1 - t0 + tensors += len(state) + native_count += len(native) + native_names |= set(native) + bitwise &= set(back) == set(state) and all( + torch.equal(back[n], state[n]) and back[n].dtype == state[n].dtype for n in state + ) + return ( + dict( + checkpoint_tensors=tensors, + native_tensors=native_count, + unique_native_names=len(native_names), + groups=len(groups), + roundtrip_bitwise=bool(bitwise), + map_seconds=t_fwd, + inverse_seconds=t_back, + ), + index, + ) + + +def load(index, name): + with safe_open(index[name], "pt") as h: + return h.get_tensor(name).cuda() + + +def timeit(fn, n=20): + for _ in range(3): + fn() + torch.cuda.synchronize() + s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + s.record() + for _ in range(n): + fn() + e.record() + torch.cuda.synchronize() + return s.elapsed_time(e) / n + + +class Lin(torch.nn.Module): + def __init__(self, w): + super().__init__() + self.weight = torch.nn.Parameter(w, requires_grad=False) + + def forward(self, x): + return F.linear(x, self.weight), None + + +def layer_rows(index, tokens=(70, 925)): + rows = [] + for blk in (0, 25, 49): + p = f"transformer_blocks.{blk}." + q, k, v = (load(index, p + f"attn.to_{s}.weight") for s in "qkv") + o = load(index, p + "attn.to_out.0.weight") + f1 = load(index, p + "ff.net.0.proj.weight") + f2 = load(index, p + "ff.net.2.weight") + native = cm.diffusers_to_native( + { + p + "attn.to_q.weight": q, + p + "attn.to_k.weight": k, + p + "attn.to_v.weight": v, + p + "attn.to_out.0.weight": o, + p + "ff.net.0.proj.weight": f1, + p + "ff.net.2.weight": f2, + } + ) + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList([torch.nn.Module()]) + b = model.blocks[0] + b.attn, b.mlp = torch.nn.Module(), torch.nn.Module() + b.attn.qkv_proj = Lin(native["blocks.%d.attn.qkv_proj.weight" % blk]) + b.attn.out_proj = Lin(native["blocks.%d.attn.out_proj.weight" % blk]) + b.mlp.fc1 = Lin(native["blocks.%d.mlp.fc1.weight" % blk]) + b.mlp.fc2 = Lin(native["blocks.%d.mlp.fc2.weight" % blk]) + for t in tokens: + gen = torch.Generator().manual_seed(blk * 1000 + t) + hidden = q.shape[1] + x = torch.randn(1, t, hidden, generator=gen).bfloat16().cuda() + xo = torch.randn(1, t, o.shape[1], generator=gen).bfloat16().cuda() + x2 = torch.randn(1, t, f2.shape[1], generator=gen).bfloat16().cuda() + ref = { + "qkv_proj": (x, b.attn.qkv_proj), + "out_proj": (xo, b.attn.out_proj), + "fc1": (x, b.mlp.fc1), + "fc2": (x2, b.mlp.fc2), + } + torch_out = {n: m(xi)[0] for n, (xi, m) in ref.items()} + torch_ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()} + with install_h3_rollout(model): + out = {n: m(xi)[0] for n, (xi, m) in ref.items()} + ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()} + again = {n: m(xi)[0] for n, (xi, m) in ref.items()} + flat = x.reshape(-1, hidden) + sep = torch.cat([g.gemm(flat, w, b_t=True) for w in (q, k, v)], dim=1) + half = f1.shape[0] // 2 + sep_fc1 = g.gemm(flat, f1, b_t=True) + fc1_native = out["fc1"].reshape(t, -1) + swapped = torch.cat([sep_fc1[:, half:], sep_fc1[:, :half]], dim=1) + for n, (xi, m) in ref.items(): + exact = xi.reshape(-1, xi.shape[-1]).double() @ m.weight.double().t() + rows.append( + dict( + block=blk, + tokens=t, + linear=n, + repeatable=bool(torch.equal(out[n], again[n])), + rel_err_vs_fp64=rel(out[n].reshape(t, -1), exact), + torch_rel_err_vs_fp64=rel(torch_out[n].reshape(t, -1), exact), + fused_equals_separate=( + bool(torch.equal(out[n].reshape(t, -1), sep)) + if n == "qkv_proj" + else None + ), + fc1_equals_diffusers_swapped=( + bool(torch.equal(fc1_native, swapped)) if n == "fc1" else None + ), + adapter_ms=ms[n], + torch_ms=torch_ms[n], + slowdown=ms[n] / torch_ms[n], + ) + ) + return rows + + +def plot(report, path): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + rows = report["layers"] + names = ("qkv_proj", "out_proj", "fc1", "fc2") + fig, ax = plt.subplots(1, 3, figsize=(18, 5)) + m = report["mapping"] + ax[0].bar( + ["checkpoint\ntensors", "native\ntensors", "round trip\nbitwise equal"], + [ + m["checkpoint_tensors"], + m["native_tensors"], + m["checkpoint_tensors"] if m["roundtrip_bitwise"] else 0, + ], + ) + ax[0].set_title("Checkpoint mapper (real 33B checkpoint)") + w = 0.38 + for j, key in enumerate(("rel_err_vs_fp64", "torch_rel_err_vs_fp64")): + vals = [max(r[key] for r in rows if r["linear"] == n) for n in names] + ax[1].bar( + [i + (j - 0.5) * w for i in range(4)], + vals, + w, + label=("adapter" if j == 0 else "torch F.linear"), + ) + ax[1].set_xticks(range(4), names) + ax[1].set_yscale("log") + ax[1].set_title("Relative error vs FP64 (worst of 3 blocks x 2 sizes)") + ax[1].legend() + for t, c in ((70, "C0"), (925, "C1")): + vals = [ + sum(r["slowdown"] for r in rows if r["linear"] == n and r["tokens"] == t) / 3 + for n in names + ] + ax[2].plot(names, vals, "o-", color=c, label=f"{t} tokens") + ax[2].axhline(1, color="gray", lw=1) + ax[2].set_title("Adapter time / torch F.linear time (mean over blocks)") + ax[2].legend() + fig.tight_layout() + fig.savefig(path, dpi=130) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--png", type=Path) + args = p.parse_args() + torch.backends.cuda.matmul.allow_tf32 = False + mapping, index = mapping_report(args.checkpoint / "transformer") + report = dict( + device=torch.cuda.get_device_name(0), + torch=torch.__version__, + mapping=mapping, + layers=layer_rows(index), + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1)) + if args.png: + plot(report, args.png) + print(json.dumps(mapping, indent=1)) + for r in report["layers"]: + print(r) + + +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_parallel_run.py b/tools/validation/models/h3_parallel_run.py new file mode 100644 index 000000000..fc5cf6dbc --- /dev/null +++ b/tools/validation/models/h3_parallel_run.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Run the H3 parallel block on a TP x CP grid under ``srun`` (one task per GPU, any node count). + + srun -N2 --ntasks-per-node=4 python tools/validation/models/h3_parallel_run.py \\ + --tp 8 --cp 1 --out evidence/ws2_tp8 + +Every rank writes its rows (and gradients) to ``--out/rank{r}.pt``; rank 0 then recomputes the +single-GPU canonical block and compares the assembled forward, input gradient and weight-gradient +shards bitwise, and writes ``report.json``. ``--local-size`` is the number of GPUs per node for a +TP group that spans nodes (selects ``hierarchical_ipc_nccl_tree``); a TP group inside one node +uses the CUDA-IPC fixed tree. +""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import time +from pathlib import Path + +import torch +import torch.distributed as dist + +from rl_engine.validation.models import h3_tp +from rl_engine.validation.models.h3_weights import load_h3_block_weights + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tp", type=int, required=True) + parser.add_argument("--cp", type=int, default=1) + parser.add_argument("--sp", action="store_true") + parser.add_argument("--seq-len", type=int, default=0) + parser.add_argument("--layout", default="tiny") + parser.add_argument("--local-size", type=int, default=4, help="GPUs per node") + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + + rank = int(os.environ["SLURM_PROCID"]) + world = int(os.environ["SLURM_NTASKS"]) + local = int(os.environ["SLURM_LOCALID"]) + torch.cuda.set_device(local) + dist.init_process_group( + "nccl", + init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ.get('MASTER_PORT', '29511')}", + rank=rank, + world_size=world, + device_id=torch.device("cuda", local), + ) + if world != args.tp * args.cp: + raise SystemExit(f"world {world} != tp {args.tp} x cp {args.cp}") + probe = torch.ones(1, device="cuda") * rank + dist.all_reduce(probe) # NCCL connectivity across nodes + spans_nodes = args.tp > args.local_size + grid = h3_tp.make_grid(args.tp, args.cp, local_size=args.local_size if spans_nodes else None) + weights = load_h3_block_weights("cuda") + kwargs = dict(sp=args.sp, seq_len=args.seq_len, layout=args.layout) + result = h3_tp.block_case(weights, grid, backward=True, **kwargs) + # timing of the forward alone + torch.cuda.synchronize() + dist.barrier() + begin = time.perf_counter() + for _ in range(args.iters): + h3_tp.block_case(weights, grid, backward=False, **kwargs) + torch.cuda.synchronize() + seconds = (time.perf_counter() - begin) / args.iters + args.out.mkdir(parents=True, exist_ok=True) + torch.save({k: v.cpu() for k, v in result.items()}, args.out / f"rank{rank}.pt") + dist.barrier() + timings = [None] * world + dist.all_gather_object( + timings, {"rank": rank, "host": socket.gethostname(), "forward_s": seconds} + ) + if rank == 0: + ref = h3_tp.block_case( + weights, h3_tp.Grid(1, 1, 0, 0), backward=True, **{**kwargs, "sp": False} + ) + ranks = [torch.load(args.out / f"rank{r}.pt") for r in range(world)] + + def assemble(key): + pieces = {} + for r in ranks: + span = tuple(int(v) for v in r["rows"]) + if span in pieces and not torch.equal(pieces[span], r[key]): + return None + pieces[span] = r[key] + return torch.cat([pieces[s] for s in sorted(pieces)], dim=1) + + out, dx = assemble("out"), assemble("dx") + checks = { + "forward_bitwise": out is not None and torch.equal(out, ref["out"].cpu()), + "input_grad_bitwise": dx is not None and torch.equal(dx, ref["dx"].cpu()), + } + if args.cp == 1 and not args.sp: + for name, axis in ( + ("attn.to_out.0.weight", 1), + ("ff.net.2.weight", 1), + ("attn.to_q.weight", 0), + ("attn.to_k.weight", 0), + ("attn.to_v.weight", 0), + ("ff.hidden", 0), + ("ff.gate", 0), + ): + shards = [r["dw." + name] for r in ranks] + checks["dw." + name] = torch.equal( + torch.cat(shards, dim=axis), ref["dw." + name].cpu() + ) + report = dict( + tp=args.tp, + cp=args.cp, + sp=args.sp, + world=world, + hosts=sorted({t["host"] for t in timings}), + collective="hierarchical_ipc_nccl_tree" if spans_nodes else "cuda_ipc_fixed_tree", + seq_len=int(ref["out"].shape[1]) if not args.seq_len else args.seq_len, + checks=checks, + all_bitwise=all(checks.values()), + forward_seconds_per_rank=timings, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + (args.out / "report.json").write_text(json.dumps(report, indent=1) + "\n") + print( + json.dumps( + { + k: report[k] + for k in ("tp", "cp", "sp", "hosts", "collective", "checks", "all_bitwise") + }, + indent=1, + ) + ) + dist.barrier() + if rank == 0: + time.sleep(3) # keep the rendezvous store up until every rank has left + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_prior_art.py b/tools/validation/models/h3_prior_art.py new file mode 100644 index 000000000..9012542e2 --- /dev/null +++ b/tools/validation/models/h3_prior_art.py @@ -0,0 +1,946 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compare one RFC #420 operator with existing implementations: batch invariance, accuracy, latency. + +The RFC #420 reuse rule asks, for every operator, whether an existing implementation is +batch-invariant (then reuse it) and for accuracy and performance comparisons. This runner +measures the operator's rl-kernel CUDA op next to the existing implementations that import in +the current environment, and writes one JSON report per operator: + + python tools/validation/models/h3_prior_art.py --op adaln_row_gather \\ + --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json \\ + [--megatron-src /path/to/Megatron-LM] [--quick] + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/adaln_row_gather.json + +Only operators whose rl-kernel op exists on the current branch can be selected. Optional +libraries (SGLang, Liger, Transformer Engine, vLLM, Megatron-LM from a source tree) are skipped +and recorded as unavailable when they do not import. Global batch-invariant modes (vLLM, +SGLang, Megatron) patch aten for the whole process, so each runs in its own subprocess, with +the environment variables it relies on set before CUDA initialises. + +Batch invariance is bitwise, for the forward, the per-row gradients, and repeatability of every +parameter and table gradient, with three checks: + +1. every row computed alone vs inside full batches of 64, 257 and 2048 rows, seeds 3-5; +2. the full workload-size batch vs sub-batches that together cover every row; +3. eight probe rows at the front, middle and back of batches of every size 1..9 and + 2^k - 1, 2^k, 2^k + 1, plus the whole batch reversed. + +Sparse probes (a few rows, a few batch sizes) miss row-specific or large-batch dependence, which +is why every row is checked. Accuracy is max|err| / max|ref| against the same computation in +FP64; latency is the median of 100 CUDA-event samples after 20 warm-ups. Run on a clean tree on +an otherwise idle GPU; the report records the commit, the tree state and every library version. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as md +import importlib.util +import json +import os +import statistics +import subprocess +import sys +from pathlib import Path +from typing import Any, Callable + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +import torch # noqa: E402 +import torch.nn.functional as F # noqa: E402 + +H, D, F1, T3, EPS = 5376, 2688, 256, 3, 1e-5 + +#: operator -> rl-kernel module that must exist on the branch +OPS = { + "timestep_sinusoid": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid", + "timestep_mlp": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp", + "adaln_projection": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection", + "adaln_row_gather": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather", + "norm_modulate": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm", + "gate_residual": "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual", + "final_adaln_out": "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out", +} +#: global batch-invariant modes compared for the GEMM operators, with their environment +MODES = { + "plain": {}, + "vllm": {"CUBLAS_WORKSPACE_CONFIG": ":16:8", "CUBLASLT_WORKSPACE_SIZE": "1"}, + "sglang": {}, + "sglang_ieee": {"TRITON_F32_DEFAULT": "ieee"}, + "megatron_te_native": {"CUBLASLT_WORKSPACE_SIZE": "0"}, + "megatron_triton": {}, + "megatron_triton_ieee": {"TRITON_F32_DEFAULT": "ieee"}, +} +GEMM_OPS = ("timestep_mlp", "adaln_projection") +LIBS = ("torch", "triton", "diffusers", "vllm", "flashinfer-python", "sglang", "liger-kernel") +LIBS += ("transformer_engine", "megatron-core") + +DEV = "cuda" + + +# --------------------------------------------------------------------------- # +# Batch-invariance checks +# --------------------------------------------------------------------------- # + + +class _Case: + """``fn(**inputs)`` with row arguments sliced together; fixed upstream gradients per seed.""" + + def __init__(self, fn, inputs, rows, grads, out_dims, seed): + self.fn, self.inputs, self.rows, self.grads, self.out_dims = ( + fn, + inputs, + rows, + grads, + out_dims, + ) + self.n = next(inputs[k].shape[d] for k, d in rows.items()) + outs = self._flat(fn(**inputs)) + gen = torch.Generator(device=DEV).manual_seed(seed) + self.ups = [torch.randn(o.shape, device=DEV, generator=gen).to(o.dtype) for o in outs] + + @staticmethod + def _flat(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + def run(self, idx=None): + args = {} + for k, v in self.inputs.items(): + if idx is not None and k in self.rows: + v = v.index_select(self.rows[k], idx) + if k in self.grads: + v = v.detach().clone().requires_grad_(True) + args[k] = v + outs = self._flat(self.fn(**args)) + ups = ( + self.ups + if idx is None + else [u.index_select(d, idx) for u, d in zip(self.ups, self.out_dims)] + ) + live = [(o, u) for o, u in zip(outs, ups) if o.requires_grad] + if live: + torch.autograd.backward([o for o, _ in live], [u for _, u in live]) + grads = {k: args[k].grad for k in self.grads} + return [o.detach() for o in outs], grads + + +def _sweep_sizes(n: int) -> list[int]: + sizes = set(range(1, min(n, 9) + 1)) + k = 16 + while k <= n: + sizes.update(v for v in (k - 1, k, k + 1) if v <= n) + k *= 2 + return sorted(sizes | {n}) + + +def batch_invariance(cand: dict[str, Any], quick: bool) -> dict[str, Any]: + fn, make, rows, grads = cand["fn"], cand["make"], cand["rows"], cand["grads"] + out_dims = cand.get("out_dims", (0,)) + params = [k for k in grads if k not in rows] + full_sizes, big, subs = cand["bi_sizes"] + if quick: + full_sizes, big, subs = (17, 64), 256, (1, 7, 64) + seeds = (3, 4, 5) + result: dict[str, Any] = {"all_rows": {}, "params_repeatable": None} + ok = True + + def same_row(a, b, d, i): + return torch.equal(a.select(d, 0), b.select(d, i)) + + for n in full_sizes: + fwd = grd = 0 + repeat = True + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + fo, fg = case.run() + for i in range(n): + o, g = case.run(torch.tensor([i], device=DEV)) + fwd += not all(same_row(a, b, d, i) for a, b, d in zip(o, fo, out_dims)) + grd += not all( + same_row(g[k], fg[k], rows[k], i) + for k in rows + if k in grads and g[k] is not None + ) + if params: + _, again = case.run() + repeat &= all(torch.equal(again[k], fg[k]) for k in params if fg[k] is not None) + result["all_rows"][str(n)] = { + "rows": len(seeds) * n, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + if params: + result["params_repeatable"] = repeat and result["params_repeatable"] is not False + ok &= fwd == 0 and grd == 0 and repeat + + fwd = grd = checked = 0 + for seed in seeds[:2]: + case = _Case(fn, make(seed, big), rows, grads, out_dims, seed) + fo, fg = case.run() + for sb in subs: + step = sb if sb >= 1024 else max(sb, big // 512) + for start in range(0, big, step): + idx = torch.arange(start, min(start + sb, big), device=DEV) + o, g = case.run(idx) + checked += 1 + fwd += not all( + torch.equal(a, b.index_select(d, idx)) for a, b, d in zip(o, fo, out_dims) + ) + grd += not all( + torch.equal(g[k], fg[k].index_select(rows[k], idx)) + for k in rows + if k in grads and g[k] is not None + ) + del case, fo, fg + torch.cuda.empty_cache() + result["full_vs_sub_batches"] = { + "full_batch": big, + "sub_batch_sizes": list(subs), + "sub_batches": checked, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + ok &= fwd == 0 and grd == 0 + + n = full_sizes[-1] + sizes = _sweep_sizes(n) + probes = sorted({0, n - 1, *list(range(n // 8 + 3, n - 1, n // 8))[:6]}) + compared = failed = 0 + first: list[dict[str, int]] = [] + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + alone = {r: case.run(torch.tensor([r], device=DEV)) for r in probes} + gen = torch.Generator().manual_seed(seed) + + def compare(idx, pos, r, seed=seed, case=case, alone=alone): + nonlocal compared, failed + o, g = case.run(idx.to(DEV)) + ao, ag = alone[r] + good = all( + torch.equal(a.select(d, pos), b.select(d, 0)) for a, b, d in zip(o, ao, out_dims) + ) + good &= all( + torch.equal(g[k].select(rows[k], pos), ag[k].select(rows[k], 0)) + for k in rows + if k in grads and g[k] is not None + ) + compared += 1 + if not good: + failed += 1 + if len(first) < 8: + first.append( + {"seed": seed, "batch": int(idx.numel()), "position": pos, "row": r} + ) + + for m in sizes: + for r in probes: + others = torch.randperm(n, generator=gen) + others = others[others != r][: m - 1] + for pos in sorted({0, (m - 1) // 2, m - 1}): + compare(torch.cat([others[:pos], torch.tensor([r]), others[pos:]]), pos, r) + for r in probes: + compare(torch.arange(n - 1, -1, -1), n - 1 - r, r) + result["batch_size_sweep"] = { + "batch_sizes": sizes, + "probe_rows": probes, + "comparisons": compared, + "differ": failed, + "first_failures": first, + } + ok &= failed == 0 + result["batch_invariant"] = bool(ok) + return result + + +# --------------------------------------------------------------------------- # +# Accuracy and latency +# --------------------------------------------------------------------------- # + + +def _time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 100) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def _rel(a: torch.Tensor, ref: torch.Tensor) -> float: + return ((a.double() - ref).abs().max() / ref.abs().max().clamp_min(1e-300)).item() + + +def accuracy_latency(cand: dict[str, Any], mode_off, quick: bool) -> dict[str, Any]: + out: dict[str, Any] = {} + for size in cand["perf_sizes"][:2] if quick else cand["perf_sizes"]: + inputs = cand["make"](11, size) + names = list(inputs) + grads = [k for k in cand["grads"] if k in names] + leaves = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + y = cand["fn"](**leaves) + y = torch.cat([t.reshape(-1) for t in y]) if isinstance(y, (tuple, list)) else y + with mode_off(): + ref_in = { + k: (v.detach().double().requires_grad_(k in grads) if v.is_floating_point() else v) + for k, v in inputs.items() + } + ref = cand["ref"](**ref_in) + ref = torch.cat([t.reshape(-1) for t in ref]) if isinstance(ref, (tuple, list)) else ref + row: dict[str, Any] = {"fwd_rel_err": _rel(y, ref.detach())} + if grads: + gen = torch.Generator(device=DEV).manual_seed(9) + up = torch.randn(y.shape, device=DEV, generator=gen).to(y.dtype) + y.backward(up) + with mode_off(): + ref.backward(up.double()) + row["grad_rel_err"] = {k: _rel(leaves[k].grad, ref_in[k].grad) for k in grads} + row["fwd_us"] = _time_us(lambda inputs=inputs: cand["fn"](**inputs)) + if grads: + + def step(inputs=inputs): + lv = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + o = cand["fn"](**lv) + outs = list(o) if isinstance(o, (tuple, list)) else [o] + torch.autograd.backward(outs, [torch.ones_like(t) for t in outs]) + + row["fwd_bwd_us"] = _time_us(step) + out[str(size)] = row + del inputs, leaves, y, ref + torch.cuda.empty_cache() + return out + + +# --------------------------------------------------------------------------- # +# Candidates +# --------------------------------------------------------------------------- # + + +def _gen(seed: int, salt: int) -> torch.Generator: + return torch.Generator(device=DEV).manual_seed(seed * 1000 + salt) + + +def _rn(seed, salt, *shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, device=DEV, generator=_gen(seed, salt)) * scale).to(dtype) + + +def _layout(n: int, seed: int): + from rl_engine.validation.models.h3_cases import h3_packed_layout + + ti, tags = h3_packed_layout(n, T3, seed=seed) + return ti.contiguous(), tags.contiguous() + + +def _version(name: str) -> str | None: + try: + return md.version(name) + except md.PackageNotFoundError: + return None + + +T_BI = ((64, 257, 2048), 4096, (1, 7, 1024, 2048)) +S_BI = ((64, 257, 2048), 131072, (1, 7, 4097, 32768, 65536)) +T_PERF = (1, 3, 4, 64, 256, 2048) +S_PERF = (4097, 32768) + + +def _sinusoid_ref(t): + half = F1 // 2 + exponent = -torch.log(torch.tensor(10000.0, dtype=torch.float64, device=DEV)) + freqs = torch.exp(exponent * torch.arange(half, device=DEV, dtype=torch.float64) / half) + e = t.double()[:, None] * freqs[None] + return torch.cat([torch.cos(e), torch.sin(e)], -1) + + +def candidates(op: str, mode: str) -> list[dict[str, Any]]: + """Candidate dicts; a factory that raises (missing library) becomes an unavailable entry.""" + + from rl_engine.validation.models import h3_provider as P + + plain = mode == "plain" + found: list[tuple[str, Callable[[], dict[str, Any]]]] = [] + + if op == "timestep_sinusoid": + base = { + "make": lambda s, n: {"t": torch.rand(n, device=DEV, generator=_gen(s, 1))}, + "rows": {"t": 0}, + "grads": [], + "ref": _sinusoid_ref, + "bi_sizes": ((64, 257, 2048), 8192, (1, 7, 2048, 4097)), + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers get_timestep_embedding (op-for-op replay)", + "fn": lambda t: P.provider_time_proj(t), + } + + def sglang(): + from sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit import ( + timestep_embedding, + ) + + return { + **base, + "source": f"SGLang {_version('sglang')} timestep_embedding", + "fn": lambda t: timestep_embedding( + t, F1, flip_sin_to_cos=True, downscale_freq_shift=0.0 + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + ) + + op_ = H3TimestepSinusoidCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepSinusoidCudaOp (check_range=False)", + "fn": lambda t: op_.forward(t, check_range=False), + } + + found = [("diffusers", diffusers), ("sglang", sglang), ("rl_kernel", ours)] + + elif op == "timestep_mlp": + + def make(s, n): + return { + "x": _rn(s, 2, n, F1), + "w1": _rn(1, 3, H, F1, scale=F1**-0.5), + "b1": _rn(1, 4, H, scale=0.02), + "w2": _rn(1, 5, D, H, scale=H**-0.5), + "b2": _rn(1, 6, D, scale=0.02), + } + + base = { + "make": make, + "rows": {"x": 0}, + "grads": ["x", "w1", "b1", "w2", "b2"], + "ref": lambda x, w1, b1, w2, b2: F.linear(F.silu(F.linear(x, w1, b1)), w2, b2), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers TimestepEmbedding, FP32 F.linear [{mode}]", + "fn": lambda x, w1, b1, w2, b2: P.provider_time_embedder(x, w1, b1, w2, b2), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import ( + H3TimestepMLPCudaOp, + ) + + op_ = H3TimestepMLPCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepMLPCudaOp", + "fn": lambda x, w1, b1, w2, b2: op_(x, w1, b1, w2, b2), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_projection": + + def make(s, n): + return { + "temb": _rn(s, 7, n, D), + "w": _rn(1, 8, 18 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 9, 18 * H, dtype=torch.bfloat16, scale=0.02), + } + + base = { + "make": make, + "rows": {"temb": 0}, + "grads": ["temb", "w", "b"], + "ref": lambda temb, w, b: F.linear(F.silu(temb), w, b), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers AdaLN projection, BF16 F.linear [{mode}]", + "fn": lambda temb, w, b: F.linear(F.silu(temb).to(w.dtype), w, b), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + + op_ = H3AdaLNProjectionCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNProjectionCudaOp", + "fn": lambda temb, w, b: op_.forward_table(temb, w, b), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_row_gather": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "table": _rn(1, 14, 3 * T3, 6 * H, dtype=torch.bfloat16, scale=0.1), + "ti": ti, + "tags": tags, + } + + def ref(table, ti, tags): + return tuple(table.index_select(0, ti * 3 + tags).chunk(6, dim=1)) + + base = { + "make": make, + "rows": {"ti": 0, "tags": 0}, + "grads": ["table"], + "out_dims": (0,) * 6, + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return {**base, "source": "diffusers six index_select calls", "fn": ref} + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + op_ = H3AdaLNRowGatherCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNRowGatherCudaOp", + "fn": lambda table, ti, tags: op_(table, ti, tags), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "norm_modulate": + + def make(s, n): + ti, tags = _layout(n, s) + w = (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 11))).bfloat16() + return { + "x": _rn(s, 10, n, H, dtype=torch.bfloat16, scale=2.0), + "w": w, + "sh": _rn(1, 12, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "sc": _rn(1, 13, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(x, w, sh, sc, idx): + n = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + return n * (1 + sc.index_select(0, idx)) + sh.index_select(0, idx) + + base = { + "make": make, + "rows": {"x": 0, "idx": 0}, + "grads": ["x", "w", "sh", "sc"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def per_token(fn): + # Liger / SGLang take per-token shift/scale rows: gather them first (index_select). + return lambda x, w, sh, sc, idx: fn( + x, w, sh.index_select(0, idx), sc.index_select(0, idx) + ) + + def diffusers(): + return { + **base, + "source": "diffusers composition (F.rms_norm + index_select modulation)", + "fn": lambda x, w, sh, sc, idx: P.provider_norm_modulate(x, w, sh, sc, idx), + } + + def liger(): + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + return { + **base, + "source": f"Liger {_version('liger-kernel')} modulated RMSNorm + index_select", + "fn": per_token( + lambda x, w, shr, scr: LigerModulatedRMSNormFunction.apply( + x, w, scr, shr, EPS, 0.0, "llama", False + ) + ), + } + + def liger_rlk_gather(): + # Same Liger kernel, with the rows gathered by rl-kernel's deterministic row gather + # (adaln_row_gather) instead of index_select: separates Liger from the gather. + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + gather = H3AdaLNRowGatherCudaOp() + + def fn(x, w, sh, sc, idx): + pad = torch.zeros_like(sh) + table = torch.cat([sh, sc, pad, pad, pad, pad], dim=1) + rows = gather(table, torch.div(idx, 3, rounding_mode="floor"), idx % 3) + return LigerModulatedRMSNormFunction.apply( + x, w, rows[1], rows[0], EPS, 0.0, "llama", False + ) + + version = _version("liger-kernel") + return { + **base, + "source": f"Liger {version} modulated RMSNorm + rl-kernel row gather", + "fn": fn, + } + + def sglang(): + from sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl import ( + fused_norm_scale_shift, + ) + + return { + **base, + "grads": [], + "source": f"SGLang {_version('sglang')} fused_norm_scale_shift (forward only)", + "fn": per_token( + lambda x, w, shr, scr: fused_norm_scale_shift( + x[None], w, None, scr[None], shr[None], "rms", EPS + )[0] + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + + op_ = H3RMSNormCudaOp() + return { + **base, + "source": "rl-kernel H3RMSNormCudaOp.forward_modulated", + "fn": lambda x, w, sh, sc, idx: op_.forward_modulated(x, w, sh, sc, idx), + } + + def norm_only(make): + def m(s, n): + d = make(s, n) + return {"x": d["x"], "w": d["w"]} + + return m + + def norm_ref(x, w): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + + def torch_norm(): + # diffusers' norm without the modulation, to show which half is batch-invariant + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x", "w"], + "ref": norm_ref, + "source": "torch F.rms_norm (no modulation)", + "fn": lambda x, w: F.rms_norm(x, (H,), w, EPS), + } + + def te_norm(): + import transformer_engine.pytorch as te + + mod = te.RMSNorm(H, eps=EPS, params_dtype=torch.bfloat16, device=DEV) + + def fn(x, w): + # TE's backward reads its own parameter, so copy w in (functional_call + # gives a wrong dx); dweight is then TE's parameter gradient, not checked. + with torch.no_grad(): + mod.weight.copy_(w) + return mod(x) + + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x"], + "ref": norm_ref, + "source": f"TE {_version('transformer_engine')} RMSNorm (no modulation)", + "fn": fn, + } + + found = [ + ("diffusers", diffusers), + ("torch_rms_norm", torch_norm), + ("transformer_engine", te_norm), + ("liger", liger), + ("liger_rlk_gather", liger_rlk_gather), + ("sglang", sglang), + ("rl_kernel", ours), + ] + + elif op == "gate_residual": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "res": _rn(s, 15, n, H, dtype=torch.bfloat16), + "y": _rn(s, 16, n, H, dtype=torch.bfloat16), + "gate": _rn(1, 17, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(res, y, gate, idx): + return res + gate.index_select(0, idx) * y + + base = { + "make": make, + "rows": {"res": 0, "y": 0, "idx": 0}, + "grads": ["res", "y", "gate"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers residual + gate.index_select(...) * y", + "fn": lambda res, y, gate, idx: P.provider_gate_residual(res, gate, idx, y), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + ) + + op_ = H3GateResidualCudaOp() + return { + **base, + "source": "rl-kernel H3GateResidualCudaOp", + "fn": lambda res, y, gate, idx: op_(res, y, gate, idx), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "final_adaln_out": + + def make(s, n): + ti, _ = _layout(n, s) + return { + "x": _rn(s, 18, n, H, dtype=torch.bfloat16, scale=2.0), + "nw": (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 19))).bfloat16(), + "temb": _rn(1, 20, T3, D), + "w": _rn(1, 21, 2 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 22, 2 * H, dtype=torch.bfloat16, scale=0.02), + "ti": ti, + } + + def ref(x, nw, temb, w, b, ti): + return P.provider_final_adaln_out(x, nw, temb, w, b, ti) + + base = { + "make": make, + "rows": {"x": 0, "ti": 0}, + "grads": ["x", "nw", "temb", "w", "b"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay)", + "fn": ref, + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + + op_ = H3FinalAdaLNOutCudaOp() + return { + **base, + "source": "rl-kernel H3FinalAdaLNOutCudaOp", + "fn": lambda x, nw, temb, w, b, ti: op_(x, nw, temb, w, b, ti), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + result = [] + for name, factory in found: + try: + result.append({"name": name, **factory()}) + except Exception as exc: # noqa: BLE001 - optional library missing or unusable + result.append({"name": name, "unavailable": f"{type(exc).__name__}: {exc}"[:300]}) + return result + + +# --------------------------------------------------------------------------- # +# Modes and driver +# --------------------------------------------------------------------------- # + + +def enable_mode(mode: str, megatron_src: str | None): + """Switch on a global batch-invariant mode; returns a context manager that turns it off.""" + + base = mode[:-5] if mode.endswith("_ieee") else mode + if base == "plain": + return contextlib.nullcontext + if base == "vllm": + from vllm.model_executor.determinism.batch_invariant import enable_batch_invariant_mode + + enable_batch_invariant_mode() + return contextlib.nullcontext + if base == "sglang": + from sglang.srt.batch_invariant_ops import batch_invariant_ops as mod + + mod.enable_batch_invariant_mode() + kwargs: dict[str, str] = {} + else: + if megatron_src: + sys.path.insert(0, megatron_src) + from megatron.core.transformer.custom_layers import batch_invariant_kernels as mod + + kwargs = {"backend": base.split("_", 1)[1]} + mod.enable_batch_invariant_mode(**kwargs) + if kwargs["backend"] != "triton": + return contextlib.nullcontext + + @contextlib.contextmanager + def off(): + # The Triton matmuls have no FP64 configuration: FP64 references run with the mode off. + mod.disable_batch_invariant_mode() + try: + yield + finally: + mod.enable_batch_invariant_mode(**kwargs) + + return off + + +def run_mode(op: str, mode: str, megatron_src: str | None, quick: bool) -> dict[str, Any]: + try: + mode_off = enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + return {"unavailable": f"{type(exc).__name__}: {exc}"[:300]} + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + out: dict[str, Any] = {} + for cand in candidates(op, mode): + name = cand.pop("name") + if "unavailable" in cand: + out[name] = cand + print(f" [{mode}] {name}: unavailable ({cand['unavailable'][:80]})", flush=True) + continue + entry: dict[str, Any] = {"source": cand["source"], "has_backward": bool(cand["grads"])} + entry["batch_invariance"] = batch_invariance(cand, quick) + entry["accuracy_latency"] = accuracy_latency(cand, mode_off, quick) + out[name] = entry + print( + f" [{mode}] {name}: batch_invariant={entry['batch_invariance']['batch_invariant']}", + flush=True, + ) + torch.cuda.empty_cache() + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument("--quick", action="store_true", help="small sizes, for a smoke test only") + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + + if args.mode is not None: # child process: one global mode + args.out.write_text(json.dumps(run_mode(args.op, args.mode, args.megatron_src, args.quick))) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + modes = list(MODES) if args.op in GEMM_OPS else ["plain"] + results: dict[str, Any] = {} + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + # The part file sits next to the report, not in /tmp: cluster job epilogs can clear a + # user's /tmp while another of their jobs on the same node is still running. + part = args.out.with_name(f".{args.out.stem}.{mode}.part.json") + part.parent.mkdir(parents=True, exist_ok=True) + part.unlink(missing_ok=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.quick: + cmd.append("--quick") + proc = subprocess.run(cmd, env={**os.environ, **MODES[mode]}) + results[mode] = ( + json.loads(part.read_text()) + if proc.returncode == 0 and part.exists() + else {"unavailable": f"subprocess exited with {proc.returncode}"} + ) + part.unlink(missing_ok=True) + megatron_commit = None + if args.megatron_src: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + report = { + "kind": "h3_prior_art", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: _version(lib) for lib in LIBS}, + "megatron_source_commit": megatron_commit, + "mode_environment": {m: MODES[m] for m in modes}, + }, + "quick": args.quick, + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_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_runtime_report.py b/tools/validation/models/h3_runtime_report.py new file mode 100644 index 000000000..b7ed41daf --- /dev/null +++ b/tools/validation/models/h3_runtime_report.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Measurements of the RFC #420 H3 integration rows (runtime modes, ablation overlay, Diffusers +adapter, training provider, LoRA gate, CPS transition): JSON report plus a PNG. + + python tools/validation/models/h3_runtime_report.py --out report.json --png figure.png \ + [--checkpoint ] [--verl-omni ] +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + +from rl_engine.models.minimax_h3 import runtime as rt +from rl_engine.models.minimax_h3 import transition as tr + + +def rel(a, b): + a, b = a.double(), b.double() + return ((a - b).norm() / b.norm().clamp_min(1e-300)).item() + + +def lora_rows(): + out = [] + for n, k, t in ((512, 384, 70), (2048, 1024, 257), (5120, 5120, 33)): + g = torch.Generator().manual_seed(n) + w = (torch.randn(n, k, generator=g) / k**0.5).bfloat16().cuda() + x = torch.randn(t, k, generator=g).bfloat16().cuda() + grad = torch.randn(t, n, generator=g).bfloat16().cuda() + gate = dict(rt.lora_gradient_gate(w, x, grad)) + # gradients against an FP64 autograd of the same function + r, scale = 8, 2.0 + gen = torch.Generator().manual_seed(0) + a0 = (torch.randn(r, k, generator=gen) / k**0.5).cuda() + b0 = (torch.randn(n, r, generator=gen) * 0.1).cuda() + a, b = a0.clone().requires_grad_(), b0.clone().requires_grad_() + xi = x.clone().requires_grad_() + rt.strict_lora_linear(xi, w, a, b, scale).backward(grad) + ad, bd, xd = (v.detach().double().requires_grad_() for v in (a0, b0, x)) + y = xd @ w.double().t() + scale * (xd @ ad.t()) @ bd.t() + (y * grad.double()).sum().backward() + gate.update( + shape=f"N{n} K{k} T{t}", + grad_a_rel_err_vs_fp64=rel(a.grad, ad.grad), + grad_b_rel_err_vs_fp64=rel(b.grad, bd.grad), + grad_x_rel_err_vs_fp64=rel(xi.grad, xd.grad), + ) + out.append(gate) + return out + + +def scheduler_cps(x, v, sigma, sigma_prev, noise, level): + s = math.sin(level * math.pi / 2) + std = sigma_prev * s + mean = (x - sigma * v) * (1 - sigma_prev) + (x + v * (1 - sigma)) * torch.sqrt( + sigma_prev**2 - std**2 + ) + prev = mean + std * noise + return prev, mean, (-((prev - mean) ** 2)).mean(dim=tuple(range(1, prev.ndim))) + + +def transition_rows(verl_omni): + sched = None + if verl_omni: + import importlib.util + + spec = importlib.util.spec_from_file_location( + "flow_match_sde", + Path(verl_omni) / "verl_omni/pipelines/schedulers/flow_match_sde.py", + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + FlowMatchSDEDiscreteScheduler = module.FlowMatchSDEDiscreteScheduler + + sched = FlowMatchSDEDiscreteScheduler(num_train_timesteps=1000, shift=1.0) + sched.sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + sched.timesteps = sched.sigmas[:-1] * 1000 + sched._step_index = 1 + rows = [] + for batch, shape in ((3, (257, 96)), (4, (4096, 96)), (2, (8192, 32))): + g = torch.Generator().manual_seed(1) + x = torch.randn(batch, *shape, generator=g).cuda() + v = torch.randn(batch, *shape, generator=g).cuda() + noise = torch.stack([tr.step_noise(7, 2, i, "video", shape, "cuda") for i in range(batch)]) + sg, sp = torch.tensor(0.75, device="cuda"), torch.tensor(0.5, device="cuda") + out = tr.cps_transition(x, v, sg, sp, noise=noise) + prev, mean, logp = scheduler_cps(x, v, sg, sp, noise, tr.NOISE_LEVEL) + row = dict( + shape=f"B{batch} x {shape[0]}x{shape[1]}", + prev_sample_bitwise=bool(torch.equal(out.prev_sample, prev)), + mean_bitwise=bool(torch.equal(out.prev_sample_mean, mean)), + logp_max_rel_diff_vs_torch_mean=((out.log_prob - logp).abs() / logp.abs()).max().item(), + ) + for i in range(batch): + one = tr.cps_transition(x[i : i + 1], v[i : i + 1], sg, sp, noise=noise[i : i + 1]) + if not ( + torch.equal(one.prev_sample, out.prev_sample[i : i + 1]) + and torch.equal(one.log_prob, out.log_prob[i : i + 1]) + ): + row["batch_invariant"] = False + break + else: + row["batch_invariant"] = True + scored = tr.cps_transition(x, v, sg, sp, prev_sample=out.prev_sample) + row["scoring_reproduces_sampling_logprob"] = bool( + torch.equal(scored.log_prob, out.log_prob) + ) + if sched is not None: + ref = sched.sample_previous_step( + sample=x.cpu(), + model_output=v.cpu(), + prev_sample=out.prev_sample.cpu(), + noise_level=tr.NOISE_LEVEL, + sde_type=tr.SDE_TYPE, + ) + row["verl_omni_mean_bitwise"] = bool(torch.equal(ref[2], out.prev_sample_mean.cpu())) + row["verl_omni_logp_max_rel_diff"] = ( + ((ref[1] - out.log_prob.cpu()).abs() / ref[1].abs()).max().item() + ) + rows.append(row) + return rows + + +def sized_inputs(device, *, text, frames, side, audio): + """A bigger packed request (``side`` x ``side`` latent grid per frame).""" + + from rl_engine.validation.models.h3_cases import h3_fl2va_layout + + layout = h3_fl2va_layout( + num_text_tokens=text, + num_latent_frames=frames, + latent_height=side, + latent_width=side, + num_audio_latents=audio, + device=device, + ) + tags = layout["token_tags"] + video, text_idx, audio_idx = ((tags == t).nonzero().flatten() for t in (0, 1, 2)) + gen = torch.Generator().manual_seed(420539) + + def rnd(*shape): + return torch.randn(shape, generator=gen).to(device).requires_grad_() + + return dict( + **layout, + hidden_states=rnd(1, len(video), 96), + audio_hidden_states=rnd(1, len(audio_idx), 32), + encoder_hidden_states=rnd(1, len(text_idx), 5120), + timestep=torch.tensor([0.0, 500.0], device=device), + video_indices=video, + text_indices=text_idx, + audio_indices=audio_idx, + return_dict=False, + ) + + +SIZES = { + "12": None, + "925": dict(text=77, frames=3, side=32, audio=40), + "4253": dict(text=77, frames=1, side=128, audio=40), + "8349": dict(text=77, frames=2, side=128, audio=40), +} + + +def model_rows(checkpoint, sizes, no_ablations=False): + from diffusers import MiniMaxH3Transformer3DModel + + from tools.validation.models.h3_full_model import packed_inputs + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(Path(checkpoint) / "transformer"), + torch_dtype=torch.bfloat16, + local_files_only=True, + low_cpu_mem_usage=True, + ).to("cuda") + model.eval().enable_gradient_checkpointing() + for p in model.parameters(): + p.register_post_accumulate_grad_hook(lambda q: setattr(q, "grad", None)) + gi = ("hidden_states", "audio_hidden_states", "encoder_hidden_states") + + def run(base, mode, spec=""): + inputs = { + k: (v.detach().clone().requires_grad_() if k in gi else v) for k, v in base.items() + } + provider = rt.H3TrainingProvider(model, mode, rt.AblationOverlay.parse(spec)) + with provider.step(), torch.autograd.graph.save_on_cpu(pin_memory=True): + video, audio = provider.forward(**inputs) + (video.float().square().mean() + audio.float().square().mean()).backward() + rep = provider.adapter.report() + trans = provider.transition( + inputs["hidden_states"].detach(), + video.detach(), + torch.tensor(0.75, device="cuda"), + torch.tensor(0.5, device="cuda"), + seed=3, + step=1, + sample_ids=[0], + modality="video", + ) + out = dict( + video=video.detach().cpu(), + audio=audio.detach().cpu(), + grad=inputs["hidden_states"].grad.detach().cpu(), + report=rep, + logp=trans.log_prob.detach().cpu(), + ) + del video, audio, inputs, trans + torch.cuda.empty_cache() + return out + + variants = ( + ("off", "off", ""), + ("strict", "strict", ""), + ("strict (repeat)", "strict", ""), + ("audit", "audit", ""), + ("strict, node:attention native", "strict", "node:attention"), + ("strict, boundary:token_refiner native", "strict", "boundary:token_refiner"), + ) + if no_ablations: + variants = tuple(v for v in variants if not v[2]) + rows, checks = [], {} + for size in sizes: + base = packed_inputs("cuda") if SIZES[size] is None else sized_inputs("cuda", **SIZES[size]) + tokens = int(base["position_ids"].shape[0]) + ref = run(base, "off") # the provider run must fit; the others may report OOM + runs = {} + for name, mode, spec in variants: + try: + r = runs[name] = run(base, mode, spec) + except torch.OutOfMemoryError: + import traceback + + traceback.print_exc() + torch.cuda.empty_cache() + rows.append(dict(tokens=tokens, variant=name, error="CUDA out of memory")) + continue + rows.append( + dict( + tokens=tokens, + variant=name, + label=r["report"]["label"], + case_id=r["report"]["case_id"], + bound=r["report"]["bound"], + video_rel_diff_vs_provider=rel(r["video"], ref["video"]), + audio_rel_diff_vs_provider=rel(r["audio"], ref["audio"]), + input_grad_rel_diff_vs_provider=rel(r["grad"], ref["grad"]), + logp_rel_diff_vs_provider=rel(r["logp"], ref["logp"]), + boundary_bindings=r["report"]["boundary_bindings"], + ) + ) + if not all(k in runs for k in ("strict", "strict (repeat)", "audit")): + continue + a, b, c = runs["strict"], runs["strict (repeat)"], runs["audit"] + keys = ("video", "audio", "grad", "logp") + checks[str(tokens)] = dict( + strict_bitwise_repeatable=all(torch.equal(a[k], b[k]) for k in keys), + audit_equals_strict=all(torch.equal(a[k], c[k]) for k in keys), + ) + del ref, runs, a, b, c + torch.cuda.empty_cache() + return rows, checks + + +def plot(report, path): + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + fig, ax = plt.subplots(1, 4, figsize=(22, 5)) + lora = report["lora_gate"] + xs = range(len(lora)) + w = 0.25 + for j, key in enumerate( + ("grad_a_rel_err_vs_fp64", "grad_b_rel_err_vs_fp64", "grad_x_rel_err_vs_fp64") + ): + ax[0].bar( + [i + (j - 1) * w for i in xs], + [r[key] for r in lora], + w, + label=key.replace("_rel_err_vs_fp64", ""), + ) + ax[0].axhline(5e-3, color="r", ls="--", lw=1, label="gate 5e-3") + ax[0].set_xticks(list(xs), [r["shape"] for r in lora], fontsize=8) + ax[0].set_yscale("log") + ax[0].set_title("LoRA gradients vs FP64 (relative error)") + ax[0].legend(fontsize=7) + + tra = report["transition"] + ax[1].bar(range(len(tra)), [max(r["logp_max_rel_diff_vs_torch_mean"], 1e-12) for r in tra]) + ax[1].set_xticks(range(len(tra)), [r["shape"] for r in tra], fontsize=8) + ax[1].set_yscale("log") + ax[1].set_title("CPS log-prob vs torch.mean (max rel diff);\nprev_sample/mean bitwise equal") + + if "model" in report: + rows = [ + r + for r in report["model"] + if "error" not in r and r["variant"] not in ("off", "strict (repeat)", "audit") + ] + for axis, key, title in ( + (ax[2], "video_rel_diff_vs_provider", "33B model: video output"), + (ax[3], "input_grad_rel_diff_vs_provider", "33B model: input gradient"), + ): + for name in dict.fromkeys(r["variant"] for r in rows): + pts = [(r["tokens"], r[key]) for r in rows if r["variant"] == name] + axis.plot(*zip(*pts), "o-", label=name) + axis.set_xscale("log") + axis.set_yscale("log") + axis.set_xlabel("tokens") + axis.set_title(title + ", relative difference to provider") + axis.legend(fontsize=6) + fig.tight_layout() + fig.savefig(path, dpi=130) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--out", type=Path, required=True) + p.add_argument("--png", type=Path) + p.add_argument("--checkpoint") + p.add_argument("--verl-omni") + p.add_argument("--no-ablations", action="store_true", help="skip the ablation variants") + p.add_argument("--sizes", default="12,925,4253", help="comma list from: " + ",".join(SIZES)) + args = p.parse_args() + modes = { + m: vars(rt.h3_runtime_policy(m)) | {"mode": m} for m in ("strict", "audit", "auto", "off") + } + report = dict( + device=torch.cuda.get_device_name(0), + torch=torch.__version__, + runtime_modes={k: {kk: str(vv) for kk, vv in v.items()} for k, v in modes.items()}, + lora_gate=lora_rows(), + transition=transition_rows(args.verl_omni), + ) + if args.checkpoint: + report["model"], report["model_checks"] = model_rows( + args.checkpoint, args.sizes.split(","), args.no_ablations + ) + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=1, default=str)) + if args.png: + plot(report, args.png) + print( + json.dumps({k: v for k, v in report.items() if k != "model"}, indent=1, default=str)[:3000] + ) + if "model" in report: + for r in report["model"]: + print({k: v for k, v in r.items() if k != "boundary_bindings"}) + print(report["model_checks"]) + + +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_charts.py b/tools/validation/models/h3_ws2_charts.py new file mode 100644 index 000000000..5bc27964d --- /dev/null +++ b/tools/validation/models/h3_ws2_charts.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Chart of the WS2 evidence: bitwise agreement and per-rank forward time per parallel layout.""" + +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(evidence: Path, out: Path) -> None: + reports = [json.loads(p.read_text()) for p in sorted(evidence.glob("h3_ws2_*.json"))] + labels = [] + times = [] + colors = [] + for r in reports: + labels.append(f"TP{r['tp']}xCP{r['cp']}{'+SP' if r.get('sp') else ''}") + ranks = r["forward_seconds_per_rank"] + times.append(sum(x["forward_s"] for x in ranks) / len(ranks)) + colors.append("tab:green" if r["all_bitwise"] else "tab:red") + fig, ax = plt.subplots(figsize=(8, 4)) + ax.bar(labels, times, color=colors) + ax.set( + ylabel="mean forward s per rank (H100)", title="WS2 layouts: green = bitwise == single GPU" + ) + for i, r in enumerate(reports): + ax.text( + i, times[i], f"{r['world']} GPU\n{len(set(r['hosts']))} node", ha="center", va="bottom" + ) + 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_ws2_evidence.py b/tools/validation/models/h3_ws2_evidence.py new file mode 100644 index 000000000..8ef5976fb --- /dev/null +++ b/tools/validation/models/h3_ws2_evidence.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ + --out reports/experiments/h3-tp-adaln-3mod-b200/report.json + python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) +from rl_engine.validation.models.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +SP_CASES = [ + {"batch": 1, "seq": 4097, "layout": "block", "seed": 1}, + {"batch": 1, "seq": 4097, "layout": "interleaved", "seed": 2}, + {"batch": 2, "seq": 4097, "layout": "block", "seed": 3}, + {"batch": 1, "seq": 32768, "layout": "block", "seed": 4}, + {"batch": 1, "seq": 32768, "layout": "interleaved", "seed": 5}, + {"batch": 1, "seq": 131072, "layout": "block", "seed": 6}, +] + + +def _sp_norm(worlds: list[int], _max_t: int) -> dict: + runs = [] + for world in worlds: + ranks = run_world(world, sp_region_sweep, {}, cases=SP_CASES) + runs.append({"sp": world, "ranks": ranks}) + equal = all(all(c["equal"].values()) for r in ranks for c in r["cases"]) + print(f"sp={world}: byte-equal to WS1 on every rank and case: {equal}") + naive = { + layout: naive_sp_mismatch(sp_case(1, 32768, layout=layout, seed=4), max(worlds)) + for layout in ("block", "interleaved") + } + region = "norm2(residual + gate_msa[row] * y) * (1 + scale_mlp[row]) + shift_mlp[row]" + return { + "region": region, + "hidden": 5376, + "runs": runs, + "naive_sp": {"sp": max(worlds), **naive}, + } + + +OPS = {"tp_adaln_3mod": _tp_adaln, "sp_norm_adaln": _sp_norm} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py new file mode 100644 index 000000000..85983fc0a --- /dev/null +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot ``h3_bitwise_reuse.py`` reports: share of forward outputs that are not bitwise equal. + +One heatmap per report. Rows are implementations, columns are the cases compared against +diffusers' path and against the rl-kernel op; 0 means bitwise identical. + + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +_REFS = ("diffusers[plain]", "diffusers", "rl_kernel") + + +def _cell(entry: dict | None) -> float | None: + if not entry: + return None + out = entry["out"][0] + if out is None or "differ" not in out: + return None + return 0.0 if out["equal"] else out["differ"] / out["numel"] + + +def _text(frac: float | None) -> str: + if frac is None: + return "—" + if frac == 0: + return "0" + return f"{frac:.1%}" if frac >= 0.001 else f"{frac:.2%}" + + +def _panel(ax, report: dict) -> None: + results = {k: v for k, v in report["results"].items() if "unavailable" not in v} + cases = report["cases"] + diffusers = next((r for r in _REFS[:2] if r in results), None) + refs = [r for r in (diffusers, "rl_kernel") if r] + rows = ["rl_kernel"] + sorted(r for r in results if r not in _REFS) + cols = [(ref, case) for ref in refs for case in cases] + grid = np.full((len(rows), len(cols)), np.nan) + for i, row in enumerate(rows): + for j, (ref, case) in enumerate(cols): + frac = _cell(results.get(row, {}).get(f"vs {ref}", {}).get(str(case))) + if frac is not None: + grid[i, j] = frac + ax.text( + j, + i, + "—" if row == ref else _text(frac), + ha="center", + va="center", + fontsize=7.5, + color="white" if frac is not None and frac > 0.3 else "black", + fontweight="bold" if frac == 0 else None, + ) + ax.imshow(np.sqrt(np.nan_to_num(grid)), cmap="Reds", vmin=0, vmax=1, aspect="auto") + ax.set_yticks(range(len(rows)), rows, fontsize=8) + ax.set_xticks(range(len(cols)), [str(c) for _, c in cols], fontsize=7, rotation=45) + transform = ax.get_xaxis_transform() + for k, ref in enumerate(refs): + if k: + ax.axvline(k * len(cases) - 0.5, color="k", lw=1.5) + label = "vs rl-kernel" if ref == "rl_kernel" else "vs diffusers" + ax.text( + k * len(cases) + (len(cases) - 1) / 2, + 1.01, + label, + ha="center", + va="bottom", + fontsize=8.5, + transform=transform, + ) + ax.set_title(f"{report['op']} @ {report['rl_kernel_commit'][:8]}", fontsize=10, pad=18) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reports", nargs="+", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + reports = [json.loads(p.read_text()) for p in args.reports] + fig, axes = plt.subplots(1, len(reports), figsize=(6.4 * len(reports), 5.8), squeeze=False) + for ax, report in zip(axes[0], reports): + _panel(ax, report) + gpu = reports[0]["environment"].get("gpu", "") + fig.suptitle( + "Forward outputs that are not bitwise equal\n" + f"(share of elements; 0 = bitwise identical), {gpu}", + fontsize=10, + ) + fig.tight_layout() + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=125) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_block.py b/tools/validation/models/plot_h3_block.py new file mode 100644 index 000000000..228f3a289 --- /dev/null +++ b/tools/validation/models/plot_h3_block.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render the RFC #420 one-block replay (``tools/validation/models/h3_block_replay.py``) as a PNG. + +Needs only the report JSON and matplotlib: + + python tools/validation/models/plot_h3_block.py \ + reports/experiments/h3-one-block-b200/report.json + +writes ``figure.png`` next to the report. Panels: the per-node invariance and +provider checks over every layout; the per-node error against the FP32 golden +for the candidate and diffusers; the same for every backward leaf. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import ListedColormap # noqa: E402 + +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +PASS = "#1baf7a" +FAIL = "#d64545" +NA = "#e4e3dd" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" +STATUS_MARK = {"row": "", "interim": "*", "reference": "†"} + + +def _style(ax, title: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=10.5, color=INK, pad=8) + ax.tick_params(colors=MUTED, labelsize=7.5, length=0) + for side in ("top", "right", "left", "bottom"): + ax.spines[side].set_visible(False) + + +def _checks_panel(ax, report: dict[str, Any]) -> None: + """Rows: nodes. Columns: forward/backward checks, ANDed over every layout.""" + + nodes = [n["node"] for n in report["nodes"]] + status = {n["node"]: n["status"] for n in report["nodes"]} + cols = [ + ("repeat", "forward", "repeat_bitwise_equal"), + ("batch row", "forward", "batch_row_bitwise_equal"), + ("= diffusers\n(isolated)", "forward", None), + ("grad\nrepeat", "backward", "grad_repeat_bitwise_equal"), + ("grad\nbatch row", "backward", "grad_batch_row_bitwise_equal"), + ] + grid = [] + for node in nodes: + row = [] + for _label, kind, key in cols: + cases = report["forward_cases" if kind == "forward" else "backward_cases"] + entries = [e for case in cases for e in case["nodes"] if e["node"] == node] + if key is None: + promised = entries[0]["provider_bitwise_isolated_promised"] + ok = all(e["isolated_vs_provider"]["bitwise_equal"] for e in entries) + row.append(2 if not promised and not ok else (1 if ok else 0)) + else: + values = [e[key] for e in entries] + row.append(2 if not values or None in values else (1 if all(values) else 0)) + grid.append(row) + ax.imshow(grid, cmap=ListedColormap([FAIL, PASS, NA]), vmin=0, vmax=2, aspect="auto") + ax.set_xticks(range(len(cols)), [c[0] for c in cols], fontsize=7.5, color=INK) + ax.set_yticks( + range(len(nodes)), [f"{n}{STATUS_MARK[status[n]]}" for n in nodes], fontsize=7.5, color=INK + ) + ax.xaxis.tick_top() + _style(ax, "Bitwise checks, all layouts") + ax.set_xlabel( + "green = holds grey = not promised (reduction differs from diffusers)\n" + "or not applicable (AdaLN tables have no batch axis)\n" + "* interim operator † provider replay (no kernel yet)", + fontsize=7, + color=MUTED, + ) + + +def _err_panel(ax, labels, cand, prov, title: str) -> None: + xs = range(len(labels)) + ax.plot( + xs, [math.log10(v) for v in cand], "o-", color=RL_KERNEL, lw=1.8, ms=5, label="RL-Kernel" + ) + ax.plot( + xs, [math.log10(v) for v in prov], "s--", color=PROVIDER, lw=1.4, ms=4, label="diffusers" + ) + ax.set_xticks(list(xs), labels, rotation=55, ha="right", fontsize=7, color=INK) + lo = math.floor(min(math.log10(v) for v in cand + prov)) + hi = math.ceil(max(math.log10(v) for v in cand + prov)) + ax.set_yticks(range(lo, hi + 1), [f"1e{v}" for v in range(lo, hi + 1)]) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_ylabel("max|err| / max|golden|", color=MUTED, fontsize=8) + ax.legend(frameon=False, fontsize=7.5, labelcolor=INK, loc="upper left") + _style(ax, title) + + +def plot(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(16, 5.6), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.0, 1.6, 1.3]} + ) + case = max(report["forward_cases"], key=lambda c: c["seq_len"]) + back = ( + max(report["backward_cases"], key=lambda c: c["seq_len"]) + if report["backward_cases"] + else None + ) + _checks_panel(axes[0], report) + _err_panel( + axes[1], + [e["node"] for e in case["nodes"]], + [e["candidate_rel_err_vs_golden"] for e in case["nodes"]], + [e["provider_rel_err_vs_golden"] for e in case["nodes"]], + f"Forward error vs FP32 golden, S = {case['seq_len']}", + ) + if back is not None: + leaves = list(back["leaves"]) + _err_panel( + axes[2], + [f"d {name}" for name in leaves], + [back["leaves"][n]["candidate_rel_err_vs_golden"] for n in leaves], + [back["leaves"][n]["provider_rel_err_vs_golden"] for n in leaves], + f"Gradient error vs FP32 golden, S = {back['seq_len']}", + ) + env = report["environment"] + sizes = ", ".join(str(c["seq_len"]) for c in report["forward_cases"]) + fig.suptitle( + "MiniMax-H3 block 0, forward and backward (RFC #420 ws1_one_h3_block)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.905, + f"{env['gpu']} | torch {env['torch']} | layouts S = {sizes} | " + f"first drift vs diffusers: {case['first_drift']} | " + f"rl-kernel {report['rl_kernel_commit'][:7]}", + fontsize=8, + color=MUTED, + ) + fig.tight_layout(rect=(0, 0, 1, 0.89)) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + out = args.report.with_name("figure.png") + plot(report).savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py new file mode 100644 index 000000000..467cb5b7d --- /dev/null +++ b/tools/validation/models/plot_h3_evidence.py @@ -0,0 +1,768 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator or WS2 report (``tools/validation/models/h3_evidence.py``, +``tools/validation/models/h3_ws2_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python tools/validation/models/plot_h3_evidence.py \\ + reports/experiments/h3-timestep-sinusoid-b200/report.json + +writes ``figure.png`` next to the report. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Callable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 +import matplotlib.pyplot as plt # noqa: E402 + +# Fixed series identity across every H3 figure (categorical slots 1-3). +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +THIRD = "#1baf7a" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" + + +def _style(ax, title: str, ylabel: str) -> None: + """Apply the shared report axis styling, title, and vertical-axis label.""" + + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_ylabel(ylabel, color=MUTED, fontsize=9) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + for side in ("top", "right", "left"): + ax.spines[side].set_visible(False) + ax.spines["bottom"].set_color(GRID) + + +def _bars(ax, labels: list[str], series: list[tuple], fmt: str, log: bool = False) -> None: + """Draw grouped series with formatted value labels and a legend; a 4th tuple item is a hatch.""" + + n = len(series) + width = 0.8 / n + for i, (name, color, values, *hatch) in enumerate(series): + xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] + bars = ax.bar( + xs, + values, + width * 0.92, + color=color if not hatch else SURFACE, + edgecolor=color, + hatch=hatch[0] if hatch else None, + linewidth=1.2, + label=name, + zorder=2, + ) + for bar, value in zip(bars, values, strict=True): + ax.annotate( + fmt.format(value), + (bar.get_x() + bar.get_width() / 2, bar.get_height()), + xytext=(0, 2), + textcoords="offset points", + ha="center", + fontsize=7, + color=INK, + ) + ax.set_xticks(range(len(labels)), labels) + top = max(v for _, _, values, *_ in series for v in values) + if log: + ax.set_yscale("log") + ax.set_ylim(min(v for _, _, values, *_ in series for v in values) / 2, top * 8) + else: + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=min(n, 4)) + + +# Exact zeros cannot sit on a log axis; draw them on this floor and label them. +LOG_FLOOR = 1e-9 + + +def _log_values(values: list[float]) -> list[float]: + """Replace nonpositive values with the plotting floor for logarithmic axes.""" + + return [v if v > 0 else LOG_FLOOR for v in values] + + +def _figure(report: dict[str, Any], suptitle: str): + """Create two report axes with a title and GPU, software, and revision provenance.""" + + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) + env = report["environment"] + fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") + fig.text( + 0.01, + 0.01, + f"{env['gpu']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · MiniMax-H3@{report['model_revision'][:7]}", + fontsize=7, + color=MUTED, + ) + return fig, axes + + +def plot_sinusoid(report: dict[str, Any]): + """Return a sinusoid report figure showing latency and errors against FP64.""" + + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") + perf = report["perf"] + _style(left, "Latency per call (lower is better)", "µs") + _bars( + left, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + ("CUDA + range check", THIRD, [row["candidate_checked_us"] for row in perf]), + ("diffusers", PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + cases = report["accuracy"]["cases"] + xs = [c["num_timesteps"] for c in cases] + _style(right, "Max |error| vs FP64 golden (log)", "max abs error") + cuda = [c["max_abs_vs_fp64"] for c in cases] + provider = [c["provider_max_abs_vs_fp64"] for c in cases] + right.plot(xs, _log_values(cuda), "o-", color=RL_KERNEL, lw=2, ms=8, label="RL-Kernel CUDA") + right.plot(xs, _log_values(provider), "s--", color=PROVIDER, lw=2, ms=5, label="diffusers") + atol = report["accuracy"]["contract_atol"] + right.axhline(atol, color=MUTED, lw=1, ls=":") + right.annotate( + "contract atol 1e-5", + (xs[-1], atol), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + for x, value in zip(xs, cuda, strict=True): + if value == 0: + right.annotate( + "exact", + (x, LOG_FLOOR), + xytext=(0, 8), + textcoords="offset points", + fontsize=7, + color=INK, + ha="center", + ) + right.set_xscale("log") + right.set_yscale("log") + right.set_ylim(LOG_FLOOR / 3, atol * 30) + right.set_xlabel("number of timesteps T", color=MUTED, fontsize=9) + equal = sum(c["bitwise_equal_to_provider"] for c in cases) + right.legend( + frameon=False, + fontsize=8, + labelcolor=INK, + loc="upper left", + ncol=2, + title=f"bitwise equal to diffusers at {equal}/{len(cases)} sizes", + title_fontsize=8, + alignment="left", + ) + return fig + + +def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + """Plot candidate and provider latency summaries in microseconds for each case.""" + + _style(ax, "Latency per call (lower is better)", "µs") + _bars( + ax, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + (provider_name, PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + + +def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], contract: float): + """Per-draw max |error| as boxes with every draw overlaid, log scale.""" + + _style(ax, title, "max abs error per draw") + for i, (_name, color, values) in enumerate(series): + ax.boxplot( + [values], + positions=[i], + widths=0.45, + showfliers=False, + patch_artist=True, + boxprops={"facecolor": color, "alpha": 0.25, "edgecolor": color}, + medianprops={"color": color, "linewidth": 2}, + whiskerprops={"color": color}, + capprops={"color": color}, + ) + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax.scatter(jitter, values, s=8, color=color, alpha=0.6, linewidths=0, zorder=3) + median = sorted(values)[len(values) // 2] + ax.annotate( + f"median {median:.1e}", + (i + 0.28, median), + fontsize=8, + color=INK, + va="center", + ) + ax.axhline(contract, color=MUTED, lw=1, ls=":") + ax.annotate( + f"contract atol {contract:g}", + (len(series) - 0.5, contract), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + ax.set_yscale("log") + ax.set_xticks(range(len(series)), [name for name, _, _ in series]) + ax.set_xlim(-0.6, len(series) - 0.1) + + +def plot_mlp(report: dict[str, Any]): + """Return an MLP report figure showing latency and per-draw FP64 error distributions.""" + + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _error_boxes( + right, + f"Error vs FP64 golden, {acc['draws']} draws of T={acc['num_timesteps']}", + [ + ("RL-Kernel CUDA", RL_KERNEL, acc["cuda_max_abs_vs_fp64"]), + ("diffusers (cuBLAS)", PROVIDER, acc["provider_max_abs_vs_fp64"]), + ], + acc["contract_atol"], + ) + return fig + + +def plot_projection(report: dict[str, Any]): + """Return an AdaLN figure with latency, BF16 rounding accuracy, and the early-cast probe.""" + + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _style( + right, + f"Outputs equal to the correctly rounded FP64 golden, {acc['draws']} draws", + "% of BF16 outputs", + ) + series = [ + ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), + ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), + ] + for i, (_name, color, values) in enumerate(series): + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax_vals = sorted(values) + median = ax_vals[len(ax_vals) // 2] + right.scatter(jitter, values, s=36, color=color, alpha=0.8, linewidths=0, zorder=3) + right.hlines(median, i - 0.25, i + 0.25, color=color, lw=2.5, zorder=4) + right.annotate( + f"median {median:.3f}%", (i + 0.28, median), fontsize=8, color=INK, va="center" + ) + right.set_xticks(range(len(series)), [name for name, _, _ in series]) + right.set_xlim(-0.6, len(series) - 0.1) + early = sorted(acc["early_cast_golden_match"])[len(acc["early_cast_golden_match"]) // 2] + right.text( + 0.98, + 0.62, + f"an early BF16 cast (probe H7)\nwould match only {100 * early:.0f}%", + transform=right.transAxes, + fontsize=8, + color=MUTED, + ha="right", + ) + return fig + + +def plot_gather(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_row_gather (T = 3, H = 5376, BF16)") + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("index_select fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "index_select bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + chain = report["accuracy"]["chain_backward"] + _style( + right, + "Whole-chain grads of the FP32 time embedder vs FP64", + "max abs error / golden max", + ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig + names = { + "candidate": ("RL-Kernel, separate ops", THIRD, "o"), + "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), + "provider": ("diffusers", PROVIDER, "s"), + } + labels = [f"T={c['num_timesteps']}\nS={c['seq_len']}" for c in chain] + for j, (mode, (name, color, marker)) in enumerate(names.items()): + leaves = [e for c in chain for e in c["leaves"].values()] + det = sum(e[mode]["repeat_bitwise_equal"] for e in leaves) + adaln = min( + e[mode]["correctly_rounded_fraction"] + for c in chain + for n, e in c["leaves"].items() + if n.startswith("transformer_blocks") + ) + label = ( + f"{name}\n repeat-bitwise {det}/{len(leaves)} · " + f"AdaLN BF16 grads ≥{100 * adaln:.1f}% correctly rounded" + ) + for i, case in enumerate(chain): + values = [ + e[mode]["max_abs_vs_golden_over_absmax"] + for n, e in case["leaves"].items() + if n.startswith("time_embedder") + ] + right.scatter( + [i + (j - 1) * 0.22] * len(values), + values, + s=30, + color=color, + marker=marker, + alpha=0.85, + linewidths=0, + zorder=3, + label=label if i == 0 else None, + ) + right.set_yscale("log") + right.set_ylim(1e-8, 1e4) + right.set_xticks(range(len(labels)), labels, fontsize=7) + right.legend(frameon=False, fontsize=7, labelcolor=INK, loc="upper left") + return fig + + +def plot_rmsnorm(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "h3_rmsnorm (block norm1 + MSA modulation, H = 5376, BF16)" + ) + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("diffusers fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "diffusers bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + acc = report["accuracy"] + bwd = acc["backward"] + keys = ["dx", "dweight", "dshift", "dscale"] + fwd_ok = acc["modulated_bitwise_vs_diffusers"] and all( + acc["plain_bitwise_vs_nn_rmsnorm"].values() + ) + _style(right, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def repeat(mode: str) -> str: + return "yes" if bwd[mode]["repeat_bitwise_equal"] else "no" + + _bars( + right, + keys, + [ + ( + f"RL-Kernel CUDA (repeat-bitwise: {repeat('cuda')})", + RL_KERNEL, + [bwd["cuda"]["rel_error"][k] for k in keys], + ), + ( + f"diffusers (repeat-bitwise: {repeat('provider')})", + PROVIDER, + [bwd["provider"]["rel_error"][k] for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + right.text( + 0.98, + 0.80, + f"forward bitwise equal to diffusers: {'yes' if fwd_ok else 'NO'}", + transform=right.transAxes, + ha="right", + fontsize=8, + color=INK, + ) + return fig + + +def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + combined = all(row.get("backward_timing_scope") == "forward_and_backward" for row in perf) + backward_label = "fwd+bwd" if combined else "bwd" + _style(ax, "Forward and backward time (log, lower is better)", "ms") + _bars( + ax, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ( + f"CUDA {backward_label}", + RL_KERNEL, + [row["candidate_backward_us"] / 1e3 for row in perf], + "////", + ), + ( + f"{provider} {backward_label}", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + + +def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: + _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + for text in ax.texts: + if text.get_text() == f"{floor:.1e}": + text.set_text("exact") + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=1) + ax.set_ylim(floor / 3, 1e4) + ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + +def plot_final(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "final_adaln_out (norm_out: projection + final norm, T = 3)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + _backward_errors( + right, + acc["backward"], + ["dx", "d_norm_w", "d_temb", "dW", "db"], + f"forward equal to diffusers on {100 * acc['equal_to_diffusers_fraction']:.2f}% " + "of elements (projection 1-ULP ties)", + ) + return fig + + +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" + ) + runs = report["runs"] + _equality_grid( + left, + "Byte-equal to WS1", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "ranks x T", + ) + # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. + num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + return fig + + +def _equality_grid(ax, title, rows, tensors, unit) -> None: + """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" + + ax.set_facecolor(SURFACE) + ax.set_title(f"{title}; checks = {unit}", loc="left", fontsize=11, color=INK, pad=10) + for y, (_, checks) in enumerate(rows): + for x, (key, _) in enumerate(tensors): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + ax.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)}" + ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") + ax.set_xlim(0, len(tensors)) + ax.set_ylim(len(rows), 0) + ax.set_xticks([i + 0.5 for i in range(len(tensors))], [label for _, label in tensors]) + ax.set_yticks([i + 0.5 for i in range(len(rows))], [name for name, _ in rows]) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + for side in ax.spines.values(): + side.set_visible(False) + + +SP_TENSORS = ( + ("out", "output rows"), + ("d_residual", "d_residual"), + ("d_y", "d_sublayer"), + ("d_weight", "d_norm_w"), + ("d_table", "d_table"), +) + + +def plot_sp_norm(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(15, 3.9), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.35, 1, 0.9]} + ) + env = report["environment"] + fig.suptitle( + "sp_norm_adaln (sequence-parallel norm2(residual + gate * y) with AdaLN modulation, " + "H = 5376, NCCL)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.01, + f"{env['gpu']} x {env['gpus']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · bf16", + fontsize=7, + color=MUTED, + ) + runs = report["runs"] + _equality_grid( + axes[0], + "Byte-equal to WS1", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "ranks x cases", + ) + + # Middle: the largest block-layout case, slowest rank. + big = max( + (c for c in runs[0]["ranks"][0]["cases"] if c["layout"] == "block"), + key=lambda c: c["seq"] * c["batch"], + ) + + def worst(run, key): + return ( + max( + c[key] + for r in run["ranks"] + for c in r["cases"] + if (c["seq"], c["batch"], c["layout"]) == (big["seq"], big["batch"], big["layout"]) + ) + / 1e3 + ) + + sp_runs = [run for run in runs if run["sp"] > 1] + ax = axes[1] + _style(ax, f"Region fwd+bwd, S = {big['seq']}, block layout", "ms (slowest rank)") + _bars( + ax, + [f"SP{run['sp']}" for run in sp_runs], + [("SP", RL_KERNEL, [worst(run, "sp_forward_backward_us") for run in sp_runs])], + "{:.2f}", + ) + ws1 = worst(runs[0], "ws1_forward_backward_us") + ax.axhline( + ws1, color=PROVIDER, linestyle="--", linewidth=1.4, zorder=1, label=f"WS1, 1 GPU: {ws1:.2f}" + ) + ax.set_ylim(0, ws1 * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper right") + + # Right: what a naive SP reduction would do instead. + naive = report["naive_sp"] + ax = axes[2] + _style(ax, f"Naive SP{naive['sp']} vs WS1 (S = 32768)", "% of elements that differ") + _bars( + ax, + ["d_norm_w", "d_table"], + [ + ("naive, block", PROVIDER, [100 * naive["block"][k] for k in ("d_weight", "d_table")]), + ( + "naive, interleaved", + PROVIDER, + [100 * naive["interleaved"][k] for k in ("d_weight", "d_table")], + "////", + ), + ("RL-Kernel SP", RL_KERNEL, [0.0, 0.0]), + ], + "{:.1f}", + ) + ax.text( + 0.98, + 0.62, + "naive: each rank's WS1 backward,\nthen a rank-order sum", + transform=ax.transAxes, + ha="right", + fontsize=7, + color=MUTED, + ) + return fig + + +PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { + "timestep_sinusoid_h3": plot_sinusoid, + "timestep_mlp_fp32": plot_mlp, + "adaln_projection_3mod": plot_projection, + "adaln_row_gather": plot_gather, + "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, + "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, +} + + +def main() -> None: + """Read an operator report and save its figure as a PNG beside it or at ``--out``.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + fig = PLOTS[report["op"]](report) + fig.tight_layout(rect=(0, 0.04, 1, 0.95)) + out = args.out or args.report.with_name("figure.png") + fig.savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_prior_art.py b/tools/validation/models/plot_h3_prior_art.py new file mode 100644 index 000000000..893612e87 --- /dev/null +++ b/tools/validation/models/plot_h3_prior_art.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot H3 prior-art latency, accuracy, and batch invariance. + + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/.json + +Writes ``.png`` next to the report. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib import transforms # noqa: E402 + + +def _entries(report): + for mode, cands in report["results"].items(): + if "unavailable" in cands: + continue + for name, entry in cands.items(): + if "unavailable" not in entry: + yield (name if name.endswith("]") or mode == "plain" else f"{name}[{mode}]"), entry + + +def _differ(bi) -> int: + rows = sum(v["fwd_differ"] + v["rowgrad_differ"] for v in bi["all_rows"].values()) + sub = bi["full_vs_sub_batches"] + return rows + sub["fwd_differ"] + sub["rowgrad_differ"] + bi["batch_size_sweep"]["differ"] + + +def main() -> None: + path = Path(sys.argv[1]) + report = json.loads(path.read_text()) + entries = list(_entries(report)) + names = [n for n, _ in entries] + fig, axes = plt.subplots(1, 3, figsize=(20, 6.5), layout="constrained") + fig.suptitle( + f"{report['op']} vs existing implementations — {report['environment']['gpu']}, " + f"commit {report['rl_kernel_commit'][:7]}", + fontsize=13, + ) + + ax = axes[0] + for name, entry in entries: + perf = entry["accuracy_latency"] + key = "fwd_bwd_us" if entry["has_backward"] else "fwd_us" + sizes = sorted(perf, key=int) + ax.plot( + [int(s) for s in sizes], + [perf[s][key] for s in sizes], + marker="o", + ls="-" if entry["has_backward"] else "--", + label=name + ("" if entry["has_backward"] else " (forward only)"), + ) + ax.set_xscale("log", base=2) + ax.set_yscale("log") + ax.set_xlabel("rows (timesteps or tokens)") + ax.set_ylabel("µs, median") + ax.set_title("forward + backward latency (dashed: forward only)") + ax.grid(True, which="both", alpha=0.3) + ax.legend(fontsize=7) + + ax = axes[1] + first = entries[0][1]["accuracy_latency"] + size = "3" if "3" in first else sorted(first, key=int)[-1] + ys = range(len(entries)) + fwd = [e["accuracy_latency"][size]["fwd_rel_err"] for _, e in entries] + grad = [ + max(e["accuracy_latency"][size].get("grad_rel_err", {}).values(), default=0) + for _, e in entries + ] + ax.barh([y - 0.2 for y in ys], fwd, 0.4, label="forward") + ax.barh([y + 0.2 for y in ys], [g or float("nan") for g in grad], 0.4, label="worst gradient") + values = [v for v in fwd + grad if v] + if values and max(values) / min(values) > 10: + ax.set_xscale("log") + inside = transforms.blended_transform_factory(ax.transAxes, ax.transData) + for y, a, b in zip(ys, fwd, grad): + ax.text( + 0.98, + y, + f"fwd {a:.1e} / grad {b:.1e}" if b else f"fwd {a:.1e}", + transform=inside, + ha="right", + va="center", + fontsize=7, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.8, "pad": 1}, + ) + ax.set_yticks(list(ys), names, fontsize=7) + ax.invert_yaxis() + ax.set_title(f"error vs FP64, max|err| / max|ref| (size {size})") + ax.legend(fontsize=7, loc="lower right") + ax.grid(True, axis="x", alpha=0.3) + + ax = axes[2] + for y, (_, e) in zip(ys, entries): + bi = e["batch_invariance"] + rep = bi["params_repeatable"] + lines = ["batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant"] + lines.append(f"{_differ(bi)} row comparisons differ") + if rep is not None: + lines.append("parameter/table grads " + ("repeatable" if rep else "NOT repeatable")) + ax.text( + 0.02, + y, + " | ".join(lines), + va="center", + fontsize=8, + color="#3aa676" if bi["batch_invariant"] else "#d14b4b", + ) + ax.set_ylim(len(entries) - 0.5, -0.5) + ax.set_xlim(0, 1) + ax.axis("off") + ax.set_title("batch invariance (bitwise)") + ax.text( + 0.02, + len(entries) - 0.3, + "every row alone vs full batches + full batch vs covering sub-batches + size sweep", + fontsize=7, + color="#555555", + ) + + out = path.with_suffix(".png") + fig.savefig(out, dpi=120) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py new file mode 100644 index 000000000..f6e32bcb5 --- /dev/null +++ b/tools/weights/prepare_h3_weights.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fetch and pin the MiniMax-H3 conditioning tensors used by the RFC #420 tests. + +Downloads only the config, the shard index and the shards that hold the +manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies +every sha256 against ``rl_engine/validation/models/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. ``--block`` also writes block 0's +attention/FFN weights (manifest ``block_tensors``, same shard) to +``/h3_block0.safetensors`` for the ``ws1_one_h3_block`` replay. + + python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 [--block] +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + BLOCK_FILE, + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, + sha256_tensor, +) + + +def _check(path: Path, expected: str) -> None: + """Verify a downloaded file's SHA-256, exiting on mismatch or printing success.""" + + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") + print(f"ok {path.name} {actual}") + + +def main() -> None: + """Download pinned model artifacts, verify tensor identities, and write extracted weights.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="output directory") + parser.add_argument( + "--download-dir", + type=Path, + default=None, + help="where the raw shard goes (default: /hf)", + ) + parser.add_argument( + "--block", action="store_true", help="also extract block 0 (manifest block_tensors)" + ) + args = parser.parse_args() + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + manifest = load_h3_manifest() + identity = manifest["model_identity"] + out_dir: Path = args.out.expanduser() + download_dir = (args.download_dir or out_dir / "hf").expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + + def fetch(filename: str) -> Path: + """Download one file from the manifest's model subfolder and pinned revision.""" + + return Path( + hf_hub_download( + identity["hf_repo"], + f"{identity['subfolder']}/{filename}", + revision=identity["revision"], + local_dir=str(download_dir), + ) + ) + + _check(fetch("config.json"), identity["config_sha256"]) + index_path = fetch(identity["index_file"]) + _check(index_path, identity["index_sha256"]) + weight_map = json.loads(index_path.read_text())["weight_map"] + + def extract(section: str, filename: str) -> None: + """Pull one manifest section's tensors from their pinned shards, verify, and save.""" + + specs = manifest[section] + tensors = {} + for shard in sorted({weight_map[name] for name in specs}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in specs: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + + target = out_dir / filename + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + + extract("tensors", EXTRACTED_FILE) + if args.block: + extract("block_tensors", BLOCK_FILE) + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main() 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()