Add experimental Shardy support. by jreiffers · Pull Request #1642 · NVIDIA/TransformerEngine

jreiffers · 2025-04-03T13:44:35Z

This is experimental and not intended for production use. If you try this, please file issues for problems you encounter.

Description

Shardy is a new partitioning system in JAX. It currently replaces the sharding propagation passes in GSPMD.

This PR adds Shardy support for TE's attention, normalization, quantization and softmax primitives.

To use this, simply enable Shardy:

jax.config.update("jax_use_shardy_partitioner", True)

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:

Adds functions to provide Shardy sharding rules to primitives.

jreiffers · 2025-04-04T07:17:00Z

/te-ci jax L1

transformer_engine/jax/quantize/scaling_modes.py

transformer_engine/jax/cpp_extensions/normalization.py

tests/jax/test_distributed_layernorm_mlp.py

jreiffers · 2025-04-09T08:10:46Z

PTAL. If this looks OK in principle, I'll do the rebase.

jreiffers · 2025-04-10T09:34:20Z

/te-ci jax L1

Production use is not yet recommended. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

for more information, see https://pre-commit.ci Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

for more information, see https://pre-commit.ci

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

test_distributed_layernorm_mlp contains some cases that require functionality that is missing from JAX. This needs to be fixed there before these tests can be enabled again. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

jreiffers · 2025-04-10T11:56:54Z

/te-ci jax L1

transformer_engine/jax/cpp_extensions/activation.py

transformer_engine/jax/quantize/scaling_modes.py

transformer_engine/jax/cpp_extensions/activation.py

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

- Add support for variable flatten_axis values - Hopefully improve the names in BlockScalingModeMetadataImpl::get_shardy_sharding_rules a bit. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

for more information, see https://pre-commit.ci

jreiffers · 2025-04-10T13:50:58Z

/te-ci jax L1

phu0ngng

The rest looks good to me!

transformer_engine/jax/cpp_extensions/activation.py

transformer_engine/jax/cpp_extensions/normalization.py

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

phu0ngng · 2025-04-10T14:00:51Z

/te-ci jax L1

tests/jax/test_distributed_layernorm_mlp.py

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

jreiffers · 2025-04-11T08:14:36Z

/te-ci jax L1

examples/jax/encoder/test_model_parallel_encoder.py

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

phu0ngng · 2025-04-14T14:47:00Z

/te-ci jax L1

phu0ngng · 2025-04-14T17:20:22Z

Pipeline #26929185.

phu0ngng self-requested a review April 7, 2025 12:45

phu0ngng reviewed Apr 7, 2025

View reviewed changes

transformer_engine/jax/quantize/scaling_modes.py Outdated Show resolved Hide resolved

transformer_engine/jax/quantize/scaling_modes.py Outdated Show resolved Hide resolved

transformer_engine/jax/cpp_extensions/normalization.py Outdated Show resolved Hide resolved

phu0ngng reviewed Apr 7, 2025

View reviewed changes

transformer_engine/jax/cpp_extensions/normalization.py Outdated Show resolved Hide resolved

phu0ngng reviewed Apr 7, 2025

View reviewed changes

tests/jax/test_distributed_layernorm_mlp.py Outdated Show resolved Hide resolved

jreiffers force-pushed the shardyv3 branch from 93cc823 to a7c205f Compare April 9, 2025 08:08

jreiffers force-pushed the shardyv3 branch 2 times, most recently from 66fb2b8 to ee15c3f Compare April 9, 2025 12:47

jreiffers and others added 13 commits April 10, 2025 11:44

Add experimental Shardy support.

3742bf0

Production use is not yet recommended. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Fix bad rebase in FusedAttnFwdPrimitive.

f075884

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Fix Lint errors.

14796ea

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

01fc03f

for more information, see https://pre-commit.ci Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Prune shardy distributed fused attn test cases.

3a3292b

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Prune shardy distributed softmax test cases.

f5b84d6

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Address review comments.

7f4b44f

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

43efbc5

for more information, see https://pre-commit.ci

More rebase fixes.

c2cdd87

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

e09ccb1

for more information, see https://pre-commit.ci

More fixes.

8e455b7

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Remove unsupported Shardy test cases.

3b94073

test_distributed_layernorm_mlp contains some cases that require functionality that is missing from JAX. This needs to be fixed there before these tests can be enabled again. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Fix Shardy verifier error.

69d2312

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

jreiffers force-pushed the shardyv3 branch from 465c406 to 296667a Compare April 10, 2025 11:55

phu0ngng reviewed Apr 10, 2025

View reviewed changes

Rebase fixes for ScalingMode names.

70282fd

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

jreiffers force-pushed the shardyv3 branch from 296667a to 70282fd Compare April 10, 2025 12:49

jreiffers and others added 2 commits April 10, 2025 13:31

Rebase fixes, BlockScalingModeMetadataImpl naming.

dcc2c0c

- Add support for variable flatten_axis values - Hopefully improve the names in BlockScalingModeMetadataImpl::get_shardy_sharding_rules a bit. Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

ad2b79e

for more information, see https://pre-commit.ci

phu0ngng requested changes Apr 10, 2025

View reviewed changes

transformer_engine/jax/cpp_extensions/activation.py Outdated Show resolved Hide resolved

transformer_engine/jax/cpp_extensions/normalization.py Outdated Show resolved Hide resolved

Address review comments.

2a3d176

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

phu0ngng reviewed Apr 10, 2025

View reviewed changes

tests/jax/test_distributed_layernorm_mlp.py Show resolved Hide resolved

Rebase fixes.

33380bc

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

phu0ngng reviewed Apr 11, 2025

View reviewed changes

examples/jax/encoder/test_model_parallel_encoder.py Outdated Show resolved Hide resolved

jreiffers added 5 commits April 14, 2025 08:26

Remove unnecessary Shardy test.

3c4f668

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Update test_model_parallel_encoder Shardy tests.

ebb5055

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Update test_multigpu_encoder Shardy tests.

52f4bc0

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Add a TODO to add an MXFP8 Shardy test to test_multiprocessing_encoder.

705386e

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

Adjust tolerance.

47e4e03

Signed-off-by: Johannes Reifferscheid <jreiffers@nvidia.com>

phu0ngng approved these changes Apr 14, 2025

View reviewed changes

phu0ngng merged commit 6117b20 into NVIDIA:main Apr 14, 2025
12 checks passed

Comments

Conversation

jreiffers commented Apr 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Uh oh!

jreiffers commented Apr 4, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jreiffers commented Apr 9, 2025

Uh oh!

jreiffers commented Apr 10, 2025

Uh oh!

jreiffers commented Apr 10, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jreiffers commented Apr 10, 2025

Uh oh!

phu0ngng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

phu0ngng commented Apr 10, 2025

Uh oh!

Uh oh!

jreiffers commented Apr 11, 2025

Uh oh!

Uh oh!

phu0ngng commented Apr 14, 2025

Uh oh!

phu0ngng commented Apr 14, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jreiffers commented Apr 3, 2025 •

edited

Loading