This document is relevant for: Trn2, Trn3

Tutorial: Pre-train GPT-OSS with TorchTitan Neuron#

This tutorial walks you through pre-training GPT-OSS with TorchTitan on AWS Trainium using TorchNeuron. When you have completed it, you will have run GPT-OSS 20B sharded across a full Trainium node, with NKI kernels handling the attention and mixture-of-experts layers.

Overview#

GPT-OSS is a mixture-of-experts model: each token is routed to a small subset of expert feed-forward networks rather than passing through all of them. The dense parts behave like any other transformer, but the MoE layers dominate both the memory footprint and the communication pattern, and they are where the Neuron-specific work shows up.

Two things make the Neuron path for GPT-OSS different from a stock TorchTitan run. Attention uses a NKI kernel with sliding-window support rather than SDPA or varlen attention, and each MoE layer is converted to a NKI implementation. Both come from the recipe; neither is something you select by hand.

If you have not run TorchTitan on Trainium before, do Tutorial: Run a basic TorchTitan training on Neuron first. It establishes the launcher, the environment variables, and the tokenizer workflow that this tutorial reuses without re-explaining.

Where the recipes live#

Neuron-specific recipes ship in the torchtitan_neuron package, installed alongside torchtitan:

torchtitan_neuron/
├── models/        Neuron Trainer recipes
└── experiments/   GraphTrainer recipes, including SimpleFSDP

Because these modules are not in TorchTitan’s built-in shorthand list, you select them by fully qualified module path rather than a short name like gpt_oss. Every launch command below does this.

Before you start#

This tutorial assumes that you have experience in the following areas:

Model details#

This tutorial trains GPT-OSS 20B. Two Neuron recipes are available, and they differ in how sharding is implemented rather than in what they train:

Recipe

Module

Sharding

gpt_oss_20b_nki_tp4_fsdp16

torchtitan_neuron.models.gpt_oss

FSDP v2

simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16

torchtitan_neuron.experiments.graph_trainer.gpt_oss

SimpleFSDP

Both build the 20b model flavor with the NKI attention backend, convert every MoE layer with to_nki_moe_config, apply the prefix-cache rotary path, turn on torch.compile with the Neuron backend, pin tensor-parallel degree 4 with sharding filling the remaining cores, set expert-parallel degree 1, and use full activation checkpointing.

Note

The upstream gpt_oss registry ships recipes of its own — gpt_oss_debugmodel, gpt_oss_debugmodel_flex, gpt_oss_20b, and gpt_oss_120b. None of them do any of the above: they are the stock TorchTitan configurations, without the NKI kernels, the Neuron compile backend, or the rotary workaround Neuron needs. They are not substitutes for the two recipes here.

You will run both in Step 2. Which one you would reach for afterwards depends on what you are optimizing for:

  • FSDP v2 wraps modules with PyTorch’s fully_shard. It is the path the basic tutorial uses.

  • SimpleFSDP instead represents sharding as collectives inside the computation graph, which keeps the whole step visible to the compiler. TorchTitan describes it as achieving memory and throughput improvements over eager FSDP2, and How to use torch.distributed on Neuron recommends it on Neuron for torch.compile workloads. The trade is compile time: a larger graph takes longer to compile. One caveat on Neuron: it is supported only with the JIT compile mode, so do not override --compile.mode. Runs print a deprecation warning for that mode; in this release that is expected rather than a problem.

Tutorial: Train Llama and DeepSeek with SimpleFSDP on Neuron covers SimpleFSDP in depth, including that limitation and where the recipes set it.

Limitations:

  • Targets a single trn2.48xlarge node (64 NeuronCores). Both recipes shard with tensor-parallel degree 4 and data-parallel sharding filling the remaining cores.

  • Expert parallelism is not used for GPT-OSS on Neuron. Both recipes scale with tensor parallelism and data-parallel sharding, and ship expert-parallel degree 1. Raising it fails — the NKI mixture-of-experts implementation shards experts with tensor parallelism only.

  • gpt_oss_120b needs more than one node and is out of scope here.

Note

The attention backend these recipes use applies causal attention but does not consume TorchTitan’s packed-document boundaries. A sequence that holds more than one document therefore lets tokens attend across the boundary into earlier, unrelated documents. The Neuron recipes do not currently offer a backend that masks this, and the backend is fixed by the recipe rather than settable on the command line. If strict cross-document isolation matters for your workload, raise it before you build on this path.

Prerequisites#

Prepare your environment#

Work from the torchtitan-neuron directory you installed in How to set up TorchTitan Neuron and verified in Tutorial: Run a basic TorchTitan training on Neuron.

Launches use the run_train.sh wrapper and the NGPU / MODULE / CONFIG environment variables introduced in the basic tutorial, under “Neuron environment variables of interest”. The only difference here is that MODULE takes a fully qualified path. Extra arguments pass straight through to torchtitan.train, which is how every override below works.

Important

Step 2 runs ten steps on c4_test, a small split bundled with the repository rather than the C4 corpus, because this tutorial is about confirming the path works, not about training a model. A loss curve from a run like that tells you the pipeline is wired correctly and little else. Step 3 shows what to change for a real run.

Step 1: Prepare the tokenizer#

The recipe reads its tokenizer from a small bundled path, so you can run Step 2 with no download at all. Do the download now only if you intend to follow Step 3 and train on real data:

python scripts/download_hf_assets.py \
  --repo_id openai/gpt-oss-20b --assets tokenizer

Assets land in ./assets/hf/gpt-oss-20b/. The dataset needs no download either: both recipes read c4_test, a small split bundled with the repository, so Step 2 needs no network access at all. Step 3 is where that changes.

Step 2: Launch the training run#

In this step, you will launch GPT-OSS 20B with NKI kernels on all 64 NeuronCores, once with each sharding implementation.

Start with FSDP v2:

NGPU=64 \
MODULE=torchtitan_neuron.models.gpt_oss \
CONFIG=gpt_oss_20b_nki_tp4_fsdp16 \
./run_train.sh --training.steps 10

Then SimpleFSDP. The model, the kernels, and the parallelism layout are the same; the sharding implementation, the recipe module, and the batch size change:

NGPU=64 \
MODULE=torchtitan_neuron.experiments.graph_trainer.gpt_oss \
CONFIG=simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16 \
./run_train.sh --training.steps 10

Both should step, and loss should drop sharply at the start and keep trending down. Expect it to bounce on the way — individual steps can rise while the trend falls, so read the trend rather than each step. Getting both recipes to run is the useful outcome here — the Neuron compile path works for this model under either implementation.

Note

The two recipes ship different batch sizes, so what you have just run is not a measurement of one against the other.

You do not pass --compile.enable, --compile.backend, an attention backend, or the MoE conversion by hand. The recipe does all of it, and reading it is the fastest way to see what the run is actually doing. In torchtitan_neuron/models/gpt_oss/config_registry.py:

def gpt_oss_20b_nki_tp4_fsdp16() -> Trainer.Config:
    model_spec = model_registry("20b", attn_backend="nki")
    for layer in model_spec.model.layers:
        layer.moe = to_nki_moe_config(layer.moe)
        # Neuron does not support the RoPE gather; use the prefix-cache path.
        layer.attention.rope.use_prefix_cache_for_positions = True
    return Trainer.Config(
        compile=CompileConfig(
            enable=True,
            components=["model", "loss", "grad_clip"],
            backend="neuron",
        ),
        # ... remaining configuration unchanged
    )

The first three lines are the Neuron-specific part: the 20b flavor built with the nki attention backend, every MoE layer converted to the NKI implementation, and the rotary workaround applied per layer. The rest of the recipe is ordinary TorchTitan configuration.

Note which components get compiled: the model, the loss, and grad_clip. Gradient-norm clipping is compiled along with the rest rather than left to run eagerly, which is what Neuron needs. The SimpleFSDP recipe inherits this, because it is built by wrapping this one.

Note

Expect the first step to take far longer than the ones after it. This recipe compiles the model with the Neuron backend, and the NKI attention and MoE kernels are compiled too. The Neuron stack caches the results on disk, so later runs of the same recipe start much faster. See TorchNeuron environment variables for the cache directories.

Step 3: Confirm success and move to a real run#

In this step, you will confirm the Step 2 runs were healthy, then launch a longer run on real data.

Check them against the Confirmation list below before going further.

Moving to real data changes three things. The dataset becomes c4, the asset path moves to the tokenizer you downloaded in Step 1, and the step count has to be set explicitly — these recipes ship a ten-step schedule, so dropping --training.steps leaves you at ten rather than handing you a longer run.

Run both recipes again, this time on real data. --training.global_batch_size 64 is the one addition: the recipes ship different per-rank batch sizes, so pinning the global batch keeps both runs training on the same number of sequences per optimizer step. A real workload would choose its own batch size rather than inherit this one.

FSDP v2:

NGPU=64 \
MODULE=torchtitan_neuron.models.gpt_oss \
CONFIG=gpt_oss_20b_nki_tp4_fsdp16 \
./run_train.sh \
  --training.steps 100 \
  --training.global_batch_size 64 \
  --dataloader.dataset c4 \
  --hf_assets_path=./assets/hf/gpt-oss-20b

SimpleFSDP:

NGPU=64 \
MODULE=torchtitan_neuron.experiments.graph_trainer.gpt_oss \
CONFIG=simple_fsdp_gpt_oss_20b_nki_tp4_fsdp16 \
./run_train.sh \
  --training.steps 100 \
  --training.global_batch_size 64 \
  --dataloader.dataset c4 \
  --hf_assets_path=./assets/hf/gpt-oss-20b

Note

You do not need to let either finish to complete the tutorial. Once you have watched the loss fall over a stretch of steps, stop the run with Ctrl-C in the terminal running run_train.sh, which signals the torchrun launcher and shuts down every rank.

Training a mixture-of-experts model on a shortened learning rate schedule carries risks to convergence, so do not read these runs as production training. The first fifty or so steps swing hard — loss can move by several points between consecutive steps — and then the curve settles and falls steadily, with the step-to-step movement shrinking to a fraction of a point. Training with a longer warmup over more steps gives more stable training. Both implementations behave the same way here.

The GPT-OSS 20B configs on Neuron do not use expert parallelism; both scale with tensor parallelism and FSDP.

Confirmation#

The Step 2 runs succeeded if all of the following hold, for each of the two recipes:

  • Training starts stepping and per-step loss lines appear.

  • Loss values are finite — never nan — and lower by the tenth step than the first.

  • neuron-top shows the NeuronCores in use during the run.

  • The second step completes far faster than the first, which is the one-time trace and compile cost being paid on step one and not again.

Note

Each recipe compiles separately, so the SimpleFSDP launch pays its own first-step cost even after the FSDP v2 one has run. Expect it to be the longer of the two.

Congratulations! You have now pre-trained GPT-OSS with TorchTitan on AWS Trainium. If you encountered any issues, see the Common issues section below.

Common issues#

  • Out of memory: the recipe already uses full activation checkpointing, so the next lever is sequence length or batch size. Halve --training.seq_len before changing the model configuration. See Best practices for TorchTitan Neuron.

  • The run hangs at startup: NGPU must match the NeuronCores actually available — 64 on a trn2.48xlarge. Confirm with neuron-ls.

  • “Config function not found”: MODULE needs the path to the Neuron recipe module, not a short name. gpt_oss resolves to the upstream registry, which has no Neuron recipes.

  • Tokenizer or dataset mismatch: the recipe points at a bundled test tokenizer by default. If you overrode --hf_assets_path, confirm it matches the directory download_hf_assets.py created.

  • Compilation fails: confirm you have not overridden the attention backend. The NKI attention path is selected by the recipe and is what the rest of the configuration assumes. For the SimpleFSDP recipe, also confirm you have not overridden --compile.mode; JIT is the only mode supported for SimpleFSDP on Neuron in this release.

  • One recipe runs and the other does not: they share the model, the kernels, and the parallelism layout, so this narrows the problem to the sharding implementation rather than to any of those.

Clean up#

Stop the training processes (Ctrl-C on the torchrun launcher). Remove the output TorchTitan wrote under ./outputs if you no longer need it, and stop or terminate the instance, since it bills while running.

Note

Running on Trn3: The instance-dependent values on this page, such as the number of ranks per node, the device memory (HBM) per rank, NEURON_RT_NUM_CORES, and the device mesh shapes, are for Trn2 (trn2.48xlarge). Trn3 servers have a different number of Trainium chips, so these values do not apply to Trn3 without change. Before you run on Trn3, calculate them again from the Amazon EC2 Trn3 architecture and Trainium3 architecture specifications. For the effect of the LNC setting on the number of visible cores, see Logical NeuronCore configuration.

Next steps#

Further reading#

This document is relevant for: Trn2, Trn3