From ed3377616024ecb835f2ba3e44b9bf0544adc4b3 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 06:08:30 +0200
Subject: [PATCH 01/14] feat(h3): WS2 parallel rows (TP/SP/CP Ulysses),
hierarchical cross-node collective, 2-node H100 evidence
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 7f1296a970efb0cfcebd9300ff92551de08cf311)
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit d56c391c6eb74c098195430cd96c8b457f14f0f0)
---
docs/operators/README.md | 1 +
docs/operators/h3-parallel.md | 87 +++++
.../model_specific/minimax_h3/parallel.py | 336 ++++++++++++++++++
.../distributed/algorithms/hierarchical.py | 112 ++++++
rl_engine/validation/models/h3_tp.py | 204 +++++++++++
tests/models/minimax_h3/test_h3_parallel.py | 299 ++++++++++++++++
tools/validation/models/h3_parallel_run.py | 115 ++++++
7 files changed, 1154 insertions(+)
create mode 100644 docs/operators/h3-parallel.md
create mode 100644 rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
create mode 100644 rl_engine/distributed/algorithms/hierarchical.py
create mode 100644 rl_engine/validation/models/h3_tp.py
create mode 100644 tests/models/minimax_h3/test_h3_parallel.py
create mode 100644 tools/validation/models/h3_parallel_run.py
diff --git a/docs/operators/README.md b/docs/operators/README.md
index c6724bf11..9d0e3a2c6 100644
--- a/docs/operators/README.md
+++ b/docs/operators/README.md
@@ -50,3 +50,4 @@ Every operator page should include:
- [MiniMax-H3 Full Transformer and Dual Trajectory](h3-full-transformer.md)
- [MiniMax-H3 Integration Rows](h3-integration.md)
- [MiniMax-H3 Trajectory Artifact](h3-trajectory-artifact.md)
+- [MiniMax-H3 Tensor/Sequence/Context Parallel Rows](h3-parallel.md)
diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md
new file mode 100644
index 000000000..b92c68342
--- /dev/null
+++ b/docs/operators/h3-parallel.md
@@ -0,0 +1,87 @@
+# MiniMax-H3 Tensor / Sequence / Context Parallel Rows (WS2)
+
+## Summary
+
+RFC #420 rows `tp_attention_7168`, `tp_ffn_14336`, `cp_full_attention` (Ulysses),
+`cp_metadata_coshard`, `cp_arbitrary_length`, `distributed_packed_io`, and closeouts
+`ws2_tp2_tp4_tp8`, `ws2_tp4_cp2`, `ws2_tp4_cp2_sp`. Contract `h3-parallel-v1`
+(`rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py`); the block that composes them is
+`rl_engine/validation/models/h3_tp.py`. Cross-node reductions use
+`rl_engine/distributed/algorithms/hierarchical.py` (`hierarchical_ipc_nccl_tree`).
+
+Result (H100, 2 nodes x 4 GPUs, block 0 of the pinned checkpoint, 264 packed rows): every rank's
+rows of the forward, the input gradient and every weight-gradient shard are **bitwise equal** to the
+single-GPU canonical block for TP2, TP4, TP8 (two nodes), TP4 x CP2 (two nodes) and TP4 x CP2 + SP.
+Evidence: `reports/experiments/h3-ws2-parallel-h100/*.json`.
+
+## The canonical reduction
+
+All cross-rank sums are defined on **8 canonical units** (the divisors of 8 are the TP sizes):
+
+| Layer | Sharding | Reduction | Rule |
+| --- | --- | --- | --- |
+| Q/K/V, gate/up forward | output rows (whole heads / halves) | none | the unsharded GEMM's columns (`h3-projection-gemm-v1`) |
+| O (7168 -> 5376), down (14336 -> 5376) forward | reduction axis, `8/tp` K-units per rank | K | each K-unit is one fixed-order FP32 GEMM; a rank adds its units as a balanced tree; the collective adds ranks as the same tree; one BF16 rounding |
+| O / down backward (`dX`, `dW`) | - | N / tokens | plain fixed-order GEMM, independent of the shard |
+| Q/K/V, gate/up backward `dX` | output rows | N | unit `j` = block `j` of every fused weight, summed in the order q,k,v (hidden,gate); tree over units; collective; one rounding |
+
+The unit of the O projection is 896 K (14 K-tiles) and of down 1792 (28 K-tiles), so a unit never
+splits a tile. FP32 partials cross ranks (`DeterministicCollective` fixed tree on a node;
+`HierarchicalCollective` across nodes). The row-parallel contract differs from the single-sweep
+K reduction of the unsharded O/down GEMM rows (claimed by another contributor): it is the
+TP-invariant definition, and TP=1 is that definition, not the sweep.
+
+## Cross-node collective
+
+`DeterministicCollective` is CUDA IPC and therefore single node. `HierarchicalCollective(local_size)`
+reduces each group of `local_size` ranks with that tree, moves the group partials with NCCL
+`all_gather` (exact copies, no NCCL arithmetic) and adds them in one fixed balanced order. For
+power-of-two sizes this is the 8-way tree; a unit test checks it equals the single-node tree of the
+same world on 4 GPUs for FP32 and BF16 (`local_size` 1 and 2), and the TP8 run above uses it across
+two nodes.
+
+## Context parallelism
+
+- `SeqShard`: contiguous row slices differing by at most one row, any length >= 0 (including fewer
+ rows than ranks): `cp_arbitrary_length`.
+- `coshard`: hidden, AdaLN row index, token tags, rotary tables are cut with one `SeqShard` so the
+ metadata follows its rows: `cp_metadata_coshard`.
+- `ulysses_attention`: all-to-all sequence <-> head (ragged split sizes), attention on `H/cp` heads
+ of the full sequence, all-to-all back; both moves are exact, so rows equal the unsharded
+ attention's if the attention is head independent (the repo's deterministic attention is). Tested at
+ lengths 1, 2, 7, 130 and in the block at 261 / 263 rows.
+- `scatter_packed` / `gather_packed`: exact packed IO, ragged gather by padding: `distributed_packed_io`.
+- `seq_sharded_linear`: weight gradient from gathered rows so it equals the unsharded token sweep.
+
+## Sequence parallelism
+
+With `sp=True` the residual stream is sharded by rows across the TP group: norm and gate-residual are
+row-local, the column-parallel input is all-gathered (exact) and the row-parallel output is
+reduce-scattered with the same tree (`rows % tp == 0` required, checked).
+
+## Tests and evidence
+
+`tests/models/minimax_h3/test_h3_parallel.py`: 70 tests on 4 H100 (sharding, canonical folding for
+TP 2/4/8, accuracy against FP64, row independence, Ulysses ragged, ten block grids with backward,
+hierarchical collective). `tools/validation/models/h3_parallel_run.py` runs a grid under `srun` on
+any node count and compares against the single-GPU block.
+
+| Grid | GPUs / nodes | Collective | Forward | dX | dW shards |
+| --- | --- | --- | --- | --- | --- |
+| TP2 | 2 / 1 | IPC tree | bitwise | bitwise | bitwise |
+| TP4 | 4 / 1 | IPC tree | bitwise | bitwise | bitwise |
+| TP8 | 8 / 2 | hierarchical | bitwise | bitwise | bitwise |
+| TP4 x CP2 | 8 / 2 | IPC tree (TP in node), NCCL all-to-all (CP across) | bitwise | bitwise | not compared (token axis sharded) |
+| TP4 x CP2 + SP | 8 / 2 | same | bitwise | bitwise | not compared |
+
+Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4xCP2+SP 104
+(tiny 264-row case, launch-bound; not a scaling claim).
+
+## Limits
+
+- Attention, SwiGLU and QK norm inside the block are the repo's interim operators; the parallel
+ rows only require them to be row/head independent, which the bitwise results confirm.
+- Pipeline parallelism is out of scope (`ws2_pipeline_parallel` is not part of this change).
+- dW under CP is exact through `seq_sharded_linear` for non-TP linears; the TP x CP block compares
+ forward and `dX` (CP shards the token axis of the TP weight gradients).
+- `HierarchicalCollective` assumes ranks are packed node-major and node counts of 1/2/4/8.
diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
new file mode 100644
index 000000000..eabef82fd
--- /dev/null
+++ b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""MiniMax-H3 tensor-, sequence- and context-parallel building blocks (contract ``h3-parallel-v1``).
+
+Rows: ``tp_attention_7168``, ``tp_ffn_14336``, ``cp_full_attention`` (Ulysses),
+``cp_metadata_coshard``, ``cp_arbitrary_length``, ``distributed_packed_io``.
+
+Every reduction that crosses ranks is defined on **8 canonical units** so TP 1/2/4/8 add the
+same numbers in the same order:
+
+* *row-parallel* (O and FFN-down, ``K`` = 7168 / 14336): the reduction axis is cut into 8
+ equal K-units; a unit is one fixed-order FP32 GEMM (``h3-projection-gemm-v1`` tile, ``K``
+ swept ascending). A rank holding ``8/tp`` consecutive units adds them in balanced-tree
+ order and the collective adds the ranks' partials in the same tree: the result is the
+ balanced tree over the 8 units for every TP size, rounded to BF16 once.
+* *column-parallel backward* (``dX`` of Q/K/V and gate/up, reduction over the output axis):
+ the same 8 units, where unit ``j`` is output block ``j`` of every weight of the fused
+ projection, added in the order given (q, k, v / hidden, gate).
+
+The forward of a column-parallel projection and the backward of a row-parallel one reduce
+over axes that do not depend on the sharding (they are plain ``h3-projection-gemm-v1``), so
+they are bitwise the unsharded result without any collective.
+
+Context parallelism (Ulysses) moves data only: sequence <-> head all-to-alls are exact, so
+attention (per head, full sequence) is the unsharded attention. Weight gradients under
+sequence sharding are computed from the gathered rows, so they equal the unsharded GEMM.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+
+from rl_engine.backends.shared.triton.gemm import h3_projection as g
+from rl_engine.distributed.algorithms.hierarchical import tree_add
+
+CONTRACT = "h3-parallel-v1"
+UNITS = 8 # canonical reduction units; TP sizes are the divisors 1, 2, 4, 8
+TP_SIZES = (1, 2, 4, 8)
+HEADS, HEAD_DIM = 56, 128
+INNER = HEADS * HEAD_DIM # 7168
+FFN = 14336
+
+
+# --- Sharding helpers ---------------------------------------------------------------------------
+
+
+def check_tp(tp: int, rank: int) -> None:
+ if tp not in TP_SIZES or not 0 <= rank < tp:
+ raise ValueError(f"tp must be one of {TP_SIZES} and rank in [0, tp); got tp={tp} rank={rank}")
+
+
+def shard_rows(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor:
+ """Output-axis shard of a column-parallel weight ``(N, K)``: rows ``[r*N/tp, (r+1)*N/tp)``."""
+
+ check_tp(tp, rank)
+ n = weight.shape[0]
+ if n % UNITS:
+ raise ValueError(f"{n} output rows do not split into {UNITS} units")
+ return weight.narrow(0, rank * n // tp, n // tp).contiguous()
+
+
+def shard_cols(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor:
+ """Reduction-axis shard of a row-parallel weight ``(N, K)``: columns of ``8/tp`` K-units."""
+
+ check_tp(tp, rank)
+ k = weight.shape[1]
+ unit = k // UNITS
+ if unit * UNITS != k or unit % g.CONFIG_BF16["BLOCK_K"]:
+ raise ValueError(f"K={k} must split into {UNITS} units of a multiple of the GEMM K tile")
+ return weight.narrow(1, rank * k // tp, k // tp).contiguous()
+
+
+def shard_heads(x: torch.Tensor, tp: int, rank: int, dim: int = -1) -> torch.Tensor:
+ """Head shard of a ``(..., 7168)`` activation or weight axis (whole 128-wide heads)."""
+
+ check_tp(tp, rank)
+ return x.narrow(dim, rank * x.shape[dim] // tp, x.shape[dim] // tp).contiguous()
+
+
+def shard_gate_up(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor:
+ """Shard of the fused ``(2*FFN, K)`` gate/up weight keeping ``[hidden; gate]`` halves.
+
+ Each half is sharded on its own, so local rows are ``[hidden shard; gate shard]`` and
+ ``chunk(2)`` still splits them.
+ """
+
+ half = weight.shape[0] // 2
+ return torch.cat(
+ [shard_rows(weight[:half], tp, rank), shard_rows(weight[half:], tp, rank)]
+ ).contiguous()
+
+
+# --- Row-parallel linear (reduction over the sharded axis) ---------------------------------------
+
+
+def _units_partial(x: torch.Tensor, w: torch.Tensor, tp: int) -> torch.Tensor:
+ """FP32 balanced-tree sum of this rank's ``8/tp`` K-units of ``x @ w.T``."""
+
+ local = UNITS // tp
+ unit = x.shape[1] // local
+ parts = [
+ g.gemm(x[:, i * unit : (i + 1) * unit], w[:, i * unit : (i + 1) * unit], b_t=True,
+ out_dtype=torch.float32)
+ for i in range(local)
+ ]
+ return tree_add(parts)
+
+
+class _RowParallel(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x, w, collective, tp, scatter):
+ ctx.save_for_backward(x, w)
+ ctx.scatter, ctx.collective = scatter, collective
+ partial = _units_partial(x, w, tp)
+ if tp > 1:
+ partial = collective.reduce_scatter(partial) if scatter else collective.all_reduce(partial)
+ return partial.to(x.dtype)
+
+ @staticmethod
+ def backward(ctx, grad):
+ x, w = ctx.saved_tensors
+ if ctx.scatter: # sequence-parallel: the rows of dY are sharded; gather them (exact)
+ grad = ctx.collective.all_gather(grad.contiguous())
+ grad = grad.contiguous()
+ dx = g.gemm(grad, w) # reduction over N: independent of the shard
+ dw = g.gemm(grad, x, a_t=True) # reduction over tokens: independent of the shard
+ return dx, dw, None, None, None
+
+
+def row_parallel_linear(x, w_shard, collective=None, tp: int = 1, scatter: bool = False):
+ """``y = x @ W.T`` with ``x (T, K/tp)`` and ``w_shard (N, K/tp)``; ``y`` bitwise TP-invariant.
+
+ ``scatter`` gives this rank's ``T/tp`` rows of ``y`` (sequence parallel) instead of all of it.
+ """
+
+ if x.dim() != 2 or w_shard.dim() != 2 or x.shape[1] != w_shard.shape[1]:
+ raise ValueError("row_parallel_linear needs x (T, K/tp) and w (N, K/tp)")
+ if tp > 1 and collective is None:
+ raise ValueError("tp > 1 needs a collective")
+ return _RowParallel.apply(x.contiguous(), w_shard, collective, tp, scatter)
+
+
+# --- Column-parallel linear (backward reduces over the sharded axis) ------------------------------
+
+
+class _ColumnParallel(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x, collective, tp, gather, *weights):
+ ctx.save_for_backward(x, *weights)
+ ctx.collective, ctx.tp, ctx.gather = collective, tp, gather
+ full_x = collective.all_gather(x) if gather else x
+ ctx.full_x = full_x if gather else None
+ return tuple(g.gemm(full_x, w, b_t=True) for w in weights)
+
+ @staticmethod
+ def backward(ctx, *grads):
+ x, *weights = ctx.saved_tensors
+ full_x = ctx.full_x if ctx.gather else x
+ local = UNITS // ctx.tp
+ grads = [t.contiguous() for t in grads]
+ units = []
+ for j in range(local):
+ acc = None
+ for grad, w in zip(grads, weights):
+ width = grad.shape[1] // local
+ piece = g.gemm(grad[:, j * width : (j + 1) * width], w[j * width : (j + 1) * width],
+ out_dtype=torch.float32)
+ acc = piece if acc is None else acc + piece
+ units.append(acc)
+ partial = tree_add(units)
+ if ctx.tp > 1:
+ partial = (ctx.collective.reduce_scatter(partial) if ctx.gather
+ else ctx.collective.all_reduce(partial))
+ dx = partial.to(x.dtype)
+ dws = tuple(g.gemm(grad, full_x, a_t=True) for grad in grads)
+ return (dx, None, None, None, *dws)
+
+
+def column_parallel_linear(x, weights: Sequence[torch.Tensor], collective=None, tp: int = 1,
+ gather: bool = False):
+ """``[x @ w.T for w in weights]`` with each ``w`` a row shard ``(N/tp, K)``.
+
+ The forward is the unsharded projection's columns (no collective). With ``gather``, ``x`` is
+ this rank's rows of a sequence-parallel activation, gathered first (exact), and ``dX`` is
+ reduce-scattered back to rows. ``dX`` is the canonical 8-unit tree for every ``tp``.
+ """
+
+ if tp > 1 and collective is None:
+ raise ValueError("tp > 1 needs a collective")
+ return _ColumnParallel.apply(x.contiguous(), collective, tp, gather, *weights)
+
+
+# --- Context parallelism (Ulysses) ---------------------------------------------------------------
+
+
+@dataclass(frozen=True)
+class SeqShard:
+ """This rank's contiguous slice of a packed sequence of ``length`` rows over ``cp`` ranks.
+
+ Slices differ by at most one row (first ``length % cp`` ranks get the extra row), so any
+ length is legal, including ``length < cp`` (some ranks hold nothing).
+ """
+
+ length: int
+ cp: int
+ rank: int
+
+ @property
+ def sizes(self) -> list[int]:
+ base, extra = divmod(self.length, self.cp)
+ return [base + (r < extra) for r in range(self.cp)]
+
+ @property
+ def lo(self) -> int:
+ return sum(self.sizes[: self.rank])
+
+ @property
+ def hi(self) -> int:
+ return self.lo + self.sizes[self.rank]
+
+
+def coshard(shard: SeqShard, *tensors: torch.Tensor, dim: int = 0) -> list[torch.Tensor]:
+ """Slice every per-row tensor (hidden, timestep_indices, token_tags, position_ids, ...) alike.
+
+ The metadata that selects AdaLN rows and rotary angles must follow the rows it describes;
+ slicing them in one call with one ``SeqShard`` is what keeps them aligned.
+ """
+
+ out = []
+ for t in tensors:
+ if t.shape[dim] != shard.length:
+ raise ValueError(f"tensor has {t.shape[dim]} rows along dim {dim}, expected {shard.length}")
+ out.append(t.narrow(dim, shard.lo, shard.hi - shard.lo).contiguous())
+ return out
+
+
+def scatter_packed(packed: torch.Tensor, shard: SeqShard, dim: int = 0) -> torch.Tensor:
+ """This rank's rows of a packed tensor that every rank holds in full."""
+
+ return coshard(shard, packed, dim=dim)[0]
+
+
+def gather_packed(local: torch.Tensor, shard: SeqShard, group=None, dim: int = 0) -> torch.Tensor:
+ """Reassemble the packed tensor from every rank's rows (exact, any length)."""
+
+ if shard.cp == 1:
+ return local
+ sizes = shard.sizes
+ widest = max(sizes)
+ moved = local.movedim(dim, 0).contiguous()
+ pad = moved.new_zeros((widest - moved.shape[0], *moved.shape[1:]))
+ chunks = [torch.empty_like(torch.cat([moved, pad])) for _ in range(shard.cp)]
+ dist.all_gather(chunks, torch.cat([moved, pad]), group=group)
+ return torch.cat([c[:n] for c, n in zip(chunks, sizes)]).movedim(0, dim)
+
+
+class _AllToAll(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x, in_sizes, out_sizes, group):
+ ctx.sizes, ctx.group = (in_sizes, out_sizes), group
+ out = x.new_empty((sum(out_sizes), *x.shape[1:]))
+ dist.all_to_all_single(out, x.contiguous(), out_sizes, in_sizes, group=group)
+ return out
+
+ @staticmethod
+ def backward(ctx, grad):
+ in_sizes, out_sizes = ctx.sizes
+ back = grad.new_empty((sum(in_sizes), *grad.shape[1:]))
+ dist.all_to_all_single(back, grad.contiguous(), in_sizes, out_sizes, group=ctx.group)
+ return back, None, None, None
+
+
+def ulysses_attention(q, k, v, attention, shard: SeqShard, group=None):
+ """Full attention of a sequence-sharded ``q, k, v (B, S_local, H, D)`` by head scatter.
+
+ All-to-all makes each rank hold ``H/cp`` heads for the whole sequence, ``attention(q, k, v)``
+ (``(B, H', S, D)`` layout, per-head independent) runs, and the result goes back. Both moves
+ are exact, so the output rows equal the unsharded attention's rows bitwise as long as
+ ``attention`` is batch/head invariant.
+ """
+
+ cp = shard.cp
+ if cp == 1:
+ return attention(q, k, v)
+ batch, local_rows, heads, dim = q.shape
+ if heads % cp:
+ raise ValueError(f"{heads} heads do not divide over cp={cp}")
+ sizes, hp = shard.sizes, heads // cp
+ if local_rows != sizes[shard.rank]:
+ raise ValueError(f"expected {sizes[shard.rank]} local rows, got {local_rows}")
+
+ def to_heads(t):
+ # peer r receives head block r of my rows; I receive every peer's rows of my head block
+ send = t.reshape(batch, local_rows, cp, hp, dim).permute(2, 1, 0, 3, 4)
+ send = send.reshape(cp * local_rows, batch, hp, dim)
+ got = _AllToAll.apply(send, [local_rows] * cp, sizes, group)
+ return got.permute(1, 0, 2, 3) # (B, S, hp, D), sequence in rank order
+
+ out = attention(to_heads(q), to_heads(k), to_heads(v)) # (B, S, hp, D)
+ send = out.permute(1, 0, 2, 3).contiguous()
+ got = _AllToAll.apply(send, sizes, [local_rows] * cp, group)
+ got = got.reshape(cp, local_rows, batch, hp, dim)
+ return got.permute(2, 1, 0, 3, 4).reshape(batch, local_rows, heads, dim)
+
+
+# --- Weight gradients under sequence sharding ----------------------------------------------------
+
+
+class _SeqLinear(torch.autograd.Function):
+ """Row-local linear on a sequence shard whose ``dW`` is computed from gathered rows."""
+
+ @staticmethod
+ def forward(ctx, x, w, shard, group):
+ ctx.save_for_backward(x, w)
+ ctx.shard, ctx.group = shard, group
+ return g.gemm(x, w, b_t=True)
+
+ @staticmethod
+ def backward(ctx, grad):
+ x, w = ctx.saved_tensors
+ grad = grad.contiguous()
+ dx = g.gemm(grad, w)
+ full_x = gather_packed(x, ctx.shard, ctx.group)
+ full_g = gather_packed(grad, ctx.shard, ctx.group)
+ return dx, g.gemm(full_g, full_x, a_t=True), None, None
+
+
+def seq_sharded_linear(x, w, shard: SeqShard, group=None):
+ """``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM's (token sweep, ascending)."""
+
+ return _SeqLinear.apply(x.contiguous(), w, shard, group) if shard.cp > 1 else g.linear_forward(x, w, None)
diff --git a/rl_engine/distributed/algorithms/hierarchical.py b/rl_engine/distributed/algorithms/hierarchical.py
new file mode 100644
index 000000000..bfc29f04d
--- /dev/null
+++ b/rl_engine/distributed/algorithms/hierarchical.py
@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Cross-node deterministic collective: node-local IPC fixed tree, then one exact gather.
+
+``DeterministicCollective`` (``cuda_ipc_fixed_tree``) reads peers' memory through CUDA IPC,
+so it is single node. This wrapper (``hierarchical_ipc_nccl_tree``) reaches the same bits on
+several nodes without relying on any NCCL reduction order:
+
+1. every group of ``local_size`` consecutive ranks reduces with the node-local fixed tree;
+2. the group partials move with NCCL ``all_gather`` (bit-exact data movement, no arithmetic);
+3. each rank adds the group partials in one fixed, balanced order.
+
+With ``world_size = nodes * local_size`` and ``local_size``, ``nodes`` powers of two, the
+tree ``((0+1)+(2+3)) + ((4+5)+(6+7))`` of an 8-way node-local reduction is exactly "the 4-way
+tree of each half, then the two halves added". The input dtype is the accumulation dtype, as in
+``DeterministicCollective`` (FP32 inputs keep every partial in FP32, which is what the H3
+parallel operators use), so the result is bitwise the single-node tree of the same world size.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.distributed as dist
+
+from rl_engine.distributed.algorithms.collectives import DeterministicCollective
+
+_GROUP_COUNTS = (1, 2, 4, 8)
+
+
+def tree_add(parts: list[torch.Tensor]) -> torch.Tensor:
+ """Balanced pairwise sum ``((p0+p1)+(p2+p3))+...`` of a power-of-two number of tensors."""
+
+ if len(parts) & (len(parts) - 1) or not parts:
+ raise ValueError(f"tree reduction needs a power-of-two part count, got {len(parts)}")
+ while len(parts) > 1:
+ parts = [parts[i] + parts[i + 1] for i in range(0, len(parts), 2)]
+ return parts[0]
+
+
+class HierarchicalCollective:
+ """TP-invariant all-reduce / reduce-scatter over ``world_size`` ranks on several nodes."""
+
+ backend_id = "hierarchical_ipc_nccl_tree"
+
+ def __init__(
+ self,
+ local_size: int,
+ group: dist.ProcessGroup | None = None,
+ device: torch.device | str | int | None = None,
+ *,
+ max_size_bytes: int = 64 << 20,
+ ) -> None:
+ if not dist.is_initialized():
+ raise RuntimeError("torch.distributed must be initialized before collectives")
+ self.group = group if group is not None else dist.group.WORLD
+ self.world_size = dist.get_world_size(group=self.group)
+ self.rank = dist.get_rank(group=self.group)
+ if self.world_size % local_size or local_size not in _GROUP_COUNTS:
+ raise ValueError(
+ f"local_size {local_size} must be 1/2/4/8 and divide world {self.world_size}"
+ )
+ self.local_size = local_size
+ self.num_groups = self.world_size // local_size
+ if self.num_groups not in _GROUP_COUNTS:
+ raise ValueError(f"{self.num_groups} node groups; need 1, 2, 4 or 8")
+ self.group_index, self.local_rank = divmod(self.rank, local_size)
+ # Every rank creates every group in the same order (a torch.distributed requirement).
+ ranks = dist.get_process_group_ranks(self.group) # global rank of each group member
+ local_groups = [
+ dist.new_group(ranks[g * local_size : (g + 1) * local_size])
+ for g in range(self.num_groups)
+ ]
+ cross_groups = [
+ dist.new_group([ranks[g * local_size + r] for g in range(self.num_groups)])
+ for r in range(local_size)
+ ]
+ self._cross = cross_groups[self.local_rank]
+ self._local = DeterministicCollective(
+ local_groups[self.group_index], device, max_size_bytes=max_size_bytes
+ )
+
+ def all_reduce(self, input: torch.Tensor) -> torch.Tensor:
+ """The fixed-tree sum of every rank's ``input`` on every rank (new tensor)."""
+
+ if input.dtype not in (torch.float32, torch.bfloat16, torch.float16):
+ raise TypeError(f"unsupported dtype {input.dtype}")
+ wide = input.contiguous() # the input dtype is the accumulation dtype, as in the node-local tree
+ partial = self._local.all_reduce(wide) if self.local_size > 1 else wide.clone()
+ if self.num_groups > 1:
+ gathered = [torch.empty_like(partial) for _ in range(self.num_groups)]
+ dist.all_gather(gathered, partial, group=self._cross)
+ partial = tree_add(gathered)
+ return partial
+
+ def reduce_scatter(self, input: torch.Tensor) -> torch.Tensor:
+ """``all_reduce`` then this rank's equal share of dimension 0."""
+
+ if input.shape[0] % self.world_size:
+ raise ValueError("dimension 0 must divide by the world size")
+ rows = input.shape[0] // self.world_size
+ return self.all_reduce(input)[self.rank * rows : (self.rank + 1) * rows].contiguous()
+
+ def all_gather(self, input: torch.Tensor) -> torch.Tensor:
+ """Rank-ordered concatenation along dimension 0 (exact data movement)."""
+
+ out = [torch.empty_like(input) for _ in range(self.world_size)]
+ dist.all_gather(out, input.contiguous(), group=self.group)
+ return torch.cat(out)
+
+ def close(self) -> None:
+ self._local.close()
diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py
new file mode 100644
index 000000000..697010300
--- /dev/null
+++ b/rl_engine/validation/models/h3_tp.py
@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""One H3 block under tensor / sequence / context parallelism (RFC #420 WS2 closeouts).
+
+``run_tp_block`` is the block graph with the attention and FFN projections replaced by the
+``h3-parallel-v1`` operators: column-parallel Q/K/V and gate/up, per-head norm, rotary and
+attention on the rank's heads, Ulysses head<->sequence all-to-all across the CP group,
+row-parallel O and down with the deterministic collective, optional sequence parallelism
+(rows of the residual stream sharded across the TP group). With ``tp=cp=1`` it is the
+canonical single-GPU block every other layout must equal bitwise.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
+from rl_engine.validation.models import h3_block as hb
+
+HIDDEN, HEAD_DIM = hb.HIDDEN, hb.HEAD_DIM
+
+
+@dataclass
+class Grid:
+ tp: int
+ cp: int
+ tp_rank: int
+ cp_rank: int
+ tp_collective: Any = None
+ cp_group: Any = None
+
+ def seq_shard(self, length: int) -> par.SeqShard:
+ return par.SeqShard(length, self.cp, self.cp_rank)
+
+
+def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid:
+ """Rank ``r`` is ``(cp_rank, tp_rank) = divmod(r, tp)``; TP groups are consecutive ranks.
+
+ ``local_size`` (ranks per node) selects the cross-node collective when the TP group spans
+ nodes; ``None`` means the TP group is on one node (CUDA IPC fixed tree).
+ """
+
+ from rl_engine.distributed.algorithms.collectives import DeterministicCollective
+ from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective
+
+ world, rank = dist.get_world_size(), dist.get_rank()
+ if world != tp * cp:
+ raise ValueError(f"world {world} != tp {tp} x cp {cp}")
+ cp_rank, tp_rank = divmod(rank, tp)
+ tp_groups = [dist.new_group(list(range(c * tp, (c + 1) * tp))) for c in range(cp)]
+ cp_groups = [dist.new_group(list(range(t, world, tp))) for t in range(tp)]
+ collective = None
+ if tp > 1:
+ mine = tp_groups[cp_rank]
+ if local_size is not None and local_size < tp:
+ collective = HierarchicalCollective(local_size, mine, torch.cuda.current_device(),
+ max_size_bytes=256 << 20)
+ else:
+ collective = DeterministicCollective(mine, torch.cuda.current_device(),
+ max_size_bytes=256 << 20)
+ return Grid(tp, cp, tp_rank, cp_rank, collective, cp_groups[tp_rank] if cp > 1 else None)
+
+
+def _attention_fn(ops):
+ def attend(q, k, v): # (B, S, H', D) -> (B, S, H', D)
+ q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, v))
+ return ops.attention(q, k, v, causal=False).transpose(1, 2).contiguous()
+
+ return attend
+
+
+def run_tp_block(
+ weights: dict[str, torch.Tensor],
+ hidden: torch.Tensor,
+ temb: torch.Tensor,
+ index: torch.Tensor,
+ rope: tuple[torch.Tensor, torch.Tensor],
+ ops: hb.CandidateOps,
+ grid: Grid,
+ *,
+ sp: bool = False,
+ seq_len: int | None = None,
+ leaves: dict[str, torch.Tensor] | None = None,
+) -> torch.Tensor:
+ """Block 0 for this rank. ``hidden/index/rope`` are the full packed sequence on every rank.
+
+ The rank slices its CP rows (``cp_metadata_coshard``: hidden, AdaLN row index and both
+ rotary tables with the same ``SeqShard``) and, with ``sp``, its TP row share of those.
+ Returns this rank's rows of the block output (CP rows, or the SP share of them). With
+ ``leaves`` the weight shards are tracked and stored there by name so the caller can read
+ their gradients.
+ """
+
+ if hidden.shape[0] != 1:
+ raise ValueError("the parallel block takes one packed sequence (batch 1)")
+ tp, w = grid.tp, (lambda name: weights[f"{hb.BLOCK}.{name}"])
+ seq = seq_len or hidden.shape[1]
+ shard = grid.seq_shard(seq)
+ h = par.coshard(shard, hidden, dim=1)[0]
+ idx = par.coshard(shard, index, dim=0)[0]
+ cos, sin = par.coshard(shard, rope[0], rope[1], dim=0)
+ rows = h.shape[1]
+ if sp:
+ if rows % tp:
+ raise ValueError(f"sequence parallelism needs {rows} local rows divisible by tp={tp}")
+ share = rows // tp
+ lo, hi = grid.tp_rank * share, (grid.tp_rank + 1) * share
+ h_res, idx_res = h[:, lo:hi].contiguous(), idx[lo:hi].contiguous()
+ else:
+ h_res, idx_res = h, idx
+ coll = grid.tp_collective
+
+ def leaf(name, tensor):
+ tensor = tensor.detach().clone()
+ if leaves is not None:
+ tensor.requires_grad_(True)
+ leaves[name] = tensor
+ return tensor
+ table = ops.projection(temb, w("adaln_proj.linear.weight"), w("adaln_proj.linear.bias"))
+
+ def proj_in(x, name_weights, gather):
+ flat = x.reshape(-1, HIDDEN)
+ return par.column_parallel_linear(flat, name_weights, coll, tp, gather=gather)
+
+ x = ops.norm.forward_modulated(h_res, w("norm1.weight"), table[0], table[1], idx_res)
+ q, k, v = proj_in(x, [leaf(f"attn.to_{n}.weight", par.shard_rows(w(f"attn.to_{n}.weight"), tp, grid.tp_rank)) for n in "qkv"], sp)
+ local_heads = par.HEADS // tp
+
+ def heads(t):
+ return t.view(1, rows, local_heads, HEAD_DIM)
+
+ q = ops.norm(heads(q), w("attn.norm_q.weight"))
+ k = ops.norm(heads(k), w("attn.norm_k.weight"))
+ from rl_engine.validation.models import h3_provider as prov
+
+ q, k = prov.provider_apply_rotary(q, cos, sin), prov.provider_apply_rotary(k, cos, sin)
+ attn = par.ulysses_attention(q, k, heads(v), _attention_fn(ops), shard, grid.cp_group)
+ o = par.row_parallel_linear(
+ attn.reshape(rows, local_heads * HEAD_DIM),
+ leaf("attn.to_out.0.weight", par.shard_cols(w("attn.to_out.0.weight"), tp, grid.tp_rank)), coll, tp, scatter=sp,
+ )
+ h1 = ops.gate(h_res, o.view(1, -1, HIDDEN), table[2], idx_res)
+ x2 = ops.norm.forward_modulated(h1, w("norm2.weight"), table[3], table[4], idx_res)
+ fused = w("ff.net.0.proj.weight")
+ half = fused.shape[0] // 2
+ hid, gate = proj_in(
+ x2,
+ [leaf("ff.hidden", par.shard_rows(fused[:half], tp, grid.tp_rank)),
+ leaf("ff.gate", par.shard_rows(fused[half:], tp, grid.tp_rank))],
+ sp,
+ )
+ act = ops.swiglu(gate.contiguous(), hid.contiguous())
+ down = par.row_parallel_linear(
+ act, leaf("ff.net.2.weight", par.shard_cols(w("ff.net.2.weight"), tp, grid.tp_rank)), coll, tp, scatter=sp
+ )
+ return ops.gate(h1, down.view(1, -1, HIDDEN), table[5], idx_res)
+
+
+def block_target(
+ collective, dummy, *, tp: int, cp: int, sp: bool = False, seq_len: int = 0,
+ backward: bool = False, local_size: int = -1, layout: str = "tiny", seed: int = 0,
+):
+ """``run_world`` target: this rank's rows of the parallel block (and gradients) for one case."""
+
+ del collective # the grid builds its own TP/CP groups
+ from rl_engine.validation.models.h3_weights import load_h3_block_weights
+
+ grid = make_grid(tp, cp, local_size=local_size if local_size > 0 else None)
+ return block_case(load_h3_block_weights("cuda"), grid, sp=sp, seq_len=seq_len,
+ backward=backward, layout=layout, seed=seed)
+
+
+def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layout="tiny", seed=0):
+ """Run block 0 on ``grid``; the global rows this rank returns are in ``rows``."""
+
+ ctx = hb.make_block_context(weights, layout=layout, seed=seed)
+ n = seq_len or ctx["seq_len"]
+ hidden = ctx["hidden"][:, :n].clone().requires_grad_(backward)
+ index = ctx["index"][:n]
+ rope = (ctx["rope"][0][:n], ctx["rope"][1][:n])
+ ops = hb.CandidateOps()
+ leaves: dict[str, torch.Tensor] | None = {} if backward else None
+ with torch.set_grad_enabled(backward):
+ out = run_tp_block(weights, hidden, ctx["temb"], index, rope, ops, grid, sp=sp,
+ seq_len=n, leaves=leaves)
+ shard = grid.seq_shard(n)
+ lo, hi = shard.lo, shard.hi
+ if sp:
+ share = (hi - lo) // grid.tp
+ lo, hi = lo + grid.tp_rank * share, lo + (grid.tp_rank + 1) * share
+ result = {"out": out.detach(), "rows": torch.tensor([lo, hi])}
+ if backward:
+ gen = torch.Generator().manual_seed(seed + 99)
+ full_grad = torch.randn(1, n, HIDDEN, generator=gen).to(out.dtype).cuda()
+ out.backward(full_grad[:, lo:hi])
+ result["dx"] = hidden.grad[:, lo:hi].detach()
+ for name, tensor in leaves.items():
+ result["dw." + name] = tensor.grad.detach()
+ return result
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
new file mode 100644
index 000000000..19b5c0516
--- /dev/null
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""RFC #420 WS2 rows: TP attention/FFN, Ulysses CP, metadata co-shard, arbitrary length, packed IO.
+
+Single-process tests pin the canonical 8-unit reduction (summing per-rank partials in tree
+order reproduces the TP=1 result for TP 2/4/8) and the sharding helpers. Multi-process tests
+(2-4 GPUs) run the whole block 0 under TP / CP / SP grids and require every rank's rows to be
+bitwise equal to the single-GPU canonical block, forward and backward.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
+from rl_engine.distributed.algorithms.hierarchical import tree_add
+from rl_engine.validation.models import h3_block as hb
+from rl_engine.validation.models import h3_tp
+from rl_engine.validation.models.h3_ws2 import run_world
+
+needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def gpus(n):
+ return pytest.mark.skipif(torch.cuda.device_count() < n, reason=f"needs {n} GPUs")
+
+
+def _triton():
+ pytest.importorskip("triton")
+
+
+class TestSharding:
+ @pytest.mark.parametrize("tp", [1, 2, 4, 8])
+ def test_row_and_column_shards_tile_once(self, tp):
+ w = torch.arange(56 * 128 * 16, dtype=torch.float32).view(56 * 128, 16)
+ assert torch.equal(torch.cat([par.shard_rows(w, tp, r) for r in range(tp)]), w)
+ w2 = torch.arange(8 * 64 * 4, dtype=torch.float32).view(4, 8 * 64)
+ assert torch.equal(torch.cat([par.shard_cols(w2, tp, r) for r in range(tp)], 1), w2)
+
+ @pytest.mark.parametrize("tp", [3, 5, 6, 16])
+ def test_illegal_tp_is_rejected(self, tp):
+ with pytest.raises(ValueError):
+ par.shard_rows(torch.zeros(56 * 128, 4), tp, 0)
+
+ def test_shards_are_whole_heads(self):
+ for tp in (1, 2, 4, 8):
+ assert (par.HEADS // tp) * par.HEAD_DIM == par.INNER // tp
+
+ def test_gate_up_shards_keep_both_halves(self):
+ w = torch.arange(2 * 64 * 3, dtype=torch.float32).view(128, 3)
+ s = par.shard_gate_up(w, 4, 1)
+ hidden, gate = s.chunk(2)
+ assert torch.equal(hidden, w[:64][16:32]) and torch.equal(gate, w[64:][16:32])
+
+ @pytest.mark.parametrize("length", [1, 2, 3, 7, 261, 264, 925])
+ @pytest.mark.parametrize("cp", [1, 2, 3, 4])
+ def test_seq_shards_tile_any_length(self, length, cp):
+ spans = [(par.SeqShard(length, cp, r).lo, par.SeqShard(length, cp, r).hi) for r in range(cp)]
+ assert spans[0][0] == 0 and spans[-1][1] == length
+ assert all(a[1] == b[0] for a, b in zip(spans, spans[1:]))
+ sizes = [hi - lo for lo, hi in spans]
+ assert max(sizes) - min(sizes) <= 1
+
+ def test_coshard_slices_every_per_row_tensor_alike(self):
+ n = 11
+ idx, tags, rows = torch.arange(n), torch.arange(n) * 3, torch.arange(n * 2).view(n, 2)
+ for r in range(3):
+ s = par.SeqShard(n, 3, r)
+ a, b, c = par.coshard(s, idx, tags, rows)
+ assert torch.equal(a * 3, b) and torch.equal(c[:, 0] // 2, a)
+
+ def test_coshard_rejects_wrong_length(self):
+ with pytest.raises(ValueError, match="rows"):
+ par.coshard(par.SeqShard(5, 2, 0), torch.zeros(4))
+
+
+@needs_gpu
+class TestCanonicalReduction:
+ """Per-rank partials added in tree order are bitwise the TP=1 result (no collective needed)."""
+
+ @pytest.mark.parametrize("k", [7168, 14336])
+ @pytest.mark.parametrize("tp", [2, 4, 8])
+ def test_row_parallel_partials_fold_to_tp1(self, k, tp):
+ _triton()
+ gen = torch.Generator().manual_seed(1)
+ x = torch.randn(33, k, generator=gen).bfloat16().cuda()
+ w = (torch.randn(5376, k, generator=gen) / k**0.5).bfloat16().cuda()
+ full = par._units_partial(x, w, 1)
+ parts = [
+ par._units_partial(par.shard_heads(x, tp, r, 1), par.shard_cols(w, tp, r), tp)
+ for r in range(tp)
+ ]
+ assert torch.equal(tree_add(parts), full)
+ assert torch.equal(
+ par.row_parallel_linear(x, w), full.bfloat16()
+ )
+
+ def test_row_parallel_close_to_fp64(self):
+ _triton()
+ gen = torch.Generator().manual_seed(2)
+ x = torch.randn(64, 7168, generator=gen).bfloat16().cuda()
+ w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda()
+ got = par.row_parallel_linear(x, w).double()
+ ref = x.double() @ w.double().t()
+ provider = torch.nn.functional.linear(x, w).double()
+ err = (got - ref).norm() / ref.norm()
+ assert err <= 1.25 * (provider - ref).norm() / ref.norm() + 1e-3
+
+ def test_each_row_is_independent_of_the_others(self):
+ _triton()
+ gen = torch.Generator().manual_seed(3)
+ x = torch.randn(70, 7168, generator=gen).bfloat16().cuda()
+ w = (torch.randn(5376, 7168, generator=gen) / 7168**0.5).bfloat16().cuda()
+ whole = par.row_parallel_linear(x, w)
+ for rows in (slice(0, 1), slice(5, 6), slice(30, 31), slice(0, 33)):
+ assert torch.equal(par.row_parallel_linear(x[rows].contiguous(), w), whole[rows])
+
+ @pytest.mark.parametrize("tp", [2, 4, 8])
+ def test_column_parallel_backward_folds_to_tp1(self, tp):
+ _triton()
+ gen = torch.Generator().manual_seed(4)
+ k, n = 5376, 7168
+ x = torch.randn(40, k, generator=gen).bfloat16().cuda()
+ ws = [(torch.randn(n, k, generator=gen) / k**0.5).bfloat16().cuda() for _ in range(3)]
+ grads = [torch.randn(40, n, generator=gen).bfloat16().cuda() for _ in range(3)]
+
+ def unit_partials(weights, gs, local):
+ units = []
+ for j in range(local):
+ acc = None
+ for grad, w in zip(gs, weights):
+ width = grad.shape[1] // local
+ piece = par.g.gemm(
+ grad[:, j * width : (j + 1) * width], w[j * width : (j + 1) * width],
+ out_dtype=torch.float32,
+ )
+ acc = piece if acc is None else acc + piece
+ units.append(acc)
+ return tree_add(units)
+
+ full = unit_partials(ws, grads, 8)
+ ranks = [
+ unit_partials(
+ [par.shard_rows(w, tp, r) for w in ws],
+ [par.shard_heads(g_, tp, r, 1) for g_ in grads],
+ 8 // tp,
+ )
+ for r in range(tp)
+ ]
+ assert torch.equal(tree_add(ranks), full)
+
+ def test_column_parallel_forward_is_the_unsharded_columns(self):
+ _triton()
+ gen = torch.Generator().manual_seed(5)
+ x = torch.randn(40, 5376, generator=gen).bfloat16().cuda()
+ w = (torch.randn(7168, 5376, generator=gen) / 5376**0.5).bfloat16().cuda()
+ (full,) = par.column_parallel_linear(x, [w])
+ for tp in (2, 4, 8):
+ for r in range(tp):
+ (part,) = par.column_parallel_linear(x, [par.shard_rows(w, tp, r)])
+ assert torch.equal(part, par.shard_heads(full, tp, r, 1))
+
+
+@needs_gpu
+class TestPackedIO:
+ @pytest.mark.parametrize("length", [1, 5, 261])
+ def test_scatter_then_gather_is_identity_single_rank(self, length):
+ t = torch.randn(length, 7).cuda()
+ shard = par.SeqShard(length, 1, 0)
+ assert torch.equal(par.gather_packed(par.scatter_packed(t, shard), shard), t)
+
+
+# --- Multi-process grids --------------------------------------------------------------------
+
+
+def _reference(seq_len=0, backward=False):
+ from rl_engine.validation.models.h3_weights import load_h3_block_weights
+
+ return h3_tp.block_case(
+ load_h3_block_weights("cuda"), h3_tp.Grid(1, 1, 0, 0), seq_len=seq_len, backward=backward
+ )
+
+
+def _assemble(ranks, tp, cp, key="out"):
+ """Global rows from the ranks; replicated ranks must agree bitwise."""
+
+ pieces: dict[tuple[int, int], torch.Tensor] = {}
+ for r in ranks:
+ span = tuple(int(v) for v in r["rows"])
+ if span in pieces:
+ assert torch.equal(pieces[span], r[key]), f"ranks disagree on rows {span}"
+ pieces[span] = r[key]
+ return torch.cat([pieces[s] for s in sorted(pieces)], dim=1), sorted(pieces)
+
+
+GRIDS = [
+ pytest.param(2, 1, False, 0, id="tp2", marks=gpus(2)),
+ pytest.param(4, 1, False, 0, id="tp4", marks=gpus(4)),
+ pytest.param(1, 2, False, 0, id="cp2", marks=gpus(2)),
+ pytest.param(2, 2, False, 0, id="tp2_cp2", marks=gpus(4)),
+ pytest.param(2, 1, True, 0, id="tp2_sp", marks=gpus(2)),
+ pytest.param(4, 1, True, 0, id="tp4_sp", marks=gpus(4)),
+ pytest.param(2, 2, True, 0, id="tp2_cp2_sp", marks=gpus(4)),
+ pytest.param(1, 2, False, 261, id="cp2_len261", marks=gpus(2)),
+ pytest.param(2, 2, False, 261, id="tp2_cp2_len261", marks=gpus(4)),
+ pytest.param(4, 1, False, 263, id="tp4_len263", marks=gpus(4)),
+]
+
+
+@needs_gpu
+@pytest.mark.parametrize(("tp", "cp", "sp", "length"), GRIDS)
+def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length):
+ _triton()
+ ranks = run_world(
+ tp * cp, h3_tp.block_target, {"dummy": torch.zeros(1)},
+ tp=tp, cp=cp, sp=sp, seq_len=length, backward=True,
+ )
+ ref = _reference(length, backward=True)
+ out, spans = _assemble(ranks, tp, cp)
+ assert out.shape == ref["out"].shape, (spans, out.shape)
+ assert torch.equal(out.cpu(), ref["out"].cpu()), "forward differs from the single-GPU block"
+ dx, _ = _assemble(ranks, tp, cp, "dx")
+ assert torch.equal(dx.cpu(), ref["dx"].cpu()), "input gradient differs"
+ if cp == 1: # weight gradients: shards of the unsharded gradient
+ for name, axis in (("attn.to_out.0.weight", 1), ("ff.net.2.weight", 1),
+ ("attn.to_q.weight", 0), ("attn.to_k.weight", 0), ("attn.to_v.weight", 0),
+ ("ff.hidden", 0), ("ff.gate", 0)):
+ if sp:
+ continue # SP shards sequence rows; dW is reduced over rows below
+ full = ref["dw." + name].cpu()
+ shards = [r["dw." + name].cpu() for r in sorted(ranks, key=lambda r: r["rank"])]
+ assert torch.equal(torch.cat(shards, dim=axis), full), f"dW {name} differs"
+
+
+@needs_gpu
+@gpus(2)
+def test_ulysses_matches_unsharded_attention_for_ragged_lengths():
+ _triton()
+ for length in (1, 2, 7, 130):
+ ranks = run_world(2, _ulysses_target, {"dummy": torch.zeros(1)}, length=length)
+ gathered = torch.cat([r["out"] for r in ranks], dim=1)
+ assert torch.equal(gathered.cpu(), ranks[0]["ref"].cpu()), f"length {length}"
+
+
+def _ulysses_target(collective, dummy, *, length):
+ import torch.distributed as dist
+
+ from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp
+
+ del collective, dummy
+ rank = dist.get_rank()
+ gen = torch.Generator().manual_seed(7)
+ q, k, v = (torch.randn(1, length, 8, 128, generator=gen).bfloat16().cuda() for _ in range(3))
+ op = DeterministicAttentionOp()
+
+ def attend(q_, k_, v_):
+ t = lambda x: x.transpose(1, 2).contiguous() # noqa: E731
+ return op(t(q_), t(k_), t(v_), causal=False).transpose(1, 2).contiguous()
+
+ shard = par.SeqShard(length, 2, rank)
+ local = [par.scatter_packed(t, shard, dim=1) for t in (q, k, v)]
+ out = par.ulysses_attention(*local, attend, shard, dist.group.WORLD)
+ return {"out": out, "ref": attend(q, k, v)}
+
+
+# --- Hierarchical (cross-node design) collective, exercised inside one node ----------------
+
+
+def _hier_target(collective, dummy, *, local_size, dtype):
+ import torch.distributed as dist
+
+ from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective
+
+ del dummy
+ h = HierarchicalCollective(local_size, device=torch.cuda.current_device(), max_size_bytes=64 << 20)
+ gen = torch.Generator().manual_seed(100 + dist.get_rank())
+ x = (torch.randn(257, 129, generator=gen) * 10 ** torch.randint(-3, 4, (257, 1), generator=gen)).to(
+ dtype
+ ).cuda()
+ flat = collective.all_reduce(x.clone())
+ mine = h.all_reduce(x)
+ rs = h.reduce_scatter(x[: 256].contiguous())
+ ag = h.all_gather(x[:2].contiguous())
+ return {"flat": flat, "hier": mine, "rs": rs, "ag": ag, "x": x}
+
+
+@needs_gpu
+@gpus(4)
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize("local_size", [1, 2])
+def test_hierarchical_equals_single_node_fixed_tree(local_size, dtype):
+ ranks = run_world(4, _hier_target, {"dummy": torch.zeros(1)}, local_size=local_size, dtype=dtype)
+ for r in ranks:
+ assert torch.equal(r["flat"], r["hier"]), f"rank {r['rank']}: hierarchical != fixed tree"
+ assert torch.equal(ranks[0]["rs"], ranks[0]["hier"][:64])
+ assert torch.equal(ranks[3]["rs"], ranks[3]["hier"][192:256])
+ assert torch.equal(ranks[1]["ag"], torch.cat([r["x"][:2] for r in ranks]))
diff --git a/tools/validation/models/h3_parallel_run.py b/tools/validation/models/h3_parallel_run.py
new file mode 100644
index 000000000..4082920b2
--- /dev/null
+++ b/tools/validation/models/h3_parallel_run.py
@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""Run the H3 parallel block on a TP x CP grid under ``srun`` (one task per GPU, any node count).
+
+ srun -N2 --ntasks-per-node=4 python tools/validation/models/h3_parallel_run.py \\
+ --tp 8 --cp 1 --out evidence/ws2_tp8
+
+Every rank writes its rows (and gradients) to ``--out/rank{r}.pt``; rank 0 then recomputes the
+single-GPU canonical block and compares the assembled forward, input gradient and weight-gradient
+shards bitwise, and writes ``report.json``. ``--local-size`` is the number of GPUs per node for a
+TP group that spans nodes (selects ``hierarchical_ipc_nccl_tree``); a TP group inside one node
+uses the CUDA-IPC fixed tree.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import socket
+import time
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+
+from rl_engine.validation.models import h3_tp
+from rl_engine.validation.models.h3_weights import load_h3_block_weights
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--tp", type=int, required=True)
+ parser.add_argument("--cp", type=int, default=1)
+ parser.add_argument("--sp", action="store_true")
+ parser.add_argument("--seq-len", type=int, default=0)
+ parser.add_argument("--layout", default="tiny")
+ parser.add_argument("--local-size", type=int, default=4, help="GPUs per node")
+ parser.add_argument("--iters", type=int, default=5)
+ parser.add_argument("--out", type=Path, required=True)
+ args = parser.parse_args()
+
+ rank = int(os.environ["SLURM_PROCID"])
+ world = int(os.environ["SLURM_NTASKS"])
+ local = int(os.environ["SLURM_LOCALID"])
+ torch.cuda.set_device(local)
+ dist.init_process_group(
+ "nccl", init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ.get('MASTER_PORT', '29511')}",
+ rank=rank, world_size=world, device_id=torch.device("cuda", local),
+ )
+ if world != args.tp * args.cp:
+ raise SystemExit(f"world {world} != tp {args.tp} x cp {args.cp}")
+ probe = torch.ones(1, device="cuda") * rank
+ dist.all_reduce(probe) # NCCL connectivity across nodes
+ spans_nodes = args.tp > args.local_size
+ grid = h3_tp.make_grid(args.tp, args.cp, local_size=args.local_size if spans_nodes else None)
+ weights = load_h3_block_weights("cuda")
+ kwargs = dict(sp=args.sp, seq_len=args.seq_len, layout=args.layout)
+ result = h3_tp.block_case(weights, grid, backward=True, **kwargs)
+ # timing of the forward alone
+ torch.cuda.synchronize(); dist.barrier()
+ begin = time.perf_counter()
+ for _ in range(args.iters):
+ h3_tp.block_case(weights, grid, backward=False, **kwargs)
+ torch.cuda.synchronize()
+ seconds = (time.perf_counter() - begin) / args.iters
+ args.out.mkdir(parents=True, exist_ok=True)
+ torch.save({k: v.cpu() for k, v in result.items()}, args.out / f"rank{rank}.pt")
+ dist.barrier()
+ timings = [None] * world
+ dist.all_gather_object(timings, {"rank": rank, "host": socket.gethostname(), "forward_s": seconds})
+ if rank == 0:
+ ref = h3_tp.block_case(
+ weights, h3_tp.Grid(1, 1, 0, 0), backward=True, **{**kwargs, "sp": False}
+ )
+ ranks = [torch.load(args.out / f"rank{r}.pt") for r in range(world)]
+
+ def assemble(key):
+ pieces = {}
+ for r in ranks:
+ span = tuple(int(v) for v in r["rows"])
+ if span in pieces and not torch.equal(pieces[span], r[key]):
+ return None
+ pieces[span] = r[key]
+ return torch.cat([pieces[s] for s in sorted(pieces)], dim=1)
+
+ out, dx = assemble("out"), assemble("dx")
+ checks = {
+ "forward_bitwise": out is not None and torch.equal(out, ref["out"].cpu()),
+ "input_grad_bitwise": dx is not None and torch.equal(dx, ref["dx"].cpu()),
+ }
+ if args.cp == 1 and not args.sp:
+ for name, axis in (("attn.to_out.0.weight", 1), ("ff.net.2.weight", 1),
+ ("attn.to_q.weight", 0), ("attn.to_k.weight", 0),
+ ("attn.to_v.weight", 0), ("ff.hidden", 0), ("ff.gate", 0)):
+ shards = [r["dw." + name] for r in ranks]
+ checks["dw." + name] = torch.equal(torch.cat(shards, dim=axis), ref["dw." + name].cpu())
+ report = dict(
+ tp=args.tp, cp=args.cp, sp=args.sp, world=world, hosts=sorted({t["host"] for t in timings}),
+ collective="hierarchical_ipc_nccl_tree" if spans_nodes else "cuda_ipc_fixed_tree",
+ seq_len=int(ref["out"].shape[1]) if not args.seq_len else args.seq_len,
+ checks=checks, all_bitwise=all(checks.values()),
+ forward_seconds_per_rank=timings,
+ torch=torch.__version__, device=torch.cuda.get_device_name(),
+ )
+ (args.out / "report.json").write_text(json.dumps(report, indent=1) + "\n")
+ print(json.dumps({k: report[k] for k in ("tp", "cp", "sp", "hosts", "collective", "checks", "all_bitwise")}, indent=1))
+ dist.barrier()
+ if rank == 0:
+ time.sleep(3) # keep the rendezvous store up until every rank has left
+ dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+ main()
From 17fe0706433f474677e1782cd888a5222d46b2d7 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 13:12:18 +0200
Subject: [PATCH 02/14] feat(h3): WS2 Hopper evidence chart and lint
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 33cbcd6ecd01c38e180b0ccb357125e0c4069823)
---
docs/operators/h3-parallel.md | 6 ++
.../model_specific/minimax_h3/parallel.py | 47 +++++++++----
.../distributed/algorithms/hierarchical.py | 4 +-
rl_engine/validation/models/h3_tp.py | 66 ++++++++++++++-----
tests/models/minimax_h3/test_h3_parallel.py | 53 ++++++++++-----
tools/validation/models/h3_parallel_run.py | 52 +++++++++++----
tools/validation/models/h3_ws2_charts.py | 38 +++++++++++
7 files changed, 207 insertions(+), 59 deletions(-)
create mode 100644 tools/validation/models/h3_ws2_charts.py
diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md
index b92c68342..dff54666b 100644
--- a/docs/operators/h3-parallel.md
+++ b/docs/operators/h3-parallel.md
@@ -85,3 +85,9 @@ Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4x
- dW under CP is exact through `seq_sharded_linear` for non-TP linears; the TP x CP block compares
forward and `dX` (CP shards the token axis of the TP weight gradients).
- `HierarchicalCollective` assumes ranks are packed node-major and node counts of 1/2/4/8.
+
+## Hopper evidence chart
+
+These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description).
+Per-layout results of the 2 nodes x 4 H100 runs (`reports/experiments/h3-ws2-parallel-h100/*.json`): every layout is bitwise equal to the single-GPU reference
+(forward and input gradient); mean per-rank forward time of block 0 shown below.
diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
index eabef82fd..f61a7dbcd 100644
--- a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
+++ b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
@@ -51,7 +51,9 @@
def check_tp(tp: int, rank: int) -> None:
if tp not in TP_SIZES or not 0 <= rank < tp:
- raise ValueError(f"tp must be one of {TP_SIZES} and rank in [0, tp); got tp={tp} rank={rank}")
+ raise ValueError(
+ f"tp must be one of {TP_SIZES} and rank in [0, tp); got tp={tp} rank={rank}"
+ )
def shard_rows(weight: torch.Tensor, tp: int, rank: int) -> torch.Tensor:
@@ -104,8 +106,12 @@ def _units_partial(x: torch.Tensor, w: torch.Tensor, tp: int) -> torch.Tensor:
local = UNITS // tp
unit = x.shape[1] // local
parts = [
- g.gemm(x[:, i * unit : (i + 1) * unit], w[:, i * unit : (i + 1) * unit], b_t=True,
- out_dtype=torch.float32)
+ g.gemm(
+ x[:, i * unit : (i + 1) * unit],
+ w[:, i * unit : (i + 1) * unit],
+ b_t=True,
+ out_dtype=torch.float32,
+ )
for i in range(local)
]
return tree_add(parts)
@@ -118,7 +124,9 @@ def forward(ctx, x, w, collective, tp, scatter):
ctx.scatter, ctx.collective = scatter, collective
partial = _units_partial(x, w, tp)
if tp > 1:
- partial = collective.reduce_scatter(partial) if scatter else collective.all_reduce(partial)
+ partial = (
+ collective.reduce_scatter(partial) if scatter else collective.all_reduce(partial)
+ )
return partial.to(x.dtype)
@staticmethod
@@ -168,21 +176,28 @@ def backward(ctx, *grads):
acc = None
for grad, w in zip(grads, weights):
width = grad.shape[1] // local
- piece = g.gemm(grad[:, j * width : (j + 1) * width], w[j * width : (j + 1) * width],
- out_dtype=torch.float32)
+ piece = g.gemm(
+ grad[:, j * width : (j + 1) * width],
+ w[j * width : (j + 1) * width],
+ out_dtype=torch.float32,
+ )
acc = piece if acc is None else acc + piece
units.append(acc)
partial = tree_add(units)
if ctx.tp > 1:
- partial = (ctx.collective.reduce_scatter(partial) if ctx.gather
- else ctx.collective.all_reduce(partial))
+ partial = (
+ ctx.collective.reduce_scatter(partial)
+ if ctx.gather
+ else ctx.collective.all_reduce(partial)
+ )
dx = partial.to(x.dtype)
dws = tuple(g.gemm(grad, full_x, a_t=True) for grad in grads)
return (dx, None, None, None, *dws)
-def column_parallel_linear(x, weights: Sequence[torch.Tensor], collective=None, tp: int = 1,
- gather: bool = False):
+def column_parallel_linear(
+ x, weights: Sequence[torch.Tensor], collective=None, tp: int = 1, gather: bool = False
+):
"""``[x @ w.T for w in weights]`` with each ``w`` a row shard ``(N/tp, K)``.
The forward is the unsharded projection's columns (no collective). With ``gather``, ``x`` is
@@ -234,7 +249,9 @@ def coshard(shard: SeqShard, *tensors: torch.Tensor, dim: int = 0) -> list[torch
out = []
for t in tensors:
if t.shape[dim] != shard.length:
- raise ValueError(f"tensor has {t.shape[dim]} rows along dim {dim}, expected {shard.length}")
+ raise ValueError(
+ f"tensor has {t.shape[dim]} rows along dim {dim}, expected {shard.length}"
+ )
out.append(t.narrow(dim, shard.lo, shard.hi - shard.lo).contiguous())
return out
@@ -331,6 +348,10 @@ def backward(ctx, grad):
def seq_sharded_linear(x, w, shard: SeqShard, group=None):
- """``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM's (token sweep, ascending)."""
+ """``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM (token sweep)."""
- return _SeqLinear.apply(x.contiguous(), w, shard, group) if shard.cp > 1 else g.linear_forward(x, w, None)
+ return (
+ _SeqLinear.apply(x.contiguous(), w, shard, group)
+ if shard.cp > 1
+ else g.linear_forward(x, w, None)
+ )
diff --git a/rl_engine/distributed/algorithms/hierarchical.py b/rl_engine/distributed/algorithms/hierarchical.py
index bfc29f04d..253b34399 100644
--- a/rl_engine/distributed/algorithms/hierarchical.py
+++ b/rl_engine/distributed/algorithms/hierarchical.py
@@ -85,7 +85,9 @@ def all_reduce(self, input: torch.Tensor) -> torch.Tensor:
if input.dtype not in (torch.float32, torch.bfloat16, torch.float16):
raise TypeError(f"unsupported dtype {input.dtype}")
- wide = input.contiguous() # the input dtype is the accumulation dtype, as in the node-local tree
+ wide = (
+ input.contiguous()
+ ) # the input dtype is the accumulation dtype, as in the node-local tree
partial = self._local.all_reduce(wide) if self.local_size > 1 else wide.clone()
if self.num_groups > 1:
gathered = [torch.empty_like(partial) for _ in range(self.num_groups)]
diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py
index 697010300..1cecab180 100644
--- a/rl_engine/validation/models/h3_tp.py
+++ b/rl_engine/validation/models/h3_tp.py
@@ -58,11 +58,13 @@ def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid:
if tp > 1:
mine = tp_groups[cp_rank]
if local_size is not None and local_size < tp:
- collective = HierarchicalCollective(local_size, mine, torch.cuda.current_device(),
- max_size_bytes=256 << 20)
+ collective = HierarchicalCollective(
+ local_size, mine, torch.cuda.current_device(), max_size_bytes=256 << 20
+ )
else:
- collective = DeterministicCollective(mine, torch.cuda.current_device(),
- max_size_bytes=256 << 20)
+ collective = DeterministicCollective(
+ mine, torch.cuda.current_device(), max_size_bytes=256 << 20
+ )
return Grid(tp, cp, tp_rank, cp_rank, collective, cp_groups[tp_rank] if cp > 1 else None)
@@ -121,6 +123,7 @@ def leaf(name, tensor):
tensor.requires_grad_(True)
leaves[name] = tensor
return tensor
+
table = ops.projection(temb, w("adaln_proj.linear.weight"), w("adaln_proj.linear.bias"))
def proj_in(x, name_weights, gather):
@@ -128,7 +131,14 @@ def proj_in(x, name_weights, gather):
return par.column_parallel_linear(flat, name_weights, coll, tp, gather=gather)
x = ops.norm.forward_modulated(h_res, w("norm1.weight"), table[0], table[1], idx_res)
- q, k, v = proj_in(x, [leaf(f"attn.to_{n}.weight", par.shard_rows(w(f"attn.to_{n}.weight"), tp, grid.tp_rank)) for n in "qkv"], sp)
+ q, k, v = proj_in(
+ x,
+ [
+ leaf(f"attn.to_{n}.weight", par.shard_rows(w(f"attn.to_{n}.weight"), tp, grid.tp_rank))
+ for n in "qkv"
+ ],
+ sp,
+ )
local_heads = par.HEADS // tp
def heads(t):
@@ -142,7 +152,10 @@ def heads(t):
attn = par.ulysses_attention(q, k, heads(v), _attention_fn(ops), shard, grid.cp_group)
o = par.row_parallel_linear(
attn.reshape(rows, local_heads * HEAD_DIM),
- leaf("attn.to_out.0.weight", par.shard_cols(w("attn.to_out.0.weight"), tp, grid.tp_rank)), coll, tp, scatter=sp,
+ leaf("attn.to_out.0.weight", par.shard_cols(w("attn.to_out.0.weight"), tp, grid.tp_rank)),
+ coll,
+ tp,
+ scatter=sp,
)
h1 = ops.gate(h_res, o.view(1, -1, HIDDEN), table[2], idx_res)
x2 = ops.norm.forward_modulated(h1, w("norm2.weight"), table[3], table[4], idx_res)
@@ -150,20 +163,35 @@ def heads(t):
half = fused.shape[0] // 2
hid, gate = proj_in(
x2,
- [leaf("ff.hidden", par.shard_rows(fused[:half], tp, grid.tp_rank)),
- leaf("ff.gate", par.shard_rows(fused[half:], tp, grid.tp_rank))],
+ [
+ leaf("ff.hidden", par.shard_rows(fused[:half], tp, grid.tp_rank)),
+ leaf("ff.gate", par.shard_rows(fused[half:], tp, grid.tp_rank)),
+ ],
sp,
)
act = ops.swiglu(gate.contiguous(), hid.contiguous())
down = par.row_parallel_linear(
- act, leaf("ff.net.2.weight", par.shard_cols(w("ff.net.2.weight"), tp, grid.tp_rank)), coll, tp, scatter=sp
+ act,
+ leaf("ff.net.2.weight", par.shard_cols(w("ff.net.2.weight"), tp, grid.tp_rank)),
+ coll,
+ tp,
+ scatter=sp,
)
return ops.gate(h1, down.view(1, -1, HIDDEN), table[5], idx_res)
def block_target(
- collective, dummy, *, tp: int, cp: int, sp: bool = False, seq_len: int = 0,
- backward: bool = False, local_size: int = -1, layout: str = "tiny", seed: int = 0,
+ collective,
+ dummy,
+ *,
+ tp: int,
+ cp: int,
+ sp: bool = False,
+ seq_len: int = 0,
+ backward: bool = False,
+ local_size: int = -1,
+ layout: str = "tiny",
+ seed: int = 0,
):
"""``run_world`` target: this rank's rows of the parallel block (and gradients) for one case."""
@@ -171,8 +199,15 @@ def block_target(
from rl_engine.validation.models.h3_weights import load_h3_block_weights
grid = make_grid(tp, cp, local_size=local_size if local_size > 0 else None)
- return block_case(load_h3_block_weights("cuda"), grid, sp=sp, seq_len=seq_len,
- backward=backward, layout=layout, seed=seed)
+ return block_case(
+ load_h3_block_weights("cuda"),
+ grid,
+ sp=sp,
+ seq_len=seq_len,
+ backward=backward,
+ layout=layout,
+ seed=seed,
+ )
def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layout="tiny", seed=0):
@@ -186,8 +221,9 @@ def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layo
ops = hb.CandidateOps()
leaves: dict[str, torch.Tensor] | None = {} if backward else None
with torch.set_grad_enabled(backward):
- out = run_tp_block(weights, hidden, ctx["temb"], index, rope, ops, grid, sp=sp,
- seq_len=n, leaves=leaves)
+ out = run_tp_block(
+ weights, hidden, ctx["temb"], index, rope, ops, grid, sp=sp, seq_len=n, leaves=leaves
+ )
shard = grid.seq_shard(n)
lo, hi = shard.lo, shard.hi
if sp:
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
index 19b5c0516..1d1e9d65e 100644
--- a/tests/models/minimax_h3/test_h3_parallel.py
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -16,7 +16,6 @@
from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
from rl_engine.distributed.algorithms.hierarchical import tree_add
-from rl_engine.validation.models import h3_block as hb
from rl_engine.validation.models import h3_tp
from rl_engine.validation.models.h3_ws2 import run_world
@@ -57,7 +56,9 @@ def test_gate_up_shards_keep_both_halves(self):
@pytest.mark.parametrize("length", [1, 2, 3, 7, 261, 264, 925])
@pytest.mark.parametrize("cp", [1, 2, 3, 4])
def test_seq_shards_tile_any_length(self, length, cp):
- spans = [(par.SeqShard(length, cp, r).lo, par.SeqShard(length, cp, r).hi) for r in range(cp)]
+ spans = [
+ (par.SeqShard(length, cp, r).lo, par.SeqShard(length, cp, r).hi) for r in range(cp)
+ ]
assert spans[0][0] == 0 and spans[-1][1] == length
assert all(a[1] == b[0] for a, b in zip(spans, spans[1:]))
sizes = [hi - lo for lo, hi in spans]
@@ -93,9 +94,7 @@ def test_row_parallel_partials_fold_to_tp1(self, k, tp):
for r in range(tp)
]
assert torch.equal(tree_add(parts), full)
- assert torch.equal(
- par.row_parallel_linear(x, w), full.bfloat16()
- )
+ assert torch.equal(par.row_parallel_linear(x, w), full.bfloat16())
def test_row_parallel_close_to_fp64(self):
_triton()
@@ -122,7 +121,6 @@ def test_column_parallel_backward_folds_to_tp1(self, tp):
_triton()
gen = torch.Generator().manual_seed(4)
k, n = 5376, 7168
- x = torch.randn(40, k, generator=gen).bfloat16().cuda()
ws = [(torch.randn(n, k, generator=gen) / k**0.5).bfloat16().cuda() for _ in range(3)]
grads = [torch.randn(40, n, generator=gen).bfloat16().cuda() for _ in range(3)]
@@ -133,7 +131,8 @@ def unit_partials(weights, gs, local):
for grad, w in zip(gs, weights):
width = grad.shape[1] // local
piece = par.g.gemm(
- grad[:, j * width : (j + 1) * width], w[j * width : (j + 1) * width],
+ grad[:, j * width : (j + 1) * width],
+ w[j * width : (j + 1) * width],
out_dtype=torch.float32,
)
acc = piece if acc is None else acc + piece
@@ -214,8 +213,14 @@ def _assemble(ranks, tp, cp, key="out"):
def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length):
_triton()
ranks = run_world(
- tp * cp, h3_tp.block_target, {"dummy": torch.zeros(1)},
- tp=tp, cp=cp, sp=sp, seq_len=length, backward=True,
+ tp * cp,
+ h3_tp.block_target,
+ {"dummy": torch.zeros(1)},
+ tp=tp,
+ cp=cp,
+ sp=sp,
+ seq_len=length,
+ backward=True,
)
ref = _reference(length, backward=True)
out, spans = _assemble(ranks, tp, cp)
@@ -224,9 +229,15 @@ def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length):
dx, _ = _assemble(ranks, tp, cp, "dx")
assert torch.equal(dx.cpu(), ref["dx"].cpu()), "input gradient differs"
if cp == 1: # weight gradients: shards of the unsharded gradient
- for name, axis in (("attn.to_out.0.weight", 1), ("ff.net.2.weight", 1),
- ("attn.to_q.weight", 0), ("attn.to_k.weight", 0), ("attn.to_v.weight", 0),
- ("ff.hidden", 0), ("ff.gate", 0)):
+ for name, axis in (
+ ("attn.to_out.0.weight", 1),
+ ("ff.net.2.weight", 1),
+ ("attn.to_q.weight", 0),
+ ("attn.to_k.weight", 0),
+ ("attn.to_v.weight", 0),
+ ("ff.hidden", 0),
+ ("ff.gate", 0),
+ ):
if sp:
continue # SP shards sequence rows; dW is reduced over rows below
full = ref["dw." + name].cpu()
@@ -274,14 +285,18 @@ def _hier_target(collective, dummy, *, local_size, dtype):
from rl_engine.distributed.algorithms.hierarchical import HierarchicalCollective
del dummy
- h = HierarchicalCollective(local_size, device=torch.cuda.current_device(), max_size_bytes=64 << 20)
+ h = HierarchicalCollective(
+ local_size, device=torch.cuda.current_device(), max_size_bytes=64 << 20
+ )
gen = torch.Generator().manual_seed(100 + dist.get_rank())
- x = (torch.randn(257, 129, generator=gen) * 10 ** torch.randint(-3, 4, (257, 1), generator=gen)).to(
- dtype
- ).cuda()
+ x = (
+ (torch.randn(257, 129, generator=gen) * 10 ** torch.randint(-3, 4, (257, 1), generator=gen))
+ .to(dtype)
+ .cuda()
+ )
flat = collective.all_reduce(x.clone())
mine = h.all_reduce(x)
- rs = h.reduce_scatter(x[: 256].contiguous())
+ rs = h.reduce_scatter(x[:256].contiguous())
ag = h.all_gather(x[:2].contiguous())
return {"flat": flat, "hier": mine, "rs": rs, "ag": ag, "x": x}
@@ -291,7 +306,9 @@ def _hier_target(collective, dummy, *, local_size, dtype):
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
@pytest.mark.parametrize("local_size", [1, 2])
def test_hierarchical_equals_single_node_fixed_tree(local_size, dtype):
- ranks = run_world(4, _hier_target, {"dummy": torch.zeros(1)}, local_size=local_size, dtype=dtype)
+ ranks = run_world(
+ 4, _hier_target, {"dummy": torch.zeros(1)}, local_size=local_size, dtype=dtype
+ )
for r in ranks:
assert torch.equal(r["flat"], r["hier"]), f"rank {r['rank']}: hierarchical != fixed tree"
assert torch.equal(ranks[0]["rs"], ranks[0]["hier"][:64])
diff --git a/tools/validation/models/h3_parallel_run.py b/tools/validation/models/h3_parallel_run.py
index 4082920b2..fc5cf6dbc 100644
--- a/tools/validation/models/h3_parallel_run.py
+++ b/tools/validation/models/h3_parallel_run.py
@@ -45,8 +45,11 @@ def main() -> None:
local = int(os.environ["SLURM_LOCALID"])
torch.cuda.set_device(local)
dist.init_process_group(
- "nccl", init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ.get('MASTER_PORT', '29511')}",
- rank=rank, world_size=world, device_id=torch.device("cuda", local),
+ "nccl",
+ init_method=f"tcp://{os.environ['MASTER_ADDR']}:{os.environ.get('MASTER_PORT', '29511')}",
+ rank=rank,
+ world_size=world,
+ device_id=torch.device("cuda", local),
)
if world != args.tp * args.cp:
raise SystemExit(f"world {world} != tp {args.tp} x cp {args.cp}")
@@ -58,7 +61,8 @@ def main() -> None:
kwargs = dict(sp=args.sp, seq_len=args.seq_len, layout=args.layout)
result = h3_tp.block_case(weights, grid, backward=True, **kwargs)
# timing of the forward alone
- torch.cuda.synchronize(); dist.barrier()
+ torch.cuda.synchronize()
+ dist.barrier()
begin = time.perf_counter()
for _ in range(args.iters):
h3_tp.block_case(weights, grid, backward=False, **kwargs)
@@ -68,7 +72,9 @@ def main() -> None:
torch.save({k: v.cpu() for k, v in result.items()}, args.out / f"rank{rank}.pt")
dist.barrier()
timings = [None] * world
- dist.all_gather_object(timings, {"rank": rank, "host": socket.gethostname(), "forward_s": seconds})
+ dist.all_gather_object(
+ timings, {"rank": rank, "host": socket.gethostname(), "forward_s": seconds}
+ )
if rank == 0:
ref = h3_tp.block_case(
weights, h3_tp.Grid(1, 1, 0, 0), backward=True, **{**kwargs, "sp": False}
@@ -90,21 +96,43 @@ def assemble(key):
"input_grad_bitwise": dx is not None and torch.equal(dx, ref["dx"].cpu()),
}
if args.cp == 1 and not args.sp:
- for name, axis in (("attn.to_out.0.weight", 1), ("ff.net.2.weight", 1),
- ("attn.to_q.weight", 0), ("attn.to_k.weight", 0),
- ("attn.to_v.weight", 0), ("ff.hidden", 0), ("ff.gate", 0)):
+ for name, axis in (
+ ("attn.to_out.0.weight", 1),
+ ("ff.net.2.weight", 1),
+ ("attn.to_q.weight", 0),
+ ("attn.to_k.weight", 0),
+ ("attn.to_v.weight", 0),
+ ("ff.hidden", 0),
+ ("ff.gate", 0),
+ ):
shards = [r["dw." + name] for r in ranks]
- checks["dw." + name] = torch.equal(torch.cat(shards, dim=axis), ref["dw." + name].cpu())
+ checks["dw." + name] = torch.equal(
+ torch.cat(shards, dim=axis), ref["dw." + name].cpu()
+ )
report = dict(
- tp=args.tp, cp=args.cp, sp=args.sp, world=world, hosts=sorted({t["host"] for t in timings}),
+ tp=args.tp,
+ cp=args.cp,
+ sp=args.sp,
+ world=world,
+ hosts=sorted({t["host"] for t in timings}),
collective="hierarchical_ipc_nccl_tree" if spans_nodes else "cuda_ipc_fixed_tree",
seq_len=int(ref["out"].shape[1]) if not args.seq_len else args.seq_len,
- checks=checks, all_bitwise=all(checks.values()),
+ checks=checks,
+ all_bitwise=all(checks.values()),
forward_seconds_per_rank=timings,
- torch=torch.__version__, device=torch.cuda.get_device_name(),
+ torch=torch.__version__,
+ device=torch.cuda.get_device_name(),
)
(args.out / "report.json").write_text(json.dumps(report, indent=1) + "\n")
- print(json.dumps({k: report[k] for k in ("tp", "cp", "sp", "hosts", "collective", "checks", "all_bitwise")}, indent=1))
+ print(
+ json.dumps(
+ {
+ k: report[k]
+ for k in ("tp", "cp", "sp", "hosts", "collective", "checks", "all_bitwise")
+ },
+ indent=1,
+ )
+ )
dist.barrier()
if rank == 0:
time.sleep(3) # keep the rendezvous store up until every rank has left
diff --git a/tools/validation/models/h3_ws2_charts.py b/tools/validation/models/h3_ws2_charts.py
new file mode 100644
index 000000000..a3e959843
--- /dev/null
+++ b/tools/validation/models/h3_ws2_charts.py
@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+"""Chart of the WS2 evidence: bitwise agreement and per-rank forward time per parallel layout."""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt # noqa: E402
+
+
+def main(evidence: Path, out: Path) -> None:
+ reports = [json.loads(p.read_text()) for p in sorted(evidence.glob("h3_ws2_*.json"))]
+ labels = []
+ times = []
+ colors = []
+ for r in reports:
+ labels.append(f"TP{r['tp']}xCP{r['cp']}{'+SP' if r.get('sp') else ''}")
+ ranks = r["forward_seconds_per_rank"]
+ times.append(sum(x["forward_s"] for x in ranks) / len(ranks))
+ colors.append("tab:green" if r["all_bitwise"] else "tab:red")
+ fig, ax = plt.subplots(figsize=(8, 4))
+ ax.bar(labels, times, color=colors)
+ ax.set(ylabel="mean forward s per rank (H100)", title="WS2 layouts: green = bitwise == single GPU")
+ for i, r in enumerate(reports):
+ ax.text(i, times[i], f"{r['world']} GPU\n{len(set(r['hosts']))} node", ha="center", va="bottom")
+ out.parent.mkdir(parents=True, exist_ok=True)
+ fig.tight_layout()
+ fig.savefig(out, dpi=120)
+
+
+if __name__ == "__main__":
+ main(Path(sys.argv[1]), Path(sys.argv[2]))
From eee3f97e4c350ea3e418c7f0e6d45853a41d4ba8 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 14:06:36 +0200
Subject: [PATCH 03/14] docs(h3): drop the reference to the moved WS2 figure
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 5f4e34c6ae3b79b643a85aded69e41fce447a204)
---
docs/operators/h3-parallel.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md
index dff54666b..3ffcacd42 100644
--- a/docs/operators/h3-parallel.md
+++ b/docs/operators/h3-parallel.md
@@ -90,4 +90,4 @@ Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4x
These outputs and figures are not part of this diff; they are on the fork's `pr_assets` branch (links in the PR description).
Per-layout results of the 2 nodes x 4 H100 runs (`reports/experiments/h3-ws2-parallel-h100/*.json`): every layout is bitwise equal to the single-GPU reference
-(forward and input gradient); mean per-rank forward time of block 0 shown below.
+(forward and input gradient); mean per-rank forward time of block 0 is in the figure.
From f0778881cadc7757a0a9cc4e2eddba78a9d6b3ca Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 16:33:14 +0200
Subject: [PATCH 04/14] style(h3): black for the WS2 chart tool
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit d086b5ce7a50e03e022334ae35b3a7b48e9c1171)
---
tools/validation/models/h3_ws2_charts.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/tools/validation/models/h3_ws2_charts.py b/tools/validation/models/h3_ws2_charts.py
index a3e959843..5bc27964d 100644
--- a/tools/validation/models/h3_ws2_charts.py
+++ b/tools/validation/models/h3_ws2_charts.py
@@ -26,9 +26,13 @@ def main(evidence: Path, out: Path) -> None:
colors.append("tab:green" if r["all_bitwise"] else "tab:red")
fig, ax = plt.subplots(figsize=(8, 4))
ax.bar(labels, times, color=colors)
- ax.set(ylabel="mean forward s per rank (H100)", title="WS2 layouts: green = bitwise == single GPU")
+ ax.set(
+ ylabel="mean forward s per rank (H100)", title="WS2 layouts: green = bitwise == single GPU"
+ )
for i, r in enumerate(reports):
- ax.text(i, times[i], f"{r['world']} GPU\n{len(set(r['hosts']))} node", ha="center", va="bottom")
+ ax.text(
+ i, times[i], f"{r['world']} GPU\n{len(set(r['hosts']))} node", ha="center", va="bottom"
+ )
out.parent.mkdir(parents=True, exist_ok=True)
fig.tight_layout()
fig.savefig(out, dpi=120)
From 72a9068d5dd7d6d5965067e47bafeddb088545bf Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 22:54:44 +0200
Subject: [PATCH 05/14] fix(h3): satisfy mypy for the TP collective and the
optional weight-shard leaves
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 5ae9f5ebbea007235325e2bce3e659ed4b561ad6)
---
rl_engine/validation/models/h3_tp.py | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py
index 1cecab180..2234b2400 100644
--- a/rl_engine/validation/models/h3_tp.py
+++ b/rl_engine/validation/models/h3_tp.py
@@ -54,7 +54,7 @@ def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid:
cp_rank, tp_rank = divmod(rank, tp)
tp_groups = [dist.new_group(list(range(c * tp, (c + 1) * tp))) for c in range(cp)]
cp_groups = [dist.new_group(list(range(t, world, tp))) for t in range(tp)]
- collective = None
+ collective: DeterministicCollective | HierarchicalCollective | None = None
if tp > 1:
mine = tp_groups[cp_rank]
if local_size is not None and local_size < tp:
@@ -235,6 +235,7 @@ def block_case(weights, grid: Grid, *, sp=False, seq_len=0, backward=False, layo
full_grad = torch.randn(1, n, HIDDEN, generator=gen).to(out.dtype).cuda()
out.backward(full_grad[:, lo:hi])
result["dx"] = hidden.grad[:, lo:hi].detach()
+ assert leaves is not None
for name, tensor in leaves.items():
result["dw." + name] = tensor.grad.detach()
return result
From 17e29e3ad4a246c5f4b7757227312a1cb2069d2b Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 18:52:12 +0200
Subject: [PATCH 06/14] feat(h3): checkpoint mapper and the vLLM-Omni rollout
adapter
RFC #420 rows h3_checkpoint_mapper (Diffusers <-> vLLM-Omni names and layouts, all 638 tensors,
bitwise round trip, fail-closed on unmapped names) and vllm_runtime_adapter_h3 (strict
fixed-order kernels behind qkv_proj / out_proj / fc1 / fc2, with read-back).
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit 7b08b63094784f9c7313cc229c1c11c9f9aa53f6)
---
docs/operators/h3-checkpoint-mapper-vllm.md | 32 ++++
.../engines/rollout/vllm/h3_omni.py | 95 ++++++++++
.../models/minimax_h3/checkpoint_mapper.py | 176 ++++++++++++++++++
.../test_h3_checkpoint_mapper_vllm.py | 132 +++++++++++++
4 files changed, 435 insertions(+)
create mode 100644 docs/operators/h3-checkpoint-mapper-vllm.md
create mode 100644 rl_engine/integrations/engines/rollout/vllm/h3_omni.py
create mode 100644 rl_engine/models/minimax_h3/checkpoint_mapper.py
create mode 100644 tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
diff --git a/docs/operators/h3-checkpoint-mapper-vllm.md b/docs/operators/h3-checkpoint-mapper-vllm.md
new file mode 100644
index 000000000..050653650
--- /dev/null
+++ b/docs/operators/h3-checkpoint-mapper-vllm.md
@@ -0,0 +1,32 @@
+# MiniMax-H3 Checkpoint Mapper and vLLM-Omni Rollout Adapter
+
+## Summary
+
+RFC #420 integration rows, validated on H100 against the installed vLLM-Omni name tables:
+
+| Row | Where | What it fixes |
+| --- | --- | --- |
+| `h3_checkpoint_mapper` | `models/minimax_h3/checkpoint_mapper.py` | Diffusers <-> vLLM-Omni native names and layouts |
+| `vllm_runtime_adapter_h3` | `integrations/engines/rollout/vllm/h3_omni.py` | strict fixed-order kernels behind vLLM-Omni's `qkv_proj` / `out_proj` / `fc1` / `fc2` |
+
+## Checkpoint mapper
+
+`transformer_blocks.N` -> `blocks.N`, `token_refiner.refiner_blocks.N` -> `token_refiner.blocks.N`;
+`to_q/k/v` -> one `attn.qkv_proj` (`[Q;K;V]` rows); `to_out.0` -> `out_proj`; `ff.net.0.proj`
+(value-first) -> `mlp.fc1` (gate-first: halves swapped); `ff.net.2` -> `mlp.fc2`; model-level names
+from vLLM-Omni's `fasth3` table. All 638 checkpoint tensors map; the round trip is bitwise;
+unmapped names and incomplete q/k/v sets fail closed. A test compares the tables with
+`vllm_omni.diffusion.models.minimax_h3.fasth3` when importable (passes with vLLM-Omni installed).
+
+## vLLM-Omni adapter
+
+`install_h3_rollout(model, collective=None, tp=1)` replaces the four fused linears' forwards
+(returning vLLM's `(output, None)`): `qkv_proj` / `fc1` use `h3-projection-gemm-v1` (a fused column
+equals the same column alone, so fused == separate bitwise, tested); `out_proj` / `fc2` use the
+canonical row-parallel reduction of [`h3-parallel-v1`](h3-parallel.md). Wrong dtype, non-CUDA
+weights or a model without those modules raise in strict mode. Attention, rotary, QK norm,
+SwiGLU and AdaLN stay vLLM-Omni's and are reported `external`. `handle.remove()` restores.
+
+## Evidence
+
+`tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py`. The name-table test compares with `vllm_omni.diffusion.models.minimax_h3.fasth3` when it is importable and is skipped otherwise.
diff --git a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
new file mode 100644
index 000000000..3d48b8caf
--- /dev/null
+++ b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""RFC #420 ``vllm_runtime_adapter_h3``: strict H3 projections inside a vLLM-Omni DiT.
+
+vLLM-Omni's ``MiniMaxH3Attention`` / MLP keep fused linears: ``attn.qkv_proj`` (``[Q;K;V]`` rows),
+``attn.out_proj`` (row parallel), ``mlp.fc1`` (gate-first), ``mlp.fc2`` (row parallel); each
+returns ``(output, bias)``. ``install_h3_rollout(model)`` replaces those four forwards with the
+fixed-order kernels so rollout reads the same bits as the strict training graph:
+
+* ``qkv_proj`` / ``fc1``: ``h3-projection-gemm-v1`` over the fused rows (a column of a fused
+ projection equals the same column computed alone, so fused == separate bitwise);
+* ``out_proj`` / ``fc2``: the canonical 8-unit row-parallel reduction of ``h3-parallel-v1``
+ (``tp`` ranks with a deterministic collective, ``tp=1`` alone), one BF16 rounding.
+
+Attention, rotary and norm stay vLLM-Omni's (their rows are separate); the report lists them as
+``external``. Everything is restored by ``handle.remove()``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import torch
+
+from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
+from rl_engine.backends.shared.triton.gemm import h3_projection as g
+
+LINEAR_NAMES = {"qkv_proj": "column", "fc1": "column", "out_proj": "row", "fc2": "row"}
+
+
+@dataclass
+class H3RolloutHandle:
+ patched: list[str] = field(default_factory=list)
+ _restore: list = field(default_factory=list)
+
+ def remove(self) -> None:
+ for module in self._restore:
+ module.__dict__.pop("forward", None)
+ self._restore.clear()
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *exc):
+ self.remove()
+
+ def report(self) -> dict:
+ return {
+ "contract": par.CONTRACT,
+ "patched": list(self.patched),
+ "external": ["attention", "rotary", "qk_norm", "swiglu", "adaln"],
+ }
+
+
+def _wrap(module, kind, collective, tp):
+ def forward(x):
+ shape = x.shape
+ flat = x.reshape(-1, shape[-1]).contiguous()
+ if kind == "column":
+ out = g.gemm(flat, module.weight, b_t=True)
+ else:
+ out = par.row_parallel_linear(flat, module.weight, collective, tp)
+ return out.view(*shape[:-1], out.shape[-1]), None
+
+ return forward
+
+
+def install_h3_rollout(
+ model: torch.nn.Module, *, collective=None, tp: int = 1, strict: bool = True
+) -> H3RolloutHandle:
+ """Patch every ``qkv_proj`` / ``out_proj`` / ``fc1`` / ``fc2`` of ``blocks`` and the refiner."""
+
+ if tp > 1 and collective is None:
+ raise ValueError("tp > 1 needs a deterministic collective")
+ handle = H3RolloutHandle()
+ for name, module in model.named_modules():
+ leaf = name.rsplit(".", 1)[-1]
+ kind = LINEAR_NAMES.get(leaf)
+ if kind is None or not hasattr(module, "weight"):
+ continue
+ weight = module.weight
+ if not weight.is_cuda or weight.dtype != torch.bfloat16:
+ if strict:
+ raise RuntimeError(
+ f"{name}: strict H3 rollout needs CUDA BF16 weights, got "
+ f"{weight.dtype} on {weight.device}"
+ )
+ continue
+ module.forward = _wrap(module, kind, collective, tp)
+ handle._restore.append(module)
+ handle.patched.append(name)
+ if strict and not handle.patched:
+ raise RuntimeError("no H3 projection modules found (qkv_proj/out_proj/fc1/fc2)")
+ return handle
diff --git a/rl_engine/models/minimax_h3/checkpoint_mapper.py b/rl_engine/models/minimax_h3/checkpoint_mapper.py
new file mode 100644
index 000000000..945900e56
--- /dev/null
+++ b/rl_engine/models/minimax_h3/checkpoint_mapper.py
@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""RFC #420 ``h3_checkpoint_mapper``: Diffusers <-> vLLM-Omni native H3 names and layouts.
+
+The training side (diffusers, FSDP2) and the rollout side (vLLM-Omni) hold the same 638
+tensors under different names and fused layouts. This module is the single, tested mapping, so
+the strict operators on both sides read bit-identical weights:
+
+* ``transformer_blocks.N`` -> ``blocks.N``;
+ ``token_refiner.refiner_blocks.N`` -> ``token_refiner.blocks.N``;
+* ``attn.to_q/to_k/to_v`` -> one fused ``attn.qkv_proj`` laid out ``[Q; K; V]`` (rows), the
+ layout ``h3_qkv_gemm.forward_fused`` is bitwise-equal to;
+* ``attn.to_out.0`` -> ``attn.out_proj``, ``attn.norm_q/k`` -> ``attn.q_norm/k_norm``;
+* ``ff.net.0.proj`` (value-first ``[hidden; gate]``) -> ``mlp.fc1`` (gate-first, halves swapped),
+ ``ff.net.2`` -> ``mlp.fc2``;
+* model-level names per vLLM-Omni's ``fasth3`` table (``proj_in`` -> ``video_patch_proj`` ...).
+
+The tables mirror ``vllm_omni.diffusion.models.minimax_h3.fasth3`` (checked by a test when that
+package is importable). Mapping is pure data movement:
+``native_to_diffusers(diffusers_to_native(x))`` is bitwise ``x``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+import torch
+
+PLAIN, SWAP_HALVES = "plain", "swap_halves"
+QKV_SLOTS = ("q", "k", "v")
+
+MODEL_LEVEL = {
+ "proj_in": "video_patch_proj",
+ "proj_out": "final_layer.video_out",
+ "audio_proj_in": "audio_patch_proj",
+ "audio_proj_out": "final_layer.audio_out",
+ "context_embedder": "condition_proj",
+ "time_embedder.linear_1": "time_embedder.proj_in",
+ "time_embedder.linear_2": "time_embedder.proj_out",
+ "norm_out.linear": "final_layer.adaln_proj.linear",
+ "norm_out.norm": "final_layer.norm",
+ "token_refiner.final_norm": "token_refiner.final_norm",
+}
+BLOCK_PREFIXES = (
+ ("token_refiner.refiner_blocks.", "token_refiner.blocks."),
+ ("transformer_blocks.", "blocks."),
+)
+BLOCK_TARGETS = {
+ "attn.to_q": ("attn.qkv_proj", "q"),
+ "attn.to_k": ("attn.qkv_proj", "k"),
+ "attn.to_v": ("attn.qkv_proj", "v"),
+ "attn.to_out.0": ("attn.out_proj", PLAIN),
+ "attn.norm_q": ("attn.q_norm", PLAIN),
+ "attn.norm_k": ("attn.k_norm", PLAIN),
+ "ff.net.0.proj": ("mlp.fc1", SWAP_HALVES),
+ "ff.net.2": ("mlp.fc2", PLAIN),
+ "adaln_proj.linear": ("adaln_proj.linear", PLAIN),
+ "norm1": ("norm1", PLAIN),
+ "norm2": ("norm2", PLAIN),
+}
+
+
+def resolve(module: str) -> tuple[str, str] | None:
+ """Diffusers module path -> ``(native module path, layout)`` or ``None``."""
+
+ if module in MODEL_LEVEL:
+ return MODEL_LEVEL[module], PLAIN
+ for source, target in BLOCK_PREFIXES:
+ if module.startswith(source):
+ index, _, suffix = module[len(source) :].partition(".")
+ mapped = BLOCK_TARGETS.get(suffix)
+ if index.isdigit() and mapped is not None:
+ return f"{target}{index}.{mapped[0]}", mapped[1]
+ return None
+
+
+def swap_halves(tensor: torch.Tensor) -> torch.Tensor:
+ first, second = tensor.chunk(2, dim=0)
+ return torch.cat((second, first), dim=0)
+
+
+def fuse_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+ """``[Q; K; V]`` rows (the serving layout)."""
+
+ return torch.cat([q, k, v], dim=0)
+
+
+def split_qkv(fused: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ if fused.shape[0] % 3:
+ raise ValueError(f"fused QKV rows {fused.shape[0]} are not a multiple of 3")
+ q, k, v = fused.chunk(3, dim=0)
+ return q, k, v
+
+
+def grouped_qkv_to_qkv(weight: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor:
+ """Native checkpoint's per-head ``[q,k,v]`` groups -> ``[Q; K; V]`` (vLLM-Omni's loader)."""
+
+ per_group = 3 * head_dim
+ if weight.shape[0] != heads * per_group:
+ raise ValueError(f"grouped QKV rows {weight.shape[0]} != {heads * per_group}")
+ grouped = weight.reshape(heads, per_group, *weight.shape[1:])
+ q, k, v = grouped.split(head_dim, dim=1)
+ return torch.cat([t.reshape(heads * head_dim, *weight.shape[1:]) for t in (q, k, v)])
+
+
+def qkv_to_grouped(fused: torch.Tensor, *, heads: int, head_dim: int) -> torch.Tensor:
+ q, k, v = split_qkv(fused)
+ parts = [t.reshape(heads, head_dim, *fused.shape[1:]) for t in (q, k, v)]
+ return torch.cat(parts, dim=1).reshape(heads * 3 * head_dim, *fused.shape[1:])
+
+
+def diffusers_to_native(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Diffusers state dict -> native names, fused QKV, gate-first ``fc1``; unmapped keys fail."""
+
+ out: dict[str, torch.Tensor] = {}
+ qkv: dict[str, dict[str, torch.Tensor]] = {}
+ for name, tensor in state.items():
+ module, _, kind = name.rpartition(".")
+ target = resolve(module)
+ if target is None or kind not in ("weight", "bias"):
+ raise ValueError(f"unmapped Diffusers H3 tensor: {name}")
+ native, layout = target
+ key = f"{native}.{kind}"
+ if layout in QKV_SLOTS:
+ qkv.setdefault(key, {})[layout] = tensor
+ elif layout == SWAP_HALVES:
+ out[key] = swap_halves(tensor)
+ else:
+ out[key] = tensor
+ for key, parts in qkv.items():
+ if set(parts) != set(QKV_SLOTS):
+ raise ValueError(f"{key} needs q, k and v, got {sorted(parts)}")
+ out[key] = fuse_qkv(parts["q"], parts["k"], parts["v"])
+ return out
+
+
+def _native_block(module: str) -> tuple[str, str, str] | None:
+ for source, target in BLOCK_PREFIXES:
+ if module.startswith(target):
+ index, _, suffix = module[len(target) :].partition(".")
+ if index.isdigit():
+ return source, index, suffix
+ return None
+
+
+def native_to_diffusers(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ """Inverse of :func:`diffusers_to_native` (splits fused QKV, swaps ``fc1`` halves back)."""
+
+ model = {native: diff for diff, native in MODEL_LEVEL.items()}
+ block = {
+ native: (diff, layout)
+ for diff, (native, layout) in BLOCK_TARGETS.items()
+ if layout not in QKV_SLOTS
+ }
+ out: dict[str, torch.Tensor] = {}
+ for name, tensor in state.items():
+ module, _, kind = name.rpartition(".")
+ if module in model:
+ out[f"{model[module]}.{kind}"] = tensor
+ continue
+ parsed = _native_block(module)
+ if parsed is None:
+ raise ValueError(f"unmapped native H3 tensor: {name}")
+ source, index, suffix = parsed
+ if suffix == "attn.qkv_proj":
+ for slot, part in zip(QKV_SLOTS, split_qkv(tensor)):
+ out[f"{source}{index}.attn.to_{slot}.{kind}"] = part
+ elif suffix in block:
+ diff, layout = block[suffix]
+ out[f"{source}{index}.{diff}.{kind}"] = (
+ swap_halves(tensor) if layout == SWAP_HALVES else tensor
+ )
+ else:
+ raise ValueError(f"unmapped native H3 tensor: {name}")
+ return out
diff --git a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
new file mode 100644
index 000000000..5686f1b10
--- /dev/null
+++ b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""RFC #420 integration rows: checkpoint mapper and the vLLM-Omni rollout adapter."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from rl_engine.models.minimax_h3 import checkpoint_mapper as cm
+from rl_engine.models.minimax_h3 import fingerprint as fp
+
+needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+# --- checkpoint mapper ----------------------------------------------------------------------
+
+
+def _names():
+ return list(fp.expected_tensor_table(fp.transformer_config()))
+
+
+def test_every_checkpoint_tensor_maps_and_native_names_are_unique():
+ names = _names()
+ assert len(names) == 638
+ state = {n: torch.zeros(6, 2) for n in names}
+ native = cm.diffusers_to_native(state)
+ # q/k/v of each of the 52 blocks (50 + 2 refiner) fuse per kind (weight): 3 -> 1
+ assert len(native) == 638 - 2 * 52
+ assert len(set(native)) == len(native)
+
+
+def test_diffusers_native_roundtrip_is_bitwise():
+ names = _names()
+ gen = torch.Generator().manual_seed(0)
+ state = {n: torch.randn(6, 2, generator=gen) for n in names}
+ back = cm.native_to_diffusers(cm.diffusers_to_native(state))
+ assert set(back) == set(state)
+ assert all(torch.equal(back[n], state[n]) for n in state)
+
+
+def test_fc1_halves_are_swapped_and_qkv_is_q_then_k_then_v():
+ w = torch.arange(8.0).view(8, 1)
+ state = {
+ "transformer_blocks.3.ff.net.0.proj.weight": w,
+ "transformer_blocks.3.attn.to_q.weight": torch.full((2, 1), 1.0),
+ "transformer_blocks.3.attn.to_k.weight": torch.full((2, 1), 2.0),
+ "transformer_blocks.3.attn.to_v.weight": torch.full((2, 1), 3.0),
+ }
+ native = cm.diffusers_to_native(state)
+ assert torch.equal(native["blocks.3.mlp.fc1.weight"], torch.cat([w[4:], w[:4]]))
+ assert native["blocks.3.attn.qkv_proj.weight"].flatten().tolist() == [1, 1, 2, 2, 3, 3]
+
+
+def test_unmapped_names_fail_closed():
+ with pytest.raises(ValueError, match="unmapped"):
+ cm.diffusers_to_native({"transformer_blocks.0.surprise.weight": torch.zeros(1)})
+ with pytest.raises(ValueError, match="unmapped"):
+ cm.native_to_diffusers({"blocks.0.surprise.weight": torch.zeros(1)})
+ with pytest.raises(ValueError, match="q, k and v"):
+ cm.diffusers_to_native({"transformer_blocks.0.attn.to_q.weight": torch.zeros(1, 1)})
+
+
+def test_grouped_checkpoint_layout_roundtrip():
+ w = torch.randn(4 * 3 * 8, 5)
+ fused = cm.grouped_qkv_to_qkv(w, heads=4, head_dim=8)
+ assert torch.equal(cm.qkv_to_grouped(fused, heads=4, head_dim=8), w)
+
+
+def test_tables_match_vllm_omni_when_importable():
+ fasth3 = pytest.importorskip("vllm_omni.diffusion.models.minimax_h3.fasth3")
+ assert dict(fasth3._MODEL_LEVEL_TARGETS) == cm.MODEL_LEVEL
+ assert {k: v[0] for k, v in fasth3._BLOCK_TARGETS.items() if k in cm.BLOCK_TARGETS} == {
+ k: v[0] for k, v in cm.BLOCK_TARGETS.items()
+ }
+
+
+class _Lin(torch.nn.Module):
+ def __init__(self, n, k):
+ super().__init__()
+ self.weight = torch.nn.Parameter(
+ (torch.randn(n, k) / 30).bfloat16().cuda(), requires_grad=False
+ )
+
+ def forward(self, x):
+ raise AssertionError("must be replaced")
+
+
+@needs_gpu
+def test_vllm_adapter_patches_projections_and_restores():
+ pytest.importorskip("triton")
+ from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
+ from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout
+
+ torch.manual_seed(0)
+ model = torch.nn.Module()
+ model.blocks = torch.nn.ModuleList([torch.nn.Module()])
+ blk = model.blocks[0]
+ blk.attn = torch.nn.Module()
+ blk.attn.qkv_proj, blk.attn.out_proj = _Lin(3 * 1024, 640), _Lin(640, 1024)
+ blk.mlp = torch.nn.Module()
+ blk.mlp.fc1, blk.mlp.fc2 = _Lin(2 * 2048, 640), _Lin(640, 2048)
+ x = torch.randn(2, 17, 640).bfloat16().cuda()
+ with install_h3_rollout(model) as handle:
+ qkv, bias = blk.attn.qkv_proj(x)
+ assert bias is None and qkv.shape == (2, 17, 3 * 1024)
+ q, k, v = (blk.attn.qkv_proj.weight[i * 1024 : (i + 1) * 1024] for i in range(3))
+ sep = [par.g.gemm(x.reshape(-1, 640), w, b_t=True) for w in (q, k, v)]
+ assert torch.equal(qkv.reshape(-1, 3072), torch.cat(sep, dim=1)) # fused == separate
+ y, _ = blk.attn.out_proj(torch.randn(2, 17, 1024).bfloat16().cuda())
+ assert y.shape == (2, 17, 640)
+ assert sorted(handle.patched) == [
+ "blocks.0.attn.out_proj",
+ "blocks.0.attn.qkv_proj",
+ "blocks.0.mlp.fc1",
+ "blocks.0.mlp.fc2",
+ ]
+ with pytest.raises(AssertionError, match="must be replaced"):
+ blk.attn.qkv_proj(x)
+
+
+@needs_gpu
+def test_vllm_adapter_strict_rejects_wrong_dtype_and_empty_models():
+ from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout
+
+ with pytest.raises(RuntimeError, match="no H3 projection"):
+ install_h3_rollout(torch.nn.Linear(2, 2).cuda())
+ holder = torch.nn.Module()
+ holder.qkv_proj = torch.nn.Linear(8, 8).cuda() # float32
+ with pytest.raises(RuntimeError, match="BF16"):
+ install_h3_rollout(holder)
From a6c4f17f375ad119616337de1c8b3e89edcdf221 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 19:15:19 +0200
Subject: [PATCH 07/14] test(h3): measurement report for the checkpoint mapper
and the vLLM-Omni adapter
Real checkpoint: 638 tensors map and round trip bitwise; fused qkv equals separate Q/K/V, fc1 halves,
FP64 error and time per projection against torch.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
(cherry picked from commit b517d13633e7fa009570f9fcb46ce8f9f32da198)
---
.../models/h3_mapper_vllm_report.py | 252 ++++++++++++++++++
1 file changed, 252 insertions(+)
create mode 100644 tools/validation/models/h3_mapper_vllm_report.py
diff --git a/tools/validation/models/h3_mapper_vllm_report.py b/tools/validation/models/h3_mapper_vllm_report.py
new file mode 100644
index 000000000..1b3e50fbe
--- /dev/null
+++ b/tools/validation/models/h3_mapper_vllm_report.py
@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Measurements of the RFC #420 ``h3_checkpoint_mapper`` and ``vllm_runtime_adapter_h3`` rows on
+the real checkpoint: JSON report plus a PNG.
+
+ python tools/validation/models/h3_mapper_vllm_report.py --checkpoint \
+ --out report.json --png figure.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import time
+from collections import defaultdict
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from safetensors import safe_open
+
+from rl_engine.backends.shared.triton.gemm import h3_projection as g
+from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout
+from rl_engine.models.minimax_h3 import checkpoint_mapper as cm
+
+GROUP = re.compile(r"^((?:token_refiner\.refiner_blocks|transformer_blocks)\.\d+)\.")
+
+
+def rel(a, b):
+ a, b = a.double(), b.double()
+ return ((a - b).norm() / b.norm()).item()
+
+
+def mapping_report(root):
+ index = {}
+ for f in sorted(root.glob("*.safetensors")):
+ with safe_open(f, "pt") as h:
+ for k in h.keys():
+ index[k] = f
+ groups = defaultdict(list)
+ for name in index:
+ m = GROUP.match(name)
+ groups[m.group(1) if m else "model-level"].append(name)
+ tensors = native_count = 0
+ bitwise, t_fwd, t_back = True, 0.0, 0.0
+ native_names = set()
+ for names in groups.values():
+ state = {}
+ for n in names:
+ with safe_open(index[n], "pt") as h:
+ state[n] = h.get_tensor(n)
+ t0 = time.perf_counter()
+ native = cm.diffusers_to_native(state)
+ t1 = time.perf_counter()
+ back = cm.native_to_diffusers(native)
+ t_back += time.perf_counter() - t1
+ t_fwd += t1 - t0
+ tensors += len(state)
+ native_count += len(native)
+ native_names |= set(native)
+ bitwise &= set(back) == set(state) and all(
+ torch.equal(back[n], state[n]) and back[n].dtype == state[n].dtype for n in state
+ )
+ return (
+ dict(
+ checkpoint_tensors=tensors,
+ native_tensors=native_count,
+ unique_native_names=len(native_names),
+ groups=len(groups),
+ roundtrip_bitwise=bool(bitwise),
+ map_seconds=t_fwd,
+ inverse_seconds=t_back,
+ ),
+ index,
+ )
+
+
+def load(index, name):
+ with safe_open(index[name], "pt") as h:
+ return h.get_tensor(name).cuda()
+
+
+def timeit(fn, n=20):
+ for _ in range(3):
+ fn()
+ torch.cuda.synchronize()
+ s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+ s.record()
+ for _ in range(n):
+ fn()
+ e.record()
+ torch.cuda.synchronize()
+ return s.elapsed_time(e) / n
+
+
+class Lin(torch.nn.Module):
+ def __init__(self, w):
+ super().__init__()
+ self.weight = torch.nn.Parameter(w, requires_grad=False)
+
+ def forward(self, x):
+ return F.linear(x, self.weight), None
+
+
+def layer_rows(index, tokens=(70, 925)):
+ rows = []
+ for blk in (0, 25, 49):
+ p = f"transformer_blocks.{blk}."
+ q, k, v = (load(index, p + f"attn.to_{s}.weight") for s in "qkv")
+ o = load(index, p + "attn.to_out.0.weight")
+ f1 = load(index, p + "ff.net.0.proj.weight")
+ f2 = load(index, p + "ff.net.2.weight")
+ native = cm.diffusers_to_native(
+ {
+ p + "attn.to_q.weight": q,
+ p + "attn.to_k.weight": k,
+ p + "attn.to_v.weight": v,
+ p + "attn.to_out.0.weight": o,
+ p + "ff.net.0.proj.weight": f1,
+ p + "ff.net.2.weight": f2,
+ }
+ )
+ model = torch.nn.Module()
+ model.blocks = torch.nn.ModuleList([torch.nn.Module()])
+ b = model.blocks[0]
+ b.attn, b.mlp = torch.nn.Module(), torch.nn.Module()
+ b.attn.qkv_proj = Lin(native["blocks.%d.attn.qkv_proj.weight" % blk])
+ b.attn.out_proj = Lin(native["blocks.%d.attn.out_proj.weight" % blk])
+ b.mlp.fc1 = Lin(native["blocks.%d.mlp.fc1.weight" % blk])
+ b.mlp.fc2 = Lin(native["blocks.%d.mlp.fc2.weight" % blk])
+ for t in tokens:
+ gen = torch.Generator().manual_seed(blk * 1000 + t)
+ hidden = q.shape[1]
+ x = torch.randn(1, t, hidden, generator=gen).bfloat16().cuda()
+ xo = torch.randn(1, t, o.shape[1], generator=gen).bfloat16().cuda()
+ x2 = torch.randn(1, t, f2.shape[1], generator=gen).bfloat16().cuda()
+ ref = {
+ "qkv_proj": (x, b.attn.qkv_proj),
+ "out_proj": (xo, b.attn.out_proj),
+ "fc1": (x, b.mlp.fc1),
+ "fc2": (x2, b.mlp.fc2),
+ }
+ torch_out = {n: m(xi)[0] for n, (xi, m) in ref.items()}
+ torch_ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()}
+ with install_h3_rollout(model):
+ out = {n: m(xi)[0] for n, (xi, m) in ref.items()}
+ ms = {n: timeit(lambda m=m, xi=xi: m(xi)) for n, (xi, m) in ref.items()}
+ again = {n: m(xi)[0] for n, (xi, m) in ref.items()}
+ flat = x.reshape(-1, hidden)
+ sep = torch.cat([g.gemm(flat, w, b_t=True) for w in (q, k, v)], dim=1)
+ half = f1.shape[0] // 2
+ sep_fc1 = g.gemm(flat, f1, b_t=True)
+ fc1_native = out["fc1"].reshape(t, -1)
+ swapped = torch.cat([sep_fc1[:, half:], sep_fc1[:, :half]], dim=1)
+ for n, (xi, m) in ref.items():
+ exact = xi.reshape(-1, xi.shape[-1]).double() @ m.weight.double().t()
+ rows.append(
+ dict(
+ block=blk,
+ tokens=t,
+ linear=n,
+ repeatable=bool(torch.equal(out[n], again[n])),
+ rel_err_vs_fp64=rel(out[n].reshape(t, -1), exact),
+ torch_rel_err_vs_fp64=rel(torch_out[n].reshape(t, -1), exact),
+ fused_equals_separate=(
+ bool(torch.equal(out[n].reshape(t, -1), sep))
+ if n == "qkv_proj"
+ else None
+ ),
+ fc1_equals_diffusers_swapped=(
+ bool(torch.equal(fc1_native, swapped)) if n == "fc1" else None
+ ),
+ adapter_ms=ms[n],
+ torch_ms=torch_ms[n],
+ slowdown=ms[n] / torch_ms[n],
+ )
+ )
+ return rows
+
+
+def plot(report, path):
+ import matplotlib
+
+ matplotlib.use("Agg")
+ import matplotlib.pyplot as plt
+
+ rows = report["layers"]
+ names = ("qkv_proj", "out_proj", "fc1", "fc2")
+ fig, ax = plt.subplots(1, 3, figsize=(18, 5))
+ m = report["mapping"]
+ ax[0].bar(
+ ["checkpoint\ntensors", "native\ntensors", "round trip\nbitwise equal"],
+ [
+ m["checkpoint_tensors"],
+ m["native_tensors"],
+ m["checkpoint_tensors"] if m["roundtrip_bitwise"] else 0,
+ ],
+ )
+ ax[0].set_title("Checkpoint mapper (real 33B checkpoint)")
+ w = 0.38
+ for j, key in enumerate(("rel_err_vs_fp64", "torch_rel_err_vs_fp64")):
+ vals = [max(r[key] for r in rows if r["linear"] == n) for n in names]
+ ax[1].bar(
+ [i + (j - 0.5) * w for i in range(4)],
+ vals,
+ w,
+ label=("adapter" if j == 0 else "torch F.linear"),
+ )
+ ax[1].set_xticks(range(4), names)
+ ax[1].set_yscale("log")
+ ax[1].set_title("Relative error vs FP64 (worst of 3 blocks x 2 sizes)")
+ ax[1].legend()
+ for t, c in ((70, "C0"), (925, "C1")):
+ vals = [
+ sum(r["slowdown"] for r in rows if r["linear"] == n and r["tokens"] == t) / 3
+ for n in names
+ ]
+ ax[2].plot(names, vals, "o-", color=c, label=f"{t} tokens")
+ ax[2].axhline(1, color="gray", lw=1)
+ ax[2].set_title("Adapter time / torch F.linear time (mean over blocks)")
+ ax[2].legend()
+ fig.tight_layout()
+ fig.savefig(path, dpi=130)
+
+
+def main():
+ p = argparse.ArgumentParser(description=__doc__)
+ p.add_argument("--checkpoint", type=Path, required=True)
+ p.add_argument("--out", type=Path, required=True)
+ p.add_argument("--png", type=Path)
+ args = p.parse_args()
+ torch.backends.cuda.matmul.allow_tf32 = False
+ mapping, index = mapping_report(args.checkpoint / "transformer")
+ report = dict(
+ device=torch.cuda.get_device_name(0),
+ torch=torch.__version__,
+ mapping=mapping,
+ layers=layer_rows(index),
+ )
+ args.out.parent.mkdir(parents=True, exist_ok=True)
+ args.out.write_text(json.dumps(report, indent=1))
+ if args.png:
+ plot(report, args.png)
+ print(json.dumps(mapping, indent=1))
+ for r in report["layers"]:
+ print(r)
+
+
+if __name__ == "__main__":
+ main()
From 4f7c5d113cdcaa9ddf69c0606ae5fe24bb9aaf85 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 22 Oct 2026 21:15:55 +0200
Subject: [PATCH 08/14] feat(h3): bind the h3_full_attention row (#553) in the
strict block and time the full model
CandidateOps attention=auto uses H3FullAttentionOp when the strict _H3_C extension is built (bounded
workspace, same #240 arithmetic) and the interim DeterministicAttentionOp otherwise; h3_full requires
it, interim forces the old op. tools/validation/models/h3_full_speed.py times one full step.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
rl_engine/validation/models/h3_block.py | 23 +++-
tools/validation/models/h3_full_speed.py | 146 +++++++++++++++++++++++
2 files changed, 167 insertions(+), 2 deletions(-)
create mode 100644 tools/validation/models/h3_full_speed.py
diff --git a/rl_engine/validation/models/h3_block.py b/rl_engine/validation/models/h3_block.py
index 0f895613e..1d5c3df2b 100644
--- a/rl_engine/validation/models/h3_block.py
+++ b/rl_engine/validation/models/h3_block.py
@@ -83,7 +83,7 @@ class CandidateOps:
deterministic on every GPU this replay runs on.
"""
- def __init__(self, registry: KernelRegistry | None = None) -> None:
+ def __init__(self, registry: KernelRegistry | None = None, attention: str = "auto") -> None:
from rl_engine.backends.cuda.activation.swiglu import SwiGLUCudaOp
from rl_engine.backends.cuda.attention.deterministic_attn import DeterministicAttentionOp
from rl_engine.backends.shared.triton.gemm.det_gemm import TritonDetGemmOp
@@ -95,11 +95,27 @@ def __init__(self, registry: KernelRegistry | None = None) -> None:
self.gemm = TritonDetGemmOp()
self.qkv = registry.get_op("h3_qkv_gemm", device="cuda")
self.attention = DeterministicAttentionOp()
+ self.attention_h3 = None
+ if attention != "interim":
+ # h3_full_attention (#553): query-chunked, bounded workspace, same #240 arithmetic
+ try:
+ from rl_engine.backends.cuda.attention.h3_full_attention import (
+ H3FullAttentionOp,
+ _extension,
+ )
+
+ _extension()
+ self.attention_h3 = H3FullAttentionOp()
+ except (ImportError, RuntimeError):
+ if attention == "h3_full":
+ raise
self.swiglu = SwiGLUCudaOp()
def backends(self) -> dict[str, str]:
return {
- name: f"{type(op).__module__}.{type(op).__name__}" for name, op in vars(self).items()
+ name: f"{type(op).__module__}.{type(op).__name__}"
+ for name, op in vars(self).items()
+ if op is not None
}
@@ -146,6 +162,9 @@ def _cand_linear(ops: CandidateOps, x: torch.Tensor, weight: torch.Tensor) -> to
def _cand_attention(ops: CandidateOps, q, k, v) -> torch.Tensor:
+ if ops.attention_h3 is not None: # BSHD, bounded workspace
+ q, k, v = (t.contiguous() for t in (q, k, _heads(v)))
+ return ops.attention_h3(q, k, v).flatten(2, 3)
q, k, v = (t.transpose(1, 2).contiguous() for t in (q, k, _heads(v)))
out = ops.attention(q, k, v, causal=False)
return out.transpose(1, 2).flatten(2, 3)
diff --git a/tools/validation/models/h3_full_speed.py b/tools/validation/models/h3_full_speed.py
new file mode 100644
index 000000000..13644c6dd
--- /dev/null
+++ b/tools/validation/models/h3_full_speed.py
@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors
+
+"""Time one full-model forward + backward step, provider against the strict RL-Kernel graph.
+
+ python tools/validation/models/h3_full_speed.py --checkpoint \
+ --sizes 12,925,4253,8349 --out speed.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import time
+from pathlib import Path
+
+import torch
+
+SIZES = {
+ "12": None,
+ "925": dict(text=77, frames=3, side=32, audio=40),
+ "4253": dict(text=77, frames=1, side=128, audio=40),
+ "8349": dict(text=77, frames=2, side=128, audio=40),
+}
+
+
+def sized_inputs(device, *, text, frames, side, audio):
+ from rl_engine.validation.models.h3_cases import h3_fl2va_layout
+
+ layout = h3_fl2va_layout(
+ num_text_tokens=text,
+ num_latent_frames=frames,
+ latent_height=side,
+ latent_width=side,
+ num_audio_latents=audio,
+ device=device,
+ )
+ tags = layout["token_tags"]
+ video, text_idx, audio_idx = ((tags == t).nonzero().flatten() for t in (0, 1, 2))
+ gen = torch.Generator().manual_seed(420539)
+
+ def rnd(*shape):
+ return torch.randn(shape, generator=gen).to(device).requires_grad_()
+
+ return dict(
+ **layout,
+ hidden_states=rnd(1, len(video), 96),
+ audio_hidden_states=rnd(1, len(audio_idx), 32),
+ encoder_hidden_states=rnd(1, len(text_idx), 5120),
+ timestep=torch.tensor([0.0, 500.0], device=device),
+ video_indices=video,
+ text_indices=text_idx,
+ audio_indices=audio_idx,
+ return_dict=False,
+ )
+
+
+def timed(model, inputs, candidate, repeats, attention="auto"):
+ from rl_engine.validation.models.h3_full_model import FullModelTrace, full_model_execution
+
+ grad_inputs = ("hidden_states", "audio_hidden_states", "encoder_hidden_states")
+ backends = {}
+
+ def once():
+ fresh = {
+ k: (v.detach().clone().requires_grad_() if k in grad_inputs else v)
+ for k, v in inputs.items()
+ }
+ trace = FullModelTrace()
+ torch.cuda.synchronize()
+ t0 = time.perf_counter()
+ ops = None
+ if candidate:
+ from rl_engine.validation.models.h3_block import CandidateOps
+
+ ops = CandidateOps(attention=attention)
+ with full_model_execution(
+ model, candidate=candidate, trace=trace, ops=ops, strict_boundaries=candidate
+ ):
+ with torch.autograd.graph.save_on_cpu(pin_memory=True):
+ video, audio = model(**fresh)
+ torch.cuda.synchronize()
+ t1 = time.perf_counter()
+ (video.float().square().mean() + audio.float().square().mean()).backward()
+ torch.cuda.synchronize()
+ backends.update(getattr(trace, "actual_backends", None) or {})
+ return t1 - t0, time.perf_counter() - t0
+
+ torch.cuda.reset_peak_memory_stats()
+ once()
+ runs = [once() for _ in range(repeats)]
+ return dict(
+ forward_s=statistics.median(r[0] for r in runs),
+ total_s=statistics.median(r[1] for r in runs),
+ peak_gib=torch.cuda.max_memory_allocated() / 2**30,
+ attention_backend=backends.get("attention_h3") or backends.get("attention"),
+ )
+
+
+def main():
+ p = argparse.ArgumentParser(description=__doc__)
+ p.add_argument("--checkpoint", type=Path, required=True)
+ p.add_argument("--sizes", default="12,925,4253,8349")
+ p.add_argument("--attention", default="auto", choices=("auto", "h3_full", "interim"))
+ p.add_argument("--repeats", type=int, default=3)
+ p.add_argument("--out", type=Path, required=True)
+ args = p.parse_args()
+ from diffusers import MiniMaxH3Transformer3DModel
+
+ from tools.validation.models.h3_full_model import packed_inputs
+
+ torch.backends.cuda.matmul.allow_tf32 = False
+ torch.backends.cudnn.allow_tf32 = False
+ model = MiniMaxH3Transformer3DModel.from_pretrained(
+ str(args.checkpoint / "transformer"),
+ torch_dtype=torch.bfloat16,
+ local_files_only=True,
+ low_cpu_mem_usage=True,
+ ).to("cuda")
+ model.eval().enable_gradient_checkpointing()
+ for q in model.parameters():
+ q.register_post_accumulate_grad_hook(lambda t: setattr(t, "grad", None))
+ result = dict(device=torch.cuda.get_device_name(0), torch=torch.__version__, sizes={})
+ for size in args.sizes.split(","):
+ inputs = (
+ packed_inputs("cuda") if SIZES[size] is None else sized_inputs("cuda", **SIZES[size])
+ )
+ entry = dict(tokens=int(inputs["position_ids"].shape[0]))
+ for name, candidate in (("provider", False), ("strict", True)):
+ try:
+ entry[name] = timed(model, inputs, candidate, args.repeats, args.attention)
+ except torch.OutOfMemoryError:
+ torch.cuda.empty_cache()
+ entry[name] = dict(error="CUDA out of memory")
+ if "total_s" in entry["provider"] and "total_s" in entry["strict"]:
+ entry["ratio"] = entry["strict"]["total_s"] / entry["provider"]["total_s"]
+ entry["fwd_ratio"] = entry["strict"]["forward_s"] / entry["provider"]["forward_s"]
+ result["sizes"][size] = entry
+ print(size, json.dumps(entry), flush=True)
+ args.out.parent.mkdir(parents=True, exist_ok=True)
+ args.out.write_text(json.dumps(result, indent=1))
+
+
+if __name__ == "__main__":
+ main()
From 0826d3d8058553620e0b98cba5c6e62c9dc91afc Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Fri, 30 Oct 2026 02:50:49 +0000
Subject: [PATCH 09/14] fix(h3): import the WS2 harness from its new location
in the parallel tests
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
tests/models/minimax_h3/test_h3_parallel.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
index 1d1e9d65e..0fed5ad3a 100644
--- a/tests/models/minimax_h3/test_h3_parallel.py
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -16,8 +16,8 @@
from rl_engine.backends.cuda.model_specific.minimax_h3 import parallel as par
from rl_engine.distributed.algorithms.hierarchical import tree_add
+from rl_engine.validation.distributed.h3_ws2 import run_world
from rl_engine.validation.models import h3_tp
-from rl_engine.validation.models.h3_ws2 import run_world
needs_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
From 6a00ba216b5bc9ab2445d987b12a30a130cb0c92 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 5 Nov 2026 10:02:49 +0000
Subject: [PATCH 10/14] test(h3): skip the parallel block grids when the pinned
weights are absent
The grid test loaded the reference block weights directly, so on a host with
two or more GPUs and no extracted weights it failed with FileNotFoundError
instead of skipping like the other real-weight tests. It now takes the
h3_block_weights_cuda fixture, which skips in that case and also supplies the
single-GPU reference weights.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
tests/models/minimax_h3/test_h3_parallel.py | 13 +++++--------
1 file changed, 5 insertions(+), 8 deletions(-)
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
index 0fed5ad3a..01a78acb2 100644
--- a/tests/models/minimax_h3/test_h3_parallel.py
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -174,12 +174,8 @@ def test_scatter_then_gather_is_identity_single_rank(self, length):
# --- Multi-process grids --------------------------------------------------------------------
-def _reference(seq_len=0, backward=False):
- from rl_engine.validation.models.h3_weights import load_h3_block_weights
-
- return h3_tp.block_case(
- load_h3_block_weights("cuda"), h3_tp.Grid(1, 1, 0, 0), seq_len=seq_len, backward=backward
- )
+def _reference(weights, seq_len=0, backward=False):
+ return h3_tp.block_case(weights, h3_tp.Grid(1, 1, 0, 0), seq_len=seq_len, backward=backward)
def _assemble(ranks, tp, cp, key="out"):
@@ -210,7 +206,8 @@ def _assemble(ranks, tp, cp, key="out"):
@needs_gpu
@pytest.mark.parametrize(("tp", "cp", "sp", "length"), GRIDS)
-def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length):
+def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length, h3_block_weights_cuda):
+ # the fixture skips without extracted weights; every rank loads the same pinned files
_triton()
ranks = run_world(
tp * cp,
@@ -222,7 +219,7 @@ def test_parallel_block_equals_single_gpu_bitwise(tp, cp, sp, length):
seq_len=length,
backward=True,
)
- ref = _reference(length, backward=True)
+ ref = _reference(h3_block_weights_cuda, length, backward=True)
out, spans = _assemble(ranks, tp, cp)
assert out.shape == ref["out"].shape, (spans, out.shape)
assert torch.equal(out.cpu(), ref["out"].cpu()), "forward differs from the single-GPU block"
From 100057f73d94522a4bbffd09628b427aff93f563 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 5 Nov 2026 15:47:08 +0000
Subject: [PATCH 11/14] fix(h3): refuse a biased projection in the vLLM-Omni
rollout adapter
The patched qkv_proj / fc1 / out_proj / fc2 forward computes x @ W.T and
returns (output, None), so a module that carries a bias would silently lose
it and rollout would diverge from both vLLM-Omni and the training graph.
H3's four projections have no bias, so instead of adding untested bias paths
the adapter now refuses such a module in strict mode and leaves it on
vLLM-Omni's own forward otherwise.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
.../integrations/engines/rollout/vllm/h3_omni.py | 11 +++++++++++
.../minimax_h3/test_h3_checkpoint_mapper_vllm.py | 6 ++++++
2 files changed, 17 insertions(+)
diff --git a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
index 3d48b8caf..e385cf425 100644
--- a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
+++ b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
@@ -13,6 +13,10 @@
* ``out_proj`` / ``fc2``: the canonical 8-unit row-parallel reduction of ``h3-parallel-v1``
(``tp`` ranks with a deterministic collective, ``tp=1`` alone), one BF16 rounding.
+H3's four projections have no bias, and the patched forward computes ``x @ W.T`` and returns
+``(output, None)``. A module that carries a bias is therefore refused in strict mode (and left
+on vLLM-Omni's own forward otherwise) instead of silently losing its bias.
+
Attention, rotary and norm stay vLLM-Omni's (their rows are separate); the report lists them as
``external``. Everything is restored by ``handle.remove()``.
"""
@@ -87,6 +91,13 @@ def install_h3_rollout(
f"{weight.dtype} on {weight.device}"
)
continue
+ if getattr(module, "bias", None) is not None:
+ if strict:
+ raise RuntimeError(
+ f"{name}: has a bias; the strict H3 rollout kernels compute the bias-free H3 "
+ "projections only"
+ )
+ continue
module.forward = _wrap(module, kind, collective, tp)
handle._restore.append(module)
handle.patched.append(name)
diff --git a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
index 5686f1b10..df0143c01 100644
--- a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
+++ b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
@@ -130,3 +130,9 @@ def test_vllm_adapter_strict_rejects_wrong_dtype_and_empty_models():
holder.qkv_proj = torch.nn.Linear(8, 8).cuda() # float32
with pytest.raises(RuntimeError, match="BF16"):
install_h3_rollout(holder)
+ biased = torch.nn.Module() # the patched forward has no bias term: refuse, never drop it
+ biased.fc2 = torch.nn.Linear(8, 8, bias=True).cuda().bfloat16()
+ with pytest.raises(RuntimeError, match="has a bias"):
+ install_h3_rollout(biased)
+ with install_h3_rollout(biased, strict=False) as handle:
+ assert handle.patched == [] and "forward" not in biased.fc2.__dict__
From 9188133b32f8a49c0a1fcc653f38562df35ff09d Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 5 Nov 2026 15:41:39 +0000
Subject: [PATCH 12/14] fix(h3): create the hierarchical collective's subgroups
in one global order
dist.new_group must be entered by every rank of the world with the same rank
lists in the same order. HierarchicalCollective derived its local and cross
subgroups from its own group only, so with several TP groups (CP > 1, e.g.
make_grid(tp=2, cp=2, local_size=1)) ranks of different TP groups made
different new_group calls and group creation could hang or mismatch.
The collective now takes group_ranks, the global ranks of every group built
together; every rank creates the subgroups of all of them in one order and
keeps its own. make_grid passes all TP groups. A subgroup that does not span
the world is refused without that list. The CP = 1 path (one TP group, as in
the measured TP8 run) creates the same groups as before.
A CPU test simulates every rank of tp/cp/local_size grids and checks that the
new_group sequences agree; there is no 4-GPU run of the CP > 1 hierarchical
path.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
docs/operators/h3-parallel.md | 4 ++
.../distributed/algorithms/hierarchical.py | 41 +++++++++++----
rl_engine/validation/models/h3_tp.py | 9 +++-
tests/models/minimax_h3/test_h3_parallel.py | 51 +++++++++++++++++++
4 files changed, 93 insertions(+), 12 deletions(-)
diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md
index 3ffcacd42..cca5a4959 100644
--- a/docs/operators/h3-parallel.md
+++ b/docs/operators/h3-parallel.md
@@ -85,6 +85,10 @@ Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4x
- dW under CP is exact through `seq_sharded_linear` for non-TP linears; the TP x CP block compares
forward and `dX` (CP shards the token axis of the TP weight gradients).
- `HierarchicalCollective` assumes ranks are packed node-major and node counts of 1/2/4/8.
+ With several TP groups (CP > 1) it is given every TP group's ranks (`group_ranks`, as
+ `make_grid` does) and every rank creates all their subgroups in one order, as
+ `torch.distributed.new_group` requires; a TP subgroup without that list is refused. The
+ hierarchical path with CP > 1 has a CPU test of the creation order only, no GPU run.
## Hopper evidence chart
diff --git a/rl_engine/distributed/algorithms/hierarchical.py b/rl_engine/distributed/algorithms/hierarchical.py
index 253b34399..6e7e35620 100644
--- a/rl_engine/distributed/algorithms/hierarchical.py
+++ b/rl_engine/distributed/algorithms/hierarchical.py
@@ -50,7 +50,13 @@ def __init__(
device: torch.device | str | int | None = None,
*,
max_size_bytes: int = 64 << 20,
+ group_ranks: list[list[int]] | None = None,
) -> None:
+ """``group`` defaults to the world. ``group_ranks`` lists the global ranks of every group
+ whose collective is built at the same time (e.g. all TP groups of a TP x CP grid, this
+ rank's included); it is required when ``group`` does not span the world.
+ """
+
if not dist.is_initialized():
raise RuntimeError("torch.distributed must be initialized before collectives")
self.group = group if group is not None else dist.group.WORLD
@@ -65,16 +71,31 @@ def __init__(
if self.num_groups not in _GROUP_COUNTS:
raise ValueError(f"{self.num_groups} node groups; need 1, 2, 4 or 8")
self.group_index, self.local_rank = divmod(self.rank, local_size)
- # Every rank creates every group in the same order (a torch.distributed requirement).
- ranks = dist.get_process_group_ranks(self.group) # global rank of each group member
- local_groups = [
- dist.new_group(ranks[g * local_size : (g + 1) * local_size])
- for g in range(self.num_groups)
- ]
- cross_groups = [
- dist.new_group([ranks[g * local_size + r] for g in range(self.num_groups)])
- for r in range(local_size)
- ]
+ # ``dist.new_group`` must be entered by every rank of the world, in the same order and
+ # with the same rank lists. A subgroup does not know its sibling groups, so every rank
+ # creates the subgroups of every listed group in one rank-independent order and keeps
+ # its own; a subgroup without the full list is refused rather than left to hang.
+ ranks = list(dist.get_process_group_ranks(self.group)) # global rank of each member
+ everyone = [list(g) for g in group_ranks] if group_ranks is not None else [ranks]
+ if sorted(r for g in everyone for r in g) != list(range(dist.get_world_size())):
+ raise ValueError(
+ "group_ranks must list every group built together, covering each world rank once"
+ )
+ if ranks not in everyone or any(len(g) != self.world_size for g in everyone):
+ raise ValueError("this rank's group must be in group_ranks, all groups the same size")
+ local_groups: list[dist.ProcessGroup] = []
+ cross_groups: list[dist.ProcessGroup] = []
+ for members in everyone:
+ local = [
+ dist.new_group(members[g * local_size : (g + 1) * local_size])
+ for g in range(self.num_groups)
+ ]
+ cross = [
+ dist.new_group([members[g * local_size + r] for g in range(self.num_groups)])
+ for r in range(local_size)
+ ]
+ if members == ranks:
+ local_groups, cross_groups = local, cross
self._cross = cross_groups[self.local_rank]
self._local = DeterministicCollective(
local_groups[self.group_index], device, max_size_bytes=max_size_bytes
diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py
index 2234b2400..1bf5f5c21 100644
--- a/rl_engine/validation/models/h3_tp.py
+++ b/rl_engine/validation/models/h3_tp.py
@@ -52,14 +52,19 @@ def make_grid(tp: int, cp: int, *, local_size: int | None = None) -> Grid:
if world != tp * cp:
raise ValueError(f"world {world} != tp {tp} x cp {cp}")
cp_rank, tp_rank = divmod(rank, tp)
- tp_groups = [dist.new_group(list(range(c * tp, (c + 1) * tp))) for c in range(cp)]
+ tp_ranks = [list(range(c * tp, (c + 1) * tp)) for c in range(cp)]
+ tp_groups = [dist.new_group(ranks) for ranks in tp_ranks]
cp_groups = [dist.new_group(list(range(t, world, tp))) for t in range(tp)]
collective: DeterministicCollective | HierarchicalCollective | None = None
if tp > 1:
mine = tp_groups[cp_rank]
if local_size is not None and local_size < tp:
collective = HierarchicalCollective(
- local_size, mine, torch.cuda.current_device(), max_size_bytes=256 << 20
+ local_size,
+ mine,
+ torch.cuda.current_device(),
+ max_size_bytes=256 << 20,
+ group_ranks=tp_ranks, # every rank creates all TP groups' subgroups (CP > 1)
)
else:
collective = DeterministicCollective(
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
index 01a78acb2..0e48c70ce 100644
--- a/tests/models/minimax_h3/test_h3_parallel.py
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -311,3 +311,54 @@ def test_hierarchical_equals_single_node_fixed_tree(local_size, dtype):
assert torch.equal(ranks[0]["rs"], ranks[0]["hier"][:64])
assert torch.equal(ranks[3]["rs"], ranks[3]["hier"][192:256])
assert torch.equal(ranks[1]["ag"], torch.cat([r["x"][:2] for r in ranks]))
+
+
+def _fake_dist(rank: int, world: int, calls: list):
+ """Just enough of ``torch.distributed`` for one rank to build a ``HierarchicalCollective``."""
+
+ from types import SimpleNamespace
+
+ everything = tuple(range(world))
+
+ def new_group(ranks):
+ calls.append(tuple(ranks))
+ return tuple(ranks)
+
+ return SimpleNamespace(
+ is_initialized=lambda: True,
+ group=SimpleNamespace(WORLD=everything),
+ get_world_size=lambda group=None: len(everything if group is None else group),
+ get_rank=lambda group=None: (everything if group is None else group).index(rank),
+ get_process_group_ranks=lambda group: list(group),
+ new_group=new_group,
+ ProcessGroup=tuple,
+ )
+
+
+@pytest.mark.parametrize(("tp", "cp", "local_size"), [(2, 2, 1), (4, 2, 2), (8, 2, 4), (8, 1, 4)])
+def test_hierarchical_subgroups_are_created_in_one_global_order(monkeypatch, tp, cp, local_size):
+ """``dist.new_group`` is entered by every rank with the same lists in the same order, also
+ when several TP groups (CP > 1) each build a hierarchical collective."""
+
+ from types import SimpleNamespace
+
+ import rl_engine.distributed.algorithms.hierarchical as hier
+
+ world = tp * cp
+ tp_ranks = [list(range(c * tp, (c + 1) * tp)) for c in range(cp)]
+ monkeypatch.setattr(hier, "DeterministicCollective", lambda g, d, **_: SimpleNamespace(g=g))
+ sequences = []
+ for rank in range(world):
+ calls: list = []
+ monkeypatch.setattr(hier, "dist", _fake_dist(rank, world, calls))
+ mine = tuple(tp_ranks[rank // tp])
+ h = hier.HierarchicalCollective(local_size, mine, group_ranks=tp_ranks)
+ sequences.append(calls)
+ assert rank in h._local.g and rank in h._cross # its own subgroups, of its own TP group
+ assert set(h._local.g) | set(h._cross) <= set(mine)
+ assert len(h._local.g) == local_size and len(h._cross) == tp // local_size
+ assert all(seq == sequences[0] for seq in sequences)
+ if cp > 1: # a TP subgroup alone would make the ranks disagree; it is refused instead
+ monkeypatch.setattr(hier, "dist", _fake_dist(0, world, []))
+ with pytest.raises(ValueError, match="group_ranks"):
+ hier.HierarchicalCollective(local_size, tuple(tp_ranks[0]))
From cca6e2fb7d8eb43ac03b98a1d44394e5c9dbfc82 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 5 Nov 2026 16:08:36 +0000
Subject: [PATCH 13/14] fix(h3): refuse short CP shards in compute paths;
correct WS2 timings and chart input
- seq_sharded_linear and the parallel block need a row on every CP rank: a
sequence shorter than cp now raises on every rank alike (same SeqShard,
before any GEMM or collective) instead of failing in the GEMM on the empty
rank while the others wait. Data movement keeps accepting any length; the
docs narrow cp_arbitrary_length to it.
- docs: the H100 per-rank times are the ones in the linked reports (31.5 to
33.3 ms, not 80 to 113 ms), stated as the time of a whole forward-only
block_case call, which also rebuilds the context, operators and weight
shards; the runner comment says the same instead of 'the forward alone'.
- h3_ws2_charts.py reads the runner's /report.json or the published
.json files and exits with a message when it finds no report,
instead of drawing an empty chart.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
docs/operators/h3-parallel.md | 10 +++++++---
.../model_specific/minimax_h3/parallel.py | 19 +++++++++++++++++--
rl_engine/validation/models/h3_tp.py | 1 +
tests/models/minimax_h3/test_h3_parallel.py | 14 ++++++++++++++
tools/validation/models/h3_parallel_run.py | 4 +++-
tools/validation/models/h3_ws2_charts.py | 8 +++++++-
6 files changed, 49 insertions(+), 7 deletions(-)
diff --git a/docs/operators/h3-parallel.md b/docs/operators/h3-parallel.md
index cca5a4959..4ebdef6ed 100644
--- a/docs/operators/h3-parallel.md
+++ b/docs/operators/h3-parallel.md
@@ -43,7 +43,9 @@ two nodes.
## Context parallelism
- `SeqShard`: contiguous row slices differing by at most one row, any length >= 0 (including fewer
- rows than ranks): `cp_arbitrary_length`.
+ rows than ranks) for the data movement (`scatter_packed` / `gather_packed`, Ulysses all-to-all):
+ `cp_arbitrary_length`. The compute paths (`seq_sharded_linear`, the parallel block) need at
+ least one row per rank; a shorter sequence is refused on every rank alike, before any collective.
- `coshard`: hidden, AdaLN row index, token tags, rotary tables are cut with one `SeqShard` so the
metadata follows its rows: `cp_metadata_coshard`.
- `ulysses_attention`: all-to-all sequence <-> head (ragged split sizes), attention on `H/cp` heads
@@ -74,8 +76,10 @@ any node count and compares against the single-GPU block.
| TP4 x CP2 | 8 / 2 | IPC tree (TP in node), NCCL all-to-all (CP across) | bitwise | bitwise | not compared (token axis sharded) |
| TP4 x CP2 + SP | 8 / 2 | same | bitwise | bitwise | not compared |
-Timing of the forward per rank (ms): TP2 80, TP4 102, TP8 110, TP4xCP2 113, TP4xCP2+SP 104
-(tiny 264-row case, launch-bound; not a scaling claim).
+Mean time per rank of one forward call (ms), from the `h3-ws2-parallel-h100/*.json` reports: TP2
+33.1, TP4 32.6, TP8 32.4, TP4xCP2 31.5, TP4xCP2+SP 33.3 (tiny 264-row case, launch-bound; not a
+scaling claim). Each call also rebuilds the block context, the operators and the weight shards, so
+this is the time of a whole `block_case(..., backward=False)` call, not of the kernels alone.
## Limits
diff --git a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
index f61a7dbcd..e8bfbbb46 100644
--- a/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
+++ b/rl_engine/backends/cuda/model_specific/minimax_h3/parallel.py
@@ -217,8 +217,11 @@ def column_parallel_linear(
class SeqShard:
"""This rank's contiguous slice of a packed sequence of ``length`` rows over ``cp`` ranks.
- Slices differ by at most one row (first ``length % cp`` ranks get the extra row), so any
- length is legal, including ``length < cp`` (some ranks hold nothing).
+ Slices differ by at most one row (first ``length % cp`` ranks get the extra row). For data
+ movement (scatter / gather / all-to-all) any length is legal, including ``length < cp`` (some
+ ranks hold nothing). The compute paths (``seq_sharded_linear``, the parallel block) need a row
+ on every rank and refuse a shorter sequence on every rank alike
+ (``require_rows_on_every_rank``), so no rank is left waiting in a collective.
"""
length: int
@@ -239,6 +242,16 @@ def hi(self) -> int:
return self.lo + self.sizes[self.rank]
+def require_rows_on_every_rank(shard: SeqShard) -> None:
+ """Raise, identically on every rank, when some rank's slice of the sequence is empty."""
+
+ if min(shard.sizes) == 0:
+ raise ValueError(
+ f"{shard.length} rows over cp={shard.cp}: the sequence-sharded compute paths need at "
+ "least one row per rank (data movement alone accepts fewer)"
+ )
+
+
def coshard(shard: SeqShard, *tensors: torch.Tensor, dim: int = 0) -> list[torch.Tensor]:
"""Slice every per-row tensor (hidden, timestep_indices, token_tags, position_ids, ...) alike.
@@ -350,6 +363,8 @@ def backward(ctx, grad):
def seq_sharded_linear(x, w, shard: SeqShard, group=None):
"""``x @ W.T`` for this rank's rows; ``dW`` equals the unsharded GEMM (token sweep)."""
+ if shard.cp > 1:
+ require_rows_on_every_rank(shard)
return (
_SeqLinear.apply(x.contiguous(), w, shard, group)
if shard.cp > 1
diff --git a/rl_engine/validation/models/h3_tp.py b/rl_engine/validation/models/h3_tp.py
index 1bf5f5c21..d298a09ec 100644
--- a/rl_engine/validation/models/h3_tp.py
+++ b/rl_engine/validation/models/h3_tp.py
@@ -108,6 +108,7 @@ def run_tp_block(
tp, w = grid.tp, (lambda name: weights[f"{hb.BLOCK}.{name}"])
seq = seq_len or hidden.shape[1]
shard = grid.seq_shard(seq)
+ par.require_rows_on_every_rank(shard) # every rank refuses alike, before any collective
h = par.coshard(shard, hidden, dim=1)[0]
idx = par.coshard(shard, index, dim=0)[0]
cos, sin = par.coshard(shard, rope[0], rope[1], dim=0)
diff --git a/tests/models/minimax_h3/test_h3_parallel.py b/tests/models/minimax_h3/test_h3_parallel.py
index 0e48c70ce..a3c7a8be1 100644
--- a/tests/models/minimax_h3/test_h3_parallel.py
+++ b/tests/models/minimax_h3/test_h3_parallel.py
@@ -362,3 +362,17 @@ def test_hierarchical_subgroups_are_created_in_one_global_order(monkeypatch, tp,
monkeypatch.setattr(hier, "dist", _fake_dist(0, world, []))
with pytest.raises(ValueError, match="group_ranks"):
hier.HierarchicalCollective(local_size, tuple(tp_ranks[0]))
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_compute_paths_refuse_a_sequence_shorter_than_cp_on_every_rank(rank):
+ """``seq_len < cp`` is fine for data movement but not for the compute paths: every rank
+ (also the one that holds a row) refuses before any GEMM or collective."""
+
+ shard = par.SeqShard(1, 2, rank)
+ with pytest.raises(ValueError, match="at least one row per rank"):
+ par.seq_sharded_linear(torch.ones(shard.hi - shard.lo, 4), torch.ones(3, 4), shard)
+ grid = h3_tp.Grid(1, 2, 0, rank)
+ with pytest.raises(ValueError, match="at least one row per rank"):
+ h3_tp.run_tp_block({}, torch.ones(1, 1, 8), None, torch.zeros(1), (None, None), None, grid)
+ par.require_rows_on_every_rank(par.SeqShard(2, 2, rank)) # one row each: accepted
diff --git a/tools/validation/models/h3_parallel_run.py b/tools/validation/models/h3_parallel_run.py
index fc5cf6dbc..db5576bb2 100644
--- a/tools/validation/models/h3_parallel_run.py
+++ b/tools/validation/models/h3_parallel_run.py
@@ -60,7 +60,9 @@ def main() -> None:
weights = load_h3_block_weights("cuda")
kwargs = dict(sp=args.sp, seq_len=args.seq_len, layout=args.layout)
result = h3_tp.block_case(weights, grid, backward=True, **kwargs)
- # timing of the forward alone
+ # forward_s: time of a whole forward-only block_case call. It includes rebuilding the block
+ # context, the operators and the weight shards on every call (kept so the recorded reports
+ # stay comparable); it is not the time of the kernels alone.
torch.cuda.synchronize()
dist.barrier()
begin = time.perf_counter()
diff --git a/tools/validation/models/h3_ws2_charts.py b/tools/validation/models/h3_ws2_charts.py
index 5bc27964d..ae9261524 100644
--- a/tools/validation/models/h3_ws2_charts.py
+++ b/tools/validation/models/h3_ws2_charts.py
@@ -15,7 +15,13 @@
def main(evidence: Path, out: Path) -> None:
- reports = [json.loads(p.read_text()) for p in sorted(evidence.glob("h3_ws2_*.json"))]
+ # h3_parallel_run.py writes /report.json; published evidence keeps one
+ # .json per layout. Read either layout, and refuse an empty chart.
+ paths = sorted(evidence.glob("*/report.json")) or sorted(evidence.glob("*.json"))
+ loaded = [json.loads(p.read_text()) for p in paths]
+ reports = [r for r in loaded if isinstance(r, dict) and "forward_seconds_per_rank" in r]
+ if not reports:
+ raise SystemExit(f"no WS2 reports under {evidence} (*/report.json or .json)")
labels = []
times = []
colors = []
From 39cb7d3dca38c3767a9bea25377e3905f22892a0 Mon Sep 17 00:00:00 2001
From: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
Date: Thu, 5 Nov 2026 16:17:23 +0000
Subject: [PATCH 14/14] fix(h3): explicit native QKV layout, exact adapter
restore and untimed op setup
- checkpoint mapper: both directions take qkv_layout. 'module' (default) is
the loaded module's [Q; K; V] rows, what install_h3_rollout and
load_state_dict see; 'checkpoint' is a serialised native checkpoint's
per-head groups, which vLLM-Omni's loader converts itself, so writing the
module layout for that loader would permute twice. native_to_diffusers
also refuses tensor kinds other than weight / bias.
- vLLM-Omni adapter: only modules at a block or refiner projection path are
patched (another module with the same leaf name is left alone), every
candidate is checked before any forward changes (a strict refusal leaves
the model untouched), and remove() restores a module's own instance
forward instead of deleting it.
- h3_full_speed.py builds CandidateOps once and starts the clock after the
binding is entered, so the strict times no longer include registry lookups,
CUDA initialisation and binding setup on every repeat.
Signed-off-by: Wenbo Ji <36562829+fusheng-ji@users.noreply.github.com>
---
.../engines/rollout/vllm/h3_omni.py | 40 +++++--
.../models/minimax_h3/checkpoint_mapper.py | 55 ++++++++--
.../test_h3_checkpoint_mapper_vllm.py | 102 +++++++++++++++++-
tools/validation/models/h3_full_speed.py | 15 +--
4 files changed, 183 insertions(+), 29 deletions(-)
diff --git a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
index e385cf425..815fa16fa 100644
--- a/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
+++ b/rl_engine/integrations/engines/rollout/vllm/h3_omni.py
@@ -23,6 +23,7 @@
from __future__ import annotations
+import re
from dataclasses import dataclass, field
import torch
@@ -31,16 +32,27 @@
from rl_engine.backends.shared.triton.gemm import h3_projection as g
LINEAR_NAMES = {"qkv_proj": "column", "fc1": "column", "out_proj": "row", "fc2": "row"}
+# The four projections of a transformer block or a token-refiner block, by module path (an outer
+# wrapper prefix is allowed). A module elsewhere that happens to share a leaf name is not touched.
+_PROJECTION_PATH = re.compile(
+ r"(?:^|\.)(?:token_refiner\.)?blocks\.\d+\.(?:attn\.(?:qkv_proj|out_proj)|mlp\.(?:fc1|fc2))$"
+)
+_ABSENT = object()
@dataclass
class H3RolloutHandle:
patched: list[str] = field(default_factory=list)
- _restore: list = field(default_factory=list)
+ _restore: list = field(default_factory=list) # (module, its own forward or _ABSENT)
def remove(self) -> None:
- for module in self._restore:
- module.__dict__.pop("forward", None)
+ """Put back exactly what each module had: its own instance ``forward`` or none."""
+
+ for module, previous in reversed(self._restore):
+ if previous is _ABSENT:
+ module.__dict__.pop("forward", None)
+ else:
+ module.forward = previous
self._restore.clear()
def __enter__(self):
@@ -73,15 +85,17 @@ def forward(x):
def install_h3_rollout(
model: torch.nn.Module, *, collective=None, tp: int = 1, strict: bool = True
) -> H3RolloutHandle:
- """Patch every ``qkv_proj`` / ``out_proj`` / ``fc1`` / ``fc2`` of ``blocks`` and the refiner."""
+ """Patch every ``qkv_proj`` / ``out_proj`` / ``fc1`` / ``fc2`` of ``blocks`` and the refiner.
+
+ Only modules at a block projection path are considered. Every candidate is checked before
+ any forward changes, so a strict refusal leaves the model untouched.
+ """
if tp > 1 and collective is None:
raise ValueError("tp > 1 needs a deterministic collective")
- handle = H3RolloutHandle()
+ selected = []
for name, module in model.named_modules():
- leaf = name.rsplit(".", 1)[-1]
- kind = LINEAR_NAMES.get(leaf)
- if kind is None or not hasattr(module, "weight"):
+ if not _PROJECTION_PATH.search(name) or not hasattr(module, "weight"):
continue
weight = module.weight
if not weight.is_cuda or weight.dtype != torch.bfloat16:
@@ -98,9 +112,13 @@ def install_h3_rollout(
"projections only"
)
continue
+ selected.append((name, module))
+ if strict and not selected:
+ raise RuntimeError("no H3 projection modules found (qkv_proj/out_proj/fc1/fc2)")
+ handle = H3RolloutHandle()
+ for name, module in selected:
+ kind = LINEAR_NAMES[name.rsplit(".", 1)[-1]]
+ handle._restore.append((module, module.__dict__.get("forward", _ABSENT)))
module.forward = _wrap(module, kind, collective, tp)
- handle._restore.append(module)
handle.patched.append(name)
- if strict and not handle.patched:
- raise RuntimeError("no H3 projection modules found (qkv_proj/out_proj/fc1/fc2)")
return handle
diff --git a/rl_engine/models/minimax_h3/checkpoint_mapper.py b/rl_engine/models/minimax_h3/checkpoint_mapper.py
index 945900e56..e2c69a429 100644
--- a/rl_engine/models/minimax_h3/checkpoint_mapper.py
+++ b/rl_engine/models/minimax_h3/checkpoint_mapper.py
@@ -19,6 +19,14 @@
The tables mirror ``vllm_omni.diffusion.models.minimax_h3.fasth3`` (checked by a test when that
package is importable). Mapping is pure data movement:
``native_to_diffusers(diffusers_to_native(x))`` is bitwise ``x``.
+
+Two QKV layouts exist on the native side, and both directions take ``qkv_layout``:
+
+* ``"module"`` (default): ``attn.qkv_proj.weight`` as the loaded module holds it, ``[Q; K; V]``
+ rows. This is what ``install_h3_rollout`` reads and what ``load_state_dict`` writes in place;
+* ``"checkpoint"``: a serialised native checkpoint, per-head ``[q, k, v]`` groups, which
+ vLLM-Omni's loader converts to ``[Q; K; V]`` itself. Writing the module layout into a file
+ for that loader would permute the rows twice.
"""
from __future__ import annotations
@@ -29,6 +37,9 @@
PLAIN, SWAP_HALVES = "plain", "swap_halves"
QKV_SLOTS = ("q", "k", "v")
+QKV_LAYOUTS = ("module", "checkpoint")
+HEAD_DIM = 128
+TENSOR_KINDS = ("weight", "bias")
MODEL_LEVEL = {
"proj_in": "video_patch_proj",
@@ -110,15 +121,33 @@ def qkv_to_grouped(fused: torch.Tensor, *, heads: int, head_dim: int) -> torch.T
return torch.cat(parts, dim=1).reshape(heads * 3 * head_dim, *fused.shape[1:])
-def diffusers_to_native(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
- """Diffusers state dict -> native names, fused QKV, gate-first ``fc1``; unmapped keys fail."""
+def _check_layout(qkv_layout: str) -> None:
+ if qkv_layout not in QKV_LAYOUTS:
+ raise ValueError(f"qkv_layout must be one of {QKV_LAYOUTS}, got {qkv_layout!r}")
+
+
+def _heads(fused: torch.Tensor) -> int:
+ if fused.shape[0] % (3 * HEAD_DIM):
+ raise ValueError(f"fused QKV rows {fused.shape[0]} are not whole {HEAD_DIM}-wide heads")
+ return fused.shape[0] // (3 * HEAD_DIM)
+
+def diffusers_to_native(
+ state: Mapping[str, torch.Tensor], *, qkv_layout: str = "module"
+) -> dict[str, torch.Tensor]:
+ """Diffusers state dict -> native names, fused QKV, gate-first ``fc1``; unmapped keys fail.
+
+ ``qkv_layout`` selects the module layout (``[Q; K; V]``) or the serialised checkpoint
+ layout (per-head groups) of the fused QKV, see the module docstring.
+ """
+
+ _check_layout(qkv_layout)
out: dict[str, torch.Tensor] = {}
qkv: dict[str, dict[str, torch.Tensor]] = {}
for name, tensor in state.items():
module, _, kind = name.rpartition(".")
target = resolve(module)
- if target is None or kind not in ("weight", "bias"):
+ if target is None or kind not in TENSOR_KINDS:
raise ValueError(f"unmapped Diffusers H3 tensor: {name}")
native, layout = target
key = f"{native}.{kind}"
@@ -131,7 +160,10 @@ def diffusers_to_native(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Te
for key, parts in qkv.items():
if set(parts) != set(QKV_SLOTS):
raise ValueError(f"{key} needs q, k and v, got {sorted(parts)}")
- out[key] = fuse_qkv(parts["q"], parts["k"], parts["v"])
+ fused = fuse_qkv(parts["q"], parts["k"], parts["v"])
+ if qkv_layout == "checkpoint":
+ fused = qkv_to_grouped(fused, heads=_heads(fused), head_dim=HEAD_DIM)
+ out[key] = fused
return out
@@ -144,9 +176,16 @@ def _native_block(module: str) -> tuple[str, str, str] | None:
return None
-def native_to_diffusers(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
- """Inverse of :func:`diffusers_to_native` (splits fused QKV, swaps ``fc1`` halves back)."""
+def native_to_diffusers(
+ state: Mapping[str, torch.Tensor], *, qkv_layout: str = "module"
+) -> dict[str, torch.Tensor]:
+ """Inverse of :func:`diffusers_to_native` (splits fused QKV, swaps ``fc1`` halves back).
+ ``qkv_layout`` must name the layout the native QKV tensors are in (``"checkpoint"`` for a
+ serialised native checkpoint). Tensor kinds other than ``weight`` / ``bias`` are refused.
+ """
+
+ _check_layout(qkv_layout)
model = {native: diff for diff, native in MODEL_LEVEL.items()}
block = {
native: (diff, layout)
@@ -156,6 +195,8 @@ def native_to_diffusers(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Te
out: dict[str, torch.Tensor] = {}
for name, tensor in state.items():
module, _, kind = name.rpartition(".")
+ if kind not in TENSOR_KINDS:
+ raise ValueError(f"unmapped native H3 tensor: {name}")
if module in model:
out[f"{model[module]}.{kind}"] = tensor
continue
@@ -164,6 +205,8 @@ def native_to_diffusers(state: Mapping[str, torch.Tensor]) -> dict[str, torch.Te
raise ValueError(f"unmapped native H3 tensor: {name}")
source, index, suffix = parsed
if suffix == "attn.qkv_proj":
+ if qkv_layout == "checkpoint":
+ tensor = grouped_qkv_to_qkv(tensor, heads=_heads(tensor), head_dim=HEAD_DIM)
for slot, part in zip(QKV_SLOTS, split_qkv(tensor)):
out[f"{source}{index}.attn.to_{slot}.{kind}"] = part
elif suffix in block:
diff --git a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
index df0143c01..3c25869df 100644
--- a/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
+++ b/tests/models/minimax_h3/test_h3_checkpoint_mapper_vllm.py
@@ -126,13 +126,105 @@ def test_vllm_adapter_strict_rejects_wrong_dtype_and_empty_models():
with pytest.raises(RuntimeError, match="no H3 projection"):
install_h3_rollout(torch.nn.Linear(2, 2).cuda())
- holder = torch.nn.Module()
- holder.qkv_proj = torch.nn.Linear(8, 8).cuda() # float32
+ holder = _model_with(attn=("qkv_proj", torch.nn.Linear(8, 8).cuda())) # float32
with pytest.raises(RuntimeError, match="BF16"):
install_h3_rollout(holder)
- biased = torch.nn.Module() # the patched forward has no bias term: refuse, never drop it
- biased.fc2 = torch.nn.Linear(8, 8, bias=True).cuda().bfloat16()
+ # the patched forward has no bias term: refuse, never drop it
+ biased = _model_with(mlp=("fc2", torch.nn.Linear(8, 8, bias=True).cuda().bfloat16()))
with pytest.raises(RuntimeError, match="has a bias"):
install_h3_rollout(biased)
+ fc2 = biased.blocks[0].mlp.fc2
with install_h3_rollout(biased, strict=False) as handle:
- assert handle.patched == [] and "forward" not in biased.fc2.__dict__
+ assert handle.patched == [] and "forward" not in fc2.__dict__
+
+
+def _model_with(**parts):
+ """``blocks.0..`` holding the given module, as in vLLM-Omni's DiT."""
+
+ model = torch.nn.Module()
+ model.blocks = torch.nn.ModuleList([torch.nn.Module()])
+ for group, (name, module) in parts.items():
+ setattr(model.blocks[0], group, torch.nn.Module())
+ setattr(getattr(model.blocks[0], group), name, module)
+ return model
+
+
+@needs_gpu
+def test_vllm_adapter_patches_block_projections_only_and_restores_exactly():
+ from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout
+
+ model = _model_with(attn=("qkv_proj", _Lin(3 * 128, 64)), mlp=("fc1", _Lin(256, 64)))
+ model.audio_head = torch.nn.Module()
+ model.audio_head.fc1 = _Lin(64, 64) # same leaf name, not a block projection
+ own = model.blocks[0].mlp.fc1.forward = lambda x: ("own forward", x)
+ with install_h3_rollout(model) as handle:
+ assert sorted(handle.patched) == ["blocks.0.attn.qkv_proj", "blocks.0.mlp.fc1"]
+ assert "forward" not in model.audio_head.fc1.__dict__
+ assert model.blocks[0].mlp.fc1.forward is own # the instance override is back
+ assert "forward" not in model.blocks[0].attn.qkv_proj.__dict__
+
+
+@needs_gpu
+def test_vllm_adapter_strict_refusal_leaves_every_module_unpatched():
+ from rl_engine.integrations.engines.rollout.vllm.h3_omni import install_h3_rollout
+
+ model = _model_with(attn=("qkv_proj", _Lin(3 * 128, 64)))
+ model.blocks.append(torch.nn.Module())
+ model.blocks[1].attn = torch.nn.Module()
+ model.blocks[1].attn.qkv_proj = torch.nn.Linear(64, 3 * 128).cuda() # float32: refused
+ with pytest.raises(RuntimeError, match="BF16"):
+ install_h3_rollout(model)
+ assert "forward" not in model.blocks[0].attn.qkv_proj.__dict__
+
+
+def test_qkv_layout_is_explicit_module_or_serialised_checkpoint():
+ """The module layout is ``[Q; K; V]``; a serialised native checkpoint holds per-head groups
+ (vLLM-Omni's loader converts them). Each direction round-trips in its own layout."""
+
+ gen = torch.Generator().manual_seed(1)
+ heads, d = 2, cm.HEAD_DIM
+ state = {
+ f"transformer_blocks.0.attn.to_{s}.weight": torch.randn(heads * d, 3, generator=gen)
+ for s in "qkv"
+ }
+ module = cm.diffusers_to_native(state)["blocks.0.attn.qkv_proj.weight"]
+ ckpt = cm.diffusers_to_native(state, qkv_layout="checkpoint")["blocks.0.attn.qkv_proj.weight"]
+ assert torch.equal(
+ module, torch.cat([state[f"transformer_blocks.0.attn.to_{s}.weight"] for s in "qkv"])
+ )
+ assert torch.equal(ckpt, cm.qkv_to_grouped(module, heads=heads, head_dim=d))
+ assert not torch.equal(ckpt, module) # writing the module layout for the loader is wrong
+ for layout, tensor in (("module", module), ("checkpoint", ckpt)):
+ back = cm.native_to_diffusers({"blocks.0.attn.qkv_proj.weight": tensor}, qkv_layout=layout)
+ assert all(torch.equal(back[n], state[n]) for n in state)
+ with pytest.raises(ValueError, match="qkv_layout"):
+ cm.diffusers_to_native(state, qkv_layout="grouped")
+
+
+def test_native_tensor_kinds_other_than_weight_and_bias_are_refused():
+ with pytest.raises(ValueError, match="unmapped"):
+ cm.native_to_diffusers({"video_patch_proj.running_mean": torch.zeros(1)})
+ with pytest.raises(ValueError, match="unmapped"):
+ cm.native_to_diffusers({"blocks.0.mlp.fc2.scale": torch.zeros(1)})
+
+
+def test_checkpoint_qkv_layout_is_what_the_vllm_omni_loader_reorders():
+ """vLLM-Omni's ``load_weights`` reorders a native ``attn.qkv_proj.weight`` from per-head groups
+ to ``[Q; K; V]``: the ``"checkpoint"`` layout must come out as the module layout."""
+
+ loader = pytest.importorskip("vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer")
+ gen = torch.Generator().manual_seed(2)
+ heads = 3
+ state = {
+ f"transformer_blocks.0.attn.to_{s}.weight": torch.randn(
+ heads * cm.HEAD_DIM, 4, generator=gen
+ )
+ for s in "qkv"
+ }
+ key = "blocks.0.attn.qkv_proj.weight"
+ module = cm.diffusers_to_native(state)[key]
+ ckpt = cm.diffusers_to_native(state, qkv_layout="checkpoint")[key]
+ loaded = loader._reorder_grouped_qkv_to_qkv(
+ ckpt, num_query_groups=heads, heads_per_group=1, head_dim=cm.HEAD_DIM
+ )
+ assert torch.equal(loaded, module)
diff --git a/tools/validation/models/h3_full_speed.py b/tools/validation/models/h3_full_speed.py
index 13644c6dd..7a7682ca4 100644
--- a/tools/validation/models/h3_full_speed.py
+++ b/tools/validation/models/h3_full_speed.py
@@ -61,6 +61,11 @@ def timed(model, inputs, candidate, repeats, attention="auto"):
grad_inputs = ("hidden_states", "audio_hidden_states", "encoder_hidden_states")
backends = {}
+ ops = None
+ if candidate: # built once and outside every timed region: registry lookups, CUDA init
+ from rl_engine.validation.models.h3_block import CandidateOps
+
+ ops = CandidateOps(attention=attention)
def once():
fresh = {
@@ -68,16 +73,12 @@ def once():
for k, v in inputs.items()
}
trace = FullModelTrace()
- torch.cuda.synchronize()
- t0 = time.perf_counter()
- ops = None
- if candidate:
- from rl_engine.validation.models.h3_block import CandidateOps
-
- ops = CandidateOps(attention=attention)
with full_model_execution(
model, candidate=candidate, trace=trace, ops=ops, strict_boundaries=candidate
):
+ # timed: forward and backward of the bound model only, not the binding setup
+ torch.cuda.synchronize()
+ t0 = time.perf_counter()
with torch.autograd.graph.save_on_cpu(pin_memory=True):
video, audio = model(**fresh)
torch.cuda.synchronize()