Add support for MXFP8 All gather #3435
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3435
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 248a403 with merge base aa21b80 ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @avizon-aws! Thank you for your pull request and welcome to our community. Action RequiredIn order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
meta-cla
bot
added
the
CLA Signed
|
looks reasonable, would you mind rebasing as we've landed some refactors to MXTensor recently? just curious, which framework are you using TP + mxfp8 from? |
|
I think I already rebased just before creating the PR (yesterday night), so it should be up to date with mainline. I fixed the ruff formating to pass the workflows, I think one of the workflow was failing due to a missing tag on the PR, i think the right tag should be "add new feature", i wasnt able to add it, I think i dont have the permission for it. I tried the TP AG using native pytorch, i created a test script for it, didnt use any other framework currently. I updated the sharding strategy for to_mx due to which it required the TP AG. |
|
lgtm! will leave approval for @vkuzo who i think is also taking a look at this. @avizon-aws to fix the code analysis ci job, please run the linter with |
danielvegamyhre
added
the
topic: improvement
|
Thanks @danielvegamyhre , i have used the commands shared above to fix the formatting. |
Update [ghstack-poisoned]
* Update [ghstack-poisoned] * Update [ghstack-poisoned] * Update [ghstack-poisoned]
Summary: taking claude code for a more thorough spin, will start with local instructions and will see what makes sense to upstream Test Plan: Reviewers: Subscribers: Tasks: Tags:
* [Reland][PT2E][X86] Add Inductor fusion passes of float8 qconv for X86Inductor backend * add torch version check for Qconv FP8 UTs * fix format issue * Skip tests for ROCm --------- Co-authored-by: Sun, Jiayi <jiayi.sun@intel.com>
* Int8Tensor migration Summary: This PR creates a new Int8Tensor and updates the configs to use the new Int8Tensor flow Test Plan: To ensure BC: ``` pytest test/quantization/test_quant_api.py ``` To test new Int8Tensor: ``` pytest test/quantization/quantize_/workflows/int8/test_int8_tensor.py ``` Reviewers: Subscribers: Tasks: Tags: * ruff fixes * add init * fix ruff again * update * wip * undo update tests * fix ruff * fix varname * fix typing * add tests * fix dtype * fix ci * address granularity cr * update _choose_quant_func_and_quantize_tensor * make block size required attribute * made dtype required as well * address nits * skip per tensor weight only test for now
… XPU (pytorch#3368) * enable test/dtypes/test_bitpacking.py on intel xpu * enable test/dtypes/test_floatx.py * enable test/dtypes/test_floatx.py * fix format issue * fix format issue * update _DEVICES
…ze_pt2e_qat} UT files to intel XPU (pytorch#3405) * add test/quantization/pt2e/test_quantize_pt2e.py * add test/quantization/pt2e/test_quantize_pt2e.py * test/quantization/pt2e/test_quantize_pt2e_qat.py * test/quantization/pt2e/test_quantize_pt2e_qat.py * fix format issue * update format * increase timeout for xpu
|
I see some issues during rebasing, will fix them and then ping for review. |
|
@vkuzo , i have fixed the rebase |
|
I think the test is failing in the CI because of version and compatibility issues (I have run the test locally and it passes on compatibility >=9)), I will have to add some conditions due to which the test might need to be skipped for certain CUDA versions. |
|
I added checks to ensure the test runs only on CUDA and also for compatibility >=9 (H100 and above as FP8 is supported well) |
|
It seems like there is only 1 H100 GPU and hence the test is failing, because the world size required for my test is 2. What are the recommendations to proceed here? There are 2 options:
I think skipping the test is risky, I would prefer option 1 |
7 participants
Currently when trying to use TP All gather with MXFP8, it gives an unimplemented error for the all gather operation. This PR adds support for MXFP8 all gather which will enable better performance as it will reduce the bandwidth required for the gather compute.