COPUS: Co-adaptive Parallelism and Batch Size Selection in Large Language Model Training

Abstract: Training LLMs requires jointly configuring two interdependent aspects of the system: the global batch size, which governs statistical efficiency, and the 3D parallelism strategy, which governs hardware throughput. Existing approaches make these decisions independently: optimization work adapts the batch size to track the evolving critical batch size while keeping parallelism fixed, and systems work selects the fastest parallelism for a given fixed batch size without anticipating that the optimal batch size could change. We show that these decisions are tightly coupled: the throughput-optimal parallelism strategy may shift as the global batch size changes, so any method that fixes one while adapting the other operates with a suboptimal configuration for part of the training run. We present COPUS, a system that adaptively tunes the global batch size, parallelism strategy, and micro-batch size as training evolves. COPUS is guided by Goodput, the product of throughput and statistical efficiency, which models both hardware and statistical effects jointly and directly measures useful convergence per unit of wall-clock time. The system combines online gradient noise scale estimation under 3D parallelism with throughput-aware evaluation of candidate configurations, and supports efficient reconfiguration of both batch size and parallelism during training. We evaluate COPUS on LLM pre-training workloads across 1-4 nodes of 8xH100 and 8xMI210 GPUs and model sizes from 3B to 32B parameters, demonstrating average time-to-convergence speedups of 3.9-8.0% over the fastest baseline across four configurations, with peak gains up to 11.1%, including system overheads.

• The paper introduces a joint optimization framework that co-adapts parallelism and batch size to maximize Goodput by integrating hardware throughput with statistical efficiency.
• It details an online reconfiguration mechanism that adjusts 3D parallelism and microbatch settings in process, significantly reducing switching overhead compared to checkpoint-restart methods.
• Empirical evaluations show up to 11.1% faster convergence over static and CBS-only methods, underscoring the benefits of dynamic system tuning in LLM pre-training.

Co-Adaptive Parallelism and Batch Size Selection in LLM Training

Motivation and Problem Formulation

The scaling of LLM pre-training highlights critical limitations in current distributed deep learning practices, where global batch size and parallelism strategy (data, tensor, pipeline) are configured independently. Empirical evidence indicates the hardware throughput-optimal parallel strategy depends on the batch size, and conversely, the optimal batch size evolves due to nonstationary gradient noise scale (GNS) throughout training. When these configuration dimensions are optimized in isolation, significant system and statistical performance is forfeited.

Industry-standard stacks (e.g., Megatron-LM, DeepSpeed) require static 3D parallel mesh configuration at job startup, preventing dynamic adaptation as batch size transitions through regimes during training—a salient inefficiency observed in OpenAI, Meta, and AI2 published pre-training records. Simultaneously, adaptive batch size protocols (e.g., GNS or CBS-based) tune global batch size to maintain statistical efficiency but assume a fixed parallelism configuration. The paper formalizes this as a joint optimization over three axes: parallelism strategy $S = (d, t, p)$, global batch size $B_g$, and microbatch size $B_m$.

The Goodput Objective

To address the above coupling, the paper proposes maximizing a Goodput metric defined as the product of hardware throughput and statistical efficiency. For a candidate configuration $(S, B_g, B_m)$ on hardware $H$,

$$\mathrm{Goodput}(S, B_g, B_m, H) = \mathrm{T}(S, B_g, B_m, H) \times SE(B_g)$$

where $\mathrm{T}$ denotes throughput and $SE$ statistical efficiency derived from GNS. Crucially, for Adam-based optimization with square-root learning rate scaling, an additional $\sqrt{B_g}$ factor is introduced:

$$\mathrm{Goodput}^{LR}(S, B_g, B_m, H) \propto \mathrm{T}(S, B_g, B_m, H) \times SE(B_g) \times \sqrt{B_g}$$

This joint metric forces the co-optimization of batch size, microbatch decomposition, and parallelism axes, directly maximizing convergence per wall-clock second. This is in contrast to CBS-only schemes where statistical and system dimensions are decoupled and to parallelism planners that ignore optimization dynamics.

System Architecture and Online Estimation

The Copus system comprises a 3D-parallelism-aware GNS estimator, a Goodput orchestrator, and an in-process trainer with support for online reconfiguration.

Figure 1: Copus system architecture: Goodput-based configuration selection via the orchestrator, informed by offline throughput profiling and online GNS measurements.

GNS (gradient noise scale) estimation is performed in a manner compatible with arbitrary data, tensor, and pipeline partitioning—using per-microbatch gradients and distributed reductions to compute the required signal and noise statistics without reconstructing full gradients on any single device.

Figure 2: Gradient noise scale estimation under 3D parallelism leverages microbatch and all-reduced gradients to provide the necessary signal and noise statistics for Goodput evaluation.

Copus's orchestrator periodically considers all memory-feasible configurations (from an offline throughput profile), computes Goodput for each candidate, and triggers the requisite system change (batch size or parallelism) if the projected improvement exceeds a specified switching margin.

Empirical Analysis of Throughput–Batch Coupling

Benchmarks across four hardware topologies and model sizes show that the throughput-optimal parallelism can shift abruptly as batch size increases due to non-uniform network interconnect speeds and the competing effects of pipeline, tensor, and data parallelism.

Figure 3: Throughput–batch size relationship for all feasible 3D parallelism strategies, illustrating that the optimal configuration shifts nonmonotonically with $B_g$0.

This establishes that any method which fixes parallelism at initialization is categorically suboptimal for realistic LLM pre-training runs, especially as the training traverses batch size regimes.

Online Reconfiguration Mechanisms

When a favorable configuration change is identified, Copus supports efficient in-process reconfiguration. For batch size or microbatch size updates, the procedure is immediate; for parallelism changes, Copus performs state resharding in host memory, avoiding the high-latency checkpoint-restart approach, and reconstructs model and optimizer shards for the target mesh. This yields $B_g$1–$B_g$2 lower switching overhead compared to checkpoint-restart.

Figure 4: Copus online reconfiguration circumvents disk I/O by staging state in host memory, significantly reducing latency compared to standard checkpoint-restart.

Training Trajectory and Comparative Evaluation

Ablation studies show that Copus attains lower loss faster than state-of-the-art CBS-style and static-parallelism baselines across all settings. The principal gains are realized when Copus transitions parallelism configuration at batch size milestones, maintaining regime-optimal throughput, while CBS baselines suffer regime mismatch and underutilization.

Figure 5: Copus training trajectory showing coordinated adaptation of batch size and parallelism to consistently minimize loss and wall-clock time.

In four model/hardware settings (3B–32B parameters, 8–32 GPUs on NVIDIA/AMD), time-to-target-loss speedups average $B_g$3–$B_g$4 over the fastest baseline including reconfiguration overheads, reaching peaks of $B_g$5. Relative benefit grows with increasing cluster scale and heterogeneity in interconnect topology.

Figure 6: Loss vs. training time shows Copus consistently achieves faster convergence across all settings compared to CBS and static baselines.

Comparisons based on processed tokens (wall-clock-invariant) show that Copus matches CBS's sample efficiency but surpasses it in real time due to superior throughput management. Decision-space Goodput decomposition demonstrates that Copus uniquely maintains both high system throughput and statistical efficiency throughout.

Figure 7: Goodput decomposition over training time for the 3B configuration, with Copus maintaining high values by co-adapting configuration.

Theoretical and Practical Implications

The central insight is that the selection of optimal parallelism and batch regime is inherently nonstationary and must be dynamically co-optimized throughout training. The use of a joint Goodput metric offers robustness to statistical estimation noise and calibration error typically endemic to GNS-based methods.

From a systems perspective, these results imply that future LLM training platforms should prioritize automated, on-the-fly configuration search over manual scheduler tuning and static launch-time mesh selection. Algorithmically, the Goodput maximization approach accommodates future directions such as optimizer-aware adaptation, data mixture curriculum learning, and joint tuning of additional axes such as memory sharding or activation checkpointing.

Future Prospects and Limitations

Copus currently optimizes over data/tensor/pipeline parallelism, batch size, and microbatch size dimensions. Integrating further axes (e.g., ZeRO, context parallel, sequence parallel) and replacing offline throughput profiles with online measurement or predictive modeling would yield a fully autonomous reconfiguration stack. Automatic GNS calibration, or more advanced online sample efficiency estimators, would further reduce reliance on empirical tuning.

A limitation is that Copus requires an offline throughput profile per model–hardware pair and smooth GNS estimation (which remains noisy and requires moderate smoothing hyperparameters). At supercomputing scale (beyond 32B parameters and 32+ GPUs), further engineering may be required for constant overhead and efficiency.

Conclusion

Copus demonstrates that optimal LLM training demands joint, dynamic tuning of parallelism strategy and batch scheduling, not only for statistical efficiency but to exploit hardware capability as batch size evolves. The system realizes consistent and measurable reductions in time-to-convergence relative to best prior baselines, confirming that static parallelism is suboptimal even when paired with adaptive batch size methods. The Goodput-guided paradigm has significant implications for the design of next-generation scalable foundation model training software and motivates further expansion to other reconfigurable dimensions.

Reference: "COPUS: Co-adaptive Parallelism and Batch Size Selection in LLM Training" (2604.26687)