This document is relevant for: Trn1

Track System Resource Utilization during Training with neuron-monitor using PyTorch Neuron#

This tutorial explains how to monitor resource utilization using neuron-monitor, Prometheus and Grafana while running a multi-layer perceptron MNIST model on Trainium using PyTorch Neuron.

Multi-layer Perceptron MNIST Model#

This tutorial is based on the MNIST example for PyTorch Neuron on Trainium. For the full tutorial, please see Multi-Layer Perceptron Training Tutorial.

The Training Job#

For this tutorial, we will make the original script do more work thus giving us more system utilization data to observe. The training loop is simply repeated 1000 times:

for run in range(0, 1000):
    print(f'Run {run}')

Save the following code as and you can run it as python3 on a Trn1 instance.

import os
import time
import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision.datasets import mnist
from torch.optim import SGD
from import DataLoader
from torchvision.transforms import ToTensor

# XLA imports
import torch_xla.core.xla_model as xm

# Declare 3-layer MLP for MNIST dataset
class MLP(nn.Module):
    def __init__(self, input_size = 28 * 28, output_size = 10, layers = [120, 84]):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_size, layers[0])
        self.fc2 = nn.Linear(layers[0], layers[1])
        self.fc3 = nn.Linear(layers[1], output_size)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

# Load MNIST train dataset
train_dataset = mnist.MNIST(root='./MNIST_DATA_train', \
                            train=True, download=True, transform=ToTensor())

def main():
    # Prepare data loader
    train_loader = DataLoader(train_dataset, batch_size=32)

    # Fix the random number generator seeds for reproducibility

    # XLA: Specify XLA device (defaults to a NeuronCore on Trn1 instance)
    device = 'xla'

    # Move model to device and declare optimizer and loss function
    model = MLP().to(device)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    loss_fn = torch.nn.NLLLoss()

    # Run the training loop
    print('----------Training ---------------')
    for run in range(0, 1000):
        print(f'Run {run}')
        for idx, (train_x, train_label) in enumerate(train_loader):
            train_x = train_x.view(train_x.size(0), -1)
            train_x =
            train_label =
            output = model(train_x)
            loss = loss_fn(output, train_label)
            xm.mark_step() # XLA: collect ops and run them in XLA runtime
            if idx < 2: # skip warmup iterations
                start = time.time()

    # Save checkpoint for evaluation
    os.makedirs("checkpoints", exist_ok=True)
    checkpoint = {'state_dict': model.state_dict()}
    # XLA: use instead of to ensure states are moved back to cpu
    # This can prevent "XRT memory handle not found" at end of execution,'checkpoints/')

    print('----------End Training ---------------')

Setting up Prometheus and Grafana#


The setup presented in the following paragraphs can be extended to monitor any number of instances running training jobs or inference workloads. For this tutorial, we will set everything up on a single Trn1 instance running Amazon Linux 2.

Setting up Prometheus#

For a more detailed guide on how to install Prometheus visit their official guide at

Download and unzip a prebuilt Prometheus binary on your Trn1 instance:

tar -xzvf prometheus-2.38.0.linux-amd64.tar.gz
cd prometheus-2.38.0.linux-amd64/

Create a config and add a scrape target:

vim prometheus.yml
- job_name:       'neuron'

# Scrape target every 5 seconds.
scrape_interval: 5s

    - targets: ['localhost:8000']

Finally, start Prometheus:

./prometheus --config.file=prometheus.yml

Setting up Grafana#

For a more detailed guide on how to install Grafana visit their official guide at

Add the Grafana repo to yum:

sudo vim /etc/yum.repos.d/grafana.repo


Install and start Grafana:

sudo yum install -y grafana
sudo /bin/systemctl start grafana-server.service

By default, Grafana will run a HTTP server on port 3000. If you need to change that, update its config and restart the service:

sudo vim /etc/grafana/grafana.ini
sudo /bin/systemctl start grafana-server.service

Using your favorite web browser, access the Grafana webpage and add a new dashboard.

The default user and password are both ‘admin’:

Image: image.png

Next, you’ll add a Prometheus data source by going to Configuration -> Data Sources:

Image: image.png

… and adding the local Prometheus server as a data source:

Image: image.png

Finally, upload the sample dashboard neuron-monitor-grafana.json <src/examples/neuron-monitor/neuron-monitor-grafana.json to Grafana:

Image: image.png

Monitoring the Training Workload#

Start the training job which, due to the artificially added complexity, will take more than 15 minutes:


On the same instance, start neuron-monitor and its companion script,

neuron-monitor |

Once they are running, you can use your web browser, access the Grafana server running on your Trn1 instance and view a timeline of the system utilization.

The upper part of the dashboard contains:
  • a list of the currently monitored instances (for this tutorial there is a single Trn1 instance)

  • aggregated metrics for stats such as NeuronCore utilization, NeuronCores in use, iteration success rates, error rates etc.

  • a timeline of execution status rates and execution latencies

Image: image.png

The lower part of the dashboard contains: - one line of charts containing a timeline of Neuron resource utilization (NeuronCore, vCPU and memory utilization) - one line of charts containing a timeline of host resource utilization (vCPU and memory utilization)

Image: image.png

This document is relevant for: Trn1