Skip to content

mxfp8 training: add TP sharding strategy for dim1 kernel#2436

Merged
vkuzo merged 25 commits into
mainfrom
gh/vkuzo/92/head
Jul 9, 2025
Merged

mxfp8 training: add TP sharding strategy for dim1 kernel#2436
vkuzo merged 25 commits into
mainfrom
gh/vkuzo/92/head

Conversation

@vkuzo

@vkuzo vkuzo commented Jun 24, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary:

Enables mxfp8 training with the dim1 triton kernel and TP, in eager mode. In detail:

  1. write a sharding strategy for the triton kernel custom op
  2. properly wrap the outputs of the custom op in DTensor
  3. add a test which enables the triton custom op, and modify the test model sizes since this kernel needs block_size 32

Note that compile does not work yet, seeing https://www.internalfb.com/phabricator/paste/view/P1851219639

Test Plan:

// TP test passes (with eager mode)
./test/prototype/mx_formats/test_mx_dtensor.sh

// unit tests still pass
pytest test/prototype/mx_formats

// torchtitan debug model loss goes down (in eager)
with-proxy CONFIG_FILE="torchtitan/models/llama3/train_configs/debug_model.toml " ./run_train.sh --model.print_after_conversion --training.steps 50 --model.converters mx --mx.recipe_name "mxfp8" --parallelism.tensor_parallel_degree=2 --mx.use_fp8_dim1_cast_triton_kernel

Reviewers:

Subscribers:

Tasks:

Tags:

vkuzo added 22 commits June 20, 2025 07:10
@vkuzo
Update
5c23c6b 
[ghstack-poisoned]
@vkuzo
Update
ad2ce62 
[ghstack-poisoned]
@vkuzo
Update
5eb2066 
[ghstack-poisoned]
@vkuzo
Update
6e3df57 
[ghstack-poisoned]
@vkuzo
Update
75e6fe7 
[ghstack-poisoned]
@vkuzo
Update
8bf42da 
[ghstack-poisoned]
@vkuzo
Update
c0080cd 
[ghstack-poisoned]
@vkuzo
Update
c6fc48b 
[ghstack-poisoned]
@vkuzo
Update
4cc1531 
[ghstack-poisoned]
@vkuzo
Update
42083e2 
[ghstack-poisoned]
@vkuzo
Update
9d171ad 
[ghstack-poisoned]
@vkuzo
Update
09c1c58 
[ghstack-poisoned]
@vkuzo
Update
e511e7b 
[ghstack-poisoned]
@vkuzo
Update
3562a5e 
[ghstack-poisoned]
@vkuzo
Update
7a0fd96 
[ghstack-poisoned]
@vkuzo
Update
2d1545f 
[ghstack-poisoned]
@vkuzo
Update
20b7db2 
[ghstack-poisoned]
@vkuzo
Update
7788412 
[ghstack-poisoned]
@vkuzo
Update
aabeb61 
[ghstack-poisoned]
@vkuzo
Update
28f32b9 
[ghstack-poisoned]
@vkuzo
Update
1001602 
[ghstack-poisoned]
@vkuzo
Update
b251829 
[ghstack-poisoned]
@vkuzo

vkuzo commented Jun 24, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Stack from ghstack (oldest at bottom):

@pytorch-bot

pytorch-bot Bot commented Jun 24, 2025 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2436

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Cancelled Job

As of commit 68b683d with merge base c57226b (image):

NEW FAILURE - The following job has failed:

CANCELLED JOB - The following job was cancelled. Please retry:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@vkuzo vkuzo mentioned this pull request Jun 24, 2025
Merged
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 24, 2025
vkuzo added a commit that referenced this pull request Jun 24, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
3737967 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 3b7f85d
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo
Update
dd36237 
[ghstack-poisoned]
@vkuzo vkuzo changed the base branch from gh/vkuzo/91/head to main June 24, 2025 19:27
vkuzo added a commit that referenced this pull request Jun 24, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
48ab874 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: ddb1f80
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo
Update
dc0803c 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Jun 25, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
3e27fca 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: b229983
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo vkuzo changed the title [wip] sharding strategy for dim1 kernel mxfp8 training: add TP sharding strategy for dim1 kernel Jun 25, 2025
vkuzo
vkuzo commented Jun 25, 2025
View reviewed changes
Comment thread test/prototype/mx_formats/test_mx_dtensor.py
)
# TODO(future PR): enable compile here, currently seeing
# https://www.internalfb.com/phabricator/paste/view/P1851219639
# _test_lowp_mlp_tensor_parallelism_base(

@vkuzo vkuzo Jun 25, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

need to uncomment this and run ./test/prototype/mx_formats/test_mx_dtensor.sh to reproduce

@vkuzo vkuzo Jul 8, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @bdhirsh

@vkuzo vkuzo added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Jun 25, 2025
@vkuzo vkuzo requested review from danielvegamyhre and drisspg June 27, 2025 13:56
danielvegamyhre
danielvegamyhre approved these changes Jul 7, 2025
View reviewed changes
Comment thread torchao/prototype/mx_formats/kernels.py Outdated
Comment thread torchao/prototype/mx_formats/kernels.py Outdated
@vkuzo
Update
68b683d 
[ghstack-poisoned]
vkuzo added a commit that referenced this pull request Jul 8, 2025
@vkuzo
[wip] sharding strategy for dim1 kernel
570770d 
Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: f66d58d
ghstack-comment-id: 3001601112
Pull Request resolved: #2436
@vkuzo vkuzo mentioned this pull request Jul 8, 2025
Merged
@vkuzo vkuzo merged commit e675ffd into main Jul 9, 2025
47 of 53 checks passed
@jithunnair-amd

jithunnair-amd commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true
cc @petrex @huydhn @drisspg

@danielvegamyhre

danielvegamyhre commented Jul 15, 2025

Copy link
Copy Markdown
Contributor

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true cc @petrex @huydhn @drisspg

Just fyi we are increasing timeout thresholds in #2549 and #2548

@petrex

petrex commented Jul 15, 2025

Copy link
Copy Markdown
Collaborator

@vkuzo This PR seems to have introduced timeouts in torch ao CI on trunk for both CUDA and ROCm nightly runs: https://hud.pytorch.org/hud/pytorch/ao/e675ffd9745e745056cd27a5f64cacad0aebd051/1?per_page=50&name_filter=regression&mergeEphemeralLF=true cc @petrex @huydhn @drisspg

Just fyi we are increasing timeout thresholds in #2549 and #2548

thanks I've labeled CI flow for rocm for #2549 and #2548.

liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
@vkuzo @liangel-02
mxfp8 training: add TP sharding strategy for dim1 kernel (#2436)
165772d 
* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]

* Update

[ghstack-poisoned]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@drisspg drisspg Awaiting requested review from drisspg drisspg is a code owner

1 more reviewer

@danielvegamyhre danielvegamyhre danielvegamyhre approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@vkuzo @jithunnair-amd @danielvegamyhre @petrex @facebook-github-bot