[rapatel]

Description

cast_master_weights_to_fp8 / quantize_master_weights with MXFP8 all-reduced every block's bf16 partial amax (2 bytes) and then turned the result into an E8M0 scale. It now computes each rank's E8M0 scales from its partial amaxes and all-reduces those (1 byte) with MAX. The scales are bit-identical and the all-reduce payload halves.

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:

In transformer_engine/pytorch/tensor/utils.py, _cast_master_weights_to_fp8_mxfp8_scaling:

• A packed uint8 buffer packed_scale_invs sits beside packed_amaxes, with per-weight rowwise and columnwise views at the same offsets.
• Step 2 runs multi_tensor_compute_scale_inv_e8m0 on the local amaxes into the packed scale views, then all-reduces packed_scale_invs with MAX. It is still one all-reduce for all weights.
• Step 3 copies the reduced scales into each weight's _rowwise_scale_inv / _columnwise_scale_inv with one torch._foreach_copy_.
• Steps 1 (partial amax) and 4 (partial cast) are unchanged. Weights without columnwise data are skipped in both lists, as before.

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

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge.

Summary

This PR changes MXFP8 master-weight updates to reduce one-byte E8M0 scales instead of two-byte BF16 maxima.

• MXFP8 master-weight casting reduces one-byte scales across ranks.

No actionable issues were found. GPU tests were not run.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Local master-weight shard] --> B[Local block maxima]
  B --> C[Packed E8M0 scales]
  C --> D[MAX all-reduce across ranks]
  D --> E[Copy scales into model weights]
  E --> F[Cast each local shard to FP8]

Loading

_{Reviews (1) · Last reviewed commit: "[PyTorch] All-reduce MXFP8 E8M0 scales i..." · Reviewed by Greptile}