This document is relevant for: Trn2, Trn3

Tutorial: Train with DDP, FSDP2, and DTensor on Neuron

Run the same training job three ways on a full trn2.48xlarge instance: DDP, FSDP2, and DTensor tensor parallelism. Each job uses 64 workers and records its outputs, gradients, and parameter updates. Run the same distributed job on CPU with Gloo, then on Neuron, and compare the collected results rank by rank. Then combine SimpleFSDP and TP inside compiled forward and backward graphs.

Overview

Use the same launch script, rank count, parallelism plan, and training step for the CPU and Neuron jobs. Change --device and --backend between runs. Both jobs use the same model, global batch, optimizer, and initial parameters. Compiled jobs also change the compiler backend for the selected device.

Mode	What each worker owns	Training data per worker
DDP	A complete model replica; gradients are averaged across workers.	Eight samples from the global batch.
FSDP2	Parameter and gradient shards; layers gather weights when needed.	Eight samples from the global batch.
DTensor / TP	A partition of each linear layer.	The same global batch of 512 samples on every worker.

Before you start

Read How to use torch.distributed on Neuron first. That guide explains process-group initialization, rank assignment, parallelism choices, EFA for multi-node execution, and HostCC overlap.

Complete the single-device check in Quickstart: Install and verify torch-neuronx. You should also be familiar with a PyTorch training loop: forward, loss, backward, and optimizer step.

Model details

• Two linear layers: 128 input features, 256 hidden features, and 128 outputs.

• ReLU activation and mean squared error loss.

• Float32 parameters and deterministic synthetic data.

• SGD with a learning rate of 0.01.

• Ten training steps, seed 0, and a global batch of 512 samples.

The comparison tool checks separately collected CPU and Neuron distributed results. The training scripts materialize parameter shards for verification. Use these checks to establish correctness before performance tuning. A per-rank PASS line confirms that the run completed. Each comparison must report PASS comparison ranks=64 steps=10 to establish numerical correctness.

Prerequisites

• Exclusive access to a trn2.48xlarge with all 64 logical NeuronCores available under logical NeuronCore configuration 2.

• A native TorchNeuron environment with torch-neuronx and PyTorch 2.14.

• A successful single-device run.

For multi-node training beyond this single-node tutorial, install EFA on every node. Follow Step 3: Install EFA in the Neuron manual setup instructions, which explicitly cover Trn2.

Prepare your environment

Activate your TorchNeuron environment, then make all cores available to the launcher. Set the logical-core configuration before starting Python.

unset NEURON_RT_VISIBLE_CORES NEURON_RT_NUM_CORES TORCH_NEURONX_ENABLE_STREAMS
export NEURON_LOGICAL_NC_CONFIG=2
export TORCH_NEURONX_ENABLE_HOST_CC=0
export OMP_NUM_THREADS=1
neuron-ls
mkdir -p neuron-distributed-tutorial
cd neuron-distributed-tutorial

Confirm that neuron-ls reports logical-core configuration 2, core IDs 0 through 63, and no other processes using the cores. Streams are enabled by default; the command above removes any previous streams override. These correctness runs keep HostCC disabled. Follow the distributed guide to enable HostCC overlap after establishing parity.

Step 1: Save the shared code and run distributed CPU references

Download these files into the current directory:

• training_common.py: model, data, training loop, and editable configuration defaults.

• distributed_train.py: distributed DDP, FSDP2, and TP launch script for both CPU and Neuron.

• compare_training.py: result comparison.

• reference_train.py: optional CPU/Gloo entry point. It calls the same distributed trainer with CPU defaults.

CPU reference execution is a separate distributed job. Each CPU and Neuron worker saves its own results, and the comparison tool checks matching ranks after training. These saved tensors are correctness artifacts, not throughput measurements.

The device, process-group backend, rank count, TP degree, seed, and training steps are configurable at the top of training_common.py:

# Edit these defaults or use the command-line options in the launch scripts.
DEVICE_TYPE = "neuron"
PROCESS_GROUP_BACKEND = "neuron"
RANKS_PER_NODE = 64
TP_DEGREE = 4
SEED = 0
STEPS = 10
GLOBAL_BATCH_SIZE = 512
INPUT_FEATURES = 128
HIDDEN_FEATURES = 256
LEARNING_RATE = 0.01

The launch scripts also accept --device, --backend, --expected-world-size, --seed, and --steps. Set device and backend options together: --device cpu --backend gloo for the reference and --device neuron --backend neuron for Neuron. Both jobs use 64 ranks. Compiled scripts also accept --compile-backend: use aot_eager on CPU and neuron on Neuron. The CPU job captures the same forward and backward graphs with PyTorch’s AOTAutograd backend.

make_model calls torch.manual_seed(seed) before creating CPU-initialized parameters. Both independent runs use seed 0, the same deterministic synthetic data, and the same optimizer. Generate a ten-step distributed CPU reference for each eager mode:

for mode in ddp fsdp tp; do
    torchrun --standalone --nproc-per-node=64 distributed_train.py \
        --parallelism "$mode" --device cpu --backend gloo \
        --seed 0 --steps 10 --output-dir "results/cpu-$mode"
done

The CPU reference uses the same DDP wrapper, FSDP2 sharding, or TP plan as its Neuron counterpart. It performs distributed forward and backward computation and collective communication. Use PyTorch 2.14 for both jobs. reference_train.py is a convenience alias with CPU/Gloo defaults, for example:

torchrun --standalone --nproc-per-node=64 reference_train.py \
    --parallelism ddp --seed 0 --steps 10 --output-dir results/cpu-ddp

The shared training loop is:

def train(parameters, optimizer, steps, train_step):
    """Run one model; train_step performs its forward and backward computation."""
    if steps < 1:
        raise ValueError("Use at least one training step.")
    records = []
    for step in range(steps):
        optimizer.zero_grad()
        before = [full_cpu(p) for p in parameters]
        output, loss = train_step(step)
        if any(p.grad is None for p in parameters):
            raise AssertionError("A model parameter has no gradient.")
        gradients = [full_cpu(p.grad) for p in parameters]
        optimizer.step()
        updated = [full_cpu(p) for p in parameters]
        deltas = [value - initial for value, initial in zip(updated, before, strict=True)]
        if not any(bool(torch.count_nonzero(delta)) for delta in deltas):
            raise AssertionError("The optimizer did not update any parameters.")
        records.append(
            {
                "output": None if output is None else full_cpu(output),
                "loss": None if loss is None else float(loss.detach().cpu()),
                "gradients": gradients,
                "parameters": updated,
                "deltas": deltas,
            }
        )
        print(f"step={step + 1}/{steps} training=completed", flush=True)
    return records

The distributed configuration is:

"""Run the same distributed training plan on CPU or Neuron."""

import argparse

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.fsdp import fully_shard
from torch.distributed.tensor import DTensor
from torch.distributed.tensor.parallel import ColwiseParallel, RowwiseParallel, parallelize_module
from torch.nn.parallel import DistributedDataParallel
from training_common import (
    DEVICE_TYPE,
    GLOBAL_BATCH_SIZE,
    HIDDEN_FEATURES,
    INPUT_FEATURES,
    LEARNING_RATE,
    PROCESS_GROUP_BACKEND,
    add_run_arguments,
    initialize,
    make_batch,
    make_model,
    save_results,
    train,
)


def main(default_device=DEVICE_TYPE, default_backend=PROCESS_GROUP_BACKEND):
    parser = argparse.ArgumentParser()
    parser.add_argument("--parallelism", choices=("ddp", "fsdp", "tp"), required=True)
    add_run_arguments(parser, device=default_device, backend=default_backend)
    args = parser.parse_args()
    rank, world_size, device = initialize(args)
    try:
        model = make_model("eager", args.seed).to(device)
        data_parallel = args.parallelism in ("ddp", "fsdp")
        if data_parallel and GLOBAL_BATCH_SIZE % world_size:
            raise ValueError("The global batch must divide evenly across DP workers.")
        if args.parallelism != "ddp" and (
            HIDDEN_FEATURES % world_size or INPUT_FEATURES % world_size
        ):
            raise ValueError("Linear dimensions must divide evenly across the mesh.")
        if args.parallelism == "ddp":
            model = DistributedDataParallel(model)
        elif args.parallelism == "fsdp":
            mesh = init_device_mesh(args.device, (world_size,), mesh_dim_names=("dp",))
            # Layers 0 and 2 are Linear modules with parameters. Layer 1 is
            # ReLU: it has no parameters to shard, so it needs no FSDP wrapper.
            fully_shard(model[0], mesh=mesh, reshard_after_forward=True)
            fully_shard(model[2], mesh=mesh, reshard_after_forward=True)
            fully_shard(model, mesh=mesh)
        else:
            mesh = init_device_mesh(args.device, (world_size,), mesh_dim_names=("tp",))
            model = parallelize_module(
                model, mesh, {"0": ColwiseParallel(), "2": RowwiseParallel()}
            )
        if args.parallelism != "ddp":
            assert isinstance(model[0].weight, DTensor)
            assert isinstance(model[2].weight, DTensor)
            assert tuple(model[0].weight.to_local().shape) == (
                HIDDEN_FEATURES // world_size,
                INPUT_FEATURES,
            )
            shape = (
                (INPUT_FEATURES // world_size, HIDDEN_FEATURES)
                if args.parallelism == "fsdp"
                else (INPUT_FEATURES, HIDDEN_FEATURES // world_size)
            )
            assert tuple(model[2].weight.to_local().shape) == shape
        parameters = list(model.parameters())
        optimizer = torch.optim.SGD(parameters, lr=LEARNING_RATE)

        def train_step(step):
            batch, target = make_batch(step)
            if data_parallel:
                batch = batch.chunk(world_size)[rank].contiguous()
                target = target.chunk(world_size)[rank].contiguous()
            output = model(batch.to(device))
            loss = torch.nn.functional.mse_loss(output, target.to(device))
            loss.backward()
            global_loss = loss.detach().clone()
            if data_parallel:
                dist.all_reduce(global_loss)
                global_loss.div_(world_size)
            return output, global_loss

        records = train(parameters, optimizer, args.steps, train_step)
        save_results(
            args,
            records,
            model="eager",
            mode=args.parallelism,
            rank=rank,
            world_size=world_size,
            dp_degree=world_size if data_parallel else 1,
            dp_rank=rank if data_parallel else 0,
            tp_degree=1 if data_parallel else world_size,
            tp_rank=0 if data_parallel else rank,
        )
        dist.barrier()
        print(f"PASS rank={rank} mode={args.parallelism} world_size={world_size}", flush=True)
    finally:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

For DDP and FSDP2, each rank receives a different eight-sample slice of the global batch. Averaging those workers’ gradients matches the gradient of the CPU loss over all 512 samples. For TP, every rank receives the complete batch; the layers partition the feature dimension.

Step 2: Train with DDP

Launch 64 workers, one per logical NeuronCore:

torchrun --standalone --nproc-per-node=64 \
    distributed_train.py --parallelism ddp --seed 0 --steps 10 \
    --output-dir results/ddp
python compare_training.py --reference-dir results/cpu-ddp \
    --actual-dir results/ddp --world-size 64

DDP(model) creates a model replica on each worker and reduces gradients during backward. The comparison tool checks the recorded averaged gradients and SGD updates against the matching distributed CPU/Gloo job.

Every worker should complete ten steps and print PASS rank=<rank> mode=ddp world_size=64. This line confirms that the run finished. The comparison is the numerical correctness check and must report PASS comparison ranks=64 steps=10 before you continue.

Step 3: Train with FSDP2

Run the same workload with sharded parameters:

torchrun --standalone --nproc-per-node=64 \
    distributed_train.py --parallelism fsdp --seed 0 --steps 10 \
    --output-dir results/fsdp
python compare_training.py --reference-dir results/cpu-fsdp \
    --actual-dir results/fsdp --world-size 64

The script applies fully_shard to each linear layer, then to the root model, before creating the optimizer. Layers 0 and 2 are linear layers with parameters; layer 1 is ReLU and has no parameters to shard. Each worker owns four rows of the first layer’s weight and two rows of the second layer’s weight.

The CPU and Neuron jobs use the same mesh and layer sharding. The script records full gradients and parameters for the separate comparison. Every worker should finish with PASS rank=<rank> mode=fsdp world_size=64.

Step 4: Train with DTensor tensor parallelism

Run the same model with partitioned linear layers:

torchrun --standalone --nproc-per-node=64 \
    distributed_train.py --parallelism tp --seed 0 --steps 10 \
    --output-dir results/tp
python compare_training.py --reference-dir results/cpu-tp \
    --actual-dir results/tp --world-size 64

ColwiseParallel partitions the first layer’s 256 output features into four features per rank. RowwiseParallel consumes that partitioned input in the second layer and returns a complete output on every worker.

Note

This step uses all 64 ranks as one TP group to check the sharding path. For training workloads, start with chip-local TP4 and use the remaining ranks for data parallelism, as in Step 5. Wider TP needs a supported rank placement and measured benefit for your model. See Best practices for training on AWS Neuron for topology-aware mesh choices, including a non-default TP8 layout.

All workers process the same global batch. The comparison tool checks their outputs, full gradients, and updates against the matching CPU TP ranks. Every worker should finish with PASS rank=<rank> mode=tp world_size=64.

Step 5: Compile SimpleFSDP and TP together

SimpleFSDP represents parameter gathering and gradient reduction as DTensor operations inside the model computation. The compiler can capture those collectives together with tensor-parallel computation.

Download simplefsdp_tp_train.py. This complete example uses a 16 x 4 mesh: 16 FSDP ranks and four TP ranks, for 64 workers total. It uses a bias-free version of the two-layer model, and records ten steps for comparison with the same distributed script on CPU. Both jobs use the same mesh and SimpleFSDP parameter getter. The shared model factory applies the same seed before initialization.

Each worker stores one sixty-fourth of both weights:

Parameter	DTensor placements (FSDP, TP)	Local shape
First weight (256 x 128)	(Shard(1), Shard(0))	(64, 8)
Second weight (128 x 256)	(Shard(0), Shard(1))	(8, 64)

The parameter getter gathers each TP partition across the FSDP dimension. Its backward operation sums and reduce-scatters gradients to the stored shards. Checkpointing the getter recomputes gathers needed by backward, rather than retaining the gathered weights.

class SimpleFSDPWeight(torch.nn.Module):
    """Gather a DP-sharded TP weight through differentiable DTensor operations."""

    def __init__(self, mesh, shape, dp_dim, tp_dim):
        super().__init__()
        self.mesh = mesh
        self.dp_mesh = mesh["dp"]
        self.tp_mesh = mesh["tp"]
        self.shape = shape
        self.dp_dim = dp_dim
        self.tp_dim = tp_dim
        tp_shape = list(shape)
        tp_shape[tp_dim] //= self.tp_mesh.size()
        self.tp_shape = tuple(tp_shape)

    def gather(self, weight):
        # Collapse to the DP submesh, gather this TP partition, then restore TP.
        dp_weight = DTensor.from_local(
            weight.to_local(),
            self.dp_mesh,
            (Shard(self.dp_dim),),
            shape=self.tp_shape,
            stride=(self.tp_shape[1], 1),
        )
        tp_local = dp_weight.redistribute(placements=(Replicate(),)).to_local(
            grad_placements=(Partial("sum"),)
        )
        return DTensor.from_local(
            tp_local,
            self.tp_mesh,
            (Shard(self.tp_dim),),
            shape=self.shape,
            stride=(self.shape[1], 1),
        )

    def forward(self, weight):
        # Recompute the gather in backward instead of retaining full weights.
        return checkpoint(self.gather, weight, use_reentrant=False)

The complete model is compiled with torch.compile(sharded, backend="neuron", fullgraph=True). The captured forward includes FSDP all-gathers, both linear layers, ReLU, and the TP output all-reduce. AOTAutograd generates the backward graph, including FSDP gradient reduce-scatter. The optimizer and result collection run outside these graphs. CPU training and result comparison run in separate jobs.

Launch the complete script:

torchrun --standalone --nproc-per-node=64 simplefsdp_tp_train.py \
    --device cpu --backend gloo --compile-backend aot_eager \
    --seed 0 --steps 10 --output-dir results/cpu-simplefsdp-tp
torchrun --standalone --nproc-per-node=64 simplefsdp_tp_train.py \
    --seed 0 --steps 10 --output-dir results/simplefsdp-tp
python compare_training.py --reference-dir results/cpu-simplefsdp-tp \
    --actual-dir results/simplefsdp-tp --world-size 64

Each data-parallel replica processes a different 32-sample slice of the global batch; its four TP workers receive the same slice. The script divides the local mean loss by 16 before backward because the parameter getter uses sum reduction across FSDP workers.

Confirm ten completed steps and PASS rank=<rank> mode=simplefsdp_tp world_size=64 on every worker, then PASS comparison ranks=64 steps=10. The separate comparison checks outputs, full gradients, updated parameters, update deltas, and global loss.

The comparison tool prints each step’s global CPU and Neuron loss after checking every rank. The bias-free model differs from the eager model in Steps 2 to 4.

This compact example follows the DTensor parameter-gather pattern in TorchTitan’s SimpleFSDP implementation. For the upstream compilation workflow, see the Graph Trainer examples.

Confirmation

For each Neuron job in Steps 2 to 5, confirm all of the following:

• The launcher exits with status 0.

• Ranks 0 through 63 each print a final PASS line.

• Each rank completes ten training steps; the comparison reports parity=passed against the matching CPU rank for all 64 ranks.

• The global losses match the corresponding CPU job within the script’s numerical tolerances. For the eager modes in Steps 2 to 4, they also follow the same trajectory across modes.

For each recorded step, the comparison checks outputs, gradients, updated parameters, update deltas, and global loss. The FSDP2 and TP setup also checks local weight shapes to confirm that parameters are sharded.

Each comparison prints a ten-step table with CPU and Neuron global losses. Use the table for the matching mode and mesh. Small numerical differences within the script’s tolerances are acceptable.

You have now run a training job across all 64 logical NeuronCores using DDP, FSDP2, and tensor parallelism.

Common issues

• A worker cannot allocate a core: Check for other Neuron processes with neuron-ls. Remove inherited core visibility limits and use logical-core configuration 2 before launching 64 workers.

• Workers hang in a collective: Check that every rank runs the same mode and step count. Keep batch shapes and collective order consistent.

• A parity assertion fails after changing the data: DDP/FSDP2 workers must receive different, equal-sized slices of the global batch. TP workers must receive the same complete batch. Use the same data and mesh in both jobs.

• A sharding assertion fails: Keep the layer dimensions divisible by the selected mesh dimensions. The defaults in this tutorial use 64 workers.

For more diagnoses, see Troubleshooting: TorchNeuron and How to use torch.distributed on Neuron.

Clean up

The script destroys its process group in finally and the workers exit. If you provisioned an instance for this tutorial, stop it when you finish.

Next steps

• Best practices for training on AWS Neuron – Choose TP, FSDP, and hybrid meshes for training workloads.

• How to use torch.distributed on Neuron – Configure FSDP prefetching, HostCC streams, pipeline parallelism, or multi-node execution.

• torch-neuronx operator compatibility table – Check supported APIs and their limits before adapting the example.

• Best practices for TorchNeuron – Apply performance and accuracy practices to a larger workload.

This document is relevant for: Trn2, Trn3