This document catalogs the main environment variables you might encounter when running Primus on AMD GPUs: distributed launchers, Primus runners and CLI, YAML substitution, libraries (NCCL/RCCL, ROCm, PyTorch, JAX), and optional integrations (Hugging Face, WandB, MLflow). It is a practical reference, not a complete list of every variable accepted by upstream libraries.
Legend
- Required: Must be set for the stated workflow; otherwise the job fails or mis-ranks.
- Optional: Has a safe default or is only needed for specific features.
- Set by: Typical source (launcher,
runner/helpers/envs/*.sh, user shell, container host).
- Used in: Representative Primus paths; many variables are also read by NVIDIA NCCL, AMD RCCL, PyTorch, or JAX without Primus wrapping them.
Set by torchrun, Slurm launchers, or runner/primus-cli-direct.sh / runner/primus-cli-slurm-entry.sh. Consumed by PyTorch distributed, RCCL, and Primus helpers.
| Variable | Default | Where set | Where used | Description |
|---|
MASTER_ADDR | localhost (direct / base_env.sh) | User, Slurm entry (primus-cli-slurm-entry.sh), or validation fallback (runner/lib/validation.sh) | primus/pretrain.py, primus/core/base_module.py, primus/core/utils/env.py, primus/tools/preflight/network/network_probe.py, PyTorch rendezvous | Rendezvous hostname or IP for process group initialization. Required for multi-node if not using Slurm auto-detection. |
MASTER_PORT | 1234 (direct), 29500 in some Python defaults | Config / CLI / user | Same as MASTER_ADDR; validation.sh enforces 1024–65535 | TCP port for the store backing torch.distributed. |
RANK | 0 if unset in helpers | torchrun | primus/tools/utils.py, primus/tools/preflight/global_vars.py, projection and profiler code | Global rank index. |
WORLD_SIZE | 1 | torchrun | Preflight, projection, primus/core/base_module.py | Total number of processes. |
LOCAL_RANK | 0 | torchrun | primus/core/base_module.py, GPU selection in benchmarks and trainers | GPU index on this node. |
LOCAL_WORLD_SIZE | 1 (Python) / 8 in benchmarks default | torchrun | primus/tools/preflight/*.py, strided_allgather_bench.py | Processes (GPUs) per node. |
NODE_RANK | 0 | primus-cli-direct / primus-cli-slurm-entry.sh | primus/pretrain.py, logging in runner/lib/common.sh | Zero-based node index in multi-node jobs. |
NNODES | 1 | Direct config (runner/.primus.yaml), primus-cli-slurm-entry.sh | primus/pretrain.py, primus/core/projection/training_config.py | Number of nodes in the job. |
GPUS_PER_NODE | 8 | runner/.primus.yaml direct section, primus-cli-slurm-entry.sh, validation.sh | primus/core/projection/module_profilers/*.py, training config helpers | GPUs per node for world-size math and binding. |
| Variable | Default | Where set | Where used | Description |
|---|
PRIMUS_PATCHES | "" / "all" | User | primus/core/patches/patch_runner.py | "all" or empty enables all patches; "none" disables; comma list enables subset. |
PRIMUS_LOG_LEVEL | INFO | User; debug paths in runner/primus-cli-*.sh set DEBUG | runner/lib/common.sh | Log verbosity: DEBUG, INFO, WARN, ERROR. |
PRIMUS_LOG_TIMESTAMP | 1 | User | runner/lib/common.sh | 1 prefixes logs with timestamps; 0 disables. |
PRIMUS_LOG_COLOR | 1 (auto-off if not a TTY) | User; tests may set 0 | runner/lib/common.sh | ANSI colors in runner logs. |
PRIMUS_DEBUG | 0 | User | runner/helpers/envs/primus-env.sh | 1 enables set -x in the env loader for shell tracing. |
PRIMUS_SKIP_VALIDATION | 0 | User / tests | runner/helpers/envs/primus-env.sh | 1 skips validate_distributed_params (not recommended). |
PRIMUS_EXPECT_IB | (unset) | User | primus/tools/preflight/network/network_standard.py | When 1, preflight treats InfiniBand as expected for validation. |
PRIMUS_CLUSTER | amd-aig-poolside (CLI default) | User | primus/tools/benchmark/rccl_bench_args.py | Cluster label for RCCL benchmark tooling. |
PRIMUS_GPU_ARCH | (auto / "mi300x" in simulators) | User / CLI | primus/core/projection/simulation_backends/origami_backend.py, sdpa_simulator.py, projection.py CLI | GPU architecture string for performance projection. |
PRIMUS_GPU_CLOCK_MHZ | (unset) | User | Same as PRIMUS_GPU_ARCH | Optional clock override for projection. |
PRIMUS_GPU_DEVICE | 0 | User | origami_backend.py | GPU index for hardware detection in projection. |
PRIMUS_GEMM_BACKEND | (unset) | User | primus/core/projection/simulation_backends/factory.py | Selects GEMM simulation backend by name. |
PRIMUS_PREFLIGHT_MIN_FREE_MEM_GB | 1 | User | primus/tools/preflight/gpu/utils.py | Minimum free GPU memory (GB) for preflight checks. |
PRIMUS_PREFLIGHT_MIN_TFLOPS | 10.0 | User | primus/tools/preflight/gpu/utils.py | Minimum TFLOPS threshold for preflight GEMM checks. |
PRIMUS_TURBO_AUTO_TUNE | (unset) | User / tests | tests/trainer/test_megatron_trainer.py (integration) | Enables Turbo auto-tuning in supported Turbo/Megatron test flows; not referenced in core primus/ Python outside tests. Optional. |
PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND | TURBO | User; hooks may set DEEP_EP | primus/backends/megatron/patches/args/rocm_arg_validation.py, examples/run_pretrain.sh, runner/helpers/hooks/05_using_uep.sh | MoE dispatch/combine backend selector. |
Parsed by primus/core/config/yaml_loader.py for patterns ${VAR} (required) and ${VAR:default} (optional). Typical experiment YAMLs under examples/ use these for sweep-friendly overrides.
| Variable | Typical default in YAML | Where set | Where used | Description |
|---|
PRIMUS_TEAM | "amd" | User | Resolved before module merge in experiment YAML | Work group / team segment in paths. |
PRIMUS_USER | "root" | User | Experiment YAML | User name segment. |
PRIMUS_EXP_NAME | per-example | User | Experiment YAML | Experiment folder name. |
PRIMUS_WORKSPACE | "./output" | User | Experiment YAML | Root workspace for artifacts. |
PRIMUS_TP | 1 | User | Megatron example YAMLs | tensor_model_parallel_size override. |
PRIMUS_PP | 1 | User | Megatron example YAMLs | pipeline_model_parallel_size override. |
PRIMUS_EP | 1 | User | Megatron example YAMLs | expert_model_parallel_size override. |
PRIMUS_SEQ_LENGTH | per-model | User | Megatron example YAMLs | Sequence length override. |
PRIMUS_MAX_POSITION_EMBEDDINGS | 4096 or 131072 | User | examples/megatron/**/*.yaml, tests | Position embedding cap override. |
PRIMUS_GLOBAL_BATCH_SIZE | per-model | User | Megatron example YAMLs | Global batch override. |
PRIMUS_NUM_LAYERS | per-model | User | Tests and MoE examples | Transformer layer count override. |
PRIMUS_MOE_LAYER_FREQ | MoE patterns | User | MoE examples / tests | MoE layer frequency pattern. |
PRIMUS_TOKENIZED_DATA_PATH | null | User | Megatron examples | Path to tokenized training data. |
PRIMUS_MODEL | per-stack | User | Megatron examples | Model preset stem (e.g. llama3_8B). |
PRIMUS_VPP | null | User | tests/trainer/test_megatron_trainer.yaml | Virtual pipeline stages override. |
Primus seeds many of these in runner/helpers/envs/base_env.sh. RCCL honors NCCL-compatible variables on AMD GPUs. See NCCL environment and RCCL environment.
| Variable | Default (Primus base) | Where set | Where used | Description |
|---|
NCCL_DEBUG | unset | User / base_env.sh empty default | Preflight reports, RCCL runtime | Log verbosity: NONE, WARN, INFO, TRACE, etc. Optional unless debugging comms. |
NCCL_SOCKET_IFNAME | derived from IP_INTERFACE | base_env.sh | primus/tools/preflight/network/*.py, GPU topology helpers | Socket NIC for host networking. |
GLOO_SOCKET_IFNAME | same as NCCL if unset | base_env.sh | Preflight | Gloo TCP backend interface. |
NCCL_IB_HCA | auto via get_nccl_ib_hca.sh if empty | base_env.sh, container passthrough | Preflight, multi-node tuning | InfiniBand HCAs to use. |
NCCL_IB_GID_INDEX | 3 | base_env.sh | RCCL | GID index for IB/RoCE; many sites use 1 for RoCE v2 (override as needed). |
NCCL_IB_TC | (unset) | User | RCCL | InfiniBand traffic class. |
NCCL_IB_FIFO_TC | (unset) | User | RCCL | InfiniBand FIFO traffic class. |
NCCL_IB_ROCE_VERSION_NUM | (unset) | User | RCCL | RoCE version selection. |
NCCL_PXN_DISABLE | 1 | base_env.sh | RCCL | Disable PXN (PCIe cross-NIC); set 0 to enable. |
NCCL_P2P_NET_CHUNKSIZE | 524288 | base_env.sh | RCCL | P2P network chunk size tuning. |
NCCL_PROTO | (unset) | User | RCCL | Protocol selection (e.g. Simple, LL, LL128). |
NCCL_CROSS_NIC | 0 | base_env.sh | RCCL | Cross-NIC communication policy. |
NCCL_IB_RETRY_CNT | (unset) | User | RCCL | IB retry count. |
NCCL_IB_TIMEOUT | (unset) | User | RCCL | IB timeout. |
NCCL_NET_GDR_LEVEL | (unset) | User | Preflight summaries | GPUDirect RDMA level. |
NCCL_IB_DISABLE | 0 | User / env | Preflight | Disable IB; use sockets only. |
NCCL_DMABUF_ENABLE | (unset) | User | RCCL | DMA-BUF registration path. |
NCCL_IGNORE_CPU_AFFINITY | (unset) | User | RCCL | Ignore CPU affinity hints. |
NCCL_IB_QPS_PER_CONNECTION | (unset) | User | RCCL | IB QPs per connection. |
NCCL_MAX_P2P_CHANNELS | (unset) | User | RCCL | Cap P2P channels. |
NCCL_GDR_FLUSH_DISABLE | (unset) | User | RCCL | Disable GDR flush. |
NCCL_IB_USE_INLINE | (unset) | User | RCCL | Inline IB sends. |
NCCL_NET_PLUGIN | (unset) | User | RCCL | Alternate network plugin (e.g. librccl-anp.so). |
RCCL_MSCCL_ENABLE | 0 | base_env.sh | RCCL | Enable MSCCL algorithms. |
RCCL_MSCCLPP_THRESHOLD | 1GiB default | base_env.sh | RCCL | MSCCL++ message-size threshold. |
RCCL_GDR_FLUSH_GPU_MEM_NO_RELAXED_ORDERING | 0 in hooks | runner/helpers/hooks/03_enable_ainic.sh | RCCL | Stricter GDR flush memory ordering; relevant for some NIC/GPU combos. |
TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK | 0 | base_env.sh | PyTorch + RCCL | Tensor allocator hook for NCCL registration. |
TORCH_NCCL_HIGH_PRIORITY | 1 | base_env.sh | PyTorch | High-priority NCCL streams. |
| Variable | Default | Where set | Where used | Description |
|---|
HSA_ENABLE_SDMA | 1 | base_env.sh | ROCm runtime | Enable SDMA engines for copies. |
HSA_NO_SCRATCH_RECLAIM | 1 | base_env.sh, container passthrough | ROCm runtime; documented for MoE stability | 1 keeps scratch allocated (often used for MoE stability). See ROCR environment. |
HIP_VISIBLE_DEVICES | 0..GPUS_PER_NODE-1 | base_env.sh | ROCm device visibility | Restricts which GPU indices ROCm exposes. |
ROCBLAS_DEFAULT_ATOMICS_MODE | (unset) | User | primus/backends/megatron/patches/args/rocm_arg_validation.py | Read for deterministic / accuracy-sensitive GEMM behavior. |
| Variable | Default | Where set | Where used | Description |
|---|
CUDA_DEVICE_MAX_CONNECTIONS | 1 | base_env.sh; Megatron patches may adjust | primus/backends/megatron/patches/env_patches.py, Megatron patches | Limits concurrent CUDA connections; often 1 for TP/PP overlap. |
TORCH_COMPILE_DISABLE | 0 | User | primus/backends/megatron/patches/args/rocm_arg_validation.py | Disable torch.compile when 1. |
| Variable | Default | Where set | Where used | Description |
|---|
NVTE_ROCM_ENABLE_MXFP8 | 1 | base_env.sh | Transformer Engine on ROCm | Enable MXFP8 paths. |
NVTE_CK_USES_BWD_V3 | 1 | base_env.sh, container passthrough | TE / CK | Use CK backward v3 kernels. |
NVTE_CK_IS_V3_ATOMIC_FP32 | (unset; examples print 0) | User / examples/run_pretrain.sh, container passthrough | TE / CK | Atomic FP32 mode for CK v3 backward. |
PATCH_TE_FLASH_ATTN | 0 | base_env.sh, container passthrough | runner/helpers/hooks/01_patch_te_flash_attn_max_version.sh | Trigger TE flash-attn patch hook when 1. |
| Variable | Default | Where set | Where used | Description |
|---|
HF_HOME | ${DATA_PATH}/huggingface | base_env.sh, primus/core/utils/env_setup.py, primus/pretrain.py | Hugging Face libraries | Cache for models and datasets. |
HF_TOKEN | (unset) | User, container passthrough | Hugging Face Hub | Auth for gated models. Required for private/gated assets. |
TORCH_HOME | under workspace | primus/core/utils/env_setup.py | PyTorch Hub | Torch Hub cache root. |
TRANSFORMERS_CACHE | aligned with HF layout | primus/core/utils/env_setup.py | transformers | Model cache for Transformers. |
WANDB_API_KEY | (unset) | User, container passthrough | Weights & Biases client, Megatron trainer checks | API key for logging. Required for Weights & Biases when enabled. |
WANDB_PROJECT | (unset) | User / TorchTitan patch | primus/backends/torchtitan/patches/wandb_patches.py | Project name. |
WANDB_RUN_NAME | (unset) | User / patches | Same | Run display name. |
WANDB_TEAM | (unset) | User | TorchTitan metrics (entity) | WandB team/entity. |
DATABRICKS_HOST | (unset) | User | mlflow client (via primus/backends/megatron/training/global_vars.py MLflow setup) | Required for Databricks-hosted MLflow when MLflow logging is enabled. |
DATABRICKS_TOKEN | (unset) | User | Databricks APIs | Auth token paired with host. |
MLFLOW_TRACKING_URI | (unset) | User | mlflow (via Megatron integrations) | MLflow tracking server URI. Optional unless using MLflow. |
MLFLOW_REGISTRY_URI | (unset) | User | MLflow | Model registry endpoint. |
NLTK_DATA | (unset) | User | runner/helpers/hooks/train/pretrain/megatron/preprocess_data.py, Megatron-LM tools | Punkt and other tokenizer data for preprocessing. |
TOKENIZED_DATA_PATH | per-hook default | User | runner/helpers/hooks/train/pretrain/megatron/prepare.py | Pre-tokenized dataset location for Megatron data prep hooks. |
| Variable | Default | Where set | Where used | Description |
|---|
PRIMUS_HIPBLASLT_TUNING | 0 | User | examples/run_pretrain.sh | Master switch for the HipBLASLt tuning flow (1 enables). Must be set before PRIMUS_HIPBLASLT_TUNING_STAGE takes effect, and is mutually exclusive with deterministic mode (PRIMUS_DETERMINISTIC=1). |
PRIMUS_HIPBLASLT_TUNING_STAGE | 0 | User | examples/run_pretrain.sh | Stages 0 off, 1 dump shapes, 2 offline tune, 3 apply tuned kernels. |
HIPBLASLT_TUNING_OVERRIDE_FILE | (unset) | User / tuning scripts | examples/run_pretrain.sh | Path to tuned-kernel override file for stage 3. |
TE_HIPBLASLT_TUNING_RUN_COUNT | varies | User | examples/run_pretrain.sh | Number of benchmark runs per shape during TE hipBLASLt tuning. |
TE_HIPBLASLT_TUNING_ALGO_COUNT | varies | User | examples/run_pretrain.sh | Transformer Engine hipBLASLt search breadth. |
TE_HIPBLASLT_TUNING_ALGO_FILE | (unset) | User | TE + HipBLASLt | Algorithm file for TE tuning flows. |
TE_HIPBLASLT_TUNING | (unset) | User | examples/run_pretrain.sh | When set, interacts with deterministic mode and tuning stages (disable conflicting modes per script comments). |
HIPBLASLT_LOG_LEVEL | (unset) | User | HipBLASLt | Library log level. |
HIPBLASLT_LOG_MASK | (unset) | User | HipBLASLt | Bitmask for log categories. |
| Variable | Default | Where set | Where used | Description |
|---|
REBUILD_PRIMUS_TURBO | 0 | User, container passthrough | runner/helpers/hooks/00_rebuild_primus_turbo.sh | 1 rebuilds Primus-Turbo on startup. |
REBUILD_BNXT | 0 | User, container passthrough | runner/helpers/hooks/02_rebuild_bnxt.sh | 1 rebuilds BNXT driver artifacts when packaged. |
USING_AINIC | (unset) | User | runner/helpers/hooks/03_enable_ainic.sh | 1 enables AINIC-oriented networking hooks. |
MAX_JOBS | (unset) | User / tooling | tools/daily/safe_wrapper.py | Parallel compile jobs for pip builds. |
BACKEND_PATH | (unset) | User | primus/pretrain.py, primus/core/backend/backend_adapter.py | Override checkout path for third-party backends (Megatron, TorchTitan, MaxText). |
runner/.primus.yaml lists names forwarded from the host into training containers (container.options.env). Primus does not assign values here; it only allowlists keys for --env forwarding.
Forwarded keys:
MASTER_ADDR, MASTER_PORT, NNODES, NODE_RANK, GPUS_PER_NODE, DOCKER_IMAGE, HF_TOKEN, WANDB_API_KEY, ENABLE_NUMA_BINDING, REBUILD_PRIMUS_TURBO, USING_AINIC, PATCH_TE_FLASH_ATTN, REBUILD_BNXT, HSA_NO_SCRATCH_RECLAIM, NVTE_CK_USES_BWD_V3, GPU_MAX_HW_QUEUES, HSA_KERNARG_POOL_SIZE, PRIMUS_TURBO_DEEPEP_TIMEOUT, NCCL_IB_HCA, NCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME, NCCL_IB_GID_INDEX, PRIMUS_TURBO_ATTN_V3_ATOMIC_FP32, NVTE_CK_IS_V3_ATOMIC_FP32, PATH_TO_BNXT_TAR_PACKAGE, ANP_HOME_DIR, RCCL_HOME_DIR, MPI_HOME_DIR, DUMP_HLO, DUMP_HLO_DIR, PRIMUS_DETERMINISTIC, PRIMUS_HIPBLASLT_TUNING, PRIMUS_HIPBLASLT_TUNING_STAGE, TE_HIPBLASLT_TUNING_RUN_COUNT, TE_HIPBLASLT_TUNING_ALGO_COUNT, HIPBLASLT_LOG_MASK, HIPBLASLT_LOG_FILE, HIPBLASLT_LOG_LEVEL, HIPBLASLT_TUNING_OVERRIDE_FILE
| Variable | Default | Where set | Where used | Description |
|---|
SLURM_NNODES / SLURM_JOB_NUM_NODES | job-dependent | Slurm | primus-cli-slurm-entry.sh (NNODES export), preflight probes | Node count for the allocation. |
SLURM_NODEID | job-dependent | Slurm | Mapped to NODE_RANK in primus-cli-slurm-entry.sh | Node index. |
SLURM_PROCID | job-dependent | Slurm | Fallback for NODE_RANK when SLURM_NODEID is unset | Process ID within the Slurm step (entry script). |
SLURM_JOB_ID | job-dependent | Slurm | primus/tools/preflight/host/host_probe.py | Job identifier string. |
| Variable | Default | Where set | Where used | Description |
|---|
DUMP_PP_DIR | output/pp_data | User | primus/backends/megatron/megatron_pretrain_trainer.py, primus/backends/megatron/patches/pp_dump_data_patches.py | Directory for pipeline-parallel debug dumps. |
DEBUG_SIMULATOR | 0 | User | primus/core/projection/performance_projection/simulator.py | 1 enables verbose projection simulator logging. |
RECORD_OFFLOAD_MEMORY_INFO | 0 | User | primus/core/pipeline_parallel/handler/offload_handler.py | Record offload memory stats when 1. |
RECORD_OFFLOAD_MEMORY_INFO_DIR | output | User | primus/core/pipeline_parallel/scheduler/scheduler.py | Output directory for offload memory logs. |
USE_PINNED_OFFLOAD | 0 | User | offload_handler.py | Use pinned host memory for offload buffers when 1. |
Primus MaxText hooks print recommended values in runner/helpers/hooks/train/pretrain/maxtext/prepare.py; MaxText and JAX read them directly.
| Variable | Default | Where set | Where used | Description |
|---|
XLA_PYTHON_CLIENT_MEM_FRACTION | e.g. .97 in prepare hook | User / hook output | JAX / XLA allocator | Fraction of GPU memory pre-allocated for JAX. |
DUMP_HLO_DIR | ${PRIMUS_PATH}/output/xla_dump_hlo (example) | User | XLA via XLA_FLAGS composition | Directory for HLO dumps when enabled. |
DUMP_HLO | 0 | User | Prepare hook → XLA flags | Gate HLO dumping (1 enables in hook samples). |
Note: MaxText also propagates many knobs through XLA_FLAGS and LIBTPU_INIT_ARGS upstream; see MaxText sources for the full list.