NeuralArmada | Multi-Node FSDP Training on SageMaker AI: Managed Spot, Checkpointing, and Warm Pools

Almost every SageMaker AI tutorial you find — including several on this site — is about inference. But the expensive, fragile part of a real project is usually the training job: a 7B or 13B model that no longer fits on one GPU, a run that dies at hour nine, and a bill nobody budgeted for. This tutorial covers the training side end to end: multi-node PyTorch FSDP on SageMaker AI training jobs, managed spot instances for the discount, checkpointing so interruptions are survivable, and warm pools so your debug loop is measured in seconds instead of minutes.

We assume you already have a working single-GPU fine-tuning script and an execution role that can read/write your S3 bucket.

When you actually need FSDP

Order of preference, cheapest first:

One GPU, LoRA/QLoRA. A 7–8B model fine-tunes on a single ml.g5.2xlarge or ml.g6.4xlarge with 4-bit quantization. If this works, stop here.
One node, many GPUs, FSDP. ml.g5.12xlarge (4× A10G) or ml.g6e.12xlarge gives you shard-across-GPUs without any network tuning.
Many nodes, FSDP + EFA. Needed for full-parameter fine-tunes of 13B+ or when your sequence length blows up activation memory. Use p4d/p5 class instances, which carry EFA; sharding across g5 nodes over standard networking is usually slower than a bigger single node.

Rough memory math for full-parameter AdamW training in bf16: ~2 bytes/param weights + 2 bytes gradients + 8 bytes optimizer state ≈ 12–16 GB per billion parameters, before activations. An 8B model is therefore ~110 GB of state — impossible on one 24 GB A10G, comfortable when sharded across 8× A100 40 GB.

Step 1: Make the training script rank-aware

SageMaker AI injects the distributed environment for you when you use the torch_distributed launcher. Your script only needs to read it:

# train.py
import os, torch, torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.llama.modeling_llama import LlamaDecoderLayer
import functools

def setup():
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    return local_rank, int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"])

local_rank, rank, world = setup()

model = AutoModelForCausalLM.from_pretrained(
    os.environ["MODEL_ID"], torch_dtype=torch.bfloat16, use_cache=False
)
model.gradient_checkpointing_enable()

wrap_policy = functools.partial(
    transformer_auto_wrap_policy, transformer_layer_cls={LlamaDecoderLayer}
)

model = FSDP(
    model,
    auto_wrap_policy=wrap_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,       # HYBRID_SHARD for multi-node
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
        buffer_dtype=torch.bfloat16,
    ),
    device_id=torch.cuda.current_device(),
    limit_all_gathers=True,
    use_orig_params=True,
)

Three settings do most of the work:

transformer_auto_wrap_policy wraps one decoder block per FSDP unit. Wrapping the whole model as a single unit gives you no memory benefit at all — this is the single most common FSDP mistake.
gradient_checkpointing_enable() trades ~20–30% step time for a large activation-memory reduction. Turn it on first when you hit OOM, before you shrink batch size.
ShardingStrategy.HYBRID_SHARD shards inside a node and replicates across nodes. On multi-node runs this cuts cross-node traffic dramatically; use FULL_SHARD only when the model genuinely does not fit within one node.

Step 2: Checkpoint to `/opt/ml/checkpoints`, and resume from it

This is the part that makes spot instances safe. SageMaker AI continuously syncs anything written to /opt/ml/checkpoints to S3, and restores it into the same path when a job restarts after an interruption. So checkpointing and resuming are the same code path:

CKPT_DIR = "/opt/ml/checkpoints"

def save_checkpoint(model, optimizer, step):
    from torch.distributed.fsdp import FullStateDictConfig
    cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
        state = model.state_dict()
    if rank == 0:
        torch.save({"step": step, "model": state}, f"{CKPT_DIR}/ckpt.pt")
    dist.barrier()

def load_checkpoint(model):
    path = f"{CKPT_DIR}/ckpt.pt"
    if not os.path.exists(path):
        return 0                      # fresh start
    blob = torch.load(path, map_location="cpu")
    from torch.distributed.fsdp import FullStateDictConfig
    cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, cfg):
        model.load_state_dict(blob["model"])
    return blob["step"]

start_step = load_checkpoint(model)

Notes from production runs:

FULL_STATE_DICT with rank0_only=True is simple and portable, but rank 0 must hold the whole model in CPU RAM. Above ~13B, switch to StateDictType.SHARDED_STATE_DICT with torch.distributed.checkpoint so every rank writes its own shard in parallel.
Also checkpoint the optimizer (FSDP.optim_state_dict) and your data-loader position, or a resumed run silently repeats samples and your loss curve develops a step you cannot explain.
Checkpoint on a time interval (every 20–30 minutes), not only on epoch boundaries. Your resume cost is bounded by the interval; your S3 cost is negligible next to a p4d hour.
Write to a temp file and rename. A checkpoint half-written when the interruption notice lands is worse than no checkpoint.

Step 3: Launch it with managed spot and warm pools

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="train.py",
    source_dir="src",
    role=role,
    framework_version="2.6",
    py_version="py312",
    instance_type="ml.g5.12xlarge",
    instance_count=2,
    distribution={"torch_distributed": {"enabled": True}},
    environment={"MODEL_ID": "meta-llama/Llama-3.1-8B-Instruct"},
    hyperparameters={"epochs": 3, "per_device_batch_size": 1, "grad_accum": 16},

    # cost + resilience
    use_spot_instances=True,
    max_run=24 * 3600,          # training time budget
    max_wait=30 * 3600,         # must be >= max_run when using spot
    checkpoint_s3_uri="s3://my-bucket/runs/llama31-8b/checkpoints",
    checkpoint_local_path="/opt/ml/checkpoints",

    # fast iteration
    keep_alive_period_in_seconds=1800,
)

estimator.fit({"train": "s3://my-bucket/data/train/", "val": "s3://my-bucket/data/val/"})

What each knob buys you:

use_spot_instances=True typically lands 60–70% off on-demand for GPU training. max_wait covers queueing plus interruption retries and must be at least max_run, or the SDK rejects the job.
checkpoint_s3_uri is what turns an interruption from a lost day into a lost 20 minutes. Never enable spot without it.
keep_alive_period_in_seconds (warm pools) keeps the cluster provisioned after the job ends, so your next fit() with the same configuration starts in ~30 seconds instead of 5–8 minutes. Warm-pool time is billed, so use ~1800 while you are actively debugging and drop it for the final run. Request a warm-pool quota increase first; the default is zero.
distribution={"torch_distributed": ...} sets RANK, LOCAL_RANK, WORLD_SIZE, and the master address on every node. Do not hand-roll torchrun.

For multi-node EFA instances, also pass FI_PROVIDER=efa and NCCL_PROTO=simple in environment, and confirm in the job log that NCCL selected EFA rather than falling back to TCP — a silent fallback is the usual cause of "why is 2 nodes slower than 1?".

Step 4: Stream data instead of downloading it

Default File mode copies the whole channel to local disk before training starts. On a multi-hundred-GB corpus that is dead billed time, twice if you get interrupted. Use fast file mode:

from sagemaker.inputs import TrainingInput

train = TrainingInput(
    "s3://my-bucket/data/train/", input_mode="FastFile"
)
estimator.fit({"train": train})

FastFile exposes S3 objects as POSIX files with lazy streaming — near-zero startup, ideal for sharded webdataset/parquet. Keep shards at 100–500 MB and read them sequentially; FastFile is bad at random access inside huge files.

Step 5: Verify you are actually scaling

Before you launch a 24-hour run, do a 15-minute sanity pass and record three numbers:

Tokens/second per GPU at 1 node vs 2 nodes. Anything below ~80% scaling efficiency means you are network- or dataloader-bound, not compute-bound. Fix that before buying more nodes.
Peak GPU memory (torch.cuda.max_memory_allocated()). Under ~70% means you can raise micro-batch size and cut wall-clock.
GPU utilization from CloudWatch. Sustained dips point at the data loader — raise num_workers, enable pin_memory, and pre-tokenize offline.

Log all three to managed MLflow or SageMaker Experiments so the comparison survives the next engineer.

Common failures and their fixes

Symptom	Likely cause	Fix
OOM immediately at step 0	model wrapped as a single FSDP unit	add `transformer_auto_wrap_policy`
OOM at ~step 50	activation growth from long sequences	gradient checkpointing, then lower micro-batch
Job restarts from step 0 after interruption	checkpoints written outside `/opt/ml/checkpoints`	use `checkpoint_local_path`
2 nodes no faster than 1	`FULL_SHARD` over non-EFA networking	`HYBRID_SHARD`, or a single larger node
`max_wait` validation error	spot enabled without `max_wait >= max_run`	set both
NCCL timeout during checkpoint save	rank 0 slow-writing while others hit the barrier	raise the process-group timeout; save sharded

Cost model, concretely

A full-parameter 8B fine-tune on 2× ml.g5.12xlarge for 18 hours is roughly 36 instance-hours. At on-demand that is real money; at spot rates with checkpointing it is 30–40% of that, and the only operational cost is that the run may take longer in wall-clock while it waits for capacity. For anything that is not deadline-critical, spot plus checkpointing is the default we recommend to clients. Reserve on-demand for the final reproducible run you intend to ship, and for HyperPod flexible training plans when you need guaranteed capacity on a schedule.

Where this fits

Once the run finishes, the artifact lands in S3 and the rest of the loop takes over: register it in the Model Registry, gate it with an evaluation step, then deploy behind inference components so idle GPUs cost you nothing. Training is the part where careless configuration costs the most per hour — and where a two-line change to a wrap policy can be the difference between "does not fit" and "ships this week".

If your team is standing up multi-node training on SageMaker AI and wants a second pair of eyes on the sharding, checkpointing, and spot strategy before the bill arrives, get in touch.

Multi-Node FSDP Training on SageMaker AI: Managed Spot, Checkpointing, and Warm Pools