### Run setup script

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-training/general/installation_guide.md

Download and execute the setup script to patch installations. Ensure the script has execute permissions.

```shell
wget https://raw.githubusercontent.com/aws-neuron/neuronx-distributed-training/master/install_setup.sh
chmod +x install_setup.sh
./install_setup.sh
```

--------------------------------

### Eager Debug Mode Setup (Example 2)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/torch/torch-neuronx/programming-guide/training/pytorch-neuron-debug.md

This is a continuation of the eager debug mode example, showing the necessary environment variable setup before importing torch-xla.

```python
# Example 2

import os
# You need to set this env variable before importing torch-xla

```

--------------------------------

### Install Package

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/_content-types/procedural-tutorial.ipynb

Use this command to install necessary packages for the environment setup.

```python
# Example setup command (Remove these comments and add the CLI commands, env variable declarations, or other operations for the user to prepare their environment.)
# pip install package_name
```

--------------------------------

### Install and Start Grafana

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/tutorials/tutorial-neuron-monitor-mnist.md

Install Grafana using dnf and start the Grafana server service. This command installs the Grafana package and begins its operation.

```bash
sudo dnf install -y grafana
sudo /bin/systemctl start grafana-server.service
```

--------------------------------

### Install LLMPerf

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/disaggregated-inference-tutorial.md

Clone the LLMPerf repository and install it using pip. This is the initial setup step for running benchmarks.

```bash
git clone https://github.com/ray-project/llmperf.git
cd llmperf
pip install -e .
```

--------------------------------

### Download Example Image and Install Pillow

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/src/examples/tensorflow/tensorflow_resnet50/resnet50.ipynb

Downloads an example image ('kitten_small.jpg') for inference and installs the Pillow library, which is necessary for image manipulation.

```bash
!curl -O https://raw.githubusercontent.com/awslabs/mxnet-model-server/master/docs/images/kitten_small.jpg
!pip install pillow  # Necessary for loading images
```

--------------------------------

### Install vLLM-Neuron Plugin and Serve Model

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/neuron-inference-overview.md

Install the vLLM-Neuron plugin and start the vLLM server for inference. This example shows a sample configuration for a Trn1 instance.

```bash
##install the vllm-neuron plugin which automatically installs the right vLLM version that is supported
git clone https://github.com/vllm-project/vllm-neuron.git
cd vllm-neuron
pip install --extra-index-url=https://pip.repos.neuron.amazonaws.com -e .

##start the vLLM server to start serving inference requests (sample config for Trn1 instance)
vllm serve meta-llama/Meta-Llama-3-8B-Instruct --tensor-parallel-size 32 --max-num-seqs 4 --max-model-len 128 --block-size 32 --num-gpu-blocks-override 256
```

--------------------------------

### Clone Repository and Setup

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/torch/torch-neuronx/tutorials/training/zero1_gpt2.md

Clones the AWS Neuron samples repository and navigates to the specific tutorial directory. It then installs the required Python packages.

```shell
git clone https://github.com/aws-neuron/aws-neuron-samples.git
cd aws-neuron-samples/torch-neuronx/training/zero1_gpt2
python3 -m pip install -r requirements.txt
```

--------------------------------

### Setup Python Virtual Environment and Install Dependencies

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/CONTRIBUTING.md

Set up a Python 3.10 virtual environment and install project dependencies using pip. This is a prerequisite for building the documentation locally.

```bash
cd .. # The root folder where you have your cloned Git repos; don't run this in the repo folder but one level up or you'll have venv files in your repo folder
python3.10 -m venv venv && . venv/bin/activate
pip install -U pip
cd private-aws-neuron-sdk-staging
pip install -r requirements.txt
```

--------------------------------

### Jupyter Notebook Server Output Example

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/notebook/setup-jupyter-notebook-steps-troubleshooting.md

Example log output when starting a Jupyter Notebook server, showing the URL and token for access.

```bash
[I 21:53:11.729 NotebookApp] Using EnvironmentKernelSpecManager...
[I 21:53:11.730 NotebookApp] Started periodic updates of the kernel list (every 3 minutes).
[I 21:53:11.867 NotebookApp] Loading IPython parallel extension
[I 21:53:11.884 NotebookApp] JupyterLab beta preview extension loaded from /home/ubuntu/anaconda3/lib/python3.6/site-packages/jupyterlab
[I 21:53:11.884 NotebookApp] JupyterLab application directory is /home/ubuntu/anaconda3/share/jupyter/lab
[I 21:53:12.002 NotebookApp] [nb_conda] enabled
[I 21:53:12.004 NotebookApp] Serving notebooks from local directory: /home/ubuntu/tutorial
[I 21:53:12.004 NotebookApp] 0 active kernels
[I 21:53:12.004 NotebookApp] The Jupyter Notebook is running at:
[I 21:53:12.004 NotebookApp] http://localhost:8888/?token=f9ad4086afd3c91f33d5587781f9fd8143b4cafbbf121a16
[I 21:53:12.004 NotebookApp] Use Control-C to stop this server and shut down all kernels (twice to skip confirmation).
[W 21:53:12.004 NotebookApp] No web browser found: could not locate runnable browser.

```

--------------------------------

### Quickstart vLLM with NxD Inference

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/developer_guides/vllm-user-guide.md

A minimal example demonstrating how to initialize and use vLLM with the NxD Inference framework for text generation. Ensure the VLLM_NEURON_FRAMEWORK environment variable is set.

```python
import os
os.environ['VLLM_NEURON_FRAMEWORK'] = "neuronx-distributed-inference"

from vllm import LLM, SamplingParams
llm = LLM(
    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    max_num_seqs=8,
    max_model_len=128,
    device="neuron",
    tensor_parallel_size=2)

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# note that top_k must be set to lower than the global_top_k defined in
# the neuronx_distributed_inference.models.config.OnDeviceSamplingConfig
sampling_params = SamplingParams(top_k=10, temperature=0.8, top_p=0.95)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
```

--------------------------------

### Install JAX with jax-neuronx (Combined Package)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/nki/get-started/setup-env.md

Installs the combined jax-neuronx package with all necessary dependencies. Use this for a straightforward setup.

```bash
pip install jax-neuronx[stable] --extra-index-url=https://pip.repos.neuron.amazonaws.com
```

--------------------------------

### Setup Python Virtual Environment

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-training/general/installation_guide.md

Create and activate a Python virtual environment for development. This is a prerequisite for installing the framework and its dependencies.

```shell
python3 -m venv env
source env/bin/activate
```

--------------------------------

### Quickstart vLLM Model Initialization and Text Generation

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/developer_guides/vllm-user-guide-v1.md

Initializes a vLLM model with specified parameters and generates text based on provided prompts. Ensure you have the vLLM library installed and the model available.

```python
from vllm import LLM, SamplingParams

if __name__ == '__main__':
    # Initialize the model
    llm = LLM(
        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        max_num_seqs=4,
        max_model_len=128,
        tensor_parallel_size=2,
        block_size=32,
        num_gpu_blocks_override=16
    )

    # Generate text
    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
    ]
    sampling_params = SamplingParams(temperature=0.0)
    outputs = llm.generate(prompts, sampling_params)

    for output in outputs:
        print(f"Prompt: {output.prompt}")
        print(f"Generated: {output.outputs[0].text}")
```

--------------------------------

### Start Server with Custom Neuron Config (vLLM V1)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/llama4-tutorial.ipynb

Starts a vLLM inference server with a custom Neuron configuration for vLLM V1. This command-line example sets necessary environment variables for vLLM V1 mode and Neuron artifacts.

```bash
%%bash
# Example server startup with custom Neuron configuration
export VLLM_USE_V1=1
export NEURON_COMPILED_ARTIFACTS="/home/ubuntu/llama4/traced_models/Llama-4-Scout-17B-16E-Instruct/"
export VLLM_RPC_TIMEOUT=100000


```

--------------------------------

### Install NeuronX Distributed and Clone Repository

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/neuronx-distributed/tutorials/training_llama_tp_pp.md

Installs the neuronx-distributed package and clones the repository containing example scripts. Ensure you have the necessary AWS infrastructure and ParallelCluster setup.

```ipython3
python -m pip install neuronx_distributed --extra-index-url https://pip.repos.neuron.amazonaws.com
git clone git@github.com:aws-neuron/neuronx-distributed.git
```

--------------------------------

### Install and Verify Docker

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/jax/dlc.md

Install Docker and add the current user to the 'docker' group. Log out and back in to apply group changes. Verify installation with 'docker run hello-world'.

```bash
sudo dnf install -y docker
sudo usermod -aG docker $USER
```

```bash
docker run hello-world
```

--------------------------------

### Start Prefill Server

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/disaggregated-inference-tutorial.md

Run these commands to set up environment variables and start the vLLM prefill server. This server is configured as a kv_producer.

```bash
sudo su - ubuntu
export MODEL="meta-llama/Llama-3.3-70B-Instruct"
export VLLM_BATCH=8
export MAX_LEN=8192
export ETCD="${HOST_IP}:8989"
export PORT=8000

# Remove old container
docker rm -f prefill-vllm-server1 2>/dev/null || true

# Start prefill server
docker run -d \
  --name prefill-vllm-server1 \
  --privileged \
  --device /dev/infiniband/uverbs0 \
  --shm-size=10g \
  -p ${PORT}:${PORT} \
  -e MODEL \
  -e VLLM_BATCH \
  -e MAX_LEN \
  -e ETCD \
  -e PORT \
  public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.9.1-neuronx-py310-sdk2.25.1-ubuntu22.04 \
  bash -c "exec python3 -m vllm.entrypoints.openai.api_server \
    --model $MODEL \
    --max-num-seqs $VLLM_BATCH \
    --max-model-len $MAX_LEN \
    --tensor-parallel-size 64 \
    --device neuron \
    --speculative-max-model-len $MAX_LEN \
    --override-neuron-config '{}' \
    --kv-transfer-config '{\"kv_connector\":\"NeuronConnector\",\"kv_role\":\"kv_producer\",\"kv_buffer_size\":2e11,\"etcd\":\"$ETCD\"}' \
    --port $PORT"
```

--------------------------------

### Run nccom-test with Slurm and Custom Setup Script

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-sys-tools/nccom-test.md

Execute collective operations across multiple nodes using Slurm. This example allocates a Slurm job with two nodes and runs a custom setup script on each node before starting the benchmark.

```bash
nccom-test -r 64 -N 2 allr --slurm-mode --slurm-setup-script path/to/my/custom-setup-script.sh
```

--------------------------------

### Replace Apex setup.py Contents

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-training/general/installation_guide.md

Replace the existing setup.py file with this content to configure a slim CPU version of Apex.

```python
import sys
import warnings
import os
from packaging.version import parse, Version

from setuptools import setup, find_packages
import subprocess

import torch
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME, load

setup(
    name="apex",
    version="0.1",
    packages=find_packages(
        exclude=("build", "csrc", "include", "tests", "dist", "docs", "tests", "examples", "apex.egg-info",)
    ),
    install_requires=["packaging>20.6",],
    description="PyTorch Extensions written by NVIDIA",
)
```

--------------------------------

### Example: Run neuron-monitor-prometheus.py with Kubernetes info

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-sys-tools/neuron-monitor-user-guide.md

An example demonstrating the combined usage of neuron-monitor-prometheus.py with Kubernetes info enabled and piping to neuron-monitor-k8s-info.py, specifying a port and a 30-second period.

```bash
neuron-monitor | neuron-monitor-prometheus.py --port 8008 --enable-k8s-info | neuron-monitor-k8s-info.py --period 30
```

--------------------------------

### Install LLMPerf from Source

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/developer_guides/llm-inference-benchmarking-guide.md

Sets up a virtual environment and installs the LLMPerf library from its source repository. Ensure you are in the desired directory before running.

```bash
python3 -m venv llmperf-env
source llmperf-env/bin/activate

git clone https://github.com/ray-project/llmperf.git ~/llmperf
cd ~/llmperf
pip install -e .
```

--------------------------------

### Example Function Python

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/_content-types/conceptual-deep-dive.rst

This is a placeholder for an example function in Python. It requires no specific setup or imports.

```python
# Code example if applicable
def example_function():
    pass
```

--------------------------------

### Install JAX, JAXlib, and NeuronX Separately

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/jax/setup/jax-setup.md

Install JAX and JAXlib first, then install `jax-neuronx`, `libneuronxla`, and `neuronx-cc`. This method offers flexibility in choosing JAX and JAXlib versions.

```bash
python3 -m pip install jax==0.6.2 jaxlib==0.6.2
```

```bash
python3 -m pip install jax-neuronx libneuronxla neuronx-cc==2.* --extra-index-url=https://pip.repos.neuron.amazonaws.com
```

--------------------------------

### Install Dependencies

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/torch/torch-neuronx/tutorials/training/zero1_gpt2.md

Installs the necessary Hugging Face libraries for the tutorial. Ensure you are in the activated virtual environment before running.

```shell
python3 -m pip install -r requirements.txt
```

--------------------------------

### Install TensorFlow Neuron and Neuron Compiler

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/src/examples/tensorflow/tensorflow_resnet50/resnet50.ipynb

Installs the necessary TensorFlow Neuron and Neuron Compiler versions for TensorFlow 1.x. Ensure your environment is set up according to the TensorFlow Installation Guide.

```python
!pip install tensorflow_neuron==1.15.5.2.8.9.0 --extra-index-url=https://pip.repos.neuron.amazonaws.com/
!pip install neuron_cc==1.13.5.0 --extra-index-url=https://pip.repos.neuron.amazonaws.com
```

--------------------------------

### Get Specific DLAMI Image ID Example

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/deploy/environments/dlami.md

Example command to retrieve the image ID for the Multi-Framework DLAMI on Ubuntu 24.04.

```bash
aws ssm get-parameter \
--region us-east-1 \
--name /aws/service/neuron/dlami/multi-framework/ubuntu-24.04/latest/image_id \
--query "Parameter.Value" \
--output text
```

--------------------------------

### Launch Prefill Server (Multi-Instance)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/disaggregated-inference-tutorial-1p1d.md

Launches the prefill server for a multi-instance disaggregated inference setup. Ensure to replace 'path/to/your/downloaded/model' and provide the correct IP addresses for prefill and decode instances.

```bash
SEND=1 ./server.sh --tp-degree 64 --batch-size 4 \
                   --model-path path/to/your/downloaded/model \
                   --compiled-model-path di_traced_model_tp64_b4/ \
                   --neuron-send-ip prefill_ip --neuron-recv-ip decode_ip
```

--------------------------------

### Verify PyTorch Neuron Installation

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/about-neuron/quick-start/training-quickstart.md

Checks if PyTorch and torch_neuronx are installed correctly by importing them and printing the PyTorch version. This confirms successful setup.

```python
import torch
import torch_neuronx
print(f'PyTorch: {torch.__version__}')
```

--------------------------------

### vLLM Client Example (Python)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/_templates/model_card_qwen3.jinja.rst

A client script to submit requests to a running vLLM server. This example is typically used after starting the server with the `vllm serve` command.

```python
import json

import requests


if __name__ == "__main__":
   # Example: generate text
   schema = {
      "prompt": "A man is eating food. A man is eating a piece of fruit, a banana. ",
      "use_beam_search": False,
      "n": 1,
      "temperature": 0.7,
      "top_p": 0.9,
      "top_k": 40,
      "stream": False,
      "stop": None,
      "ignore_eos": False,
      "max_tokens": 16,
      "logprobs": None,
      "echo": False,
   }

   response = requests.post(
      "http://localhost:8080/generate",
      json=schema,
   )

   if response.status_code == 200:
      result = response.json()["text"] [0]
      print(f"Generated text: {result}")
   else:
      print(f"Error: {response.status_code}")

   # Example: get engine status
   response = requests.get("http://localhost:8080/health/ready")
   if response.status_code == 200:
      print("Server is ready.")
   else:
      print("Server is not ready.")
```

--------------------------------

### Example: Run neuron-monitor-prometheus.py

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-sys-tools/neuron-monitor-user-guide.md

An example showing how to run the neuron-monitor-prometheus.py script, setting a custom port.

```bash
neuron-monitor | neuron-monitor-prometheus.py --port 8008
```

--------------------------------

### Compile C Example with Neuron Runtime

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/neuron-runtime/guides/nrt-developer-guide.md

Compile the C example using gcc, linking against the Neuron Runtime library and including necessary headers. Ensure the Neuron SDK is installed.

```bash
gcc run_neff.c -o run_neff -lnrt -pthread -I/opt/aws/neuron/include -L/opt/aws/neuron/lib
```

--------------------------------

### Launch Prefill Server (Single-Instance)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/disaggregated-inference-tutorial-1p1d.md

Launches the prefill server for a single-instance disaggregated inference setup. Uses '127.0.0.1' for IP addresses as prefill and decode run on the same instance.

```bash
SEND=1 SINGLE_INSTANCE=1 ./server.sh --tp-degree 32 --batch-size 4 \
                                     --model-path path/to/your/downloaded/model \
                                     --compiled-model-path di_traced_model_tp32_b4/ \
                                     --neuron-send-ip 127.0.0.1 --neuron-recv-ip 127.0.0.1
```

--------------------------------

### Install PyTorch Neuron 2.18.0 on Amazon Linux 2

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/torch/torch-neuronx/setup/pytorch-install-prev-al2.md

Installs Python venv, Jupyter kernel, and Neuron Compiler/Framework for Neuron version 2.18.0. Ensure you have the necessary permissions and environment setup.

```text
# Install Python venv 
sudo dnf install -y python3.8-venv gcc-c++ 

# Create Python venv
python3.8 -m venv aws_neuron_venv_pytorch 

# Activate Python venv 
source aws_neuron_venv_pytorch/bin/activate 
python -m pip install -U pip 

# Install Jupyter notebook kernel
pip install ipykernel 
python3.8 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)"
pip install jupyter notebook
pip install environment_kernels

# Set pip repository pointing to the Neuron repository 
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com

# Install wget, awscli 
python -m pip install wget 
python -m pip install awscli 

# Install Neuron Compiler and Framework
python -m pip install neuronx-cc==2.13.66.0 torch-neuronx==1.13.1.1.14.0
```

--------------------------------

### Example: Run neuron-monitor-cloudwatch.py

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-sys-tools/neuron-monitor-user-guide.md

An example demonstrating how to execute the neuron-monitor-cloudwatch.py script with specific namespace and region values.

```bash
neuron-monitor | neuron-monitor-cloudwatch.py --namespace neuron_monitor_test --region us-west-2
```

--------------------------------

### Base Ubuntu Image and Environment Setup

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/deploy/docker-examples/inference/Dockerfile-inference-dlc.md

Establishes the base Ubuntu image and sets up essential environment variables for the container, including language, library paths, and serving module.

```dockerfile
FROM ubuntu:24.04

#SDK 1.17.1 has version 1. We skipped 1.18.0.
LABEL dlc_major_version="2"
LABEL maintainer="Amazon AI"
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true

ARG PYTHON=python3.7
ARG PYTHON_VERSION=3.7.10
ARG TS_VERSION=0.5.2
ARG MAMBA_VERSION=4.12.0-0

# See http://bugs.python.org/issue19846
ENV LANG C.UTF-8
ENV LD_LIBRARY_PATH /lib/x86_64-linux-gnu:/opt/conda/lib/:$LD_LIBRARY_PATH
ENV PATH /opt/conda/bin:$PATH
ENV SAGEMAKER_SERVING_MODULE sagemaker_pytorch_serving_container.serving:main
ENV TEMP=/home/model-server/tmp
```

--------------------------------

### SQL Query Example

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-explorer/overview-database-viewer.md

Use this to filter data by specifying conditions in standard SQL. The query starts with SELECT.

```sql
SELECT field_name FROM table_name WHERE condition
```

--------------------------------

### Print Tensor Debugging Example in PyTorch NeuronX

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/frameworks/torch/torch-neuronx/programming-guide/training/pytorch-neuron-debug.md

This example demonstrates how to set up a basic PyTorch model with NeuronX device and define tensors. It serves as a starting point for debugging by allowing inspection of tensor values and model structure.

```python
import os
import torch
import torch_xla
import torch_xla.core.xla_model as xm

device = xm.xla_device()
input1 = torch.randn(2,10).to(device)
# Defining 2 linear layers
linear1 = torch.nn.Linear(10,30).to(device)
linear2 = torch.nn.Linear(30,20).to(device)

```

--------------------------------

### Set up JAX Environment with NKI (Separate Packages)

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/nki/get-started/setup-env.md

Installs Python 3.12 venv, activates it, and installs jax, jaxlib, libneuronxla, neuronx-cc, and nki separately. jax-neuronx is an optional addition.

```bash
# Install Python venv
sudo apt-get install -y python3.12-venv g++

# Create Python venv
python3.12 -m venv aws_neuron_venv_jax

# Activate Python venv
source aws_neuron_venv_jax/bin/activate
python -m pip install -U pip

pip install jax==0.7.0 jaxlib==0.7.0
pip install jax-neuronx libneuronxla neuronx-cc==2.* nki --extra-index-url=https://pip.repos.neuron.amazonaws.com

```

--------------------------------

### Full JAX Example: Distributed Matmul with AllGather Profiling

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/tools/neuron-explorer/how-to-profile-workload.md

A comprehensive example demonstrating distributed matrix multiplication with AllGather in JAX, including JAX profiler setup and custom trace annotations. Ensure NEURON_RT_INSPECT_DEVICE_PROFILE and NEURON_RT_INSPECT_OUTPUT_DIR are set.

```python
from functools import partial
import os
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from jax.experimental.shard_map import shard_map
from time import sleep

os.environ["NEURON_RT_INSPECT_DEVICE_PROFILE"] = "1"
os.environ["NEURON_RT_INSPECT_OUTPUT_DIR"] = "./output"

jax.config.update("jax_default_prng_impl", "rbg")
mesh = Mesh(jax.devices(), ('i',))

def device_put(x, pspec):
    return jax.device_put(x, NamedSharding(mesh, pspec))

lhs_spec = P('i', None)
lhs = device_put(jax.random.normal(jax.random.key(0), (128, 128)), lhs_spec)
rhs_spec = P('i', None)
rhs = device_put(jax.random.normal(jax.random.key(1), (128, 16)), rhs_spec)

@jax.jit
@partial(shard_map, mesh=mesh, in_specs=(lhs_spec, rhs_spec), out_specs=rhs_spec)
def matmul_allgather(lhs_block, rhs_block):
    rhs = jax.lax.all_gather(rhs_block, 'i', tiled=True)
    return lhs_block @ rhs

with jax.profiler.trace(os.environ["NEURON_RT_INSPECT_OUTPUT_DIR"]):
    out = matmul_allgather(lhs, rhs)
    for i in range(10):
        with jax.profiler.TraceAnnotation("my_label" + str(i)):
            out = matmul_allgather(lhs, rhs)
        sleep(0.001)

expected = lhs @ rhs
with jax.default_device(jax.devices('cpu')[0]):
    equal = jnp.allclose(jax.device_get(out), jax.device_get(expected), atol=1e-3, rtol=1e-3)
print("Tensors are the same") if equal else print("Tensors are different")
```

--------------------------------

### Launch TensorBoard for Monitoring

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-training/tutorials/megatron_gpt_pretraining.md

Navigate to the experiment directory and launch TensorBoard to visualize training progress.

```bash
cd ~/neuronx-distributed-training/examples/nemo_experiments/megatron_gpt/
tensorboard --logdir ./
```

--------------------------------

### Execute vLLM Server Start Script

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-inference/tutorials/trn2-llama3.3-70b-dp-tutorial.ipynb

These commands execute the `start_vllm.sh` script to launch vLLM servers. The first command starts a server on port 8000 using neuron cores 0-31, and the second starts another on port 8001 using cores 32-63. This setup is for a distributed inference scenario with two servers.

```bash
!chmod +x ./start_vllm.sh
!./start_vllm.sh -p 8000 -c 0-31
```

```bash
!./start_vllm.sh -p 8001 -c 32-63
```

--------------------------------

### Install Neuron SDK and PyTorch

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/pytorch/update-manual.md

This snippet covers the complete process of setting up a Python virtual environment, configuring pip to use the Neuron repository, and installing the Neuron compiler and PyTorch framework. Ensure you have the necessary build tools installed.

```bash
sudo dnf install -y libxcrypt-compat
sudo dnf install python3.12
sudo dnf install -y gcc-c++
python3.12 -m venv aws_neuron_venv_pytorch
source aws_neuron_venv_pytorch/bin/activate
python -m pip install -U pip
pip install ipykernel
python3.12 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python (torch-neuronx)"
pip install jupyter notebook
pip install environment_kernels
python -m pip config set global.extra-index-url https://pip.repos.neuron.amazonaws.com
python -m pip install wget
python -m pip install awscli
python -m pip install neuronx-cc==2.22.12471.0 torch-neuronx==2.8.0.2.11.19912
```

--------------------------------

### Setup Single-Worker Inference

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/src/examples/mxnet/data_parallel/data_parallel_tutorial.ipynb

Configures and starts a single-worker inference setup using NeuronSimpleDataParallel. This involves setting the number of Neuron cores to use and initializing the parallel model. It also includes a warm-up phase to ensure cores are ready for inference.

```python
from parallel import NeuronSimpleDataParallel
from benchmark_utils import Results
import time
import functools
import os
import numpy as np
import warnings

num_cores = 1
batch_size=1

# Each worker process should use one core, hence we set
#    os.environ['NEURON_RT_NUM_CORES'] = "1"
os.environ["NEURON_RT_NUM_CORES"] = "1"

#Result aggregation class (code in bert_benchmark_utils.py)
results = Results(batch_size, num_cores)
def result_handler(output, start, end):
    elapsed = end - start
    results.add_result([elapsed], [end], [start])

inputs = get_sample_inputs(batch_size, seq_len)
parallel_neuron_model = NeuronSimpleDataParallel(compiled_model_path, num_cores, inputs)

#Starting the inference threads
parallel_neuron_model.start_continuous_inference()

# Warm up the cores
for _ in range(num_cores*4):
    parallel_neuron_model.warmup(inputs)
```

--------------------------------

### Install and Verify Docker

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/jax/dlc.md

Installs the Docker engine and adds the current user to the 'docker' group for necessary permissions. Includes a verification step using 'docker run hello-world'.

```bash
sudo apt-get install -y docker.io
sudo usermod -aG docker $USER

# Log out and log back in to refresh group membership, then verify:
docker run hello-world
```

--------------------------------

### Troubleshooting Jupyter Notebook Startup

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/notebook/setup-jupyter-notebook-steps-troubleshooting.md

Commands to reset Jupyter configuration, install necessary kernels and packages, and restart the notebook server when it fails to start.

```default
mv ~/.jupyter ~/.jupyter.old
mkdir -p ~/.jupyter
echo "c.NotebookApp.iopub_data_rate_limit = 10000000000" > ~/.jupyter/jupyter_notebook_config.py

# Instal Jupyter notebook kernel
 pip install ipykernel
 python3 -m ipykernel install --user --name aws_neuron_venv_pytorch --display-name "Python Neuronx"
 pip install jupyter notebook
 pip install environment_kernels

jupyter notebook
```

--------------------------------

### Set Up JAX Environment for NKI

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/nki/get-started/setup-env.md

Installs Python, GCC, and creates a Python virtual environment for JAX development with NKI.

```bash
# Install External Dependency
    sudo dnf install -y libxcrypt-compat

    # Install Python
    sudo dnf install -y python3.11

    # Install GCC
    sudo dnf install -y gcc-c++

    # Create Python venv
    python3.11 -m venv aws_neuron_venv_jax

    # Activate Python venv
    source aws_neuron_venv_jax/bin/activate
    pip install -U pip
    
```

--------------------------------

### Run a Quick PyTorch Example on Neuron

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/setup/pytorch/dlami.md

This snippet demonstrates a simple tensor operation and model compilation for Neuron using PyTorch. Ensure PyTorch and torch-neuronx are installed before running.

```python
import torch
import torch_neuronx

# Simple tensor operation on Neuron
x = torch.randn(3, 3)
model = torch.nn.Linear(3, 3)

# Compile for Neuron
trace = torch_neuronx.trace(model, x)
print(trace(x))
```

--------------------------------

### Download Tutorial Source Code

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/neuron-customops/tutorials/customop-mlp-training.md

Clone the repository containing the tutorial's source code and navigate to the customop_mlp directory.

```bash
git clone https://github.com/aws-neuron/aws-neuron-samples.git
cd aws-neuron-samples/torch-neuronx/training/customop_mlp
```

--------------------------------

### Launch TensorBoard for Monitoring

Source: https://github.com/aws-neuron/aws-neuron-sdk/blob/master/libraries/nxd-training/tutorials/hf_llama3_8B_pretraining.md

Navigate to the experiment directory and launch TensorBoard to visualize training progress. Ensure the logdir points to the correct experiment directory.

```bash
cd ~/neuronx-distributed-training/examples/nemo_experiments/hf_llama3_8B/
tensorboard --logdir ./
```