Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

[PyTorch] Request a migration path for downstream `_GroupedLinear` integrations after the 2.17 signature change

Aperta
  • Qualcosa non va?
#3,191 2 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
5/5
Tempo stimato
Più di una settimana
Idoneità per principianti
35/100
Tipo di issue
Bug
Chiarezza
Abbastanza chiara
Stato di attività
Tranquilla
Stack tecnologico
python, pytorch
Ambito
backend, machine-learning

Direzione di ricerca

Inizia da transformer_engine/pytorch/module/grouped_linear.py ed esamina la firma di _GroupedLinear.forward e il dispatch di linear_fn intorno al commit f8bda5d0. Esegui il target pytest di Megatron-Bridge indicato per riprodurre il fallimento len(False), quindi esamina gli esempi di compatibilità collegati di ModelOpt e Megatron-Bridge. Il lavoro è completo quando il comportamento di migrazione o transizione supportato è documentato e coperto da un test di regressione per le chiamate a wrapped linear_fn.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

Describe the bug

Transformer Engine 2.17 intentionally changed the private PyTorch _GroupedLinear autograd call layout while making GroupedLinear graph-safe.

TE 2.16:

forward(ctx, inp, non_tensor_args, *weights_and_biases)
# non_tensor_args[0] is m_splits

TE 2.17:

forward(ctx, inp, m_splits, non_tensor_args, *weights_and_biases)
# m_splits is now a separate int64 tensor
# non_tensor_args[0] is use_bias

The change was introduced by f8bda5d0. The public GroupedLinear.forward accepts both list and tensor split inputs, while the internal linear_fn dispatch boundary changed positionally.

We understand that _GroupedLinear is private and that downstream users cannot assume its signature will remain stable. In practice, however, there does not appear to be a public functional or interception API for these use cases, so some downstream integrations currently use this boundary to intercept grouped GEMMs or invoke them with externally owned weights. Upgrading TE alone therefore causes deterministic runtime failures in those integrations.

Two downstream examples are affected:

  1. NVIDIA ModelOpt intercepts the grouped-linear function to quantize its input and weights. Its compatibility logic finds non_tensor_args in the signature and reads non_tensor_args[0] as m_splits. With TE 2.17, that value is use_bias, so calibration fails with:

    File ".../transformer_engine/pytorch/module/grouped_linear.py", line 1788, in forward
    out, new_workspaces = linear_fn(
File ".../modelopt/torch/quantization/plugins/transformer_engine.py", line 178, in te_grouped_quantized_linear_fn
    num_gemms = len(args[sig_params.index("non_tensor_args") - ctx_offset][0])
TypeError: object of type 'bool' has no len()

    Tracking issue: NVIDIA/Model-Optimizer#1940

  2. Megatron-Bridge directly calls _GroupedLinear.apply/forward with externally owned grouped adapter weights. The TE 2.16 layout passes x, non_tensor_args, weights..., biases.... Under TE 2.17, the same call is shifted: the old tuple is interpreted as m_splits, the first weight is interpreted as non_tensor_args, and the remaining weights and biases are mispartitioned. The downstream compatibility fix is NVIDIA-NeMo/Megatron-Bridge#4721.

Steps/Code to reproduce bug

The failure is reproduced by changing only the TE pin in Megatron-Bridge PR #4696:

  • Passing baseline: 2.16.0+d64bc14d at d64bc14dc87eb658ab98839e4b7687595ee53e2d
  • Failing version: 2.17.0+2e559f06 at 2e559f062497bef768dfbe9d7e45548fadeca80a
  • ModelOpt remains pinned at nvidia-modelopt==0.44.0rc5

Run:

uv run python -m pytest -s -x \
  tests/functional_tests/test_groups/quantization/models/qwen/test_qwen3_moe_quantization_workflow.py::TestQwen3MoeQuantizationWorkflow::test_qwen3_moe_quantization_and_generation_with_expert_parallelism

The first calibration forward fails at the len(False) exception shown above. The same failure occurs on both GB200 and H100 runners. Full failing job: Megatron-Bridge GitHub Actions.

The signature transition can also be confirmed directly:

import inspect
from transformer_engine.pytorch.module.grouped_linear import _GroupedLinear

print(inspect.signature(_GroupedLinear.forward))

Requested guidance / compatibility support

We will update the affected downstream integrations to handle the new layout. To make that migration less disruptive, would the TE maintainers consider publishing a TE 2.17.x patch release with a short-lived transition path for both grouped-linear call layouts?

# Legacy layout used through TE 2.16
(ctx, inp, non_tensor_args, *weights_and_biases)

# Graph-safe layout introduced in TE 2.17
(ctx, inp, m_splits, non_tensor_args, *weights_and_biases)

If practical, the compatibility path could cover both grad-enabled _GroupedLinear.apply and direct no-grad _GroupedLinear.forward dispatch, preserve the correct backward arity, and emit a deprecation warning for the legacy form before it is removed in a later feature release. We are not asking for indefinite compatibility for the private API.

Because wrappers such as ModelOpt intercept the linear_fn boundary before _GroupedLinear.forward executes, compatibility may need to be provided at that dispatch/hook boundary rather than only inside the autograd function. A regression test using a wrapped linear_fn would cover this downstream use case.

If dual-layout support at this private boundary is not feasible, guidance on the intended supported approach—or an equivalent stable functional/interception API—would be equally helpful. Public GroupedLinear owns its parameters, so it is not currently a drop-in replacement for callers that need to supply externally owned grouped weights.

Environment overview (please complete the following information)

  • Environment location: GitHub Actions, Docker, GCP GPU runners
  • Transformer Engine install: Git source pin resolved by uv
  • Passing TE: 2.16.0+d64bc14d
  • Failing TE: 2.17.0+2e559f06
  • Downstream packages: Megatron-Bridge PR #4696 and nvidia-modelopt==0.44.0rc5

Environment details

  • OS: Linux container
  • Python: 3.12.3
  • Transformer Engine: versions and exact commits listed above
  • The full container and runner information is recorded in the linked GitHub Actions job.

Device details

  • Reproduced on GB200 and H100 CI runners.

Additional context

The old and new contracts are straightforward for downstream code to distinguish when it owns the call: prefer the explicit m_splits parameter when present, otherwise use non_tensor_args[0]. The harder compatibility case is an interceptor that TE calls through the changed positional boundary, which is why a short TE-side transition window would be valuable.

This report concerns only the grouped-linear signature transition. It is independent of other TE 2.17 packaging or import behavior changes.

Lingua principale
Python
Stelle
3.6k
Fork
851
Merge medio
4g 20h
PR unite (30g)
58

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/TransformerEngine

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.