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..cea040ccf 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -32,5 +32,15 @@ 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) - [H3 Flash Attention (optional fast backend)](h3-flash-attention.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/h3-adaln-gate-residual.md b/docs/operators/h3-adaln-gate-residual.md new file mode 100644 index 000000000..ffd200b0a --- /dev/null +++ b/docs/operators/h3-adaln-gate-residual.md @@ -0,0 +1,143 @@ +# MiniMax-H3 AdaLN Gated Residual + +## Summary + +`adaln_gate_residual` adds each sublayer's output back to the residual stream through the +per-row AdaLN gate (RFC #420, WS1 step 6). In every H3 block this happens twice, after attention +with `gate_msa` and after the FFN with `gate_mlp`: + +```text +hidden = residual + gate[adaln_indices] * sublayer_output # RFC #420 §4 order +``` + +`gate` is an `(R, H)` row view of the [AdaLN projection](h3-adaln-projection.md) table. The op +gathers the gate row inside the kernel, so the gathered `(S, H)` gate is never materialised. + +## Entry Point + +```python +from rl_engine.runtime.registry import kernel_registry + +op = kernel_registry.get_op("adaln_gate_residual", device="cuda") +hidden = op(residual, attn_output, gate_msa, adaln_indices) +``` + +## Backends + +| Backend | Wrapper | Native symbols | Status | +| --- | --- | --- | --- | +| CUDA | `rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp` | `rl_engine._C.h3_gate_residual_{forward,backward}` | Forward and `d_sublayer` bitwise equal to diffusers | +| PyTorch reference | `rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp` | n/a | Eager forward with a deterministic gate gradient; `forward_fp32`: FP64 golden | +| ROCm | n/a | n/a | Falls back to the PyTorch reference | + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `residual`, `sublayer_output` | `(..., S, H)` | bf16 / fp16 / fp32 | Same shape and dtype | +| `gate` | `(R, H)` row view | same | Unit column stride | +| `index` | `(S,)` | int64 | In `[0, R)`; shared across the batch | + +## Numerics + +- **Forward.** Each element computes `p = round(gate * y)` and then `out = round(residual + p)`. + `__fmul_rn`/`__fadd_rn` keep FP32 from contracting the two operations into one FMA. The + roundings fall exactly where the eager expression rounds, so the forward is **bitwise equal to + diffusers** in BF16, FP16 and FP32. Rows are independent. +- **Backward.** + - `d_residual` is the incoming gradient. + - `d_sublayer = round(grad * gate)`. The product of two 16-bit values is exact in FP32, so this + is bitwise equal to the eager VJP. + - `d_gate` is an FP32 segmented sum of `grad * sublayer_output` over positions sorted stably + by table row, using the same tiles as [`adaln_row_gather`](h3-adaln-row-gather.md), rounded + once. + + Diffusers' `d_gate` uses `index_select`'s BF16 atomic scatter-add instead. It is + non-deterministic and 13–49× further from the FP64 golden (its error varies from run to run). + +## Performance Notes + +```bash +python benchmarks/models/benchmark_h3_conditioning.py --op adaln_gate_residual +``` + +B200, B = 1, H = 5376, BF16. Backward timings exclude input and forward-graph setup: + +| S | CUDA fwd | diffusers fwd | CUDA bwd | diffusers bwd | +| --- | --- | --- | --- | --- | +| 4097 | 0.07 ms | 0.09 ms | 1.12 ms | 0.58 ms | +| 32768 | 0.24 ms | 0.42 ms | 1.12 ms | 2.24 ms | +| 131072 | 0.88 ms | 1.61 ms | 2.22 ms | 8.92 ms | + +The forward is one pass over `residual` and `sublayer_output`, using 16-byte vectors with one +block per row. At small S the backward pays a fixed cost of about 1 ms for the stable sort and +tile setup. + +## Evidence + +![adaln_gate_residual on B200: latency and backward accuracy](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/figure.png) + +The data is in [`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-adaln-gate-residual-b200/report.json), +written by `tools/validation/models/h3_evidence.py` from a clean tree at commit `bde1d8d`. It also records +forward bitwise equality with diffusers in bf16, fp16 and fp32, and row invariance. + +## Existing implementations (RFC #420 reuse rule) + +![gate_residual vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers residual + gate.index_select(...) * y | **no** (param/table grads not repeatable) | 3.0e-03 / 4.3e-02 / 637 µs | 2.9e-03 / 1.4e-01 / 2502 µs | +| rl-kernel H3GateResidualCudaOp | yes | 3.0e-03 / 2.6e-03 / 1210 µs | 2.9e-03 / 2.6e-03 / 1596 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-row batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-prior-art-b200/gate_residual.json) +was written from a clean tree at `4010854` by + +```bash +python tools/validation/models/h3_prior_art.py --op gate_residual --out reports/experiments/h3-prior-art-b200/gate_residual.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/gate_residual.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![gate_residual bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.png) + +There is no reuse candidate beyond diffusers' residual + gate.index_select(...) * y. The forward, d_residual and d_sublayer are bitwise identical to it at every S (1 to 32768). The gate-table gradient matches at S = 1 and differs from S = 257 on, where index_select accumulates with atomics and this kernel in a fixed FP32 order. +[`gate_residual.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr488/reports/experiments/h3-bitwise-reuse-b200/gate_residual.json) was written from a clean tree at `a746bff0` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op gate_residual \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/gate_residual.json \ + --out reports/experiments/h3-bitwise-reuse-b200/gate_residual.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_gate_residual.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, incl. the gated residual +python tools/validation/operators/check_operator.py --op adaln_gate_residual --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_gate_residual \ + --out reports/experiments/h3-adaln-gate-residual-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-gate-residual-b200/report.json +``` + +## Known Limitations + +- Until the attention and FFN rows land, the end-to-end chain feeds a seeded stand-in for the + sublayer output. The op itself is exercised on random and pinned-table inputs. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-adaln-projection.md b/docs/operators/h3-adaln-projection.md new file mode 100644 index 000000000..526e2b4cf --- /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 `579775f1` on an +otherwise idle B200 (regenerated after the golden's declared cast was changed to round the FP32 +SiLU, as the model does; the accuracy fields are unchanged on the report's 20 draws). The report also +records that a timestep's 18 modulation rows are bitwise identical whether it runs alone +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 `50d1e924` 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..923ad2f77 --- /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 `29b07c5d` with backward-only +timings and alternating execution order, and with the op-backward check's rows and gradients +drawn from a seeded generator. Before, they came from the process's default CUDA generator, +so the fraction depended on how many draws the runner made first (99.99965% at `d2a6546`). +`chain_replay.json` was written from a clean tree at commit `61d23ade`. Both were regenerated +after the AdaLN projection golden's declared cast was changed to round the FP32 SiLU, as the +model does; every CUDA pass/fail flag is unchanged. + +## 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 `62f8009e` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op adaln_row_gather \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.json \ + --out reports/experiments/h3-bitwise-reuse-b200/adaln_row_gather.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_adaln_row_gather.py tests/models/minimax_h3/test_h3_adaln_modulation.py -v +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # whole chain, forward + backward +python tools/validation/operators/check_operator.py --op adaln_row_gather --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 1365 --normalized-dim 5376 --check-grad +python tools/validation/models/h3_evidence.py --op adaln_row_gather \ + --out reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-adaln-row-gather-b200/report.json +python tools/validation/models/h3_chain_replay.py --timesteps 1,2,3,4 --seq-lens 3,257,4097,32768 \ + --backward --out reports/experiments/h3-adaln-row-gather-b200/chain_replay.json +``` + +## Known Limitations + +- Like diffusers, the op materialises six `(S, H)` tensors, which is 8 GB at + S = 131072. Fusing the gather into the norm/modulate and gate-residual consumers + (`h3_rmsnorm`, `adaln_gate_residual`) avoids this. Those are separate rows. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-final-adaln-out.md b/docs/operators/h3-final-adaln-out.md new file mode 100644 index 000000000..a46954f06 --- /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 `ddca8f95` on an otherwise idle B200, with FP64 leaves and upstream gradients, after + the golden's declared cast was changed to round the FP32 SiLU as the model does. The + backward reference keeps only the SiLU and table roundings, so the errors include the BF16 + rounding of `norm(x)` and `1 + scale`. +- [`chain_replay.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-final-adaln-out-b200/chain_replay.json): the + whole conditioning chain (timestep → … → norm_out) replayed stage by stage over + T in {1, 2, 3, 4} × S in {3, 257, 4097, 32768}, with the backward replay. Written from a clean + tree at `a4bc99f9` (#489). Every CUDA pass/fail flag matches the previous replay; the + `h3_rmsnorm` stage's chained error against the golden is lower in all 16 cases (worst 0.0796 + before, 0.0671 now). + +## Existing implementations (RFC #420 reuse rule) + +![final_adaln_out vs existing implementations](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.png) + +| Implementation | Batch-invariant | size 4097: fwd err / worst grad err / fwd+bwd | size 32768: fwd err / worst grad err / fwd+bwd | +|---|---|---|---| +| diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay) | **no** (param/table grads not repeatable) | 8.3e-03 / 7.3e-02 / 1021 µs | 8.9e-03 / 2.4e-01 / 4525 µs | +| rl-kernel H3FinalAdaLNOutCudaOp | yes | 8.3e-03 / 5.2e-03 / 1642 µs | 8.9e-03 / 5.7e-03 / 2959 µs | + +Errors are max|err| / max|ref| against the same computation in FP64; latency is the median +forward + backward time on an otherwise idle B200. Batch invariance is bitwise and covers +three checks: every row computed alone vs inside full batches of 64, 257 and 2048 rows; the full +131072-token batch vs sub-batches that together cover every row; and a dense batch-size sweep. A +"no" means that at least one row, sub-batch or gradient differed. [`final_adaln_out.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-prior-art-b200/final_adaln_out.json) +was written from a clean tree at `03da729` by + +```bash +python tools/validation/models/h3_prior_art.py --op final_adaln_out --out reports/experiments/h3-prior-art-b200/final_adaln_out.json +python tools/validation/models/plot_h3_prior_art.py reports/experiments/h3-prior-art-b200/final_adaln_out.json +``` + +Libraries that do not import are skipped and recorded as unavailable in the report. + +### Direct bitwise comparison + +Equal FP64 errors do not show equal bits, so the candidates are also compared byte for byte. +`tools/validation/models/h3_bitwise_reuse.py` runs every candidate above on identical inputs +(each global batch-invariant mode in its own process) and compares every output and leaf +gradient with diffusers' path and with this kernel. + +![final_adaln_out bitwise comparison](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr489/reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png) + +There is no reuse candidate beyond diffusers' MiniMaxH3AdaLayerNormOut. This op does not promise bitwise equality: its shift/scale projection uses #485's fixed-order linear, not cuBLAS. The output differs from diffusers in 0.05-0.07% of BF16 elements at every S (1 to 32768), and 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 `8cb4abed` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op final_adaln_out \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json +python tools/validation/models/plot_h3_bitwise_reuse.py reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.json \ + --out reports/experiments/h3-bitwise-reuse-b200/final_adaln_out.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_final_adaln_out.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end: timestep -> ... -> norm_out +python tools/validation/operators/check_operator.py --op final_adaln_out --candidate cuda --device cuda \ + --dtype bf16 --batch 3 --seq 4097 --normalized-dim 5376 --check-grad # also 257, 1024 +python tools/validation/models/h3_evidence.py --op final_adaln_out \ + --out reports/experiments/h3-final-adaln-out-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-final-adaln-out-b200/report.json +``` + +## Known Limitations + +- **FP32 gtest at larger S.** H3 runs this op with BF16 activations and weights, and the BF16 + gtest passes at S = 257, 1024 and 4097. With `--dtype fp32` (every input FP32), the FP32 + reduction tolerance (atol = rtol = 1e-4) misses `d_temb` from S = 257 and `dW` from + S = 1024. Both gradients are FP32 sums over 3·S rows. Where large terms cancel, the + summation error exceeds the 1e-4 absolute floor. The PyTorch reference misses the same + gradients at the same S (`--batch 3`, seed 123, max abs error): + + | S | `d_temb` CUDA | `d_temb` PyTorch | `dW` CUDA | `dW` PyTorch | + | --- | --- | --- | --- | --- | + | 64 | 9.2e-5 | 1.5e-4 | 1.5e-4 | 9.2e-5 | + | 257 | 3.1e-4 ✗ | 5.9e-4 ✗ | 8.5e-4 | 2.4e-4 | + | 1024 | 1.5e-3 ✗ | 2.4e-3 ✗ | 2.0e-3 ✗ | 1.5e-3 ✗ | + | 4097 | 5.9e-3 ✗ | 9.3e-3 ✗ | 5.9e-3 ✗ | 3.9e-3 ✗ | + + ✗ fails the gtest. Unmarked values above 1e-4 pass through the rtol term. In a PyTorch + emulation, accumulating only the projection backward in FP64 clears S = 257 but not + S ≥ 1024, because the segment sums contribute as well. Making every reduction FP64 would + change the `h3_det_linear` kernels shared with `timestep_mlp_fp32` and + `adaln_projection_3mod`. The alternative is an FP32 gradient + tolerance that scales with reduction length. Which of the two to take is an open question + for the maintainers. +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. diff --git a/docs/operators/h3-one-block.md b/docs/operators/h3-one-block.md new file mode 100644 index 000000000..b2f1d4ce0 --- /dev/null +++ b/docs/operators/h3-one-block.md @@ -0,0 +1,183 @@ +# MiniMax-H3 One Block, Forward and Backward + +## Summary + +`ws1_one_h3_block` (RFC #420, WS1 closeout) runs **block 0 of the pinned MiniMax-H3 checkpoint** +node by node, forward and backward, and reports where it first departs from diffusers. The block +is a fixed graph of 17 nodes, from the AdaLN projection to the second gated residual. It runs on +packed FL2VA layouts that are **bitwise equal to the pinned diffusers pipeline's** `position_ids` +and modality tags. + +On one B200, at S = 264, 925 and 3160: + +- **Determinism.** Every node, every gradient reaching a node, and all 14 leaf gradients are + bitwise repeatable. +- **Batch invariance.** Batch row 0 of a two-row batch has the single-row bytes at every node, + forward and backward. +- **Agreement with diffusers.** Every node that promises it is bitwise equal to diffusers on + diffusers' inputs. The first drift from diffusers is `adaln_projection`, the first reduction + whose tree differs from cuBLAS. +- **Accuracy.** The block output is exactly as close to the FP32 golden as diffusers' (relative + error 5.2e-3 / 7.9e-3 / 9.0e-3). + +Only three of the 17 nodes are bound to this stack's row operators. The others use **interim** +deterministic operators until their own rows land, or the provider replay where no kernel exists +yet (see Bindings). The block-level promises above hold with whatever a node is bound to. Speed +and the exact rounding of interim nodes belong to those operators, not to this row. + +## Entry Point + +```python +from rl_engine.validation.models.h3_block import CandidateOps, run_block_backward, run_block_case +from rl_engine.validation.models.h3_weights import load_h3_block_weights + +weights = load_h3_block_weights("cuda") # conditioning + block-0 tensors, sha256-checked +ops = CandidateOps() +forward = run_block_case(weights, layout="small", ops=ops) # per node, first_*_drift +backward = run_block_backward(weights, layout="small", ops=ops) +``` + +The CLI writes the JSON evidence and the plot script renders it: + +```bash +python tools/weights/prepare_h3_weights.py --out --block # shard 1 of 14, ~4.8 GB +export RL_KERNEL_H3_WEIGHTS= +python tools/validation/models/h3_block_replay.py --layouts tiny,small,medium --out report.json +python tools/validation/models/plot_h3_block.py report.json +``` + +## Graph and Bindings + +`rl_engine/validation/models/h3_block.py` `NODES` (graph order): + +| Node | Owning RFC row | Binding today | +|---|---|---| +| `adaln_projection` | `adaln_projection_3mod` | **row**: `H3AdaLNProjectionCudaOp` (registry) | +| `norm1`, `norm2` | `h3_rmsnorm` | **row**: `H3RMSNormCudaOp.forward_modulated` (registry) | +| `q_proj`, `k_proj`, `v_proj` | `h3_qkv_gemm` | interim: Triton FP32-accumulate GEMM, one BF16 rounding | +| `q_norm`, `k_norm` | `h3_qk_rmsnorm_d128` | interim: `H3RMSNormCudaOp.forward` on the 128-wide heads | +| `rope_q`, `rope_k` | `h3_mm_rope_3axis_partial` | reference: the provider replay (no kernel exists yet) | +| `attention` | `h3_full_attention` | interim: `DeterministicAttentionOp(causal=False)` | +| `o_proj` | `h3_attention_o_gemm` | interim: Triton FP32-accumulate GEMM | +| `residual_attn`, `residual_mlp` | `adaln_gate_residual` | **row**: `H3GateResidualCudaOp` (registry) | +| `ffn_gate_up` | `h3_ffn_gate_up_gemm` | interim: Triton FP32-accumulate GEMM | +| `swiglu` | `h3_swiglu` | interim: `SwiGLUCudaOp` (FP32 math, one rounding) | +| `ffn_down` | `h3_ffn_down_gemm` | interim: Triton FP32-accumulate GEMM | + +When a row lands, replacing its node's `candidate` callable is the whole integration. The node's +report entries then become that operator's block-level evidence. + +The interim GEMMs use `_triton_gemm_fp32`, one fixed-order FP32 K loop with no split-K, and round +to BF16 once. They do not use `TritonDetGemmOp.__call__`, whose K tree has BF16 nodes. Measured on +B200: + +| Shape | Tree path err | FP32 K loop → BF16 | cuBLAS | Tree | FP32 K loop | +|---|---|---|---|---|---| +| 264 × 5376 → 7168 | 7.0e-3 | 2.4e-3 | 2.4e-3 | 10.4 ms | 1.2 ms | +| 925 × 14336 → 5376 | 7.5e-3 | 2.6e-3 | 2.6e-3 | 47 ms | 2.7 ms | + +Both paths are bitwise repeatable and row-invariant. Error is max|err| / max|ref| against FP64. + +## What Each Node Reports + +Every node runs three ways: the provider replay (`h3_provider`, diffusers@f53d552 op for op), the +candidate (chained, and *isolated* on the provider's inputs), and the FP32 golden +(`rl_engine/reference/minimax_h3/block.py`: FP32 math, TF32 off, no BF16 storage between nodes). + +- **Forward:** `repeat_bitwise_equal` and `batch_row_bitwise_equal` (row 0 of a B = 2 batch whose + row 1 has a different input). Also chained and isolated agreement with the provider (bitwise, + max abs, first differing packed row), the candidate's and provider's `max|err| / max|golden|`, + and the sha256 of the node's output. +- **Backward:** for the block input, `temb` and the 12 block parameters: repeat bitwise, and the + candidate's and provider's error against the golden. For the gradient reaching each node's + output, in reverse graph order: repeat bitwise and batch-row bitwise. Parameter gradients sum + over the batch, so only activation gradients take part in the batch check. +- **First drift:** `first_drift`, `first_isolated_drift`, `first_repeat_drift`, + `first_batch_drift`, `first_grad_repeat_drift` and `first_grad_batch_drift` each name the first + node where that comparison stops being bitwise. +- **Permutation probe (control H-C3, reported, not promised):** permute the packed rows together + with positions and AdaLN indices, and compare the output to the permuted original. + +Inputs: `temb` comes from the provider's conditioning path (the conditioning rows are checked by +`h3_chain`). The block input is a seeded BF16 stand-in, because the packed input projections are +other rows. Text rows use AdaLN table row 0 and media rows use row 1. + +## Evidence + +[`report.json`](https://github.com/fusheng-ji/RL-Kernel/blob/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/report.json) was generated from a clean tree +at `41e4293` on an otherwise idle B200 (torch 2.13.0+cu130): + +![one-block evidence](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr539/reports/experiments/h3-one-block-b200/figure.png) + +| Layout | S | First drift / first isolated drift | Repeat / batch drift (fwd, grad) | Output err: RL-Kernel / diffusers | Fwd ms: RL-Kernel / diffusers | Fwd+bwd ms | +|---|---|---|---|---|---|---| +| tiny | 264 | `adaln_projection` / `adaln_projection` | none | 5.21e-3 / 5.21e-3 | 7.4 / 0.61 | 22.7 / 2.8 | +| small | 925 | `adaln_projection` / `adaln_projection` | none | 7.87e-3 / 7.87e-3 | 30.1 / 1.12 | 81.7 / 5.0 | +| medium | 3160 | `adaln_projection` / `adaln_projection` | none | 8.97e-3 / 8.97e-3 | 225 / 3.8 | 577 / 11.5 | + +- **Forward accuracy.** Per node, the candidate's error is within 0.93–1.20× diffusers'. The + largest is `attention` at S = 3160, which is the interim attention operator's. +- **Gradient accuracy.** The candidate is better than diffusers on `adaln_proj.weight`, + `adaln_proj.bias` and `temb`, by up to 7×: diffusers' `index_select` backward accumulates the + table gradient in BF16 atomics, while the row operators use FP32 segment sums. The candidate is + within 1.16× on the other leaves, except `attn.norm_q.weight` at S = 3160 (9.5e-3 vs 5.5e-3). + At S = 925 the same leaf is 2× *more* accurate than diffusers (1.3e-2 vs 2.4e-2), so this + 128-element gradient's error moves with BF16 rounding rather than favouring one side. +- **Permutation probe (H-C3).** The output changes in 79 / 450 / 2199 of the rows, by at most 8. + Attention sums over keys in packed order, so a consistent permutation is not byte-preserving with + the interim attention operator. This is for `h3_full_attention` to settle. +- **Speed.** The candidate is 12–60× slower than diffusers forward. Almost all of it is the + interim GEMMs and attention (FP32 scores materialised). The three row operators are each + faster than or equal to diffusers in their own evidence. + +## Existing implementations (RFC #420 reuse rule) + +- **Is there an existing block-level harness?** diffusers' `MiniMaxH3TransformerBlock` is the + executable reference. It is replayed op for op as the provider, so diffusers is not imported and + the replay pins the reference's dtypes and ordering. The Qwen3-8B chain gate from #315 + (`rl_engine/validation/models/chain_gate.py`) has per-node token digests and first-drift + localisation, but its graph, shapes and acceptance are Qwen3-specific (causal attention, no + AdaLN). This row reuses its method (per-node digests, graph-ordered first drift) on the H3 + `h3_chain` stage pattern. +- **Interim operators.** These are reused from RL-Kernel and are not new code. The only choice + made here is the FP32-accumulate entry of the Triton GEMM instead of its BF16 tree, for the + accuracy and speed measured above. + +## Tests + +`tests/models/minimax_h3/test_h3_one_block.py`: + +- **CPU:** + - the three layouts hash to the pinned diffusers `build_packed_sequence` output; + - the graph is well formed; + - the partial MM-RoPE leaves channels 96–127 unchanged. +- **GPU, `tiny` and `small` forward:** + - every node repeats and is batch-row bitwise; + - promised nodes are isolated-bitwise to diffusers; + - node and output errors are at most 1.25 × diffusers' + 1e-3. +- **GPU, `tiny` backward:** + - every leaf repeats; + - dX is batch-row bitwise; + - no gradient drift; + - leaf errors are at most 1.5 × diffusers' + 5e-3, or, for a BF16 leaf, at most 6 ulps + (of its largest golden element) above diffusers'. A stored gradient's worst element moves + in whole ulps, so the relative rule alone is seed- and torch-version-sensitive; + `tools/validation/models/h3_block_ulp_sweep.py` measured at most +4.15 ulps over + tiny/small × 10 seeds on torch 2.13 and 2.9.1, and +5.15 ulps over 34 single-node + candidate/provider swaps (reports: `reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json`). + +```bash +RL_KERNEL_H3_WEIGHTS= CUDA_VISIBLE_DEVICES=0,1 python -m pytest tests/models/minimax_h3 -q +``` + +## Known Limitations + +- **Single platform.** CUDA only. `DeterministicAttentionOp` and the Triton GEMM also build for + ROCm, but the block has not been run there. +- **Interim nodes.** Eleven of the 17 nodes are interim or reference bindings. Their speed and + rounding are not the owning rows' final ones, and the MM-RoPE nodes are the provider replay + itself, so they carry no RL-Kernel evidence until `h3_mm_rope_3axis_partial` lands. +- **Stand-in inputs.** The block input is a seeded stand-in and `temb` is the provider's. + Real packed projections arrive with the input-projection rows. +- **Not covered here.** Block 0 only (`ws1_full_h3_transformer` covers 50 blocks and the refiner), + and one sequence per batch row (H3's batch is a replication axis). diff --git a/docs/operators/h3-rmsnorm.md b/docs/operators/h3-rmsnorm.md new file mode 100644 index 000000000..ac625ae35 --- /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 `d2a6546`. 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 `83797833` 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..eeb90f72f --- /dev/null +++ b/docs/operators/h3-sp-norm-adaln.md @@ -0,0 +1,117 @@ +# MiniMax-H3 Sequence-Parallel Norm, AdaLN Modulation and Gated Residual + +## Summary + +`sp_norm_adaln` runs the row-wise part of an H3 block with the packed sequence split across +sequence-parallel ranks (RFC #420, WS2). It covers the [RMSNorm and AdaLN modulation](h3-rmsnorm.md) +and the [gated residual](h3-adaln-gate-residual.md). Every rank's output and row gradients are the +**WS1 bytes of its rows**. The gradients of the replicated tensors (`d_norm_weight`, `d_shift`, +`d_scale`, `d_gate`) are **the WS1 bytes, identical on every rank**. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + +collective = DeterministicCollective(group=sp_group, device=local_rank) +op = H3SPNormAdaLNCudaOp(collective, seq_len=S, batch=B) +rows = slice(op.layout.lo, op.layout.hi) # this rank's packed positions +hidden = op.gate_residual(residual[:, rows], attn_out_local, gate_msa, adaln_indices) +normed = op.norm_modulated(hidden, norm2_weight, shift_mlp, scale_mlp, adaln_indices) +final = op.norm_modulated(h, norm_out_weight, shift, scale, timestep_indices) # norm_out +op.readback() # sp, rank, positions, collective backend, reduction order, fallback +``` + +The table views, the norm weight and the full `(S,)` row index are replicated. Activations are +the rank's `(B, S_local, H)` rows. + +## Ownership + +Rank `r` of `sp` holds packed positions `[r * S // sp, (r + 1) * S // sp)` of every batch item, +together with the matching rows of the metadata. Any `S >= sp` works, including `S` not divisible +by `sp`. Shapes that do not match the rank's slice, a sharded index, an out-of-range index, CPU +tensors and malformed tables all fail closed. + +## Numerics + +- **Forward and row gradients.** `out`, `dx`, `d_sublayer` and `d_residual` are the WS1 kernels on + the local rows, with the local slice of the index. Rows are independent, so they are byte-equal + to WS1. +- **Cross-row gradients.** In WS1, `d_norm_weight` sums `d_n * x * rstd` over fixed 256-row tiles + of the `B x S` rows. `d_shift`/`d_scale`/`d_gate` sum over 256-element tiles of each table row's + positions, sorted stably. In both cases the tile partials are then folded in ascending order. SP + keeps exactly that two-level order: + 1. Every rank builds the global WS1 tile list from the replicated index (`sp_plan`). + 2. Each tile is computed by the rank that holds its first row. Rows of that tile that live on + other ranks are all-gathered first. Only tiles that cross a shard boundary move rows. + 3. The tile partials (`h3_rmsnorm_backward_partials`, `h3_gate_grad_partials`, the WS1 partial + kernels over explicit row lists) are all-gathered and put back in WS1 tile order. + 4. Every rank runs the WS1 folds (`h3_rmsnorm_fold_partials`, `h3_gate_grad_fold`). + + The only collectives are rank-ordered all-gathers, which are copies. No collective reduction + arithmetic is used. + +A naive SP backward runs the WS1 backward on each rank's rows and sums the per-rank results. It +changes the reduction tree, so it is not WS1. The evidence measures how often it differs. + +The plan depends only on the layout and the index. `H3SPNormAdaLNCudaOp` builds it once per index +object and reuses it for every norm and gated residual that uses the same index, which in a forward +pass is every block. + +## Rows Exchanged + +How many rows move depends on how the table rows interleave along the sequence: + +| Packing | Rows sent per rank (S = 32768, SP8, 4096 rows per rank) | +| --- | --- | +| block (H3's layout: each timestep's text, video and audio tokens contiguous) | 0 – 224 | +| interleaved (modality and timestep random per position; stress case) | 0 – 1536 | + +## Evidence + +![sp_norm_adaln on 8 x B200: byte equality, time, naive SP](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr494/reports/experiments/h3-sp-norm-adaln-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `ecde5bb`, with +real NCCL processes on one 8 x B200 node, one GPU per rank. The region is +`norm2(residual + gate_msa[row] * y)` with `shift_mlp`/`scale_mlp` modulation, H = 5376, BF16. It +covers six cases: S = 4097 (block, interleaved, B = 2), 32768 (block, interleaved) and 131072 +(block). For SP 1, 2, 4 and 8, every rank's output rows, `d_residual` and `d_sublayer`, and its +`d_norm_w` and `d_table`, are byte-equal to WS1 computed on the same GPU. + +Forward + backward of the region (slowest rank): + +| S, packing | WS1 (1 GPU) | SP2 | SP4 | SP8 | +| --- | --- | --- | --- | --- | +| 4097, block | 2.04 ms | 1.73 ms | 1.88 ms | 2.43 ms | +| 32768, block | 3.72 ms | 2.85 ms | 2.40 ms | 2.72 ms | +| 32768, interleaved | 3.72 ms | 3.45 ms | 3.88 ms | 6.82 ms | +| 131072, block | 10.73 ms | 7.27 ms | 4.66 ms | 3.92 ms | + +At long sequences the rows split and the gain grows: 2.74x at SP8 for S = 131072. At short +sequences, each backward's handful of host-synchronising all-gathers (about 0.1 ms each) +dominates. Interleaved packing also moves more rows. A naive SP8 backward (each rank's WS1 +backward, then a rank-order sum) differs from WS1 on 40% of `d_norm_w` elements and on 6% (block) +or 21% (interleaved) of `d_table` elements. + +## Tests + +```bash +python -m pytest tests/models/minimax_h3/test_h3_sp_norm_adaln.py -v # layout, rank backward, NCCL SP2/4/8 +python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-sp-norm-adaln-b200/report.json +``` + +The rank-backward tests run SP 2–8 on one GPU, with ranks as threads around the plain backward +functions. They cover odd `S`, `B = 2`, shards smaller than a tile, block and interleaved packing, +and bf16 and fp32, and compare every rank against WS1. The NCCL tests run the autograd region end +to end and need as many GPUs as ranks. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- With interleaved packing, up to a full shard of rows can move in the backward. The result is + still exact. +- Attention and FFN are other rows. Until they land, the region uses a seeded stand-in for the + sublayer output. diff --git a/docs/operators/h3-timestep-mlp.md b/docs/operators/h3-timestep-mlp.md new file mode 100644 index 000000000..b65b19fca --- /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 `418998c`. It also records that a +timestep's row is bitwise identical whether it runs alone or in a batch of 9. + +## Existing implementations (RFC #420 reuse rule) + +![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 `73939f2d` 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..ff6579cd8 --- /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 `be6a94ae` by + +```bash +python tools/validation/models/h3_bitwise_reuse.py --op timestep_sinusoid --fast-math-variant \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json +python tools/validation/models/plot_h3_bitwise_reuse.py \ + reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.json \ + --out reports/experiments/h3-bitwise-reuse-b200/timestep_sinusoid.png +``` + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_timestep_sinusoid.py -v # operator +python -m pytest tests/models/minimax_h3/test_h3_conditioning_e2e.py -v # end to end, pinned weights +python tools/validation/operators/check_operator.py --op timestep_sinusoid_h3 --candidate cuda --device cuda \ + --dtype fp32 --batch 7 --check-grad +python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-timestep-sinusoid-b200/report.json +``` + +## Known Limitations + +- There is no ROCm kernel. ROCm dispatches the PyTorch reference. +- `max_period` is fixed at 10000, the only value H3 uses. +- `KERNEL_ALIGN_USE_FAST_MATH=1` builds leave out every H3 kernel (`H3_BITWISE_SOURCES` in + `build_tools/extensions.py`): fast math replaces `expf`/`sinf`/`cosf` with approximations and + breaks bitwise parity. Such builds dispatch the PyTorch reference. diff --git a/docs/operators/h3-tp-adaln-3mod.md b/docs/operators/h3-tp-adaln-3mod.md new file mode 100644 index 000000000..868df445a --- /dev/null +++ b/docs/operators/h3-tp-adaln-3mod.md @@ -0,0 +1,113 @@ +# MiniMax-H3 Tensor-Parallel AdaLN Projection + +## Summary + +`tp_adaln_3mod` shards the [AdaLN projection](h3-adaln-projection.md) (`2688 -> 96768`, 520 MB of +BF16 weight per block) across tensor-parallel ranks (RFC #420, WS2). Every rank ends with the +**same bytes as the single-GPU (WS1) op**: the six modulation tensors, the `3T` modality rows, +`d_temb`, and its shard of `dW`/`db`. + +## Entry Point + +```python +from rl_engine.distributed.algorithms.collectives import DeterministicCollective +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, +) + +collective = DeterministicCollective(group=tp_group, device=local_rank) +op = H3TPAdaLNProjectionCudaOp(collective, n_total=96768) +w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) +shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = op(temb, w_shard, b_shard) +op.readback() # tp, rank, owned columns/slots, collective backend, kernel ids, fallback +``` + +## Ownership + +Rank `r` of `tp` owns the contiguous output columns `[r * N / tp, (r + 1) * N / tp)` of the +`N = 3 x 6 x 5376` table, which are the same rows of `adaln_proj.linear.weight` and `bias`. +`N / tp` must be a multiple of 64, the `d_input` chunk of contract `h3-det-linear-v1`, so a shard +never splits a WS1 partial sum. The shard constructor accepts positive TP sizes satisfying +this alignment, including 1, 2, 3, 4, 6, 7, 8 and 9 (TP=9 gives 10752 columns per rank). +The `DeterministicCollective` backend supports only 1, 2, 4 and 8 ranks; shard alignment +alone does not imply backend support. + +| TP | Columns per rank | Rank 0 owns | +| --- | --- | --- | +| 2 | 48384 | video: all six chunks; text: shift/scale/gate_msa | +| 4 | 24192 | video: shift_msa .. shift_mlp, half of scale_mlp | +| 8 | 12096 | video: shift_msa, scale_msa, 1344 channels of gate_msa | + +`AdaLNColumnShard.slots()` lists each rank's `(modality, chunk, h_begin, h_end)` pieces. The +all-gathered table is the WS1 table, so the six tensors and the `t * 3 + modality` rows are views +of it exactly as in WS1. Nothing downstream needs to know the TP size. + +## Numerics + +Every floating-point operation is the WS1 kernel itself. The only collective is a rank-ordered +all-gather, which is a copy. No collective reduction arithmetic is used. + +- **Forward.** An output column depends only on `silu(temb)` and its own weight row, so the local + tensor-core GEMV yields exactly the WS1 columns. The all-gather rebuilds the `(T, N)` table. +- **`dW`, `db`.** These are rows of the shard, computed locally with the WS1 ascending-`t` fold. +- **`d_temb`.** WS1 sums `N` in 64-row chunks and then left-folds the chunk partials in ascending + order. Each rank computes the partials of its own chunks + (`h3_det_linear_backward_input_partials`), the all-gather puts them in global chunk order, and + every rank runs the WS1 fold (`h3_det_linear_fold_chunks`) and the FP32 SiLU VJP. + +The backward reads its own columns of the table gradient. It therefore requires that gradient to +be identical on every TP rank, as it is when every rank applies the modulation to the same rows. + +## Communication + +| Direction | Payload per rank | At T = 3, TP = 8 | +| --- | --- | --- | +| forward | `T x N / tp` BF16 | 72.6 KB | +| backward | `(N / 64 / tp) x T x 2688` FP32 | 6.1 MB | + +`DeterministicCollective` copies all-gathers whose output is at most 256 KiB with a single-block +byte loop, at about 4 µs per KiB on B200. `ws2_comm.gather_rows` pads such shards just past that +size, where the multi-block path takes about 45 µs. This matters for the forward at small `T`. + +## Evidence + +![tp_adaln_3mod on 8 x B200: byte equality and time](https://raw.githubusercontent.com/fusheng-ji/RL-Kernel/6d6193d94a83266d5f5f62bb9d90fb3215716ee3/pr493/reports/experiments/h3-tp-adaln-3mod-b200/figure.png) + +The report is written by `tools/validation/models/h3_ws2_evidence.py` from a clean tree at commit `507c9ad`, with +real NCCL processes on one 8 x B200 node, one GPU per rank, on the pinned block-0 weights. For +TP 1, 2, 4 and 8 and T = 1..4, every rank's table and `d_temb`, and its `dW`/`db` shard, are +byte-equal to WS1 computed on the same GPU. + +| T | | WS1 (1 GPU) | TP2 | TP4 | TP8 | +| --- | --- | --- | --- | --- | --- | +| 1 | forward | 0.10 ms | 0.12 ms | 0.11 ms | 0.12 ms | +| 1 | forward + backward | 1.90 ms | 1.49 ms | 1.19 ms | 1.04 ms | +| 3 | forward | 0.10 ms | 0.11 ms | 0.11 ms | 0.12 ms | +| 3 | forward + backward | 2.25 ms | 1.77 ms | 1.48 ms | 1.32 ms | +| 4 | forward + backward | 2.09 ms | 1.83 ms | 1.57 ms | 1.45 ms | + +Times are for the slowest rank. The WS1 forward already streams the 520 MB weight at about +5 TB/s, so the TP forward is bounded by the all-gather. `DeterministicCollective` synchronises the +host on every call, which costs about 0.1 ms. The backward is dominated by the shard-local `dW` +and gains 1.7-1.8x at TP8. Each rank also holds only `1 / tp` of the weight and its gradient. + +## Tests + +```bash +export RL_KERNEL_H3_WEIGHTS= +python -m pytest tests/models/minimax_h3/test_h3_tp_adaln.py -v # ownership, shard arithmetic, NCCL TP2/4/8 +python tools/validation/models/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \ + --out reports/experiments/h3-tp-adaln-3mod-b200/report.json +python tools/validation/models/plot_h3_evidence.py reports/experiments/h3-tp-adaln-3mod-b200/report.json +``` + +The NCCL tests need as many GPUs as ranks and skip otherwise. The shard-arithmetic tests check +the same property on one GPU: each rank's columns and chunk partials equal the matching slice of +the WS1 call, and the fold of the gathered partials equals WS1's `d_input`. + +## Known Limitations + +- CUDA only. `DeterministicCollective` supports world sizes 1, 2, 4 and 8. +- `norm_out.linear` (`2688 -> 10752`) is not sharded here. It is small and stays replicated. +- There is no ROCm path. diff --git a/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/reference-manifest.json b/examples/models/minimax_h3/reference-manifest.json new file mode 100644 index 000000000..acfe57486 --- /dev/null +++ b/examples/models/minimax_h3/reference-manifest.json @@ -0,0 +1,1641 @@ +{ + "source_commit": "2681cadb4dd14c8980f6d760ba11751be657f261", + "note": "Regenerated on 2026-10-16 with the original commands at source_commit after the first run's artifacts were deleted. 1614 of the first run's 1635 files are byte-identical; the others are metadata (report.json, the thresholds file that embeds their hashes) and the video.mp4 / audio.wav containers, which differ even between two runs of one job.", + "files": { + "fp32-reference/report.json": "c82f94617e01fd97920cd890bca0225d476a92543afdd5634e08f40b337b1b8d", + "fp32-reference/gradients.json": "1e933436ca157876aa7dd5b6f1f4b17d1a51329ad52ac9c20107033be16292e2", + "fp32-reference/tensors.pt": "3944d33e8331efcd6b6ef3a0aa7a7d5528bde3cb7a51cf3bbe27d47bbf5264c6", + "fp32-reference/gradients/0637.pt": "44f5bbaed42fff5acf99e8c204882a0c574de0b29d65fef631383e2dd5098d74", + "fp32-reference/gradients/0636.pt": "254731e038755b142a49f7751856177c062266a3a6c3b7bd59efd4187f0f1158", + "fp32-reference/gradients/0635.pt": "a48608cbefc9a5aa90a1981238b8aacef860f7225f0ed6b6e0b217e2fa89effa", + "fp32-reference/gradients/0634.pt": "a7909e90edad95fc8c1721f7e8dd5dff6b326bfdd4c525eb93e373b4dc1efbc7", + "fp32-reference/gradients/0631.pt": "0d75a0b1311b45f7d2d74622a61bfcbcdbd59b32f455fb16a45cb5d9a444f2aa", + "fp32-reference/gradients/0633.pt": "665c07fb50f760eaf86cb0de99a9b0550fff06a4f0133c34991ad08d99d7fb4d", + "fp32-reference/gradients/0632.pt": "54ca30b8ff8c43b8b41e49e145d526c8b9a678ac2d80b4cb645293fd4f4e8c2b", + "fp32-reference/gradients/0628.pt": "e853596f30bf518fccd1508de3682c15b89fe823f7a2916ede958b54efd6c2e8", + "fp32-reference/gradients/0627.pt": "90f3f9cb0982f472ccc91fd6e696c42bb204c1faef5de058bdf01ebdfdff2976", + "fp32-reference/gradients/0626.pt": "bc42369ec3de7ea1166281989736b960ee1bd92dfc34d28b3e5b8bc1e9c22570", + "fp32-reference/gradients/0625.pt": "10beec67d5f8a74e03451f12e1ce7867c3b9f8740287d2ee6a64962bb290570d", + "fp32-reference/gradients/0624.pt": "decd792208d9969a670b04bf357c0bf7a08061a8d44fbc401a378d0f2665be49", + "fp32-reference/gradients/0623.pt": "b70e894832b85566ff63c954bc817f2d22ec8edb0a1591f89fbd9986c194ee7b", + "fp32-reference/gradients/0622.pt": "b74e7fe0e27483e4bc02030c72f9bd8d9cf13fe3e9b36981d5e7c766520a2f72", + "fp32-reference/gradients/0621.pt": "3a21a0af227127a43604e0dc12ec4019d34ce597481ea019e85f5d2382c1bb2c", + "fp32-reference/gradients/0620.pt": "31679f9ddaebeecc04ed14f0566dabb117ee185e8665500232a691e95decb1e9", + "fp32-reference/gradients/0619.pt": "e90a3f84a10a51357abf17e880cf5f26ec4875423cfab369acc9837d5950a24c", + "fp32-reference/gradients/0630.pt": "00c5f09323896833fdb4e8f7e472ccc92a87e2c8d77ba635f5e3e9655fe1f67f", + "fp32-reference/gradients/0629.pt": "05658b8175118bdd51cdf1703f585e62381593314a06a9bcf296ce698dbf593a", + "fp32-reference/gradients/0616.pt": "dc2d855b1c2a8c0689a31f26dbcc49b52f59342d8abbfd02c53555653c5a7a99", + "fp32-reference/gradients/0615.pt": "273f5f9e8f5f285402b951ceaf434d44ba8fa594b32e57b6af586880c33b929d", + "fp32-reference/gradients/0614.pt": "8efca892a52bdfffa7abc41f2bdee17397e73e0c4bcdc28be888dd0959579c28", + "fp32-reference/gradients/0613.pt": "37e41946f8b7dc1de33eed0275f37562249fe0d55ce037e13f1c35a5683d211c", + "fp32-reference/gradients/0612.pt": "3dc925c6fdba3cd5033a637ff8ca873f9ae655eb0e58ab1360a864c22c122995", + "fp32-reference/gradients/0611.pt": "fb09ec9e84b92033dc0d4644f6ec72e0868830ecc362670782f198f3795dfc55", + "fp32-reference/gradients/0610.pt": "7529a771ef1c37b3415ca8de67e454c88c83952047457f6334b0d8ea0124d3ed", + "fp32-reference/gradients/0609.pt": "22e05312b01f606c9481481686665ec98bd68dd23bb522098769726c0081c0a3", + "fp32-reference/gradients/0608.pt": "180abc034fbed4b13c7558c20e5ba3d2a5ccca4cf1647e3e0d2addc25a01accb", + "fp32-reference/gradients/0607.pt": "1a34a0d1bb63396ee46373f80034e76cd06dcb7c56582f3eea15715eeaa681ca", + "fp32-reference/gradients/0618.pt": "bf8e30cfd617e225b5aa4beaa44a44be00e89e3510a0f46d3ed946b0817fb9ae", + "fp32-reference/gradients/0617.pt": "b2814b4164c79a314e50b652d97c7dde7bd18033a9e7af86ad3a6ce1be6bc839", + "fp32-reference/gradients/0604.pt": "33b5102eceb55ce3ef632e08e5f0d258d060833e9d34ab7aced4e1a4513f804f", + "fp32-reference/gradients/0603.pt": "bf302df2dc6a5fcdd4c94587eafdf84bec36dd716d902a0c35d01449a474ca45", + "fp32-reference/gradients/0602.pt": "3fe3279733a24f54602fc0b76e0413dbc967cbcf4a6890f67b3df16bc27ef631", + "fp32-reference/gradients/0601.pt": "a720456c6e189915151fe83a96c991c4dfec8327c83ff97c9f52627e54077273", + "fp32-reference/gradients/0600.pt": "baf9bd4336b597dcf9775959c63be60eb094687945e88a5702ede50f9c226a42", + "fp32-reference/gradients/0599.pt": "f9536b7ade5e21e340c122636b7761964d863e5abb25c97bf26ca26b37221087", + "fp32-reference/gradients/0598.pt": "d14af50b4470ca9245f1a2998949e09531af4d06f0d56b71b65fcff925ea3a93", + "fp32-reference/gradients/0597.pt": "bb9016a2949975d81fe4cbe92f4b3600078c360ad170715054596fc9876fbf06", + "fp32-reference/gradients/0596.pt": "33e3f6636ce02508301b8891be3551908e82da23db45b48b2ac461deb51b0abc", + "fp32-reference/gradients/0595.pt": "197e55d14ca2f0b0a59c16a8532add4336471be3efdeb16d67608cf512fed9c3", + "fp32-reference/gradients/0606.pt": "ca5e63333d5e4c48af48fb4c2aefa3420701d67b9d7c5cfde98aa306fc90faf7", + "fp32-reference/gradients/0605.pt": "646f515364bfc1ab4372208d51d405066c62e541a86f8e6021d8d07f30a2fe56", + "fp32-reference/gradients/0592.pt": "41f2515a72db0c0cbfc31691bdcfc2f683d67ab5877618c83480afa36a40fcfe", + "fp32-reference/gradients/0591.pt": "f453e824a82d73817f087d9a18dc6702d18075212b2280ede5268c49fe91c0df", + "fp32-reference/gradients/0590.pt": "2535ff23d17149d142ad967251ec4a659a281f0624488cf514f9ed0e658c98ce", + "fp32-reference/gradients/0589.pt": "8a3d3f193b56a3ede7ea543ea59a890f923a4fd91927667a02672744c249ebdf", + "fp32-reference/gradients/0588.pt": "c780e2182cc311a799ce3dd6399fff3a03fc1a4e3efa807ce22c25a5ba1508d6", + "fp32-reference/gradients/0587.pt": "562fc28ed2dce596e281aec9ed03f5a0757cbd8ddba265940bfb723866113bed", + "fp32-reference/gradients/0586.pt": "921043d6bf593d325aa59acf2d54e6e1a68c3effdf9b129293945f3477366794", + "fp32-reference/gradients/0585.pt": "c972070823d14a2f3a93b508204d9b69dd34b974a4bc408b5bc704de971b0553", + "fp32-reference/gradients/0584.pt": "d583d795b1478b779ade820ad6d5d578afd3c57d225237a7022926ae90ee0b53", + "fp32-reference/gradients/0583.pt": "6193692f06308374c4b957e22eac7903d1b3ed2ef8ea3bc9ede70db7440b23d5", + "fp32-reference/gradients/0594.pt": "74e34ec8619044dda13074a19b9142e24154a8172b8f0b942f005f29f2c7be0d", + "fp32-reference/gradients/0593.pt": "bae9313909674635cf9f33352f58f1b67bd568aca424ae9b25e6c5db2aecf2a9", + "fp32-reference/gradients/0580.pt": "c04fe8669c0ecac74a1b733fe1fd7debb0ee0d046528ad533913b5e3089963f0", + "fp32-reference/gradients/0579.pt": "fa351a8a3fe71204357f9e802709beba5fd9efffe0d5e2b849faa8da300e99b1", + "fp32-reference/gradients/0578.pt": "1f40e8332afa6a18a9ca57add6205fbf08e8d1bb7c6eeedcf62b3f89e77a4dcb", + "fp32-reference/gradients/0577.pt": "aea66f63e383a5d071ce9a20061b17fbeb9807604529431f6522f9dedbf07d65", + "fp32-reference/gradients/0576.pt": "3075eb7cd56feab500821f11a8097809949095c2a31e3941d189a0375507731e", + "fp32-reference/gradients/0575.pt": "e47d9e80455259ebe88c8991d38b07b0a7d8617625a62a536874e42dd9edb5e2", + "fp32-reference/gradients/0574.pt": "b8e196fff622d71b28ff84c19e70c78bea5b803c4eb43c7f5e2f1dce6a6278e7", + "fp32-reference/gradients/0573.pt": "dac0a85d7bc8ae9b0c08ad76374ed7881285b1d552552ddf92bc426127a02d3b", + "fp32-reference/gradients/0572.pt": "6aa08bc98cbcca883cb53fa12f80f5f0bdc5be204d33a46734697709c40d288c", + "fp32-reference/gradients/0571.pt": "88eef77acf516a8cc66a709b738746bb732f33ee9d20e5a98035ce1b21b8ee62", + "fp32-reference/gradients/0582.pt": "f0a6f8cf798c92ce501a2ecb43727bbc0cb656dcf67af11fc0bc17f96518e576", + "fp32-reference/gradients/0581.pt": "c9f1cc11bbd50411948b1785c09c7a1b0b2bc0f43dcc39ae0f6a968f80d1a0b2", + "fp32-reference/gradients/0568.pt": "25eba65a0b9ab0c565008aafebc65f3e0ad5ca8e7d73a39c504f88ef92d30bee", + "fp32-reference/gradients/0567.pt": "c0a0a485d1e25c2c275bc4614b5b1251ad04bc07a0e1fc682fcce8a0a33ec056", + "fp32-reference/gradients/0566.pt": "12c60adc2b6d12959d1d7d6bc6e959ccdf47a31d5fb657203736b2721474b382", + "fp32-reference/gradients/0565.pt": "abe6e3af6bccd98f5418ff0994a99f38f55b71b849a6a6bd48300162a092cc16", + "fp32-reference/gradients/0564.pt": "f6c11d204ae2f6a1f8cf53b3d59f0cf8618ec0f518650784ac12367e9aaf785b", + "fp32-reference/gradients/0563.pt": "13a275ed3a508c43e5e8116a75c6db44ea27c26fa3f68bc8b65bf1e43f389659", + "fp32-reference/gradients/0562.pt": "9b3839bbe302539a908f053a0b96662f657313a74bea71f56460411f32a1c1bb", + "fp32-reference/gradients/0561.pt": "c1614ca2877742c4bc056e33e1b7e47e7f525ece5ec2f9f045875ef2688ef2a2", + "fp32-reference/gradients/0560.pt": "bd1f4e9f86449803a63fde68633487f426d2503208fa6a2a74354440138254c8", + "fp32-reference/gradients/0559.pt": "069e800a99f054f11ffb3709c942e76049e7b05314cf2a1cf370ac62c17ae99f", + "fp32-reference/gradients/0570.pt": "ae867dd8f4e994bc4275508ae2c1caf8f1f8a35c6485a12b02121257471d3586", + "fp32-reference/gradients/0569.pt": "efd950800dd04a89d271131c81a0715c53b2fc98240e84948eb7c996ac14a952", + "fp32-reference/gradients/0556.pt": "f5096f499f47380174c2b551409907df82f929c2225cb37c1466344b9ecb7e3a", + "fp32-reference/gradients/0555.pt": "54d1f6946665e11ae7c8a9f6a8f24c34f53da005f58bbe1a49e83325b2b473de", + "fp32-reference/gradients/0554.pt": "e266c4857acb712b3f34394d24982838318358afcf3422e6638270677fa7b3c7", + "fp32-reference/gradients/0553.pt": "de2138d1665e1e5fb59a830d82602010e22f9ecc70958bc51028e330d3dfa92f", + "fp32-reference/gradients/0552.pt": "0d478c3f601c1d04769663777ec496a96d095e936006d5995e348b1fd74993fc", + "fp32-reference/gradients/0551.pt": "cf98d824f0f80bf32329cb91290685c5246a590a477477157942a3ba9f8444dc", + "fp32-reference/gradients/0550.pt": "ca7a07a9ad353f2b5f334333c661501c782f35c24a7bb2ee1bbf24b2849740f2", + "fp32-reference/gradients/0549.pt": "942b6b6d0fd0e9c81d84705bf7aa9c9da6a63f269f2a0fdeb466a77742b7b6df", + "fp32-reference/gradients/0548.pt": "96f2fcddc5d2ce4bf5199432f70f08dbeab0fc5f8910bd0667fa7a6e1cf1d841", + "fp32-reference/gradients/0547.pt": "1250c5a774a0837114d2927fe25f9abcbd29f2541a06cb6fc711862a50f86c06", + "fp32-reference/gradients/0558.pt": "98d6c6c7b9cf59fc506ee12a825da8d9818091c3e2ac3dbc2571372b77223ada", + "fp32-reference/gradients/0557.pt": "b10177d042eded9e6ab9905cdb666be6741a462bbc68983f98bc9482d2a580f7", + "fp32-reference/gradients/0544.pt": "50edc8355df398821f71866d565d5315410bb7e315c38793a2bba56a4777362d", + "fp32-reference/gradients/0543.pt": "8fa3752eff89183e99bb8e2cad6580752007ec066a50c0befdae0cf56edd80a5", + "fp32-reference/gradients/0542.pt": "2f3305e9de3502d8c4eb39f6ec7edd88d095bc89a3c7564e653e11aac36478bb", + "fp32-reference/gradients/0541.pt": "fbd3c1d13b2b3c77c669423deb4b6bab16f8533b83ab21157fa8e7faf4cbf23c", + "fp32-reference/gradients/0540.pt": "c3345fc3b40045674a6056adb5a8b950c3d68809169f4b9f869b107c2be20d11", + "fp32-reference/gradients/0539.pt": "6a2034e9f9fc7f97545623164c6294dc5d8ad144dd47496757e9fa2c9a4b1754", + "fp32-reference/gradients/0538.pt": "9b69cd5698548300b19f15a6f47bf2da4c67cfe7fdcb1c8cd98ffff7237e4311", + "fp32-reference/gradients/0537.pt": "c501a28230fff172649e2582d00bd0045245be2579cba72a10c1bafc8ee71a60", + "fp32-reference/gradients/0536.pt": "e608e623f2f61c554f61f4ea3dba5d3a19652190609623150e30a0555d6b6540", + "fp32-reference/gradients/0535.pt": "f46efdb2f5aefd044af14259bc8aabb2401f805a7ee0040880133c128d260770", + "fp32-reference/gradients/0546.pt": "85c068aacfed9e8192ff677e7737d061da7d856223270bcc1c03f289a1304585", + "fp32-reference/gradients/0545.pt": "85f48b3f6304874c95a9809f6985ba2945a5a6030347c6dae41372001a1c1bad", + "fp32-reference/gradients/0532.pt": "943695429503883526618ca2bf73208139d79471fa5bb741f91ee921723af1a5", + "fp32-reference/gradients/0531.pt": "77a97794ff01fa1888a3d06512e4103da118d821fb255d81876455a0ea058fb7", + "fp32-reference/gradients/0530.pt": "28a5f17c058e33274d209ca9ad4ea9b0d9156804a99dc3efe19641800217f54c", + "fp32-reference/gradients/0529.pt": "90a288c366c9cdf0aa20c4ac71820ecf9c97ddc5ea2a3a73cba8bf61786eb85b", + "fp32-reference/gradients/0528.pt": "e90a01fb080473bebb60405512d74810ea034960d9e48c14a1b2375ccbeb5011", + "fp32-reference/gradients/0527.pt": "70d05a36b918f5413b63b8a09034a64d7b7c6fb223f6aaa6abe171724ae5ee96", + "fp32-reference/gradients/0526.pt": "ce091cef2919bcc951a8ef1cfd8128fdd2c46f628dd4dd1aff13bb2141e5bf2f", + "fp32-reference/gradients/0525.pt": "1207b3d5a9a08f6fd525365a405529c67a9f86f959e0c361881310e38037be2a", + "fp32-reference/gradients/0524.pt": "1c5058f32db55c0615d936c8ff4ee8b2da5e745fac42e88dcd01e8f77a0358b2", + "fp32-reference/gradients/0523.pt": "1796a50450a530a31f47ed133de263736b0034c0c8ced0981ed572fd882f6f58", + "fp32-reference/gradients/0534.pt": "4d7bf480f2b928a86e909594ae851461d42bdef139e1cf9a62cbdad51dad8ca4", + "fp32-reference/gradients/0533.pt": "670ea78d38faca647db06d2c2c6d1e53a19280e4d55bc1f5f851ed38d26a5537", + "fp32-reference/gradients/0520.pt": "43cc5f2f0e18e1f9ec61a750c87bd9db3cf097f453bfe6f45fa9c67a65edb647", + "fp32-reference/gradients/0519.pt": "2e72ba1a0ea33725378e093d98bf853cdc68886ce67cfe0d03169709b83e7197", + "fp32-reference/gradients/0518.pt": "da36f02ee00005f0702d5a0a01687bc4bfe87ee617ed4e7ede9ebfe1370cfa34", + "fp32-reference/gradients/0517.pt": "7ec01591947add93c475af44b6e1d8728e0d0fea292f998b2e92001da5a8ff22", + "fp32-reference/gradients/0516.pt": "8b36ed32494a658f7d6037e57bb90f5b0f2c44d10c86c63d3188cbe2701970e4", + "fp32-reference/gradients/0515.pt": "73794c97a169cec101c036e22b867b1f80c665eaae74454119b08a6cabf228cb", + "fp32-reference/gradients/0514.pt": "21ac37ba854658a299d18635b581b1a731fe7a873090b6bf2527609744c2f706", + "fp32-reference/gradients/0513.pt": "c45611708b5dcfd3f37fc561e29706b20552566dfe07ec802639a9fb53dc76cc", + "fp32-reference/gradients/0512.pt": "f17e67901af8cee77e500cd7ccc816658254d0e9f0e7f97ac966e37eec414066", + "fp32-reference/gradients/0511.pt": "72ec8e7ff122c89b9ad5f081bb38780a05798baf5db199ec65b0bb9531861a03", + "fp32-reference/gradients/0522.pt": "006b34aa59e38dd230e0deb54437a4d865f21dd5973904e65180fce39e10d9ac", + "fp32-reference/gradients/0521.pt": "6568995b15407cbd88fb3acf9d8d53b66a73ece10977a45a130a9e63f61a8fe1", + "fp32-reference/gradients/0508.pt": "b06e0d4f4ec0c738ffa9aec689ab50ccfb9f1f5b21347ed722779a9a9270c6c8", + "fp32-reference/gradients/0507.pt": "543df7123d3a2a58123ab577238b08f4ae9a996550ac427f9b4e67b0c410beae", + "fp32-reference/gradients/0506.pt": "c1281ee12c79c65278789b3f3088a3878cd180bee4b356eee859f1e15f521662", + "fp32-reference/gradients/0505.pt": "28d4541906dca31a95663cd8871ad0d1bd8f2f9a9b0ec4fec2fbd131e929678c", + "fp32-reference/gradients/0504.pt": "3e0a93c47d2e1aeecf7b88523df493ed221452dca89ad7a7303350ac7e97b49b", + "fp32-reference/gradients/0503.pt": "6e5574a0ac9f3ebf84943be5e0ea709dfff1d14736162babd9dc6822a64cf236", + "fp32-reference/gradients/0502.pt": "2a58c4eccb5676c4a62438925c86a5b4cddb1fcac574428979dbbe0ccc7e6e49", + "fp32-reference/gradients/0501.pt": "13d5eac783b4d9fd547ac31888e1a81187e810d13de9d4edb7f3a74a9579c490", + "fp32-reference/gradients/0500.pt": "9db2150d5ffa5de04cec3cea3c918955b56ebcd2abdccf012553b47b1e09c52a", + "fp32-reference/gradients/0499.pt": "a6ce0ef449017393dd6d0f193f43a390c1694d50cfa234c5ba87575db6c255d0", + "fp32-reference/gradients/0510.pt": "dc6b7b5902b4fdd6cc8f95a4664ef737148c4f658a37e672f599aa49aa878800", + "fp32-reference/gradients/0509.pt": "f06a43d2844738a53cc4465b536eda78a51363c42c4eb301191451ddb72649bd", + "fp32-reference/gradients/0496.pt": "40e8e219f53a6fd32cd13f25f016b5fee910167ea3693d277088df2e30aae283", + "fp32-reference/gradients/0495.pt": "fa09566b9cc7aa8282c5c633a6e020818489d151dc860771a7aeba5130569182", + "fp32-reference/gradients/0494.pt": "a8e1bacaf41297dd7a02e4c0694683410db0cc83f74289f48dc4ef54d3881c20", + "fp32-reference/gradients/0493.pt": "38826546a90bfe5288484fbbc7e16fbfcdb6f99b73afda65f0ae2319b1cf62e8", + "fp32-reference/gradients/0492.pt": "206d571b87c88493d5e302d424c1e21c9b4f8fefb1c5d5b676deafa9dab5c943", + "fp32-reference/gradients/0491.pt": "2ff197a7f8d707fece3e2471469cbbe1e82936f3ee36822c6f95a87ab2cdd726", + "fp32-reference/gradients/0490.pt": "0390cd25a8adf59249cda2d30c3d06aad0d7bade0eb0f76ae049481fd8d42532", + "fp32-reference/gradients/0489.pt": "f4ea4a3182afcff0da048f60ecd1bca38def8ed141ae3e0f38ce36c63c95aa91", + "fp32-reference/gradients/0488.pt": "400662f9763b4ae3850e2d38ab235b2ed26041215462c2c2d8908312465e7aa0", + "fp32-reference/gradients/0487.pt": "30a3ed098eb7eb5f14732fd40665cf8688b246cd8c0f473b6e714525b2c15fdc", + "fp32-reference/gradients/0498.pt": "d825eee41db8392a5a5732534a5ad28a235d0de217b3f3e53b5a752d2cfd2a92", + "fp32-reference/gradients/0497.pt": "704279a36d508ded75e63d94e52be769c79853cc9dd890a72b644962413ccede", + "fp32-reference/gradients/0484.pt": "9e0cb0aa6de8230b78bf2e6c3ce1e545b697ae5953c5fd40497f5cd291f10673", + "fp32-reference/gradients/0483.pt": "e661fe7eb319d398ccf230f331f638c6ec766e1a9da286450ab52cdc0a320658", + "fp32-reference/gradients/0482.pt": "dcb47d66beb78c80a357abdbc490ab890fdf10a1c3ba5585a8f07637c42ab6d3", + "fp32-reference/gradients/0481.pt": "326ecce695eb2154fc11dcb8183c6d6d3f2c539492c7dda9cface28f4a830f10", + "fp32-reference/gradients/0480.pt": "6a1bc745f3205e625de2982a281dbfa8fe9e64a0720dda995b39a4f79159ea05", + "fp32-reference/gradients/0479.pt": "c261a24a15163c2363e7a3a22b8a25dbe3612ac588c1ad8881a9a6f80365a323", + "fp32-reference/gradients/0478.pt": "14b50f0bb5050805c84ad92dc420254d1ec031c4e8574f37f4a40e4f80f76603", + "fp32-reference/gradients/0477.pt": "4c3a42f151a26de43655da95b7987019eff5a4f041e2da6dcc67fac8d3fae685", + "fp32-reference/gradients/0476.pt": "713271643a14adf133f189ebbb69ba97aea5761a975572a8be93d1dd8dc4d9ba", + "fp32-reference/gradients/0475.pt": "5a399f516b0e67d766153961a6f4031f82d1ad223bc169a1d851c7e5321c5f6b", + "fp32-reference/gradients/0486.pt": "46fdef3e713423cb93df30ee7285674bc26d14e5e7060fb3407987a763be2583", + "fp32-reference/gradients/0485.pt": "d8975b3bd49a354ed8a856863bb38cce8090082bca012f7bd6238958567f2b1d", + "fp32-reference/gradients/0472.pt": "cbbb64c0c541c893a28a09427a0d1ebfc3892f4d48dd648fa2ed48d16b39dbe2", + "fp32-reference/gradients/0471.pt": "322be960b040310dc6037828a1ab58439059daf66fe1a36a1b5436e710624cd0", + "fp32-reference/gradients/0470.pt": "ef81d5a07a2ba847257b5cbef97233b346cdce42417a4df78748802ad0251fa8", + "fp32-reference/gradients/0469.pt": "59a1240c5f65c6f87a658123d7c25db803a29a30839503649fa5f75bbb7382c8", + "fp32-reference/gradients/0468.pt": "957a78f6d9061c3f7f4a0e2c8dce6cf958425c63178dec09a9789fd20944211f", + "fp32-reference/gradients/0467.pt": "4aee1927a6078a296a97226231b13dabc98d6e79ed5584f462dbe60e1187e4db", + "fp32-reference/gradients/0466.pt": "8dd58690d2b120c1f9a6519ec7f030351206f11abc01b49e1c513faf3c9d7a22", + "fp32-reference/gradients/0465.pt": "f21a5c5a9cc17a30f1d095e64d9e2d0e6442c80ab3f4c6fe185fe4b129ac8019", + "fp32-reference/gradients/0464.pt": "f425133d2df900b0fdb4c0308fcb0bb624fd8b4e5d82c8bfd0e6e2a85cf53b6a", + "fp32-reference/gradients/0463.pt": "3e2b380eeeaec92c66c2042d35c5ac491dc0870025eb0f578d7070db0324a80d", + "fp32-reference/gradients/0474.pt": "930be3a0d6fa35fe355628d8a3e203604c5174baf0f56e91de547cd36ebb9595", + "fp32-reference/gradients/0473.pt": "5e9f2affab05ae19018a73d6317635b76bca52196226d4b36f50081a3acbe1e4", + "fp32-reference/gradients/0460.pt": "c6a71ba438f79c7e7cc15f581aceebab346bf2b3ea3aa97bbd28cff471c28205", + "fp32-reference/gradients/0459.pt": "35267c9a1f22e7ea83e08a7e1b47760de3a119f173a941ea58b5e03d5fa2564e", + "fp32-reference/gradients/0458.pt": "7b3cc3bfc0c6c03c9694830532949aba765d6c8eac45843532328973f7a9e878", + "fp32-reference/gradients/0457.pt": "c6ff9b6a3ad24ba2ed0dcbf80a8c42bbe3d23e96d2f64daef69f8cfc07b6ffe9", + "fp32-reference/gradients/0456.pt": "603609358dbbb1edf4d134e4edf8f8e2eafb7980a6d9ec4a1847e3c92c1abb6c", + "fp32-reference/gradients/0455.pt": "8c8805ecad1a9281fb9aff6f25c8239ac38bd35fc10e7afa5055d607dd18bcf2", + "fp32-reference/gradients/0454.pt": "984b2d3a175b67d0b2b3ee91dc6a72c0e2074faa0958386c2f9dfab2ed65bf7f", + "fp32-reference/gradients/0453.pt": "deab52b241bfbe5a72a47bfd8115fa91df25f0d3732cb8d8de8db17afdd42345", + "fp32-reference/gradients/0452.pt": "f2c89000d816ae526c707d1806862160a80f69cc9932a131f28b07453c10a302", + "fp32-reference/gradients/0451.pt": "dd33693e1fc442b2779b7192aa95e6607b396fd8eb34133fb1259dc0adc4bdfa", + "fp32-reference/gradients/0462.pt": "98cfb654e544a673ca011163aed6f1e9d9f9c0e551ba96147ddd8729470f3d06", + "fp32-reference/gradients/0461.pt": "2c03f7c8e551207b401d849529c1861edad03f1c18a7087f9b0dd80116348054", + "fp32-reference/gradients/0448.pt": "81cd65ae67a59af052232c49299b08cd30d669c216602207bab48f49182d8828", + "fp32-reference/gradients/0447.pt": "af8e51ba5a0d5e4bbecb1dc7b107a77a14a1ac487da2ee16185573f7be5e3455", + "fp32-reference/gradients/0446.pt": "677a8910044036fce0c62f5bcb84b7a372da4c3d7181250c0ddb0b8a63b13e09", + "fp32-reference/gradients/0445.pt": "487db7dcc022a52162cd98cf26dd8147355c608a21d13c6af16155bb82318070", + "fp32-reference/gradients/0444.pt": "df3201fb9a082cdb9d93c8066e74a5318681a60f5f924ab7e064e61030ee8974", + "fp32-reference/gradients/0443.pt": "97af11e81de3f569e077cb031cd98f894e6c5d3d4588867302e776def0f5d1c2", + "fp32-reference/gradients/0442.pt": "127a82048329e8f6784e0774a03c1f19397b9a69feaeab3e6b7b3fc244c84435", + "fp32-reference/gradients/0441.pt": "6a76887e80fc3135b73a551ce908e64e23e78b67dfaa5f4f4988b7a8a4d3ba83", + "fp32-reference/gradients/0440.pt": "7a6edef7f1ee390b8c52a06161c7fe053acc7e2f38c178f28203643715627021", + "fp32-reference/gradients/0439.pt": "993197a391f39470dd4d0f5868ac165767b54b7abec0465e77ef4a2e62a7f6f8", + "fp32-reference/gradients/0450.pt": "973638a3a1b1673a22a62cbdabec4617f0aee24a103978c2cd6aba5ee672047d", + "fp32-reference/gradients/0449.pt": "2d54af03139f375dede9804624777a461dcbb4e65be05c0e4e34fa53c0c475ad", + "fp32-reference/gradients/0436.pt": "8ccb65336ec154b5ceec97e7c7e77771959094c970304cbc9e669eedafbc984b", + "fp32-reference/gradients/0435.pt": "a06840af610d2d695f155cb579496efc71a48eec470900a9367e4b4d2d6d0589", + "fp32-reference/gradients/0434.pt": "498768b33a2ddfe59e2d05643197b65d88532805e5ff2424bb13ea34183b335c", + "fp32-reference/gradients/0433.pt": "402b8c894be2381a1b3effb145f8d7c2b315fe7d861a2b554171f1dd5c5932f5", + "fp32-reference/gradients/0432.pt": "281948db17d5cba6b3147be99d40d83a61a0eb584afb64a5a3559931cbb1f1d0", + "fp32-reference/gradients/0431.pt": "45f30f12ed0ed990450967f11f708e01131c7010ea8a911cc30dfceac68c54ed", + "fp32-reference/gradients/0430.pt": "aba14a507adb1457d5bf2f1bd78252a402195b9b79eca8daa88b5ed5b8af64bf", + "fp32-reference/gradients/0429.pt": "c9da0de7771e8ad19613ce08bb8de95adc5b4c9577ba57946f13d7c1ff3e8832", + "fp32-reference/gradients/0428.pt": "4979d33eb8e70d2c8951a180cab912154ef015b4d7808aeddf061043ee5c4331", + "fp32-reference/gradients/0427.pt": "1e4030f3ca19efea712ac4f1939f3522fb3dc630cc79766b86f0fbd9ef543b23", + "fp32-reference/gradients/0438.pt": "81fd57d56eb2af15ce2d8e45c56250cbc5e9b2ece6227bf71ecb12c73277b7be", + "fp32-reference/gradients/0437.pt": "eb3313558184ac8a8dc339736ff5c5658a2ed3d865d43cd585072c66fa8ff2fa", + "fp32-reference/gradients/0424.pt": "0cc8f805a9744b814159f4e6b272e70a1b92662f85d604db8d46acca1128c38a", + "fp32-reference/gradients/0423.pt": "5530b03cf876b7fb8a5450f1f0709eb6e91ab4b5e23ac76d5fed093714375eaa", + "fp32-reference/gradients/0422.pt": "690554a5d80a0b98ba6378f3c6ed6d41357e4899bf9bc88cebb32b8bae9abfa7", + "fp32-reference/gradients/0421.pt": "9879ec70748ad0f7c274d709b1d8a24335398aaef2db6e09a2cd7d8f912ae3d1", + "fp32-reference/gradients/0420.pt": "c2e67c2a62f16d8f19068d3b9bad00bf1dd9ca62cd2ae6e54232f6dfb71e3b66", + "fp32-reference/gradients/0419.pt": "31be6eff12f5ad14cdcbb43dabe3c493911a953f32dc1c60f06306fa6310e2c7", + "fp32-reference/gradients/0418.pt": "ed3d0feb5ed1eb58bdb4e516aee3ee1caa866d0bf9d9c04b37f25e464317e108", + "fp32-reference/gradients/0417.pt": "0740fdd7326f0ea2c1bb4ea7140359b1f329341d3f20db1bc1f8c3d0b5d927d8", + "fp32-reference/gradients/0416.pt": "556ac4f1e52dcdb880c4288242a7d1761021b7cfe7d3f9c3c7671c49d0ab1a65", + "fp32-reference/gradients/0415.pt": "46a96495be56bb52abdfd7267938fd55044658cbb2894767f4f9ca014f698fb3", + "fp32-reference/gradients/0426.pt": "e070b09cf801ed3a8a65e5b2ec63ecc1642b2389bbf8f7486a9e792af4562c3e", + "fp32-reference/gradients/0425.pt": "55c1d8439a2cd53dc49bc35c43f76ca4a3f76a2d2054cd85e025bea41cd3991b", + "fp32-reference/gradients/0412.pt": "61dcac766920bc765df7fc69fd30e77fa8a07f0874bfa9ed4889164a60e37cd3", + "fp32-reference/gradients/0411.pt": "99c45965ccd1495cc6e7821ff2030354791c435be19706609dffc55f52197429", + "fp32-reference/gradients/0410.pt": "81b4d92fbcef872dc767575bd04fc04d0a99e7d73c9bd7b534753277bb4acf20", + "fp32-reference/gradients/0409.pt": "8304957c813471f601c71ea518fe5a49249638d0413a3f14a20e20df00a17ae1", + "fp32-reference/gradients/0408.pt": "ca69b37d38be66fafa21ef99922f3e6ce737701ff17c862db68779f15d9eda15", + "fp32-reference/gradients/0407.pt": "b927fc05062ecbae856da3b7bb845d5f46988a2b11763d02720a60abffa2de35", + "fp32-reference/gradients/0406.pt": "b2188a96cd251ab54b15c7f05be6168823cdec771bc8a5563cc2ed0ee713c57c", + "fp32-reference/gradients/0405.pt": "cd009b42b8316e84b06670df68c3f1056058ce1b3c6ea5eebcf2e39ab7fdc7a4", + "fp32-reference/gradients/0404.pt": "0c092e2448b5f7c815bbd67c06052019551ce387d77000e192841c8aab18878f", + "fp32-reference/gradients/0403.pt": "1c10e4c4aee8e5232206d8b8268d2e98f768ea5075122e50d99f7273d13c282c", + "fp32-reference/gradients/0414.pt": "a1a5c98c122831855aaadffe0f55a746acb3562b5cbacbe983640a880edf74ab", + "fp32-reference/gradients/0413.pt": "d226b6558e5ef7bb333da5318dfc2c998674634c47c59333fb4812bff8df1565", + "fp32-reference/gradients/0400.pt": "babdf8d72aa4ebadf25177e8755a63e216a639ca599c868583d71158aa4b2ad1", + "fp32-reference/gradients/0399.pt": "f4c1f48734bfc6a92d095bc358b0bca49e29058ec8bb0bb312d4662e8fe4a6e0", + "fp32-reference/gradients/0398.pt": "993c0beb924ca34bd796c74ff2ace0e3b6b7b63189e6822fed76ae022e40c151", + "fp32-reference/gradients/0397.pt": "cb1b248337dd8869f80f077de96241ee6a3e0f1ad89ee70ca32797bcebd7d2b1", + "fp32-reference/gradients/0396.pt": "479869b785c2e46f81e047701031624a34ddd0daa1fce350a4c4c13f5708f046", + "fp32-reference/gradients/0395.pt": "02bf140d6ca2c7b05b627c671771d58a12525b556fea24d0e6b3cb1da8a19894", + "fp32-reference/gradients/0394.pt": "6ea8d33e05f192a3904d0df8238b4d3dc06d60885170d677504109b0b2f42f7b", + "fp32-reference/gradients/0393.pt": "497dd32db340a8602a5667f0967854fd47efd8b0eedb0d78a1b1892dae69b0c2", + "fp32-reference/gradients/0392.pt": "e241fdb37146d5df78d6a5cd2a386af505a314af8cce700d4b31b858412aa07e", + "fp32-reference/gradients/0391.pt": "4cd2bc752abfee65b34d656a719e3d41ccfaaca69b83252dbdf441b0ced575cc", + "fp32-reference/gradients/0402.pt": "b269ddd15d6b3d1c7eb93abaa14eefb096f11a07ae2e1a7b40b00f8479969364", + "fp32-reference/gradients/0401.pt": "c3f791af8aa2f3dd31b82faf4e24be613ee1af339fb3dd0b72307e21d905b3f6", + "fp32-reference/gradients/0388.pt": "055d0f42b1d00d33caafb3b36e0a191d77e664ac9fa6d1622fb73de6ed29cf8f", + "fp32-reference/gradients/0387.pt": "09eadfc79f8d76fc9fc1527980e87fb32cdabbf47820fba6a2bc9a96e417e112", + "fp32-reference/gradients/0386.pt": "bcbac108b3e79a5920c78755848acc74d3ea88ab83c7d4cfb66a3ddae5e49701", + "fp32-reference/gradients/0385.pt": "dab9514a1d728bed815212d68f49ccd086f65685636de4f6428d3a4b62aa25ca", + "fp32-reference/gradients/0384.pt": "61c28bf78e73158fa599ae5e4fca1fff10f631fcb242fbb19603dc063f42e5b6", + "fp32-reference/gradients/0383.pt": "f98993334332d6b3fb603bff01157ef4121f68e83ac4ffe866d64ce7e7f79e74", + "fp32-reference/gradients/0382.pt": "da0dfb211b031fac433ddc0f22bc5837e944ea9724e8912866b8085f4337c6f1", + "fp32-reference/gradients/0381.pt": "5c4b234b4780b7aaaa601c4cfe152bba2369655900fdb39f4989c8e8f3899c2e", + "fp32-reference/gradients/0380.pt": "0bcafb5cf2c1c7521b15dfaadd32d254675adb96e9b0bdf2c2350f60ebf1fda2", + "fp32-reference/gradients/0379.pt": "fc1adaad6c208cf4ed78501d6ca918624f9b5abe395267703a9b449c3a391d89", + "fp32-reference/gradients/0390.pt": "926fc65ea61633bc44e4d42e1e32eec564dce7b3043d617512e0fabb5f90b087", + "fp32-reference/gradients/0389.pt": "d7f9ef3c7b1aab9cf97eba892e19793aaca7b4270d2a04411eed323c608b6fd4", + "fp32-reference/gradients/0376.pt": "fca479e9c367975ec12aef6bcdcbbff639faca4441503a7c4a7388ddbd1e81cf", + "fp32-reference/gradients/0375.pt": "8fedbec498e53baadf9b7921de73e53fc3b4ef06ac7a39a7717b09511909eacf", + "fp32-reference/gradients/0374.pt": "cb14cd46f5564f1f558ca0ed65cdb0782157ba931887de4efdc973d0c4d79f94", + "fp32-reference/gradients/0373.pt": "6deadc0364719fe6f8fc6c5dcf3c308bc540ed97a46da3b17e2304eb9a197dfc", + "fp32-reference/gradients/0372.pt": "6d87e6164df3fb9e066120461671f47be986b399d2a97890fb86e1ccc250dcc4", + "fp32-reference/gradients/0371.pt": "26acf2f7ba06424123a879fc2c08251fc1afcc72fe71ce13764fb14f21a9e247", + "fp32-reference/gradients/0370.pt": "21a95af2c977ede2b79e6a14e83e70f44b5eea58fae021bb3ba189505823f6d5", + "fp32-reference/gradients/0369.pt": "5040b554bd3a9e027cdb0c3a53b609b6eb2b90636932c575a00b5f9867244a29", + "fp32-reference/gradients/0368.pt": "c5f6a9a097c9fd6aa9c978ffc1d9b5d223370be963f293e2d9ebdff9acffe1ef", + "fp32-reference/gradients/0367.pt": "141fc2a9740cca0e795932932f082b5f0a6fd2f4bd235f5ba88b3d3550703767", + "fp32-reference/gradients/0378.pt": "2a5160f53f07d9545f97166fbc156e0a998557d251201100b5c8ad04e72c3216", + "fp32-reference/gradients/0377.pt": "4b891cea1186ea29963536e51bff3f61f4b131e9902396104eab8fe0883a6301", + "fp32-reference/gradients/0364.pt": "297e433afd8df2f4e363d2deb099d38e41c13ddc4acc0c5ac4d5beae92353886", + "fp32-reference/gradients/0363.pt": "830c8fa1f764f5d3c1a89351070f1ad76e4340273e46b38fdcc068b2f1a4fe04", + "fp32-reference/gradients/0362.pt": "b71b38fbcc78f6d31641241633e001f20db6f57529c570c8071e1147169fa42f", + "fp32-reference/gradients/0361.pt": "cfc0ea845a88b63b1e22c922564d39fe64a6d7e9eb0f049905bf513245625f5d", + "fp32-reference/gradients/0360.pt": "c1b8a40a0859cdf37efd0133c492c043a5231f12f616cb7acf46601ac3af6899", + "fp32-reference/gradients/0359.pt": "6da7d1c92ea12e9a8562f985157967e36f840e07f5868d41f1d9873941d289ca", + "fp32-reference/gradients/0358.pt": "98381516e776803417d7441f825aeef0cce21a357310943c3851a2bb914790c1", + "fp32-reference/gradients/0357.pt": "d230e6f3b0f30ad25069564c7f6fe583acf3ca88840c86bac1b404c30f590caa", + "fp32-reference/gradients/0356.pt": "f142bbc7005fd0daad1cf4f2b9e5e698752eda421b75b18bee14a1f3fb0f0ae3", + "fp32-reference/gradients/0355.pt": "5622fee6dcaeae422325f614b9feb9d809d1cb0d86dc35f90735bb015baf51ef", + "fp32-reference/gradients/0366.pt": "be69db3dcc33b8749c8b21bcf9da4afcccfb68fb671884772afb623c57b2fd99", + "fp32-reference/gradients/0365.pt": "4c254dd353ebc1b09b7eab21170765e489b8d9ddd798dd99addf772285fe0cdc", + "fp32-reference/gradients/0352.pt": "d0ea48a2b67f353d9f3f8148e93069b85182cb3f306c5c7749e989f9776be585", + "fp32-reference/gradients/0351.pt": "664941612ea2d708d020617616bdecc811efc94c6baf4b3d8e8fa060a3f322f2", + "fp32-reference/gradients/0350.pt": "a46c006861b58b0959738f1a7e4af5385b9e9ff72bf8f6e1085856c1cef6adcf", + "fp32-reference/gradients/0349.pt": "01646aee7e5ec1f0c3e0b270a89dd93955c977981b0a835d88018e771f191078", + "fp32-reference/gradients/0348.pt": "dad5413edd983c8a8bfde79146d99dfe938564fa8d0029cf1279e28b48690186", + "fp32-reference/gradients/0347.pt": "8e67e8e8331059794a4da7d91441617ebaeadde4879a1acfd23d61c1fc2703f7", + "fp32-reference/gradients/0346.pt": "7c44095544782f10c97dbee924dd647e6297553676ab13543a3d832f2c32aadf", + "fp32-reference/gradients/0345.pt": "3b91220e343c52b71eac538752c6e2661b4756ecdb3a5d314b8546e5d0c7a77f", + "fp32-reference/gradients/0344.pt": "18529e38c196dac0c5fe86c3f7bbfbfb02db2875d9fd89d0a722e6c699e12573", + "fp32-reference/gradients/0343.pt": "0ebb81d928088e9a98ae13c93db5964545e382beee8cbf9c58d366c59e19f97e", + "fp32-reference/gradients/0354.pt": "7f28f9f54e3efc7e43c30f4beb199b0784eaf365e86cb343f8579f2cb31cceef", + "fp32-reference/gradients/0353.pt": "e9bf6f6ec4046508264dc20f914f629f77e9d6e3010a26c47334cf22cf22a0a7", + "fp32-reference/gradients/0340.pt": "30c1ca0449204c49165ee0ae647516732b457697e4a0c46dbbeff6bb9bfe4103", + "fp32-reference/gradients/0339.pt": "bbb21d4990a263e4bd00bc6a685d963beb4e8ead1397e6710678804a78a697a3", + "fp32-reference/gradients/0338.pt": "1ec28f042a16d01964213fe873cd8f4358fa8b84341763307d0393fd7ab7f20c", + "fp32-reference/gradients/0337.pt": "31cbfb86edcd646f8e67050b218451fa2696bb22ba4599a6a3b4e32849c57236", + "fp32-reference/gradients/0336.pt": "26e3d7b7d6ea02900d252c202288afdb36b34b580393ddfd90f877ac1d158792", + "fp32-reference/gradients/0335.pt": "ad309b703b42366a9bd5df22967d2663240da28c9056a32f7a039f164b439c32", + "fp32-reference/gradients/0334.pt": "e6140fd83b826824ebe2dce3131947772fb0d95210a110b101b54ad8d99d32d5", + "fp32-reference/gradients/0333.pt": "91061cdce387adfbadecd2b52e31be797d093ccff436cf7c337483d26df30f9b", + "fp32-reference/gradients/0332.pt": "6784885842c58905df6382ba505a76d4271d5faae11d8486203ad93db3b8e086", + "fp32-reference/gradients/0331.pt": "f87214d58ec9e88064e24e472391264823ae025a939c9c0cff9408b63d72798b", + "fp32-reference/gradients/0342.pt": "c80472d05577d563baeeb3da54f51511a4d588b29e95aac30872ee23d617265d", + "fp32-reference/gradients/0341.pt": "1b97fe915bab23a8f2acdcb9a4e14e12c1bde2b495c19cfad61a5d19dd52a0f0", + "fp32-reference/gradients/0328.pt": "a46f533054a001ff28f32b915cf85a05d8b95bd4cd13ea4205842b611b1e6f9b", + "fp32-reference/gradients/0327.pt": "9e85dcbfb51f3d492c3c93b87d5cae0bc961110b3be6c94a6281fbb19608d151", + "fp32-reference/gradients/0326.pt": "61446f3b0ab7612c75cd61a325ec49977d827985ef78c0c9d4407073599eb586", + "fp32-reference/gradients/0325.pt": "f30573b2637ee6c35c97042a913a2d48112169bdcd6b31ab5155aabe4febdd0b", + "fp32-reference/gradients/0324.pt": "9ab6259674209f759dc294e731a10a80205973b98c201e8a2cec04a63a51bfbf", + "fp32-reference/gradients/0323.pt": "e7bc8e825b3b558a0740fd26425d103647650801e3cc1103fd0d39c4cf3fbfda", + "fp32-reference/gradients/0322.pt": "cac33dfe30e43034a3d814b379cd9ec6390e3307941ad5bed24a1c0a79d0f2f7", + "fp32-reference/gradients/0321.pt": "932fa66b71a1bdf1e8ababcfd9ecae9ce9ed1bd8589791b7342140ecb4cc5bc0", + "fp32-reference/gradients/0320.pt": "94d781055abb8fa924731009a5ac4c6f23b27d69364b92a367eecc5917ac9536", + "fp32-reference/gradients/0319.pt": "6047d1b068cf0252ef47033e1111f070a19919ae4deba4e95e6486b2fdd1367d", + "fp32-reference/gradients/0330.pt": "43b519b8c6d8a59c735918b465fb461b8026356941cfa3fc0fb777f4f8439f4d", + "fp32-reference/gradients/0329.pt": "c31311bb28b3f451096ee4017f227c457f0f9f4d9af8044ee906c20af5e37516", + "fp32-reference/gradients/0316.pt": "aa535cfd1f5e50eb830a4451df49b5b4a017b4a5ffa0a9cc3bd0ce7b39c7af06", + "fp32-reference/gradients/0315.pt": "0e3a47131d9328e1f4d9eb6d202ef20decab26ff352d84b3bd8c16c232cb998e", + "fp32-reference/gradients/0314.pt": "c900945c0b6af837c551a71df7edbafed618c66d15f8ac8be9340608c3bcfa49", + "fp32-reference/gradients/0313.pt": "c98d8cc2c560a21ca866f8d4a0b3ad3b52ec9997e8329e61f18db9aa8af1de09", + "fp32-reference/gradients/0312.pt": "529b655af8096c58d58fb1b3c3aff067b97a68729ea84e9604748c4ac54a7844", + "fp32-reference/gradients/0311.pt": "07369812478dbcca4562706fa2523f4b0262f8c2b8ca5e1ea5f6b10649ca599e", + "fp32-reference/gradients/0310.pt": "7d810f9a97fb1a1a5f3f0ad9c4919b6cc05ce2439b1989a2cd5ff28a90aad19f", + "fp32-reference/gradients/0309.pt": "8aae11e69cea48219c2994364b21b5939b6d72a082393881cb50ddbb4658fc4d", + "fp32-reference/gradients/0308.pt": "89bec3cbe3662565b5a735087e2468b2284b3be9bcd6efa646dfc783803f0747", + "fp32-reference/gradients/0307.pt": "3e5769ec13b31cc8dfb80caa425e24defb8c245815f3b04bcb2a021bba2356b9", + "fp32-reference/gradients/0318.pt": "6a58c2cd5662cf5f600b467b827730026bb8e7fefaf1ed0dd59dc31ac1b50b03", + "fp32-reference/gradients/0317.pt": "8349ec6d6d454cbee6161a21fe93381a31d99f041f545c0420df84511e3ed59e", + "fp32-reference/gradients/0304.pt": "56927f97e433dc8c4a4ac52ff00747a03e7c8f4c9e3014540ff256fdaaccfc2c", + "fp32-reference/gradients/0303.pt": "7eca12e84511596c21ffb80a2b25d80cded57e5c8ad6b0de1e785a6b9501f590", + "fp32-reference/gradients/0302.pt": "86bc7606eca062317ec2bcfefa4eda0cd3e53c5f2b5d4175281b4ad55a53672e", + "fp32-reference/gradients/0301.pt": "e826563d9b100cd3e3e22b7538d4961642fbfd1381b3973664dd96e5d7e50755", + "fp32-reference/gradients/0300.pt": "f1a98109766ef78c2eecb238dec287d4ff647fb8dce0c2caad229f853aaff82d", + "fp32-reference/gradients/0299.pt": "f948e050467b4d540dc4a9339f551c317e8ac12f0cdcc7b1ddac35bb50778d71", + "fp32-reference/gradients/0298.pt": "36d0016c7ca4e22235e2e676b5ad3b82678375744e2a3da7f86a63dc39506099", + "fp32-reference/gradients/0297.pt": "2e7e293b196e88cfcdda3313ee7cb9f2586af250b2dfd64894c52a71380636e6", + "fp32-reference/gradients/0296.pt": "44fc8f38aad56d152da5931db0c48a1d1adb09cf93da010feed7e718a6542d53", + "fp32-reference/gradients/0295.pt": "d428d1ae88d44f78d842e2f86c83a55c8614dd98d9146615c5bb86c76c5839da", + "fp32-reference/gradients/0306.pt": "47977dee08ee2be3dc9b3457e67fbc668b0748879f61f320e37e3f4802a0a60f", + "fp32-reference/gradients/0305.pt": "83a717de444b9e8c67f092a1486764954c44df1c5a3169994f02491ae3e2febb", + "fp32-reference/gradients/0292.pt": "e6f26acfa525135c68d7aa2a621385489211f29234ffaa1391f71d9136ba1abd", + "fp32-reference/gradients/0291.pt": "6568c1c9332be4f04d93f7c9df5403705244be1680065951f8d3111c10a76311", + "fp32-reference/gradients/0290.pt": "ac9d5e5ee0f40b4da80cc7b250c479e98d16fd24080344313298ff263978c342", + "fp32-reference/gradients/0289.pt": "9ccf2de8ca240f9ae55c7d54b517c71c39c9eec1e6e328cb0309c72255878cd5", + "fp32-reference/gradients/0288.pt": "b44208f51754dcf8dc83e82a48b3bc9390c5fbbe8ddb8d20c05c556ddcc0fd1c", + "fp32-reference/gradients/0287.pt": "4e9ee7340584f74e2769eaf5397214455a126302c5c9c769dda55c5429a9fdf1", + "fp32-reference/gradients/0286.pt": "e958e029137c0eaf4321f4e2becd71fcc591d7e41fc56f003019ab6af5cbea44", + "fp32-reference/gradients/0285.pt": "2eb3c1cf558bd3efbf1ba135be3309ea7f4625a69ba53d8ef7b7f7d540e5ab77", + "fp32-reference/gradients/0284.pt": "c9f18d0ba332d302fe1db2ae201c3159759650f697db7583df26fde6f1450460", + "fp32-reference/gradients/0283.pt": "57d52095ba5b541d7bf1e96ab47b16d12071ac63ad48386c120ff458492f9ca4", + "fp32-reference/gradients/0294.pt": "7f54a9b76f96e3c5a15d276f057d7215db701fbaff25104b7461a6a4e534d992", + "fp32-reference/gradients/0293.pt": "445dab90ea8620c59f2ede7de5ee4a79a9da227aab31767d45ab433116e2ceaa", + "fp32-reference/gradients/0280.pt": "bb173c55f854dc97199e915da49e6c61f763decf02e3af4fd2a9229932f600ca", + "fp32-reference/gradients/0279.pt": "09918cb3e216463cfe2229411b587edf64e184db29ac4449d1db6924241a93c8", + "fp32-reference/gradients/0278.pt": "24ead138531983bdc29613177e2a12d1e7c5a99521100e444930ffd84779ee23", + "fp32-reference/gradients/0277.pt": "0791e037b3685357f850875d4eaa239127eeb3170ff72985fd509950c7b1926e", + "fp32-reference/gradients/0276.pt": "bb7689d72d84a0d2ad58677ca28a9747c33fcb1f5e50a94843aa0425f5380270", + "fp32-reference/gradients/0275.pt": "f8b425462d1239aa42a7d33410233d39ac02a9df052d3d0cc1d4aca8d4594ab1", + "fp32-reference/gradients/0274.pt": "9e3a878665c09d337a3a349a2536f4e2dfceb3dc68f44dc9be4e2863a925d55f", + "fp32-reference/gradients/0273.pt": "c3f20c91bc0b816d38ec0c46219e0776fdb1041b59b5c3df7c3708052974ed7c", + "fp32-reference/gradients/0272.pt": "f9f149bc940f7cd388b1d2303857732ca510494c392a303d2404ef598e61a408", + "fp32-reference/gradients/0271.pt": "89c8029ad82790e44abbdbe41c813e520967b96cdcebb45e98362cb419c03d4d", + "fp32-reference/gradients/0282.pt": "01461d30708557a310882c612cdaea64dcf2f4db11372f6c7c932ff110bb8236", + "fp32-reference/gradients/0281.pt": "83a5000db45c95e38d4be269495bc82fc456ad2ce7294fe6728ce65e49bcabc9", + "fp32-reference/gradients/0268.pt": "74b1618f5397c50676856eb6ec11be0dfd0934599535f89fc2a2709791247620", + "fp32-reference/gradients/0267.pt": "48aa3a0fb63758a5ab02cceb6a0b2a3d6d339b8420085091b5a05e572819ae72", + "fp32-reference/gradients/0266.pt": "a06bf21815dbd6ef3b6b7f08022ad523ea583af27ccdd82b6a58221274a57f16", + "fp32-reference/gradients/0265.pt": "87d805f43b18dac430d0d63ca05dadd55a5ca32e4852856f0310381a281dcab8", + "fp32-reference/gradients/0264.pt": "a469ad817abaeec10adeebaffb9aea07c8befb7279634d2d788b29c40821a6a3", + "fp32-reference/gradients/0263.pt": "08a25e62a6f8680731ae589381bae0429081c03cfa255cafbe06d04b4ea08675", + "fp32-reference/gradients/0262.pt": "3904581f9ef042d39c459e0c62070846962d3d5b7200920080d815af59307c94", + "fp32-reference/gradients/0261.pt": "015443a3746e4b298cea82ff383735370802948a008fa478d831ae47fde942de", + "fp32-reference/gradients/0260.pt": "62672784add9ad400444f1ce37f8cdc5c05a2c04275920a2408fd4818a7aaa99", + "fp32-reference/gradients/0259.pt": "16cfb35a77918631f4affa6bf494d7b7ef28de3e3cb3ba078112bb380117f008", + "fp32-reference/gradients/0270.pt": "a798d90c7b1e37acfc75eb3c8fb9609ec51a607cf02350ddce86fdd1f96c1c65", + "fp32-reference/gradients/0269.pt": "4042dde7e9a5f8a2137dd9f3b21adf3444cbbe0eee4094dde3532a31e2616bbf", + "fp32-reference/gradients/0256.pt": "5d9ba1422fc50d1ced3a8e7845cf09dfae1e2034f62cb1646887000a2aa2278c", + "fp32-reference/gradients/0255.pt": "0f6ed31fe4b00d96ce35c3754d1439f08f4201a23869164840ad921fc7aa523a", + "fp32-reference/gradients/0254.pt": "13bfc45f4ff6d72e810ace53e61656914ea406da30913629b7dce9e6b7fcdd22", + "fp32-reference/gradients/0253.pt": "957610f17e94a964b3ce98682830f9d5f033338888d431e0fcdf4636192478ae", + "fp32-reference/gradients/0252.pt": "cdefa2367ac2f76cb315427623e7751731247cfc463fb6961e830209a8f36565", + "fp32-reference/gradients/0251.pt": "ff6c7251f77ffa0380a5b216f474e5cb89b2d0c65e78df0fadb9b782a2a8a7b6", + "fp32-reference/gradients/0250.pt": "331a77e539eb03f25188dbbba071e4890eeb5ee0353efa2d6e04fc3a0e4791fc", + "fp32-reference/gradients/0249.pt": "58cb63a6e35cb65fbb897bd7fb45cfd59a5a1fde8be6ea089e76f7428f4937c8", + "fp32-reference/gradients/0248.pt": "6573b84ca394825d338b9b3d7b3f6e3a329ca505ebd81480ada17009071343da", + "fp32-reference/gradients/0247.pt": "53f565fc6d5264df52f5b0947e017d072cbd57f578932041e2506e6290a15820", + "fp32-reference/gradients/0258.pt": "9710d5632b54a5317ff9563a004544de71fd583e2b328293e3977d42d46a333e", + "fp32-reference/gradients/0257.pt": "805d4b08afd97ace82c902e0aaee865d417b1999cf52f445219deeafd8b70d7d", + "fp32-reference/gradients/0244.pt": "15da11791244a0b7a82e9c4b8fd4a7ca450d999b8296ff93767bb52f13c09321", + "fp32-reference/gradients/0243.pt": "f52c438fc875bf0fa379b8f89a49bf45a3871a967290445066ff196a9005eba4", + "fp32-reference/gradients/0242.pt": "68fbe79a87f5151b3d6aac46b7915788eb0f06e8df15ee7d4d98ab7a97c473fa", + "fp32-reference/gradients/0241.pt": "32018a1c61616fe55d4eb061f4fe97691dd20384c47cc91ae78c9874e9c8c540", + "fp32-reference/gradients/0240.pt": "fe17142d8523a3e82c108c252aee0ef93968eeebf2b39109651cd22b11a2d1c7", + "fp32-reference/gradients/0239.pt": "52eeb688c6bde3a35c304e8b3723e80ca13ec8bdf0fe9a650c564ddc23016b5a", + "fp32-reference/gradients/0238.pt": "90523ef79ff4d4822ec45016d7d033d848220f9a215b54a2ac376aaad2ee73d2", + "fp32-reference/gradients/0237.pt": "4c844c2ddf60ebc6711e4f4df72bcfcdb2985016ff30affa67f1226c92f028b1", + "fp32-reference/gradients/0236.pt": "3cce77166f9b8a42cbcb57bca033688f6b6bff865884449454437acb361b1859", + "fp32-reference/gradients/0235.pt": "63dd5be6524df313c429b7c64049a8a710b5e0b8e3e7c67e456e99ea110e96b9", + "fp32-reference/gradients/0246.pt": "5c74a456b051370097290c6aae8c654804ef4f5d819ff5f36f21a0758855b6ef", + "fp32-reference/gradients/0245.pt": "f34ddd4fdc299207f52de538234b5071ff28a6f0f7ffe817916af075dd1f0f12", + "fp32-reference/gradients/0232.pt": "63a832e049d3a280a339d9667adc479fa409493e08f7911244114a128cd3e854", + "fp32-reference/gradients/0231.pt": "53a28778723343eeec78750c57c7ff4d69d610ff906a4d2915a50cb3b16d9af2", + "fp32-reference/gradients/0230.pt": "2a01daadbf34e70ecc1b5cc1fa8d44f6add72ad9a7f80c0e7f6310810afba0b3", + "fp32-reference/gradients/0229.pt": "ec95345a5bb65c8b09f7bfa56a6f18f10b693b81878dae1b13bf157b27f5f9f9", + "fp32-reference/gradients/0228.pt": "88893e25ac13fbb7f61549868f5efd41e0a310fad2ef92efe3a67f6a9bbd0f50", + "fp32-reference/gradients/0227.pt": "ece7264aeb17f5cbccfebe102f64ca822a9d36d786958126da75454f61d3a8ca", + "fp32-reference/gradients/0226.pt": "737db907ea2d2bc477b0f713340f6b6a3968ee2c4c29e199bd6343d96cf74ada", + "fp32-reference/gradients/0225.pt": "780d7021878199e7385d612bc2736ef6afa30bab555c9a5bd1d3551751a0094b", + "fp32-reference/gradients/0224.pt": "f31570780d9507fe089fc08ea55913db84390b42eb433ab42e0b95f771d2fe0b", + "fp32-reference/gradients/0223.pt": "2ffdf58e282f2776a4e13b1fea9aa737dc7dc7ea0486bc14f0489a5f6c34fa81", + "fp32-reference/gradients/0234.pt": "e16320be47fef92d457a1282ee5bacbd32d8406e99840016852c942dcb9191e7", + "fp32-reference/gradients/0233.pt": "ab8ffb2e37b4fd738831b1e0d14532175235cc9563ca59e6127fbf8816e4a011", + "fp32-reference/gradients/0220.pt": "b79aaeaef0a4df18ab2e8363ae21d175185f78b088b79232a8b1ee4a11b6f566", + "fp32-reference/gradients/0219.pt": "f7f8ed60e67a53edbaf38770c80f4dac3db2683381c7df9d2ce346283645cf67", + "fp32-reference/gradients/0218.pt": "d1919d831be54c10c1fff2a78f604774d280f7bc6e33d42a4e38a3a18ffd975d", + "fp32-reference/gradients/0217.pt": "a62dc9c713e64abf9433f7d2baa7bd0209649f44009d6a3a3d114808bbe1beef", + "fp32-reference/gradients/0216.pt": "1b0d2e81da8cc84a5297d622b0ce1760f83a125319b32ca04461d092608946e6", + "fp32-reference/gradients/0215.pt": "4181323351f413d3b8d96352c742893d072f220d5dec5c6cc2894fa049dc6e76", + "fp32-reference/gradients/0214.pt": "52aec816316351f1cf09e3d2f8a1fac7ddf60c2be166f442a344958ec465e915", + "fp32-reference/gradients/0213.pt": "cc75dcf14ed94d45a8f993f1e36e73124ec946134ff88e28d02ee11ee08087c9", + "fp32-reference/gradients/0212.pt": "7a0ab54bfd179305ef69f6555cdd724c2c0e017be27385fa6966088717c32757", + "fp32-reference/gradients/0211.pt": "1a4525da9aa3cc3ebbb07f10edb4c03a22eecb299868e253a4df142c64f49a87", + "fp32-reference/gradients/0222.pt": "7034b0dccf92c1fea7caba281b139c41527a853db8a60e560616b149a92a686e", + "fp32-reference/gradients/0221.pt": "e422fc8c1978ab439c4920978ec8148784135544af2d9abda46c2e317e8a4da9", + "fp32-reference/gradients/0208.pt": "17afd90ee3fd66673308cbca0d91138d6f63637f73bceb1e0e734b74bfb61961", + "fp32-reference/gradients/0207.pt": "73284b75bfa35fb6f4453632ad0021175c0ffb2a8cfb0192f9598d67d4b7abcc", + "fp32-reference/gradients/0206.pt": "be3853417df32c614a41562ba6ce11e2a3d12ba341b0113acae8fb27a7dd25e3", + "fp32-reference/gradients/0205.pt": "392c3545647a01f50144be8c0908bb5974cc995d8f416521c269fce3b24fe426", + "fp32-reference/gradients/0204.pt": "5c76c7f71014a0abeaf47753763483186bb4ee1d9b83db15a886fe9b56a4d128", + "fp32-reference/gradients/0203.pt": "ed5588d737e6ca4f5d5c71708516fe43ab90afd49f966308201cef6db624a4d6", + "fp32-reference/gradients/0202.pt": "0b1b89065e42478eb33ce7c71f580f528dda416b65c997b75202e4fe27e913b9", + "fp32-reference/gradients/0201.pt": "e95e304dd8bbf323ba3a644314d8901dfa3f5cf6f65ac14ee1b6446463029a85", + "fp32-reference/gradients/0200.pt": "21fe5ea00c1eb8648dcefc6a352d04c56d47b3d3975ad4f8339dd1091b570c54", + "fp32-reference/gradients/0199.pt": "1f22dac9532faa8d265e9678f8e63f759b4162ebe3fd236aa72a01e78889535a", + "fp32-reference/gradients/0210.pt": "3962095170b0967eeee1c86e6fc27256225160abf987133fcba67ee5e0bb88ba", + "fp32-reference/gradients/0209.pt": "77231e7a1d4b5b3a9c99ed04d022864ccd5a321e290515c3111764ec9408fdde", + "fp32-reference/gradients/0196.pt": "9973b017c9345214248b04363544ef1d3ece268ffd383d505ff510a6c8dc824b", + "fp32-reference/gradients/0195.pt": "9a960ae7c115ffc2bdc2458017ed9c33d2e1746f0a615d725b08f5f72e920388", + "fp32-reference/gradients/0194.pt": "e14d1ba9ab68ede0a055ac665761f62bfaa5acfb12fba96fec91fcfaf165c3b2", + "fp32-reference/gradients/0193.pt": "c3394f2d71d8bf6572ae31239fdfc7a7d925a93894911efbd778fbdee4b5889f", + "fp32-reference/gradients/0192.pt": "0191597afc0d3ad11d02adc401028a9967b5d3cf1984f2137161b8c1219b4b90", + "fp32-reference/gradients/0191.pt": "50a83e66da7b52aa46618325411680efdf7a008df939e82ece17aa7d4eff36a0", + "fp32-reference/gradients/0190.pt": "05e90464879175abba5316f3e852091d11748e6c1bedc53c102ab7e4b9721b16", + "fp32-reference/gradients/0189.pt": "7ac70df33e3b5a715eeebed46db148f3c12060ba3453dc97426fba2b8afd3fc0", + "fp32-reference/gradients/0188.pt": "5158f69cd770cf5db251a96b2b1634fd36dc781ed1249e59880caad14f1a107f", + "fp32-reference/gradients/0187.pt": "2bf8c64a232aa5c0259287f76214eac70cba6282dfb06db4b361cd02e1064adc", + "fp32-reference/gradients/0198.pt": "78ef8746e59eab1471ef933617d8aec8c4dea3803a132e047b1d2fc44509753b", + "fp32-reference/gradients/0197.pt": "408b3bdbb4f384d2459866122b893d3810f0b07191b790a6d1cef3041b76f372", + "fp32-reference/gradients/0184.pt": "b3e02e5674aa74e6371679edc8485ee573ee3d861b205ab8c6db88858ba23c68", + "fp32-reference/gradients/0183.pt": "dde7a8bc4595275db1c6fb26bfb62a802c9dddde2807a3bf0981d08de2febf2b", + "fp32-reference/gradients/0182.pt": "be56c87888d39287928f45af5d828976a34ec73a4a9437eef6d4a881b61b51cd", + "fp32-reference/gradients/0181.pt": "7777e4eb88e1e2312c8dbe1fdd7210d2fc9086b687300526ef0fe1f5418c5681", + "fp32-reference/gradients/0180.pt": "f4532dbd854c7b95ad09dc80841dcf11410b6e495f6710f5126cf04bb7a91434", + "fp32-reference/gradients/0179.pt": "54be67e63f3159371053dc664a051b60fb52bc42ed8a883087cd3b60cace98e3", + "fp32-reference/gradients/0178.pt": "09360ac4602a57af14b53ed431f50f7dac1d41bd65c2d396a4c991fdd2c8b6c0", + "fp32-reference/gradients/0177.pt": "330cd7a890d573598367fa63e37f56ac86c1b416eb26cbf7125137ef6b4518e0", + "fp32-reference/gradients/0176.pt": "9b71236e3685061d9a4bf032ba06ec5faef50b82ece1b7255208a37578c5e5b1", + "fp32-reference/gradients/0175.pt": "520dfc2f5c73201a764a36b927e88cc259708d7c791a84974b94f26875f9bfc4", + "fp32-reference/gradients/0186.pt": "4437758b852edc34f8cfad3abfd5a9a0fcca6092a3d6cee9bec375c189d969c0", + "fp32-reference/gradients/0185.pt": "020c5c3262063b7de91284e2c0cea160ac1ea2aa2f8cd49ddec81fe720d5713f", + "fp32-reference/gradients/0172.pt": "882862c9b0142712b1c4df5601ea8c6f169cc68bc9c4500901d84715ceb3fd19", + "fp32-reference/gradients/0171.pt": "096ceea720f53c78afff13b465a9640bfac01d1c6a3b5b361ccc07b033c299d0", + "fp32-reference/gradients/0170.pt": "d2e0784fc06524ffecdb4a3482e986fd07621ccebe1d6cbf56a25c98351d9255", + "fp32-reference/gradients/0169.pt": "f3dddfc6ae5b77c9a1d39cb902cc856f3ce8ab85f512061046c0e8f146f69806", + "fp32-reference/gradients/0168.pt": "7e6fbf962e9166dbd5c271eafdddc8071a40a446cac930b77c676afb4731fbe3", + "fp32-reference/gradients/0167.pt": "ee003a6246d66822acca80d7e4a6135cac6e6e61215738a978ac6fa23a48474b", + "fp32-reference/gradients/0166.pt": "dad720f418df89b6a9076ccfda4c4454f0b921a48589043a4932318cfab7ae6c", + "fp32-reference/gradients/0165.pt": "d3c36343db437f2ec2d69fc6182256271d153a68f98a25ce01a7d09dd82555bd", + "fp32-reference/gradients/0164.pt": "c4a8a78c37e9f71bc4bf273d8af0b428ad3cc400142dec2cefa8c1f064a4ebc1", + "fp32-reference/gradients/0163.pt": "4a99b33fb4182893bea399da2f290b312b3cb252683292aae11d2e96ebddb581", + "fp32-reference/gradients/0174.pt": "cc2ff6280cae269b3d0f817c88f8a0a715b4a50d5fb17d764c79b3382325a30d", + "fp32-reference/gradients/0173.pt": "4591792c167fccfa8630c57547f1c0f01ba2182e4edbe93f488a6018dc537289", + "fp32-reference/gradients/0160.pt": "7f1f4106d8097ee5c5531c8e38e2b7bd16c81d81aebf60a9ff056889fc11bf31", + "fp32-reference/gradients/0159.pt": "befc5346742535b4b009d4800bc69403bbd9e14b21a640b7f58c9267e23f5c14", + "fp32-reference/gradients/0158.pt": "19df487803b9e7534d28c91fcd2aae4abd47eacc59c3bf7398645d273cba5c53", + "fp32-reference/gradients/0157.pt": "c5d25f9369eb7b457e4b708dd8703adacf76a15b206004d1b34dd7345421f1a5", + "fp32-reference/gradients/0156.pt": "3166f23d79b4ed7aadd48d6532ca02e8090521d14e35159f11f60ccca493db3e", + "fp32-reference/gradients/0155.pt": "077f1bd6d24a041db7f4ae9032d897821608eff5a61db7a53f468bf2eac1eedc", + "fp32-reference/gradients/0154.pt": "4c4575ac1020a8efb8ea4925507ac2876423c053ab8823271fac7bd4eb2a10d7", + "fp32-reference/gradients/0153.pt": "0c9efab1fe6e2eb94d425d1c575803bc7225dac499b9d1210b8e7abae56913bf", + "fp32-reference/gradients/0152.pt": "f8be0b6b44bfe1b67bebeeca2bbc8f1ab8145b857136115982572b680d77eed0", + "fp32-reference/gradients/0151.pt": "1c8be24f986eefed4dbe1959588edf1422ad1756ffb2283125c5cb6548b0f766", + "fp32-reference/gradients/0162.pt": "d4ac31763f94cefac0e7ffdf139a9b1b0a1f834ba8041a2b2111ae9ae8c15afd", + "fp32-reference/gradients/0161.pt": "d2ff019717053fc34e5f0206174a770866f699ba21962a03cbb924a7ba1cbb6d", + "fp32-reference/gradients/0148.pt": "f597968fbe2d77bee43d5559d1431d5b7da81b18874f932559393406eaaf7177", + "fp32-reference/gradients/0147.pt": "fe3fc60bc9005807ff5c21cdc3ceea0973edc448ac5b44dce4604e9d18b4043c", + "fp32-reference/gradients/0146.pt": "20e4bfe11ba93df308bc986e6c4a45c36b7ae8813fd5b2e1006e7c4820af1dfe", + "fp32-reference/gradients/0145.pt": "4ab42890c9082364010b4df131170f989a2538c486d89907de1104bba2fe4ef6", + "fp32-reference/gradients/0144.pt": "cd60b2a8446b40d6a7a6f0a889ca5c421bbd73bf539cf73fc3ce2041699b63f9", + "fp32-reference/gradients/0143.pt": "4f495c888e6cf02d397f37339c50aa0ab1d098d57f2068aa100c5ecd3541e30c", + "fp32-reference/gradients/0142.pt": "577149715278de3eb1e73c6875e6840ef7d8329baf33ffc424c6006a703221f5", + "fp32-reference/gradients/0141.pt": "c90e0e8ede5f868e097ac2fb04e8215cb4b34659c3114205f2a96db9ee400e34", + "fp32-reference/gradients/0140.pt": "505f0ce6e6462c38703163cf6733dbf28990f0985ebf6e9c1dfac35aeefc069c", + "fp32-reference/gradients/0139.pt": "dcb2c21276b9bfaa597605352caa86b66019441b19785058d0d8b9f3ae6f49a6", + "fp32-reference/gradients/0150.pt": "91ecd6deb6402a6b9654562418a8dd8f231d19735576982188ce886a0173bfc2", + "fp32-reference/gradients/0149.pt": "f4b515d08ef82923144e9ea06e82ddfbae76d7e017277746f7307912af6f90f3", + "fp32-reference/gradients/0136.pt": "4db5a0dc8f6127148d007e28d163968f820508625e2d8cba0c30adb245b367df", + "fp32-reference/gradients/0135.pt": "ab09687d9b1df75194c431a43f041784e1da0dc3e0279b9001f208a0269e90a0", + "fp32-reference/gradients/0134.pt": "e82046c52150ba940b6fb1a17bef72518e96923a5e358920ad0d3d5b99bccc19", + "fp32-reference/gradients/0133.pt": "b96f4bfcf59e9660377e7f4a09c3d716c6a310bac91c3108e439c7a716b374a4", + "fp32-reference/gradients/0132.pt": "9b8a47fa6fda0b78fc3230473f5e3397c66d286508faf91121e33aa1b8eb0a0d", + "fp32-reference/gradients/0131.pt": "197e9a7ab040df42415317b624ae6a8179d2cb897a048e8ec8b1585253d987d1", + "fp32-reference/gradients/0130.pt": "a82998d9cb3f69b29f02e973d4f6b05d7fd5b71a8b3274c6e01d22d7d083283d", + "fp32-reference/gradients/0129.pt": "c603922d13f5d498e416addfd24cdc692f0c88826463857985e91f51b6242d61", + "fp32-reference/gradients/0128.pt": "eac06846dde403b093c414e46211ae5420c95e6676028cc0c19e0fb820221a2f", + "fp32-reference/gradients/0127.pt": "584dd5c4dafe5526ada168b79564d2d5a2d6ac9885475235c79b7c6db5f7dc17", + "fp32-reference/gradients/0138.pt": "31b9f1e938dd2b239e872c6d83599349f210bb52c9236e84765c04ea4ceda70a", + "fp32-reference/gradients/0137.pt": "6af6a3eaaa857bd1b34a075c77bdecd0bc232448bb762c49c95202a15b771219", + "fp32-reference/gradients/0124.pt": "b5846f779c74725d02a6cedcd676eb43734d4d20bcde782f3a3012fd25daf113", + "fp32-reference/gradients/0123.pt": "1e130cfd1e73fdf877088f9c998f8f5754ecdda5ecdf6ed42ba96faee3bf2dac", + "fp32-reference/gradients/0122.pt": "78ca0aa4d199db71f297f54c2b52a0aac683760322cfdc61476857a5661e85fd", + "fp32-reference/gradients/0121.pt": "20231c6a4f9db313eb50f34a38583e245fd0c32f52257bbce19b3483fde2569b", + "fp32-reference/gradients/0120.pt": "cf26d994ef0dd0b6077a9f79699287d032d04e196d7ac942433d78af6ee1b7b6", + "fp32-reference/gradients/0119.pt": "3ff84f127c5bf5e89ff6a817d4b3be0b4f457ca915ac666f38d1f9aeba2ecb76", + "fp32-reference/gradients/0118.pt": "d744872d12bbba24546dd2ba89c4109e921367a974937d6c8628e4adfc58065d", + "fp32-reference/gradients/0117.pt": "e2f4c477bd9e8d41f05b433baa278277b1888797e4d40ab7eb6cce8814649e29", + "fp32-reference/gradients/0116.pt": "4e0a4f1866f621f004de2b20a08eb631f572833b5232bc0acfad693cda1a67f4", + "fp32-reference/gradients/0115.pt": "44fe298e7d2d8f0a9ac3013172486e1c7db7862e5510c4e278f4e7c8f0c3c42a", + "fp32-reference/gradients/0126.pt": "bb5b91b0fc8f5505dd13c59279fb36545020ac46543afba7d1029eabff7d5a18", + "fp32-reference/gradients/0125.pt": "4e605ef9b9b5a6dd2cf3c0d097f07264e8c86fb5315173386a52907d1444a283", + "fp32-reference/gradients/0112.pt": "1d81d90143fdee2c5bbceba8b10e18f5460c91c79892f70f7bd452ebe853b4d0", + "fp32-reference/gradients/0111.pt": "1f4cf67047321d713c447c5182377ac7e70cd792e2185b2ba8f498affa30caa1", + "fp32-reference/gradients/0110.pt": "af8b9e1ecbc48d297314e2886a80be8196c3ad50b62780891f47ba833c9f0d2f", + "fp32-reference/gradients/0109.pt": "f28d836783666399312c65d077a44478f890459bdb8b7f62bba025b05a0f791c", + "fp32-reference/gradients/0108.pt": "ad1506ed4e9d06d43381af842d2e603771e39d8ea35cbffcc169f7fff6584226", + "fp32-reference/gradients/0107.pt": "ee3522f5f906c621463224feb56f7f75021586c4c78d02aee28defc866954c19", + "fp32-reference/gradients/0106.pt": "8ed67af3aba1a42c0173f5fee9da3d65fd63683568ec679bf7ff0c556b7d6da1", + "fp32-reference/gradients/0105.pt": "fe1228cda80642b2d3e541909b26e3e53e99bb3ad5254aa9b80d79c8f32af793", + "fp32-reference/gradients/0104.pt": "c60e6980543b98ebd0dd271687cbd4bebd75f5dfb2244267b4c5c7ed7bbe0aea", + "fp32-reference/gradients/0103.pt": "11dac88c8b802a30dcf3e77a08400e7fbe3eecf3564538966df03c6ede4f9c34", + "fp32-reference/gradients/0114.pt": "493169ce6544ee958d39d54f7ae2c8b68aea1b998987794c6a2fd9d6d5231535", + "fp32-reference/gradients/0113.pt": "98ad3557a69c4218a09241a3476ff8b5049d6056ee29ad0fb574d504facaf968", + "fp32-reference/gradients/0100.pt": "70f66af21bb9f5d5f868eb3c064c7b810377123da2d57c253e70f9dd2505c44d", + "fp32-reference/gradients/0099.pt": "7a4e1dd9733c55b88d5275e40325b8ffefbf9c2c70ecd19d0d8d9b0ecd3f02c4", + "fp32-reference/gradients/0098.pt": "4dbc44dc7e03ad611ca70fb20179749afc75dc296809b9ece2a762729411032b", + "fp32-reference/gradients/0097.pt": "b769dd57b82e763c726e908f275998bb3854901ae6278695ef797668c2ebc8b4", + "fp32-reference/gradients/0096.pt": "a7f51e770c17c1a83fb39ca6561268efe315421e7aa3a1da7be8ca1740857487", + "fp32-reference/gradients/0095.pt": "edb1849e30227917f2c234c0e808499493d25866b14cb523379f62285458c349", + "fp32-reference/gradients/0094.pt": "eee1eb0387492d2eeef00eb8224c9a524e55b3fd555264e3d1034291988194d2", + "fp32-reference/gradients/0093.pt": "f21d6b20698c4daf711e87917064e7ce2291d13cb775f90b4bdf0a852d4a5fdb", + "fp32-reference/gradients/0092.pt": "82000f2b130e0585f1609129b09d8ecb5cef98a3412ce47b7107c2b48bb6cfd5", + "fp32-reference/gradients/0091.pt": "75b0f1d8c206663516eb65d47f94b1fbc5d20c91b4853ea3544e02449077d70a", + "fp32-reference/gradients/0102.pt": "b625032731794390d28521187feac7c92745d2ed56789d821540b73b64f854e2", + "fp32-reference/gradients/0101.pt": "fbaa0c0ff79a307daa72da00b9f3d0d16f60e17a872621d2c0931f4813085816", + "fp32-reference/gradients/0088.pt": "6630823b8faa6cda492942db91e73524e2dc9d5760175e9858b7bca512b77690", + "fp32-reference/gradients/0087.pt": "7001edfd5dc9854b2c6c55031590d9cacb1db146b3d148b6b62f334ab2c5b81c", + "fp32-reference/gradients/0086.pt": "a5c226ce3763a07fb29e3b2656bbb69f6ee4505b232830b0ccfc14772e2a6be2", + "fp32-reference/gradients/0085.pt": "2b66738ad8814a1ba47dcd9776029b3c27b19cfdec8d018720622e764e223708", + "fp32-reference/gradients/0084.pt": "83bb85db5b6579a0bb7a7b4af6f276b8130991a32805051755c91a2e69f8e308", + "fp32-reference/gradients/0083.pt": "9722fb5d0e3560b5cd1dee15e382ae36c85077219aefe1d6d9cb86faad7a6079", + "fp32-reference/gradients/0082.pt": "40d597f6fe90c7f9c0482008f02b3d3674ff62e24fc1a09d825d8737cdaabbea", + "fp32-reference/gradients/0081.pt": "5152cc9686464f57f751f8f50c9ec3a28dab0869ba3450e33c3ab3188319587e", + "fp32-reference/gradients/0080.pt": "1c410fb3553d74495ae5621ad82060807a463ce82732701c5e7919345a4a1cc5", + "fp32-reference/gradients/0079.pt": "2a5b84b3868ae17c3fa593118ebb50f3010aa0c80366e76e8dc638f503eecc50", + "fp32-reference/gradients/0090.pt": "dfa763e291bd9fc454ab24d0afb6bb554d485cd309eb691459484e3ae5d0d71f", + "fp32-reference/gradients/0089.pt": "cdddaa734b91a56d3982037ddc06b31b3de1b127da0c18d24d418c4a22e5aba5", + "fp32-reference/gradients/0076.pt": "728eaaf85f2fe00fca2709fc058ec4251489905d0b80e0c125d7b5f9aa91cfe5", + "fp32-reference/gradients/0075.pt": "20013f7b225182f41902427568763ed16c57a37ee8c6f3fccbcb6dc2db0a0240", + "fp32-reference/gradients/0074.pt": "4bd48f9b26a284f88094ee2d1b5e25d7a8551dd0adcd2a1022e61e7033cc1e2e", + "fp32-reference/gradients/0073.pt": "2d8c58d640fff1c73e8ff384be388e516ff7c141707ca5fabb63058783d51eaa", + "fp32-reference/gradients/0072.pt": "8f465620b1bc98917d2afa8afa9df1d7b598b620b9b43ab3bab980b672ba59d5", + "fp32-reference/gradients/0071.pt": "7c3b0806b92db3c8b232d47bf69b780c79d818d7a18a34d59547cc933ed52369", + "fp32-reference/gradients/0070.pt": "61ef4c0077c8eeb8425fc99afc26fbb80bd305e508528770434aec19477a4424", + "fp32-reference/gradients/0069.pt": "f3cccc0dd4318165e0a57c5ccc6e350396f175e88fb225535041cc57b6766814", + "fp32-reference/gradients/0068.pt": "6d2ad0be502f304db305602b70e10ff016717a25cff20bee9c21cf7d6870d11a", + "fp32-reference/gradients/0067.pt": "97522b486c15336bbcad70a081771d91bd6b0c496b9e7b68580b180338f3c32d", + "fp32-reference/gradients/0078.pt": "1aba7ec1b18b858cab2a82553b1b5b22ad8a1ec8d5ac619c46fc2e652330e4b8", + "fp32-reference/gradients/0077.pt": "7a8a4c1eb4380d436f070aeb4b062c56e3e389ec20b8e0ec75c777252140e729", + "fp32-reference/gradients/0064.pt": "9d487ed2abd411d54da33a569dd4ab64f37f126f76c5327b72794c0d0b77ad2c", + "fp32-reference/gradients/0063.pt": "53f2f6cd0c8985d68e2ba59c4374ee61ee308047cc41e6c70d9e515eb6a5fc9b", + "fp32-reference/gradients/0062.pt": "06c1a9ad1b28e22cee769d513b771784c6ace0130ac97497c1c2a8070b36ba67", + "fp32-reference/gradients/0061.pt": "7071b184165b47a92d5f415a47db3ee8705b9187af41bc5b57a49a188961f1ce", + "fp32-reference/gradients/0060.pt": "4ef21ba0a0a59600d8a33c4b8bfad87463d0988343a30add0827b7715945e7c5", + "fp32-reference/gradients/0059.pt": "fe2eca1fe0316fd0917bd5abf263e19d1324513f53b11ee94f1372c089fecf94", + "fp32-reference/gradients/0058.pt": "c4332f8a32817da839859cfdba35217a3d48ee77ef99bf776a807f6d1c92479e", + "fp32-reference/gradients/0057.pt": "d788d2c6d6d2136cde78c9b2a87328f7f1f473d6b8170414c579c40cc7a44bdc", + "fp32-reference/gradients/0056.pt": "b18fe46471fa08b76922c5c8469b45e32f4cce5b667fb140b5c9ee13850e2a62", + "fp32-reference/gradients/0055.pt": "7f5f5d80d2d0413c175d5d7f1fbb4e979a4a5e75d66eddaf269e7a5a15406c18", + "fp32-reference/gradients/0066.pt": "674e9a06e7d0a44fe1c680cbadb8b694714cf8988c68e625ce82a2cf1078e3e5", + "fp32-reference/gradients/0065.pt": "2170eb1d748979d92954df19908340928198ad5637faadc399ffe4356b0789c3", + "fp32-reference/gradients/0052.pt": "65eb7545e27ff9452acde808b36b588a2f9ebdda83c9cb34fa29ef2c55a72438", + "fp32-reference/gradients/0051.pt": "1b1345f6ac31172a92562a5f0ba8b82d94ceb7f476bfc8c679e203799b1a660a", + "fp32-reference/gradients/0050.pt": "55e84dd9a515859978ae9a5a96ce178baf057da1cc14cc49864dc56897539727", + "fp32-reference/gradients/0049.pt": "10a7c34b3b1997aee50246b7fb5a51b87866153f57e54d3faa75cac826cca248", + "fp32-reference/gradients/0048.pt": "d4c84c0b1e38e077e69dfbb2eaef2dbfffb022d38bc6284514aeecfd36444b83", + "fp32-reference/gradients/0047.pt": "ec1a717e2a51887118ab8f4c62af44ce5176a16cffff35e25af6a75b955b8845", + "fp32-reference/gradients/0046.pt": "d6670b2182ffd8f77f953328dcfea550cf00ee03122ab2e804b629b516302814", + "fp32-reference/gradients/0045.pt": "a42872fcc9b0203c79329e0f731c62e2cffaddfffc3dc105ee136157d0f777d1", + "fp32-reference/gradients/0044.pt": "e68139314f366b7c18709f828e37f7c063c5d44d198808f59c1dcd042c711f00", + "fp32-reference/gradients/0043.pt": "8ac7a0cab906637797d6919c6a2e26a536babee044e564deeddb16e18cee45cc", + "fp32-reference/gradients/0054.pt": "4dc8058a898db51a5310c51838d3500feb29e00a833ce8c244b714c92691bdab", + "fp32-reference/gradients/0053.pt": "5148c5396672518eec5c4fc641fa5301231971fa70d3c15db300ac230bfa7158", + "fp32-reference/gradients/0040.pt": "50f4495cde40013a908e8110c47fc15e6de4c11254649b154573187b3131d8da", + "fp32-reference/gradients/0039.pt": "8a0b2d7962117ea6f70ca44abd2d23748c71bad596388e397a780780c318b0e6", + "fp32-reference/gradients/0038.pt": "e816f0674c99299f65154c3a43fc21ff40b0fb05b7a3a5f49462526b2fed2fab", + "fp32-reference/gradients/0037.pt": "40605f89ea7f0dd0fdae830b978d06dd216240beb9d93bdb0306fe9f7e3271b3", + "fp32-reference/gradients/0036.pt": "99ccc3a8e9739279978cf612bf1703f9ea64874fd0a83f9a39fb8a289d43361c", + "fp32-reference/gradients/0035.pt": "733a0981c95af7b1e5a4f14f05bb576f905796e9d09ba065193fcafc4b91c7aa", + "fp32-reference/gradients/0034.pt": "2ad4201fe5b25d1ee3b13a32c794f94a85e73962d614ec4f9b91a6e2cf5cdfb5", + "fp32-reference/gradients/0033.pt": "3906ebb90fd6a62739df341571e2c64375bdfeedb47536513c4e52433c9c6aa6", + "fp32-reference/gradients/0032.pt": "06fc4c4d064a07037b709e41059bac406daa7cf468f371f1ce5c9b635482e400", + "fp32-reference/gradients/0031.pt": "c0aae5701c42c167c4b65ac48f0f55086bd3afe390a2fcce405655ae2be7ed39", + "fp32-reference/gradients/0042.pt": "3984142de083c8d2c12c2ff72e672a3ef2ab16ac182c80b2e4765c444f29c1fb", + "fp32-reference/gradients/0041.pt": "f0e5722180894e058da84b2bec560b937e35e59378ca095d53a4d6745935649d", + "fp32-reference/gradients/0009.pt": "97d4a4c63a4d19a156195d96cd0489d2c94b8fc05b9ba5b07f5cecd083263924", + "fp32-reference/gradients/0008.pt": "ff74c7140174f036a208d51310f75fcd7c7aa8d10b03d0c172d58eb245b94157", + "fp32-reference/gradients/0007.pt": "a9ed52ea232dc4559e4fbcf024bad48a5aaafb79ee6e399d68b5844931c3b24e", + "fp32-reference/gradients/0006.pt": "ea30e4874686fd68cdeb741a2f0fc7dab4704a250b3b46b808da46fdaf1ca392", + "fp32-reference/gradients/0030.pt": "2559c36a9f9532b818bfa734baefdaa5bd6ab465a7456f3306153377715011b5", + "fp32-reference/gradients/0029.pt": "77d6954cfd6fb5b4dce2fdb75f75bc24b6e20bbe61f61eba6d381929450438e1", + "fp32-reference/gradients/0028.pt": "3a4bc8d1864ce3048424bb7760800882aa41aff44acfa94b405d1d053fac9ba5", + "fp32-reference/gradients/0027.pt": "0831a649430c7fccdfe0d09f9c2836f151ed1ad9a2d812440bea3dbc4877e4a4", + "fp32-reference/gradients/0026.pt": "539679efb96862b5792975b15de270da58ac14a45d6dac71c824dda9fa1d1723", + "fp32-reference/gradients/0025.pt": "e69675232c8ac02fb7874276df98c44fe1b570ff3dd217ad92c39c4a0cde15de", + "fp32-reference/gradients/0024.pt": "27a8fc9823a9d73b3f2fda8a39fa87af6d7368b9ca4f8a036a23967257f3b035", + "fp32-reference/gradients/0023.pt": "cdbee4034ed3595a82c19c2e128019a9b54624e1b904e9f8b98cecd8f1f422eb", + "fp32-reference/gradients/0022.pt": "ba98182fec10ffb5faf61fe825fe24348d47524a6afb9e6b73861164dda17e44", + "fp32-reference/gradients/0021.pt": "3ac846501b778b7ac829a350173d68e24c4655901027d9d99858c00126e5543d", + "fp32-reference/gradients/0020.pt": "6373d0f0a778e4c4e76ae7910ac98ec5cad79716d870a593f8956fca3d22f690", + "fp32-reference/gradients/0019.pt": "8d5a5f8d5d4f6721c6a4447ea9d91b57c5a16e57d0f8c333a9cfe359202cdb68", + "fp32-reference/gradients/0018.pt": "260dba92837138bfd2c58a4dd4c2e23f224a32c1b607146d88d0c9770dedbc03", + "fp32-reference/gradients/0017.pt": "f4e2784dc14cc048acc47cd0e1448c87af3a2b94656e555b68124ff6edc59c60", + "fp32-reference/gradients/0016.pt": "ee3e74c0568d935973f196a7db97c83437d9c923713f84665ad2187017dbba8f", + "fp32-reference/gradients/0015.pt": "11b82267677ec4b6ec327c5c34b578fec10e55e73d50b0d16e51bf7f9ad43422", + "fp32-reference/gradients/0014.pt": "a06fe168ab2e23c279b8aa1d3facde62a7f764423d3195e0366f71884f33b4df", + "fp32-reference/gradients/0013.pt": "bb5e1e290216b6ff0747caf7e5f9dc44f7ab386d71d7a5cbd439d13a5358c11a", + "fp32-reference/gradients/0012.pt": "630b5c9bf756301a9a4a14ee045202b463fdbc9bc8c297d90a61dabdbae536fd", + "fp32-reference/gradients/0011.pt": "60fcc18f64a16e2f0b29fefd7d2c248abe2ef3587a0b3ff6b82a8323f856d66c", + "fp32-reference/gradients/0010.pt": "319cc63c3a9858022c03b06f513070c6ace8615a5d144a58bd0be3c9cc386a9a", + "fp32-reference/gradients/0005.pt": "d454aff83dcd6e14cbf9685128797a913c73539ddb470fd286c47e7ad5e14c3f", + "fp32-reference/gradients/0004.pt": "2d20102884ed602eb8ec075554844fc5840bf9464a4bc353be5f993ad741074f", + "fp32-reference/gradients/0003.pt": "53c91edfe9637cc99b14ff5e1503e0e6a21996c73fdb94f3149099a52a384137", + "fp32-reference/gradients/0002.pt": "0419ee6d6ad48307830cc52cdeb3421584ce0e43175cdf5ef832b13f10ac02ab", + "fp32-reference/gradients/0001.pt": "a13fa20784d73f4073e29ccc80a63d6cd8d1d3e99127c939d1f3e0532d74aadc", + "fp32-reference/gradients/0000.pt": "82f4f19f10c9540f0f9e7daf898b793723970846f0e2af0f19e5cb6137ddc335", + "bf16-provider/report.json": "d49a80d6e438eb1f7d91193ad3f3ef5b5b3692b2c86a7532b86f23359ad33d79", + "bf16-provider/gradients.json": "4aa8c3e141a8649ba4384ef1d7852cef2d1ea1aabcc3a06f64acd17efed61e9f", + "bf16-provider/tensors.pt": "0946fbd60bef1996693d2f97d69f13e01165c9094fca9fe988cfc95ec64773db", + "bf16-provider/gradients/0637.pt": "e3844a24c49f10a55cb8cd4fcd97f71d814bccebad3edd81988a49502bdf3b86", + "bf16-provider/gradients/0636.pt": "4f0171f5de3f271bdbab29b9250b2e7b9415209be2c43fde44314b75be9082ff", + "bf16-provider/gradients/0635.pt": "0a93c6cebf6871f3a8872fca367c5b26a23800fc924d9ed5543df2f2a335f0bf", + "bf16-provider/gradients/0634.pt": "bd513cac9f098d35c3b83e59d5dba5400e32865958783fd2e21e151ecad331c1", + "bf16-provider/gradients/0631.pt": "51598eb239d21a6e2cd3079821ea65660515f1468b8df864c3a5e566b93c6214", + "bf16-provider/gradients/0633.pt": "33806dce327b8335d65d3b898fa90056564828a353b823b625f11593e73d5514", + "bf16-provider/gradients/0632.pt": "0e60adf736145369074cf76f434d7cf9eb221d6f49a67da2f4978e24485a92bc", + "bf16-provider/gradients/0628.pt": "4ef0834886c1f5916a6f6305766e1024b1203483797a21460c28b06bb5beba6a", + "bf16-provider/gradients/0627.pt": "913665b9659bdb29a70bc3d02eb2d445c85a5009c83169a4147330676b21e264", + "bf16-provider/gradients/0626.pt": "28c56ce68138d83acaa74e7b2a26e4f7bbcf5f76551df90ef6c01ccd1bb6f113", + "bf16-provider/gradients/0625.pt": "ff8c9ec07cec56bb2d12787fffbea708f8b734291f62730e0019dd0baf15b1fa", + "bf16-provider/gradients/0624.pt": "3fc9857d0ddac770e071e5d998908b0988d5ef7405eec8138f8f45c462082ef2", + "bf16-provider/gradients/0623.pt": "a0c593df280bde9d01eeef67decdd64a51828b36d51046c114cf9bd5bb590a1d", + "bf16-provider/gradients/0622.pt": "5cb69d6bc4671abe1a18e59031c26a746e6dd585c6bbe3cc2be15aa813cd54d8", + "bf16-provider/gradients/0621.pt": "851a411f41d4b9e9a278c97a2a92043aab6366a7a4bc4a9cbcb364979dcb18fb", + "bf16-provider/gradients/0620.pt": "9db3e1ffb4b9ae818bdf8bfe0f06ecf09cca5dcaea67e7d901d5db79e53c0b10", + "bf16-provider/gradients/0619.pt": "d1babab5f9512b93122dd571e60da19cb76da15dfadf504509d67d03a94f68d8", + "bf16-provider/gradients/0630.pt": "ca7f24a16c16afeda5d876e29ef3b7e3293ec77645bb33cc7b1b8fda6925cc64", + "bf16-provider/gradients/0629.pt": "38ec65bd40e3f255dbbf49da0cadc63becbe53a5d5c6c89f2770fca50bfbbc8f", + "bf16-provider/gradients/0616.pt": "fa914d3a6b6cbecc3a38cccbcddbff1fa34aca0cbc2775cb91055e43a76d1b65", + "bf16-provider/gradients/0615.pt": "2b3075e68c3d6f92b541d642810ad7b11f505bb5d33b67637a44968688093ac0", + "bf16-provider/gradients/0614.pt": "ea0d5c550b383dab6225d6a7790239a77c0a90b6e5fbe4a5f3df8f680bccdf9a", + "bf16-provider/gradients/0613.pt": "62d8bca5994bd2ca25a853951ff7449554c7c1389644deef17da1b72ae5d17c1", + "bf16-provider/gradients/0612.pt": "7a7229e523859a3d0582b2ca62d54dcfaea25d6c4757be3bd8895515b724261c", + "bf16-provider/gradients/0611.pt": "bdf44f033071dd300e9e55b20e5db52b7838b67e2db3561532b642b3c3765f8a", + "bf16-provider/gradients/0610.pt": "d3fdd2b2236b3bc9f013f9efa2fa36226b6aefe2c9f2c6701ed99070b7d7605a", + "bf16-provider/gradients/0609.pt": "e8d0bfbd34d5ebe63d08c85c2eb61f040425e8885dea6dbe89282861358fac32", + "bf16-provider/gradients/0608.pt": "2b2fc55e2d8f80a9a81354b259bc49cbaaa109b1e509eb75521d7ac68fe00760", + "bf16-provider/gradients/0607.pt": "7fb3904db359e9aade77abd89a6adceccc3bbe3c781bc1c21b99b0efe5ae6245", + "bf16-provider/gradients/0618.pt": "018aa39a33eb4c1b40782f074212f339aae5b9fa0b3e892838d90abdef0a1bc5", + "bf16-provider/gradients/0617.pt": "016dc3451f14408615f3ea626beb19d7572df1c340c76fd12e839a9a5f69bb6e", + "bf16-provider/gradients/0604.pt": "bcefb03c9741e6625375bbe3bf23a4e183e15b94e04ec9f17937b490b6913643", + "bf16-provider/gradients/0603.pt": "0ed996bb8bf4df3dd76ec08056f76f03aab6eedd219f13a2f160cc782222a0f6", + "bf16-provider/gradients/0602.pt": "54e929bb4912a7ea5c0df1d971c6d7f2592608e03882d34337e2d440627b1434", + "bf16-provider/gradients/0601.pt": "6a9e96d568953e43e0cfdad1bf21f0c7af94ab58ec20d9f7455e4e0897b7f39f", + "bf16-provider/gradients/0600.pt": "fbdbc4a352ba654f2c3a9fcf102e2a4b297876e9e5bdacc0ebd2cf9b4535e83d", + "bf16-provider/gradients/0599.pt": "88dd012e28d652906053a614309260f8dff041b49f5bd25acc735694016ee9a2", + "bf16-provider/gradients/0598.pt": "221d3d71a7f2ae8df65643da20c1598506dd329486e48bfcefd2ae90fa15a1d2", + "bf16-provider/gradients/0597.pt": "8ae935a939083e830c030362f857d70dc3b403f31aa60b069c8502fe4c78d92c", + "bf16-provider/gradients/0596.pt": "9888c82184122e0b7f1b24eb98d886545a11871e6f765be63c4250d58d25b6a9", + "bf16-provider/gradients/0595.pt": "890446a2799cd772077e1b3a96b9c28b44bd67a238835915352c72d69bbc9b62", + "bf16-provider/gradients/0606.pt": "d878219a9c322bdc610b840e657905ebece4bcf489d4f60e20ecb8a634cc9b27", + "bf16-provider/gradients/0605.pt": "ad169fbd40159df0b2a04665eccf1ede575c60d63c1c7837fd484f0aa38de5a3", + "bf16-provider/gradients/0592.pt": "8bbdd22256b837d712b1ee07cd7638feacebdb19cb05d123a0b14aef46f7dda0", + "bf16-provider/gradients/0591.pt": "fd3c0839a61f3c8e8b3387b0d64aa57c0cdc8314c1cab47b2af6fcf0de18e6ea", + "bf16-provider/gradients/0590.pt": "689f0c9b67e7e3fcd0fa372b130626ec3b3912940caf34b39c2f35a79ae54d1a", + "bf16-provider/gradients/0589.pt": "0a8199f7a5c05ac343dec7fff537a821690c2eab374e9bb4b72e985920637188", + "bf16-provider/gradients/0588.pt": "bd960d7750933208a3a0deaf760806e8f41d1b809fc952ca4721efb30112e196", + "bf16-provider/gradients/0587.pt": "0ecedec069314be2c5c804ac9c080aed436a26bb617e00e16143482e17abb9ca", + "bf16-provider/gradients/0586.pt": "c68f6b73c56a0364d8726b63b5014cb85ab0d44fe4a6b97d28cb881acb77c11e", + "bf16-provider/gradients/0585.pt": "8cb2b4527d51586788fce5328f2046fdd02fb3dff61b3d4db1bca938797cbb82", + "bf16-provider/gradients/0584.pt": "fd9e2170bfdaee78cbad2af809c172ed1cc011c39fe5ce5e81126b5e21e2d485", + "bf16-provider/gradients/0583.pt": "945e8f576364faffaaab6ddde5e24d26b159369d124c04510da98221f40f82af", + "bf16-provider/gradients/0594.pt": "8602f38b3b8d337dff24593156a4adecf274da9a80cfeb66c477c131547f240e", + "bf16-provider/gradients/0593.pt": "46dbd687badab03b9018c139d00adc6bb4b4e8dbb577bf1eebeec2458ab871fe", + "bf16-provider/gradients/0580.pt": "1491b47c7978d9ea8359a56292fc44c30b89a89ea59f208715c78daf9d6ac2d6", + "bf16-provider/gradients/0579.pt": "c63604ee2ffc3ac441857498a5cc8e563ae0a1a68d79d378ebdbb654726c9985", + "bf16-provider/gradients/0578.pt": "8573eed9b76c6d67e917a9568684fbf073e4967f1b47aa2f1b3ab6f330fe08cc", + "bf16-provider/gradients/0577.pt": "c5222d4c103c9860ea6fddcfb2bb08cf1dd0986784e3abeab51e0fedba6e54ec", + "bf16-provider/gradients/0576.pt": "23aef65cf4f34badb6f54f79937f785907a60361f8ff8f376d38ec4f9325ff3c", + "bf16-provider/gradients/0575.pt": "a5a3e293f026a3ccfc1f3a9f69eaf9c2b48fb36b307e4e5c63d4bcb273150402", + "bf16-provider/gradients/0574.pt": "ad03e712eae857ae84a4390c896367ade6323952d7ced6eed7e4a8d7eb40da31", + "bf16-provider/gradients/0573.pt": "57ee0afb6f94214756c9c54a490823b27abdb9005a849b0fdc5b640747a32a89", + "bf16-provider/gradients/0572.pt": "8768cade03c13d12225068b4809d7cb13d89a387055fc3ece71bd1cbd01a04d1", + "bf16-provider/gradients/0571.pt": "7be2bc4784f7e9950a73628a98feb51c287665f55c8d471a6e63f27c73df57db", + "bf16-provider/gradients/0582.pt": "54baeabca7ef89c73e80f7380f99f0de89601beafe54afde037852fb0eb6f7ab", + "bf16-provider/gradients/0581.pt": "9359cbd8a00cb77ccd6dcd98303949a3d661bdd1f5194be1c3302565d0a67ab7", + "bf16-provider/gradients/0568.pt": "be46d5c4351a616510ab17d76b5e806bb29dbcd48b00554ca91d91effa60ecb4", + "bf16-provider/gradients/0567.pt": "5770498c485f2adba56c536b11ec8e6b310b125f35e3d161c9114b50e48c08fd", + "bf16-provider/gradients/0566.pt": "984de7b29f1ec4a997c47ac9cde37713c5318e38a03974f225e435980322e5c5", + "bf16-provider/gradients/0565.pt": "0edd6067054df8df147fab513af9d769bd5ffcbcae67417b1fd14a4cf7d3af2f", + "bf16-provider/gradients/0564.pt": "a7a9edb3ab3d9d41d05df54033c91f211eb4ff312ef936f7cb55c83f679cd673", + "bf16-provider/gradients/0563.pt": "4a025606f6cfe00c2de7eaf2b80cd472133eebf6108e7d64f71588453b4339fd", + "bf16-provider/gradients/0562.pt": "9fecc5de4630537cf580421569821b71bafcfe65dcc38ad81633c08134045ed2", + "bf16-provider/gradients/0561.pt": "533684836475cff6d2bc5d95a923ebc2ba5e0a313f91b98d597f5f6a7c80c47b", + "bf16-provider/gradients/0560.pt": "27ed55d0e7e8f8ccb0d3e8d5cd88f31df7cba976ee71537919c6106cf3fdf908", + "bf16-provider/gradients/0559.pt": "53a2f2a9f5ff38cd615ea1d0b716554685ccc76b1c1c0ed574a419e3a61ab48b", + "bf16-provider/gradients/0570.pt": "17ac192b08997134047773026351be1d1db636cc7f6f8b7932db24e40031ccfb", + "bf16-provider/gradients/0569.pt": "86c813be3a9da399fc2b9ef799f5362f2113a34cb8dfa09759004d819adc70de", + "bf16-provider/gradients/0556.pt": "5d38387ad6a82948ce726466261d5080e0be20ba316cdd2a733063555bc6dbad", + "bf16-provider/gradients/0555.pt": "22cc0d5842e8410830f2c4bfa8aece9afd10cf946cb8bec33c76fdf8fac329a7", + "bf16-provider/gradients/0554.pt": "aefb9d245e82e0489b09a907b22cbce33049b23a6d2eae2288a4ce8d6d791414", + "bf16-provider/gradients/0553.pt": "a7a46058da0b6bbba09abc1eb9ba6476c8e6b79b1a1b9c44a46640337ac7a0d8", + "bf16-provider/gradients/0552.pt": "4f5217e95931f275f7f8f123e11bc785e993ae3e1666097fcea3e3d9a3d8b04e", + "bf16-provider/gradients/0551.pt": "a4316761352fa318b00679b1ac45ac02e367d12e1e3cb684d76fd9f06fccede3", + "bf16-provider/gradients/0550.pt": "5146dc1f6469a40d648a442cc26722c6a936ec0d61ab9f68d5a4582ae5c0dfc8", + "bf16-provider/gradients/0549.pt": "fac434ba5e865bd6692d8a6547e2884ef72aa51784ccfa5b2534cc30599f8a1f", + "bf16-provider/gradients/0548.pt": "5c102395455ce8168e4be4c20b397924957ed1d5eed6fb190ed706188f57c7cc", + "bf16-provider/gradients/0547.pt": "ffc8141ac3d86477b39d0866968032db4e834753e6280a38439383705ba296e2", + "bf16-provider/gradients/0558.pt": "ef70685243de893d4f8a4ca105ee66d81ff4f65b791a4ad971cca9b1069689f7", + "bf16-provider/gradients/0557.pt": "eb10aeab6ce010ee9836d3fa25e99d0847483c30be9b5215e01712306a1081cf", + "bf16-provider/gradients/0544.pt": "f5812b13181808ded4de9c07802f9b30b5181d1a1d3724177a4e90803bff9c60", + "bf16-provider/gradients/0543.pt": "ccc434434e34eccaf5744e962d07bc25e461f57f8e1e862f2da4bd5b9eca6c16", + "bf16-provider/gradients/0542.pt": "f065482e4865b0847c0a4c677177a9db80c52fbde3315460a52b82d2b01cae32", + "bf16-provider/gradients/0541.pt": "c87bd58dc250a8d960ab75703713325c8d64d944ac80cf59cb6185625bf54e7d", + "bf16-provider/gradients/0540.pt": "9d8a6ab3d9f4f746760ea4e590fb645062a13d0da36bf49ebb90ee254aa33328", + "bf16-provider/gradients/0539.pt": "1db2f753557e53896f7a3b55594405472f91dc51389e3fc77f1f74dadd9a2ffe", + "bf16-provider/gradients/0538.pt": "a37e2b4389b788567f0748fa937358f1d92decdcbbee2aa30f5be376081f227d", + "bf16-provider/gradients/0537.pt": "63e82dcd38eb8f7e9bfd752299bd89642ee4650111b083b048127d8c276dbd2f", + "bf16-provider/gradients/0536.pt": "178f0d0b39cecb680d755ebba612815be2bca21b2b31adf43ab52a3ba8e9e382", + "bf16-provider/gradients/0535.pt": "b1cbc58016470fb40996db3f6ee8e4d49c6872edf3cf1c0ec748d1d9986385dc", + "bf16-provider/gradients/0546.pt": "8aba3a6f188e499801e293ea99e6a24872a0b3da322a0f2686472fbadb97bcd0", + "bf16-provider/gradients/0545.pt": "13cbdc90c566f913377702a0579a5ba75c5517ce23a68e3120fb9d5a4ccacb3e", + "bf16-provider/gradients/0532.pt": "2c615efd2923d609bf1cf6f77d44dcb3e351cc7dbf0283fe843be0198c04c77c", + "bf16-provider/gradients/0531.pt": "6e045c469beefbc41e97702d42c908ccb777c3682e278ade0ba47bef890580c6", + "bf16-provider/gradients/0530.pt": "b703f285f960fcf2f4e0984ee64f3119ce26820ce9204db657d0c21043ec4fd0", + "bf16-provider/gradients/0529.pt": "eaef32f9eeb96ba92f46432c8114feb21dd6d1d4fc3bde96d6300d84bd8665ed", + "bf16-provider/gradients/0528.pt": "e0592d062f6a01396ccb01cc53deed7a84edafa914af63241a8d06c16cc3b60d", + "bf16-provider/gradients/0527.pt": "8e0ff55072cf1798c0963f4a00b0b9ec25e20052cc8cd79b3ed633c826e24d96", + "bf16-provider/gradients/0526.pt": "18d9af8ebcd68e06a3dc9a0cea36ca6d6efd2231d7289a3e385e0fbe88d4cb22", + "bf16-provider/gradients/0525.pt": "e441d5ec143c790b19762e2f217028ab7266b8d8b92864e61c07749499a19158", + "bf16-provider/gradients/0524.pt": "3ae5eaca1a440c8373b5b61584644db4d650c30a024d446c06d631f310ef19f1", + "bf16-provider/gradients/0523.pt": "92f7e448dd792f3cf7ae0b006ccb190edc0310fd08028cd388f194ae2f98f9a2", + "bf16-provider/gradients/0534.pt": "b0b7832cfecec6c81a0f1910d5141a8a8b5b13bef75599abdf588ebd80bc39df", + "bf16-provider/gradients/0533.pt": "5cb82cd6cd1f7709dd1402d3029f28703e7eb599dec844e8e3b422bbbfbc3c04", + "bf16-provider/gradients/0520.pt": "dc6193ebf4b7e1f9acf742f6c049b44914dee467f09b4ec704de28547f773962", + "bf16-provider/gradients/0519.pt": "fde3ad4471b0504fe2fbf6795e385394831e9eb98822ecc390922e8f3cc43b1b", + "bf16-provider/gradients/0518.pt": "306f556e6b7df14938884265422e149b638a6ed52888432e9c7a598ef35f7a97", + "bf16-provider/gradients/0517.pt": "0e6ffec908b0b43cada76470fea711bddac5c6104a3db49412e357eec175769c", + "bf16-provider/gradients/0516.pt": "835066eb37267c91e1fa65874e5700976430603a6e7a8fbe99699c402642196a", + "bf16-provider/gradients/0515.pt": "90043d814566b6a61e2ea7888f15d681101165a8833a51f503d58a2a6d90ca1b", + "bf16-provider/gradients/0514.pt": "f3f633bdeabcf26df4bb70dcde9f5e10fc50c4479f190cf09ceb4395e569e67e", + "bf16-provider/gradients/0513.pt": "211944e3666590a18abef957222a5999fe7c9895cde6f34d02ba5d1a20f0c8e9", + "bf16-provider/gradients/0512.pt": "f7d95606e6c293746562cda0a7697ccc9a73b20cb2d2062681ae73802730513d", + "bf16-provider/gradients/0511.pt": "7dcbb1587b8a88bdf4df3016df63b7db60ec82e350c01d8d1b60324b76f879f2", + "bf16-provider/gradients/0522.pt": "3c623cfae38c555d279e4ebc55dce1580241c46d9fd80e7d79b9a7f3bb3aceba", + "bf16-provider/gradients/0521.pt": "c7fef0cd677c2c2a5c0906590d3fe37f28459f1b1a9124ae8813a60260523b39", + "bf16-provider/gradients/0508.pt": "b9962e684b42a46ea85472f2677b753424c202c4fc47dfc402fb7bedb36db7c1", + "bf16-provider/gradients/0507.pt": "a4ab90249ece77aaf07ba8ada94c057904c1892d54873aa3461aed852724c6f7", + "bf16-provider/gradients/0506.pt": "cc059e7e82af46b0fb5fca3377ac0ec774411bf8fe79224d86350023b49fcf8d", + "bf16-provider/gradients/0505.pt": "35e5ad1f88d24fbda6cb315a64c991646c22a9101cd84bee0866b730e25261cf", + "bf16-provider/gradients/0504.pt": "01f2cf395c8e0a322301ab737620f025ff4c51241497bfb5b156f48b39fd0f0b", + "bf16-provider/gradients/0503.pt": "77cff9445ba8f1019aa7b9828f7b8d38c3bcdb1225fa28c06b296f822c29e64f", + "bf16-provider/gradients/0502.pt": "7130564a5ed1b02ea68e72aee594c2bff4302db8ab5e0f45fbaad2b75c19c9bf", + "bf16-provider/gradients/0501.pt": "8ff81f867922902edb9453e61fd21d7896b12c63b1b226c691bdfd7171b0f9b3", + "bf16-provider/gradients/0500.pt": "94263e7b0e65c6b51ee84fe45c61873723e89d320b08d5a4e1ca78af04e04756", + "bf16-provider/gradients/0499.pt": "663aaff4d99ec8b2f58085cc3bd5cf6f4062467b980aabdb0099d0c04bad9b07", + "bf16-provider/gradients/0510.pt": "5658660b4191fe7a0b7b9ae7c3ecc5ceef8b0576be6a83c270785dd129094a46", + "bf16-provider/gradients/0509.pt": "817adb51b2e64d38622d83fd59e360243862be28863311a019792583f3f7f5c2", + "bf16-provider/gradients/0496.pt": "1ae60be582155b48fe58a92a2a9a1e25c1c97b0d58b07be1334b70882795f0af", + "bf16-provider/gradients/0495.pt": "b7f691c0beae28cc99ecb050dae1154ca27d5a83ebf759bdf27b785e69593514", + "bf16-provider/gradients/0494.pt": "2f441517cc9479f5cf538ac4e0dcc86f37df201b80e9bdeab69b1daa7a64b615", + "bf16-provider/gradients/0493.pt": "f20240e0b40e95b910eb4a262948fe8ca21d4d2dab6f0236d50662619d94dcd0", + "bf16-provider/gradients/0492.pt": "f56d7184fa1fc74eab5e8ffe910151b7df772170d0beba7ff099ee9e4e6cd9d0", + "bf16-provider/gradients/0491.pt": "03ebd4d2a58ebdbc22ebf2e6bdcdee6533c256c7c3c1184cdb0e8a2fef1feeb6", + "bf16-provider/gradients/0490.pt": "4d796d39f7b8bb74e311c777d5ba00ba259c4556fb8fb0d54e917e6901cdcd2a", + "bf16-provider/gradients/0489.pt": "84d87af311c8b91664283ab5b2aae4299313b999d188df6798b309a9cd1c1fae", + "bf16-provider/gradients/0488.pt": "78b4e4da3bb7a0b2b0c68d203f70631143beb55af9e38e55eb73731426dc1915", + "bf16-provider/gradients/0487.pt": "6c52d0b39d7d7d8521743459744aac609ad2310464333b7635e2307bf96bbce0", + "bf16-provider/gradients/0498.pt": "ac4a9f164643ddb499717218c034112daf3419c7d5355fd160b8522843c9eb6b", + "bf16-provider/gradients/0497.pt": "ed7fd638cc5f89b9d794f9851ac983683ba48f54d1bd4a6fca55c09e93d10fbe", + "bf16-provider/gradients/0484.pt": "d677b348266831ebbc4c523e3d84ec83885bb851a00e266b0a51fa2ab5493621", + "bf16-provider/gradients/0483.pt": "22a12b0ef8de6043639f920c524e8c930d2a02b1d20451109a8b6fca0a919529", + "bf16-provider/gradients/0482.pt": "c5bced8939f0d1005c917c2d4d3f1e42eb115b4cf95f6e9ec086e9103f5d3a37", + "bf16-provider/gradients/0481.pt": "9c6cb5e1b11bf8f01a34aa7955cbd31b0b02d4eabc29250f11edb94eff21f0fc", + "bf16-provider/gradients/0480.pt": "5da555815a3fe105b85057c507eb8ff2498f3102d13028f65969daf0ffa90195", + "bf16-provider/gradients/0479.pt": "c59623111a24369ef542337593b4244453ec11121f932831d845a09fa10767eb", + "bf16-provider/gradients/0478.pt": "51b3a8f563657f5fb55bec5774d9e3b431a6fca3be71a95c789e0838ad4bd431", + "bf16-provider/gradients/0477.pt": "891d737107ca018bb54c89fb5c919af92709baf5e55281108cb6d1b62f5b577e", + "bf16-provider/gradients/0476.pt": "03c78db510bb19dae85acb75d11e5f5ad65d6e54d7aa6ede0c5f5ae1e598ff29", + "bf16-provider/gradients/0475.pt": "bf623529e4399a3b5f45d58fbaa3e994e101e35d4e7ef4bf174d11038216eb03", + "bf16-provider/gradients/0486.pt": "844838a0d4a67972838a0fa63c0d19c52924a10bcf7e64a6a16631b501325c84", + "bf16-provider/gradients/0485.pt": "d34cf4a12d9f02908d36a7f7775411de7e3a1e3e6539ee98c6cf88e6b812078d", + "bf16-provider/gradients/0472.pt": "cffca0566f6916795e9a852648b53c5a12ecc297141bb2d28dbcb1cea1f94140", + "bf16-provider/gradients/0471.pt": "47789d466380fa6c27a12fbed584875b4adfbecdb92be3ab8bf86996097d0fe8", + "bf16-provider/gradients/0470.pt": "335519ee69f6d2b2732de1f686500bb6cf1c0e1989f0b13e6f05aee7107f84dc", + "bf16-provider/gradients/0469.pt": "852f5894d5bd0e2e876b2dbcbc818348aa4b8828cbecf8992bedc826460d56ba", + "bf16-provider/gradients/0468.pt": "3b8e415537faefd522f822b0def895db00fda3fa56eb1a9e099b084ceb8a76cd", + "bf16-provider/gradients/0467.pt": "00de339693e9e42a50d2d68053f506fb74dd2edd0cbcb5e36c600c05386516f7", + "bf16-provider/gradients/0466.pt": "7c4d30defcd08033f135e5b18b8755561b05cf5d483b96459a5cbe571245a065", + "bf16-provider/gradients/0465.pt": "3b78037be58521e544113886f171d23c1373f2737a67045e19718f3d2419199e", + "bf16-provider/gradients/0464.pt": "b051ca08bef62117ec3460c5401114293d18e85e7425bcf4865998c3846105aa", + "bf16-provider/gradients/0463.pt": "05e46077a1beaca7d3424fdc71b7caa74f7819351b07c567a994374a4b01bc13", + "bf16-provider/gradients/0474.pt": "71be156216ba221a97e1423a03217a7a7167c88f150f603efd2d8c3a11f38503", + "bf16-provider/gradients/0473.pt": "bf55e7d09d5c5b856d6feca13155c7c4e82ecca6a4520856365e706ab00f0d5b", + "bf16-provider/gradients/0460.pt": "7fabc4c35ccc0907d2ecd5b49a57c1efb64cfef04c8828dc96de9c22529cbde9", + "bf16-provider/gradients/0459.pt": "459398bfa18a8c7021ea269175b9d639fdf138ee51dd84cc9d1101f52bbb6298", + "bf16-provider/gradients/0458.pt": "c1c19be6ddf5f7d84700f5bb13f9154b39cfd949f51a730d946bb66a8677c2bd", + "bf16-provider/gradients/0457.pt": "30dd1380984026f13f99f8d9a3df21ceeba1b3dc526684db5fd5acee94c64715", + "bf16-provider/gradients/0456.pt": "830b6ce8c09e1e0d1d51acd72d2ad4ad8595d23072ceae4efc3970bed5570cb3", + "bf16-provider/gradients/0455.pt": "d243fbca25f5f6aa29ad48552f38fec86f8a01c8534eae18f5fefa3187210d99", + "bf16-provider/gradients/0454.pt": "8d6507b3747f7224ea8d168b034911ba1a7f1837a52fc070faa696186b8c9c02", + "bf16-provider/gradients/0453.pt": "4f1524673d3334f5b0884719e71620c95df2719d8d68e990bb8fbe89c99827fe", + "bf16-provider/gradients/0452.pt": "81a00bb1487bdc2dd595e0e3f7ee75e7d9148eb943f566d7bf3777c97c5e915b", + "bf16-provider/gradients/0451.pt": "d31dd660f1c1e8ae3fe121978d12fe12d006a3fe4ea06b43b0ca948c43c56261", + "bf16-provider/gradients/0462.pt": "99cee9ccb61645d1dbdfb9abdc91a677fda2c0c9d67c654f4ee265e6d3df614d", + "bf16-provider/gradients/0461.pt": "0c0bcf12ef370674906bdf31844ee8cbc460516270da6b924faa8eb22354a7b4", + "bf16-provider/gradients/0448.pt": "d669605d10e1954f66a509a33d4cbe8f01820a0aae89accb562a946699bd0a43", + "bf16-provider/gradients/0447.pt": "f0991a36dc1b3f2de055077396835da8cd918b592291f189b11b21968fedfe6f", + "bf16-provider/gradients/0446.pt": "8fee8797705a82f026c1d4369af14836a07113a3000396c9afdec448c149e539", + "bf16-provider/gradients/0445.pt": "fead9a7e8c04fc2624058097546e874eb151025d806206c7b8c0b78583111fa0", + "bf16-provider/gradients/0444.pt": "f94353bd4cdfac7d0d14be7193cf401a5bf920d0790deaa7841459878614bfff", + "bf16-provider/gradients/0443.pt": "a9e90619d6e5521aeb2334f3f6be6c803274580eec082d91c209219b5d5d51b3", + "bf16-provider/gradients/0442.pt": "8e21e37aaf689d5ac7a935c6a00b271ae54ac323d891970c8f0c7e2f3503b6b8", + "bf16-provider/gradients/0441.pt": "5a0ab2c9b9163f3a568e2bfdffa89227dcd6986f41f295cc03a570f24e748d41", + "bf16-provider/gradients/0440.pt": "36a3c16a220274d2bc149cf95b75770ffb0660b21f893dae6c75e20d2ab5c9d5", + "bf16-provider/gradients/0439.pt": "b171415dc97e7661edf1ad12a2674373d5dd9f12f1a7ae45c6fb2c7143464d1e", + "bf16-provider/gradients/0450.pt": "8e6f07c28eae7a00bc37f8d6f9d92a8c0b06d6de8231b0361bf4bb7b3d9bdee5", + "bf16-provider/gradients/0449.pt": "7b2a0f5393f7c4b6342495cae101f5d0fa69194bf5c190bed72bf28fc75003e6", + "bf16-provider/gradients/0436.pt": "834ec333f78bd28821895823aca1b30ed8fe08da07b718b05dce0ac5e6dc8bf0", + "bf16-provider/gradients/0435.pt": "34cb3acae0d595bca5bdafa674db93e07af5f2b622d1d522c51891886cc6baca", + "bf16-provider/gradients/0434.pt": "73068876b071591f76196e1c2c8ec8056a40b8c76e47a1791cfbb6a164ddecf8", + "bf16-provider/gradients/0433.pt": "dac1afd8bcd9a2aff4726ab8ae02fc31a549f990e957584b2f1f3f943f355c67", + "bf16-provider/gradients/0432.pt": "eda8de51540f5547b3bf15019803b443f0c53a34bfae2c69c7895bf0b603fb72", + "bf16-provider/gradients/0431.pt": "9595b4baac45a409602faa397f33dcc17748636ce674617fb5fc4e7334c1bce9", + "bf16-provider/gradients/0430.pt": "940334deb170ddfed4b1a31e8c229628aa009f745d708ab6104bea92c8fc1e0f", + "bf16-provider/gradients/0429.pt": "e4d4a1c4cc701cd55c69c00c508ac3bf76925cf2b31d66d4e07a5ee929485a97", + "bf16-provider/gradients/0428.pt": "14d3e7302550ae94a58643aed0a2511a3bfa63daecb3cede0d62eaa682a8377c", + "bf16-provider/gradients/0427.pt": "3d104208990f1ea302885f316d2b08c2774a3d5e028ddae08c94eda048d88f32", + "bf16-provider/gradients/0438.pt": "1baed369a2ce360a0de3eb5863270a5f257117221db024f0f98932a86fc26405", + "bf16-provider/gradients/0437.pt": "d6dc7e416eb2fdb75eaac86bd4f1c8b78efd500cfafae9a8f065429d56ad530b", + "bf16-provider/gradients/0424.pt": "1b898b8bc3b14d42334488f9df4165690876762702441029edf32b6e753d9035", + "bf16-provider/gradients/0423.pt": "f87f47ba3ad6649007e9020f8706aa331b98513294ea61f62b3420f382d7aff1", + "bf16-provider/gradients/0422.pt": "10288c7ee2c446c385058a64c8c095fdf56bda47f389999a513c382a63748e12", + "bf16-provider/gradients/0421.pt": "50b3c863396334d9f8ea937c0c716052a4edefaa2bfd31f70e314d61b5d9e04a", + "bf16-provider/gradients/0420.pt": "e00669d8f3e71b772e5a0f7fa8911c7cda6e59195f36df868659c5e2bf1bfd8f", + "bf16-provider/gradients/0419.pt": "2c2f4bead308123ce434bdb8ba3ad24a06e4b45770b50066885d02484b470006", + "bf16-provider/gradients/0418.pt": "dc69f98edaa1b6b26ea2c0df259f60f097129f3cba6ff943ea8a8277ff3d05b7", + "bf16-provider/gradients/0417.pt": "e723311caaf808dc9c0459b41f94618b7c2e2bcf806cd1691a14875eb239f41b", + "bf16-provider/gradients/0416.pt": "68d79df4f2a59dbb98ea8758ea1fe1a5b9a649909531df08e3ef5b19bee21790", + "bf16-provider/gradients/0415.pt": "c657bfb1ec98bd4b6e56e8a8e6dae63e716cf4b91688bde1f39066ca70222c64", + "bf16-provider/gradients/0426.pt": "7cf0e8f7f7b12fbc639cef97ace25ca674d93f1d88d4bf2597bf64653679f483", + "bf16-provider/gradients/0425.pt": "0500fb0d766b281637ba24e2050de2e3fc10601d564c674a8640dfe3c6f373fc", + "bf16-provider/gradients/0412.pt": "4bfc5c032d26964efa3fd2a84e1d8968b769d40eec8bbe9718d392855c626f71", + "bf16-provider/gradients/0411.pt": "2f4ae0f3c0e4096397e7786b3ea0259099187206e5c1b7b1b6b3c5fa86b25d69", + "bf16-provider/gradients/0410.pt": "42f419cbd8544eda7be7fa2f667097d8bd88e006538e9d67d1bce3961b11225a", + "bf16-provider/gradients/0409.pt": "eef39b3b6b1d834cf0e506ca09c0de4f3843fc891e2995ebf09174704b4371dd", + "bf16-provider/gradients/0408.pt": "b3adae34f2641f63b9d9429a286631d22ce93a8b119572b47594f3f04439e6ab", + "bf16-provider/gradients/0407.pt": "ce7cf574a586123e9cc0f8d0117e409b1411663269383a346af027ce19db29cf", + "bf16-provider/gradients/0406.pt": "a499999820e53a0f239336712c66ee76ead36582c149c9d3d6f905f65f6653ec", + "bf16-provider/gradients/0405.pt": "3111bb92fd28d64f382b4aa0767397cd2514d5602b0a2de635f47009b7133861", + "bf16-provider/gradients/0404.pt": "e751f5cc45343da810a307d273962c735c8edbe8d28ec19b1957e80c9c392c96", + "bf16-provider/gradients/0403.pt": "6336ce3fbc6f62af1a59f40650229c6a72791e15c6a78c730a5445029e9edc23", + "bf16-provider/gradients/0414.pt": "983932ac4f61e0a7c2a592a58ddeea75dea31f7aba0ced0b922e213064e72ed3", + "bf16-provider/gradients/0413.pt": "3c3e33a6f32818988c23262eed7ef1a226d69f5a5f68ef5c43582bae43f99420", + "bf16-provider/gradients/0400.pt": "6e20fa0e1e3756137131e6744ffd2d8ee67d285accc58638ed14b56f641daa2b", + "bf16-provider/gradients/0399.pt": "199826583b91b8a1e4ab99510e120815e327d477876a48cd06436f80568061ab", + "bf16-provider/gradients/0398.pt": "ad1c0c75231bb2ccf4dddbf7252ba64c7f85011718730af27f15ab91866f42ae", + "bf16-provider/gradients/0397.pt": "6126e25b6ba7c4f39aa4169f82efb2edd8866a0f7ebab22901fe4806280bc99d", + "bf16-provider/gradients/0396.pt": "b9401ae1183b64c57979e10df3fd738439960d17687bd10fcfbe23fd74c72d6d", + "bf16-provider/gradients/0395.pt": "00cb1f8b781f7df14ccbeaccdfe13d0f184258680d09e5b121ed4f62d4b50b43", + "bf16-provider/gradients/0394.pt": "7c41ba4078643c2c007c57a88417a1d82ba8159083c979b518f4ff440c48ce05", + "bf16-provider/gradients/0393.pt": "dee0364f82594a2d9d3880ea80a46c85322f38155f6b45b6fa46b4a8728da63c", + "bf16-provider/gradients/0392.pt": "7d2210cf08a9090efe94c510c1ea16cb1f547511f5dad337d55cbdfa0c72d10d", + "bf16-provider/gradients/0391.pt": "a8ad0998896ee8ffe5972b3f214746319e4e52733984e8e1df0424bf5ee799a4", + "bf16-provider/gradients/0402.pt": "3f8dbe7fadc5ea145906611363bf22173a47368175f76ab266384018ea60017a", + "bf16-provider/gradients/0401.pt": "de62139743fb85546f4a2c5a71e9c1ecdeb53c48a49b7d97bc257748623203e4", + "bf16-provider/gradients/0388.pt": "c01a779947fb30a631e357fd880decdd520c4caa90f31473d4b58b546b596eef", + "bf16-provider/gradients/0387.pt": "1af9a5d52305c25c7ef28239e6b73cc6b69aed344f9c2c677d8580c189405401", + "bf16-provider/gradients/0386.pt": "83f94e71e1ac221cb1e56132ade7704b2856cf3d126227a8b217c0300e5144ca", + "bf16-provider/gradients/0385.pt": "93daf48ecb6e04cf0dd96764571546a3d8b0cbc08eebc39978a87dbbc360eb3f", + "bf16-provider/gradients/0384.pt": "4c889512e4ba10b671cfd286a629b3d19f965584a50b7f5916593109eb05dd2c", + "bf16-provider/gradients/0383.pt": "83a4d89b8d6fbdd88a4f3224565faf09180735e7da060fd6a44395b78b4131f4", + "bf16-provider/gradients/0382.pt": "736fdc79e54f720023033cf411b5a56eed176d259cda852d6fac8e4fe0070962", + "bf16-provider/gradients/0381.pt": "e1650693f27a890330a91b847ff669f2845d962df460ff9f18e7706b0b8f0ced", + "bf16-provider/gradients/0380.pt": "6af3e26fcaf59b4ca8d600dffdded8191cfa7d31773b4bad11dd689d8fe1891c", + "bf16-provider/gradients/0379.pt": "7a9026ef060cb6d7ef16718bac6f3aab381f38bb4afb1bed63cf75213cd45ae5", + "bf16-provider/gradients/0390.pt": "9d01f008020aa5f7bab6bd45d2321186ae97ecb3f747f21c38d56d13cd08563a", + "bf16-provider/gradients/0389.pt": "fed763e3838a95b12a4a4683b1a863eaf4f8a17e1ca7ff596b8f1f1824e80fb8", + "bf16-provider/gradients/0376.pt": "80d991324b2a369e92703b4d330cf9f95a2625c05a6db6e2a1b33cad19ef6cc8", + "bf16-provider/gradients/0375.pt": "9a2548d35c9168c4b76fa30fae9a755e8b2401130261a645fab732f9ba58b631", + "bf16-provider/gradients/0374.pt": "366e7a501ca2023befd3a1f1a6599a72f59f49f9af9cd52895a9e1d245b67d47", + "bf16-provider/gradients/0373.pt": "2f7113d23b7e40fd5f5c3940bb0b54b8caa829e8b278ea720ba99ccbd4533640", + "bf16-provider/gradients/0372.pt": "fb22006855623f04e5823a46ce0ba4e0ac7ae939bfe5e81efbb25a812ec196c1", + "bf16-provider/gradients/0371.pt": "c04413582df24961b0f65b52f7a0480f922105e877c465f574dd9ec4bdb66df8", + "bf16-provider/gradients/0370.pt": "df53702eb87cd51d31af9d714012b7fdfc031b0b6f334903c9d02507324e5db5", + "bf16-provider/gradients/0369.pt": "5cd145903aeaddfd0523565eabc48ed2ec3a9f9e2de099ca02d3085d76387672", + "bf16-provider/gradients/0368.pt": "442b10d48479cd3bb6637e77028c7f37275c0498e0b12c657461af6f9e479dfd", + "bf16-provider/gradients/0367.pt": "2820a3ddd747ee37c3f92ef04921cfb014f1bbcafbd47304f48aa4cd64557401", + "bf16-provider/gradients/0378.pt": "7bc6ac4c284eef24a85062aeab27d2db52d9f52294f8249473186a2b9619b406", + "bf16-provider/gradients/0377.pt": "0fb26d3d09283c9c93ee5ff99e8d015058afa2759f59678297a05465d586d941", + "bf16-provider/gradients/0364.pt": "7086c5889af19c994a10c9fee802da9a0d3505be3ca936a39bc819379635b3c1", + "bf16-provider/gradients/0363.pt": "b10767c3b65bd0440f5374265e67bdf2998868413214e1d5348b7c7a6496af7b", + "bf16-provider/gradients/0362.pt": "7540995f2696f4fd1fd660595c81d78d2ca64ae36553ac4eb0cac167c3f3fd36", + "bf16-provider/gradients/0361.pt": "cfd05b67644a4c309d6a72c22f399f36f977156a150d66ffa6e2586d9a49cbd9", + "bf16-provider/gradients/0360.pt": "c587df6077f08da4b7e0619960cca589f5248db0c32bb82cdb781602cebc59c5", + "bf16-provider/gradients/0359.pt": "92b58b50957eaa4c9bde642ae32635a31f1fab21f5979b1a4b98ba09edbde855", + "bf16-provider/gradients/0358.pt": "10c4b56975dd96a39c64a0fc1296f89e0c2ba0b63d363bdb65d333eed53648c7", + "bf16-provider/gradients/0357.pt": "16f230ee1e7ea6f01259b29c1747d9126043a8bf5939d8a44dacd673cd2ef072", + "bf16-provider/gradients/0356.pt": "e8ec4e58e9d1df18453191536147588efc7a9d7734c81cb748997e6234ee26a5", + "bf16-provider/gradients/0355.pt": "215feb49ffa57c65ead59544bf40b64cddbc2b408f5650684218810e9f9428ab", + "bf16-provider/gradients/0366.pt": "8dc4a8cb1ef2ef03b7261b9057a2d1b3415925b8e102dd713d7952324bac9906", + "bf16-provider/gradients/0365.pt": "048dfe529664b2f752583875838bf01c6382bd1356ba5c6ad4181b533aba765c", + "bf16-provider/gradients/0352.pt": "73b885e92b77e1a87a2ef06a0b9ba5b6106294192b8bfa925fd1a7ceb0b0bd46", + "bf16-provider/gradients/0351.pt": "8ada9649ccd6e091746424be7c0e9f3c4d7f4b5de9d0bf1624d6c94616638c61", + "bf16-provider/gradients/0350.pt": "b897d88524957c46b4b5c829385b1ded1c97a5c60995b1de494e047fa542c019", + "bf16-provider/gradients/0349.pt": "d4f9dda2f1686a3daf3cb6303d07bf76b14f2d162a6ec68edd2826c26bfffab3", + "bf16-provider/gradients/0348.pt": "114a6555be1fea1a2411da97bd39ee547da87ff8401f7c884b38111d285769af", + "bf16-provider/gradients/0347.pt": "d99e95c384ac1e1a79269ab365e7a641fb51cec7e53f48f19023b591162ed321", + "bf16-provider/gradients/0346.pt": "c8196dfd759a551f9073919e1184d9ae2c1e7df7f679b53ff72e5fb323828c90", + "bf16-provider/gradients/0345.pt": "061fbb5691890cb52bb3d4c4d523c30f19441631617f27f52a2156ab1e713d41", + "bf16-provider/gradients/0344.pt": "bc7afb0a9f9b5bf678de90d4e32760c9dd2133b2ff7665542f68bd10c73ca3ac", + "bf16-provider/gradients/0343.pt": "c89979544fadc5881e4db713ef85ba71fa97f3875b750ab6cfa1f61d6c8cbda5", + "bf16-provider/gradients/0354.pt": "9d28f231db5552c04e92baeeed5c1039326d284d896354b39151fb13059d9acd", + "bf16-provider/gradients/0353.pt": "2cdfa89c6fa01d30190c3561c3dcfa1a99cf345d252352b3d36a1d85d3b40f7c", + "bf16-provider/gradients/0340.pt": "16e736c422313236fa0687931f559902adbbc030a71a6a9265f0e0d48a744ca7", + "bf16-provider/gradients/0339.pt": "c44c4f8a6e6b690c9a676688a102a57a557d6b0fadaa2e5ab3914c37831d514e", + "bf16-provider/gradients/0338.pt": "432eb47149d2699aa0938b44e8311cf02bd578b2e4e3d9026e1b1ccecd078af7", + "bf16-provider/gradients/0337.pt": "9d2996d3e10e01646574bb35f9e4ea39d32041dda99842ba88443f6b1e90a187", + "bf16-provider/gradients/0336.pt": "5ac82f7000cebb3952880b1624f7c822661dd1e80e36272205c130eaf2c7de9a", + "bf16-provider/gradients/0335.pt": "5f96493dd608cb99d6d6b55638f872230e4074eefcaeb2ad19e1f74ccd52f775", + "bf16-provider/gradients/0334.pt": "ab438df3ff6b41104fbb062d14be1d32db9acd9bcce90586eff6d28af0952db7", + "bf16-provider/gradients/0333.pt": "c0b604625ceea9e2be57f6712581231b1c8a7fb19a7466a57203e1ef9b55fc95", + "bf16-provider/gradients/0332.pt": "89fa2ab19f9530652a5589e5dec63dbc5658d6292728abfceb6e971a7361cd68", + "bf16-provider/gradients/0331.pt": "a28d779aa3c7ea69749cfe92f72150ab4253b3620a25fd26157bc50eba51249f", + "bf16-provider/gradients/0342.pt": "dbe29d903f6f33fb03d6e57fc57066a77783484ac332a7001aec997aac07f2d3", + "bf16-provider/gradients/0341.pt": "15db439a8cc038776a700c98ff1af3dd5535d5b9a81f1f14d33102c3697266d3", + "bf16-provider/gradients/0328.pt": "9c2096499920a1f2ab43be7737404f8c7ca77df77ee8da8c4d175ffe79b1b8e7", + "bf16-provider/gradients/0327.pt": "6c209b6e47c1dfec4b212471dcc803ae83376d08853117a860d4d4800481ee4d", + "bf16-provider/gradients/0326.pt": "9c65c872e4526ca10382a1a58cf2f917aaf21e12ab9e23ac78a6f711b1c83c79", + "bf16-provider/gradients/0325.pt": "7f2fca93366b019fa9cf2399d3614722bf4026d2c62be5998cfa4f4ae8436e71", + "bf16-provider/gradients/0324.pt": "a383beb4c6dc0a9b81668c4ebd111b5a6e527c576253ccbbcc574967277fb386", + "bf16-provider/gradients/0323.pt": "bd611186afd49b0f5ba6c530c7c2307f30bd03e0a1fbfeeb4d15ad01818a6c27", + "bf16-provider/gradients/0322.pt": "9f7a464f463566fc87aa8800106925523a6e651f42595abc6969936574bb9845", + "bf16-provider/gradients/0321.pt": "07f5921da45d7d5a3d10203bb49465d56b92106aa9799baa0bf5f30947a287cf", + "bf16-provider/gradients/0320.pt": "dab1df649ff666c8e37eee09779a8b39eefd171a8f99da0102af1590bf4cb0d8", + "bf16-provider/gradients/0319.pt": "a706aca7f87492191c7c7434ee5d185764eb1f59d94ae7c4299d316aaeef51d3", + "bf16-provider/gradients/0330.pt": "134317c0c9e4b21682d825d0c35c9436e269c9517e42dff6440f73aba83c7867", + "bf16-provider/gradients/0329.pt": "46b37f67a291d05e35ededf73764f5763d68dab2ffb115ffe5d85afb8b8f7fe4", + "bf16-provider/gradients/0316.pt": "9390fd6efd83be43ee17922a778e5be65806fc0a88fd127c250d3139bf9a4eb6", + "bf16-provider/gradients/0315.pt": "d56e52357f154ae94ce733262f7284c38bbcb2cb4ac06079bc8c7008a485055a", + "bf16-provider/gradients/0314.pt": "bf43b0a5eab07b921a7c7d587f90fae7abca00158cce1e089ed0acc182ccec13", + "bf16-provider/gradients/0313.pt": "6136f702dce8910065120cd350ebf5cf3cef47ed4fe158f87d8421cfb3da669d", + "bf16-provider/gradients/0312.pt": "d9a8b18e267e976a7346bf311cfd79fa4625f091d03c104b74f0724a011231d1", + "bf16-provider/gradients/0311.pt": "5b295970ff06425a7b72ba8021c74ce779e14c49b4d3fe2be355e05d5655ff6b", + "bf16-provider/gradients/0310.pt": "0d653d6145f8ced6c4fa61a187bb1773503d188b51729343434c153d0c954877", + "bf16-provider/gradients/0309.pt": "e010c4c69dadd88e4b434a3c39f4fd87b84ae264659ce400b6343e935aadf4a0", + "bf16-provider/gradients/0308.pt": "da3c2b653c9a47bd73805db2bd693b463e8ef20beb16273b11484df3c76b19b4", + "bf16-provider/gradients/0307.pt": "5023bfd386b924ee3b89c869ef73022bc9a77fc7bdbe3d1c90a636874ddf071f", + "bf16-provider/gradients/0318.pt": "cfccb4ffc0d753db41273a3fc30c11773a5b36a3c9ad58e575a2f9f3e2d6328c", + "bf16-provider/gradients/0317.pt": "47aed8bf8daa5dcd663f8e5bdfa0ccf9942fdc886baf5b374b383adb86500462", + "bf16-provider/gradients/0304.pt": "97739c64dced790b37ae7c068e9c6d53f3e20d30e6e54c6a179a1b9f84453d71", + "bf16-provider/gradients/0303.pt": "3304aa31f19d4af1973f91e58484b35ebeed42f1ccbfa5979be754e9f2e11149", + "bf16-provider/gradients/0302.pt": "ca8c7d81f786b9ce66cb5889306f513d6e03024376bdb3d44acdacae6dcbbfeb", + "bf16-provider/gradients/0301.pt": "787675e431dc8cbea94c09decc3149fca55cb1e70293365770e2214615b47451", + "bf16-provider/gradients/0300.pt": "668ef60e15794443875978942371c0e340578c7991eceb3ee832830ec0c82130", + "bf16-provider/gradients/0299.pt": "6eef743f79d8f93aa8e010987b292b0f4a72c7d4a5203cb0c74f37da09c88ef4", + "bf16-provider/gradients/0298.pt": "233c25284295be84a203fe10833fa9801ab197a468a9b7bc1153e7aa4bc3d8ab", + "bf16-provider/gradients/0297.pt": "7f2bb4b77e5380e5caa61cb8781b66a7b68e20679497e7cc334dbd4e48af6ce5", + "bf16-provider/gradients/0296.pt": "49b573ab96c800002052bf1d470d39ab94810c9ee7a7135d7b03e9f1ce62f4bf", + "bf16-provider/gradients/0295.pt": "f88f6163f48a80473a21d129f619ab42a4d88334f1ecfa74005d1627e4057333", + "bf16-provider/gradients/0306.pt": "796f1ecc30c0895b2645a17d2820d9c0c3c437f0c1629f086869d764236758d3", + "bf16-provider/gradients/0305.pt": "fb7e07b6d951996018c55a750b4de1a7587e50a9b1e8b1ace03d381423b5166d", + "bf16-provider/gradients/0292.pt": "d774e15a385a6e2affe56fd4bf7565b8a62f90a7984dc1121cace9520140d5c2", + "bf16-provider/gradients/0291.pt": "0f295a53ea57aa75fa7607fb78bec5b95247494f29651269b54385f9f3d63f31", + "bf16-provider/gradients/0290.pt": "5e4ae29bfeb14e558df98e786d1a10bc280afa49d5f2ccdb89ded768f3f2bec8", + "bf16-provider/gradients/0289.pt": "0e50704ebdab957ad87686cf78306b329410a668bd177568538388e5cf8c9e38", + "bf16-provider/gradients/0288.pt": "3f5034461985f9f4a3d3c03d67dca680b3693b7bedc4e98abcb661baf72c9762", + "bf16-provider/gradients/0287.pt": "86c0f75adb9d77407f34f834b3123a66113a0e4433efad21dc6d644919897d02", + "bf16-provider/gradients/0286.pt": "35913b01eb81cf78be88754a54e9d3001561f33363145b6df93492f39bef5486", + "bf16-provider/gradients/0285.pt": "afc40a230f35770796db2500ae780c5f3be23bb0ae7f25f77fc0bc64cc807665", + "bf16-provider/gradients/0284.pt": "26e2bec40c56e38cc2ff769a5668677cfc914242de5d996fb799d0bf525dbbb7", + "bf16-provider/gradients/0283.pt": "4417139e2f7fa76dbc6bd9234fab69923b8ec6dfb2ab8006e2d4e97d5c0225eb", + "bf16-provider/gradients/0294.pt": "808a1412bd7e2daef88ea1c267c0135b823be7e1038336a884721c56cabc4bb6", + "bf16-provider/gradients/0293.pt": "76ba8b80c0d1b68afd1d6b058511a7b61d2672d1fb1f265d53d1d750f43d64d4", + "bf16-provider/gradients/0280.pt": "dc6b8ea09a1e08f2ced13ff7112a2ed17c3c137b50a6cf8dbbd962eb0650c421", + "bf16-provider/gradients/0279.pt": "3f29a1ebd7a98e252bad3ff33801d197d9e955445217e878cd5dbe032336cb90", + "bf16-provider/gradients/0278.pt": "9f0d2d5267f3011ac21138c731cbc0099ba70eccb81d11a5c8e2c45e73021c4c", + "bf16-provider/gradients/0277.pt": "8d9315411992cf348e28707747169acf8084122ef7e78839a51aa615ab33dcde", + "bf16-provider/gradients/0276.pt": "2bc86872c0427eef27c8d41828ed65185cd20d43c97faa02faa276b435112a6d", + "bf16-provider/gradients/0275.pt": "d18373120bfc7561ba43364cc7ef2952c1412a934b11bf19aaa43446b6319b6e", + "bf16-provider/gradients/0274.pt": "d03638a9a7cdc8097cdb71710c5f6fde6a4cb0c72e9f904ed29bcc3be379b993", + "bf16-provider/gradients/0273.pt": "d2a9a7ea8eaffffe8c698ca8eec6c4e3ca404eb00ec77d808fc3834c62c3b331", + "bf16-provider/gradients/0272.pt": "6db399f8af1c9925068b79ddd174c9a3a8707db1f6d33b6c89e53ff2d43d673f", + "bf16-provider/gradients/0271.pt": "a186a972fad6db0a157ccb39a88f0abc5926f4129b582090077aafd84f22d74f", + "bf16-provider/gradients/0282.pt": "9ee552757a1496ec847ba7c77445854cf51d5b857f5a869e7c779883e0154605", + "bf16-provider/gradients/0281.pt": "0d0252a452b5a900e0e1cb3211fc88b44f2b9f6aa38b183d55e19a37f3a2b9ad", + "bf16-provider/gradients/0268.pt": "eda4537d480a47f9b8c854ad5da8c51c14e11e9ef79445db161bb87ead56034c", + "bf16-provider/gradients/0267.pt": "a543f6e2eef7495a8ed588339eabc7df314634ab816e08e8b1e430a9dcba66e5", + "bf16-provider/gradients/0266.pt": "e759cff02bd6196fe5de0c1028b9ad2668d467354bdb16cd416708cdb15ae185", + "bf16-provider/gradients/0265.pt": "042bb32b5a2a98e22fb7f4e7c9b27867c70d2eedb1198a637dc703c5cb3703f0", + "bf16-provider/gradients/0264.pt": "35b9ad420c48646419a247e113f0319c3b5b28dd85ce64c009b83943e254db70", + "bf16-provider/gradients/0263.pt": "c7964735b2b3bc3b030a62306d2f99c3ffa91487c92a8fee0dcc4832a7173faa", + "bf16-provider/gradients/0262.pt": "f9ff815ad5e7325b80892fb584703cd6e8735e17efce777cff2e1144fc65bbdf", + "bf16-provider/gradients/0261.pt": "4ad2f2b3328ddfa1828d3ebd55a5f3605231fe5b919e0c5d17594aa0151e5210", + "bf16-provider/gradients/0260.pt": "53f7d3922bbea80523b20f065de6c4d58161a69404bda4010828e119bc14c5a3", + "bf16-provider/gradients/0259.pt": "deb98987e5cdc89e76beb9fc92c52409bf8beb9e3fd66f29b5c557ff944a898c", + "bf16-provider/gradients/0270.pt": "331af297b425fca0d9f996e2afe83661f98533818922f3bc7ef35ecd90c7458f", + "bf16-provider/gradients/0269.pt": "21e1849f0dc4290caa005a924f36dee7aa4a5ca43718f7b2aafa8fd1652f976e", + "bf16-provider/gradients/0256.pt": "d9a3f3ad49d2ec14f3e1c98939daac81d748bd5ce3223209fd0f79c50e4ccb40", + "bf16-provider/gradients/0255.pt": "8827e510654d2727d8d9a24316e2332c30c43d31db69dd0b988df016fcfc5824", + "bf16-provider/gradients/0254.pt": "0e3bfcea4d4857a0affab4319659a6de26be790c102e76a1c142de78a8e877b3", + "bf16-provider/gradients/0253.pt": "e182a9865b8f41da2e50b5c625a502aa4b69afd3a7f3e027f90bc7ac22b6b95d", + "bf16-provider/gradients/0252.pt": "64044c1e4986c3f95008f294d1dd807b27e1aece6ed48745422c13cbd3de32fc", + "bf16-provider/gradients/0251.pt": "995d73480e764bfeca642a3e98aad010e02aeb1f382bb6a43b7968f23becd61d", + "bf16-provider/gradients/0250.pt": "e4d24a090de0e3b44a7dc0a0e7b37eb776cf600ad4f7fbf47f0e56a31ca10e92", + "bf16-provider/gradients/0249.pt": "2b47350687d676b10fe5b477f8448b396276b40618645b866c5292ea907d28d9", + "bf16-provider/gradients/0248.pt": "3eee607af2409ff8291636a53d8653616c2155633d69f07f35a72192721bfb88", + "bf16-provider/gradients/0247.pt": "7b4921e404de60678f39763f3f01dc9aa0525b3700f6a666868175d4ed6d2a70", + "bf16-provider/gradients/0258.pt": "7fa9e2361f768b3c11a05912ffad5db282a8c92627cb5c1328015bfa1ed66b16", + "bf16-provider/gradients/0257.pt": "35a0be6b96d695c4bace48f372f6d6da79d529917241a0fc408f542a92da8f79", + "bf16-provider/gradients/0244.pt": "2edda58756a5909e00df680ba20935aaf4b2560219d6795271c7f07cf63fff32", + "bf16-provider/gradients/0243.pt": "0d5dd5c1bc41d326425e58aa87814b63c524da5202202f65ed38e5f955a6e236", + "bf16-provider/gradients/0242.pt": "5e0d121d7d055e4c666de5ea54194efff27fb38124dc46efdf88379fbc952282", + "bf16-provider/gradients/0241.pt": "97ff757508a1f23ea0a5037c65db9afc86d6438abe1b613ebd036f9d0e674918", + "bf16-provider/gradients/0240.pt": "4c994d82c228bf5a5c39d3eaba7fb8290a4593176396979cd05732a34b02ef2b", + "bf16-provider/gradients/0239.pt": "81390d41da237daaf46382ae3d9281a258228ef0d7fa88b2f1bb90abe1f6a63b", + "bf16-provider/gradients/0238.pt": "cce6e8691c4e078d1bc2f6a660fb3b68cdbbad9cd605a062c43f1f4735021dba", + "bf16-provider/gradients/0237.pt": "ababa08f47c4071fa17302e9029a220d69c1d076dd62680ea60aa78ed560724a", + "bf16-provider/gradients/0236.pt": "5c2b4c74861056484b085664f3be32a3dfeb088f4a792876bdc6474c9faa35fa", + "bf16-provider/gradients/0235.pt": "e0818c2664fe961bda4170d740112c554245340796bfc88af79a2465e0aaa1ab", + "bf16-provider/gradients/0246.pt": "96769ebad77b2012bcf50570d35720984269ad4886376203150a0315b2e77b03", + "bf16-provider/gradients/0245.pt": "856296f7d8d4c371d7e6496c781b73b753d90fd1bb93933a308d9e0c3c426969", + "bf16-provider/gradients/0232.pt": "138e34674452dacd239e811f4d0fd95a03ee33c4e2e4ffc82ddc3ab81fb0747a", + "bf16-provider/gradients/0231.pt": "ef62f36d3acd46a543240d3e4eb2457c078041e0d6d7c247266495c32688423d", + "bf16-provider/gradients/0230.pt": "028f6ec9ac2acb4596391ab04bed596121ba3590e8ea542025b3a7b36ba76233", + "bf16-provider/gradients/0229.pt": "f6a26bdf2356aecce69b55a10ed5f757376a572b4c1bf9adf1ea5b48c60889e5", + "bf16-provider/gradients/0228.pt": "3337413bb7b2476214add605a710ddad5b52b0b884c0909625e6aed1fd2fa903", + "bf16-provider/gradients/0227.pt": "f8e336ec4317ab0959462b1514b00488be782431a251f8226d5aeb327bb95ddd", + "bf16-provider/gradients/0226.pt": "e5acb44c93e7eba97d8766da911f7f2cec1d9a81e8494bc0fcee36ac39628392", + "bf16-provider/gradients/0225.pt": "fd46ba7fe63dd918a34eb099b720e3a2cbea2fe8d2a32ce710487dd7e8825249", + "bf16-provider/gradients/0224.pt": "66feef305f8e73b239261d70d32282cc6a4862fc6389284e667b2f9cb1bcdc5c", + "bf16-provider/gradients/0223.pt": "4c00f3c929248a54acb02ebcb2691f3c98eea5dba89b836f641bfe87deeddd4e", + "bf16-provider/gradients/0234.pt": "3a00acaf15200bda3868ebf8657e48a66d743fe15b4a84ad46b8bbbbe6c3ce1e", + "bf16-provider/gradients/0233.pt": "f3713c4b0601d73011409a8a7786a3a2c5005575040d7b7b67e64c3b9db2d951", + "bf16-provider/gradients/0220.pt": "9b9a49fe27a7d4885a41c31084cda29cbfe820861ceec6931c9fb9f272cc2900", + "bf16-provider/gradients/0219.pt": "43a94aa6cda2a92ac3d0b458f1c6b6aab9f8f94149f08fce0d6290179160492c", + "bf16-provider/gradients/0218.pt": "231c88e71b80ba4b5ee9658309e540e75dbae979fd203d5f91f6570bb8c01545", + "bf16-provider/gradients/0217.pt": "02ea797c01ff789685a3daaaccd55b4e4b3da1bc69afc9d015732890d27d2d51", + "bf16-provider/gradients/0216.pt": "a68db9b9da5c22cc5015d4db7f5eff44ea8048ce7aa74a3d5c1948a3d23c4de8", + "bf16-provider/gradients/0215.pt": "e6a8c4f881f31c148bf2424cdcd3f62a89492385659db4cf638d1768f6c3e7d4", + "bf16-provider/gradients/0214.pt": "6268b4b6c92da5faadfeaf8a5268f20f0783fb127d64f25a10916fab0d296e93", + "bf16-provider/gradients/0213.pt": "d1ab32df4a95afc81a053f3087a05f0c6d8dd88eca14ebac174d415bbbb1aaec", + "bf16-provider/gradients/0212.pt": "accce9d5f93aa2cea3c290073756588ae67a72e875031660a5d711ce8744669d", + "bf16-provider/gradients/0211.pt": "c7bb5e8a8094d7addf1260486a3768ce296f18fa7a3790b44afb734d9dc4c6eb", + "bf16-provider/gradients/0222.pt": "bb262babfe633986374c150d32a6c2c4918008707e38d79118cfb3138d0a4a8a", + "bf16-provider/gradients/0221.pt": "367d15082c9c3443e1d302044f0ba89d48a6c4607c1e75a7e00f2bc66af6f081", + "bf16-provider/gradients/0208.pt": "ec69a37acf4fd5ec8b48896c0984b2f36ff580d3c847f0d3f821bc4ddd399d5f", + "bf16-provider/gradients/0207.pt": "ef7b8ad9b92f94a04799bf5f0e470499ad35e382baa15f4c8fa7e9ef402ac66d", + "bf16-provider/gradients/0206.pt": "7ea4af6d3be7edf6d195cacaf5360ab95b4256f4f41fb4cc10a8ed4bb1771d6f", + "bf16-provider/gradients/0205.pt": "3da5315785f276d61cef3108d748074a5a16cc9c6a94e625a059755113e523a4", + "bf16-provider/gradients/0204.pt": "c860184bf7f61c1b423b7a76c496923b20eaa48104265b04aa3da36f4bcd4662", + "bf16-provider/gradients/0203.pt": "142f97eeb0d5376046d69e5e325a0b8a4797adca669687e6f393a64e6bd168fb", + "bf16-provider/gradients/0202.pt": "1f070a3c00ea16a81f7bb03fdc123837f9a6f72eca06c4bed0e7d74d05fe5ea5", + "bf16-provider/gradients/0201.pt": "583c07b39f496470dabd088eee80f0b2e4b0dcb0f1c5da9a796025bd80303955", + "bf16-provider/gradients/0200.pt": "8d2b63e23ee19dd7f9cbfeafcd61d4d5a4052dc7dcbe1344762c7e4473d7b05e", + "bf16-provider/gradients/0199.pt": "a68321b3801cf41e9fa06723164feb62e8d285986cbb7bb3422c85bb4b4cce1f", + "bf16-provider/gradients/0210.pt": "3d1b6ef3a7257378e4c0745b0e8140ab68b1a83b513acf1507417c74429aaf08", + "bf16-provider/gradients/0209.pt": "8da497487256c7cebbeb3acc10d59bf7c438108fa2d0b75fede9c3c6d482e922", + "bf16-provider/gradients/0196.pt": "d56d202cdd76c0302322836571ab2463c93479fe288bfed1857826694fa42742", + "bf16-provider/gradients/0195.pt": "57b41605efeadb40d6aa9108db0e7a5aa7e1f46e5b0f41551fdff305bf2a7f07", + "bf16-provider/gradients/0194.pt": "5d23277607ee99ce66886822f8914ce125a695c1c4a507d9b94d91ce529f9aa0", + "bf16-provider/gradients/0193.pt": "8b7104e81471686c7793303c31404cc359890fb701afedfdc1f0825e8ff03b4d", + "bf16-provider/gradients/0192.pt": "a972ea28abae2e9606fa3ec4e30bf39052e18aff339b66c2de785f6104e8f4ef", + "bf16-provider/gradients/0191.pt": "cad419f09a97deb1912989597ca98c29148705ebb4768f6ede89a5de7d0a4f20", + "bf16-provider/gradients/0190.pt": "dc61442a3ae246cf85e03efed12725f4ae033237e3b92760b81bff1a61b40b05", + "bf16-provider/gradients/0189.pt": "22bc489135f9d4c1ff0f56d8c52977857b32922900555ec1f5bd097ebd9c1a3e", + "bf16-provider/gradients/0188.pt": "b2614dc807f92b54e0ea8589b3fc30241d572e90b8351d5d1db9e467bfc9cab3", + "bf16-provider/gradients/0187.pt": "daba5072a2576fce5fb57cf32e3ee225e4c323ecb5ca67890b7f0e59351658d2", + "bf16-provider/gradients/0198.pt": "45f6987de952c6c13e8379d76188cd83bc648c7a2a05f4292f2f206338ba2211", + "bf16-provider/gradients/0197.pt": "dd4d032f2330d6d263bb9e8ec402ccf5b8103e5ca453286637f67cf373009625", + "bf16-provider/gradients/0184.pt": "7ffd4d1b9fe71b0f514abe5b564a7bc60bdc08551c5fd76ca643589f40279e6e", + "bf16-provider/gradients/0183.pt": "9975bac895e20c1ddc9937b17c851960c2b35a9e5975e9529a36fd246dfe366d", + "bf16-provider/gradients/0182.pt": "a9c98f1544a7c4d2ec3c412b01250ebb273779d6e7d1cfcd76c51533a2bf1b1f", + "bf16-provider/gradients/0181.pt": "6bc9d0c23aae1a4377728a5f80fe0159283d0cc5ed475e8524e2c6427ab4f1c6", + "bf16-provider/gradients/0180.pt": "c5739199f54a40ce008e8500c1dab2b3bd12b7e4c467796152957c652eb3b800", + "bf16-provider/gradients/0179.pt": "e2c76c1bbe74302a29de725df3c8b2e0f14611df60c94a0a2dd6d5c7d8089309", + "bf16-provider/gradients/0178.pt": "cb668f84ff41f5caab5b717d5ca608228d770e208f842de229d71d4c8b550e69", + "bf16-provider/gradients/0177.pt": "7b3dc7301357f3702cd3fd5694d2e5cd456f07b43c439c9e22a2548d30a2de8c", + "bf16-provider/gradients/0176.pt": "6d7ebff37525d991c1879f82dee3e5fa2dbe4dc6373a0c709108c661b2768701", + "bf16-provider/gradients/0175.pt": "a4eeab0695a22cf78fd777301257a2589261390e1cc6f9a14c4d87f6a883b559", + "bf16-provider/gradients/0186.pt": "8434eeded31e0c7ea8b047c94ff5f164c7b50062b574dc0939dd18cbd248d8a8", + "bf16-provider/gradients/0185.pt": "1c1cc9022678617e00a2e806b8e18dafa8b85e6165fcbf78e0590be822441d43", + "bf16-provider/gradients/0172.pt": "016a36a9d3567779fe23d91f99ff855e1a9af8066e373eef9f23299bfffc5837", + "bf16-provider/gradients/0171.pt": "b0c8b6b6aa50924a8139da67dee1efce799e4a8cba77c994ee4087a15ceeca04", + "bf16-provider/gradients/0170.pt": "dc0221a72663a02aa2c447e624e9f5a81d3c8324d1c0918329b72fa5b9ea5682", + "bf16-provider/gradients/0169.pt": "ccc96b95774d18294008e74868b40e27006ccc8c2bd914b5386346524080a97d", + "bf16-provider/gradients/0168.pt": "505b0a5d3af2787f104dda0b8ef704250f737816488c52fbea7d5d7d18a53330", + "bf16-provider/gradients/0167.pt": "711d2fd41c74af125cec317c7cd3b2fda5722f7572124486e2b4913b94591225", + "bf16-provider/gradients/0166.pt": "9c74031094116729122f0f311dbc62f65f56a36e4c2df75ab6a96afff189d3a6", + "bf16-provider/gradients/0165.pt": "58804c807ce3e933f23acc19a9f7298fe75a921a9ea41bd148db05e1bdcf94bc", + "bf16-provider/gradients/0164.pt": "69db65b8d18eeed754cc5896693ad5f42ef0a1b8187416cb761de299cce77ee0", + "bf16-provider/gradients/0163.pt": "ef62de4a59ed2b7e58a66dd933affab1263387e3798e07cfa3490b84f6e11b83", + "bf16-provider/gradients/0174.pt": "09ede8866a958b2ee9e8126369fef6ce7e4323409fe3c00563d8360d9d42101f", + "bf16-provider/gradients/0173.pt": "6a1096389ef27a462216c6b69f687496093ac17f42981176afb6dc45fb6413a0", + "bf16-provider/gradients/0160.pt": "9bde85c82bffb4798e008201297dfd29dbfa0724eb98d499ee1ade1e9f43a549", + "bf16-provider/gradients/0159.pt": "1c559ae77b1bd082da5e9c5e54f2b50c4e58f9e1b49fc37f2ecc303642d7ddba", + "bf16-provider/gradients/0158.pt": "018d969699505bdc6bf01a63d9d49a24004ca43d4e00558fbfede39ad7f7731a", + "bf16-provider/gradients/0157.pt": "9f18ab851c0c1d9c73813dc72daec46bc826635fc0818bdf48df0295b8386e81", + "bf16-provider/gradients/0156.pt": "cd2a439b79d958254592d7f5960001b3f3e6ac4bd9930f7d7774a57f06b957f8", + "bf16-provider/gradients/0155.pt": "c4e09428b8557906d650939be1935fa41c04757640720cd1e8d82351d458821d", + "bf16-provider/gradients/0154.pt": "13b9ecfc3c2af43fde4b48e02654633c6fe0ddd7423190ee23fe785d46c90e56", + "bf16-provider/gradients/0153.pt": "29c671db274b7200b9a57226cc85ab5241cd7c1a30a2ce5c9461470c27484e17", + "bf16-provider/gradients/0152.pt": "3a7377c5348ea6165e9b7b82683e7a84d73e3047bfa84275598c7c2135abd8a4", + "bf16-provider/gradients/0151.pt": "593708635544f6403cf94e36c59324a516cc8cfed11530c16da1dc54ab90e5d4", + "bf16-provider/gradients/0162.pt": "afae5321af111190bbe3a27f8633ca05f7a95b7c195178682c5b3d18062eda0f", + "bf16-provider/gradients/0161.pt": "15bb8d7e4ef8e2b4bff6ab0ac4641aec9faf6f9e6bcac66c62cba59b851315b0", + "bf16-provider/gradients/0148.pt": "e23cd4f3ae6d0f9b19affbd1c48dc80a47d16ef05c624ca126577e0d0495cb7b", + "bf16-provider/gradients/0147.pt": "85f26343f8242e0ff98a6ea9faf60cffaafe9076d4e415476121ae9c28547a68", + "bf16-provider/gradients/0146.pt": "815508d3a5721fb5850dc275553142796cc47f9fcbcc3c6d97a003c66d95df28", + "bf16-provider/gradients/0145.pt": "a759e15ec38a04639058443c47a2a1fa5660eb0d439524f6b986ee652f6936fc", + "bf16-provider/gradients/0144.pt": "acb1925dc915c8cc874d66abd51c5e944b98c629dd334639ff34cc0a97ec2245", + "bf16-provider/gradients/0143.pt": "10da5b0274b175a6223cf48b88ddf37ebf16e0083abfcf7c1a7e96b948f26283", + "bf16-provider/gradients/0142.pt": "162739213b2cd276a285823e42b40e49b737108a5802cce1132cb3f5b812222d", + "bf16-provider/gradients/0141.pt": "cd4c8c4f39356449b76bfa1ca2e7126f6b5e843345ce221cc34cc590524dea3b", + "bf16-provider/gradients/0140.pt": "ec04f8fe998018da5fbd1820a5856d53c690fa17206212057a59e33f321d935e", + "bf16-provider/gradients/0139.pt": "53a5cd74f51b7995c2c0418ae70d94c27d24a05ef83e5dcd1a27bebe5471770a", + "bf16-provider/gradients/0150.pt": "07932e7ee2dbc633b8ab0f8e213bc5f793058f7b59a4b8ebd8c6d5f98485e2d3", + "bf16-provider/gradients/0149.pt": "4f66852763363e3cc77cdc62cfe13ded71b03f9b607849f32443b50b5b520931", + "bf16-provider/gradients/0136.pt": "1d712a653fcf0981214b1924d30d3e8c8c04c50f57e075d70c39aaf65df3eaba", + "bf16-provider/gradients/0135.pt": "c75dc1ed5222c4fa46a564d6032b151bdd36c46cd789f9d75593bbcd14ffb6a2", + "bf16-provider/gradients/0134.pt": "826cc623be29f797ecf4546c34870dbe86bcad0a20b6a61043c7771851153ede", + "bf16-provider/gradients/0133.pt": "1eae2273a5afb2b3acdae724666600f103af57b3d072cd4967bfd0107442d840", + "bf16-provider/gradients/0132.pt": "534b45c8e194ec24845c3ab685390f5d6a8194b2058c85f1e2f91f4bd5210312", + "bf16-provider/gradients/0131.pt": "087678ebe2ed7fb82735eda24a4631aa309a2f40546753fda03d6d691e1b6d00", + "bf16-provider/gradients/0130.pt": "b33178c0f5a3921fba3d9adf920f222a4d20c33f1aeeb206975bdc258c92dc98", + "bf16-provider/gradients/0129.pt": "b701fbc13ec95e85222495fec6c6028b1a1e93aff213eca68c3e7105e450b6c4", + "bf16-provider/gradients/0128.pt": "4c7621976840a1f8459727e71edd47c42ca19f20475c49597355c25344e4f842", + "bf16-provider/gradients/0127.pt": "274dc09a775da8f4a61b32173d2178e8df21ea1b83bbe765c99baa2714f64870", + "bf16-provider/gradients/0138.pt": "877c7a18047c4e3ceadaf73b508d90877081685c75ba1102249a6939ec1d4c25", + "bf16-provider/gradients/0137.pt": "18b570db82f4deeecbf75fb8dc029cca837e6a8d0da23a5fe45f73b663e41892", + "bf16-provider/gradients/0124.pt": "cf19e165eaf709bfe0359d8e486e78497e2b0e00db7c024df03fadf42ebea5f5", + "bf16-provider/gradients/0123.pt": "ac8722337c50c2fb9798afb3afb098e71fc5bf036dc081b6479525d12481959d", + "bf16-provider/gradients/0122.pt": "774e8af6593a0dea7dd67593f6485c293625ff5645fcc251f25d841ea572d7a7", + "bf16-provider/gradients/0121.pt": "9161e0f787bfcc9cbf4e1108de217bc77d25d37811eb4b641c95fe4443e047de", + "bf16-provider/gradients/0120.pt": "89e5bbe7aa9f2dd6dbc0bf76c5e405fb784ced9c68b1de06133a43cc59a65fd3", + "bf16-provider/gradients/0119.pt": "55b11262ad988c34f26974d3cd40a86ae59543f911c8c591517ee27f85efd640", + "bf16-provider/gradients/0118.pt": "0a098c417551c7e109535459167c92d26294628cf3fcbca6f3f17cc22442c2e3", + "bf16-provider/gradients/0117.pt": "fe8c8212c7b9673111ec777f46013df5fd477d8bf00b86c7e4bdab7bfc5e2c93", + "bf16-provider/gradients/0116.pt": "8c3033828275a2771b6bf640047642bebc6c05c0f7369723e8e9e2195dca89bb", + "bf16-provider/gradients/0115.pt": "31cc3f5422a20c671f2cce549e93cc8c9ca1da510f16e46a17a3b4e5aa32bc63", + "bf16-provider/gradients/0126.pt": "3d5e9096ac850b1f195a13943969b7f9ee28a241943c1f70850473ca4672d4d0", + "bf16-provider/gradients/0125.pt": "144130e8520ffbcca2a2ab0f41f6f84807f181c1d1077c3435206965b75b53ab", + "bf16-provider/gradients/0112.pt": "5c7af8d04a634bededab66333074d7626ad411ac306d693dc95ed6def0459cd8", + "bf16-provider/gradients/0111.pt": "f59db7383925bd7ea207b917235a06ede25592c747ee02a38c721677a6003314", + "bf16-provider/gradients/0110.pt": "ac80f3091f2725f92ba02eeb582c806eb863da016ff92af423bf777814181f0d", + "bf16-provider/gradients/0109.pt": "f0a2afc6c0fc3844b3529e2be8ffb7255a2ca864a8ae2d0a9e3f669a20c14ca8", + "bf16-provider/gradients/0108.pt": "9bb5c3eb397301d96fd1490756031a1f8d689f9b6242e5f1d6de836a34292d1c", + "bf16-provider/gradients/0107.pt": "7c613e45624c62ae2e08c0c477923fd46fdb9092cb733edd403976c4a04b4b6b", + "bf16-provider/gradients/0106.pt": "63c7b0a3b11ccf9e9d6011884d9cb36b6284bbe67d160f14028fbfb1de61bbac", + "bf16-provider/gradients/0105.pt": "aecae0cef955e9e1e1fb66d901716569c8e5d74dd3db01eee052e81ea3c6dc42", + "bf16-provider/gradients/0104.pt": "d6d8d40daeac243f6fd9975d4ccac10bcf30354f35173a562ba6e05f6ad6a7a1", + "bf16-provider/gradients/0103.pt": "0db7dbf32f62f6659ee1205f91b81ca1805b529ae75e725eb5ca7b1909ff44cd", + "bf16-provider/gradients/0114.pt": "6da9b3527255cd82ac91ea421bb0802a9474902783191eb0f5d65b0ac062dde4", + "bf16-provider/gradients/0113.pt": "10630f0a39d2d0abfb49aa3ac08d173f34ab6eb11ce15bf3672f207a887ead14", + "bf16-provider/gradients/0100.pt": "2832a9cc77db09443da100feb1553606e3ea625218f5ef4574724c909cc54a9e", + "bf16-provider/gradients/0099.pt": "5e37cd9e586a650413176939381063b26530930fa26daabb00e60d50b1353d2c", + "bf16-provider/gradients/0098.pt": "8bbf8b390b4441065205c70a7e5e8b011bd1b8f35f5f2da2a827562190c3d9c7", + "bf16-provider/gradients/0097.pt": "6ebc32a32feb456765d1e3ad7bd96987e6ff8cad906c1b7d3e352e08e410f711", + "bf16-provider/gradients/0096.pt": "4bee3136156a423a43b4ff236a72d3a2a56932233ab25e89a984bbaaf2c8e4f4", + "bf16-provider/gradients/0095.pt": "bb34ac908f20a239d98c1f5c0416d91cefa95f5c105ecd241f24de8bebacfa3a", + "bf16-provider/gradients/0094.pt": "17d4322dd958c3e096b66230312b58cf14646e24b500234e3c6f302b70f2568d", + "bf16-provider/gradients/0093.pt": "dd70dcf21fb8ccf93a51991e89415380bf421b2f550245076f7fd4c877d786b0", + "bf16-provider/gradients/0092.pt": "904a0e7949f730a8257cbb4fb7afe68e05bf491ad40bf36e51e55aacf2438eee", + "bf16-provider/gradients/0091.pt": "15ce62a31c2871e75845a26cb32f8668c13c6d544e0ee941e090eddf48af1f41", + "bf16-provider/gradients/0102.pt": "2105ddf33f8b5d864df4e9c8f005099fda2d241cea290e0c8c0edc7f7f9d1f44", + "bf16-provider/gradients/0101.pt": "3f6b9693e70bc787745771556135dfed5ccf4265de5eef6da488622de0517577", + "bf16-provider/gradients/0088.pt": "2cb9cf8ce9dd39c7f7f57f4ad225f9a5bc99e9113b243095905f5edcedde9cbd", + "bf16-provider/gradients/0087.pt": "435fa661238b6c1b4551c5491d18099d7468149a4d48a1da7a45b3cdd8a116e8", + "bf16-provider/gradients/0086.pt": "b9d949c5cadedcfe6d46f8192de6f307fff9f7b05c9faa0d73b165f349f8c05f", + "bf16-provider/gradients/0085.pt": "4c2a4273bd9defa879591794a590eab1242047245433c618cf23e95f41cb87e8", + "bf16-provider/gradients/0084.pt": "f35f309d595786a6e929928b5b50c9a93bf0a91d8b9a2e744836ab53978eea6c", + "bf16-provider/gradients/0083.pt": "166cfdaf62b8893a83d92dcf8bbdd162edfa88bc48f0fdfd4d2e5aa8fab1d8d0", + "bf16-provider/gradients/0082.pt": "143aae86d17c6a1935980f0b2d1eff88d1aab8c942d28fe982a8091f3f9a0031", + "bf16-provider/gradients/0081.pt": "5f473c39fc3374181c611d1eb265679ddcdda92ca1fab207f75a4cf5762bb459", + "bf16-provider/gradients/0080.pt": "4d42a16f74698f5bcee7ae76b980878506cb23b4b8f3801bd1f4bd0cbc014ba4", + "bf16-provider/gradients/0079.pt": "c61dd2aa8d7b9c462fbb35803ca2232ebafe460a83147eb4aa8db48922035b2b", + "bf16-provider/gradients/0090.pt": "0141bc5cf23698c2f891fe0824f2f7d1bea7400b01f1ffb38c908574260fe5b9", + "bf16-provider/gradients/0089.pt": "938b32ab210aa1ce7cff95866361c0230371d44cf54c67dd945324613926972a", + "bf16-provider/gradients/0076.pt": "293d60dffa0f358385c13b7e3d456b5598dd9966e3fd121c6ff82455c318f374", + "bf16-provider/gradients/0075.pt": "f36571750844e76325093b7c37ac3bc916ce5374afec6f260b7d046c03c3a70a", + "bf16-provider/gradients/0074.pt": "9038d73341d28c607d92439ff89ca0bda9c372953ad4c6ea5251e7e748c9cb09", + "bf16-provider/gradients/0073.pt": "339cd2fcee85b12299e2cb9548f33ceaff90765e4eedfb086edb6c13a9ecd65f", + "bf16-provider/gradients/0072.pt": "1ab591630cbed9bdc50b6efc206608f81ff91db1398c80174d6b941f181ecdea", + "bf16-provider/gradients/0071.pt": "d09d2cdc8fa64016333be8b54e51b169944e4a4a408e216152760d79f2ada1dd", + "bf16-provider/gradients/0070.pt": "5d91967d801a1adc4cfcc5a317fad0c0fce6aed02f7b406ba30883f493b9d33f", + "bf16-provider/gradients/0069.pt": "d793217a0b139aa0130f525a81bf5ca4d70b91ccf920782d2f6d0ed49e199640", + "bf16-provider/gradients/0068.pt": "f4d4e7601a3d80129447decf647ead76d36587b33ffa68cf135b54c68395a202", + "bf16-provider/gradients/0067.pt": "2d646468bfd531007d63d06c684387e9a86b9c634746ba8c20790a9e348f58c5", + "bf16-provider/gradients/0078.pt": "1b8c659f0574cf4febdf7f8fa8df027d81d411244705639760530716780647de", + "bf16-provider/gradients/0077.pt": "e48b2388258704879784d4b0bd4e589e85f84285c333fb04d0d23e3c1af332b6", + "bf16-provider/gradients/0064.pt": "b7d99f27a3c4b724fc6b7cc74903657b04465c0432e52c92121ff65ffef82eb7", + "bf16-provider/gradients/0063.pt": "4133fb7012619a578275b9011dfc5eb37c3c1619d22a18379d69d7808f7d3e08", + "bf16-provider/gradients/0062.pt": "0fa4df1da173e0fbc15ede70c252eed208b1b3e6c1b4a6ef77f064bc160be685", + "bf16-provider/gradients/0061.pt": "39b847e7d283acdd0acada980458ea5272dacf5d9a89902228ae58106f596de2", + "bf16-provider/gradients/0060.pt": "97872df7510332413337eca50c394132423e2c26c231339d1328261332892216", + "bf16-provider/gradients/0059.pt": "152c079726dd0ae673badbc4c878db1ba144bbdffa781947a9018dc3a901530d", + "bf16-provider/gradients/0058.pt": "17996bdb44c46cdba536edcb3271ac25a6e914c11f6f8b9392702475012e19da", + "bf16-provider/gradients/0057.pt": "e6a2d5f495f9dfa98402fcfe5292a85b31927351a00fb0aa0c01e139f21ed7c8", + "bf16-provider/gradients/0056.pt": "6c7eb067b1ef06be780b031196b052f2065f7646798962af6ee5910251ab814a", + "bf16-provider/gradients/0055.pt": "ee511ef6481d96e4178843075ea5d3c65d5fdb5a0640b3fb80b83215a01d79b0", + "bf16-provider/gradients/0066.pt": "71fa3630d746d844ef8332249afde4763182fb108b603d90d32807e60ac085ea", + "bf16-provider/gradients/0065.pt": "3a5690dfac45adec64f184ee74a2e514bd435c2ab6bb7eea078e924e7b3b9975", + "bf16-provider/gradients/0052.pt": "55f1a640d777a587dae8ecba98331694099dcee1ec05f2801e3e34a890c675ec", + "bf16-provider/gradients/0051.pt": "10d6da1d756568758347c08247e50184ada2aafaf5c3ddbdf73d661d672f71e6", + "bf16-provider/gradients/0050.pt": "56bed646e4f450fa1d6f8a8bd5ec45484a69c874efb7ba9b2fd7d6149e1b1318", + "bf16-provider/gradients/0049.pt": "dbbd40845e50516541f796388f91efaaf805d1594880bdb63292f97bd6a9f025", + "bf16-provider/gradients/0048.pt": "2bc9cb41dd36ffe8d4681a45e69969a29b97a30ecd717cf1fe6f776dc239a76e", + "bf16-provider/gradients/0047.pt": "4c14231efb114fea51cd04115a861cb9d1e3acc56257c84cac2c3a8d3ff63a95", + "bf16-provider/gradients/0046.pt": "1f6dfc8143744f1bdfc3b6dd2a42e041ec1dcf863369a02b4431da4c2b0e0b46", + "bf16-provider/gradients/0045.pt": "e42525b6185b5070da925ba8ecea153090ba1d08edf3f9d89f0bd7dddd105bf7", + "bf16-provider/gradients/0044.pt": "911f5953d0c2a5f4746549c1e9703f1f596c02496198d2831797a3e37c4f5d71", + "bf16-provider/gradients/0043.pt": "44d0fda8dc1ce5e177bbd028228a60f3c71fe96689c0811ad700476ebec370b4", + "bf16-provider/gradients/0054.pt": "1ffec37aa9359996c99b32f7d9c785d89bb107806e8fa1ba3bbb1f5d6bf60402", + "bf16-provider/gradients/0053.pt": "d874a0151cfe590905adcd4b089cb0cf8d0193ec4073da7dcac82b3f3999eb6d", + "bf16-provider/gradients/0040.pt": "4895de0e001cb9602018260946f4b3c5f508d526aaac4659b8b18b30e677c441", + "bf16-provider/gradients/0039.pt": "12eebd64cd38985f12a7f1d39ca002b35206c2d37a00c94846209fe96b2b9811", + "bf16-provider/gradients/0038.pt": "44d110c789407ada35b4b984fed865ade732f26d79c78762afcba528897ed0e1", + "bf16-provider/gradients/0037.pt": "24fb83adcaa5f42f51ea86c546221a85ab7937605150460652459d44c6dd29e4", + "bf16-provider/gradients/0036.pt": "ba28120bda4a87b324e20c4e03e9497e9df18e489e1b89df82eb5b9233ec5537", + "bf16-provider/gradients/0035.pt": "c5f967b3ef117e01c047e6b578deaf86154b44619055e82f10378833356c571e", + "bf16-provider/gradients/0034.pt": "fcf2313216dd2b8fdc8ed31a389b459d8704e281685f83b8be8893cbbdf8236a", + "bf16-provider/gradients/0033.pt": "48ed901ebc30a8eb15c8aa00bec6bee7fd589c70031fa84227a34bd830dc51c9", + "bf16-provider/gradients/0032.pt": "9f23a16e0d67d0d5781e55fe6d90230f2c7a69f668867d084dc064843f861e6a", + "bf16-provider/gradients/0031.pt": "89b10eef54ff52932a865f92174f3d97306ced96759aed3a069e8c11a6eee8ac", + "bf16-provider/gradients/0042.pt": "080dafd3d8a7c0c8077931a5a7f3a160e4b1add36f84cd70ce7d359b0313853b", + "bf16-provider/gradients/0041.pt": "07d1ba8a4a7d26df9da7ca67280d374596a320739659b21500814befcfce439d", + "bf16-provider/gradients/0009.pt": "de47e82f2f3500d93f56e4c028f0a4fa8d2790f4e62479754dd2412a99fa7b7d", + "bf16-provider/gradients/0008.pt": "1dfc2144fa6d768523cb6898d6a20b99cca90b46e3515ea4fb1ec738310a5b74", + "bf16-provider/gradients/0007.pt": "17f995c829b58ad1bb6a584222727be244251458f564b7fef5c7dd88d4c198dd", + "bf16-provider/gradients/0006.pt": "62bcb914fc74eff342f30a65eb2cf3d1509d550912f1bc58bf684250b7017534", + "bf16-provider/gradients/0030.pt": "e88add9ff07e2e340f831ddeb0d6cdb7234da718dbeacf483e96fcddb340b3dc", + "bf16-provider/gradients/0029.pt": "565007293ed7d6ccd577c0879757e489f63aa81ee5846f268a55831610b14adf", + "bf16-provider/gradients/0028.pt": "ae8bf5675c319d904ad3a47466ac56616eea4cd76dcad9f21a860eaea6702cbe", + "bf16-provider/gradients/0027.pt": "5b3b2dba09b7e4811283d4797c97eef57378dfee7c4f0bb0bc89b3732f3c3d79", + "bf16-provider/gradients/0026.pt": "ed74b7bca7380650b552056563196958929d6d130ebff126550df717eef15123", + "bf16-provider/gradients/0025.pt": "d15c5a45bb7858ce7a2033573d3bd8d15b3d502db00ddc9335ac2880d5f20ef6", + "bf16-provider/gradients/0024.pt": "6970eb0e2285784f4c5495009a87dfbb36ac612c5570e04d220c7484669e727d", + "bf16-provider/gradients/0023.pt": "de6c1a0c38b17d5054a85664dd4587a942e036305754bf3796bead97ce2b9957", + "bf16-provider/gradients/0022.pt": "55761c3189e84cc1bc2de51c4f270e679ffbd532145994deaf253930117d6e7d", + "bf16-provider/gradients/0021.pt": "b232701911cda69913bff55f8c305b3d6da1b08f3ddd1eef63c1b3883a4e2313", + "bf16-provider/gradients/0020.pt": "b113fe0092d38ff7c209ca0cae8bb438e437e1d9e80cc59b7f1933a84fdf6be1", + "bf16-provider/gradients/0019.pt": "73c9019f4869b3233e55b8d3ddcf641f27a2c398e6d843cf50040b6c0215ffa8", + "bf16-provider/gradients/0018.pt": "94a2d5447adf7806a9b1c012b4ba95e180dec535f9b658b8fea5c160dc108e0f", + "bf16-provider/gradients/0017.pt": "d721ce7e8891d57d9a5e34c06211d2e3e5aafeb610fb4b00c23ed21d61a17591", + "bf16-provider/gradients/0016.pt": "6ffaf1248e2c88315943bdfbece1abc8e24dc655634d801a6203ea6016572e85", + "bf16-provider/gradients/0015.pt": "d1d638a00f6a41657f8073129fd3f7c3d57cf58915aab47a24d79f66f69a6aec", + "bf16-provider/gradients/0014.pt": "b12ab91a2aa7392166ad5d2090a26a3ca1f62b98ccdabab9594d685806ecfa42", + "bf16-provider/gradients/0013.pt": "ba45fd0ff6f5600ecb2a9ec6babda8836c863225430c5430d5b78c16b0ccb75a", + "bf16-provider/gradients/0012.pt": "816a2593b6597196ee735b7909ba63ff33f8003e98e0411a834daad804333047", + "bf16-provider/gradients/0011.pt": "a0b24512b80d5a7de9c1c5d636fc977c263d9507d30352e537f4c4af04d5b8ac", + "bf16-provider/gradients/0010.pt": "aa8e3825d6987a4ae2ebe5af3e3d7326d2102f1a98ac4fc4516327897e5b6a7b", + "bf16-provider/gradients/0005.pt": "5a661a8f6aa437196011babdafe14d8c754ad2e4be75900786e6507cb134b700", + "bf16-provider/gradients/0004.pt": "4beae74ac56b4eaed53c2b17a7d31f98d045e64e01d067d5067851304a0efcb3", + "bf16-provider/gradients/0003.pt": "b3427ad172c897d9ca3f684d70286d98f99f9074b37a369e7aad5e2fb69eb3df", + "bf16-provider/gradients/0002.pt": "1c56aa9e0a604d21e52062270e0c6ce5e319b79694986338da13dbaa8dd81f08", + "bf16-provider/gradients/0001.pt": "f575932d4d780dcf5da597987286d0c9cccb1bc589113a38d0a0b010b56a5e87", + "bf16-provider/gradients/0000.pt": "369cf6e3c0a16b175a5eb8d95e43b339b96a1f7f01b1b765f05eb8a38048897a", + "generation-reference/42/report.json": "e55e877570fa3f089b158be76a4423f8b25443af5c85fca45ba3dd378974c123", + "generation-reference/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference/42/audio.wav": "373efac29bd3073f1dfebf0cba30c51848866351fb6d2c8e55921047c82cd296", + "generation-reference/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference/42/video.mp4": "560d134d198e725880f308be6b8ebafd71718d01f323ba7b5f3fbee379302d01", + "generation-reference/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference/43/report.json": "7e181902ac5984cd04c8773a2c0bae77ba10618b63b0e5b2a3b56599c4ac5525", + "generation-reference/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference/43/audio.wav": "eb404d2f89fbf3eb9795ef2704c245618e47e552a3b4d782b52602762569de6a", + "generation-reference/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference/43/video.mp4": "14dc8307aab50e41849f02e3726612a5e2b5d197dc04987ae9196d2b63c2e5c7", + "generation-reference/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference/44/report.json": "50f571441bad710594cb53fb861cc73fd6b8b9e81a2fbe82de1f9dea4feaddfc", + "generation-reference/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference/44/audio.wav": "0cab0190275bb143bbd997f50928cea3e2c0208f90cd71b9633046ec6e4ddf0c", + "generation-reference/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference/44/video.mp4": "e68e18a0fd4a1ea9c5da84969c70299d9d21db6ce9aa8592bd649f84c6b9b595", + "generation-reference/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "generation-reference-repeat/42/report.json": "e1f254991fa8c8e72f2dffc4e30a056d38e9e302fd2f684d8c593895afc020e4", + "generation-reference-repeat/42/audio.npy": "9861a472ef9c8c58d6cc5a2664685b43ebb0d0e88dd6551dac28611aaa0e41b2", + "generation-reference-repeat/42/audio.wav": "b9891b3a34e097361f06e63a9e890dc4551600a9ec575a08dc07ef9493c19fe8", + "generation-reference-repeat/42/frame-000.png": "ff73cab289183fc6590cdde6e051ff377f9ecbfa812d42e8d18f2dca110a4cff", + "generation-reference-repeat/42/frame-062.png": "f1b56ce226dd3291a4f28f985b3863d52c56da5adf9291caf0d140a8414d1107", + "generation-reference-repeat/42/frame-123.png": "7d2c577c5d0b31ae20d6e67442b01d5b0029c80037a1db39aa5f6e971117bd48", + "generation-reference-repeat/42/initial.pt": "287e2c9ba5dd5c3fb760db8343ef8c0d55cd042c823fc34bc2d4b856d3d7b79e", + "generation-reference-repeat/42/step-000.pt": "b56edd8630d8638a412f97fffc75b3fd79cc4e75479bdc38eedc2c42530be4e6", + "generation-reference-repeat/42/step-001.pt": "7dcd9eb3bac5358290534b3313f189848655209c7f57a4032e5dd95338f7d2b4", + "generation-reference-repeat/42/step-002.pt": "cc4661b4b188dcd3e8f7921cf69ce98061cf269baa2486c81d93cc6c3f792c6f", + "generation-reference-repeat/42/step-003.pt": "e8d4a33633da852ecde3b11b946016b01013ecdde772d02888ed12d112e72af9", + "generation-reference-repeat/42/step-004.pt": "b052400f11dc1654b8ea543a5d1cc1cff7d3454ac66761eac1ffeaf30c978973", + "generation-reference-repeat/42/step-005.pt": "e9a9a75ed4bb103dfb65361222b01595fa8815faafbf050bfc82fa01c8562b09", + "generation-reference-repeat/42/step-006.pt": "b682b9852a4e38e08f1877084d7c36f4da7cbd8a5ffe487c7ccb8afba34446f9", + "generation-reference-repeat/42/step-007.pt": "2113018dc67dd4cefc49b7d043375f07b6b18266faf29e072ebb3775beba568e", + "generation-reference-repeat/42/step-008.pt": "fbc231e9c11996ac287eb063c232dfc6c4705ca1f54f50025034403091ff0a8d", + "generation-reference-repeat/42/step-009.pt": "1033bf656d2f82f2ae4d93ea7b9e71ad581de0d019ae0459f01cba56ed09e350", + "generation-reference-repeat/42/step-010.pt": "85a9e8f0bfa1338a5f4ace82a845075a960091e712804ae6ea6ed1d4c67a88c0", + "generation-reference-repeat/42/step-011.pt": "53de40d1408176c6bf996e2574915a1eec4bd64db1fdc72199195c98839a53f1", + "generation-reference-repeat/42/step-012.pt": "994d4d8f680c29552429b086794a4cdc746efcf0a64f8e1501ebcf92cbd79c82", + "generation-reference-repeat/42/step-013.pt": "8bdb4c48aca1a9c04838755d61a44337b5cb0706a028e51257e5682f8d4a0998", + "generation-reference-repeat/42/step-014.pt": "6196f0fa34815f28e19ffbaa459269f9ca73296a802fe36be69623f378664806", + "generation-reference-repeat/42/step-015.pt": "33ae1dac3c11a2edfbded966645819017e0c58689b50cc52e9edfaa36ba9da2e", + "generation-reference-repeat/42/step-016.pt": "f4bcc7d6cd2cf962b24216f6ddcb880f26bb669fe8a41b337fc6a01d5efb4a1b", + "generation-reference-repeat/42/step-017.pt": "795e90deedc17502edd183dc96abb674f4798b30e7b6e333aa7497b74bf10b76", + "generation-reference-repeat/42/step-018.pt": "c2dc7abecf43b225b7c7e5e550dfa3749347bc895672f16b3825e49c1055d556", + "generation-reference-repeat/42/step-019.pt": "99be79dcdc48fa0c3634b1e454ec3d860e164672f26a3ff535ed092e7ffce465", + "generation-reference-repeat/42/step-020.pt": "20ab02d76387107a8407621819abb7dbd06d5e49cf9fac6d49b37a603ce54de2", + "generation-reference-repeat/42/step-021.pt": "c7b7081913001681eee3efdd9378fd6ef2f4052082d54f2e6799cf439543e7ac", + "generation-reference-repeat/42/step-022.pt": "d960567c973e7d8cddc036c09c6516aa2f0bbc3f44a0a401cb3875370c5ba531", + "generation-reference-repeat/42/step-023.pt": "abd38eb7e2e1b745776767ff6fd81ad4b3ef16a4e6fa4ce88488bd9c6f07abdd", + "generation-reference-repeat/42/step-024.pt": "57b9f578e76a371ff0b0e32171a0cb8bc26cc753e5919e459f3750c94cb00846", + "generation-reference-repeat/42/step-025.pt": "8a831855bc9951a5e1864b767b60857e5712c1d893f73b960b7ce434940a3577", + "generation-reference-repeat/42/step-026.pt": "39cb28c468015fdc1646fb20a4b31243bc3e765b3d7c0908f741cadd6648847c", + "generation-reference-repeat/42/step-027.pt": "ee7408d62850bef88ef29f4ecb397b1d25437a3907ada3f7e51a7366096a59be", + "generation-reference-repeat/42/step-028.pt": "c1814b768cdb2aa4d0d92a4fe6e9d3a1172677efaa1576e2f3189c3f174945ee", + "generation-reference-repeat/42/step-029.pt": "4b9bec6be67576a86d4641f3f62e3cdb2f9a62cec6b97012f3e6604645da30c2", + "generation-reference-repeat/42/step-030.pt": "ae39379b8b7a0bdefb69a2f9e621404d2cd1fd78215069ae42b0c82253e42fb4", + "generation-reference-repeat/42/step-031.pt": "f9a8fd41bab27657bb9fbca79bbce6563cac339869b0576f04b3641f38e01bc2", + "generation-reference-repeat/42/step-032.pt": "c79a65580d3dae0e0ffca66fbaec9a1d0f7df1cdfd99d375daca09145c8cd6a9", + "generation-reference-repeat/42/step-033.pt": "1c8f9fb5a19d3f4fdae293f03883fea0ccf69773e4797194f69a901c5ea3fb5b", + "generation-reference-repeat/42/step-034.pt": "8ea65bf600017b2127acd93927d9cca8e17065b676bdc313307a6e001791a03e", + "generation-reference-repeat/42/step-035.pt": "d02dd61c201e9f7e772126364ee4ff27a02ecd46b7cb913391e6784cc8fab5b3", + "generation-reference-repeat/42/step-036.pt": "a1f554effd959e5b29ca553fe85476ce9af4fca7af6f3d8bc60da6130fc725e3", + "generation-reference-repeat/42/step-037.pt": "fefbcafa6f95850770ec913386b56fe5d4d05f45bd9ba4e1980a83af62ca6fae", + "generation-reference-repeat/42/step-038.pt": "a993216cb68d1a5945005a89812392e5f4a3d749df2cb554010e7771bb83210d", + "generation-reference-repeat/42/step-039.pt": "7d4fe8032d4b6e014add67fcd8494377605029cd0d79dd0d16203c4350e8055b", + "generation-reference-repeat/42/step-040.pt": "e857886ed4a6c5ac04b04d67bbbf8fd028669fdd6f8b2f5779bb6741988d5cc6", + "generation-reference-repeat/42/step-041.pt": "7d246e8a0f4565a09a55239d9811187846a914c060c3e695cb06251931da5d64", + "generation-reference-repeat/42/step-042.pt": "d9fc437f5b867e3e07c7ba04a573bfa43b512b75975af10d9cc05f1672d49143", + "generation-reference-repeat/42/step-043.pt": "a48098517f70da1d338ea774efc47ac1cef943f78db24ceef66c07eb0b85a3cc", + "generation-reference-repeat/42/step-044.pt": "62276b3b774c8ad428b82f1a868004f48c6d80224e8aa300e6c5ad4565b8baa1", + "generation-reference-repeat/42/step-045.pt": "7e2e5ba2d4206146bed14c9bcbfb98cc9d5c6931630f0f7e0b0f89f357bfc7dc", + "generation-reference-repeat/42/step-046.pt": "5f21d828458f092ff63e2cac8b031575e3d58c1c67539057b01e040601c910a2", + "generation-reference-repeat/42/step-047.pt": "f9b5881172bbfd966c6880e2d8f6fae4bde6381fbe8397d13895f71773293dac", + "generation-reference-repeat/42/step-048.pt": "70e5c1209285253ffc94d52fb1649ce3a5fbe6b5731b3e7c0626beab51f1deae", + "generation-reference-repeat/42/video.mp4": "f97c122eada187d8805bb088a2326774af2524e598091446c1302924ee043d9e", + "generation-reference-repeat/42/video.npy": "0638e202fcb58d0d4cd428e600758e8f0d8e60aa7b64aba5e234240dcc2554b2", + "generation-reference-repeat/43/report.json": "914a6e564097b23305d6bce5ba4b06a30603d5dd3c0866e3688f3a8a8bd76987", + "generation-reference-repeat/43/audio.npy": "fce1f4b53406f651cd707291c34db7ef21bffa7711d996d732a2300dbba68d09", + "generation-reference-repeat/43/audio.wav": "1cc055d49792fa9760701ad092f1f9e9dba691503ca7a4205702fc3e0c6d63e1", + "generation-reference-repeat/43/frame-000.png": "548f629a956b028568430820899bcaf373b06035dcab83bed10076c0ac113696", + "generation-reference-repeat/43/frame-062.png": "0f08c8cf856d14f9f593e5511c07a7695f687a198c0b25c57511f9e123e55011", + "generation-reference-repeat/43/frame-123.png": "50a3a11f302e3c491e683f4553c6e21f131a3433a4faf589d7a6eeb5b17b1a0a", + "generation-reference-repeat/43/initial.pt": "01ca4819f3d9260f926bd8a229be922b9d35e45009c1deb864841e9a20e2bd9a", + "generation-reference-repeat/43/step-000.pt": "ef51b8c119b9efd613f0bec53b4f705682f676c76857db50a88895b5c2a72d60", + "generation-reference-repeat/43/step-001.pt": "6d6b1f0963aac327fbcf92ca0b803154d992155de222cef6fc04396e38941ad0", + "generation-reference-repeat/43/step-002.pt": "436745a64e2aa37963420f54b852ebb96167e5c2058cd31d8af29c96d51f06a1", + "generation-reference-repeat/43/step-003.pt": "f3c67f8e4d07b214bc19d1f4e173b7b43d734bb061d8b89893caedea8dccde38", + "generation-reference-repeat/43/step-004.pt": "ee2538a44ec15a678da6700633f9a40cbccc66747e83a774632a370b1704c3bf", + "generation-reference-repeat/43/step-005.pt": "4cc7db11be0b53d7186f8937944be68ef412c22728974abe05186c3c73e7c062", + "generation-reference-repeat/43/step-006.pt": "132c641748425a57daa62ced3dc8502fc2bd9c33018e8fa4ab574d8ccc3e9d9e", + "generation-reference-repeat/43/step-007.pt": "06707ecc8795663787af4fb0cb93361f536070d2b93ffad78f39f1cc2e4e3baa", + "generation-reference-repeat/43/step-008.pt": "991dbf6eafdf2907a010acb608486b9e9815164263ed33d84ad9add145a60dd3", + "generation-reference-repeat/43/step-009.pt": "f8ec55fd35a0cab8efd5d53ad917c1f378d6973446081ce3f2232700e8b7481b", + "generation-reference-repeat/43/step-010.pt": "5045c6c68df597fd69a735b441e52da2781972a72cf4e95345f55c1548441564", + "generation-reference-repeat/43/step-011.pt": "6949be69f34f3d4e7bc7b01985994a220074de2ab9209f7e2f49726ae6c77486", + "generation-reference-repeat/43/step-012.pt": "a5ed2159307174e7dfe61f640eaef03299be4fc388de396ed85f91e33e6a17f2", + "generation-reference-repeat/43/step-013.pt": "da21ea765c52e626cba277c375b784cfa6d546aad4f412bd0a9c80988dac5023", + "generation-reference-repeat/43/step-014.pt": "e468641367340d527f403a6aacffe1497efbcdfbe2ec9d510c968dcd9c746a67", + "generation-reference-repeat/43/step-015.pt": "1e94f53ddaa365a11b03792aab8e984b88672c450d75c83fe3a729b7a20e953f", + "generation-reference-repeat/43/step-016.pt": "2198dadba30badaad090f5e712ae73ca83fadccb150d9dcbfcc6ec154e83951f", + "generation-reference-repeat/43/step-017.pt": "19cf82887c5f3da113ff3f0bf49d026997d3e0fe83787075cc98aa5d49f723a2", + "generation-reference-repeat/43/step-018.pt": "397589d9e597c473cbd6ead9909de24b9051f7b2439c268f0508ef87da65e131", + "generation-reference-repeat/43/step-019.pt": "d8633928769d57e739874d6c71e49a668148bd0241530a712ed343d4cf3163c7", + "generation-reference-repeat/43/step-020.pt": "4153c74c8110a74c12b129f58e080680621360f97e2b98f172ec2951f91aa6ce", + "generation-reference-repeat/43/step-021.pt": "5a1f0a1f855966b26b4419a67b3c76842f0da84b924d8a795cef28066a6c0d5b", + "generation-reference-repeat/43/step-022.pt": "4367a0eaae2018676ced9e84675eaa577ac1c8425ea6594769322ae2753b2ec2", + "generation-reference-repeat/43/step-023.pt": "7369f00e8ecc039e674872d81a7f1115e101b85ec00e8ae00ea429515a348677", + "generation-reference-repeat/43/step-024.pt": "323558190222094fa43c30f6ee4e20787526edebf4357ec45adfa65162c2fc08", + "generation-reference-repeat/43/step-025.pt": "c7ca33fb2a2a818a30109e937f101d9ea666926a61ae2c30ea4a5ffa8ab32398", + "generation-reference-repeat/43/step-026.pt": "a39a9c80bafe0aab4427c4c2f58a53b6e9a56cf6f45ec313ad39b0ff8610e3ba", + "generation-reference-repeat/43/step-027.pt": "c86b948a332c3e514187de662cf8912a154cc8b0980eeffb5cd938d8f05fb660", + "generation-reference-repeat/43/step-028.pt": "5600cd67ae665aa920c189e340b3b875322ff0ad8c3c11bba97e24c765eb85a6", + "generation-reference-repeat/43/step-029.pt": "aee0410558627c84b26baf469939ec33bb4069b30d5ca72a06a7e5ef2cacc662", + "generation-reference-repeat/43/step-030.pt": "f3976d7195bd38a0a2e9850b81488db46270a8cd08cfb9c73b05f1582c8be6e8", + "generation-reference-repeat/43/step-031.pt": "414e39885b68a22e3199cc5a5eaa3a71f6e53a274946843b35fc7b3367f6c652", + "generation-reference-repeat/43/step-032.pt": "103c0230a28d5b6ffa51ae56aec59f67748f5bc6a7240fcba746f40d51a83193", + "generation-reference-repeat/43/step-033.pt": "dc7926c6d76d4659dca824a716dab42363a535f7f35be82d5235966229d9bec6", + "generation-reference-repeat/43/step-034.pt": "7366c3b774c6205236a085ae008c85a7003905500003f8e994d313519a474ae6", + "generation-reference-repeat/43/step-035.pt": "1e88250fb28907bbc04a20e78306c8608fccf36de137f4a90e877bff36460c78", + "generation-reference-repeat/43/step-036.pt": "288c116c1be130d98e00348e956a6e36d516721b52aed894e90f0eae5e7e9c48", + "generation-reference-repeat/43/step-037.pt": "db5021e694a07743105ebd9b95210f4202cca222911c2f51c8d720f63dbfd330", + "generation-reference-repeat/43/step-038.pt": "5fc3d0f40929f12f29f5ed7babcd83626766ef77a2c6798aaa784d96659e6b6a", + "generation-reference-repeat/43/step-039.pt": "c01234332b204eb15900dd162b5f1d22942741f4c9a1dc86a1c52993789e6a94", + "generation-reference-repeat/43/step-040.pt": "a082fca297bab52c472b072a8ed5218b62620ae501032d374aea3616627392b7", + "generation-reference-repeat/43/step-041.pt": "a745bfd042e09d47fe95830ee14d8486d325cda9ae555075f5707031b9b9112c", + "generation-reference-repeat/43/step-042.pt": "63f96775fa57046d3af24025aa2814e56bc0ee9a21ee3e8e503ce202fec9a463", + "generation-reference-repeat/43/step-043.pt": "ad8dd8024a254d9ae1b804023f16a28913ec5f9dbb63990bcf0990e624065d8b", + "generation-reference-repeat/43/step-044.pt": "5cdf133961f925e236ccc0d6d315a163a675cba36e876d98c118d144a4abc1c8", + "generation-reference-repeat/43/step-045.pt": "1c8063587517124623fefdf5e4340b93ad4bb0c8310c06eec38899dae9e3f453", + "generation-reference-repeat/43/step-046.pt": "d343169046b8ee02981bd11c7af9d73034270b922637927b7fc92537c619c53e", + "generation-reference-repeat/43/step-047.pt": "85fc7d772a5308c4f1e1baf06ec8f76e3d3e805d5ca4a46906160686922b80ce", + "generation-reference-repeat/43/step-048.pt": "1f84391342e9d755093d87ca5a8c636ed73d2bbe8e482eff4807b22632418cb5", + "generation-reference-repeat/43/video.mp4": "51af10f6386bbe0b12db43d7c5d4727fe5e6c1aa91e5ee7502ff4b9e71c96f3a", + "generation-reference-repeat/43/video.npy": "578fbebcd4d38c8dbc2e7dc39c865032be1f7642940566b7df79af084aee5577", + "generation-reference-repeat/44/report.json": "d72c8d8cdc680c7af21481cd1fd35e821fc231020caa8c310497e25833433100", + "generation-reference-repeat/44/audio.npy": "a1972b7945babed4e3ce0b80854de8dc23bcdd5b4a5952bfa31298ee7ac79713", + "generation-reference-repeat/44/audio.wav": "edda8f495c4288cd872d2a585bfa17d9a4f7e6d0df79893a0cf1b08f3adaff74", + "generation-reference-repeat/44/frame-000.png": "65b7aab7c9029e1138fdf2d3958848974b5e1895e3347f91cae7b0f6388d6281", + "generation-reference-repeat/44/frame-062.png": "3bc9bea327c82767b06b1369d6b61cf00903b34b984e63a24d4b7725b1ac4992", + "generation-reference-repeat/44/frame-123.png": "1af5e294a949fa96c0afee0cf9a8fdd740d017628fcaa319ea55a755f9a768d5", + "generation-reference-repeat/44/initial.pt": "302ec05aa9bbd09182d6a6962bb8d75bb076ad75bbf84446f95f34e597c9bcd0", + "generation-reference-repeat/44/step-000.pt": "39c7619bfc5ff7eaf3c10aa4ca14e293ad496388966d9713de81945f9b743f33", + "generation-reference-repeat/44/step-001.pt": "45471fa2ee1d4ab1b1b357dd2931239db6cdea7799d4c7326d9bc35d0f26b5bd", + "generation-reference-repeat/44/step-002.pt": "077e41311a91277a157ae11bd66df9d28ee5c2a3ed6c6c1fcfade53f394d36cb", + "generation-reference-repeat/44/step-003.pt": "67432ad99b311154d2e90f950431d1ff371599bb9701daf5266bdc535a2e4f14", + "generation-reference-repeat/44/step-004.pt": "8859f733cc283308c340d3d44f1a5099ee3ae71cdb77761f0fbddf1b02cdf577", + "generation-reference-repeat/44/step-005.pt": "efbee8b13a2dae639192fc3c47d78f51cd5134331e85defc7351de21401a6ce0", + "generation-reference-repeat/44/step-006.pt": "a66b7db3f5c61750387a06aae13981aa6eabb0015dbb64815452584fbb980480", + "generation-reference-repeat/44/step-007.pt": "057bbbe657c12632bf7c8ead1c3111b57ae043c9a751c4e267f79e5f14bacde5", + "generation-reference-repeat/44/step-008.pt": "084c2b4a00e2420bf18b460aeab85cd099391b7f6908450e1f20627c47df03e3", + "generation-reference-repeat/44/step-009.pt": "d6881b3574378f8a7ac6f689dc3bf13d8a69f6a39a0ce8ee97198e9a9c46c213", + "generation-reference-repeat/44/step-010.pt": "749245a766ea9fd299720636ea27cfd14ee89eabf188db94ac4c4ebf0f07c81a", + "generation-reference-repeat/44/step-011.pt": "ed5fcc97a4bacbc4c62d668ac1aa7c9da00abd1540a9df3ea1b5be343777419f", + "generation-reference-repeat/44/step-012.pt": "f65e97731dc484ab2828f1a4ba439d5f6d3502576fc3870b3bcece5958d04a71", + "generation-reference-repeat/44/step-013.pt": "6e223fb0c3db20cadd7b06c1b3dfec5b3328d0c8df5e9c9559c9612c805b648a", + "generation-reference-repeat/44/step-014.pt": "58b8fc6672d6564b6c2fb306608986cbfdbf3887e5a93b3881a2ff880005a25b", + "generation-reference-repeat/44/step-015.pt": "3264e1989950fb0b0c89cb6ae1536dfb605316c73df66973e6cb5cadcb99535c", + "generation-reference-repeat/44/step-016.pt": "acbdfb98b66b28cc59b58c3bd9fa745821fd026e456feafb169fc6d2f44c3c94", + "generation-reference-repeat/44/step-017.pt": "661e2cfe17f866a62f09793956a95ab02096b167a20a8503336c1d0ad31d3436", + "generation-reference-repeat/44/step-018.pt": "19bd1378c573fcee28c27866acbe2c53be88c8c30ff850f919d9f7790f9d5c8a", + "generation-reference-repeat/44/step-019.pt": "57d7da03e567cb1c1bdea77cd8768144f32af6b556cdd652b0557fe39531ef2e", + "generation-reference-repeat/44/step-020.pt": "dd883aa09a70cde3e17b082ffdc8fabda9043b2b460818e4b382632ff5eb2e35", + "generation-reference-repeat/44/step-021.pt": "c644f23f16052b66990b3b25e76fc05b518f954d2e4192b959f7964cd88f7ea2", + "generation-reference-repeat/44/step-022.pt": "0c467aeccc6f4ba2e281593b42730f344bb93b13813c098d4d2decafaf9745eb", + "generation-reference-repeat/44/step-023.pt": "34f5bf5e004fb748529112b0374b5dcf452a5aa1159a1e85397f47e8b300b697", + "generation-reference-repeat/44/step-024.pt": "dbebf9825df277fe4ee8eaa6f59c6354843b5a9a076e9377f5a20f20a6523c04", + "generation-reference-repeat/44/step-025.pt": "90d980b8f51fa0a1b86d14b01232f6f3f3db22b5a2110885dbbc9b59daf142a3", + "generation-reference-repeat/44/step-026.pt": "ddfa270b092a6520b9662bafdf4bf2c8a8afc43b7052c5619f38aeb2441be93f", + "generation-reference-repeat/44/step-027.pt": "c290ca897f228f9eed50e829f64734514f19d5f6e033fba138720f52c52796a2", + "generation-reference-repeat/44/step-028.pt": "06039abebd9edf5c4447469f9618c3b9e9a5ba3664ef4078e4cd473b8ae1c352", + "generation-reference-repeat/44/step-029.pt": "47a6dffb08d186ab23572dab55c99f86bf26a33856f0fcc91afa953949c24c1d", + "generation-reference-repeat/44/step-030.pt": "a25e3ab3fbc73f94b2376eae117f447edf82c52cd4530633c0b8ce70fc967eb4", + "generation-reference-repeat/44/step-031.pt": "b08ed619094c076663c31aaecf5163c50c364cb03a14d4f9fc5b2032b12e8164", + "generation-reference-repeat/44/step-032.pt": "7b481348cf67e19d659f86bbcb32c564c40715feec3695491e58329d201a79cb", + "generation-reference-repeat/44/step-033.pt": "c523dd9c31bea576f03954a5b9079548417c7f148d4509f19865b5b4938f79a3", + "generation-reference-repeat/44/step-034.pt": "dc357ac4a20a217794221bbb29c039d9b764ba5bcaf66a787e8aa303cf20853b", + "generation-reference-repeat/44/step-035.pt": "6ce4198a268546acbbf441d2725b6076fde8bf8a395aaccf021b9b9f503fb850", + "generation-reference-repeat/44/step-036.pt": "a1d0f3dd1018ddefd6537d6ecd76ebf55134df19515da5232ce1c2be04dac12c", + "generation-reference-repeat/44/step-037.pt": "e434291c8c360f34c11544b8ac6b133615468e765f4c638bd614946ad18825e4", + "generation-reference-repeat/44/step-038.pt": "020c3f7be0fd1d911a39c4666bf0ac97a77df67c19b15fa8452fc1042988d6bc", + "generation-reference-repeat/44/step-039.pt": "4b853f31936e94678196424ab4af30842f432e035000bbfb5391a8ec9e418d42", + "generation-reference-repeat/44/step-040.pt": "96f0f02148a625f85d522a1dcf040785fbbc4d4f1ca809046eedb61e278c8c1a", + "generation-reference-repeat/44/step-041.pt": "37a1b65a39c1ff4892a2f311f7cecafaa48b050a575f8d9cfa14c294f5c53c2e", + "generation-reference-repeat/44/step-042.pt": "1411b77734ecad2978af0e7beea8be6295c9308c13db926f9b34af6e3959f012", + "generation-reference-repeat/44/step-043.pt": "7b7d4b8f7282334f6db2fd298b9deec81ff9c4b55b84eb5e35dbfcb4864ec80a", + "generation-reference-repeat/44/step-044.pt": "83a841307323c17020fab91c72ea68d7c83d92e80987d9945903477173e84718", + "generation-reference-repeat/44/step-045.pt": "f02464efaa745fa373469a2e3421a6c92f13a0ff8cd65f43a7cb5c4ae963b9aa", + "generation-reference-repeat/44/step-046.pt": "14ed97cbc60ea26a5fb3885092ab185d71920f1c2d12cd42e8a743590dc4af5a", + "generation-reference-repeat/44/step-047.pt": "3bf7c22072a4f84e35e6d366825e314ae7745423dc6a261a762ba938ff5685cb", + "generation-reference-repeat/44/step-048.pt": "189d20f8817b6e01b6d0120e02d22a83e070b254634ab2bb5b753dc19362533a", + "generation-reference-repeat/44/video.mp4": "929285c63024073b07eab1f6b7c3fbf5d2e49cc0f07b2e426dff5fb006d09196", + "generation-reference-repeat/44/video.npy": "7853db70a7373356bcadfac87447001da27e05409e1cb4a5a4ceee6d746f45a0", + "full-model-thresholds.json": "9542bce25d394e72f24c2cc02012fc47dc16100a37d5f86feda59c132132eb30", + "reference-repeat-quality.json": "3817569746ac3a1c99cffe32bdf183fbe12de2817c3e60fd6c21c17711bdd691", + "commit.txt": "7c9a380b0ba89a2288ca661bd2fda1e102a02606521681ab0bcf3077e5d2d600", + "dependencies.txt": "6139443eb38b47a805548b0d892808f5b6a26c7bc3b3796ff96ef6b9b0185cbf", + "extra-dependencies.json": "d2655f41325039e28fd4aae62d5637de0df49d08ea8dc9545295cdbbf1f47ce6" + } +} diff --git a/examples/models/minimax_h3/reference.sbatch b/examples/models/minimax_h3/reference.sbatch new file mode 100644 index 000000000..ce7429710 --- /dev/null +++ b/examples/models/minimax_h3/reference.sbatch @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Site options (partition, account, job name) are given at submission; see README.md. +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gres=gpu:b200:1 +#SBATCH --cpus-per-task=16 +#SBATCH --mem=512G +#SBATCH --time=24:00:00 +set -euo pipefail +source_dir=$1 +python_bin=$2 +checkpoint=$3 +result_root=$4 +reference_image=$5 +provider_source=$6 +extra_dependencies=$7 +cd "$source_dir" +[[ -z $(git status --porcelain) ]] +result_dir="$result_root/job-$SLURM_JOB_ID-attempt-${SLURM_RESTART_COUNT:-0}" +mkdir "$result_dir" +export PYTHONPATH="$source_dir:$provider_source/src:$extra_dependencies" +export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 USE_TF=0 USE_FLAX=0 +export HF_HUB_OFFLINE=1 TOKENIZERS_PARALLELISM=false +scontrol show job "$SLURM_JOB_ID" > "$result_dir/slurm.txt" +nvidia-smi -q > "$result_dir/device.txt" +git rev-parse HEAD > "$result_dir/commit.txt" +"$python_bin" -m pip freeze > "$result_dir/dependencies.txt" +"$python_bin" -m pip list --path "$extra_dependencies" --format=json > "$result_dir/extra-dependencies.json" +for mode in fp32-reference bf16-provider; do + "$python_bin" -m tools.validation.models.h3_full_model \ + --checkpoint "$checkpoint" --manifest rl_engine/validation/models/h3_full_manifest.json \ + --mode "$mode" --output "$result_dir/$mode" > "$result_dir/$mode.log" 2>&1 +done +"$python_bin" -m tools.validation.models.h3_full_compare \ + --reference "$result_dir/fp32-reference" --other "$result_dir/bf16-provider" \ + --output "$result_dir/full-model-thresholds.json" > "$result_dir/calibration.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --preflight --output "$result_dir/generation-preflight" \ + > "$result_dir/generation-preflight.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference" \ + > "$result_dir/generation-reference.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_full_generation --checkpoint "$checkpoint" \ + --reference-image "$reference_image" --output "$result_dir/generation-reference-repeat" \ + > "$result_dir/generation-reference-repeat.log" 2>&1 +"$python_bin" -m tools.validation.models.h3_generation_compare --reference "$result_dir/generation-reference" \ + --candidate "$result_dir/generation-reference-repeat" --provider-repeat \ + --contract examples/models/minimax_h3/generation-contract.json --output "$result_dir/reference-repeat-quality.json" \ + > "$result_dir/reference-repeat-quality.log" 2>&1 +# Candidate jobs are submitted only after these immutable reference artifacts +# are reviewed and the numerical thresholds have been committed. diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index fb3c4379c..40fb56213 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -325,3 +325,153 @@ def deterministic_collective_rocm_ipc_all_gather_input( input: torch.Tensor, output: torch.Tensor, ) -> None: ... +def h3_timestep_sinusoid_forward( + timestep: torch.Tensor, + num_channels: int = 256, + max_period: float = 10000.0, + check_range: bool = True, +) -> torch.Tensor: + """Return FP32 CUDA cosine-then-sine features of shape ``(T, num_channels)``. + + Require a nonempty 1-D FP32 CUDA timestep tensor, positive even channels + and positive ``max_period``. Check finite values in ``[0, 1]`` when enabled; + the Python operator supplies the analytic backward for this native forward. + """ + ... + +def h3_det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: int = 0, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Return a deterministic CUDA projection and optional FP32 pre-activation. + + Require same-device/dtype FP32 or BF16 contiguous, 16-byte-aligned matrices + ``x (T, K)`` and ``weight (N, K)``, ``T > 0`` and ``K`` divisible by the + 16-byte vector width. Optional contiguous bias is ``(N,)`` in that dtype. + Activation 0 is identity and 1 is SiLU; BF16 forward requires SM80 or newer. + Return ``[out]`` in the input dtype or ``[out, pre]`` when saving activation. + """ + ... + +def h3_det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Return a deterministic ``(T, K)`` CUDA input gradient in ``out_dtype``. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``weight (N, K)``. Output dtype + must be FP32 or BF16; output-column chunks accumulate in fixed order. + """ + ... + +def h3_det_linear_backward_input_partials( + grad: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: ... +def h3_det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: ... +def h3_det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, with_bias: bool = True +) -> list[torch.Tensor]: + """Return CUDA weight and optional bias gradients using an ascending row fold. + + Require contiguous FP32 ``grad (T, N)`` with ``T > 0`` and same-device + contiguous, 16-byte-aligned FP32 or BF16 ``x (T, K)``. Return ``[dweight]`` + with shape ``(N, K)``, or ``[dweight, dbias]`` with ``dbias (N,)`` when + requested, in FP32 or BF16 as specified by ``w_dtype``. + """ + ... + +def h3_adaln_row_gather_forward( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + chunks: int = 6, + modality_num: int = 3, +) -> torch.Tensor: ... +def h3_adaln_row_gather_backward( + grad: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: ... +def h3_rmsnorm_forward( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, + sorted_pos: torch.Tensor | None = None, + tile_begin: torch.Tensor | None = None, + tile_end: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_residual_forward( + residual: torch.Tensor, y: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... +def h3_gate_residual_backward( + grad: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + sorted_pos: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, + seg_first_tile: torch.Tensor, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_backward_partials( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None, + scale: torch.Tensor | None, + index: torch.Tensor | None, + dw_rows: torch.Tensor, + dw_begin: torch.Tensor, + dw_end: torch.Tensor, + seg_rows: torch.Tensor | None = None, + seg_begin: torch.Tensor | None = None, + seg_end: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_rmsnorm_fold_partials( + dw_partial: torch.Tensor, + weight: torch.Tensor, + seg_partial: torch.Tensor | None = None, + seg_first_tile: torch.Tensor | None = None, +) -> list[torch.Tensor]: ... +def h3_gate_grad_partials( + grad: torch.Tensor, + y: torch.Tensor, + rows: torch.Tensor, + tile_begin: torch.Tensor, + tile_end: torch.Tensor, +) -> torch.Tensor: ... +def h3_gate_grad_fold( + partial: torch.Tensor, seg_first_tile: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: ... +def h3_rmsnorm_backward_dx( + grad: torch.Tensor, + x: torch.Tensor, + weight: torch.Tensor, + rstd: torch.Tensor, + shift: torch.Tensor | None = None, + scale: torch.Tensor | None = None, + index: torch.Tensor | None = None, +) -> torch.Tensor: ... +def h3_gate_residual_backward_dy( + grad: torch.Tensor, gate: torch.Tensor, index: torch.Tensor +) -> torch.Tensor: ... diff --git a/rl_engine/backends/cuda/model_specific/__init__.py b/rl_engine/backends/cuda/model_specific/__init__.py new file mode 100644 index 000000000..3b09d3c11 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA implementations with model-specific arithmetic contracts.""" diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py new file mode 100644 index 000000000..86cf4c9da --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py new file mode 100644 index 000000000..17a3a3884 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused H3 AdaLN modulation: ``adaln_projection_3mod`` followed by ``adaln_row_gather``. + +The forward is exactly the two RFC #420 ops back to back (same kernels, same +bits). The difference is the backward. With two separate ops, autograd hands +the gather's table gradient to the projection in the table's dtype (BF16), +which rounds a long FP32 segment sum once more before it reaches the +projection weights and ``temb``. Here the segment sum stays FP32 into the +projection backward, so the parameter gradients carry no extra BF16 rounding. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + adaln_row_gather_available, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import validate_h3_adaln_projection +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_adaln_row_gather_forward" +) +BACKWARD_IMPL = "fp32_segment_sum_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_adaln_row_gather_backward[fp32]" + "+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNModulationCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias, timestep_indices, token_tags, hidden): + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias) + rows = table.view(-1, H3_ADALN_CHUNKS * hidden) + packed = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(temb, act, weight, timestep_indices, token_tags) + ctx.table_shape = table.shape + return packed + + @staticmethod + def backward(ctx, grad_packed): + temb, act, weight, timestep_indices, token_tags = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad[:3] + num_timesteps, width = ctx.table_shape + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_timesteps * H3_MODALITY_NUM) + d_rows = _C.h3_adaln_row_gather_backward(grad_packed.contiguous(), *tiles, torch.float32) + d_table = d_rows.view(num_timesteps, width) + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + if need_temb: + d_temb = silu_backward_fp32( + det_linear_backward_input(d_table, weight, torch.float32), temb + ) + record_backward( + "adaln_modulation_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return ( + d_temb, + dw if need_weight else None, + db if need_bias else None, + None, + None, + None, + ) + + +class H3AdaLNModulationCudaOp: + """``adaln_projection_3mod`` + ``adaln_row_gather`` with an FP32 table gradient. + + Returns the six ``(S, H)`` modulation tensors for every packed position, + bitwise equal to running the two ops separately. + """ + + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and adaln_row_gather_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_adaln_row_gather_* ops") + + def __call__(self, temb, weight, bias, timestep_indices, token_tags, *, check_range=True): + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNModulationCudaOp needs CUDA tensors") + _validate_indices(timestep_indices, token_tags, temb.shape[0], check_range) + packed = _H3AdaLNModulationCuda.apply( + temb, + weight, + bias, + timestep_indices.contiguous(), + token_tags.contiguous(), + hidden, + ) + return tuple(packed.unbind(0)) + + +def _validate_indices(timestep_indices, token_tags, num_timesteps, check_range): + # Reuse the gather's checks against a stand-in of the right row count. + stand_in = torch.empty( + (num_timesteps * H3_MODALITY_NUM, H3_ADALN_CHUNKS), device=timestep_indices.device + ) + validate_h3_adaln_row_gather(stand_in, timestep_indices, token_tags, check_range=check_range) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..f02431903 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 three-modality AdaLN projection (RFC #420 ``adaln_projection_3mod``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.adaln_projection import ( + split_adaln_table, + validate_h3_adaln_projection, +) + +KERNEL_ID = "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3AdaLNProjectionCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight, bias): + """Project FP32 ``(T, D)`` embeddings to a weight-dtype ``(T, 18H)`` table. + + Save the embedding, rounded SiLU activation and weight for the VJP. + """ + + temb = temb.contiguous() + # Declared mixed-precision boundary: SiLU at temb's FP32 precision, + # then exactly one rounding to the projection dtype. These are the + # provider's own elementwise ops, so the activation is bitwise equal. + act = F.silu(temb).to(weight.dtype) + (table,) = det_linear_forward(act, weight, bias, activation=ACT_NONE) + ctx.save_for_backward(temb, act, weight) + return table + + @staticmethod + def backward(ctx, grad_table): + """Return requested embedding, weight and bias gradients in their input dtypes. + + Accumulate projection gradients deterministically and evaluate the + straight-through cast and SiLU VJP in FP32 for the embedding gradient. + """ + + temb, act, weight = ctx.saved_tensors + need_temb, need_weight, need_bias = ctx.needs_input_grad + grad_table = grad_table.float().contiguous() + d_temb = dw = db = None + if need_weight or need_bias: + dw, db = det_linear_backward_weight(grad_table, act, weight.dtype) + if need_temb: + # The cast's VJP is the identity; the SiLU VJP runs in FP32. + d_act = det_linear_backward_input(grad_table, weight, torch.float32) + d_temb = silu_backward_fp32(d_act, temb) + record_backward( + "adaln_projection_3mod", + kernel_id=BACKWARD_KERNEL_ID, + impl=BACKWARD_IMPL, + family="cuda", + ) + return d_temb, dw if need_weight else None, db if need_bias else None + + +class H3AdaLNProjectionCudaOp: + """CUDA candidate: FP32 SiLU, one BF16 cast, deterministic GEMV, six views. + + The projection weight (96768 x 2688 BF16, 520 MB per block) is streamed + once per call; every output element follows contract ``h3-det-linear-v1``, + so a timestep's 3 x 6 modulation rows do not depend on the other + timesteps in the call. The six outputs are views of one ``(T, 6H*3)`` + table, laid out exactly like diffusers' ``view(-1, 6H).chunk(6)``. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` modulation views for CUDA inputs.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six ``(3T, H)`` views after FP32 SiLU and a deterministic projection. + + Require CUDA FP32 ``temb`` of shape ``(T, D)`` with ``T > 0`` and + same-device FP32 or BF16 weight/bias of shapes ``(18H, D)``/``(18H,)``. + Outputs share one table in the weight dtype; autograd covers all inputs. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + table = _H3AdaLNProjectionCuda.apply(temb, weight, bias) + return split_adaln_table(table, hidden) + + def forward_table(self, temb, weight, bias) -> torch.Tensor: + """The raw ``(T, 6H*3)`` projection, for consumers that gather it directly.""" + + validate_h3_adaln_projection(temb, weight, bias) + if not temb.is_cuda: + raise ValueError("H3AdaLNProjectionCudaOp needs CUDA tensors") + return _H3AdaLNProjectionCuda.apply(temb, weight, bias) + + def forward_fp32(self, temb, weight, bias): + """Run the CUDA projection with its declared weight-dtype output boundary.""" + + return self.forward(temb, weight, bias) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..ee681454d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 AdaLN row gather (RFC #420 ``adaln_row_gather``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_row_gather import ( + h3_adaln_indices, + validate_h3_adaln_row_gather, +) + +KERNEL_ID = "rl_engine._C.h3_adaln_row_gather_forward" +BACKWARD_IMPL = "stable_sorted_segment_tiles_fp32_ordered_fold" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_adaln_row_gather_backward" +# Positions per backward tile; part of the fixed summation order. +BACKWARD_TILE = 256 + + +def adaln_row_gather_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_adaln_row_gather_forward") + and hasattr(_C, "h3_adaln_row_gather_backward") + ) + + +def _segment_tiles(index: torch.Tensor, num_rows: int): + """Sorted positions and fixed-size tiles per row segment (all integer ops).""" + + sorted_pos = torch.argsort(index, stable=True) + counts = torch.bincount(index, minlength=num_rows) + seg_start = torch.cumsum(counts, 0) - counts + tiles_per_seg = (counts + BACKWARD_TILE - 1) // BACKWARD_TILE + seg_first_tile = torch.zeros(num_rows + 1, dtype=torch.int64, device=index.device) + seg_first_tile[1:] = torch.cumsum(tiles_per_seg, 0) + tile_seg = torch.repeat_interleave(torch.arange(num_rows, device=index.device), tiles_per_seg) + tile_local = torch.arange(tile_seg.numel(), device=index.device) - seg_first_tile[tile_seg] + tile_begin = seg_start[tile_seg] + tile_local * BACKWARD_TILE + tile_end = torch.minimum(tile_begin + BACKWARD_TILE, seg_start[tile_seg] + counts[tile_seg]) + return sorted_pos, tile_begin.contiguous(), tile_end.contiguous(), seg_first_tile + + +class _H3AdaLNRowGatherCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, rows, timestep_indices, token_tags): + out = _C.h3_adaln_row_gather_forward( + rows, timestep_indices, token_tags, H3_ADALN_CHUNKS, H3_MODALITY_NUM + ) + ctx.save_for_backward(timestep_indices, token_tags) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return out + + @staticmethod + def backward(ctx, grad_out): + timestep_indices, token_tags = ctx.saved_tensors + num_rows, rows_dtype = ctx.rows_meta + index = h3_adaln_indices(timestep_indices, token_tags) + tiles = _segment_tiles(index, num_rows) + grad_rows = _C.h3_adaln_row_gather_backward(grad_out.contiguous(), *tiles, rows_dtype) + record_backward( + "adaln_row_gather", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad_rows, None, None + + +class H3AdaLNRowGatherCudaOp: + """CUDA candidate: one launch gathers all six modulation tensors. + + The forward is a byte copy (bitwise equal to six ``index_select`` calls). + The backward replaces ``index_select``'s atomic BF16 scatter-add with a + deterministic FP32 segmented sum in stable sorted order, cast once. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not adaln_row_gather_available(): + raise RuntimeError( + "rl_engine._C lacks h3_adaln_row_gather_*; rebuild with " + "csrc/cuda/h3/adaln_row_gather.cu" + ) + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + if not rows.is_cuda: + raise ValueError("H3AdaLNRowGatherCudaOp needs CUDA tensors") + if rows.stride(1) != 1: + rows = rows.contiguous() + packed = _H3AdaLNRowGatherCuda.apply( + rows, timestep_indices.contiguous(), token_tags.contiguous() + ) + return tuple(packed.unbind(0)) + + def gather_chunks(self, chunks, timestep_indices, token_tags, **kwargs): + """Drop-in for diffusers' six ``(3T, H)`` tensors: concatenate, then gather.""" + + if len(chunks) != H3_ADALN_CHUNKS: + raise ValueError(f"expected {H3_ADALN_CHUNKS} modulation tensors, got {len(chunks)}") + return self.forward(torch.cat(list(chunks), dim=1), timestep_indices, token_tags, **kwargs) + + def forward_fp32(self, rows, timestep_indices, token_tags): + return tuple(out.float() for out in self.forward(rows, timestep_indices, token_tags)) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py new file mode 100644 index 000000000..1274364f2 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Python surface of the H3 deterministic row-wise linear (contract h3-det-linear-v1). + +See ``csrc/cuda/h3/det_linear.cu`` for the reduction order. Every helper +raises instead of falling back to cuBLAS when the extension is missing. +""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE + +ACT_NONE = 0 +ACT_SILU = 1 +CONTRACT = "h3-det-linear-v1" +DINPUT_CHUNK = 64 # N rows per d_input partial (kDInputChunk) +_SYMBOLS = ( + "h3_det_linear_forward", + "h3_det_linear_backward_input", + "h3_det_linear_backward_weight", + "h3_det_linear_backward_input_partials", + "h3_det_linear_fold_chunks", +) + + +def det_linear_available() -> bool: + """Return whether the extension exposes all three H3 deterministic linear APIs.""" + + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +def _require() -> None: + """Raise when a deterministic linear entry point is unavailable.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks the h3_det_linear_* symbols; rebuild the CUDA extension " + "with csrc/cuda/h3/det_linear.cu (no cuBLAS fallback is allowed)" + ) + + +def det_linear_forward( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + *, + activation: int = ACT_NONE, + save_pre_activation: bool = False, +) -> list[torch.Tensor]: + """Project CUDA ``x`` of shape ``(T, K)`` with ``(N, K)`` weights and optional bias. + + Inputs share an FP32 or BF16 dtype/device, ``T > 0``, and ``K`` supports + 16-byte vector loads; BF16 forward needs SM80 or newer. Return ``[out]`` + in the input dtype, plus an FP32 pre-activation when requested. Activation + is identity or SiLU; contiguous copies preserve the reduction contract. + """ + + _require() + return _C.h3_det_linear_forward( + x.contiguous(), + weight.contiguous(), + None if bias is None else bias.contiguous(), + int(activation), + bool(save_pre_activation), + ) + + +def det_linear_backward_input( + grad: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype +) -> torch.Tensor: + """Compute a deterministic ``(T, K)`` input VJP on the inputs' CUDA device. + + Convert nonempty ``(T, N)`` gradients to FP32 and combine them with FP32 + or BF16 ``(N, K)`` weights; return FP32 or BF16 as specified by ``out_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_input( + grad.float().contiguous(), weight.contiguous(), out_dtype + ) + + +def det_linear_backward_input_partials(grad: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """``(ceil(N / 64), T, K)`` FP32 chunk sums of ``grad @ weight``, before the fold.""" + + _require() + return _C.h3_det_linear_backward_input_partials(grad.float().contiguous(), weight.contiguous()) + + +def det_linear_fold_chunks(partial: torch.Tensor, out_dtype: torch.dtype) -> torch.Tensor: + """Ascending left fold of the chunk partials, cast once (the second half of d_input).""" + + _require() + return _C.h3_det_linear_fold_chunks(partial.contiguous(), out_dtype) + + +def det_linear_backward_weight( + grad: torch.Tensor, x: torch.Tensor, w_dtype: torch.dtype, *, with_bias: bool = True +) -> list[torch.Tensor]: + """Fold nonempty CUDA ``(T, N)`` gradients and ``(T, K)`` inputs in row order. + + Convert gradients to FP32; ``x`` is FP32 or BF16 on the same device. + Return ``[dweight]`` of shape ``(N, K)``, optionally followed by ``dbias`` + of shape ``(N,)``, both in the requested FP32 or BF16 ``w_dtype``. + """ + + _require() + return _C.h3_det_linear_backward_weight( + grad.float().contiguous(), x.contiguous(), w_dtype, bool(with_bias) + ) + + +def silu_backward_fp32(grad: torch.Tensor, pre_activation: torch.Tensor) -> torch.Tensor: + """Elementwise SiLU VJP in FP32: g * s * (1 + z * (1 - s)), s = sigmoid(z).""" + + z = pre_activation.float() + s = torch.sigmoid(z) + return grad.float() * (s * (1.0 + z * (1.0 - s))) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..741a0a85e --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 final AdaLN output (RFC #420 ``final_adaln_out``).""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available +from rl_engine.backends.extension import _C +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.final_adaln_out import validate_h3_final_adaln_out +from rl_engine.reference.minimax_h3.rmsnorm import H3_NORM_EPS + +KERNEL_ID = ( + "torch.silu[fp32]+cast[weight_dtype]+rl_engine._C.h3_det_linear_forward" + "+rl_engine._C.h3_rmsnorm_forward[modulated]" +) +BACKWARD_IMPL = "fp32_table_grad_into_h3_det_linear_backward" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_rmsnorm_backward+rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3FinalAdaLNOutCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, norm_weight, temb, weight, bias, timestep_indices, eps): + hidden = x.shape[-1] + temb = temb.contiguous() + act = F.silu(temb).to(weight.dtype) # declared boundary: one cast after the FP32 SiLU + (table,) = det_linear_forward(act, weight, bias) + shift, scale = table.chunk(2, dim=-1) + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward( + x2, norm_weight.contiguous(), float(eps), shift, scale, timestep_indices + ) + ctx.save_for_backward(x2, norm_weight, rstd, temb, act, weight, table, timestep_indices) + ctx.x_shape = x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, norm_weight, rstd, temb, act, weight, table, timestep_indices = ctx.saved_tensors + shift, scale = table.chunk(2, dim=-1) + positions = timestep_indices.repeat(x2.shape[0] // timestep_indices.shape[0]) + tiles = _segment_tiles(positions, table.shape[0]) + dx, d_norm_weight, d_shift, d_scale = _C.h3_rmsnorm_backward( + grad.contiguous().view_as(x2), + x2, + norm_weight.contiguous(), + rstd, + shift, + scale, + timestep_indices, + *tiles, + ) + # The table gradient stays FP32 into the projection backward. + d_table = torch.cat([d_shift, d_scale], dim=1) + dw, db = det_linear_backward_weight(d_table, act, weight.dtype) + d_temb = silu_backward_fp32(det_linear_backward_input(d_table, weight, torch.float32), temb) + record_backward( + "final_adaln_out", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), d_norm_weight, d_temb, dw, db, None, None + + +class H3FinalAdaLNOutCudaOp: + """CUDA candidate: tensor-core shift/scale projection + fused norm/modulation. + + The projection is ``adaln_projection_3mod``'s deterministic GEMV + (``h3-det-linear-bf16-mma-v1``) on ``norm_out.linear``; the norm and + modulation are ``h3_rmsnorm``'s kernel indexed by ``timestep_indices``. + One autograd node, so the table gradient reaches the projection in FP32. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not (det_linear_available() and h3_rmsnorm_available()): + raise RuntimeError("rl_engine._C lacks the h3_det_linear_* / h3_rmsnorm_* ops") + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + if weight.dtype != x.dtype: + raise TypeError("norm_out.linear must share the activations' dtype") + if not x.is_cuda: + raise ValueError("H3FinalAdaLNOutCudaOp needs CUDA tensors") + return _H3FinalAdaLNOutCuda.apply( + x, norm_weight, temb, weight, bias, timestep_indices.contiguous(), eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py new file mode 100644 index 000000000..3d040e9e8 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 gated residual (RFC #420 ``adaln_gate_residual``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual + +KERNEL_ID = "rl_engine._C.h3_gate_residual_forward" +BACKWARD_IMPL = "exact_dy_sorted_segment_dgate" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_gate_residual_backward" + + +def h3_gate_residual_available() -> bool: + return bool( + _EXT_AVAILABLE + and hasattr(_C, "h3_gate_residual_forward") + and hasattr(_C, "h3_gate_residual_backward") + ) + + +class _H3GateResidualCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, index): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, index) + ctx.save_for_backward(y2, gate, index) + ctx.shape = residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + positions = index.repeat(y2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, gate.shape[0]) + dy, dgate = _C.h3_gate_residual_backward(g2, y2, gate, index, *tiles) + record_backward( + "adaln_gate_residual", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return grad, dy.view(ctx.shape), dgate, None + + +class H3GateResidualCudaOp: + """CUDA candidate: ``residual + gate[index] * y`` in one elementwise pass. + + The gate row is gathered in the kernel from the AdaLN table view, and the + two roundings match the eager expression, so the output is bitwise equal + to diffusers. ``dy`` is bitwise equal to the eager VJP; ``dgate`` is a + deterministic FP32 segment sum instead of ``index_select``'s BF16 atomics. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_gate_residual_available(): + raise RuntimeError( + "rl_engine._C lacks h3_gate_residual_*; rebuild with csrc/cuda/h3/gate_residual.cu" + ) + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + if not residual.is_cuda: + raise ValueError("H3GateResidualCudaOp needs CUDA tensors") + if gate.stride(1) != 1: + gate = gate.contiguous() + return _H3GateResidualCuda.apply(residual, y, gate, index.contiguous()) + + def forward_fp32(self, residual, y, gate, index): + return self.forward(residual, y, gate, index).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..d9266e759 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 RMSNorm with fused AdaLN modulation (RFC #420 ``h3_rmsnorm``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +KERNEL_ID = "rl_engine._C.h3_rmsnorm_forward" +BACKWARD_IMPL = "row_local_dx_tiled_dweight_sorted_table_grads" +BACKWARD_KERNEL_ID = "torch.argsort[stable]+rl_engine._C.h3_rmsnorm_backward" + + +def h3_rmsnorm_available() -> bool: + return bool( + _EXT_AVAILABLE and hasattr(_C, "h3_rmsnorm_forward") and hasattr(_C, "h3_rmsnorm_backward") + ) + + +class _H3RMSNormCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, index, eps): + hidden = x.shape[-1] + x2 = x.contiguous().view(-1, hidden) + y, rstd = _C.h3_rmsnorm_forward(x2, weight.contiguous(), float(eps), shift, scale, index) + ctx.save_for_backward(x2, weight, rstd, shift, scale, index) + ctx.x_shape = x.shape + ctx.modulated = index is not None + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale, index = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + if not ctx.modulated: + dx, dweight = _C.h3_rmsnorm_backward(g2, x2, weight.contiguous(), rstd) + d_shift = d_scale = None + else: + positions = index.repeat(x2.shape[0] // index.shape[0]) + tiles = _segment_tiles(positions, shift.shape[0]) + dx, dweight, d_shift, d_scale = _C.h3_rmsnorm_backward( + g2, x2, weight.contiguous(), rstd, shift, scale, index, *tiles + ) + d_shift, d_scale = d_shift.to(shift.dtype), d_scale.to(scale.dtype) + record_backward( + "h3_rmsnorm", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None + + +class H3RMSNormCudaOp: + """CUDA candidate. + + The statistics replay PyTorch's own RMSNorm reduction order, so the plain + norm is bitwise equal to ``nn.RMSNorm``; the modulated form gathers its + ``shift``/``scale`` rows in the kernel (no ``(S, H)`` intermediates) and + rounds each step where the eager expression does, so it is bitwise equal to + diffusers' ``n * (1.0 + scale[i]) + shift[i]`` as well. Rows are + independent of batch size and position. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + if not h3_rmsnorm_available(): + raise RuntimeError( + "rl_engine._C lacks h3_rmsnorm_*; rebuild with csrc/cuda/h3/rmsnorm_modulate.cu" + ) + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + return _H3RMSNormCuda.apply(x, weight, None, None, None, eps) + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + if not x.is_cuda: + raise ValueError("H3RMSNormCudaOp needs CUDA tensors") + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + return _H3RMSNormCuda.apply(x, weight, shift, scale, index.contiguous(), eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + return self.forward(x, weight, eps).float() diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py new file mode 100644 index 000000000..37cea620b --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sequence-parallel H3 RMSNorm / AdaLN modulation / gated residual (RFC #420 ``sp_norm_adaln``). + +Ownership: rank ``r`` of ``sp`` holds packed positions ``[r * S // sp, (r + 1) * S // sp)`` +of every batch item. The full ``(S,)`` row index is replicated; it is small, and every +rank derives the same reduction plan from it. + +Forward and the row-local gradients (``dx``, ``d_sublayer``, ``d_residual``) are the WS1 +kernels on the local rows. Rows are independent, so they are byte-equal by construction. + +The cross-row reductions (``d_norm_weight``, ``d_shift``/``d_scale``, ``d_gate``) are WS1 +two-level folds: FP32 sums over fixed 256-element tiles, then an ascending fold of the +tiles. Every rank builds the global WS1 tile list. A tile is computed by the rank that +holds its first row; the tile's rows from other ranks are all-gathered beforehand (only +rows of tiles that straddle a shard boundary move). The tile partials are all-gathered, +put back in WS1 tile order, and every rank runs the WS1 fold. The partial and fold kernels +are the WS1 kernels, so every gradient equals WS1 on every rank. The only collectives are +rank-ordered all-gathers (copies). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + BACKWARD_TILE, + _segment_tiles, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.gate_residual import validate_h3_gate_residual +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + validate_h3_modulation, + validate_h3_rmsnorm, +) + +_SYMBOLS = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward_dx", + "h3_rmsnorm_backward_partials", + "h3_rmsnorm_fold_partials", + "h3_gate_residual_forward", + "h3_gate_residual_backward_dy", + "h3_gate_grad_partials", + "h3_gate_grad_fold", +) +BACKWARD_IMPL = "row_local_ws1+gather_straddling_rows+ws1_tile_partials+gather+ws1_fold" + + +def sp_norm_adaln_available() -> bool: + return bool(_EXT_AVAILABLE and all(hasattr(_C, name) for name in _SYMBOLS)) + + +@dataclass(frozen=True) +class SPRowLayout: + """Which packed positions each sequence-parallel rank holds.""" + + seq_len: int + batch: int + sp: int + rank: int + + def bounds(self, rank: int) -> tuple[int, int]: + return rank * self.seq_len // self.sp, (rank + 1) * self.seq_len // self.sp + + @property + def lo(self) -> int: + return self.bounds(self.rank)[0] + + @property + def hi(self) -> int: + return self.bounds(self.rank)[1] + + @property + def local_len(self) -> int: + return self.hi - self.lo + + +def sp_row_layout(seq_len: int, batch: int, sp: int, rank: int) -> SPRowLayout: + if sp < 1 or not 0 <= rank < sp: + raise ValueError(f"need 0 <= rank < sp, got rank={rank}, sp={sp}") + if batch < 1 or seq_len < sp: + raise ValueError(f"every rank needs a row: seq_len={seq_len}, sp={sp}, batch={batch}") + return SPRowLayout(seq_len, batch, sp, rank) + + +def _ordinal_by_group(group: torch.Tensor, groups: int) -> tuple[torch.Tensor, torch.Tensor]: + """Per-group counts, and each element's position within its group (in input order).""" + + counts = torch.bincount(group, minlength=groups) + order = torch.argsort(group, stable=True) + ordinal = torch.empty_like(group) + first = torch.cumsum(counts, 0) - counts + ordinal[order] = torch.arange(group.numel(), device=group.device) - first[group[order]] + return counts, ordinal + + +class _Plan: + """Global WS1 tiles of one backward, and this rank's part in computing them. + + ``families`` are ``(rows, tile_begin, tile_end)`` over global flattened rows + ``b * S + s``; consecutive tiles cover ``rows`` in order. Every rank builds the same + plan from replicated metadata, so collective calls and shapes agree. + """ + + def __init__(self, layout: SPRowLayout, families, device) -> None: + self.layout, sp = layout, layout.sp + self._his = torch.tensor([layout.bounds(r)[1] for r in range(sp)], device=device) + + # Rows of a tile computed by another rank: the only activation rows that move. + tile_owner, shared = [], [] + for rows, begin, end in families: + owner = self.owner(rows[begin]) + tile_owner.append(owner) + shared.append(rows[torch.repeat_interleave(owner, end - begin) != self.owner(rows)]) + shared = torch.unique(torch.cat(shared)) + shared_owner = self.owner(shared) + counts, slot = _ordinal_by_group(shared_owner, sp) + self.max_send = int(counts.max()) if shared.numel() else 0 + + # Buffer row of every global row this rank reads: its own rows first, then + # every rank's send block (padded to max_send) in rank order. + own = self.global_rows(layout.rank) + buffer_index = torch.full((layout.batch * layout.seq_len,), -1, device=device) + buffer_index[own] = torch.arange(own.numel(), device=device) + foreign = shared_owner != layout.rank + buffer_index[shared[foreign]] = ( + own.numel() + shared_owner[foreign] * self.max_send + slot[foreign] + ) + mine = ~foreign + self.send_local = self.local_index(shared[mine][torch.argsort(slot[mine])]) + self.recv_global = torch.zeros(sp * self.max_send, dtype=torch.long, device=device) + self.recv_global[shared_owner * self.max_send + slot] = shared # padding: never read + + self.families = [ + self._my_tiles(rows, begin, end, owner, buffer_index) + for (rows, begin, end), owner in zip(families, tile_owner) + ] + + def owner(self, rows: torch.Tensor) -> torch.Tensor: + return torch.searchsorted(self._his, rows % self.layout.seq_len, right=True) + + def global_rows(self, rank: int) -> torch.Tensor: + lo, hi = self.layout.bounds(rank) + s = torch.arange(lo, hi, device=self._his.device) + b = torch.arange(self.layout.batch, device=self._his.device) + return (b[:, None] * self.layout.seq_len + s[None, :]).reshape(-1) + + def local_index(self, rows: torch.Tensor) -> torch.Tensor: + b, s = rows // self.layout.seq_len, rows % self.layout.seq_len + return b * self.layout.local_len + s - self.layout.lo + + def _my_tiles(self, rows, begin, end, owner, buffer_index): + mine = torch.nonzero(owner == self.layout.rank).flatten() + lengths = (end - begin)[mine] + my_end = torch.cumsum(lengths, 0) + my_begin = my_end - lengths + elems = torch.repeat_interleave(begin[mine] - my_begin, lengths) + elems = elems + torch.arange(elems.numel(), device=rows.device) + counts, ordinal = _ordinal_by_group(owner, self.layout.sp) + width = int(counts.max()) if owner.numel() else 0 + return { + "tiles": (buffer_index[rows[elems]].contiguous(), my_begin, my_end), + "width": width, # tiles per rank in the padded partial gather + "gather_index": owner * width + ordinal, # WS1 tile order within that gather + } + + def exchange(self, collective, *tensors: torch.Tensor) -> list[torch.Tensor]: + """For each ``(M_local, ...)`` tensor: local rows, then the gathered send blocks.""" + + out = [] + for t in tensors: + if self.max_send == 0: + out.append(t) + continue + send = t.new_zeros((self.max_send, *t.shape[1:])) + send[: self.send_local.numel()] = t[self.send_local] + out.append(torch.cat([t, gather_rows(collective, send)])) + return out + + def gather_partials(self, collective, family: int, partial: torch.Tensor) -> torch.Tensor: + """This rank's tile partials -> every tile's partial, in WS1 tile order.""" + + fam = self.families[family] + padded = partial.new_zeros((fam["width"], *partial.shape[1:])) + padded[: partial.shape[0]] = partial + return gather_rows(collective, padded).index_select(0, fam["gather_index"]).contiguous() + + +def _dweight_tiles(total: int, device): + begin = torch.arange(0, total, BACKWARD_TILE, device=device) + return torch.arange(total, device=device), begin, torch.clamp(begin + BACKWARD_TILE, max=total) + + +@dataclass(frozen=True) +class SPPlan: + """Everything an SP backward needs that depends only on the layout and the row index. + + Built once per ``(row index, table rows)`` and shared by every norm and gated + residual that uses that index (every block of a forward pass). + """ + + layout: SPRowLayout + plan: _Plan + local_index: torch.Tensor | None # (S_local,) table row of each local position + index_buf: torch.Tensor | None # table row of each exchange-buffer row + seg_first_tile: torch.Tensor | None + num_rows: int | None + + +def sp_plan( + layout: SPRowLayout, index_full=None, num_rows: int | None = None, *, device=None +) -> SPPlan: + """The reduction plan for one row index (``None``: the plain norm's ``dweight`` only).""" + + device = index_full.device if index_full is not None else torch.device(device or "cuda") + families = [_dweight_tiles(layout.batch * layout.seq_len, device)] + if index_full is None: + return SPPlan(layout, _Plan(layout, families, device), None, None, None, None) + if num_rows is None: + raise ValueError("indexed SP plans require num_rows") + *seg_tiles, seg_first_tile = _segment_tiles(index_full.repeat(layout.batch), num_rows) + plan = _Plan(layout, [*families, tuple(seg_tiles)], device) + local_index = index_full[layout.lo : layout.hi].clone(memory_format=torch.contiguous_format) + received = index_full[plan.recv_global % layout.seq_len] + index_buf = torch.cat([local_index.repeat(layout.batch), received]).contiguous() + return SPPlan(layout, plan, local_index, index_buf, seg_first_tile, num_rows) + + +def sp_rmsnorm_backward(grad, x, weight, rstd, shift, scale, sp: SPPlan, collective): + """One rank's backward on its ``(B * S_local, H)`` rows: ``(dx, dweight, d_shift, d_scale)``. + + ``dx`` is row-local; the other gradients are the full WS1 reductions, identical on + every rank. ``shift``/``scale`` are ``None`` for the plain norm. + """ + + weight, plan, modulated = weight.contiguous(), sp.plan, sp.local_index is not None + dx = _C.h3_rmsnorm_backward_dx(grad, x, weight, rstd, shift, scale, sp.local_index) + g_buf, x_buf, rstd_buf = plan.exchange(collective, grad, x, rstd) + seg = plan.families[1]["tiles"] if modulated else (None, None, None) + partials = _C.h3_rmsnorm_backward_partials( + g_buf, x_buf, weight, rstd_buf, shift, scale, sp.index_buf, + *plan.families[0]["tiles"], *seg, + ) # fmt: skip + gathered = [plan.gather_partials(collective, i, p) for i, p in enumerate(partials)] + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_rmsnorm_backward_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + if not modulated: + (dweight,) = _C.h3_rmsnorm_fold_partials(gathered[0], weight) + return dx, dweight, None, None + dweight, d_shift, d_scale = _C.h3_rmsnorm_fold_partials( + gathered[0], weight, gathered[1], sp.seg_first_tile + ) + return dx, dweight, d_shift.to(shift.dtype), d_scale.to(scale.dtype) + + +def sp_gate_residual_backward(grad, y, gate, sp: SPPlan, collective): + """One rank's backward on its rows: ``(d_sublayer, d_gate)``; ``d_residual`` is ``grad``.""" + + plan = sp.plan + dy = _C.h3_gate_residual_backward_dy(grad, gate, sp.local_index) + g_buf, y_buf = plan.exchange(collective, grad, y) + partial = _C.h3_gate_grad_partials(g_buf, y_buf, *plan.families[1]["tiles"]) + dgate = _C.h3_gate_grad_fold( + plan.gather_partials(collective, 1, partial), sp.seg_first_tile, gate.dtype + ) + record_backward( + "sp_norm_adaln", + kernel_id="rl_engine._C.h3_gate_grad_partials", + impl=BACKWARD_IMPL, + family="cuda", + ) + return dy, dgate + + +class _SPRMSNorm(torch.autograd.Function): + @staticmethod + def forward(ctx, x, weight, shift, scale, sp, collective, eps): + x2 = x.contiguous().view(-1, x.shape[-1]) + y, rstd = _C.h3_rmsnorm_forward( + x2, weight.contiguous(), float(eps), shift, scale, sp.local_index + ) + ctx.save_for_backward(x2, weight, rstd, shift, scale) + ctx.sp, ctx.collective, ctx.x_shape = sp, collective, x.shape + return y.view(x.shape) + + @staticmethod + def backward(ctx, grad): + x2, weight, rstd, shift, scale = ctx.saved_tensors + g2 = grad.contiguous().view_as(x2) + dx, dweight, d_shift, d_scale = sp_rmsnorm_backward( + g2, x2, weight, rstd, shift, scale, ctx.sp, ctx.collective + ) + return dx.view(ctx.x_shape), dweight, d_shift, d_scale, None, None, None + + +class _SPGateResidual(torch.autograd.Function): + @staticmethod + def forward(ctx, residual, y, gate, sp, collective): + hidden = residual.shape[-1] + res2 = residual.contiguous().view(-1, hidden) + y2 = y.contiguous().view(-1, hidden) + out = _C.h3_gate_residual_forward(res2, y2, gate, sp.local_index) + ctx.save_for_backward(y2, gate) + ctx.sp, ctx.collective, ctx.shape = sp, collective, residual.shape + return out.view(residual.shape) + + @staticmethod + def backward(ctx, grad): + y2, gate = ctx.saved_tensors + g2 = grad.contiguous().view_as(y2) + dy, dgate = sp_gate_residual_backward(g2, y2, gate, ctx.sp, ctx.collective) + return grad, dy.view(ctx.shape), dgate, None, None + + +class H3SPNormAdaLNCudaOp: + """Sequence-parallel ``h3_rmsnorm`` (+ modulation) and ``adaln_gate_residual``. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` and + ``backend_id`` (e.g. ``DeterministicCollective``). Activations are this rank's + ``(B, S_local, H)`` rows; the norm weight, the table views and the ``(S,)`` row + index are the full, replicated tensors. Gradients of the replicated tensors come + back identical on every rank, equal to WS1. + """ + + op_class = "reduction" + backward_impl = BACKWARD_IMPL + + def __init__(self, collective, seq_len: int, batch: int) -> None: + if not sp_norm_adaln_available(): + raise RuntimeError("rl_engine._C lacks the H3 SP symbols; rebuild the CUDA extension") + self.collective = collective + self.layout = sp_row_layout(seq_len, batch, collective.world_size, collective.rank) + self._plans: list[tuple[object, int, int | None, torch.device, SPPlan]] = [] + + def plan(self, index: torch.Tensor | None, num_rows: int | None, *, device=None) -> SPPlan: + """The cached plan for this index object (rebuilt if it was modified in place).""" + + device = index.device if index is not None else torch.device(device or "cuda") + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + version = -1 if index is None else index._version + for obj, ver, rows, cached_device, plan in self._plans: + if obj is index and ver == version and rows == num_rows and cached_device == device: + return plan + plan = sp_plan(self.layout, index, num_rows, device=device) + # Holding the index keeps its storage alive, so a new tensor can never alias it. + self._plans = [(index, version, num_rows, device, plan), *self._plans[:3]] + return plan + + def _check(self, x: torch.Tensor, index: torch.Tensor | None, num_rows: int | None) -> None: + expected = (self.layout.batch, self.layout.local_len) + if not isinstance(x, torch.Tensor) or x.dim() != 3 or tuple(x.shape[:2]) != expected: + shape = tuple(x.shape) if isinstance(x, torch.Tensor) else type(x).__name__ + raise ValueError(f"activations must be this rank's {expected} + (H,) rows, got {shape}") + if not x.is_cuda: + raise ValueError("H3SPNormAdaLNCudaOp needs CUDA tensors") + if index is None: + return + if index.dim() != 1 or index.shape[0] != self.layout.seq_len: + raise ValueError( + f"index must be the full ({self.layout.seq_len},) row index, " + f"got {tuple(index.shape)}" + ) + if num_rows is None: + raise ValueError("indexed SP validation requires num_rows") + if index.numel() and (int(index.min()) < 0 or int(index.max()) >= num_rows): + raise IndexError(f"index must be in [0, {num_rows})") + + def norm(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + self._check(x, None, None) + validate_h3_rmsnorm(x, weight, eps) + return _SPRMSNorm.apply( + x, weight, None, None, self.plan(None, None, device=x.device), self.collective, eps + ) + + def norm_modulated(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + self._check(x, index, shift.shape[0]) + validate_h3_rmsnorm(x, weight, eps) + validate_h3_modulation(x, shift, scale, index[self.layout.lo : self.layout.hi]) + if shift.stride(1) != 1 or scale.stride(1) != 1 or shift.stride(0) != scale.stride(0): + shift, scale = shift.contiguous(), scale.contiguous() + sp = self.plan(index, shift.shape[0]) + return _SPRMSNorm.apply(x, weight, shift, scale, sp, self.collective, eps) + + def gate_residual(self, residual, y, gate, index) -> torch.Tensor: + self._check(residual, index, gate.shape[0]) + validate_h3_gate_residual(residual, y, gate, index[self.layout.lo : self.layout.hi]) + if gate.stride(1) != 1: + gate = gate.contiguous() + sp = self.plan(index, gate.shape[0]) + return _SPGateResidual.apply(residual, y, gate, sp, self.collective) + + def readback(self) -> dict: + return { + "op": "sp_norm_adaln", + "backward_impl": BACKWARD_IMPL, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_256_element_tiles_then_ascending_fold", + "sp": self.layout.sp, + "rank": self.layout.rank, + "positions": [self.layout.lo, self.layout.hi], + "batch": self.layout.batch, + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..4d58f6d9f --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 FP32 timestep MLP (RFC #420 ``timestep_mlp_fp32``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + ACT_NONE, + ACT_SILU, + CONTRACT, + det_linear_available, + det_linear_backward_input, + det_linear_backward_weight, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3.timestep_mlp import validate_h3_timestep_mlp + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[silu]+rl_engine._C.h3_det_linear_forward" +BACKWARD_IMPL = "h3_det_linear_v1_chunked_dinput_ascending_row_dweight" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input" + "+rl_engine.backends.cuda.model_specific.minimax_h3.det_linear.silu_backward_fp32" +) + + +class _H3TimestepMLPCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w1, b1, w2, b2): + """Evaluate the FP32 CUDA linear-SiLU-linear graph and save VJP inputs.""" + + x = x.contiguous() + hidden, pre = det_linear_forward(x, w1, b1, activation=ACT_SILU, save_pre_activation=True) + (out,) = det_linear_forward(hidden, w2, b2, activation=ACT_NONE) + ctx.save_for_backward(x, w1, w2, pre, hidden) + return out + + @staticmethod + def backward(ctx, grad_out): + """Return requested FP32 input and parameter VJPs with fixed reduction order.""" + + x, w1, w2, pre, hidden = ctx.saved_tensors + need_x, need_w1, need_b1, need_w2, need_b2 = ctx.needs_input_grad + grad_out = grad_out.float().contiguous() + dw2 = db2 = dw1 = db1 = dx = None + if need_w2 or need_b2: + dw2, db2 = det_linear_backward_weight(grad_out, hidden, torch.float32) + d_pre = None + if need_x or need_w1 or need_b1: + d_hidden = det_linear_backward_input(grad_out, w2, torch.float32) + d_pre = silu_backward_fp32(d_hidden, pre).contiguous() + if need_w1 or need_b1: + dw1, db1 = det_linear_backward_weight(d_pre, x, torch.float32) + if need_x: + dx = det_linear_backward_input(d_pre, w1, torch.float32) + record_backward( + "timestep_mlp_fp32", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return ( + dx, + dw1 if need_w1 else None, + db1 if need_b1 else None, + dw2 if need_w2 else None, + db2 if need_b2 else None, + ) + + +class H3TimestepMLPCudaOp: + """CUDA candidate: two warp-per-column GEMVs, SiLU fused into the first. + + The weights dominate the traffic (5.5 MB + 57.8 MB FP32), the rows are the + handful of distinct timesteps, and every output element follows contract + ``h3-det-linear-v1``, so a timestep's ``temb`` bytes do not depend on which + other timesteps share the call. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self) -> None: + """Require all deterministic linear symbols in the CUDA extension.""" + + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` timestep embeddings for same-device CUDA inputs.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Apply two deterministic FP32 CUDA projections with an intervening SiLU. + + Require same-device FP32 tensors: ``x`` is nonempty ``(T, K)``, + weights are ``(H, K)`` and ``(D, H)``, and biases are ``(H,)`` and + ``(D,)``. Return ``(T, D)`` with autograd support for all five inputs. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + if not x.is_cuda: + raise ValueError("H3TimestepMLPCudaOp needs CUDA tensors") + return _H3TimestepMLPCuda.apply(x, w1, b1, w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Use the standard CUDA MLP, whose inputs and output are already FP32.""" + + return self.forward(x, w1, b1, w2, b2) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..65934e1e9 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CUDA H3 sinusoidal timestep features (RFC #420 ``timestep_sinusoid_h3``).""" + +from __future__ import annotations + +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD +from rl_engine.reference.minimax_h3.fixed_order import tree_sum_lastdim_fp32 +from rl_engine.reference.minimax_h3.timestep_sinusoid import ( + NativeH3TimestepSinusoidOp, + validate_h3_timesteps, +) + +KERNEL_ID = "rl_engine._C.h3_timestep_sinusoid_forward" +BACKWARD_IMPL = "row_local_fp32_analytic_tree_sum" + + +def h3_sinusoid_cuda_available() -> bool: + """Return whether the extension provides the CUDA H3 sinusoid entry point.""" + + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +class _H3TimestepSinusoidCuda(torch.autograd.Function): + @staticmethod + def forward(ctx, timestep: torch.Tensor, num_channels: int, check_range: bool): + """Compute FP32 CUDA features and save them for the analytic timestep VJP.""" + + t32 = timestep.detach().float().contiguous() + out = _C.h3_timestep_sinusoid_forward( + t32, int(num_channels), float(H3_MAX_PERIOD), check_range + ) + ctx.save_for_backward(out) + ctx.num_channels = int(num_channels) + ctx.timestep_dtype = timestep.dtype + return out + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + """Reduce channel derivatives with a fixed FP32 tree into the timestep dtype.""" + + (out,) = ctx.saved_tensors + half = ctx.num_channels // 2 + freq = NativeH3TimestepSinusoidOp.frequencies_fp32(ctx.num_channels, out.device) + cos_part, sin_part = out[:, :half], out[:, half:] + g = grad_out.float() + # d cos(t f)/dt = -f sin(t f); d sin(t f)/dt = f cos(t f). Row-local, + # summed over channels with a fixed pairwise tree. + contrib = torch.cat( + [-(g[:, :half] * sin_part) * freq, (g[:, half:] * cos_part) * freq], dim=-1 + ) + grad_t = tree_sum_lastdim_fp32(contrib).to(ctx.timestep_dtype) + record_backward( + "timestep_sinusoid_h3", + kernel_id="rl_engine.reference.minimax_h3.fixed_order.tree_sum_lastdim_fp32", + impl=BACKWARD_IMPL, + family="cuda", + ) + return grad_t, None, None + + +class H3TimestepSinusoidCudaOp: + """CUDA candidate: one thread per (timestep, frequency), FP32 output. + + Construction fails when the extension lacks the symbol, so the registry + falls back to the PyTorch reference instead of silently mis-dispatching. + """ + + op_class = "elementwise" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + + def __init__(self) -> None: + """Require the native H3 CUDA sinusoid symbol before creating this operator.""" + + if not h3_sinusoid_cuda_available(): + raise RuntimeError( + "rl_engine._C.h3_timestep_sinusoid_forward is unavailable; rebuild the " + "CUDA extension with csrc/cuda/h3/timestep_sinusoid.cu (it is omitted " + "from KERNEL_ALIGN_USE_FAST_MATH=1 builds to keep bitwise parity)" + ) + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 CUDA features for timesteps in ``[0, 1]`` with gradients.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on CUDA. + + Require nonempty 1-D FP32, FP16 or BF16 timesteps and a positive even + channel count. Native validation rejects nonfinite/out-of-range values + when ``check_range`` is true; backward returns the original input dtype. + """ + + # Native code owns value validation, avoiding a second host sync here. + validate_h3_timesteps(timestep, num_channels, check_range=False) + if not timestep.is_cuda: + raise ValueError("H3TimestepSinusoidCudaOp needs a CUDA timestep tensor") + return _H3TimestepSinusoidCuda.apply(timestep, num_channels, check_range) + + def forward_fp32(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Use the standard CUDA sinusoid path, which already returns FP32 features.""" + + # The op is FP32 end to end; the FP32 path is the op itself. + return self.forward(timestep, num_channels=num_channels) diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py new file mode 100644 index 000000000..c4ceed334 --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Tensor-parallel H3 AdaLN projection (RFC #420 ``tp_adaln_3mod``). + +Ownership: rank ``r`` of ``tp`` owns the contiguous output columns +``[r * N / tp, (r + 1) * N / tp)`` of the ``N = 3 * 6 * H`` projection, i.e. the +same rows of ``adaln_proj.linear.weight``/``bias``. ``N / tp`` must be a +multiple of the 64-row ``d_input`` chunk of contract ``h3-det-linear-v1``. + +Every floating-point operation is the WS1 kernel itself, so every rank ends +with the WS1 bytes: + +* forward: each column depends only on ``x`` and its own weight row, so the + local GEMV produces the WS1 columns; a rank-ordered all-gather (a copy) + rebuilds the ``(T, N)`` table, and the six tensors and ``3T`` modality rows + are views of it exactly as in WS1; +* ``dW``/``db``: rows of the shard, computed locally (no cross-rank sum); +* ``d_temb``: each rank computes the WS1 chunk partials of its columns, the + all-gather puts them in global chunk order, and every rank runs the WS1 + ascending fold. No collective reduction arithmetic is used. + +The backward expects the table gradient to be identical on every TP rank +(the modulation consumes the full replicated table); it reads its own columns. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + CONTRACT, + DINPUT_CHUNK, + det_linear_available, + det_linear_backward_input_partials, + det_linear_backward_weight, + det_linear_fold_chunks, + det_linear_forward, + silu_backward_fp32, +) +from rl_engine.backends.cuda.model_specific.minimax_h3.ws2_comm import gather_rows +from rl_engine.ops.autograd.backward_runtime import record_backward +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import split_adaln_table + +KERNEL_ID = "rl_engine._C.h3_det_linear_forward[column_shard]+all_gather[rank_order]" +BACKWARD_IMPL = "local_dweight+all_gather_chunk_partials+ascending_fold" +BACKWARD_KERNEL_ID = ( + "rl_engine._C.h3_det_linear_backward_weight" + "+rl_engine._C.h3_det_linear_backward_input_partials" + "+all_gather+rl_engine._C.h3_det_linear_fold_chunks" +) + + +@dataclass(frozen=True) +class AdaLNColumnShard: + """The output columns a TP rank owns, and the WS1 objects they map to.""" + + tp: int + rank: int + n_total: int + begin: int + end: int + + @property + def hidden(self) -> int: + return self.n_total // (H3_MODALITY_NUM * H3_ADALN_CHUNKS) + + def slots(self) -> list[tuple[int, int, int, int]]: + """``(modality, chunk, h_begin, h_end)`` pieces covered by this shard.""" + + out, col = [], self.begin + while col < self.end: + slot, h0 = divmod(col, self.hidden) + h1 = min(self.hidden, h0 + self.end - col) + out.append((*divmod(slot, H3_ADALN_CHUNKS), h0, h1)) + col += h1 - h0 + return out + + +def adaln_column_shard(n_total: int, tp: int, rank: int) -> AdaLNColumnShard: + if tp < 1 or not 0 <= rank < tp: + raise ValueError(f"need 0 <= rank < tp, got rank={rank}, tp={tp}") + if n_total <= 0: + raise ValueError(f"N must be positive, got {n_total}") + if n_total % (H3_MODALITY_NUM * H3_ADALN_CHUNKS): + raise ValueError(f"N={n_total} is not 3 modalities x 6 chunks x H") + if n_total % (tp * DINPUT_CHUNK): + # A shard boundary inside a d_input chunk would split one WS1 partial. + raise ValueError( + f"tp={tp} does not split N={n_total} on {DINPUT_CHUNK}-column chunk boundaries" + ) + width = n_total // tp + return AdaLNColumnShard(tp, rank, n_total, rank * width, (rank + 1) * width) + + +def shard_adaln_projection(weight: torch.Tensor, bias: torch.Tensor, tp: int, rank: int): + """This rank's rows of the full projection weight and bias (views).""" + + shard = adaln_column_shard(weight.shape[0], tp, rank) + return weight[shard.begin : shard.end], bias[shard.begin : shard.end] + + +def _gather_columns(collective, local: torch.Tensor) -> torch.Tensor: + """``(T, N / tp)`` column shards -> the ``(T, N)`` table, in rank order (a copy).""" + + num_t, width = local.shape + rows = gather_rows(collective, local) # (tp * T, N / tp), rank-major + return rows.view(-1, num_t, width).transpose(0, 1).reshape(num_t, -1) + + +def _validate_shard_inputs(temb, weight, bias, shard: AdaLNColumnShard) -> None: + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if not tensor.is_cuda or tensor.device != temb.device: + raise ValueError(f"{name} must be on this rank's CUDA device, got {tensor.device}") + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in (torch.bfloat16, torch.float32) or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share bfloat16 or float32, got {weight.dtype}/{bias.dtype}" + ) + width = shard.end - shard.begin + if weight.dim() != 2 or weight.shape[0] != width or bias.shape != (width,): + raise ValueError( + f"rank {shard.rank} of {shard.tp} owns {width} rows, got weight " + f"{tuple(weight.shape)} and bias {tuple(bias.shape)}" + ) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + + +class _TPAdaLNProjection(torch.autograd.Function): + @staticmethod + def forward(ctx, temb, weight_shard, bias_shard, collective, shard): + temb = temb.contiguous() + act = F.silu(temb).to(weight_shard.dtype) + (local,) = det_linear_forward(act, weight_shard, bias_shard) + ctx.save_for_backward(temb, act, weight_shard) + ctx.collective, ctx.shard = collective, shard + return _gather_columns(collective, local) + + @staticmethod + def backward(ctx, grad_table): + temb, act, weight = ctx.saved_tensors + shard = ctx.shard + grad_local = grad_table[:, shard.begin : shard.end].float().contiguous() + dw, db = det_linear_backward_weight(grad_local, act, weight.dtype) + partial = det_linear_backward_input_partials(grad_local, weight) + d_act = det_linear_fold_chunks(gather_rows(ctx.collective, partial), torch.float32) + record_backward( + "tp_adaln_3mod", kernel_id=BACKWARD_KERNEL_ID, impl=BACKWARD_IMPL, family="cuda" + ) + return silu_backward_fp32(d_act, temb), dw, db, None, None + + +class H3TPAdaLNProjectionCudaOp: + """Column-parallel ``adaln_projection_3mod``, byte-equal to WS1 on every rank. + + ``collective`` is a rank-ordered all-gather with ``rank``, ``world_size`` + and ``backend_id`` (e.g. ``DeterministicCollective``). Inputs are the replicated + FP32 ``temb`` and this rank's weight/bias rows. + """ + + op_class = "reduction" + kernel_id = KERNEL_ID + backward_impl = BACKWARD_IMPL + contract = CONTRACT + + def __init__(self, collective, n_total: int) -> None: + if not det_linear_available(): + raise RuntimeError( + "rl_engine._C lacks h3_det_linear_*; rebuild with csrc/cuda/h3/det_linear.cu" + ) + self.collective = collective + self.shard = adaln_column_shard(n_total, collective.world_size, collective.rank) + + def __call__(self, temb, weight_shard, bias_shard): + return self.forward(temb, weight_shard, bias_shard) + + def forward(self, temb, weight_shard, bias_shard) -> tuple[torch.Tensor, ...]: + table = self.forward_table(temb, weight_shard, bias_shard) + return split_adaln_table(table, self.shard.hidden) + + def forward_table(self, temb, weight_shard, bias_shard) -> torch.Tensor: + _validate_shard_inputs(temb, weight_shard, bias_shard, self.shard) + return _TPAdaLNProjection.apply(temb, weight_shard, bias_shard, self.collective, self.shard) + + def readback(self) -> dict: + """Runtime identity of this rank's projection, for evidence and strict traces.""" + + return { + "op": "tp_adaln_3mod", + "kernel_id": KERNEL_ID, + "backward_kernel_id": BACKWARD_KERNEL_ID, + "contract": CONTRACT, + "collective_backend": getattr( + self.collective, "backend_id", type(self.collective).__name__ + ), + "collective_ops": ["all_gather"], + "reduction_order": "ws1_ascending_chunk_fold_after_rank_ordered_gather", + "tp": self.shard.tp, + "rank": self.shard.rank, + "columns": [self.shard.begin, self.shard.end], + "slots": [list(s) for s in self.shard.slots()], + "fallback": None, + } diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py new file mode 100644 index 000000000..f03597b7d --- /dev/null +++ b/rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Rank-ordered row gather shared by the H3 WS2 ops (a copy: no arithmetic).""" + +from __future__ import annotations + +import torch + +# DeterministicCollective sends all-gathers whose output is at most this many bytes +# through a single-block byte-copy kernel (csrc/cuda/distributed/ +# deterministic_collective.cu, kSingleBlockFastPathMaxBytes) that costs about 4 us per +# KiB on B200; larger outputs take the multi-block path (~45 us). Padding the shard +# past the threshold is cheaper for every non-trivial size. +_SINGLE_BLOCK_GATHER_MAX_BYTES = 256 * 1024 + + +def gather_rows(collective, shard: torch.Tensor) -> torch.Tensor: + """``(world * rows, ...)``: every rank's ``(rows, ...)`` shard, in rank order.""" + + shard = shard.contiguous() + rows, world = shard.shape[0], collective.world_size + row_bytes = shard[0].numel() * shard.element_size() if rows else 0 + if world == 1: + return shard + if row_bytes == 0: + return collective.all_gather(shard) + padded = max(rows, _SINGLE_BLOCK_GATHER_MAX_BYTES // (world * row_bytes) + 1) + if padded == rows: + return collective.all_gather(shard) + staged = torch.cat([shard, shard.new_zeros((padded - rows, *shard.shape[1:]))]) + out = collective.all_gather(staged).view(world, padded, *shard.shape[1:]) + return out[:, :rows].reshape(world * rows, *shard.shape[1:]) diff --git a/rl_engine/reference/minimax_h3/__init__.py b/rl_engine/reference/minimax_h3/__init__.py new file mode 100644 index 000000000..a01f9218b --- /dev/null +++ b/rl_engine/reference/minimax_h3/__init__.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""MiniMax-H3 (RFC #420) PyTorch goldens. + +Shape constants are pinned to ``MiniMaxAI/MiniMax-H3@42ed227`` (see +``rl_engine/validation/models/h3_manifest.json``). The ops accept other sizes so that +tiny synthetic shapes can be tested, but the H3 layout rules (three modality +rows per timestep, six modulation chunks) are fixed. +""" + +from __future__ import annotations + +H3_FREQ_DIM = 256 +H3_MAX_PERIOD = 10000 +H3_TIME_EMBED_HIDDEN_DIM = 5376 +H3_TIME_EMBED_DIM = 2688 +H3_HIDDEN_SIZE = 5376 +# Modality tags of the packed sequence: 0 video, 1 text, 2 audio. +H3_MODALITY_NUM = 3 +# shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp +H3_ADALN_CHUNKS = 6 +H3_ADALN_CHUNK_NAMES = ( + "shift_msa", + "scale_msa", + "gate_msa", + "shift_mlp", + "scale_mlp", + "gate_mlp", +) diff --git a/rl_engine/reference/minimax_h3/adaln_projection.py b/rl_engine/reference/minimax_h3/adaln_projection.py new file mode 100644 index 000000000..be83239c5 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_projection.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 three-modality AdaLN projection (RFC #420 row ``adaln_projection_3mod``). + +One ``MiniMaxH3AdaLayerNormModulation`` per transformer block:: + + act = silu(temb).to(weight.dtype) SiLU in FP32, one cast to BF16 + table = act @ W.T + b (T, 6 * H * 3), BF16 + rows = table.view(3 * T, 6 * H) row t * 3 + m, m in {video, text, audio} + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = rows.chunk(6, -1) + +Output channel ``o = m * 6H + c * H + h`` of the projection is chunk ``c``, hidden +index ``h`` of modality ``m``. Each of the six outputs has shape ``(3T, H)``, +and row ``t * 3 + m`` is what ``timestep_indices * 3 + token_tags`` addresses. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_WEIGHT_DTYPES = (torch.bfloat16, torch.float32) + + +def adaln_hidden_size(weight: torch.Tensor) -> int: + """Infer ``H`` from a 2-D projection weight whose row count is divisible by 18.""" + + rows_per_hidden = H3_ADALN_CHUNKS * H3_MODALITY_NUM + if weight.dim() != 2 or weight.shape[0] % rows_per_hidden != 0: + raise ValueError( + f"AdaLN weight must be (6 * H * 3, D); got {tuple(weight.shape)}, whose first " + f"dim is not a multiple of {rows_per_hidden}" + ) + return weight.shape[0] // rows_per_hidden + + +def validate_h3_adaln_projection(temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor): + """Check the AdaLN input shapes, shared device and precision boundary; return ``H``. + + Require nonempty FP32 ``temb`` of shape ``(T, D)`` and FP32 or BF16 weight + of shape ``(18H, D)`` with matching-dtype bias of shape ``(18H,)``. + """ + + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.device != temb.device: + raise ValueError(f"{name} is on {tensor.device}, temb is on {temb.device}") + if temb.dtype != torch.float32: + # RFC probe H7: the activation must run at the timestep-embedding + # precision; a BF16 temb means the cast happened before the SiLU. + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if weight.dtype not in _WEIGHT_DTYPES or bias.dtype != weight.dtype: + raise TypeError( + f"weight/bias must share a dtype in {_WEIGHT_DTYPES}, got {weight.dtype}/{bias.dtype}" + ) + hidden = adaln_hidden_size(weight) + if temb.dim() != 2 or temb.shape[0] == 0 or temb.shape[1] != weight.shape[1]: + raise ValueError( + f"temb must be a non-empty (T, {weight.shape[1]}) matrix, got {tuple(temb.shape)}" + ) + if bias.shape != (weight.shape[0],): + raise ValueError(f"bias must be ({weight.shape[0]},), got {tuple(bias.shape)}") + return hidden + + +def split_adaln_table(table: torch.Tensor, hidden: int) -> tuple[torch.Tensor, ...]: + """(T, 6H*3) projection output -> six (3T, H) views, diffusers order.""" + + return table.view(-1, H3_ADALN_CHUNKS * hidden).chunk(H3_ADALN_CHUNKS, dim=-1) + + +class NativeH3AdaLNProjectionOp: + """PyTorch reference for the H3 AdaLN projection. + + ``forward`` is the provider path (``F.silu`` in FP32, cast, ``F.linear``). + ``forward_fp32`` is the golden: the SiLU in FP64, rounded once to the + weight dtype at the declared boundary (the cast is model semantics, not + an implementation detail), then the projection in FP64. The six outputs are + returned in FP32 without the final BF16 rounding. Its gradient is FP64 + end to end: the cast is applied straight-through (identity VJP). + """ + + op_class = "reduction" + + def __call__(self, temb, weight, bias): + """Return six weight-dtype ``(3T, H)`` views using the PyTorch provider path.""" + + return self.forward(temb, weight, bias) + + def forward(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Apply FP32 SiLU, cast to the weight dtype and project on the input device. + + Validate ``(T, D)`` embeddings and ``(18H, D)``/``(18H,)`` parameters; + return six weight-dtype ``(3T, H)`` views with PyTorch autograd support. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + table = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + return split_adaln_table(table, hidden) + + def forward_fp32(self, temb, weight, bias) -> tuple[torch.Tensor, ...]: + """Return six FP32 golden outputs from FP64 SiLU and projection arithmetic. + + Preserve the declared activation rounding to the weight dtype while + using an identity VJP at that cast boundary for the FP64 golden graph. + """ + + hidden = validate_h3_adaln_projection(temb, weight, bias) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # The declared cast rounds the FP32 SiLU, as the model (and the kernel) + # computes it: an FP64 SiLU can round to the other weight-dtype neighbour + # near a tie and shift a whole output column. Its VJP is the identity, + # written straight-through so autograd does not round the FP64 gradient + # to the weight dtype on its way back (``.to(bf16)`` would). + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + return tuple(chunk.float() for chunk in split_adaln_table(table, hidden)) diff --git a/rl_engine/reference/minimax_h3/adaln_row_gather.py b/rl_engine/reference/minimax_h3/adaln_row_gather.py new file mode 100644 index 000000000..2cd713d45 --- /dev/null +++ b/rl_engine/reference/minimax_h3/adaln_row_gather.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 AdaLN row gather (RFC #420 row ``adaln_row_gather``). + +H3 computes one row index per packed-sequence position and uses it to pick +every position's six modulation vectors:: + + adaln_indices = timestep_indices * 3 + token_tags (S,) + shift_msa[adaln_indices], scale_msa[adaln_indices], ... six (S, H) + +``rows`` is the ``(3T, 6H)`` view of one block's AdaLN projection +(``table.view(-1, 6H)``); its six column blocks are the six modulation +tensors. ``timestep_indices`` and ``token_tags`` are semantic inputs +(RFC #420 §4): out-of-range values are an error, never clamped. +""" + +from __future__ import annotations + +import torch + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM + +_INDEX_DTYPES = (torch.int64, torch.int32) + + +def validate_h3_adaln_row_gather( + rows: torch.Tensor, + timestep_indices: torch.Tensor, + token_tags: torch.Tensor, + *, + check_range: bool = True, +) -> int: + """Return the hidden size; raise on anything H3 does not define.""" + + if not isinstance(rows, torch.Tensor) or rows.dim() != 2: + raise ValueError("rows must be a 2-D (3T, 6H) tensor") + if not rows.is_floating_point(): + raise TypeError(f"rows must be floating point, got {rows.dtype}") + num_rows, width = rows.shape + if num_rows == 0 or num_rows % H3_MODALITY_NUM != 0: + raise ValueError(f"rows must hold 3 rows per timestep, got {num_rows}") + if width == 0 or width % H3_ADALN_CHUNKS != 0: + raise ValueError(f"rows width must be 6 * H, got {width}") + for name, index in (("timestep_indices", timestep_indices), ("token_tags", token_tags)): + if not isinstance(index, torch.Tensor) or index.dim() != 1: + raise ValueError(f"{name} must be a 1-D (S,) tensor") + if index.dtype not in _INDEX_DTYPES: + raise TypeError(f"{name} must be int64 or int32, got {index.dtype}") + if index.device != rows.device: + raise ValueError(f"{name} is on {index.device}, rows are on {rows.device}") + if timestep_indices.shape != token_tags.shape: + raise ValueError( + f"timestep_indices {tuple(timestep_indices.shape)} and token_tags " + f"{tuple(token_tags.shape)} must have the same length" + ) + if timestep_indices.numel() == 0: + raise ValueError("the packed sequence must not be empty") + if timestep_indices.dtype != token_tags.dtype: + raise TypeError("timestep_indices and token_tags must share an integer dtype") + if check_range: + num_timesteps = num_rows // H3_MODALITY_NUM + bad = ( + (timestep_indices < 0) + | (timestep_indices >= num_timesteps) + | (token_tags < 0) + | (token_tags >= H3_MODALITY_NUM) + ) + if bool(bad.any()): # one host sync + raise IndexError( + f"timestep_indices must lie in [0, {num_timesteps}) and token_tags in " + f"[0, {H3_MODALITY_NUM}) (0 video, 1 text, 2 audio)" + ) + return width // H3_ADALN_CHUNKS + + +def h3_adaln_indices(timestep_indices: torch.Tensor, token_tags: torch.Tensor) -> torch.Tensor: + return timestep_indices.long() * H3_MODALITY_NUM + token_tags.long() + + +class _DeterministicRowGather(torch.autograd.Function): + """``index_select`` forward; FP32 per-row segment sums for the backward. + + ``index_select``'s own backward is an atomic scatter-add in the grad dtype + (BF16 for H3), which is neither deterministic nor accurate. There are + only 3T table rows, so the reference sums each row's positions with one + FP32 ``torch.sum`` per row (a fixed reduction for a fixed packing) and + rounds once. + """ + + @staticmethod + def forward(ctx, rows, index): + ctx.save_for_backward(index) + ctx.rows_meta = (rows.shape[0], rows.dtype) + return torch.stack( + [chunk.index_select(0, index) for chunk in rows.chunk(H3_ADALN_CHUNKS, dim=-1)] + ) + + @staticmethod + def backward(ctx, grad): + (index,) = ctx.saved_tensors + num_rows, dtype = ctx.rows_meta + grad32 = grad.float() # (6, S, H) + out = grad32.new_zeros((num_rows, grad32.shape[0], grad32.shape[2])) + for row in range(num_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + out[row] = grad32.index_select(1, positions).sum(dim=1) + return out.reshape(num_rows, -1).to(dtype), None + + +class NativeH3AdaLNRowGatherOp: + """PyTorch reference: ``index_select`` on each of the six column blocks. + + The forward is byte-identical to diffusers; the backward is deterministic + (see ``_DeterministicRowGather``). The raw diffusers path, including its + atomic backward, is ``rl_engine.validation.models.h3_provider.provider_adaln_row_gather``. + ``forward_fp32`` returns the same rows in FP32 with an FP64 gradient. + + The forward is a copy, but the VJP is a segmented reduction over the + packed sequence, so the op is judged with the ``reduction`` tolerance + class; forward bitwise equality is asserted separately. + """ + + op_class = "reduction" + + def __call__(self, rows, timestep_indices, token_tags): + return self.forward(rows, timestep_indices, token_tags) + + def forward(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple(_DeterministicRowGather.apply(rows, index).unbind(0)) + + def forward_fp32(self, rows, timestep_indices, token_tags, *, check_range: bool = True): + validate_h3_adaln_row_gather(rows, timestep_indices, token_tags, check_range=check_range) + index = h3_adaln_indices(timestep_indices, token_tags) + return tuple( + chunk.index_select(0, index).float() + for chunk in rows.double().chunk(H3_ADALN_CHUNKS, dim=-1) + ) diff --git a/rl_engine/reference/minimax_h3/block.py b/rl_engine/reference/minimax_h3/block.py new file mode 100644 index 000000000..427d7de2e --- /dev/null +++ b/rl_engine/reference/minimax_h3/block.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""FP32 golden of one MiniMax-H3 transformer block (RFC #420 ``ws1_one_h3_block``). + +Every node takes FP32 inputs (BF16 weights and activations upcast exactly), +computes in FP32 with TF32 off, and returns FP32 without rounding back to +BF16, so the golden is the block's arithmetic without the BF16 storage the +provider and the candidate both have. It is the yardstick their errors are +measured against, not a bitwise target. The one declared cast kept from the +checkpoint contract is the AdaLN input: ``silu(temb)`` is rounded to BF16 +before the projection, as every implementation must do it. + +The partial MM-RoPE uses exact FP32 ``cos``/``sin`` tables on the leading 96 +of 128 head channels (the provider rounds the tables to BF16 first). +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + +H3_HEADS = 56 +H3_HEAD_DIM = 128 +H3_ROTARY_DIM = 96 +H3_NORM_EPS = 1e-5 + + +def _f32(t: torch.Tensor) -> torch.Tensor: + return t.to(torch.float32) + + +def golden_adaln_tables(temb, weight, bias, hidden_size: int = 5376) -> tuple[torch.Tensor, ...]: + """Six ``(3T, H)`` FP32 tables; ``silu(temb)`` rounded to the weight dtype first.""" + + act = F.silu(_f32(temb)).to(weight.dtype) + table = F.linear(_f32(act), _f32(weight), _f32(bias)) + return table.view(-1, 6 * hidden_size).chunk(6, dim=-1) + + +def golden_rmsnorm(x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + x = _f32(x) + rstd = torch.rsqrt(x.square().mean(dim=-1, keepdim=True) + eps) + return x * rstd * _f32(weight) + + +def golden_norm_modulate(x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + n = golden_rmsnorm(x, weight, eps) + return n * (1.0 + _f32(scale).index_select(0, index)) + _f32(shift).index_select(0, index) + + +def golden_linear(x, weight) -> torch.Tensor: + return F.linear(_f32(x), _f32(weight)) + + +def golden_rope_tables(position_ids, rope_freq_dim: int = 16, rope_theta: float = 10000.0): + """FP32 ``(S, 96)`` cos/sin, the same angles as the provider's rotary module.""" + + exponent = torch.arange( + 0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device + ) + inv_freq = 1.0 / (rope_theta ** (exponent / (2 * rope_freq_dim))) + freqs = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs = torch.cat(freqs.unbind(dim=1), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def golden_apply_rotary(x, cos, sin) -> torch.Tensor: + """Rotate-half on the leading ``cos.shape[-1]`` channels of ``(B, S, H, D)``; FP32 tables.""" + + x = _f32(x) + rot = cos.shape[-1] + head, tail = x[..., :rot], x[..., rot:] + x1, x2 = head.chunk(2, dim=-1) + rotated = torch.cat((-x2, x1), dim=-1) + head = head * cos[None, :, None, :] + rotated * sin[None, :, None, :] + return torch.cat((head, tail), dim=-1) + + +def golden_attention(q, k, v) -> torch.Tensor: + """Non-causal full attention over one packed document, ``(B, S, H, D)`` -> ``(B, S, H*D)``.""" + + q, k, v = (_f32(t).permute(0, 2, 1, 3) for t in (q, k, v)) + scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(q.shape[-1]) + out = torch.matmul(torch.softmax(scores, dim=-1), v) + return out.permute(0, 2, 1, 3).flatten(2, 3) + + +def golden_swiglu(projected) -> torch.Tensor: + hidden, gate = _f32(projected).chunk(2, dim=-1) + return hidden * F.silu(gate) + + +def golden_gate_residual(residual, gate, index, y) -> torch.Tensor: + return _f32(residual) + _f32(gate).index_select(0, index) * _f32(y) diff --git a/rl_engine/reference/minimax_h3/final_adaln_out.py b/rl_engine/reference/minimax_h3/final_adaln_out.py new file mode 100644 index 000000000..a21a5cbdb --- /dev/null +++ b/rl_engine/reference/minimax_h3/final_adaln_out.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 final AdaLN output (RFC #420 row ``final_adaln_out``). + +``MiniMaxH3AdaLayerNormOut`` runs after the 50 blocks:: + + shift, scale = linear(silu(temb).to(bf16)).chunk(2) # (T, H) each, shift first + out = norm(x) * (1.0 + scale[timestep_indices]) + shift[timestep_indices] + +``linear`` is ``norm_out.linear`` (``2688 -> 2 * 5376``, BF16), ``norm`` is +``norm_out.norm`` (``nn.RMSNorm(5376, eps=1e-5)``), and the table is indexed +by ``timestep_indices`` (one row per distinct timestep), not by +``adaln_indices``. Diffusers then casts the result to the FP32 output heads' +dtype, an exact upcast left to the caller. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.rmsnorm import ( + H3_NORM_EPS, + NativeH3RMSNormOp, + validate_h3_rmsnorm, +) + + +def validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, *, same_dtype: bool = True +) -> int: + hidden = x.shape[-1] + if same_dtype: + validate_h3_rmsnorm(x, norm_weight, eps) + if temb.dtype != torch.float32: + raise TypeError(f"temb must be float32 (the SiLU runs before the cast), got {temb.dtype}") + if temb.dim() != 2 or temb.shape[0] == 0: + raise ValueError(f"temb must be a non-empty (T, D) matrix, got {tuple(temb.shape)}") + if weight.shape != (2 * hidden, temb.shape[1]) or bias.shape != (2 * hidden,): + raise ValueError( + f"norm_out.linear must be ({2 * hidden}, {temb.shape[1]}) with a ({2 * hidden},) " + f"bias, got {tuple(weight.shape)} / {tuple(bias.shape)}" + ) + if bias.dtype != weight.dtype: + raise TypeError("norm_out.linear weight and bias must share a dtype") + for name, tensor in (("temb", temb), ("weight", weight), ("bias", bias)): + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if timestep_indices.dtype != torch.int64 or timestep_indices.dim() != 1: + raise TypeError("timestep_indices must be a 1-D int64 tensor") + if x.dim() < 2 or x.shape[-2] != timestep_indices.shape[0]: + raise ValueError("x must be (..., S, H) with S = len(timestep_indices)") + if bool(((timestep_indices < 0) | (timestep_indices >= temb.shape[0])).any()): + raise IndexError(f"timestep_indices must lie in [0, {temb.shape[0]})") + return hidden + + +class NativeH3FinalAdaLNOutOp: + """PyTorch reference: ``forward`` replays diffusers' ``norm_out``; + ``forward_fp32`` is the FP64 golden, returned in FP32. It rounds only where + the model itself stores a value in its own dtype (the SiLU cast, the BF16 + table, the ``norm_out.norm`` output and ``1 + scale``), straight-through for + the gradient, and computes everything else in FP64.""" + + op_class = "reduction" + + def __call__(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + return self.forward(x, norm_weight, temb, weight, bias, timestep_indices, eps) + + def forward(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out(x, norm_weight, temb, weight, bias, timestep_indices, eps) + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + return NativeH3RMSNormOp().forward_modulated( + x, norm_weight, shift, scale, timestep_indices, eps + ) + + def forward_fp32(self, x, norm_weight, temb, weight, bias, timestep_indices, eps=H3_NORM_EPS): + validate_h3_final_adaln_out( + x, norm_weight, temb, weight, bias, timestep_indices, eps, same_dtype=False + ) + t64 = temb.double() + act = t64 * torch.sigmoid(t64) + # Declared cast of the FP32 SiLU, as the model computes it (an FP64 SiLU can + # round to the other neighbour near a tie); identity VJP. + declared = F.silu(temb.float()).to(weight.dtype).double() + act = act + (declared - act).detach() + table = F.linear(act, weight.double(), bias.double()) + # norm_out.linear is a BF16 module: its (T, 2H) output is rounded to the + # weight dtype, and every position of a timestep shares that row, so the + # rounding is model semantics (straight-through for the gradient). + table = table + (table.to(weight.dtype).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * norm_weight.double() + # norm_out.norm is an x.dtype module and diffusers forms 1 + scale in + # x.dtype, so both are stored rounded. 1 + scale is shared by every + # position of a timestep: left unrounded, its error is systematic and + # grows with S in d_norm_weight; n's rounding enters d_scale (and so + # dW, d_temb) summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.index_select(0, timestep_indices) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.index_select(0, timestep_indices) + return out.float() diff --git a/rl_engine/reference/minimax_h3/fixed_order.py b/rl_engine/reference/minimax_h3/fixed_order.py new file mode 100644 index 000000000..e95a66054 --- /dev/null +++ b/rl_engine/reference/minimax_h3/fixed_order.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shape-independent FP32 reductions shared by the H3 goldens and backwards. + +``torch.sum`` picks its reduction tree from the tensor shape, so a row's +result can change with the number of rows next to it. These helpers spell +the order out with elementwise ops only, which keeps every row's bytes a +function of that row alone. +""" + +from __future__ import annotations + +import torch + + +def tree_sum_lastdim_fp32(x: torch.Tensor) -> torch.Tensor: + """Pairwise tree over the last dim: pad to a power of two, halve until 1. + + Level ``l`` adds element ``i`` and ``i + width/2`` for every ``i``. The + tree depends only on the last-dim size. + """ + + x = x.float() + width = x.shape[-1] + if width == 0: + return x.new_zeros(x.shape[:-1]) + padded = 1 << (width - 1).bit_length() + if padded != width: + x = torch.nn.functional.pad(x, (0, padded - width)) + while x.shape[-1] > 1: + half = x.shape[-1] // 2 + x = x[..., :half] + x[..., half:] + return x[..., 0] + + +def fold_rows_fp32(rows: torch.Tensor) -> torch.Tensor: + """Left fold over dim 0 in ascending row order: ((r0 + r1) + r2) + ...""" + + rows = rows.float() + if rows.shape[0] == 0: + return rows.new_zeros(rows.shape[1:]) + acc = rows[0].clone() + for index in range(1, rows.shape[0]): + acc = acc + rows[index] + return acc diff --git a/rl_engine/reference/minimax_h3/gate_residual.py b/rl_engine/reference/minimax_h3/gate_residual.py new file mode 100644 index 000000000..f68a0fdf4 --- /dev/null +++ b/rl_engine/reference/minimax_h3/gate_residual.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 gated residual (RFC #420 row ``adaln_gate_residual``). + +After each sublayer of a block, diffusers adds the sublayer output back to +the residual stream through the per-row AdaLN gate:: + + hidden = residual + gate.index_select(0, adaln_indices) * sublayer_output + +``gate`` is ``gate_msa`` (after attention) or ``gate_mlp`` (after the FFN), +an ``(R, H)`` row view of the AdaLN table; ``residual`` and the sublayer +output are ``(..., S, H)``. The order is fixed by RFC #420 §4: +``residual + gate[row] * sublayer_output``. +""" + +from __future__ import annotations + +import torch + +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_gate_residual( + residual: torch.Tensor, + y: torch.Tensor, + gate: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + if residual.dtype not in _DTYPES: + raise TypeError(f"residual must be fp32, bf16 or fp16, got {residual.dtype}") + if y.shape != residual.shape or y.dtype != residual.dtype or y.device != residual.device: + raise ValueError("y must match the residual's shape, dtype and device") + if residual.dim() < 2 or residual.numel() == 0: + raise ValueError( + f"residual must be a non-empty (..., S, H) tensor, got {tuple(residual.shape)}" + ) + hidden = residual.shape[-1] + if gate.dim() != 2 or gate.shape[1] != hidden or gate.device != residual.device: + raise ValueError(f"gate must be (R, {hidden}) on the residual's device") + if same_dtype and gate.dtype != residual.dtype: + raise TypeError(f"gate dtype {gate.dtype} must match the residual dtype {residual.dtype}") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != residual.device: + raise TypeError("index must be a 1-D int64 tensor on the residual's device") + if residual.shape[-2] != index.shape[0]: + raise ValueError(f"residual must be (..., S, H) with S = len(index) = {index.shape[0]}") + if check_range and bool(((index < 0) | (index >= gate.shape[0])).any()): + raise IndexError(f"index must lie in [0, {gate.shape[0]})") + + +class _DeterministicGateResidual(torch.autograd.Function): + """Eager forward; the gate gradient as FP32 per-row sums instead of + ``index_select``'s BF16 atomic scatter-add (``d_residual``/``d_y`` are the + eager VJPs).""" + + @staticmethod + def forward(ctx, residual, y, gate, index): + gathered = gate.index_select(0, index) + ctx.save_for_backward(y, gathered, index) + ctx.gate_rows = gate.shape[0] + return residual + gathered * y + + @staticmethod + def backward(ctx, grad): + y, gathered, index = ctx.saved_tensors + contrib = (grad.float() * y.float()).reshape(-1, index.shape[0], y.shape[-1]) + per_position = contrib.sum(dim=0) # (S, H): fixed reduction over the batch + d_gate = per_position.new_zeros((ctx.gate_rows, y.shape[-1])) + for row in range(ctx.gate_rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_gate[row] = per_position.index_select(0, positions).sum(dim=0) + return grad, grad * gathered, d_gate.to(gathered.dtype), None + + +class NativeH3GateResidualOp: + """PyTorch reference: ``forward`` is the eager provider expression (with a + deterministic gate gradient), ``forward_fp32`` the FP64 golden returned in + FP32. The raw diffusers path, atomic backward included, is + ``rl_engine.validation.models.h3_provider.provider_gate_residual``.""" + + op_class = "reduction" # forward is elementwise; the gate VJP sums positions + + def __call__(self, residual, y, gate, index): + return self.forward(residual, y, gate, index) + + def forward(self, residual, y, gate, index, *, check_range: bool = True): + validate_h3_gate_residual(residual, y, gate, index, check_range=check_range) + return _DeterministicGateResidual.apply(residual, y, gate, index) + + def forward_fp32(self, residual, y, gate, index): + validate_h3_gate_residual(residual, y, gate, index, same_dtype=False) + out = residual.double() + gate.double().index_select(0, index) * y.double() + return out.float() diff --git a/rl_engine/reference/minimax_h3/rmsnorm.py b/rl_engine/reference/minimax_h3/rmsnorm.py new file mode 100644 index 000000000..1727fd1d0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/rmsnorm.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 RMSNorm and its fused AdaLN modulation (RFC #420 row ``h3_rmsnorm``). + +Every H3 RMSNorm is ``nn.RMSNorm(5376, eps=1e-5)`` with an affine weight in the +block dtype: the block ``norm1``/``norm2``, the token-refiner norms and the +final ``norm_out.norm``. Inside a transformer block (and in ``norm_out``) the +normalised rows are immediately modulated by per-row AdaLN parameters:: + + n = rms_norm(x, weight, eps) + out = n * (1.0 + scale[index]) + shift[index] # eager, rounded after every op + +``index`` is ``adaln_indices`` in a block and ``timestep_indices`` in +``norm_out``; ``shift``/``scale`` are ``(R, H)`` rows (views of the AdaLN +table). ``x`` is ``(..., S, H)`` and ``index`` has one entry per position ``S``. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + +H3_NORM_EPS = 1e-5 +_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> int: + for name, tensor in (("x", x), ("weight", weight)): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype not in _DTYPES: + raise TypeError(f"{name} must be fp32, bf16 or fp16, got {tensor.dtype}") + if weight.dtype != x.dtype: + raise TypeError(f"weight dtype {weight.dtype} must match x dtype {x.dtype}") + if weight.device != x.device: + raise ValueError(f"weight is on {weight.device}, x is on {x.device}") + if x.dim() < 2 or x.numel() == 0: + raise ValueError(f"x must be a non-empty (..., N) tensor, got {tuple(x.shape)}") + if weight.shape != (x.shape[-1],): + raise ValueError(f"weight must be ({x.shape[-1]},), got {tuple(weight.shape)}") + if not eps > 0: + raise ValueError(f"eps must be positive, got {eps}") + return x.shape[-1] + + +def validate_h3_modulation( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + index: torch.Tensor, + *, + check_range: bool = True, + same_dtype: bool = True, +) -> None: + hidden = x.shape[-1] + for name, tensor in (("shift", shift), ("scale", scale)): + if tensor.device != x.device or (same_dtype and tensor.dtype != x.dtype): + raise TypeError(f"{name} must match x's dtype and device") + if tensor.dim() != 2 or tensor.shape[1] != hidden: + raise ValueError(f"{name} must be (R, {hidden}), got {tuple(tensor.shape)}") + if shift.shape != scale.shape: + raise ValueError("shift and scale must have the same shape") + if index.dtype != torch.int64 or index.dim() != 1 or index.device != x.device: + raise TypeError("index must be a 1-D int64 tensor on x's device") + if x.dim() < 2 or x.shape[-2] != index.shape[0]: + raise ValueError( + f"x must be (..., S, {hidden}) with S = len(index) = {index.shape[0]}, " + f"got {tuple(x.shape)}" + ) + if check_range and bool(((index < 0) | (index >= shift.shape[0])).any()): + raise IndexError(f"index must lie in [0, {shift.shape[0]})") + + +class _DeterministicModulation(torch.autograd.Function): + """Eager ``n * (1.0 + scale[i]) + shift[i]``; the ``shift``/``scale`` + gradients as FP32 per-row sums instead of ``index_select``'s atomic + scatter-add (``d_n`` is the eager VJP).""" + + @staticmethod + def forward(ctx, n, shift, scale, index): + t1 = 1.0 + scale.index_select(0, index) + ctx.save_for_backward(n, t1, index) + ctx.rows = shift.shape[0] + return n * t1 + shift.index_select(0, index) + + @staticmethod + def backward(ctx, grad): + n, t1, index = ctx.saved_tensors + hidden = n.shape[-1] + g = grad.float().reshape(-1, index.shape[0], hidden) + per_shift = g.sum(dim=0) + per_scale = (g * n.float().reshape(-1, index.shape[0], hidden)).sum(dim=0) + d_shift = per_shift.new_zeros((ctx.rows, hidden)) + d_scale = per_scale.new_zeros((ctx.rows, hidden)) + for row in range(ctx.rows): + positions = torch.nonzero(index == row).flatten() + if positions.numel(): + d_shift[row] = per_shift.index_select(0, positions).sum(dim=0) + d_scale[row] = per_scale.index_select(0, positions).sum(dim=0) + return grad * t1, d_shift.to(t1.dtype), d_scale.to(t1.dtype), None + + +class NativeH3RMSNormOp: + """PyTorch reference. + + ``forward`` / ``forward_modulated`` are the provider path (``F.rms_norm`` and + the eager modulation expression, with a deterministic shift/scale gradient; + the raw diffusers path is ``h3_provider.provider_norm_modulate``). ``forward_fp32`` / + ``forward_modulated_fp32`` are the golden: the same math in FP64, returned in + FP32. The modulated golden rounds, straight-through, only the two values the + model stores in ``x``'s dtype: ``norm(x)`` and ``1 + scale``. + """ + + op_class = "reduction" + + def __call__(self, x, weight, eps: float = H3_NORM_EPS): + return self.forward(x, weight, eps) + + def forward(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + hidden = validate_h3_rmsnorm(x, weight, eps) + return F.rms_norm(x, (hidden,), weight, eps) + + def forward_fp32(self, x, weight, eps: float = H3_NORM_EPS) -> torch.Tensor: + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + rstd = torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) + return (x64 * rstd * weight.double()).float() + + def forward_modulated( + self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS, *, check_range=True + ): + validate_h3_modulation(x, shift, scale, index, check_range=check_range) + n = self.forward(x, weight, eps) + return _DeterministicModulation.apply(n, shift, scale, index) + + def forward_modulated_fp32(self, x, weight, shift, scale, index, eps: float = H3_NORM_EPS): + # The golden may be fed FP32 golden modulation rows for a BF16 x. + validate_h3_modulation(x, shift, scale, index, same_dtype=False) + validate_h3_rmsnorm(x, weight, eps) + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(dim=-1, keepdim=True) + eps) * weight.double() + # The norm is an x.dtype module and diffusers forms 1 + scale in x.dtype. + # 1 + scale is shared by every position of a row, so leaving it unrounded + # skews d_norm_weight systematically with S; n's rounding enters d_scale + # summed over S. Identity in FP32. + n = n + (n.to(x.dtype).double() - n).detach() + one_plus_scale = 1.0 + scale.double().index_select(0, index) + one_plus_scale = ( + one_plus_scale + (one_plus_scale.to(x.dtype).double() - one_plus_scale).detach() + ) + out = n * one_plus_scale + shift.double().index_select(0, index) + return out.float() diff --git a/rl_engine/reference/minimax_h3/timestep_mlp.py b/rl_engine/reference/minimax_h3/timestep_mlp.py new file mode 100644 index 000000000..51ef48db7 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_mlp.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 FP32 timestep MLP (RFC #420 row ``timestep_mlp_fp32``). + +``TimestepEmbedding(in_channels=256, time_embed_dim=5376, out_dim=2688)``:: + + temb = linear_2(silu(linear_1(features))) 256 -> 5376 -> 2688 + +``time_embedder`` is in ``_keep_in_fp32_modules``: weights, biases, +activations and the output are FP32, and ``temb`` stays FP32 because every +AdaLN projection applies its own SiLU to it before casting (RFC #420 §4). +A BF16 argument is a contract violation (RFC probe H7), not a lower-precision +variant, so it is rejected. +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def validate_h3_timestep_mlp( + x: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> None: + """Check a same-device FP32 ``(T, K) -> (T, H) -> (T, D)`` MLP contract. + + Require ``T > 0``, weights of shapes ``(H, K)``/``(D, H)`` and biases + of shapes ``(H,)``/``(D,)``; reject lower-precision parameters or inputs. + """ + + named = {"x": x, "w1": w1, "b1": b1, "w2": w2, "b2": b2} + for name, tensor in named.items(): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise TypeError( + f"{name} must be float32: the H3 time_embedder is an FP32 module " + f"(got {tensor.dtype})" + ) + if tensor.device != x.device: + raise ValueError(f"{name} is on {tensor.device}, x is on {x.device}") + if x.dim() != 2 or x.shape[0] == 0: + raise ValueError(f"x must be a non-empty (T, K) matrix, got {tuple(x.shape)}") + hidden, k_in = w1.shape if w1.dim() == 2 else (None, None) + if k_in != x.shape[1]: + raise ValueError(f"w1 must be (H, {x.shape[1]}), got {tuple(w1.shape)}") + if b1.shape != (hidden,): + raise ValueError(f"b1 must be ({hidden},), got {tuple(b1.shape)}") + if w2.dim() != 2 or w2.shape[1] != hidden: + raise ValueError(f"w2 must be (D, {hidden}), got {tuple(w2.shape)}") + if b2.shape != (w2.shape[0],): + raise ValueError(f"b2 must be ({w2.shape[0]},), got {tuple(b2.shape)}") + + +class NativeH3TimestepMLPOp: + """PyTorch reference for the H3 timestep MLP. + + ``forward`` is the provider path (``nn.Linear`` / ``F.silu`` in FP32; run + it with TF32 disabled, as the RL-Kernel contract requires). + ``forward_fp32`` is the independent golden: the same graph in FP64, + rounded once to FP32, so its own rounding error (~1e-16 relative) is far + below the FP32 contract and its reduction order does not matter. + """ + + op_class = "reduction" + + def __call__(self, x, w1, b1, w2, b2): + """Return FP32 ``(T, D)`` embeddings using the PyTorch provider graph.""" + + return self.forward(x, w1, b1, w2, b2) + + def forward(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Validate and evaluate linear-SiLU-linear in FP32 on the input device. + + Inputs follow the ``(T, K) -> (T, H) -> (T, D)`` contract and retain + PyTorch autograd support; CUDA callers must disable TF32 for this path. + """ + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + return F.linear(F.silu(F.linear(x, w1, b1)), w2, b2) + + def forward_fp32(self, x, w1, b1, w2, b2) -> torch.Tensor: + """Evaluate the validated MLP in FP64 and round its ``(T, D)`` output to FP32.""" + + validate_h3_timestep_mlp(x, w1, b1, w2, b2) + z = F.linear(x.double(), w1.double(), b1.double()) + h = z * torch.sigmoid(z) + return F.linear(h, w2.double(), b2.double()).float() diff --git a/rl_engine/reference/minimax_h3/timestep_sinusoid.py b/rl_engine/reference/minimax_h3/timestep_sinusoid.py new file mode 100644 index 000000000..0d95e58f0 --- /dev/null +++ b/rl_engine/reference/minimax_h3/timestep_sinusoid.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""H3 sinusoidal timestep features (RFC #420 row ``timestep_sinusoid_h3``). + +H3 builds ``Timesteps(num_channels=256, flip_sin_to_cos=True, +downscale_freq_shift=0)`` and feeds it the *distinct* timestep values of the +packed sequence, unscaled in ``[0, 1]`` (``t = 1 - sigma``):: + + freq[k] = exp(-ln(10000) * k / 128) k = 0..127, FP32 + arg[t, k] = t * freq[k] FP32 + out[t] = [cos(arg[t]) | sin(arg[t])] (T, 256) FP32 + +The output is always FP32: ``time_embedder`` is a ``_keep_in_fp32_modules`` +module in the checkpoint, so the features never pass through BF16. +""" + +from __future__ import annotations + +import math + +import torch + +from rl_engine.reference.minimax_h3 import H3_FREQ_DIM, H3_MAX_PERIOD + +_TIMESTEP_DTYPES = (torch.float32, torch.bfloat16, torch.float16) + + +def validate_h3_timesteps( + timestep: torch.Tensor, num_channels: int, *, check_range: bool = True +) -> None: + """Fail closed on inputs the H3 convention does not define. + + ``check_range`` rejects values outside ``[0, 1]``: H3 consumes + ``t = 1 - sigma`` unscaled, so a ``t * 1000`` caller (RFC probe H10) is a + convention error, not a different valid input. It reads the tensor back + to the host, which is cheap for the handful of distinct timesteps. + """ + + if not isinstance(timestep, torch.Tensor): + raise TypeError("timestep must be a torch.Tensor") + if timestep.dim() != 1: + raise ValueError(f"timestep must be 1-D (num_timesteps,), got {tuple(timestep.shape)}") + if timestep.numel() == 0: + raise ValueError("timestep must hold at least one timestep") + if timestep.dtype not in _TIMESTEP_DTYPES: + raise TypeError(f"timestep must be fp32, bf16 or fp16, got {timestep.dtype}") + if num_channels <= 0 or num_channels % 2 != 0: + raise ValueError(f"num_channels must be a positive even number, got {num_channels}") + if check_range: + t32 = timestep.detach().float() + # One host sync: NaN fails both comparisons, so this also rejects it. + if not bool(((t32 >= 0) & (t32 <= 1)).all()): + raise ValueError( + "timestep must be finite and lie in [0, 1]: H3 consumes t = 1 - sigma " + f"unscaled (got {t32.cpu().tolist()[:8]})" + ) + + +class NativeH3TimestepSinusoidOp: + """PyTorch reference for the H3 sinusoidal timestep features. + + ``forward`` replays diffusers' ``get_timestep_embedding`` op for op on the + input's device (the provider path). ``forward_fp32`` is the independent + golden: the same formula evaluated in FP64 and rounded once to FP32. + """ + + op_class = "elementwise" + + def __call__(self, timestep: torch.Tensor, *, num_channels: int = H3_FREQ_DIM): + """Return FP32 sinusoidal features for nonempty timesteps in ``[0, 1]``.""" + + return self.forward(timestep, num_channels=num_channels) + + def forward( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Return FP32 ``(T, num_channels)`` cosine-then-sine features on the input device. + + Require nonempty 1-D FP32, BF16 or FP16 timesteps and positive even + channels. Validate finite values in ``[0, 1]`` when ``check_range`` is + true and preserve PyTorch autograd through the provider formula. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + # Same op sequence and dtypes as diffusers get_timestep_embedding with + # flip_sin_to_cos=True, downscale_freq_shift=0, scale=1. + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / half + freq = torch.exp(exponent) + arg = timestep[:, None].float() * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1) + + def forward_fp32( + self, + timestep: torch.Tensor, + *, + num_channels: int = H3_FREQ_DIM, + check_range: bool = True, + ) -> torch.Tensor: + """Evaluate the validated sinusoid formula in FP64 and round once to FP32. + + Return ``(T, num_channels)`` on the timestep device with cosine channels + followed by sine channels and gradients through the FP64 golden graph. + """ + + validate_h3_timesteps(timestep, num_channels, check_range=check_range) + half = num_channels // 2 + k = torch.arange(half, dtype=torch.float64, device=timestep.device) + freq = torch.exp(-math.log(H3_MAX_PERIOD) * k / half) + arg = timestep.double()[:, None] * freq[None, :] + return torch.cat([torch.cos(arg), torch.sin(arg)], dim=-1).float() + + @staticmethod + def frequencies_fp32(num_channels: int, device: torch.device | str) -> torch.Tensor: + """The FP32 frequency table the provider path multiplies by.""" + + half = num_channels // 2 + exponent = -math.log(H3_MAX_PERIOD) * torch.arange( + start=0, end=half, dtype=torch.float32, device=device + ) + return torch.exp(exponent / half) diff --git a/rl_engine/runtime/registry.py b/rl_engine/runtime/registry.py index 528c77b98..a733e95ce 100644 --- a/rl_engine/runtime/registry.py +++ b/rl_engine/runtime/registry.py @@ -204,6 +204,43 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_SM90_LM_HEAD = "rl_engine.kernels.ops.cuda.linear.lm_head.SM90LMHeadOp" CUDA_SM90_EMBEDDING = "rl_engine.kernels.ops.cuda.linear.embedding.SM90EmbeddingOp" + # MiniMax-H3 (RFC #420) conditioning path + PYTORCH_H3_TIMESTEP_SINUSOID = ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ) + CUDA_H3_TIMESTEP_SINUSOID = ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ) + PYTORCH_H3_TIMESTEP_MLP = "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp" + CUDA_H3_TIMESTEP_MLP = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp.H3TimestepMLPCudaOp" + ) + PYTORCH_H3_ADALN_PROJECTION = ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ) + CUDA_H3_ADALN_PROJECTION = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection.H3AdaLNProjectionCudaOp" + ) + PYTORCH_H3_ADALN_ROW_GATHER = ( + "rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp" + ) + CUDA_H3_ADALN_ROW_GATHER = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather.H3AdaLNRowGatherCudaOp" + ) + PYTORCH_H3_RMSNORM = "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp" + CUDA_H3_RMSNORM = "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp" + PYTORCH_H3_GATE_RESIDUAL = "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp" + CUDA_H3_GATE_RESIDUAL = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual.H3GateResidualCudaOp" + ) + PYTORCH_H3_FINAL_ADALN_OUT = ( + "rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp" + ) + CUDA_H3_FINAL_ADALN_OUT = ( + "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out.H3FinalAdaLNOutCudaOp" + ) + def _default_semantic_descriptors() -> tuple[OperatorBackendDescriptor, ...]: return ( @@ -502,6 +539,7 @@ class KernelRegistry: """ def __init__(self): + """Initialize backend caches, capability contracts, and device-specific priorities.""" self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() self.semantic = SemanticOperatorCatalog(_default_semantic_descriptors()) @@ -673,6 +711,31 @@ def __init__(self): OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE, ], + "timestep_sinusoid_h3": [ + OpBackend.CUDA_H3_TIMESTEP_SINUSOID, + OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID, + ], + "timestep_mlp_fp32": [ + OpBackend.CUDA_H3_TIMESTEP_MLP, + OpBackend.PYTORCH_H3_TIMESTEP_MLP, + ], + "adaln_projection_3mod": [ + OpBackend.CUDA_H3_ADALN_PROJECTION, + OpBackend.PYTORCH_H3_ADALN_PROJECTION, + ], + "adaln_row_gather": [ + OpBackend.CUDA_H3_ADALN_ROW_GATHER, + OpBackend.PYTORCH_H3_ADALN_ROW_GATHER, + ], + "h3_rmsnorm": [OpBackend.CUDA_H3_RMSNORM, OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [ + OpBackend.CUDA_H3_GATE_RESIDUAL, + OpBackend.PYTORCH_H3_GATE_RESIDUAL, + ], + "final_adaln_out": [ + OpBackend.CUDA_H3_FINAL_ADALN_OUT, + OpBackend.PYTORCH_H3_FINAL_ADALN_OUT, + ], }, "rocm": { "logp": [ @@ -714,6 +777,13 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "musa": { "logp": [ @@ -763,6 +833,13 @@ def __init__(self): ], "silu": [OpBackend.TRITON_SILU, OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.TRITON_SWIGLU, OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], @@ -789,6 +866,13 @@ def __init__(self): "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], "silu": [OpBackend.PYTORCH_NATIVE_SILU], "swiglu": [OpBackend.PYTORCH_NATIVE_SWIGLU], + "timestep_sinusoid_h3": [OpBackend.PYTORCH_H3_TIMESTEP_SINUSOID], + "timestep_mlp_fp32": [OpBackend.PYTORCH_H3_TIMESTEP_MLP], + "adaln_projection_3mod": [OpBackend.PYTORCH_H3_ADALN_PROJECTION], + "adaln_row_gather": [OpBackend.PYTORCH_H3_ADALN_ROW_GATHER], + "h3_rmsnorm": [OpBackend.PYTORCH_H3_RMSNORM], + "adaln_gate_residual": [OpBackend.PYTORCH_H3_GATE_RESIDUAL], + "final_adaln_out": [OpBackend.PYTORCH_H3_FINAL_ADALN_OUT], }, # Ascend NPU: op types without an entry fall back to their CPU # candidates (see the runtime override below), so only diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py new file mode 100644 index 000000000..b20e885a1 --- /dev/null +++ b/rl_engine/validation/models/h3_block.py @@ -0,0 +1,718 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 transformer block, node by node (RFC #420 ``ws1_one_h3_block``). + +The block is a fixed graph of 17 nodes (``NODES``), from the AdaLN projection +through attention and the FFN to the second gated residual, on block 0's +pinned weights and a packed FL2VA layout. Every node runs three ways: + +* provider: the diffusers op sequence (``h3_provider``), fed provider inputs; +* candidate: the RL-Kernel binding of the node, fed candidate inputs + (chained) and, separately, provider inputs (isolated); +* golden: the FP32 reference (``rl_engine.reference.minimax_h3.block``). + +Each node also names the RFC row that owns its arithmetic and how it is bound +today (``Binding.status``): + +* ``row`` -- that row's operator, in this PR stack; +* ``interim`` -- an existing deterministic RL-Kernel operator standing in + until the owning row lands (its arithmetic, speed and rounding are that + operator's, not the row's); +* ``reference`` -- no kernel exists yet; the node runs the provider replay + itself, so it is trivially bitwise and contributes no RL-Kernel evidence. + +The report gives, per node: repeat and batch-row bitwise checks, chained and +isolated agreement with the provider, error against the golden for both the +candidate and the provider, and the backend that ran. ``first_*_drift`` name +the first node (graph order) where a comparison stops being bitwise; the +backward report does the same for the gradients reaching each node's output, +in reverse graph order. +""" + +from __future__ import annotations + +import hashlib +import time +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any, Callable + +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout, h3_timesteps + +HIDDEN = 5376 +HEADS = 56 +HEAD_DIM = 128 +BLOCK = "transformer_blocks.0" +NUM_TIMESTEPS = 2 + + +@dataclass(frozen=True) +class Binding: + rfc_row: str + status: str # "row" | "interim" | "reference" + backend: str + + +@dataclass(frozen=True) +class Node: + name: str + inputs: tuple[str, ...] + binding: Binding + # (ops, ctx, *inputs) -> output + candidate: Callable[..., Any] + # (ctx, *inputs) -> output + provider: Callable[..., Any] + golden: Callable[..., Any] + # Promise checked by the tests: the isolated candidate equals the provider bitwise. + provider_bitwise_isolated: bool + + +class CandidateOps: + """The operators bound to the block's nodes, resolved once per run. + + The three rows of this stack are resolved through the registry, so the + report records what dispatch actually returned; the interim operators are + named directly because their registry entries (``det_gemm``, + ``attention``, ``swiglu``) prefer backends that are not available or not + deterministic on every GPU this replay runs on. + """ + + def __init__(self, registry: KernelRegistry | None = None) -> None: + from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp + from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp + from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp + + registry = registry or KernelRegistry() + self.projection = registry.get_op("adaln_projection_3mod", device="cuda") + self.norm = registry.get_op("h3_rmsnorm", device="cuda") + self.gate = registry.get_op("adaln_gate_residual", device="cuda") + self.gemm = TritonDetGemmOp() + self.attention = DeterministicAttentionOp() + self.swiglu = SwiGLUCudaOp() + + def backends(self) -> dict[str, str]: + return { + name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items() + } + + +def _w(ctx: dict[str, Any], name: str) -> torch.Tensor: + return ctx["weights"][f"{BLOCK}.{name}"] + + +def _heads(x: torch.Tensor) -> torch.Tensor: + return x.unflatten(-1, (HEADS, HEAD_DIM)) + + +class _InterimLinear(torch.autograd.Function): + """``y = x @ W^T`` on the Triton FP32-accumulate kernel, rounded to BF16 once. + + Interim binding for the block's GEMM rows. ``TritonDetGemmOp.__call__`` + reduces K as a tree with BF16 nodes (about 3x cuBLAS's error at these + shapes); ``_triton_gemm_fp32`` runs one fixed-order FP32 K loop (no + split-K), so each output row depends only on its own input row. The + backward uses the same kernel for ``dX`` (reduction over the output + features, row-local) and ``dW`` (reduction over tokens). + """ + + @staticmethod + def forward(ctx, x2, weight): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + ctx.save_for_backward(x2, weight) + return _triton_gemm_fp32(x2, weight.t()).to(x2.dtype) + + @staticmethod + def backward(ctx, grad): + from rl_engine.backends.shared.triton.gemm.det_gemm import _triton_gemm_fp32 + + x2, weight = ctx.saved_tensors + grad = grad.contiguous() + dx = _triton_gemm_fp32(grad, weight).to(x2.dtype) + dw = _triton_gemm_fp32(grad.t(), x2).to(weight.dtype) + return dx, dw + + +def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + flat = x.reshape(-1, x.shape[-1]) + return _InterimLinear.apply(flat, weight).view(*x.shape[:-1], weight.shape[0]) + + +def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor: + q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v))) + out = ops.attention(q, k, v, causal=False) + return out.transpose(1, 2).flatten(2, 3) + + +def _cand_swiglu(ops: CandidateOps, projected: torch.Tensor) -> torch.Tensor: + hidden, gate = projected.chunk(2, dim=-1) + return ops.swiglu(gate.contiguous(), hidden.contiguous()) + + +def _linear_node(name: str, source: str, weight: str, row: str, status: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding( + row, status, "triton _det_gemm_fp32_kernel (FP32 K loop, one BF16 rounding)" + ), + candidate=lambda ops, ctx, x: _cand_linear(ops, x, _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_linear(x, _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_linear(x, _w(ctx, weight)), + provider_bitwise_isolated=False, # tree reduction vs cuBLAS + ) + + +def _norm_modulate_node(name: str, source: str, weight: str, shift: int, scale: int) -> Node: + return Node( + name=name, + inputs=(source, "adaln_projection"), + binding=Binding("h3_rmsnorm", "row", "H3RMSNormCudaOp.forward_modulated"), + candidate=lambda ops, ctx, x, t: ops.norm.forward_modulated( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider=lambda ctx, x, t: prov.provider_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + golden=lambda ctx, x, t: gold.golden_norm_modulate( + x, _w(ctx, weight), t[shift], t[scale], ctx["index"] + ), + provider_bitwise_isolated=True, + ) + + +def _head_norm_node(name: str, source: str, weight: str) -> Node: + return Node( + name=name, + inputs=(source,), + # The row is someone else's; until it lands, this stack's plain H3 RMSNorm + # (bitwise equal to nn.RMSNorm at any width) runs on the 128-wide heads. + binding=Binding("h3_qk_rmsnorm_d128", "interim", "H3RMSNormCudaOp.forward"), + candidate=lambda ops, ctx, x: ops.norm(_heads(x), _w(ctx, weight)), + provider=lambda ctx, x: prov.provider_head_rmsnorm(_heads(x), _w(ctx, weight)), + golden=lambda ctx, x: gold.golden_rmsnorm(_heads(x), _w(ctx, weight)), + provider_bitwise_isolated=True, + ) + + +def _rope_node(name: str, source: str) -> Node: + return Node( + name=name, + inputs=(source,), + binding=Binding("h3_mm_rope_3axis_partial", "reference", "provider_apply_rotary"), + candidate=lambda ops, ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + provider=lambda ctx, x: prov.provider_apply_rotary(x, *ctx["rope"]), + golden=lambda ctx, x: gold.golden_apply_rotary(x, *ctx["rope_fp32"]), + provider_bitwise_isolated=True, + ) + + +def _gate_node(name: str, residual: str, sublayer: str, gate: int) -> Node: + return Node( + name=name, + inputs=(residual, sublayer, "adaln_projection"), + binding=Binding("adaln_gate_residual", "row", "H3GateResidualCudaOp"), + candidate=lambda ops, ctx, r, y, t: ops.gate(r, y, t[gate], ctx["index"]), + provider=lambda ctx, r, y, t: prov.provider_gate_residual(r, t[gate], ctx["index"], y), + golden=lambda ctx, r, y, t: gold.golden_gate_residual(r, t[gate], ctx["index"], y), + provider_bitwise_isolated=True, + ) + + +def _adaln_params(ctx): + return _w(ctx, "adaln_proj.linear.weight"), _w(ctx, "adaln_proj.linear.bias") + + +NODES: list[Node] = [ + Node( + name="adaln_projection", + inputs=("temb",), + binding=Binding("adaln_projection_3mod", "row", "H3AdaLNProjectionCudaOp"), + candidate=lambda ops, ctx, temb: ops.projection(temb, *_adaln_params(ctx)), + provider=lambda ctx, temb: prov.provider_adaln_modulation(temb, *_adaln_params(ctx)), + golden=lambda ctx, temb: gold.golden_adaln_tables(temb, *_adaln_params(ctx)), + provider_bitwise_isolated=False, # tensor-core tree vs cuBLAS + ), + _norm_modulate_node("norm1", "hidden", "norm1.weight", shift=0, scale=1), + _linear_node("q_proj", "norm1", "attn.to_q.weight", "h3_qkv_gemm", "interim"), + _linear_node("k_proj", "norm1", "attn.to_k.weight", "h3_qkv_gemm", "interim"), + _linear_node("v_proj", "norm1", "attn.to_v.weight", "h3_qkv_gemm", "interim"), + _head_norm_node("q_norm", "q_proj", "attn.norm_q.weight"), + _head_norm_node("k_norm", "k_proj", "attn.norm_k.weight"), + _rope_node("rope_q", "q_norm"), + _rope_node("rope_k", "k_norm"), + Node( + name="attention", + inputs=("rope_q", "rope_k", "v_proj"), + binding=Binding("h3_full_attention", "interim", "DeterministicAttentionOp(causal=False)"), + candidate=lambda ops, ctx, q, k, v: _cand_attention(ops, q, k, v), + provider=lambda ctx, q, k, v: prov.provider_attention(q, k, _heads(v)), + golden=lambda ctx, q, k, v: gold.golden_attention(q, k, _heads(v)), + provider_bitwise_isolated=False, + ), + _linear_node("o_proj", "attention", "attn.to_out.0.weight", "h3_attention_o_gemm", "interim"), + _gate_node("residual_attn", "hidden", "o_proj", gate=2), + _norm_modulate_node("norm2", "residual_attn", "norm2.weight", shift=3, scale=4), + _linear_node("ffn_gate_up", "norm2", "ff.net.0.proj.weight", "h3_ffn_gate_up_gemm", "interim"), + Node( + name="swiglu", + inputs=("ffn_gate_up",), + binding=Binding("h3_swiglu", "interim", "SwiGLUCudaOp"), + candidate=lambda ops, ctx, x: _cand_swiglu(ops, x), + provider=lambda ctx, x: prov.provider_swiglu(x), + golden=lambda ctx, x: gold.golden_swiglu(x), + provider_bitwise_isolated=False, # one FP32 rounding vs the eager two + ), + _linear_node("ffn_down", "swiglu", "ff.net.2.weight", "h3_ffn_down_gemm", "interim"), + _gate_node("residual_mlp", "residual_attn", "ffn_down", gate=5), +] +NODE_NAMES = tuple(node.name for node in NODES) +OUTPUT = NODES[-1].name + +# Gradient leaves of the backward replay: the block input, temb and every block parameter. +PARAM_LEAVES = ( + "adaln_proj.linear.weight", + "adaln_proj.linear.bias", + "norm1.weight", + "attn.to_q.weight", + "attn.to_k.weight", + "attn.to_v.weight", + "attn.norm_q.weight", + "attn.norm_k.weight", + "attn.to_out.0.weight", + "norm2.weight", + "ff.net.0.proj.weight", + "ff.net.2.weight", +) + + +# --- Inputs --------------------------------------------------------------------------------- + + +def make_block_context( + weights: dict[str, torch.Tensor], *, layout: str, seed: int = 0, batch: int = 1 +) -> dict[str, Any]: + """Block-0 weights, a packed FL2VA layout, the timestep embedding and a seeded input. + + ``temb`` comes from the provider's conditioning path (``h3_chain`` covers + those rows); the block input is a seeded BF16 stand-in for the packed + projections, which are other rows. Batch rows ``b > 0`` get their own + seeded inputs; row 0 is the same tensor at every ``batch``. + """ + + spec = H3_BLOCK_LAYOUTS[layout] + packed = h3_fl2va_layout( + num_text_tokens=spec["num_text_tokens"], + num_latent_frames=spec["num_latent_frames"], + latent_height=spec["latent_height"], + latent_width=spec["latent_width"], + num_audio_latents=spec["num_audio_latents"], + ) + seq_len = packed["position_ids"].shape[0] + timestep = h3_timesteps(NUM_TIMESTEPS, seed=seed) + w = weights + with torch.no_grad(): + features = prov.provider_time_proj(timestep) + temb = prov.provider_time_embedder( + features, + w["time_embedder.linear_1.weight"], + w["time_embedder.linear_1.bias"], + w["time_embedder.linear_2.weight"], + w["time_embedder.linear_2.bias"], + ) + rows = [ + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2 + b)) + for b in range(batch) + ] + hidden = torch.cat(rows).to(torch.bfloat16).cuda() + return { + "layout": layout, + "seq_len": seq_len, + "batch": batch, + "weights": weights, + "temb": temb, + "hidden": hidden, + "index": packed["timestep_indices"] * 3 + packed["token_tags"], + "position_ids": packed["position_ids"], + "rope": prov.provider_rope(packed["position_ids"]), + "rope_fp32": gold.golden_rope_tables(packed["position_ids"]), + } + + +@contextmanager +def strict_fp32_matmul(): + """Golden matmuls in full FP32: no TF32, no reduced-precision reductions.""" + + saved = (torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32) + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = saved + + +# --- Graph execution ------------------------------------------------------------------------ + + +def run_graph(mode: str, ctx: dict[str, Any], ops: CandidateOps | None = None) -> dict[str, Any]: + """Every node's output for one mode ("candidate", "provider" or "golden").""" + + values: dict[str, Any] = {"hidden": ctx["hidden"], "temb": ctx["temb"]} + # Diagnostic hybrids only: these candidate nodes run their provider replay instead. + provider_nodes = ctx.get("provider_nodes", ()) + with strict_fp32_matmul(): + for node in NODES: + args = [values[name] for name in node.inputs] + if mode == "candidate" and node.name in provider_nodes: + values[node.name] = node.provider(ctx, *args) + elif mode == "candidate": + values[node.name] = node.candidate(ops, ctx, *args) + elif mode == "provider": + values[node.name] = node.provider(ctx, *args) + else: + values[node.name] = node.golden(ctx, *args) + observer = ctx.get("node_observer") + if observer is not None: + observer(node.name) + return values + + +def run_isolated( + ctx: dict[str, Any], ops: CandidateOps, provider: dict[str, Any] +) -> dict[str, Any]: + """Each candidate node fed the provider's inputs to that node.""" + + return { + node.name: node.candidate(ops, ctx, *[provider[name] for name in node.inputs]) + for node in NODES + } + + +def _flat(value: Any) -> list[torch.Tensor]: + if isinstance(value, torch.Tensor): + return [value] + return [t for item in value for t in _flat(item)] + + +def _batch_row(value: Any, row: int, batch: int) -> Any: + """Row ``row`` of a batched activation; AdaLN tables have no batch axis.""" + + if isinstance(value, torch.Tensor): + return value[row : row + 1] if value.dim() >= 3 and value.shape[0] == batch else value + return tuple(_batch_row(v, row, batch) for v in value) + + +def bitwise(lhs: Any, rhs: Any) -> bool: + return all(torch.equal(a, b) for a, b in zip(_flat(lhs), _flat(rhs), strict=True)) + + +def max_abs(lhs: Any, rhs: Any) -> float: + return max( + float((a.float() - b.float()).abs().max()) if a.numel() else 0.0 + for a, b in zip(_flat(lhs), _flat(rhs), strict=True) + ) + + +def first_differing_row(lhs: Any, rhs: Any) -> int | None: + """First packed row (sequence position) where two node outputs differ, or None.""" + + for a, b in zip(_flat(lhs), _flat(rhs), strict=True): + if a.dim() < 3: + if not torch.equal(a, b): + return -1 # a table row, not a sequence position + continue + diff = (a != b).reshape(a.shape[0], a.shape[1], -1).any(dim=-1).any(dim=0) + rows = diff.nonzero() + if rows.numel(): + return int(rows[0]) + return None + + +def rel_err(value: Any, golden: Any) -> float: + """``max|value - golden| / max|golden|`` over a node's tensors.""" + + num = den = 0.0 + for a, g in zip(_flat(value), _flat(golden), strict=True): + g32 = g.float() + num = max(num, float((a.float() - g32).abs().max())) + den = max(den, float(g32.abs().max())) + return num / max(den, 1e-30) + + +def ulp_err(value: torch.Tensor, golden: torch.Tensor) -> float: + """``max|value - golden|`` in units of ``value.dtype``'s ulp at ``max|golden|``. + + ``rel_err`` of a stored gradient moves in steps of this ulp: its worst + element is a single rounded number, so two equally valid reduction orders + can differ by whole ulps there while every other element agrees. + """ + + g32 = golden.float() + top = float(g32.abs().max()) + if top == 0.0: + return 0.0 + ulp = torch.finfo(value.dtype).eps * 2.0 ** torch.floor(torch.log2(torch.tensor(top))).item() + return float((value.float() - g32).abs().max()) / ulp + + +def digest(value: Any) -> str: + """sha256 over a node's tensors' bytes, for comparing runs across processes or hosts.""" + + h = hashlib.sha256() + for t in _flat(value): + h.update(t.detach().contiguous().view(torch.uint8).cpu().numpy().tobytes()) + return h.hexdigest() + + +def _time_ms(fn: Callable[[], Any], repeats: int) -> float: + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize() + samples.append((time.perf_counter() - start) * 1e3) + return sorted(samples)[len(samples) // 2] + + +# --- Forward report ------------------------------------------------------------------------- + + +def run_block_case( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Forward replay of block 0: per-node agreement, invariance, accuracy and first drift.""" + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + with torch.no_grad(): + provider = run_graph("provider", ctx) + golden = run_graph("golden", ctx) + chained = run_graph("candidate", ctx, ops) + repeat = run_graph("candidate", ctx, ops) + batched = run_graph("candidate", ctx2, ops) + isolated = run_isolated(ctx, ops, provider) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "backends": ops.backends(), + "nodes": [], + } + firsts: dict[str, str | None] = {"first_drift": None, "first_isolated_drift": None} + firsts.update({"first_repeat_drift": None, "first_batch_drift": None}) + for node in NODES: + name = node.name + row0 = _batch_row(batched[name], 0, 2) + entry: dict[str, Any] = { + "node": name, + "rfc_row": node.binding.rfc_row, + "status": node.binding.status, + "backend": node.binding.backend, + "provider_bitwise_isolated_promised": node.provider_bitwise_isolated, + "repeat_bitwise_equal": bitwise(repeat[name], chained[name]), + "batch_row_bitwise_equal": bitwise(row0, chained[name]), + "chained_vs_provider": { + "bitwise_equal": bitwise(chained[name], provider[name]), + "max_abs": max_abs(chained[name], provider[name]), + "first_differing_row": first_differing_row(chained[name], provider[name]), + }, + "isolated_vs_provider": { + "bitwise_equal": bitwise(isolated[name], provider[name]), + "max_abs": max_abs(isolated[name], provider[name]), + }, + "candidate_rel_err_vs_golden": rel_err(chained[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "candidate_sha256": digest(chained[name]), + } + checks = { + "first_drift": entry["chained_vs_provider"]["bitwise_equal"], + "first_isolated_drift": entry["isolated_vs_provider"]["bitwise_equal"], + "first_repeat_drift": entry["repeat_bitwise_equal"], + "first_batch_drift": entry["batch_row_bitwise_equal"], + } + for key, ok in checks.items(): + if firsts[key] is None and not ok: + firsts[key] = name + report["nodes"].append(entry) + report.update(firsts) + report["output"] = { + "candidate_rel_err_vs_golden": rel_err(chained[OUTPUT], golden[OUTPUT]), + "provider_rel_err_vs_golden": rel_err(provider[OUTPUT], golden[OUTPUT]), + "candidate_vs_provider_max_abs": max_abs(chained[OUTPUT], provider[OUTPUT]), + } + report["permutation_probe"] = permutation_probe(ctx, ops, chained[OUTPUT], seed) + if timing_repeats: + with torch.no_grad(): + report["timing_ms"] = { + "candidate_forward": _time_ms( + lambda: run_graph("candidate", ctx, ops), timing_repeats + ), + "provider_forward": _time_ms(lambda: run_graph("provider", ctx), timing_repeats), + } + return report + + +def permutation_probe( + ctx: dict[str, Any], ops: CandidateOps, base_output: torch.Tensor, seed: int +) -> dict[str, Any]: + """RFC #420 control H-C3, reported and not promised. + + Permute the packed rows together with every per-row tensor (positions, + AdaLN indices); the block output should be the same permutation of the + original. Row-local nodes keep each row's bytes, but attention sums over + keys in packed order, so a different key order may round differently. + """ + + perm = torch.randperm(ctx["seq_len"], generator=torch.Generator().manual_seed(seed + 7)).cuda() + permuted = { + **ctx, + "hidden": ctx["hidden"][:, perm], + "index": ctx["index"][perm], + "position_ids": ctx["position_ids"][perm], + "rope": tuple(t[perm] for t in ctx["rope"]), + "rope_fp32": tuple(t[perm] for t in ctx["rope_fp32"]), + } + with torch.no_grad(): + out = run_graph("candidate", permuted, ops)[OUTPUT] + expected = base_output[:, perm] + return { + "bitwise_equal_after_permutation": bool(torch.equal(out, expected)), + "max_abs": max_abs(out, expected), + "differing_rows": int((out != expected).reshape(out.shape[1], -1).any(dim=-1).sum()), + } + + +# --- Backward report ------------------------------------------------------------------------ + + +def _backward( + mode: str, ctx: dict[str, Any], ops: CandidateOps | None, upstream: torch.Tensor +) -> tuple[dict[str, torch.Tensor], dict[str, list[torch.Tensor]]]: + """Gradients of every leaf, and the gradient reaching each node's output, for one mode. + + The golden runs on FP32 copies of the leaves, so its gradients are not + rounded to the BF16 parameter dtype. + """ + + def leaf(t: torch.Tensor) -> torch.Tensor: + t = t.detach().float() if mode == "golden" else t.detach() + return t.clone().requires_grad_(True) + + leaves = {name: leaf(ctx["weights"][f"{BLOCK}.{name}"]) for name in PARAM_LEAVES} + leaves["hidden"] = leaf(ctx["hidden"]) + leaves["temb"] = leaf(ctx["temb"]) + weights = {**ctx["weights"], **{f"{BLOCK}.{k}": leaves[k] for k in PARAM_LEAVES}} + run_ctx = {**ctx, "weights": weights, "hidden": leaves["hidden"], "temb": leaves["temb"]} + values = run_graph(mode, run_ctx, ops) + for node in NODES: + for t in _flat(values[node.name]): + if t.requires_grad: + t.retain_grad() + out = values[OUTPUT] + torch.autograd.backward(out, upstream.to(out.dtype)) + node_grads = { + node.name: [ + t.grad if t.grad is not None else torch.zeros_like(t) for t in _flat(values[node.name]) + ] + for node in NODES + } + return {name: t.grad for name, t in leaves.items()}, node_grads + + +def run_block_backward( + weights: dict[str, torch.Tensor], + *, + layout: str, + seed: int = 0, + ops: CandidateOps | None = None, + timing_repeats: int = 0, +) -> dict[str, Any]: + """Backward replay of block 0. + + Per leaf (block input, ``temb`` and the 12 block parameters): repeat + bitwise, and the candidate's and provider's error against the FP32 + golden. Per node, in reverse graph order: whether the gradient reaching + the node's output repeats bitwise, and whether batch row 0 of a two-row + batch gets the same bytes as the batch of one (rows are separate packed + documents, so every activation gradient is row-local). Parameter + gradients sum over the batch and are not part of the batch check. + """ + + ops = ops or CandidateOps() + ctx = make_block_context(weights, layout=layout, seed=seed, batch=1) + ctx2 = make_block_context(weights, layout=layout, seed=seed, batch=2) + generator = torch.Generator().manual_seed(seed + 11) + up_row0 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + up_row1 = torch.randn((1, ctx["seq_len"], HIDDEN), generator=generator) + upstream = up_row0.bfloat16().cuda() + upstream2 = torch.cat([up_row0, up_row1]).bfloat16().cuda() + + golden, _ = _backward("golden", ctx, None, upstream) + provider, _ = _backward("provider", ctx, None, upstream) + first, first_nodes = _backward("candidate", ctx, ops, upstream) + second, second_nodes = _backward("candidate", ctx, ops, upstream) + batched, batched_nodes = _backward("candidate", ctx2, ops, upstream2) + + report: dict[str, Any] = { + "layout": layout, + "seq_len": ctx["seq_len"], + "seed": seed, + "leaves": {}, + "nodes": [], + } + for name in (*PARAM_LEAVES, "hidden", "temb"): + report["leaves"][name] = { + "repeat_bitwise_equal": bool(torch.equal(first[name], second[name])), + "candidate_rel_err_vs_golden": rel_err(first[name], golden[name]), + "provider_rel_err_vs_golden": rel_err(provider[name], golden[name]), + "storage_dtype": str(first[name].dtype).removeprefix("torch."), + "candidate_ulps_vs_golden": ulp_err(first[name], golden[name]), + "provider_ulps_vs_golden": ulp_err(provider[name], golden[name]), + } + report["leaves"]["hidden"]["batch_row_bitwise_equal"] = bool( + torch.equal(batched["hidden"][:1], first["hidden"]) + ) + first_repeat = first_batch = None + for node in reversed(NODES): + name = node.name + row0 = [_batch_row(g, 0, 2) for g in batched_nodes[name]] + entry: dict[str, Any] = { + "node": name, + "grad_repeat_bitwise_equal": bitwise(first_nodes[name], second_nodes[name]), + "grad_batch_row_bitwise_equal": ( + bitwise(row0, first_nodes[name]) if node.name != "adaln_projection" else None + ), + } + if first_repeat is None and not entry["grad_repeat_bitwise_equal"]: + first_repeat = name + if first_batch is None and entry["grad_batch_row_bitwise_equal"] is False: + first_batch = name + report["nodes"].append(entry) + report["first_grad_repeat_drift"] = first_repeat + report["first_grad_batch_drift"] = first_batch + if timing_repeats: + report["timing_ms"] = { + "candidate_forward_backward": _time_ms( + lambda: _backward("candidate", ctx, ops, upstream), timing_repeats + ), + "provider_forward_backward": _time_ms( + lambda: _backward("provider", ctx, None, upstream), timing_repeats + ), + } + return report diff --git a/rl_engine/validation/models/h3_cases.py b/rl_engine/validation/models/h3_cases.py new file mode 100644 index 000000000..92455bae0 --- /dev/null +++ b/rl_engine/validation/models/h3_cases.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Deterministic MiniMax-H3 (RFC #420) test inputs shared by tests and scripts.""" + +from __future__ import annotations + +import torch + + +def h3_timesteps(num: int, *, seed: int = 0, device: str = "cuda") -> torch.Tensor: + """Return seeded FP32 timesteps of shape ``(num,)`` on ``device``. + + Requires positive ``num``; the first value is 0 and the last is 1 when + ``num > 1``, with all remaining values sampled uniformly in [0, 1]. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + t = torch.rand(num, generator=generator) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(device) + + +def h3_packed_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """Return seeded int64 timestep indices and modality tags, each shape ``(seq_len,)``. + + Each row independently samples a timestep in ``[0, num_timesteps)`` and + a video/text/audio tag in ``[0, 3)``. The first three rows cover all tags + when available; both tensors are returned on ``device``. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + token_tags = torch.randint(0, 3, (seq_len,), generator=generator) + timestep_indices = torch.randint(0, num_timesteps, (seq_len,), generator=generator) + token_tags[: min(3, seq_len)] = torch.arange(min(3, seq_len)) + return timestep_indices.to(device), token_tags.to(device) + + +def h3_block_layout( + seq_len: int, num_timesteps: int, *, seed: int = 0, device: str = "cuda" +) -> tuple[torch.Tensor, torch.Tensor]: + """(timestep_indices, token_tags) for block-structured packing, as H3 packs requests. + + Each timestep owns one contiguous run of text, video and audio blocks (about + 5% / 80% / 15% of its tokens), so table rows form contiguous position runs. + ``h3_packed_layout`` is the interleaved stress case. + """ + + generator = torch.Generator(device="cpu").manual_seed(seed) + cuts = torch.sort(torch.randperm(seq_len - 1, generator=generator)[: num_timesteps - 1] + 1) + edges = [0, *cuts.values.tolist(), seq_len] + timestep_indices = torch.empty(seq_len, dtype=torch.long) + token_tags = torch.empty(seq_len, dtype=torch.long) + for t, (lo, hi) in enumerate(zip(edges, edges[1:])): + text, audio = (hi - lo) // 20, (hi - lo) * 3 // 20 + timestep_indices[lo:hi] = t + token_tags[lo:hi] = 0 # video + token_tags[lo : lo + text] = 1 + token_tags[hi - audio : hi] = 2 + return timestep_indices.to(device), token_tags.to(device) + + +# diffusers@f53d552 modular_pipelines/minimax_h3/before_denoise.py constants. +_ROPE_FRAME_RESCALE = 5.0 / 3.0 +_ROPE_FRAMES_PER_LATENT = (1, 4, 4, 4, 4) +_ROPE_SPATIAL_SCALE = 32 +VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2 + + +def _spatial_grid(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + """One aspect-normalised spatial axis, as numpy's ``linspace(endpoint=False)`` builds it.""" + + import numpy as np + + ratio = dim / sqrt_area + left = (1.0 - ratio) / 2.0 + grid = np.linspace(left, left + ratio, dim // patch, endpoint=False) * _ROPE_SPATIAL_SCALE + return torch.from_numpy(grid).to(torch.float64) + + +def h3_fl2va_layout( + *, + num_text_tokens: int, + num_latent_frames: int, + latent_height: int, + latent_width: int, + num_audio_latents: int, + audio_channels: int = 2, + patch: tuple[int, int] = (2, 2), + device: str = "cuda", +) -> dict[str, torch.Tensor]: + """Packed text/audio/video rows of a text-only FL2VA request, laid out as diffusers packs them. + + Replays ``position_ids``/``token_tags`` of the pinned pipeline's FL2VA + layout step (no keyframe anchors): text rows on the time axis at their row + index, channel-major audio rows pinned to the width extremes, then the video + frames on the non-uniform ``5/3 * (1, 4, 4, 4, 4)`` clock. ``position_ids`` + is ``(S, 3)`` float64 as the pipeline makes it; the rotary module casts it. + ``timestep_indices`` puts text on table row 0 and the generated media on + row 1, a two-timestep stand-in for the conditioning rows. + """ + + import numpy as np + + patch_h, patch_w = patch + rows_per_frame = (latent_height // patch_h) * (latent_width // patch_w) + num_audio_rows = num_audio_latents * audio_channels + num_video_rows = num_latent_frames * rows_per_frame + seq_len = num_text_tokens + num_audio_rows + num_video_rows + audio_start = num_text_tokens + video_start = audio_start + num_audio_rows + + position_ids = torch.zeros(seq_len, 3, dtype=torch.float64) + position_ids[:num_text_tokens, 0] = torch.arange(num_text_tokens, dtype=torch.float64) + sqrt_area = np.sqrt(latent_height * latent_width) + height_grid = _spatial_grid(latent_height, patch_h, sqrt_area) + width_grid = _spatial_grid(latent_width, patch_w, sqrt_area) + grids = torch.meshgrid(height_grid, width_grid, indexing="ij") + frame_grid = torch.stack([grid.reshape(-1) for grid in grids], dim=-1) + + audio_time = float(num_text_tokens) + torch.arange(num_audio_latents, dtype=torch.float64) + position_ids[audio_start:video_start, 0] = audio_time.repeat(audio_channels) + position_ids[audio_start:video_start, 2] = torch.cat( + [ + torch.full((num_audio_latents,), float(width_grid[0]), dtype=torch.float64), + torch.full( + (num_audio_rows - num_audio_latents,), float(width_grid[-1]), dtype=torch.float64 + ), + ] + ) + spans = torch.tensor( + [ + _ROPE_FRAME_RESCALE * _ROPE_FRAMES_PER_LATENT[i % len(_ROPE_FRAMES_PER_LATENT)] + for i in range(num_latent_frames) + ], + dtype=torch.float64, + ) + frame_time = float(num_text_tokens) + torch.cat( + [torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)] + ) + video = torch.empty(num_latent_frames, rows_per_frame, 3, dtype=torch.float64) + video[:, :, 0] = frame_time[:, None] + video[:, :, 1:] = frame_grid[None] + position_ids[video_start:] = video.reshape(-1, 3) + + token_tags = torch.full((seq_len,), VIDEO_TAG, dtype=torch.long) + token_tags[:num_text_tokens] = TEXT_TAG + token_tags[audio_start:video_start] = AUDIO_TAG + timestep_indices = torch.ones(seq_len, dtype=torch.long) + timestep_indices[:num_text_tokens] = 0 + return { + "position_ids": position_ids.to(device), + "token_tags": token_tags.to(device), + "timestep_indices": timestep_indices.to(device), + } + + +# Named FL2VA layouts for the one-block replay: (text, latent frames, latent H, latent W, audio). +H3_BLOCK_LAYOUTS = { + # 32 + 2*20 + 2*(8*12) = 264 rows: the CI-sized case. + "tiny": dict( + num_text_tokens=32, + num_latent_frames=2, + latent_height=16, + latent_width=24, + num_audio_latents=20, + ), + # 77 + 2*40 + 3*(16*16) = 925 rows: a non-power-of-two, three-frame case. + "small": dict( + num_text_tokens=77, + num_latent_frames=3, + latent_height=32, + latent_width=32, + num_audio_latents=40, + ), + # 120 + 2*80 + 5*(18*32) = 3160 rows: a 16:9 five-frame case. + "medium": dict( + num_text_tokens=120, + num_latent_frames=5, + latent_height=36, + latent_width=64, + num_audio_latents=80, + ), +} diff --git a/rl_engine/validation/models/h3_chain.py b/rl_engine/validation/models/h3_chain.py new file mode 100644 index 000000000..3aa1aa191 --- /dev/null +++ b/rl_engine/validation/models/h3_chain.py @@ -0,0 +1,455 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Stage-by-stage replay of the MiniMax-H3 conditioning chain (RFC #420). + +Each stage runs three ways on the same device: + +* provider: the diffusers op sequence (``rl_engine.validation.models.h3_provider``); +* candidate: the dispatched RL-Kernel op, fed the previous candidate output + (chained) and, separately, the previous provider output (isolated); +* golden: the FP32/FP64 reference, fed the previous golden output. + +A stage also declares what it promises, so the end-to-end test and the +evidence JSON check the same thing: whether its isolated output is bitwise +equal to the provider, and its tolerance against the golden (rows of +``tolerance_contract.json``). ``first_drift`` is the first stage whose chained +output differs from the provider; ``first_isolated_drift`` the first stage +that differs even on provider inputs. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) + +REPO_ROOT = Path(__file__).resolve().parents[3] +HIDDEN = 5376 + + +@dataclass(frozen=True) +class Stage: + name: str + op_type: str + # (op, ctx, upstream) -> outputs; ``upstream`` is the previous stage's output. + candidate: Callable[[Any, dict, Any], Any] + provider: Callable[[dict, Any], Any] + golden: Callable[[Any, dict, Any], Any] + # Promises checked by tests/models/minimax_h3/test_h3_conditioning_e2e.py. + provider_bitwise_isolated: bool + golden_atol: float + golden_rtol: float + + +def mlp_params(ctx: dict[str, Any]) -> list[torch.Tensor]: + """Return the context's MLP weights and biases in linear_1, linear_2 call order.""" + + weights = ctx["weights"] + return [weights[f"time_embedder.linear_{i}.{p}"] for i in (1, 2) for p in ("weight", "bias")] + + +def adaln_params(ctx: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: + """Return the context's block-0 AdaLN projection weight and bias.""" + + prefix = "transformer_blocks.0.adaln_proj.linear" + return ctx["weights"][f"{prefix}.weight"], ctx["weights"][f"{prefix}.bias"] + + +STAGES: list[Stage] = [ + Stage( + name="timestep_sinusoid_h3", + op_type="timestep_sinusoid_h3", + candidate=lambda op, ctx, _up: op(ctx["timestep"]), + provider=lambda ctx, _up: provider_time_proj(ctx["timestep"]), + golden=lambda op, ctx, _up: op.forward_fp32(ctx["timestep"]), + provider_bitwise_isolated=True, + golden_atol=1e-5, # elementwise / float32 + golden_rtol=1e-5, + ), + Stage( + name="timestep_mlp_fp32", + op_type="timestep_mlp_fp32", + candidate=lambda op, ctx, up: op(up, *mlp_params(ctx)), + provider=lambda ctx, up: provider_time_embedder(up, *mlp_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *mlp_params(ctx)), + provider_bitwise_isolated=False, # reduction: different tree from cuBLAS + golden_atol=1e-4, # reduction / float32 + golden_rtol=1e-4, + ), + Stage( + name="adaln_projection_3mod", + op_type="adaln_projection_3mod", + candidate=lambda op, ctx, up: op(up, *adaln_params(ctx)), + provider=lambda ctx, up: provider_adaln_modulation(up, *adaln_params(ctx)), + golden=lambda op, ctx, up: op.forward_fp32(up, *adaln_params(ctx)), + provider_bitwise_isolated=False, # reduction: tensor-core tree differs from cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ), + Stage( + name="adaln_row_gather", + op_type="adaln_row_gather", + candidate=lambda op, ctx, up: op.gather_chunks( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + provider=lambda ctx, up: provider_adaln_row_gather( + up, ctx["timestep_indices"], ctx["token_tags"] + ), + golden=lambda op, ctx, up: op.forward_fp32( + torch.cat(list(up), dim=1), ctx["timestep_indices"], ctx["token_tags"] + ), + provider_bitwise_isolated=True, # a copy + golden_atol=5e-2, # carries the projection's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ), +] + + +def adaln_indices(ctx: dict[str, Any]) -> torch.Tensor: + return ctx["timestep_indices"] * 3 + ctx["token_tags"] + + +def _norm1_modulated(ctx: dict[str, Any], modulation, call): + shift_msa, scale_msa = modulation[0], modulation[1] + weight = ctx["weights"]["transformer_blocks.0.norm1.weight"] + return call(ctx["hidden"], weight, shift_msa, scale_msa, adaln_indices(ctx)) + + +STAGES.append( + Stage( + name="h3_rmsnorm", + op_type="h3_rmsnorm", + # Block norm1 followed by the MSA shift/scale of the projection stage. + candidate=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated + ), + provider=lambda ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], provider_norm_modulate + ), + golden=lambda op, ctx, _up: _norm1_modulated( + ctx, ctx["history"]["adaln_projection_3mod"], op.forward_modulated_fp32 + ), + provider_bitwise_isolated=True, # replays nn.RMSNorm's reduction order + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +STAGES.append( + Stage( + name="adaln_gate_residual", + op_type="adaln_gate_residual", + # residual + gate_msa[row] * attention_output (the stand-in sublayer output). + candidate=lambda op, ctx, _up: op( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider=lambda ctx, _up: provider_gate_residual( + ctx["hidden"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ctx["sublayer"], + ), + golden=lambda op, ctx, _up: op.forward_fp32( + ctx["hidden"], + ctx["sublayer"], + ctx["history"]["adaln_projection_3mod"][2], + adaln_indices(ctx), + ), + provider_bitwise_isolated=True, # eager rounding order, elementwise + golden_atol=5e-2, # carries the gate's BF16 rounding (reduction / bfloat16) + golden_rtol=2e-2, + ) +) + + +def _final(ctx: dict[str, Any], call): + weights = ctx["weights"] + history = ctx["history"] + return call( + history["adaln_gate_residual"], # the residual stream after the gated sublayer + weights["norm_out.norm.weight"], + history["timestep_mlp_fp32"], # temb + weights["norm_out.linear.weight"], + weights["norm_out.linear.bias"], + ctx["timestep_indices"], + ) + + +STAGES.append( + Stage( + name="final_adaln_out", + op_type="final_adaln_out", + candidate=lambda op, ctx, _up: _final(ctx, op), + provider=lambda ctx, _up: _final(ctx, provider_final_adaln_out), + golden=lambda op, ctx, _up: _final(ctx, op.forward_fp32), + provider_bitwise_isolated=False, # shift/scale projection: tensor-core tree vs cuBLAS + golden_atol=5e-2, # reduction / bfloat16 + golden_rtol=2e-2, + ) +) + +# The backward replay covers the conditioning chain up to the row gather. +CONDITIONING_STAGES = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", +) + +# Parameters whose gradients the backward replay reports, in chain order. +GRAD_LEAVES = ( + "time_embedder.linear_1.weight", + "time_embedder.linear_1.bias", + "time_embedder.linear_2.weight", + "time_embedder.linear_2.bias", + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +) +BACKWARD_MODES = ("candidate", "candidate_fused", "provider") + + +def _flat(value: Any) -> list[torch.Tensor]: + """Collect tensors from nested sequences in their original order.""" + + if isinstance(value, torch.Tensor): + return [value] + return [tensor for item in value for tensor in _flat(item)] + + +def compare(lhs: Any, rhs: Any, atol: float = 0.0, rtol: float = 0.0) -> dict[str, Any]: + """Report tensor equality, maximum FP32 absolute error, and tolerance checks.""" + + lhs_t, rhs_t = _flat(lhs), _flat(rhs) + bitwise = all(torch.equal(a, b) for a, b in zip(lhs_t, rhs_t, strict=True)) + max_abs = 0.0 + within = True + for a, b in zip(lhs_t, rhs_t, strict=True): + diff = (a.float() - b.float()).abs() + if diff.numel(): + max_abs = max(max_abs, float(diff.max())) + within = within and bool((diff <= atol + rtol * b.float().abs()).all()) + return {"bitwise_equal": bitwise, "max_abs": max_abs, "within_tolerance": within} + + +def make_context(weights, *, num_timesteps: int, seq_len: int, seed: int) -> dict[str, Any]: + """Build seeded CUDA timesteps and packed row labels with the supplied weights.""" + + timestep_indices, token_tags = h3_packed_layout(seq_len, num_timesteps, seed=seed) + return { + "timestep": h3_timesteps(num_timesteps, seed=seed), + "timestep_indices": timestep_indices, + "token_tags": token_tags, + "weights": weights, + # Stand-in for the packed hidden states (the patch/text projections are + # other RFC rows): BF16 (1, S, H), seeded. + "hidden": torch.randn( + (1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 2) + ) + .to(torch.bfloat16) + .cuda(), + # Stand-in for the attention output the gate multiplies (the attention + # rows belong to other contributors): BF16 (1, S, H), seeded. + "sublayer": ( + torch.randn((1, seq_len, HIDDEN), generator=torch.Generator().manual_seed(seed + 3)) + * 3.0 + ) + .to(torch.bfloat16) + .cuda(), + } + + +def golden_op(registry: KernelRegistry, op_type: str): + """Return the PyTorch reference at the end of the operator's CPU priority list.""" + + return registry._get_or_create_backend(registry._priority_map["cpu"][op_type][-1]) + + +def run_case( + registry: KernelRegistry, + weights, + *, + num_timesteps: int, + seq_len: int, + seed: int = 0, + stages: list[Stage] | None = None, +) -> dict[str, Any]: + """Replay CUDA candidate, provider, and golden paths and report each stage's drift. + + Supplied stages must include their prerequisites in execution order. + Omitting ``stages`` runs the full registered chain without gradients. + """ + + stages = STAGES if stages is None else stages + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + report: dict[str, Any] = { + "num_timesteps": num_timesteps, + "seq_len": seq_len, + "seed": seed, + "stages": [], + } + chained: Any = None + provider: Any = None + golden: Any = None + first_drift = first_isolated = None + # Each mode sees its own earlier outputs (a stage may read any earlier stage). + history: dict[str, dict[str, Any]] = {"candidate": {}, "provider": {}, "golden": {}} + ctx_c = {**ctx, "history": history["candidate"]} + ctx_p = {**ctx, "history": history["provider"]} + ctx_g = {**ctx, "history": history["golden"]} + for stage in stages: + op = registry.get_op(stage.op_type, device="cuda") + gold = golden_op(registry, stage.op_type) + with torch.no_grad(): + isolated = stage.candidate(op, ctx_p, provider) + chained_input = chained + chained = stage.candidate(op, ctx_c, chained_input) + repeat = stage.candidate(op, ctx_c, chained_input) + golden = stage.golden(gold, ctx_g, golden) + provider = stage.provider(ctx_p, provider) + history["candidate"][stage.name] = chained + history["provider"][stage.name] = provider + history["golden"][stage.name] = golden + entry = { + "stage": stage.name, + "backend": type(op).__name__, + "kernel_id": getattr(op, "kernel_id", type(op).__name__), + "repeat_bitwise_equal": compare(repeat, chained)["bitwise_equal"], + "chained_vs_provider": compare(chained, provider), + "isolated_vs_provider": compare(isolated, provider), + "chained_vs_golden": compare(chained, golden, stage.golden_atol, stage.golden_rtol), + "provider_vs_golden": compare(provider, golden, stage.golden_atol, stage.golden_rtol), + } + if first_drift is None and not entry["chained_vs_provider"]["bitwise_equal"]: + first_drift = stage.name + if first_isolated is None and not entry["isolated_vs_provider"]["bitwise_equal"]: + first_isolated = stage.name + report["stages"].append(entry) + report["first_drift"] = first_drift + report["first_isolated_drift"] = first_isolated + return report + + +def chain_grads(mode: str, registry: KernelRegistry, ctx, upstream) -> list[torch.Tensor]: + """Parameter gradients of the full chain for one execution mode. + + ``candidate`` runs the four RL-Kernel ops separately, ``candidate_fused`` + runs projection + gather as ``H3AdaLNModulationCudaOp``, ``provider`` the + diffusers path, ``golden`` the FP64 references. + """ + + leaves = { + name: ctx["weights"][name].detach().clone().requires_grad_(True) for name in GRAD_LEAVES + } + run_ctx = {**ctx, "weights": {**ctx["weights"], **leaves}, "history": {}} + value = None + if mode == "candidate_fused": + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + + for stage in STAGES[:2]: + value = stage.candidate(registry.get_op(stage.op_type, device="cuda"), run_ctx, value) + value = H3AdaLNModulationCudaOp()( + value, *adaln_params(run_ctx), ctx["timestep_indices"], ctx["token_tags"] + ) + else: + for stage in (s for s in STAGES if s.name in CONDITIONING_STAGES): + if mode == "provider": + value = stage.provider(run_ctx, value) + elif mode == "candidate": + op = registry.get_op(stage.op_type, device="cuda") + value = stage.candidate(op, run_ctx, value) + else: + value = stage.golden(golden_op(registry, stage.op_type), run_ctx, value) + if value is None: + raise ValueError("conditioning chain produced no output") + outputs = list(value) + torch.autograd.backward( + outputs, [g.to(out.dtype) for g, out in zip(upstream, outputs, strict=True)] + ) + return [leaves[name].grad for name in GRAD_LEAVES] + + +def run_backward_case( + registry: KernelRegistry, weights, *, num_timesteps: int, seq_len: int, seed: int = 0 +) -> dict[str, Any]: + """Determinism and accuracy of the chain's parameter gradients, per mode.""" + + ctx = make_context(weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=seed) + generator = torch.Generator(device="cuda").manual_seed(seed + 1) + hidden = weights[GRAD_LEAVES[4]].shape[0] // 18 + upstream = [ + torch.randn(seq_len, hidden, device="cuda", generator=generator).bfloat16() + for _ in range(6) + ] + golden = chain_grads("golden", registry, ctx, upstream) + runs = { + mode: [chain_grads(mode, registry, ctx, upstream) for _ in range(2)] + for mode in BACKWARD_MODES + } + report: dict[str, Any] = {"num_timesteps": num_timesteps, "seq_len": seq_len, "leaves": {}} + for index, name in enumerate(GRAD_LEAVES): + gold = golden[index].float() + entry: dict[str, Any] = {"golden_absmax": float(gold.abs().max())} + for mode, (first, second) in runs.items(): + grad = first[index] + err = (grad.float() - gold).abs().max() + entry[mode] = { + "repeat_bitwise_equal": bool(torch.equal(first[index], second[index])), + "max_abs_vs_golden": float(err), + "max_abs_vs_golden_over_absmax": float(err / gold.abs().max().clamp_min(1e-30)), + "correctly_rounded_fraction": float( + (grad == golden[index].to(grad.dtype)).float().mean() + ), + } + report["leaves"][name] = entry + return report + + +def git_state() -> dict[str, Any]: + """Return HEAD and tracked-tree dirtiness, defaulting to unknown/dirty if Git fails.""" + + def run(*args: str) -> str: + """Run Git in the repository root and return stripped stdout.""" + + return subprocess.check_output(["git", *args], cwd=REPO_ROOT, text=True).strip() + + try: + sha = run("rev-parse", "HEAD") + dirty = bool(run("status", "--porcelain", "--untracked-files=no")) + except Exception: # noqa: BLE001 + sha, dirty = "unknown", True + return {"rl_kernel_commit": sha, "tracked_tree_dirty": dirty} + + +def environment() -> dict[str, Any]: + """Record the active CUDA device, software versions, and matmul TF32 setting.""" + + return { + "gpu": torch.cuda.get_device_name(), + "capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + "tf32": bool(torch.backends.cuda.matmul.allow_tf32), + } diff --git a/rl_engine/validation/models/h3_full_manifest.json b/rl_engine/validation/models/h3_full_manifest.json new file mode 100644 index 000000000..6cc2e5b5a --- /dev/null +++ b/rl_engine/validation/models/h3_full_manifest.json @@ -0,0 +1,311 @@ +{ + "repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "files": [ + { + "path": "audio_scheduler/scheduler_config.json", + "size": 96, + "sha256": "804780f7133477067bd6bbfbc02dc8b3cf9feeb400f97c08f5b1d5f6cbab3840" + }, + { + "path": "audio_vae/config.json", + "size": 2271, + "sha256": "9a3c645ff892b376c6f5f4c8685964cd75474731af594ff058492a0000caabb6" + }, + { + "path": "audio_vae/diffusion_pytorch_model.safetensors", + "size": 605429340, + "sha256": "52c59e67ba8de5477c81bfbced0327aabf500f1bfdeefd5ee754529241cb26cb" + }, + { + "path": "model_index.json", + "size": 2936, + "sha256": "5a587fe13b2371427415ac892463142683aefcd8d322e274a3a095eac37ac7d2" + }, + { + "path": "modular_model_index.json", + "size": 2935, + "sha256": "a2b6a210e482ffb78e613b553f570c44e101afce6741bd4ed91429d0559af031" + }, + { + "path": "processor/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "processor/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "processor/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "processor/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "processor/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "processor/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "processor/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "scheduler/scheduler_config.json", + "size": 97, + "sha256": "8fa6c3aa70dc9e691e1a6df899fd1b6f75f70481a27cee6e18a303817075c304" + }, + { + "path": "text_encoder/chat_template.json", + "size": 5499, + "sha256": "5c72a170d2a4a1a3bc5adad2e689ae28138a9700e5b8c96c0266331e86c0acce" + }, + { + "path": "text_encoder/config.json", + "size": 1474, + "sha256": "d2dd0c60d01b9e195d9447c52da61c7302d28828524914c044d9c6e1b81d0427" + }, + { + "path": "text_encoder/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "text_encoder/model-00001-of-00014.safetensors", + "size": 4932328944, + "sha256": "6b9dfbc930e505402ae9d7e5091a9d7d656cda5f34614f01cfe70bfb0cca27cb" + }, + { + "path": "text_encoder/model-00002-of-00014.safetensors", + "size": 4875990528, + "sha256": "d8bb44b4ff303fe76fe9e894022fb3dc71b15a2e716592790fe0e3c3e60478fa" + }, + { + "path": "text_encoder/model-00003-of-00014.safetensors", + "size": 4875990552, + "sha256": "54f22e8b3168f8dc962fac0d313607ebf52a12b433d4cf3098a0d82d9f042940" + }, + { + "path": "text_encoder/model-00004-of-00014.safetensors", + "size": 4875990584, + "sha256": "ad09c74d3c13ee29b5d0d84548fd8a3424a651564eaccd519946c296e59c557f" + }, + { + "path": "text_encoder/model-00005-of-00014.safetensors", + "size": 4875990584, + "sha256": "fc993c8a0e2a5b0570f383e1a95dc3a1281d1b224b6f3ee908f4827941e1dfc2" + }, + { + "path": "text_encoder/model-00006-of-00014.safetensors", + "size": 4875990584, + "sha256": "82f05620d1f718a90c362b221d6a184ff1a0f53301d706882d3df49695fa1974" + }, + { + "path": "text_encoder/model-00007-of-00014.safetensors", + "size": 4875990584, + "sha256": "fb91da8cb01ff4de3eef0eab1c3e769a734b3a1aafc61734068638a0d6c86934" + }, + { + "path": "text_encoder/model-00008-of-00014.safetensors", + "size": 4875990584, + "sha256": "431ca56535c8781944ce3801f5eb61c45531e853ecc5846d936ebaf4761b764f" + }, + { + "path": "text_encoder/model-00009-of-00014.safetensors", + "size": 4875990584, + "sha256": "3825e3f4302f4d2f7d76aa7430d2ce0864fde6b9e540a5806bf0d8e38e4d9f47" + }, + { + "path": "text_encoder/model-00010-of-00014.safetensors", + "size": 4875990584, + "sha256": "aded5a4d1d5e22dbd8b6f79266b6eb88c840411b09527c53917a1419ace22e2f" + }, + { + "path": "text_encoder/model-00011-of-00014.safetensors", + "size": 4875990584, + "sha256": "3820ffe8d8d6477f6fe8d614ef3c87abb264ee39accebf43a1507b970d80946f" + }, + { + "path": "text_encoder/model-00012-of-00014.safetensors", + "size": 4875990584, + "sha256": "05ad2d08ce71963121c9b03f1d9ec5d7641052f4b23c6c12b80d71065eb8e98e" + }, + { + "path": "text_encoder/model-00013-of-00014.safetensors", + "size": 4875990584, + "sha256": "b64f2289871261fdd1abbd3b78bcd66011b341de3dc8eeb2ed1a473ee7c8d95c" + }, + { + "path": "text_encoder/model-00014-of-00014.safetensors", + "size": 3270697008, + "sha256": "e45b6c9998c77ee5a6577f9f47bc76416c1d4d387169e50c4c9d3134ea51b13b" + }, + { + "path": "text_encoder/model.safetensors.index.json", + "size": 97831, + "sha256": "06c952c569285870b811989b794b9766493e280fb77fbcb957fc4e5fcf25403a" + }, + { + "path": "text_encoder/preprocessor_config.json", + "size": 390, + "sha256": "27225450ac9c6529872ee1924fcb0962ff5634834f817040f444118116f4e516" + }, + { + "path": "text_encoder/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "text_encoder/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "text_encoder/video_preprocessor_config.json", + "size": 385, + "sha256": "7768af27c1fafa9cc9011c1dc20067e03f8915e03b63504550e11d5066986d13" + }, + { + "path": "text_encoder/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "tokenizer/merges.txt", + "size": 1671839, + "sha256": "599bab54075088774b1733fde865d5bd747cbcc7a547c5bc12610e874e26f5e3" + }, + { + "path": "tokenizer/tokenizer.json", + "size": 7032403, + "sha256": "a5d85b6dcc535e6b93115a9ef287e6132fdbf30270da6218194ba742261173c7" + }, + { + "path": "tokenizer/tokenizer_config.json", + "size": 11003, + "sha256": "a07e942ac874baa13758de8d1fbdb186683cc03416b5589e1b6671c6b3057c68" + }, + { + "path": "tokenizer/vocab.json", + "size": 2776833, + "sha256": "ca10d7e9fb3ed18575dd1e277a2579c16d108e32f27439684afa0e10b1440910" + }, + { + "path": "transformer/config.json", + "size": 546, + "sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00014.safetensors", + "size": 4825958704, + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00014.safetensors", + "size": 4702158032, + "sha256": "2c4d362eddd2802180ac9c744849eb9ba8d9c8b984bdf9822cb02ed004b29184" + }, + { + "path": "transformer/diffusion_pytorch_model-00003-of-00014.safetensors", + "size": 4933368192, + "sha256": "949c5aafbbfa5654da730a6a7fafd75adb164d0857b095a30e8bb6d390887d69" + }, + { + "path": "transformer/diffusion_pytorch_model-00004-of-00014.safetensors", + "size": 4567069608, + "sha256": "eef7616790105ee839766bb2027203bf2c0d87c6aa038dca84145a8675f5ce28" + }, + { + "path": "transformer/diffusion_pytorch_model-00005-of-00014.safetensors", + "size": 4702158080, + "sha256": "43fdf42d638e8bc6745f713fae80c93bb301807a1a5ae7249344ce28e202a494" + }, + { + "path": "transformer/diffusion_pytorch_model-00006-of-00014.safetensors", + "size": 4933368232, + "sha256": "6442510b34d173653f0cce5c964b935395a8f7accf0b9cc0aa31aec59805239d" + }, + { + "path": "transformer/diffusion_pytorch_model-00007-of-00014.safetensors", + "size": 4567069608, + "sha256": "29f48f535c91dac76496ca821eeb16ca24bc4caf3f0cae8b920a89b1f966da6d" + }, + { + "path": "transformer/diffusion_pytorch_model-00008-of-00014.safetensors", + "size": 4702158080, + "sha256": "c711b096c764bd60f0b8b6ad49518bfab6d614fb788c725add8741c0674a4cd8" + }, + { + "path": "transformer/diffusion_pytorch_model-00009-of-00014.safetensors", + "size": 4933368232, + "sha256": "44428defe3976cbb87635ad200b958199e739986697cd29fdf27aeb7294b5944" + }, + { + "path": "transformer/diffusion_pytorch_model-00010-of-00014.safetensors", + "size": 4567069608, + "sha256": "3d44939c374c9da382e9c6877e1946adf7b84e08c7a881c068f228d6849411c9" + }, + { + "path": "transformer/diffusion_pytorch_model-00011-of-00014.safetensors", + "size": 4702158080, + "sha256": "224d24430b58127a5577721084e0e704a0e74ec96dd7c35bc6fc0994ebd87c33" + }, + { + "path": "transformer/diffusion_pytorch_model-00012-of-00014.safetensors", + "size": 4933368232, + "sha256": "48fa2bd8fe134eef565ab2464f1c2589a6657cba0d14283dfc06b532f8961f3c" + }, + { + "path": "transformer/diffusion_pytorch_model-00013-of-00014.safetensors", + "size": 4567069608, + "sha256": "be5b4b1809f9d546ffd4b3fcf41e5c1e02b819125caa6bc105c109b04c051bd3" + }, + { + "path": "transformer/diffusion_pytorch_model-00014-of-00014.safetensors", + "size": 4644161920, + "sha256": "8fbd5e6c1fb1df7ce988ca90f3d59e7610e465c7517e4b344eda4a214ba4b97d" + }, + { + "path": "transformer/diffusion_pytorch_model.safetensors.index.json", + "size": 64488, + "sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e" + }, + { + "path": "vae/config.json", + "size": 2011, + "sha256": "78f67deec3d63aae807f2bfe7154bc1e26f6372cb20b63265fcbae1b62bb5745" + }, + { + "path": "vae/diffusion_pytorch_model-00001-of-00003.safetensors", + "size": 5061033024, + "sha256": "72f4c6be84ac0674f27398cde991dd9d719762f3952c4921aa66b2ce542f6374" + }, + { + "path": "vae/diffusion_pytorch_model-00002-of-00003.safetensors", + "size": 4955986528, + "sha256": "2e05e8bc23fa4071043e17fd242be8acd0685e781a43987432b2eae925be4198" + }, + { + "path": "vae/diffusion_pytorch_model-00003-of-00003.safetensors", + "size": 398539336, + "sha256": "c05d6ac4b1a33de372799d708531da6320f6a3ce6d1ce6d895e770988e004a39" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors.index.json", + "size": 74228, + "sha256": "15f6d44553c3c616b0dc999920aa784f92ecee7e4201f1f99ac405cfbf3061ca" + } + ] +} diff --git a/rl_engine/validation/models/h3_full_model.py b/rl_engine/validation/models/h3_full_model.py new file mode 100644 index 000000000..1219ef89f --- /dev/null +++ b/rl_engine/validation/models/h3_full_model.py @@ -0,0 +1,242 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Instrument the full pinned H3 provider with the actual RL-Kernel block graph. + +Input/output projections and both text-refiner layers retain provider arithmetic. +Each of the 50 distinct blocks executes the existing candidate graph with its own +parameters. No weights are copied from block zero. This is integration evidence, +not qualification of the reference/interim bindings as final RFC kernels. +""" + +from __future__ import annotations + +from collections import Counter +from contextlib import contextmanager +from types import MethodType + +import torch + +from rl_engine.validation.models import h3_block + +MODEL_REVISION = "42ed227ee7df40d41602854ae760620d6eb651fe" +PROVIDER_REVISION = "f53d552036a0d1bd5570782a39cd40cfabf112bc" + + +def validate_full_structure(model): + expected = dict( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + for key, value in expected.items(): + if getattr(model.config, key) != value: + raise ValueError(f"full H3 requires {key}={value}") + blocks = model.transformer_blocks + if len(blocks) != 50 or len({id(b) for b in blocks}) != 50: + raise ValueError("50 distinct transformer blocks are required") + owned = set() + for i, block in enumerate(blocks): + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError(f"block {i} has missing or shared parameters") + owned.update(id(p) for p in params) + required = ( + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ) + if any(not hasattr(model, name) for name in required): + raise ValueError("full-model boundary components are missing") + refiners = model.token_refiner.refiner_blocks + if len(refiners) != 2 or len({id(b) for b in refiners}) != 2: + raise ValueError("two distinct real refiner blocks are required") + for block in refiners: + params = list(block.parameters()) + if not params or any(id(p) in owned for p in params): + raise ValueError("refiner parameters are missing or shared") + owned.update(id(p) for p in params) + return required + + +class FullModelTrace: + def __init__(self): + self.calls = Counter() + self.backward_calls = Counter() + self.snapshots = {} + self.capture = False + self.actual_backends = {} + self.candidate = False + self.candidate_calls = Counter() + self.forward_order = [] + self.backward_order = [] + self.execution_order = [] + self._events = set() + + def record_event(self, name): + if name not in self._events: + self._events.add(name) + self.execution_order.append(name) + + def observe(self, name, value): + self.calls[name] += 1 + self.record_event("nodes." + name) + if name not in self.forward_order: + self.forward_order.append(name) + if not isinstance(value, torch.Tensor): + return + if self.capture: + self.snapshots[name] = value.detach().cpu() + if value.requires_grad: + + def record(grad): + self.backward_calls[name] += 1 + self.record_event("nodes." + name + ".grad") + if name not in self.backward_order: + self.backward_order.append(name) + if self.capture: + self.snapshots[name + ".grad"] = grad.detach().cpu() + + value.register_hook(record) + + def validate(self, *, backward=False): + required = ( + [f"block.{i}" for i in range(50)] + + ["refiner.0", "refiner.1"] + + [ + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ] + ) + missing = [name for name in required if not self.calls[name]] + if backward: + missing += [name + ".backward" for name in required if not self.backward_calls[name]] + if missing: + raise ValueError(f"incomplete full-model execution: {missing}") + if self.candidate: + missing_nodes = [ + f"block.{i}.{node.name}" + for i in range(50) + for node in h3_block.NODES + if not self.candidate_calls[f"block.{i}.{node.name}"] + ] + if missing_nodes: + raise ValueError(f"candidate nodes did not execute: {missing_nodes}") + return dict( + candidate_calls=dict(self.candidate_calls), + execution_order=self.execution_order, + forward_order=self.forward_order, + backward_order=self.backward_order, + actual_backends=self.actual_backends, + calls=dict(self.calls), + backward_calls=dict(self.backward_calls), + bindings={ + n.name: dict( + rfc_row=n.binding.rfc_row, status=n.binding.status, backend=n.binding.backend + ) + for n in h3_block.NODES + }, + boundary_binding="provider/reference", + full_transformer_blocks=50, + refiner_layers=2, + ) + + +def require_candidate_cuda(ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + if ".cuda." not in backends[name]: + raise ValueError(f"candidate {name} fell back from its CUDA implementation") + return backends + + +@contextmanager +def full_model_execution( + model, *, candidate: bool, trace: FullModelTrace, ops=None, provider_nodes=() +): + """Temporarily bind each real block, preserving module parameters and restoring on error. + + ``provider_nodes`` (diagnostic hybrids only) names candidate graph nodes that run + their provider replay in every block; such a run localises drift and is never a + candidate result. + """ + unknown = set(provider_nodes) - {n.name for n in h3_block.NODES} + if unknown: + raise ValueError(f"unknown block nodes: {sorted(unknown)}") + boundaries = validate_full_structure(model) + originals, handles = [], [] + trace.candidate = candidate + if candidate: + ops = ops or h3_block.CandidateOps() + if hasattr(ops, "backends"): + trace.actual_backends = require_candidate_cuda(ops) + try: + for i, block in enumerate(model.transformer_blocks): + if candidate: + original = block.forward + originals.append((block, original)) + + def forward( + self, + hidden_states, + temb, + adaln_indices, + rotary_emb, + attention_mask=None, + _block_index=i, + ): + if attention_mask is not None: + raise ValueError( + "the qualified H3 candidate requires an unmasked packed sequence" + ) + weights = {h3_block.BLOCK + "." + n: p for n, p in self.named_parameters()} + ctx = dict( + weights=weights, + hidden=hidden_states, + temb=temb, + index=adaln_indices, + rope=rotary_emb, + provider_nodes=frozenset(provider_nodes), + node_observer=lambda name: trace.candidate_calls.update( + [f"block.{_block_index}.{name}"] + ), + ) + return h3_block.run_graph("candidate", ctx, ops)[h3_block.OUTPUT] + + block.forward = MethodType(forward, block) + + def hook(module, args, output, name=f"block.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for i, block in enumerate(model.token_refiner.refiner_blocks): + + def hook(module, args, output, name=f"refiner.{i}"): + trace.observe(name, output) + + handles.append(block.register_forward_hook(hook)) + for name in boundaries: + + def hook(module, args, output, name=name): + trace.observe(name, output) + + handles.append(getattr(model, name).register_forward_hook(hook)) + yield trace + finally: + for handle in handles: + handle.remove() + for block, original in originals: + block.forward = original diff --git a/rl_engine/validation/models/h3_manifest.json b/rl_engine/validation/models/h3_manifest.json new file mode 100644 index 000000000..e590a8473 --- /dev/null +++ b/rl_engine/validation/models/h3_manifest.json @@ -0,0 +1,137 @@ +{ + "version": "h3-conditioning-v1", + "rfc": "RL-Align/RL-Kernel#420", + "model_identity": { + "hf_repo": "MiniMaxAI/MiniMax-H3", + "revision": "42ed227ee7df40d41602854ae760620d6eb651fe", + "subfolder": "transformer", + "architecture": "MiniMaxH3Transformer3DModel", + "config_sha256": "74c11bff524336576096993cbfcdcdc2ef4fa2fa4409df693bdcbc6c666282ae", + "index_file": "diffusion_pytorch_model.safetensors.index.json", + "index_sha256": "ac30a3b58963f2e735d493475fbb81853a5735ec947619648b3e045acda6783e", + "config_fingerprint": { + "hidden_size": 5376, + "num_attention_heads": 56, + "attention_head_dim": 128, + "num_layers": 50, + "num_refiner_layers": 2, + "ffn_dim": 14336, + "freq_dim": 256, + "time_embed_hidden_dim": 5376, + "time_embed_dim": 2688, + "norm_eps": 1e-05, + "final_norm_eps": 1e-05 + } + }, + "weight_shards": { + "diffusion_pytorch_model-00001-of-00014.safetensors": { + "sha256": "2d847200c45c09dd7f973c1b096663068408ef851ee0b3711d059b6dc5dcd028", + "size_bytes": 4825958704 + } + }, + "tensors": { + "time_embedder.linear_1.weight": { + "dtype": "float32", "shape": [5376, 256], + "sha256": "4c50eabe3c434dc05422e6501cdf7baa5cecbb1a7706ae9599c7a5b0e34e2913" + }, + "time_embedder.linear_1.bias": { + "dtype": "float32", "shape": [5376], + "sha256": "64763a5e39cb078eb0d8c5be5d56ad36cac3dec37be93a87e106d82ea6c071b7" + }, + "time_embedder.linear_2.weight": { + "dtype": "float32", "shape": [2688, 5376], + "sha256": "697819dd9413e0fec3a50987a964f1fd5db568ca545267414238f452294a8154" + }, + "time_embedder.linear_2.bias": { + "dtype": "float32", "shape": [2688], + "sha256": "096a3c7539f0248698cc65376c8da260e7365bf4ccef847c09bb9a36782a743c" + }, + "transformer_blocks.0.adaln_proj.linear.weight": { + "dtype": "bfloat16", "shape": [96768, 2688], + "sha256": "a0000193b30ea36a1996f53fc0d6c4c37ce5fd4585284f0f9bae570d25c6b8cf" + }, + "transformer_blocks.0.adaln_proj.linear.bias": { + "dtype": "bfloat16", "shape": [96768], + "sha256": "d59f9de9e07e6b4e3b57009627e276a26971ce379555fc99dbb06d3a237c8fc0" + }, + "transformer_blocks.0.norm1.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "5e32e4d7379a9f9302755af395ed50b5b998352a221dc2c99204a85dd9f9b49f" + }, + "transformer_blocks.0.norm2.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "a2c513ec459414e211a3d37f77c1f556f11117f56516b22eb5912f8ff5f98872" + }, + "token_refiner.final_norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "cf10a56a11370a216a4ac3afe9e8b8a3133206bccd6011a92b9aa791a9196b93" + }, + "norm_out.norm.weight": { + "dtype": "bfloat16", "shape": [5376], + "sha256": "91ac17792929e0a84533b28caf71ac831ec3d666c07a473393639be20392f5ce" + }, + "norm_out.linear.weight": { + "dtype": "bfloat16", "shape": [10752, 2688], + "sha256": "38063ae4b492aacd320e71d4a38a32325c0d47b29f59978f0f4979ca080fe673" + }, + "norm_out.linear.bias": { + "dtype": "bfloat16", "shape": [10752], + "sha256": "3b12cdc7bbda772e7e31539ffa9a32f529dba7f4148273d2128bdc621fd27438" + } + }, + "block_tensors": { + "transformer_blocks.0.attn.to_q.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "196f9120ad018784bcefd4224f0e08fae88ba1adf079ecdca18410a414a6a2a4" + }, + "transformer_blocks.0.attn.to_k.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "8e9171f3b8cf894904d1ddbd96510739d7291d2992310de6a52e1d733ec895cd" + }, + "transformer_blocks.0.attn.to_v.weight": { + "dtype": "bfloat16", "shape": [7168, 5376], + "sha256": "275bbc993b1ba738831e04d86f290c9988108d96fe8c97c3b576a043c7ff3fe1" + }, + "transformer_blocks.0.attn.norm_q.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "34a3f13557a2c0a9af260e48d431041e381779a0f61616cfe9800ec839693519" + }, + "transformer_blocks.0.attn.norm_k.weight": { + "dtype": "bfloat16", "shape": [128], + "sha256": "8991949ddea3bc2136a0fcf0b980fa63bd0df760207e3116f7ce96df4e6ca058" + }, + "transformer_blocks.0.attn.to_out.0.weight": { + "dtype": "bfloat16", "shape": [5376, 7168], + "sha256": "822cb426aaebe57d3d219b3c4f6855ca5c13821cb84ebe0a7474fe3a87869fd6" + }, + "transformer_blocks.0.ff.net.0.proj.weight": { + "dtype": "bfloat16", "shape": [28672, 5376], + "sha256": "49587044f325aae73f046c4cda9637052670227ec3c2c289359d9d1ebf29f95d" + }, + "transformer_blocks.0.ff.net.2.weight": { + "dtype": "bfloat16", "shape": [5376, 14336], + "sha256": "e19d4e1649e4b55b1931842e5cc4bb46892c2e635b9315b2996630a86284e88f" + } + }, + "reference_implementation": { + "repo": "huggingface/diffusers", + "commit": "f53d552036a0d1bd5570782a39cd40cfabf112bc", + "file": "src/diffusers/models/transformers/transformer_minimax_h3.py", + "file_sha256": "435873d9bc26b8770342961e18cb20324aeeaa4d045e85e0941638d4857d7db3", + "timestep_embedding_file": "src/diffusers/models/embeddings.py" + }, + "conventions": { + "timestep": "distinct timesteps t = 1 - sigma, unscaled in [0, 1], shape (T,)", + "modality_tags": {"video": 0, "text": 1, "audio": 2}, + "adaln_row": "timestep_index * 3 + token_tag", + "adaln_table_rows": "[t0_mod0, t0_mod1, t0_mod2, t1_mod0, ...]", + "adaln_chunk_order": ["shift_msa", "scale_msa", "gate_msa", "shift_mlp", "scale_mlp", "gate_mlp"], + "precision": { + "timestep_sinusoid": "float32", + "timestep_mlp": "float32 weights, float32 activations", + "adaln_projection": "SiLU in float32, cast to bfloat16, bfloat16 weight/bias/output", + "rmsnorm": "nn.RMSNorm(5376, eps=1e-5), bfloat16 affine weight, statistics in float32", + "final_modulation": "norm_out.linear(silu(temb).bfloat16) -> (shift, scale), indexed by timestep_indices" + } + } +} diff --git a/rl_engine/validation/models/h3_provider.py b/rl_engine/validation/models/h3_provider.py new file mode 100644 index 000000000..3eec1818f --- /dev/null +++ b/rl_engine/validation/models/h3_provider.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The diffusers MiniMax-H3 conditioning path, op for op, without diffusers. + +Each function replays the tensor ops of the pinned reference +(``huggingface/diffusers@f53d552``, see ``h3_manifest.json``) in the same +order and dtypes, so it runs the same kernels the provider does on a given +device. This is the *provider* side of the RFC #420 comparisons; the FP32/FP64 +goldens live in ``rl_engine.reference.minimax_h3``. +""" + +from __future__ import annotations + +import math + +import torch +import torch.nn.functional as F + + +def provider_time_proj(timestep: torch.Tensor, num_channels: int = 256) -> torch.Tensor: + """Replay provider FP32 cosine/sine features for ``(T,)`` timesteps. + + Matches ``Timesteps(256, flip_sin_to_cos=True, downscale_freq_shift=0)`` + with the default channel count, returning shape ``(T, 256)``. + """ + + half_dim = num_channels // 2 + exponent = -math.log(10000) * torch.arange( + start=0, end=half_dim, dtype=torch.float32, device=timestep.device + ) + exponent = exponent / (half_dim - 0) + emb = torch.exp(exponent) + emb = timestep[:, None].float() * emb[None, :] + emb = 1 * emb + emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) + emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) + return emb + + +def provider_time_embedder( + features: torch.Tensor, + w1: torch.Tensor, + b1: torch.Tensor, + w2: torch.Tensor, + b2: torch.Tensor, +) -> torch.Tensor: + """Replay linear_1, SiLU, and linear_2 using the checkpoint weight dtype. + + H3's FP32 weights map features of shape ``(T, 256)`` to ``(T, 2688)``. + """ + + sample = F.linear(features.to(w1.dtype), w1, b1) + sample = F.silu(sample) + return F.linear(sample, w2, b2) + + +def provider_adaln_modulation( + temb: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, hidden_size: int = 5376 +) -> tuple[torch.Tensor, ...]: + """Replay provider SiLU, the weight-dtype cast, and the AdaLN projection. + + Return six ``(3T, hidden_size)`` views in timestep-major modality order + for checkpoint weights shaped ``(18 * hidden_size, 2688)``. + """ + + temb = F.linear(F.silu(temb).to(weight.dtype), weight, bias) + temb = temb.view(-1, 6 * hidden_size) + return temb.chunk(6, dim=-1) + + +def provider_adaln_row_gather( + modulation: tuple[torch.Tensor, ...], timestep_indices: torch.Tensor, token_tags: torch.Tensor +) -> tuple[torch.Tensor, ...]: + """``adaln_indices = timestep_indices * 3 + token_tags`` then six ``index_select``.""" + + adaln_indices = timestep_indices * 3 + token_tags + return tuple(tensor.index_select(0, adaln_indices) for tensor in modulation) + + +def provider_norm_modulate( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``norm(x) * (1.0 + scale[i]) + shift[i]`` with ``index_select``; block and norm_out.""" + + norm_hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return norm_hidden_states * (1.0 + scale.index_select(0, indices)) + shift.index_select( + 0, indices + ) + + +def provider_gate_residual( + residual: torch.Tensor, gate: torch.Tensor, indices: torch.Tensor, sublayer_output: torch.Tensor +) -> torch.Tensor: + """``residual + gate.index_select(0, adaln_indices) * sublayer_output`` (block).""" + + return residual + gate.index_select(0, indices) * sublayer_output + + +def provider_final_adaln_out( + hidden_states: torch.Tensor, + norm_weight: torch.Tensor, + temb: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + timestep_indices: torch.Tensor, + eps: float = 1e-5, +) -> torch.Tensor: + """``MiniMaxH3AdaLayerNormOut.forward``: shift first, indexed by timestep.""" + + shift, scale = F.linear(F.silu(temb).to(weight.dtype), weight, bias).chunk(2, dim=-1) + hidden_states = F.rms_norm(hidden_states, (hidden_states.shape[-1],), norm_weight, eps) + return hidden_states * (1.0 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +# --- One transformer block (RFC #420 ``ws1_one_h3_block``) ------------------------------------ + + +def provider_rope( + position_ids: torch.Tensor, rope_freq_dim: int = 16, rope_theta: float = 10000.0 +) -> tuple[torch.Tensor, torch.Tensor]: + """``MiniMaxH3RotaryPosEmbed``: FP32 ``(S, 96)`` cos/sin over the ``(t, h, w)`` axes.""" + + inv_freq = 1.0 / ( + rope_theta + ** ( + torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32, device=position_ids.device) + / (2 * rope_freq_dim) + ) + ) + position_ids = position_ids.to(torch.float32) + freqs = position_ids.unsqueeze(-1) * inv_freq.view(1, 1, -1) + freqs_t, freqs_h, freqs_w = freqs.unbind(dim=1) + freqs = torch.cat((freqs_t, freqs_h, freqs_w), dim=-1) + freqs = torch.cat((freqs, freqs), dim=-1) + return freqs.cos(), freqs.sin() + + +def provider_apply_rotary( + hidden_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> torch.Tensor: + """``_apply_rotary_emb``: rotate-half on the leading 96 of 128 channels of ``(B, S, H, D)``. + + The FP32 tables are cast to the activation dtype (BF16) before the rotation. + """ + + rotary_dim = cos.shape[-1] + hidden_states_rotary = hidden_states[..., :rotary_dim] + hidden_states_pass = hidden_states[..., rotary_dim:] + cos = cos.to(hidden_states.dtype)[None, :, None, :] + sin = sin.to(hidden_states.dtype)[None, :, None, :] + x1, x2 = hidden_states_rotary.chunk(2, dim=-1) + hidden_states_rotated = torch.cat((-x2, x1), dim=-1) + hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin + return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous() + + +def provider_linear(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """A bias-free ``nn.Linear`` of the block (Q/K/V, O and both FFN projections).""" + + return F.linear(x, weight) + + +def provider_head_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-5) -> torch.Tensor: + """``norm_q``/``norm_k``: ``nn.RMSNorm(128)`` on ``(B, S, H, D)``.""" + + return F.rms_norm(x, (x.shape[-1],), weight, eps) + + +def provider_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + """The native diffusers backend: SDPA on ``(B, H, S, D)`` views, non-causal, no mask. + + Returns ``(B, S, H*D)`` after ``flatten(2, 3).type_as(query)``. + """ + + q, k, v = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + return out.permute(0, 2, 1, 3).flatten(2, 3).type_as(query) + + +def provider_swiglu(projected: torch.Tensor) -> torch.Tensor: + """diffusers ``SwiGLU`` after its projection: ``hidden * silu(gate)`` on the two halves.""" + + hidden_states, gate = projected.chunk(2, dim=-1) + return hidden_states * F.silu(gate) diff --git a/rl_engine/validation/models/h3_report.py b/rl_engine/validation/models/h3_report.py new file mode 100644 index 000000000..16b212e6d --- /dev/null +++ b/rl_engine/validation/models/h3_report.py @@ -0,0 +1,831 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Per-operator performance and accuracy measurements for the H3 evidence. + +``PERF_CASES[op]`` builds the timed cases that ``benchmarks/models/benchmark_h3_conditioning.py`` +prints and ``tools/validation/models/h3_evidence.py`` stores. ``ACCURACY[op]`` measures the +op against the provider path and its golden. Each RFC #420 row registers both. +""" + +from __future__ import annotations + +import statistics +from typing import Any, Callable + +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models.h3_cases import h3_packed_layout, h3_timesteps +from rl_engine.validation.models.h3_provider import ( + provider_adaln_modulation, + provider_adaln_row_gather, + provider_final_adaln_out, + provider_gate_residual, + provider_norm_modulate, + provider_time_embedder, + provider_time_proj, +) +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_conditioning_weights, +) + +# Keys of a perf case that hold a timed callable, in display order. +TIMED_KEYS = ( + "candidate", + "candidate_checked", + "provider", + "candidate_backward", + "provider_backward", +) + + +def _prepared_call( + fn: Callable[..., Any], setup: Callable[[], Any] | None = None +) -> Callable[[], Any]: + if setup is None: + return fn + prepared = setup() + return lambda: fn(prepared) + + +def _sample_us(fn: Callable[..., Any], setup: Callable[[], Any] | None = None) -> float: + run = _prepared_call(fn, setup) + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + run() + end.record() + end.synchronize() + return start.elapsed_time(end) * 1e3 + + +def time_us( + fn: Callable[..., Any], + warmup: int = 20, + iters: int = 200, + *, + setup: Callable[[], Any] | None = None, +) -> float: + """Median CUDA-event microseconds; optional setup runs outside the timed region. + + When supplied, ``setup`` runs before every call and its return value is passed + to ``fn``. Backward measurements use it to build a fresh forward graph. + """ + + for _ in range(warmup): + _prepared_call(fn, setup)() + torch.cuda.synchronize() + return statistics.median(_sample_us(fn, setup) for _ in range(iters)) + + +def peak_mib(fn: Callable[..., Any], *, setup: Callable[[], Any] | None = None) -> float: + """Incremental peak CUDA memory of the call, excluding optional setup.""" + + run = _prepared_call(fn, setup) + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + run() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - base) / 2**20 + + +def measure(case: dict[str, Any], warmup: int = 20, iters: int = 200) -> dict[str, Any]: + """Interleave timed keys, reversing their execution order on alternate iterations. + + Record raw CUDA-event samples, executed orders, median microseconds, + decimal GB/s from ``case['bytes']``, and peak allocated MiB per callable. + """ + + row: dict[str, Any] = { + key: case[key] for key in ("op", "case", "backend", "bytes") if key in case + } + keys = [key for key in TIMED_KEYS if key in case] + if not keys: + return row + orders = (keys, list(reversed(keys))) + row["execution_order"] = { + "policy": "alternating", + "iteration_0": orders[0], + "iteration_1": orders[1], + } + if any(key.endswith("_backward") for key in keys): + row["backward_timing_scope"] = "backward_only" + for iteration in range(warmup): + for key in orders[iteration % 2]: + _prepared_call(case[key], case.get(f"{key}_setup"))() + torch.cuda.synchronize() + samples: dict[str, list[float]] = {key: [] for key in keys} + for iteration in range(iters): + for key in orders[iteration % 2]: + samples[key].append(_sample_us(case[key], case.get(f"{key}_setup"))) + row["timing_samples_us"] = samples + row["timing_order"] = [orders[iteration % 2] for iteration in range(iters)] + for key in keys: + us = statistics.median(samples[key]) + row[f"{key}_us"] = us + row[f"{key}_gbps"] = case["bytes"] / (us * 1e-6) / 1e9 + row[f"{key}_peak_mib"] = peak_mib(case[key], setup=case.get(f"{key}_setup")) + return row + + +# --------------------------------------------------------------------------- # +# timestep_sinusoid_h3 +# --------------------------------------------------------------------------- # + + +def _sinusoid_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA sinusoid timing cases with checked, unchecked, and provider calls.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + cases = [] + for num in (1, 2, 4, 64): + t = h3_timesteps(num) + cases.append( + { + "op": "timestep_sinusoid_h3", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": num * 4 + num * 256 * 4, + "candidate": lambda t=t: op.forward(t, check_range=False), + "candidate_checked": lambda t=t: op.forward(t), + "provider": lambda t=t: provider_time_proj(t), + } + ) + return cases + + +def _sinusoid_accuracy(registry: KernelRegistry) -> dict[str, Any]: + """Compare sinusoid outputs to provider and FP64 golden across timestep counts.""" + + op = registry.get_op("timestep_sinusoid_h3", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["timestep_sinusoid_h3"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + rows = [] + for num in (1, 2, 3, 7, 64, 1000, 4097): + t = h3_timesteps(num, seed=num) + out, ref, gold = op(t), provider_time_proj(t), golden.forward_fp32(t) + rows.append( + { + "num_timesteps": num, + "bitwise_equal_to_provider": bool(torch.equal(out, ref)), + "max_abs_vs_fp64": float((out - gold).abs().max()), + "provider_max_abs_vs_fp64": float((ref - gold).abs().max()), + } + ) + return {"contract_atol": 1e-5, "cases": rows} + + +def h3_params(names: list[str], shapes: list[tuple[int, ...]]) -> list[torch.Tensor]: + """Return CUDA checkpoint parameters, or seeded random tensors when weights are unset.""" + + if h3_weights_dir() is not None: + weights = load_h3_conditioning_weights("cuda", names) + return [weights[name] for name in names] + generator = torch.Generator(device="cuda").manual_seed(0) + return [ + torch.randn(shape, device="cuda", generator=generator) * shape[-1] ** -0.5 + for shape in shapes + ] + + +# --------------------------------------------------------------------------- # +# timestep_mlp_fp32 +# --------------------------------------------------------------------------- # + +MLP_NAMES = [f"time_embedder.linear_{i}.{p}" for i in (1, 2) for p in ("weight", "bias")] +MLP_SHAPES = [(5376, 256), (5376,), (2688, 5376), (2688,)] + + +def _mlp_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA MLP candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + weight_bytes = sum(p.numel() * p.element_size() for p in params) + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + x = torch.rand(num, 256, device="cuda", generator=gen) * 2 - 1 + cases.append( + { + "op": "timestep_mlp_fp32", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight_bytes, + "candidate": lambda x=x: op(x, *params), + "provider": lambda x=x: provider_time_embedder(x, *params), + } + ) + return cases + + +def _mlp_accuracy(registry: KernelRegistry, draws: int = 200) -> dict[str, Any]: + """Measure per-draw MLP error against FP64 and test single-row/batched equality.""" + + op = registry.get_op("timestep_mlp_fp32", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["timestep_mlp_fp32"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + sinusoid = registry.get_op("timestep_sinusoid_h3", device="cuda") + params = h3_params(MLP_NAMES, MLP_SHAPES) + cuda_err, provider_err = [], [] + for seed in range(draws): + x = sinusoid(h3_timesteps(4, seed=seed)) + gold = golden.forward_fp32(x, *params) + cuda_err.append(float((op(x, *params) - gold).abs().max())) + provider_err.append(float((provider_time_embedder(x, *params) - gold).abs().max())) + # Row invariance: each timestep alone equals its row in the batch of 9. + x = sinusoid(h3_timesteps(9, seed=1)) + full = op(x, *params) + invariant = all(torch.equal(op(x[i : i + 1], *params)[0], full[i]) for i in range(9)) + return { + "contract_atol": 1e-4, + "draws": draws, + "num_timesteps": 4, + "cuda_max_abs_vs_fp64": cuda_err, + "provider_max_abs_vs_fp64": provider_err, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_projection_3mod +# --------------------------------------------------------------------------- # + +ADALN_NAMES = [ + "transformer_blocks.0.adaln_proj.linear.weight", + "transformer_blocks.0.adaln_proj.linear.bias", +] + + +def _adaln_params() -> tuple[torch.Tensor, torch.Tensor]: + """Return block-0 AdaLN projection weights and bias as BF16 CUDA tensors.""" + + weight, bias = h3_params(ADALN_NAMES, [(96768, 2688), (96768,)]) + return weight.bfloat16(), bias.bfloat16() + + +def _projection_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + """Build CUDA AdaLN candidate/provider timing cases for one to four timesteps.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + weight, bias = _adaln_params() + cases = [] + for num in (1, 2, 3, 4): + gen = torch.Generator(device="cuda").manual_seed(num) + temb = torch.randn(num, 2688, device="cuda", generator=gen) + cases.append( + { + "op": "adaln_projection_3mod", + "case": f"T={num}", + "backend": type(op).__name__, + "bytes": weight.numel() * weight.element_size(), + "candidate": lambda temb=temb: op(temb, weight, bias), + "provider": lambda temb=temb: provider_adaln_modulation(temb, weight, bias), + } + ) + return cases + + +def _flat_cat(outputs) -> torch.Tensor: + """Concatenate flattened modulation tensors in their output order.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +def _projection_accuracy(registry: KernelRegistry, draws: int = 20) -> dict[str, Any]: + """Report per-draw BF16 equality to FP64 goldens and single-row/batch invariance.""" + + op = registry.get_op("adaln_projection_3mod", device="cuda") + golden = registry._get_or_create_backend( + registry._priority_map["cpu"]["adaln_projection_3mod"][-1] + ) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + weight, bias = _adaln_params() + cuda_frac, provider_frac, early_frac = [], [], [] + for seed in range(draws): + g = torch.Generator(device="cuda").manual_seed(seed) + temb = torch.randn(4, 2688, device="cuda", generator=g) * 2 + gold = _flat_cat(golden.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat_cat(golden.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + ours = _flat_cat(op(temb, weight, bias)) + theirs = _flat_cat(provider_adaln_modulation(temb, weight, bias)) + cuda_frac.append(float((ours == gold).float().mean())) + provider_frac.append(float((theirs == gold).float().mean())) + early_frac.append(float((ours == early).float().mean())) + gen = torch.Generator(device="cuda").manual_seed(draws) + temb = torch.randn(9, 2688, device="cuda", generator=gen) + full = op(temb, weight, bias) + invariant = all( + all( + torch.equal(s, f[3 * i : 3 * i + 3]) + for s, f in zip(op(temb[i : i + 1], weight, bias), full, strict=True) + ) + for i in range(9) + ) + return { + "draws": draws, + "num_timesteps": 4, + "cuda_correctly_rounded": cuda_frac, + "provider_correctly_rounded": provider_frac, + "early_cast_golden_match": early_frac, + "rows_batch_invariant": invariant, + } + + +# --------------------------------------------------------------------------- # +# adaln_row_gather +# --------------------------------------------------------------------------- # + +GATHER_SEQ_LENS = (4097, 32768, 131072) + + +def _gather_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_row_gather", device="cuda") + num_timesteps, hidden = 3, 5376 + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(3 * num_timesteps, 6 * hidden, device="cuda", generator=gen).bfloat16() + chunks = rows.chunk(6, dim=-1) + cases = [] + for seq in GATHER_SEQ_LENS: + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + grads = [ + torch.randn(seq, hidden, device="cuda", generator=gen).bfloat16() for _ in range(6) + ] + + def prepare_backward(fn, ti=ti, tags=tags): + leaf = rows.detach().requires_grad_(True) + return list(fn(leaf, ti, tags)) + + def backward(outputs, grads=grads): + torch.autograd.backward(outputs, grads) + + cases.append( + { + "op": "adaln_row_gather", + "case": f"S={seq}", + "backend": type(op).__name__, + # bytes written (six (S, H) BF16 outputs); the 3T-row table stays in L2 + "bytes": 6 * seq * hidden * 2, + "candidate": lambda ti=ti, tags=tags: op.forward(rows, ti, tags, check_range=False), + "provider": lambda ti=ti, tags=tags: provider_adaln_row_gather(chunks, ti, tags), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: op.forward(r, indices, row_tags, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward: b( + lambda r, indices, row_tags: provider_adaln_row_gather( + r.chunk(6, dim=-1), indices, row_tags + ) + ), + } + ) + return cases + + +def _gather_accuracy(registry: KernelRegistry) -> dict[str, Any]: + from rl_engine.validation.models.h3_chain import run_backward_case + + op = registry.get_op("adaln_row_gather", device="cuda") + golden = registry._get_or_create_backend(registry._priority_map["cpu"]["adaln_row_gather"][-1]) + if golden is None: + raise RuntimeError("required H3 golden backend is unavailable") + # Seeded so the op-backward fractions below are reproducible from this commit. + gen = torch.Generator(device="cuda").manual_seed(0) + rows = torch.randn(9, 6 * 5376, device="cuda", generator=gen).bfloat16() + forward_bitwise = {} + for seq in (1, 257, 4097, 32768): + ti, tags = h3_packed_layout(seq, 3, seed=seq) + ours = op(rows, ti, tags) + theirs = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + forward_bitwise[str(seq)] = all( + torch.equal(a, b) for a, b in zip(ours, theirs, strict=True) + ) + + # Op-level backward: correctly rounded FP32 segment sums vs the atomic BF16 scatter-add. + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = [torch.randn(4097, 5376, device="cuda", generator=gen).bfloat16() for _ in range(6)] + + def grad_of(fn): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(fn(leaf)), grads) + return leaf.grad + + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward(list(golden.forward_fp32(ref, ti, tags)), [g.float() for g in grads]) + gold = ref.grad.bfloat16() + backward = {} + for name, fn in ( + ("cuda", lambda leaf: op(leaf, ti, tags)), + ("provider", lambda leaf: provider_adaln_row_gather(leaf.chunk(6, dim=-1), ti, tags)), + ): + first, second = grad_of(fn), grad_of(fn) + backward[name] = { + "repeat_bitwise_equal": bool(torch.equal(first, second)), + "correctly_rounded_fraction": float((first == gold).float().mean()), + } + + result: dict[str, Any] = { + "forward_bitwise_vs_index_select": forward_bitwise, + "op_backward": backward, + "chain_backward": [], + } + if h3_weights_dir() is None: + result["chain_backward_skipped"] = ( + f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py for chain measurements" + ) + else: + weights = load_h3_conditioning_weights("cuda") + result["chain_backward"] = [ + run_backward_case(registry, weights, num_timesteps=t, seq_len=s) + for t, s in ((1, 257), (3, 257), (1, 4097), (3, 4097), (4, 32768)) + ] + return result + + +# --------------------------------------------------------------------------- # +# h3_rmsnorm (block norm1 + MSA modulation) +# --------------------------------------------------------------------------- # + +NORM_SEQ_LENS = (4097, 32768, 131072) + + +def _norm_inputs(seq: int, seed: int = 0): + weight = h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0].bfloat16() + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).bfloat16() + shift, scale = table.view(9, 6, 5376)[:, 0], table.view(9, 6, 5376)[:, 1] + ti, tags = h3_packed_layout(seq, 3, seed=seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + return x, weight, shift, scale, ti * 3 + tags + + +def _norm_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("h3_rmsnorm", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, weight, shift, scale, index = _norm_inputs(seq) + grad = torch.randn_like(x) + + def prepare_backward(fn, x=x, weight=weight, shift=shift, scale=scale): + leaves = [t.detach().requires_grad_(True) for t in (x, weight, shift, scale)] + return fn(*leaves) + + def backward(output, grad=grad): + output.backward(grad) + + cases.append( + { + "op": "h3_rmsnorm", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 2 * seq * 5376 * 2, # read x, write y + "candidate": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + op.forward_modulated(x, w, sh, sc, i, check_range=False) + ), + "provider": lambda x=x, w=weight, sh=shift, sc=scale, i=index: ( + provider_norm_modulate(x, w, sh, sc, i) + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: op.forward_modulated( + x_, w_, sh_, sc_, i, check_range=False + ) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda x_, w_, sh_, sc_: provider_norm_modulate(x_, w_, sh_, sc_, i) + ), + } + ) + return cases + + +def _norm_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("h3_rmsnorm", device="cuda") + weight_source = "pinned_checkpoint" if h3_weights_dir() is not None else "synthetic" + names = [ + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", + ] + weights = { + name: weight.bfloat16() + for name, weight in zip(names, h3_params(names, [(5376,)] * len(names)), strict=True) + } + gen = torch.Generator(device="cuda").manual_seed(0) + plain = {} + for name in names: + x = torch.randn(2, 777, 5376, device="cuda", generator=gen).bfloat16() + ref = torch.nn.functional.rms_norm(x, (5376,), weights[name], 1e-5) + plain[name] = bool(torch.equal(op(x, weights[name]), ref)) + x, weight, shift, scale, index = _norm_inputs(4097, seed=1) + modulated = bool( + torch.equal( + op.forward_modulated(x, weight, shift, scale, index), + provider_norm_modulate(x, weight, shift, scale, index), + ) + ) + invariant = bool( + torch.equal( + op.forward_modulated(x[:, 100:140], weight, shift, scale, index[100:140])[0], + op.forward_modulated(x, weight, shift, scale, index)[0, 100:140], + ) + ) + + # Backward against FP64, for the CUDA op and the diffusers expression. + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, weight, shift, scale)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden(x_, w_, sh_, sc_): + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * w_ + return n * (1 + sc_.index_select(0, index)) + sh_.index_select(0, index) + + ref = grads(golden, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op.forward_modulated(*t, index)), + ("provider", lambda *t: provider_norm_modulate(t[0], t[1], t[2], t[3], index)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "dweight", "dshift", "dscale"), first, ref, strict=True) + }, + } + return { + "weight_source": weight_source, + "plain_bitwise_vs_nn_rmsnorm": plain, + "modulated_bitwise_vs_diffusers": modulated, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# adaln_gate_residual (gate_msa after attention) +# --------------------------------------------------------------------------- # + + +def _gate_inputs(seq: int, seed: int = 0, dtype=torch.bfloat16): + g = torch.Generator(device="cuda").manual_seed(seed) + table = (torch.randn(9, 6 * 5376, device="cuda", generator=g) * 0.5).to(dtype) + residual = torch.randn(1, seq, 5376, device="cuda", generator=g).to(dtype) + y = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 3).to(dtype) + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate_view(table): + return table.view(table.shape[0], 6, -1)[:, 2] + + +def _gate_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("adaln_gate_residual", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + table, residual, y, index = _gate_inputs(seq) + gate = _gate_view(table) + grad = torch.randn_like(residual) + + def prepare_backward(fn, residual=residual, y=y, table=table, grad=grad): + leaves = [t.detach().requires_grad_(True) for t in (residual, y, table)] + return fn(leaves[0], leaves[1], _gate_view(leaves[2])), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "adaln_gate_residual", + "case": f"S={seq}", + "backend": type(op).__name__, + "bytes": 3 * seq * 5376 * 2, # read residual and y, write out + "candidate": lambda r=residual, y=y, g=gate, i=index: op.forward( + r, y, g, i, check_range=False + ), + "provider": lambda r=residual, y=y, g=gate, i=index: provider_gate_residual( + r, g, i, y + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: op.forward(r, y_, g, i, check_range=False) + ), + "provider_backward_setup": lambda b=prepare_backward, i=index: b( + lambda r, y_, g: provider_gate_residual(r, g, i, y_) + ), + } + ) + return cases + + +def _gate_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("adaln_gate_residual", device="cuda") + forward_bitwise = {} + for dtype in (torch.bfloat16, torch.float16, torch.float32): + table, residual, y, index = _gate_inputs(777, seed=1, dtype=dtype) + gate = _gate_view(table) + forward_bitwise[str(dtype).removeprefix("torch.")] = bool( + torch.equal( + op(residual, y, gate, index), provider_gate_residual(residual, gate, index, y) + ) + ) + table, residual, y, index = _gate_inputs(4097, seed=2) + invariant = bool( + torch.equal( + op(residual[:, 10:60], y[:, 10:60], _gate_view(table), index[10:60])[0], + op(residual, y, _gate_view(table), index)[0, 10:60], + ) + ) + grad = torch.randn( + residual.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(residual.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate_view(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ref = grads(lambda r, y_, g: provider_gate_residual(r, g, index, y_), torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda r, y_, g: op(r, y_, g, index)), + ("provider", lambda r, y_, g: provider_gate_residual(r, g, index, y_)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all( + torch.equal(a, b) for a, b in zip(first, second, strict=True) + ), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip( + ("d_residual", "d_sublayer", "d_gate"), first, ref, strict=True + ) + }, + } + return { + "forward_bitwise_vs_diffusers": forward_bitwise, + "rows_batch_invariant": invariant, + "backward": backward, + } + + +# --------------------------------------------------------------------------- # +# final_adaln_out (norm_out) +# --------------------------------------------------------------------------- # + +FINAL_NAMES = ["norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias"] + + +def _final_inputs(seq: int, seed: int = 0): + norm_weight, weight, bias = h3_params(FINAL_NAMES, [(5376,), (10752, 2688), (10752,)]) + g = torch.Generator(device="cuda").manual_seed(seed) + x = (torch.randn(1, seq, 5376, device="cuda", generator=g) * 2).bfloat16() + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, _ = h3_packed_layout(seq, 3, seed=seed) + return x, norm_weight.bfloat16(), temb, weight.bfloat16(), bias.bfloat16(), ti + + +def _final_perf(registry: KernelRegistry) -> list[dict[str, Any]]: + op = registry.get_op("final_adaln_out", device="cuda") + cases = [] + for seq in NORM_SEQ_LENS: + x, nw, temb, w, b, ti = _final_inputs(seq) + grad = torch.randn_like(x) + + def prepare_backward(fn, tensors=(x, nw, temb, w, b), grad=grad): + leaves = [t.detach().requires_grad_(True) for t in tensors] + return fn(*leaves), grad + + def backward(state): + output, grad = state + output.backward(grad) + + cases.append( + { + "op": "final_adaln_out", + "case": f"S={seq}", + "backend": type(op).__name__, + # read x + norm_out.linear, write out + "bytes": 2 * seq * 5376 * 2 + w.numel() * 2, + "candidate": lambda t=(x, nw, temb, w, b), i=ti: op(*t, i), + "provider": lambda t=(x, nw, temb, w, b), i=ti: provider_final_adaln_out( + t[0], t[1], t[2], t[3], t[4], i + ), + "candidate_backward": backward, + "provider_backward": backward, + "candidate_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: op(*t, i) + ), + "provider_backward_setup": lambda bw=prepare_backward, i=ti: bw( + lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], i) + ), + } + ) + return cases + + +def _final_accuracy(registry: KernelRegistry) -> dict[str, Any]: + op = registry.get_op("final_adaln_out", device="cuda") + golden_op = registry._get_or_create_backend( + registry._priority_map["cpu"]["final_adaln_out"][-1] + ) + if golden_op is None: + raise RuntimeError("required H3 golden backend is unavailable") + x, nw, temb, w, b, ti = _final_inputs(4097, seed=1) + ours = op(x, nw, temb, w, b, ti) + theirs = provider_final_adaln_out(x, nw, temb, w, b, ti) + golden = golden_op.forward_fp32(x, nw, temb, w, b, ti) + invariant = bool( + torch.equal(op(x[:, 100:160], nw, temb, w, b, ti[100:160])[0], ours[0, 100:160]) + ) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (x, nw, temb, w, b)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + def golden_backward(x_, nw_, t_, w_, b_): + # Keep the declared BF16 boundaries with FP64 leaves and identity VJPs. + act = t_ * torch.sigmoid(t_) + # The declared cast rounds the FP32 SiLU, as the model computes it. + declared = torch.nn.functional.silu(t_.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = torch.nn.functional.linear(act, w_, b_) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x_ * torch.rsqrt(x_.square().mean(-1, keepdim=True) + 1e-5) * nw_ + return n * (1.0 + scale.index_select(0, ti)) + shift.index_select(0, ti) + + ref = grads(golden_backward, torch.float64) + backward = {} + for name, fn in ( + ("cuda", lambda *t: op(*t, ti)), + ("provider", lambda *t: provider_final_adaln_out(t[0], t[1], t[2], t[3], t[4], ti)), + ): + first, second = grads(fn), grads(fn) + backward[name] = { + "repeat_bitwise_equal": all(torch.equal(a, c) for a, c in zip(first, second)), + "rel_error": { + key: float((g.double() - r).abs().max() / r.abs().max()) + for key, g, r in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), first, ref) + }, + } + return { + "equal_to_diffusers_fraction": float((ours == theirs).float().mean()), + "max_abs_vs_golden": float((ours.float() - golden).abs().max()), + "provider_max_abs_vs_golden": float((theirs.float() - golden).abs().max()), + "rows_batch_invariant": invariant, + "backward": backward, + } + + +PERF_CASES: dict[str, Callable[[KernelRegistry], list[dict[str, Any]]]] = { + "timestep_sinusoid_h3": _sinusoid_perf, + "timestep_mlp_fp32": _mlp_perf, + "adaln_projection_3mod": _projection_perf, + "adaln_row_gather": _gather_perf, + "h3_rmsnorm": _norm_perf, + "adaln_gate_residual": _gate_perf, + "final_adaln_out": _final_perf, +} +ACCURACY: dict[str, Callable[[KernelRegistry], dict[str, Any]]] = { + "timestep_sinusoid_h3": _sinusoid_accuracy, + "timestep_mlp_fp32": _mlp_accuracy, + "adaln_projection_3mod": _projection_accuracy, + "adaln_row_gather": _gather_accuracy, + "h3_rmsnorm": _norm_accuracy, + "adaln_gate_residual": _gate_accuracy, + "final_adaln_out": _final_accuracy, +} diff --git a/rl_engine/validation/models/h3_weights.py b/rl_engine/validation/models/h3_weights.py new file mode 100644 index 000000000..f587f2726 --- /dev/null +++ b/rl_engine/validation/models/h3_weights.py @@ -0,0 +1,108 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Pinned MiniMax-H3 conditioning weights for the RFC #420 tests. + +Only the tensors listed in ``h3_manifest.json`` are used. They all live in +shard 1 of 14, so ``tools/weights/prepare_h3_weights.py`` downloads that shard, +checks it against the manifest and extracts the tensors into +``$RL_KERNEL_H3_WEIGHTS/h3_conditioning.safetensors``. With ``--block`` it also +extracts block 0's attention and FFN weights (manifest ``block_tensors``, +about 1.1 GB) into ``h3_block0.safetensors`` for ``ws1_one_h3_block``. +""" + +from __future__ import annotations + +import hashlib +import json +import os +from pathlib import Path +from typing import Any + +import torch + +MANIFEST_PATH = Path(__file__).with_name("h3_manifest.json") +WEIGHTS_ENV = "RL_KERNEL_H3_WEIGHTS" +EXTRACTED_FILE = "h3_conditioning.safetensors" +BLOCK_FILE = "h3_block0.safetensors" + + +def load_h3_manifest() -> dict[str, Any]: + """Read the bundled checkpoint identity and pinned tensor contracts.""" + + return json.loads(MANIFEST_PATH.read_text()) + + +def h3_weights_dir() -> Path | None: + """Expand the configured extraction directory, or return None when unset or blank.""" + + value = os.environ.get(WEIGHTS_ENV, "").strip() + return Path(value).expanduser() if value else None + + +def sha256_file(path: Path, chunk: int = 1 << 24) -> str: + """Return a file SHA256 while reading bounded chunks.""" + + digest = hashlib.sha256() + with path.open("rb") as handle: + while block := handle.read(chunk): + digest.update(block) + return digest.hexdigest() + + +def sha256_tensor(tensor: torch.Tensor) -> str: + """Hash a CPU tensor's contiguous bytes without safetensors metadata.""" + + tensor_bytes = tensor.contiguous().view(torch.uint8).numpy() + return hashlib.sha256(memoryview(tensor_bytes)).hexdigest() + + +def _load_pinned( + filename: str, section: str, device: torch.device | str, names: list[str] | None +) -> dict[str, torch.Tensor]: + """Load one extracted file and check every tensor of a manifest section against it.""" + + from safetensors.torch import load_file + + root = h3_weights_dir() + if root is None: + raise FileNotFoundError(f"set {WEIGHTS_ENV} to the prepare_h3_weights.py output dir") + path = root / filename + if not path.is_file(): + flag = " --block" if section == "block_tensors" else "" + raise FileNotFoundError(f"{path} missing; run tools/weights/prepare_h3_weights.py{flag}") + specs = load_h3_manifest()[section] + tensors = load_file(str(path), device="cpu") + wanted = names if names is not None else list(specs) + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise ValueError(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise ValueError(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise ValueError(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + out: dict[str, torch.Tensor] = {} + for name in wanted: + if name not in specs: + raise KeyError(name) + out[name] = tensors[name].to(device=device) + return out + + +def load_h3_conditioning_weights( + device: torch.device | str = "cpu", names: list[str] | None = None +) -> dict[str, torch.Tensor]: + """Load the extracted pinned tensors; raise if they are missing or off-manifest.""" + + return _load_pinned(EXTRACTED_FILE, "tensors", device, names) + + +def load_h3_block_weights(device: torch.device | str = "cpu") -> dict[str, torch.Tensor]: + """Conditioning tensors plus block 0's attention/FFN weights, every one sha-checked.""" + + return { + **load_h3_conditioning_weights(device), + **_load_pinned(BLOCK_FILE, "block_tensors", device, None), + } diff --git a/rl_engine/validation/models/h3_ws2.py b/rl_engine/validation/models/h3_ws2.py new file mode 100644 index 000000000..9ec7b22c3 --- /dev/null +++ b/rl_engine/validation/models/h3_ws2.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS1 vs WS2 comparisons for the H3 conditioning path (tests and evidence). + +The per-rank helpers run inside real multi-process groups (tests and +``tools/validation/models/h3_ws2_evidence.py``); the WS1 side runs on one GPU. +""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def ws1_projection(temb, weight, bias, grad) -> dict[str, torch.Tensor]: + """The single-GPU ``adaln_projection_3mod`` table and its three gradients.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + table = H3AdaLNProjectionCudaOp().forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight": leaves[1].grad, + "d_bias": leaves[2].grad, + } + + +def tp_projection_rank(collective, temb, weight, bias, grad) -> dict[str, Any]: + """One TP rank: shard the full weight, run the TP op, return its view of the results.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + op = H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]) + w_shard, b_shard = shard_adaln_projection(weight, bias, collective.world_size, collective.rank) + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, w_shard, b_shard)] + table = op.forward_table(*leaves) + table.backward(grad) + return { + "table": table.detach(), + "d_temb": leaves[0].grad, + "d_weight_shard": leaves[1].grad, + "d_bias_shard": leaves[2].grad, + "readback": op.readback(), + } + + +def tp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Byte equality of every rank's replicated outputs and of the concatenated shards.""" + + return { + "table": all(torch.equal(r["table"], ws1["table"]) for r in ranks), + "d_temb": all(torch.equal(r["d_temb"], ws1["d_temb"]) for r in ranks), + "d_weight": torch.equal(torch.cat([r["d_weight_shard"] for r in ranks]), ws1["d_weight"]), + "d_bias": torch.equal(torch.cat([r["d_bias_shard"] for r in ranks]), ws1["d_bias"]), + } + + +def tp_projection_sweep(collective, temb, weight, bias, grad, iters: int = 100) -> dict[str, Any]: + """Evidence for one rank: per-T byte equality against WS1 on this GPU, timing, readback. + + ``temb``/``grad`` hold the largest T; every T in ``1..T`` uses their first + rows. Each rank compares its own replicated outputs and its own shard of + ``dW``/``db`` with the WS1 result computed here, so nothing large is shipped. + """ + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + adaln_column_shard, + ) + from rl_engine.validation.models.h3_report import time_us + + shard = adaln_column_shard(weight.shape[0], collective.world_size, collective.rank) + cols = slice(shard.begin, shard.end) + tp_op, ws1_op = ( + H3TPAdaLNProjectionCudaOp(collective, weight.shape[0]), + H3AdaLNProjectionCudaOp(), + ) + equality, perf = [], [] + for num_t in range(1, temb.shape[0] + 1): + t, g = temb[:num_t], grad[:num_t] + ws1 = ws1_projection(t, weight, bias, g) + mine = tp_projection_rank(collective, t, weight, bias, g) + equality.append( + { + "num_timesteps": num_t, + "table": torch.equal(mine["table"], ws1["table"]), + "d_temb": torch.equal(mine["d_temb"], ws1["d_temb"]), + "d_weight_shard": torch.equal(mine["d_weight_shard"], ws1["d_weight"][cols]), + "d_bias_shard": torch.equal(mine["d_bias_shard"], ws1["d_bias"][cols]), + } + ) + leaves = [x.detach().clone().requires_grad_(True) for x in (t, weight[cols], bias[cols])] + full = [x.detach().clone().requires_grad_(True) for x in (t, weight, bias)] + + def step(op, args): + for x in args: + x.grad = None + op.forward_table(*args).backward(g) + + perf.append( + { + "num_timesteps": num_t, + "tp_forward_us": time_us( + lambda: tp_op.forward_table(t, weight[cols], bias[cols]), 30, iters + ), + "tp_forward_backward_us": time_us(lambda: step(tp_op, leaves), 10, iters), + "ws1_forward_us": time_us(lambda: ws1_op.forward_table(t, weight, bias), 30, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_op, full), 10, iters), + } + ) + return {"equality": equality, "perf": perf, "readback": tp_op.readback()} + + +def _region(gate_fn, norm_fn, residual, y, weight, table): + """One block's MLP-side SP region: ``norm2(residual + gate_msa * y)`` with modulation.""" + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = table.chunk(6, dim=-1) + hidden = gate_fn(residual, y, gate_msa) + return norm_fn(hidden, weight, shift_mlp, scale_mlp) + + +def ws1_sp_region(residual, y, weight, table, index, grad) -> dict[str, torch.Tensor]: + """The SP region on one GPU with the WS1 ops (the reference for ``sp_norm_adaln``).""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import H3GateResidualCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + leaves = [t.detach().clone().requires_grad_(True) for t in (residual, y, weight, table)] + out = _region( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad) + names = ("d_residual", "d_y", "d_weight", "d_table") + return {"out": out.detach(), **{n: leaf.grad for n, leaf in zip(names, leaves)}} + + +def sp_region_rank(collective, residual, y, weight, table, index, grad) -> dict[str, Any]: + """One SP rank: its rows of the region and its (replicated) parameter gradients.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + op = H3SPNormAdaLNCudaOp(collective, seq_len=residual.shape[1], batch=residual.shape[0]) + rows = slice(op.layout.lo, op.layout.hi) + local = [t[:, rows].contiguous() for t in (residual, y)] + leaves = [t.detach().clone().requires_grad_(True) for t in (*local, weight, table)] + out = _region( + lambda r, yy, g: op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: op.norm_modulated(h, w, sh, sc, index), + *leaves, + ) + out.backward(grad[:, rows].contiguous()) + names = ("d_residual", "d_y", "d_weight", "d_table") + return { + "out": out.detach(), + **{n: leaf.grad for n, leaf in zip(names, leaves)}, + "readback": op.readback(), + } + + +def sp_matches_ws1(ws1: dict[str, torch.Tensor], ranks: list[dict[str, Any]]) -> dict[str, bool]: + """Rows: each rank's slice equals WS1's; parameters: every rank's copy equals WS1's.""" + + def rows(key): + return all( + torch.equal(r[key], ws1[key][:, slice(*r["readback"]["positions"])]) for r in ranks + ) + + return { + "out": rows("out"), + "d_residual": rows("d_residual"), + "d_y": rows("d_y"), + "d_weight": all(torch.equal(r["d_weight"], ws1["d_weight"]) for r in ranks), + "d_table": all(torch.equal(r["d_table"], ws1["d_table"]) for r in ranks), + } + + +def sp_case(batch, seq, hidden=5376, num_t=3, *, layout="block", dtype=torch.bfloat16, seed=0): + """Inputs of the SP region; identical on every rank (CPU generator).""" + + from rl_engine.validation.models.h3_cases import h3_block_layout, h3_packed_layout + + g = torch.Generator(device="cpu").manual_seed(seed) + make = h3_block_layout if layout == "block" else h3_packed_layout + ti, tags = make(seq, num_t, seed=seed) + return { + "residual": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + "y": (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda(), + "weight": (torch.rand(hidden, generator=g) + 0.5).to(dtype).cuda(), + "table": (torch.randn(3 * num_t, 6 * hidden, generator=g) * 0.5).to(dtype).cuda(), + "index": (ti * 3 + tags).cuda(), + "grad": torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda(), + } + + +def naive_sp_mismatch(case: dict[str, torch.Tensor], sp: int) -> dict[str, float]: + """Fraction of parameter-gradient elements where a naive SP backward differs from WS1. + + Naive: each rank runs the WS1 backward on its own rows, then the per-rank results + are summed in rank order (what a plain all-reduce of local gradients does). + """ + + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_row_layout + + ws1 = ws1_sp_region(**case) + parts = [] + for rank in range(sp): + lay = sp_row_layout(case["residual"].shape[1], case["residual"].shape[0], sp, rank) + part = slice(lay.lo, lay.hi) + local = {k: (v[:, part] if k in ("residual", "y", "grad") else v) for k, v in case.items()} + local["index"] = case["index"][part] + parts.append(ws1_sp_region(**local)) + out = {} + for key in ("d_weight", "d_table"): + total = parts[0][key].float() + for p in parts[1:]: + total = total + p[key].float() + out[key] = (total.to(ws1[key].dtype) != ws1[key]).float().mean().item() + return out + + +def sp_region_sweep(collective, cases: list[dict[str, Any]], iters: int = 20) -> dict[str, Any]: + """Evidence for one SP rank: byte equality, rows exchanged and time, per case.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import _segment_tiles + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + _dweight_tiles, + _Plan, + sp_row_layout, + ) + + results = [] + for spec in cases: + case = sp_case(**spec) + batch, seq = case["residual"].shape[:2] + ws1 = ws1_sp_region(**case) + mine = sp_region_rank(collective, **case) + equal = sp_matches_ws1(ws1, [mine]) + lay = sp_row_layout(seq, batch, collective.world_size, collective.rank) + tiles = _segment_tiles(case["index"].repeat(batch), case["table"].shape[0]) + plan = _Plan(lay, [_dweight_tiles(batch * seq, "cuda"), tuple(tiles[:3])], "cuda") + del ws1, mine + torch.cuda.empty_cache() + results.append( + { + **spec, + "equal": equal, + "local_rows": batch * lay.local_len, + "rows_sent": int(plan.send_local.numel()), + "rows_gathered_per_rank": plan.max_send, + **_region_times(collective, case, lay, iters), + } + ) + return {"cases": results} + + +def _region_times(collective, case, lay, iters) -> dict[str, float]: + """Forward + backward time of the region (leaves prepared once, grads reset per call).""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import H3GateResidualCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + from rl_engine.validation.models.h3_report import time_us + + index, rows = case["index"], slice(lay.lo, lay.hi) + sp_op = H3SPNormAdaLNCudaOp(collective, seq_len=lay.seq_len, batch=lay.batch) + gate_op, norm_op = H3GateResidualCudaOp(), H3RMSNormCudaOp() + + def leaves(local): + acts = [case[k][:, rows] if local else case[k] for k in ("residual", "y")] + return [ + t.detach().contiguous().requires_grad_(True) + for t in (*acts, case["weight"], case["table"]) + ] + + def step(fns, args, grad): + for leaf in args: + leaf.grad = None + _region(*fns, *args).backward(grad) + + sp_fns = ( + lambda r, yy, g: sp_op.gate_residual(r, yy, g, index), + lambda h, w, sh, sc: sp_op.norm_modulated(h, w, sh, sc, index), + ) + ws1_fns = ( + lambda r, yy, g: gate_op(r, yy, g, index), + lambda h, w, sh, sc: norm_op.forward_modulated(h, w, sh, sc, index), + ) + sp_args, ws1_args = leaves(True), leaves(False) + sp_grad = case["grad"][:, rows].contiguous() + return { + "sp_forward_backward_us": time_us(lambda: step(sp_fns, sp_args, sp_grad), 3, iters), + "ws1_forward_backward_us": time_us(lambda: step(ws1_fns, ws1_args, case["grad"]), 3, iters), + } + + +# --- real multi-process groups ------------------------------------------------- + + +def _bootstrap(rank, world, init_method, queue, inputs_path, target, kwargs): + import traceback + + import torch.distributed as dist + + from rl_engine.distributed.algorithms.collectives import DeterministicCollective + + try: + torch.cuda.set_device(rank) + dist.init_process_group( + "nccl", init_method=init_method, rank=rank, world_size=world, device_id=rank + ) + inputs = {k: v.cuda() for k, v in torch.load(inputs_path).items()} + collective = DeterministicCollective(device=rank, max_size_bytes=256 << 20) + out = target(collective, **inputs, **kwargs) + # Results go through a file: tensors sent over the queue are shared by + # file descriptor and vanish if this process exits before the parent reads. + torch.save( + {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in out.items()}, + inputs_path.with_name(f"rank{rank}.pt"), + ) + collective.close() + dist.destroy_process_group() + queue.put({"rank": rank}) + except Exception: # forwarded to the parent + queue.put({"rank": rank, "error": traceback.format_exc()}) + + +def run_world(world: int, target, inputs: dict[str, torch.Tensor], timeout: float = 900, **kwargs): + """Run ``target(collective, **inputs, **kwargs)`` on ranks ``0..world-1`` (one GPU each). + + ``target`` must be importable (it is pickled into spawned processes). Each + rank gets a NCCL group and a ``DeterministicCollective``; results come back + on the CPU, ordered by rank. Any rank's exception is re-raised here. + """ + + import tempfile + import time + from pathlib import Path + from queue import Empty + + import torch.multiprocessing as mp + + if torch.cuda.device_count() < world: + raise RuntimeError(f"needs {world} GPUs, found {torch.cuda.device_count()}") + ctx = mp.get_context("spawn") + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "inputs.pt" + torch.save({k: v.cpu() for k, v in inputs.items()}, path) + queue = ctx.Queue() + init = (Path(tmp) / "init").as_uri() + procs = [ + ctx.Process(target=_bootstrap, args=(r, world, init, queue, path, target, kwargs)) + for r in range(world) + ] + started = [] + deadline = time.monotonic() + timeout + status: dict[int, object] = {} + try: + for proc in procs: + proc.start() + started.append(proc) + while len(status) < world or any(proc.is_alive() for proc in procs): + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError(f"rank results timed out after {timeout}s") + try: + result = queue.get(timeout=min(0.1, remaining)) + except Empty: + exited = [ + (rank, proc.exitcode) + for rank, proc in enumerate(procs) + if proc.exitcode is not None and (proc.exitcode != 0 or rank not in status) + ] + if not exited: + continue + # A worker may have flushed its final message between the + # timed read and the exit-code check. + try: + result = queue.get_nowait() + except Empty: + rank, code = exited[0] + raise RuntimeError( + f"rank {rank} exited with code {code} without completing" + ) from None + if "error" in result: + raise RuntimeError(f"rank {result['rank']} failed:\n{result['error']}") + status[result["rank"]] = result + for rank, proc in enumerate(procs): + proc.join() + if proc.exitcode != 0: + raise RuntimeError(f"rank {rank} exited with code {proc.exitcode}") + return [{"rank": r, **torch.load(path.with_name(f"rank{r}.pt"))} for r in range(world)] + finally: + for proc in started: + if proc.is_alive(): + proc.terminate() + for proc in started: + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join() + queue.close() + queue.join_thread() diff --git a/rl_engine/validation/operators/operator_inputs.py b/rl_engine/validation/operators/operator_inputs.py index aaebc279f..15d7adc91 100644 --- a/rl_engine/validation/operators/operator_inputs.py +++ b/rl_engine/validation/operators/operator_inputs.py @@ -16,6 +16,10 @@ DEFAULT_VOCAB = 151936 DEFAULT_ROPE_THETA = 1.0e6 DEFAULT_RMS_EPS = 1.0e-6 +H3_FREQ_DIM = 256 +H3_TIME_HIDDEN = 5376 +H3_TIME_EMBED = 2688 +H3_HIDDEN = 5376 def make_operator_inputs( @@ -24,6 +28,7 @@ def make_operator_inputs( dtype: torch.dtype, device: torch.device, ) -> dict[str, Any]: + """Build operator keyword inputs on the requested device from CLI shape and seed options.""" builders = { "rms_norm": _make_rms_norm_inputs, "qk_norm": _make_qk_norm_inputs, @@ -44,6 +49,13 @@ def make_operator_inputs( "embedding": _make_embedding_inputs, "lm_head": _make_lm_head_inputs, "kv_cache_attention": _make_kv_cache_attention_inputs, + "timestep_sinusoid_h3": _make_timestep_sinusoid_h3_inputs, + "timestep_mlp_fp32": _make_timestep_mlp_fp32_inputs, + "adaln_projection_3mod": _make_adaln_projection_3mod_inputs, + "adaln_row_gather": _make_adaln_row_gather_inputs, + "h3_rmsnorm": _make_h3_rmsnorm_inputs, + "adaln_gate_residual": _make_adaln_gate_residual_inputs, + "final_adaln_out": _make_final_adaln_out_inputs, } try: return builders[op_name](args, dtype, device) @@ -52,6 +64,7 @@ def make_operator_inputs( def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: + """Return the operator's dimension label for benchmark and evidence reports.""" if op_name == "video_patch_pack_unpack": shape = "x".join(str(dim) for dim in _video_patch_shape(args)) return f"{_arg_str(args, 'patch_direction', 'pack')}-{shape}" @@ -78,6 +91,16 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "embedding": f"{batch}x{seq}x{vocab}x{_normalized_dim(args)}", "lm_head": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "kv_cache_attention": f"{batch}x{DEFAULT_N_HEADS}x1x{seq + 1}x{DEFAULT_HEAD_DIM}", + "timestep_sinusoid_h3": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}", + "timestep_mlp_fp32": f"{_h3_num_timesteps(args)}x{H3_FREQ_DIM}x{H3_TIME_HIDDEN}" + f"x{H3_TIME_EMBED}", + "adaln_projection_3mod": f"{_h3_num_timesteps(args)}x{H3_TIME_EMBED}" + f"x{6 * 3 * _h3_hidden(args)}", + "adaln_row_gather": f"{3 * _h3_num_timesteps(args)}x{6 * _h3_hidden(args)}" + f"->{batch * seq}", + "h3_rmsnorm": f"{batch}x{seq}x{_h3_hidden(args)}", + "adaln_gate_residual": f"{batch}x{seq}x{_h3_hidden(args)}", + "final_adaln_out": f"{batch}x{seq}x{_h3_hidden(args)}", } try: return names[op_name] @@ -370,6 +393,151 @@ def _make_kv_cache_attention_inputs( } +def _h3_num_timesteps(args: argparse.Namespace) -> int: + """Read the packed H3 timestep count, falling back to the CLI batch size.""" + # H3 packs a handful of distinct timesteps; reuse --batch as their count. + return _arg_int(args, "num_timesteps", _arg_int(args, "batch", 2)) + + +def _h3_timesteps( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> torch.Tensor: + """Build seeded timesteps in [0, 1] with endpoint cases in the requested storage dtype.""" + num = _h3_num_timesteps(args) + mode = _arg_str(args, "input_mode", "random") + if mode == "constant": + t = torch.full((num,), 0.5, device=device) + else: + t = torch.rand((num,), generator=_generator(args, device, offset=7), device=device) + t[0] = 0.0 + if num > 1: + t[-1] = 1.0 + return t.to(dtype) + + +def _make_timestep_sinusoid_h3_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Supply a one-dimensional packed timestep tensor to the H3 sinusoid operator.""" + return {"timestep": _h3_timesteps(args, dtype, device)} + + +def _make_timestep_mlp_fp32_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build [T, 256] sinusoid features and seeded FP32 256→5376→2688 MLP parameters.""" + # The H3 time_embedder is declared FP32; ``dtype`` does not apply to it. + del dtype + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + features = NativeH3TimestepSinusoidOp().forward(_h3_timesteps(args, torch.float32, device)) + scale1, scale2 = H3_FREQ_DIM**-0.5, H3_TIME_HIDDEN**-0.5 + return { + "x": features, + "w1": _floating_tensor((H3_TIME_HIDDEN, H3_FREQ_DIM), args, torch.float32, device, 1) + * scale1, + "b1": _floating_tensor((H3_TIME_HIDDEN,), args, torch.float32, device, 2) * 0.1, + "w2": _floating_tensor((H3_TIME_EMBED, H3_TIME_HIDDEN), args, torch.float32, device, 3) + * scale2, + "b2": _floating_tensor((H3_TIME_EMBED,), args, torch.float32, device, 4) * 0.1, + } + + +def _h3_hidden(args: argparse.Namespace) -> int: + """Read the AdaLN channel width, defaulting to the checkpoint's 5376 channels.""" + return _arg_int(args, "normalized_dim", H3_HIDDEN) + + +def _make_adaln_projection_3mod_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + """Build FP32 [T, 2688] embeddings and 18H projection parameters in the chosen dtype.""" + # temb is FP32 by contract (the SiLU runs before the cast); ``dtype`` is the + # projection's weight dtype, BF16 in the checkpoint. + n_out = 6 * 3 * _h3_hidden(args) + temb = _floating_tensor( + (_h3_num_timesteps(args), H3_TIME_EMBED), args, torch.float32, device, 0 + ) + weight = _floating_tensor((n_out, H3_TIME_EMBED), args, torch.float32, device, 1) + bias = _floating_tensor((n_out,), args, torch.float32, device, 2) + return { + "temb": temb * 2.0, + "weight": (weight * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (bias * 0.1).to(dtype), + } + + +def _make_adaln_row_gather_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # T distinct timesteps (--batch) and a packed sequence of batch * seq rows + # mixing all three modalities. + batch, seq = _batch_seq(args) + num_timesteps = _h3_num_timesteps(args) + generator = _generator(args, device, offset=29) + packed = batch * seq + token_tags = torch.randint(0, 3, (packed,), generator=generator, device=device) + timestep_indices = torch.randint( + 0, num_timesteps, (packed,), generator=generator, device=device + ) + rows = _floating_tensor((3 * num_timesteps, 6 * _h3_hidden(args)), args, dtype, device, 0) + return {"rows": rows, "timestep_indices": timestep_indices, "token_tags": token_tags} + + +def _make_h3_rmsnorm_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + weight = _floating_tensor((hidden,), args, torch.float32, device, 1).abs() + 0.5 + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "weight": weight.to(dtype), + "eps": 1e-5, + } + + +def _make_adaln_gate_residual_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # gate rows: 3 modality rows for each of --num-timesteps (default --batch) timesteps. + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + rows = 3 * _h3_num_timesteps(args) + generator = _generator(args, device, offset=31) + return { + "residual": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "y": _floating_tensor((batch, seq, hidden), args, dtype, device, 1), + "gate": (_floating_tensor((rows, hidden), args, torch.float32, device, 2) * 0.5).to(dtype), + "index": torch.randint(0, rows, (seq,), generator=generator, device=device), + } + + +def _make_final_adaln_out_inputs( + args: argparse.Namespace, dtype: torch.dtype, device: torch.device +) -> dict[str, Any]: + # temb is FP32 by contract; norm_out.linear/norm/x use ``dtype`` (BF16 in H3). + batch, seq = _batch_seq(args) + hidden = _h3_hidden(args) + num_timesteps = _h3_num_timesteps(args) + # Realistic scales (norm weight in [0.5, 1], unit activations): with unscaled + # inputs the parameter gradients sum hundreds of rows into values ~1e3, and + # their cancelling elements cannot meet a per-element FP32 atol. + linear = _floating_tensor((2 * hidden, H3_TIME_EMBED), args, torch.float32, device, 3) + norm = 0.5 + 0.5 * torch.sigmoid(_floating_tensor((hidden,), args, torch.float32, device, 4)) + generator = _generator(args, device, offset=37) + return { + "x": _floating_tensor((batch, seq, hidden), args, dtype, device, 0), + "norm_weight": norm.to(dtype), + "temb": _floating_tensor((num_timesteps, H3_TIME_EMBED), args, torch.float32, device, 1), + "weight": (linear * H3_TIME_EMBED**-0.5).to(dtype), + "bias": (_floating_tensor((2 * hidden,), args, torch.float32, device, 2) * 0.1).to(dtype), + "timestep_indices": torch.randint( + 0, num_timesteps, (seq,), generator=generator, device=device + ), + } + + def _floating_tensor( shape: tuple[int, ...], args: argparse.Namespace, diff --git a/rl_engine/validation/operators/operator_specs.py b/rl_engine/validation/operators/operator_specs.py index 92995b938..d5e6a7c85 100644 --- a/rl_engine/validation/operators/operator_specs.py +++ b/rl_engine/validation/operators/operator_specs.py @@ -272,6 +272,110 @@ def _load_object(path: str) -> Any: }, grad_input_names=("logits",), ), + # MiniMax-H3 (RFC #420) conditioning path. The op is FP32 end to end, so + # its FP32 rows are the declared contract; other dtypes only change the + # timestep input's storage dtype. + "timestep_sinusoid_h3": OperatorSpec( + name="timestep_sinusoid_h3", + op_class="elementwise", + gold_path="rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.timestep_sinusoid.NativeH3TimestepSinusoidOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_sinusoid.H3TimestepSinusoidCudaOp" + ), + }, + grad_input_names=("timestep",), + ), + "timestep_mlp_fp32": OperatorSpec( + name="timestep_mlp_fp32", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.timestep_mlp.NativeH3TimestepMLPOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "timestep_mlp.H3TimestepMLPCudaOp" + ), + }, + grad_input_names=("x", "w1", "b1", "w2", "b2"), + ), + "adaln_projection_3mod": OperatorSpec( + name="adaln_projection_3mod", + op_class="reduction", + gold_path=("rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp"), + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ( + "rl_engine.reference.minimax_h3.adaln_projection.NativeH3AdaLNProjectionOp" + ), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_projection.H3AdaLNProjectionCudaOp" + ), + }, + grad_input_names=("temb", "weight", "bias"), + ), + # The gather's forward is a copy (asserted bitwise in tests/models/minimax_h3), but its VJP + # sums every packed position of a table row, so it is judged as a reduction. + "adaln_row_gather": OperatorSpec( + name="adaln_row_gather", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.adaln_row_gather.NativeH3AdaLNRowGatherOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "adaln_row_gather.H3AdaLNRowGatherCudaOp" + ), + }, + grad_input_names=("rows",), + ), + "h3_rmsnorm": OperatorSpec( + name="h3_rmsnorm", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.rmsnorm.NativeH3RMSNormOp", + "cuda": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm.H3RMSNormCudaOp", + }, + grad_input_names=("x", "weight"), + ), + "adaln_gate_residual": OperatorSpec( + name="adaln_gate_residual", + op_class="reduction", # forward is elementwise; the gate VJP sums positions + gold_path="rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.reference.minimax_h3.gate_residual.NativeH3GateResidualOp", + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "gate_residual.H3GateResidualCudaOp" + ), + }, + grad_input_names=("residual", "y", "gate"), + ), + "final_adaln_out": OperatorSpec( + name="final_adaln_out", + op_class="reduction", + gold_path="rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": ("rl_engine.reference.minimax_h3.final_adaln_out.NativeH3FinalAdaLNOutOp"), + "cuda": ( + "rl_engine.backends.cuda.model_specific.minimax_h3." + "final_adaln_out.H3FinalAdaLNOutCudaOp" + ), + }, + grad_input_names=("x", "norm_weight", "temb", "weight", "bias"), + ), "pack": OperatorSpec( name="pack", op_class="elementwise", diff --git a/tests/build/test_build_platform_collectives.py b/tests/build/test_build_platform_collectives.py index b813506b9..820347673 100644 --- a/tests/build/test_build_platform_collectives.py +++ b/tests/build/test_build_platform_collectives.py @@ -10,7 +10,12 @@ import torch from torch.utils import cpp_extension -from build_tools.extensions import _resolve_package_data_source, get_cmdclass +from build_tools.extensions import ( + H3_BITWISE_SOURCES, + H3_KERNELS_DEFINE, + _resolve_package_data_source, + get_cmdclass, +) def _load_extension_config(monkeypatch, *, hip: str | None) -> dict[str, Any]: @@ -83,3 +88,28 @@ def test_package_data_real_json_is_not_reinterpreted(tmp_path) -> None: assert _resolve_package_data_source(source, project_root=tmp_path) == source assert "build_py" in get_cmdclass() + + +def test_cuda_build_compiles_h3_kernels_without_fast_math(monkeypatch) -> None: + monkeypatch.delenv("KERNEL_ALIGN_USE_FAST_MATH", raising=False) + extension = _load_extension_config(monkeypatch, hip=None) + + assert set(H3_BITWISE_SOURCES) <= set(extension["sources"]) + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE in extension["extra_compile_args"]["nvcc"] + assert "--use_fast_math" not in extension["extra_compile_args"]["nvcc"] + + +def test_fast_math_build_excludes_bitwise_h3_kernels(monkeypatch) -> None: + # --use_fast_math swaps expf/sinf/cosf/rsqrtf and division for approximations + # and enables FTZ, which breaks the H3 kernels' bitwise parity promise; nvcc + # exposes no macro to reject it in-source, so the build drops the kernels + # and their bindings. + monkeypatch.setenv("KERNEL_ALIGN_USE_FAST_MATH", "1") + extension = _load_extension_config(monkeypatch, hip=None) + + assert "--use_fast_math" in extension["extra_compile_args"]["nvcc"] + assert not set(H3_BITWISE_SOURCES) & set(extension["sources"]) + assert not [s for s in extension["sources"] if s.startswith("csrc/cuda/h3/")] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["cxx"] + assert H3_KERNELS_DEFINE not in extension["extra_compile_args"]["nvcc"] diff --git a/tests/models/minimax_h3/conftest.py b/tests/models/minimax_h3/conftest.py new file mode 100644 index 000000000..f9c2b000f --- /dev/null +++ b/tests/models/minimax_h3/conftest.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Shared fixtures for the MiniMax-H3 (RFC #420) operator tests. + +Real-weight tests read the pinned tensors from ``$RL_KERNEL_H3_WEIGHTS`` +(written by ``tools/weights/prepare_h3_weights.py``) and skip when it is unset. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.extension import _C, _EXT_AVAILABLE +from rl_engine.validation.models.h3_weights import ( + WEIGHTS_ENV, + h3_weights_dir, + load_h3_block_weights, + load_h3_conditioning_weights, +) + + +@pytest.fixture(scope="session") +def h3_weights_cpu() -> dict[str, torch.Tensor]: + """Load pinned CPU weights once per session, skipping unavailable artifacts.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py") + try: + return load_h3_conditioning_weights("cpu") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(scope="session") +def h3_block_weights_cuda() -> dict[str, torch.Tensor]: + """Conditioning plus block-0 tensors on CUDA, skipping when they were not extracted.""" + + if h3_weights_dir() is None: + pytest.skip(f"{WEIGHTS_ENV} not set; run tools/weights/prepare_h3_weights.py --block") + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + try: + return load_h3_block_weights("cuda") + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +def h3_kernels_built() -> bool: + """False for KERNEL_ALIGN_USE_FAST_MATH=1 builds, which omit every H3 kernel.""" + return bool(_EXT_AVAILABLE and hasattr(_C, "h3_timestep_sinusoid_forward")) + + +@pytest.fixture +def require_h3_kernels() -> None: + if not h3_kernels_built(): + pytest.skip("rl_engine._C has no H3 kernels (not built, or a fast-math build)") diff --git a/tests/models/minimax_h3/test_h3_adaln_gate_residual.py b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py new file mode 100644 index 000000000..e4c303e10 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_gate_residual.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_gate_residual``: ``residual + gate[row] * sublayer_output``. + +* forward and ``d_sublayer`` are bitwise equal to diffusers' eager + expression in every dtype (the gate row is gathered in the kernel); +* ``d_gate`` is a deterministic FP32 segment sum, checked against FP64; +* rows are elementwise-independent; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.gate_residual import NativeH3GateResidualOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_gate_residual + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + h3_gate_residual_available, + ) + + if not h3_gate_residual_available(): + pytest.skip("rl_engine._C lacks h3_gate_residual_*") + return H3GateResidualCudaOp() + + +def _case(seq, hidden=5376, dtype=torch.bfloat16, seed=0, batch=2): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(9, 6 * hidden, generator=g) * 0.5).to(dtype).cuda() + residual = torch.randn(batch, seq, hidden, generator=g).to(dtype).cuda() + y = (torch.randn(batch, seq, hidden, generator=g) * 3).to(dtype).cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seed) + return table, residual, y, ti * 3 + tags + + +def _gate(table, chunk=2): + return table.view(table.shape[0], 6, -1)[:, chunk] # gate_msa (2) / gate_mlp (5) view + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3GateResidualOp() + res, y, gate = torch.zeros(4, 8), torch.zeros(4, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op(res, y, gate, torch.tensor([0, 1, 2, 3])) + with pytest.raises(ValueError): + op(res, y[:, :4], gate, torch.tensor([0, 1, 2, 0])) + with pytest.raises(TypeError): + op(res, y, gate.bfloat16(), torch.tensor([0, 1, 2, 0])) + with pytest.raises(ValueError): # S mismatch + op(res, y, gate, torch.tensor([0, 1])) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("seq, hidden", [(4097, 5376), (33, 12), (17, 13)]) + def test_forward_bitwise_equal_to_diffusers(self, dtype, seq, hidden): + table, residual, y, index = _case(seq, hidden, dtype, seed=seq) + for chunk in (2, 5): + ours = _cuda_op()(residual, y, _gate(table, chunk), index) + assert torch.equal( + ours, provider_gate_residual(residual, _gate(table, chunk), index, y) + ) + + def test_rows_are_batch_and_position_invariant(self): + table, residual, y, index = _case(500, seed=1) + full = _cuda_op()(residual, y, _gate(table), index) + part = _cuda_op()(residual[1:2, 100:150], y[1:2, 100:150], _gate(table), index[100:150]) + assert torch.equal(part[0], full[1, 100:150]) + + def test_backward(self): + table, residual, y, index = _case(4097, seed=2) + grad = torch.randn_like(residual) + + def grads(fn, dtype=None): + tensors = [t if dtype is None else t.to(dtype) for t in (residual, y, table)] + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(leaves[0], leaves[1], _gate(leaves[2])).backward( + grad if dtype is None else grad.to(dtype) + ) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + again = grads(lambda r, yy, gg: _cuda_op()(r, yy, gg, index)) + theirs = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy)) + golden = grads(lambda r, yy, gg: provider_gate_residual(r, gg, index, yy), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + assert torch.equal(ours[0], theirs[0]) # d_residual + assert torch.equal(ours[1], theirs[1]) # d_sublayer: exact product, one rounding + rel = ((ours[2].double() - golden[2]).abs().max() / golden[2].abs().max()).item() + assert rel < 5e-3 # only the BF16 rounding of the table gradient + assert torch.count_nonzero(ours[2].view(9, 6, -1)[:, [0, 1, 3, 4, 5]]) == 0 + + def test_rejects_cpu_and_bad_index(self): + table, residual, y, index = _case(10, hidden=16) + with pytest.raises(IndexError): + _cuda_op()(residual, y, _gate(table), index + 9) + with pytest.raises(ValueError): + _cuda_op()(residual.cpu(), y.cpu(), _gate(table).cpu(), index.cpu()) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_gate_residual", device="cuda") + assert type(op).__name__ == "H3GateResidualCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_modulation.py b/tests/models/minimax_h3/test_h3_adaln_modulation.py new file mode 100644 index 000000000..066445fc0 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_modulation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fused AdaLN modulation (projection + row gather) with an FP32 table gradient. + +The forward must be bitwise equal to the two RFC #420 ops run separately; the +backward must be deterministic and must not re-round the table gradient to +BF16 (so ``d_temb`` reaches FP64-golden accuracy). +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models.h3_cases import h3_packed_layout + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _ops(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_modulation import ( + H3AdaLNModulationCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks the H3 kernels") + return H3AdaLNModulationCudaOp(), H3AdaLNProjectionCudaOp(), H3AdaLNRowGatherCudaOp() + + +def _params(hidden, dim, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = 18 * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).bfloat16().cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).bfloat16().cuda() + temb = (torch.randn(3, dim, generator=g) * 2).cuda() + return temb, weight, bias + + +def _grads(seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + +def _leaf_grads(fn, temb, weight, bias, grads): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(fn(*leaves)), grads) + return [leaf.grad for leaf in leaves] + + +@requires_cuda +class TestFusedModulation: + def test_forward_bitwise_equal_to_separate_ops(self): + fused, proj, gather = _ops() + temb, weight, bias = _params(64, 2688) + ti, tags = h3_packed_layout(1000, 3) + ours = fused(temb, weight, bias, ti, tags) + separate = gather.gather_chunks(proj(temb, weight, bias), ti, tags) + assert all(torch.equal(a, b) for a, b in zip(ours, separate)) + + def test_backward_keeps_table_gradient_fp32(self, h3_weights_cpu): + fused, proj, gather = _ops() + weight = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.weight"].cuda() + bias = h3_weights_cpu["transformer_blocks.0.adaln_proj.linear.bias"].cuda() + g = torch.Generator(device="cuda").manual_seed(0) + temb = torch.randn(3, 2688, device="cuda", generator=g) * 2 + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = _grads(4097, 5376, seed=1) + ours = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + separate = _leaf_grads( + lambda t, w, b: gather.gather_chunks(proj(t, w, b), ti, tags), + temb, + weight, + bias, + grads, + ) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column (about 1 temb draw in 30 dropped below 0.9999 that way). + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # identity-VJP cast + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376) + index = ti * 3 + tags + torch.autograd.backward( + [chunk.index_select(0, index) for chunk in table.chunk(6, dim=-1)], + [g.double() for g in grads], + ) + + def rel(grad, golden): + return ((grad.double() - golden).abs().max() / golden.abs().max()).item() + + # d_temb: no BF16 rounding anywhere on its path, so it reaches FP32 accuracy. + assert rel(ours[0], ref[0].grad) < 1e-5 + assert rel(ours[0], ref[0].grad) < rel(separate[0], ref[0].grad) / 100 + # dW, db: only their own final BF16 rounding remains. + for grad, golden in zip(ours[1:], (ref[1].grad, ref[2].grad)): + assert (grad == golden.to(torch.bfloat16)).float().mean() > 0.9999 + + def test_backward_repeat_bitwise(self): + fused, _, _ = _ops() + temb, weight, bias = _params(64, 2688, seed=3) + ti, tags = h3_packed_layout(3000, 3, seed=3) + grads = _grads(3000, 64, seed=3) + first = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + again = _leaf_grads(lambda t, w, b: fused(t, w, b, ti, tags), temb, weight, bias, grads) + assert all(torch.equal(a, b) for a, b in zip(first, again)) + + def test_rejects_bad_indices(self): + fused, _, _ = _ops() + temb, weight, bias = _params(8, 16) + ti, tags = h3_packed_layout(10, 3) + with pytest.raises(IndexError): + fused(temb, weight, bias, ti + 3, tags) + with pytest.raises(TypeError): + fused(temb.bfloat16(), weight, bias, ti, tags) diff --git a/tests/models/minimax_h3/test_h3_adaln_projection.py b/tests/models/minimax_h3/test_h3_adaln_projection.py new file mode 100644 index 000000000..5fb10a73f --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_projection.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_projection_3mod``: SiLU + 2688 -> 96768 projection, 6 x 3 table. + +* layout: the six outputs and the three modality rows per timestep land + exactly where diffusers puts them (RFC probe H4); +* mixed precision: SiLU in FP32, one cast to BF16; an early cast (probe H7) + is rejected at the API and is measurably different numerically; +* accuracy: forward and gradients against the declared-cast golden on the + pinned block-0 weights; +* invariance and determinism: a timestep's 18 modulation rows do not depend + on the other timesteps, and repeated runs are bitwise equal. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3 import H3_ADALN_CHUNKS, H3_MODALITY_NUM +from rl_engine.reference.minimax_h3.adaln_projection import NativeH3AdaLNProjectionOp +from rl_engine.validation.models.h3_provider import provider_adaln_modulation + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _cuda_op(): + """Construct the CUDA projection operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3AdaLNProjectionCudaOp() + + +def _temb(num, dim=2688, seed=0, device="cuda"): + """Generate seeded FP32 timestep embeddings on CPU and move them to the test device.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(num, dim, generator=g) * 2).to(device) + + +def _synthetic(hidden, dim, dtype=torch.bfloat16, seed=0, device="cuda"): + """Create seeded projection weights for six chunks and three modalities per channel.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + weight = (torch.randn(n_out, dim, generator=g) / dim**0.5).to(dtype) + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype) + return weight.to(device), bias.to(device) + + +def _flat(outputs): + """Flatten and concatenate all modulation chunks for elementwise comparisons.""" + + return torch.cat([out.reshape(-1) for out in outputs]) + + +class TestReference: + def test_provider_layout_matches_diffusers(self): + """Match the provider values and timestep-major modality layout on CPU.""" + + temb = _temb(2, dim=16, device="cpu") + weight, bias = _synthetic(8, 16, device="cpu") + ours = NativeH3AdaLNProjectionOp().forward(temb, weight, bias) + theirs = provider_adaln_modulation(temb, weight, bias, hidden_size=8) + assert len(ours) == 6 + for a, b in zip(ours, theirs, strict=True): + assert a.shape == (6, 8) and torch.equal(a, b) + + def test_rejects_bf16_temb(self): + # RFC probe H7: casting temb before the SiLU. + """Reject an early BF16 cast so SiLU receives the declared FP32 embeddings.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(TypeError, match="float32"): + NativeH3AdaLNProjectionOp().forward(_temb(2, 16, device="cpu").bfloat16(), weight, bias) + + def test_rejects_bad_shapes(self): + """Reject invalid projection dimensions, empty batches, and mismatched bias dtype.""" + + weight, bias = _synthetic(8, 16, device="cpu") + op = NativeH3AdaLNProjectionOp() + with pytest.raises(ValueError): # not a multiple of 6 * 3 rows + op.forward(_temb(2, 16, device="cpu"), weight[:-1], bias[:-1]) + with pytest.raises(ValueError): + op.forward(_temb(2, 15, device="cpu"), weight, bias) + with pytest.raises(ValueError): + op.forward(_temb(0, 16, device="cpu"), weight, bias) + with pytest.raises(TypeError): + op.forward(_temb(2, 16, device="cpu"), weight, bias.float()) + + +@requires_cuda +class TestLayout: + def test_channel_identity_lands_in_its_chunk_and_modality_row(self): + """Probe H4: bias o = m*6H + c*H + h tagged (m, c); zero weight passes it through.""" + + hidden, dim, num = 16, 32, 3 + n_out = H3_ADALN_CHUNKS * H3_MODALITY_NUM * hidden + o = torch.arange(n_out) + modality, chunk = o // (H3_ADALN_CHUNKS * hidden), (o // hidden) % H3_ADALN_CHUNKS + bias = (modality * 10 + chunk).to(torch.bfloat16).cuda() + weight = torch.zeros(n_out, dim, dtype=torch.bfloat16, device="cuda") + outputs = _cuda_op()(_temb(num, dim), weight, bias) + assert len(outputs) == 6 + for c, out in enumerate(outputs): + assert out.shape == (num * 3, hidden) + for t in range(num): + for m in range(3): + expected = torch.full((hidden,), float(m * 10 + c), device="cuda") + assert torch.equal(out[t * 3 + m].float(), expected), (c, t, m) + + def test_outputs_are_views_of_one_table_like_diffusers(self): + """Require all six modulation chunks to share storage with the expected offsets.""" + + weight, bias = _synthetic(16, 32) + outputs = _cuda_op()(_temb(2, 32), weight, bias) + base = outputs[0].untyped_storage().data_ptr() + assert all(out.untyped_storage().data_ptr() == base for out in outputs) + assert outputs[1].data_ptr() - outputs[0].data_ptr() == 16 * outputs[0].element_size() + + +@pytest.fixture(scope="module") +def block0(h3_weights_cpu): + """Move the pinned block-zero projection weight and bias to CUDA for this module.""" + + return h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4]) + def test_forward_is_correctly_rounded_declared_golden(self, block0, num): + """Check checkpoint accuracy and near-exact BF16 rounding of the declared golden.""" + + weight, bias = block0 + temb = _temb(num, seed=num) + outputs = _cuda_op()(temb, weight, bias) + golden = NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias) + for out, gold in zip(outputs, golden, strict=True): + assert out.dtype == torch.bfloat16 and out.shape == (3 * num, 5376) + # tolerance_contract.json forward_accuracy / reduction / bfloat16. + torch.testing.assert_close(out.float(), gold, atol=5e-2, rtol=2e-2) + # Beyond the contract: almost every element is the correctly rounded + # golden value. The rest come from the FP32 tensor-core accumulation over + # K = 2688 (ties, and a few ULPs where near-zero sums cancel), so exact + # equality is not reachable: over 300 temb draws per T the worst match was + # 0.999618 (T = 1), 0.999690, 0.999728 and 0.999731 (T = 4). + declared = _flat(golden).bfloat16() + assert (_flat(outputs) == declared).float().mean().item() >= 0.9995 + + def test_early_cast_is_numerically_detectable(self, block0): + """Probe H7: a BF16 cast before the SiLU moves about half the outputs.""" + + weight, bias = block0 + temb = _temb(3, seed=7) + op = NativeH3AdaLNProjectionOp() + ours = _flat(_cuda_op()(temb, weight, bias)) + declared = _flat(op.forward_fp32(temb, weight, bias)).bfloat16() + early = _flat(op.forward_fp32(temb.bfloat16().float(), weight, bias)).bfloat16() + assert (ours == declared).float().mean() >= 0.9995 + assert (ours == early).float().mean() < 0.8 + + def test_rows_are_batch_and_position_invariant(self, block0): + """Keep each timestep modulation row identical alone and under batch reordering.""" + + weight, bias = block0 + op = _cuda_op() + temb = _temb(5, seed=3) + full = op(temb, weight, bias) + for i in range(5): + single = op(temb[i : i + 1], weight, bias) + for s, f in zip(single, full, strict=True): + assert torch.equal(s, f[3 * i : 3 * i + 3]) + swapped = op(temb.flip(0), weight, bias) + for s, f in zip(swapped, full, strict=True): + assert torch.equal(s.view(5, 3, -1).flip(0), f.view(5, 3, -1)) + + def test_backward_against_golden(self, block0): + """Check embedding and parameter gradients against a high-precision declared-cast VJP.""" + + weight, bias = block0 + temb = _temb(2, seed=5) + g = torch.Generator(device="cuda").manual_seed(0) + grads = [torch.randn(6, 5376, device="cuda", generator=g).bfloat16() for _ in range(6)] + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + torch.autograd.backward(list(_cuda_op()(*leaves)), grads) + + ref = [t.detach().double().requires_grad_(True) for t in (temb, weight, bias)] + act = ref[0] * torch.sigmoid(ref[0]) + # The declared cast rounds the FP32 SiLU, as the kernel and diffusers do: an + # FP64 SiLU can round to the other BF16 neighbour near a tie, which shifts a + # whole dW column. + act_bf16 = F.silu(temb.float()).to(torch.bfloat16).double() + act = act + (act_bf16 - act).detach() # declared cast, identity VJP + table = F.linear(act, ref[1], ref[2]).view(-1, 6 * 5376).chunk(6, dim=-1) + torch.autograd.backward(list(table), [gr.double() for gr in grads]) + # tolerance_contract.json gradient_accuracy / reduction: FP32 temb grad + # and BF16 weight/bias grads (rounded once from an FP32 accumulation). + torch.testing.assert_close(leaves[0].grad.double(), ref[0].grad, atol=1e-4, rtol=1e-4) + for leaf, r in zip(leaves[1:], ref[1:], strict=True): + torch.testing.assert_close(leaf.grad.double(), r.grad, atol=1e-1, rtol=2e-2) + # With T = 2 each dW/db element sums two products, and the one BF16 cast + # of the FP32 sum is the correctly rounded value of every element + # (checked on 500 of 500 temb draws). + assert torch.equal(leaf.grad, r.grad.to(torch.bfloat16)) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) + @pytest.mark.parametrize("hidden, dim", [(24, 16), (40, 72), (8, 2688)]) + def test_other_sizes(self, dtype, hidden, dim): + """Match the projection golden at non-checkpoint dimensions in FP32 and BF16.""" + + weight, bias = _synthetic(hidden, dim, dtype=dtype, seed=hidden) + temb = _temb(3, dim, seed=hidden) + tol = dict(atol=5e-2, rtol=2e-2) if dtype == torch.bfloat16 else dict(atol=1e-4, rtol=1e-4) + for out, gold in zip( + _cuda_op()(temb, weight, bias), + NativeH3AdaLNProjectionOp().forward_fp32(temb, weight, bias), + strict=True, + ): + torch.testing.assert_close(out.float(), gold, **tol) + + def test_repeat_forward_backward_bitwise(self): + """Require repeated projection outputs and all input gradients to match bitwise.""" + + weight, bias = _synthetic(64, 2688, seed=9) + temb = _temb(3, seed=9) + grads = [torch.randn(9, 64, device="cuda").bfloat16() for _ in range(6)] + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (temb, weight, bias)] + outs = _cuda_op()(*leaves) + torch.autograd.backward(list(outs), grads) + runs.append([*[o.detach() for o in outs], *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_temb_grad_rows_are_batch_invariant(self): + """Keep embedding gradients bitwise equal in full batches and single-row calls.""" + + weight, bias = _synthetic(64, 2688, seed=11) + temb = _temb(4, seed=11) + grads = [torch.randn(12, 64, device="cuda").bfloat16() for _ in range(6)] + + def d_temb(rows): + """Compute selected embedding-row gradients with their matching modality gradients.""" + + leaf = temb[rows].detach().clone().requires_grad_(True) + outs = _cuda_op()(leaf, weight, bias) + torch.autograd.backward( + list(outs), [g.view(4, 3, -1)[rows].reshape(-1, 64) for g in grads] + ) + return leaf.grad + + full = d_temb(slice(0, 4)) + for i in range(4): + assert torch.equal(d_temb(slice(i, i + 1))[0], full[i]) + + @pytest.mark.parametrize("num", [8, 9, 17]) + def test_rows_invariant_across_tensor_core_row_tiles(self, num): + """The BF16 path computes rows in tiles of 8; a row's bits must not depend on its tile.""" + + weight, bias = _synthetic(64, 2688, seed=13) + temb = _temb(num, seed=num) + full = _cuda_op()(temb, weight, bias) + for i in (0, 7, num - 1): + single = _cuda_op()(temb[i : i + 1], weight, bias) + for s_out, f_out in zip(single, full, strict=True): + assert torch.equal(s_out, f_out[3 * i : 3 * i + 3]) + + def test_rejects_cpu(self): + """Reject CPU tensors at the CUDA projection boundary.""" + + weight, bias = _synthetic(8, 16, device="cpu") + with pytest.raises(ValueError): + _cuda_op()(_temb(2, 16, device="cpu"), weight, bias) + + def test_registry_dispatches_cuda(self): + """Resolve the projection registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_projection_3mod", device="cuda") + assert type(op).__name__ == "H3AdaLNProjectionCudaOp" diff --git a/tests/models/minimax_h3/test_h3_adaln_row_gather.py b/tests/models/minimax_h3/test_h3_adaln_row_gather.py new file mode 100644 index 000000000..0f9fec6c3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_adaln_row_gather.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``adaln_row_gather``: ``timestep_index * 3 + token_tag`` lookup of six tensors. + +* forward is a copy: bitwise equal to diffusers' six ``index_select`` calls, + for every packing, length and index dtype; +* semantic indices fail closed (out-of-range tags or timesteps, RFC probes + H2/H3), and flipping a tag or offsetting a timestep picks a different row; +* backward is a deterministic FP32 segmented sum: repeat-bitwise and + correctly rounded, unlike the atomic BF16 ``index_select`` backward; +* batch invariance: a position's output does not depend on which other + positions share the call, and a row's gradient does not depend on the + positions that reference other rows. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest +import torch + +from rl_engine.reference.minimax_h3.adaln_row_gather import NativeH3AdaLNRowGatherOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_adaln_row_gather + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + adaln_row_gather_available, + ) + + if not adaln_row_gather_available(): + pytest.skip("rl_engine._C lacks h3_adaln_row_gather_*") + return H3AdaLNRowGatherCudaOp() + + +def _rows(num_timesteps, hidden=5376, dtype=torch.bfloat16, seed=0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(3 * num_timesteps, 6 * hidden, generator=g).to(dtype).to(device) + + +class TestReference: + def test_matches_diffusers_index_select(self): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = h3_packed_layout(11, 2, device="cpu") + index = ti * 3 + tags + chunks = rows.chunk(6, dim=-1) + for ours, chunk in zip(NativeH3AdaLNRowGatherOp().forward(rows, ti, tags), chunks): + assert torch.equal(ours, chunk.index_select(0, index)) + + @pytest.mark.parametrize( + "mutate, error", + [ + (lambda ti, tags: (ti, tags.clone().fill_(3)), IndexError), # H2: no 4th modality + (lambda ti, tags: (ti, tags - 1), IndexError), + (lambda ti, tags: (ti + 2, tags), IndexError), # H3: timestep offset past T + (lambda ti, tags: (ti, tags[:-1]), ValueError), + (lambda ti, tags: (ti[:0], tags[:0]), ValueError), + (lambda ti, tags: (ti.float(), tags), TypeError), + (lambda ti, tags: (ti.int(), tags), TypeError), # mixed index dtypes + ], + ) + def test_rejects_invalid_indices(self, mutate, error): + rows = _rows(2, hidden=4, device="cpu") + ti, tags = mutate(*h3_packed_layout(9, 2, device="cpu")) + with pytest.raises(error): + NativeH3AdaLNRowGatherOp().forward(rows, ti, tags) + + def test_rejects_bad_rows(self): + ti, tags = h3_packed_layout(4, 1, device="cpu") + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(4, 24), ti, tags) # not 3 per t + with pytest.raises(ValueError): + NativeH3AdaLNRowGatherOp().forward(torch.zeros(3, 25), ti, tags) # not 6 * H + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("seq", [1, 2, 3, 257, 4097, 32768]) + @pytest.mark.parametrize("num_timesteps", [1, 3]) + def test_bitwise_equal_to_index_select(self, seq, num_timesteps): + rows = _rows(num_timesteps, seed=seq) + ti, tags = h3_packed_layout(seq, num_timesteps, seed=seq) + ours = _cuda_op()(rows, ti, tags) + ref = provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + assert len(ours) == 6 + for a, b in zip(ours, ref): + assert a.shape == (seq, 5376) and a.is_contiguous() and torch.equal(a, b) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [24, 5, 5376]) + def test_dtypes_and_unaligned_hidden(self, dtype, hidden): + rows = _rows(2, hidden=hidden, dtype=dtype) + ti, tags = h3_packed_layout(33, 2) + for a, b in zip( + _cuda_op()(rows, ti.int(), tags.int()), NativeH3AdaLNRowGatherOp()(rows, ti, tags) + ): + assert torch.equal(a, b) + + def test_strided_rows_view(self): + table = _rows(4, hidden=16) + rows = table[::2] # row stride 2 * 6H, still 3 rows per timestep after slicing 12 -> 6 + ti, tags = h3_packed_layout(20, 2) + for a, b in zip(_cuda_op()(rows, ti, tags), NativeH3AdaLNRowGatherOp()(rows, ti, tags)): + assert torch.equal(a, b) + + def test_tag_flip_and_timestep_offset_select_other_rows(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(30, 2) + base = _cuda_op()(rows, ti, tags) + flipped = _cuda_op()(rows, ti, (tags + 1) % 3) # H2 + shifted = _cuda_op()(rows, 1 - ti, tags) # H3 within range + assert not torch.equal(base[0], flipped[0]) + assert not torch.equal(base[0], shifted[0]) + # Every output row is exactly the addressed table row, nothing else. + index = ti * 3 + tags + for c, out in enumerate(base): + assert torch.equal(out, rows[index, c * 16 : (c + 1) * 16]) + + def test_rejects_out_of_range_on_device(self): + rows = _rows(1, hidden=16) + ti, tags = h3_packed_layout(8, 1) + with pytest.raises(IndexError): + _cuda_op()(rows, ti, tags + 3) + + @pytest.mark.parametrize("entrypoint", ["native", "unchecked_wrapper"]) + @pytest.mark.parametrize( + "timestep, tag, index_dtype", + [ + pytest.param(-1, 0, "int32", id="negative-timestep"), + pytest.param(2, 0, "int64", id="timestep-past-table"), + pytest.param(0, 3, "int32", id="tag-past-modality-valid-row"), + pytest.param(1, -1, "int64", id="negative-tag-valid-row"), + # Multiplication by 3 wraps this value to row 2 in signed int64. + pytest.param((2**64 + 2) // 3, 0, "int64", id="timestep-overflow-valid-row"), + pytest.param(0, 2**63 - 1, "int64", id="int64-max-tag"), + ], + ) + def test_native_bounds_assertions(self, entrypoint, timestep, tag, index_dtype): + """Reject semantic index violations without poisoning pytest's CUDA context.""" + + _cuda_op() + probe = textwrap.dedent( + """ + import sys + import torch + from rl_engine.backends.extension import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + entrypoint, timestep, tag, index_dtype = sys.argv[1:] + rows = torch.zeros((6, 6), dtype=torch.float32, device="cuda") + dtype = getattr(torch, index_dtype) + ti = torch.tensor([int(timestep)], dtype=dtype, device="cuda") + tags = torch.tensor([int(tag)], dtype=dtype, device="cuda") + if entrypoint == "native": + _C.h3_adaln_row_gather_forward(rows, ti, tags, 6, 3) + else: + H3AdaLNRowGatherCudaOp().forward(rows, ti, tags, check_range=False) + torch.cuda.synchronize() + """ + ) + completed = subprocess.run( + [sys.executable, "-c", probe, entrypoint, str(timestep), str(tag), index_dtype], + cwd=Path(__file__).resolve().parents[3], + env={**os.environ, "CUDA_LAUNCH_BLOCKING": "1"}, + capture_output=True, + text=True, + timeout=60, + ) + assert completed.returncode != 0, "native gather accepted invalid semantic indices" + assert ( + "device-side assert triggered" in completed.stderr + ), f"expected a CUDA bounds assertion, got:\n{completed.stdout}\n{completed.stderr}" + + def test_gather_chunks_drop_in(self): + rows = _rows(2, hidden=16) + ti, tags = h3_packed_layout(10, 2) + chunks = rows.chunk(6, dim=-1) + for a, b in zip(_cuda_op().gather_chunks(chunks, ti, tags), _cuda_op()(rows, ti, tags)): + assert torch.equal(a, b) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, seq, hidden, seed=0): + g = torch.Generator(device="cuda").manual_seed(seed) + return [torch.randn(seq, hidden, device="cuda", generator=g).bfloat16() for _ in range(6)] + + def _cuda_grad(self, rows, ti, tags, grads): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(_cuda_op()(leaf, ti, tags)), grads) + return leaf.grad + + def test_correctly_rounded_fp32_segment_sum(self): + rows = _rows(3, seed=1) + ti, tags = h3_packed_layout(4097, 3, seed=1) + grads = self._grads(4097, 5376) + ours = self._cuda_grad(rows, ti, tags, grads) + ref = rows.detach().double().requires_grad_(True) + torch.autograd.backward( + list(NativeH3AdaLNRowGatherOp().forward_fp32(ref, ti, tags)), [g.float() for g in grads] + ) + assert ours.dtype == torch.bfloat16 + # tolerance_contract.json gradient_accuracy / elementwise / bfloat16. + torch.testing.assert_close(ours.double(), ref.grad, atol=2e-2, rtol=1.6e-2) + assert (ours == ref.grad.bfloat16()).float().mean() > 0.9999 + + def test_repeat_bitwise(self): + rows = _rows(2, hidden=512, seed=2) + ti, tags = h3_packed_layout(5000, 2, seed=2) + grads = self._grads(5000, 512, seed=2) + first = self._cuda_grad(rows, ti, tags, grads) + for _ in range(3): + assert torch.equal(self._cuda_grad(rows, ti, tags, grads), first) + + def test_forward_rows_invariant_to_batch_size_and_position(self): + rows = _rows(3, hidden=512, seed=5) + ti, tags = h3_packed_layout(4097, 3, seed=5) + full = _cuda_op()(rows, ti, tags) + perm = torch.randperm(4097, generator=torch.Generator().manual_seed(5)).cuda() + for pick in ( + torch.tensor([0]), + torch.tensor([2048]), + torch.arange(100, 357), + torch.arange(4000, 4097), + perm, + ): + pick = pick.cuda() + part = _cuda_op()(rows, ti[pick], tags[pick]) + for a, b in zip(part, full): + assert torch.equal(a, b[pick]) + + def test_backward_row_gradient_independent_of_other_rows_tokens(self): + # Each row's segment is tiled on its own, so a row's gradient from the + # full packing equals the one from a packing of only its own positions + # (same order), even though ~455 positions per row cross tile boundaries. + rows = _rows(3, hidden=512, seed=6) + ti, tags = h3_packed_layout(4097, 3, seed=6) + grads = self._grads(4097, 512, seed=6) + full = self._cuda_grad(rows, ti, tags, grads) + flat = ti * 3 + tags + for r in range(rows.shape[0]): + own = (flat == r).nonzero().squeeze(1) + assert own.numel() > 256 + alone = self._cuda_grad(rows, ti[own], tags[own], [g[own] for g in grads]) + assert torch.equal(alone[r], full[r]) + assert torch.count_nonzero(alone[torch.arange(rows.shape[0], device="cuda") != r]) == 0 + + def test_unreferenced_rows_get_zero_grad(self): + rows = _rows(3, hidden=16) + ti = torch.zeros(7, dtype=torch.long, device="cuda") + tags = torch.zeros(7, dtype=torch.long, device="cuda") # only row 0 is used + grad = self._cuda_grad(rows, ti, tags, self._grads(7, 16)) + assert torch.count_nonzero(grad[1:]) == 0 + assert torch.count_nonzero(grad[0]) > 0 + + def test_segments_longer_than_one_tile(self): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import BACKWARD_TILE + + seq = 3 * BACKWARD_TILE + 17 + rows = _rows(1, hidden=8, dtype=torch.float32) + ti = torch.zeros(seq, dtype=torch.long, device="cuda") + tags = torch.zeros(seq, dtype=torch.long, device="cuda") + grads = [torch.ones(seq, 8, device="cuda") for _ in range(6)] + grad = self._cuda_grad(rows, ti, tags, grads) + assert torch.equal(grad[0], torch.full((48,), float(seq), device="cuda")) + + def test_reference_backward_is_deterministic_and_accurate(self): + rows = _rows(2, hidden=256, seed=4) + ti, tags = h3_packed_layout(3000, 2, seed=4) + grads = self._grads(3000, 256, seed=4) + + def ref_grad(): + leaf = rows.detach().clone().requires_grad_(True) + torch.autograd.backward(list(NativeH3AdaLNRowGatherOp()(leaf, ti, tags)), grads) + return leaf.grad + + first = ref_grad() + assert torch.equal(ref_grad(), first) + torch.testing.assert_close( + first.float(), self._cuda_grad(rows, ti, tags, grads).float(), atol=1e-1, rtol=2e-2 + ) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("adaln_row_gather", device="cuda") + assert type(op).__name__ == "H3AdaLNRowGatherCudaOp" diff --git a/tests/models/minimax_h3/test_h3_benchmark_timing.py b/tests/models/minimax_h3/test_h3_benchmark_timing.py new file mode 100644 index 000000000..bd3a44363 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_benchmark_timing.py @@ -0,0 +1,267 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU regressions for benchmark ordering and backward timing boundaries.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_report + + +@pytest.fixture +def cuda_clock(monkeypatch): + clock = SimpleNamespace(time=0.0, active=False, log=[], allocated=0, peak=0, events=0) + + class Event: + def __init__(self, *, enable_timing): + assert enable_timing + self.start = clock.events % 2 == 0 + clock.events += 1 + + def record(self): + clock.active = self.start + clock.log.append("start" if self.start else "end") + self.time = clock.time + + def synchronize(self): + assert not clock.active + + def elapsed_time(self, end): + return end.time - self.time + + def reset_peak(): + clock.log.append("reset_peak") + clock.peak = clock.allocated + + def allocated(): + clock.log.append("memory_baseline") + return clock.allocated + + monkeypatch.setattr(torch.cuda, "Event", Event) + monkeypatch.setattr(torch.cuda, "synchronize", lambda: clock.log.append("synchronize")) + monkeypatch.setattr(torch.cuda, "reset_peak_memory_stats", reset_peak) + monkeypatch.setattr(torch.cuda, "memory_allocated", allocated) + monkeypatch.setattr(torch.cuda, "max_memory_allocated", lambda: clock.peak) + return clock + + +def test_setup_is_outside_backward_events_and_memory_baseline(cuda_clock): + states = [] + + def setup(): + assert not cuda_clock.active + cuda_clock.log.append("setup") + cuda_clock.time += 50.0 + cuda_clock.allocated += 64 * 2**20 + return object() + + def backward(state): + states.append(state) + cuda_clock.log.append("backward") + cuda_clock.time += 2.0 + cuda_clock.allocated += 3 * 2**20 + cuda_clock.peak = cuda_clock.allocated + + assert h3_report.time_us(backward, warmup=1, iters=2, setup=setup) == 2000.0 + assert len({id(state) for state in states}) == 3 + assert cuda_clock.log.count("start") == 2 + for index, entry in enumerate(cuda_clock.log): + if entry == "start": + assert cuda_clock.log[index - 1] == "setup" + assert cuda_clock.log[index + 1 : index + 3] == ["backward", "end"] + + cuda_clock.log.clear() + assert h3_report.peak_mib(backward, setup=setup) == 3.0 + assert cuda_clock.log == [ + "setup", + "synchronize", + "reset_peak", + "memory_baseline", + "backward", + "synchronize", + ] + + +@pytest.mark.parametrize("keys", [("candidate", "provider"), h3_report.TIMED_KEYS]) +def test_measure_interleaves_and_records_both_orders(cuda_clock, keys): + calls = [] + case = {"op": "test", "case": "tiny", "backend": "cpu", "bytes": 4096} + for position, key in enumerate(keys, start=1): + + def run(key=key, duration=position): + calls.append(key) + cuda_clock.time += duration + + case[key] = run + + row = h3_report.measure(case, warmup=2, iters=4) + first, second = list(keys), list(reversed(keys)) + assert calls == (first + second) * 3 + first # warmup, samples, then memory calls + assert row["execution_order"] == { + "policy": "alternating", + "iteration_0": first, + "iteration_1": second, + } + for position, key in enumerate(keys, start=1): + assert row[f"{key}_us"] == position * 1000.0 + assert row[f"{key}_gbps"] == pytest.approx(4096 / (position * 1e-3) / 1e9) + assert row[f"{key}_peak_mib"] == 0.0 + assert not any(key.endswith("_setup") for key in row) + + +@pytest.mark.parametrize("operator", ["norm", "gather", "gate", "final"]) +def test_backward_perf_cases_prepare_fresh_cpu_graphs_before_events( + monkeypatch, cuda_clock, operator +): + prepared_leaves = [] + previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def track_backward(outputs): + tensors = outputs if isinstance(outputs, tuple) else (outputs,) + + def backward_hook(grad): + cuda_clock.time += 2.0 / len(tensors) + return grad + + for tensor in tensors: + tensor.register_hook(backward_hook) + return outputs + + def track(fn, leaves, *args): + if leaves[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([leaf if leaf.is_leaf else leaf._base for leaf in leaves]) + cuda_clock.time += 1.0 + outputs = fn(*leaves, *args) + return track_backward(outputs) if leaves[0].requires_grad else outputs + + if operator == "norm": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 4), + torch.randn(3, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_norm_modulate + + def forward(x, weight, shift, scale, index, **kwargs): + return track(provider, [x, weight, shift, scale], index) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_norm_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_norm_modulate", forward) + op = SimpleNamespace(forward_modulated=forward) + build = h3_report._norm_perf + elif operator == "gate": + + def inputs(seq): + return ( + torch.randn(3, 24), + torch.randn(1, seq, 4), + torch.randn(1, seq, 4), + torch.arange(seq) % 3, + ) + + provider = h3_report.provider_gate_residual + + def forward(residual, y, gate, index, **kwargs): + return track(lambda r, y_, g: provider(r, g, index, y_), [residual, y, gate]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_gate_inputs", inputs) + monkeypatch.setattr( + h3_report, "provider_gate_residual", lambda r, g, i, y: forward(r, y, g, i) + ) + op = SimpleNamespace(forward=forward) + build = h3_report._gate_perf + elif operator == "final": + + def inputs(seq): + return ( + torch.randn(1, seq, 4), + torch.ones(4), + torch.randn(3, 2), + torch.randn(8, 2), + torch.randn(8), + torch.arange(seq) % 3, + ) + + def forward(x, nw, temb, w, b, ti): + def compute(x, nw, temb, w, b): + rows = torch.nn.functional.linear(torch.nn.functional.silu(temb), w, b) + shift, scale = rows.chunk(2, dim=-1) + norm = torch.nn.functional.rms_norm(x, (4,), nw) + return norm * (1 + scale[ti]) + shift[ti] + + return track(compute, [x, nw, temb, w, b]) + + monkeypatch.setattr(h3_report, "NORM_SEQ_LENS", (2, 3)) + monkeypatch.setattr(h3_report, "_final_inputs", inputs) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", forward) + op = forward + build = h3_report._final_perf + else: + randn = torch.randn + cpu_generators = {} + + def cpu_randn(*args, **kwargs): + kwargs["device"] = "cpu" + # A seeded CUDA generator cannot draw a CPU tensor: continue the same seed on + # one CPU generator per CUDA generator. + gen = kwargs.get("generator") + if gen is not None and gen.device.type != "cpu": + if id(gen) not in cpu_generators: + cpu_generators[id(gen)] = torch.Generator().manual_seed(gen.initial_seed()) + kwargs["generator"] = cpu_generators[id(gen)] + return randn(*args, **kwargs) + + provider = h3_report.provider_adaln_row_gather + + def forward(rows, ti, tags, **kwargs): + def gather(leaf, ti, tags): + return provider(leaf.chunk(6, dim=-1), ti, tags) + + return track(gather, [rows], ti, tags) + + def provider_forward(chunks, ti, tags): + if chunks[0].requires_grad: + assert not cuda_clock.active + prepared_leaves.append([chunks[0]._base]) + cuda_clock.time += 1.0 + outputs = provider(chunks, ti, tags) + return track_backward(outputs) if chunks[0].requires_grad else outputs + + monkeypatch.setattr(torch, "randn", cpu_randn) + monkeypatch.setattr(h3_report, "GATHER_SEQ_LENS", (2, 3)) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: (torch.arange(seq) % num, torch.arange(seq) % 3), + ) + monkeypatch.setattr(h3_report, "provider_adaln_row_gather", provider_forward) + op = SimpleNamespace(forward=forward) + build = h3_report._gather_perf + + try: + registry = SimpleNamespace(get_op=lambda *args, **kwargs: op) + for case in build(registry): + row = h3_report.measure(case, warmup=1, iters=2) + assert row["backward_timing_scope"] == "backward_only" + assert row["candidate_backward_us"] == pytest.approx(2000.0) + assert row["provider_backward_us"] == pytest.approx(2000.0) + finally: + torch.set_num_threads(previous_threads) + + assert prepared_leaves + assert len({id(leaves[0]) for leaves in prepared_leaves}) == len(prepared_leaves) + assert all(leaf.grad is not None for leaves in prepared_leaves for leaf in leaves) diff --git a/tests/models/minimax_h3/test_h3_chain_replay_cli.py b/tests/models/minimax_h3/test_h3_chain_replay_cli.py new file mode 100644 index 000000000..b9baea84d --- /dev/null +++ b/tests/models/minimax_h3/test_h3_chain_replay_cli.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Reject incomplete chain dependencies before checking CUDA or loading weights.""" + +from __future__ import annotations + +import sys + +import pytest + +from tools.validation.models import h3_chain_replay + +STAGE_NAMES = [stage.name for stage in h3_chain_replay.h3_chain.STAGES] + + +@pytest.fixture +def no_cuda_check(monkeypatch): + """Fail if invalid CLI arguments reach a CUDA availability check.""" + + def unexpected_check(): + """Report an environment check performed before argument rejection.""" + + pytest.fail("invalid arguments must be rejected before checking CUDA") + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", unexpected_check) + + +@pytest.mark.parametrize("stages, message", [("", "unknown stages"), ("unknown", "unknown stages")]) +def test_reject_invalid_stages(monkeypatch, capsys, no_cuda_check, stages, message): + """Reject unknown stages before any device check.""" + + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert message in capsys.readouterr().err + + +@pytest.mark.parametrize("count", range(1, len(STAGE_NAMES))) +def test_backward_requires_all_stages(monkeypatch, capsys, no_cuda_check, count): + """Require the full chain for parameter-gradient replay.""" + + monkeypatch.setattr( + sys, + "argv", + ["h3_chain_replay.py", "--stages", ",".join(STAGE_NAMES[:count]), "--backward"], + ) + with pytest.raises(SystemExit) as exc: + h3_chain_replay.main() + assert exc.value.code == 2 + assert "--backward requires every stage" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "stages", + [ + *(",".join(STAGE_NAMES[:count]) for count in range(1, len(STAGE_NAMES) + 1)), + STAGE_NAMES[1], + STAGE_NAMES[-1], + ",".join(STAGE_NAMES[:1] + STAGE_NAMES[2:]), + ",".join(reversed(STAGE_NAMES)), + ",".join([STAGE_NAMES[0], STAGE_NAMES[0]]), + ], +) +def test_accept_any_known_selection(monkeypatch, stages): + """Any selection of known stages reaches the device requirement; the replay + runs its prerequisites (tests/models/minimax_h3/test_h3_cli.py).""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", "--stages", stages]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() + + +@pytest.mark.parametrize("options", [[], ["--backward"]]) +def test_accept_all_stages_by_default(monkeypatch, options): + """Keep the complete chain as the default for forward and backward replay.""" + + monkeypatch.setattr(h3_chain_replay.torch.cuda, "is_available", lambda: False) + monkeypatch.setattr(sys, "argv", ["h3_chain_replay.py", *options]) + with pytest.raises(SystemExit, match="the chain replay needs a CUDA device"): + h3_chain_replay.main() diff --git a/tests/models/minimax_h3/test_h3_cli.py b/tests/models/minimax_h3/test_h3_cli.py new file mode 100644 index 000000000..8e97fc6a2 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_cli.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""The H3 CLIs preserve chain dependencies and require pinned report weights.""" + +from __future__ import annotations + +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_chain +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV +from tools.validation.models import h3_chain_replay, h3_evidence + + +@pytest.fixture +def cpu_chain_cli(monkeypatch): + """Replace CUDA chain stages with CPU stubs that track dependencies and provider drift.""" + + executed = [] + state = {"provider_shift": 0} + + class CpuOp: + def __call__(self, value): + """Advance the synthetic chain value by one for candidate and golden replay.""" + + return value + 1 + + forward_fp32 = __call__ + + stages = [] + for index, stage in enumerate(h3_chain.STAGES): + + def stage_input(ctx, upstream, index=index): + """Use raw timesteps for stage zero and require predecessor outputs thereafter.""" + + if index == 0: + return ctx["timestep"] + assert isinstance(upstream, torch.Tensor), "a preceding stage must run first" + return upstream + + stages.append( + replace( + stage, + candidate=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + golden=lambda op, ctx, up, stage_input=stage_input: op(stage_input(ctx, up)), + provider=lambda ctx, up, stage_input=stage_input, index=index: ( + stage_input(ctx, up) + 1 + (state["provider_shift"] if index == 0 else 0) + ), + ) + ) + + def get_op(op_type, *, device): + """Record each dispatched stage and return its deterministic CPU stub.""" + + executed.append(op_type) + return CpuOp() + + monkeypatch.setattr(h3_chain, "STAGES", stages) + monkeypatch.setattr(h3_chain, "make_context", lambda *a, **kw: {"timestep": torch.zeros(1)}) + monkeypatch.setattr(h3_chain, "golden_op", lambda *a: CpuOp()) + monkeypatch.setattr(h3_chain_replay, "KernelRegistry", lambda: SimpleNamespace(get_op=get_op)) + monkeypatch.setattr(h3_chain_replay, "load_h3_conditioning_weights", lambda *a: {}) + monkeypatch.setattr(h3_chain, "environment", lambda: {}) + monkeypatch.setattr(h3_chain, "git_state", lambda: {}) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + return executed, state + + +@pytest.mark.parametrize("provider_shift", [0, 1]) +@pytest.mark.parametrize( + "requested, execution_count", + [ + ("timestep_mlp_fp32", 2), + ("adaln_projection_3mod", 3), + ("timestep_sinusoid_h3,adaln_projection_3mod", 3), + ("timestep_mlp_fp32,timestep_sinusoid_h3", 2), + ], +) +def test_chain_selection_runs_dependencies( + cpu_chain_cli, monkeypatch, tmp_path, capsys, requested, execution_count, provider_shift +): + """Execute prerequisite stages but report requested stages and preserve first drift.""" + + executed, state = cpu_chain_cli + state["provider_shift"] = provider_shift + out = tmp_path / "chain.json" + monkeypatch.setattr( + sys, + "argv", + [ + "h3_chain_replay.py", + "--stages", + requested, + "--timesteps", + "1", + "--seq-lens", + "3", + "--out", + str(out), + ], + ) + h3_chain_replay.main() + + report = json.loads(out.read_text()) + stage_names = [stage.name for stage in h3_chain.STAGES] + expected_report = [name for name in stage_names if name in requested.split(",")] + assert executed == stage_names[:execution_count] + assert report["executed_stages"] == executed + assert report["stages"] == expected_report + case = report["cases"][0] + assert [entry["stage"] for entry in case["stages"]] == expected_report + assert all(entry["chained_vs_golden"]["bitwise_equal"] for entry in case["stages"]) + assert case["first_drift"] == (stage_names[0] if provider_shift else None) + assert case["first_isolated_drift"] == (stage_names[0] if provider_shift else None) + summary = capsys.readouterr().out.split("; first_drift=", 1)[0] + for name in stage_names: + assert (f"{name}=" in summary) == (name in expected_report) + + +@pytest.mark.parametrize("op", ["timestep_mlp_fp32", "adaln_projection_3mod"]) +@pytest.mark.parametrize("weights_env", [None, " "]) +def test_weighted_evidence_requires_pinned_weights(monkeypatch, tmp_path, op, weights_env): + """Reject weighted evidence without a configured checkpoint before writing a report.""" + + if weights_env is None: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, weights_env) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + with pytest.raises(SystemExit, match=WEIGHTS_ENV): + h3_evidence.main() + assert not out.exists() + + +@pytest.mark.parametrize( + "op, weight_source", + [ + ("timestep_sinusoid_h3", "not_applicable"), + ("timestep_mlp_fp32", "pinned_checkpoint"), + ("adaln_projection_3mod", "pinned_checkpoint"), + ], +) +def test_evidence_records_weight_source(monkeypatch, tmp_path, op, weight_source): + """Label weighted evidence as pinned and sinusoid evidence as weight-independent.""" + + if weight_source == "not_applicable": + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + else: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False) + monkeypatch.setattr(torch.backends.cudnn, "allow_tf32", False) + monkeypatch.setattr(h3_evidence, "KernelRegistry", lambda: object()) + monkeypatch.setattr( + h3_evidence, + "git_state", + lambda: {"rl_kernel_commit": "test_commit", "tracked_tree_dirty": False}, + ) + monkeypatch.setattr(h3_evidence, "environment", lambda: {}) + monkeypatch.setitem(h3_evidence.ACCURACY, op, lambda registry: {}) + monkeypatch.setitem(h3_evidence.PERF_CASES, op, lambda registry: []) + out = tmp_path / "report.json" + monkeypatch.setattr(sys, "argv", ["h3_evidence.py", "--op", op, "--out", str(out)]) + h3_evidence.main() + assert json.loads(out.read_text())["weight_source"] == weight_source diff --git a/tests/models/minimax_h3/test_h3_conditioning_e2e.py b/tests/models/minimax_h3/test_h3_conditioning_e2e.py new file mode 100644 index 000000000..07cd49776 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_conditioning_e2e.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""End to end: the H3 conditioning chain on the pinned checkpoint, through the registry. + +Runs every stage registered in ``rl_engine.validation.models.h3_chain.STAGES`` (each +RFC #420 row adds its own), chained exactly as the model calls them, and +checks every stage's promises: + +* the registry dispatched the CUDA backend; +* a repeated run is bitwise equal; +* the chained output is within the stage's contract tolerance of the golden; +* stages that promise it are bitwise equal to diffusers on diffusers' inputs, + so the first isolated drift can only be a stage that declares a reduction. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.runtime.registry import KernelRegistry +from rl_engine.validation.models import h3_chain + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# (T distinct timesteps, S packed rows): single row, short packed, realistic-size packed. +CASES = [(1, 3), (3, 257), (4, 4097)] + + +@pytest.fixture(scope="module") +def chain_weights(h3_weights_cpu): + """Move all pinned conditioning tensors to CUDA once for the end-to-end module.""" + + return {name: tensor.cuda() for name, tensor in h3_weights_cpu.items()} + + +@pytest.fixture(scope="module") +def registry(): + """Provide the production registry used to dispatch every conditioning stage.""" + + return KernelRegistry() + + +@pytest.mark.parametrize("num_timesteps, seq_len", CASES) +def test_chain_forward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Check CUDA dispatch, golden tolerance, repeatability, and declared provider parity.""" + + report = h3_chain.run_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + stages = {stage.name: stage for stage in h3_chain.STAGES} + for entry in report["stages"]: + stage = stages[entry["stage"]] + assert entry["backend"].endswith("CudaOp"), entry + assert entry["repeat_bitwise_equal"], entry["stage"] + assert entry["chained_vs_golden"]["within_tolerance"], entry + if stage.provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], entry["stage"] + drift = report["first_isolated_drift"] + assert drift is None or not stages[drift].provider_bitwise_isolated + + +@pytest.mark.parametrize("num_timesteps, seq_len", [(1, 257), (3, 4097)]) +def test_chain_backward(require_h3_kernels, registry, chain_weights, num_timesteps, seq_len): + """Parameter gradients of the whole chain: deterministic, and FP32-accurate when fused.""" + + report = h3_chain.run_backward_case( + registry, chain_weights, num_timesteps=num_timesteps, seq_len=seq_len + ) + for name, entry in report["leaves"].items(): + for mode in ("candidate", "candidate_fused"): + assert entry[mode]["repeat_bitwise_equal"], (mode, name) + fused = entry["candidate_fused"] + if name.startswith("time_embedder"): + # FP32 parameters: no BF16 rounding anywhere on the fused path. + assert fused["max_abs_vs_golden_over_absmax"] < 1e-5, (name, fused) + else: + # BF16 parameters: only their own final rounding remains. + assert fused["correctly_rounded_fraction"] > 0.99, (name, fused) diff --git a/tests/models/minimax_h3/test_h3_det_linear.py b/tests/models/minimax_h3/test_h3_det_linear.py new file mode 100644 index 000000000..ca11c4e53 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_det_linear.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Native input validation for the H3 deterministic linear kernels.""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available +from rl_engine.backends.extension import _C + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device"), + pytest.mark.skipif(not det_linear_available(), reason="rl_engine._C lacks h3_det_linear_*"), +] + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_backward_input_rejects_zero_rows(dtype): + """Reject empty gradient batches in native input-gradient kernels for both dtypes.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + weight = torch.empty(16, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_input(grad, weight, dtype) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_rejects_zero_rows(dtype, with_bias): + """Reject empty gradient batches before parameter-gradient or bias-gradient reduction.""" + + grad = torch.empty(0, 16, device="cuda", dtype=torch.float32) + x = torch.empty(0, 8, device="cuda", dtype=dtype) + with pytest.raises(RuntimeError, match="grad must have at least one row"): + _C.h3_det_linear_backward_weight(grad, x, dtype, with_bias) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +def test_backward_input_handles_empty_columns(dtype, out_dtype, k_in, n_out): + """Match the input VJP of empty linear dimensions in the requested output dtype.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + weight = torch.empty(n_out, k_in, device="cuda", dtype=dtype) + expected = (grad @ weight.float()).to(out_dtype) + actual = _C.h3_det_linear_backward_input(grad, weight, out_dtype) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("k_in,n_out", [(0, 16), (8, 0), (0, 0)]) +@pytest.mark.parametrize("with_bias", [False, True]) +def test_backward_weight_handles_empty_columns(dtype, w_dtype, k_in, n_out, with_bias): + """Return empty weight gradients while preserving any nonempty bias reduction.""" + + grad = torch.arange(3 * n_out, device="cuda", dtype=torch.float32).reshape(3, n_out) + x = torch.arange(3 * k_in, device="cuda", dtype=dtype).reshape(3, k_in) + expected_weight = (grad.T @ x.float()).to(w_dtype) + actual = _C.h3_det_linear_backward_weight(grad, x, w_dtype, with_bias) + assert len(actual) == (2 if with_bias else 1) + torch.testing.assert_close(actual[0], expected_weight, atol=0, rtol=0) + if with_bias: + torch.testing.assert_close(actual[1], grad.sum(dim=0).to(w_dtype), atol=0, rtol=0) diff --git a/tests/models/minimax_h3/test_h3_final_adaln_out.py b/tests/models/minimax_h3/test_h3_final_adaln_out.py new file mode 100644 index 000000000..5fb964677 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_final_adaln_out.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``final_adaln_out``: norm_out shift/scale projection + final norm/modulation. + +* layout: ``shift`` is the first half of ``norm_out.linear``'s output, rows + are indexed by ``timestep_indices`` (not ``adaln_indices``); +* forward matches diffusers' ``norm_out`` to the projection's 1-ULP ties + (the norm and modulation are bitwise; the GEMV tree differs from cuBLAS); +* backward is deterministic, with the table gradient kept in FP32; +* rows are batch/position invariant; malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.final_adaln_out import NativeH3FinalAdaLNOutOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_final_adaln_out + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import h3_rmsnorm_available + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3FinalAdaLNOutCudaOp() + + +@pytest.fixture(scope="module") +def norm_out(h3_weights_cpu): + return [ + h3_weights_cpu[name].cuda() + for name in ("norm_out.norm.weight", "norm_out.linear.weight", "norm_out.linear.bias") + ] + + +def _inputs(seq, num_timesteps=3, seed=0, batch=1): + g = torch.Generator(device="cpu").manual_seed(seed) + x = (torch.randn(batch, seq, 5376, generator=g) * 2).bfloat16().cuda() + temb = (torch.randn(num_timesteps, 2688, generator=g) * 2).cuda() + ti, _ = h3_packed_layout(seq, num_timesteps, seed=seed) + return x, temb, ti + + +def _golden(x, norm_weight, temb, weight, bias, timestep_indices): + act = temb * torch.sigmoid(temb) + # Declared cast of the FP32 SiLU, as the model computes it; identity VJP. + declared = F.silu(temb.detach().float()).to(torch.bfloat16).double() + act = act + (declared - act).detach() + table = F.linear(act, weight, bias) + table = table + (table.to(torch.bfloat16).double() - table).detach() + shift, scale = table.chunk(2, dim=-1) + n = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + 1e-5) * norm_weight + return n * (1 + scale.index_select(0, timestep_indices)) + shift.index_select( + 0, timestep_indices + ) + + +class TestReference: + def test_rejects_bad_inputs(self): + op = NativeH3FinalAdaLNOutOp() + x, nw = torch.zeros(1, 4, 8), torch.ones(8) + temb, w, b = torch.zeros(2, 16), torch.zeros(16, 16), torch.zeros(16) + ti = torch.tensor([0, 1, 1, 0]) + with pytest.raises(TypeError): # BF16 temb (probe H7) + op(x, nw, temb.bfloat16(), w, b, ti) + with pytest.raises(ValueError): # not 2H rows + op(x, nw, temb, w[:10], b[:10], ti) + with pytest.raises(IndexError): + op(x, nw, temb, w, b, ti + 2) + + +@requires_cuda +class TestCuda: + def test_shift_first_layout_and_timestep_indexing(self): + """Zero weights: the bias halves are (shift, scale); rows follow timestep_indices.""" + + hidden, seq = 16, 6 + x = torch.randn(1, seq, hidden, device="cuda").bfloat16() + nw = torch.ones(hidden, device="cuda", dtype=torch.bfloat16) + w = torch.zeros(2 * hidden, 32, device="cuda", dtype=torch.bfloat16) + b = torch.cat([torch.full((hidden,), 1.0), torch.full((hidden,), -1.0)]).bfloat16().cuda() + ti = torch.tensor([0, 1, 0, 1, 1, 0], device="cuda") + out = _cuda_op()(x, nw, torch.zeros(2, 32, device="cuda"), w, b, ti) + # scale = -1 -> norm(x) * 0 + shift = 1 everywhere + assert torch.equal(out, torch.ones_like(out)) + + @pytest.mark.parametrize("seq, num_timesteps", [(1, 1), (257, 2), (4097, 3)]) + def test_forward_against_diffusers_and_golden(self, norm_out, seq, num_timesteps): + x, temb, ti = _inputs(seq, num_timesteps, seed=seq) + ours = _cuda_op()(x, *norm_out[:1], temb, *norm_out[1:], ti) + theirs = provider_final_adaln_out(x, norm_out[0], temb, *norm_out[1:], ti) + golden = NativeH3FinalAdaLNOutOp().forward_fp32(x, norm_out[0], temb, *norm_out[1:], ti) + assert (ours == theirs).float().mean() > 0.999 + torch.testing.assert_close(ours.float(), golden, atol=5e-2, rtol=2e-2) + + def test_rows_are_batch_and_position_invariant(self, norm_out): + x, temb, ti = _inputs(600, seed=1, batch=2) + full = _cuda_op()(x, norm_out[0], temb, *norm_out[1:], ti) + part = _cuda_op()(x[1:2, 200:260], norm_out[0], temb, *norm_out[1:], ti[200:260]) + assert torch.equal(part[0], full[1, 200:260]) + + def test_backward_deterministic_and_fp32_table_gradient(self, norm_out): + x, temb, ti = _inputs(4097, seed=2) + grad = torch.randn_like(x) + + def grads(fn, dtype=None): + tensors = [x, norm_out[0], temb, *norm_out[1:]] + leaves = [ + (t if dtype is None else t.to(dtype)).detach().clone().requires_grad_(True) + for t in tensors + ] + fn(*leaves).backward(grad if dtype is None else grad.to(dtype)) + return [leaf.grad for leaf in leaves] + + ours = grads(lambda *t: _cuda_op()(*t, ti)) + again = grads(lambda *t: _cuda_op()(*t, ti)) + ref = grads(lambda *t: _golden(*t, ti), torch.float64) + assert all(torch.equal(a, b) for a, b in zip(ours, again)) + for name, g, r in zip(("dx", "d_norm_weight", "d_temb", "dW", "db"), ours, ref): + rel = ((g.double() - r).abs().max() / r.abs().max()).item() + assert rel < 1e-2, (name, rel) # BF16 output rounding and round(1 + scale) only + + def test_report_backward_uses_fp64_golden(self, monkeypatch): + from types import SimpleNamespace + + from rl_engine.validation.models import h3_report + + rng = torch.Generator(device="cuda").manual_seed(3) + x = torch.randn(1, 160, 8, device="cuda", generator=rng).bfloat16() + nw = torch.randn(8, device="cuda", generator=rng).bfloat16() + temb = torch.randn(3, 4, device="cuda", generator=rng) + w = torch.randn(16, 4, device="cuda", generator=rng).bfloat16() + b = torch.randn(16, device="cuda", generator=rng).bfloat16() + ti = torch.arange(160, device="cuda") % 3 + inputs = (x, nw, temb, w, b) + native = NativeH3FinalAdaLNOutOp() + registry = SimpleNamespace( + get_op=lambda *args, **kwargs: native, + _get_or_create_backend=lambda _: native, + _priority_map={"cpu": {"final_adaln_out": [None]}}, + ) + monkeypatch.setattr(h3_report, "_final_inputs", lambda *args, **kwargs: (*inputs, ti)) + monkeypatch.setattr(h3_report, "provider_final_adaln_out", native) + report = h3_report._final_accuracy(registry) + grad = torch.randn( + x.shape, device="cuda", generator=torch.Generator(device="cuda").manual_seed(7) + ).to(x.dtype) + leaves = [t.detach().clone().requires_grad_(True) for t in inputs] + native(*leaves, ti).backward(grad) + ref = [t.double().detach().requires_grad_(True) for t in inputs] + _golden(*ref, ti).backward(grad.double()) + for name, leaf, reference in zip(("dx", "d_norm_w", "d_temb", "dW", "db"), leaves, ref): + assert reference.grad.dtype == torch.float64 + expected = float( + (leaf.grad.double() - reference.grad).abs().max() / reference.grad.abs().max() + ) + for backend in ("cuda", "provider"): + assert report["backward"][backend]["rel_error"][name] == pytest.approx(expected) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("final_adaln_out", device="cuda") + assert type(op).__name__ == "H3FinalAdaLNOutCudaOp" + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_bf16_gradients_meet_gtest_contract(self, seq): + # The check_operator path; S=257 used to miss d_norm_weight before the + # golden stored norm(x) and 1 + scale in BF16 like the model does. + import argparse + + from rl_engine.validation.operators import run_operator_suite + from rl_engine.validation.operators.operator_specs import make_candidate, make_operator_case + + _cuda_op() + args = argparse.Namespace( + op="final_adaln_out", + candidate="cuda", + arch_key=None, + batch=3, + seq=seq, + normalized_dim=5376, + seed=123, + input_mode="random", + ) + report = run_operator_suite( + "final_adaln_out", + candidates=[make_candidate(args)], + cases=[make_operator_case(args, torch.bfloat16, torch.device("cuda"))], + check_grad=True, + ) + assert report.passed, report.candidates[0].cases[0].outputs diff --git a/tests/models/minimax_h3/test_h3_full_model.py b/tests/models/minimax_h3/test_h3_full_model.py new file mode 100644 index 000000000..b10898105 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_full_model.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.validation.models import h3_block +from rl_engine.validation.models.h3_full_model import ( + FullModelTrace, + full_model_execution, + validate_full_structure, +) + + +class Block(torch.nn.Module): + def __init__(self, value): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor(float(value))) + + def forward(self, x, *args): + return x + self.weight + + +class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace( + num_layers=50, + num_refiner_layers=2, + hidden_size=5376, + num_attention_heads=56, + attention_head_dim=128, + ) + self.transformer_blocks = torch.nn.ModuleList([Block(i) for i in range(50)]) + for name in [ + "proj_in", + "audio_proj_in", + "context_embedder", + "time_proj", + "time_embedder", + "token_refiner", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + setattr(self, name, torch.nn.Identity()) + self.token_refiner.refiner_blocks = torch.nn.ModuleList([Block(50), Block(51)]) + + +def test_full_stack_uses_each_layers_own_parameters_and_keeps_gradients(monkeypatch): + model = Model() + seen = [] + + def run(mode, ctx, ops): + p = ctx["weights"][h3_block.BLOCK + ".weight"] + seen.append(p.item()) + return {h3_block.OUTPUT: ctx["hidden"] + p} + + monkeypatch.setattr(h3_block, "run_graph", run) + original = model.transformer_blocks[0].forward + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + x = torch.ones(1, requires_grad=True) + for block in model.transformer_blocks: + x = block(x, None, None, None) + x.sum().backward() + assert x.item() == 1 + sum(range(50)) + assert seen == list(range(50)) + assert all(b.weight.grad.item() == 1 for b in model.transformer_blocks) + assert model.transformer_blocks[0].forward == original + + +def test_repeated_block_zero_is_not_full_model(): + model = Model() + model.transformer_blocks[49] = model.transformer_blocks[0] + with pytest.raises(ValueError, match="distinct"): + validate_full_structure(model) + + +def test_partial_execution_cannot_publish_full_model_pass(): + trace = FullModelTrace() + trace.observe("block.0", torch.ones(1)) + with pytest.raises(ValueError, match="incomplete"): + trace.validate() + + +def test_binding_restored_after_failure(): + model = Model() + original = model.transformer_blocks[12].forward + with pytest.raises(RuntimeError, match="injected"): + with full_model_execution(model, candidate=True, trace=FullModelTrace(), ops=object()): + raise RuntimeError("injected") + assert model.transformer_blocks[12].forward == original + + +def test_candidate_execution_requires_every_node(): + trace = FullModelTrace() + trace.candidate = True + for name in [f"block.{i}" for i in range(50)] + [ + "refiner.0", + "refiner.1", + "proj_in", + "audio_proj_in", + "context_embedder", + "token_refiner", + "time_embedder", + "norm_out", + "proj_out", + "audio_proj_out", + ]: + trace.calls[name] = 1 + with pytest.raises(ValueError, match="candidate nodes did not execute"): + trace.validate() + + +def test_drift_order_follows_execution_instead_of_lexical_names(): + from tools.validation.models.h3_full_compare import drift_order + + report = { + "trace": { + "calls": {"block.2": 1, "block.10": 1}, + "backward_calls": {"block.10": 1, "block.2": 1}, + } + } + names = [ + "video", + "nodes.block.10", + "nodes.block.2", + "nodes.block.2.grad", + "nodes.block.10.grad", + ] + assert drift_order(report, names) == [ + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "nodes.block.2.grad", + "video", + ] + + +def test_reference_inventory_rejects_tampered_tensor(tmp_path): + import hashlib + import json + + from tools.validation.models.h3_evidence_integrity import verify_inventory + + tensor = tmp_path / "tensor.pt" + tensor.write_bytes(b"original") + manifest = tmp_path / "manifest.json" + manifest.write_text( + json.dumps({"files": {"tensor.pt": hashlib.sha256(b"original").hexdigest()}}) + ) + verify_inventory(tmp_path, manifest) + tensor.write_bytes(b"changed") + with pytest.raises(ValueError, match="changed"): + verify_inventory(tmp_path, manifest) + + +def test_cuda_guard_rejects_one_fallback_among_native_bindings(): + from rl_engine.validation.models.h3_full_model import require_candidate_cuda + + native = "rl_engine.backends.cuda.model_specific.minimax_h3.NativeOp" + ops = SimpleNamespace( + backends=lambda: {"projection": native, "norm": "torch.nn.RMSNorm", "gate": native} + ) + with pytest.raises(ValueError, match="candidate norm fell back"): + require_candidate_cuda(ops) + + +def test_drift_order_interleaves_parameter_and_activation_backward_hooks(): + from tools.validation.models.h3_full_compare import drift_order + + trace = FullModelTrace() + for name in ( + "nodes.block.2", + "nodes.block.10", + "nodes.block.10.grad", + "parameter.block.10.weight", + "nodes.block.2.grad", + ): + trace.record_event(name) + trace.record_event("nodes.block.2") # checkpoint recomputation is not a new first event + names = list(reversed(trace.execution_order)) + assert drift_order({"trace": {"execution_order": trace.execution_order}}, names) == ( + trace.execution_order + ) diff --git a/tests/models/minimax_h3/test_h3_native_validation.py b/tests/models/minimax_h3/test_h3_native_validation.py new file mode 100644 index 000000000..03a70f1a3 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_native_validation.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Exercise H3 validation at the native entrypoints, bypassing Python guards.""" + +from __future__ import annotations + +import pytest +import torch + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +@pytest.fixture +def native_extension(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + names = ( + "h3_rmsnorm_forward", + "h3_rmsnorm_backward", + "h3_gate_residual_forward", + "h3_gate_residual_backward", + ) + if not _EXT_AVAILABLE or not all(hasattr(_C, name) for name in names): + pytest.skip("rl_engine._C lacks native H3 entrypoints") + return _C + + +@pytest.fixture +def native_case(): + device = "cuda:0" + x = torch.arange(16, device=device, dtype=torch.float32).view(2, 8) / 10 + return { + "x": x, + "weight": torch.ones(8, device=device), + "shift": torch.zeros(2, 8, device=device), + "scale": torch.zeros(2, 8, device=device), + "gate": torch.ones(2, 8, device=device), + "grad": torch.ones_like(x), + "rstd": torch.rsqrt(x.square().mean(-1) + 1e-5), + "index": torch.arange(2, device=device, dtype=torch.int64), + "sorted_pos": torch.arange(2, device=device, dtype=torch.int64), + "tile_begin": torch.arange(2, device=device, dtype=torch.int64), + "tile_end": torch.arange(1, 3, device=device, dtype=torch.int64), + "seg_first_tile": torch.arange(3, device=device, dtype=torch.int64), + } + + +def _native_call(extension, operation, boundary, case): + x, index = case["x"], case["index"] + tiles = tuple(case[name] for name in ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile")) + if operation == "rmsnorm": + modulation = (case["shift"], case["scale"], index) + if boundary == "forward": + return extension.h3_rmsnorm_forward(x, case["weight"], 1e-5, *modulation) + return extension.h3_rmsnorm_backward( + case["grad"], x, case["weight"], case["rstd"], *modulation, *tiles + ) + if boundary == "forward": + return extension.h3_gate_residual_forward(x, x, case["gate"], index) + return extension.h3_gate_residual_backward(case["grad"], x, case["gate"], index, *tiles) + + +def _other_cuda(tensor): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + return tensor.to("cuda:1") + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("boundary", ["forward", "backward"]) +@pytest.mark.parametrize("index_case", ["empty", "cpu", "other_cuda"]) +def test_native_rejects_invalid_row_index( + native_extension, native_case, operation, boundary, index_case +): + index = native_case["index"] + if index_case == "empty": + native_case["index"] = index[:0] + elif index_case == "cpu": + native_case["index"] = index.cpu() + else: + native_case["index"] = _other_cuda(index) + with pytest.raises( + RuntimeError, match="index must be a non-empty contiguous int64 tensor on cuda:0" + ): + _native_call(native_extension, operation, boundary, native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("device", ["cpu", "other_cuda"]) +def test_native_rejects_mixed_device_tiles( + native_extension, native_case, operation, metadata, device +): + tensor = native_case[metadata] + native_case[metadata] = tensor.cpu() if device == "cpu" else _other_cuda(tensor) + with pytest.raises( + RuntimeError, match="tile metadata must be contiguous int64 tensors on cuda:0" + ): + _native_call(native_extension, operation, "backward", native_case) + + +@pytest.mark.parametrize("operation", ["rmsnorm", "gate_residual"]) +def test_native_accepts_same_device_inputs(native_extension, native_case, operation): + x = native_case["x"].detach().clone().requires_grad_(True) + if operation == "rmsnorm": + weight = native_case["weight"].detach().clone().requires_grad_(True) + shift = native_case["shift"].detach().clone().requires_grad_(True) + scale = native_case["scale"].detach().clone().requires_grad_(True) + norm = torch.nn.functional.rms_norm(x, (8,), weight, 1e-5) + ref = norm * (1 + scale) + shift + ref.backward(native_case["grad"]) + expected_grads = (x.grad, weight.grad, shift.grad, scale.grad) + got, _ = _native_call(native_extension, operation, "forward", native_case) + else: + gate = native_case["gate"].detach().clone().requires_grad_(True) + ref = native_case["x"] + gate * x + ref.backward(native_case["grad"]) + expected_grads = (x.grad, gate.grad) + got = _native_call(native_extension, operation, "forward", native_case) + torch.testing.assert_close(got, ref) + actual_grads = _native_call(native_extension, operation, "backward", native_case) + for actual, expected in zip(actual_grads, expected_grads, strict=True): + torch.testing.assert_close(actual, expected) + + +TILE_NAMES = ("sorted_pos", "tile_begin", "tile_end", "seg_first_tile") + + +@pytest.fixture +def native(): + from rl_engine.backends.extension import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not all( + hasattr(_C, name) for name in ("h3_rmsnorm_forward", "h3_rmsnorm_backward") + ): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return _C + + +@pytest.fixture +def inputs(native): + x = torch.randn(4, 8, device="cuda") + weight = torch.ones(8, device=x.device) + _, rstd = native.h3_rmsnorm_forward(x, weight, 1e-5) + shift, scale = torch.zeros(2, 16, device=x.device).chunk(2, dim=1) + return { + "grad": torch.ones_like(x), + "x": x, + "weight": weight, + "rstd": rstd, + "shift": shift, + "scale": scale, + "index": torch.tensor([0, 1], dtype=torch.int64, device=x.device), + "sorted_pos": torch.tensor([0, 2, 1, 3], dtype=torch.int64, device=x.device), + "tile_begin": torch.tensor([0, 2], dtype=torch.int64, device=x.device), + "tile_end": torch.tensor([2, 4], dtype=torch.int64, device=x.device), + "seg_first_tile": torch.tensor([0, 1, 2], dtype=torch.int64, device=x.device), + } + + +def _call(native, inputs, entrypoint): + if entrypoint == "forward": + return native.h3_rmsnorm_forward( + inputs["x"], + inputs["weight"], + 1e-5, + inputs["shift"], + inputs["scale"], + inputs["index"], + ) + return native.h3_rmsnorm_backward(**inputs) + + +def _malformed(tensor, defect): + if defect == "cpu": + return tensor.cpu() + if defect == "dtype": + dtype = torch.float16 if tensor.dtype == torch.float32 else torch.int32 + return tensor.to(dtype) + if defect == "rank": + return tensor.unsqueeze(0) + if defect == "stride": + return torch.stack((tensor, tensor), dim=1)[:, 0] + if defect == "length": + return tensor[:-1] + raise ValueError(f"unknown defect: {defect}") + + +@pytest.mark.cuda_only +def test_valid_plain_and_modulated_backward(native, inputs): + plain = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + assert len(native.h3_rmsnorm_backward(**plain)) == 2 + dx, dweight, dshift, dscale = native.h3_rmsnorm_backward(**inputs) + assert dx.shape == inputs["x"].shape + assert dweight.shape == inputs["weight"].shape + assert dscale.shape == inputs["scale"].shape + torch.testing.assert_close(dshift, torch.full_like(inputs["shift"], 2.0)) + torch.cuda.synchronize(inputs["x"].device) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_empty_row_index(native, inputs, entrypoint): + inputs["index"] = inputs["index"][:0] + with pytest.raises(RuntimeError, match="row index must be a non-empty"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("entrypoint", ["forward", "backward"]) +def test_rejects_zero_hidden_size(native, inputs, entrypoint): + inputs["x"] = inputs["x"][:, :0].contiguous() + inputs["weight"] = inputs["weight"][:0] + with pytest.raises(RuntimeError, match="at least one column"): + _call(native, inputs, entrypoint) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +@pytest.mark.parametrize("mode", ["plain", "modulated"]) +def test_rejects_invalid_rstd(native, inputs, defect, mode): + inputs["rstd"] = _malformed(inputs["rstd"], defect) + if mode == "plain": + inputs = {name: inputs[name] for name in ("grad", "x", "weight", "rstd")} + with pytest.raises(RuntimeError, match="rstd must be contiguous float32"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +@pytest.mark.parametrize("defect", ["cpu", "dtype", "rank", "stride", "length"]) +def test_rejects_invalid_segment_tiles(native, inputs, name, defect): + inputs[name] = _malformed(inputs[name], defect) + message = "tile_begin and tile_end" if defect == "length" and name.startswith("tile_") else name + with pytest.raises(RuntimeError, match=message): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_rejects_missing_segment_tiles(native, inputs, name): + inputs[name] = None + with pytest.raises(RuntimeError, match="needs the sorted segment tiles"): + native.h3_rmsnorm_backward(**inputs) + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("name", TILE_NAMES) +def test_plain_backward_rejects_segment_tiles(native, inputs, name): + plain = {key: inputs[key] for key in ("grad", "x", "weight", "rstd", name)} + with pytest.raises(RuntimeError, match="sorted segment tiles require modulation"): + native.h3_rmsnorm_backward(**plain) + + +@pytest.mark.cuda_only +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize( + ("entrypoint", "name"), + [ + ("forward", "weight"), + ("forward", "index"), + ("backward", "weight"), + ("backward", "index"), + ("backward", "rstd"), + ("backward", "sorted_pos"), + ("backward", "tile_begin"), + ("backward", "tile_end"), + ("backward", "seg_first_tile"), + ], +) +def test_rejects_tensor_on_another_cuda_device(native, inputs, entrypoint, name): + other_device = (inputs["x"].device.index + 1) % torch.cuda.device_count() + inputs[name] = inputs[name].to(f"cuda:{other_device}") + message = "row index" if name == "index" else name + with pytest.raises(RuntimeError, match=message): + _call(native, inputs, entrypoint) + + +@pytest.mark.parametrize( + "operation", ["gate_forward", "gate_dy", "norm_forward", "norm_dx", "norm_partials"] +) +@pytest.mark.parametrize("bad_index", [-1, 2]) +def test_native_table_bounds_fail_in_isolated_process(native_extension, operation, bad_index): + import subprocess + import sys + import textwrap + + script = textwrap.dedent( + """ + import sys + import torch + from rl_engine import _C + op, bad = sys.argv[1], int(sys.argv[2]) + x = torch.ones(2, 8, device="cuda") + w = torch.ones(8, device="cuda") + table = torch.ones(2, 8, device="cuda") + index = torch.tensor([bad, 0], device="cuda") + rstd = torch.ones(2, device="cuda") + rows = torch.arange(2, device="cuda") + begin = torch.tensor([0], device="cuda") + end = torch.tensor([2], device="cuda") + try: + if op == "gate_forward": + _C.h3_gate_residual_forward(x, x, table, index) + elif op == "gate_dy": + _C.h3_gate_residual_backward_dy(x, table, index) + elif op == "norm_forward": + _C.h3_rmsnorm_forward(x, w, 1e-5, table, table, index) + elif op == "norm_dx": + _C.h3_rmsnorm_backward_dx(x, x, w, rstd, table, table, index) + else: + _C.h3_rmsnorm_backward_partials( + x, x, w, rstd, table, table, index, rows, begin, end, rows, begin, end) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid index reached native table lookup without an assertion") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, operation, str(bad_index)], + capture_output=True, + text=True, + timeout=60, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +def test_row_gather_rejects_metadata_on_other_device(native_extension, metadata): + if torch.cuda.device_count() < 2: + pytest.skip("needs two CUDA devices") + grad = torch.ones(6, 2, 8, device="cuda:0") + tensors = dict( + sorted_pos=torch.tensor([0, 1], device="cuda:0"), + tile_begin=torch.tensor([0], device="cuda:0"), + tile_end=torch.tensor([2], device="cuda:0"), + seg_first_tile=torch.tensor([0, 1], device="cuda:0"), + ) + tensors[metadata] = tensors[metadata].to("cuda:1") + with pytest.raises(RuntimeError, match="grad device"): + native_extension.h3_adaln_row_gather_backward(grad, **tensors, out_dtype=torch.float32) + + +@pytest.mark.parametrize("metadata", ["sorted_pos", "tile_begin", "tile_end", "seg_first_tile"]) +@pytest.mark.parametrize("bound", ["negative", "past_end"]) +def test_row_gather_metadata_bounds_in_isolated_process(native_extension, metadata, bound): + import subprocess + import sys + import textwrap + + script = textwrap.dedent( + """ + import sys + import torch + from rl_engine import _C + name, bound = sys.argv[1:] + grad = torch.ones(6, 2, 8, device="cuda") + metadata = dict(sorted_pos=torch.tensor([0, 1], device="cuda"), + tile_begin=torch.tensor([0], device="cuda"), + tile_end=torch.tensor([2], device="cuda"), + seg_first_tile=torch.tensor([0, 1], device="cuda")) + metadata[name][0] = -1 if bound == "negative" else 3 + try: + _C.h3_adaln_row_gather_backward(grad, **metadata, out_dtype=torch.float32) + torch.cuda.synchronize() + except RuntimeError as exc: + if "device-side assert" in str(exc): + print("BOUNDS_ASSERT_CONFIRMED") + sys.exit(0) + raise + raise AssertionError("invalid metadata was not rejected") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, metadata, bound], capture_output=True, text=True, timeout=60 + ) + assert result.returncode == 0, result.stdout + result.stderr + assert "BOUNDS_ASSERT_CONFIRMED" in result.stdout diff --git a/tests/models/minimax_h3/test_h3_one_block.py b/tests/models/minimax_h3/test_h3_one_block.py new file mode 100644 index 000000000..8f64a4508 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_one_block.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""One real MiniMax-H3 block, forward and backward (RFC #420 ``ws1_one_h3_block``). + +GPU tests run block 0 of the pinned checkpoint on packed FL2VA layouts and +check the block-level promises of ``rl_engine.validation.models.h3_block``: + +* every node is bitwise repeatable, and batch row 0 of a two-row batch has the + bytes of the single-row run (forward, and the gradient reaching every node); +* nodes that promise it are bitwise equal to diffusers on diffusers' inputs; +* every node, the block output and every gradient are about as close to the + FP32 golden as diffusers itself. + +CPU tests pin the packed layouts to the pinned diffusers pipeline and check +the graph's structure. +""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.reference.minimax_h3 import block as gold +from rl_engine.validation.models import h3_block +from rl_engine.validation.models import h3_provider as prov +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS, h3_fl2va_layout + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 8, + reason="needs an SM80+ CUDA device", +) + +# sha256(position_ids float64 bytes + token_tags int64 bytes) of diffusers@f53d552's +# ``MiniMaxH3PrepareLayoutStep.build_packed_sequence`` for each layout (text tags 1, patch +# (1, 2, 2), two audio channels, audio tag 2, video tag 0, no keyframe anchors). +PINNED_LAYOUTS = { + "tiny": "5329a0a13a9996c4fea7a0e676b65aa24302d65bea2e2248ba68d12184cc7c37", + "small": "64e4e79d8fb91d5a68ae7d15bb29049e08add38da2639030e089c6ef80968931", + "medium": "16e448800695eedfb85e0aacbca9bd5415ed53fbbe25f701bcb60df5f4754a31", +} + +# Accuracy: the candidate's error against the FP32 golden may exceed diffusers' by this +# factor plus an absolute floor (relative to the golden's max magnitude). +FORWARD_FACTOR, FORWARD_FLOOR = 1.25, 1e-3 +BACKWARD_FACTOR, BACKWARD_FLOOR = 1.5, 5e-3 +# A BF16 leaf gradient may instead exceed diffusers' error by this many ulps of its largest +# golden element. Its max|err| / max|golden| is set by one rounded element, so it moves in +# whole ulps (0.004 relative for attn.norm_q.weight), and the relative rule alone failed on +# seed 2 with torch 2.13 and on seed 0 with torch 2.9.1. Sized from +# tools/validation/models/h3_block_ulp_sweep.py over tiny/small x seeds 0-9 on both torch +# versions (max excess +4.15 ulp) and 34 single-node candidate/provider swaps (max +5.15 ulp); +# reports in reports/experiments/h3-one-block-b200/ulp-sweep-torch*.json. +BACKWARD_BF16_EXCESS_ULPS = 6.0 + + +@pytest.mark.parametrize("layout", sorted(PINNED_LAYOUTS)) +def test_layout_matches_pinned_pipeline(layout): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS[layout], device="cpu") + data = packed["position_ids"].numpy().tobytes() + packed["token_tags"].numpy().tobytes() + assert hashlib.sha256(data).hexdigest() == PINNED_LAYOUTS[layout] + + +def test_graph_is_well_formed(): + seen = {"hidden", "temb"} + for node in h3_block.NODES: + assert set(node.inputs) <= seen, node.name + assert node.name not in seen, node.name + assert node.binding.status in ("row", "interim", "reference"), node.name + seen.add(node.name) + assert h3_block.OUTPUT == "residual_mlp" + rows = {n.binding.rfc_row for n in h3_block.NODES if n.binding.status == "row"} + assert rows == {"adaln_projection_3mod", "h3_rmsnorm", "adaln_gate_residual"} + + +def test_partial_rope_rotates_96_of_128_channels(): + packed = h3_fl2va_layout(**H3_BLOCK_LAYOUTS["tiny"], device="cpu") + cos, sin = prov.provider_rope(packed["position_ids"]) + assert cos.shape == (packed["position_ids"].shape[0], 96) and cos.dtype == torch.float32 + gold_cos, gold_sin = gold.golden_rope_tables(packed["position_ids"]) + assert torch.equal(cos, gold_cos) and torch.equal(sin, gold_sin) + x = torch.randn(1, cos.shape[0], 2, 128).bfloat16() + for out in (prov.provider_apply_rotary(x, cos, sin), gold.golden_apply_rotary(x, cos, sin)): + assert torch.equal(out[..., 96:].float(), x[..., 96:].float()) + assert not torch.equal(out[..., :96].float(), x[..., :96].float()) + + +@pytest.fixture(scope="module") +def ops(): + return h3_block.CandidateOps() + + +@needs_gpu +def test_stack_rows_dispatch_to_cuda(require_h3_kernels, ops): + backends = ops.backends() + for name in ("projection", "norm", "gate"): + assert backends[name].endswith("CudaOp"), backends + + +@needs_gpu +@pytest.mark.parametrize("layout", ["tiny", "small"]) +def test_block_forward(h3_block_weights_cuda, ops, layout): + report = h3_block.run_block_case(h3_block_weights_cuda, layout=layout, ops=ops) + nodes = {node.name: node for node in h3_block.NODES} + for entry in report["nodes"]: + name = entry["node"] + assert entry["repeat_bitwise_equal"], name + assert entry["batch_row_bitwise_equal"], name + if nodes[name].provider_bitwise_isolated: + assert entry["isolated_vs_provider"]["bitwise_equal"], name + bound = FORWARD_FACTOR * entry["provider_rel_err_vs_golden"] + FORWARD_FLOOR + assert entry["candidate_rel_err_vs_golden"] <= bound, entry + assert report["first_repeat_drift"] is None + assert report["first_batch_drift"] is None + drift = report["first_isolated_drift"] + assert drift is None or not nodes[drift].provider_bitwise_isolated + out = report["output"] + assert out["candidate_rel_err_vs_golden"] <= ( + FORWARD_FACTOR * out["provider_rel_err_vs_golden"] + FORWARD_FLOOR + ) + + +@needs_gpu +def test_block_backward(h3_block_weights_cuda, ops): + report = h3_block.run_block_backward(h3_block_weights_cuda, layout="tiny", ops=ops) + for name, entry in report["leaves"].items(): + assert entry["repeat_bitwise_equal"], name + bound = BACKWARD_FACTOR * entry["provider_rel_err_vs_golden"] + BACKWARD_FLOOR + within_rel = entry["candidate_rel_err_vs_golden"] <= bound + within_ulps = entry["storage_dtype"] == "bfloat16" and ( + entry["candidate_ulps_vs_golden"] + <= entry["provider_ulps_vs_golden"] + BACKWARD_BF16_EXCESS_ULPS + ) + assert within_rel or within_ulps, (name, entry) + assert report["leaves"]["hidden"]["batch_row_bitwise_equal"] + assert report["first_grad_repeat_drift"] is None + assert report["first_grad_batch_drift"] is None diff --git a/tests/models/minimax_h3/test_h3_report.py b/tests/models/minimax_h3/test_h3_report.py new file mode 100644 index 000000000..ed74d1901 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_report.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Optional checkpoint handling preserves standalone operator measurements, and +comparative timing balances backend order and preserves its raw evidence.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.validation.models import h3_chain, h3_report +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import ( + provider_adaln_row_gather, + provider_norm_modulate, +) +from rl_engine.validation.models.h3_weights import WEIGHTS_ENV + + +@pytest.fixture +def small_cpu_report(monkeypatch): + """Exercise actual report/autograd logic with small CPU tensors and provider ops.""" + + def dimension(size): + return {5376: 4, 6 * 5376: 24, 257: 160, 777: 160, 4097: 160, 32768: 160}.get(size, size) + + class CpuTorch: + nn = SimpleNamespace( + functional=SimpleNamespace( + rms_norm=lambda x, _shape, weight, eps: F.rms_norm(x, (x.shape[-1],), weight, eps) + ) + ) + + def __getattr__(self, name): + return getattr(torch, name) + + def Generator(self, device): + return torch.Generator(device="cpu") + + def randn(self, *shape, **kwargs): + if len(shape) == 1 and isinstance(shape[0], (tuple, list, torch.Size)): + shape = shape[0] + kwargs["device"] = "cpu" + return torch.randn(*(dimension(size) for size in shape), **kwargs) + + class Gather: + def __call__(self, rows, ti, tags): + return provider_adaln_row_gather(rows.chunk(6, dim=-1), ti, tags) + + forward_fp32 = __call__ + + class Norm: + def __call__(self, x, weight): + return F.rms_norm(x, (x.shape[-1],), weight, 1e-5) + + forward_modulated = staticmethod(provider_norm_modulate) + + gather = Gather() + ops = {"adaln_row_gather": gather, "h3_rmsnorm": Norm()} + registry = SimpleNamespace( + get_op=lambda name, device: ops[name], + _priority_map={"cpu": {"adaln_row_gather": ["gather"]}}, + _get_or_create_backend=lambda _backend: gather, + ) + monkeypatch.setattr(h3_report, "torch", CpuTorch()) + monkeypatch.setattr( + h3_report, + "h3_packed_layout", + lambda seq, num, seed: h3_packed_layout(dimension(seq), num, seed=seed, device="cpu"), + ) + + def norm_inputs(seq, seed=0): + generator = torch.Generator().manual_seed(seed) + weight = h3_report.h3_params(["transformer_blocks.0.norm1.weight"], [(5376,)])[0] + shift, scale = [torch.randn(9, 4, generator=generator).bfloat16() for _ in range(2)] + x = torch.randn(1, dimension(seq), 4, generator=generator).bfloat16() + ti, tags = h3_packed_layout(dimension(seq), 3, seed=seed, device="cpu") + return x, weight.bfloat16(), shift, scale, ti * 3 + tags + + monkeypatch.setattr(h3_report, "_norm_inputs", norm_inputs) + return registry + + +def test_gather_without_weights_preserves_op_results(small_cpu_report, monkeypatch): + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + + def unexpected_load(*args, **kwargs): + pytest.fail("an unset checkpoint must not be loaded") + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", unexpected_load) + result = h3_report._gather_accuracy(small_cpu_report) + assert all(result["forward_bitwise_vs_index_select"].values()) + assert set(result["op_backward"]) == {"cuda", "provider"} + assert all(row["repeat_bitwise_equal"] for row in result["op_backward"].values()) + assert result["chain_backward"] == [] + assert WEIGHTS_ENV in result["chain_backward_skipped"] + + +@pytest.mark.parametrize("configured", [False, True]) +def test_norm_reports_weight_source(small_cpu_report, monkeypatch, tmp_path, configured): + if configured: + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + else: + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + calls = [] + + def load_weights(device, names): + assert configured, "an unset checkpoint must not be loaded" + calls.append(names) + return {name: torch.ones(4, dtype=torch.bfloat16) for name in names} + + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", load_weights) + result = h3_report._norm_accuracy(small_cpu_report) + assert result["weight_source"] == ("pinned_checkpoint" if configured else "synthetic") + assert len(result["plain_bitwise_vs_nn_rmsnorm"]) == 4 + assert all(result["plain_bitwise_vs_nn_rmsnorm"].values()) + assert result["modulated_bitwise_vs_diffusers"] + assert result["rows_batch_invariant"] + assert set(result["backward"]) == {"cuda", "provider"} + assert bool(calls) is configured + + +def test_gather_with_weights_runs_chain(small_cpu_report, monkeypatch, tmp_path): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + weights = {"weight": torch.ones(1)} + monkeypatch.setattr(h3_report, "load_h3_conditioning_weights", lambda device: weights) + + def run_chain(registry, loaded, **case): + assert registry is small_cpu_report + assert loaded is weights + return case + + monkeypatch.setattr(h3_chain, "run_backward_case", run_chain) + result = h3_report._gather_accuracy(small_cpu_report) + assert len(result["chain_backward"]) == 5 + assert result["chain_backward"][-1] == {"num_timesteps": 4, "seq_len": 32768} + assert "chain_backward_skipped" not in result + + +@pytest.mark.parametrize("accuracy", [h3_report._gather_accuracy, h3_report._norm_accuracy]) +def test_configured_missing_weights_raise(small_cpu_report, monkeypatch, tmp_path, accuracy): + monkeypatch.setenv(WEIGHTS_ENV, str(tmp_path)) + with pytest.raises(FileNotFoundError, match="missing; run tools/weights/prepare_h3_weights.py"): + accuracy(small_cpu_report) + + +def test_plot_gather_without_chain(small_cpu_report, monkeypatch, tmp_path): + pytest.importorskip("matplotlib") + monkeypatch.delenv(WEIGHTS_ENV, raising=False) + script = ( + Path(__file__).resolve().parents[3] + / "tools" + / "validation" + / "models" + / "plot_h3_evidence.py" + ) + spec = importlib.util.spec_from_file_location("plot_h3_evidence_test", script) + plot = importlib.util.module_from_spec(spec) + spec.loader.exec_module(plot) + accuracy = h3_report._gather_accuracy(small_cpu_report) + report = { + "environment": {"gpu": "CPU fixture", "torch": torch.__version__, "cuda": "none"}, + "rl_kernel_commit": "testcommit", + "model_revision": "pinned_revision", + "accuracy": accuracy, + "perf": [ + { + "case": "S=160", + "candidate_us": 1, + "provider_us": 2, + "candidate_backward_us": 3, + "provider_backward_us": 4, + } + ], + } + figure = plot.plot_gather(report) + output = tmp_path / "figure.png" + figure.savefig(output) + assert output.stat().st_size > 0 + assert figure.axes[1].texts[0].get_text() == accuracy["chain_backward_skipped"] + plot.plt.close(figure) + + +def test_measure_interleaves_backends_and_records_samples(monkeypatch): + """Balance candidate/provider timing order and preserve samples and derived bandwidth.""" + + elapsed_us = 0.0 + calls = [] + + class Event: + def __init__(self, *, enable_timing): + """Require timing-enabled events in the deterministic timing stub.""" + + assert enable_timing + + def record(self): + """Capture the synthetic elapsed clock at this event.""" + + self.timestamp = elapsed_us + + def synchronize(self): + """Leave synchronization inert because the timing stub has no pending GPU work.""" + + pass + + def elapsed_time(self, end): + """Return elapsed synthetic event time in CUDA milliseconds.""" + + return (end.timestamp - self.timestamp) / 1e3 + + def candidate(): + """Record a candidate call and advance the synthetic clock by two microseconds.""" + + nonlocal elapsed_us + calls.append("candidate") + elapsed_us += 2.0 + + def provider(): + """Record a provider call and advance the synthetic clock by four microseconds.""" + + nonlocal elapsed_us + calls.append("provider") + elapsed_us += 4.0 + + monkeypatch.setattr(h3_report.torch.cuda, "Event", Event) + monkeypatch.setattr(h3_report.torch.cuda, "synchronize", lambda: None) + monkeypatch.setattr(h3_report, "peak_mib", lambda fn, **kwargs: 0.0) + report = h3_report.measure( + {"bytes": 1000, "candidate": candidate, "provider": provider}, warmup=1, iters=4 + ) + assert calls == [ + "candidate", + "provider", # warmup + "candidate", + "provider", + "provider", + "candidate", + "candidate", + "provider", + "provider", + "candidate", + ] + assert report["timing_order"] == [ + ["candidate", "provider"], + ["provider", "candidate"], + ["candidate", "provider"], + ["provider", "candidate"], + ] + assert report["timing_samples_us"] == {"candidate": [2.0] * 4, "provider": [4.0] * 4} + assert report["candidate_us"] == pytest.approx(2.0) + assert report["provider_us"] == pytest.approx(4.0) + assert report["candidate_gbps"] == pytest.approx(0.5) + assert report["provider_gbps"] == pytest.approx(0.25) diff --git a/tests/models/minimax_h3/test_h3_rmsnorm.py b/tests/models/minimax_h3/test_h3_rmsnorm.py new file mode 100644 index 000000000..ced9a9377 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_rmsnorm.py @@ -0,0 +1,302 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``h3_rmsnorm``: RMSNorm (eps 1e-5) and its fused AdaLN modulation. + +* forward is bitwise equal to ``nn.RMSNorm`` and to diffusers' + ``norm(x) * (1.0 + scale[i]) + shift[i]`` (block: ``adaln_indices``; + ``norm_out``: ``timestep_indices``) on the pinned norm weights; +* rows are independent of batch size and position; +* backward is deterministic and checked against an FP64 golden; +* malformed inputs fail closed. +""" + +from __future__ import annotations + +import pytest +import torch + +from rl_engine.reference.minimax_h3.rmsnorm import NativeH3RMSNormOp +from rl_engine.validation.models.h3_cases import h3_packed_layout +from rl_engine.validation.models.h3_provider import provider_norm_modulate + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NORMS = ( + "transformer_blocks.0.norm1.weight", + "transformer_blocks.0.norm2.weight", + "token_refiner.final_norm.weight", + "norm_out.norm.weight", +) + + +def _cuda_op(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import ( + H3RMSNormCudaOp, + h3_rmsnorm_available, + ) + + if not h3_rmsnorm_available(): + pytest.skip("rl_engine._C lacks h3_rmsnorm_*") + return H3RMSNormCudaOp() + + +def _x(shape, dtype=torch.bfloat16, seed=0, scale=2.0, device="cuda"): + g = torch.Generator(device="cpu").manual_seed(seed) + return (torch.randn(shape, generator=g) * scale).to(device=device, dtype=dtype) + + +def _table(rows, hidden, chunks, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + table = (torch.randn(rows, chunks * hidden, generator=g) * 0.5).to(dtype).cuda() + return table.chunk(chunks, dim=-1) # strided (R, H) views, like the AdaLN outputs + + +def _fp64_modulated(x, weight, shift, scale, index, eps=1e-5): + x64 = x.double() + n = x64 * torch.rsqrt(x64.square().mean(-1, keepdim=True) + eps) * weight.double() + return n * (1 + scale.double().index_select(0, index)) + shift.double().index_select(0, index) + + +class TestReference: + def test_golden_close_to_provider(self): + x, w = _x((3, 64), torch.float32, device="cpu"), torch.rand(64) + 0.5 + op = NativeH3RMSNormOp() + torch.testing.assert_close(op.forward(x, w), op.forward_fp32(x, w), atol=1e-6, rtol=1e-6) + + def test_modulated_with_prevalidated_indices(self): + x, w = _x((2, 8), torch.float32, device="cpu"), torch.ones(8) + shift, scale = torch.randn(3, 8), torch.randn(3, 8) + index = torch.tensor([2, 0]) + op = NativeH3RMSNormOp() + expected = op.forward(x, w) * (1 + scale[index]) + shift[index] + torch.testing.assert_close( + op.forward_modulated(x, w, shift, scale, index, check_range=False), + expected, + ) + + def test_rejects_bad_inputs(self): + op = NativeH3RMSNormOp() + x, w = torch.randn(2, 8), torch.ones(8) + with pytest.raises(TypeError): + op.forward(x, w.double()) + with pytest.raises(TypeError): + op.forward(x, w.bfloat16()) + with pytest.raises(ValueError): + op.forward(x, torch.ones(7)) + with pytest.raises(ValueError): + op.forward(x, w, eps=0.0) + shift, scale = torch.zeros(3, 8), torch.zeros(3, 8) + with pytest.raises(IndexError): + op.forward_modulated(x, w, shift, scale, torch.tensor([0, 3])) + with pytest.raises(ValueError): # one index per position S + op.forward_modulated(x, w, shift, scale, torch.tensor([0])) + + +@requires_cuda +class TestCudaForward: + @pytest.mark.parametrize("name", NORMS) + def test_bitwise_equal_to_nn_rmsnorm_on_pinned_weights(self, h3_weights_cpu, name): + weight = h3_weights_cpu[name].cuda() + module = torch.nn.RMSNorm(5376, eps=1e-5).cuda().bfloat16() + module.weight.data.copy_(weight) + x = _x((2, 777, 5376), seed=1) + assert torch.equal(_cuda_op()(x, weight), module(x)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) + @pytest.mark.parametrize("hidden", [12, 256, 4096, 5376]) + def test_bitwise_equal_to_f_rms_norm(self, dtype, hidden): + x = _x((5, 33, hidden), dtype, seed=hidden, scale=30.0) + w = (torch.rand(hidden, device="cuda") * 2).to(dtype) + ref = torch.nn.functional.rms_norm(x, (hidden,), w, 1e-5) + assert torch.equal(_cuda_op()(x, w), ref) + + def test_block_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["transformer_blocks.0.norm1.weight"].cuda() + shift, scale, *_ = _table(9, 5376, 6, seed=2) # (3T, H) rows, T = 3 + ti, tags = h3_packed_layout(4097, 3, seed=2) + index = ti * 3 + tags + x = _x((2, 4097, 5376), seed=2) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, index)) + + def test_norm_out_modulation_bitwise_equal_to_diffusers(self, h3_weights_cpu): + weight = h3_weights_cpu["norm_out.norm.weight"].cuda() + shift, scale = _table(3, 5376, 2, seed=3) # (T, H) rows indexed by timestep + ti, _ = h3_packed_layout(1000, 3, seed=3) + x = _x((1, 1000, 5376), seed=3) + ours = _cuda_op().forward_modulated(x, weight, shift, scale, ti) + assert torch.equal(ours, provider_norm_modulate(x, weight, shift, scale, ti)) + + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_modulation_bitwise_in_every_dtype(self, dtype): + """FP32 must not contract the separately rounded eager ops into an FMA.""" + + w = (torch.rand(5376, device="cuda") + 0.5).to(dtype) + shift, scale, *_ = _table(9, 5376, 6, dtype=dtype, seed=9) + ti, tags = h3_packed_layout(333, 3, seed=9) + index = ti * 3 + tags + x = _x((2, 333, 5376), dtype, seed=9) + ours = _cuda_op().forward_modulated(x, w, shift, scale, index) + assert torch.equal(ours, provider_norm_modulate(x, w, shift, scale, index)) + + def test_rows_are_batch_and_position_invariant(self): + w = (torch.rand(5376, device="cuda") + 0.5).bfloat16() + shift, scale, *_ = _table(9, 5376, 6, seed=4) + ti, tags = h3_packed_layout(300, 3, seed=4) + index = ti * 3 + tags + x = _x((3, 300, 5376), seed=4) + full = _cuda_op().forward_modulated(x, w, shift, scale, index) + part = _cuda_op().forward_modulated(x[2:3, 50:90], w, shift, scale, index[50:90]) + assert torch.equal(part[0], full[2, 50:90]) + + def test_rejects_unsupported(self): + w = torch.ones(6, device="cuda", dtype=torch.bfloat16) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(_x((2, 6)), w) + with pytest.raises(ValueError): + _cuda_op()(_x((2, 8)).cpu(), torch.ones(8).bfloat16()) + + +_DEVICE_MESSAGES = { + "weight": "with x's dtype and device", + "index": r"row index must be a non-empty contiguous int64 tensor on cuda:0", + "rstd": "on x's device", +} + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +class TestCudaNativeDeviceValidation: + @pytest.mark.parametrize("argument", ["weight", "index"]) + def test_forward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + if argument == "weight": + weight = weight.to("cuda:1") + else: + index = index.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + @pytest.mark.parametrize("argument", ["weight", "index", "rstd"]) + def test_backward_rejects_different_device(self, argument): + from rl_engine.backends.extension import _C + + _cuda_op() + x = torch.ones(2, 8, device="cuda:0", dtype=torch.bfloat16) + weight = torch.ones(8, device="cuda:0", dtype=x.dtype) + shift, scale = torch.zeros(2, 16, device="cuda:0", dtype=x.dtype).chunk(2, dim=-1) + index = torch.tensor([0, 1], device="cuda:0") + _, rstd = _C.h3_rmsnorm_forward(x, weight, 1e-5, shift, scale, index) + if argument == "weight": + weight = weight.to("cuda:1") + elif argument == "index": + index = index.to("cuda:1") + else: + rstd = rstd.to("cuda:1") + with pytest.raises(RuntimeError, match=_DEVICE_MESSAGES[argument]): + _C.h3_rmsnorm_backward(torch.ones_like(x), x, weight, rstd, shift, scale, index) + torch.cuda.synchronize(0) + torch.cuda.synchronize(1) + + +@requires_cuda +class TestCudaBackward: + def _grads(self, fn, tensors, grad): + leaves = [t.detach().clone().requires_grad_(True) for t in tensors] + fn(*leaves).backward(grad) + return [leaf.grad for leaf in leaves] + + def test_modulated_backward_against_fp64(self): + g = torch.Generator(device="cpu").manual_seed(5) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.5).bfloat16().cuda() + ti, tags = h3_packed_layout(2049, 3, seed=5) + index = ti * 3 + tags + x = _x((2, 2049, 5376), seed=5) + grad = _x((2, 2049, 5376), seed=6, scale=1.0) + op = _cuda_op() + + def ours(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return op.forward_modulated(x_, w_, shift, scale, index) + + def golden(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return _fp64_modulated(x_, w_, shift, scale, index) + + got = self._grads(ours, (x, w, table), grad) + ref = self._grads(golden, (x.double(), w.double(), table.double()), grad.double()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + rel = ((g_.double() - r).abs().max() / r.abs().max()).item() + # Two BF16 roundings remain: the eager VJP uses round(1 + scale), and + # every gradient is rounded once to BF16 (2^-9 relative each). + assert rel < 1e-2, (name, rel) + again = self._grads(ours, (x, w, table), grad) + assert all(torch.equal(a, b) for a, b in zip(got, again)) + + @pytest.mark.parametrize("seq", [257, 1024, 4097]) + def test_modulated_gradients_meet_contract_against_golden(self, seq): + # Against forward_modulated_fp32 before it stored norm(x) and 1 + scale + # in BF16, d_weight missed from S=1024 and d_scale at S=4097. + from rl_engine.contracts.numerical import load_contract, resolve_tolerance + + spec = resolve_tolerance( + load_contract(), + judgment="gradient_accuracy", + op_class="reduction", + dtype=torch.bfloat16, + ) + g = torch.Generator(device="cpu").manual_seed(seq) + w = (torch.rand(5376, generator=g) + 0.5).bfloat16().cuda() + table = (torch.randn(9, 2 * 5376, generator=g) * 0.1).bfloat16().cuda() + ti, tags = h3_packed_layout(seq, 3, seed=seq) + index = ti * 3 + tags + x = _x((1, seq, 5376), seed=seq) + grad = _x((1, seq, 5376), seed=seq + 1, scale=1.0) + op, native = _cuda_op(), NativeH3RMSNormOp() + + def run(fn, upstream): + def call(x_, w_, t_): + shift, scale = t_.chunk(2, dim=-1) + return fn(x_, w_, shift, scale, index) + + return self._grads(call, (x, w, table), upstream) + + got = run(op.forward_modulated, grad) + ref = run(native.forward_modulated_fp32, grad.float()) + for name, g_, r in zip(("dx", "dweight", "dtable"), got, ref): + torch.testing.assert_close( + g_.float(), r.float(), atol=spec.atol, rtol=spec.rtol, msg=lambda m: f"{name}: {m}" + ) + + def test_plain_backward_against_fp64_and_row_local_dx(self): + w = (torch.rand(5376, generator=torch.Generator().manual_seed(7)) + 0.5).bfloat16().cuda() + x = _x((4, 64, 5376), seed=7) + grad = _x((4, 64, 5376), seed=8, scale=1.0) + op = _cuda_op() + got = self._grads(lambda a, b: op(a, b), (x, w), grad) + + def golden(a, b): + return a * torch.rsqrt(a.square().mean(-1, keepdim=True) + 1e-5) * b + + ref = self._grads(golden, (x.double(), w.double()), grad.double()) + for g, r in zip(got, ref): + assert ((g.double() - r).abs().max() / r.abs().max()).item() < 5e-3 + part = self._grads(lambda a, b: op(a, b), (x[1:2, 10:20], w), grad[1:2, 10:20]) + assert torch.equal(part[0][0], got[0][1, 10:20]) + + def test_registry_dispatches_cuda(self): + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + assert ( + type(KernelRegistry().get_op("h3_rmsnorm", device="cuda")).__name__ == "H3RMSNormCudaOp" + ) diff --git a/tests/models/minimax_h3/test_h3_sp_norm_adaln.py b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py new file mode 100644 index 000000000..a60489a04 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_sp_norm_adaln.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``sp_norm_adaln``: sequence-parallel norm / modulation / gated residual. + +* ownership: contiguous position shards per batch item; bad layouts fail closed; +* backward (one GPU, ranks as threads around the plain backward functions): every + rank's row gradients equal WS1's rows, and its ``dweight``/``d_shift``/``d_scale``/ + ``d_gate`` equal WS1's, for SP 2..8, odd S, B > 1, interleaved modality rows and + bf16/fp32; +* end to end (real NCCL, autograd): the MLP-side block region + ``norm2(residual + gate_msa * y)`` with modulation is byte-equal to WS1 on every rank. +""" + +from __future__ import annotations + +import os +import threading +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + SPRowLayout, + sp_row_layout, +) +from rl_engine.validation.models.h3_ws2 import ( + run_world, + sp_case, + sp_matches_ws1, + sp_region_rank, + ws1_sp_region, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +ALL_TRUE = dict.fromkeys(("out", "d_residual", "d_y", "d_weight", "d_table"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_norm_adaln_available, + ) + + if not sp_norm_adaln_available(): + pytest.skip("rl_engine._C lacks the H3 SP symbols") + + +class _ThreadGroup: + """Ranks as threads on one GPU; ``all_gather`` is the rank-ordered concatenation. + + Only for the plain backward functions: autograd runs every backward of a device + on one worker thread, so autograd-driven ranks need real processes. + """ + + def __init__(self, world: int) -> None: + self.barrier, self.slots = threading.Barrier(world), [None] * world + + def handle(self, rank: int): + def all_gather(t): + self.slots[rank] = t + self.barrier.wait() + out = torch.cat(self.slots) + self.barrier.wait() + return out + + return SimpleNamespace(rank=rank, world_size=len(self.slots), all_gather=all_gather) + + def run(self, fn): + results, errors = [None] * len(self.slots), [] + + def body(rank): + try: + results[rank] = fn(self.handle(rank)) + except BaseException as exc: # re-raised below + errors.append(exc) + self.barrier.abort() + + threads = [threading.Thread(target=body, args=(r,)) for r in range(len(self.slots))] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise next( + (e for e in errors if not isinstance(e, threading.BrokenBarrierError)), errors[0] + ) + return results + + +class TestLayout: + @pytest.mark.parametrize("seq, sp", [(8, 8), (257, 2), (1000, 3), (4097, 8)]) + def test_shards_cover_every_position_once(self, seq, sp): + bounds = [SPRowLayout(seq, 1, sp, r).bounds(r) for r in range(sp)] + assert bounds[0][0] == 0 and bounds[-1][1] == seq + assert all(a[1] == b[0] and a[0] < a[1] for a, b in zip(bounds, bounds[1:])) + + @pytest.mark.parametrize("seq, batch, sp, rank", [(3, 1, 4, 0), (8, 0, 2, 0), (8, 1, 2, 2)]) + def test_rejects(self, seq, batch, sp, rank): + with pytest.raises(ValueError): + sp_row_layout(seq, batch, sp, rank) + + +@requires_cuda +class TestBackwardRanks: + @pytest.mark.parametrize( + "sp, batch, seq, hidden, layout, dtype", + [ + (2, 1, 4097, 5376, "block", torch.bfloat16), + (2, 1, 4097, 5376, "interleaved", torch.bfloat16), + (3, 2, 1000, 256, "interleaved", torch.bfloat16), + (4, 2, 777, 512, "block", torch.float32), + (5, 1, 1301, 128, "interleaved", torch.bfloat16), + (8, 2, 2049, 256, "block", torch.bfloat16), + (8, 1, 40, 64, "interleaved", torch.float32), # shards far smaller than a tile + ], + ) + def test_ranks_reproduce_ws1(self, sp, batch, seq, hidden, layout, dtype): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + _segment_tiles, + ) + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_gate_residual_backward, + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(batch, seq, hidden, layout=layout, dtype=dtype, seed=sp * 100 + seq) + chunks = c["table"].chunk(6, dim=-1) + shift, scale, gate = chunks[3], chunks[4], chunks[2] # shift/scale_mlp, gate_msa + x2, g2, y2 = (c[k].view(-1, hidden) for k in ("residual", "grad", "y")) + idx, rows = c["index"], shift.shape[0] + tiles = _segment_tiles(idx.repeat(batch), rows) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5, shift, scale, idx) + ws1_norm = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd, shift, scale, idx, *tiles) + ws1_gate = _C.h3_gate_residual_backward(g2, y2, gate, idx, *tiles) + + def rank(coll): + lay = sp_row_layout(seq, batch, sp, coll.rank) + part = slice(lay.lo, lay.hi) + local = { + k: c[k][:, part].reshape(-1, hidden).contiguous() for k in ("residual", "grad", "y") + } + rs = rstd.view(batch, seq)[:, part].reshape(-1).contiguous() + plan = sp_plan(lay, idx, rows) + norm = sp_rmsnorm_backward( + local["grad"], local["residual"], c["weight"], rs, shift, scale, plan, coll + ) + gate_grads = sp_gate_residual_backward(local["grad"], local["y"], gate, plan, coll) + return lay, norm, gate_grads + + for lay, (dx, dw, dsh, dsc), (dy, dgate) in _ThreadGroup(sp).run(rank): + part = slice(lay.lo, lay.hi) + assert torch.equal( + dx.view(batch, -1, hidden), ws1_norm[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dw, ws1_norm[1]) + assert torch.equal(dsh, ws1_norm[2].to(dtype)) and torch.equal( + dsc, ws1_norm[3].to(dtype) + ) + assert torch.equal( + dy.view(batch, -1, hidden), ws1_gate[0].view(batch, seq, hidden)[:, part] + ) + assert torch.equal(dgate, ws1_gate[1]) + + def test_plain_norm_ranks_reproduce_ws1(self): + _require_ext() + from rl_engine import _C + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + sp_plan, + sp_rmsnorm_backward, + ) + + c = sp_case(2, 1500, 256, seed=5) + x2, g2 = c["residual"].view(-1, 256), c["grad"].view(-1, 256) + _, rstd = _C.h3_rmsnorm_forward(x2, c["weight"], 1e-5) + ws1 = _C.h3_rmsnorm_backward(g2, x2, c["weight"], rstd) + + def rank(coll): + lay = sp_row_layout(1500, 2, 4, coll.rank) + part = slice(lay.lo, lay.hi) + loc = [c[k][:, part].reshape(-1, 256).contiguous() for k in ("grad", "residual")] + rs = rstd.view(2, 1500)[:, part].reshape(-1).contiguous() + return sp_rmsnorm_backward(*loc, c["weight"], rs, None, None, sp_plan(lay), coll) + + assert all(torch.equal(dw, ws1[1]) for _, dw, _, _ in _ThreadGroup(4).run(rank)) + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import ( + H3SPNormAdaLNCudaOp, + ) + + c = sp_case(1, 100, 64) + shift, scale = c["table"].chunk(6, dim=-1)[3:5] + op = H3SPNormAdaLNCudaOp(SimpleNamespace(rank=1, world_size=2), seq_len=100, batch=1) + local = c["residual"][:, 50:] + with pytest.raises(ValueError): # the full sequence instead of this rank's rows + op.norm_modulated(c["residual"], c["weight"], shift, scale, c["index"]) + with pytest.raises(ValueError): # a shard of the index instead of the full index + op.norm_modulated(local, c["weight"], shift, scale, c["index"][50:]) + with pytest.raises(IndexError): + op.gate_residual(local, c["y"][:, 50:], shift, c["index"] + 9) + with pytest.raises(ValueError): + op.norm(local.cpu(), c["weight"].cpu()) + assert op.readback()["positions"] == [50, 100] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_region_byte_equal_to_ws1(world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + case = sp_case(2, 4097, layout="interleaved", seed=world) + ws1 = {k: v.cpu() for k, v in ws1_sp_region(**case).items()} + torch.cuda.empty_cache() + ranks = run_world(world, sp_region_rank, {k: v.cpu() for k, v in case.items()}) + assert sp_matches_ws1(ws1, ranks) == ALL_TRUE + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@requires_cuda +def test_plan_snapshots_caller_index(): + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import sp_plan + + index = torch.tensor([0, 1, 0, 1], device="cuda") + plan = sp_plan(sp_row_layout(4, 1, 1, 0), index, 2) + local, exchanged = plan.local_index.clone(), plan.index_buf.clone() + index.fill_(99) + assert torch.equal(plan.local_index, local) + assert torch.equal(plan.index_buf, exchanged) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +def test_plain_norm_uses_activation_device_and_separate_cache(): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.sp_norm_adaln import H3SPNormAdaLNCudaOp + + collective = SimpleNamespace(rank=0, world_size=1, all_gather=lambda x: x) + op = H3SPNormAdaLNCudaOp(collective, 5, 1) + with torch.cuda.device(0): + for device in ("cuda:1", "cuda:0"): + x = torch.randn(1, 5, 8, device=device, requires_grad=True) + weight = torch.randn(8, device=device, requires_grad=True) + ref_x = x.detach().clone().requires_grad_(True) + ref_w = weight.detach().clone().requires_grad_(True) + result = op.norm(x, weight) + expected = torch.nn.functional.rms_norm(ref_x, (8,), ref_w, 1e-5) + result.sum().backward() + expected.sum().backward() + torch.testing.assert_close(result, expected) + torch.testing.assert_close(x.grad, ref_x.grad) + torch.testing.assert_close(weight.grad, ref_w.grad) + assert len(op._plans) == 2 + assert op.plan(None, None, device="cuda:1").plan.send_local.device.index == 1 + assert op.plan(None, None, device="cuda:0").plan.send_local.device.index == 0 + + +def _worker_exit_without_result(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + os._exit(kwargs["exit_code"]) + time.sleep(60) + + +def _worker_reports_error(rank, world, init, queue, path, target, kwargs): + import time + + if rank == 0: + queue.put({"rank": rank, "error": "intentional worker failure"}) + time.sleep(60) + + +def _worker_never_reports(rank, world, init, queue, path, target, kwargs): + import time + + time.sleep(60) + + +@pytest.mark.parametrize("mode", ["clean_exit", "crash", "reported_error", "timeout"]) +def test_run_world_cleans_up_failed_workers(monkeypatch, mode): + import time + + import torch.multiprocessing as mp + + from rl_engine.validation.models import h3_ws2 + + workers = { + "clean_exit": _worker_exit_without_result, + "crash": _worker_exit_without_result, + "reported_error": _worker_reports_error, + "timeout": _worker_never_reports, + } + monkeypatch.setattr(h3_ws2, "_bootstrap", workers[mode]) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + expected = TimeoutError if mode == "timeout" else RuntimeError + with pytest.raises(expected): + run_world( + 2, + None, + {}, + timeout=1 if mode == "timeout" else 30, + exit_code=7 if mode == "crash" else 0, + ) + assert time.monotonic() - start < 20 + assert {p.pid for p in mp.active_children()} <= before + + +def _worker_reports_result(rank, world, init, queue, path, target, kwargs): + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + + +def test_run_world_returns_rank_ordered_results(monkeypatch): + from rl_engine.validation.models import h3_ws2 + + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker_reports_result) + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + results = run_world(2, None, {}, timeout=30) + assert [result["rank"] for result in results] == [0, 1] + assert [result["value"].item() for result in results] == [0, 1] diff --git a/tests/models/minimax_h3/test_h3_timestep_mlp.py b/tests/models/minimax_h3/test_h3_timestep_mlp.py new file mode 100644 index 000000000..555a28ebe --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_mlp.py @@ -0,0 +1,296 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_mlp_fp32``: FP32 256 -> 5376 -> 2688 timestep MLP. + +* accuracy: forward and every gradient against an FP64 golden, on the pinned + checkpoint weights and on synthetic non-H3 sizes; +* invariance: a timestep's ``temb`` bytes (and its row-local ``dx``) do not + depend on how many timesteps share the call or where it sits; +* determinism: repeated forward/backward runs are bitwise equal; +* contract: BF16 inputs (RFC probe H7) and malformed shapes fail closed. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.reference.minimax_h3.timestep_mlp import NativeH3TimestepMLPOp +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +NAMES = ("x", "w1", "b1", "w2", "b2") + + +def _cuda_op(): + """Construct the CUDA MLP operator, skipping missing native linear kernels.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import H3TimestepMLPCudaOp + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + return H3TimestepMLPCudaOp() + + +def _features(num, device="cuda", seed=0): + """Generate seeded H3 sinusoidal features on the requested test device.""" + + return NativeH3TimestepSinusoidOp().forward(h3_timesteps(num, seed=seed, device=device)) + + +def _synthetic_params(k_in, hidden, out, device="cuda", seed=0): + """Create seeded FP32 parameters for a two-layer SiLU MLP of the requested sizes.""" + + g = torch.Generator(device="cpu").manual_seed(seed) + w1 = torch.randn(hidden, k_in, generator=g) / k_in**0.5 + b1 = torch.randn(hidden, generator=g) * 0.1 + w2 = torch.randn(out, hidden, generator=g) / hidden**0.5 + b2 = torch.randn(out, generator=g) * 0.1 + return [t.to(device) for t in (w1, b1, w2, b2)] + + +def _real_params(weights): + """Return pinned timestep MLP weights and biases on CUDA in layer order.""" + + return [ + weights[f"time_embedder.linear_{i}.{p}"].cuda() for i in (1, 2) for p in ("weight", "bias") + ] + + +def _fp64_grads(x, params, grad_out): + """Compute gradients of both linear layers and SiLU using FP64 autograd.""" + + leaves = [t.detach().double().requires_grad_(True) for t in (x, *params)] + z = F.linear(leaves[0], leaves[1], leaves[2]) + F.linear(z * torch.sigmoid(z), leaves[3], leaves[4]).backward(grad_out.double()) + return [leaf.grad for leaf in leaves] + + +class TestReference: + def test_provider_and_golden_agree(self): + """Keep the FP32 provider close to the high-precision MLP replay on CPU.""" + + x = _features(5, device="cpu") + params = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + torch.testing.assert_close( + op.forward(x, *params), op.forward_fp32(x, *params), atol=1e-5, rtol=1e-5 + ) + + @pytest.mark.parametrize("which", range(5)) + def test_rejects_bf16_anywhere(self, which): + # RFC probe H7: casting a declared FP32 path to BF16 early. + """Reject BF16 in every argument of the declared FP32 MLP path.""" + + args = [_features(2, device="cpu"), *_synthetic_params(256, 64, 32, device="cpu")] + args[which] = args[which].bfloat16() + with pytest.raises(TypeError, match="float32"): + NativeH3TimestepMLPOp().forward(*args) + + def test_rejects_shape_mismatches(self): + """Reject incompatible layer dimensions, malformed biases, and empty feature batches.""" + + x = _features(2, device="cpu") + w1, b1, w2, b2 = _synthetic_params(256, 64, 32, device="cpu") + op = NativeH3TimestepMLPOp() + with pytest.raises(ValueError): + op.forward(x[:, :128], w1, b1, w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1[:10], w2, b2) + with pytest.raises(ValueError): + op.forward(x, w1, b1, w2[:, :10], b2) + with pytest.raises(ValueError): + op.forward(x[:0], w1, b1, w2, b2) + + +@requires_cuda +class TestCudaRealWeights: + @pytest.mark.parametrize("num", [1, 2, 3, 4, 7]) + def test_forward_against_fp64_golden(self, h3_weights_cpu, num): + """Check checkpoint forward accuracy, FP32 dtype, and the timestep output shape.""" + + params = _real_params(h3_weights_cpu) + x = _features(num, seed=num) + out = _cuda_op()(x, *params) + gold = NativeH3TimestepMLPOp().forward_fp32(x, *params) + assert out.dtype == torch.float32 and out.shape == (num, 2688) + # tolerance_contract.json forward_accuracy / reduction / float32. + torch.testing.assert_close(out, gold, atol=1e-4, rtol=1e-4) + assert (out - gold).abs().max().item() < 1e-5 + + def test_rows_are_batch_and_position_invariant(self, h3_weights_cpu): + """Keep each checkpoint embedding row bitwise unchanged across batch sizes and positions.""" + + op = _cuda_op() + params = _real_params(h3_weights_cpu) + x_all = _features(9, seed=1) + full = op(x_all, *params) + for i in range(9): + assert torch.equal(op(x_all[i : i + 1], *params)[0], full[i]) + for num in (2, 3, 4, 5, 8): + for pos in (0, num - 1): + x = _features(num, seed=10 + num) + x[pos] = x_all[4] + assert torch.equal(op(x, *params)[pos], full[4]), (num, pos) + + def test_backward_against_fp64_golden(self, h3_weights_cpu): + """Match all five checkpoint input gradients to the FP64 autograd golden.""" + + params = _real_params(h3_weights_cpu) + x = _features(3, seed=2) + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + grad_out = torch.randn( + 3, 2688, device="cuda", generator=torch.Generator("cuda").manual_seed(0) + ) + _cuda_op()(*leaves).backward(grad_out) + for name, leaf, ref in zip(NAMES, leaves, _fp64_grads(x, params, grad_out), strict=True): + # tolerance_contract.json gradient_accuracy / reduction / float32. + torch.testing.assert_close(leaf.grad.double(), ref, atol=1e-4, rtol=1e-4, msg=name) + + +@requires_cuda +class TestCudaSynthetic: + @pytest.mark.parametrize("k_in, hidden, out", [(8, 24, 16), (12, 40, 20), (256, 5376, 2688)]) + def test_other_sizes(self, k_in, hidden, out): + """Match the high-precision forward golden for small and checkpoint-sized MLPs.""" + + params = _synthetic_params(k_in, hidden, out, seed=k_in) + x = torch.randn(5, k_in, device="cuda") + torch.testing.assert_close( + _cuda_op()(x, *params), + NativeH3TimestepMLPOp().forward_fp32(x, *params), + atol=1e-4, + rtol=1e-4, + ) + + def test_zero_hidden_size_matches_cpu_forward_backward(self): + """Match CPU forward values and all input gradients with a zero-width hidden layer.""" + + generator = torch.Generator(device="cpu").manual_seed(0) + cpu_inputs = [ + torch.randn(3, 8, generator=generator), + torch.empty(0, 8), + torch.empty(0), + torch.empty(16, 0), + torch.linspace(-1.0, 1.0, 16), + ] + cpu_inputs = [tensor.requires_grad_(True) for tensor in cpu_inputs] + cuda_inputs = [tensor.detach().cuda().requires_grad_(True) for tensor in cpu_inputs] + cpu_out = NativeH3TimestepMLPOp()(*cpu_inputs) + cuda_out = _cuda_op()(*cuda_inputs) + torch.testing.assert_close(cuda_out.cpu(), cpu_out, atol=0, rtol=0) + grad = torch.randn(3, 16, generator=generator) + cpu_out.backward(grad) + cuda_out.backward(grad.cuda()) + for name, cpu_input, cuda_input in zip(NAMES, cpu_inputs, cuda_inputs, strict=True): + torch.testing.assert_close( + cuda_input.grad.cpu(), cpu_input.grad, atol=0, rtol=0, msg=name + ) + + def test_rejects_unaligned_k(self): + """Reject inner dimensions that violate the native FP32 four-element alignment.""" + + params = _synthetic_params(6, 8, 8) + with pytest.raises(RuntimeError, match="multiple of 4"): + _cuda_op()(torch.randn(2, 6, device="cuda"), *params) + + def test_rejects_cpu_and_mixed_devices(self): + """Reject CPU-only and mixed-device inputs at the CUDA MLP boundary.""" + + params = _synthetic_params(8, 24, 16) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *[p.cpu() for p in params]) + with pytest.raises(ValueError): + _cuda_op()(torch.randn(2, 8), *params) + + def test_repeat_and_backward_are_deterministic(self): + """Require repeated MLP outputs and all input gradients to be bitwise equal.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=3) + x = _features(4, seed=3) + grad_out = torch.randn(4, 2688, device="cuda") + runs = [] + for _ in range(3): + leaves = [t.detach().clone().requires_grad_(True) for t in (x, *params)] + out = op(*leaves) + out.backward(grad_out) + runs.append([out.detach(), *[leaf.grad for leaf in leaves]]) + for later in runs[1:]: + for a, b in zip(runs[0], later, strict=True): + assert torch.equal(a, b) + + def test_dx_rows_are_batch_invariant(self): + """Keep feature gradients bitwise unchanged when rows are evaluated independently.""" + + op = _cuda_op() + params = _synthetic_params(256, 5376, 2688, seed=4) + x = _features(6, seed=4) + grad_out = torch.randn(6, 2688, device="cuda") + + def dx(rows): + """Compute selected feature-row gradients using matching upstream rows.""" + + leaf = x[rows].detach().clone().requires_grad_(True) + op(leaf, *params).backward(grad_out[rows]) + return leaf.grad + + full = dx(slice(0, 6)) + for i in range(6): + assert torch.equal(dx(slice(i, i + 1))[0], full[i]) + + def test_registry_dispatches_cuda(self): + """Resolve the timestep MLP registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_mlp_fp32", device="cuda") + assert type(op).__name__ == "H3TimestepMLPCudaOp" + + +@requires_cuda +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="needs two CUDA devices") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +class TestCudaDetLinearDeviceContract: + def test_forward_rejects_bias_on_another_device(self, dtype): + """Reject a projection bias on a different CUDA device from its weight.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_forward + + _cuda_op() + x = torch.randn(2, 8, device="cuda:0", dtype=dtype) + weight = torch.randn(16, 8, device="cuda:0", dtype=dtype) + bias = torch.randn(16, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="bias and weight must be on the same device"): + det_linear_forward(x, weight, bias) + + def test_backward_input_rejects_weight_on_another_device(self, dtype): + """Reject native input-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_input, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + weight = torch.randn(16, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and weight must be on the same device"): + det_linear_backward_input(grad, weight, dtype) + + def test_backward_weight_rejects_input_on_another_device(self, dtype): + """Reject native parameter-gradient operands on different CUDA devices.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import ( + det_linear_backward_weight, + ) + + _cuda_op() + grad = torch.randn(2, 16, device="cuda:0") + x = torch.randn(2, 8, device="cuda:1", dtype=dtype) + with pytest.raises(RuntimeError, match="grad and x must be on the same device"): + det_linear_backward_weight(grad, x, dtype) diff --git a/tests/models/minimax_h3/test_h3_timestep_sinusoid.py b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py new file mode 100644 index 000000000..277beadea --- /dev/null +++ b/tests/models/minimax_h3/test_h3_timestep_sinusoid.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``timestep_sinusoid_h3``: FP32 [cos | sin] timestep features. + +Three properties, each checked separately: + +* provider parity: the CUDA kernel is bitwise equal to the diffusers + ``get_timestep_embedding`` replay on the same device (an elementwise op + with the same FP32 operation sequence); +* accuracy: within the FP32 elementwise contract of an FP64 golden; +* invariance: a row's bytes do not depend on how many timesteps share the + launch, their order, or repeated runs. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp +from rl_engine.validation.models.h3_cases import h3_timesteps +from rl_engine.validation.models.h3_provider import provider_time_proj + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") + + +def _cuda_op(): + """Construct the CUDA sinusoid operator, skipping a missing native forward entrypoint.""" + + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + h3_sinusoid_cuda_available, + ) + + if not h3_sinusoid_cuda_available(): + pytest.skip("rl_engine._C lacks h3_timestep_sinusoid_forward") + return H3TimestepSinusoidCudaOp() + + +def _dense_unit_sweep() -> torch.Tensor: + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + edges = torch.tensor([0x00000001, 0x007FFFFF, 0x00800000, 0x3F800000], dtype=torch.int32) + return torch.cat([bits, edges.cuda()]).view(torch.float32) + + +class TestReference: + def test_layout_is_cos_then_sin(self): + """Check FP32 cosine-then-sine layout and the unit-frequency boundary values.""" + + t = torch.tensor([0.0, 1.0]) + out = NativeH3TimestepSinusoidOp().forward(t) + assert out.shape == (2, 256) and out.dtype == torch.float32 + assert torch.equal(out[0, :128], torch.ones(128)) + assert torch.equal(out[0, 128:], torch.zeros(128)) + # k = 0 has frequency 1: cos(1), sin(1) at t = 1. + assert out[1, 0].item() == pytest.approx(math.cos(1.0), abs=1e-7) + assert out[1, 128].item() == pytest.approx(math.sin(1.0), abs=1e-7) + + def test_provider_path_matches_diffusers_cpu(self): + """Require CPU sinusoid values to match the diffusers provider bitwise.""" + + t = h3_timesteps(9, device="cpu") + ours = NativeH3TimestepSinusoidOp().forward(t) + assert torch.equal(ours, provider_time_proj(t, 256)) + + def test_golden_close_to_provider_path(self): + """Keep the CPU provider within FP32 error of the high-precision sinusoid golden.""" + + t = h3_timesteps(33, device="cpu") + op = NativeH3TimestepSinusoidOp() + torch.testing.assert_close(op.forward(t), op.forward_fp32(t), atol=1e-6, rtol=0) + + @pytest.mark.parametrize( + "bad, error", + [ + (torch.tensor([[0.5]]), ValueError), # not 1-D + (torch.empty(0), ValueError), # empty + (torch.tensor([1, 0]), TypeError), # integer + (torch.tensor([0.5, 1.5]), ValueError), # H10: t outside [0, 1] + (torch.tensor([500.0]), ValueError), # H10: t * 1000 convention + (torch.tensor([float("nan")]), ValueError), + ], + ) + def test_rejects_invalid_timesteps(self, bad, error): + """Reject empty, malformed, non-floating, non-finite, or out-of-range timesteps.""" + + with pytest.raises(error): + NativeH3TimestepSinusoidOp().forward(bad) + + def test_rejects_odd_channels(self): + """Reject odd feature counts that cannot split into equal cosine and sine halves.""" + + with pytest.raises(ValueError): + NativeH3TimestepSinusoidOp().forward(torch.tensor([0.5]), num_channels=255) + + +@requires_cuda +class TestCuda: + @pytest.mark.parametrize("num", [1, 2, 3, 7, 64, 1000, 4097]) + def test_bitwise_equal_to_diffusers_cuda_path(self, num): + """Require CUDA sinusoid values to match both provider replays for varied batch sizes.""" + + t = h3_timesteps(num, seed=num) + out = _cuda_op()(t) + assert out.dtype == torch.float32 and out.shape == (num, 256) + assert torch.equal(out, provider_time_proj(t, 256)) + assert torch.equal(out, NativeH3TimestepSinusoidOp().forward(t)) + + def test_bitwise_over_dense_fp32_sweep_of_unit_interval(self): + # Every 997th FP32 bit pattern in [0, 1] plus the subnormals' edge: a + # fast-math build (approximate __expf/__sinf/__cosf, FTZ) fails here. + t = _dense_unit_sweep() + assert torch.equal(_cuda_op().forward(t, check_range=False), provider_time_proj(t, 256)) + + def test_sglang_timestep_embedding_is_bitwise_equal(self): + # RFC #420 reuse rule: parity with SGLang is checked bit for bit, not + # through matching FP64 error summaries. + jit = pytest.importorskip("sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit") + t = torch.cat([_dense_unit_sweep(), h3_timesteps(4097, seed=3)]) + theirs = jit.timestep_embedding(t, 256, flip_sin_to_cos=True, downscale_freq_shift=0.0) + assert torch.equal(theirs, provider_time_proj(t, 256)) + assert torch.equal(theirs, _cuda_op().forward(t, check_range=False)) + + @pytest.mark.parametrize("num_channels", [2, 8, 96, 256, 320]) + def test_other_channel_counts_match_reference(self, num_channels): + """Preserve bitwise reference parity across supported even channel counts.""" + + t = h3_timesteps(5) + out = _cuda_op()(t, num_channels=num_channels) + ref = NativeH3TimestepSinusoidOp().forward(t, num_channels=num_channels) + assert torch.equal(out, ref) + + def test_fp32_contract_against_fp64_golden(self): + """Bound CUDA sinusoid error against the high-precision golden by FP32 precision.""" + + t = h3_timesteps(257) + out = _cuda_op()(t) + gold = NativeH3TimestepSinusoidOp().forward_fp32(t) + # tolerance_contract.json forward_accuracy / elementwise / float32. + torch.testing.assert_close(out, gold, atol=1e-5, rtol=1e-5) + assert (out - gold).abs().max().item() <= 2.0**-23 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) + def test_low_precision_timesteps_are_upcast_first(self, dtype): + """Match provider FP32 upcasting for BF16 and FP16 timesteps.""" + + t = h3_timesteps(6).to(dtype) + assert torch.equal(_cuda_op()(t), provider_time_proj(t, 256)) + + def test_batch_position_and_repeat_invariance(self): + """Keep row bits unchanged across batch positions, repeated calls, and permutations.""" + + op = _cuda_op() + base = torch.tensor(0.3141592, device="cuda") + single = op(base.view(1)) + for num in (2, 5, 64, 1023): + for pos in (0, num // 2, num - 1): + t = h3_timesteps(num, seed=num) + t[pos] = base + assert torch.equal(op(t)[pos], single[0]), (num, pos) + t = h3_timesteps(129, seed=11) + first = op(t) + for _ in range(3): + assert torch.equal(op(t), first) + perm = torch.randperm(129, generator=torch.Generator().manual_seed(0)).cuda() + assert torch.equal(op(t[perm]), first[perm]) + + def test_non_contiguous_input(self): + """Treat strided timesteps identically to their contiguous copies.""" + + t = h3_timesteps(20) + strided = t[::2] + assert not strided.is_contiguous() + assert torch.equal(_cuda_op()(strided), _cuda_op()(strided.contiguous())) + + def test_rejects_cpu_input(self): + """Reject CPU timesteps at the CUDA sinusoid boundary.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.5])) + + def test_rejects_out_of_range(self): + """Reject scaled timesteps outside the declared unit interval.""" + + with pytest.raises(ValueError): + _cuda_op()(torch.tensor([0.0, 999.0], device="cuda")) + + @pytest.mark.parametrize("bad", [-0.1, 1.1, float("nan"), float("inf"), -float("inf")]) + def test_native_entrypoint_rejects_invalid_values(self, bad): + """Reject invalid values in native and wrapper calls even when neighbors are valid.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + # Valid neighbors must not hide an invalid timestep in the native call. + t = torch.tensor([0.0, bad, 1.0], device="cuda") + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _C.h3_timestep_sinusoid_forward(t) + with pytest.raises(ValueError, match=r"finite and lie in \[0, 1\]"): + _cuda_op()(t) + + def test_native_entrypoint_accepts_boundaries(self): + """Accept both interval endpoints and match provider values in the native call.""" + + from rl_engine.backends.extension import _C + + _cuda_op() + t = torch.tensor([0.0, 0.5, 1.0], device="cuda") + assert torch.equal(_C.h3_timestep_sinusoid_forward(t), provider_time_proj(t, 256)) + + def test_trusted_path_supports_cuda_graph(self): + """Capture and replay the prevalidated sinusoid path without changing its output.""" + + op = _cuda_op() + t = h3_timesteps(7) + expected = op(t) + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + op.forward(t, check_range=False) + torch.cuda.current_stream().wait_stream(warmup) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = op.forward(t, check_range=False) + graph.replay() + assert torch.equal(out, expected) + + def test_backward_matches_autograd_reference(self): + """Match timestep gradients to the high-precision autograd sinusoid replay.""" + + t = h3_timesteps(17).requires_grad_(True) + t_ref = t.detach().clone().requires_grad_(True) + grad = torch.randn(17, 256, device="cuda", generator=torch.Generator("cuda").manual_seed(1)) + _cuda_op()(t).backward(grad) + NativeH3TimestepSinusoidOp().forward_fp32(t_ref).backward(grad) + torch.testing.assert_close(t.grad, t_ref.grad, atol=1e-4, rtol=1e-4) + + def test_backward_is_row_invariant(self): + """Keep a timestep gradient bitwise unchanged when unrelated rows join the batch.""" + + op = _cuda_op() + grad_row = torch.randn(256, device="cuda", generator=torch.Generator("cuda").manual_seed(2)) + + def grad_of(t, pos): + """Differentiate one selected timestep using a fixed upstream feature-gradient row.""" + + t = t.clone().requires_grad_(True) + grad = torch.zeros(t.shape[0], 256, device="cuda") + grad[pos] = grad_row + op(t).backward(grad) + return t.grad[pos] + + single = grad_of(torch.tensor([0.7], device="cuda"), 0) + t = h3_timesteps(40) + t[13] = 0.7 + assert torch.equal(grad_of(t, 13), single) + + def test_registry_dispatches_cuda(self): + """Resolve the sinusoid registry entry to its dedicated CUDA implementation.""" + + from rl_engine.runtime.registry import KernelRegistry + + _cuda_op() + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda") + assert type(op).__name__ == "H3TimestepSinusoidCudaOp" diff --git a/tests/models/minimax_h3/test_h3_tp_adaln.py b/tests/models/minimax_h3/test_h3_tp_adaln.py new file mode 100644 index 000000000..c9b511153 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_tp_adaln.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""RFC #420 ``tp_adaln_3mod``: column-parallel AdaLN projection, byte-equal to WS1. + +* ownership: contiguous 64-aligned column shards cover the 3 x 6 x H table + exactly once; unsupported TP sizes and wrong shards fail closed; +* shard arithmetic (one GPU): a rank's forward columns and ``d_input`` chunk + partials are bitwise the matching slice of the full WS1 call, and folding + the rank-ordered partials reproduces WS1's ``d_input``; +* real NCCL (2/4/8 GPUs): every rank's table and ``d_temb`` and the + concatenated ``dW``/``db`` shards are bitwise equal to WS1. +""" + +from __future__ import annotations + +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import adaln_column_shard +from rl_engine.validation.models.h3_ws2 import ( + run_world, + tp_matches_ws1, + tp_projection_rank, + ws1_projection, +) + +requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" +N_H3 = 3 * 6 * 5376 +ALL_TRUE = dict.fromkeys(("table", "d_temb", "d_weight", "d_bias"), True) + + +def _require_ext(): + from rl_engine.backends.cuda.model_specific.minimax_h3.det_linear import det_linear_available + + if not det_linear_available(): + pytest.skip("rl_engine._C lacks h3_det_linear_*") + + +def _inputs(num_t, n_out, k, dtype=torch.bfloat16, seed=0): + g = torch.Generator(device="cpu").manual_seed(seed) + temb = (torch.randn(num_t, k, generator=g) * 2).cuda() + weight = (torch.randn(n_out, k, generator=g) / k**0.5).to(dtype).cuda() + bias = (torch.randn(n_out, generator=g) * 0.1).to(dtype).cuda() + grad = torch.randn(num_t, n_out, generator=g).to(dtype).cuda() + return temb, weight, bias, grad + + +class TestOwnership: + @pytest.mark.parametrize("tp", [1, 2, 3, 4, 6, 7, 8, 9]) + def test_shards_tile_the_table_once(self, tp): + shards = [adaln_column_shard(N_H3, tp, r) for r in range(tp)] + assert [s.begin for s in shards] == [0, *[s.end for s in shards[:-1]]] + assert shards[-1].end == N_H3 + assert sum(h1 - h0 for s in shards for (_, _, h0, h1) in s.slots()) == N_H3 + + def test_tp2_slots(self): + # rank 0: video's six chunks and text's first three; rank 1: the rest. + assert adaln_column_shard(N_H3, 2, 0).slots() == [ + *[(0, c, 0, 5376) for c in range(6)], + *[(1, c, 0, 5376) for c in range(3)], + ] + + @pytest.mark.parametrize("tp, rank", [(5, 0), (16, 0), (0, 0), (2, 2), (2, -1)]) + def test_rejects_unsupported(self, tp, rank): + with pytest.raises(ValueError): + adaln_column_shard(N_H3, tp, rank) + + +@requires_cuda +class TestShardArithmetic: + @pytest.mark.parametrize("tp", [2, 4, 8]) + def test_pinned_shards_reproduce_ws1(self, h3_weights_cpu, tp): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3 import det_linear as dl + + weight, bias = h3_weights_cpu[WEIGHT].cuda(), h3_weights_cpu[BIAS].cuda() + g = torch.Generator(device="cpu").manual_seed(tp) + act = F.silu(torch.randn(3, weight.shape[1], generator=g).cuda() * 2).bfloat16() + grad = torch.randn(3, weight.shape[0], generator=g).cuda() + (full,) = dl.det_linear_forward(act, weight, bias) + full_partials = dl.det_linear_backward_input_partials(grad, weight) + pieces = [] + for rank in range(tp): + s = adaln_column_shard(weight.shape[0], tp, rank) + (cols,) = dl.det_linear_forward(act, weight[s.begin : s.end], bias[s.begin : s.end]) + assert torch.equal(cols, full[:, s.begin : s.end]) + part = dl.det_linear_backward_input_partials( + grad[:, s.begin : s.end].contiguous(), weight[s.begin : s.end] + ) + c0 = s.begin // dl.DINPUT_CHUNK + assert torch.equal(part, full_partials[c0 : c0 + part.shape[0]]) + pieces.append(part) + folded = dl.det_linear_fold_chunks(torch.cat(pieces), torch.float32) + assert torch.equal(folded, dl.det_linear_backward_input(grad, weight, torch.float32)) + + def test_tp1_op_is_ws1(self): + _require_ext() + one = SimpleNamespace(rank=0, world_size=1, backend_id="none", all_gather=lambda t: t) + temb, weight, bias, grad = _inputs(3, 3 * 6 * 256, 136) + ws1 = ws1_projection(temb, weight, bias, grad) + assert tp_matches_ws1(ws1, [tp_projection_rank(one, temb, weight, bias, grad)]) == ALL_TRUE + + def test_fails_closed(self): + _require_ext() + from rl_engine.backends.cuda.model_specific.minimax_h3.tp_adaln_projection import ( + H3TPAdaLNProjectionCudaOp, + shard_adaln_projection, + ) + + temb, weight, bias, _ = _inputs(2, 3 * 6 * 256, 136) + rank1 = SimpleNamespace(rank=1, world_size=2, backend_id="unused", all_gather=None) + op = H3TPAdaLNProjectionCudaOp(rank1, weight.shape[0]) + w, b = shard_adaln_projection(weight, bias, 2, 1) + with pytest.raises(TypeError): # RFC probe H7: cast before the SiLU + op(temb.bfloat16(), w, b) + with pytest.raises(ValueError): # the full weight instead of this rank's rows + op(temb, weight, bias) + with pytest.raises(ValueError): + op(temb.cpu(), w, b) + with pytest.raises(ValueError): # 4608 / 16 is not a multiple of 64 + H3TPAdaLNProjectionCudaOp(SimpleNamespace(rank=0, world_size=16), weight.shape[0]) + assert op.readback()["columns"] == [2304, 4608] + + +@pytest.mark.parametrize("world", [2, 4, 8]) +@pytest.mark.skipif(int(os.environ.get("WORLD_SIZE", "1")) != 1, reason="owns its processes") +def test_nccl_byte_equal_to_ws1(h3_weights_cpu, world): + if torch.cuda.device_count() < world: + pytest.skip(f"needs {world} GPUs") + _require_ext() + weight, bias = h3_weights_cpu[WEIGHT], h3_weights_cpu[BIAS] + g = torch.Generator(device="cpu").manual_seed(world) + inputs = { + "temb": torch.randn(3, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(3, weight.shape[0], generator=g).bfloat16(), + } + ws1 = { + k: v.cpu() for k, v in ws1_projection(**{k: v.cuda() for k, v in inputs.items()}).items() + } + torch.cuda.empty_cache() + ranks = run_world(world, tp_projection_rank, inputs) + assert tp_matches_ws1(ws1, ranks) == ALL_TRUE + assert [r["readback"]["rank"] for r in ranks] == list(range(world)) + assert {r["readback"]["collective_backend"] for r in ranks} == {"cuda_ipc_fixed_tree"} + + +@pytest.mark.parametrize("n_total", [0, -1152]) +def test_rejects_nonpositive_projection(n_total): + with pytest.raises(ValueError, match="N must be positive"): + adaln_column_shard(n_total, 1, 0) diff --git a/tests/models/minimax_h3/test_h3_weights.py b/tests/models/minimax_h3/test_h3_weights.py new file mode 100644 index 000000000..13b9b2489 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_weights.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Extracted H3 weights must retain the pinned tensor contents.""" + +from __future__ import annotations + +import hashlib + +import pytest +import torch + +from rl_engine.validation.models import h3_weights + + +@pytest.mark.parametrize( + "dtype,raw_bytes", + [ + (torch.float32, bytes.fromhex("000000000000803f")), + (torch.bfloat16, bytes.fromhex("0000803f")), + ], +) +def test_tensor_checksum_hashes_raw_bytes(dtype, raw_bytes): + """Hash the exact FP32 and BF16 bit patterns without numerical conversion.""" + + tensor = torch.tensor([0.0, 1.0], dtype=dtype) + assert h3_weights.sha256_tensor(tensor) == hashlib.sha256(raw_bytes).hexdigest() + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +@pytest.mark.parametrize("names", [None, ["bias"]]) +def test_loading_rejects_modified_weight_contents(tmp_path, monkeypatch, dtype, mismatch, names): + """Reject altered shape, dtype, or bytes even outside the requested tensor subset.""" + + from safetensors.torch import save_file + + tensors = { + "weight": torch.tensor([[0.0, 1.0], [2.0, 3.0]], dtype=dtype), + "bias": torch.tensor([0.0, 1.0], dtype=dtype), + } + manifest = { + "tensors": { + name: { + "dtype": str(tensor.dtype).removeprefix("torch."), + "shape": list(tensor.shape), + "sha256": h3_weights.sha256_tensor(tensor), + } + for name, tensor in tensors.items() + } + } + if mismatch == "dtype": + tensors["weight"] = tensors["weight"].to(torch.float64) + elif mismatch == "shape": + tensors["weight"] = tensors["weight"].reshape(4) + elif mismatch == "sha256": + tensors["weight"][0, 0] = 1.0 + # Tensor checksums do not depend on serializer metadata or tensor key order. + save_file( + dict(reversed(list(tensors.items()))), + str(tmp_path / h3_weights.EXTRACTED_FILE), + metadata={"revision": "a different serializer metadata value"}, + ) + monkeypatch.setenv(h3_weights.WEIGHTS_ENV, str(tmp_path)) + monkeypatch.setattr(h3_weights, "load_h3_manifest", lambda: manifest) + + if mismatch: + with pytest.raises(ValueError, match=f"weight: {mismatch} .* != manifest"): + h3_weights.load_h3_conditioning_weights("cpu", names=names) + else: + actual = h3_weights.load_h3_conditioning_weights("cpu", names=names) + assert list(actual) == (names if names is not None else list(tensors)) + for name, tensor in actual.items(): + assert torch.equal(tensor, tensors[name]) diff --git a/tests/models/minimax_h3/test_h3_ws2_processes.py b/tests/models/minimax_h3/test_h3_ws2_processes.py new file mode 100644 index 000000000..f448ba580 --- /dev/null +++ b/tests/models/minimax_h3/test_h3_ws2_processes.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""CPU process lifecycle regressions for the multi-GPU evidence runner.""" + +import os +import time + +import pytest +import torch +import torch.multiprocessing as mp + +from rl_engine.validation.models import h3_ws2 + + +def _worker(rank, world, init, queue, path, target, kwargs): + mode = kwargs["mode"] + if mode == "success": + torch.save({"value": torch.tensor(rank)}, path.with_name(f"rank{rank}.pt")) + queue.put({"rank": rank}) + elif rank == 0 and mode == "error": + queue.put({"rank": rank, "error": "original rank traceback"}) + elif rank == 0 and mode == "crash": + os._exit(7) + else: + time.sleep(120) + + +@pytest.mark.parametrize("mode", ["success", "error", "crash", "timeout"]) +def test_run_world_reaps_workers(monkeypatch, mode): + monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) + monkeypatch.setattr(h3_ws2, "_bootstrap", _worker) + before = {p.pid for p in mp.active_children()} + start = time.monotonic() + if mode == "success": + ranks = h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert [r["rank"] for r in ranks] == [0, 1] + assert [r["value"].item() for r in ranks] == [0, 1] + elif mode == "timeout": + with pytest.raises(TimeoutError, match="timed out"): + h3_ws2.run_world(2, None, {}, timeout=0.2, mode=mode) + else: + message = "original rank traceback" if mode == "error" else "rank 0 exited with code 7" + with pytest.raises(RuntimeError, match=message): + h3_ws2.run_world(2, None, {}, timeout=30, mode=mode) + assert time.monotonic() - start < 25 + assert {p.pid for p in mp.active_children()} == before diff --git a/tests/models/minimax_h3/test_prepare_h3_weights.py b/tests/models/minimax_h3/test_prepare_h3_weights.py new file mode 100644 index 000000000..5c9a966b9 --- /dev/null +++ b/tests/models/minimax_h3/test_prepare_h3_weights.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Weight extraction rejects off-manifest tensors even under Python optimization.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest +import torch + +from rl_engine.validation.models.h3_weights import EXTRACTED_FILE, sha256_file, sha256_tensor + + +@pytest.mark.parametrize("optimize", [0, 1]) +@pytest.mark.parametrize("mismatch", [None, "dtype", "shape", "sha256"]) +def test_manifest_validation_before_write(tmp_path, monkeypatch, optimize, mismatch): + """Enforce tensor contracts before extraction writes, including under Python optimization.""" + + import huggingface_hub + from safetensors.torch import load_file, save_file + + tensor = torch.tensor([0.0, 1.0]) + config = tmp_path / "config.json" + config.write_text("{}") + index = tmp_path / "index.json" + index.write_text(json.dumps({"weight_map": {"weight": "shard.safetensors"}})) + shard = tmp_path / "shard.safetensors" + save_file({"weight": tensor}, str(shard)) + manifest = { + "model_identity": { + "hf_repo": "test/h3", + "subfolder": "dit", + "revision": "pinned", + "config_sha256": sha256_file(config), + "index_file": index.name, + "index_sha256": sha256_file(index), + }, + "weight_shards": {shard.name: {"sha256": sha256_file(shard)}}, + "tensors": { + "weight": { + "dtype": "bfloat16" if mismatch == "dtype" else "float32", + "shape": [3] if mismatch == "shape" else [2], + "sha256": "0" * 64 if mismatch == "sha256" else sha256_tensor(tensor), + } + }, + } + monkeypatch.setattr( + huggingface_hub, + "hf_hub_download", + lambda repo, filename, **kwargs: str(tmp_path / Path(filename).name), + ) + out_dir = tmp_path / "extracted" + monkeypatch.setattr(sys, "argv", ["prepare_h3_weights.py", "--out", str(out_dir)]) + script = Path(__file__).resolve().parents[3] / "tools" / "weights" / "prepare_h3_weights.py" + namespace = {"__name__": "prepare_h3_weights_test", "__file__": str(script)} + # Compiling the actual script at optimize=1 reproduces python -O's assert removal. + exec(compile(script.read_text(), str(script), "exec", optimize=optimize), namespace) + namespace["load_h3_manifest"] = lambda: manifest + target = out_dir / EXTRACTED_FILE + if mismatch: + with pytest.raises(SystemExit, match=f"weight: {mismatch} .* != manifest"): + namespace["main"]() + assert not target.exists() + else: + namespace["main"]() + assert torch.equal(load_file(str(target))["weight"], tensor) diff --git a/tests/runtime/test_dispatch.py b/tests/runtime/test_dispatch.py index 3c47bbbcf..647053ce1 100644 --- a/tests/runtime/test_dispatch.py +++ b/tests/runtime/test_dispatch.py @@ -4,6 +4,7 @@ import sys from types import ModuleType +import pytest import torch import rl_engine.platforms.device as device_module @@ -231,6 +232,44 @@ def fake_load_backend(backend): ] +H3_OPS = ( + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", +) + + +@pytest.mark.parametrize("op_name", H3_OPS) +def test_h3_ops_prefer_cuda_and_fall_back_to_pytorch_reference(op_name): + """Prefer H3 CUDA backends on CUDA and use their PyTorch fallbacks elsewhere.""" + + registry = KernelRegistry() + cuda_chain = registry._priority_map["cuda"][op_name] + assert cuda_chain[0].name.startswith("CUDA_H3_") + assert cuda_chain[-1].name.startswith("PYTORCH_H3_") + for platform in ("rocm", "musa", "cpu", "npu"): + chain = registry._priority_map[platform][op_name] + assert [backend.name for backend in chain] == [cuda_chain[-1].name] + + +def test_h3_timestep_falls_back_when_fast_math_build_omits_kernel(monkeypatch): + # KERNEL_ALIGN_USE_FAST_MATH=1 builds leave out the bitwise timestep kernel + # and its binding; dispatch must then pick the PyTorch reference. + from rl_engine.backends.cuda.model_specific.minimax_h3 import timestep_sinusoid + from rl_engine.reference.minimax_h3.timestep_sinusoid import NativeH3TimestepSinusoidOp + + monkeypatch.setattr(timestep_sinusoid, "h3_sinusoid_cuda_available", lambda: False) + monkeypatch.setattr(torch.version, "hip", None, raising=False) + + op = KernelRegistry().get_op("timestep_sinusoid_h3", device="cuda:0") + + assert isinstance(op, NativeH3TimestepSinusoidOp) + + def test_npu_available_handles_runtime_failure(monkeypatch): class BrokenNPU: @staticmethod diff --git a/tests/validation/operators/test_ws1_gtest_gpu.py b/tests/validation/operators/test_ws1_gtest_gpu.py index bcaebc2f5..e7d1e5275 100644 --- a/tests/validation/operators/test_ws1_gtest_gpu.py +++ b/tests/validation/operators/test_ws1_gtest_gpu.py @@ -33,6 +33,8 @@ def _run(script: str, *args: str, timeout: int = 300) -> None: def test_all_ws1_single_ops_are_registered(): + """Require gtest specs for all WS1 single operators and the H3 conditioning stages.""" + names = set(operator_names()) assert { "rms_norm", @@ -49,6 +51,13 @@ def test_all_ws1_single_ops_are_registered(): "pack", "linear_logp", "prefix_shared_attention", + "timestep_sinusoid_h3", + "timestep_mlp_fp32", + "adaln_projection_3mod", + "adaln_row_gather", + "h3_rmsnorm", + "adaln_gate_residual", + "final_adaln_out", } <= names diff --git a/tools/migration/README.md b/tools/migration/README.md index 4a6cee300..48491e383 100644 --- a/tools/migration/README.md +++ b/tools/migration/README.md @@ -8,3 +8,6 @@ Historical workload/operator identifiers and archived evidence intentionally retain their original values. New source imports should use the mapped canonical module. Remove compatibility paths only in a separately announced migration once downstream launchers, patches, imports and saved objects have migrated. + +The MiniMax-H3 entries also record the open RFC #420 PR stack migration to +`test-h3`; those H3 files were not part of the original main baseline. diff --git a/tools/migration/layout.json b/tools/migration/layout.json index cb8fe469a..6c7fc9898 100644 --- a/tools/migration/layout.json +++ b/tools/migration/layout.json @@ -347,6 +347,7 @@ "benchmarks/benchmark_deterministic_attention.py": "benchmarks/operators/attention/benchmark_deterministic_attention.py", "benchmarks/benchmark_grpo_loss.py": "benchmarks/operators/loss/benchmark_grpo_loss.py", "benchmarks/benchmark_grpo_op.py": "benchmarks/operators/loss/benchmark_grpo_op.py", + "benchmarks/benchmark_h3_conditioning.py": "benchmarks/models/benchmark_h3_conditioning.py", "benchmarks/benchmark_linear_logp.py": "benchmarks/operators/logprob/benchmark_linear_logp.py", "benchmarks/benchmark_pack.py": "benchmarks/operators/packing/benchmark_pack.py", "benchmarks/benchmark_qwen_ffn_layout.py": "benchmarks/layers/benchmark_qwen_ffn_layout.py", @@ -466,6 +467,42 @@ "docs/design/ws2-attention-single-gpu-harness.md": "docs/archive/design/ws2-attention-single-gpu-harness.md", "docs/design/ws2-cp-attention-contract.md": "docs/contracts/ws2-cp-attention-contract.md", "docs/design/ws2-module-debug-matrix.md": "docs/archive/design/ws2-module-debug-matrix.md", + "docs/usage/evidence/h3-adaln-gate-residual-b200/figure.png": "reports/experiments/h3-adaln-gate-residual-b200/figure.png", + "docs/usage/evidence/h3-adaln-gate-residual-b200/report.json": "reports/experiments/h3-adaln-gate-residual-b200/report.json", + "docs/usage/evidence/h3-adaln-projection-b200/figure.png": "reports/experiments/h3-adaln-projection-b200/figure.png", + "docs/usage/evidence/h3-adaln-projection-b200/report.json": "reports/experiments/h3-adaln-projection-b200/report.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/chain_replay.json": "reports/experiments/h3-adaln-row-gather-b200/chain_replay.json", + "docs/usage/evidence/h3-adaln-row-gather-b200/figure.png": "reports/experiments/h3-adaln-row-gather-b200/figure.png", + "docs/usage/evidence/h3-adaln-row-gather-b200/report.json": "reports/experiments/h3-adaln-row-gather-b200/report.json", + "docs/usage/evidence/h3-final-adaln-out-b200/chain_replay.json": "reports/experiments/h3-final-adaln-out-b200/chain_replay.json", + "docs/usage/evidence/h3-final-adaln-out-b200/figure.png": "reports/experiments/h3-final-adaln-out-b200/figure.png", + "docs/usage/evidence/h3-final-adaln-out-b200/report.json": "reports/experiments/h3-final-adaln-out-b200/report.json", + "docs/usage/evidence/h3-one-block-b200/figure.png": "reports/experiments/h3-one-block-b200/figure.png", + "docs/usage/evidence/h3-one-block-b200/report.json": "reports/experiments/h3-one-block-b200/report.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.json": "reports/experiments/h3-prior-art-b200/adaln_projection.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_projection.png": "reports/experiments/h3-prior-art-b200/adaln_projection.png", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.json": "reports/experiments/h3-prior-art-b200/adaln_row_gather.json", + "docs/usage/evidence/h3-prior-art-b200/adaln_row_gather.png": "reports/experiments/h3-prior-art-b200/adaln_row_gather.png", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.json": "reports/experiments/h3-prior-art-b200/final_adaln_out.json", + "docs/usage/evidence/h3-prior-art-b200/final_adaln_out.png": "reports/experiments/h3-prior-art-b200/final_adaln_out.png", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.json": "reports/experiments/h3-prior-art-b200/gate_residual.json", + "docs/usage/evidence/h3-prior-art-b200/gate_residual.png": "reports/experiments/h3-prior-art-b200/gate_residual.png", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.json": "reports/experiments/h3-prior-art-b200/norm_modulate.json", + "docs/usage/evidence/h3-prior-art-b200/norm_modulate.png": "reports/experiments/h3-prior-art-b200/norm_modulate.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.json": "reports/experiments/h3-prior-art-b200/timestep_mlp.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_mlp.png": "reports/experiments/h3-prior-art-b200/timestep_mlp.png", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.json": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.json", + "docs/usage/evidence/h3-prior-art-b200/timestep_sinusoid.png": "reports/experiments/h3-prior-art-b200/timestep_sinusoid.png", + "docs/usage/evidence/h3-rmsnorm-b200/figure.png": "reports/experiments/h3-rmsnorm-b200/figure.png", + "docs/usage/evidence/h3-rmsnorm-b200/report.json": "reports/experiments/h3-rmsnorm-b200/report.json", + "docs/usage/evidence/h3-sp-norm-adaln-b200/figure.png": "reports/experiments/h3-sp-norm-adaln-b200/figure.png", + "docs/usage/evidence/h3-sp-norm-adaln-b200/report.json": "reports/experiments/h3-sp-norm-adaln-b200/report.json", + "docs/usage/evidence/h3-timestep-mlp-b200/figure.png": "reports/experiments/h3-timestep-mlp-b200/figure.png", + "docs/usage/evidence/h3-timestep-mlp-b200/report.json": "reports/experiments/h3-timestep-mlp-b200/report.json", + "docs/usage/evidence/h3-timestep-sinusoid-b200/figure.png": "reports/experiments/h3-timestep-sinusoid-b200/figure.png", + "docs/usage/evidence/h3-timestep-sinusoid-b200/report.json": "reports/experiments/h3-timestep-sinusoid-b200/report.json", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/figure.png": "reports/experiments/h3-tp-adaln-3mod-b200/figure.png", + "docs/usage/evidence/h3-tp-adaln-3mod-b200/report.json": "reports/experiments/h3-tp-adaln-3mod-b200/report.json", "envs.py": "build_tools/envs.py", "examples/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json": "configs/experiments/cross_config/cross_config_qwen3_8b_megatron_tp2_cp2_vllm.json", "examples/cross_config_s0_cpu_smoke.json": "configs/experiments/cross_config/cross_config_s0_cpu_smoke.json", @@ -675,6 +712,19 @@ "rl_engine/kernels/ops/cuda/attention/prefix_shared_attn.py": "rl_engine/backends/cuda/attention/prefix_shared_attn.py", "rl_engine/kernels/ops/cuda/attention/strict_runtime.py": "rl_engine/backends/cuda/attention/strict_runtime.py", "rl_engine/kernels/ops/cuda/ffn.py": "rl_engine/backends/cuda/ffn/ffn.py", + "rl_engine/kernels/ops/cuda/h3/__init__.py": "rl_engine/backends/cuda/model_specific/minimax_h3/__init__.py", + "rl_engine/kernels/ops/cuda/h3/adaln_modulation.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_modulation.py", + "rl_engine/kernels/ops/cuda/h3/adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/cuda/h3/adaln_row_gather.py": "rl_engine/backends/cuda/model_specific/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/cuda/h3/det_linear.py": "rl_engine/backends/cuda/model_specific/minimax_h3/det_linear.py", + "rl_engine/kernels/ops/cuda/h3/final_adaln_out.py": "rl_engine/backends/cuda/model_specific/minimax_h3/final_adaln_out.py", + "rl_engine/kernels/ops/cuda/h3/gate_residual.py": "rl_engine/backends/cuda/model_specific/minimax_h3/gate_residual.py", + "rl_engine/kernels/ops/cuda/h3/rmsnorm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/cuda/h3/sp_norm_adaln.py": "rl_engine/backends/cuda/model_specific/minimax_h3/sp_norm_adaln.py", + "rl_engine/kernels/ops/cuda/h3/timestep_mlp.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_mlp.py", + "rl_engine/kernels/ops/cuda/h3/timestep_sinusoid.py": "rl_engine/backends/cuda/model_specific/minimax_h3/timestep_sinusoid.py", + "rl_engine/kernels/ops/cuda/h3/tp_adaln_projection.py": "rl_engine/backends/cuda/model_specific/minimax_h3/tp_adaln_projection.py", + "rl_engine/kernels/ops/cuda/h3/ws2_comm.py": "rl_engine/backends/cuda/model_specific/minimax_h3/ws2_comm.py", "rl_engine/kernels/ops/cuda/linear/__init__.py": "rl_engine/backends/cuda/embedding/__init__.py", "rl_engine/kernels/ops/cuda/linear/embedding.py": "rl_engine/backends/cuda/embedding/embedding.py", "rl_engine/kernels/ops/cuda/linear/lm_head.py": "rl_engine/backends/cuda/embedding/lm_head.py", @@ -705,6 +755,15 @@ "rl_engine/kernels/ops/pytorch/attention/stateful_kv.py": "rl_engine/reference/attention/stateful_kv.py", "rl_engine/kernels/ops/pytorch/ffn/__init__.py": "rl_engine/ops/ffn/__init__.py", "rl_engine/kernels/ops/pytorch/ffn/ffn.py": "rl_engine/ops/ffn/qwen3.py", + "rl_engine/kernels/ops/pytorch/h3/__init__.py": "rl_engine/reference/minimax_h3/__init__.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_projection.py": "rl_engine/reference/minimax_h3/adaln_projection.py", + "rl_engine/kernels/ops/pytorch/h3/adaln_row_gather.py": "rl_engine/reference/minimax_h3/adaln_row_gather.py", + "rl_engine/kernels/ops/pytorch/h3/final_adaln_out.py": "rl_engine/reference/minimax_h3/final_adaln_out.py", + "rl_engine/kernels/ops/pytorch/h3/fixed_order.py": "rl_engine/reference/minimax_h3/fixed_order.py", + "rl_engine/kernels/ops/pytorch/h3/gate_residual.py": "rl_engine/reference/minimax_h3/gate_residual.py", + "rl_engine/kernels/ops/pytorch/h3/rmsnorm.py": "rl_engine/reference/minimax_h3/rmsnorm.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_mlp.py": "rl_engine/reference/minimax_h3/timestep_mlp.py", + "rl_engine/kernels/ops/pytorch/h3/timestep_sinusoid.py": "rl_engine/reference/minimax_h3/timestep_sinusoid.py", "rl_engine/kernels/ops/pytorch/linear/__init__.py": "rl_engine/reference/embedding/__init__.py", "rl_engine/kernels/ops/pytorch/linear/embedding.py": "rl_engine/reference/embedding/embedding.py", "rl_engine/kernels/ops/pytorch/linear/lm_head.py": "rl_engine/reference/embedding/lm_head.py", @@ -776,6 +835,13 @@ "rl_engine/testing/__init__.py": "rl_engine/validation/common/__init__.py", "rl_engine/testing/attention_comparison.py": "rl_engine/validation/common/attention_comparison.py", "rl_engine/testing/distributed_logprob_comparison.py": "rl_engine/validation/common/distributed_logprob_comparison.py", + "rl_engine/testing/h3_cases.py": "rl_engine/validation/models/h3_cases.py", + "rl_engine/testing/h3_chain.py": "rl_engine/validation/models/h3_chain.py", + "rl_engine/testing/h3_manifest.json": "rl_engine/validation/models/h3_manifest.json", + "rl_engine/testing/h3_provider.py": "rl_engine/validation/models/h3_provider.py", + "rl_engine/testing/h3_report.py": "rl_engine/validation/models/h3_report.py", + "rl_engine/testing/h3_weights.py": "rl_engine/validation/models/h3_weights.py", + "rl_engine/testing/h3_ws2.py": "rl_engine/validation/models/h3_ws2.py", "rl_engine/testing/logprob_comparison.py": "rl_engine/validation/common/logprob_comparison.py", "rl_engine/testing/logprob_drift.py": "rl_engine/validation/common/logprob_drift.py", "rl_engine/testing/reference_ops.py": "rl_engine/reference/definitions.py", @@ -790,6 +856,13 @@ "scripts/check_rocm_env.py": "tools/env/check_rocm_env.py", "scripts/check_stateful_kv.py": "tools/validation/operators/check_stateful_kv.py", "scripts/ci_smoke.py": "tools/checks/ci_smoke.py", + "scripts/h3_chain_replay.py": "tools/validation/models/h3_chain_replay.py", + "scripts/h3_evidence.py": "tools/validation/models/h3_evidence.py", + "scripts/h3_prior_art.py": "tools/validation/models/h3_prior_art.py", + "scripts/h3_ws2_evidence.py": "tools/validation/models/h3_ws2_evidence.py", + "scripts/plot_h3_evidence.py": "tools/validation/models/plot_h3_evidence.py", + "scripts/plot_h3_prior_art.py": "tools/validation/models/plot_h3_prior_art.py", + "scripts/prepare_h3_weights.py": "tools/weights/prepare_h3_weights.py", "scripts/prepare_ws1_weights.py": "tools/weights/prepare_ws1_weights.py", "scripts/run_perf.py": "tools/benchmarking/run_perf.py", "scripts/run_profile_suite.py": "tools/benchmarking/run_profile_suite.py", @@ -820,6 +893,27 @@ "tests/distributed/test_rocm_attention_transport.py": "tests/distributed/cp/test_rocm_attention_transport.py", "tests/distributed/test_rocm_strict_attention_cp.py": "tests/distributed/cp/test_rocm_strict_attention_cp.py", "tests/distributed/test_transport_deterministic_collective.py": "tests/distributed/collectives/test_transport_deterministic_collective.py", + "tests/h3/conftest.py": "tests/models/minimax_h3/conftest.py", + "tests/h3/test_h3_adaln_gate_residual.py": "tests/models/minimax_h3/test_h3_adaln_gate_residual.py", + "tests/h3/test_h3_adaln_modulation.py": "tests/models/minimax_h3/test_h3_adaln_modulation.py", + "tests/h3/test_h3_adaln_projection.py": "tests/models/minimax_h3/test_h3_adaln_projection.py", + "tests/h3/test_h3_adaln_row_gather.py": "tests/models/minimax_h3/test_h3_adaln_row_gather.py", + "tests/h3/test_h3_benchmark_timing.py": "tests/models/minimax_h3/test_h3_benchmark_timing.py", + "tests/h3/test_h3_chain_replay_cli.py": "tests/models/minimax_h3/test_h3_chain_replay_cli.py", + "tests/h3/test_h3_cli.py": "tests/models/minimax_h3/test_h3_cli.py", + "tests/h3/test_h3_conditioning_e2e.py": "tests/models/minimax_h3/test_h3_conditioning_e2e.py", + "tests/h3/test_h3_det_linear.py": "tests/models/minimax_h3/test_h3_det_linear.py", + "tests/h3/test_h3_final_adaln_out.py": "tests/models/minimax_h3/test_h3_final_adaln_out.py", + "tests/h3/test_h3_native_validation.py": "tests/models/minimax_h3/test_h3_native_validation.py", + "tests/h3/test_h3_report.py": "tests/models/minimax_h3/test_h3_report.py", + "tests/h3/test_h3_rmsnorm.py": "tests/models/minimax_h3/test_h3_rmsnorm.py", + "tests/h3/test_h3_sp_norm_adaln.py": "tests/models/minimax_h3/test_h3_sp_norm_adaln.py", + "tests/h3/test_h3_timestep_mlp.py": "tests/models/minimax_h3/test_h3_timestep_mlp.py", + "tests/h3/test_h3_timestep_sinusoid.py": "tests/models/minimax_h3/test_h3_timestep_sinusoid.py", + "tests/h3/test_h3_tp_adaln.py": "tests/models/minimax_h3/test_h3_tp_adaln.py", + "tests/h3/test_h3_weights.py": "tests/models/minimax_h3/test_h3_weights.py", + "tests/h3/test_h3_ws2_processes.py": "tests/models/minimax_h3/test_h3_ws2_processes.py", + "tests/h3/test_prepare_h3_weights.py": "tests/models/minimax_h3/test_prepare_h3_weights.py", "tests/linear_logp_tp.py": "tests/distributed/tp/linear_logp_tp.py", "tests/test_alignment_model_wrappers.py": "tests/models/qwen3/test_alignment_model_wrappers.py", "tests/test_alignment_wrapper_interfaces.py": "tests/models/qwen3/test_alignment_wrapper_interfaces.py", diff --git a/tools/validation/models/h3_bitwise_reuse.py b/tools/validation/models/h3_bitwise_reuse.py new file mode 100644 index 000000000..d43e5c70e --- /dev/null +++ b/tools/validation/models/h3_bitwise_reuse.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Direct bitwise comparison of one RFC #420 operator with its reuse candidates. + +``h3_prior_art.py`` reports each implementation's batch invariance and its error against FP64. +Matching error summaries do not show that two implementations produce the same bits, so before +a candidate is adopted this runner compares outputs directly. Every candidate that +``h3_prior_art.candidates()`` defines runs on the same inputs (each global batch-invariant mode +in its own process, as in ``h3_prior_art.py``), and every forward output and leaf gradient is +compared byte for byte with diffusers' path and with the rl-kernel op: + + python tools/validation/models/h3_bitwise_reuse.py --op norm_modulate \\ + --out reports/experiments/h3-bitwise-reuse-b200/norm_modulate.json \\ + [--megatron-src /path/to/Megatron-LM] [--fast-math-variant] + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png + +For ``timestep_sinusoid`` the inputs are three timestep sets: 2^20 uniform draws, every 997th +FP32 bit pattern in [0, 1] (a dense sweep, mostly tiny values), and a 100001-point grid. +``--fast-math-variant`` also compiles ``csrc/cuda/h3/timestep_sinusoid.cu`` with +``--use_fast_math`` and compares it, to show what the fast-math build guard prevents. Other +operators use each candidate's own input generator at four sizes. + +Intermediate tensors are written next to the report (not to /tmp, which cluster job epilogs +can clear) and removed afterwards. Run on a clean tree; the report records the commit, the tree +state and every library version. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import json +import os +import shutil +import subprocess +import sys +from pathlib import Path +from typing import Any + +import h3_prior_art as PA +import torch + +SEED = 7 +SIZES = { + "timestep_mlp": (1, 4, 64, 2048), + "adaln_projection": (1, 4, 64, 2048), + "adaln_row_gather": (1, 257, 4097, 32768), + "norm_modulate": (1, 257, 4097, 32768), + "gate_residual": (1, 257, 4097, 32768), + "final_adaln_out": (1, 257, 4097, 32768), +} +SINUSOID_SETS = ("random_2^20", "dense_fp32_sweep", "grid_100001") +_FAST_MATH_BINDING = """ +#include +torch::Tensor h3_timestep_sinusoid_forward(torch::Tensor, int64_t, double, bool); +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &h3_timestep_sinusoid_forward); } +""" + + +def _sinusoid_inputs(name: str) -> dict[str, torch.Tensor]: + if name == "random_2^20": + t = torch.rand(1 << 20, device="cuda", generator=torch.Generator("cuda").manual_seed(0)) + elif name == "dense_fp32_sweep": + bits = torch.arange(0, 0x3F800001, 997, dtype=torch.int32, device="cuda") + t = bits.view(torch.float32) + else: + t = torch.linspace(0, 1, 100001, device="cuda") + return {"t": t} + + +def _fast_math_sinusoid(build_dir: Path): + from torch.utils.cpp_extension import load + + build_dir.mkdir(parents=True, exist_ok=True) + binding = build_dir / "binding.cpp" + binding.write_text(_FAST_MATH_BINDING) + root = Path(__file__).resolve().parents[3] + mod = load( + name="h3_timestep_sinusoid_fast_math", + sources=[str(root / "csrc/cuda/h3/timestep_sinusoid.cu"), str(binding)], + build_directory=str(build_dir), + extra_cuda_cflags=["-O3", "--use_fast_math"], + ) + return { + "source": "rl-kernel timestep_sinusoid.cu compiled with --use_fast_math", + "grads": [], + "fn": lambda t: mod.forward(t, PA.F1, 10000.0, False), + } + + +def _run(cand: dict[str, Any], inputs: dict[str, torch.Tensor]): + args = {} + for k, v in inputs.items(): + args[k] = v.detach().clone().requires_grad_(True) if k in cand["grads"] else v + out = cand["fn"](**args) + outs = list(out) if isinstance(out, (tuple, list)) else [out] + live = [o for o in outs if o.requires_grad] + if live: + gen = torch.Generator("cuda").manual_seed(SEED) + ups = [torch.randn(o.shape, device="cuda", generator=gen).to(o.dtype) for o in live] + torch.autograd.backward(live, ups) + grads = {k: args[k].grad.cpu() if args[k].grad is not None else None for k in cand["grads"]} + return [o.detach().cpu() for o in outs], grads + + +def dump(op: str, mode: str, parts: Path, megatron_src: str | None, fast_math: bool) -> None: + """Child process: run every candidate of one mode and save its outputs.""" + + try: + PA.enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + (parts / f"{mode}.unavailable").write_text(f"{type(exc).__name__}: {exc}"[:300]) + return + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + cands = PA.candidates(op, mode) + if fast_math and op == "timestep_sinusoid" and mode == "plain": + cands.append({"name": "rl_kernel_fast_math", **_fast_math_sinusoid(parts / "build")}) + for cand in cands: + name = cand.pop("name") + if "unavailable" in cand: + (parts / f"{name}.unavailable").write_text(cand["unavailable"][:300]) + continue + cases = SINUSOID_SETS if op == "timestep_sinusoid" else SIZES[op] + for case in cases: + inputs = ( + _sinusoid_inputs(case) if op == "timestep_sinusoid" else cand["make"](SEED, case) + ) + outs, grads = _run(cand, inputs) + torch.save( + {"source": cand["source"], "out": outs, "grads": grads}, + parts / f"{name}__{case}.pt", + ) + print(f" [{mode}] {name}: done", flush=True) + torch.cuda.empty_cache() + + +def _same(a: torch.Tensor | None, b: torch.Tensor | None) -> dict[str, Any] | None: + if a is None or b is None: + return None + if a.shape != b.shape or a.dtype != b.dtype: + return {"equal": False, "note": f"{a.dtype}{tuple(a.shape)} vs {b.dtype}{tuple(b.shape)}"} + width = a.element_size() + differ = ( + a.contiguous().view(torch.uint8).view(-1, width) + != b.contiguous().view(torch.uint8).view(-1, width) + ).any(1) + return { + "equal": not bool(differ.any()), + "differ": int(differ.sum()), + "numel": a.numel(), + "max_abs": float((a.double() - b.double()).abs().max()), + } + + +def compare(parts: Path) -> dict[str, Any]: + runs: dict[str, dict[str, Any]] = {} + for f in sorted(parts.glob("*.pt")): + name, case = f.stem.rsplit("__", 1) + runs.setdefault(name, {})[case] = torch.load(f) + refs = [r for r in ("diffusers[plain]", "diffusers", "rl_kernel") if r in runs] + out: dict[str, Any] = {} + for name, by_case in runs.items(): + entry: dict[str, Any] = {"source": next(iter(by_case.values()))["source"]} + for ref in refs: + if ref == name: + continue + entry[f"vs {ref}"] = { + case: { + "out": [_same(a, b) for a, b in zip(run["out"], runs[ref][case]["out"])], + "grads": { + g: _same(run["grads"].get(g), runs[ref][case]["grads"].get(g)) + for g in run["grads"] + }, + } + for case, run in by_case.items() + } + out[name] = entry + for f in parts.glob("*.unavailable"): + out[f.stem] = {"unavailable": f.read_text()} + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(PA.OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument( + "--fast-math-variant", + action="store_true", + help="timestep_sinusoid only: also compare a --use_fast_math build of the kernel", + ) + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(PA.OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {PA.OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + parts = args.out.with_name(f".{args.out.stem}.parts") + + if args.mode is not None: # child process: one global mode + dump(args.op, args.mode, parts, args.megatron_src, args.fast_math_variant) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + shutil.rmtree(parts, ignore_errors=True) + parts.mkdir(parents=True) + modes = list(PA.MODES) if args.op in PA.GEMM_OPS else ["plain"] + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(args.out), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.fast_math_variant: + cmd.append("--fast-math-variant") + proc = subprocess.run(cmd, env={**os.environ, **PA.MODES[mode]}) + if proc.returncode != 0: + (parts / f"{mode}.unavailable").write_text(f"subprocess exited with {proc.returncode}") + results = compare(parts) + shutil.rmtree(parts, ignore_errors=True) + report = { + "kind": "h3_bitwise_reuse", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: PA._version(lib) for lib in PA.LIBS}, + "mode_environment": {m: PA.MODES[m] for m in modes}, + }, + "cases": list(SINUSOID_SETS if args.op == "timestep_sinusoid" else SIZES[args.op]), + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_diagnostic.py b/tools/validation/models/h3_block_diagnostic.py new file mode 100644 index 000000000..19e5b111a --- /dev/null +++ b/tools/validation/models/h3_block_diagnostic.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Localize a failed block-0 gradient gate without changing its numerical contract.""" +import argparse +import json +import os +import subprocess +from dataclasses import replace +from pathlib import Path + +import torch + +from rl_engine.backends.extension import _C +from rl_engine.validation.models import h3_block as block +from rl_engine.validation.models.h3_weights import load_h3_block_weights +from tools.validation.models.h3_full_model import digest + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--failed-job", required=True) + parser.add_argument("--causal-sweep", action="store_true") + args = parser.parse_args() + args.output.mkdir(parents=True, exist_ok=False) + if subprocess.check_output(["git", "status", "--porcelain"], text=True).strip(): + raise ValueError("diagnostic requires clean source") + ops = block.CandidateOps() + weights = load_h3_block_weights("cuda") + ctx = block.make_block_context(weights, layout="tiny", seed=0, batch=1) + upstream = ( + torch.randn((1, ctx["seq_len"], block.HIDDEN), generator=torch.Generator().manual_seed(11)) + .bfloat16() + .cuda() + ) + report = dict( + status="diagnostic_only_not_acceptance", + failed_job=args.failed_job, + job_id=os.environ.get("SLURM_JOB_ID"), + commit=subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip(), + extension_sha256=digest(Path(_C.__file__)), + torch=torch.__version__, + device=torch.cuda.get_device_name(), + backends=ops.backends(), + default_matmul_tf32=torch.backends.cuda.matmul.allow_tf32, + bf16_reduced_precision_reduction=( + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + ), + leaves={}, + nodes=[], + ) + # Keep FP32 arithmetic strict through backward as well as forward. + with block.strict_fp32_matmul(): + golden, golden_nodes = block._backward("golden", ctx, None, upstream) + provider, provider_nodes = block._backward("provider", ctx, None, upstream) + candidate, candidate_nodes = block._backward("candidate", ctx, ops, upstream) + for name in golden: + report["leaves"][name] = dict( + candidate=block.rel_err(candidate[name], golden[name]), + provider=block.rel_err(provider[name], golden[name]), + ) + for node in reversed(block.NODES): + report["nodes"].append( + dict( + node=node.name, + errors=[ + dict(candidate=block.rel_err(c, g), provider=block.rel_err(p, g)) + for g, p, c in zip( + golden_nodes[node.name], + provider_nodes[node.name], + candidate_nodes[node.name], + strict=True, + ) + ], + ) + ) + # Same input and upstream gradient isolate the RMSNorm implementation from + # error accumulated through other candidate nodes. + with torch.no_grad(): + values = block.run_graph("provider", ctx) + isolated = {} + for name, source, weight_name in ( + ("q_norm", "q_proj", "attn.norm_q.weight"), + ("k_norm", "k_proj", "attn.norm_k.weight"), + ): + grads = {} + node = next(n for n in block.NODES if n.name == name) + for mode in ("golden", "provider", "candidate"): + x = values[source].detach().clone() + w = weights[block.BLOCK + "." + weight_name].detach().clone() + if mode == "golden": + x, w = x.float(), w.float() + x.requires_grad_() + w.requires_grad_() + local = {**ctx, "weights": {**weights, block.BLOCK + "." + weight_name: w}} + y = ( + node.candidate(ops, local, x) + if mode == "candidate" + else getattr(node, mode)(local, x) + ) + torch.autograd.backward(y, provider_nodes[name][0].to(y.dtype)) + grads[mode] = w.grad + isolated[name] = dict( + candidate=block.rel_err(grads["candidate"], grads["golden"]), + provider=block.rel_err(grads["provider"], grads["golden"]), + ) + report["isolated_norm_weight_gradients"] = isolated + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + if args.causal_sweep: + # These hybrids intentionally use provider nodes for localization only. + # They cannot be submitted to the candidate acceptance comparator. + report["causal_sweep"] = [] + original_nodes = block.NODES + try: + for node in original_nodes: + for baseline in ("candidate", "provider"): + mixed = [] + for item in original_nodes: + use_provider = (item.name == node.name) == (baseline == "candidate") + mixed.append( + replace( + item, + candidate=lambda ops, ctx, *xs, item=item: item.provider(ctx, *xs), + ) + if use_provider + else item + ) + block.NODES = tuple(mixed) + with block.strict_fp32_matmul(): + leaves, _ = block._backward("candidate", ctx, ops, upstream) + report["causal_sweep"].append( + dict( + baseline=baseline, + replaced_node=node.name, + errors={ + name: block.rel_err(leaves[name], golden[name]) for name in golden + }, + ) + ) + del leaves + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + finally: + block.NODES = original_nodes + report["causal_sweep_complete"] = True + (args.output / "report.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_replay.py b/tools/validation/models/h3_block_replay.py new file mode 100644 index 000000000..7d8230294 --- /dev/null +++ b/tools/validation/models/h3_block_replay.py @@ -0,0 +1,114 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay one real MiniMax-H3 block forward and backward and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_block.py`` for what is compared. + + python tools/weights/prepare_h3_weights.py --out --block + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_replay.py --out block_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_cases import H3_BLOCK_LAYOUTS # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_block_weights, + load_h3_manifest, +) + + +def main() -> None: + """Run the forward (and optionally backward) block replay per layout and save the report.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--layouts", + default="tiny,small,medium", + help=f"comma list of packed FL2VA layouts: {', '.join(H3_BLOCK_LAYOUTS)}", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--no-backward", action="store_true", help="forward replay only") + parser.add_argument("--timing-repeats", type=int, default=5) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + + layouts = args.layouts.split(",") + unknown = sorted(set(layouts) - set(H3_BLOCK_LAYOUTS)) + if unknown: + parser.error(f"unknown layouts: {unknown}") + if not torch.cuda.is_available(): + raise SystemExit("the block replay needs a CUDA device") + + weights = load_h3_block_weights("cuda") + manifest = load_h3_manifest() + ops = h3_block.CandidateOps() + + forward, backward = [], [] + for layout in layouts: + case = h3_block.run_block_case( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + forward.append(case) + print( + f"{layout} S={case['seq_len']}: first_drift={case['first_drift']} " + f"first_isolated_drift={case['first_isolated_drift']} " + f"repeat_drift={case['first_repeat_drift']} batch_drift={case['first_batch_drift']} " + f"out err cand {case['output']['candidate_rel_err_vs_golden']:.2e} " + f"prov {case['output']['provider_rel_err_vs_golden']:.2e}" + ) + if not args.no_backward: + case = h3_block.run_block_backward( + weights, layout=layout, seed=args.seed, ops=ops, timing_repeats=args.timing_repeats + ) + backward.append(case) + det = all(v["repeat_bitwise_equal"] for v in case["leaves"].values()) + print( + f"{layout} backward: leaves {'det' if det else 'NONDET'}, " + f"grad_repeat_drift={case['first_grad_repeat_drift']} " + f"grad_batch_drift={case['first_grad_batch_drift']} " + f"dX batch-row {case['leaves']['hidden']['batch_row_bitwise_equal']}" + ) + torch.cuda.empty_cache() + + evidence = { + "kind": "h3_one_block_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "nodes": [ + { + "node": n.name, + "inputs": list(n.inputs), + "rfc_row": n.binding.rfc_row, + "status": n.binding.status, + "backend": n.binding.backend, + } + for n in h3_block.NODES + ], + "forward_cases": forward, + "backward_cases": backward, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_block_ulp_sweep.py b/tools/validation/models/h3_block_ulp_sweep.py new file mode 100644 index 000000000..7fe66646b --- /dev/null +++ b/tools/validation/models/h3_block_ulp_sweep.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Sweep the one-block backward over layouts and seeds and record leaf errors in ulps. + +Backs the ulp allowance of ``tests/models/minimax_h3/test_h3_one_block.py``: +for every leaf gradient it records the candidate's and diffusers' error against +the FP32 golden both as ``max|err| / max|golden|`` and in ulps of the leaf's +storage dtype at ``max|golden|``, and the per-leaf excess +``candidate_ulps - provider_ulps``. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_block_ulp_sweep.py --layouts tiny,small --seeds 0,1,2,3,4 \\ + --out ulp_sweep.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models import h3_block, h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import load_h3_block_weights # noqa: E402 + + +def main() -> None: + """Run the backward replay for every (layout, seed) and write per-leaf ulp errors.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--layouts", default="tiny,small") + parser.add_argument("--seeds", default="0,1,2,3,4") + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + if not torch.cuda.is_available(): + raise SystemExit("the sweep needs a CUDA device") + + weights = load_h3_block_weights("cuda") + ops = h3_block.CandidateOps() + cases = [] + for layout in args.layouts.split(","): + for seed in (int(s) for s in args.seeds.split(",")): + report = h3_block.run_block_backward(weights, layout=layout, seed=seed, ops=ops) + leaves = { + name: { + key: entry[key] + for key in ( + "storage_dtype", + "candidate_rel_err_vs_golden", + "provider_rel_err_vs_golden", + "candidate_ulps_vs_golden", + "provider_ulps_vs_golden", + ) + } + | { + "excess_ulps": entry["candidate_ulps_vs_golden"] + - entry["provider_ulps_vs_golden"] + } + for name, entry in report["leaves"].items() + } + worst = max(leaves, key=lambda n: leaves[n]["excess_ulps"]) + print( + f"{layout} S={report['seq_len']} seed={seed}: worst excess " + f"{leaves[worst]['excess_ulps']:+.2f} ulp ({worst})" + ) + cases.append( + {"layout": layout, "seq_len": report["seq_len"], "seed": seed, "leaves": leaves} + ) + torch.cuda.empty_cache() + + bf16 = [ + leaf["excess_ulps"] + for case in cases + for leaf in case["leaves"].values() + if leaf["storage_dtype"] == "bfloat16" + ] + evidence = { + "kind": "h3_one_block_backward_ulp_sweep", + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "cases": cases, + "max_bf16_excess_ulps": max(bf16), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"max BF16 excess: {max(bf16):+.2f} ulp; wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_candidate_prepare.py b/tools/validation/models/h3_candidate_prepare.py new file mode 100644 index 000000000..ae17b6600 --- /dev/null +++ b/tools/validation/models/h3_candidate_prepare.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fail closed on missing native H3 kernels and pin the candidate/reference identities.""" +import argparse +import json +import os +from pathlib import Path + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +from rl_engine.backends.extension import _C +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import require_candidate_cuda +from rl_engine.validation.models.h3_weights import load_h3_manifest, sha256_tensor +from tools.validation.models.h3_evidence_integrity import verify_inventory +from tools.validation.models.h3_full_model import digest, verify_checkpoint + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--checkpoint", type=Path, required=True) + p.add_argument("--reference-root", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + a = p.parse_args() + a.output.mkdir() + manifest = Path("examples/models/minimax_h3/reference-manifest.json") + verify_inventory(a.reference_root, manifest) + verify_checkpoint( + a.checkpoint, Path("rl_engine/validation/models/h3_full_manifest.json"), all_files=True + ) + if _C is None or not hasattr(_C, "h3_timestep_sinusoid_forward"): + raise ValueError("candidate extension lacks H3 kernels") + ops = CandidateOps() + backends = require_candidate_cuda(ops) + # Inject a real eager CUDA fallback and require the same guard used by the + # full-model path to reject it before it can count as candidate execution. + original = ops.norm + try: + ops.norm = torch.nn.RMSNorm(128).cuda() + try: + require_candidate_cuda(ops) + except ValueError as exc: + if "candidate norm fell back" not in str(exc): + raise + else: + raise AssertionError("eager CUDA norm fallback was accepted") + finally: + ops.norm = original + require_candidate_cuda(ops) + extension = Path(_C.__file__).resolve() + if not extension.is_relative_to(Path.cwd().resolve()): + raise ValueError("candidate loaded an extension outside its clean source snapshot") + spec = load_h3_manifest() + index = json.loads( + (a.checkpoint / "transformer/diffusion_pytorch_model.safetensors.index.json").read_text() + )["weight_map"] + weights = a.output / "weights" + weights.mkdir() + for section, filename in [ + ("tensors", "h3_conditioning.safetensors"), + ("block_tensors", "h3_block0.safetensors"), + ]: + values = {} + for shard in sorted({index[name] for name in spec[section]}): + with safe_open(str(a.checkpoint / "transformer" / shard), framework="pt") as handle: + for name, row in spec[section].items(): + if index[name] == shard: + value = handle.get_tensor(name).contiguous() + if sha256_tensor(value) != row["sha256"]: + raise ValueError(f"block weight identity changed: {name}") + values[name] = value + save_file(values, str(weights / filename)) + report = dict( + extension=str(extension), + extension_sha256=digest(extension), + torch=torch.__version__, + backends=backends, + negative_controls={"eager_cuda_norm_fallback_rejected": True}, + reference_manifest_sha256=digest(manifest), + build={ + key: os.environ.get(key) + for key in ("CUDA_HOME", "MAX_JOBS", "KERNEL_ALIGN_USE_FAST_MATH") + }, + ) + (a.output / "identity.json").write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_chain_replay.py b/tools/validation/models/h3_chain_replay.py new file mode 100644 index 000000000..11068259c --- /dev/null +++ b/tools/validation/models/h3_chain_replay.py @@ -0,0 +1,137 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the MiniMax-H3 conditioning chain and write JSON evidence (RFC #420). + +See ``rl_engine/validation/models/h3_chain.py`` for what is compared. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_chain_replay.py --out chain_replay.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models import h3_chain # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) + + +def _ints(text: str) -> list[int]: + """Parse a comma-separated list of integer case sizes.""" + + return [int(v) for v in text.split(",")] + + +def main() -> None: + """Replay requested CUDA stages with prerequisites and optionally save JSON evidence.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--stages", + default=",".join(s.name for s in h3_chain.STAGES), + help="comma list of stages to report; preceding stages run automatically", + ) + parser.add_argument("--timesteps", type=_ints, default=[1, 2, 4], help="comma list of T") + parser.add_argument("--seq-lens", type=_ints, default=[3, 257, 4097], help="comma list of S") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--out", type=Path, default=None) + parser.add_argument( + "--backward", + action="store_true", + help="also replay the chain backward (needs every stage): parameter-gradient " + "determinism and accuracy for the separate-op, fused and diffusers chains", + ) + args = parser.parse_args() + + wanted = args.stages.split(",") + stage_names = [s.name for s in h3_chain.STAGES] + unknown = sorted(set(wanted) - set(stage_names)) + if unknown: + parser.error(f"unknown stages: {unknown}") + stages = [s for s in h3_chain.STAGES if s.name in wanted] + if args.backward and len(stages) != len(h3_chain.STAGES): + parser.error("--backward requires every stage of the chain") + last_stage = max(i for i, stage in enumerate(h3_chain.STAGES) if stage.name in wanted) + execution_stages = h3_chain.STAGES[: last_stage + 1] + + if not torch.cuda.is_available(): + raise SystemExit("the chain replay needs a CUDA device") + + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + weights = load_h3_conditioning_weights("cuda") + registry = KernelRegistry() + manifest = load_h3_manifest() + + cases = [] + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_case( + registry, + weights, + num_timesteps=num_timesteps, + seq_len=seq_len, + seed=args.seed, + stages=execution_stages, + ) + case["stages"] = [entry for entry in case["stages"] if entry["stage"] in wanted] + cases.append(case) + summary = ", ".join( + f"{e['stage']}=" + f"{'eq' if e['chained_vs_provider']['bitwise_equal'] else 'drift'}" + f"(gold {e['chained_vs_golden']['max_abs']:.3e})" + for e in case["stages"] + ) + print(f"T={num_timesteps} S={seq_len}: {summary}; first_drift={case['first_drift']}") + + backward_cases = [] + if args.backward: + for num_timesteps in args.timesteps: + for seq_len in args.seq_lens: + case = h3_chain.run_backward_case( + registry, weights, num_timesteps=num_timesteps, seq_len=seq_len, seed=args.seed + ) + backward_cases.append(case) + parts = [] + for mode in h3_chain.BACKWARD_MODES: + entries = case["leaves"].values() + det = all(e[mode]["repeat_bitwise_equal"] for e in entries) + worst = max(e[mode]["max_abs_vs_golden_over_absmax"] for e in entries) + parts.append(f"{mode}: {'det' if det else 'NONDET'} {worst:.1e}") + summary = " | ".join(parts) + print(f"backward T={num_timesteps} S={seq_len}: {summary}") + + evidence = { + "kind": "h3_conditioning_chain_replay", + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weights_sha256": manifest["weight_shards"], + "reference_commit": manifest["reference_implementation"]["commit"], + **h3_chain.git_state(), + "environment": h3_chain.environment(), + "stages": [s.name for s in stages], + "executed_stages": [s.name for s in execution_stages], + "cases": cases, + "backward_cases": backward_cases, + } + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(evidence, indent=2) + "\n") + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence.py b/tools/validation/models/h3_evidence.py new file mode 100644 index 000000000..213d5beb3 --- /dev/null +++ b/tools/validation/models/h3_evidence.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write one RFC #420 operator's evidence report (performance + accuracy) as JSON. + +Run it on a clean tree; the report records the commit and whether the tree +was dirty. ``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_evidence.py --op timestep_sinusoid_h3 \\ + --out reports/experiments/h3-timestep-sinusoid-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.runtime.registry import KernelRegistry # noqa: E402 +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_report import ACCURACY, PERF_CASES, measure # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + WEIGHTS_ENV, + h3_weights_dir, + load_h3_manifest, +) + + +def main() -> None: + """Collect CUDA accuracy and performance evidence, requiring weights for weighted ops.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(ACCURACY)) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + uses_weights = args.op != "timestep_sinusoid_h3" + if uses_weights and h3_weights_dir() is None: + raise SystemExit(f"set {WEIGHTS_ENV} to the tools/weights/prepare_h3_weights.py output dir") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + registry = KernelRegistry() + manifest = load_h3_manifest() + report = { + "kind": "h3_operator_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + "weight_source": "pinned_checkpoint" if uses_weights else "not_applicable", + "reference_commit": manifest["reference_implementation"]["commit"], + **git_state(), + "environment": environment(), + "accuracy": ACCURACY[args.op](registry), + "perf": [measure(case) for case in PERF_CASES[args.op](registry)], + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_evidence_integrity.py b/tools/validation/models/h3_evidence_integrity.py new file mode 100644 index 000000000..42d9b383d --- /dev/null +++ b/tools/validation/models/h3_evidence_integrity.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Validate immutable reference inventories before consuming any tensor evidence.""" +from __future__ import annotations + +import json +from pathlib import Path + +from tools.validation.models.h3_full_model import digest + + +def verify_inventory(root: Path, manifest: Path): + records = json.loads(manifest.read_text())["files"] + for name, expected in records.items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()): + raise ValueError("evidence path escapes root") + if digest(path) != expected: + raise ValueError(f"reference evidence changed: {name}") + + +def verify_generation(root): + report = json.loads((root / "report.json").read_text()) + if report["status"] != "generation_complete_comparison_pending": + raise ValueError("generation did not complete") + required = {"initial.pt", "video.npy", "audio.npy"} | { + f"step-{i:03}.pt" for i in range(report["steps"]) + } + if not required.issubset(report["artifacts"]): + raise ValueError("incomplete generation artifact inventory") + for name, expected in report["artifacts"].items(): + path = root / name + if not path.resolve().is_relative_to(root.resolve()) or digest(path) != expected: + raise ValueError(f"generation artifact changed: {name}") + return report diff --git a/tools/validation/models/h3_first_step_diagnostic.py b/tools/validation/models/h3_first_step_diagnostic.py new file mode 100644 index 000000000..e433b8481 --- /dev/null +++ b/tools/validation/models/h3_first_step_diagnostic.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Replay the first denoising step of a generation under one execution configuration. + +The first transformer call of a generation sees the saved packed inputs (``initial.pt``) +and nothing has fed back yet, so the difference to the provider's saved ``step-000.pt`` +is the model's own error on a real generation input. A configuration is the provider +(optionally with one forced SDPA backend or cuBLASLt), the candidate, or the candidate +with chosen block nodes run as their provider replay (``--provider-nodes``). Replaying +the provider default must reproduce ``step-000.pt`` bitwise; ``baseline_bitwise`` records +whether it did, which validates the replay itself. Diagnostic only: nothing here accepts +or rejects a candidate. + + python -m tools.validation.models.h3_first_step_diagnostic --checkpoint \\ + --generation --seed 43 --mode candidate \\ + --provider-nodes attention --output out.json +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import torch + +from rl_engine.validation.models.h3_block import CandidateOps +from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution +from tools.validation.models.h3_full_generation import QueryChunkAttention +from tools.validation.models.h3_full_model import PROVIDER_VARIANTS, provider_variant_context + + +def errors(value: torch.Tensor, reference: torch.Tensor) -> dict[str, float]: + """max_abs and relative_l2 of ``value`` against ``reference`` (float64 norms).""" + + delta = value.double() - reference.double() + return dict( + max_abs=float(delta.abs().max()), + relative_l2=float(delta.norm() / reference.double().norm().clamp_min(1e-12)), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--generation", type=Path, required=True) + parser.add_argument("--contract", type=Path, required=True) + parser.add_argument("--seed", type=int, required=True) + parser.add_argument("--mode", choices=["provider", "candidate"], required=True) + parser.add_argument("--provider-variant", choices=PROVIDER_VARIANTS, default="default") + parser.add_argument("--provider-nodes", default="") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + if args.mode == "candidate" and args.provider_variant != "default": + parser.error("--provider-variant applies to --mode provider only") + if args.mode == "provider" and args.provider_nodes: + parser.error("--provider-nodes applies to --mode candidate only") + nodes = tuple(n for n in args.provider_nodes.split(",") if n) + + from diffusers import MiniMaxH3Transformer3DModel + + directory = args.generation / str(args.seed) + inputs = torch.load(directory / "initial.pt", map_location="cpu", weights_only=True) + saved = torch.load(directory / "step-000.pt", map_location="cpu", weights_only=True) + contract = json.loads(args.contract.read_text()) + + started = time.perf_counter() + model = MiniMaxH3Transformer3DModel.from_pretrained( + str(args.checkpoint / "transformer"), torch_dtype=torch.bfloat16, local_files_only=True + ) + model = model.to("cuda").eval() + loaded = time.perf_counter() - started + kwargs = {k: v.to("cuda") for k, v in inputs.items()} + ops = None + if args.mode == "candidate": + ops = CandidateOps() + ops.attention = QueryChunkAttention(ops.attention, contract["attention_query_chunk"]) + + trace = FullModelTrace() + started = time.perf_counter() + with ( + torch.inference_mode(), + full_model_execution( + model, + candidate=args.mode == "candidate", + trace=trace, + ops=ops, + provider_nodes=nodes, + ), + provider_variant_context(args.provider_variant), + ): + video, audio = model(**kwargs, attention_kwargs=None, return_dict=False) + torch.cuda.synchronize() + forward_seconds = time.perf_counter() - started + video, audio = video.float().cpu(), audio.float().cpu() + result = dict( + seed=args.seed, + mode=args.mode, + provider_variant=args.provider_variant, + provider_nodes=list(nodes), + video=errors(video, saved[0]), + audio=errors(audio, saved[1]), + baseline_bitwise=bool(torch.equal(video, saved[0]) and torch.equal(audio, saved[1])), + load_seconds=loaded, + forward_seconds=forward_seconds, + torch=torch.__version__, + device=torch.cuda.get_device_name(), + ) + args.output.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_compare.py b/tools/validation/models/h3_full_compare.py new file mode 100644 index 000000000..6f1645e67 --- /dev/null +++ b/tools/validation/models/h3_full_compare.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Freeze full-H3 reference error bounds, or evaluate a candidate without changing them.""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from tools.validation.models.h3_full_model import digest + + +def tensors(root): + report = json.loads((root / "report.json").read_text()) + if digest(root / "tensors.pt") != report["tensors_sha256"]: + raise ValueError("full-model tensor evidence changed") + payload = torch.load(root / "tensors.pt", map_location="cpu", weights_only=True) + yield "video", payload["video"] + yield "audio", payload["audio"] + for section in ("nodes", "input_gradients"): + for name, value in sorted(payload[section].items()): + yield section + "." + name, value + for name, row in sorted(json.loads((root / "gradients.json").read_text()).items()): + path = root / row["path"] + if digest(path) != row["sha256"]: + raise ValueError(f"gradient artifact changed: {path}") + yield "parameter." + name, torch.load(path, map_location="cpu", weights_only=True) + + +def compare(left, right): + a, b = tensors(left), tensors(right) + from itertools import zip_longest + + for x, y in zip_longest(a, b): + if x is None or y is None or x[0] != y[0]: + raise ValueError("different full-model tensor inventories") + name, value = x + other = y[1] + if ( + value.shape != other.shape + or not torch.isfinite(value).all() + or not torch.isfinite(other).all() + ): + raise ValueError(f"nonfinite/shape mismatch: {name}") + delta = value.float() - other.float() + yield name, dict( + max_abs=delta.abs().max().item(), + relative_l2=(delta.double().norm() / value.double().norm().clamp_min(1e-12)).item(), + ) + + +def drift_order(report, names): + trace = report["trace"] + if "execution_order" in trace: + observed = trace["execution_order"] + return [name for name in observed if name in names] + [ + name for name in names if name not in observed + ] + forward = trace.get("forward_order", list(trace["calls"])) + backward = trace.get("backward_order", list(trace["backward_calls"])) + ordered = ["nodes." + name for name in forward] + ordered += ["nodes." + name + ".grad" for name in backward] + ordered += [name for name in names if name not in ordered] + return [name for name in ordered if name in names] + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--reference", type=Path, required=True) + p.add_argument("--other", type=Path, required=True) + p.add_argument("--output", type=Path, required=True) + p.add_argument("--thresholds", type=Path) + p.add_argument("--reference-manifest", type=Path) + p.add_argument("--reference-root", type=Path) + args = p.parse_args() + if args.output.exists(): + raise ValueError("never overwrite calibration or failed evidence") + from tools.validation.models.h3_evidence_integrity import verify_inventory + + if args.thresholds and not (args.reference_root and args.reference_manifest): + raise ValueError("candidate comparison requires pinned reference inventory") + if args.reference_manifest: + verify_inventory(args.reference_root, args.reference_manifest) + ref = json.loads((args.reference / "report.json").read_text()) + other = json.loads((args.other / "report.json").read_text()) + if ref["mode"] != "fp32-reference" or ref["status"] != "execution_complete_acceptance_pending": + raise ValueError("independent completed FP32 reference required") + if other["status"] != "execution_complete_acceptance_pending": + raise ValueError("other execution is incomplete") + reference_sha = digest(args.reference / "report.json") + if args.thresholds: + contract = json.loads(args.thresholds.read_text()) + if other["mode"] != "candidate" or contract["reference_report_sha256"] != reference_sha: + raise ValueError("wrong candidate/reference contract") + if other.get("thresholds_sha256") != digest(args.thresholds): + raise ValueError("thresholds were not frozen before candidate execution") + elif other["mode"] != "bf16-provider": + raise ValueError("only independent provider can establish thresholds") + rows = dict(compare(args.reference, args.other)) + report = dict( + reference_report_sha256=reference_sha, + other_report_sha256=digest(args.other / "report.json"), + ) + if args.thresholds: + if rows.keys() != contract["bounds"].keys(): + raise ValueError("threshold tensor inventory differs") + failures = [ + name + for name in drift_order(other if "execution_order" in other["trace"] else ref, rows) + if any(rows[name][k] > contract["bounds"][name][k] for k in rows[name]) + ] + hybrid = other.get("diagnostic_provider_nodes") or [] + report.update( + status="diagnostic_only" if hybrid else ("failed" if failures else "passed"), + diagnostic_provider_nodes=hybrid, + errors=rows, + failures=failures, + first_drift=failures[0] if failures else None, + thresholds_sha256=digest(args.thresholds), + ) + else: + # Committed before candidate execution. Forward and gradient tensors use + # separate floors; the existing block-0 golden contract is unchanged. + bounds = { + name: { + key: 1.5 * value + + ( + 5e-3 + if ".grad" in name + or name.startswith("parameter.") + or name.startswith("input_gradients.") + else 1e-3 + ) + for key, value in row.items() + } + for name, row in rows.items() + } + report.update( + status="frozen_reference_bounds", + rule="1.5 * independent BF16 provider error + fixed floor", + bounds=bounds, + errors=rows, + ) + args.output.write_text(json.dumps(report, indent=2) + "\n") + if args.thresholds and failures and not report.get("diagnostic_provider_nodes"): + raise SystemExit(1) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_ensemble_rule.py b/tools/validation/models/h3_full_ensemble_rule.py new file mode 100644 index 000000000..72c397533 --- /dev/null +++ b/tools/validation/models/h3_full_ensemble_rule.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Full-model acceptance from an ensemble of independent BF16 provider configurations. + +The single-configuration rule (``1.5 * one provider's error + floor``) accepts only the +configuration it was frozen from: providers that differ only in their SDPA backend or +BLAS library fail it on tens of tensors. This rule is calibrated on several provider +configurations, compared with the same frozen FP32 reference, and contains no candidate +data: + +* E[t] is the largest ``relative_l2`` error any configuration has on tensor ``t``; +* the bound is ``k * E[t] + floor`` (floors as in the single-configuration rule); +* ``k`` is the smallest factor, rounded up to 0.1, for which every configuration held + out passes against the maximum of the others (leave-one-out); +* ``max_abs``, set by one element, is reported but not gated: held-out configurations + exceed the others' maximum by several times on it. + + h3_full_ensemble_rule.py calibrate --provider NAME=comparison.json ... \\ + --output rule.json + h3_full_ensemble_rule.py check --rule rule.json --candidate comparison.json \\ + --output out.json +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math +from pathlib import Path + + +def _floor(name: str) -> float: + gradient = ( + ".grad" in name or name.startswith("parameter.") or name.startswith("input_gradients.") + ) + return 5e-3 if gradient else 1e-3 + + +def _digest(path: Path) -> str: + return hashlib.sha256(path.read_bytes()).hexdigest() + + +def calibrate(args) -> None: + """Build the ensemble bounds and the leave-one-out factor from provider comparisons.""" + + providers, sources = {}, {} + for item in args.provider: + name, _, path = item.partition("=") + report = json.loads(Path(path).read_text()) + if report.get("diagnostic_provider_nodes"): + raise ValueError(f"{name} is a hybrid diagnostic, not a provider configuration") + providers[name] = report["errors"] + sources[name] = { + "sha256": _digest(Path(path)), + "reference": report["reference_report_sha256"], + } + if len({s["reference"] for s in sources.values()}) != 1: + raise ValueError("provider comparisons use different FP32 references") + # Configurations whose errors are identical to an earlier one are one configuration. + distinct = {} + for name in sorted(providers): + if not any(providers[name] == providers[kept] for kept in distinct): + distinct[name] = providers[name] + if len(distinct) < 3: + raise ValueError("need at least three distinct provider configurations") + names = sorted(next(iter(distinct.values()))) + loo = {} + for held in distinct: + others = [distinct[o] for o in distinct if o != held] + worst = 0.0 + for t in names: + top = max(o[t]["relative_l2"] for o in others) + if top > 0: + worst = max(worst, (distinct[held][t]["relative_l2"] - _floor(t)) / top) + loo[held] = worst + factor = math.ceil(max(loo.values()) * 10) / 10 + ensemble = {t: max(d[t]["relative_l2"] for d in distinct.values()) for t in names} + rule = { + "status": "frozen_ensemble_rule", + "rule": "relative_l2 <= k * max over provider configurations + floor; max_abs reported", + "factor_k": factor, + "leave_one_out_worst_ratio": loo, + "providers": sorted(distinct), + "duplicates_dropped": sorted(set(providers) - set(distinct)), + "sources": sources, + "reference_report_sha256": next(iter(sources.values()))["reference"], + "ensemble_relative_l2": ensemble, + "ensemble_max_abs": {t: max(d[t]["max_abs"] for d in distinct.values()) for t in names}, + "bounds_relative_l2": {t: factor * ensemble[t] + _floor(t) for t in names}, + } + args.output.write_text(json.dumps(rule, indent=2) + "\n") + print( + f"k = {factor} from leave-one-out {loo}; " + f"{len(distinct)} configurations, {len(names)} tensors" + ) + + +def check(args) -> None: + """Apply a frozen ensemble rule to one candidate comparison; exit 1 on any failure.""" + + rule = json.loads(args.rule.read_text()) + report = json.loads(args.candidate.read_text()) + if report["reference_report_sha256"] != rule["reference_report_sha256"]: + raise ValueError("candidate was compared with a different FP32 reference") + if report.get("diagnostic_provider_nodes"): + raise ValueError("a hybrid diagnostic is not a candidate result") + errors = report["errors"] + if errors.keys() != rule["bounds_relative_l2"].keys(): + raise ValueError("tensor inventory differs from the rule") + failures = [ + t for t in sorted(errors) if errors[t]["relative_l2"] > rule["bounds_relative_l2"][t] + ] + ratios = {t: errors[t]["relative_l2"] / rule["bounds_relative_l2"][t] for t in errors} + max_abs_ratio = { + t: errors[t]["max_abs"] / rule["ensemble_max_abs"][t] + for t in errors + if rule["ensemble_max_abs"][t] + } + out = { + "status": "failed" if failures else "passed", + "rule_sha256": _digest(args.rule), + "candidate_comparison_sha256": _digest(args.candidate), + "failures": failures, + "worst_relative_l2_over_bound": sorted(ratios.items(), key=lambda x: -x[1])[:10], + "outputs": {t: errors[t] for t in ("video", "audio") if t in errors}, + "max_abs_over_ensemble_max_reported": sorted(max_abs_ratio.items(), key=lambda x: -x[1])[ + :10 + ], + } + args.output.write_text(json.dumps(out, indent=2) + "\n") + print( + f"{out['status']}: {len(failures)} failures; worst {out['worst_relative_l2_over_bound'][0]}" + ) + if failures: + raise SystemExit(1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + sub = parser.add_subparsers(dest="cmd", required=True) + cal = sub.add_parser("calibrate") + cal.add_argument("--provider", action="append", required=True, help="name=comparison.json") + cal.add_argument("--output", type=Path, required=True) + chk = sub.add_parser("check") + chk.add_argument("--rule", type=Path, required=True) + chk.add_argument("--candidate", type=Path, required=True) + chk.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + calibrate(args) if args.cmd == "calibrate" else check(args) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_full_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_prior_art.py b/tools/validation/models/h3_prior_art.py new file mode 100644 index 000000000..9012542e2 --- /dev/null +++ b/tools/validation/models/h3_prior_art.py @@ -0,0 +1,946 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Compare one RFC #420 operator with existing implementations: batch invariance, accuracy, latency. + +The RFC #420 reuse rule asks, for every operator, whether an existing implementation is +batch-invariant (then reuse it) and for accuracy and performance comparisons. This runner +measures the operator's rl-kernel CUDA op next to the existing implementations that import in +the current environment, and writes one JSON report per operator: + + python tools/validation/models/h3_prior_art.py --op adaln_row_gather \\ + --out reports/experiments/h3-prior-art-b200/adaln_row_gather.json \\ + [--megatron-src /path/to/Megatron-LM] [--quick] + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/adaln_row_gather.json + +Only operators whose rl-kernel op exists on the current branch can be selected. Optional +libraries (SGLang, Liger, Transformer Engine, vLLM, Megatron-LM from a source tree) are skipped +and recorded as unavailable when they do not import. Global batch-invariant modes (vLLM, +SGLang, Megatron) patch aten for the whole process, so each runs in its own subprocess, with +the environment variables it relies on set before CUDA initialises. + +Batch invariance is bitwise, for the forward, the per-row gradients, and repeatability of every +parameter and table gradient, with three checks: + +1. every row computed alone vs inside full batches of 64, 257 and 2048 rows, seeds 3-5; +2. the full workload-size batch vs sub-batches that together cover every row; +3. eight probe rows at the front, middle and back of batches of every size 1..9 and + 2^k - 1, 2^k, 2^k + 1, plus the whole batch reversed. + +Sparse probes (a few rows, a few batch sizes) miss row-specific or large-batch dependence, which +is why every row is checked. Accuracy is max|err| / max|ref| against the same computation in +FP64; latency is the median of 100 CUDA-event samples after 20 warm-ups. Run on a clean tree on +an otherwise idle GPU; the report records the commit, the tree state and every library version. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib.metadata as md +import importlib.util +import json +import os +import statistics +import subprocess +import sys +from pathlib import Path +from typing import Any, Callable + +REPO_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(REPO_ROOT)) + +import torch # noqa: E402 +import torch.nn.functional as F # noqa: E402 + +H, D, F1, T3, EPS = 5376, 2688, 256, 3, 1e-5 + +#: operator -> rl-kernel module that must exist on the branch +OPS = { + "timestep_sinusoid": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid", + "timestep_mlp": "rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp", + "adaln_projection": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection", + "adaln_row_gather": "rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather", + "norm_modulate": "rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm", + "gate_residual": "rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual", + "final_adaln_out": "rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out", +} +#: global batch-invariant modes compared for the GEMM operators, with their environment +MODES = { + "plain": {}, + "vllm": {"CUBLAS_WORKSPACE_CONFIG": ":16:8", "CUBLASLT_WORKSPACE_SIZE": "1"}, + "sglang": {}, + "sglang_ieee": {"TRITON_F32_DEFAULT": "ieee"}, + "megatron_te_native": {"CUBLASLT_WORKSPACE_SIZE": "0"}, + "megatron_triton": {}, + "megatron_triton_ieee": {"TRITON_F32_DEFAULT": "ieee"}, +} +GEMM_OPS = ("timestep_mlp", "adaln_projection") +LIBS = ("torch", "triton", "diffusers", "vllm", "flashinfer-python", "sglang", "liger-kernel") +LIBS += ("transformer_engine", "megatron-core") + +DEV = "cuda" + + +# --------------------------------------------------------------------------- # +# Batch-invariance checks +# --------------------------------------------------------------------------- # + + +class _Case: + """``fn(**inputs)`` with row arguments sliced together; fixed upstream gradients per seed.""" + + def __init__(self, fn, inputs, rows, grads, out_dims, seed): + self.fn, self.inputs, self.rows, self.grads, self.out_dims = ( + fn, + inputs, + rows, + grads, + out_dims, + ) + self.n = next(inputs[k].shape[d] for k, d in rows.items()) + outs = self._flat(fn(**inputs)) + gen = torch.Generator(device=DEV).manual_seed(seed) + self.ups = [torch.randn(o.shape, device=DEV, generator=gen).to(o.dtype) for o in outs] + + @staticmethod + def _flat(out): + return list(out) if isinstance(out, (tuple, list)) else [out] + + def run(self, idx=None): + args = {} + for k, v in self.inputs.items(): + if idx is not None and k in self.rows: + v = v.index_select(self.rows[k], idx) + if k in self.grads: + v = v.detach().clone().requires_grad_(True) + args[k] = v + outs = self._flat(self.fn(**args)) + ups = ( + self.ups + if idx is None + else [u.index_select(d, idx) for u, d in zip(self.ups, self.out_dims)] + ) + live = [(o, u) for o, u in zip(outs, ups) if o.requires_grad] + if live: + torch.autograd.backward([o for o, _ in live], [u for _, u in live]) + grads = {k: args[k].grad for k in self.grads} + return [o.detach() for o in outs], grads + + +def _sweep_sizes(n: int) -> list[int]: + sizes = set(range(1, min(n, 9) + 1)) + k = 16 + while k <= n: + sizes.update(v for v in (k - 1, k, k + 1) if v <= n) + k *= 2 + return sorted(sizes | {n}) + + +def batch_invariance(cand: dict[str, Any], quick: bool) -> dict[str, Any]: + fn, make, rows, grads = cand["fn"], cand["make"], cand["rows"], cand["grads"] + out_dims = cand.get("out_dims", (0,)) + params = [k for k in grads if k not in rows] + full_sizes, big, subs = cand["bi_sizes"] + if quick: + full_sizes, big, subs = (17, 64), 256, (1, 7, 64) + seeds = (3, 4, 5) + result: dict[str, Any] = {"all_rows": {}, "params_repeatable": None} + ok = True + + def same_row(a, b, d, i): + return torch.equal(a.select(d, 0), b.select(d, i)) + + for n in full_sizes: + fwd = grd = 0 + repeat = True + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + fo, fg = case.run() + for i in range(n): + o, g = case.run(torch.tensor([i], device=DEV)) + fwd += not all(same_row(a, b, d, i) for a, b, d in zip(o, fo, out_dims)) + grd += not all( + same_row(g[k], fg[k], rows[k], i) + for k in rows + if k in grads and g[k] is not None + ) + if params: + _, again = case.run() + repeat &= all(torch.equal(again[k], fg[k]) for k in params if fg[k] is not None) + result["all_rows"][str(n)] = { + "rows": len(seeds) * n, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + if params: + result["params_repeatable"] = repeat and result["params_repeatable"] is not False + ok &= fwd == 0 and grd == 0 and repeat + + fwd = grd = checked = 0 + for seed in seeds[:2]: + case = _Case(fn, make(seed, big), rows, grads, out_dims, seed) + fo, fg = case.run() + for sb in subs: + step = sb if sb >= 1024 else max(sb, big // 512) + for start in range(0, big, step): + idx = torch.arange(start, min(start + sb, big), device=DEV) + o, g = case.run(idx) + checked += 1 + fwd += not all( + torch.equal(a, b.index_select(d, idx)) for a, b, d in zip(o, fo, out_dims) + ) + grd += not all( + torch.equal(g[k], fg[k].index_select(rows[k], idx)) + for k in rows + if k in grads and g[k] is not None + ) + del case, fo, fg + torch.cuda.empty_cache() + result["full_vs_sub_batches"] = { + "full_batch": big, + "sub_batch_sizes": list(subs), + "sub_batches": checked, + "fwd_differ": fwd, + "rowgrad_differ": grd, + } + ok &= fwd == 0 and grd == 0 + + n = full_sizes[-1] + sizes = _sweep_sizes(n) + probes = sorted({0, n - 1, *list(range(n // 8 + 3, n - 1, n // 8))[:6]}) + compared = failed = 0 + first: list[dict[str, int]] = [] + for seed in seeds: + case = _Case(fn, make(seed, n), rows, grads, out_dims, seed) + alone = {r: case.run(torch.tensor([r], device=DEV)) for r in probes} + gen = torch.Generator().manual_seed(seed) + + def compare(idx, pos, r, seed=seed, case=case, alone=alone): + nonlocal compared, failed + o, g = case.run(idx.to(DEV)) + ao, ag = alone[r] + good = all( + torch.equal(a.select(d, pos), b.select(d, 0)) for a, b, d in zip(o, ao, out_dims) + ) + good &= all( + torch.equal(g[k].select(rows[k], pos), ag[k].select(rows[k], 0)) + for k in rows + if k in grads and g[k] is not None + ) + compared += 1 + if not good: + failed += 1 + if len(first) < 8: + first.append( + {"seed": seed, "batch": int(idx.numel()), "position": pos, "row": r} + ) + + for m in sizes: + for r in probes: + others = torch.randperm(n, generator=gen) + others = others[others != r][: m - 1] + for pos in sorted({0, (m - 1) // 2, m - 1}): + compare(torch.cat([others[:pos], torch.tensor([r]), others[pos:]]), pos, r) + for r in probes: + compare(torch.arange(n - 1, -1, -1), n - 1 - r, r) + result["batch_size_sweep"] = { + "batch_sizes": sizes, + "probe_rows": probes, + "comparisons": compared, + "differ": failed, + "first_failures": first, + } + ok &= failed == 0 + result["batch_invariant"] = bool(ok) + return result + + +# --------------------------------------------------------------------------- # +# Accuracy and latency +# --------------------------------------------------------------------------- # + + +def _time_us(fn: Callable[[], Any], warmup: int = 20, iters: int = 100) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + samples = [] + for _ in range(iters): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1e3) + return statistics.median(samples) + + +def _rel(a: torch.Tensor, ref: torch.Tensor) -> float: + return ((a.double() - ref).abs().max() / ref.abs().max().clamp_min(1e-300)).item() + + +def accuracy_latency(cand: dict[str, Any], mode_off, quick: bool) -> dict[str, Any]: + out: dict[str, Any] = {} + for size in cand["perf_sizes"][:2] if quick else cand["perf_sizes"]: + inputs = cand["make"](11, size) + names = list(inputs) + grads = [k for k in cand["grads"] if k in names] + leaves = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + y = cand["fn"](**leaves) + y = torch.cat([t.reshape(-1) for t in y]) if isinstance(y, (tuple, list)) else y + with mode_off(): + ref_in = { + k: (v.detach().double().requires_grad_(k in grads) if v.is_floating_point() else v) + for k, v in inputs.items() + } + ref = cand["ref"](**ref_in) + ref = torch.cat([t.reshape(-1) for t in ref]) if isinstance(ref, (tuple, list)) else ref + row: dict[str, Any] = {"fwd_rel_err": _rel(y, ref.detach())} + if grads: + gen = torch.Generator(device=DEV).manual_seed(9) + up = torch.randn(y.shape, device=DEV, generator=gen).to(y.dtype) + y.backward(up) + with mode_off(): + ref.backward(up.double()) + row["grad_rel_err"] = {k: _rel(leaves[k].grad, ref_in[k].grad) for k in grads} + row["fwd_us"] = _time_us(lambda inputs=inputs: cand["fn"](**inputs)) + if grads: + + def step(inputs=inputs): + lv = { + k: (v.detach().clone().requires_grad_(True) if k in grads else v) + for k, v in inputs.items() + } + o = cand["fn"](**lv) + outs = list(o) if isinstance(o, (tuple, list)) else [o] + torch.autograd.backward(outs, [torch.ones_like(t) for t in outs]) + + row["fwd_bwd_us"] = _time_us(step) + out[str(size)] = row + del inputs, leaves, y, ref + torch.cuda.empty_cache() + return out + + +# --------------------------------------------------------------------------- # +# Candidates +# --------------------------------------------------------------------------- # + + +def _gen(seed: int, salt: int) -> torch.Generator: + return torch.Generator(device=DEV).manual_seed(seed * 1000 + salt) + + +def _rn(seed, salt, *shape, dtype=torch.float32, scale=1.0): + return (torch.randn(*shape, device=DEV, generator=_gen(seed, salt)) * scale).to(dtype) + + +def _layout(n: int, seed: int): + from rl_engine.validation.models.h3_cases import h3_packed_layout + + ti, tags = h3_packed_layout(n, T3, seed=seed) + return ti.contiguous(), tags.contiguous() + + +def _version(name: str) -> str | None: + try: + return md.version(name) + except md.PackageNotFoundError: + return None + + +T_BI = ((64, 257, 2048), 4096, (1, 7, 1024, 2048)) +S_BI = ((64, 257, 2048), 131072, (1, 7, 4097, 32768, 65536)) +T_PERF = (1, 3, 4, 64, 256, 2048) +S_PERF = (4097, 32768) + + +def _sinusoid_ref(t): + half = F1 // 2 + exponent = -torch.log(torch.tensor(10000.0, dtype=torch.float64, device=DEV)) + freqs = torch.exp(exponent * torch.arange(half, device=DEV, dtype=torch.float64) / half) + e = t.double()[:, None] * freqs[None] + return torch.cat([torch.cos(e), torch.sin(e)], -1) + + +def candidates(op: str, mode: str) -> list[dict[str, Any]]: + """Candidate dicts; a factory that raises (missing library) becomes an unavailable entry.""" + + from rl_engine.validation.models import h3_provider as P + + plain = mode == "plain" + found: list[tuple[str, Callable[[], dict[str, Any]]]] = [] + + if op == "timestep_sinusoid": + base = { + "make": lambda s, n: {"t": torch.rand(n, device=DEV, generator=_gen(s, 1))}, + "rows": {"t": 0}, + "grads": [], + "ref": _sinusoid_ref, + "bi_sizes": ((64, 257, 2048), 8192, (1, 7, 2048, 4097)), + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers get_timestep_embedding (op-for-op replay)", + "fn": lambda t: P.provider_time_proj(t), + } + + def sglang(): + from sglang.kernels.ops.diffusion.modulate.timestep_embedding_jit import ( + timestep_embedding, + ) + + return { + **base, + "source": f"SGLang {_version('sglang')} timestep_embedding", + "fn": lambda t: timestep_embedding( + t, F1, flip_sin_to_cos=True, downscale_freq_shift=0.0 + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_sinusoid import ( + H3TimestepSinusoidCudaOp, + ) + + op_ = H3TimestepSinusoidCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepSinusoidCudaOp (check_range=False)", + "fn": lambda t: op_.forward(t, check_range=False), + } + + found = [("diffusers", diffusers), ("sglang", sglang), ("rl_kernel", ours)] + + elif op == "timestep_mlp": + + def make(s, n): + return { + "x": _rn(s, 2, n, F1), + "w1": _rn(1, 3, H, F1, scale=F1**-0.5), + "b1": _rn(1, 4, H, scale=0.02), + "w2": _rn(1, 5, D, H, scale=H**-0.5), + "b2": _rn(1, 6, D, scale=0.02), + } + + base = { + "make": make, + "rows": {"x": 0}, + "grads": ["x", "w1", "b1", "w2", "b2"], + "ref": lambda x, w1, b1, w2, b2: F.linear(F.silu(F.linear(x, w1, b1)), w2, b2), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers TimestepEmbedding, FP32 F.linear [{mode}]", + "fn": lambda x, w1, b1, w2, b2: P.provider_time_embedder(x, w1, b1, w2, b2), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.timestep_mlp import ( + H3TimestepMLPCudaOp, + ) + + op_ = H3TimestepMLPCudaOp() + return { + **base, + "source": "rl-kernel H3TimestepMLPCudaOp", + "fn": lambda x, w1, b1, w2, b2: op_(x, w1, b1, w2, b2), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_projection": + + def make(s, n): + return { + "temb": _rn(s, 7, n, D), + "w": _rn(1, 8, 18 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 9, 18 * H, dtype=torch.bfloat16, scale=0.02), + } + + base = { + "make": make, + "rows": {"temb": 0}, + "grads": ["temb", "w", "b"], + "ref": lambda temb, w, b: F.linear(F.silu(temb), w, b), + "bi_sizes": T_BI, + "perf_sizes": T_PERF, + } + + def diffusers(): + return { + **base, + "source": f"diffusers AdaLN projection, BF16 F.linear [{mode}]", + "fn": lambda temb, w, b: F.linear(F.silu(temb).to(w.dtype), w, b), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_projection import ( + H3AdaLNProjectionCudaOp, + ) + + op_ = H3AdaLNProjectionCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNProjectionCudaOp", + "fn": lambda temb, w, b: op_.forward_table(temb, w, b), + } + + found = [(f"diffusers[{mode}]", diffusers)] + ([("rl_kernel", ours)] if plain else []) + + elif op == "adaln_row_gather": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "table": _rn(1, 14, 3 * T3, 6 * H, dtype=torch.bfloat16, scale=0.1), + "ti": ti, + "tags": tags, + } + + def ref(table, ti, tags): + return tuple(table.index_select(0, ti * 3 + tags).chunk(6, dim=1)) + + base = { + "make": make, + "rows": {"ti": 0, "tags": 0}, + "grads": ["table"], + "out_dims": (0,) * 6, + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return {**base, "source": "diffusers six index_select calls", "fn": ref} + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + op_ = H3AdaLNRowGatherCudaOp() + return { + **base, + "source": "rl-kernel H3AdaLNRowGatherCudaOp", + "fn": lambda table, ti, tags: op_(table, ti, tags), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "norm_modulate": + + def make(s, n): + ti, tags = _layout(n, s) + w = (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 11))).bfloat16() + return { + "x": _rn(s, 10, n, H, dtype=torch.bfloat16, scale=2.0), + "w": w, + "sh": _rn(1, 12, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "sc": _rn(1, 13, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(x, w, sh, sc, idx): + n = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + return n * (1 + sc.index_select(0, idx)) + sh.index_select(0, idx) + + base = { + "make": make, + "rows": {"x": 0, "idx": 0}, + "grads": ["x", "w", "sh", "sc"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def per_token(fn): + # Liger / SGLang take per-token shift/scale rows: gather them first (index_select). + return lambda x, w, sh, sc, idx: fn( + x, w, sh.index_select(0, idx), sc.index_select(0, idx) + ) + + def diffusers(): + return { + **base, + "source": "diffusers composition (F.rms_norm + index_select modulation)", + "fn": lambda x, w, sh, sc, idx: P.provider_norm_modulate(x, w, sh, sc, idx), + } + + def liger(): + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + return { + **base, + "source": f"Liger {_version('liger-kernel')} modulated RMSNorm + index_select", + "fn": per_token( + lambda x, w, shr, scr: LigerModulatedRMSNormFunction.apply( + x, w, scr, shr, EPS, 0.0, "llama", False + ) + ), + } + + def liger_rlk_gather(): + # Same Liger kernel, with the rows gathered by rl-kernel's deterministic row gather + # (adaln_row_gather) instead of index_select: separates Liger from the gather. + from liger_kernel.ops.modulated_rms_norm import LigerModulatedRMSNormFunction + + from rl_engine.backends.cuda.model_specific.minimax_h3.adaln_row_gather import ( + H3AdaLNRowGatherCudaOp, + ) + + gather = H3AdaLNRowGatherCudaOp() + + def fn(x, w, sh, sc, idx): + pad = torch.zeros_like(sh) + table = torch.cat([sh, sc, pad, pad, pad, pad], dim=1) + rows = gather(table, torch.div(idx, 3, rounding_mode="floor"), idx % 3) + return LigerModulatedRMSNormFunction.apply( + x, w, rows[1], rows[0], EPS, 0.0, "llama", False + ) + + version = _version("liger-kernel") + return { + **base, + "source": f"Liger {version} modulated RMSNorm + rl-kernel row gather", + "fn": fn, + } + + def sglang(): + from sglang.kernels.ops.diffusion.norm.scale_residual_norm_cutedsl import ( + fused_norm_scale_shift, + ) + + return { + **base, + "grads": [], + "source": f"SGLang {_version('sglang')} fused_norm_scale_shift (forward only)", + "fn": per_token( + lambda x, w, shr, scr: fused_norm_scale_shift( + x[None], w, None, scr[None], shr[None], "rms", EPS + )[0] + ), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.rmsnorm import H3RMSNormCudaOp + + op_ = H3RMSNormCudaOp() + return { + **base, + "source": "rl-kernel H3RMSNormCudaOp.forward_modulated", + "fn": lambda x, w, sh, sc, idx: op_.forward_modulated(x, w, sh, sc, idx), + } + + def norm_only(make): + def m(s, n): + d = make(s, n) + return {"x": d["x"], "w": d["w"]} + + return m + + def norm_ref(x, w): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + EPS) * w + + def torch_norm(): + # diffusers' norm without the modulation, to show which half is batch-invariant + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x", "w"], + "ref": norm_ref, + "source": "torch F.rms_norm (no modulation)", + "fn": lambda x, w: F.rms_norm(x, (H,), w, EPS), + } + + def te_norm(): + import transformer_engine.pytorch as te + + mod = te.RMSNorm(H, eps=EPS, params_dtype=torch.bfloat16, device=DEV) + + def fn(x, w): + # TE's backward reads its own parameter, so copy w in (functional_call + # gives a wrong dx); dweight is then TE's parameter gradient, not checked. + with torch.no_grad(): + mod.weight.copy_(w) + return mod(x) + + return { + **base, + "make": norm_only(make), + "rows": {"x": 0}, + "grads": ["x"], + "ref": norm_ref, + "source": f"TE {_version('transformer_engine')} RMSNorm (no modulation)", + "fn": fn, + } + + found = [ + ("diffusers", diffusers), + ("torch_rms_norm", torch_norm), + ("transformer_engine", te_norm), + ("liger", liger), + ("liger_rlk_gather", liger_rlk_gather), + ("sglang", sglang), + ("rl_kernel", ours), + ] + + elif op == "gate_residual": + + def make(s, n): + ti, tags = _layout(n, s) + return { + "res": _rn(s, 15, n, H, dtype=torch.bfloat16), + "y": _rn(s, 16, n, H, dtype=torch.bfloat16), + "gate": _rn(1, 17, 3 * T3, H, dtype=torch.bfloat16, scale=0.1), + "idx": ti * 3 + tags, + } + + def ref(res, y, gate, idx): + return res + gate.index_select(0, idx) * y + + base = { + "make": make, + "rows": {"res": 0, "y": 0, "idx": 0}, + "grads": ["res", "y", "gate"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers residual + gate.index_select(...) * y", + "fn": lambda res, y, gate, idx: P.provider_gate_residual(res, gate, idx, y), + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.gate_residual import ( + H3GateResidualCudaOp, + ) + + op_ = H3GateResidualCudaOp() + return { + **base, + "source": "rl-kernel H3GateResidualCudaOp", + "fn": lambda res, y, gate, idx: op_(res, y, gate, idx), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + elif op == "final_adaln_out": + + def make(s, n): + ti, _ = _layout(n, s) + return { + "x": _rn(s, 18, n, H, dtype=torch.bfloat16, scale=2.0), + "nw": (0.5 + torch.rand(H, device=DEV, generator=_gen(1, 19))).bfloat16(), + "temb": _rn(1, 20, T3, D), + "w": _rn(1, 21, 2 * H, D, dtype=torch.bfloat16, scale=D**-0.5), + "b": _rn(1, 22, 2 * H, dtype=torch.bfloat16, scale=0.02), + "ti": ti, + } + + def ref(x, nw, temb, w, b, ti): + return P.provider_final_adaln_out(x, nw, temb, w, b, ti) + + base = { + "make": make, + "rows": {"x": 0, "ti": 0}, + "grads": ["x", "nw", "temb", "w", "b"], + "ref": ref, + "bi_sizes": S_BI, + "perf_sizes": S_PERF, + } + + def diffusers(): + return { + **base, + "source": "diffusers MiniMaxH3AdaLayerNormOut (op-for-op replay)", + "fn": ref, + } + + def ours(): + from rl_engine.backends.cuda.model_specific.minimax_h3.final_adaln_out import ( + H3FinalAdaLNOutCudaOp, + ) + + op_ = H3FinalAdaLNOutCudaOp() + return { + **base, + "source": "rl-kernel H3FinalAdaLNOutCudaOp", + "fn": lambda x, nw, temb, w, b, ti: op_(x, nw, temb, w, b, ti), + } + + found = [("diffusers", diffusers), ("rl_kernel", ours)] + + result = [] + for name, factory in found: + try: + result.append({"name": name, **factory()}) + except Exception as exc: # noqa: BLE001 - optional library missing or unusable + result.append({"name": name, "unavailable": f"{type(exc).__name__}: {exc}"[:300]}) + return result + + +# --------------------------------------------------------------------------- # +# Modes and driver +# --------------------------------------------------------------------------- # + + +def enable_mode(mode: str, megatron_src: str | None): + """Switch on a global batch-invariant mode; returns a context manager that turns it off.""" + + base = mode[:-5] if mode.endswith("_ieee") else mode + if base == "plain": + return contextlib.nullcontext + if base == "vllm": + from vllm.model_executor.determinism.batch_invariant import enable_batch_invariant_mode + + enable_batch_invariant_mode() + return contextlib.nullcontext + if base == "sglang": + from sglang.srt.batch_invariant_ops import batch_invariant_ops as mod + + mod.enable_batch_invariant_mode() + kwargs: dict[str, str] = {} + else: + if megatron_src: + sys.path.insert(0, megatron_src) + from megatron.core.transformer.custom_layers import batch_invariant_kernels as mod + + kwargs = {"backend": base.split("_", 1)[1]} + mod.enable_batch_invariant_mode(**kwargs) + if kwargs["backend"] != "triton": + return contextlib.nullcontext + + @contextlib.contextmanager + def off(): + # The Triton matmuls have no FP64 configuration: FP64 references run with the mode off. + mod.disable_batch_invariant_mode() + try: + yield + finally: + mod.enable_batch_invariant_mode(**kwargs) + + return off + + +def run_mode(op: str, mode: str, megatron_src: str | None, quick: bool) -> dict[str, Any]: + try: + mode_off = enable_mode(mode, megatron_src) + except Exception as exc: # noqa: BLE001 + return {"unavailable": f"{type(exc).__name__}: {exc}"[:300]} + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + out: dict[str, Any] = {} + for cand in candidates(op, mode): + name = cand.pop("name") + if "unavailable" in cand: + out[name] = cand + print(f" [{mode}] {name}: unavailable ({cand['unavailable'][:80]})", flush=True) + continue + entry: dict[str, Any] = {"source": cand["source"], "has_backward": bool(cand["grads"])} + entry["batch_invariance"] = batch_invariance(cand, quick) + entry["accuracy_latency"] = accuracy_latency(cand, mode_off, quick) + out[name] = entry + print( + f" [{mode}] {name}: batch_invariant={entry['batch_invariance']['batch_invariant']}", + flush=True, + ) + torch.cuda.empty_cache() + return out + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--out", type=Path, required=True) + parser.add_argument("--megatron-src", default=None, help="Megatron-LM source tree (optional)") + parser.add_argument("--quick", action="store_true", help="small sizes, for a smoke test only") + parser.add_argument("--mode", default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + if importlib.util.find_spec(OPS[args.op]) is None: + raise SystemExit(f"{args.op}: {OPS[args.op]} is not on this branch") + if not torch.cuda.is_available(): + raise SystemExit("needs a CUDA device") + + if args.mode is not None: # child process: one global mode + args.out.write_text(json.dumps(run_mode(args.op, args.mode, args.megatron_src, args.quick))) + return + + from rl_engine.validation.models.h3_chain import environment, git_state + + modes = list(MODES) if args.op in GEMM_OPS else ["plain"] + results: dict[str, Any] = {} + for mode in modes: + print(f"[{args.op}] mode {mode}", flush=True) + # The part file sits next to the report, not in /tmp: cluster job epilogs can clear a + # user's /tmp while another of their jobs on the same node is still running. + part = args.out.with_name(f".{args.out.stem}.{mode}.part.json") + part.parent.mkdir(parents=True, exist_ok=True) + part.unlink(missing_ok=True) + cmd = [sys.executable, __file__, "--op", args.op, "--out", str(part), "--mode", mode] + if args.megatron_src: + cmd += ["--megatron-src", args.megatron_src] + if args.quick: + cmd.append("--quick") + proc = subprocess.run(cmd, env={**os.environ, **MODES[mode]}) + results[mode] = ( + json.loads(part.read_text()) + if proc.returncode == 0 and part.exists() + else {"unavailable": f"subprocess exited with {proc.returncode}"} + ) + part.unlink(missing_ok=True) + megatron_commit = None + if args.megatron_src: + megatron_commit = ( + subprocess.run( + ["git", "rev-parse", "--short", "HEAD"], + cwd=args.megatron_src, + capture_output=True, + text=True, + ).stdout.strip() + or None + ) + report = { + "kind": "h3_prior_art", + "rfc": "RL-Align/RL-Kernel#420", + "op": args.op, + **git_state(), + "environment": { + **environment(), + "libraries": {lib: _version(lib) for lib in LIBS}, + "megatron_source_commit": megatron_commit, + "mode_environment": {m: MODES[m] for m in modes}, + }, + "quick": args.quick, + "results": results, + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + print( + f"wrote {args.out} (commit {report['rl_kernel_commit'][:7]}, " + f"dirty={report['tracked_tree_dirty']})" + ) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/h3_ws2_evidence.py b/tools/validation/models/h3_ws2_evidence.py new file mode 100644 index 000000000..8ef5976fb --- /dev/null +++ b/tools/validation/models/h3_ws2_evidence.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Write an RFC #420 WS2 evidence report (byte equality vs WS1, timing, readback). + +Each TP size runs as real NCCL processes, one GPU per rank, with the +deterministic CUDA collective. Run it on a clean tree on one 8-GPU node; +``tools/validation/models/plot_h3_evidence.py`` turns the report into a figure. + + export RL_KERNEL_H3_WEIGHTS= + python tools/validation/models/h3_ws2_evidence.py --op tp_adaln_3mod --worlds 1,2,4,8 \\ + --out reports/experiments/h3-tp-adaln-3mod-b200/report.json + python tools/validation/models/h3_ws2_evidence.py --op sp_norm_adaln --worlds 1,2,4,8 \\ + --out reports/experiments/h3-sp-norm-adaln-b200/report.json +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[3])) + +import torch # noqa: E402 + +from rl_engine.validation.models.h3_chain import environment, git_state # noqa: E402 +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + load_h3_conditioning_weights, + load_h3_manifest, +) +from rl_engine.validation.models.h3_ws2 import ( # noqa: E402 + naive_sp_mismatch, + run_world, + sp_case, + sp_region_sweep, + tp_projection_sweep, +) + +WEIGHT = "transformer_blocks.0.adaln_proj.linear.weight" +BIAS = "transformer_blocks.0.adaln_proj.linear.bias" + + +def _tp_adaln(worlds: list[int], max_t: int) -> dict: + weights = load_h3_conditioning_weights("cpu") + weight, bias = weights[WEIGHT], weights[BIAS] + g = torch.Generator(device="cpu").manual_seed(0) + inputs = { + "temb": torch.randn(max_t, weight.shape[1], generator=g) * 2, + "weight": weight, + "bias": bias, + "grad": torch.randn(max_t, weight.shape[0], generator=g).bfloat16(), + } + runs = [] + for world in worlds: + ranks = run_world(world, tp_projection_sweep, inputs) + runs.append({"tp": world, "ranks": ranks}) + equal = all( + v for r in ranks for e in r["equality"] for k, v in e.items() if k != "num_timesteps" + ) + print(f"tp={world}: byte-equal to WS1 on every rank and T: {equal}") + return {"weights": [WEIGHT, BIAS], "runs": runs} + + +SP_CASES = [ + {"batch": 1, "seq": 4097, "layout": "block", "seed": 1}, + {"batch": 1, "seq": 4097, "layout": "interleaved", "seed": 2}, + {"batch": 2, "seq": 4097, "layout": "block", "seed": 3}, + {"batch": 1, "seq": 32768, "layout": "block", "seed": 4}, + {"batch": 1, "seq": 32768, "layout": "interleaved", "seed": 5}, + {"batch": 1, "seq": 131072, "layout": "block", "seed": 6}, +] + + +def _sp_norm(worlds: list[int], _max_t: int) -> dict: + runs = [] + for world in worlds: + ranks = run_world(world, sp_region_sweep, {}, cases=SP_CASES) + runs.append({"sp": world, "ranks": ranks}) + equal = all(all(c["equal"].values()) for r in ranks for c in r["cases"]) + print(f"sp={world}: byte-equal to WS1 on every rank and case: {equal}") + naive = { + layout: naive_sp_mismatch(sp_case(1, 32768, layout=layout, seed=4), max(worlds)) + for layout in ("block", "interleaved") + } + region = "norm2(residual + gate_msa[row] * y) * (1 + scale_mlp[row]) + shift_mlp[row]" + return { + "region": region, + "hidden": 5376, + "runs": runs, + "naive_sp": {"sp": max(worlds), **naive}, + } + + +OPS = {"tp_adaln_3mod": _tp_adaln, "sp_norm_adaln": _sp_norm} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--op", required=True, choices=sorted(OPS)) + parser.add_argument("--worlds", default="1,2,4,8") + parser.add_argument("--max-timesteps", type=int, default=4) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + worlds = [int(w) for w in args.worlds.split(",")] + if torch.cuda.device_count() < max(worlds): + raise SystemExit(f"needs {max(worlds)} GPUs, found {torch.cuda.device_count()}") + + manifest = load_h3_manifest() + report = { + "kind": "h3_ws2_report", + "op": args.op, + "rfc": manifest["rfc"], + "model_revision": manifest["model_identity"]["revision"], + **git_state(), + "environment": {**environment(), "gpus": torch.cuda.device_count()}, + **OPS[args.op](worlds, args.max_timesteps), + } + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(report, indent=2) + "\n") + commit, dirty = report["rl_kernel_commit"][:7], report["tracked_tree_dirty"] + print(f"wrote {args.out} (commit {commit}, dirty={dirty})") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_bitwise_reuse.py b/tools/validation/models/plot_h3_bitwise_reuse.py new file mode 100644 index 000000000..85983fc0a --- /dev/null +++ b/tools/validation/models/plot_h3_bitwise_reuse.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot ``h3_bitwise_reuse.py`` reports: share of forward outputs that are not bitwise equal. + +One heatmap per report. Rows are implementations, columns are the cases compared against +diffusers' path and against the rl-kernel op; 0 means bitwise identical. + + python tools/validation/models/plot_h3_bitwise_reuse.py \\ + reports/experiments/h3-bitwise-reuse-b200/*.json --out figure.png +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +_REFS = ("diffusers[plain]", "diffusers", "rl_kernel") + + +def _cell(entry: dict | None) -> float | None: + if not entry: + return None + out = entry["out"][0] + if out is None or "differ" not in out: + return None + return 0.0 if out["equal"] else out["differ"] / out["numel"] + + +def _text(frac: float | None) -> str: + if frac is None: + return "—" + if frac == 0: + return "0" + return f"{frac:.1%}" if frac >= 0.001 else f"{frac:.2%}" + + +def _panel(ax, report: dict) -> None: + results = {k: v for k, v in report["results"].items() if "unavailable" not in v} + cases = report["cases"] + diffusers = next((r for r in _REFS[:2] if r in results), None) + refs = [r for r in (diffusers, "rl_kernel") if r] + rows = ["rl_kernel"] + sorted(r for r in results if r not in _REFS) + cols = [(ref, case) for ref in refs for case in cases] + grid = np.full((len(rows), len(cols)), np.nan) + for i, row in enumerate(rows): + for j, (ref, case) in enumerate(cols): + frac = _cell(results.get(row, {}).get(f"vs {ref}", {}).get(str(case))) + if frac is not None: + grid[i, j] = frac + ax.text( + j, + i, + "—" if row == ref else _text(frac), + ha="center", + va="center", + fontsize=7.5, + color="white" if frac is not None and frac > 0.3 else "black", + fontweight="bold" if frac == 0 else None, + ) + ax.imshow(np.sqrt(np.nan_to_num(grid)), cmap="Reds", vmin=0, vmax=1, aspect="auto") + ax.set_yticks(range(len(rows)), rows, fontsize=8) + ax.set_xticks(range(len(cols)), [str(c) for _, c in cols], fontsize=7, rotation=45) + transform = ax.get_xaxis_transform() + for k, ref in enumerate(refs): + if k: + ax.axvline(k * len(cases) - 0.5, color="k", lw=1.5) + label = "vs rl-kernel" if ref == "rl_kernel" else "vs diffusers" + ax.text( + k * len(cases) + (len(cases) - 1) / 2, + 1.01, + label, + ha="center", + va="bottom", + fontsize=8.5, + transform=transform, + ) + ax.set_title(f"{report['op']} @ {report['rl_kernel_commit'][:8]}", fontsize=10, pad=18) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reports", nargs="+", type=Path) + parser.add_argument("--out", type=Path, required=True) + args = parser.parse_args() + reports = [json.loads(p.read_text()) for p in args.reports] + fig, axes = plt.subplots(1, len(reports), figsize=(6.4 * len(reports), 5.8), squeeze=False) + for ax, report in zip(axes[0], reports): + _panel(ax, report) + gpu = reports[0]["environment"].get("gpu", "") + fig.suptitle( + "Forward outputs that are not bitwise equal\n" + f"(share of elements; 0 = bitwise identical), {gpu}", + fontsize=10, + ) + fig.tight_layout() + args.out.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(args.out, dpi=125) + print(f"wrote {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_block.py b/tools/validation/models/plot_h3_block.py new file mode 100644 index 000000000..228f3a289 --- /dev/null +++ b/tools/validation/models/plot_h3_block.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render the RFC #420 one-block replay (``tools/validation/models/h3_block_replay.py``) as a PNG. + +Needs only the report JSON and matplotlib: + + python tools/validation/models/plot_h3_block.py \ + reports/experiments/h3-one-block-b200/report.json + +writes ``figure.png`` next to the report. Panels: the per-node invariance and +provider checks over every layout; the per-node error against the FP32 golden +for the candidate and diffusers; the same for every backward leaf. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path +from typing import Any + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.colors import ListedColormap # noqa: E402 + +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +PASS = "#1baf7a" +FAIL = "#d64545" +NA = "#e4e3dd" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" +STATUS_MARK = {"row": "", "interim": "*", "reference": "†"} + + +def _style(ax, title: str) -> None: + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=10.5, color=INK, pad=8) + ax.tick_params(colors=MUTED, labelsize=7.5, length=0) + for side in ("top", "right", "left", "bottom"): + ax.spines[side].set_visible(False) + + +def _checks_panel(ax, report: dict[str, Any]) -> None: + """Rows: nodes. Columns: forward/backward checks, ANDed over every layout.""" + + nodes = [n["node"] for n in report["nodes"]] + status = {n["node"]: n["status"] for n in report["nodes"]} + cols = [ + ("repeat", "forward", "repeat_bitwise_equal"), + ("batch row", "forward", "batch_row_bitwise_equal"), + ("= diffusers\n(isolated)", "forward", None), + ("grad\nrepeat", "backward", "grad_repeat_bitwise_equal"), + ("grad\nbatch row", "backward", "grad_batch_row_bitwise_equal"), + ] + grid = [] + for node in nodes: + row = [] + for _label, kind, key in cols: + cases = report["forward_cases" if kind == "forward" else "backward_cases"] + entries = [e for case in cases for e in case["nodes"] if e["node"] == node] + if key is None: + promised = entries[0]["provider_bitwise_isolated_promised"] + ok = all(e["isolated_vs_provider"]["bitwise_equal"] for e in entries) + row.append(2 if not promised and not ok else (1 if ok else 0)) + else: + values = [e[key] for e in entries] + row.append(2 if not values or None in values else (1 if all(values) else 0)) + grid.append(row) + ax.imshow(grid, cmap=ListedColormap([FAIL, PASS, NA]), vmin=0, vmax=2, aspect="auto") + ax.set_xticks(range(len(cols)), [c[0] for c in cols], fontsize=7.5, color=INK) + ax.set_yticks( + range(len(nodes)), [f"{n}{STATUS_MARK[status[n]]}" for n in nodes], fontsize=7.5, color=INK + ) + ax.xaxis.tick_top() + _style(ax, "Bitwise checks, all layouts") + ax.set_xlabel( + "green = holds grey = not promised (reduction differs from diffusers)\n" + "or not applicable (AdaLN tables have no batch axis)\n" + "* interim operator † provider replay (no kernel yet)", + fontsize=7, + color=MUTED, + ) + + +def _err_panel(ax, labels, cand, prov, title: str) -> None: + xs = range(len(labels)) + ax.plot( + xs, [math.log10(v) for v in cand], "o-", color=RL_KERNEL, lw=1.8, ms=5, label="RL-Kernel" + ) + ax.plot( + xs, [math.log10(v) for v in prov], "s--", color=PROVIDER, lw=1.4, ms=4, label="diffusers" + ) + ax.set_xticks(list(xs), labels, rotation=55, ha="right", fontsize=7, color=INK) + lo = math.floor(min(math.log10(v) for v in cand + prov)) + hi = math.ceil(max(math.log10(v) for v in cand + prov)) + ax.set_yticks(range(lo, hi + 1), [f"1e{v}" for v in range(lo, hi + 1)]) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_ylabel("max|err| / max|golden|", color=MUTED, fontsize=8) + ax.legend(frameon=False, fontsize=7.5, labelcolor=INK, loc="upper left") + _style(ax, title) + + +def plot(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(16, 5.6), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.0, 1.6, 1.3]} + ) + case = max(report["forward_cases"], key=lambda c: c["seq_len"]) + back = ( + max(report["backward_cases"], key=lambda c: c["seq_len"]) + if report["backward_cases"] + else None + ) + _checks_panel(axes[0], report) + _err_panel( + axes[1], + [e["node"] for e in case["nodes"]], + [e["candidate_rel_err_vs_golden"] for e in case["nodes"]], + [e["provider_rel_err_vs_golden"] for e in case["nodes"]], + f"Forward error vs FP32 golden, S = {case['seq_len']}", + ) + if back is not None: + leaves = list(back["leaves"]) + _err_panel( + axes[2], + [f"d {name}" for name in leaves], + [back["leaves"][n]["candidate_rel_err_vs_golden"] for n in leaves], + [back["leaves"][n]["provider_rel_err_vs_golden"] for n in leaves], + f"Gradient error vs FP32 golden, S = {back['seq_len']}", + ) + env = report["environment"] + sizes = ", ".join(str(c["seq_len"]) for c in report["forward_cases"]) + fig.suptitle( + "MiniMax-H3 block 0, forward and backward (RFC #420 ws1_one_h3_block)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.905, + f"{env['gpu']} | torch {env['torch']} | layouts S = {sizes} | " + f"first drift vs diffusers: {case['first_drift']} | " + f"rl-kernel {report['rl_kernel_commit'][:7]}", + fontsize=8, + color=MUTED, + ) + fig.tight_layout(rect=(0, 0, 1, 0.89)) + return fig + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + out = args.report.with_name("figure.png") + plot(report).savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_evidence.py b/tools/validation/models/plot_h3_evidence.py new file mode 100644 index 000000000..467cb5b7d --- /dev/null +++ b/tools/validation/models/plot_h3_evidence.py @@ -0,0 +1,768 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Render an RFC #420 operator or WS2 report (``tools/validation/models/h3_evidence.py``, +``tools/validation/models/h3_ws2_evidence.py``) as a PNG. + +Needs only the report JSON and matplotlib (not torch or a GPU), so it can run +anywhere: + + python tools/validation/models/plot_h3_evidence.py \\ + reports/experiments/h3-timestep-sinusoid-b200/report.json + +writes ``figure.png`` next to the report. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any, Callable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.patches # noqa: E402 +import matplotlib.pyplot as plt # noqa: E402 + +# Fixed series identity across every H3 figure (categorical slots 1-3). +RL_KERNEL = "#2a78d6" +PROVIDER = "#eb6834" +THIRD = "#1baf7a" +INK = "#2b2b29" +MUTED = "#6f6e69" +GRID = "#e4e3dd" +SURFACE = "#fcfcfb" + + +def _style(ax, title: str, ylabel: str) -> None: + """Apply the shared report axis styling, title, and vertical-axis label.""" + + ax.set_facecolor(SURFACE) + ax.set_title(title, loc="left", fontsize=11, color=INK, pad=10) + ax.set_ylabel(ylabel, color=MUTED, fontsize=9) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + ax.grid(axis="y", color=GRID, linewidth=0.8) + ax.set_axisbelow(True) + for side in ("top", "right", "left"): + ax.spines[side].set_visible(False) + ax.spines["bottom"].set_color(GRID) + + +def _bars(ax, labels: list[str], series: list[tuple], fmt: str, log: bool = False) -> None: + """Draw grouped series with formatted value labels and a legend; a 4th tuple item is a hatch.""" + + n = len(series) + width = 0.8 / n + for i, (name, color, values, *hatch) in enumerate(series): + xs = [x + (i - (n - 1) / 2) * width for x in range(len(labels))] + bars = ax.bar( + xs, + values, + width * 0.92, + color=color if not hatch else SURFACE, + edgecolor=color, + hatch=hatch[0] if hatch else None, + linewidth=1.2, + label=name, + zorder=2, + ) + for bar, value in zip(bars, values, strict=True): + ax.annotate( + fmt.format(value), + (bar.get_x() + bar.get_width() / 2, bar.get_height()), + xytext=(0, 2), + textcoords="offset points", + ha="center", + fontsize=7, + color=INK, + ) + ax.set_xticks(range(len(labels)), labels) + top = max(v for _, _, values, *_ in series for v in values) + if log: + ax.set_yscale("log") + ax.set_ylim(min(v for _, _, values, *_ in series for v in values) / 2, top * 8) + else: + ax.set_ylim(0, top * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=min(n, 4)) + + +# Exact zeros cannot sit on a log axis; draw them on this floor and label them. +LOG_FLOOR = 1e-9 + + +def _log_values(values: list[float]) -> list[float]: + """Replace nonpositive values with the plotting floor for logarithmic axes.""" + + return [v if v > 0 else LOG_FLOOR for v in values] + + +def _figure(report: dict[str, Any], suptitle: str): + """Create two report axes with a title and GPU, software, and revision provenance.""" + + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8), facecolor=SURFACE) + env = report["environment"] + fig.suptitle(suptitle, x=0.01, ha="left", fontsize=12, color=INK, fontweight="bold") + fig.text( + 0.01, + 0.01, + f"{env['gpu']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · MiniMax-H3@{report['model_revision'][:7]}", + fontsize=7, + color=MUTED, + ) + return fig, axes + + +def plot_sinusoid(report: dict[str, Any]): + """Return a sinusoid report figure showing latency and errors against FP64.""" + + fig, (left, right) = _figure(report, "timestep_sinusoid_h3") + perf = report["perf"] + _style(left, "Latency per call (lower is better)", "µs") + _bars( + left, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + ("CUDA + range check", THIRD, [row["candidate_checked_us"] for row in perf]), + ("diffusers", PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + cases = report["accuracy"]["cases"] + xs = [c["num_timesteps"] for c in cases] + _style(right, "Max |error| vs FP64 golden (log)", "max abs error") + cuda = [c["max_abs_vs_fp64"] for c in cases] + provider = [c["provider_max_abs_vs_fp64"] for c in cases] + right.plot(xs, _log_values(cuda), "o-", color=RL_KERNEL, lw=2, ms=8, label="RL-Kernel CUDA") + right.plot(xs, _log_values(provider), "s--", color=PROVIDER, lw=2, ms=5, label="diffusers") + atol = report["accuracy"]["contract_atol"] + right.axhline(atol, color=MUTED, lw=1, ls=":") + right.annotate( + "contract atol 1e-5", + (xs[-1], atol), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + for x, value in zip(xs, cuda, strict=True): + if value == 0: + right.annotate( + "exact", + (x, LOG_FLOOR), + xytext=(0, 8), + textcoords="offset points", + fontsize=7, + color=INK, + ha="center", + ) + right.set_xscale("log") + right.set_yscale("log") + right.set_ylim(LOG_FLOOR / 3, atol * 30) + right.set_xlabel("number of timesteps T", color=MUTED, fontsize=9) + equal = sum(c["bitwise_equal_to_provider"] for c in cases) + right.legend( + frameon=False, + fontsize=8, + labelcolor=INK, + loc="upper left", + ncol=2, + title=f"bitwise equal to diffusers at {equal}/{len(cases)} sizes", + title_fontsize=8, + alignment="left", + ) + return fig + + +def _latency_panel(ax, perf: list[dict[str, Any]], provider_name: str) -> None: + """Plot candidate and provider latency summaries in microseconds for each case.""" + + _style(ax, "Latency per call (lower is better)", "µs") + _bars( + ax, + [row["case"] for row in perf], + [ + ("RL-Kernel CUDA", RL_KERNEL, [row["candidate_us"] for row in perf]), + (provider_name, PROVIDER, [row["provider_us"] for row in perf]), + ], + "{:.0f}", + ) + + +def _error_boxes(ax, title: str, series: list[tuple[str, str, list[float]]], contract: float): + """Per-draw max |error| as boxes with every draw overlaid, log scale.""" + + _style(ax, title, "max abs error per draw") + for i, (_name, color, values) in enumerate(series): + ax.boxplot( + [values], + positions=[i], + widths=0.45, + showfliers=False, + patch_artist=True, + boxprops={"facecolor": color, "alpha": 0.25, "edgecolor": color}, + medianprops={"color": color, "linewidth": 2}, + whiskerprops={"color": color}, + capprops={"color": color}, + ) + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax.scatter(jitter, values, s=8, color=color, alpha=0.6, linewidths=0, zorder=3) + median = sorted(values)[len(values) // 2] + ax.annotate( + f"median {median:.1e}", + (i + 0.28, median), + fontsize=8, + color=INK, + va="center", + ) + ax.axhline(contract, color=MUTED, lw=1, ls=":") + ax.annotate( + f"contract atol {contract:g}", + (len(series) - 0.5, contract), + xytext=(0, 3), + textcoords="offset points", + fontsize=7, + color=MUTED, + ha="right", + ) + ax.set_yscale("log") + ax.set_xticks(range(len(series)), [name for name, _, _ in series]) + ax.set_xlim(-0.6, len(series) - 0.1) + + +def plot_mlp(report: dict[str, Any]): + """Return an MLP report figure showing latency and per-draw FP64 error distributions.""" + + fig, (left, right) = _figure(report, "timestep_mlp_fp32 (pinned time_embedder weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _error_boxes( + right, + f"Error vs FP64 golden, {acc['draws']} draws of T={acc['num_timesteps']}", + [ + ("RL-Kernel CUDA", RL_KERNEL, acc["cuda_max_abs_vs_fp64"]), + ("diffusers (cuBLAS)", PROVIDER, acc["provider_max_abs_vs_fp64"]), + ], + acc["contract_atol"], + ) + return fig + + +def plot_projection(report: dict[str, Any]): + """Return an AdaLN figure with latency, BF16 rounding accuracy, and the early-cast probe.""" + + fig, (left, right) = _figure(report, "adaln_projection_3mod (pinned block-0 AdaLN weights)") + _latency_panel(left, report["perf"], "diffusers (cuBLAS)") + acc = report["accuracy"] + _style( + right, + f"Outputs equal to the correctly rounded FP64 golden, {acc['draws']} draws", + "% of BF16 outputs", + ) + series = [ + ("RL-Kernel CUDA", RL_KERNEL, [100 * v for v in acc["cuda_correctly_rounded"]]), + ("diffusers (cuBLAS)", PROVIDER, [100 * v for v in acc["provider_correctly_rounded"]]), + ] + for i, (_name, color, values) in enumerate(series): + jitter = [i + ((k * 0.6180339) % 1 - 0.5) * 0.3 for k in range(len(values))] + ax_vals = sorted(values) + median = ax_vals[len(ax_vals) // 2] + right.scatter(jitter, values, s=36, color=color, alpha=0.8, linewidths=0, zorder=3) + right.hlines(median, i - 0.25, i + 0.25, color=color, lw=2.5, zorder=4) + right.annotate( + f"median {median:.3f}%", (i + 0.28, median), fontsize=8, color=INK, va="center" + ) + right.set_xticks(range(len(series)), [name for name, _, _ in series]) + right.set_xlim(-0.6, len(series) - 0.1) + early = sorted(acc["early_cast_golden_match"])[len(acc["early_cast_golden_match"]) // 2] + right.text( + 0.98, + 0.62, + f"an early BF16 cast (probe H7)\nwould match only {100 * early:.0f}%", + transform=right.transAxes, + fontsize=8, + color=MUTED, + ha="right", + ) + return fig + + +def plot_gather(report: dict[str, Any]): + fig, (left, right) = _figure(report, "adaln_row_gather (T = 3, H = 5376, BF16)") + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("index_select fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "index_select bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + chain = report["accuracy"]["chain_backward"] + _style( + right, + "Whole-chain grads of the FP32 time embedder vs FP64", + "max abs error / golden max", + ) + if not chain: + right.text( + 0.5, + 0.5, + report["accuracy"].get("chain_backward_skipped", "No chain measurements available"), + transform=right.transAxes, + ha="center", + va="center", + color=INK, + fontsize=9, + wrap=True, + ) + right.set_axis_off() + return fig + names = { + "candidate": ("RL-Kernel, separate ops", THIRD, "o"), + "candidate_fused": ("RL-Kernel, fused modulation", RL_KERNEL, "D"), + "provider": ("diffusers", PROVIDER, "s"), + } + labels = [f"T={c['num_timesteps']}\nS={c['seq_len']}" for c in chain] + for j, (mode, (name, color, marker)) in enumerate(names.items()): + leaves = [e for c in chain for e in c["leaves"].values()] + det = sum(e[mode]["repeat_bitwise_equal"] for e in leaves) + adaln = min( + e[mode]["correctly_rounded_fraction"] + for c in chain + for n, e in c["leaves"].items() + if n.startswith("transformer_blocks") + ) + label = ( + f"{name}\n repeat-bitwise {det}/{len(leaves)} · " + f"AdaLN BF16 grads ≥{100 * adaln:.1f}% correctly rounded" + ) + for i, case in enumerate(chain): + values = [ + e[mode]["max_abs_vs_golden_over_absmax"] + for n, e in case["leaves"].items() + if n.startswith("time_embedder") + ] + right.scatter( + [i + (j - 1) * 0.22] * len(values), + values, + s=30, + color=color, + marker=marker, + alpha=0.85, + linewidths=0, + zorder=3, + label=label if i == 0 else None, + ) + right.set_yscale("log") + right.set_ylim(1e-8, 1e4) + right.set_xticks(range(len(labels)), labels, fontsize=7) + right.legend(frameon=False, fontsize=7, labelcolor=INK, loc="upper left") + return fig + + +def plot_rmsnorm(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "h3_rmsnorm (block norm1 + MSA modulation, H = 5376, BF16)" + ) + perf = report["perf"] + _style(left, "Forward and backward time (log, lower is better)", "ms") + _bars( + left, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + ("diffusers fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ("CUDA bwd", RL_KERNEL, [row["candidate_backward_us"] / 1e3 for row in perf], "////"), + ( + "diffusers bwd", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + acc = report["accuracy"] + bwd = acc["backward"] + keys = ["dx", "dweight", "dshift", "dscale"] + fwd_ok = acc["modulated_bitwise_vs_diffusers"] and all( + acc["plain_bitwise_vs_nn_rmsnorm"].values() + ) + _style(right, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def repeat(mode: str) -> str: + return "yes" if bwd[mode]["repeat_bitwise_equal"] else "no" + + _bars( + right, + keys, + [ + ( + f"RL-Kernel CUDA (repeat-bitwise: {repeat('cuda')})", + RL_KERNEL, + [bwd["cuda"]["rel_error"][k] for k in keys], + ), + ( + f"diffusers (repeat-bitwise: {repeat('provider')})", + PROVIDER, + [bwd["provider"]["rel_error"][k] for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + right.text( + 0.98, + 0.80, + f"forward bitwise equal to diffusers: {'yes' if fwd_ok else 'NO'}", + transform=right.transAxes, + ha="right", + fontsize=8, + color=INK, + ) + return fig + + +def _latency_fwd_bwd(ax, perf: list[dict[str, Any]], provider: str) -> None: + combined = all(row.get("backward_timing_scope") == "forward_and_backward" for row in perf) + backward_label = "fwd+bwd" if combined else "bwd" + _style(ax, "Forward and backward time (log, lower is better)", "ms") + _bars( + ax, + [row["case"] for row in perf], + [ + ("CUDA fwd", RL_KERNEL, [row["candidate_us"] / 1e3 for row in perf]), + (f"{provider} fwd", PROVIDER, [row["provider_us"] / 1e3 for row in perf]), + ( + f"CUDA {backward_label}", + RL_KERNEL, + [row["candidate_backward_us"] / 1e3 for row in perf], + "////", + ), + ( + f"{provider} {backward_label}", + PROVIDER, + [row["provider_backward_us"] / 1e3 for row in perf], + "////", + ), + ], + "{:.2f}", + log=True, + ) + + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + + +def _backward_errors(ax, backward: dict[str, Any], keys: list[str], note: str) -> None: + _style(ax, "Backward vs FP64 golden (log, lower is better)", "max abs error / golden max") + + def label(name: str, mode: str) -> str: + return ( + f"{name} (repeat-bitwise: {'yes' if backward[mode]['repeat_bitwise_equal'] else 'no'})" + ) + + floor = 1e-9 + _bars( + ax, + keys, + [ + ( + label("RL-Kernel CUDA", "cuda"), + RL_KERNEL, + [max(backward["cuda"]["rel_error"][k], floor) for k in keys], + ), + ( + label("diffusers", "provider"), + PROVIDER, + [max(backward["provider"]["rel_error"][k], floor) for k in keys], + ), + ], + "{:.1e}", + log=True, + ) + for text in ax.texts: + if text.get_text() == f"{floor:.1e}": + text.set_text("exact") + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=1) + ax.set_ylim(floor / 3, 1e4) + ax.text(0.98, 0.80, note, transform=ax.transAxes, ha="right", fontsize=8, color=INK) + + +def plot_gate_residual(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "adaln_gate_residual (residual + gate_msa[row] * y, H = 5376)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + ok = all(acc["forward_bitwise_vs_diffusers"].values()) + _backward_errors( + right, + acc["backward"], + ["d_residual", "d_sublayer", "d_gate"], + f"forward bitwise equal to diffusers (bf16/fp16/fp32): {'yes' if ok else 'NO'}", + ) + return fig + + +def plot_final(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "final_adaln_out (norm_out: projection + final norm, T = 3)" + ) + _latency_fwd_bwd(left, report["perf"], "diffusers") + acc = report["accuracy"] + _backward_errors( + right, + acc["backward"], + ["dx", "d_norm_w", "d_temb", "dW", "db"], + f"forward equal to diffusers on {100 * acc['equal_to_diffusers_fraction']:.2f}% " + "of elements (projection 1-ULP ties)", + ) + return fig + + +TP_TENSORS = ( + ("table", "table (6 x 3T rows)"), + ("d_temb", "d_temb"), + ("d_weight_shard", "dW shard"), + ("d_bias_shard", "db shard"), +) + + +def plot_tp_adaln(report: dict[str, Any]): + fig, (left, right) = _figure( + report, "tp_adaln_3mod (column-parallel AdaLN projection, 2688 -> 96768, NCCL)" + ) + runs = report["runs"] + _equality_grid( + left, + "Byte-equal to WS1", + [(f"TP{run['tp']}", [e for r in run["ranks"] for e in r["equality"]]) for run in runs], + TP_TENSORS, + "ranks x T", + ) + # Right: slowest rank's time at the largest T; WS1 on one GPU as reference lines. + num_t = max(p["num_timesteps"] for p in runs[0]["ranks"][0]["perf"]) + + def worst(run, key): + perf = [p for r in run["ranks"] for p in r["perf"] if p["num_timesteps"] == num_t] + return max(p[key] for p in perf) / 1e3 + + tp_runs = [run for run in runs if run["tp"] > 1] + _style(right, f"Time per call, T = {num_t} (slowest rank; lower is better)", "ms") + _bars( + right, + [f"TP{run['tp']}" for run in tp_runs], + [ + ("TP fwd", RL_KERNEL, [worst(run, "tp_forward_us") for run in tp_runs]), + ( + "TP fwd+bwd", + RL_KERNEL, + [worst(run, "tp_forward_backward_us") for run in tp_runs], + "////", + ), + ], + "{:.2f}", + log=True, + ) + for key, label, style in ( + ("ws1_forward_us", "WS1 fwd", ":"), + ("ws1_forward_backward_us", "WS1 fwd+bwd", "--"), + ): + value = worst(runs[0], key) + right.axhline( + value, + color=PROVIDER, + linestyle=style, + linewidth=1.4, + zorder=1, + label=f"{label}, 1 GPU: {value:.2f}", + ) + right.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper left", ncol=2) + return fig + + +def _equality_grid(ax, title, rows, tensors, unit) -> None: + """One cell per (configuration, tensor): how many checks were byte-equal to WS1.""" + + ax.set_facecolor(SURFACE) + ax.set_title(f"{title}; checks = {unit}", loc="left", fontsize=11, color=INK, pad=10) + for y, (_, checks) in enumerate(rows): + for x, (key, _) in enumerate(tensors): + ok = sum(bool(e[key]) for e in checks) + good = ok == len(checks) + ax.add_patch( + matplotlib.patches.FancyBboxPatch( + (x + 0.06, y + 0.08), + 0.88, + 0.84, + boxstyle="round,pad=0,rounding_size=0.06", + facecolor=RL_KERNEL if good else PROVIDER, + edgecolor=SURFACE, + linewidth=2, + ) + ) + label = f"{'equal' if good else 'DIFF'}\n{ok}/{len(checks)}" + ax.text(x + 0.5, y + 0.5, label, ha="center", va="center", fontsize=8, color="white") + ax.set_xlim(0, len(tensors)) + ax.set_ylim(len(rows), 0) + ax.set_xticks([i + 0.5 for i in range(len(tensors))], [label for _, label in tensors]) + ax.set_yticks([i + 0.5 for i in range(len(rows))], [name for name, _ in rows]) + ax.tick_params(colors=MUTED, labelsize=8, length=0) + for side in ax.spines.values(): + side.set_visible(False) + + +SP_TENSORS = ( + ("out", "output rows"), + ("d_residual", "d_residual"), + ("d_y", "d_sublayer"), + ("d_weight", "d_norm_w"), + ("d_table", "d_table"), +) + + +def plot_sp_norm(report: dict[str, Any]): + fig, axes = plt.subplots( + 1, 3, figsize=(15, 3.9), facecolor=SURFACE, gridspec_kw={"width_ratios": [1.35, 1, 0.9]} + ) + env = report["environment"] + fig.suptitle( + "sp_norm_adaln (sequence-parallel norm2(residual + gate * y) with AdaLN modulation, " + "H = 5376, NCCL)", + x=0.01, + ha="left", + fontsize=12, + color=INK, + fontweight="bold", + ) + fig.text( + 0.01, + 0.01, + f"{env['gpu']} x {env['gpus']} · torch {env['torch']} · CUDA {env['cuda']} · " + f"commit {report['rl_kernel_commit'][:7]} · bf16", + fontsize=7, + color=MUTED, + ) + runs = report["runs"] + _equality_grid( + axes[0], + "Byte-equal to WS1", + [ + (f"SP{run['sp']}", [c["equal"] for r in run["ranks"] for c in r["cases"]]) + for run in runs + ], + SP_TENSORS, + "ranks x cases", + ) + + # Middle: the largest block-layout case, slowest rank. + big = max( + (c for c in runs[0]["ranks"][0]["cases"] if c["layout"] == "block"), + key=lambda c: c["seq"] * c["batch"], + ) + + def worst(run, key): + return ( + max( + c[key] + for r in run["ranks"] + for c in r["cases"] + if (c["seq"], c["batch"], c["layout"]) == (big["seq"], big["batch"], big["layout"]) + ) + / 1e3 + ) + + sp_runs = [run for run in runs if run["sp"] > 1] + ax = axes[1] + _style(ax, f"Region fwd+bwd, S = {big['seq']}, block layout", "ms (slowest rank)") + _bars( + ax, + [f"SP{run['sp']}" for run in sp_runs], + [("SP", RL_KERNEL, [worst(run, "sp_forward_backward_us") for run in sp_runs])], + "{:.2f}", + ) + ws1 = worst(runs[0], "ws1_forward_backward_us") + ax.axhline( + ws1, color=PROVIDER, linestyle="--", linewidth=1.4, zorder=1, label=f"WS1, 1 GPU: {ws1:.2f}" + ) + ax.set_ylim(0, ws1 * 1.35) + ax.legend(frameon=False, fontsize=8, labelcolor=INK, loc="upper right") + + # Right: what a naive SP reduction would do instead. + naive = report["naive_sp"] + ax = axes[2] + _style(ax, f"Naive SP{naive['sp']} vs WS1 (S = 32768)", "% of elements that differ") + _bars( + ax, + ["d_norm_w", "d_table"], + [ + ("naive, block", PROVIDER, [100 * naive["block"][k] for k in ("d_weight", "d_table")]), + ( + "naive, interleaved", + PROVIDER, + [100 * naive["interleaved"][k] for k in ("d_weight", "d_table")], + "////", + ), + ("RL-Kernel SP", RL_KERNEL, [0.0, 0.0]), + ], + "{:.1f}", + ) + ax.text( + 0.98, + 0.62, + "naive: each rank's WS1 backward,\nthen a rank-order sum", + transform=ax.transAxes, + ha="right", + fontsize=7, + color=MUTED, + ) + return fig + + +PLOTS: dict[str, Callable[[dict[str, Any]], Any]] = { + "timestep_sinusoid_h3": plot_sinusoid, + "timestep_mlp_fp32": plot_mlp, + "adaln_projection_3mod": plot_projection, + "adaln_row_gather": plot_gather, + "h3_rmsnorm": plot_rmsnorm, + "adaln_gate_residual": plot_gate_residual, + "final_adaln_out": plot_final, + "tp_adaln_3mod": plot_tp_adaln, + "sp_norm_adaln": plot_sp_norm, +} + + +def main() -> None: + """Read an operator report and save its figure as a PNG beside it or at ``--out``.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("report", type=Path) + parser.add_argument("--out", type=Path, default=None) + args = parser.parse_args() + report = json.loads(args.report.read_text()) + fig = PLOTS[report["op"]](report) + fig.tight_layout(rect=(0, 0.04, 1, 0.95)) + out = args.out or args.report.with_name("figure.png") + fig.savefig(out, dpi=150, facecolor=SURFACE) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/models/plot_h3_prior_art.py b/tools/validation/models/plot_h3_prior_art.py new file mode 100644 index 000000000..893612e87 --- /dev/null +++ b/tools/validation/models/plot_h3_prior_art.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Plot H3 prior-art latency, accuracy, and batch invariance. + + python tools/validation/models/plot_h3_prior_art.py \\ + reports/experiments/h3-prior-art-b200/.json + +Writes ``.png`` next to the report. +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib import transforms # noqa: E402 + + +def _entries(report): + for mode, cands in report["results"].items(): + if "unavailable" in cands: + continue + for name, entry in cands.items(): + if "unavailable" not in entry: + yield (name if name.endswith("]") or mode == "plain" else f"{name}[{mode}]"), entry + + +def _differ(bi) -> int: + rows = sum(v["fwd_differ"] + v["rowgrad_differ"] for v in bi["all_rows"].values()) + sub = bi["full_vs_sub_batches"] + return rows + sub["fwd_differ"] + sub["rowgrad_differ"] + bi["batch_size_sweep"]["differ"] + + +def main() -> None: + path = Path(sys.argv[1]) + report = json.loads(path.read_text()) + entries = list(_entries(report)) + names = [n for n, _ in entries] + fig, axes = plt.subplots(1, 3, figsize=(20, 6.5), layout="constrained") + fig.suptitle( + f"{report['op']} vs existing implementations — {report['environment']['gpu']}, " + f"commit {report['rl_kernel_commit'][:7]}", + fontsize=13, + ) + + ax = axes[0] + for name, entry in entries: + perf = entry["accuracy_latency"] + key = "fwd_bwd_us" if entry["has_backward"] else "fwd_us" + sizes = sorted(perf, key=int) + ax.plot( + [int(s) for s in sizes], + [perf[s][key] for s in sizes], + marker="o", + ls="-" if entry["has_backward"] else "--", + label=name + ("" if entry["has_backward"] else " (forward only)"), + ) + ax.set_xscale("log", base=2) + ax.set_yscale("log") + ax.set_xlabel("rows (timesteps or tokens)") + ax.set_ylabel("µs, median") + ax.set_title("forward + backward latency (dashed: forward only)") + ax.grid(True, which="both", alpha=0.3) + ax.legend(fontsize=7) + + ax = axes[1] + first = entries[0][1]["accuracy_latency"] + size = "3" if "3" in first else sorted(first, key=int)[-1] + ys = range(len(entries)) + fwd = [e["accuracy_latency"][size]["fwd_rel_err"] for _, e in entries] + grad = [ + max(e["accuracy_latency"][size].get("grad_rel_err", {}).values(), default=0) + for _, e in entries + ] + ax.barh([y - 0.2 for y in ys], fwd, 0.4, label="forward") + ax.barh([y + 0.2 for y in ys], [g or float("nan") for g in grad], 0.4, label="worst gradient") + values = [v for v in fwd + grad if v] + if values and max(values) / min(values) > 10: + ax.set_xscale("log") + inside = transforms.blended_transform_factory(ax.transAxes, ax.transData) + for y, a, b in zip(ys, fwd, grad): + ax.text( + 0.98, + y, + f"fwd {a:.1e} / grad {b:.1e}" if b else f"fwd {a:.1e}", + transform=inside, + ha="right", + va="center", + fontsize=7, + bbox={"facecolor": "white", "edgecolor": "none", "alpha": 0.8, "pad": 1}, + ) + ax.set_yticks(list(ys), names, fontsize=7) + ax.invert_yaxis() + ax.set_title(f"error vs FP64, max|err| / max|ref| (size {size})") + ax.legend(fontsize=7, loc="lower right") + ax.grid(True, axis="x", alpha=0.3) + + ax = axes[2] + for y, (_, e) in zip(ys, entries): + bi = e["batch_invariance"] + rep = bi["params_repeatable"] + lines = ["batch-invariant" if bi["batch_invariant"] else "NOT batch-invariant"] + lines.append(f"{_differ(bi)} row comparisons differ") + if rep is not None: + lines.append("parameter/table grads " + ("repeatable" if rep else "NOT repeatable")) + ax.text( + 0.02, + y, + " | ".join(lines), + va="center", + fontsize=8, + color="#3aa676" if bi["batch_invariant"] else "#d14b4b", + ) + ax.set_ylim(len(entries) - 0.5, -0.5) + ax.set_xlim(0, 1) + ax.axis("off") + ax.set_title("batch invariance (bitwise)") + ax.text( + 0.02, + len(entries) - 0.3, + "every row alone vs full batches + full batch vs covering sub-batches + size sweep", + fontsize=7, + color="#555555", + ) + + out = path.with_suffix(".png") + fig.savefig(out, dpi=120) + print(f"wrote {out}") + + +if __name__ == "__main__": + main() diff --git a/tools/weights/prepare_h3_weights.py b/tools/weights/prepare_h3_weights.py new file mode 100644 index 000000000..f6e32bcb5 --- /dev/null +++ b/tools/weights/prepare_h3_weights.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Fetch and pin the MiniMax-H3 conditioning tensors used by the RFC #420 tests. + +Downloads only the config, the shard index and the shards that hold the +manifest tensors (shard 1 of 14, ~4.8 GB) at the pinned revision, verifies +every sha256 against ``rl_engine/validation/models/h3_manifest.json`` and writes the +tensors to ``/h3_conditioning.safetensors``. Point the tests at it with +``export RL_KERNEL_H3_WEIGHTS=``. ``--block`` also writes block 0's +attention/FFN weights (manifest ``block_tensors``, same shard) to +``/h3_block0.safetensors`` for the ``ws1_one_h3_block`` replay. + + python tools/weights/prepare_h3_weights.py --out ~/.cache/rl-kernel/h3 [--block] +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT)) + +from rl_engine.validation.models.h3_weights import ( # noqa: E402 + BLOCK_FILE, + EXTRACTED_FILE, + load_h3_manifest, + sha256_file, + sha256_tensor, +) + + +def _check(path: Path, expected: str) -> None: + """Verify a downloaded file's SHA-256, exiting on mismatch or printing success.""" + + actual = sha256_file(path) + if actual != expected: + raise SystemExit(f"sha256 mismatch for {path.name}: {actual} != {expected}") + print(f"ok {path.name} {actual}") + + +def main() -> None: + """Download pinned model artifacts, verify tensor identities, and write extracted weights.""" + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--out", type=Path, required=True, help="output directory") + parser.add_argument( + "--download-dir", + type=Path, + default=None, + help="where the raw shard goes (default: /hf)", + ) + parser.add_argument( + "--block", action="store_true", help="also extract block 0 (manifest block_tensors)" + ) + args = parser.parse_args() + + from huggingface_hub import hf_hub_download + from safetensors import safe_open + from safetensors.torch import save_file + + manifest = load_h3_manifest() + identity = manifest["model_identity"] + out_dir: Path = args.out.expanduser() + download_dir = (args.download_dir or out_dir / "hf").expanduser() + out_dir.mkdir(parents=True, exist_ok=True) + + def fetch(filename: str) -> Path: + """Download one file from the manifest's model subfolder and pinned revision.""" + + return Path( + hf_hub_download( + identity["hf_repo"], + f"{identity['subfolder']}/{filename}", + revision=identity["revision"], + local_dir=str(download_dir), + ) + ) + + _check(fetch("config.json"), identity["config_sha256"]) + index_path = fetch(identity["index_file"]) + _check(index_path, identity["index_sha256"]) + weight_map = json.loads(index_path.read_text())["weight_map"] + + def extract(section: str, filename: str) -> None: + """Pull one manifest section's tensors from their pinned shards, verify, and save.""" + + specs = manifest[section] + tensors = {} + for shard in sorted({weight_map[name] for name in specs}): + if shard not in manifest["weight_shards"]: + raise SystemExit(f"shard {shard} is not pinned in the manifest") + shard_path = fetch(shard) + _check(shard_path, manifest["weight_shards"][shard]["sha256"]) + with safe_open(str(shard_path), framework="pt") as handle: + for name in specs: + if weight_map[name] == shard: + tensors[name] = handle.get_tensor(name).contiguous() + + for name, spec in specs.items(): + tensor = tensors[name] + if str(tensor.dtype).removeprefix("torch.") != spec["dtype"]: + raise SystemExit(f"{name}: dtype {tensor.dtype} != manifest {spec['dtype']}") + if list(tensor.shape) != spec["shape"]: + raise SystemExit(f"{name}: shape {list(tensor.shape)} != manifest {spec['shape']}") + actual = sha256_tensor(tensor) + if actual != spec["sha256"]: + raise SystemExit(f"{name}: sha256 {actual} != manifest {spec['sha256']}") + + target = out_dir / filename + save_file(tensors, str(target), metadata={"revision": identity["revision"]}) + print(f"wrote {len(tensors)} tensors to {target}") + + extract("tensors", EXTRACTED_FILE) + if args.block: + extract("block_tensors", BLOCK_FILE) + print(f"export RL_KERNEL_H3_WEIGHTS={out_dir}") + + +if __name__ == "__main__": + main()