Fix FP8 block scaling with sequence parallel by cuichenx · Pull Request #2637 · NVIDIA/TransformerEngine

cuichenx · 2026-01-31T00:24:50Z

Description

Problem

Using Float8BlockQuantizer with sequence parallel fails with AssertionError: All-gather requires quantizable tensor for quantizer Float8BlockQuantizer when local tensor dimensions aren't divisible by 128.

Solution

Skip the assert_dim_for_all_gather check for Float8BlockQuantizer since gather_along_first_dim already has a fallback path
Fix the fallback in _start_all_gather_fp8_blockwise to handle already-quantized inputs by dequantizing before high-precision all-gather

###Note
The fallback path (high-precision all-gather → quantize) may increase the communication overhead.

Verification

The code change does not alter convergence behavior

When SP is True, the previous code did not run. When SP is False, this change doesn't affect anything.

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Change A
Change B

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: Chen Cui <chcui@nvidia.com>

greptile-apps · 2026-01-31T00:28:20Z

Greptile Overview

Greptile Summary

This PR fixes an AssertionError when using Float8BlockQuantizer with sequence parallel on tensors whose dimensions aren't divisible by 128 (the block length).

Key Changes:

Removed assert_dim_for_all_gather checks from layernorm_linear, layernorm_mlp, and linear modules, since the fallback path can handle non-quantizable dimensions
Enhanced fallback paths in _start_all_gather_fp8_blockwise, _all_gather_nvfp4, and _all_gather_mxfp8 to dequantize already-quantized inputs before performing high-precision all-gather
Fixed dequantize() calls to pass explicit dtype parameter for consistency
Added warning messages when falling back to high-precision all-gather

Issues Found:

Three new fallback paths are missing .contiguous() calls before all_gather_into_tensor, which can cause runtime errors with non-contiguous tensors (e.g., after transpose or slicing). Other all-gather paths in the codebase consistently use .contiguous().

Confidence Score: 3/5

This PR is mostly safe but has a critical bug in the fallback paths that needs to be fixed
The core logic correctly removes unnecessary assertions and adds proper fallback handling for non-quantizable dimensions. However, three new all-gather fallback paths are missing .contiguous() calls, which can cause runtime errors with non-contiguous tensors. This is a well-understood pattern violation (other paths use .contiguous()), making it straightforward to fix.
Pay close attention to transformer_engine/pytorch/distributed.py - the three missing .contiguous() calls must be added before merging

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Adds fallback path for FP8 block quantization when dimensions aren't divisible by 128. Fixed dequantize calls to pass dtype parameter. Missing .contiguous() calls in three all-gather fallback paths.
transformer_engine/pytorch/module/layernorm_linear.py	Removed assert_dim_for_all_gather check for Float8BlockQuantizer, allowing fallback path to handle non-quantizable dimensions
transformer_engine/pytorch/module/layernorm_mlp.py	Removed assert_dim_for_all_gather check for Float8BlockQuantizer, allowing fallback path to handle non-quantizable dimensions
transformer_engine/pytorch/module/linear.py	Removed assert_dim_for_all_gather check for Float8BlockQuantizer, allowing fallback path to handle non-quantizable dimensions
transformer_engine/pytorch/utils.py	Removed assert_dim_for_all_gather function that was checking quantizer compatibility - no longer needed with fallback paths

Sequence Diagram

sequenceDiagram
    participant Module as LayerNorm/Linear Module
    participant Utils as utils.py
    participant Gather as gather_along_first_dim
    participant FP8Block as _start_all_gather_fp8_blockwise
    participant Quantizer as Float8BlockQuantizer
    participant NCCL as torch.distributed

    Module->>Module: Remove assert_dim_for_all_gather check
    Module->>Gather: Call with sequence_parallel=True
    Gather->>FP8Block: Route to FP8 blockwise handler
    FP8Block->>Quantizer: is_quantizable(inp)?
    
    alt Tensor dimensions divisible by 128
        Quantizer-->>FP8Block: True
        FP8Block->>FP8Block: Quantize to FP8 blockwise
        FP8Block->>NCCL: all_gather FP8 data + scales
        NCCL-->>FP8Block: Gathered FP8 tensor
    else Dimensions not divisible by 128
        Quantizer-->>FP8Block: False
        FP8Block->>FP8Block: Check if already quantized
        alt Already quantized
            FP8Block->>FP8Block: Dequantize to high precision
        end
        FP8Block->>NCCL: all_gather in high precision
        NCCL-->>FP8Block: Gathered high-precision tensor
        FP8Block->>Quantizer: Quantize output
        Quantizer-->>FP8Block: FP8 blockwise tensor
    end
    
    FP8Block-->>Module: Return gathered tensor

transformer_engine/pytorch/utils.py

Perform all-gather in high-precision if the input tensor is too small to quantize. Signed-off-by: Tim Moon <tmoon@nvidia.com>

greptile-apps

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

timmoon10

I've tweaked the module logic so that it no longer checks that the local tensor can be quantized. In addition to FP8 block scaling, MXFP8 and NVFP4 will also support high-precision all-gathers. That said, this may not be fully robust for LayerNormLinear and LayerNormMLP, since the norm kernel may still expect to output to a quantized local tensor.

timmoon10 · 2026-02-03T20:20:44Z

/te-ci pytorch L1

greptile-apps

_{2 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-03T20:25:23Z

transformer_engine/pytorch/distributed.py

-        out = torch.empty(out_shape, dtype=dtype, device=device)
+        warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.")
+        if isinstance(inp, QuantizedTensorStorage):
+            inp = inp.dequantize()  # Dequantize if needed


The dequantize() method defaults to dtype=torch.float32. Consider whether this is always appropriate for the fallback path, especially when the original tensor might have been in a different precision (e.g., bfloat16).

greptile-apps · 2026-02-03T20:25:25Z

transformer_engine/pytorch/distributed.py

    ):
+        warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.")
+        if isinstance(inp, QuantizedTensorStorage):
+            inp = inp.dequantize()  # Dequantize if needed


Same consideration as in _start_all_gather_fp8_blockwise: the dequantize() method defaults to dtype=torch.float32, which may not match the original tensor's precision.

greptile-apps · 2026-02-03T20:25:26Z

transformer_engine/pytorch/distributed.py

    ):
+        warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.")
+        if isinstance(inp, QuantizedTensorStorage):
+            inp = inp.dequantize()  # Dequantize if needed


Same consideration as in _start_all_gather_fp8_blockwise: the dequantize() method defaults to dtype=torch.float32, which may not match the original tensor's precision.

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

greptile-apps

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

ptrendx · 2026-02-10T23:02:59Z

/te-ci pytorch L1

greptile-apps

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-10T23:10:25Z

transformer_engine/pytorch/distributed.py

    if not quantizer.is_quantizable(inp) or quantizer.block_scaling_dim != 1:
+        warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.")
+        if isinstance(inp, QuantizedTensorStorage):
+            inp = inp.dequantize(dtype=dtype)  # Dequantize if needed
        out = torch.empty(out_shape, dtype=dtype, device=device)
        torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False)
        out = quantizer(out)


Non-contiguous gather input

In the new high-precision fallback (if not quantizer.is_quantizable(inp) ...), all_gather_into_tensor(out, inp, ...) passes inp directly. Elsewhere in this same module the plain-tensor path uses inp.contiguous() (distributed.py:1737-1742) and the FP8 path uses _data.contiguous() (distributed.py:1031-1035), which strongly suggests the collective expects contiguous inputs. If inp is a non-contiguous view (common after transpose/slicing), this fallback can raise at runtime. This same issue also appears in the NVFP4 and MXFP8 high-precision fallbacks (distributed.py:1353 and :1523).

ptrendx · 2026-02-12T19:20:00Z

/te-ci pytorch

greptile-apps

_{5 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

greptile-apps · 2026-02-12T19:23:06Z

transformer_engine/pytorch/distributed.py

+        if isinstance(inp, QuantizedTensorStorage):
+            inp = inp.dequantize(dtype=dtype)  # Dequantize if needed
        out = torch.empty(out_shape, dtype=dtype, device=device)
        torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False)


Missing .contiguous() call on inp before all-gather

Other all-gather paths in this file use .contiguous() (lines 1739, 1033). Non-contiguous tensors (from transpose/slicing) can cause runtime errors.

Suggested change

torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False)

torch.distributed.all_gather_into_tensor(out, inp.contiguous(), group=process_group, async_op=False)

greptile-apps · 2026-02-12T19:23:16Z

Additional Comments (2)

transformer_engine/pytorch/distributed.py
Missing .contiguous() call on inp before all-gather

        torch.distributed.all_gather_into_tensor(out, inp.contiguous(), group=process_group)

transformer_engine/pytorch/distributed.py
Missing .contiguous() call on inp before all-gather

        torch.distributed.all_gather_into_tensor(out, inp.contiguous(), group=process_group)

fix subchannel fp8 + sp

3ba991b

Signed-off-by: Chen Cui <chcui@nvidia.com>

This comment was marked as outdated.

Sign in to view

cyanguwa requested a review from timmoon10 February 2, 2026 18:48

Merge branch 'main' into chcui/fix_subchannel_fp8+sp

390b2e1

This comment was marked as outdated.

Sign in to view

timmoon10 self-requested a review February 2, 2026 19:38

timmoon10 requested changes Feb 2, 2026

View reviewed changes

transformer_engine/pytorch/utils.py Outdated Show resolved Hide resolved

Support sequence-parallel all-gather with small inputs

637ba0f

Perform all-gather in high-precision if the input tensor is too small to quantize. Signed-off-by: Tim Moon <tmoon@nvidia.com>

greptile-apps bot reviewed Feb 2, 2026

View reviewed changes

timmoon10 previously approved these changes Feb 2, 2026

View reviewed changes

This comment was marked as outdated.

Sign in to view

Merge branch 'main' into chcui/fix_subchannel_fp8+sp

9fe572b

greptile-apps bot reviewed Feb 3, 2026

View reviewed changes

Fix lint error

19fd927

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

ptrendx dismissed timmoon10’s stale review via 19fd927 February 10, 2026 22:21

greptile-apps bot reviewed Feb 10, 2026

View reviewed changes

Keep the previous behavior with dtype

adfe33b

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

greptile-apps bot reviewed Feb 10, 2026

View reviewed changes

ptrendx approved these changes Feb 10, 2026

View reviewed changes

Merge branch 'main' into chcui/fix_subchannel_fp8+sp

e6d1559

greptile-apps bot reviewed Feb 12, 2026

View reviewed changes

	torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False)
	torch.distributed.all_gather_into_tensor(out, inp.contiguous(), group=process_group, async_op=False)

Conversation

cuichenx commented Jan 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Problem

Solution

Verification

Type of change

Changes

Checklist:

Uh oh!

greptile-apps bot commented Jan 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Sequence Diagram

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

This comment was marked as outdated.

timmoon10 commented Feb 3, 2026

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

ptrendx commented Feb 10, 2026

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 10, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendx commented Feb 12, 2026

Uh oh!

greptile-apps bot left a comment

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot Feb 12, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps bot commented Feb 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cuichenx commented Jan 31, 2026 •

edited

Loading

greptile-apps bot commented Jan 31, 2026 •

edited

Loading