Basis mannequin coaching has reached an inflection level the place conventional checkpoint-based restoration strategies have gotten a bottleneck to effectivity and cost-effectiveness. As fashions develop to trillions of parameters and coaching clusters increase to 1000’s of AI accelerators, even minor disruptions may end up in vital prices and delays.
On this publish, we introduce checkpointless coaching on Amazon SageMaker HyperPod, a paradigm shift in mannequin coaching that reduces the want for conventional checkpointing by enabling peer-to-peer state restoration. Outcomes from production-scale validation present 80–93% discount in restoration time (from 15–half-hour or extra to underneath 2 minutes) and permits as much as 95% coaching goodput on cluster sizes with 1000’s of AI accelerators.
Understanding goodput
Basis mannequin coaching is without doubt one of the most resource-intensive processes in AI, usually involving thousands and thousands of {dollars} in compute spend throughout 1000’s of AI accelerators working for days to months. Due to the inherent all-or-none distributed synchrony throughout all ranks, even a lack of a single rank due to software program or {hardware} faults brings the coaching workloads to a whole halt. To mitigate such localized faults, the trade has relied on checkpoint-based restoration; periodically saving coaching states (checkpoints) to a sturdy retailer based mostly on a user-defined checkpoint interval. When a fault happens, the coaching workload resumes by restoring from the most recent saved checkpoint. This conventional restart-to-recover mannequin has turn out to be more and more untenable as mannequin sizes develop from billions to trillions of parameters and coaching workloads develop from lots of to 1000’s of AI accelerators.
This problem of sustaining environment friendly coaching operations at scale has led to the idea of goodput—the precise helpful work completed in an AI coaching system in comparison with its theoretical most capability. In basis mannequin coaching, goodput is impacted by system failures and restoration overhead. The hole between the system’s theoretical most throughput and its precise productive output (goodput) grows bigger with: elevated frequency of failures (which rises with cluster measurement), longer restoration instances (which scale with mannequin measurement and cluster measurement), and better prices of idle assets throughout restoration. This definition helps body why measuring and optimizing goodput turns into more and more essential as AI coaching scales to bigger clusters and extra complicated fashions, the place even small inefficiencies may end up in vital monetary and time prices.
A pre-training workload on a HyperPod cluster with 256 P5 situations, checkpointing each 20 minutes, faces two challenges when disrupted: 10 minutes of misplaced work plus 10 minutes for restoration. With ml.p5.24xlarge situations costing $55 per hour, every disruption prices $4,693 in compute time. For a month-long coaching, each day disruptions would accumulate to $141,000 in further prices and delay completion by 10 hours.
As cluster sizes develop, the chance and frequency of failures can improve.
Because the coaching spans throughout 1000’s of nodes, disruptions attributable to faults turn out to be more and more frequent. In the meantime, restoration turns into slower as a result of the workload reinitialization overhead grows linearly with cluster measurement. The cumulative impression of large-scale AI coaching failures can attain thousands and thousands of {dollars} yearly and translate on to delayed time-to-market, slower mannequin iteration cycles, and aggressive drawback. Each hour of idle GPU time is an hour not spent advancing mannequin capabilities.
Checkpoint-based restoration
Checkpoint-based restoration in distributed coaching is much extra complicated and time-consuming than generally understood. When a failure happens in conventional distributed coaching, the restart course of includes way over loading the final checkpoint. Understanding what occurs throughout restoration reveals why it takes so lengthy and why the whole cluster should sit idle.
The all-or-none cascade
A single failure—one GPU error, one community timeout, or one {hardware} fault—can set off a whole coaching cluster shutdown. As a result of distributed coaching treats all processes as tightly coupled, any single failure necessitates a whole restart. When any course of fails, the orchestration system (for instance, TorchElastic or Kubernetes) should terminate each course of throughout the job and restart from scratch. Every restart requires navigating a posh, multi-stage restoration course of the place each stage is sequential and blocking:
- Stage 1: Coaching job restart – The coaching job orchestrator detects a failure, terminates all processes in all nodes adopted by a cluster-wide restart or the coaching job.
- Stage 2: Course of and community initialization – Each course of should re-execute the coaching script from the start. That features rank initialization, loading of Python modules from sturdy retailer akin to Community File System (NFS) or object storage, establishing the coaching topology and communication backend by way of peer discovery and course of teams creation. The method group initialization alone can take tens of minutes on massive clusters.
- Stage 3: Checkpoint retrieval – Every course of should first establish the final fully saved checkpoint, then retrieve it from persistent storage (for instance, NFS or object storage) and cargo a number of state dictionaries: the mannequin’s parameters and buffers, the optimizer’s inner state (momentum, variance, and so forth), the training price scheduler, and coaching loop metadata (epoch, batch quantity). This step can take tens of minutes or longer relying on cluster and mannequin measurement.
- Stage 4: Knowledge loader initialization – The info-loading ranks have further duty to initialize the info buffers. That features retrieving the info checkpoint from sturdy storage akin to Amazon FSx or Amazon Easy Storage Service (Amazon S3) and prefetching the coaching knowledge to begin the coaching loop. Knowledge checkpointing is a necessary step to keep away from processing the identical knowledge samples a number of instances or skipping samples upon coaching disruption. Relying on the info combine technique, knowledge locality, and bandwidth, the method can take a couple of minutes.
- Stage 5: First step overhead – After checkpoint and coaching knowledge are retrieved and loaded, there’s further overhead to run the primary coaching step, we name it first step overhead (FSO). Throughout this primary step, there’s sometimes time spent in reminiscence allocation, creating and establishing the CUDA context for communication with GPUs, and compilation a part of the CUDA graph, and so forth.
- Stage 6: Misplaced steps overhead – Solely in any case earlier phases full efficiently can the coaching loop resume its common progress. As a result of the coaching resumes from the final saved mannequin checkpoint, all of the steps computed between the checkpoint and the fault encountered are misplaced. These misplaced steps should be recomputed, we name this misplaced steps overhead (LSO). Following the recomputation section, the coaching job resumes productive work that straight contributes to goodput.
How checkpointless coaching eliminates these bottlenecks
The 5 phases outlined above—termination and restart, course of discovery and community setup, checkpoint retrieval, GPU context reinitialization, and coaching loop resumption—symbolize the elemental bottlenecks in checkpoint-based restoration. Every stage is sequential and blocking, and coaching restoration can take minutes to a number of hours for big fashions. Critically, the whole cluster should wait for each stage to finish earlier than coaching can resume.
Checkpointless coaching eliminates this cascade. Checkpointless coaching preserves mannequin state coherence throughout the distributed cluster, eliminating the necessity for periodic snapshots. When failures happen, the system rapidly recovers through the use of wholesome friends, avoiding each storage I/O operations and full course of restarts sometimes required by conventional checkpointing approaches.
Checkpointless coaching structure
Checkpointless coaching is constructed on 5 elements that work collectively to get rid of the standard checkpoint-restart bottlenecks. Every part addresses a selected bottleneck within the restoration course of, and collectively they allow automated detection and restoration of infrastructure faults in minutes with zero guide intervention, even with 1000’s of AI accelerators.
Element 1: TCPStore-less/root-less NCCL and Gloo initialization (optimizing stage 2)
In a typical distributed coaching setup (for instance, utilizing torch.distributed), all ranks should initialize a course of group. The method group creates a communication layer, permitting all processes (or ranks, that’s, particular person nodes) to pay attention to one another and trade info. A TCPStore is commonly used as a rendezvous level the place all ranks test in to find one another’s connection info. When 1000’s of ranks attempt to contact a designated root server (sometimes rank 0) concurrently, it turns into a bottleneck. This results in a flood of simultaneous community requests to a single root server that may trigger community congestion, improve latency by tens of minutes, and additional sluggish the communication course of.
Checkpointless coaching eliminates this centralized dependency. As a substitute of funneling all connection requests by way of a single root server, the system makes use of a symmetric deal with sample the place every rank independently computes peer connection info utilizing a world group counter. Ranks join straight to one another utilizing predetermined port assignments, avoiding the TCPStore bottleneck. Course of group initialization drops from tens of minutes to seconds, even on clusters with 1000’s of nodes. The system additionally eliminates the single-point-of-failure danger inherent in root-based initialization.
Element 2: Reminiscence-mapped knowledge loading (optimizing stage 4)
One of many hidden prices in conventional restoration is reloading coaching knowledge. When a course of restarts, it should reload batches from disk, rebuild knowledge loader state, and punctiliously place itself to keep away from processing duplicate samples or skipping knowledge. On large-scale coaching runs, this knowledge loading can add minutes to each restoration cycle.
Checkpointless coaching makes use of memory-mapped knowledge loading to keep up cached knowledge throughout accelerators. Coaching knowledge is mapped into shared reminiscence areas that persist even when particular person processes fail. When a node recovers, it doesn’t reload knowledge from disk however reconnects to the prevailing memory-mapped cache. The info loader state is preserved, serving to to make sure that coaching continues from the proper place with out duplicate or skipped samples. MMAP additionally reduces host CPU reminiscence utilization by sustaining just one copy of information per node (in comparison with eight copies with conventional knowledge loaders on 8-GPU nodes), and coaching can resume instantly utilizing cached batches whereas the info loader concurrently prefetches the subsequent knowledge within the background.
Reminiscence-mapped knowledge loading workflow
Element 3: In-process restoration (optimizing stage 1, 2, and 5)
Conventional checkpoint-based restoration treats failures as job-level occasions: a single GPU error triggers termination of the whole distributed coaching job. Each course of throughout the cluster should be killed and restarted, despite the fact that just one part failed.
Checkpointless coaching makes use of in-process restoration to isolate failures on the course of stage. When a GPU or course of fails, solely the failed course of executes an in-process restoration to rejoin the coaching loop inside seconds, overcoming recoverable or transient errors. Wholesome processes proceed working with out interruption. The failed course of stays alive (avoiding full course of teardown), preserving the CUDA context, compiler cache, and GPU state, therefore eliminating minutes of reinitialization overhead. In instances the place the error is non-recoverable (akin to {hardware} failure), the system routinely swaps the defective part with a pre-warmed scorching spare, enabling coaching to proceed with out disruptions.
This eliminates the necessity for full cluster termination and restart, dramatically decreasing restoration overhead.
Element 4: Peer-to-peer state replication (optimizing stage 3 and 6)
Checkpoint-based restoration requires loading mannequin and optimizer state from persistent storage (akin to Amazon S3 or FSx for Lustre). For fashions with billions to trillions of parameters, this implies transferring tens to lots of of gigabytes over the community, deserializing state dictionaries, and reconstructing optimizer buffers which might take tens of minutes and create a large I/O bottleneck.
Essentially the most vital innovation in checkpointless coaching is steady peer-to-peer state replication. As a substitute of periodically saving mannequin state to centralized storage, every GPU maintains redundant copies of its mannequin shards on peer GPUs. When a failure happens, the recovering course of doesn’t load from Amazon S3. It copies state straight from a wholesome peer over the high-speed Elastic Cloth Adapter (EFA) community interconnect. This peer-to-peer structure eliminates the I/O bottleneck that dominates conventional checkpoint restoration. State switch occurs in seconds, in comparison with minutes for loading multi-gigabyte checkpoints from storage. The recovering node pulls solely the precise shards it wants, additional decreasing switch time.
Element 5: SageMaker HyperPod coaching operator (optimizing all phases)
The SageMaker HyperPod coaching operator orchestrates the checkpointless coaching elements, serving because the coordination layer that ties collectively initialization, knowledge loading, checkpointless restoration, and checkpoint fallback mechanisms. It maintains a centralized management aircraft with a world view of coaching course of well being throughout the whole cluster, coordinating fault detection, restoration choices, and cluster-wide synchronization.
The operator implements clever restoration escalation: it first makes an attempt in-process restart for failed elements, and if that’s not possible (for instance, due to container crashes or node failures), it escalates to process-level restoration. Throughout a process-level restoration, as an alternative of restarting the whole job when failures happen, the operator restarts solely coaching processes, retaining the containers alive. Consequently, the restoration instances are quicker than a job-level restart, which requires tearing down and recreating the coaching infrastructure, involving pod rescheduling, container pulls, surroundings initialization, and re-loading from checkpoints. When failures happen, the operator broadcasts coordinated cease alerts to stop cascading timeouts and integrates with the SageMaker HyperPod health-monitoring agent to routinely detect {hardware} points and set off restoration with out guide intervention.
Getting began with checkpointless coaching
This part guides you thru establishing and configuring checkpointless coaching on SageMaker HyperPod to scale back fault restoration from hours to minutes.
Conditions
Earlier than integrating checkpointless coaching into your coaching workload, confirm that your surroundings meets the next necessities:
Infrastructure necessities:
Software program necessities:
- Supported frameworks: Nemo, PyTorch, PyTorch Lightning
- Coaching knowledge codecs: JSON, JSONGZ (compressed JSON), or ARROW
- Amazon Elastic Container Registry (Amazon ECR) repository for container photographs. Use the HyperPod checkpointless coaching container—required for rootless NCCL initialization (Tier 1) and peer-to-peer checkpointless restoration (Tier 4)
Checkpointless coaching workflow
Checkpointless coaching is designed for incremental adoption. You can begin with primary capabilities and progressively allow superior options as your coaching scales. The mixing is organized into 4 tiers, every constructing on the earlier one:
Tier 1: NCCL initialization optimization
NCCL initialization optimization eliminates the centralized root course of bottleneck throughout initialization. Nodes uncover and connect with friends independently utilizing infrastructure alerts. This permits quicker course of group initialization (seconds as an alternative of minutes) and elimination of single-point-of-failure throughout startup.
Integration steps: Allow an surroundings variable as a part of the job specification and confirm that the job runs with the checkpointless coaching container.
Tier 2: Reminiscence-mapped knowledge loading
Reminiscence mapped knowledge loading retains coaching knowledge cached in shared reminiscence throughout course of restarts, eliminating knowledge reload overhead throughout restoration. This permits prompt knowledge entry throughout restoration. No must reload or re-shuffle knowledge when a course of restarts.
Integration steps: Increase the prevailing knowledge loader with a reminiscence mapped cache
Tier 3: In-process restoration
In-process restoration isolates failures to particular person processes as an alternative of requiring full job restarts. Failed processes recuperate independently whereas wholesome processes proceed coaching. It permits sub-minute restoration from process-level failures. Wholesome processes keep alive, whereas failed processes recuperate independently.
Integration steps:
Tier 4: Checkpointless (peer-to-peer restoration) (NeMo integration)
Checkpointless restoration permits full peer-to-peer state replication and restoration. Failed processes recuperate mannequin and optimizer state straight from wholesome friends with out loading from storage. This step permits elimination of checkpoint loading. Failed processes recuperate mannequin and optimizer state from wholesome replicas over the high-speed EFA interconnect.
Integration steps:
wait_rank: All ranks will await the rank info from the Hyperod coaching operator infrastructure.
HPWrapper: Python perform wrapper that permits restart capabilities for a restart code block (RCB). The implementation makes use of a context supervisor as an alternative of a Python decorator as a result of the decision wrapper lacks details about the variety of RCBs it ought to monitor.
CudaHealthCheck: Helps make sure that the CUDA context for the present course of is in a wholesome state. It synchronizes with the GPU and makes use of the system similar to LOCAL_RANK surroundings variable, or the principle thread’s default CUDA system if LOCAL_RANK was not specified within the surroundings.
HPAgentK8sAPIFactory: That is the API that checkpointless coaching will use to know the coaching standing from the opposite pods in a K8s coaching cluster. It additionally supplies an infrastructure-level barrier, which makes certain each rank can efficiently carry out the abort and restart.
CheckpointManager: Manages in-memory checkpoints and peer-to-peer restoration for checkpointless fault tolerance.
We advocate beginning with Tier 1 and validating it in your surroundings. Add Tier 2 when knowledge loading overhead turns into a bottleneck. Undertake Tier 3 and Tier 4 for max resilience on the most important coaching clusters.
For NeMo customers and HyperPod recipe customers, Tier 4 is out there out-of-the-box with minimal configuration modifications for Llama and GPT open supply recipes. NeMo examples for Llama and GPT open supply fashions could be present in SageMaker HyperPod checkpointless coaching.
Efficiency outcomes
Checkpointless coaching has been validated at manufacturing scale throughout a number of cluster configurations. The most recent Amazon Nova fashions had been educated utilizing this expertise on tens of 1000’s of AI accelerators.
On this part, we exhibit outcomes from intensive testing throughout a variety of cluster sizes, spanning 16 GPUs to 2,304 GPUs. Checkpointless coaching demonstrated vital enhancements in restoration time, constantly decreasing downtime by 80–93% in comparison with conventional checkpoint-based restoration.
| Cluster (H100s) | Mannequin | Conventional restoration | Checkpointless restoration | Enchancment |
|---|---|---|---|---|
| 2,304 GPUs | Inner mannequin | 15–half-hour | Lower than 2 minutes | ~87–93% quicker |
| 256 GPUs | Llama-3 70B (pre-training) | 4 min, 52 sec | 47 seconds | ~84% quicker |
| 16 GPUs | Llama-3 70B (fine-tuning) | 5 min 10 sec | 50 seconds | ~84% quicker |
These restoration time enhancements have a direct relationship to ML goodput, outlined as the proportion of time your cluster spends making ahead progress on coaching somewhat than sitting idle throughout failures. As clusters scale to 1000’s of nodes, failure frequency will increase proportionally. On the identical time, conventional checkpoint-based restoration instances additionally improve with cluster measurement on account of rising coordination overhead. This creates a compounding downside: extra frequent failures mixed with longer restoration instances quickly erode goodput at scale.
Checkpointless coaching makes optimizations throughout the whole restoration stack, enabling greater than 95% goodput even on clusters with 1000’s of AI accelerators. Based mostly on our inner research, we constantly noticed goodput upwards of 95% throughout massive-scale deployments that exceeded 2,300 GPUs.
We additionally verified that mannequin coaching accuracy isn’t impacted by checkpointless coaching. Particularly, we measured checksum matching for conventional checkpoint-based coaching and checkpointless coaching, and at each coaching step verified a bit-wise match on coaching loss. The next is a plot for the coaching loss for a Llama-3 70B pre-training workload on 32 x ml.p5.48xlarge situations for each conventional checkpointing versus checkpointless coaching.
Conclusion
Basis mannequin coaching has reached an inflection level. As clusters scale to 1000’s of AI accelerators and coaching runs lengthen to months, the standard checkpoint-based restoration paradigm is more and more turning into a bottleneck. A single GPU failure that beforehand would have induced minutes of downtime now triggers tens of minutes of cluster-wide idle time on 1000’s of AI accelerators, with cumulative prices reaching thousands and thousands of {dollars} yearly.
Checkpointless coaching rethinks this paradigm fully by treating failures as native, recoverable occasions somewhat than cluster-wide catastrophes. Failed processes recuperate state from wholesome friends in seconds, enabling the remainder of the cluster to proceed making ahead progress. The shift is key: from How can we restart rapidly? to How can we keep away from stopping in any respect?
This expertise has enabled greater than 95% goodput when coaching on SageMaker HyperPod. Our inner research on 2,304 GPUs present restoration instances dropped from 15–half-hour to underneath 90 seconds, translating to over 80% discount in idle GPU time per failure.
To get began, discover What’s Amazon SageMaker AI?. Pattern implementations and recipes can be found within the AWS GitHub HyperPod checkpointless coaching and SageMaker HyperPod recipes repositories.
Concerning the Authors
Anirudh Viswanathan is a Senior Product Supervisor, Technical, at AWS with the SageMaker crew, the place he focuses on Machine Studying. He holds a Grasp’s in Robotics from Carnegie Mellon College and an MBA from the Wharton Faculty of Enterprise. Anirudh is a named inventor on greater than 50 AI/ML patents. He enjoys long-distance working, exploring artwork galleries, and attending Broadway reveals. You possibly can join with Anirudh on LinkedIn.
Roy Allela is a Senior AI/ML Specialist Options Architect at AWS. He helps AWS clients, from small startups to massive enterprises to coach and deploy basis fashions effectively on AWS. He has a background in Microprocessor Engineering enthusiastic about computational optimization issues and bettering the efficiency of AI workloads. You possibly can join with Roy on LinkedIn.
Fei Wu is a Senior Software program Developer at AWS with Sagemaker crew. Fei’s focus is on ML system and distributed coaching strategies. He holds a PhD in Electrical Engineering from StonyBrook College. When outdoors of labor, Fei enjoys enjoying basketball and watching motion pictures. You possibly can join with Fei on LinkedIn.
Trevor Harvey is a Principal Specialist in Generative AI at Amazon Net Companies (AWS) and an AWS Licensed Options Architect – Skilled. At AWS, Trevor works with clients to design and implement machine studying options and leads go-to-market methods for generative AI companies.
Anirban Roy is a Principal Engineer at AWS with the SageMaker crew, primarily focussing on AI coaching infra, resiliency and observability. He holds a Grasp’s in Laptop Science from Indian Statistical Institute in Kolkata. Anirban is a seasoned distributed software program system builder with greater than 20 years of expertise and a number of patents and publications. He enjoys street biking, studying non-fiction, gardening and nature touring. You possibly can join with Anirban on LinkedIn
Arun Nagarajan is a Principal Engineer on the Amazon SageMaker AI crew, the place he at the moment focuses on distributed coaching throughout the whole stack. Since becoming a member of the SageMaker crew throughout its launch yr, Arun has contributed to a number of merchandise inside SageMaker AI, together with real-time inference and MLOps options. When he’s not engaged on machine studying infrastructure, he enjoys exploring the outside within the Pacific Northwest and hitting the slopes for snowboarding.
