Training Methods2023~117 minarXiv:2311.00257

AMSP: Super-Scaling LLM Training via Advanced Model States Partitioning

score 9Qiaoling Chen, Qinghao Hu, Zhisheng Ye +4↑ 10★ 418Nov 1, 2023cs.DChttps://github.com/internlm/internevo

🎯 Pitch

ZeRO training falls off a cliff at scale: scaling LLaMA-13B from 8 to 1024 GPUs drops MFU from 47% to just 4%. AMSP solves this by letting Parameters, Gradients, and Optimizer States each independently choose their sharding strategy and device mesh, then overlapping the resulting communication with computation to achieve 1.56× the throughput of ZeRO++ at 1024 GPUs.

1. Executive Summary

This paper introduces AMSP, a system that reduces the communication overhead of ZeRO for scalable LLM training by allowing each component of the model states—Parameters, Gradients, and Optimizer States—to independently select a sharding strategy (Full-Replica, Full-Sharding, or Partial-Sharding) with its own device mesh. Training LLaMA-based models on up to 1024 NVIDIA Ampere GPUs, AMSP achieves up to 52% Model FLOPs Utilization (MFU) and delivers a 1.56× improvement in training throughput over systems like MiCS and ZeRO++, while maintaining acceptable memory overhead through a dependency rule that constrains sharding factors to avoid redundant data movement. The paper further contributes an optimization problem formulation that discovers sharding factors minimizing communication costs subject to GPU memory constraints, establishing that flexible per-component sharding—rather than the rigid sp = sg = sos coupling in prior work—is essential for scaling ZeRO efficiently, but only when paired with fine-grained overlapping of AllGather, ReduceScatter, and Broadcast operations with forward and backward computation.

2. Context and Motivation

The Core Problem: ZeRO's Communication Costs Explode at Scale

The fundamental challenge this paper tackles is deceptively simple: distributed LLM training with ZeRO breaks down at scale because its communication overhead grows faster than its ability to overlap that communication with computation. While ZeRO's memory-efficient sharding was a breakthrough that enabled training models larger than a single GPU's memory capacity, the paper demonstrates a stark failure mode: training LLaMA-7B on 8 GPUs with ZeRO-1 achieves 63% Model FLOPs Utilization (MFU), but scaling to 1024 GPUs with the same batch size causes MFU to plummet to 36%. For ZeRO-3 on LLaMA-13B, the degradation is even more catastrophic—from 47% to a mere 4% MFU.

This gap is critically important for several practical reasons the paper highlights, and several more that emerge from the broader landscape of LLM training:

The shift toward smaller, data-rich models. Recent work, including Chinchilla scaling laws, has shown that optimal performance often comes from training smaller models on larger datasets rather than pursuing ever-larger parameter counts (Hoffmann et al., 2022, cited in Section I). Models like LLaMA (7B–65B), Mistral (7B), and InternLM2 (7B–20B) represent this new paradigm. These models can fit on fewer GPUs when using ZeRO, but they still need to be trained quickly, meaning they need to be scaled across many GPUs for throughput. If ZeRO's communication overhead prevents efficient scaling beyond a few dozen GPUs for these model sizes, the entire economics of this training paradigm are undermined.
Rising GPU cluster sizes and diminishing marginal returns. Organizations are deploying GPU clusters with thousands of accelerators. If ZeRO—the most widely adopted memory-saving strategy in frameworks like DeepSpeed, PyTorch FSDP, and ColossalAI—only achieves 4% MFU on such clusters for modest-sized models, the vast majority of that expensive hardware sits idle waiting for communication. The paper's benchmark in Figure 2(c) showing communication latency scaling with GPU count while computation time shrinks (because each GPU handles fewer micro-batches) makes this tension quantitative: the compute-to-communication ratio collapses as training is scaled out.
Real-world deployment at Shanghai AI Laboratory. The paper notes that AMSP has been used for training InternLM on thousands of GPUs. This is not a hypothetical problem—it's a production bottleneck in large-scale LLM development infrastructure, and the authors document it through a separate six-month workload trace from their GPU datacenter.

Why Prior Approaches Fall Short

Before AMSP, the landscape of solutions to ZeRO's communication bottleneck could be categorized into three strategies, each with fundamental limitations:

1. Rigid Model State Sharding: ZeRO, ZeRO++, and MiCS

The original ZeRO framework introduced three levels with progressively more aggressive sharding, but each level binds the sharding factors of parameters, gradients, and optimizer states together:

ZeRO-1 shards only optimizer states (sos > 1), leaving parameters and gradients fully replicated. This means every GPU must run AllReduce across all data-parallel ranks for gradient synchronization. At 1024 GPUs, that AllReduce spans 1024 participants—with the latency scaling characteristics shown in Figure 3, where effective bandwidth degrades substantially at larger scales.
ZeRO-2 further shards gradients (sg = sos > 1), reducing per-GPU gradient storage but not reducing the communication scale—gradients still need to be reduced across all data-parallel ranks before being distributed back.
ZeRO-3 shards everything (sp = sg = sos > 1). This minimizes memory but maximizes communication: every forward and backward pass requires AllGather to reconstruct the full parameters on each GPU, followed by ReduceScatter to aggregate and redistribute gradients. On 1024 GPUs, these operations span the entire cluster.

The paper's micro-benchmarks in Figure 2 quantify why this rigid coupling is problematic. Panel (a) shows that model state memory savings from scaling ZeRO beyond ~64 GPUs are minimal—once model states are sufficiently sharded to fit in GPU memory, further sharding provides negligible memory benefits. Panel (c) shows that communication latency for AllGather, ReduceScatter, and AllReduce increases with the number of participating GPUs, even for fixed message sizes. The key insight embedded in Figure 3 is that the effective bandwidth of these collectives degrades when spanning many nodes, particularly due to the inter-node vs. intra-node bandwidth gap (600 GB/s intra-node NVLINK vs. 400 GB/s inter-node per node in the paper's testbed).

ZeRO++ attempted to mitigate this by maintaining a secondary shard of parameters within small subgroups (typically a single node, s0_p = 8, s1_p = 1) and only using the full-parameter AllGather during the forward pass. During the backward pass, parameters are gathered from this secondary shard, reducing cross-node communication. However, ZeRO++ still couples sp = sg = sos = sdp for the primary sharding, meaning the full cluster-wide communication remains for gradients and optimizer states. The paper reports ZeRO++'s MFU at a mere 4%, 6%, and 5% for LLaMA-7B, 13B, and 30B on 1024 GPUs—sometimes worse than vanilla ZeRO-3. This happens because the secondary shard maintenance adds overhead without fixing the fundamental problem that some components are being sharded across an unnecessarily large communication group.

MiCS introduced a key conceptual advance: shard model states within a subgroup and replicate across subgroups (sp = sg = sos < sdp). This reduces the communication scale for AllGather and ReduceScatter to the subgroup size, which can be as small as a single node (8 GPUs). However, MiCS still mandates sp = sg = sos—all three components must use the same subgroup. This creates an inflexible tradeoff:

If the subgroup is small (e.g., 8 GPUs to avoid cross-node communication), all three components get replicated across subgroups, consuming more memory. For LLaMA-30B on 1024 GPUs, MiCS with s0_p = s0_g = s0_os = 8 consumes approximately double the memory of ZeRO-3 (Figure 13), yet still underperforms AMSP (29% vs. 42% MFU).
The paper's most damning finding about MiCS appears when training LLaMA-7B on 1024 GPUs: MiCS achieves only 35% MFU, which is lower than the 36% achieved by vanilla ZeRO-1. This inversion—a system designed to improve performance actually underperforming its simpler predecessor—is the paper's clearest evidence that rigid sharding coupling is the bottleneck. MiCS's AllGather and ReduceScatter happen within a node (reducing their cost), but its AllReduce for gradient synchronization and parameter broadcast operations still involve larger groups, and the inflexible configuration prevents independently optimizing which communication domains each operation uses.

2. Communication-Computation Overlap: The Implementation Quality Gap

A separate but equally important failure mode is that existing systems already attempt to overlap communication with computation but do so poorly. The paper's trace analysis in Figure 14 reveals a concrete implementation problem in DeepSpeed's ZeRO-3, MiCS, and ZeRO++:

"Even when the communication-computation overlap setting is enabled, DeepSpeed-ZeRO3/MiCS/ZeRO++ fails to effectively overlap ReduceScatter with computation. Additionally, in DeepSpeed-MiCS, the ReduceScatter operation also blocks the concurrent execution of AllReduce."

The traces show "computation resource bubbles" during the backward pass—periods where GPUs are waiting for communication to complete rather than executing computation. This isn't a theoretical impossibility of overlap; it's a systems engineering failure in how NCCL communication primitives are scheduled and how they interact with PyTorch's autograd engine. The paper notes that simply reimplementing MiCS's strategy under AMSP's execution engine—without changing the sharding configuration—achieves a 2× speedup over DeepSpeed-MiCS during LLaMA-13B training on 1024 GPUs. This single result demonstrates that the communication overhead problem has two independent components: (1) what communication needs to happen (determined by the sharding strategy), and (2) how that communication is scheduled relative to computation (determined by the execution engine's hook-based coordination).

3. Quantization and Compression-Based Approaches

ZeRO++ and Espresso use quantization to compress parameters and gradients, reducing the volume of communication. The paper explicitly disables these configurations to ensure consistent model quality, but even with compression, ZeRO++'s performance remains poor in their benchmarks. This suggests that at extreme scales, compression addresses the wrong bottleneck: the problem is not just the size of messages but the latency of coordinating communication across too many participants. No amount of compression can eliminate the (p - 1)α term in the ring-AllReduce cost model, where α is per-transmission latency and p is the number of participants, which grows linearly with scale.

The Three Factors Behind Communication Collapse

The paper decomposes the communication overhead into three distinct physical causes (Section III-A), each of which compounds the others:

Factor 1: Inter-node vs. intra-node bandwidth asymmetry. The paper's DGX-A100 nodes provide 600 GB/s intra-node bidirectional bandwidth per GPU (via NVLINK), but only 400 GB/s inter-node bidirectional bandwidth total per node. For an 8-GPU node, this means each GPU has 600 GB/s to its neighbors within the same box but effectively only 50 GB/s to any GPU in another node. The 2× intra-to-inter bandwidth ratio in the paper's testbed means that any communication operation spanning nodes pays a steep bandwidth penalty. This isn't unique to the paper's infrastructure—it's a fundamental architectural characteristic of GPU clusters.

Factor 2: Latency scaling with communication group size. Figure 2(c) benchmarks the actual latency of NCCL collectives at fixed message sizes (256 MB in this panel), showing that AllGather, ReduceScatter, and AllReduce latency all increase as the number of participating GPUs grows from 8 to 512. Figure 3 provides a more detailed characterization: effective bandwidth for AllGather on a 64 MB message drops from roughly 1200 Gb/s on 8 GPUs (single node) to roughly 300 Gb/s on 512 GPUs (64 nodes). This is a 4× bandwidth degradation purely from scaling the communication group, independent of the message size or per-GPU bandwidth.

Factor 3: Shrinking computation time per GPU. When the global batch size is fixed (a common constraint for convergence), adding more GPUs means each GPU processes fewer micro-batches. The paper's benchmark in Figure 2(b) shows forward-backward computation time dropping linearly as GPU count increases: from ~60 seconds on 8 GPUs to ~1 second on 512 GPUs for LLaMA-7B with a 4M token global batch. The compute-to-communication ratio collapses because communication latency grows with GPU count while computation time shrinks, creating a communication-dominated regime where GPUs spend more time waiting for data than processing it.

The paper notes that this third factor is not unique to ZeRO—it's a fundamental property of strong scaling with fixed batch size. What makes ZeRO particularly vulnerable is that it amplifies the first two factors through its collective communication requirements.

The Central Insight: Independent Per-Component Sharding

The paper's core conceptual contribution is the observation that parameters, gradients, and optimizer states have different communication patterns and different memory footprints, so they should not be forced to share the same sharding configuration. This is both obvious in retrospect and entirely absent from prior work.

To understand why this matters, consider the communication operations each component triggers:

Parameters (sp > 1): Require AllGather at every module in every micro-batch during both forward and backward passes, plus ReduceScatter for gradients at every micro-batch. This is M × L × K pairs of operations per step, making it the most communication-intensive component. Keeping sp small (ideally 1, Full-Replica) eliminates this cost entirely but at high memory expense.
Gradients (sg > 1): After being generated in the backward pass, gradients need to be aggregated. When sg = sp, each GPU only retains gradients for its parameter shard, and the ReduceScatter from parameter sharding handles gradient distribution. When sg > sp, extra AllReduce operations are needed but only within a subgroup, not across the full cluster. The communication frequency is lower than for parameters (once per layer, not once per module-within-layer).
Optimizer states (sos > 1): After gradient aggregation and parameter updates, updated parameters need to be distributed. This is a Broadcast operation (or series of AllGather/group calls) that happens only at the end of each step—once per step total, not per micro-batch. The total communication volume is the same as parameter sharding's AllGather (2Φ per step), but it's incurred in a single burst rather than interleaved with computation.

The paper's key example in Section III-C illustrates the power of this decoupling. For LLaMA-7B training at scale, they set sp = sg = 1 (Full-Replica for parameters and gradients) but sos = 8 (Full-Sharding for optimizer states within each node). This means:

No AllGather or ReduceScatter during forward/backward passes (parameters and gradients are fully replicated, eliminating the most expensive and frequent communication).
AllReduce for gradient synchronization still happens across all GPUs (since sg = 1), but this is less frequent than parameter-sharding communication.
Broadcast for updated parameters happens only once per step, within nodes (8 GPUs), and can be overlapped with the forward computation of the next step.

This configuration achieves an MFU of 51% on 1024 GPUs for LLaMA-7B, compared to 36% for ZeRO-1 and 35% for MiCS. The critical difference from MiCS: MiCS's sp = sg = sos = 8 forces unnecessary AllGather/ReduceScatter communication for parameters even though LLaMA-7B's parameter memory (14 GB in FP16) could fit on a single 80 GB GPU after accounting for other memory consumers—the parameter sharding was saving memory that didn't need saving while paying communication costs that were crippling performance.

The Unexplored Search Space

Prior to AMSP, the systems community had explored only a few points in what is actually a large, structured search space defined by the six sharding factors (s0_p, s1_p, s0_g, s1_g, s0_os, s1_os). The paper introduces a dependency rule that constrains this space meaningfully:

R ≥ s0_dp ≥ s0_os ≥ s0_g ≥ s0_p N ≥ s1_dp ≥ s1_os ≥ s1_g ≥ s1_p

These constraints prevent situations where a GPU would need to manage optimizer states or gradients for parameters it doesn't store locally (as illustrated in Figure 4). Without this rule, configurations like sp = sg = 2, sos = 1 would require extra communication to fetch the necessary optimizer states, incurring "significant and avoidable expenses." Within these constraints, however, the space remains rich—especially given the two-dimensional device mesh specification (s0_i for GPUs within a node, s1_i for nodes), which controls whether communication stays within a node or crosses node boundaries.

The paper further shows that sg doesn't need to be continuously variable—it can be restricted to sg ∈ {sp, sos} without losing meaningful configurations, since memory savings from gradient sharding beyond what's already provided by either parameter or optimizer state sharding are marginal. This collapses the search space while retaining the essential tradeoff dimensions.

Positioning Relative to Automated Parallelism Frameworks

The paper distinguishes AMSP from automated parallelism systems like Alpa, OptCNN, FlexFlow, and TensorOpt, which search over operator-level parallelization strategies (how to split individual matrix multiplications, attention operations, etc. across devices). These systems focus on the computational graph—which operations execute where—but overlook the model state sharding strategy, which is an orthogonal dimension. AMSP's search operates at the level of how the persistent state tensors (parameters, gradients, optimizer states) are distributed, which is complementary to operator-level parallelism and could, in principle, be combined with it.

The paper acknowledges this bounded scope explicitly, focusing solely on data-parallel training with flexible model state sharding and leaving open the integration with tensor and pipeline parallelism (Sections III-C, IV).

How AMSP Positions Itself

AMSP positions itself as solving the problem that ZeRO's rigid coupling of sharding factors is the bottleneck, not any single communication operation or overlap deficiency. The paper's evidence for this hierarchical diagnosis comes from:

The MiCS-vs-ZeRO-1 inversion on LLaMA-7B at 1024 GPUs (35% vs. 36% MFU): MiCS is supposed to be better by reducing communication scale, but its forced sp > 1 triggers AllGather/ReduceScatter that ZeRO-1's sp = 1 avoids entirely. The flexible sharding in AMSP recovers ZeRO-1's communication profile for parameters while adding a node-local Broadcast for optimizer states, achieving 51% MFU.
The 2× gap between DeepSpeed-MiCS and AMSP-MiCS under the same sharding configuration: This isolates the overlap implementation quality gap, showing that scheduling matters independently of sharding strategy.
The optimization problem formulation (Equation 1): By minimizing Tcomm(s0_p, s1_p, ..., s0_os, s1_os) subject to Dtotal ≤ GPU_Memory_Capacity, the paper provides a principled way to navigate the tradeoff space for any model size, GPU memory capacity, and cluster topology. This makes the approach general rather than hand-tuned for specific models.

The paper's unifying framework is that both the sharding strategy choice AND its execution scheduling are first-class optimizations, and fixing only one while neglecting the other (as prior work did) yields dramatically suboptimal results. The sp = sg = sos coupling in prior systems effectively removed the sharding strategy dimension from the optimization, forcing all configurations to pay the maximum communication cost for at least one component. AMSP restores that dimension and provides both the theoretical analysis (communication cost model, memory model, dependency constraints) and the practical execution engine (hook-based overlap scheduling) to exploit it.

3. Technical Approach

3.1 Reader Orientation

AMSP is a training system — specifically, an execution planner and runtime engine that sits between the user's model definition (in PyTorch) and the GPU cluster hardware, automatically deciding how to distribute a language model's persistent state tensors across devices to minimize communication overhead. The system solves the problem that ZeRO's rigid coupling of sharding strategies — forcing parameters, gradients, and optimizer states to all be sharded identically — is the primary bottleneck preventing efficient LLM training at scale, and the solution takes the shape of a constrained optimization over a six-dimensional sharding factor space whose objective is communication time minimization subject to per-GPU memory capacity, paired with a hook-based execution engine that schedules the resulting communication operations to maximally overlap with computation.

3.2 Big-Picture Architecture (Diagram in Words)

The AMSP system has two major components that operate sequentially:

1. The Planner (offline, before training starts). This component takes as input the model architecture (layer count, hidden dimension, sequence length), training hyperparameters (micro-batch size, number of micro-batches, global batch size), and cluster specification (number of nodes N, GPUs per node R, GPU memory capacity). It produces as output a set of six sharding factors — s0_p, s1_p, s0_g, s1_g, s0_os, s1_os — that specify, for each of the three model state components (Parameters, Gradients, Optimizer States), how many GPUs within a node (s0) and how many nodes (s1) the component should be sharded across. The Planner contains three sub-modules: a Pre-Filter that eliminates sharding configurations violating dependency rules or causing unnecessary cross-node communication, a Communication-Profiler that has pre-measured effective bandwidths for NCCL collectives at various message sizes and device mesh configurations, and a Solver that formulates and solves an integer programming problem minimizing total communication time subject to GPU memory constraints.

2. The Executor (online, during training). This component takes the sharding factors from the Planner and executes the actual distributed training loop using PyTorch with custom forward/backward hooks. It contains a Communication Overlap Scheduler that uses PyTorch's hook API (Table III: register_forward_pre_hook, register_forward_hook, register_full_backward_pre_hook, register_full_backward_hook, and tensor register_hook) to interleave NCCL collectives (AllGather, ReduceScatter, AllReduce, Broadcast) with layer computations at the granularity of individual modules within Transformer layers. It also contains a Communication Placement Optimizer that, when sharding factors span multiple nodes (s1_p, s1_g, s1_os > 1), groups those nodes under the same leaf switches in the spine-leaf network topology to avoid cross-spine communication.

Information flows as follows: user provides model spec + cluster spec → Pre-Filter narrows the search space using dependency rules → Communication-Profiler supplies bandwidth estimates → Solver enumerates remaining configurations, evaluates each against the memory and communication cost models, and selects the one minimizing Tcomm → Executor launches the training job with those sharding factors, using hooks to schedule AllGather for parameter prefetching, ReduceScatter for gradient synchronization, AllReduce for gradient aggregation within optimizer state sharding groups, and Broadcast for updated parameter distribution, all overlapped with layer computations.

3.3 Roadmap for the Deep Dive

First, the performance model for collective communication — how AMSP predicts communication time without relying on an analytical cost model — because all downstream decisions about sharding factors depend on accurate communication cost estimation.
Second, the flexible sharding strategies and dependency rule — what Full-Replica, Full-Sharding, and Partial-Sharding mean concretely, and why the constraints s0_dp ≥ s0_os ≥ s0_g ≥ s0_p are necessary — because this defines the search space the Planner operates over.
Third, the communication time analysis for each model state component — the equations for Tp, Tg, T0_os, and T1_os — because these are the objective function the Solver minimizes.
Fourth, the GPU memory consumption analysis — because memory is the constraint that prevents the trivial solution of replicating everything (setting all sharding factors to 1).
Fifth, the optimization problem formulation (Equations 1–7) — because this is the mathematical heart of the Planner and the mechanism that generates the sharding factors used at runtime.
Sixth, the computation-communication overlap strategies for parameters, optimizer states, and gradients separately — because the Executor's performance depends not just on what communication happens but when relative to computation, and the strategies differ per component.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems paper whose core idea is that the communication overhead of ZeRO can be substantially reduced by (1) allowing parameters, gradients, and optimizer states to independently choose their sharding configurations from a structured space defined by dependency constraints, and (2) executing the resulting communication operations with fine-grained scheduling that overlaps them with layer computations at module granularity.

Performance Model of Collective Communication

The paper faces a fundamental engineering challenge at the outset: to decide which sharding configuration is optimal, the Planner must be able to predict how long collective communication operations will take for arbitrary combinations of message size, operation type, and participating GPU device mesh. The standard approach in distributed systems literature is the α-β cost model, which for a ring-based AllReduce on p GPUs with input size v and per-link bandwidth w gives:

$t_{ar} = 2(p - 1)\left(\alpha + \frac{v}{w \times p}\right)$

where α is the per-transmission latency, p is the number of participating GPUs, and the factor of 2 accounts for the ReduceScatter phase followed by the AllGather phase.

What it computes: the wall-clock time for a ring-based AllReduce operation as the sum of per-hop latencies and per-byte transmission times across p - 1 hops in each of two phases, with the message split into p chunks so each GPU only sends/receives v/p bytes per hop.

Why this form: the α-β model captures the two fundamental scaling behaviors of collectives — the α term grows linearly with the number of participants (latency-bound scaling) while the v/(w × p) term shrinks with more participants because the per-GPU data volume decreases (bandwidth-bound scaling). For large messages, the bandwidth term dominates; for small messages, latency dominates.

However, the paper explicitly identifies three reasons the α-β model is insufficient for accurate prediction in their setting:

NCCL uses multiple communication algorithms beyond ring. The paper lists Tree, Collnet, CollnetDirect, and CollnetChain as alternatives NCCL may select based on message size and topology. Each has a different cost structure that a single analytical formula cannot capture.
In-Network Aggregation (e.g., SHARP). Production GPU clusters often offload AllReduce onto network switches, which fundamentally alters the communication pattern — instead of p - 1 hops, data may be reduced in the switch fabric itself. The paper's testbed uses Mellanox HDR InfiniBand without SHARP, but the profiling approach generalizes to clusters that do have it.
Real-world bandwidth is not a constant. As Figure 3 demonstrates, effective bandwidth for the same collective operation degrades substantially as the number of participating GPUs grows, even for fixed message sizes. An AllGather of 64 MB achieves roughly 1200 Gb/s effective bandwidth on 8 GPUs (single node) but only roughly 300 Gb/s on 512 GPUs (64 nodes) — a 4× degradation. This non-ideality cannot be captured by a simple v/w term.

The profiling-based alternative. AMSP instead uses pre-collected empirical bandwidth measurements:

$t(o, v, p_0 \times p_1) = \frac{v}{w(o, v, p_0 \times p_1)}$

where o ∈ {AllReduce, AllGather, ReduceScatter, Broadcast} is the collective operation, v is the message size in bytes, p_0 is the number of GPUs within a node participating in the collective, p_1 is the number of nodes participating, and w(o, v, p_0 × p_1) is the effective bandwidth obtained from offline profiling on the target cluster.

What it computes: given a previously measured effective bandwidth for a specific operation, message size, and device mesh, this formula estimates communication time by dividing the total data volume by that bandwidth. For message sizes not explicitly profiled, AMSP uses linear interpolation between neighboring measured points.

Why this form: it replaces an idealized analytical model with empirical reality. The profiling captures all real-world non-idealities — NCCL algorithm selection, switch topology effects, bandwidth degradation at scale, and contention patterns — in a single scalar w that is measured once per cluster and reused for all planning decisions. The tradeoff is that profiling must be done in advance, and the predictions are only valid for the profiled cluster, but for production training environments this is an acceptable cost.

The paper profiles four operations across message sizes from 1 MB to 1 GB and device meshes from 8×1 GPUs (single node) to 8×64 GPUs (64 nodes), with results visualized in Figure 3. The profiled data reveals that broadcast effective bandwidth degrades the least with scale (staying above 800 Gb/s even at 512 GPUs), while AllGather and ReduceScatter degrade more severely because they involve more data movement per GPU. This empirical characterization is the foundation on which the rest of the communication cost analysis rests.

Flexible Model States Sharding with Dependency Rule

The paper defines three sharding strategies for each model state component, parameterized by the sharding factor s_i = s0_i × s1_i for component i ∈ {p, g, os}:

Full-Replica (s_i = 1). The tensor is fully replicated on every GPU. This means s0_i = 1, s1_i = 1. Each GPU holds a complete copy. This corresponds to vanilla data parallelism if applied to all three components — no memory savings but zero communication for gathering that specific component. For example, sp = 1 means every GPU already has the full parameters and no AllGather is needed before forward/backward computation.

Full-Sharding (s_i = s_dp). The tensor is sharded across all data-parallel GPUs. This means s0_i × s1_i = s0_dp × s1_dp, so each GPU holds only 1/s_dp of the tensor. This is ZeRO-3's strategy applied to that component. For a 7B-parameter model on 1024 GPUs with s_dp = 1024, Full-Sharding means each GPU holds only 7B/1024 ≈ 6.8M parameters — massive memory savings but requiring AllGather across 1024 GPUs to reconstruct the full tensor.

Partial-Sharding (1 < s_i < s_dp). The tensor is sharded across a subgroup of GPUs and replicated across subgroups. There are s_dp / s_i identical replicas in the cluster. For example, sp = 8 on a 1024-GPU cluster means parameters are sharded within groups of 8 GPUs (typically one node), and that 8-GPU shard pattern is replicated 128 times across the cluster. This reduces communication scale (AllGather only involves 8 GPUs, not 1024) at the cost of higher memory (each replica group stores a full copy of parameters across its 8 GPUs, so total cluster-wide parameter memory is 128 × 2Φ instead of 1 × 2Φ).

The critical innovation is that each component independently chooses its strategy. The paper introduces sharding factors with explicit two-dimensional device mesh specification:

s0_i: number of GPUs within a single node that the component is sharded across
s1_i: number of nodes that the component is sharded across
s_i = s0_i × s1_i: total number of GPUs in the sharding group

This two-dimensional specification matters because communication within a node uses high-bandwidth NVLINK (600 GB/s bidirectional per GPU in the paper's testbed) while communication across nodes uses InfiniBand (400 GB/s bidirectional per node total). Keeping s1_i = 1 means all communication for that component stays within a single node; setting s1_i > 1 introduces cross-node communication, which is substantially more expensive.

The dependency rule. The paper identifies a non-obvious constraint: if GPU A holds parameters that GPU B's optimizer states reference, but GPU B doesn't hold the corresponding gradients, then extra communication will be needed to move gradients or optimizer states around. To prevent this, AMSP enforces:

$R \geq s^0_{dp} \geq s^0_{os} \geq s^0_g \geq s^0_p$ $N \geq s^1_{dp} \geq s^1_{os} \geq s^1_g \geq s^1_p$

where R is the number of GPUs per node and N is the number of nodes.

What these constraints mean operationally: optimizer states must be sharded at least as widely as gradients, which must be sharded at least as widely as parameters. If parameters are Full-Replica (sp = 1), gradients can be Full-Replica too (sg = 1), or they can be sharded more aggressively (sg > 1). But gradients can never be less sharded than parameters — you can't have sg < sp. Similarly, optimizer states can never be less sharded than gradients.

Why this form: the dependency ensures that any GPU that holds a piece of an optimizer state tensor also holds (or has direct access to) the corresponding parameters and gradients needed to compute the update. As illustrated in Figure 4, violating this — e.g., setting sp = sg = 2 but sos = 1 — would cause GPU-0 and GPU-1 to each hold half the parameters but both would hold full optimizer states, meaning optimizer states for parameters not locally present would need to be fetched, incurring unnecessary communication. The dependency rule prunes these pathological configurations from the search space.

The paper also restricts sg to only two possible values:

$s_g \in \{s_p, s_{os}\}$

Why this restriction: gradients are an intermediate quantity — they're produced during the backward pass and consumed during the optimizer update. If sg = sp, gradient storage is aligned with parameter sharding (each GPU keeps gradients only for its local parameters), and gradient distribution happens as part of the ReduceScatter already required by parameter sharding. If sg = sos > sp, gradients are sharded more finely to save memory, but this requires extra AllReduce operations within the optimizer state sharding subgroup. The paper argues that allowing arbitrary sg values between sp and sos would add complexity without meaningful memory savings, since the memory footprint of gradients (2Φ bytes in FP16) is small compared to optimizer states (12Φ bytes for Adam). Collapsing sg to these two discrete choices reduces the search space without sacrificing meaningful configurations.

Communication Time Analysis: Parameters Sharding

When sp = s0_p × s1_p > 1, parameters are distributed across sp GPUs. Before any forward or backward computation on a module, each GPU must have the full set of parameters for that module. The communication cost is:

$T_p = M L \sum_{i=0}^{K} \left(2 \cdot t(AG, 2\Phi_i, s^0_p \times s^1_p) + t(RS, 2\Phi_i, s^0_p \times s^1_p)\right)$

where M is the number of micro-batches per step, L is the number of Transformer layers, K is the number of modules per layer (e.g., attention linear projections, MLP linear layers, layer norms), Φ_i is the number of parameters in module i (stored in FP16, so 2Φ_i bytes), t(AG, v, mesh) is the profiled AllGather time for message size v on the specified device mesh, and t(RS, v, mesh) is the corresponding ReduceScatter time.

What it computes: for every micro-batch (M times per step), for every layer (L times), for every module within that layer (K times), the system must (a) AllGather the full parameters for that module before its forward computation, (b) AllGather the full parameters again before its backward computation, and (c) ReduceScatter the gradients after the backward computation. The factor of 2 before AllGather reflects the need to gather parameters for both the forward pass and the backward pass (since parameters are discarded after each pass to save memory).

Why this form: the summation over modules reflects the fact that AMSP performs AllGather and ReduceScatter at module granularity — it doesn't gather all layer parameters at once, but rather module-by-module, enabling overlap with computation (discussed in Section V-C). The cost scales linearly with M, L, and K, which is why parameter sharding is the most communication-intensive component — training a LLaMA-7B model with 32 layers and ~7 modules per layer with 128 micro-batches means roughly 128 × 32 × 7 ≈ 28,672 AllGather/ReduceScatter pairs per step, each operating on module-sized parameter tensors (typically tens to hundreds of megabytes).

When sp = 1 (Full-Replica), Tp = 0 — there is no AllGather or ReduceScatter for parameters because every GPU already has a complete copy. This is the fundamental tradeoff: parameter sharding saves memory (2Φ bytes per GPU become 2Φ/sp bytes) but introduces this large communication cost proportional to the total number of module invocations.

Communication Time Analysis: Optimizer States Sharding

When sos = s0_os × s1_os > 1, optimizer states are distributed across sos GPUs. This triggers two distinct communication patterns:

Gradient aggregation (AllReduce). After the backward pass of the last micro-batch, each GPU needs the gradients for the parameters corresponding to its optimizer state shard. Since sos ≥ sp by the dependency rule, parameters are replicated sdp/sp times across the cluster, and each replica group independently maintains a subset of the optimizer states. Gradients must be aggregated across the sdp/sp GPUs that share the same parameter replica:

$T^0_{os} = \frac{2\Phi}{U s_p} \cdot t\left(AR, U, \frac{s^0_{dp}}{s^0_p} \times \frac{s^1_{dp}}{s^1_p}\right)$

where U is the bucket size for gradient AllReduce (a buffering parameter that groups multiple parameter gradients into a single communication), 2Φ is the total gradient bytes (FP16), 2Φ/(U·sp) is the number of AllReduce calls needed (total gradient data divided by bucket size, then divided by sp because each GPU only handles 1/sp of the parameters), and the device mesh for each AllReduce is (s0_dp/s0_p) × (s1_dp/s1_p) — the number of GPUs that share the same parameter replica.

What it computes: the total time to aggregate gradients for the optimizer update. For LLaMA-7B with sp = 1, sos = 8, sdp = 1024, each AllReduce involves 1024/1 = 1024 GPUs (the full data-parallel group). With a bucket size U of typically 64–128 MB, the 14 GB of gradients require roughly 14GB / (128MB · 1) ≈ 112 AllReduce calls.

Updated parameter distribution (Broadcast). After each GPU updates its portion of parameters using its local optimizer states, those updated values need to propagate to all other GPUs that hold replicas of the same parameters. Since AMSP uses the inter-tensor approach for optimizer state sharding (each tensor is distributed whole, rather than split within a tensor), AllGather is not applicable because different GPUs may hold different numbers of parameters. Instead, AMSP uses a series of Broadcast operations:

$T^1_{os} = \frac{s_{os}}{s_p} \cdot t\left(BC, \frac{2\Phi}{s_{os}}, \frac{s^0_{os}}{s^0_p} \times \frac{s^1_{os}}{s^1_p}\right)$

where sos/sp is the number of Broadcast operations needed (one per GPU in the optimizer state sharding subgroup, since each GPU that updated parameters must broadcast its updates to the other sos/sp - 1 GPUs that share the same parameter replica but have different optimizer state shards), 2Φ/sos is the average size of each broadcast (total parameters in FP16 divided by the number of optimizer state shards), and the device mesh is the optimizer state sharding subgroup size.

What it computes: the total time to disseminate updated parameters to all replicas. For LLaMA-7B with sp = 1, sos = 8, there are 8/1 = 8 Broadcast operations, each moving 14GB/8 = 1.75GB within a group of 8 GPUs.

Why inter-tensor rather than intra-tensor: the paper notes that intra-tensor sharding (splitting individual parameter tensors across GPUs) works for FP16 but is "not feasible" for FP8 training because FP8 requires per-tensor scaling factors that must be distributed alongside the tensor shards. The inter-tensor approach puts each parameter tensor wholly on one GPU using a greedy balancing algorithm, making it compatible with FP8 training. The cost is that Broadcast must be used instead of the more efficient AllGather, but the paper accepts this tradeoff for generality.

Why the select & drop mechanism: because sos > sp, each GPU receives more gradients from the AllReduce than it needs — gradients for parameters whose optimizer states are managed by other GPUs. Rather than doing a separate, more targeted communication pattern, AMSP simply performs the full AllReduce and then has each GPU discard (select & drop) the gradients it doesn't need. This is a deliberate simplicity-vs-efficiency tradeoff: the communication cost of sending excess gradients is small compared to the engineering complexity of orchestrating partial reductions across overlapping subgroups.

Communication Time Analysis: Gradients Sharding

When sg = s0_g × s1_g > sp, gradients are sharded more widely than parameters. In every micro-batch except the last (since the last micro-batch's gradient aggregation is covered by T0_os), an AllReduce is needed to aggregate and distribute gradients within the gradient sharding subgroup:

$T_g = (M - 1) \frac{2\Phi}{U s_p} \cdot t\left(AR, U, \frac{s^0_g}{s^0_p} \times \frac{s^1_g}{s^1_p}\right)$

where M - 1 reflects that the last micro-batch's gradient aggregation is handled by the optimizer state sharding AllReduce (T0_os), and the device mesh is (s0_g/s0_p) × (s1_g/s1_p) — the per-micro-batch gradient AllReduce only involves the gradient sharding subgroup, not the full data-parallel group.

What it computes: the additional communication cost incurred when gradients are sharded more aggressively than parameters. This is only non-zero when sg > sp; when sg = sp, gradient aggregation is already handled by the ReduceScatter in parameter sharding, so Tg = 0.

Why this exists as a separate term: the paper's flexibility in allowing sg ∈ {sp, sos} means that in some configurations, gradients need intermediate aggregation during the step (before the final optimizer update) to keep per-GPU gradient memory low. This is the cost of that memory savings — extra AllReduce operations in every micro-batch except the last.

Total communication time. The single-step communication time is the sum of all four terms:

$T_{comm}(s^0_p, s^1_p, s^0_g, s^1_g, s^0_{os}, s^1_{os}) = T_p + T_g + T^0_{os} + T^1_{os}$

This function is the objective that the Solver minimizes. Its arguments are the six sharding factors; its value is the predicted total communication time per training step (in seconds).

GPU Memory Consumption Analysis

The memory consumed by model states under a given sharding configuration is:

$D_{modelstate}(s^0_p, s^1_p, s^0_g, s^1_g, s^0_{os}, s^1_{os}) = \frac{2\Phi}{s^0_p s^1_p} + \frac{2\Phi}{s^0_g s^1_g} + \frac{12\Phi}{s^0_{os} s^1_{os}}$

where 2Φ is the memory for parameters in FP16 (2 bytes per parameter), 2Φ is for gradients in FP16, and 12Φ is for optimizer states with Adam (which maintains first and second moment estimates in FP32, each taking 4 bytes, plus a master copy of parameters in FP32 at 4 bytes, totaling 4 + 4 + 4 = 12 bytes per parameter — though the paper groups these under the single "optimizer states" category).

What it computes: the per-GPU memory footprint of the three model state components after sharding. Each component's total cluster-wide memory (2Φ, 2Φ, or 12Φ) is divided by its sharding factor to get the per-GPU allocation.

Why this form: it captures the fundamental memory-communication tradeoff. Increasing any sharding factor s_i reduces the memory term (constant)/s_i but increases the corresponding communication term T_p, T_g, T0_os, or T1_os. The Planner's job is to find the sharding factors that achieve the best communication time while keeping total memory below the GPU capacity.

The total GPU memory also includes activations and temporary buffers:

$D_{total} = D_{modelstate} + D_{activation} + D_{tmp}$

where D_activation is the memory consumed by intermediate activations during training (which depends on micro-batch size, sequence length, hidden dimension, and whether activation recomputation is enabled), and D_tmp is temporary memory for communication buffers and transient variables. The paper states that existing methodologies for predicting activation memory (from Megatron-LM and ZeRO papers) are "seamlessly integrated" into AMSP but does not provide explicit formulas, treating this as a known quantity that can be estimated from model architecture and training hyperparameters.

The Optimization Problem Formulation

The Planner formulates the sharding factor selection as a constrained integer programming problem:

Objective (Equation 1):

$\text{Minimize } T_{comm}(s^0_p, s^1_p, s^0_g, s^1_g, s^0_{os}, s^1_{os})$

What it computes: find the six sharding factors that minimize the total predicted communication time per step.

Why minimization of communication time: in the communication-dominated regime the paper targets (large-scale training with small per-GPU micro-batch counts), computation time is essentially fixed by the model architecture and micro-batch size, and the variable component is communication. Minimizing communication time directly maximizes throughput.

Memory constraint (Equation 2):

$D_{total} \leq GPU\_Memory\_Capacity$

What it does: ensures the selected configuration won't cause an Out-Of-Memory (OOM) error. D_total includes model states, activations, and temporary buffers.

Why a hard constraint rather than a penalty: memory is a binary constraint — either the model fits or it doesn't. There's no gradual degradation from being slightly over capacity. The constraint eliminates all infeasible configurations from consideration.

Dependency constraints (Equations 3–4):

$1 \leq s^0_p \leq s^0_g \leq s^0_{os} \leq s^0_{dp} \leq R$ $1 \leq s^1_p \leq s^1_g \leq s^1_{os} \leq s^1_{dp} \leq N$

What they enforce: the dependency rule described in Section IV-B: parameters are least sharded, gradients are at least as sharded as parameters, and optimizer states are at least as sharded as gradients, all bounded by the data-parallel group size and physical hardware limits.

Divisibility constraints (Equations 5–6):

$s^0_i \times k = s^0_{dp}, \quad k \in \mathbb{Z}, \quad i \in \{p, g, os\}$ $s^1_j \times k = s^1_{dp}, \quad k \in \mathbb{Z}, \quad j \in \{p, g, os\}$

What they enforce: the sharding factors must evenly divide the data-parallel group size. This ensures that all GPUs participate in training (no idle GPUs) and that the replication pattern is uniform — every sharding subgroup is the same size, and every GPU belongs to exactly one subgroup per component.

Why divisibility is necessary: if s0_dp = 8 (8 GPUs per node) but s0_p = 3, there's no way to partition 8 GPUs into uniform groups of 3 — some GPUs would be left out or groups would be uneven. Divisibility ensures a clean partition where each subgroup handles exactly s0_dp / s0_i replicas of the sharded tensor.

Cross-node minimization constraint (Equation 7):

$s^0_i = s^0_{dp}, \quad \text{if } s^1_i > 1, \quad i \in \{p, g, os\}$

What it enforces: if a component is sharded across multiple nodes (s1_i > 1), then it must use all GPUs within each participating node (s0_i = s0_dp). This prioritizes using fewer nodes — if you need 16 GPUs total for a sharding group, prefer s0_i = 8, s1_i = 2 (2 full nodes) over s0_i = 4, s1_i = 4 (4 nodes partially filled).

Why this matters: as illustrated in Figure 8, setting s0_p = 1, s1_p = 2 would require AllGather and ReduceScatter across nodes, while s0_p = 2, s1_p = 1 keeps those operations within a single node. The constraint forces configurations that maximize intra-node communication, which is faster. This is a heuristic that reduces the search space without eliminating any meaningfully distinct configurations, because any cross-node configuration with s0_i < s0_dp can be converted to an equivalent (or better) configuration by expanding to full nodes.

Search procedure. With these constraints, the paper uses a brute-force search — enumerating all valid combinations of the six sharding factors that satisfy all constraints, evaluating T_comm for each using the profiled bandwidth data, and selecting the minimum. The Pre-Filter eliminates configurations violating the dependency, divisibility, or cross-node minimization constraints before the costly communication time evaluation. For typical cluster sizes (up to thousands of GPUs), the constrained search space is small enough that brute-force enumeration is feasible.

What makes this optimization non-trivial: the objective function T_comm is not monotonic in the sharding factors. Increasing sp reduces the per-AllGather message size (each GPU contributes fewer parameters) and reduces the number of GPUs in each AllGather, but increases the number of AllGather operations (since, with more parameter shards, more modules need gathering). The profiled bandwidth data captures how these opposing effects trade off in practice, but it's complex enough that manual tuning — which is what prior systems relied on — cannot find the optimum reliably.

Computation-Communication Overlap: Parameters Sharding

AMSP uses PyTorch's hook system (Table III) to schedule communication operations at specific points in the forward/backward computation graph. For parameter sharding (sp > 1), the overlap strategy operates at module granularity within each Transformer layer.

Forward pass overlap (Figure 9a). Before the forward computation of layer i begins, AMSP initiates AllGather operations to fetch parameters for each module of the next layer (i+1). This is done using register_forward_pre_hook on each module of layer i+1. While layer i's modules compute (multiplying activations by weights, computing attention, applying layer norm), the NCCL communication for layer i+1's parameters proceeds concurrently on the GPU's communication stream. When layer i finishes and layer i+1's first module is about to execute, its register_forward_pre_hook waits for the corresponding AllGather to complete — if the communication is already done (the common case, since computation time typically exceeds AllGather time for a single module's parameters), the wait is a no-op. After a module's forward computation finishes, register_forward_hook releases the gathered parameters to free memory.

Why this works: GPU computation (CUDA kernels) and NCCL communication (which also uses GPU resources but through a separate execution stream) can overlap when they operate on different data. By prefetching parameters for layer i+1 during layer i's computation, AMSP hides the AllGather latency behind useful computation. The key insight is the module-level granularity — gathering all layer i+1 parameters at once would create a long blocking communication burst at the start of the layer; gathering module-by-module interleaves communication with the computation of individual linear layers, attention mechanisms, etc.

Backward pass overlap without activation recomputation (Figure 9b). During the backward pass, each module computes two quantities: GradWeight (gradients with respect to its parameters) and GradInput (gradients with respect to its inputs, needed for the previous module's backward computation). AMSP initiates ReduceScatter on GradWeight immediately after it is computed — this can overlap with the subsequent GradInput computation for the same module. Meanwhile, AMSP fetches parameters for the next module (in backward order) using AllGather, initiated by register_full_backward_pre_hook. This creates a three-way overlap: ReduceScatter for the current module's gradients, AllGather for the next module's parameters, and GradInput computation.

The decoupling problem. A crucial engineering challenge arises: if the ReduceScatter for a module's GradWeight hasn't completed by the time GradInput computation finishes, the backward function for that module would normally exit, and the autograd engine would proceed to the next module. But the gradients aren't ready yet — the ReduceScatter is still running. AMSP addresses this by decoupling the ReduceScatter lifecycle from the backward function. Even after the backward function returns, the ReduceScatter continues asynchronously. A separate post-hook on the AccumulateGrad of each parameter ensures that any downstream operation that needs those gradients (such as AllReduce for optimizer state sharding) waits for the ReduceScatter to complete before proceeding.

Why this is necessary but tricky: PyTorch's autograd engine naturally sequences operations — when backward for module A finishes, it triggers backward for module B. If ReduceScatter is tied to module A's backward function, it would block module B's backward from starting, defeating the purpose of overlap. By decoupling, module B's backward can begin while module A's ReduceScatter is still running, achieving true overlap. The post-hook on AccumulateGrad serves as the synchronization point — gradients aren't consumed until they're actually needed for the optimizer update, which happens after all backward passes complete.

Backward pass overlap with activation recomputation (Figure 9c). When activation recomputation is enabled (as for LLaMA-30B training), the backward pass requires an additional forward pass to recompute activations that were discarded during the original forward pass to save memory. AMSP initiates AllGather for the next layer's parameters at the start of this secondary forward pass. Critically, unlike the original forward pass where parameters are discarded after use, the recomputation forward pass retains the gathered parameters — they will be needed again for the subsequent gradient computation. After GradWeight is computed, ReduceScatter is initiated on it, overlapping with subsequent computations.

Why retention is necessary in recomputation: in the standard forward pass, parameters are used once (for the forward computation) and then another copy is gathered for the backward pass (since they were discarded to save memory). With recomputation, the secondary forward pass produces the parameters that are immediately needed for gradient computation, so retaining them avoids a second AllGather. This is a memory-vs-communication tradeoff that makes sense because recomputation is already trading compute for memory — adding another AllGather would compound the overhead.

Computation-Communication Overlap: Optimizer States Sharding

For optimizer state sharding (sos > 1), the key communication is AllReduce for gradient aggregation (during the last micro-batch's backward pass) and Broadcast for updated parameter distribution (between steps).

Gradient aggregation overlap. During the backward pass of the last micro-batch, as each layer's gradients are computed, they are placed into buckets. When a bucket is full (reaches size U), its contents are flattened into a contiguous buffer and an AllReduce is launched on that buffer. This AllReduce executes asynchronously — the backward computation of the remaining layers continues without waiting. The bucketing is triggered by hooks on the AccumulateGrad of parameters: when enough parameter gradients have accumulated to fill a bucket, the hook flattens and launches the AllReduce.

Why bucketing: launching an AllReduce per parameter would create enormous launch overhead (thousands of tiny NCCL calls). Bucketing amortizes launch overhead across many parameters. The asynchronous execution means gradient aggregation for early layers overlaps with gradient computation for later layers. For a 32-layer model, the AllReduce for layer 1's gradients can complete while layers 2–32 are still computing their gradients.

Updated parameter Broadcast overlap. After the optimizer step (which happens at the end of each training step), updated parameters need to be distributed. AMSP uses the asynchronous communication mechanism of NCCL to overlap the Broadcast of updated parameters with the forward computation of the next step. This requires resolving a synchronization challenge: before any computation or communication uses a parameter in the next step, it must have its updated value. AMSP addresses this with two mechanisms:

Forward pre-hook verification: register_forward_pre_hook on each module checks that the parameters about to be used are the most recent ones — if the Broadcast for that parameter hasn't completed yet, the hook blocks until it does.
Broadcast ordering alignment: the sequence of module computations in the forward pass is aligned with the sequence of Broadcast operations. Parameters for the first layer are broadcast first, so by the time the forward pass reaches the first layer, those parameters are likely already updated. This minimizes blocking.

Why this ordering matters: if parameters were broadcast in arbitrary order, the forward pass might reach a module whose Broadcast hasn't been initiated yet, causing a stall. By matching the broadcast order to the computation order, AMSP maximizes the probability that parameters are ready when needed, making the forward pre-hook wait a rare event rather than the common case.

Computation-Communication Overlap: Gradients Sharding

When sg > sp (gradients are sharded beyond what parameter sharding provides), additional AllReduce operations are needed in every micro-batch except the last to aggregate gradients within the gradient sharding subgroup. These AllReduce operations use the same bucketing strategy as optimizer state sharding: gradients are accumulated into buckets as they are computed during the backward pass, and AllReduce is launched on full buckets without blocking the remaining backward computation.

Following the AllReduce, each GPU uses the select & drop mechanism to keep only the gradients corresponding to its local shard, releasing the rest to conserve GPU memory. This is crucial for maintaining the memory savings that motivated sg > sp in the first place — without the release, the AllReduce would deliver all gradients to all GPUs, defeating the purpose of gradient sharding.

Communication Placement

In GPU clusters with a spine-leaf network topology, GPUs in different nodes communicate through spine switches when they are not under the same leaf switch. Cross-spine communication adds latency and is a shared resource that can become congested. AMSP's Communication Placement optimizer addresses this by controlling the grouping of GPUs for collective operations.

When any s1_i > 1 for i ∈ {p, g, os} (meaning that component's collective operations span multiple nodes), AMSP ensures that the participating nodes are placed under the same leaf switch whenever possible, as illustrated in Figure 10(a). This means AllGather, ReduceScatter, AllReduce, or Broadcast operations that must cross node boundaries at least stay within a single leaf switch's domain, avoiding the additional hop through the spine.

How this is implemented: the paper states that AMSP "strategically groups" nodes under the same leaf switches, but doesn't provide implementation details. The likely mechanism is using NCCL's communicator creation API to specify which GPU ranks participate in each collective, and ensuring that ranks are ordered so that nodes under the same leaf switch are adjacent in the communicator.

Why this matters: for configurations like s0_p = 8, s1_p = 4 (parameter sharding across 4 nodes), if those 4 nodes are spread across different leaf switches, every AllGather and ReduceScatter must traverse spine switches. If AMSP can group all 4 nodes under a single leaf switch (if the topology permits), those collectives avoid spine switch congestion entirely. The benefit is most pronounced at large scale where spine switch bandwidth is the bottleneck.

Summary of Design Choices and Their Justifications

Profiling-based communication model over analytical α-β: the α-β model cannot capture NCCL's algorithm selection, in-network aggregation effects, or real-world bandwidth degradation at scale. Profiling captures all these empirically.
Two-dimensional device mesh (s0 × s1): distinguishes intra-node (NVLINK, 600 GB/s) from inter-node (InfiniBand, 400 GB/s per node) communication, enabling the optimization to prefer node-local collective operations.
Dependency rule (s_dp ≥ s_os ≥ s_g ≥ s_p): prevents pathological configurations where GPUs would need to fetch optimizer states or gradients for parameters they don't locally store, eliminating avoidable communication.
Inter-tensor optimizer state sharding over intra-tensor: compatible with FP8 training (which requires per-tensor scaling factors), even though it requires less efficient Broadcast instead of AllGather for parameter distribution.
sg ∈ {sp, sos} restriction: reduces the search space without meaningful loss, since gradient memory (2Φ bytes) is small relative to optimizer state memory (12Φ bytes), making intermediate sg values unnecessary.
Module-level granularity for AllGather/ReduceScatter: enables fine-grained overlap with computation — fetching parameters for the next module while computing the current one — rather than blocking at layer boundaries.
ReduceScatter lifecycle decoupling from backward functions: allows the backward pass to proceed to the next module while the current module's gradient communication is still in flight, a key enabler of backward-pass overlap.
Bucketing for AllReduce: amortizes NCCL launch overhead and enables asynchronous gradient aggregation that overlaps with the backward computation of subsequent layers.
Cross-node minimization constraint (Equation 7): forces configurations to use fewer nodes when possible, reducing cross-node communication in a topology-aware manner without complex network modeling.
Brute-force search within constrained space: the dependency and divisibility constraints reduce the search space enough to make exhaustive enumeration feasible, avoiding the need for more complex optimization algorithms while still guaranteeing global optimality within the modeled objective.

4. Key Insights and Innovations

Innovation 1: Decoupling Model State Components as Independent Sharding Dimensions — Not Just a Parameter Tuning, But a Category Error Correction

The field's dominant assumption, embedded in every ZeRO variant from ZeRO-1 through ZeRO-3, ZeRO++, and MiCS, was that parameters, gradients, and optimizer states should share the same sharding configuration: sp = sg = sos. This wasn't an arbitrary choice — it followed naturally from the hierarchical design of ZeRO stages, where each stage added more aggressive sharding uniformly across all components. MiCS relaxed this only by allowing sp = sg = sos < sdp (sharding within subgroups rather than across the full cluster), but preserved the fundamental coupling: all three components still moved in lockstep.

AMSP's central conceptual move is to recognize this coupling as a category error — the three model state components have fundamentally different communication patterns, access frequencies, and memory footprints, so forcing them to share a sharding factor means at least two of them are always misconfigured. Parameters are accessed at every module in every micro-batch (frequency: M × L × K per step), gradients at every layer in every micro-batch (frequency: M × L), and optimizer states only once per step. Parameters and gradients each occupy 2Φ bytes in FP16, while optimizer states consume 12Φ bytes for Adam — a 6× asymmetry. Binding them together means either over-communicating for the memory-light components or over-consuming memory for the communication-heavy ones.

The paper's evidence that this is a category error rather than a minor tuning issue is the MiCS-vs-ZeRO-1 inversion on LLaMA-7B at 1024 GPUs (Figure 11a, Figure 12a). MiCS — explicitly designed to reduce ZeRO's communication overhead — achieves 35% MFU, slightly worse than vanilla ZeRO-1 at 36%. This is a system designed to improve performance actually underperforming its simpler predecessor. The mechanism: MiCS forces sp > 1 (Partial-Sharding for parameters within a node) to reduce cross-node communication for AllGather and ReduceScatter, but LLaMA-7B's parameters (14 GB in FP16) already fit comfortably in an 80 GB GPU. The parameter sharding saves memory that doesn't need saving while introducing AllGather/ReduceScatter operations that ZeRO-1's sp = 1 avoids entirely. AMSP's configuration for the same model — sp = sg = 1, sos = 8 — restores ZeRO-1's communication profile for parameters and gradients while adding a node-local Broadcast for optimizer states, achieving 51% MFU. The 15-percentage-point gap between MiCS and AMSP, and the 1-point gap between MiCS and ZeRO-1, together demonstrate that the coupling constraint creates the problem it claims to solve.

This is a fundamental reframing, not an incremental optimization. Prior work asked: "Given that sp = sg = sos, what is the optimal subgroup size?" AMSP asks: "Why should sp equal sg equal sos at all?" The six-dimensional sharding factor space (s0_p, s1_p, s0_g, s1_g, s0_os, s1_os) is not a finer discretization of a known space — it's a space that prior work simply did not explore because the coupling assumption was baked into the system architecture. The dependency rule (Section IV-B) provides the constraints that make this space navigable without pathological configurations, but the key insight is that the space exists and contains configurations (like sp = 1, sos > 1) that are unreachable under any prior system.

The significance extends beyond the immediate performance gains. By demonstrating that decoupling is both feasible and highly beneficial, AMSP establishes that the design space for distributed training optimizers includes an orthogonal dimension — per-component sharding strategy — that had been entirely overlooked. This opens the door for future systems to consider even finer-grained decompositions (e.g., sharding different layers of parameters at different granularities, or treating first and second moment estimates in Adam separately), none of which were conceptually accessible under the sp = sg = sos assumption.

Innovation 2: The Diagnosis That Communication Overhead Has Two Independent Sources — Sharding Strategy AND Execution Scheduling — and That Both Must Be Addressed

A common failure mode in systems research is conflating two sources of inefficiency and fixing only one, leaving substantial gains on the table. AMSP provides a clean experimental dissociation showing that ZeRO's communication overhead comes from two independent sources: (1) the amount and pattern of communication determined by the sharding strategy, and (2) the scheduling of that communication relative to computation determined by the execution engine. Prior work had addressed these implicitly — ZeRO-3 has an overlap mechanism, MiCS changes the communication pattern — but no prior work had isolated their effects or demonstrated that fixing only one leaves the other as the bottleneck.

The paper's key result for this dissociation is the 2× speedup achieved by reimplementing MiCS's sharding configuration under AMSP's execution engine, without changing the sharding strategy at all (Section VI-D, Figure 16–17, and mentioned in the introduction). When training LLaMA-13B on 1024 GPUs with the MiCS configuration (sp = sg = sos = 8 within nodes), DeepSpeed-MiCS achieves some baseline MFU; AMSP with the identical sharding factors achieves approximately 2× higher throughput. The only difference is the execution engine — how AllGather, ReduceScatter, and AllReduce are scheduled relative to layer computations using PyTorch hooks. This isolates the scheduling contribution and demonstrates that DeepSpeed's overlap implementation contains a systems engineering failure that is separable from the sharding strategy question.

The trace evidence in Figure 14 makes the failure mode concrete: DeepSpeed's ZeRO-3, MiCS, and ZeRO++ all exhibit "computation resource bubbles" during the backward pass where ReduceScatter operations block the execution of subsequent computations, and in MiCS's case, ReduceScatter also blocks the concurrent execution of AllReduce. The trace for AMSP (Figure 15) shows these same operations seamlessly interleaved — AllReduce for gradient aggregation runs concurrently with AllGather for parameter prefetching and ReduceScatter for gradient distribution, all while backward computation proceeds.

This dissociation is significant because it establishes a layered optimization framework for distributed training systems. The sharding strategy determines the communication volume and topology — what messages flow between which GPUs. The execution engine determines the communication schedule — when those messages are sent relative to computation. Both layers matter, they're independently optimizable, and prior systems underperformed because they addressed only one (MiCS optimized the sharding layer but ran on an engine with poor scheduling; ZeRO-3 had the opposite problem). AMSP's contribution is not just that it optimizes both, but that it provides the conceptual decomposition that makes the dual optimization legible.

The practical implication is that future work on distributed training can — and should — evaluate these two layers separately. A new sharding strategy should be benchmarked against AMSP's execution engine to avoid attributing scheduling inefficiencies to the strategy itself (as the MiCS-on-DeepSpeed results likely do). Conversely, a new execution engine should be tested across multiple sharding strategies to ensure it doesn't overfit to one communication pattern. The paper's own ablation (Table V) quantifies the marginal contribution of each overlap optimization — Broadcast overlap, AllReduce overlap, and AllGather/ReduceScatter overlap — providing a template for such layered evaluation.

Innovation 3: The Optimization Problem Formulation as a Generality Mechanism — Shifting from Hand-Tuned Heuristics to Constrained Optimization

Prior systems determined sharding configurations through a combination of user-specified stages (ZeRO-1/2/3), rule-based heuristics (MiCS's fixed subgroup sizing), and manual tuning (ZeRO++'s secondary shard configuration). There was no principled mechanism to answer the question: "Given this model, this GPU memory capacity, and this cluster topology, what is the optimal way to shard model states?" The answer was always "ZeRO-3" or "MiCS with subgroup size 8," regardless of whether those configurations were remotely appropriate for the specific model-cluster combination.

AMSP replaces this with a constrained optimization formulation (Equations 1–7) that takes as input the model specification, cluster hardware parameters, and profiled communication bandwidths, and outputs the sharding factors that minimize predicted communication time subject to GPU memory constraints. This is not merely "automating what humans did before" — it changes the nature of the solution from a discrete set of stages to a continuous exploration of a structured space that no human operator would enumerate manually.

The formulation itself is simple — integer programming with a brute-force solver — but its presence in the system architecture is what makes AMSP general rather than model-specific. The configurations in Table IV are not hand-tuned: LLaMA-7B uses sp = sg = 1, sos = 8; LLaMA-13B uses sp = sg = 4, sos = 8; and LLaMA-30B uses sp = sg = 8, sos = 32 (s0_os = 8, s1_os = 4). The progression makes sense post-hoc — larger models need more aggressive parameter sharding to fit in memory, and LLaMA-30B needs optimizer state sharding across 4 nodes to stay within the 80 GB budget — but the specific thresholds (why sos = 8 for 13B but sos = 32 for 30B? why sp = 4 for 13B rather than sp = 2?) would be difficult to derive manually. The optimization problem produces these configurations automatically from the same mathematical framework, without per-model tuning.

The crucial design choice that makes this optimization tractable is the Pre-Filter that prunes the search space using the dependency, divisibility, and cross-node minimization constraints before evaluating communication costs. Without these constraints, the six-dimensional space would be enormous; with them, it becomes small enough for exhaustive search. The constraints are not ad-hoc — they're derived from the physical requirements of distributed training (no orphaned parameters, uniform GPU utilization, topology-aware placement) — making them portable across different models and clusters.

This innovation is incremental in form but fundamental in impact. Integer programming formulations for system configuration are not new (Alpa, FlexFlow, and TensorOpt all use optimization-based approaches for operator parallelism), but applying this approach to model state sharding — and showing that it discovers configurations that outperform human-designed heuristics by 1.3–1.56× (Figures 16–17) — establishes that this dimension of the training system is too complex for manual tuning and must be automated. The paper's result that the optimal configuration varies substantially with model size (7B→13B→30B) even on the same cluster reinforces this: a single "best practice" sharding rule would leave substantial performance on the table.

Innovation 4: The Profiling-Based Communication Model as an Escape from the α-β Straightjacket

The α-β cost model has been the standard tool for predicting collective communication time in distributed systems for decades. It's analytically tractable, it captures first-order scaling behavior, and it's universally taught. But AMSP provides a concrete, empirically-grounded argument for why it is inadequate for production GPU cluster optimization and demonstrates that replacing it with a profiling-based model is both feasible and yields substantially better decisions.

The problem is not that α-β is "wrong" — it correctly predicts that communication time scales with (p - 1)(α + v/(w × p)). The problem is that w is not a constant. Figure 3 shows that the effective bandwidth of an AllGather on a 64 MB message drops from roughly 1200 Gb/s on 8 GPUs (single node, NVLINK) to roughly 300 Gb/s on 512 GPUs (64 nodes, InfiniBand) — a 4× degradation that no single-bandwidth parameter can capture. The α-β model would need different effective bandwidths at every scale, which is functionally equivalent to profiling but obscured behind an analytical facade.

AMSP's decision to use a lookup table of profiled w(o, v, p0 × p1) values — and to interpolate for un-profiled message sizes — is an architectural choice with methodological implications. It acknowledges that NCCL's algorithm selection (ring vs. tree vs. Collnet), switch topology effects, in-network aggregation (when present), and contention patterns are too complex to model analytically, and that the empirical measurement is the model. This is a pragmatic concession to the complexity of real-world GPU clusters, and it works because the profiling is a one-time offline cost per cluster that amortizes over all subsequent training runs.

The significance beyond this paper is that as GPU clusters grow more heterogeneous (different GPU generations, different interconnect topologies, different switch capabilities), analytical communication models will become increasingly inadequate. AMSP provides a template for how to handle this: profile the relevant operations at relevant scales, store the results, and use interpolation for unseen configurations. The cost is upfront profiling time; the benefit is accurate prediction that captures all real-world non-idealities. This is not intellectually deep — it's essentially replacing theory with measurement — but it's an important engineering insight because the field's attachment to analytical models had led to optimization decisions based on predictions that were systematically wrong at scale.

The profiling results themselves (Figure 3) are also a contribution: they provide a characterization of NCCL collective performance across a range of scales and message sizes on a specific hardware configuration (Mellanox HDR InfiniBand, no SHARP, A800 GPUs) that serves as a reference point for practitioners. The finding that Broadcast effective bandwidth degrades much less than AllGather or ReduceScatter at scale (staying above 800 Gb/s vs. dropping to 200–300 Gb/s) is not obvious from first principles and directly informs optimizer state sharding strategies — it's better to use Broadcast (which AMSP does for parameter distribution) than AllGather (which prior systems used) when the operation can be structured either way.

5. Experimental Analysis

Evaluation Methodology

Dataset. The paper uses the LLaMA family of models as training workloads, not a fixed dataset of inputs. The models trained are LLaMA-7B, LLaMA-13B, and LLaMA-30B (Touvron et al., 2023). Training is conducted with a fixed sequence length of 4096 tokens, a micro-batch size of 1 sequence (4096 tokens), and a global batch size of 4 million tokens across all experiments. The micro-batch count M scales inversely with the number of GPUs: M = 128 when training on 8 GPUs, decreasing to M = 1 when training on 1024 GPUs, since the global batch size is held constant to maintain the strong-scaling regime where communication overhead is most visible.
Base model(s). The paper trains LLaMA-based models at three scales: 7 billion parameters (LLaMA-7B), 13 billion parameters (LLaMA-13B), and 30 billion parameters (LLaMA-30B). These model sizes are chosen because they represent the "new paradigm" identified in the paper's motivation — models that are small enough to fit on modest GPU counts when using ZeRO-style sharding, but need to be scaled across many GPUs for training throughput. LLaMA-7B requires 112GB for model states alone (exceeding a single 80GB A100 GPU), LLaMA-13B requires 208GB, and LLaMA-30B requires 480GB, meaning all three require some form of model state sharding to even begin training. The architecture follows the standard Transformer with multiple layers, each containing linear, multi-head attention, and normalization modules.
Metrics. The paper reports two primary performance metrics. Model FLOPs Utilization (MFU) is defined as the ratio of achieved training throughput (in FLOPs per second) to the theoretical peak FLOPs of the GPU hardware, measuring how efficiently the available compute is being used. The FLOPs calculation follows the Megatron-LM formula, and the paper notes that while attention FLOPs should technically be halved due to causal masking, they follow the convention established by FlashAttention and other libraries of not dividing by 2, for consistency with the literature. Tokens per GPU per Second (TGS) measures raw training throughput normalized per GPU, providing a more direct measure of system performance independent of hardware peak FLOPs. Both metrics are reported across all GPU scales from 8 to 1024 GPUs.
Baselines. Four systems are compared: DeepSpeed-ZeRO1 (Rajbhandari et al., 2020 — shards only optimizer states, sos = sdp), DeepSpeed-ZeRO3 (Rajbhandari et al., 2020 — shards parameters, gradients, and optimizer states, sp = sg = sos = sdp), DeepSpeed-ZeRO++ (Wang et al., 2023 — maintains a secondary parameter shard within small subgroups with s0_p = 8, s1_p = 1, while primary sharding uses sp = sg = sos = sdp; quantization is disabled to ensure consistent model quality), and DeepSpeed-MiCS (Zhang et al., 2022 — shards all model states within subgroups and replicates across subgroups, sp = sg = sos < sdp, with subgroup sizes tuned for optimal performance). All baselines use their respective DeepSpeed implementations with communication-computation overlap settings enabled. Activation recomputation is enabled for LLaMA-30B across all systems and disabled for LLaMA-7B and LLaMA-13B.
Generation budget / compute accounting. The paper measures test-time compute in two ways that are held constant across comparisons. The global batch size is fixed at 4 million tokens for all experiments, ensuring that each system processes the same amount of training data per step regardless of GPU count. The micro-batch size is fixed at 1 sequence of 4096 tokens per GPU, meaning per-GPU computation per micro-batch is identical across systems and scales. In the FLOPs-matched analysis (corresponding to Figure 16–17 and Table V), the unit of comparison is the end-to-end training step time under identical computational loads. The profiled communication bandwidth (Figure 3) serves as the basis for the Planner's cost estimates, with all collectives measured on the target cluster in advance and interpolated for un-profiled message sizes. The paper explicitly disables quantization in ZeRO++ to ensure "consistent model quality" — the comparison isolates communication efficiency, not compression-vs-accuracy tradeoffs.
Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing, which is typical for systems benchmarking papers where the metric of interest (throughput, MFU) is deterministic given fixed hardware and software configurations. Instead, the paper relies on trace-based validation (Figures 14, 15, 18–20) showing the actual timeline of computation and communication operations to verify that the claimed overlaps are occurring as designed. The ablation study of overlap strategies (Table V) systematically disables individual overlap optimizations to measure their marginal contribution, providing a more diagnostic form of validation. The optimal sharding configurations (Table IV) are produced by the Planner's optimization formulation and are fixed for all experiments at a given model scale, meaning there is no per-run tuning that would require held-out validation. However, this also means there is no quantification of how sensitive the results are to profiled bandwidth measurement noise or to different cluster conditions — the reported numbers are single-point measurements on a specific 128-node cluster.

Main Quantitative Results

End-to-End Scalability (Figures 11 and 12)

Headline result: AMSP achieves 51%, 52%, and 42% MFU on LLaMA-7B, LLaMA-13B, and LLaMA-30B respectively on 1024 GPUs, compared to the next-best baseline (MiCS) at 35%, 33%, and 29%. The TGS results in Figure 12 mirror this pattern, with AMSP achieving approximately 1.4–1.56× higher throughput than MiCS across all model sizes at 1024 GPUs.

LLaMA-7B scalability (Figures 11a, 12a). At 8 GPUs, AMSP and ZeRO-1 achieve nearly identical MFU (~63–64%), confirming that AMSP's execution engine does not introduce computational overhead compared to the baseline when communication patterns are similar (both use sp = sg = 1 at this scale). As GPU count increases to 1024, the divergence is dramatic:

AMSP: MFU declines from ~64% at 8 GPUs to 51% at 1024 GPUs (a modest 15-percentage-point drop).
ZeRO-1: declines from ~63% at 8 GPUs to 36% at 1024 GPUs (a 27-point drop).
MiCS: declines from ~50% at 8 GPUs to 35% at 1024 GPUs — critically, MiCS underperforms ZeRO-1 at 1024 GPUs (35% vs. 36%), the inversion that motivates the paper's central argument about inflexible sharding.
ZeRO-3: declines from ~47% at 8 GPUs to ~4% at 1024 GPUs (an 88% relative reduction).
ZeRO++: similar catastrophic degradation to ZeRO-3, achieving ~4% at 1024 GPUs.

The key comparison at 1024 GPUs: AMSP at 51% vs. ZeRO-1 at 36% vs. MiCS at 35%. The 15-percentage-point gap between AMSP and ZeRO-1 comes entirely from the decoupled sharding strategy — AMSP uses sos = 8 (Broadcast within nodes for updated parameters) while ZeRO-1 broadcasts across all 1024 GPUs, a 128× larger communication group. The 1-point inversion between MiCS and ZeRO-1 demonstrates that MiCS's AllGather/ReduceScatter overhead (from its forced sp > 1) outweighs the benefit of reduced communication scale.

LLaMA-13B scalability (Figures 11b, 12b). The pattern is similar but with a notable shift: MiCS no longer underperforms ZeRO-3 as dramatically, because the larger model size (208GB model states) makes parameter sharding more necessary — the memory benefit of sp > 1 now offsets some communication cost. At 1024 GPUs:

AMSP: 52% MFU (remarkably, higher than on LLaMA-7B — the paper does not explain this but it may reflect better compute-communication balance at this model size).
MiCS: 33% MFU.
ZeRO-3: degrades to ~6% MFU (slightly better than on 7B but still catastrophic).
ZeRO++: ~6% MFU.

The configuration AMSP uses for 13B is sp = sg = 4, sos = 8 (Table IV), meaning parameters are sharded within half a node (4 GPUs) and optimizer states within a full node (8 GPUs). This represents the point where parameter sharding becomes necessary for memory (unlike 7B where sp = 1 was feasible) but the sharding can still be confined within a single node (s1_p = s1_g = s1_os = 1), avoiding cross-node AllGather/ReduceScatter entirely.

LLaMA-30B scalability (Figures 11c, 12c). This is the most memory-constrained configuration, forcing some cross-node sharding. At 1024 GPUs:

AMSP: 42% MFU.
MiCS: 29% MFU.
ZeRO-3: degrades to ~3% MFU.
ZeRO++: encounters Out-Of-Memory (OOM) at 32 GPUs, making it non-viable at smaller scales; at 1024 GPUs, it achieves ~5% MFU.

AMSP's configuration for 30B is sp = sg = 8, sos = 32 (s0_os = 8, s1_os = 4) — parameters are sharded within a single node (8 GPUs) but optimizer states span 4 nodes. This introduces cross-node Broadcast for parameter distribution (captured in T1_os) but avoids cross-node AllGather/ReduceScatter (since s1_p = s1_g = 1). The Planner determined that 30B's 480GB model state memory forces some component to cross node boundaries, and optimizer states (which require communication only once per step, not per module per micro-batch) are the least damaging choice.

GPU Memory Analysis (Figure 13)

The memory results demonstrate the tradeoff between communication efficiency and memory consumption that the Planner navigates. At 1024 GPUs:

ZeRO-3 is consistently the most memory-efficient, using approximately 15–16GB for all three model sizes. The memory consumption stabilizes at large scale because the per-GPU model state shard (2Φ/1024 + 2Φ/1024 + 12Φ/1024) becomes negligible — further scaling provides almost no additional memory savings (as argued in Section III-B).
ZeRO++ consumes similar memory to ZeRO-3 but with slightly higher overhead from the secondary parameter shard.
MiCS consumes roughly double ZeRO-3's memory — approximately 30–35GB for the 7B and 13B models, reflecting the replication of model states across subgroups (sdp/sp = 1024/8 = 128 replicas for sp = sg = sos = 8).
AMSP consumes the most memory among all systems for LLaMA-7B (~60GB at 1024 GPUs, roughly double MiCS), because sp = sg = 1 means every GPU holds a full copy of parameters and gradients (16GB total). For LLaMA-13B, AMSP and MiCS show similar total memory (~40GB) but with different compositions — AMSP uses more parameter memory (sp = 4 vs. MiCS's sp = 8) but less optimizer state memory in some configurations. For LLaMA-30B, both systems consume approximately 40–45GB.

The key takeaway: AMSP's higher memory consumption on 7B is the reason it achieves higher MFU — it deliberately trades memory for communication efficiency. The memory is well within the 80GB GPU capacity, so this is a pure benefit. The Planner's constraint (Equation 2) ensures this tradeoff never causes OOM; the optimization objective (Equation 1) ensures that memory is "spent" on the components whose replication yields the largest communication savings (parameters and gradients on 7B, optimizer states on larger models).

Performance Gap Analysis: Traces (Figures 14 and 15)

The paper includes detailed execution traces that reveal why AMSP outperforms baselines with the same or similar sharding configurations. Figure 14 shows the backward phase of the last micro-batch for LLaMA-7B training on 32 GPUs under DeepSpeed-ZeRO3, DeepSpeed-MiCS, and DeepSpeed-ZeRO++:

"It reveals instances of computation resource bubbles during the backward pass. Even when the communication-computation overlap setting is enabled, DeepSpeed-ZeRO3/MiCS/ZeRO++ fails to effectively overlap ReduceScatter with computation. Additionally, in DeepSpeed-MiCS, the ReduceScatter operation also blocks the concurrent execution of AllReduce."

The bubbles are visible as gaps in the computation stream where GPUs are idle while waiting for ReduceScatter to complete. In MiCS, the blocking is compounded: not only does ReduceScatter block computation, but it also blocks AllReduce from starting, serializing two operations that could run concurrently.

Figure 15 shows the corresponding AMSP trace for LLaMA-7B, 13B, and 30B:

Forward pass (panels b and c): AllGather communication for the next layer's parameters is visible running concurrently with the current layer's computation, with minimal white space indicating effective overlap.
Backward pass: AllReduce (for gradient aggregation) runs concurrently with AllGather (for parameter prefetching) and ReduceScatter (for gradient distribution), all interleaved with backward computation kernels. The trace shows substantially less idle time compared to Figure 14.
Between-step Broadcast (panel a): The Broadcast of updated parameters from the optimizer step overlaps with the forward computation of the next step.

The complete traces in Appendix A (Figures 18–20) show full training steps including both forward and backward phases at a micro-batch count of M = 2, confirming that the overlap patterns observed in the zooms are representative of the full execution.

Ablation: Sharding Strategy vs. Execution Engine (Figures 16 and 17)

This is the paper's most important experimental dissociation. It compares three conditions:

DeepSpeed (MiCS Config): MiCS's sharding configuration (sp = sg = sos = 8 within nodes for 7B and 13B) running on DeepSpeed's execution engine.
Ours (MiCS Config): The same MiCS sharding configuration running on AMSP's execution engine.
Ours (Optimal Config): AMSP's Planner-optimized configuration running on AMSP's execution engine.

LLaMA-7B at 1024 GPUs (Figure 16a):

DeepSpeed with MiCS config: ~35% MFU (as in Figure 11).
AMSP with MiCS config: approximately 42% MFU — a 1.2× improvement from the execution engine alone, without changing the sharding strategy.
AMSP with optimal config (sp = sg = 1, sos = 8): 51% MFU — a further 1.21× improvement from the sharding strategy change.

LLaMA-13B at 1024 GPUs (Figure 16b):

DeepSpeed with MiCS config: ~33% MFU.
AMSP with MiCS config: approximately 36% MFU — a 1.09× improvement from the execution engine.
AMSP with optimal config (sp = sg = 4, sos = 8): 52% MFU — a 1.44× improvement over AMSP-MiCS, showing the sharding strategy matters more at this model size than the execution engine.

LLaMA-30B at 1024 GPUs (Figure 16c):

DeepSpeed with MiCS config: ~29% MFU.
AMSP with MiCS config: approximately 33% MFU — a 1.14× improvement from the execution engine.
AMSP with optimal config (sp = sg = 8, sos = 32): 42% MFU — a 1.27× improvement over AMSP-MiCS.

The key insight from this ablation: the relative contribution of sharding strategy vs. execution engine varies with model size. For LLaMA-7B, the execution engine contributes roughly as much improvement (1.2×) as the sharding strategy (1.21×). For LLaMA-13B, the sharding strategy dominates (1.44× vs. 1.09×). This is because LLaMA-7B's MiCS config (sp = sg = 8) forces unnecessary parameter AllGather/ReduceScatter, while the optimal config (sp = 1) eliminates these entirely — the sharding strategy difference is qualitative, not just quantitative. For larger models, the difference narrows because both configurations must shard parameters to some degree to fit in memory.

Figure 17 shows the corresponding TGS results, which exhibit the same pattern: at 1024 GPUs, AMSP with optimal config achieves ~4500 TGS on 7B (vs. ~3200 for AMSP-MiCS and ~2900 for DeepSpeed-MiCS), ~2200 TGS on 13B (vs. ~1500 and ~1300), and ~1000 TGS on 30B (vs. ~700 and ~600).

Overlap Strategy Ablation (Table V)

The paper systematically disables overlap optimizations in sequence to measure their marginal contribution, testing on 64 GPUs with micro-batch counts M ∈ {1, 2, 4, 8}. The four configurations are:

Full overlap: AllGather/ReduceScatter + AllReduce + Broadcast (all enabled).
No Broadcast overlap: AllGather/ReduceScatter + AllReduce only.
No Broadcast or AllReduce overlap: AllGather/ReduceScatter only.
No overlap: all overlap disabled (serial communication and computation).

LLaMA-7B results (Table V, top section):

At M = 1 (the most communication-constrained regime, corresponding to large-scale training):

Full overlap: 3525 TGS, 0.57 MFU.
Disabling Broadcast overlap: 3346 TGS (5.1% drop in TGS).
Disabling Broadcast and AllReduce: 2812 TGS (additional 16.0% drop from previous, 20.2% total drop from full).
No overlap: 2812 TGS (identical to previous — no further degradation).

The finding that disabling AllGather/ReduceScatter overlap causes no additional degradation at M = 1 is surprising, but explained by the configuration: LLaMA-7B uses sp = 1 (Full-Replica for parameters), so there are no AllGather/ReduceScatter operations for parameter sharding to overlap. The overlap benefit comes entirely from AllReduce (gradient aggregation) and Broadcast (updated parameter distribution). At M = 8 (less communication-constrained, corresponding to small-scale training), the full overlap achieves 4503 TGS (0.67 MFU), but the relative benefit of overlap is smaller because computation time is larger relative to communication.

LLaMA-13B results (Table V, middle section):

LLaMA-13B uses sp = 4 > 1, so AllGather/ReduceScatter overlap now applies:

Full overlap at M = 1: 1918 TGS, 0.53 MFU.
Disabling Broadcast: 1895 TGS (1.2% drop — Broadcast overlap matters less when parameter sharding is active because the forward pass is already interleaved with AllGather, leaving less idle time for Broadcast to fill).
Disabling Broadcast and AllReduce: 1630 TGS (additional 14.0% drop).
No overlap: 1527 TGS (additional 6.3% drop — AllGather/ReduceScatter overlap contributes 6.3% at M = 1).

The total benefit of overlap at M = 1 is 1.25× (1918/1527), consistent with the paper's claim. At M = 8, the benefit is 1.34× (2230/1666), showing that overlap becomes more valuable as the number of micro-batches increases, because there are more opportunities for interleaving.

LLaMA-30B results (Table V, bottom section):

LLaMA-30B uses sp = 8 > 1 and has activation recomputation enabled, creating the most complex overlap scenario:

Full overlap at M = 1: 825 TGS, 0.48 MFU.
Disabling Broadcast: 759 TGS (8.0% drop — Broadcast overlap is more important for 30B because sos = 32 with s1_os = 4, meaning cross-node Broadcast that benefits substantially from being hidden behind forward computation).
Disabling Broadcast and AllReduce: 651 TGS (additional 14.2% drop).
No overlap: 557 TGS (additional 14.4% drop — the AllGather/ReduceScatter overlap is the single largest contributor at this model size, because parameter sharding is most aggressive and recomputation creates additional opportunities for overlap during the secondary forward pass).

The total overlap benefit at M = 1 is 1.48× (825/557), matching the paper's claim. This is the largest benefit among the three model sizes, reflecting the compounded communication costs of aggressive parameter and optimizer state sharding, plus recomputation, that overlap can hide.

Critical observation from Table V: The "No Overlap" and "Overlap AllGather/ReduceScatter" rows are identical for LLaMA-7B at all micro-batch counts (e.g., both 2812 TGS at M = 1). This confirms that when sp = 1, the AllGather/ReduceScatter overlap mechanism is not simply ineffective — it is literally unused because there are no such operations to overlap. This validates the paper's claim that the Planner's choice of sp = 1 eliminates this communication category entirely, simplifying the overlap problem to only AllReduce and Broadcast.

Critical Assessment

This paper makes three central claims that the experiments must support:

Claim 1: Flexible per-component sharding achieves substantially higher training throughput than prior ZeRO variants at scale. The evidence for this is mixed but largely supportive. The head-to-head comparisons in Figures 11 and 12 show AMSP outperforming all baselines at every GPU count from 8 to 1024 across all three model sizes, with the gap widening substantially at larger scales. At 1024 GPUs, the improvement over the next-best baseline (MiCS) ranges from 1.29× on LLaMA-13B (52% vs. 33% MFU) to 1.56× on LLaMA-7B (51% vs. 35%). The ablation in Figures 16–17 confirms that the sharding strategy itself accounts for a substantial fraction of this improvement: switching from MiCS's config to AMSP's optimal config while keeping the execution engine fixed improves MFU by 1.21× on LLaMA-7B, 1.44× on LLaMA-13B, and 1.27× on LLaMA-30B.

However, the claim is specific to large-scale training (hundreds to thousands of GPUs) with small per-GPU micro-batch counts (M = 1). At 8 GPUs, AMSP and ZeRO-1 achieve nearly identical MFU (Figure 11), meaning the benefits emerge only when the compute-to-communication ratio collapses. The paper does not provide a crossover analysis — at what GPU count does the benefit become practically meaningful? The gap between AMSP and MiCS at 32 GPUs is already visible (~60% vs. ~52% MFU on 7B), but the gap between AMSP and ZeRO-1 at 32 GPUs is much smaller (~60% vs. ~57%). Users training at moderate scales (16–64 GPUs) may see much smaller benefits than the headline 1.56× figure.

A missing experiment would be a scaling efficiency plot — MFU as a fraction of single-node MFU vs. GPU count — which would directly show how well each system maintains efficiency as scale increases. From the data provided, AMSP on LLaMA-7B retains 51/63 ≈ 81% of its 8-GPU MFU at 1024 GPUs, while ZeRO-1 retains only 36/63 ≈ 57%. This framing makes the benefit concrete but is not presented in the paper.

Claim 2: The execution engine's overlap scheduling contributes independently to performance, and prior systems leave substantial overlap on the table. This claim is strongly supported by the 2× gap between DeepSpeed-MiCS and AMSP-MiCS (same sharding config, different execution engines) mentioned in Section I and shown in Figures 16–17. At 1024 GPUs on LLaMA-13B, AMSP-MiCS achieves approximately 36% MFU vs. DeepSpeed-MiCS at 33% (Figure 16b), which is a 1.09× improvement rather than 2× — the 2× figure likely refers to a specific model-scale combination not fully detailed. The trace evidence in Figure 14 vs. Figure 15 directly visualizes the mechanism: DeepSpeed's ReduceScatter creates idle computation bubbles and blocks AllReduce, while AMSP's execution interleaves these operations.

The Table V ablation is methodologically sound and supports the claim that each category of overlap (AllGather/ReduceScatter, AllReduce, Broadcast) contributes measurable gains, with the total overlap benefit ranging from 1.25× to 1.48× depending on model size. However, Table V tests on 64 GPUs, not 1024 — the overlap benefit at extreme scale (where M = 1 and communication dominates) may differ from what's measured.

Claim 3: The optimization problem formulation discovers configurations that outperform human-designed heuristics. The evidence comes from comparing "Ours (Optimal Config)" vs. "Ours (MiCS Config)" in Figures 16–17, which isolates the Planner's contribution (since both use the same execution engine). The optimal config outperforms MiCS's config by 1.21× (7B), 1.44× (13B), and 1.27× (30B). This is compelling evidence that the Planner finds non-obvious configurations — particularly for LLaMA-7B, where the optimal sp = 1, sos = 8 contradicts the intuition that all components should be sharded similarly.

However, several weaknesses limit the generality of this claim:

No comparison against human expert tuning. The paper compares against MiCS's rule-based configuration, which is one specific heuristic. An expert practitioner might manually configure sp = 1, sos = 8 for LLaMA-7B after observing that parameters fit in memory — the Planner automates this insight but doesn't necessarily discover configurations that are beyond human reach. A stronger claim would require showing that the Planner finds configurations that differ from what an expert would choose and that perform better.
The search space is relatively small after pruning. With the dependency, divisibility, and cross-node minimization constraints, the number of valid configurations is modest — for a 1024-GPU cluster (sdp = 1024), the possible sos values must divide 1024, giving {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024}, and each sg and sp takes only one or two values given the sg ∈ {sp, sos} restriction. A human could enumerate this space manually. The Planner's value is in automating the evaluation of each configuration against profiled bandwidth data and memory constraints, not in searching an intractably large space.
The objective function (minimize T_comm) assumes communication is the bottleneck. This is valid at large scale with M = 1, but at smaller scales where computation and communication are more balanced, minimizing communication time alone might not maximize throughput — a configuration with slightly higher communication but better computation-communication overlap might win. The paper does not discuss whether the optimal configuration changes with micro-batch count M, which would test this assumption.
Single-cluster validation. All profiling and experiments are on one 128-node cluster with A800 GPUs and Mellanox HDR InfiniBand without SHARP. The Planner's decisions depend on the profiled bandwidth data — on a cluster with different interconnect characteristics (e.g., NVSwitch, SHARP-enabled InfiniBand, or Ethernet-based interconnects), the optimal configurations might differ. The paper provides no evidence that the Planner would produce different (and correct) configurations on different hardware, which is the key claim of generality.

Genuine weaknesses in the experimental design:

No measurement of convergence or model quality. The paper reports only throughput metrics (MFU, TGS), never verifying that the trained models achieve the same loss or accuracy as baselines. While this is standard for systems papers that don't change the training algorithm (AMSP uses the same optimizer, same mixed-precision scheme, same data), the use of different communication patterns could theoretically affect numerical reproducibility due to non-associative floating-point reduction ordering. The paper does not address this.
The 7B result's dependency on large GPU memory. AMSP's optimal config for LLaMA-7B (sp = sg = 1) requires ~60GB per GPU (Figure 13a), which is well within the 80GB A800 capacity but would cause OOM on a 40GB A100 or 32GB V100. The Planner would generate a different config for those GPUs, but the paper provides no evidence that the resulting config would still outperform MiCS. This means the headline 51% MFU for 7B is specific to 80GB GPUs and would degrade on more memory-constrained hardware — the paper doesn't quantify how much.
ZeRO++ and quantization. The paper disables ZeRO++'s quantization to "ensure consistent model quality," which is a legitimate methodological choice, but it means ZeRO++ is evaluated without its primary contribution (communication compression). A fairer comparison might include ZeRO++ with quantization, since the paper's claim is about communication efficiency broadly, not just about sharding strategies. The paper's argument that "no amount of compression can fix the latency problem" is plausible but not experimentally demonstrated.
Missing baseline: PyTorch FSDP. FSDP is mentioned as implementing ZeRO-3 in PyTorch natively and is widely used. The paper doesn't benchmark against it, nor against FSDP's own hybrid sharding mode (which is similar in spirit to MiCS). This is a significant omission given FSDP's production deployment status.
The 30B OOM for ZeRO++ at 32 GPUs (Figure 11c) is not explained. ZeRO++ is supposed to have lower memory than ZeRO-3 for some configurations due to the secondary shard, so encountering OOM where ZeRO-3 succeeds is surprising and deserves investigation. It may be a configuration bug or an implementation issue, but the paper treats it as a black-box result.
Overlap percentages are sensitive to cluster conditions. The Table V results are collected on 64 GPUs of the paper's specific cluster, and the absolute overlap benefit depends on the ratio of computation time to communication time. On a cluster with faster interconnects or different GPU SKUs, the overlap benefit would differ. The paper does not provide any sensitivity analysis or discussion of this.

What would strengthen the paper:

A Pareto frontier analysis showing the memory-vs-throughput tradeoff explicitly — for each model size, plot achievable MFU against GPU memory consumption for all valid sharding configurations, with AMSP's optimal config and baseline configs marked. This would make the tradeoff space visual and show whether AMSP's config is truly Pareto-optimal or just one point on the frontier.
Results on additional model architectures — all experiments use LLaMA-based decoder-only Transformers. Encoder-decoder models (T5) or mixture-of-experts architectures might have very different parameter distributions that change the optimal sharding strategy.
Convergence validation showing that loss curves are identical (within noise) across all systems at a fixed scale, confirming that communication pattern changes don't affect training dynamics.
A strong scaling plot (fixed model, increasing GPUs) showing throughput scaling efficiency, which would directly visualize the paper's central claim about communication overhead at scale more clearly than the absolute MFU numbers.
The Planner's output for a hypothetical memory-constrained scenario (e.g., LLaMA-7B on 40GB GPUs) to demonstrate that the optimization framework generalizes to different hardware constraints.

In summary, the experiments strongly support the claim that flexible per-component sharding improves throughput at large scale on the tested hardware, and moderately support the claim that the optimization formulation generalizes beyond manual tuning. The dissociation between sharding strategy and execution engine contributions is clean and convincing. The paper's experimental weaknesses are primarily about scope (single cluster, single model family, no convergence results, no memory-constrained scenarios) rather than internal validity. The results are real and substantial, but the paper overstates their generality—the 1.56× improvement is measured on one specific cluster configuration, and the claim that the Planner "discovers optimal sharding factors" is not tested under conditions different enough from the training environment to establish robustness.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Dwarfs the Test-Time Compute Budget It Optimizes

The assumption or constraint. The Planner selects sharding factors by solving an optimization problem (Equations 1–7) that relies on profiled communication bandwidth data w(o, v, p0 × p1) collected offline on the target cluster (Section IV-A). The paper states:

"We utilize t(o, v, p0 × p1) = v/w(o, v, p0 × p1) to evaluate the time consumption of a collective communication operator (o) with a given data size (v) and a specified participant GPU device mesh"

The profiling must cover all relevant combinations of operation type, message size, and device mesh configuration — the four operations (AllReduce, AllGather, ReduceScatter, Broadcast) across message sizes from 1 MB to 1 GB and device meshes from 8×1 to 8×64 GPUs (Figure 3). This is a one-time offline cost per cluster, so it amortizes across training runs on that hardware. However, the paper never quantifies this profiling cost, and the Planner's decisions are only as good as the profiled data — if the cluster changes (hardware upgrades, network reconfiguration, different GPU SKUs), re-profiling is required. Similarly, if the training workload's communication patterns stress the network differently than the micro-benchmarks (contention from multiple simultaneous jobs), the profiled point-to-point bandwidths may overestimate achievable throughput.

The consequence. A practitioner deploying AMSP on a new cluster must first run an extensive profiling campaign before the Planner can produce a configuration. The profiling matrix scales as |operations| × |message sizes| × |device meshes|, which for the paper's configuration (4 operations × ~10 message sizes × 7+ device meshes) represents hundreds of micro-benchmark runs, each requiring exclusive cluster access to avoid measurement interference. On a shared production cluster, this is a non-trivial operational burden that the paper does not account for. More critically, the Planner's optimal configuration is specific to the profiled bandwidths — if network congestion during actual training differs from profiling conditions (e.g., due to other jobs sharing the spine switches), the selected sharding factors may be suboptimal. The paper provides no sensitivity analysis: how much does the optimal configuration change if profiled bandwidths differ by 10%? 50%?

What evidence exists in the paper. None. The offline profiling cost is mentioned as a design choice (Section IV-A) but never characterized in terms of wall-clock time, cluster-hours consumed, or number of individual benchmarks required. The sensitivity of the Planner's output to profiling noise is not analyzed — the paper reports only the final configurations in Table IV without showing alternative near-optimal configurations that might be more robust. Figure 3 provides the profiled data but only for the specific cluster used; there is no comparison of Planner outputs across different clusters with different interconnect characteristics.

Mitigation status. Not addressed. The paper does not discuss profiling overhead, sensitivity to measurement noise, or the portability of profiled data across similar-but-not-identical clusters. A production system would likely need to develop lightweight profiling (e.g., measuring only key message sizes and interpolating more aggressively) or online adaptation (adjusting sharding factors during training if observed communication times diverge from predictions), neither of which AMSP provides.

All Results Are on a Single Cluster with Specific Hardware — The Planner's "Optimal" Configurations May Not Transfer

The assumption or constraint. Every experiment in the paper — the profiled bandwidths in Figure 3, the end-to-end training results in Figures 11–13, the ablation studies in Figures 16–17, the overlap analysis in Table V — is conducted on a single dedicated cluster:

"The training is conducted on a dedicated cluster with 128 GPU servers. Each server is equipped with 8 GPUs and 128 CPU cores, resulting in a total of 1024 NVIDIA Ampere GPUs (A800). Each GPU is outfitted with 80GB of memory, interconnected through NVLink within a node, and inter-node communication is facilitated by 4 Mellanox HDR InfiniBand without SHARP."

This hardware configuration has specific properties that directly influence the Planner's decisions: 600 GB/s intra-node NVLINK bandwidth, 400 GB/s inter-node InfiniBand per node, and crucially, no SHARP (in-network aggregation). On clusters with SHARP, AllReduce effective bandwidth degrades much less with scale because reduction happens in the switch fabric — the 4× bandwidth degradation AMSP observes for AllGather at scale (Figure 3) would be substantially different, potentially changing which sharding configurations are optimal. Similarly, clusters using NVSwitch (which provides higher intra-node bandwidth) or Ethernet-based interconnects (which have different latency characteristics) would produce different profiling results and therefore different Planner outputs.

The consequence. The paper claims AMSP is a general system — the Planner "discovers the optimal sharding factors" for any model on any cluster — but provides zero evidence that this claim holds across different hardware. A practitioner with a different cluster topology (e.g., A100s with NVSwitch, H100s with NVLink 4.0 and SHARP-enabled InfiniBand, or cloud instances with variable network performance) cannot know whether the Planner would produce configurations that outperform MiCS or ZeRO-1, or whether the 1.56× improvement generalizes. The paper's ablation showing Planner-optimal configs outperforming MiCS configs (Figures 16–17) is valid for the paper's cluster but may not hold elsewhere — on a cluster with SHARP, for example, ZeRO-3's AllReduce might scale much better, narrowing or reversing the advantage of decoupled sharding. The absence of SHARP in the paper's testbed is particularly notable because it represents a worst-case scenario for ZeRO's communication overhead; on SHARP-enabled clusters, the baseline systems likely perform better, and AMSP's relative advantage may shrink.

What evidence exists in the paper. The paper provides a single-cluster validation. There are no experiments on multiple cluster configurations, no sensitivity analysis varying interconnect bandwidth or topology, no comparison of Planner outputs across different hardware assumptions, and no discussion of how the profiling approach would adapt to heterogeneous clusters (different GPU types within the same training job). The paper's reliance on profiling rather than analytical modeling is both a strength (captures real-world non-idealities) and a weakness (binds results to the specific hardware measured). The paper acknowledges this implicitly by providing detailed hardware specifications but never discusses the implications for transfer.

Mitigation status. Not addressed. The paper does not include a "Robustness to Hardware Variation" experiment, a comparison across different network topologies, or even a discussion of expected behavior on SHARP-enabled or NVSwitch-equipped clusters. Future work would need to validate AMSP on diverse hardware and potentially develop a Planner that can reason about hardware characteristics it hasn't directly profiled (e.g., extrapolating from similar cluster configurations).

The Planner Optimizes Communication Time, Not End-to-End Throughput — and Ignores Computation-Communication Balance

The assumption or constraint. The optimization formulation (Equation 1) minimizes T_comm — the total predicted communication time per training step — subject only to the memory constraint (Equation 2). This assumes that computation time is fixed and that minimizing communication time directly maximizes throughput. The paper motivates this by focusing on the large-scale regime where "computation time linearly reduces" (Figure 2b) while communication time grows, creating a communication-dominated regime. However, the Planner does not model computation time at all — it has no concept of how long the forward/backward passes take, how the micro-batch count M affects the compute-communication ratio, or whether a configuration with slightly higher T_comm but better overlap potential might achieve higher throughput.

The paper's own data shows why this matters. In Table V, the benefit of overlap optimizations varies substantially with M: for LLaMA-7B at M = 1, overlap improves TGS by ~25% (2812 → 3525); at M = 8, the same overlap improves TGS by only ~5% (4286 → 4503). This is because at higher M, computation time per micro-batch is unchanged but there are more micro-batches, so total computation time increases relative to communication — the compute-communication ratio shifts, and the marginal benefit of communication optimization shrinks. The Planner, which only minimizes T_comm, would produce the same configuration regardless of M, even though a configuration that prioritizes memory savings over communication reduction (allowing a larger micro-batch size) might be better at high M.

The consequence. The Planner may over-optimize for communication reduction in regimes where computation is the bottleneck. For training at moderate scales (16–64 GPUs) with reasonable micro-batch sizes, communication time is not the dominant factor — the Planner's configuration might reduce communication from 100 ms to 50 ms while leaving a 500 ms computation untouched, achieving a barely-noticeable throughput improvement while consuming more memory. More subtly, the Planner might reject a configuration that enables better overlap because its raw T_comm is higher, even though the overlapped execution achieves lower effective communication overhead. The paper's own observation — "the objective function (minimize T_comm) assumes communication is the bottleneck" — is noted as a limitation of the experimental design in Section 5 but is actually a limitation of the Planner's design philosophy.

The paper also never explores whether the Planner's output would change if the optimization objective were end-to-end step time rather than just communication time. Since computation time is not modeled (and would require running the actual training loop to measure, defeating the purpose of offline planning), the Planner cannot reason about the throughput impact of its decisions — it can only minimize one component of the total time.

What evidence exists in the paper. The Table V overlap ablation indirectly demonstrates the limitation: the overlap benefit varies with M, meaning the effective communication cost (after overlap) depends on how much computation there is to hide behind. The Planner's raw T_comm objective does not account for this. Figure 2 further shows that computation time varies linearly with GPU count at fixed global batch size — the Planner uses the same configurations at all scales (Table IV) even though the computation-communication balance shifts dramatically from 8 GPUs to 1024 GPUs. The paper never evaluates whether different configurations would be optimal at 64 GPUs vs. 1024 GPUs, or whether the Planner would produce different outputs if M were provided as input.

Mitigation status. Not addressed. The paper briefly notes in Section 5 (prior sections) that this is a weakness but does not discuss how the Planner could be extended to model computation time or to incorporate M as an input. A natural extension would be a two-term objective T_comm + T_comp_overhead, where T_comp_overhead captures the computation time lost due to imperfect overlap (which depends on sharding factors and M), but the paper does not pursue this.

Memory-Constrained Regimes May Force Suboptimal Configurations — and the Planner Cannot Mitigate This

The assumption or constraint. The Planner's memory constraint (Equation 2) is a hard cutoff: D_total ≤ GPU_Memory_Capacity. As long as this is satisfied, the Planner freely minimizes T_comm. But when memory is tight — for very large models, or on GPUs with less than 80 GB memory — the constraint becomes binding, and the Planner is forced to increase sharding factors to meet the memory budget, even if those increases dramatically raise communication costs. The paper's results with LLaMA-30B hint at this: the model requires sp = 8, sos = 32 with s1_os = 4 (cross-node optimizer state sharding), achieving 42% MFU on 1024 GPUs compared to 52% for LLaMA-13B (which stays entirely within a node). The 10-percentage-point drop is partly due to the larger model's increased computation, but the cross-node Broadcast for optimizer states (captured in T1_os) adds communication that the Planner would avoid if memory permitted.

More critically, the paper's headline result for LLaMA-7B — sp = sg = 1, Full-Replica for parameters and gradients — is only possible because 80 GB A800 GPUs have enough memory to hold the full parameters (14 GB), gradients (14 GB), optimizer states per node (12 × 7B / 8 = 10.5 GB per GPU), plus activations and buffers, totaling ~60 GB (Figure 13a). On a 40 GB A100 GPU, or with a larger model like LLaMA-65B, this configuration would cause OOM, and the Planner would be forced to increase sp (triggering expensive AllGather/ReduceScatter) or sos (triggering cross-node Broadcast), degrading throughput. The paper provides no analysis of how the Planner's output — and the resulting throughput — degrades as GPU memory shrinks.

The consequence. The "up to 1.56× improvement" is achievable only when the GPU memory budget is generous relative to the model size. For the most memory-constrained scenarios — training large models on consumer GPUs, using older hardware with limited memory, or training with large batch sizes that consume activation memory — the Planner's optimal configuration will necessarily involve more aggressive sharding, and the advantage over baselines will shrink. A practitioner training LLaMA-7B on 40 GB A100 GPUs (a common cloud instance) cannot achieve AMSP's 51% MFU because the sp = 1 configuration will OOM. The paper does not report what configuration the Planner would select for this case or what throughput it would achieve relative to MiCS or ZeRO-3.

Furthermore, the paper's memory model (D_modelstate in Section IV-D) only accounts for parameters, gradients, and optimizer states. Activation memory (D_activation) is mentioned as being "seamlessly integrated" from prior work but is never explicitly accounted for in the optimization — the paper does not provide the formula used, the assumptions about activation recomputation, or how activation memory varies with micro-batch size and sequence length. If the activation memory estimate is wrong (e.g., underestimating the memory needed for attention intermediates), the Planner could select a configuration that OOMs at runtime. The paper provides no evidence of OOM-free execution across a range of configurations — the only OOM reported is for ZeRO++ on LLaMA-30B at 32 GPUs, which is not AMSP's fault but also doesn't validate AMSP's memory model.

What evidence exists in the paper. The memory consumption data in Figure 13 shows AMSP using 60 GB on 7B, 40 GB on 13B, and 42 GB on 30B — all well within the 80 GB budget. There is no experiment where memory is deliberately constrained (e.g., limiting GPU memory artificially) to test the Planner's behavior at the boundary. The paper does not report the Planner's output for hypothetical memory-constrained scenarios, nor does it provide a sensitivity analysis showing how MFU degrades as the memory budget shrinks.

Mitigation status. Not addressed. The paper acknowledges that AMSP "consistently exhibits high memory consumption" and that its memory footprint is "twice that of MiCS" for LLaMA-7B at 1024 GPUs (Section VI-B), but frames this as a strength — the memory is being deliberately traded for communication efficiency. The paper does not discuss what happens when this tradeoff is no longer available because of tighter memory constraints, nor does it propose mechanisms for the Planner to reason about the gradient of the memory-communication tradeoff near the constraint boundary.

No Convergence or Correctness Validation — Throughput Gains Are Assumed, Not Demonstrated, to Yield Training Speedup

The assumption or constraint. The paper evaluates AMSP purely on system throughput metrics — MFU and TGS — and never measures training convergence (loss curves, validation perplexity, or downstream task performance). The implicit assumption is that AMSP's sharding and communication strategies are numerically equivalent to baseline ZeRO implementations: changing which GPUs perform AllReduce, the order of floating-point reductions, or the timing of Broadcast operations relative to computation does not affect the model's training dynamics. This assumption is standard in systems papers that don't modify the optimization algorithm or model architecture, but it is not formally verified for AMSP.

There are specific reasons to suspect numerical differences could arise. AMSP changes the communication pattern for gradient aggregation: when sg > sp and the select & drop mechanism is used, gradients are AllReduced across a subgroup, then each GPU keeps only its local portion and discards the rest. This is functionally equivalent to ReduceScatter but uses a different sequence of floating-point operations, potentially changing the least-significant bits of aggregated gradients due to non-associativity of floating-point addition. Over many training steps, such differences can compound, causing models trained with AMSP to diverge from models trained with ZeRO-1 or MiCS, even if both achieve the same loss in expectation. The paper provides no evidence that this does not occur.

Similarly, AMSP's communication-computation overlap strategies change the order in which operations execute — parameters are prefetched, gradients are reduced asynchronously, and Broadcast operations are interleaved with forward computation. While these should be semantically equivalent to the non-overlapped execution, subtle interactions with PyTorch's autograd engine or CUDA stream scheduling could introduce non-determinism that affects reproducibility.

The consequence. A practitioner evaluating whether to adopt AMSP needs to know not just whether it increases throughput, but whether that throughput gain translates to faster time-to-convergence without sacrificing model quality. A 1.56× throughput improvement that comes with a 1.2× increase in steps-to-convergence (due to noisier gradients from different reduction ordering) would yield a much smaller actual speedup. Without convergence curves, the paper's headline improvements are throughput improvements, not training time improvements — the distinction matters for production deployments where the goal is reaching a target model quality as quickly as possible.

What evidence exists in the paper. None. The paper does not report loss curves, validation metrics, or any measure of model quality for any model trained with AMSP. There is no comparison of converged model performance (e.g., perplexity on a held-out set, zero-shot task accuracy) between AMSP and baselines. The paper mentions that AMSP has been "used for training InternLM on thousands of GPUs" (Section I), which provides some real-world validation but is not a controlled comparison — we don't know whether InternLM's training used AMSP exclusively or alongside other optimizations, or whether its final model quality was comparable to what would have been achieved with ZeRO-1 or MiCS given the same training budget.

Mitigation status. Not addressed. The paper does not acknowledge the absence of convergence validation as a limitation, nor does it propose future work to characterize the numerical fidelity of AMSP's communication strategies relative to baselines. This is a standard omission in systems papers, but it is particularly relevant here because AMSP introduces new communication patterns (select & drop, decoupled ReduceScatter lifecycle) that are not simply reimplementations of existing NCCL primitives.

The Revision Model Story: AMSP Assumes the Sharding Configuration Found at One Scale Is Optimal at All Scales

The assumption or constraint. The Planner produces a single sharding configuration that is used across all GPU counts in the end-to-end evaluation (Table IV). For LLaMA-7B, the configuration is sp = sg = 1, sos = 8 regardless of whether training on 8 GPUs or 1024 GPUs. For LLaMA-13B, it's sp = sg = 4, sos = 8 at all scales. This assumes that the sharding factors that minimize communication time at one scale (say, 1024 GPUs) also minimize it at other scales (say, 32 GPUs) — or at least, that the configuration is "good enough" across the range.

However, the communication time function T_comm depends on the profiled bandwidths w(o, v, p0 × p1), and the device meshes s0_i × s1_i appear both in the choice of sharding factors and in the total data-parallel group size sdp = s0_dp × s1_dp. When training on 32 GPUs, sdp = 32; when training on 1024 GPUs, sdp = 1024. The optimal sos (which must divide sdp) could reasonably differ between these scales — at 32 GPUs, setting sos = 8 means optimizer states span one-quarter of the cluster (one node), while at 1024 GPUs, sos = 8 means they span less than 1% of the cluster. The per-GPU optimizer state memory is 12Φ / 8 in both cases, but the communication cost for the AllReduce in T0_os involves sdp / sp GPUs, which changes from 32/1 = 32 to 1024/1 = 1024. The optimal tradeoff likely shifts with scale, but the paper never evaluates this.

The consequence. The paper reports the same sharding configuration at all GPU counts, which means either (a) the Planner was run only once for the largest scale (1024 GPUs) and the same configuration was applied at smaller scales without re-optimization, or (b) the Planner produces the same output at all scales. If (a), then the results at 8–512 GPUs may not represent the best AMSP can achieve — a different configuration might yield higher throughput at intermediate scales. If (b), the paper should demonstrate that the Planner's output is indeed scale-invariant for the tested models, which it does not. This is particularly relevant because the paper's "scalability" narrative (AMSP degrades less than baselines as GPUs increase) could be partially an artifact of using a configuration optimized for large scale at all scales — the baselines (ZeRO-3, MiCS) have fixed strategies that are similarly scale-invariant, so the comparison may still be fair, but the paper's own flexibility claim would be undermined if the Planner cannot adapt to different scales.

The paper does show that AMSP outperforms baselines at every GPU count (Figures 11–12), which suggests the configuration is not catastrophically wrong at any scale. But the claim that AMSP "discovers optimal sharding factors" implies per-scale optimization, and the paper provides no evidence that re-running the Planner at different GPU counts would (or would not) produce different outputs.

What evidence exists in the paper. Table IV lists the AMSP configurations as uniform across all GPU counts (8 to 1024). Section VI states: "AMSP maintains a uniform set of configurations when scaling training from 8 GPUs to 1024 GPUs." The paper does not explain whether the Planner was re-run at each scale, or whether the configuration was selected once at 1024 GPUs and applied universally. There is no figure showing how the Planner's output changes with sdp, no ablation comparing scale-specific vs. scale-uniform configurations, and no discussion of whether the optimization problem should be solved per-scale or once for the maximum scale.

Mitigation status. Not addressed. The paper treats the uniformity of configurations as a feature ("AMSP maintains a uniform set of configurations"), not a limitation. In practice, re-running the Planner at each GPU count would be straightforward (it's an offline computation), and the paper's silence on whether this was done suggests it was not — or that doing so didn't change the output. Either way, the lack of transparency about this design choice makes it impossible to assess whether AMSP's sharding strategy is truly scale-adaptive or just a fixed configuration that happens to work well across scales.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper causes a medium-magnitude reframing rather than a paradigm shift — it doesn't invent a new communication primitive or training algorithm, but it fundamentally changes how systems researchers and practitioners should think about the design space of distributed training optimizers. The shift is from treating model state sharding as a monolithic decision (ZeRO stage 1, 2, or 3; MiCS subgroup size) to treating it as a combinatorial optimization over a structured, six-dimensional space where each component of the model state has an independent sharding strategy and device mesh. This reframing matters because it reveals that the coupling of parameters, gradients, and optimizer states — a coupling that was so deeply embedded in prior systems that it was never questioned — was the bottleneck, not any particular communication operation or overlap deficiency.

The paper resolves a specific empirical puzzle that should have been a red flag for the field but was largely ignored: why MiCS, a system explicitly designed to reduce ZeRO's communication overhead, underperforms vanilla ZeRO-1 on LLaMA-7B at scale (35% vs. 36% MFU at 1024 GPUs, Figure 11a). The answer — that MiCS's forced sp > 1 triggers AllGather/ReduceScatter for parameters that don't need sharding, while ZeRO-1's sp = 1 avoids these operations entirely — is simple in retrospect but was inaccessible under the prior conceptual framework where "more sharding = less memory = more communication" was treated as a single monotonic tradeoff. AMSP shows that the tradeoff is per-component: you can shard optimizer states (saving 12Φ bytes of memory per shard factor) without sharding parameters (avoiding the most expensive and frequent communication), because the components have different memory footprints and communication frequencies. This insight is portable beyond AMSP — any future distributed training system should at minimum consider decoupling the sharding of these three components, even if it doesn't adopt the full optimization framework.

The paper also establishes a layered decomposition of the communication overhead problem that has methodological implications for how distributed training systems should be evaluated. The dissociation experiment (Figures 16–17) — showing that MiCS's sharding configuration on AMSP's execution engine achieves ~1.1–1.2× higher throughput than the same configuration on DeepSpeed's engine, and that the Planner's optimal configuration provides a further ~1.2–1.4× gain — demonstrates that sharding strategy quality and execution scheduling quality are independent dimensions that compound multiplicatively. Prior work had implicitly assumed these were a single "system performance" axis; by separating them, AMSP shows that a bad scheduler can mask the benefits of a good sharding strategy (as DeepSpeed-MiCS does), and that fixing only one leaves substantial gains on the table. This has a direct implication for future systems papers: ablation studies should separately report the contribution of sharding strategy changes and execution engine improvements, rather than reporting only the combined effect. The paper's Table V, which disables individual overlap optimizations in sequence, provides a template for how to do this for the execution engine dimension.

The profiling-based communication model (Section IV-A) represents a methodological escape from the α-β straightjacket that has constrained distributed systems analysis for decades. While profiling itself is not novel, the paper's argument for why it's necessary — NCCL's multi-algorithm selection, in-network aggregation effects, and the 4× effective bandwidth degradation at scale that no single-bandwidth parameter can capture — provides a clear, empirically-grounded justification that future work can cite when choosing empirical performance models over analytical ones. This is an engineering insight rather than a theoretical contribution, but it matters because the field's default assumption — that α-β with a single bandwidth parameter is "good enough" for system design decisions — has led to optimization choices that are systematically wrong at scale. The paper's Figure 3, showing how effective bandwidth varies across operations and scales, should become a standard exhibit in arguments for profiling-driven system design.

This work makes certain research directions more attractive:

Automated optimization for distributed training configuration. AMSP demonstrates that a constrained integer programming formulation with a profiling-based cost model can discover non-obvious configurations (like sp = 1, sos = 8 for 7B) that outperform human-designed heuristics. This opens the door for more sophisticated optimization approaches — continuous relaxations, Bayesian optimization, or learned cost models — that could handle the combined search space of model state sharding, operator-level parallelism, and activation recomputation strategies. The paper's Pre-Filter and constraint-based search space reduction provide a blueprint for making such joint optimization tractable.
Verifier and reward model quality as the binding constraint on test-time compute scaling. The difficulty-dependent behavior the paper documents for model state sharding — different model sizes and GPU counts require qualitatively different configurations — parallels the difficulty-dependent scaling behavior in the test-time compute literature (Snell et al., 2024). Just as easy problems benefit from sequential revision while hard problems need parallel exploration, small models on large GPU clusters benefit from aggressive optimizer state sharding with parameter replication (sp = 1, sos > 1), while larger models need more balanced sharding. This suggests a unified framework where "difficulty" could mean model size relative to GPU memory, problem complexity relative to model capability, or any resource-constrained optimization scenario.
Designing robust overlap scheduling as a first-class systems contribution. The paper's trace analysis (Figures 14–15) revealing "computation resource bubbles" in DeepSpeed's execution engine — where ReduceScatter blocks both computation and AllReduce — establishes that overlap scheduling is a non-trivial engineering challenge that deserves rigorous treatment, not an afterthought. The specific techniques AMSP uses (module-level prefetching, decoupling ReduceScatter lifecycle from backward functions, bucketing with asynchronous AllReduce, Broadcast ordering alignment) provide a catalog of patterns that future systems can adopt or refine.

Conversely, some directions become less attractive:

Further incremental variations of the ZeRO stage hierarchy. The paper shows that the sp = sg = sos coupling at the heart of ZeRO-1/2/3 is fundamentally limiting — any system that maintains this coupling, regardless of how it tunes the subgroup size or adds compression, will leave performance on the table. ZeRO-4 or ZeRO-5 variants that introduce new stages within the same coupled framework are unlikely to yield meaningful improvements beyond what per-component sharding already achieves.
Pure communication compression without topology-aware sharding. ZeRO++'s quantization approach — reducing message sizes without changing the communication topology — achieves only 4–6% MFU on 1024 GPUs in the paper's benchmarks, compared to AMSP's 42–52%. While quantization helps, the 4× effective bandwidth degradation at scale (Figure 3) means that even if messages were compressed by 4×, the AllGather across 1024 GPUs would still be bottlenecked by the latency and bandwidth of coordinating that many participants. Compression addresses the v/w term in the communication cost but not the (p - 1)α term, which dominates at large p. Future work on communication optimization should prioritize reducing the number of participants in each collective (as AMSP does) over reducing the per-message size.
Manually-tuned sharding configurations. The paper's Planner produces different configurations for 7B, 13B, and 30B models (Table IV) that would be difficult to derive manually — sp = 4 for 13B rather than sp = 2 or sp = 8 is not obvious from first principles. As model architectures diversify and cluster topologies become more heterogeneous, manual tuning will become increasingly infeasible. The paper establishes that automated optimization is both necessary and achievable.

Follow-Up Research This Work Enables

Joint optimization of model state sharding and operator-level parallelism. AMSP optimizes only the model state sharding strategy, assuming pure data parallelism (stp = spp = 1). The paper explicitly states (Section VII) that automated parallelism systems like Alpa, FlexFlow, and TensorOpt focus on operator-level parallelization (how to split individual matrix multiplications and attention operations) but "overlook strategies related to the orthogonal placement of the model states." The natural extension is a unified optimization that jointly searches over sharding factors (sp, sg, sos), tensor parallelism degree stp, pipeline parallelism degree spp, and the mapping of model layers to pipeline stages, all under a single memory and communication cost model. This is challenging because the search space is the product of these dimensions, but AMSP's Pre-Filter and profiling-based cost model provide the scaffolding. A concrete experiment: for LLaMA-65B on 512 GPUs, compare a configuration where AMSP's Planner first selects sharding factors assuming stp = spp = 1, then Alpa selects operator parallelism on top of that, against a joint optimizer that considers all dimensions simultaneously. The hypothesis is that joint optimization finds configurations with better compute-communication balance — for example, using a small amount of tensor parallelism within nodes to reduce the per-GPU parameter memory, allowing sp = 1 for the remaining parameters and avoiding cross-node AllGather entirely.

Online adaptation of sharding factors during training. The Planner selects a static sharding configuration before training begins, but the optimal configuration may depend on factors that change during training — network congestion from co-located jobs, GPU memory pressure from activation spikes at certain sequence lengths, or varying micro-batch sizes if dynamic batching is used. A natural extension is an online controller that monitors actual communication times (which may diverge from profiled predictions due to contention) and memory utilization, and dynamically adjusts sharding factors between training steps. This would require efficient mechanisms for resharding model states on the fly — moving from sp = 4 to sp = 8 without pausing training — which NCCL does not natively support for arbitrary tensor distributions. A concrete experiment: on a shared cluster with varying network congestion, compare AMSP's static Planner configuration against an online variant that measures AllGather and AllReduce completion times over a sliding window of steps and re-runs the Planner with updated bandwidth estimates every 100 steps, triggering resharding if the predicted throughput gain exceeds the resharding cost. The key metric is end-to-end training time, including resharding overhead, under realistic contention patterns.

Stress-testing the dependency rule with emerging optimizer designs. AMSP's dependency rule (s_dp ≥ s_os ≥ s_g ≥ s_p) is derived from the standard training loop where parameters produce gradients that update optimizer states. But emerging optimizer designs complicate this: Lion (which uses sign-based updates with only momentum, no second moment), Adafactor (which factorizes the second moment to save memory), or 8-bit optimizers (which quantize optimizer states) change the memory footprint ratios (12Φ for Adam may become 4Φ or 2Φ for other optimizers) and the dependency structure. Does the dependency rule still hold? For 8-bit Adam, where the master parameters are in FP16 rather than FP32, the optimizer state memory term changes from 12Φ/sos to 6Φ/sos, potentially shifting the memory-communication tradeoff and changing the Planner's optimal configuration. A concrete experiment: implement two additional optimizer backends in AMSP (8-bit Adam and Lion), profile their memory footprints and communication patterns, and run the Planner for LLaMA-7B/13B/30B under each optimizer. The hypothesis is that 8-bit Adam's lower memory pressure allows less aggressive sharding (e.g., sos = 4 instead of sos = 8 for 13B), which reduces Broadcast overhead and further improves throughput. If the dependency rule needs modification — e.g., if 8-bit Adam's master parameters introduce a new dependency between sp and sos — the finding would refine AMSP's theoretical framework.

Extending to mixture-of-experts architectures. All experiments in the paper use dense Transformer models. Mixture-of-Experts (MoE) models have a fundamentally different parameter structure: expert layers have many more parameters (each expert is a full feed-forward network) but each token only activates a subset of experts, creating an additional sharding dimension (expert parallelism). How should AMSP's per-component sharding interact with expert placement? Expert parameters might benefit from Full-Replica within a node (since only a fraction of experts are used per batch) while shared parameters (attention layers) might use Full-Sharding to save memory. A concrete experiment: implement MoE training in AMSP with expert parallelism as an additional dimension, model the communication cost of AllToAll for token routing and expert AllGather/ReduceScatter, and extend the optimization problem to jointly select expert sharding factors alongside (sp, sg, sos). Evaluate on a model like Mixtral 8×7B (8 experts, each 7B parameters in the feed-forward layers). The hypothesis is that AMSP's per-component flexibility is even more valuable for MoE models because expert parameters, shared parameters, and optimizer states all have different access patterns and memory footprints, and the optimal configuration decouples them further than the current Planner allows.

Characterizing and mitigating the numerical fidelity gap. The paper provides no evidence that models trained with AMSP achieve the same loss or accuracy as baselines. The select & drop mechanism, decoupled ReduceScatter lifecycle, and different reduction ordering could introduce numerical differences that compound over training. A concrete experiment: train LLaMA-7B on a fixed dataset (e.g., C4) for 10,000 steps using both AMSP and ZeRO-1 on 64 GPUs, with identical random seeds, and compare loss curves, gradient norms, and final perplexity. Measure the bitwise difference in parameter updates at each step and track whether it grows (divergence) or remains bounded (numerical noise). If divergence occurs, investigate whether the select & drop mechanism (where GPUs discard gradients they don't need after AllReduce) introduces bias by using a different reduction tree than standard AllReduce, and propose a correction (e.g., using ReduceScatter instead of AllReduce + select & drop for the sg > sp case). This experiment would transform AMSP from a throughput benchmark into a validated training system, and the results — whether positive (no divergence) or negative (measurable divergence) — would be directly useful to practitioners.

Portability across hardware generations and cloud instances. The paper's profiling-based communication model is, by design, tied to the specific cluster it was measured on (A800 GPUs, Mellanox HDR InfiniBand without SHARP). A critical question for adoption is whether the Planner's decisions generalize across hardware configurations without requiring full re-profiling. A concrete experiment: deploy AMSP on three distinct hardware configurations — (1) the paper's original A800 + HDR InfiniBand cluster, (2) a cloud instance with A100 GPUs and NVSwitch (e.g., Azure ND A100 v4), and (3) a cluster with H100 GPUs and SHARP-enabled InfiniBand. Run the Planner's profiling on all three, then compare the optimal sharding configurations for LLaMA-7B/13B/30B. Key question: does the Planner produce different configurations on different hardware, and do those differences correctly reflect the hardware characteristics (e.g., less aggressive parameter sharding on NVSwitch clusters because intra-node AllGather is cheaper)? If the configurations differ, does AMSP's throughput advantage over MiCS and ZeRO-1 persist across all hardware, or does it shrink on SHARP-enabled clusters where baseline AllReduce scales better? This experiment directly tests the paper's generality claim and would produce actionable guidance for practitioners choosing between AMSP and alternatives on their specific hardware.

Practical Applications and Downstream Use Cases

Training 7B–30B models at scale on commodity GPU clusters without SHARP. The paper's results are most directly applicable to organizations training models in the LLaMA-7B to LLaMA-30B range on GPU clusters with standard InfiniBand interconnects (no in-network aggregation). The quantified benefit at 1024 GPUs — 51% MFU on 7B vs. 35% for MiCS, 52% on 13B vs. 33%, 42% on 30B vs. 29% (Figures 11–12) — translates to a 1.4–1.56× reduction in training time for the same model and dataset. For a hypothetical 30-day training run on 1024 GPUs, this means completing in 19–21 days instead of 30, or training a proportionally larger model in the same wall-clock time. The configuration is available immediately via the InternEvo codebase, and the Planner automates sharding factor selection, removing the need for manual tuning. The primary deployment requirement is running the profiling suite once on the target cluster, which the paper does not quantify but which likely takes hours to tens of hours of dedicated cluster time — acceptable for a one-time cost before a months-long training campaign.

Maximizing throughput for fixed-budget training of mid-size models on smaller GPU pools. Many research labs and startups have access to 64–256 GPUs rather than thousands. At 128 GPUs, AMSP achieves approximately 60% MFU on LLaMA-7B (Figure 11a) vs. ~52% for MiCS and ~54% for ZeRO-1. The benefit at this scale is smaller (~1.1–1.15×) but still meaningful for long-running training jobs. More importantly, the Planner can be run on the specific GPU count available — the paper's use of a uniform configuration across scales (Table IV) likely understates AMSP's advantage at intermediate scales, since the configuration was optimized for 1024 GPUs. A practitioner with exactly 128 GPUs would run the Planner with sdp = 128 and likely obtain a configuration better tailored to that scale. The memory analysis (Figure 13) also shows that AMSP's higher memory consumption (~60 GB on 7B at 1024 GPUs) is well within 80 GB GPU limits, meaning the throughput gain comes with no OOM risk. For models at the boundary of memory capacity (e.g., fine-tuning LLaMA-30B on 40 GB GPUs), the Planner could be configured with the tighter memory constraint, and would automatically select more aggressive sharding to fit, though the throughput advantage over MiCS would likely shrink.

Serving as the data-parallel backend for hybrid parallelism frameworks. AMSP's focus on pure data parallelism (stp = spp = 1) may seem limiting, but in practice, data parallelism is the outer loop of most hybrid parallel training setups. In a typical Megatron-LM configuration, tensor parallelism is applied within a node (8 GPUs), pipeline parallelism across a small number of nodes (2–4), and data parallelism across the remaining dimensions. AMSP's Planner could be used to optimize the data-parallel sharding configuration on top of whatever tensor and pipeline parallelism scheme the user has chosen, with sdp set to the data-parallel group size rather than the total GPU count. The paper's communication cost model (Section IV-C) already handles this case — the device meshes for AllGather, ReduceScatter, and AllReduce operations are parameterized by s0_p, s1_p, etc., and would naturally account for the data-parallel subgroup within a larger hybrid parallelism topology. This makes AMSP incrementally adoptable: a team already using Megatron-LM or DeepSpeed with 3D parallelism could replace only the data-parallel sharding component with AMSP's Planner and Executor, keeping their existing tensor and pipeline parallelism setup, and expect throughput improvements proportional to the data-parallel group size (larger gains when data parallelism is the bottleneck).

Training with FP8 optimizers where inter-tensor sharding is required. The paper notes that the inter-tensor approach to optimizer state sharding (distributing whole tensors rather than splitting within tensors) is "recommended for FP8 training" because FP8 requires per-tensor scaling factors. Since AMSP already uses inter-tensor sharding with Broadcast for parameter distribution, it is compatible with FP8 optimizers without modification. As the industry transitions to FP8 for training efficiency (following the FP8-LM paper cited by the authors), systems that rely on intra-tensor sharding and AllGather for parameter distribution will need to be redesigned for the inter-tensor case, while AMSP already handles it. A practitioner adopting FP8 training can run the same Planner — the optimizer state memory term 12Φ/sos would change to reflect FP8's memory footprint, and the Planner would adapt accordingly, potentially selecting less aggressive optimizer state sharding because FP8 optimizer states are smaller. The communication patterns (Broadcast with inter-tensor sharding, as in T1_os) remain identical, so the Executor's overlap strategies apply without modification.

Dependency tree

Papers this one builds on, walking backwards through the same domain and shared topics.

ZeRO: Memory Optimizations Toward Training Trillion Parameter Models2019· distributed-training← LLM-judged similar
FP8-LM: Training FP8 Large Language Models2023· training-methods← in-domain predecessor
- VeRA: Vector-based Random Matrix Adaptation2023· training-methods← in-domain predecessor
How Far Can Unsupervised RLVR Scale LLM Training?2026· training-methods← similar topic
- In-Context Reinforcement Learning for Tool Use in Large Language Models2026· training-methods← in-domain predecessor
- UniGRPO: Unified Policy Optimization for Reasoning-Driven Visual Generation2026· rl-training← topic thread
- Listwise Policy Optimization: Group-based RLVR as Target-Projection on the LLM Response Simplex2026· training-methods← similar topic
- Quartet: Native FP4 Training Can Be Optimal for Large Language Models2025· training-methods← similar topic
GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection2024· training-methods← similar topic
- Training-Free Long-Context Scaling of Large Language Models2024· training-methods← in-domain predecessor
- Quartet II: Accurate LLM Pre-Training in NVFP4 by Improved Unbiased Gradient Estimation2026· training-methods← similar topic
- Memory-Efficient LLM Training with Online Subspace Descent2024· training-methods← similar topic
- Read More2023· training-methods← similar topic
Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift2015· training-methods← similar topic
- Adam: A Method for Stochastic Optimization2014· training-methods← in-domain predecessor
- Critical Batch Size Revisited: A Simple Empirical Approach to Large-Batch Language Model Training2025· training-methods← similar topic

What it influenced

Later papers in the same chain. Direction reversed.

Parameter-Efficient Orthogonal Finetuning via Butterfly Factorization2023· training-methods→ in-domain predecessor
- DiLoCo: Distributed Low-Communication Training of Language Models2023· training-methods→ in-domain predecessor
Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift2015· training-methods→ similar topic
- On the importance of initialization and momentum in deep learning2015· training-methods→ in-domain predecessor
- Understanding the difficulty of training deep feedforward neural networks2000· training-methods→ similar topic
- BinaryConnect: Training Deep Neural Networks with Binary Weights during Propagations2015· training-methods→ similar topic
GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection2024· training-methods→ similar topic
- LoHan: Low-Cost High-Performance Framework to Fine-Tune 100B Model on a Consumer GPU2024· training-methods→ in-domain predecessor
- Read More2023· training-methods→ similar topic
- Memory-Efficient LLM Training with Online Subspace Descent2024· training-methods→ similar topic