From e70fe29b895d2c85090bb7c0f5989a5f52bfd85f Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 22:10:42 +0000 Subject: [PATCH 1/8] docs: focus README on AWS Neuron --- README.md | 217 +++++++++++++++--------------------------------------- 1 file changed, 58 insertions(+), 159 deletions(-) diff --git a/README.md b/README.md index b0496fa816c..e81538fc0a3 100644 --- a/README.md +++ b/README.md @@ -1,191 +1,90 @@
-# torchtitan +# TorchTitan on AWS Neuron -#### A PyTorch native platform for training generative AI models - -[![8 GPU Feature Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml?query=branch%3Amain) -[![8 GPU Model Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml?query=branch%3Amain) -[![arXiv](https://img.shields.io/badge/arXiv-2410.06511-b31b1b.svg)](https://arxiv.org/abs/2410.06511) -[![ICLR](https://img.shields.io/badge/ICLR-2025-violet.svg)](https://iclr.cc/virtual/2025/poster/29620) -[![forum](https://img.shields.io/badge/pytorch-forum-DE3412.svg)](https://discuss.pytorch.org/c/distributed/torchtitan/44) -[![license](https://img.shields.io/badge/license-BSD_3--Clause-lightgrey.svg)](./LICENSE) -[![pip](https://img.shields.io/pypi/v/torchtitan?color=blue)](https://pypi.org/project/torchtitan/) -[![conda](https://img.shields.io/conda/vn/conda-forge/torchtitan?color=green)](https://anaconda.org/conda-forge/torchtitan) +### Train generative AI models at scale on AWS Trainium +[Overview](#overview) | [Quick start](#quick-start) | [Documentation](#documentation) | [Testing](#testing)
-`torchtitan` is under extensive development. To use the latest features of `torchtitan`, we recommend using the most recent PyTorch nightly. - - -## Latest News -- [2025/11] AMD released an [optimized fork](https://github.com/AMD-AGI/torchtitan-amd/tree/main) of `torchtitan` for AMD GPUs. -- [2025/10] We released `torchtitan` [v0.2.0](https://github.com/pytorch/torchtitan/releases). -- [2025/10] SkyPilot now supports `torchtitan`! See the tutorial [here](https://docs.skypilot.co/en/latest/examples/training/torchtitan.html). -- [2025/07] We published [instructions](/torchtitan/models/README.md) on how to add a model to `torchtitan`. -- [2025/04] Our paper was accepted by [ICLR 2025](https://iclr.cc/virtual/2025/poster/29620). -- [2024/12] GPU MODE [lecture](https://www.youtube.com/watch?v=VYWRjcUqW6w) on torchtitan. -- [2024/07] [Presentation](https://pytorch2024.sched.com/event/1fHn3) at PyTorch Conference 2024. - - ## Overview -`torchtitan` is a PyTorch native platform designed for **rapid experimentation and large-scale training** of generative AI models. As a minimal clean-room implementation of PyTorch native scaling techniques, `torchtitan` provides a flexible foundation for developers to build upon. With `torchtitan` [extension points](docs/extension.md), one can easily create custom extensions tailored to specific needs. - -Our mission is to accelerate innovation in the field of generative AI by empowering researchers and developers to explore new modeling architectures and infrastructure techniques. - -The Guiding Principles when building `torchtitan` -* Designed to be easy to understand, use and extend for different training purposes. -* Minimal changes to the model code when applying multi-dimensional parallelism. -* Bias towards a clean, minimal codebase while providing basic reusable / swappable components. - -`torchtitan` has been showcasing PyTorch's latest distributed training features, via support for pretraining Llama 3.1 LLMs of various sizes. - -## Contributing - -We look forward to your contributions! - -* To accelerate contributions to and innovations around torchtitan, we host an [`experiments`](torchtitan/experiments) folder. New ideas should start there. To contribute, follow the [`experiments guidelines`](torchtitan/experiments/README.md). -* For fixes and contributions to core, follow these [`guidelines`](CONTRIBUTING.md). +TorchTitan on AWS Neuron provides Neuron-specific components for training +generative AI models on AWS Trainium. The `torchtitan_neuron` package is +installed alongside `torchtitan` and uses the same configuration and launch +interfaces. -## Llama 3.1 training +The package includes: -### Key features available +- Neuron training recipes for Llama3, Qwen3, DeepSeek-V3, and GPT-OSS. +- Dense and mixture-of-experts training with eager mode, `torch.compile`, + FSDP2, and SimpleFSDP configurations. +- Neuron Kernel Interface (NKI) attention and MoE kernels. +- MXFP8 training support on Trainium3. +- Neuron-specific configuration builders, tests, and documentation. -1. Multi-dimensional composable parallelisms - - [FSDP2](docs/fsdp.md) with per-parameter sharding - - [Tensor Parallel](https://pytorch.org/docs/stable/distributed.tensor.parallel.html) (including [async TP](https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487)) - - [Pipeline Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) - - [Context Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) -2. [Meta device](https://pytorch.org/docs/stable/meta.html) initialization -3. Per-op selective and full activation checkpointing -4. [Distributed checkpointing](https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250) (including async checkpointing) - - [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning -5. `torch.compile` support -6. [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323) support ([how-to](torchtitan/components/quantization/float8.md)) -7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8.md) on Blackwell GPUs. -8. Supervised Fine-Tuning (SFT) with chat-formatted datasets -9. DDP and HSDP -10. [TorchFT](https://github.com/pytorch/torchft) integration -11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for [custom datasets](docs/datasets.md) -12. Gradient accumulation, enabled by giving an additional `--training.global_batch_size` argument on the CLI -13. Flexible learning rate scheduler (warmup-stable-decay) -14. [BF16 optimizer states](docs/bf16_optimizer_states.md) for reduced memory usage -15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via [Tensorboard or Weights & Biases](/docs/metrics.md) -16. [Debugging tools](docs/debugging.md) including CPU/GPU profiling, memory profiling, Flight Recorder, etc. -17. All options easily configured via [Python config registry](torchtitan/models/llama3/config_registry.py) with `--module` and `--config` CLI flags -18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)) -19. [Helper scripts](scripts/) to - - download tokenizers from Hugging Face - - convert original Llama 3 checkpoints into the expected DCP format - - estimate FSDP/HSDP memory usage without materializing the model - - run distributed inference with Tensor Parallel +See the [`torchtitan_neuron` package overview](torchtitan_neuron/README.md) for +the package boundary, layout, included components, and recipe selection. -We report [performance](benchmarks/llama3_h100_202412_torchtitan.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques. +## Quick start -### Dive into the code - -You may want to see how the model is defined or how parallelism techniques are applied. For a guided tour, see these files first: -* [torchtitan/train.py](torchtitan/train.py) - the main training loop and high-level setup code -* [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - the Llama 3.1 model definition -* [torchtitan/models/llama3/parallelize.py](torchtitan/models/llama3/parallelize.py) - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and `torch.compile` to the model -* [torchtitan/distributed/pipeline_parallel.py](torchtitan/distributed/pipeline_parallel.py) - helpers for applying Pipeline Parallel to the model -* [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - utils for saving/loading distributed checkpoints -* [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - utils for applying Float8 techniques - - -## Installation - -One can directly run the source code, or install `torchtitan` from a nightly build, or a stable release. - -### From source - -This method requires the nightly build of PyTorch, or the latest PyTorch built [from source](https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source). +Install the repository in a configured AWS Neuron environment: ```bash -git clone https://github.com/pytorch/torchtitan -cd torchtitan +cd /workspace/torchtitan_neuron pip install -r requirements.txt -pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu -``` - -> **Note:** The nightly build of `torchdata` is required when using a PyTorch nightly. Install it from the nightly index as shown above. - -### Nightly builds - -This method requires the nightly build of PyTorch. You can replace `cu130` with another version of cuda or an AMD GPU (e.g. `rocm6.3`). - -```sh -pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 --force-reinstall -pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130 -``` - -### Stable releases -One can install the latest [stable release](https://github.com/pytorch/torchtitan/releases) of `torchtitan` via `pip` or `conda`. -```sh -pip install torchtitan +pip install -e . ``` -```sh -conda install conda-forge::torchtitan -``` -Note that each stable release pins the nightly versions of `torch` and `torchao`. Please see [release.md](docs/release.md) for more details. - -### Downloading a tokenizer - -`torchtitan` currently supports training Llama 3.1 (8B, 70B, 405B) out of the box. To get started training these models, we need to download the tokenizer. Follow the instructions on the official [meta-llama](https://huggingface.co/meta-llama/Llama-3.1-8B) repository to ensure you have access to the Llama model weights. -Once you have confirmed access, you can run the following command to download the Llama 3.1 tokenizer to your local machine. +Select a Neuron recipe by passing its registry module and configuration name to +the TorchTitan training entrypoint. For example: ```bash -# Get your HF token from https://huggingface.co/settings/tokens - -# Llama 3.1 tokenizer -python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=... +NGPU=64 MODULE=torchtitan_neuron.models.qwen3 \ + CONFIG=qwen3_8b_tp4_fsdp16_sac ./run_train.sh ``` -### Start a training run -Llama 3 8B model locally on 8 GPUs - -```bash -MODULE=llama3 CONFIG=llama3_8b ./run_train.sh +Continue with the [dense model](torchtitan_neuron/docs/dense-tutorials.md) or +[MoE model](torchtitan_neuron/docs/moe-tutorials.md) tutorial for model assets, +topology settings, and complete launch commands. + +## Documentation + +| Topic | Guide | +| --- | --- | +| Dense Qwen3 and Llama3 training | [Dense model tutorials](torchtitan_neuron/docs/dense-tutorials.md) | +| Qwen3, DeepSeek-V3, and GPT-OSS MoE training | [MoE model tutorials](torchtitan_neuron/docs/moe-tutorials.md) | +| SDPA and NKI attention behavior | [Neuron attention backends](torchtitan_neuron/docs/neuron_attention_backends.md) | +| MXFP8 training on Trainium3 | [Neuron MXFP8](torchtitan_neuron/docs/neuron_mxfp8.md) | +| Qwen3 NKI mixture-of-experts recipes | [Qwen3 NKI MoE](torchtitan_neuron/docs/qwen3_nki_moe.md) | +| Full Qwen3-235B-A22B training | [Four-node Trainium3 UltraServer guide](torchtitan_neuron/docs/qwen3_235b_full_ep128.md) | +| Qwen3 inference with `torch.compile` | [Qwen3 inference](torchtitan_neuron/docs/qwen3_inference.md) | +| Neuron extension test suite | [Testing `torchtitan_neuron`](torchtitan_neuron/tests/README.md) | + +## Package layout + +```text +torchtitan_neuron/ +|-- models/ Neuron modules and Trainer configuration registries +|-- experiments/ GraphTrainer configuration registries +|-- components/ Neuron-specific reusable components +|-- docs/ AWS Neuron documentation +`-- tests/ Neuron kernel, Trainer, and GraphTrainer tests ``` -### Multi-Node Training -For training on ParallelCluster/Slurm type configurations, you can use the `multinode_trainer.slurm` file to submit your sbatch job. +## Testing -To get started adjust the number of nodes and GPUs -``` -#SBATCH --ntasks=2 -#SBATCH --nodes=2 -``` +Run the Neuron extension test suite from the repository root: -Then start a run where `nnodes` is your total node count, matching the sbatch node count above. - -``` -srun torchrun --nnodes 2 -``` - -If your gpu count per node is not 8, adjust `--nproc_per_node` in the torchrun command and `#SBATCH --gpus-per-task` in the SBATCH command section. - - -## Citation - -We provide a detailed look into the parallelisms and optimizations available in `torchtitan`, along with summary advice on when to use various techniques. - -[TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://openreview.net/forum?id=SFN6Wm7YBI) -``` -@inproceedings{ - liang2025torchtitan, - title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining}, - author={Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos}, - booktitle={The Thirteenth International Conference on Learning Representations}, - year={2025}, - url={https://openreview.net/forum?id=SFN6Wm7YBI} -} +```bash +pytest -s torchtitan_neuron/tests/ ``` +Hardware-specific tests require AWS Trainium and the corresponding AWS Neuron +dependencies. Tests that do not require Neuron hardware skip when those +dependencies are unavailable. ## License -Source code is made available under a [BSD 3 license](./LICENSE), however you may have other legal obligations that govern your use of other content linked in this repository, such as the license or terms of service for third-party data and models. +This project is licensed under the [BSD 3-Clause License](LICENSE). From ee19136b369241374e913ae56fb74dd645348ca4 Mon Sep 17 00:00:00 2001 From: Arash Jamalian <179860940+aws-arashj@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:07:01 +0000 Subject: [PATCH 2/8] docs: add TorchTitan Neuron overview to README (cherry picked from commit 80e000eb509a7f286b26b870de6fd9efd814841b) --- README.md | 68 +++++++++++++++++++++++++++---------------------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index e81538fc0a3..f19594e1d03 100644 --- a/README.md +++ b/README.md @@ -1,56 +1,50 @@ -
+# TorchTitan Neuron -# TorchTitan on AWS Neuron +TorchTitan Neuron is the AWS Neuron fork of [pytorch/torchtitan](https://github.com/pytorch/torchtitan). It adds AWS Trainium support through [TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html), the native PyTorch backend for Neuron. TorchTitan's models, parallelism strategies, and `torch.compile` integration run on the `neuron` device. The Neuron-specific recipes, NKI kernel integrations, and docs are in [`torchtitan_neuron/`](torchtitan_neuron/README.md). -### Train generative AI models at scale on AWS Trainium +## Support matrix -[Overview](#overview) | [Quick start](#quick-start) | [Documentation](#documentation) | [Testing](#testing) - -
- -## Overview - -TorchTitan on AWS Neuron provides Neuron-specific components for training -generative AI models on AWS Trainium. The `torchtitan_neuron` package is -installed alongside `torchtitan` and uses the same configuration and launch -interfaces. - -The package includes: - -- Neuron training recipes for Llama3, Qwen3, DeepSeek-V3, and GPT-OSS. -- Dense and mixture-of-experts training with eager mode, `torch.compile`, - FSDP2, and SimpleFSDP configurations. -- Neuron Kernel Interface (NKI) attention and MoE kernels. -- MXFP8 training support on Trainium3. -- Neuron-specific configuration builders, tests, and documentation. - -See the [`torchtitan_neuron` package overview](torchtitan_neuron/README.md) for -the package boundary, layout, included components, and recipe selection. +| Component | Supported version | +|---|---| +| Neuron SDK | 2.33.0 (TorchNeuron public beta) | +| PyTorch / `torch-neuronx` | 2.14.0 / 2.14.3.0.* | +| Python | 3.12, 3.13, 3.14 (Linux x86_64) | +| Instances | Trn2, Trn3 | +| Execution | Eager, and `torch.compile` with `--compile.backend neuron` (static shapes only) | +| Models | Llama3, Qwen3, DeepSeek-V3, GPT-OSS | +| Parallelism | FSDP2, SimpleFSDP, TP, CP, EP, PP, DDP, activation checkpointing | ## Quick start -Install the repository in a configured AWS Neuron environment: +Before you start, install and verify `torch-neuronx` on a Trn2 or Trn3 instance by following [Quickstart: Install and verify torch-neuronx](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/get-started/quickstart-torch-neuron.html). Then install this fork: ```bash -cd /workspace/torchtitan_neuron +git clone https://github.com/aws-neuron/torchtitan-neuron.git +cd torchtitan-neuron pip install -r requirements.txt pip install -e . ``` -Select a Neuron recipe by passing its registry module and configuration name to -the TorchTitan training entrypoint. For example: +Run a short verification on one `trn2.48xlarge` (64 logical NeuronCores, TP4 × FSDP16) using the bundled test tokenizer and dataset: ```bash -NGPU=64 MODULE=torchtitan_neuron.models.qwen3 \ - CONFIG=qwen3_8b_tp4_fsdp16_sac ./run_train.sh +NGPU=64 MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_tp4_fsdp16_sac \ + ./run_train.sh --training.steps 10 --metrics.log_freq 1 \ + --dataloader.dataset c4_test --hf_assets_path=./tests/assets/tokenizer ``` -Continue with the [dense model](torchtitan_neuron/docs/dense-tutorials.md) or -[MoE model](torchtitan_neuron/docs/moe-tutorials.md) tutorial for model assets, -topology settings, and complete launch commands. - ## Documentation +- [What is TorchTitan Neuron?](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/about-torchtitan-neuron.html): architecture, supported models, distributed strategies, and known limitations +- [Set up TorchTitan Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/get-started/setup-torchtitan-neuron.html) +- [TorchTitan Neuron tutorials](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/tutorials/index.html) +- [Release notes](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchtitan-neuron.html) + +### Repository guides + +Start with the [Neuron package overview](torchtitan_neuron/README.md) for the +package boundary, layout, included components, and recipe selection. + | Topic | Guide | | --- | --- | | Dense Qwen3 and Llama3 training | [Dense model tutorials](torchtitan_neuron/docs/dense-tutorials.md) | @@ -62,6 +56,10 @@ topology settings, and complete launch commands. | Qwen3 inference with `torch.compile` | [Qwen3 inference](torchtitan_neuron/docs/qwen3_inference.md) | | Neuron extension test suite | [Testing `torchtitan_neuron`](torchtitan_neuron/tests/README.md) | +## Issues + +Report problems with running on Neuron or Trainium in this repository's [issues](https://github.com/aws-neuron/torchtitan-neuron/issues). Report problems that also occur on upstream torchtitan to [pytorch/torchtitan](https://github.com/pytorch/torchtitan/issues). + ## Package layout ```text From b7227cd4fcd7ab19590644ebe3b7973af93629f9 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:11:03 +0000 Subject: [PATCH 3/8] docs: remove testing section --- README.md | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/README.md b/README.md index f19594e1d03..2a43423ef9a 100644 --- a/README.md +++ b/README.md @@ -71,18 +71,6 @@ torchtitan_neuron/ `-- tests/ Neuron kernel, Trainer, and GraphTrainer tests ``` -## Testing - -Run the Neuron extension test suite from the repository root: - -```bash -pytest -s torchtitan_neuron/tests/ -``` - -Hardware-specific tests require AWS Trainium and the corresponding AWS Neuron -dependencies. Tests that do not require Neuron hardware skip when those -dependencies are unavailable. - ## License This project is licensed under the [BSD 3-Clause License](LICENSE). From 0dc2f6e1efae0135295762b489fc22ab7eca5893 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:19:33 +0000 Subject: [PATCH 4/8] docs: check out Neuron release branch in quick start --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 2a43423ef9a..37a5f778861 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ Before you start, install and verify `torch-neuronx` on a Trn2 or Trn3 instance ```bash git clone https://github.com/aws-neuron/torchtitan-neuron.git cd torchtitan-neuron +git checkout release_2.33.0 pip install -r requirements.txt pip install -e . ``` From 36356fd472b6a0a63ffd27149953d678e28ccbc5 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:28:23 +0000 Subject: [PATCH 5/8] docs: retain upstream README below Neuron guide --- README.md | 199 ++++++++++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 195 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 37a5f778861..8c4216294f9 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # TorchTitan Neuron -TorchTitan Neuron is the AWS Neuron fork of [pytorch/torchtitan](https://github.com/pytorch/torchtitan). It adds AWS Trainium support through [TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html), the native PyTorch backend for Neuron. TorchTitan's models, parallelism strategies, and `torch.compile` integration run on the `neuron` device. The Neuron-specific recipes, NKI kernel integrations, and docs are in [`torchtitan_neuron/`](torchtitan_neuron/README.md). +TorchTitan Neuron is the AWS Neuron's extension of [pytorch/torchtitan](https://github.com/pytorch/torchtitan). It adds AWS Trainium support through [TorchNeuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/torch/about-pytorch-neuron.html), the native PyTorch backend for Neuron. The Neuron-specific recipes, NKI kernel integrations, and docs are in [`torchtitan_neuron/`](torchtitan_neuron/README.md). ## Support matrix @@ -36,7 +36,7 @@ NGPU=64 MODULE=torchtitan_neuron.models.qwen3 CONFIG=qwen3_8b_tp4_fsdp16_sac \ ## Documentation -- [What is TorchTitan Neuron?](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/about-torchtitan-neuron.html): architecture, supported models, distributed strategies, and known limitations +- [What is TorchTitan Neuron?](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/about-torchtitan-neuron.html) - [Set up TorchTitan Neuron](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/get-started/setup-torchtitan-neuron.html) - [TorchTitan Neuron tutorials](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/training/torchtitan/tutorials/index.html) - [Release notes](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/release-notes/components/torchtitan-neuron.html) @@ -50,7 +50,6 @@ package boundary, layout, included components, and recipe selection. | --- | --- | | Dense Qwen3 and Llama3 training | [Dense model tutorials](torchtitan_neuron/docs/dense-tutorials.md) | | Qwen3, DeepSeek-V3, and GPT-OSS MoE training | [MoE model tutorials](torchtitan_neuron/docs/moe-tutorials.md) | -| SDPA and NKI attention behavior | [Neuron attention backends](torchtitan_neuron/docs/neuron_attention_backends.md) | | MXFP8 training on Trainium3 | [Neuron MXFP8](torchtitan_neuron/docs/neuron_mxfp8.md) | | Qwen3 NKI mixture-of-experts recipes | [Qwen3 NKI MoE](torchtitan_neuron/docs/qwen3_nki_moe.md) | | Full Qwen3-235B-A22B training | [Four-node Trainium3 UltraServer guide](torchtitan_neuron/docs/qwen3_235b_full_ep128.md) | @@ -72,6 +71,198 @@ torchtitan_neuron/ `-- tests/ Neuron kernel, Trainer, and GraphTrainer tests ``` +--- + +# Upstream TorchTitan README + +
+ +# torchtitan + +#### A PyTorch native platform for training generative AI models + +[![8 GPU Feature Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml?query=branch%3Amain) +[![8 GPU Model Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml?query=branch%3Amain) +[![arXiv](https://img.shields.io/badge/arXiv-2410.06511-b31b1b.svg)](https://arxiv.org/abs/2410.06511) +[![ICLR](https://img.shields.io/badge/ICLR-2025-violet.svg)](https://iclr.cc/virtual/2025/poster/29620) +[![forum](https://img.shields.io/badge/pytorch-forum-DE3412.svg)](https://discuss.pytorch.org/c/distributed/torchtitan/44) +[![license](https://img.shields.io/badge/license-BSD_3--Clause-lightgrey.svg)](./LICENSE) +[![pip](https://img.shields.io/pypi/v/torchtitan?color=blue)](https://pypi.org/project/torchtitan/) +[![conda](https://img.shields.io/conda/vn/conda-forge/torchtitan?color=green)](https://anaconda.org/conda-forge/torchtitan) + + +
+ +`torchtitan` is under extensive development. To use the latest features of `torchtitan`, we recommend using the most recent PyTorch nightly. + + +## Latest News +- [2025/11] AMD released an [optimized fork](https://github.com/AMD-AGI/torchtitan-amd/tree/main) of `torchtitan` for AMD GPUs. +- [2025/10] We released `torchtitan` [v0.2.0](https://github.com/pytorch/torchtitan/releases). +- [2025/10] SkyPilot now supports `torchtitan`! See the tutorial [here](https://docs.skypilot.co/en/latest/examples/training/torchtitan.html). +- [2025/07] We published [instructions](/torchtitan/models/README.md) on how to add a model to `torchtitan`. +- [2025/04] Our paper was accepted by [ICLR 2025](https://iclr.cc/virtual/2025/poster/29620). +- [2024/12] GPU MODE [lecture](https://www.youtube.com/watch?v=VYWRjcUqW6w) on torchtitan. +- [2024/07] [Presentation](https://pytorch2024.sched.com/event/1fHn3) at PyTorch Conference 2024. + + +## Overview + +`torchtitan` is a PyTorch native platform designed for **rapid experimentation and large-scale training** of generative AI models. As a minimal clean-room implementation of PyTorch native scaling techniques, `torchtitan` provides a flexible foundation for developers to build upon. With `torchtitan` [extension points](docs/extension.md), one can easily create custom extensions tailored to specific needs. + +Our mission is to accelerate innovation in the field of generative AI by empowering researchers and developers to explore new modeling architectures and infrastructure techniques. + +The Guiding Principles when building `torchtitan` +* Designed to be easy to understand, use and extend for different training purposes. +* Minimal changes to the model code when applying multi-dimensional parallelism. +* Bias towards a clean, minimal codebase while providing basic reusable / swappable components. + +`torchtitan` has been showcasing PyTorch's latest distributed training features, via support for pretraining Llama 3.1 LLMs of various sizes. + +## Contributing + +We look forward to your contributions! + +* To accelerate contributions to and innovations around torchtitan, we host an [`experiments`](torchtitan/experiments) folder. New ideas should start there. To contribute, follow the [`experiments guidelines`](torchtitan/experiments/README.md). +* For fixes and contributions to core, follow these [`guidelines`](CONTRIBUTING.md). + +## Llama 3.1 training + +### Key features available + +1. Multi-dimensional composable parallelisms + - [FSDP2](docs/fsdp.md) with per-parameter sharding + - [Tensor Parallel](https://pytorch.org/docs/stable/distributed.tensor.parallel.html) (including [async TP](https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487)) + - [Pipeline Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) + - [Context Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) +2. [Meta device](https://pytorch.org/docs/stable/meta.html) initialization +3. Per-op selective and full activation checkpointing +4. [Distributed checkpointing](https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250) (including async checkpointing) + - [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning +5. `torch.compile` support +6. [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323) support ([how-to](torchtitan/components/quantization/float8.md)) +7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8.md) on Blackwell GPUs. +8. Supervised Fine-Tuning (SFT) with chat-formatted datasets +9. DDP and HSDP +10. [TorchFT](https://github.com/pytorch/torchft) integration +11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for [custom datasets](docs/datasets.md) +12. Gradient accumulation, enabled by giving an additional `--training.global_batch_size` argument on the CLI +13. Flexible learning rate scheduler (warmup-stable-decay) +14. [BF16 optimizer states](docs/bf16_optimizer_states.md) for reduced memory usage +15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via [Tensorboard or Weights & Biases](/docs/metrics.md) +16. [Debugging tools](docs/debugging.md) including CPU/GPU profiling, memory profiling, Flight Recorder, etc. +17. All options easily configured via [Python config registry](torchtitan/models/llama3/config_registry.py) with `--module` and `--config` CLI flags +18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)) +19. [Helper scripts](scripts/) to + - download tokenizers from Hugging Face + - convert original Llama 3 checkpoints into the expected DCP format + - estimate FSDP/HSDP memory usage without materializing the model + - run distributed inference with Tensor Parallel + +We report [performance](benchmarks/llama3_h100_202412_torchtitan.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques. + +### Dive into the code + +You may want to see how the model is defined or how parallelism techniques are applied. For a guided tour, see these files first: +* [torchtitan/train.py](torchtitan/train.py) - the main training loop and high-level setup code +* [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - the Llama 3.1 model definition +* [torchtitan/models/llama3/parallelize.py](torchtitan/models/llama3/parallelize.py) - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and `torch.compile` to the model +* [torchtitan/distributed/pipeline_parallel.py](torchtitan/distributed/pipeline_parallel.py) - helpers for applying Pipeline Parallel to the model +* [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - utils for saving/loading distributed checkpoints +* [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - utils for applying Float8 techniques + + +## Installation + +One can directly run the source code, or install `torchtitan` from a nightly build, or a stable release. + +### From source + +This method requires the nightly build of PyTorch, or the latest PyTorch built [from source](https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source). + +```bash +git clone https://github.com/pytorch/torchtitan +cd torchtitan +pip install -r requirements.txt +pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu +``` + +> **Note:** The nightly build of `torchdata` is required when using a PyTorch nightly. Install it from the nightly index as shown above. + +### Nightly builds + +This method requires the nightly build of PyTorch. You can replace `cu130` with another version of cuda or an AMD GPU (e.g. `rocm6.3`). + +```sh +pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 --force-reinstall +pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130 +``` + +### Stable releases +One can install the latest [stable release](https://github.com/pytorch/torchtitan/releases) of `torchtitan` via `pip` or `conda`. +```sh +pip install torchtitan +``` +```sh +conda install conda-forge::torchtitan +``` +Note that each stable release pins the nightly versions of `torch` and `torchao`. Please see [release.md](docs/release.md) for more details. + +### Downloading a tokenizer + +`torchtitan` currently supports training Llama 3.1 (8B, 70B, 405B) out of the box. To get started training these models, we need to download the tokenizer. Follow the instructions on the official [meta-llama](https://huggingface.co/meta-llama/Llama-3.1-8B) repository to ensure you have access to the Llama model weights. + +Once you have confirmed access, you can run the following command to download the Llama 3.1 tokenizer to your local machine. + +```bash +# Get your HF token from https://huggingface.co/settings/tokens + +# Llama 3.1 tokenizer +python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=... +``` + +### Start a training run +Llama 3 8B model locally on 8 GPUs + +```bash +MODULE=llama3 CONFIG=llama3_8b ./run_train.sh +``` + +### Multi-Node Training +For training on ParallelCluster/Slurm type configurations, you can use the `multinode_trainer.slurm` file to submit your sbatch job. + +To get started adjust the number of nodes and GPUs +``` +#SBATCH --ntasks=2 +#SBATCH --nodes=2 +``` + +Then start a run where `nnodes` is your total node count, matching the sbatch node count above. + +``` +srun torchrun --nnodes 2 +``` + +If your gpu count per node is not 8, adjust `--nproc_per_node` in the torchrun command and `#SBATCH --gpus-per-task` in the SBATCH command section. + + +## Citation + +We provide a detailed look into the parallelisms and optimizations available in `torchtitan`, along with summary advice on when to use various techniques. + +[TorchTitan: One-stop PyTorch native solution for production ready LLM pre-training](https://openreview.net/forum?id=SFN6Wm7YBI) +``` +@inproceedings{ + liang2025torchtitan, + title={TorchTitan: One-stop PyTorch native solution for production ready {LLM} pretraining}, + author={Wanchao Liang and Tianyu Liu and Less Wright and Will Constable and Andrew Gu and Chien-Chin Huang and Iris Zhang and Wei Feng and Howard Huang and Junjie Wang and Sanket Purandare and Gokul Nadathur and Stratos Idreos}, + booktitle={The Thirteenth International Conference on Learning Representations}, + year={2025}, + url={https://openreview.net/forum?id=SFN6Wm7YBI} +} +``` + + ## License -This project is licensed under the [BSD 3-Clause License](LICENSE). +Source code is made available under a [BSD 3 license](./LICENSE), however you may have other legal obligations that govern your use of other content linked in this repository, such as the license or terms of service for third-party data and models. From fb3bffbb4f86edb26a06a625105681e430834618 Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:35:27 +0000 Subject: [PATCH 6/8] docs: retain upstream citation and license notice --- README.md | 175 ------------------------------------------------------ 1 file changed, 175 deletions(-) diff --git a/README.md b/README.md index 8c4216294f9..b465423c98d 100644 --- a/README.md +++ b/README.md @@ -71,181 +71,6 @@ torchtitan_neuron/ `-- tests/ Neuron kernel, Trainer, and GraphTrainer tests ``` ---- - -# Upstream TorchTitan README - -
- -# torchtitan - -#### A PyTorch native platform for training generative AI models - -[![8 GPU Feature Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_features.yaml?query=branch%3Amain) -[![8 GPU Model Tests](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml/badge.svg?branch=main)](https://github.com/pytorch/torchtitan/actions/workflows/integration_test_8gpu_models.yaml?query=branch%3Amain) -[![arXiv](https://img.shields.io/badge/arXiv-2410.06511-b31b1b.svg)](https://arxiv.org/abs/2410.06511) -[![ICLR](https://img.shields.io/badge/ICLR-2025-violet.svg)](https://iclr.cc/virtual/2025/poster/29620) -[![forum](https://img.shields.io/badge/pytorch-forum-DE3412.svg)](https://discuss.pytorch.org/c/distributed/torchtitan/44) -[![license](https://img.shields.io/badge/license-BSD_3--Clause-lightgrey.svg)](./LICENSE) -[![pip](https://img.shields.io/pypi/v/torchtitan?color=blue)](https://pypi.org/project/torchtitan/) -[![conda](https://img.shields.io/conda/vn/conda-forge/torchtitan?color=green)](https://anaconda.org/conda-forge/torchtitan) - - -
- -`torchtitan` is under extensive development. To use the latest features of `torchtitan`, we recommend using the most recent PyTorch nightly. - - -## Latest News -- [2025/11] AMD released an [optimized fork](https://github.com/AMD-AGI/torchtitan-amd/tree/main) of `torchtitan` for AMD GPUs. -- [2025/10] We released `torchtitan` [v0.2.0](https://github.com/pytorch/torchtitan/releases). -- [2025/10] SkyPilot now supports `torchtitan`! See the tutorial [here](https://docs.skypilot.co/en/latest/examples/training/torchtitan.html). -- [2025/07] We published [instructions](/torchtitan/models/README.md) on how to add a model to `torchtitan`. -- [2025/04] Our paper was accepted by [ICLR 2025](https://iclr.cc/virtual/2025/poster/29620). -- [2024/12] GPU MODE [lecture](https://www.youtube.com/watch?v=VYWRjcUqW6w) on torchtitan. -- [2024/07] [Presentation](https://pytorch2024.sched.com/event/1fHn3) at PyTorch Conference 2024. - - -## Overview - -`torchtitan` is a PyTorch native platform designed for **rapid experimentation and large-scale training** of generative AI models. As a minimal clean-room implementation of PyTorch native scaling techniques, `torchtitan` provides a flexible foundation for developers to build upon. With `torchtitan` [extension points](docs/extension.md), one can easily create custom extensions tailored to specific needs. - -Our mission is to accelerate innovation in the field of generative AI by empowering researchers and developers to explore new modeling architectures and infrastructure techniques. - -The Guiding Principles when building `torchtitan` -* Designed to be easy to understand, use and extend for different training purposes. -* Minimal changes to the model code when applying multi-dimensional parallelism. -* Bias towards a clean, minimal codebase while providing basic reusable / swappable components. - -`torchtitan` has been showcasing PyTorch's latest distributed training features, via support for pretraining Llama 3.1 LLMs of various sizes. - -## Contributing - -We look forward to your contributions! - -* To accelerate contributions to and innovations around torchtitan, we host an [`experiments`](torchtitan/experiments) folder. New ideas should start there. To contribute, follow the [`experiments guidelines`](torchtitan/experiments/README.md). -* For fixes and contributions to core, follow these [`guidelines`](CONTRIBUTING.md). - -## Llama 3.1 training - -### Key features available - -1. Multi-dimensional composable parallelisms - - [FSDP2](docs/fsdp.md) with per-parameter sharding - - [Tensor Parallel](https://pytorch.org/docs/stable/distributed.tensor.parallel.html) (including [async TP](https://discuss.pytorch.org/t/distributed-w-torchtitan-introducing-async-tensor-parallelism-in-pytorch/209487)) - - [Pipeline Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) - - [Context Parallel](https://discuss.pytorch.org/t/distributed-w-torchtitan-breaking-barriers-training-long-context-llms-with-1m-sequence-length-in-pytorch-using-context-parallel/215082) -2. [Meta device](https://pytorch.org/docs/stable/meta.html) initialization -3. Per-op selective and full activation checkpointing -4. [Distributed checkpointing](https://discuss.pytorch.org/t/distributed-w-torchtitan-optimizing-checkpointing-efficiency-with-pytorch-dcp/211250) (including async checkpointing) - - [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning -5. `torch.compile` support -6. [Float8](https://discuss.pytorch.org/t/distributed-w-torchtitan-enabling-float8-all-gather-in-fsdp2/209323) support ([how-to](torchtitan/components/quantization/float8.md)) -7. [MXFP8 training for dense and MoE models](torchtitan/components/quantization/mxfp8.md) on Blackwell GPUs. -8. Supervised Fine-Tuning (SFT) with chat-formatted datasets -9. DDP and HSDP -10. [TorchFT](https://github.com/pytorch/torchft) integration -11. Checkpointable data-loading, with the C4 dataset pre-configured (144M entries) and support for [custom datasets](docs/datasets.md) -12. Gradient accumulation, enabled by giving an additional `--training.global_batch_size` argument on the CLI -13. Flexible learning rate scheduler (warmup-stable-decay) -14. [BF16 optimizer states](docs/bf16_optimizer_states.md) for reduced memory usage -15. Loss, GPU memory, throughput (tokens/sec), TFLOPs, and MFU displayed and logged via [Tensorboard or Weights & Biases](/docs/metrics.md) -16. [Debugging tools](docs/debugging.md) including CPU/GPU profiling, memory profiling, Flight Recorder, etc. -17. All options easily configured via [Python config registry](torchtitan/models/llama3/config_registry.py) with `--module` and `--config` CLI flags -18. Structured logging: per-rank trace of key training phases; (see [`torchtitan/observability/structured_logger/README.md`](torchtitan/observability/structured_logger/README.md)) -19. [Helper scripts](scripts/) to - - download tokenizers from Hugging Face - - convert original Llama 3 checkpoints into the expected DCP format - - estimate FSDP/HSDP memory usage without materializing the model - - run distributed inference with Tensor Parallel - -We report [performance](benchmarks/llama3_h100_202412_torchtitan.md) on up to 512 GPUs, and verify [loss converging](docs/converging.md) correctness of various techniques. - -### Dive into the code - -You may want to see how the model is defined or how parallelism techniques are applied. For a guided tour, see these files first: -* [torchtitan/train.py](torchtitan/train.py) - the main training loop and high-level setup code -* [torchtitan/models/llama3/model.py](torchtitan/models/llama3/model.py) - the Llama 3.1 model definition -* [torchtitan/models/llama3/parallelize.py](torchtitan/models/llama3/parallelize.py) - helpers for applying Data Parallel, Tensor Parallel, activation checkpointing, and `torch.compile` to the model -* [torchtitan/distributed/pipeline_parallel.py](torchtitan/distributed/pipeline_parallel.py) - helpers for applying Pipeline Parallel to the model -* [torchtitan/components/checkpoint.py](torchtitan/components/checkpoint.py) - utils for saving/loading distributed checkpoints -* [torchtitan/components/quantization/float8.py](torchtitan/components/quantization/float8.py) - utils for applying Float8 techniques - - -## Installation - -One can directly run the source code, or install `torchtitan` from a nightly build, or a stable release. - -### From source - -This method requires the nightly build of PyTorch, or the latest PyTorch built [from source](https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source). - -```bash -git clone https://github.com/pytorch/torchtitan -cd torchtitan -pip install -r requirements.txt -pip install --pre torchdata --index-url https://download.pytorch.org/whl/nightly/cpu -``` - -> **Note:** The nightly build of `torchdata` is required when using a PyTorch nightly. Install it from the nightly index as shown above. - -### Nightly builds - -This method requires the nightly build of PyTorch. You can replace `cu130` with another version of cuda or an AMD GPU (e.g. `rocm6.3`). - -```sh -pip3 install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu130 --force-reinstall -pip install --pre torchtitan --index-url https://download.pytorch.org/whl/nightly/cu130 -``` - -### Stable releases -One can install the latest [stable release](https://github.com/pytorch/torchtitan/releases) of `torchtitan` via `pip` or `conda`. -```sh -pip install torchtitan -``` -```sh -conda install conda-forge::torchtitan -``` -Note that each stable release pins the nightly versions of `torch` and `torchao`. Please see [release.md](docs/release.md) for more details. - -### Downloading a tokenizer - -`torchtitan` currently supports training Llama 3.1 (8B, 70B, 405B) out of the box. To get started training these models, we need to download the tokenizer. Follow the instructions on the official [meta-llama](https://huggingface.co/meta-llama/Llama-3.1-8B) repository to ensure you have access to the Llama model weights. - -Once you have confirmed access, you can run the following command to download the Llama 3.1 tokenizer to your local machine. - -```bash -# Get your HF token from https://huggingface.co/settings/tokens - -# Llama 3.1 tokenizer -python scripts/download_hf_assets.py --repo_id meta-llama/Llama-3.1-8B --assets tokenizer --hf_token=... -``` - -### Start a training run -Llama 3 8B model locally on 8 GPUs - -```bash -MODULE=llama3 CONFIG=llama3_8b ./run_train.sh -``` - -### Multi-Node Training -For training on ParallelCluster/Slurm type configurations, you can use the `multinode_trainer.slurm` file to submit your sbatch job. - -To get started adjust the number of nodes and GPUs -``` -#SBATCH --ntasks=2 -#SBATCH --nodes=2 -``` - -Then start a run where `nnodes` is your total node count, matching the sbatch node count above. - -``` -srun torchrun --nnodes 2 -``` - -If your gpu count per node is not 8, adjust `--nproc_per_node` in the torchrun command and `#SBATCH --gpus-per-task` in the SBATCH command section. - - ## Citation We provide a detailed look into the parallelisms and optimizations available in `torchtitan`, along with summary advice on when to use various techniques. From 8e03a9dcde05d5cb78a3f8c64fce94d539995d9c Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:36:57 +0000 Subject: [PATCH 7/8] docs: remove MXFP8 repository guide --- README.md | 1 - 1 file changed, 1 deletion(-) diff --git a/README.md b/README.md index b465423c98d..8ad05f5c627 100644 --- a/README.md +++ b/README.md @@ -50,7 +50,6 @@ package boundary, layout, included components, and recipe selection. | --- | --- | | Dense Qwen3 and Llama3 training | [Dense model tutorials](torchtitan_neuron/docs/dense-tutorials.md) | | Qwen3, DeepSeek-V3, and GPT-OSS MoE training | [MoE model tutorials](torchtitan_neuron/docs/moe-tutorials.md) | -| MXFP8 training on Trainium3 | [Neuron MXFP8](torchtitan_neuron/docs/neuron_mxfp8.md) | | Qwen3 NKI mixture-of-experts recipes | [Qwen3 NKI MoE](torchtitan_neuron/docs/qwen3_nki_moe.md) | | Full Qwen3-235B-A22B training | [Four-node Trainium3 UltraServer guide](torchtitan_neuron/docs/qwen3_235b_full_ep128.md) | | Qwen3 inference with `torch.compile` | [Qwen3 inference](torchtitan_neuron/docs/qwen3_inference.md) | From 6e0eb005d2c0668a12f2aca7a5af6407d719e5ed Mon Sep 17 00:00:00 2001 From: Ryan Li Date: Mon, 19 Oct 2026 23:41:22 +0000 Subject: [PATCH 8/8] docs: clarify Neuron package layout scope --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 8ad05f5c627..ec05834a0cd 100644 --- a/README.md +++ b/README.md @@ -61,6 +61,9 @@ Report problems with running on Neuron or Trainium in this repository's [issues] ## Package layout +The Neuron-specific recipes, NKI kernel integrations, and docs are in +[`torchtitan_neuron/`](torchtitan_neuron/README.md): + ```text torchtitan_neuron/ |-- models/ Neuron modules and Trainer configuration registries