AWS Trainium 50 Exercises #1: Let’s Get Started with Trainium

This article is Chapter 1 of the six-part series “AWS Trainium 50 Exercises,” designed to help you gain practical knowledge for performing distributed LLM training on AWS Trainium — by doing it hands-on.

0. Introduction
1. Let’s Get Started with Trainium) (Exercises 1–7)　← This article
2. Fundamentals of Trainium (Exercises 8–10)
3. Building a Compute Cluster (Exercises 11–18)
4. Training NxD-Ready Models (Exercises 19–32)
5. Fundamentals of Distributed Training (Exercises 33–37)
6. Re-porting Llama 3 to Trainium (Exercises 38–50)

The content of this “50 Exercises” series is the responsibility of Karakuri. We have taken great care to ensure accuracy^[1], but please refer to official AWS documentation for the most reliable and up-to-date information.

Chapter 1: Let’s Get Started with Trainium

In this chapter, we assume the following prerequisites:

• You have an AWS account
• You understand basic shell operations
• You have a basic understanding of PyTorch

Exercises (1–7)

To begin, let’s get a feel for what Trainium is by actually trying it out.

1. Launch an EC2 instance of type trn1.2xlarge. Be sure to select the following AMI (machine image):

   • Deep Learning AMI Neuron (Ubuntu 22.04)

     

     • When searching for AMIs that include Neuron in their names, you may also find an Amazon Linux 2023 version in addition to the Ubuntu 22.04 version shown above. The choice will lead to slight differences in subsequent steps (this tutorial assumes the Ubuntu 22.04 version).

   Instances with names starting with trn are equipped with Trainium chips. The trn1.2xlarge instance is the most affordable among them and costs $1.34/hour (as of 2025-05-09, on-demand pricing in us-east-2). Available regions are limited (as of that date, Japan region is not yet supported).

2. Log in to the instance via SSH.

   • The username is ubuntu (you won’t be able to connect using ec2-user).

3. To confirm that your instance is indeed equipped with a Trainium chip, run the neuron-top command.

   If successful, you’ll see a screen like this. It allows you to monitor the real-time activity of the Trainium cores. The machine shown here has two cores, NC0 and NC1, whose utilization is displayed. Press q to exit.

   

4. Activate a Python virtual environment using one of the following methods (either is fine).

   Deep Learning AMI Neuron comes preinstalled with the latest Neuron drivers and virtual environments, so you can use an environment tailored to your framework or workload.

   • For training (PyTorch 2.7 + NxD Training library):

     • Activate the environment and run the setup script setup_nxdt.sh:

       source /opt/aws_neuronx_venv_pytorch_2_7_nxd_training/bin/activate
       setup_nxdt.sh
       

   • For inference (PyTorch 2.7 + NxD Inference library):

     source /opt/aws_neuronx_venv_pytorch_2_7_nxd_inference/bin/activate
     

   Deep Learning AMI Neuron is periodically updated to support the latest Neuron SDK releases. Right after a new SDK release, please check the AMI release date to ensure it’s already up-to-date.

   To verify that the installed drivers and libraries are the latest versions, compare the versions listed by the commands below with those in the Neuron documentation.

   dpkg -l | grep neuron
   pip list | grep -e neuron -e torch
   

5. Once inside the Python virtual environment, launch the Python interactive console and execute the following code.

   First, let’s perform the computation on the CPU.

   >>> import torch
   >>> x1 = torch.arange(6).reshape(2, 3).to(dtype=torch.bfloat16)
   >>> y1 = (x1 @ x1.T).flatten()
   >>> z1 = y1[1:]
   >>> z1
   

   You should see the output: tensor([14., 14., 50.], dtype=torch.bfloat16)
   Next, let’s perform the same computation on Trainium. Continue by running the following:

   >>> import torch_xla
   >>> x2 = x1.to("xla")
   >>> y2 = (x2 @ x2.T).flatten()
   >>> z2 = y2[1:]
   >>> z2
   

   • If you see the following warnings, you can safely ignore them. (REF: Official Docs)
     Show warning content

     2025-08-21 07:25:18.566356: W neuron/nrt_adaptor.cc:53] nrt_tensor_write_hugepage() is not available, will fall back to nrt_tensor_write().
     2025-08-21 07:25:18.566388: W neuron/nrt_adaptor.cc:62] nrt_tensor_read_hugepage() is not available, will fall back to nrt_tensor_read().
     2025-Aug-21 07:25:18.0568 3598:4732 [1] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):213 CCOM WARN NET/OFI Failed to initialize sendrecv protocol
     2025-Aug-21 07:25:18.0573 3598:4732 [1] int nccl_net_ofi_create_plugin(nccl_net_ofi_plugin_t**):354 CCOM WARN NET/OFI aws-ofi-nccl initialization failed
     2025-Aug-21 07:25:18.0577 3598:4732 [1] ncclResult_t nccl_net_ofi_init_no_atexit_fini_v6(ncclDebugLogger_t):183 CCOM WARN NET/OFI Initializing plugin failed
     2025-Aug-21 07:25:18.0582 3598:4732 [1] net_plugin.cc:97 CCOM WARN OFI plugin initNet() failed is EFA enabled?
     

   The call .to("xla") moves x2 onto the Trainium device, and the same applies to y2 and z2. After compilation messages appear, there will be a short delay before the final result of z2 is shown. (The first run may take several minutes instead of a few seconds.)

   2025-05-15 07:32:34.000963:  42146  INFO ||NEURON_CC_WRAPPER||: Call compiler with cmd: neuronx-cc compile --framework=XLA /tmp/ubuntu/neuroncc_compile_workdir/d98fe98d-019a-4d6d-b127-0efdd43b88f2/model.MODULE_6654042255826386093+e30acd3a.hlo_module.pb --output /tmp/ubuntu/neuroncc_compile_workdir/d98fe98d-019a-4d6d-b127-0efdd43b88f2/model.MODULE_6654042255826386093+e30acd3a.neff --target=trn1 --verbose=35
   .Completed run_backend_driver.
   
   Compiler status PASS
   tensor([14., 14., 50.], device='xla:0', dtype=torch.bfloat16)
   

   Computation on Trainium is performed lazily.
   This means the actual computation is deferred until the value is needed. The operations y2 = (x2 @ x2.T).flatten() and z2 = y2[1:] don’t execute the computation immediately; instead, the computation graph is recorded. When you finally print z2, the graph is compiled (i.e., translated into executable instructions for the Trainium chip), and once compilation succeeds (Compiler status PASS), the computation is executed.

6. Understanding when compilation and lazy evaluation occur is crucial when porting models to Trainium. Try the following and observe when the compiler messages appear.

   >>> import torch
   >>> import torch_xla
   >>> a = torch.randn(4, dtype=torch.bfloat16).to("xla")
   >>> b = a * a * a
   >>> b
   >>> c = torch.randn(4, dtype=torch.bfloat16).to("xla")
   >>> d = c * c * c
   >>> d
   >>> e = torch.randn(4, dtype=torch.bfloat16).to("xla")
   >>> f = e * e
   >>> f
   >>> g = torch.randn(5, dtype=torch.bfloat16).to("xla")
   >>> h = g * g
   >>> h
   

   You’ll notice compilation occurs when displaying b, f, and h, but not when displaying d.
   This illustrates:

   • d is computed by cubing c, which matches the computation graph used for a → b. Since that graph is already compiled, no recompilation occurs.
   • f squares e, a different computation from a → b, so it must be compiled.
   • h also squares a tensor, but because g has a different shape from e, it’s treated as a different computation graph, triggering compilation again.

   Note: For more on lazy mode vs eager mode, see:

   • https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/programming-guide/training/pytorch-neuron-programming-guide.html#understand-the-lazy-mode-in-pytorch-neuronx
   • https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/programming-guide/inference/trace-vs-xla-lazytensor.html

   Despite this “lazy” behavior, why use Trainium at all? Because its parallel compute cost-performance is outstanding.

7. Stop and terminate your trn1.2xlarge instance.

   Forgetting to do this will keep incurring costs ($1.34/hour in us-east-2 as of 2025-05-15), so please make sure to shut it down.

脚注

1. We would like to thank Mr. Tokoyo from AWS for his supervision of this material.