fix(peft): follow TE's _GroupedLinear contract instead of pinning one layout - #24
Merged
Merged
Conversation
… 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
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
PEFT._forward_te_grouped_linearcallstransformer_engine's private_GroupedLineardirectly. It was written against TE ≤ 2.14, and TE has moved thatcontract twice since:
non_tensor_argsforward(ctx, inp, non_tensor_args, *weights_and_biases)m_splitsfirst,modulenear the endTensormodule→weight_workspaces+cache_weightTensorforward(ctx, inp, **m_splits**, non_tensor_args, *weights_and_biases)m_splitspromoted 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_argstuple binds to them_splitsparameter andweight[0]binds to
non_tensor_args. TE then unpacks that weight tensor, iterating itsfirst dimension, and the failure surfaces far from its cause:
16is the weight tensor's dim-0, not any tuple length — which is what makes thisone 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 ism_splitspositional? frominspect.signature, and the tuple field names from the= non_tensor_argsunpack via
inspect.getsource.these releases asks for, and raises a named
RuntimeErrorlisting the unknownfield if TE adds one.
Tensorand(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_weightare passed empty/False: this shim alwayspasses
is_first_microbatch=None, so no fp8 workspace caching happens on either theold (
module) or new (explicit workspaces) contract.Verification
Checked offline against the 2.14.0 / 2.15.0 / 2.16.0 / 2.17.0 wheels:
and the supplied field dict covers every field in all four (no gaps). Repo
pre-commitpasses.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.