NCCL environment generator

#!/usr/bin/env bash
# Denpex — NCCL/PyTorch environment for AWS (p4d / p5, EFA) + PyTorch FSDP
# Every value is explained at denpex.com/nccl-env. Values marked VERIFY depend on
# your hardware — check them before trusting them.

# ── REQUIRED ────────────────────────────────────────────────────
# Makes EFA safe across fork(). PyTorch DataLoader workers fork, so without this you get intermittent, un-debuggable corruption or hangs that correlate with num_workers.
export FI_EFA_FORK_SAFE=1

# Enables device RDMA so the NIC reads GPU memory directly. Required on p4d; already the default on p5 and later, where setting it is harmless.
export FI_EFA_USE_DEVICE_RDMA=1

# Selects the EFA libfabric provider. Without it libfabric picks tcp or sockets and every byte crosses the kernel network stack instead of the fabric — often a 10x collective slowdown that looks like a model problem.
export FI_PROVIDER=efa

# EFA and its OpenMPI live outside the default loader path. Missing this is why FI_PROVIDER=efa "does nothing".
export LD_LIBRARY_PATH="/opt/amazon/efa/lib:/opt/amazon/openmpi/lib:$LD_LIBRARY_PATH"

# The bootstrap interface prefix (NOT the data path — that is EFA). AWS instances name interfaces ens5/ens32; the danger is NCCL picking docker0 or a virtual bridge and failing rendezvous.
# VERIFY: this value depends on your hardware.
export NCCL_SOCKET_IFNAME=ens

# Makes a failed collective raise instead of hanging forever. Without it a dead rank turns into a job that burns GPU-hours doing nothing until the wall clock kills it.
export TORCH_NCCL_ASYNC_ERROR_HANDLING=1

# ── RECOMMENDED ─────────────────────────────────────────────────
# Loads the checkpoint on rank 0 and broadcasts, instead of every rank materializing the full model in host RAM. On a 70B model this is the difference between loading and being OOM-killed at startup.
export FSDP_CPU_RAM_EFFICIENT_LOADING=1

# The LL and LL128 protocols are not supported over EFA in several aws-ofi-nccl versions; leaving the choice to NCCL can produce hangs or wrong results. Pin it to simple unless AWS docs for your exact version say otherwise.
export NCCL_PROTO=simple

# Lets the caching allocator grow segments instead of fragmenting. This is the single highest-value setting for OOM-at-step-N failures where memory is available but not contiguous.
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# FSDP's all-gather/reduce-scatter pattern interacts badly with recorded streams; this reduces peak memory.
export TORCH_NCCL_AVOID_RECORD_STREAMS=1

# Keep the watchdog asynchronous. Setting this to 1 serializes the wait and hides which rank actually stalled.
export TORCH_NCCL_BLOCKING_WAIT=0