1 Introduction↩︎

Large language models (LLM), based on the transformer architecture, have emerged as foundational models driving diverse applications. However, training these models incurs substantial computational and memory demands. To achieve state-of-the-art performance in quality and sample efficiency [1], contemporary LLMs are progressively scaled up in model size. The immense computational demands and prolonged execution times of LLM training necessitate deployment on shared large-scale clusters, where hybrid-parallel training jobs are submitted to the cluster job scheduler.

In shared clusters, the GPU resources visible to a long-running training job do not remain static. Job arrivals, job completions, quota reclaiming, and scheduler-driven reallocation can continuously reshape the device set assigned to a running job. Therefore, elastic hybrid-parallel training should be prepared to migrate model state whenever the available device set changes, rather than assuming a fixed cluster partition throughout the entire training process. To achieve better resource utilization, a training job should scale in/out when fluctuating GPU resources become available due to cluster load change, and the scheduler should frequently reconcile the GPU allocation among jobs in a proactive manner (i.e., preemption and scaling), especially on shared clusters [2]–[4] that co-locates inference jobs and training jobs. Therefore, state migration—the process of transferring model states when moving a training job from device set \(A\) to device set \(B\)—is essential for hybrid-parallel training jobs, especially those deployed in shared clusters.

Standard solutions of state migration for a hybrid-parallel training job predominantly rely on the checkpointing mechanism, which checkpoints complete states (parameters, optimizer states, etc.) to persistent storage before migration and then resumes from the states after the restart. If the parallel configuration (i.e., pipeline/tensor-parallel dimension) should be altered according to dynamic resources, the state should be resharded to match the new parallel configuration before resuming. Unfortunately, checkpointing and resharding reside on the critical path of the training loop. Consequently, frequent migrations triggered by elastic training force all participating GPUs to stall computation during the lengthy checkpointing, leading to significant resource under-utilization.

Recently, many efforts have attempted to reduce checkpointing overhead. A common practice is to leverage DRAM for intermediate persistence, where the checkpoint is then transferred from DRAM to storage asynchronously [5]–[9]. Based on that, PCcheck [6] splits the checkpoint into chunks to enable concurrent chunk persistence with multiple threads; Gemini [9] presents a placement strategy to maximize failure recovery and a traffic scheduling algorithm to avoid the interference of checkpoint traffic on model training; Tenplex [10] focuses on dynamic parallelism transformation through an in-memory file system and RDMA. Besides, Check-N-Run [11] performs differential checkpointing on values and leverages quantization to reduce checkpoint size. Oobleck [12] exploits inherent redundancy of the replicated model states through pipeline templates for fast recovery. Despite these optimization efforts, checkpoint-based migration remains fundamentally costly. Even the state-of-the-art techniques still incur substantial latency (often tens to hundreds of seconds [6], [10]) because they inevitably involve moving massive amounts of state tensors across the memory hierarchy (e.g., GPU memory\(\to\)host DRAM\(\to\)disk). As demonstrated in Figure 1 ,this process suffers from the bandwidth limitations of storage devices, and disregards the crucial state locality residing in GPU memory, necessitating the expensive transfer and persistence of the entire model state.

To overcome above issues, we propose ETC, a checkpoint-free state migration framework with minimal data movement guarantee for elastic hybrid-parallel LLM training. ETC can migrate the model states of a hybrid-parallel configuration (e.g., data/pipeline/tensor parallel) to the states of another configuration (denoted as reconfiguration). Unlike other checkpoint-based migrations, all migrations in ETC are completely through inter-device communication with no need of any heavy I/O persistence, considering the common practice of equipping state-of-the-art GPUs (or NPUs) with high-bandwidth interconnects. The key idea of ETC is exploiting the model state locality to keep more states in-place, so as to minimize the data movement during migration. Note that ETC can be integrated with automatic parallel frameworks (e.g., Alpa [13], Unity [14], AutoDDL [15]), as ETC switches model states according to the timely generated configurations efficiently, which together optimize the holistic process of elastic hybrid-parallel training.

ETC provides a collective communication primitive for the state migration of hybrid-parallel reconfiguration. The primitive takes the current & previous parallel configuration, and model state as input (§3), and then derives the cost matrix to describe data movement (i.e., quantity of model states requiring transmission) between current ranks and previous ranks (§3.2 for pipeline parallel and §3.3 for tensor/data parallel) and derive coarse-grained migration sketch, and then generates the fine-grained device-to-device communication instructions for migration (§3.4), and then performs resource de-fragmentation through additional migration operators passed to the execution engine (§3.5). We also coalesce the migration instructions to avoid the redundant indirect communications. Specifically, the key contributions are as follows:

• We formalize a cost matrix to represent the data movement across devices, and employ a cost-matrix-driven approach to minimize the data movement globally.

• We design the execution engine to transform the coarse-grained migration sketch into fine-grained Send, Recv, and Refer instructions.

• We propose a virtual rank allocator, which decouples the virtual and physical process ranks to eliminate resource fragmentation. It guarantees sustained training efficiency, resource occupancy, and minimized data movement.

• We deploy ETC in Megatron-LM to enable the elastic training of hybrid-parallel transformer-based models. Evaluation results show that ETC can accelerate the model state migration under various model sizes and parallel configurations.

2 Motivation↩︎

2.1 Hybrid-Parallel LLM Training↩︎

Efficient LLM training requires hybrid-parallelism—the combination of data, tensor, pipeline, and other parallelisms [13], [16]–[19]. A large-scale training job should be broken down into sub-tasks and then assigned to multiple devices for parallel execution, thereby balancing computation, memory, and communication among devices. Currently, the LLM training jobs are executed in fixed degree of parallelism (DoP), and they follow the fixed parallel configuration during training. 1) Data parallel (DP) [16], [17]: It splits the training data into multiple mini-batches. Each device loads a complete model replica, computes gradients of mini-batch, and synchronizes gradients with Allreduce for model update. 2) Pipeline parallel (PP) [20], [21]: It partitions the model layers into multiple stages, where each device is responsible for one or more stages. And a mini-batch is split into micro-batches and then processed in a pipeline manner. In PP, each device sends/receives intermediate tensors to/from its neighbor using peer-to-peer (P2P) communications. 3) Tensor parallel (TP) [18], [22]: It shards the model parameters at tensor level, and places each shard on a device. When a full parameter or activation tensor is required, TP uses collective communications to reassemble the full tensor from corresponding shards.

2.2 Demands of State Migration under Resource Fluctuation↩︎

LLM training relies on large-scale GPU clusters and features a training cycle that spans several days to several months. During this process, the resource fluctuation is common on shared clusters, and dynamic resource changes for jobs are inevitable. Public traces from Aurora supercomputer show that these disturbances are frequent at the job-event level. In January 2026 alone, the cluster records 32,256/32,249 GPU job starts/ends. Figure 2 provides the minute-level start/end events of GPU jobs (i.e., 1.4 events/min on average) and shows that job fluctuation affects long-running training jobs persistently. Besides, the training/development clusters also exhibit diurnal patterns, and the cluster schedulers usually adopt the minute-level scheduling cycles (e.g., 60s) [27]. For hybrid-parallel jobs, such fluctuation can force migrations.

These cluster-level fluctuations reach training jobs through several common triggers: 1) Preemption. High-priority jobs can take GPU resources from low-priority running jobs [28]. For instance, when co-locating inference jobs and training jobs in a shared cluster [2]–[4], training jobs can be preempted by inference jobs to meet strict service-level objectives (SLO) [29], [30]. Even in training clusters [25], [31], a training job can be easily preempted due to constant resource revocation. Because the cluster scheduler tends to allocate the currently idle GPUs (but belonging to other users’ quota) to the running jobs, when other users submit jobs, these GPUs are reclaimed and the corresponding jobs are preempted. As for spot instances [12], [24], [32], the jobs are more likely to be preempted. For example, a 24-hour trace on AWS EC2 suffers from 127 preemptions [32]. 2) Scaling. To guarantee the strict latency SLO, GPU resources are commonly over-provisioned for inference jobs when tackling bursting inference requests (up to 50\(\times\) than the average [33]). It leads to resource under-utilization during off-peak hours [34]–[36]. To saturate GPU resources, training jobs should scale out swiftly to utilize idle GPUs, although they would be preempted or proactively scale in to avoid interference against inference jobs.

2.3 Opportunity: Minimize Data Movement in State Migration↩︎

The change of device number must modify the parallel configuration and thus disrupts the collaborative logic among devices, which causes the state migration of training jobs. Therefore, to adapt to stage migration, standard solutions usually rely on the checkpointing mechanism (the process of checkpoint–reshard–resume) [37], [38].

However, as shown in Table ¿tbl:tab:method95comparison?, they suffer from huge data movement overhead over migration, hindering their adaptation to frequent-migration scenarios. Specifically, 1) Checkpointing enforces the persistence through I/O, including both directly persisting to disk and redirecting to disk from CPU DRAM. In that scenario, the data movement can be the full model state (i.e., a replica of parameters and optimizer states). While the approaches targeting migration within the limited parallel configurations (e.g., Varuna focus on pipeline dimension) should also persist the partial state. 2) When changing hybrid-parallel configurations, existing approaches perform state transformation (i.e., reshard) on persisted checkpointing files, leading to additional access to the full model state. 3) The state locality is usually ignored because of unawareness of the device ranks. Consequently, they fail to transmit the differential states. Therefore, to enable efficient state migration, we should try to minimize the data movement. We turn to a checkpoint-free approach by keeping model state in GPU memory and performing reconfiguration with inter-GPU connections.

Closely related to our study is Oobleck [12] and Tenplex [10]. However, Oobleck only support migration among pipeline templates, while it sacrifices the job’s fine-grained control over the parallelism, rendering it inadequate for general-purpose migration across arbitrary device sets; and Tenplex employs an in-memory file system for persistence and differential resharding, but still requires transferring full model states with CPU DRAM.

3 ETC Design↩︎

3.1 Overview↩︎

ETC is a model state migration framework containing the collective communication primitive for various hybrid-parallel training frameworks. ETC can migrate model states during the hybrid-parallel reconfiguration. The design principle of ETC is to exploit the locality of model states and eliminate the redundant data movement across devices. Unlike all the other works, ETC does not require persistent storage and can directly migrate from one hybrid-parallel state to another, which enables rapid model state migration. In addition, ETC can retain the original process context and reuse the context, enabling rapid resumption.

As illustrated in Figure 3, a user submits a LLM training job to the cluster. Once execution, the job will obtain a hybrid-parallel training state. When a migration occurs, this training state is denoted as the prev state, and ETC is notified with the target state after migration, denoted as curr state. Both states are mainly determined by the hybrid-parallel configurations. In step 1, we calculate the migration cost matrix between two states. Through the cost matrix, we can obtain the pairwise data movement overhead for any process pair in prev state and curr state. And we can then derive the optimal migration sketch with minimal data movement. In step 2, we leverage the execution engine to emit the coarse-grained migration operators of the migration sketch. It ultimately executes the direct state migration across GPUs through the fine-grained inter-GPU communication instructions we designed. In step 3, we employ a virtual rank allocator to provide us more flexibility to adjust the roles of devices. It can help perform node de-fragmentation that appears in checkpoint-free migration. The de-fragmentation is also represented through migration operators and passed to the execution engine. Additionally, the execution engine proactively performs communication coalescing to further minimize the data movement. ETC is implemented based on Megatron-LM, and it can be ported to other training frameworks, benefiting from the abstracted workload, migration sketch, virtual rank, and coarse-grained migration operators.

3.2 Cost-Matrix Driven Migration↩︎

Cost Matrix: Quantity the Data Movement between Parallel Configurations.   During state migration, the states need to be redistributed across devices. Specifically, the states are resharded and then transferred from one device to another, where we denote the inter-device communication volume as data movement. Therefore, we propose the cost matrix to quantify the data movement. Cost matrix is derived from the difference of workloads among all process pairs \(\langle\)a,b\(\rangle\), where \(a\) / \(b\) represents the process rank of the hybrid-parallel configuration before / after migration.

As shown in Figure 4, we first capture the state information (including prev state and curr state) residing on each process (). According to the captured states, we express the logical workload of the processes in both states. Note that workload is the abstracted representation of tensors, and each item in workload represents a tensor in each process (e.g., [0] represents the first transformer layer, and [0.a] represents the first TP slice of the first transformer layer). Then we build the cost matrix \(C\in\mathbb{Z}^{M\times N}\), where the row indices represent the ranks in the curr state (\(M\) ranks in total), and the column indices represent the ranks in the prev state (\(N\) ranks in total) (). Note that in hybrid-parallel scenarios, the ranks can also be expressed as tuple of ranks in all parallel dimensions (e.g., P0T1 represents rank 0 in PP and rank 1 in TP). Each element \(c_{i,j}\) in the cost matrix is derived through \(c_{i,j} = len(curr_i) - len(curr_i \cap prev_j)\) , where \(curr_i\) represents the workload list in curr state and \(prev_j\) represents the workload list in prev state. Therefore, \(c_{i,j}\) can quantify the data movement required to transform previous rank \(j\) to current rank \(i\).

We provide a example of reconfiguring a model (containing 16 transformer layers) from PP4 to PP2 in Figure 4. At prev state, there are 4 devices participating, where each holds 4 transformer layers, the workload of rank 0 is [0,1,2,3]. When migration, only 2 devices are needed to continue the training. In this scenario, each device should hold 8 transformer layers and the workload of rank 0 is [0,1,2,3,4,5,6,7]. The \(c_{0,0}\) can be derived by comparing the workloads, and its value is \(4\).

Note that the cost matrix should be synchronized across all devices, so that each device can recognize its responsibility and perform migration in a decentralized manner. Besides, in hybrid-parallel training frameworks that follow the static model state partition rules (e.g., Megatron-LM partitions tensors evenly in PP and TP dimension), each device can derive its workload following these rules and calculate global cost matrix locally, which helps to avoid synchronization overhead.

Migration Sketch: Pairing Current and Previous Processes through Cost Matrix.   By exploiting the cost matrix, we pair the processes in prev state with the processes in curr state, so as to maximize the reuse of the previous process context and thus minimize the data movement. The migration sketch is a series of migration operators on process pairs (e.g., \(\langle\)a,b\(\rangle\)), which can be translated into fine-grained migration instructions and executed by the execution engine (§3.4). Note that the migration operator for process pair \(\langle\)a,b\(\rangle\) should transform process \(b\) of prev state into process \(a\) in curr state. The process pairing can be formulated as an assignment problem. \(c_{ij}\) denotes the cost when the \(i\)-th process in curr state is associated with the \(j\)-th process in prev state. We define the indicator variable \(x_{ij}\) as: \[x_{ij}= \begin{cases} 1, & \text{if the } i\text{-th rank is paired to the } j\text{-th rank} \\ 0, & \text{otherwise} \end{cases}\]

We try to minimize the total data movement among all process pairs, so that the objective function is as follows:

\[\label{equ:objective} \min \quad \sum_{i=1}^{m} \sum_{j=1}^{n} c_{ij} x_{ij}\tag{1}\]

When the process number of curr state is larger than that of prev state (i.e., \(M > N\)), there are two constraints (Equation 2 ). a) Each process in the prev state must be paired with one process in the curr state; and b) Each process in curr state can be paired with at most one process in prev state.

\[\label{equ:constraint-1} \begin{cases} \sum_{i=1}^{m} x_{ij} = 1, & \forall j = 1,2,\dots,n \quad \triangleright\text{Constraint \textit{a)}} \\ \sum_{j=1}^{n} x_{ij} \leq 1, & \forall i = 1,2,\dots,m \quad \triangleright\text{Constraint \textit{b)}} \end{cases}\tag{2}\]

When the process number of curr state is not larger than that of prev state (i.e., \(M \leq N\)), there are two constraints (Equation 3 ). a) Each process in prev state must be paired with at most one process in curr state; and b) Each process in curr state must be paired with one process in prev state.

\[\label{equ:constraint-2} \begin{cases} \sum_{i=1}^{m} x_{ij} \leq 1, & \forall j = 1,2,\dots,n \quad \triangleright\text{Constraint \textit{a)}} \\ \sum_{j=1}^{n} x_{ij} = 1, & \forall i = 1,2,\dots,m \quad \triangleright\text{Constraint \textit{b)}} \end{cases}\tag{3}\]

The assignment problem of process pairing can be efficiently solved using a non-standard Hungarian algorithm [39], as illustrated in Algorithm 5. The input to the algorithm is the cost matrix \(C\in\mathbb{Z}^{M\times N}\), and the output is the migration sketch (i.e., a series of process pairs). In the hybrid-parallel reconfiguration scenarios, \(M\) and \(N\) are usually not equal. To adopt the Hungarian algorithm, the cost matrix is padded to a square matrix with zeros (line [line:padding]). Then, we call the Hungarian algorithm in the SciPy¹ package to obtain the process pairs (line [line:solving]). Its time complexity is \(O(k^3)\) for a \(k \times k\) cost matrix. However, in our cost matrix, \(k\) is usually less than 1,024, because ETC just need to cover the PP dimension (\(\leq 16\)) and TP dimension (\(\leq 16\)), and thus the solving time is just tens of milliseconds on a standard server (refer to §3.3 for details). Note that the solving time is much less than a training iteration. Finally, we remove the processes involved in the padding part from the process pairs (line [line:remove95start]–[line:remove95end]). The output (process pairs) is the optimal migration sketch.

3.3 Extend Cost Matrix to Hybrid-Parallel↩︎

Workload Abstraction and Cost Matrix of TP.   A one-dimensional (e.g., PP) workload can be represented by a single integer (e.g., [0]), and a multi-dimensional (e.g., PP + TP) workload should describe the logical workloads of multiple dimensions in a similar manner. Unlike PP, TP does not have an inherent layer structure for partitioning and indexing. Therefore, we further extend the workload for TP. We use the least common multiple (LCM) of TP size in prev state and TP size in curr state to determine the total workload in TP dimension, so that the workload of a TP rank (e.g., rank 0 in a TP=4 group) can be represented by a single alphabet (e.g., [a] of [a,b,c,d]). Then the cost matrix of TP can be derived by differentiating the workloads in prev state and curr state through \(c_{i,j} = len(curr_i) - len(curr_i \cap prev_j)\), which is similar to PP (refer to §3.2).

Merge the Cost Matrices of PP and TP.   Under the PP+TP scenario, as a process has individual ranks in PP group and TP group (e.g., P0T0), the workload is represented by the nested PP workload and TP workload (e.g., [0.a, 1.a]). Once we derive the individual cost matrices of PP and TP, we should merge them into a monolithic cost matrix. For each cost matrix, there is a maximum cost value in the process pairs where the two workloads have no intersection at all. Two processes can intersect in the TP dimension only if they also intersect in the PP dimension; equivalently, no TP intersection implies no PP intersection. For brevity, we set the priority of PP to be higher than that of TP during the merging. We can merge cost values dimension by dimension. 1) If the cost of the high-priority dimension equals the maximum value of the cost matrix for that dimension (denoted as max), it indicates that there is no intersection at all between the process pair, so that the merged cost is equal to max. 2) If we reach the cost matrix of the lowest-priority dimension, a coefficient is returned, which equals the cost divided by the maximum cost of this dimension. This coefficient is then passed back dimension by dimension for multiplication. An example is illustrated in Figure 6.

DP Dimension.   There is no need to capture DP in the merged cost matrix. Different from PP/TP, DP replicates the model and the replicas are either kept or discarded during reconfiguration, which has no data movement. Besides, we isolate the PP/TP processes under different DP ranks. If two processes have different DP ranks, even if their workloads on PP/TP are the same, we regard the difference between the workloads to be infinite. Because there is inherent load balance among different DP replicas, where the migration of PP/TP processes is completely identical. If we create data movement across DP replicas, the load balance tends to be broken. When increasing DP size during reconfiguration, we insert a series of process-to-process migration operators to the migration sketch, where each PP/TP process in existing DP groups sends all model states to its counterpart in new DP groups. And the source groups are selected in a round-robin manner to ensure the load balance among existing DP groups.

Fault-tolerance through DP.   Since each DP group contains a complete and synchronized model state, DP groups can be regarded as replicas of each other for fault tolerance. When the dimension of DP is \(d\), ETC can accommodate at least \(d-1\) hardware failures simultaneously. When a certain device fails, the workload on the failed device can be obtained from other DP groups via the the PP/TP ranks. And we append the corresponding migration operator to the migration sketch.

Hybrid-Parallel Migration Sketch.   After obtaining the multi-dimensional cost matrix, we can derive the migration sketch, the solving procedure is similar to that in §3.2. Unlike the one-dimensional migration sketch, the cost matrix index of the hybrid-parallel migration sketch is combinations of multiple dimensions (e.g., P0T0). Meanwhile, the migration sketch within existing DP group is identical.

3.4 Execution Engine of Migration Sketches↩︎

The execution engine is responsible for executing the migration sketch generated above. Basically, it transforms the coarse-grained migration operators in the sketch into fine-grained migration instructions, and then executes the instructions for the final migration.

Migration Instructions.   We design three migration instructions including Send, Recv, and Refer. Each instruction takes the source process id sid, destination process id did, workload type wt, and workload id wid as the input. Each process pair in the migration sketch can be transformed into one or more instructions. Specifically, as for a process pair \(\langle\)curr, prev\(\rangle\), 1) if a workload is not in curr process but in prev state, it will be packed into a Send instruction; 2) if a workload is in curr process but not in prev state, it will be packed into a Recv instruction; 3) if a workload is both in curr process and in prev process, it will be packed into Refer instruction and stay in place. Note that Refer only binds the workload to the new state, requiring no data movement.

It is necessary to migrate model parameters, optimizer states and other states. Therefore, we provide the workload type wt variable in migration instructions, and maintain the mapping relationship between the logical workload to different workload types, so that the instructions are compatible to migrate these workload types. The states can be packed to the instructions when Send and then unpacked to update the corresponding state after migration.

Meta Model for Indexing TP Slices.   Given a nested workload (e.g., [0.a]), we can locate the transformer layer with the PP index (e.g., \(0\)), but we cannot locate the TP slices, because we lack the information that which dimension of the high-dimension tensor is sliced by TP. Therefore, we initialize a meta model in PyTorch’s meta device, which maintains only the tensor shapes without actually loading the actual tensors into memory. The TP size of the meta model is set to 1. By differentiating the tensor shapes of the meta model and that in prev state, we can help identifying how the TP dimension is partitioned. For example, if there is a tensor dimension satisfying meta_model_dim / prev_model_dim == prev_TP_size, this dimension is sliced, and the slice shape corresponding to the logical workload can calculated as meta_model_dim / LCM(curr_TP_size, prev_TP_size).

Execution Flow.   The complete execution flow is illustrated in Figure 7. For each coarse-grained migration operator \(\langle\)curr:j, prev:i\(\rangle\) in the migration sketch, the execution engine compares the logical workloads of process \(i\) in the prev state and process \(j\) in the curr state. Then the migration instructions including Send, Recv, and Refer are generated with above three rules 1)2)3). Then the device derives the index of the tensor slice corresponding to the workload based on the meta model and the workload type \(wt\). At this time, the migration instructions are recorded into an instruction list on each device. With all migration instructions in the migration sketch are ready, all devices emit their instruction lists simultaneously till empty.

To prevent prev states and curr states from coexisting, we use state chunking to migrate states in sequential chunks with on-demand buffer allocation, and eager release of source-side state to promptly free source-side tensors after each round.

3.5 Virtual Rank Allocator for Resource De-fragmentation↩︎

When performing hybrid-parallel reconfiguration with ETC, we find the issue of resource fragmentation. Because ETC reuses the process contexts for better locality, the devices are not restarted with a new placement but only assigned with new ranks. There are two categories of resource fragmentation.

1) Occupancy-level. As shown in Figure 8 (a), when scaling in the PP size, some devices are occupied while others stay idle. From the cluster scheduler’s perspective, incoming jobs are infeasible to be assigned to nodes with good occupancy.

2) Efficiency-level. As shown in Figure 8 (b), when the PP size and TP size of a job change simultaneously or the TP size scales out, the cross-node TP groups may arise. This leads to inefficiency in collective communications.

To avoid resource fragmentation, we propose virtual rank allocator. It isolates the virtual rank within the hybrid-parallel grid and the physical rank within the allocated devices, and manages the mapping between virtual ranks and physical ranks. Through remapping physical devices with new virtual ranks, it can directly change the roles of physical devices. The remapping can also be represented as migration operators.

Placement-Aware Remapping.   The key idea to resolve node fragmentation is to adjust the workload on the fragmented nodes through placement-aware remapping. The details are shown in Algorithm 9.

For occupancy-level de-fragmentation, the remapping is illustrated by defragmentation_for_occupancy function. Its input includes the world size of prev state \(pw\), the world size of curr state \(cw\), and the pairing relationship \(p2c\) between the previous processes and the current processes generated from the cost matrix. And the device number in a node is \(n\). First, we obtain the placement of processes, \(pNodes\), \(cNodes\), in both prev state and curr state (line [line:pp95placement95prev]–[line:pp95placement95curr]). If a process is dropped after migration, its \(id\) is set to -1. All indices in \(curr\_nodes\) corresponding to -1 represent idle devices (line [line:pp95placement95curr]), and the rest are occupied devices. We pair the idle devices in order with the occupied device in reverse order, continuing this process sequentially for all remaining devices (line [line:pp95fragement95exec95start]–[line:pp95fragement95exec95end]), so that we can populate the migration operators through pairing.

For efficiency-level de-fragmentation, the remapping is illustrated by defragmentation_for_efficiency function. In common practice, the TP group should be placed inside a node for optimal collective communication. First, we iterate through all process ids in the curr state, which is the target placement recording the optimal placement on physical devices (line [line:device95nodes]). The \(loads\) records the placement after state migration (line [line:load95nodes]), which contains the efficiency-level fragmentation. The inconsistent node pairs of \(target\) and \(loads\) are marked as fragmentation (line [line:tp95frag95nodes]). The process pairs are generated though the processes in fragmentation nodes, which are then passed to the migration operators (line [line:tp95fragement95exec95start]–[line:tp95fragement95exec95end]).

Communication Coalescing.   To eliminate the redundant data movement, the migration operators of state migration (§3.2 and §3.3) and de-fragmentation can be coalesced when the execution engine generates instructions. As shown in Figure 8, the state migration is marked by arrow , the de-fragmentation is marked by arrow . If sends the workload \(w\) from device \(i\) to device \(k\), and sends the workload \(w\) from process \(k\) to process \(j\), then the execution engine coalesces these two instructions into sending the workload \(w\) from process \(i\) to process \(j\). The coalesced instructions are marked by arrow . The Recv instructions are coalesced in the same way.

4 Evaluation↩︎

4.1 Experimental Setup↩︎

Clusters.   We use two different clusters for evaluation. a) Cluster-A is equipped with 32 nodes connected with \(4\times 100\) Gb/s RoCE network, and each node contains 4 Nvidia GPUs (40GB, without NVLink). b) Cluster-B is equipped with 4 nodes connected with \(4\times 200\) Gb/s InfiniBand network, and each node contains 8 Nvidia GPUs (80GB, with NVLink).

Models and Datasets.   We evaluate ETC on GPT3 models. We carefully tune the hyper-parameters to saturate the GPU resources: 1) the number of transformer layers is set to be divisible by the PP dimensions before and after migration; 2) the hidden size is set to be divisible by TP dimensions before and after migration. We train the models using OpenWebText dataset.

Baselines.   We compare ETC with four checkpoint-based approaches (i.e., Megatron, Megatron-Dist [37], EasyCkpt [8], and Fastest–their fastest combination) and a dynamic state management library (i.e., Tenplex [10]).

1. Megatron – the default checkpointing approach. As for reconfiguration, an extra convert script is used to reshard the checkpoint to another hybrid-parallel configuration.

2. Megatron-Dist – an improved version of 1). It enables distributed asynchronous saving and loading of the checkpoint. Moreover, it performs state resharding when loading the checkpoint, and the convert script in 1) is omitted.

3. EasyCkpt – an on-premise in-memory checkpointing. It caches the state in host memory and persists it to disk after many training steps. The convert script in 1) is required.

4. Fastest-baseline (FB) – To represent the fastest checkpoint-based baseline, we pick the least times of reshard, and save/load phases from above baselines.

5. Tenplex – a state management library designed for LLM training. It supports dynamic parallelism transformation through an in-memory file system.

4.2 Comparison with Checkpoint-based Approaches↩︎

Under various scenarios of hybrid-parallel reconfiguration, we evaluate the model state migration time of ETC and the checkpoint-based approaches. We train GPT3-32B model on Cluster-A. The hidden size is 3072, and the transformer layer number is 36. In the scenario of reconfiguring PP8TP1 to PP16TP1, since neither 8 nor 16 is divisible by 36, we set the transformer layer number to 32. Experimental results are shown in Table ¿tbl:tab:checkpoint95stats?. Since EasyCkpt does not optimize the reshard and load phases, we have merged EasyCkpt and Megatron in these columns. Compared with checkpoint-based baselines , ETC achieves superior performance across all reconfigurations.

In the save phase, Megatron-Dist and EasyCkpt achieve 6.52\(\times\) and 11.24\(\times\) on average compared to Megatron, because Megatron-Dist adopts distributed I/O to save, and EasyCkpt just offloads the checkpoint to host memory. In the reshard and load phase, Megatron-Dist shows 21.34\(\times\)–60.22\(\times\) speedup compared to Megatron and EasyCkpt. This is because it fuses the two phases to eliminate redundant I/O, and it adopts distributed I/O to load the checkpoint to host memory. With "Fastest-baseline (FB)", we combine EasyCkpt’s advantage in save and Megatron-Dist’s advantage in reshard and load. However, ETC still achieves 2.33\(\times\)–6.37\(\times\) speedup than FB. This is because: 1) with ETC, the model state is completed within the GPU memory, the checkpoint persistence and offloading is completely eliminated. 2) ETC directly sends and receives the model state with inter-GPU communication, which usually has higher bandwidth than offloading and I/O. 3) the model state resharding of ETC is inherently parallel and leverages locality to minimize the data movement.

4.3 Comparison with the State Management Library Tenplex↩︎

We compare ETC with Tenplex on Cluster-B. For fair comparison, the model configurations (i.e., GPT3-2.7B and GPT3-6.7B) and the scenarios of hybrid-parallel reconfiguration are consistent with the open-source repository of Tenplex. The reconfigurations involve the scaling in TP, PP, and DP. We compare the Tenplex and Tenplex-Central modes, where Tenplex has distributed state management between State Transformer instances on different nodes and Tenplex-Central performs all model state resharding at the central node.

As shown in Figure 10, ETC achieves 4.79\(\times\) and 10.74\(\times\) speedup compared to Tenplex and Tenplex-Central, respectively. This is attributed to two reasons. First, ETC maintains the cost matrix to represent the data movement and solves the cost matrix to generate the migration instructions with minimal data movement overhead. Second, ETC transfers the model state from GPU to GPU directly, while Tenplex still adopts a checkpoint-based design.

4.4 Benefits of Communication Coalescing↩︎

We evaluate the benefits brought by coalescing the communication in migration and de-fragmentation stages. We adopt the hybrid-parallel reconfigurations, including 1) scale in PP, 2) scale out TP, and 3) scale them simultaneously. We also set different transformer layers (\(l\)) and hidden sizes (\(h\)).

As shown in Figure 11, the coalesced communication achieves 2.43\(\times\) speedup on average compared to the two-stage communication. In the first three scenarios, we keep the same parallel configuration and increase the transformer layer number from 32 to 128. With the total parameters increase, the time required for two stages is on the rise. However, the communication time after coalescing is almost the same as the time of state migration. This is because the primary ability of communication coalescing is to change the destination to which the workload is sent. According to our observation, the communication coalescing does not alter the maximum communication volume across devices throughout the whole hybrid-parallel reconfiguration.

4.5 Performance Improvement After De-fragmentation↩︎

To demonstrate the utilization after resource de-fragmentation, we design an experiment on Cluster-A, where each node has 4 GPUs. The training job of GPT3-64B changes from PP16TP1 to PP8TP1, where the number of GPUs is reduced by half. As shown in Figure 12, before de-fragmentation, 2 GPUs in each node are occupied and the others are idle, making it impossible to release resources on a node-by-node basis. After the node de-fragmentation of ETC, the occupied GPUs are concentrated in two nodes.

To demonstrate the performance improvement after de-fragmentation, we design an experiment on Cluster-B. The training job of GPT3-64B changes from PP4TP2 to PP4TP4. From PP4TP2 to PP4TP4 reconfiguration, the number of GPUs doubles. However, before de-fragmentation, the per-iteration time of the model increases. After de-fragmentation of ETC, the per-iteration time decreases by 26.7% by leveraging the NVLink in TP communication.

4.6 Scalability↩︎

To demonstrate the scalability of ETC, we analyze time breakdown of ETC under various hybrid-parallel reconfigurations on Cluster-A. We divide the entire elastic migration into 5 parts: 1) cost matrix, which includes the generation of cost matrix and the coarse-grained migration sketch; 2) execution, which performs the fine-grained migration instructions through the execution engine; 3) build group, which represents reconstructing a global communication group with re-assigned devices; and 4) others, which includes the other modules, such as the meta model. Among them, 1)2)4) belong to the state migration of ETC. And build group relies on the training framework implementation (e.g., PyTorch and Megatron) and the hardware configuration (e.g., network and GPU topology), and its time varies significantly. For example, on Cluster-B, the time to construct a communication group with 4 nodes (32 GPUs) is around 1.5s, while on Cluster-A, the time with the 32 GPUs is around 3.0s.

As shown in Figure 13, we perform experiments where PP \(\times\) TP ranges from 8 to 64, which covers the common ranges in Megatron-LM repository. It is obvious that the time of cost matrix is almost negligible (7.54ms on average). And execution accounts for a larger proportion of time, because of communication. The time of execution is affected by three factors: 1) the size of the model parameters; 2) the network topology; and 3) the complexity of migration sketch. Factors 1) and 2) are quite intuitive, and factor 3) explains why the execution time for PP4TP2/PP6TP4 is slightly longer. In this migration, both PP size and TP size are changed. It leads to a more complex migration sketch, where some devices need to sequentially execute Recv and Send instructions.

4.7 Memory Footprint and Accuracy Analysis↩︎

We conduct a case study on a GPT3-6.7B model. The initial configuration is PP8TP1 on 2 nodes of Cluster-A. At iteration 20, elastic migration is triggered to expand the training job to 4 nodes with a configuration PP16TP1. We profile the GPU memory allocation on rank 0 throughout the migration process and record the training loss across 100 iterations.

Figure 14 shows the normalized GPU memory during migration, with all values normalized to the PP8 steady-state. The memory footprint exhibits three distinct phases. 1) Peak phase: Upon triggering migration, communication buffers are allocated for receiving resharded model parameters (1.11\(\times\)). As optimizer states are migrated via ETC in a pipelined manner, the co-residence of previous and current states causes a transient peak of 1.25\(\times\). ETC reduces the peak to one optimizer component at a time, avoiding doubling all states. 2) Cleanup phase: When communication completes, the previous PP8 states are freed (0.59\(\times\)), and the new PP16 model/optimizer is reconstructed (0.82\(\times\)). (3) PP16 steady state: The final memory is 33% lower than PP8, as doubling the pipeline depth halves the per-rank layer count.

Figure 15 presents the training loss over 100 iterations with elastic migration at iteration 20. Two observations confirm the correctness of our approach. First, the loss curve shows no spike or discontinuity at the migration boundary, the loss transitions seamlessly from PP8 to PP16, Second, after transitioning to PP16, the loss continues to decrease, following the same convergence trend. By ensuring that all optimizer states and model parameters are transferred with bit-exact fidelity, and keeping the global batch size and learning rate consistent during migration, ETC achieves elastic migration without any significant impact on loss.

5 Related Works↩︎

Elastic Training.   It allows a job to continue its training procedure under varying resources. TorchElastic [23] and Horovod Elastic [40] support elastic training and fault tolerance mainly on the number of participating worker processes (e.g., one worker per GPU), and thus their native implementation only supports elasticity on data parallel. Based on them, hybrid-parallel training frameworks, such as Megatron-LM [18] and DeepSpeed [19], can also support elastic training, but need developers to tackle the job states with either checkpoint-based approach (e.g., PcCheck [6], ByteCheckpoint [26]) or checkpoint-free migration approach (e.g., this paper). KungFu [41] and Pollux [42] support adjusting the hyper-parameters (e.g., batch sizes) of elastic training jobs to enable various scheduling policies. EasyScale [25] enables efficient and accuracy-consistent elastic training but only focuses on data parallel. Varuna [24] allows elasticity on pipeline dimension while requiring users to specify the cutting points on models. Oobleck [12] focuses on fault tolerance but provides limited elasticity by pre-defined pipeline templates. In contrast, ETC focuses on the time-consuming state migration of hybrid-parallel reconfiguration, which is frequent in the elastic training scenario.

Auto Parallelization.   With the given resources and target models, searching for the optimal hybrid-parallel configuration is difficult. Alpa [13] partitions the compute cluster into meshes, assigns model stages to meshes, and orchestrates the inter- and intra-operator operation passes. Unity [14] represents the parallel optimizations and algebraical transformations as graph substitutions in the computation graph of distributed training, and adopts a hierarchical search algorithm. And ETC can help the training job to efficiently migrate to the searched hybrid-parallel configurations.

6 Conclusion↩︎

With ETC, we demonstrate the success of checkpoint-free model state migration for elastic hybrid-parallel LLM training. Compared to the checkpoint-based migrations, ETC performs inter-GPU communications without any I/O persistence. We try the best to keep model state in place, exploit a unique cost matrix to minimize the data movement during migration, and avoid resource fragmentation through communication coalescing, enabling second-level migrations especially on clusters co-locating inference and training jobs. In the future, we hope that ETC will make the elastic training more practical in production environments for cluster utilization.

References↩︎