Environment variables reference

Environment variables reference#

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.

1. PyTorch distributed#

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.

2. Primus core#

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.

3. Primus YAML substitution#

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.

4. NCCL / RCCL#

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.

5. ROCm / HSA / HIP#

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.

6. CUDA / PyTorch#

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`.

7. Transformer engine#

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`.

8. Caches and authentication#

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.

9. hipBLASLt tuning#

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.

10. Build and rebuild#

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).

11. Container passthrough#

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

12. Slurm#

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.

13. Debug and pipeline#

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`.

14. JAX / XLA (MaxText)#

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.