Skip to content

[PyTorch] Fix MXFP8 master weight cast when a rank owns no shard. - #3348

Open
rapatel wants to merge 2 commits into
NVIDIA:mainfrom
rapatel:ripatel/handle_empty_shards
Open

[PyTorch] Fix MXFP8 master weight cast when a rank owns no shard.#3348
rapatel wants to merge 2 commits into
NVIDIA:mainfrom
rapatel:ripatel/handle_empty_shards

Conversation

@rapatel

@rapatel rapatel commented Aug 11, 2026

Copy link
Copy Markdown

Description

Fix MXFP8 master weight cast when a rank owns no shard.

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:

  • Fix MXFP8 master weight cast when a rank owns no shard.
  • Adds tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py.

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: Ritesh Patel <ripatel@nvidia.com>
@rapatel
rapatel requested a review from ksivaman as a code owner August 11, 2026 20:38
@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Aug 11, 2026
@greptile-apps

greptile-apps Bot commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes MXFP8 master-weight casting when a rank owns no shard by deriving the fallback amax-buffer dtype from the quantized model weight. It also adds regression coverage for empty and populated shards and safely tears down a process group created by the test fixture.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/pytorch/tensor/utils.py Initializes the master-weight dtype from the model weight so empty-shard ranks participate in amax reduction with the same dtype as populated ranks.
tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py Adds empty-shard regression coverage and conditionally destroys only the distributed process group created by the fixture.

Reviews (2): Last reviewed commit: "[PyTorch] Destroy the process group crea..." | Re-trigger Greptile

Comment thread tests/pytorch/mxfp8/test_mxfp8_master_weight_empty_shard.py
@negvet negvet self-assigned this Aug 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants