huggingface/transformers · error · ValueError

_aten_grouped_mm: number of experts (mat_b.shape[0]) must be

Error message

_aten_grouped_mm: number of experts (mat_b.shape[0]) must be static at translation time

What it means

The ONNX translation of torch._grouped_mm (used by MoE experts) unrolls the grouped matmul into per-group Slice+MatMul at translation time. It needs the number of experts G = mat_b.shape[0] as a concrete int; if the axis is symbolic (marked dynamic, or an unbacked SymInt from data-dependent slicing) it cannot unroll and raises this ValueError.

Source

Thrown at src/transformers/exporters/exporter_onnx.py:847

    one_hot = op.OneHot(self, depth, op.Constant(value_ints=[0, 1]), axis=-1)
    return op.ReduceSum(one_hot, op.Constant(value_ints=[0]), keepdims=0)


def _aten_grouped_mm(mat_a: TReal, mat_b: TReal, offs: INT64, bias=None, out_dtype=None) -> TReal:
    """ONNX implementation of `aten._grouped_mm.default`.

    `_grouped_mm(mat_a: (M, K), mat_b: (G, K, N), offs: (G,))` computes `out[r] =
    mat_a[r] @ mat_b[group(r)]` where rows are sorted by group and `offs` holds the
    cumulative end index per group.

    Per-group `Slice + MatMul + Concat`. `G` (number of experts) is static for any
    concrete model, so unroll at translation time: emit one `Slice + MatMul` triple
    per group and a final `Concat`. Avoids the `(M, K, N)` materialisation a naive
    `weight[group_idx]` gather would emit — peak memory is `O(M·N + max(n_g)·K + K·N)`.
    """
    G = mat_b.shape[0]
    if not isinstance(G, int):
        raise ValueError("_aten_grouped_mm: number of experts (mat_b.shape[0]) must be static at translation time")

    offs_i64 = op.Cast(offs, to=7)
    axes_0 = op.Constant(value_ints=[0])
    zero_1d = op.Constant(value_ints=[0])

    outputs = []
    prev_end = zero_1d
    for g in range(G):
        g_lo = op.Constant(value_ints=[g])
        g_hi = op.Constant(value_ints=[g + 1])
        end = op.Slice(offs_i64, g_lo, g_hi, axes_0)  # (1,) — offs[g]
        a_g = op.Slice(mat_a, prev_end, end, axes_0)  # (n_g, K)
        w_g = op.Squeeze(op.Slice(mat_b, g_lo, g_hi, axes_0), axes_0)  # (K, N)
        outputs.append(op.MatMul(a_g, w_g))  # (n_g, N)
        prev_end = end

    return op.Concat(*outputs, axis=0)  # (M, N)

View on GitHub (pinned to a597f97485)

Solutions

  1. Pass explicit dynamic_shapes that leave the experts (mat_b.shape[0]) axis static; only mark truly varying axes (batch, sequence) dynamic.
  2. If using dynamic=True alone, switch to dynamic=True plus dynamic_shapes with Dim.AUTO only on batch/seq.
  3. For ONNX specifically, prefer fully static shapes for MoE weights — expert count is fixed per checkpoint.

Example fix

# before
cfg = DynamoConfig(dynamic=True)  # experts dim becomes symbolic -> ValueError at ONNX translation

# after
from torch.export import Dim
seq = Dim("seq", min=1, max=4096)
cfg = DynamoConfig(
    dynamic=True,
    dynamic_shapes={"hidden_states": {0: Dim.AUTO, 1: seq}},  # experts dim stays static
)
Defensive patterns

Strategy: validation

Validate before calling

# Before ONNX export of a MoE model, ensure the experts dim is static.
from torch.export import Dim

cfg = DynamoConfig(
    dynamic=True,
    dynamic_shapes={
        "hidden_states": {0: Dim.AUTO, 1: Dim("seq")},  # expert weights left static
        # do NOT mark any dim of router_logits/expert weights dynamic
    },
)

Type guard

def experts_dim_is_static(mat_b: torch.Tensor) -> bool:
    return isinstance(mat_b.shape[0], int)

Try / catch

try:
    OnnxExporter().export(model, inputs, cfg)
except ValueError as e:
    if "number of experts" in str(e) and "static" in str(e):
        cfg = DynamoConfig(dynamic=False)  # fully static retry — expert count is fixed per checkpoint
        OnnxExporter().export(model, inputs, cfg)
    else:
        raise

Prevention

When it happens

Trigger: Exporting a Mixture-of-Experts model to ONNX while dynamic_shapes marks the experts dimension dynamic (e.g. blanket Dim.AUTO from dynamic=True), leaving mat_b.shape[0] a SymInt instead of an int.

Common situations: Using DynamoConfig(dynamic=True) with no explicit dynamic_shapes — the exporter even warns that this marks every axis symbolic; exporting splinter/MoE models where G accidentally falls into the dynamic set.

Related errors


AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/35ac827bf35aa8ec. Report an issue: GitHub.