Add ROCm support #1112
Add ROCm support #1112
Conversation
akasharidas
commented
Apr 16, 2025
akasharidas
commented
- Add a Transformer Engine wrapper to enable calling Flash Attention on ROCm.
- Add a config and launch script to demonstrate training Fuji-v2-70B on a single MI300 node.
| "importlab==0.8.1", # breaks pytype on 0.8 | ||
| "jax==0.4.38", | ||
| "jaxlib==0.4.38", | ||
| "jax==0.4.35", |
There was a problem hiding this comment.
Is this intentional to downgrade the jax version?
There was a problem hiding this comment.
I am guessing they are relying on a fixed version on transformer_engine, which has to be compatible with Jax version.
There was a problem hiding this comment.
This is because 0.4.35 is the closest available ROCm JAX release (https://github.com/ROCm/jax/releases). The next version available is 0.5 which we are open to upgrading to whenever axlearn decides to use it.
| _allow_explicit_bias = True | ||
|
|
||
|
|
||
| class ROCmTransformerEngineFlashAttention(BaseFlashAttention): |
There was a problem hiding this comment.
Add a unittest for this layer?
| --module=text.gpt.c4_trainer --config=fuji-70B-v2-flash-single-host \ | ||
| --trainer_dir=/tmp/gpt_c4_test --data_dir=gs://axlearn-public/tensorflow_datasets \ | ||
| --jax_backend=gpu \ | ||
| --mesh_selector="amd-mi300-single-node" \ No newline at end of file |
There was a problem hiding this comment.
How can we reproduce your experiments?
| from jax.ad_checkpoint import checkpoint_name | ||
| from jax.experimental import pallas as pl | ||
| try: | ||
| from transformer_engine.jax.flax.transformer import DotProductAttention |
There was a problem hiding this comment.
Is it possible to extract the lower level API from transformer engine? We waited to integrate cudnn from Jax native exactly to avoid dependency on transformer_engine, which make Jax upgrade and version control much more difficult.
There was a problem hiding this comment.
Fwiw, we are also open to plumb through xla customized call
| @@ -0,0 +1,208 @@ | |||
| # Copyright © 2023 Apple Inc. | |||
There was a problem hiding this comment.
| # Copyright © 2023 Apple Inc. | |
| # Copyright © 2025 Apple Inc. |
There was a problem hiding this comment.
And feel free to make it AMD if you intend to maintain it.
| gpu=[ | ||
| GPUDecoding, | ||
| # For GPU, prefer cuDNN (without bias) whenever possible, as it's the fastest. | ||
| ROCmTransformerEngineFlashAttention, |
There was a problem hiding this comment.
how is ROCmTransformerEngineFlashAttention code path selected here?
There was a problem hiding this comment.
added a check in is_supported which will throw unsupported if not on ROCm backend. Let me know if this approach is fine
| Currently tested on MI300. To run tests in parallel on a multi-GPU machine, use this: | ||
| ``` | ||
| PARALLEL_GPU_TEST=1 pytest -n 8 axlearn/common/flash_attention/gpu_attention_test.py |
There was a problem hiding this comment.
This file path needs to be revised to the current file.
|
Closing this PR due to inactivity. Feel free to reopen if you would like to continue the work. |
5 participants
