Optimize fp8 block scaling Allgather for FSDP2

[vthumbe1503]

Description

Eliminate Columnwise allgather for fp8_model_init with fsdp2. For weights when FP8 blockscaling is used, we typically use 2d. And in such a case, columnwise data and scale inv is just the transpose of the rowwise data and scale inverse. And so allgathering the rowwise data/scales are enough

Fixes # (issue)

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

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

Greptile Summary

This PR introduces a communication-halving optimization for FP8 block-scaled weights under PyTorch FSDP2: instead of all-gathering both rowwise and columnwise tensors, only rowwise data and scales are all-gathered, and the columnwise view is derived locally via fp8_transpose (a cheap on-device kernel) in fsdp_post_all_gather. This is valid because 2D block scaling (128×128 tiles) makes the columnwise tensor a mathematical transpose of the rowwise tensor, so all_gather(shards) → transpose gives the same result as transpose(shards) → all_gather. The PR also applies a consistent defensive improvement across all three tensor types (Float8BlockwiseQTensor, Float8Tensor, MXFP8Tensor): caching fsdp_state._fsdp_param_group into a local variable and raising a descriptive RuntimeError when it is None, instead of letting an AttributeError propagate.

Key changes:

• float8_blockwise_tensor.py: fsdp_pre_all_gather now reads _training_state to determine whether the call is for a forward or backward pass (when reshard_after_forward=True), encodes rowwise_usage/columnwise_usage in metadata, and always ships only (rowwise_data, rowwise_scale_inv). fsdp_post_all_gather unconditionally extracts these two tensors, then calls _create_columnwise() only when columnwise_usage=True, followed by update_usage() to release whichever form is not needed.
• float8_tensor.py / mxfp8_tensor.py: Guard added for _fsdp_param_group is None; param_group._training_state access deduped through the local variable.
• No test coverage was added for the new Float8BlockwiseQTensor FSDP2 path (checklist unchecked).

Confidence Score: 4/5

• PR is safe to merge; optimization is mathematically correct and prior review concerns have been addressed.
• The core optimization — all-gathering only rowwise tensors and deriving columnwise via local fp8_transpose — is mathematically sound for 2D block scaling. The forward/backward pass detection via TrainingState.PRE_BACKWARD mirrors the existing pattern in float8_tensor.py and mxfp8_tensor.py. The buffer-reuse path in _create_columnwise (passing out=self._columnwise_data) correctly avoids per-iteration GPU allocations on subsequent iterations. Prior concerns from the review thread (param_group None-guard, stale columnwise fields, allocation reuse) have all been addressed. The only remaining items are two minor P2 style notes (redundant _create_columnwise call path and a trailing semicolon in a comment) plus the absence of a dedicated test for the new blockwise FSDP2 optimization.
• transformer_engine/pytorch/tensor/float8_blockwise_tensor.py — contains the main optimization; fsdp_pre_all_gather and fsdp_post_all_gather logic should be validated with an integration test covering both reshard_after_forward=True and reshard_after_forward=False.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	Core optimization: eliminates columnwise all-gather for 2D block-scaled weights by all-gathering only rowwise data and deriving columnwise locally via fp8_transpose in post_all_gather. Logic is mathematically sound (transpose distributes over all-gather for dim-0 sharding). Minor: explicit _create_columnwise() call before update_usage is slightly redundant with update_usage's own fallback; trailing semicolon in a comment.
transformer_engine/pytorch/tensor/float8_tensor.py	Defensive improvement only: caches fsdp_state._fsdp_param_group in a local variable and adds a None-guard with a descriptive RuntimeError, matching the new pattern in float8_blockwise_tensor.py. No functional logic change.
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Same defensive improvement as float8_tensor.py: None-guard added for _fsdp_param_group with a clear RuntimeError. No algorithmic changes.

Sequence Diagram

sequenceDiagram
    participant FSDP2
    participant fsdp_pre_all_gather
    participant AllGather
    participant fsdp_post_all_gather
    participant _create_columnwise

    FSDP2->>fsdp_pre_all_gather: call (module, mesh, …)
    fsdp_pre_all_gather->>fsdp_pre_all_gather: check _fsdp_param_group != None
    fsdp_pre_all_gather->>fsdp_pre_all_gather: read reshard_after_forward & training_state
    alt reshard_after_forward=True, forward pass
        fsdp_pre_all_gather-->>FSDP2: sharded=(rowwise_data, rowwise_scale_inv), usage=(row=T, col=F)
    else reshard_after_forward=True, PRE_BACKWARD
        fsdp_pre_all_gather-->>FSDP2: sharded=(rowwise_data, rowwise_scale_inv), usage=(row=F, col=T)
    else reshard_after_forward=False
        fsdp_pre_all_gather-->>FSDP2: sharded=(rowwise_data, rowwise_scale_inv), usage=(row=T, col=quantizer.cw)
    end

    FSDP2->>AllGather: all-gather rowwise_data + rowwise_scale_inv
    AllGather-->>FSDP2: full rowwise_data, full rowwise_scale_inv

    FSDP2->>fsdp_post_all_gather: call (all_gather_outputs, metadata, param_dtype, out)
    fsdp_post_all_gather->>fsdp_post_all_gather: extract rowwise_data, rowwise_scale_inv
    alt out is None (first iteration)
        fsdp_post_all_gather->>fsdp_post_all_gather: construct Float8BlockwiseQTensor (cw=None)
    else out is not None (subsequent iterations)
        fsdp_post_all_gather->>fsdp_post_all_gather: update out._rowwise_data / scale_inv in-place
    end
    alt columnwise_usage=True
        fsdp_post_all_gather->>_create_columnwise: derive columnwise via fp8_transpose (reuse buffer)
        _create_columnwise-->>fsdp_post_all_gather: _columnwise_data, _columnwise_scale_inv set
    end
    fsdp_post_all_gather->>fsdp_post_all_gather: update_usage(row, col) — clears unused form
    fsdp_post_all_gather-->>FSDP2: (Float8BlockwiseQTensor, all_gather_outputs)

Loading

_{Reviews (5): Last reviewed commit: "Merge branch 'main' into optimize_fp8_bl..." | Re-trigger Greptile}

[greptile-apps]

Unguarded access to _fsdp_param_group

fsdp_state._fsdp_param_group is typed as Optional[FSDPParamGroup] in PyTorch's FSDP2 internals — it is None for any FSDP module that does not directly manage parameters (e.g. a container module whose children are individually sharded). Accessing ._reshard_after_forward on it unconditionally will raise AttributeError: 'NoneType' object has no attribute '_reshard_after_forward' in that case.

While in practice fsdp_pre_all_gather is only called for tensors managed by a param group, this assumption is implicit. A guard makes the failure mode explicit and easier to diagnose:

fsdp_state = _get_module_fsdp_state(module)
param_group = fsdp_state._fsdp_param_group
if param_group is None:
    raise RuntimeError(
        "FSDP state for this module has no parameter group; "
        "cannot determine reshard_after_forward."
    )
reshard_after_forward = param_group._reshard_after_forward

[vthumbe1503]

/te-ci L1 pytorch

[greptile-apps]

Inconsistent import style for TrainingState

TrainingState is imported at the module level (line 10) in float8_tensor.py and at line 13 in mxfp8_tensor.py, but here it's imported lazily inside fsdp_pre_all_gather. While the inline comment about the private API and PyTorch version is valuable, the inconsistency across the three sibling files may confuse readers.

Consider either:

• Moving the TrainingState import to the module level and placing the version comment there (matching the other two files), or
• Adding the same lazy-import pattern and version comment to float8_tensor.py and mxfp8_tensor.py for symmetry.

Suggested change

	# PyTorch FSDP2 private API – tested with PyTorch 2.5+;
	from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
	# PyTorch FSDP2 private API – tested with PyTorch 2.5+;
	from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
	from transformer_engine.pytorch.distributed import _get_module_fsdp_state

_{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

[greptile-apps]

columnwise_usage not derived from training state in non-resharded path

When reshard_after_forward=False, the same all-gathered weight is reused through both forward and backward passes. The code sets:

rowwise_usage = True
columnwise_usage = self._quantizer.columnwise_usage

This means whether columnwise data gets derived locally (and kept) is entirely controlled by the sharded quantizer's setting, not the actual pass. The comment in the previous code explicitly noted that both forms were needed when not resharding. If self._quantizer.columnwise_usage is False (e.g. on an architecture that doesn't need the transpose), columnwise data won't be created and won't be available for the backward pass GEMM.

This matches the pre-existing float8_tensor.py behavior (same pattern there), so it's presumably already validated by the existing usage assumptions — but it would be worth a brief comment here documenting that self._quantizer.columnwise_usage must be True whenever the backward GEMM needs columnwise access for the non-resharding path.

[jomitchellnv]

Btw I don't see any test files that were updated. I'd expect a test under tests/pytorch/fsdp/ or similar validating that the locally-derived columnwise output matches the old all-gathered columnwise output.

[jomitchellnv]

LGTM just hoping theres some test coverage around this new implementation. i think i wrote some last time but not sure

[vthumbe1503]

/te-ci L1 pytorch

[ksivaman]

LGTM, CI pending, consistent with other recipes