Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

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

Abierto
  • ¿Algo no está bien?
#3,191 2 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
5/5
Tiempo estimado
Más de una semana
Aptitud para principiantes
35/100
Tipo de issue
Error
Claridad
Bastante claro
Estado de actividad
Tranquilo
Stack tecnológico
python, pytorch
Área
backend, machine-learning

Línea de trabajo

Empieza por transformer_engine/pytorch/module/grouped_linear.py e inspecciona la firma de _GroupedLinear.forward y el dispatch de linear_fn alrededor del commit f8bda5d0. Ejecuta el objetivo de pytest de Megatron-Bridge indicado para reproducir el fallo len(False) y, después, examina los ejemplos de compatibilidad enlazados de ModelOpt y Megatron-Bridge. Se considera terminado cuando el comportamiento de migración o transición compatible esté documentado y cubierto por una prueba de regresión para llamadas a wrapped linear_fn.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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.

Lenguaje dominante
Python
Estrellas
3.6k
Forks
851
Merge medio
4 d 20 h
PR fusionados (30 d)
58

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.