This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3

What is the Neuron Graph Compiler (neuronx-cc)?#

The Neuron Graph Compiler (neuronx-cc) is the XLA-based compiler that turns a machine learning model into optimized code for AWS Trainium and Inferentia accelerators. It analyzes the model graph, applies hardware-specific optimizations, and produces a NEFF (Neuron Executable File Format) file that the Neuron Runtime loads and runs on NeuronCores.

neuronx-cc targets the NeuronCores v2 to v4 architectures (Trn1, Trn1n, Inf2, and Trn2). It is one of the two Neuron compilers; the older neuron-cc targets NeuronCores v1 (Inf1). For the difference between them, see neuron-cc vs. neuronx-cc.

Applies to#

This concept is applicable to:

  • Training models on Trainium (Trn1, Trn1n, Trn2)

  • Running inference on Inferentia and Trainium (Inf2, Trn1, Trn2)

How neuronx-cc fits into the Neuron workflow#

Most users never call neuronx-cc directly. A Neuron framework plugin (for example, PyTorch NeuronX) invokes the compiler for you and forwards options through the NEURON_CC_FLAGS environment variable. The compilation flow is:

  1. The framework traces your model and emits an XLA HLO graph (hlo.pb).

  2. neuronx-cc compiles that graph for a target instance family.

  3. The compiler writes a NEFF archive.

  4. The Neuron Runtime loads the NEFF and runs it on one or more NeuronCores.

Calling the compiler on the command line is useful when you do not use a framework, when you customize a framework, or when you want to inspect or reproduce a build.

Commands#

neuronx-cc has the following commands. Run neuronx-cc <command> --help for full usage.

  • compile: Compile a model graph into a NEFF file for a target instance family.

  • list-operators: Print the operators the compiler supports for a given framework.

For the full set of flags for these compiler commands, see the Neuron Graph Compiler command reference.

Files neuronx-cc works with#

  • Model input (HLO): an XLA HLO file (hlo.pb) generated by the framework. This is the graph the compiler reads. --framework XLA selects this input type.

  • NEFF output: the compiled artifact, written to file.neff by default and set with --output. The Neuron Runtime loads this file to run the model.

  • Log file: compiler messages, written to log-neuron-cc.txt by default and set with --logfile. Use --verbose to control how much detail is logged.

Key features#

You control the compiler through compile options. The most common ones:

  • Target selection (--target): the instance family to compile for — trn1, trn1n, inf2, or trn2.

  • Model-type optimizations (--model-type): generic (default), transformer, or unet-inference. These apply optimizations tuned to a model family.

  • Mixed precision (--auto-cast and --auto-cast-type): trade accuracy for performance by casting FP32 operations to bf16, fp16, tf32, or fp8_e4m3. See Mixed Precision and Performance-accuracy Tuning (neuronx-cc) for guidance.

  • Mixed precision accumulation (--enable-mixed-precision-accumulation): enabled by default; accumulates operators such as softmax and layernorm in FP32 for better accuracy. Disable with --disable-mixed-precision-accumulation.

  • Optimization level (--optlevel / -O): 1, 2 (default), or 3, trading compile time against runtime performance.

  • Graph sharding (--logical-nc-config / -lnc): on trn2, shard the graph across 1 or 2 physical NeuronCores (default 2).

  • Distributed training (--distribution-strategy llm-training): enable optimizations for LLM training that shards parameters, gradients, and optimizer state across data-parallel workers.

  • Runtime tuning: --enable-fast-context-switch (faster model switching), --enable-fast-loading-neuron-binaries (larger but faster-loading NEFFs), and --enable-saturate-infinity (avoid NaNs on trn1).

For every option and its valid values, see the Neuron Compiler CLI reference.

Quickstart#

Compile an XLA HLO file into a NEFF for a Trainium instance, applying transformer-specific optimizations:

neuronx-cc compile bert-model.hlo \
  --framework XLA \
  --target trn1 \
  --model-type transformer \
  --output bert.neff

List the operators the compiler supports:

neuronx-cc list-operators --framework XLA

The compiler returns exit status 0 on success and a non-zero status on error.

Note

When you compile through a framework, you do not run these commands directly. Pass the same options through NEURON_CC_FLAGS instead, and the framework forwards them to neuronx-cc.

Further reading#

This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3