Repository navigation
Add a guide for training on AWS Trainium - #7729
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
github-actions
Bot
added
the
📚 documentation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 376bd88151
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| ## Data parallelism | ||
|
|
||
| A Trainium2 device exposes 4 logical NeuronCores. Train on all of them with [Accelerate](https://huggingface.co/docs/accelerate): | ||
|
|
||
| ```bash | ||
| accelerate launch --num_processes 4 train.py | ||
| ``` | ||
|
|
||
| SFT scales to 3.6x on 4 NeuronCores. To pin a run to some NeuronCores with `NEURON_RT_VISIBLE_CORES`, use the pairs `0,1` or `2,3`, or all four: the Neuron runtime fails to set up the collectives on `1,2`. |
There was a problem hiding this comment.
Is it the only form of parallelism supported in TRL for now?
There was a problem hiding this comment.
It's the only one I tested on Trainium. TP/FSDP through the transformers distributed API (DistributedConfig) should work since TRL takes the model as is, but I haven't tried it yet.
Co-authored-by: Michael Benayoun <mickbenayoun@gmail.com>
…nd the numbers, tidy the wording
michaelbenayoun
left a comment
michaelbenayoun
left a comment
There was a problem hiding this comment.
LGTM, thanks a lot for this!
Resolves #7661.
Adds a "Training on AWS Trainium" how-to guide: installation, SFT (with the
pad_to_multiple_of="auto"default from #7724),torch_compilewith the static dynamo shapes workaround (#7660), GRPO/RLOO with the static cache (#7664), data parallel over NeuronCores, and memory. Every number comes from runs on a trn2 instance, and the SFT, compile and GRPO examples were run there (capped to a few steps).Merge after #7724, since the page describes its default.
cc @michaelbenayoun
Note
Low Risk
Documentation and marketing copy only; no library or training runtime code changes in this diff.
Overview
Adds a Training on AWS Trainium how-to (
aws_trainium.md) and wires it into the docs nav and What's New on the README and docs index (replacing the prior fused LM head highlight).The guide documents running TRL trainers on Neuron via
torch-neuronxwith no code changes, and focuses on static shapes to limit slow recompiles: defaultpad_to_multiple_ofbehavior on Neuron, disabling Dynamo automatic dynamic shapes when usingtorch_compile=True, andcache_implementation="static"plus bf16 for GRPO/RLOO without vLLM. It also covers Accelerate data parallel across logical NeuronCores andNEURON_RT_VISIBLE_CORES, with benchmark numbers fromtrn2runs.Reviewed by Cursor Bugbot for commit 928802b. Bugbot is set up for automated code reviews on this repo. Configure here.