Skip to content

fix(peft): follow TE's _GroupedLinear contract instead of pinning one layout - #24

Merged
yueming-yuan merged 2 commits into
bridgefrom
radixark/bridge/te-grouped-linear-2.17
Jul 27, 2026
Merged

fix(peft): follow TE's _GroupedLinear contract instead of pinning one layout#24
yueming-yuan merged 2 commits into
bridgefrom
radixark/bridge/te-grouped-linear-2.17

Conversation

@yueming-yuan

Copy link
Copy Markdown

Problem

PEFT._forward_te_grouped_linear calls transformer_engine's private
_GroupedLinear directly. It was written against TE ≤ 2.14, and TE has moved that
contract twice since:

TE signature non_tensor_args returns
≤ 2.14 forward(ctx, inp, non_tensor_args, *weights_and_biases) 21 fields, m_splits first, module near the end Tensor
2.15 / 2.16 same 22 — moduleweight_workspaces + cache_weight Tensor
≥ 2.17 forward(ctx, inp, **m_splits**, non_tensor_args, *weights_and_biases) 21 — m_splits promoted out of the tuple (Tensor, list)

On TE 2.17 the old call passes only two leading positionals, so everything shifts by
one: the non_tensor_args tuple binds to the m_splits parameter and weight[0]
binds to non_tensor_args. TE then unpacks that weight tensor, iterating its
first dimension, and the failure surfaces far from its cause:

File "transformer_engine/pytorch/module/grouped_linear.py", line 405, in forward
ValueError: not enough values to unpack (expected 21, got 16)

16 is the weight tensor's dim-0, not any tuple length — which is what makes this
one confusing to chase.

Multi-LoRA on grouped MoE experts is therefore broken on TE ≥ 2.17, and was already
wrong on 2.15/2.16 (they want 22 fields). Hit while bumping miles to
TransformerEngine 2.17.0.

Approach

A version table would silently misalign again the next time TE moves — the same
failure mode, just later. Instead this derives the contract from the TE that is
actually installed:

  • _te_grouped_linear_contract() reads is m_splits positional? from
    inspect.signature, and the tuple field names from the = non_tensor_args
    unpack via inspect.getsource.
  • The call site builds the tuple by name from a dict covering every field any of
    these releases asks for, and raises a named RuntimeError listing the unknown
    field if TE adds one.
  • The result is normalised so both the Tensor and (Tensor, list) return forms work.

Net effect: a future TE change becomes an explicit error naming the field, instead of
a misaligned tuple surfacing as an unrelated unpack error inside TE.

Note weight_workspaces / cache_weight are passed empty/False: this shim always
passes is_first_microbatch=None, so no fp8 workspace caching happens on either the
old (module) or new (explicit workspaces) contract.

Verification

Checked offline against the 2.14.0 / 2.15.0 / 2.16.0 / 2.17.0 wheels:

TE 2.14.0: m_splits positional=False  fields=21  first=m_splits
TE 2.15.0: m_splits positional=False  fields=22  first=m_splits
TE 2.16.0: m_splits positional=False  fields=22  first=m_splits
TE 2.17.0: m_splits positional=True   fields=21  first=use_bias

and the supplied field dict covers every field in all four (no gaps). Repo
pre-commit passes.

I do not have a GPU box to run the grouped-MoE LoRA path end to end — the arity and
ordering are verified statically against the real wheels, but a functional run on
TE 2.17 before merge would be worth it.

… layout

`_forward_te_grouped_linear` calls `transformer_engine`'s private
`_GroupedLinear` directly, and was written against TE <= 2.14. TE has moved
that contract twice since:

  TE <= 2.14   forward(ctx, inp,           non_tensor_args, *weights_and_biases)
               21 fields, `m_splits` first, a `module` handle near the end,
               returns a Tensor
  TE 2.15/2.16 same signature; `module` replaced by `weight_workspaces` +
               `cache_weight` (22 fields)
  TE >= 2.17   forward(ctx, inp, m_splits, non_tensor_args, *weights_and_biases)
               `m_splits` promoted to its own positional parameter (21 fields),
               returns (out, new_workspaces)

On TE 2.17 the old call shifts every positional argument by one: the
non_tensor_args tuple lands in the `m_splits` slot and `weight[0]` lands in
`non_tensor_args`. TE then unpacks that weight tensor, iterating its first
dimension, and the failure surfaces far from the cause as

    File transformer_engine/pytorch/module/grouped_linear.py, line 405, in forward
    ValueError: not enough values to unpack (expected 21, got 16)

where 16 is simply the tensor's dim-0, not any tuple length. Multi-LoRA on
grouped MoE experts is dead on TE >= 2.17 (and was already wrong on 2.15/2.16,
which want 22 fields).

Rather than add a version table -- which would silently misalign again on the
next TE change -- this reads the field order off the installed TE:
`_te_grouped_linear_contract()` takes `m_splits`-is-positional from
`inspect.signature` and the tuple field names from the `= non_tensor_args`
unpack in `inspect.getsource`. The call site then builds the tuple by name from
a dict covering every field any of these releases asks for, and raises a named
RuntimeError if TE introduces one we do not supply. The return is normalised so
both the Tensor and (Tensor, list) forms work.

Verified offline against the 2.14.0 / 2.15.0 / 2.16.0 / 2.17.0 wheels:
detection yields (m_splits positional, field count) = (False, 21), (False, 22),
(False, 22), (True, 21), and the supplied dict covers every field in all four.
yueming-yuan added a commit to radixark/miles that referenced this pull request Jul 26, 2026
…ge#24

Validates the TE 2.17 grouped-linear contract fix end to end. The
`not enough values to unpack (expected 21, got 16)` failures in the
test_glm5_*_lora_ci tests come from megatron-bridge's
`_forward_te_grouped_linear` calling TE's private `_GroupedLinear` with TE
<=2.14's positional layout; miles picked that up when it bumped
TransformerEngine 2.12 -> 2.17 in #1781.

Revert to `@bridge` -- ideally pinned to the merged commit rather than a branch,
as line 82 does for mbridge -- once #24 lands.

Note this does not address `window_size_left`, which is a separate TE 2.17 vs
flash-attn skew inside TE's own context-parallel path.
TE 2.17 did not only hoist m_splits out of non_tensor_args into its own
positional parameter -- it also retyped it. _GroupedLinear.forward annotates
it torch.Tensor and calls m_splits.tolist() on it, so passing the Python list
raises AttributeError: 'list' object has no attribute 'tolist'.

TE's public wrapper normalises the argument (GroupedLinear.forward: list ->
torch.int64 CPU tensor) before invoking the autograd function. This shim calls
the autograd function directly and so has to normalise it the same way, with
the same dtype and device.

Verified against the installed TE 2.17 by calling _GroupedLinear.forward with
both shapes: the list reproduces the AttributeError, the CPU int64 tensor
returns the expected (64, 512) output.

Signed-off-by: yueming-yuan <yym022502@gmail.com>
@yueming-yuan
yueming-yuan merged commit 7f0fb34 into bridge Jul 27, 2026
3 checks passed
@yueming-yuan
yueming-yuan deleted the radixark/bridge/te-grouped-linear-2.17 branch July 27, 2026 04:33
yueming-yuan added a commit to radixark/miles that referenced this pull request Jul 27, 2026
radixark/Megatron-Bridge#24 landed on @bridge as 7f0fb345 (squash). Its
src/megatron/bridge/peft/utils.py is byte-identical to 36d20288, the revision
this PR already exercised, so the pin moves off the PR branch without changing
what runs.

Kept as a SHA rather than @bridge. Line 82 pins mbridge the same way, and the
alternative is the failure this PR already hit once: buildkit keys the layer on
the instruction text, so a moving branch silently reuses the cached layer and
ships the previous revision.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant