[PyTorch] 2.21 regression: NVFP4 sequence-parallel all-gather of 3-D grad_output with B=1 crashes in swap_first_dims ("Input is not allocated"); please backport fix to release_v2.21

Read _all_gather_nvfp4 in transformer_engine/pytorch/distributed.py and compare its columnwise shape calculation with NVFP4Quantizer::create_tensor and the get_columnwise_shape change in #3628. Start by reproducing the two-rank sequence-parallel case with a [S, 1, H] input. Done means backward completes and the gathered columnwise shape is correct; add the regression test described in the issue and verify it fails before and passes after the fix.

Describe the bug

In TE 2.21.0, backward through a row-parallel te.Linear with sequence_parallel=True crashes under the NVFP4 recipe. It happens when grad_output is 3-D [S/TP, B, H] with B == 1, which is the usual Megatron [seq, micro_batch, hidden] layout at micro-batch size 1:

File ".../transformer_engine/pytorch/module/linear.py", line 1259, in _linear_backward_impl
    ) = TransformerEngineBaseModule.grad_output_preprocess(
File ".../transformer_engine/pytorch/module/base.py", line 1803, in grad_output_preprocess
    grad_output, _ = gather_along_first_dim(
File ".../transformer_engine/pytorch/distributed.py", line 1795, in gather_along_first_dim
    return _all_gather_nvfp4(
File ".../transformer_engine/pytorch/distributed.py", line 1536, in _all_gather_nvfp4
    _post_process_nvfp4_gather(out, out_columnwise_data, out_scale_inv, world_size, handle)
File ".../transformer_engine/pytorch/distributed.py", line 1313, in _post_process_nvfp4_gather
    out._columnwise_scale_inv = _swap_first_dims(columnwise_scale_inv_interleaved, world_size)
File ".../transformer_engine/pytorch/distributed.py", line 1296, in _swap_first_dims
    tensor = tex.swap_first_dims(tensor, out=None)
RuntimeError: .../transformer_engine/common/transpose/swap_first_dims.cu:130 in function swap_first_dims: Assertion failed: input.data.dptr != nullptr. Input is not allocated.

The same training configuration works on TE 2.20.2 (6ea2a74a).

Root cause (source at d77ef88c, the current release_v2.21 HEAD)

1. #3523 (ac544de1, "Remove the view() calls from Linear/LNLinear/LNMLP modules") removed grad_output = grad_output.reshape((-1, grad_output.shape[-1])) from TransformerEngineBaseModule.grad_output_preprocess. The SP gather now receives the 3-D gradient.
2. _all_gather_nvfp4 has not changed since 2.20. It computes the columnwise shape for a high-precision input with in_shape_t = convert_shape_for_fp4(get_columnwise_shape(inp.size())).
   • On 2.21, get_columnwise_shape keeps the input rank ([S, B, H] -> [H, S, B]), and convert_shape_for_fp4 halves the last dim.
   • For B == 1 this gives [H, S, 0], so unpadded_dim1 = 0 and the gathered columnwise scale-inv is a zero-element tensor with a null data pointer. tex.swap_first_dims then asserts.
   • For odd B > 1 the scale width is silently wrong. For (4096, 3, 8192) it is 512 instead of 768.
3. The C++ NVFP4Quantizer::create_tensor already forces this shape to 2-D for exactly this case. The bug is only in the Python all-gather shape bookkeeping.

Fix status

#3628 (f0b10961, merged to main on 2026-10-20) changes NVFP4Quantizer.get_columnwise_shape to return (shape[-1], prod(shape[:-1])). For the _all_gather_nvfp4 call this gives the same shapes as the 2-D flatten tested below. It was motivated by torch.compile (#3626), though, and is not on release_v2.21. Please backport a fix for this path to 2.21. Either cherry-pick the get_columnwise_shape change from #3628 or apply this minimal one-hunk patch:

--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -1365,8 +1365,12 @@ def _all_gather_nvfp4(
     if isinstance(inp, torch.Tensor) and not isinstance(inp, NVFP4TensorStorage):
         # High-precision tensor.
         in_shape = NVFP4Quantizer.convert_shape_for_fp4(inp.size())
+        # Columnwise NVFP4 data is always 2D [K, M/2] with M = prod(leading dims),
+        # matching NVFP4Quantizer::create_tensor. Using the N-D shape here would
+        # give [K, d0, ..., d_{n-2}//2], which is zero-width for e.g. [S, 1, H].
+        flat_shape = (math.prod(inp.shape[:-1]), inp.shape[-1])
         in_shape_t = NVFP4Quantizer.convert_shape_for_fp4(
-            NVFP4Quantizer.get_columnwise_shape(inp.size())
+            NVFP4Quantizer.get_columnwise_shape(flat_shape)
         )

A regression test with te.Linear(parallel_mode="row", sequence_parallel=True), NVFP4, and a [S, 1, H] input would cover this path.

Validation of the one-hunk patch on 2.21.0 (d77ef88c)

I applied the patch to the installed 2.21.0 package in an otherwise unchanged container. That is the same image that crashes without the patch. Results:

• Llama 3 70B (Megatron-Bridge NVFP4 pretraining: TP=2, PP=4, VP=5, SP, micro-batch 1, seq 8192, 64x GB200): the crash is gone. All 30 steps completed with no skipped or NaN iterations. Steady-state step time was 8.50 s, versus 8.52 s on TE 2.20.2. The loss at step 1 is identical to 2.20.2 (13.402) and stays close afterwards.
• Nemotron-3 Super (NVFP4, TP=2, SP, micro-batch 1, 64x GB200): the crash is gone and all 30 steps completed.
• Caveat: both runs then hit a separate segfault at process exit. It occurs in cudnn-frontend 1.31.0 while TE's fused-attention GraphCache is being destroyed. That is the known teardown issue NVIDIA/cudnn-frontend#1486 and is unrelated to this bug.

Steps/Code to reproduce bug

The full-scale reproduction is the Megatron-Bridge Llama 3 70B NVFP4 pretraining run above. TP=1 NVFP4 runs on the same container pass, because they have no SP gather.

Minimal sketch, run with torchrun --nproc_per_node=2. It has not been executed standalone:

import torch, torch.distributed as dist
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

dist.init_process_group("nccl")
torch.cuda.set_device(dist.get_rank())
tp = dist.new_group([0, 1])
S, B, H = 4096, 1, 8192
lin = te.Linear(H, H, parallel_mode="row", sequence_parallel=True, tp_group=tp, tp_size=2,
                params_dtype=torch.bfloat16, device="cuda")
x = torch.randn(S, B, H // 2, dtype=torch.bfloat16, device="cuda", requires_grad=True)
with te.autocast(recipe=recipe.NVFP4BlockScaling()):
    y = lin(x)
y.sum().backward()   # RuntimeError: swap_first_dims ... Input is not allocated

Expected behavior

Backward of an SP row-parallel NVFP4 Linear should work for any leading-dimension layout, as it does in 2.20.2.

Actual behavior

On TE 2.21.0 every rank fails in the first backward pass with RuntimeError: ... swap_first_dims.cu:130 in function swap_first_dims: Assertion failed: input.data.dptr != nullptr. Input is not allocated. (stack above). The run aborts before the first optimizer step.

Environment overview

• Environment location: Docker (NVIDIA PyTorch-based container), Slurm multi-node
• Method of Transformer Engine install: preinstalled in the NVIDIA PyTorch container

Environment details

• OS version: Ubuntu 24.04
• PyTorch version: 2.15.0a0 (NVIDIA internal nightly)
• Python version: 3.12.3
• Transformer Engine version: 2.21.0+d77ef88c (bad); 2.20.2+6ea2a74a (good)
• CUDA version: 13.5
• CUDNN version: 9.28.0

Device details

• GPU model: GB200 NVL (driver 580.196)

Additional context

• Regression bracket: last good 2.20.2 (6ea2a74a), first bad 2.21.0 (d77ef88c).
• ac544de1 (#3523) is the only change between those versions that alters the shape reaching _all_gather_nvfp4. _all_gather_nvfp4 and convert_shape_for_fp4 are unchanged.
• Related: #3626 / #3628.

---

This issue was drafted with assistance from the opus AI model.