Skip to content

[docs] AWS Trainium: rewrite guide for native TorchNeuron support - #15138

Merged
JingyaHuang merged 11 commits into
huggingface:mainfrom
JingyaHuang:refactor-neuron-doc
Oct 28, 2026
Merged

JingyaHuang merged 11 commits into
huggingface:mainfrom
JingyaHuang:refactor-neuron-doc

Conversation

@JingyaHuang

Copy link
Copy Markdown
Contributor

What does this PR do?

Rewrites the Neuron guide (docs/source/en/optimization/neuron.md) for native TorchNeuron support and gives it the same layout as the TPU guide.

AWS Trainium page

  • Replaces the old Optimum Neuron export flow with the native torch_neuronx backend: pipe.to("neuron") and torch.compile(backend="neuron").
  • Sections:
    • Eager mode (FLUX.2-klein-4B, single NeuronCore)
    • Compiled mode (static shapes, compile_repeated_blocks)
    • Tensor parallelism (Qwen-Image-2.1 with TensorParallelConfig; text encoder stays on CPU; fixed-seed CPU generator)
  • Notes the NEFF compile cache and TORCH_NEURONX_NEFF_CACHE_DIR.
  • Keeps a short pointer to Optimum Neuron for ahead-of-time exported pipelines.

Google TPU page

  • Adds the optimized eager mode (EagerMode.DEFER_AND_FUSE, enabled with TORCH_TPU_DEFER_AND_FUSE=1) to the execution-mode table.

Both pages

  • Pages are titled by hardware ("Google TPU", "AWS Trainium"). The intro explains that TorchTPU / TorchNeuron is the PyTorch backend that runs Diffusers on that hardware.
  • Sidebar entries in _toctree.yml renamed to match.

Validation

All Neuron snippets were run unmodified on trn2 (Neuron SDK 2.33 / torch-neuronx 2.14.3), and each produced a correct 1024² image:

Snippet Cores Result
Eager, FLUX.2-klein-4B, 4 steps 1 ✅
Compiled, FLUX.2-klein-4B, 4 steps 1 ✅
Tensor parallel, Qwen-Image-2.1, 40 steps 8 ✅
Tensor parallel + compile_repeated_blocks 8 ✅

The TPU snippets weren't run; only the mode table changed there.

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@github-actions github-actions Bot added documentation Improvements or additions to documentation models tests modular-pipelines hooks size/L PR with diff > 200 LOC labels Oct 27, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@github-actions github-actions Bot added size/M PR with diff < 200 LOC and removed models tests modular-pipelines hooks size/L PR with diff > 200 LOC labels Oct 27, 2026
@JingyaHuang
JingyaHuang marked this pull request as ready for review October 27, 2026 15:43

@stevhliu stevhliu left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

thanks for the docs! 🚀

Comment thread docs/source/en/optimization/neuron.md Resolved
Comment thread docs/source/en/optimization/neuron.md Resolved
Comment thread docs/source/en/optimization/neuron.md Resolved
Comment thread docs/source/en/optimization/neuron.md Resolved
Comment thread docs/source/en/optimization/neuron.md Resolved
`pipe.transformer.compile_repeated_blocks(backend="neuron", fullgraph=True, dynamic=False)` compiles one block and
reuses it for the others, which cuts the warmup time considerably.

## Tensor parallelism

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

We should refer the readers to the dedicated TP section and show Neuron-relevant stuff here?

| Mode | Constant | How to activate | Notes |
|---|---|---|---|
| Strict eager (default) | `EagerMode.DEFER_NEVER` | `pipe.to("tpu")` | Operations dispatched one at a time, asynchronous |
| Optimized eager | `EagerMode.DEFER_AND_FUSE` | `TORCH_TPU_DEFER_AND_FUSE=1` | Operations deferred and fused into larger graphs; faster than strict eager, longer compile |

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Have we tested for this mode?

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.

Yes we did, actually there are 3 eager modes... (there's a debug mode as well), I did not put in the beginning since I found it confusing and one should use torch.compile for deployment anyway, but indeed the semi-compiled mode is faster in between strict compile and compiled...

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Nice!

JingyaHuang and others added 5 commits October 28, 2026 11:58
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com>
@JingyaHuang
JingyaHuang requested a review from sayakpaul October 28, 2026 11:38
@JingyaHuang
JingyaHuang merged commit 40e551a into huggingface:main Oct 28, 2026
3 checks passed
@JingyaHuang
JingyaHuang deleted the refactor-neuron-doc branch October 28, 2026 16:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation size/M PR with diff < 200 LOC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants