Skip to main content

Running on SageMaker HyperPod with Slurm

Cluster Assumptions​

  • A SageMaker HyperPod cluster with Slurm orchestration.
  • At least 2 compute nodes, each with multiple NVIDIA GPUs (e.g., p5.48xlarge with 8× H100).
  • Shared filesystem (FSx for Lustre) mounted at /fsx.
  • NCCL and EFA drivers installed (HyperPod handles this).

Setup​

Place your training script and dependencies on the shared filesystem:

# On the head node
mkdir -p /fsx/training/ddp-demo
cp train_ddp.py /fsx/training/ddp-demo/

Writing the Sbatch Job Script​

Save as submit_ddp.sh:

#!/bin/bash
#SBATCH --job-name=ddp-pretrain
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1        # one task per node (torchrun handles GPUs)
#SBATCH --gpus-per-node=8          # all GPUs on each node
#SBATCH --cpus-per-task=96         # all CPUs for data loading workers
#SBATCH --partition=gpu
#SBATCH --output=/fsx/training/ddp-demo/logs/%x_%j.out
#SBATCH --error=/fsx/training/ddp-demo/logs/%x_%j.err
#SBATCH --exclusive

# ─── Environment ──────────────────────────────────────────────────────────────
export NCCL_DEBUG=INFO                  # Useful for debugging connectivity
export FI_EFA_USE_DEVICE_RDMA=1        # Enable EFA RDMA for faster comms
export FI_PROVIDER=efa                  # Use EFA provider

# Note: NCCL_PROTO=simple was recommended for older EFA versions.
# On p5/H100 instances with EFA 2.x, omit this or test both settings.
# export NCCL_PROTO=simple

# ─── Resolve master address from the first node ──────────────────────────────
# Re-derive from live job info (handles node replacement after auto-resume)
NODE_LIST=$(scontrol show jobid=$SLURM_JOBID | awk -F= '/NodeList=/{print $2}' | grep -v Exc)
MASTER_ADDR=$(scontrol show hostname $NODE_LIST | head -n 1)
MASTER_PORT=29500  # Safe with --exclusive; for shared nodes use: $((29500 + SLURM_JOB_ID % 1000))
export MASTER_ADDR MASTER_PORT

# ─── Calculate total processes ────────────────────────────────────────────────
GPUS_PER_NODE=8
NNODES=$SLURM_NNODES
WORLD_SIZE=$((NNODES * GPUS_PER_NODE))

echo "Master: $MASTER_ADDR:$MASTER_PORT | Nodes: $NNODES | World Size: $WORLD_SIZE"

# ─── Launch with srun + torchrun ──────────────────────────────────────────────
srun --auto-resume=1 torchrun \
    --nnodes=$NNODES \
    --nproc_per_node=$GPUS_PER_NODE \
    --rdzv_id=$SLURM_JOB_ID \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT \
    /fsx/training/ddp-demo/train_ddp.py

Submitting and Monitoring​

# Create log directory
mkdir -p /fsx/training/ddp-demo/logs

# Submit
sbatch submit_ddp.sh

# Monitor
squeue -u $USER
tail -f /fsx/training/ddp-demo/logs/ddp-pretrain_*.out

Verifying Multi-Node Communication​

Check the logs for NCCL initialization messages:

NCCL INFO Connected all rings
NCCL INFO comm 0x... rank 0 nranks 16 ...

You should see nranks equal to your total world size (nodes × GPUs_per_node).

Leveraging HyperPod Auto-Resume​

SageMaker HyperPod's key differentiator is automatic node health monitoring and replacement. When a node fails:

  1. HyperPod detects the unhealthy node and replaces it.
  2. With --auto-resume=1 on srun, HyperPod automatically requeues the job on healthy nodes.
  3. The job receives SIGTERM before termination, triggering a checkpoint save.
  4. On restart, master address is re-derived from the updated node list.

The training script above already handles SIGTERM (see the sigterm_handler). The --auto-resume=1 flag on srun tells HyperPod's Slurm plugin to requeue the job automatically when a node is replaced — no separate --requeue directive needed.

Troubleshooting Common Issues​

SymptomLikely CauseFix
Hang at init_process_groupFirewall/security group blocking portEnsure nodes can reach each other on MASTER_PORT
NCCL WARN ... no EFA deviceEFA not configuredVerify EFA is installed: fi_info -p efa
OOM on some ranksUneven batch sizesUse drop_last=True in DataLoader
Loss divergesLR not scaled for effective batch sizeApply linear scaling: lr = base_lr * world_size
Different losses per rankForgot sampler.set_epoch()Always call sampler.set_epoch(epoch)
Silent hang mid-trainingUnused parameters in modelTry find_unused_parameters=True in DDP wrapper
Timeout after 10minDefault NCCL timeout (10 min) too shortSet timeout=timedelta(minutes=30) or higher in init_process_group