Checkpointless coaching on Amazon SageMaker HyperPod: Manufacturing-scale coaching with quicker fault restoration

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 develop to 1000’s of AI accelerators, even minor disruptions can lead to important prices and delays.

On this submit, 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 below 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, typically 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 newest 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 a whole lot 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 dimension), longer restoration instances (which scale with mannequin dimension and cluster dimension), 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 advanced fashions, the place even small inefficiencies can lead to important monetary and time prices.

A pre-training workload on a HyperPod cluster with 256 P5 cases, checkpointing each 20 minutes, faces two challenges when disrupted: 10 minutes of misplaced work plus 10 minutes for restoration. With ml.p5.24xlarge cases 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 additional prices and delay completion by 10 hours.

As cluster sizes develop, the likelihood and frequency of failures can improve.

Because the coaching spans throughout 1000’s of nodes, disruptions brought on by 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 dimension. 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 way extra advanced and time-consuming than generally understood. When a failure happens in conventional distributed coaching, the restart course of entails excess of loading the final checkpoint. Understanding what occurs throughout restoration reveals why it takes so lengthy and why your complete 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 an entire coaching cluster shutdown. As a result of distributed coaching treats all processes as tightly coupled, any single failure necessitates an entire 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 corresponding to Community File System (NFS) or object storage, establishing the coaching topology and communication backend by peer discovery and course of teams creation. The method group initialization alone can take tens of minutes on giant clusters.
• Stage 3: Checkpoint retrieval – Every course of should first establish the final utterly 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 educational price scheduler, and coaching loop metadata (epoch, batch quantity). This step can take tens of minutes or longer relying on cluster and mannequin dimension.
• Stage 4: Knowledge loader initialization – The information-loading ranks have extra accountability to initialize the information buffers. That features retrieving the information checkpoint from sturdy storage corresponding 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 information 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 may be extra overhead to run the primary coaching step, we name it first step overhead (FSO). Throughout this primary step, there may be usually 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 levels 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 immediately contributes to goodput.

How checkpointless coaching eliminates these bottlenecks

The 5 levels 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 giant fashions. Critically, your complete 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 by utilizing wholesome friends, avoiding each storage I/O operations and full course of restarts usually required by conventional checkpointing approaches.

Checkpointless coaching structure

Checkpointless coaching is constructed on 5 elements that work collectively to remove the normal checkpoint-restart bottlenecks. Every element addresses a selected bottleneck within the restoration course of, and collectively they permit automated detection and restoration of infrastructure faults in minutes with zero handbook intervention, even with 1000’s of AI accelerators.

Part 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 concentrate on one another and alternate 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 (usually 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 gradual the communication course of.

Checkpointless coaching eliminates this centralized dependency. As a substitute of funneling all connection requests by 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 worldwide group counter. Ranks join immediately 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.

Part 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 take care of 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 present memory-mapped cache. The information loader state is preserved, serving to to make sure that coaching continues from the right place with out duplicate or skipped samples. MMAP additionally reduces host CPU reminiscence utilization by sustaining just one copy of knowledge 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 information loader concurrently prefetches the subsequent knowledge within the background.

Reminiscence-mapped knowledge loading workflow

Part 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 your complete distributed coaching job. Each course of throughout the cluster should be killed and restarted, despite the fact that just one element 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 circumstances the place the error is non-recoverable (corresponding to {hardware} failure), the system mechanically swaps the defective element with a pre-warmed sizzling spare, enabling coaching to proceed with out disruptions.

This eliminates the necessity for full cluster termination and restart, dramatically decreasing restoration overhead.

Part 4: Peer-to-peer state replication (optimizing stage 3 and 6)

Checkpoint-based restoration requires loading mannequin and optimizer state from persistent storage (corresponding to Amazon S3 or FSx for Lustre). For fashions with billions to trillions of parameters, this implies transferring tens to a whole lot of gigabytes over the community, deserializing state dictionaries, and reconstructing optimizer buffers which may take tens of minutes and create an enormous I/O bottleneck.

Essentially the most crucial 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 immediately from a wholesome peer over the high-speed Elastic Material 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.

Part 5: SageMaker HyperPod coaching operator (optimizing all levels)

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 worldwide view of coaching course of well being throughout your complete 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 a substitute of restarting your complete job when failures happen, the operator restarts solely coaching processes, preserving the containers alive. Because of this, the restoration instances are quicker than a job-level restart, which requires tearing down and recreating the coaching infrastructure, involving pod rescheduling, container pulls, atmosphere initialization, and re-loading from checkpoints. When failures happen, the operator broadcasts coordinated cease indicators to stop cascading timeouts and integrates with the SageMaker HyperPod health-monitoring agent to mechanically detect {hardware} points and set off restoration with out handbook intervention.

Getting began with checkpointless coaching

This part guides you thru establishing and configuring checkpointless coaching on SageMaker HyperPod to cut back fault restoration from hours to minutes.

Conditions

Earlier than integrating checkpointless coaching into your coaching workload, confirm that your atmosphere 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 pictures. Use the HyperPod checkpointless coaching container—required for rootless NCCL initialization (Tier 1) and peer-to-peer checkpointless restoration (Tier 4)

658645717510.dkr.ecr..amazonaws.com/sagemaker-hyperpod/pytorch-training:2.3.0-checkpointless

Checkpointless coaching workflow

Checkpointless coaching is designed for incremental adoption. You can begin with fundamental 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 indicators. This allows quicker course of group initialization (seconds as a substitute of minutes) and elimination of single-point-of-failure throughout startup.

Integration steps: Allow an atmosphere variable as a part of the job specification and confirm that the job runs with the checkpointless coaching container.

# kubernetes job spec
env:
  - identify: HPCT_USE_CONN_DATA # Allow Rootless
    worth: "1"
  - identify: TORCH_SKIP_TCPSTORE # Allow TCPStore Removing
    worth: "1"

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 allows on the spot knowledge entry throughout restoration. No have to reload or re-shuffle knowledge when a course of restarts.

Integration steps: Increase the present knowledge loader with a reminiscence mapped cache

from hyperpod_checkpointless_training.dataloader.mmap_data_module import MMAPDataModule
from hyperpod_checkpointless_training.dataloader.config import CacheResumeMMAPConfig

base_data_module = MY_DATA_MODULE(...). # Buyer's personal datamodule

mmap_config = CacheResumeMMAPConfig(
    cache_dir=self.cfg.mmap.cache_dir,
)

mmap_dm = MMAPDataModule(
    data_module=base_data_module,
    mmap_config=CacheResumeMMAPConfig(
        cache_dir=self.cfg.mmap.cache_dir,
    ),
)

Tier 3: In-process restoration

In-process restoration isolates failures to particular person processes as a substitute of requiring full job restarts. Failed processes get better independently whereas wholesome processes proceed coaching. It permits sub-minute restoration from process-level failures. Wholesome processes keep alive, whereas failed processes get better independently.

Integration steps:

from hyperpod_checkpointless_training.inprocess.health_check import CudaHealthCheck
from hyperpod_checkpointless_training.inprocess.wrap import HPCallWrapper, HPWrapper
from hyperpod_checkpointless_training.inprocess.train_utils import HPAgentK8sAPIFactory
@HPWrapper(
    health_check=CudaHealthCheck(),
    hp_api_factory=HPAgentK8sAPIFactory(),
    abort_timeout=60.0,
)
def re_executable_codeblock(): # The re-executable codeblock outlined by person, normally it is foremost operate or prepare loop
    ...

Tier 4: Checkpointless (peer-to-peer restoration) (NeMo integration)

Checkpointless restoration permits full peer-to-peer state replication and restoration. Failed processes get better mannequin and optimizer state immediately from wholesome friends with out loading from storage. This step permits elimination of checkpoint loading. Failed processes get better mannequin and optimizer state from wholesome replicas over the high-speed EFA interconnect.

Integration steps:

from hyperpod_checkpointless_training.inprocess.train_utils import wait_rank
    wait_rank() 
    
def foremost():   
    @HPWrapper(
        health_check=CudaHealthCheck(),
        hp_api_factory=HPAgentK8sAPIFactory(),
        abort_timeout=60.0,
        checkpoint_manager=PEFTCheckpointManager(enable_offload=True),
        abort=CheckpointlessAbortManager.get_default_checkpointless_abort(),
        finalize=CheckpointlessFinalizeCleanup(),
    )
    def run_main(cfg, caller: Elective[HPCallWrapper] = None):
        ...
        coach = Coach(
            technique=CheckpointlessMegatronStrategy(...,
                num_distributed_optimizer_instances=2),
            callbacks=[..., CheckpointlessCallback(...)],
            )
        coach.fresume = resume
        coach._checkpoint_connector = CheckpointlessCompatibleConnector(coach)
        coach.wrapper = caller

wait_rank: All ranks will look forward to the rank info from the Hyperod coaching operator infrastructure.

HPWrapper: Python operate wrapper that permits restart capabilities for a restart code block (RCB). The implementation makes use of a context supervisor as a substitute of a Python decorator as a result of the decision wrapper lacks details about the variety of RCBs it ought to monitor.

CudaHealthCheck: Helps be certain that the CUDA context for the present course of is in a wholesome state. It synchronizes with the GPU and makes use of the gadget similar to LOCAL_RANK atmosphere variable, or the primary thread’s default CUDA gadget if LOCAL_RANK was not specified within the atmosphere.

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 positive 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 atmosphere. Add Tier 2 when knowledge loading overhead turns into a bottleneck. Undertake Tier 3 and Tier 4 for optimum resilience on the biggest coaching clusters.

For NeMo customers and HyperPod recipe customers, Tier 4 is on the market out-of-the-box with minimal configuration modifications for Llama and GPT open supply recipes. NeMo examples for Llama and GPT open supply fashions will be present in SageMaker HyperPod checkpointless coaching.

Efficiency outcomes

Checkpointless coaching has been validated at manufacturing scale throughout a number of cluster configurations. The newest Amazon Nova fashions have been educated utilizing this know-how on tens of 1000’s of AI accelerators.

On this part, we show outcomes from intensive testing throughout a variety of cluster sizes, spanning 16 GPUs to 2,304 GPUs. Checkpointless coaching demonstrated important enhancements in restoration time, persistently decreasing downtime by 80–93% in comparison with conventional checkpoint-based restoration.

These restoration time enhancements have a direct relationship to ML goodput, outlined as the share of time your cluster spends making ahead progress on coaching fairly than sitting idle throughout failures. As clusters scale to 1000’s of nodes, failure frequency will increase proportionally. On the similar time, conventional checkpoint-based restoration instances additionally improve with cluster dimension attributable to 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 your complete restoration stack, enabling greater than 95% goodput even on clusters with 1000’s of AI accelerators. Based mostly on our inner research, we persistently noticed goodput upwards of 95% throughout massive-scale deployments that exceeded 2,300 GPUs.

We additionally verified that mannequin coaching accuracy just 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 cases 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 prolong to months, the normal checkpoint-based restoration paradigm is more and more turning into a bottleneck. A single GPU failure that beforehand would have brought about 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 solely by treating failures as native, recoverable occasions fairly than cluster-wide catastrophes. Failed processes get better state from wholesome friends in seconds, enabling the remainder of the cluster to proceed making ahead progress. The shift is prime: from How can we restart rapidly? to How can we keep away from stopping in any respect?

This know-how 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 below 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.

In regards to 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 College 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 may join with Anirudh on LinkedIn.

Roy Allela is a Senior AI/ML Specialist Options Architect at AWS. He helps AWS prospects, from small startups to giant enterprises to coach and deploy basis fashions effectively on AWS. He has a background in Microprocessor Engineering enthusiastic about computational optimization issues and enhancing the efficiency of AI workloads. You may 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 taking part in basketball and watching films. You may 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 prospects 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 may 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 your complete stack. Since becoming a member of the SageMaker crew throughout its launch 12 months, 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.