This document is relevant for: Trn2, Trn3

Guidelines for memory management#

This page explains how torch-neuronx manages host and device memory and recommends practices for managing it. Use it to understand what the memory counters mean, tune the caching allocator, and keep host memory local to each Neuron device.

Applies to#

  • torch-neuronx (native TorchNeuron), eager mode and torch.compile

  • Trn2 and Trn3 Trainium instances

How torch-neuronx manages device memory#

Device memory is managed by a caching segment allocator. It allocates large segments of device memory from the Neuron runtime (NRT), then places tensors as blocks inside those segments. Each block is an NRT slice of its segment, so it does not need a new device memory allocation. When a tensor is freed, its block stays in the segment for reuse and is not returned to NRT.

The allocator serves memory allocations from two pools:

  • Requests smaller than 1 MB go to the small pool.

  • Requests of 1 MB or larger go to the large pool.

This keeps small, short-lived allocations from fragmenting the large segments that hold weights and activations.

Requests are also rounded differently by size:

  • Requests smaller than 16 MB are rounded up to a size class, with a minimum of 256 bytes.

  • Requests of 16 MB or larger are not rounded.

The memory APIs (see How to profile memory usage in torch-neuronx) report three quantities, which always satisfy reserved >= allocated >= active:

  • Reserved — device memory the allocator holds from NRT, including cached free blocks inside segments. reserved - allocated is cached free memory, available for immediate reuse.

  • Allocated — the total of the rounded block sizes of live tensors.

  • Active — the total bytes that live tensors requested, without rounding. allocated - active is internal fragmentation from rounding.

The allocator design also helps in the following behaviors:

  • Asynchronous allocations/deallocations. Allocations/Deallocations run in stream order on a runtime worker thread. When a tensor is freed, a deallocation is queued on its stream and runs in the order it was submitted, after the work already queued on that stream. The block then returns to the shared pool, where any stream can reuse it. If other streams also used the tensor (for example through record_stream), a deallocation is queued on each of them, and the block returns to the pool only after the last one runs. Because the pool is shared, memory freed on one stream can serve allocations on any other stream. This keeps the reserved memory lower, and reusing a cached block is faster than allocating a new segment. Call torch.neuron.synchronize() before you read a memory statistic.

  • OOM triggers automatic recovery in two steps. First, if an allocation fails and work is pending on the same stream, the allocator waits for that work to finish and free memory, then retries. Second, if it still fails, it releases cached free segments and retries once more before it raises an error. Both steps are on by default. Set NEURON_OOM_RECOVERY_ENABLED=0 to disable the second step.

Call torch_neuronx.empty_cache() to release the cached free segments that the allocator holds, for example between phases with very different memory footprints. It synchronizes the devices first.

Tune the Neuron caching allocator#

torch-neuronx also provides ways to tune the caching allocator behavior. The defaults are tuned from training-workload experiments. Override them only if you know the effects of these knobs. To check fragmentation, use segment_pool_analysis_summary() (see How to profile memory usage in torch-neuronx).

All tuning knobs are set through one environment variable. The allocator reads it once, the first time it is used:

TORCH_NEURONX_SEGMENT_ALLOCATOR_CONFIG="key1=value1,key2=value2"

Example adjustments:

  • Do not split free blocks larger than 512 MB, and let such a block serve a request up to 64 MB smaller than the block: kMaxSplitSizeBytes=536870912,kMaxNonSplitRoundingBytes=67108864. By default, every free block can be split.

  • Round requests smaller than 16 MB in finer steps to reduce internal fragmentation: kRoundupPower2Divisions=8. The default is 4 steps between each power of 2.

See TorchNeuron environment variables for the full key table, value format, and defaults.

Understand 64-bit dtype memory behavior#

int64 and float64 tensors keep their full 8 bytes per element in device memory.

How int64 ops run depends on TORCH_NEURONX_INT64_MODE:

  • native (default): each op is handled separately. Ops that support int64 natively run on the device as 64-bit. Other ops use the packed view (below). Matrix multiplication runs on the CPU.

  • packed: every int64 op uses the packed view. The device stores each int64 element as two int32 words, and arithmetic uses only the low 32 bits. Set TORCH_NEURONX_PACKED_INT64_CHECK_OVERFLOW=1 to raise an error on values outside the int32 range.

  • native_strict: ops that support int64 natively run on the device; all other int64 ops run on the CPU.

  • cpu_bounce: int64 values are cast to int32 through a CPU round trip, and the op runs on the device.

  • cpu_fallback: every int64 op runs on the CPU.

Neuron lacks native support for float64. Conversions between float64 and float32 run through the CPU.

Pin host memory to the local NUMA node#

On multi-rank Trainium workloads (for example, 64 ranks on a trn2.48xlarge), a Neuron device’s CPU threads may run on a different NUMA node without pinning, causing host memory accesses to cross the interconnect.

torch_neuronx.numa provides NUMA CPU affinity: it pins all CPU threads in each rank’s process to the same NUMA node as its assigned Neuron device. The API follows the same patterns as torch.numa.binding (PyTorch 2.11+), adapted for Neuron’s PrivateUse1 device model. Topology is discovered dynamically via neuron-ls — there are no hardcoded per-instance-type tables.

NUMA CPU affinity is on by default and needs no code changes. It activates when torch.distributed.init_process_group creates the Neuron process group. After pinning, new memory allocations land on the local node via Linux’s first-touch policy. However, memory allocated before the pin (Python interpreter startup, imported modules, early torch allocations) may already reside on the wrong node. A single process that does not create a process group is not pinned.

To also cover the allocations made before the pin, wrap each rank with neuron-numa-launch (installed with torch-neuronx; requires numactl). This changes only the launch command:

# Before:
torchrun --nproc_per_node 64 ... ./train.py --args

# After:
torchrun --no-python --nproc_per_node 64 ... neuron-numa-launch ./train.py --args

--no-python tells torchrun to launch neuron-numa-launch as a binary. The wrapper then runs Python under numactl --membind --cpunodebind, so all allocations — including Python startup — land on the local NUMA node. If numactl is not installed, the wrapper falls back to plain python, and the automatic pinning still applies.

Set NEURON_NUMA_PIN=0 to disable NUMA CPU affinity.

To check pinning status at runtime:

from torch_neuronx.numa import get_numa_status
print(get_numa_status())
# {'numa_pin_enabled': True, 'numa_pinned': True,
#  'current_affinity': [...], 'topology_available': True, 'num_devices': 16}

Place caches deliberately#

Compilation caches trade disk for recompilation time and affect shared filesystems:

  • TORCH_NEURONX_NEFF_CACHE_DIR (default /tmp/neff_cache) — point at a shared mount (for example, FSx) to reuse compiled NEFFs across runs and instances.

  • TORCH_NEURONX_NEFF_LOCAL_CACHE_DIR (default /tmp/local_cache) — keep it on a local filesystem. It holds the cache’s lock files, and file locks are not reliable on network filesystems such as NFS or Lustre.

See TorchNeuron environment variables for the complete cache variable list.

Note

Running on Trn3: The instance-dependent values on this page, such as the number of ranks per node, the device memory (HBM) per rank, NEURON_RT_NUM_CORES, and the device mesh shapes, are for Trn2 (trn2.48xlarge). Trn3 servers have a different number of Trainium chips, so these values do not apply to Trn3 without change. Before you run on Trn3, calculate them again from the Amazon EC2 Trn3 architecture and Trainium3 architecture specifications. For the effect of the LNC setting on the number of visible cores, see Logical NeuronCore configuration.