Source-linked AI summary
Beyond the Trial-and-Error Loop: Hybrid Projection and Automated Tuning for Distributed Training
Anshu Raina, Peyman Razaghi, Yuankai Chen, Cheng Yao, Yao Fu, Vidushi Goyal, Devang Patel, Wen Xie, Zhenyu Gu, Emad Barsoum
TL;DR
Distributed LLM training configuration requires costly exploration across interacting choices, and existing tools do not close that search loop. Primus Projection combines kernel measurements and analytical reconstruction, and its tuning agent uses the resulting estimates to search configurations. The projections remain close to measured throughput, and the agent reports a higher-throughput Mixtral recipe than the published reference.
Problem
Distributed LLM training has a rapidly growing configuration space, making exhaustive trial and error infeasible and multi-node iterations costly.
Source: §1 · p. 1Method
Primus Projection benchmarks kernels on a small GPU budget and analytically reconstructs out-of-scope communication and pipeline behavior, while the Primus Tuning Agent searches configurations using its projection oracle.
Source: §1 · p. 2Results
+27% throughput over AMD’s published 4-node BF16 Mixtral reference was found in under 30 minutes, while validated projections stayed within ∼10% of measured throughput.
Source: §6 · p. 10Takeaways & Limitations
The framework supports multi-node configuration projection and tuning through single-node exploration without requiring multi-node hardware.
Source: §6 · p. 10The framework has been validated on AMD Instinct GPUs and ROCm, but not yet empirically validated on NVIDIA or TPU; network contention is not modeled.
Source: §4.12 · p. 9
Abstract
from arXiv · showhide
Configuring distributed LLM training across hundreds of GPUs requires jointly optimizing tensor, pipeline, expert, context, and data parallelism alongside micro-batch size, recomputation, sharding, and overlap, a combinatorial space where memory legality can hinge on a single byte of activation accounting. Today, teams navigate it through expensive multi-node trial and error, burning GPU-hours on crashes and suboptimal recipes. We present Primus Projection, which closes this loop with two coupled components. First, a hybrid projection engine that measures what it can and simulates what it cannot: GPU microbenchmarks on as few as a single device capture compute kernels and any in-scope communication; a sub-node downscale-measure-upscale workflow analytically restores out-of-scope collectives, FSDP overlap, and the full pipeline schedule at the multi-node target. A pure simulation mode runs entirely on CPU, enabling pre-silicon planning without any GPU. Second, the Primus Tuning Agent treats the projection engine as a fast scoring oracle (seconds per candidate instead of tens of minutes) and uses it to drive an intelligent search over the configuration space. A deterministic seed planner sweeps single axes; an LLM investigation loop proposes cross-axis combinations the planner cannot reach, with an analytic memory pre-filter that rejects infeasible candidates before any GPU time is spent. Validated on Llama 3.1 (dense) and Mixtral 8x22B (MoE) across two GPU generations, all projections stay within ~10% of measured multi-node throughput, and fidelity holds at production scale: 4.7% error on DeepSeek V3 671B at 1,024 GPUs and 1.9% on Qwen3-235B-A22B at 256 GPUs. In a Mixtral 8x22B case study, the Primus Tuning Agent discovers a configuration delivering +27% throughput over the published 4-node BF16 reference at the same scale, in under 30 minutes of single-node exploration.
1 Introduction
Distributed LLM training configuration spans many interacting choices, making exhaustive trial and error infeasible. Primus Projection combines measured kernels with analytical reconstruction, while the Primus Tuning Agent searches configurations using the projection as a fast scoring oracle.
1 Introduction
- Teams face an expanding configuration space in which multi-node trial and error consumes hours per iteration, while existing modeling tools miss important workload effects or require extensive calibration.The passage identifies frequency throttling, software-stack variation, and MoE grouped-GEMM behavior among the effects that complicate existing approaches.Source: §1 · p. 1
- Projections stay within ∼10% of measured throughput across validated dense and MoE recipes, with 4.7% error at 1,024 GPUs and 1.9% at 256 GPUs.In the Mixtral case study, the agent found +27% throughput over AMD’s published 4-node BF16 reference in under 30 minutes.Source: §1 · p. 2
- The hybrid engine benchmarks compute on as few as one GPU and analytically restores out-of-scope communication and the target pipeline schedule, with a CPU-only simulation option.The tuning agent uses the projection oracle to explore configurations; its Mixtral case study reports +27% throughput over the published 4-node BF16 reference in under 30 minutes.Source: §1 · p. 2
2 Background and Motivation
Distributed training configuration combines many interacting choices, while manual exploration and purely analytical performance models have practical and accuracy limits. Primus Projection pairs deployed-stack kernel measurements with analytical models to project multi-node configurations.
2 Background and Motivation
- A Mixtral 8×22B run on 64 GPUs has roughly 160 ordered parallelism factorizations and an estimated O(10^4) topologically valid configurations after additional choices are included.Only a small fraction are both memory-legal and high-throughput.Source: §2 · p. 2
- A 4-byte activation-memory miscount per element across 56 MoE layers can turn a legal configuration into an out-of-memory run, wasting queued cluster time.The passage describes the feasible boundary as razor-thin.Source: §2 · p. 2
- Kernel benchmarks capture deployed-stack effects that analytical performance models can miss, while analytical communication, pipeline, and FSDP models transfer measurements to multi-node targets.The passage cites throttling, software-stack changes, MoE grouped GEMMs, and optimizer or allocator overhead as modeling challenges.Source: §2 · p. 2
3 Methodology
Primus Projection combines analytical memory and communication models, layer profiling, pipeline simulation, and multi-node scaling in a projection oracle. The Primus Tuning Agent uses that oracle to search for high-throughput, memory-legal configurations.
3 Methodology
- The modeling pipeline combines analytical memory projection, layer profiling, communication modeling, pipeline schedule simulation, and multi-node scaling into a projection oracle.The oracle takes a model card, hardware profile, and cluster specification and returns memory and throughput estimates.Source: §3 · p. 2, Figure 1 · p. 3
3.1 Modeling
- Memory projection tracks sharded parameters, gradients, optimizer states, and per-microbatch activations across parallelism, precision, sharding, schedule, and recomputation choices.Full recomputation reduces each layer’s activation footprint to a checkpointed input.Source: §3.1 · p. 2
- Per-iteration time combines the simulated pipeline makespan with exposed communication and optimizer-step costs.In-scope communication can be captured during benchmarking, while communication beyond benchmark scope is reconstructed analytically.Source: §3.1 · p. 3
- The downscale-measure-upscale workflow reduces PP, then EP, then TP to fit the benchmark budget, preserving per-rank MoE workload when reducing EP and never reducing CP.PP is reduced first because its schedule effect can be simulated; TP is last because reducing it changes tensor slice shapes.Source: §3.1 · pp. 3–4
- Layer profiling times forward and backward for each distinct layer type, decomposes MoE compute and All-to-All, and feeds per-stage timings to the pipeline simulator.With PP=1, per-type times scale by layer count; with PP>1, the simulator accounts for layer assignment and pipeline bubbles.Source: §3.1 · p. 4
- The upscale stage restores pipeline schedules by simulation, adds analytical All-to-All deltas for EP, and rescales compute and AllReduce for TP.The restoration order PP → EP → TP reflects decreasing fidelity, with TP relying on a compute-scaling assumption.Source: §3.1 · p. 4
- The projector offers CPU-only simulation for planning and models communication across intra-node, inter-node, and inter-pod topology domains.Its communication and pipeline models select among applicable collective algorithms and scheduling strategies for the target configuration.Source: §3.1 · pp. 4–5
3.2 Primus Tuning Agent
- The Primus Tuning Agent searches configurations across parallelism, schedules, micro-batch size, recomputation, precision, MoE backend, and overlap, scoring candidates with the projection oracle.A deterministic planner seeds single-axis sweeps, then an LLM loop proposes cross-axis combinations; an analytic memory pre-filter rejects infeasible candidates before GPU time is spent.Source: §3.2 · p. 5
4 Evaluation
Primus Projection’s evaluations show close agreement with measured throughput, including at production scale, while the Tuning Agent finds faster Mixtral configurations than the published reference. The evaluation also examines communication and pipeline models, memory accounting, ranking reliability, MoE load balance, and scope boundaries.
4.2 Communication Model Correlation; 4.3 Pipeline Schedule Algorithm Comparison
- For training workloads with messages ≥64 MB, the communication model averages 4.2% error; Zero-bubble scheduling reaches +13% throughput over 1F1B with 2% memory overhead.Zero-bubble schedules weight gradients into pipeline bubbles; small-message latency-dominated cases are outside the training regime.Source: §4.2 · p. 6, §4.3 · p. 6
4.4 Memory Projection: Mixtral 4-Node Reference
- 198 GB of activation memory exceeds Mixtral’s parameter-plus-optimizer budget by over 2× in the evaluated no-recompute configuration, motivating memory-aware search.At EP=8, one expert MLP is retained per rank, and no recomputation keeps microbatch activations live simultaneously.Source: §4.4 · p. 6, Table 5 · p. 6
4.5 End-to-End Analysis
- All projections stay within ∼10% of measured throughput across benchmark GPU counts, model families, precision modes, GPU generations, and tested parallelism configurations.All 16 cells in Figure 2 are within ∼10% of measured throughput.Source: Figure 2 · p. 6, §4.5 · p. 7
4.6 Large-Scale Validation (256 and 1,024 GPUs)
- 4.7% error on DeepSeek V3 671B at 1,024 GPUs and 1.9% on Qwen3-235B-A22B at 256 GPUs show fidelity at production scale.Both recipes were selected end-to-end by the agent and confirmed on hardware.Source: §4.6 · p. 7, Table 7 · p. 7
4.7 Primus Tuning Agent Case Study: Mixtral 8×22B
- +27% measured throughput over AMD’s 4-node BF16 reference was achieved on Mixtral 8×22B; the workflow rejected 7 of 30 trials before GPU benchmarking.The top two non-redundant LLM proposals measured 4,402 and 4,377 tok/s/GPU, versus the 3,475 tok/s/GPU reference.Source: §4.7 · pp. 7–8, Table 9 · p. 8
4.8 Rank Preservation and the Calibration Offset
- Selection regret was 0% across three production MoE models, while the comparison with Calculon found 6.6% average absolute error for Primus versus 28.9%.Calculon’s opposite-direction errors inflated the measured 1.5× FP8–BF16 gap to 2.9×; Primus’s benchmark oracle and agent framework also outperformed the tested alternatives on Mixtral.Source: §4.8 · p. 8, §4.9 · pp. 8–9, Table 10 · p. 8
4.10 MoE Load-Balancing Robustness
- A busiest-rank correction cuts worst-case MoE load-balance error from 53–71% to under 10%; on trained DeepSeek V3, corrected error is ≈2% at EP=8.The real-model step-level pooled imbalance is 1.05–1.12 despite higher per-layer medians.Source: §4.10 · p. 9
4.12 Scope and Limitations
- Primus Projection has been validated on AMD Instinct GPUs and ROCm, but not empirically on NVIDIA or TPU; network contention is not modeled.ZeRO-2 is excluded, and FP8 recipes require convergence validation before production.Source: §4.12 · p. 9
5 Related Work
Prior approaches use analytical models, target-cluster profiling, or trace-driven simulation, each with limitations for broad configuration exploration. Primus Projection combines benchmark-anchored projection with an LLM-driven search over cross-axis configurations.
Analytical performance models
- Analytical models and profiling tools search configurations or characterize target systems, but do not capture all workload effects or avoid target-cluster profiling.Closed-form models miss clock throttling, kernel-library drift, and MoE grouped-GEMM wave quantization; FlexFlow and Alpa require target-cluster profiling.Source: §5 · p. 10
Trace-driven and simulation-based approaches
- Trace-driven simulators can provide higher fidelity but require target-hardware traces or extensive calibration, while Primus Projection offers CPU-only pre-silicon exploration.Its hybrid mode instead needs only a single-node benchmark to project multi-node throughput.Source: §5 · p. 10
Hyperparameter optimization: from classical search to LLM agents
- Bayesian optimization is a poor fit for this configuration problem because feasibility can change with small memory-accounting errors and candidates differ in cost and cross-axis structure.Primus uses an analytical memory check followed by benchmark projection, and its LLM agent proposes candidates in natural language.Source: §5 · p. 10
Configuration tuning for distributed training
- Primus Tuning Agent couples a benchmark-anchored projection oracle with an LLM investigation loop to propose cross-axis combinations without requiring a multi-node cluster.The authors describe it as the first such coupling, to their knowledge.Source: §5 · p. 10
6 Conclusion
Primus Projection measures compute and in-scope communication on as few as one device, then analytically reconstructs out-of-scope communication and the target pipeline schedule. Its tuning agent searches the joint configuration space, with projections within ∼10% of measured throughput and a Mixtral configuration delivering +27% over AMD’s 4-node BF16 reference.
6 Conclusion
- Primus Tuning Agent delivered +27% over AMD’s Mixtral 4-node BF16 reference in ∼30 minutes of single-node exploration; Primus Projection’s validated-recipe throughput projections were within ∼10% of measured throughput (4.7% at 1,024 GPUs).Source: §6 · p. 10
- Primus Projection benchmarks compute and in-scope communication on as few as one device, then analytically reconstructs out-of-scope communication and the full target-scale pipeline schedule.The Primus Tuning Agent uses the projection framework as a fast surrogate for end-to-end configuration search.Source: §6 · p. 10
Appendices
The appendices describe the hardware and software environment and the analytical models underlying Primus Projection. Memory projection accounts for sharding, precision, activation sizes, pipeline schedules, and recomputation.
Appendix A: Hardware and Software Environment
- Measurements use AMD Instinct MI300X, MI325X, and MI355X GPUs with ROCm and Megatron-LM; custom hardware profiles can describe arbitrary cluster topologies and hypothetical pre-silicon architectures.Benchmarking reflects installed kernel performance and GPU clock behavior under power-cap throttling.Source: Appendix A · p. 13
Appendix B: Analytical Models
- Primus Projection’s analytical models are organized by component, including a hierarchical memory profiler that accounts for per-GPU parameter, gradient, optimizer-state, and activation memory.TP, PP, EP, CP, and FSDP/ZeRO sharding affect per-GPU parameter counts.Source: §B.1 · p. 13
B.1 Memory Projection
- For BF16 mixed-precision Adam, parameter, gradient, and optimizer-state costs are (2, 2, 10) bytes, giving bytes_per_param = 4 + 10/DP; FP8 reduces the weight term to 1 byte.The distributed-optimizer flag further shards optimizer state across DP ranks.Source: §B.1 · p. 13
- A large MoE layer can require 16.19 GB for BF16 saved activations, over 30× the attention activation, while FP8 saved activations reduce it to about 8.1 GB.Full recomputation reduces the cited 18.20 GB per MoE layer to about 0.25 GB at BF16 or 0.13 GB at FP8.Source: §B.1 · pp. 13–14
B.2 Performance Projection (Sub-Node Methodology)
Primus restores sub-node measurements to target-scale performance through analytical communication and schedule models, then searches configurations with a seeded planner and LLM investigation loop. A memory pre-filter rejects infeasible trials before GPU evaluation.
B.2 Performance Projection (Sub-Node Methodology)
- The workflow restores EP with an All-to-All delta, adjusts TP compute and AllReduce, and re-simulates the target PP schedule.When measured AllReduce data is available, its target contribution is scaled by the ratio of analytical target and benchmark estimates.Source: §B.2 · p. 14
B.3 Communication Modeling
- The communication model selects the minimum-time applicable collective algorithm across intra-node, inter-node, and inter-pod topology domains.For MoE with EP>1, expert and non-expert gradient AllReduce are modeled separately, with synchronization barriers placing the full time on the critical path.Source: §B.3 · p. 14
B.4 FSDP Overlap Model
- FSDP overlap hides portions of forward and recompute AllGather and ReduceScatter, leaving only exposed communication on the critical path.The overlap ceilings are 0.95, 0.93, and 0.92, respectively.Source: §B.4 · p. 14
B.5 Pipeline Schedule Algorithms
- The pipeline simulator supports six scheduling algorithms, including Zero-Bubble variants that schedule weight-gradient work into pipeline bubbles.For VPP=2, ZBV Formatted uses deterministic phase transitions and ZBV Greedy uses configurable greedy placement.Source: §B.5 · p. 15, Table 14 · p. 15
C.1 Seed Planner Axis Schedule
- The seed planner sweeps recomputation, MoE communication backend, pipeline schedule, micro-batch size, and distributed-optimizer settings from a user baseline.An analytic memory pre-filter rejects over-cap candidates before GPU evaluation; surviving trials are recorded in a shared ledger.Source: §C.1 · p. 15
C.2 LLM Investigation Loop
- The LLM phase reads the trial ledger and proposes cross-axis combinations that the single-axis seed planner cannot reach.An RLM evaluates candidates through the projection oracle; the cross-product search is reported as the source of a +21.6% advantage over the seed plan.Source: §C.2 · p. 15
C.3 Memory Pre-Filter and Trial Selection
- Every configuration projected above the memory cap OOMed on the real cluster, while every configuration at or below it ran successfully.Trials also require valid axis tiling and parseable throughput; the winner is the legal configuration with the highest tok/s/GPU.Source: §C.3 · p. 16
Appendix D: Large-Scale Search Traces
Large-scale traces show Primus scaling DeepSeek V3's global batch to the memory ceiling and selecting a Qwen3 configuration at 256 GPUs. Both production-scale projections closely match measured throughput.
Appendix D: Large-Scale Search Traces
- 2,808 tok/s/GPU versus 1,869 followed scaling DeepSeek V3's global batch to the memory ceiling; projection was within 4.7% of the measured 2,946.The final selection was validated at equal GBS=32,768 on 1,024 GPUs.Source: Appendix D · p. 16, Table 16 · p. 16, Table 17 · p. 16
- 4.7% error on DeepSeek V3 at 1,024 GPUs and 1.9% on Qwen3-235B at 256 GPUs show close production-scale projection fidelity.Figure 5 reports projected versus measured throughput at production scale.Source: Figure 5 · p. 16
- 4,860 projected versus 4,952 measured (1.9%) for Qwen3-235B-A22B at 256 GPUs; the selected configuration ranked ahead of a deeper PP=16 alternative.The selected shape used TP=1/PP=4/EP=8/CP=1/VPP=4/DP=64, MBS=4, and GBS=8,192.Source: Appendix D · p. 16
Appendix E: Rank Preservation: Full Per-Config Tables
Across measured configuration sweeps, projected throughput generally preserves ranking, and the top-1 selection matches the measured optimum. The documented GPT-OSS reversals cluster among near-tied configurations.
Appendix E: Rank Preservation: Full Per-Config Tables
- ρ=τ=1.000 across all seven measured Mixtral configurations, with offsets k=1.08–1.18 and no pair flips.Table 18 covers the 4-node/32-GPU FP8 recomputation/MBS sweep.Source: Table 18 · p. 17
- ρ=τ=1.000 across all seven measured Qwen3 configurations despite offset drift from 1.01 to 0.69; no pair flips.The low offset occurs for the MBS=1, GBS=1,024 configuration in the 4-node/32-GPU FP8 sweep.Source: Table 19 · p. 17
- GPT-OSS's three reversals involve mbs16, the baseline, and mbs2, while selection regret remains 0%.Those three measured configurations lie within ∼9% of one another; pairs separated by more than ∼9% are ordered correctly.Source: Appendix E · p. 17, Table 21 · p. 18
- Top-1 projected matched top-1 measured for GPT-OSS-120B, with +20.4% over the reference baseline versus +17.3% projected.Table 20 reports the 8-node/64-GPU BF16 sweep with no prior AMD anchor.Source: Table 20 · p. 18
- Qwen3 is perfectly monotone (ρ=τ=1.000), while GPT-OSS has ρ=0.933 and 0% selection regret.Figure 6 identifies the GPT-OSS reversals as one over-projected trio around mbs16; its top-1 pick and extremes remain ordered exactly.Source: Figure 6 · p. 18
Appendix F: MoE Load-Balancing Analysis
The load-balancing analysis separates forced-balanced benchmarks, trained workloads, and stress tests. A busiest-rank correction substantially reduces projection error under skew, while pooled step-level imbalance stays near uniform across tested EP sizes.
Appendix F: MoE Load-Balancing Analysis
- Forced-balanced benchmark runs set executed load to I=1 by construction; in real training, balancing and large-batch averaging keep steady-state I approximately 1.1–1.3.The benchmark-only router flag is never used in real training.Source: Appendix F · p. 19
- 53–71% worst-case error from a uniform-load model under disabled balancing fell below 10% with the busiest-rank correction across three models.The stress test swept imbalance from 1× to 4×; I≥2 lies outside the training regime.Source: Appendix F · p. 19, Table 22 · p. 19
- At EP=8, the trained checkpoint's per-layer median imbalance is 1.38 but step-level pooled imbalance is 1.05 because the busiest device rotates across layers.In pipelined, overlapped execution, pooled per-device work is the quantity that governs throughput.Source: Appendix F · p. 19, Table 23 · p. 19
- Pooled imbalance stays at 1.05–1.12 across EP=8/16/32, and the correction transfers across EP on Mixtral with ±6% validation.The reported imbalance order is forced-balanced 1.00 < trained 1.38 < random-init 2.27.Source: Appendix F · pp. 19–20
Appendix G: Baselines: Calculon and Galvatron-Style Search
Primus outperforms the supplied baselines in tuning and prediction accuracy, while the appendix describes a planner and iterative search loop for exploring configurations. The comparison also shows that the projection oracle affects which configurations the search identifies.
Appendix G: Baselines: Calculon and Galvatron-Style Search
- ∼4× more accurate: Primus’s Llama-3.1-70B throughput prediction beats Calculon’s, whose opposite-direction errors inflate the FP8–BF16 gap from measured 1.5× to 2.9×.Calculon models no experts, All-to-All, or ROCm path, so Mixtral and DeepSeek V3 cannot be expressed.Source: Appendix G · p. 20, Figure 8 · p. 20
- 28% higher: Primus Tuning Agent+Benchmark reaches 4,698 tok/s/GPU versus 3,656 for the exhaustive Galvatron-style FP8 sweep.The sweep was granted Primus’s simulate oracle, but only hardware benchmarking captures the superposed FP8 tensor-core, RCCL A2A, and grouped-GEMM effects.Source: Figure 8 · p. 20, Appendix G · p. 20, Table 25 · p. 21
Appendix H: Tuning-Agent LLM Backbone and Prompts
- The planner prompt asks for an ordered, numbered search plan based on model and cluster characteristics, legal axes, memory headroom, and known trade-offs.It requests hypotheses and candidate configurations across parallelism, micro-batch, recomputation, scheduling, and overlap choices.Source: Appendix H · p. 21
- The iterative search evaluates planned candidates with GPU benchmark projections, records findings, and stops when its budget is exhausted or five trials fail to improve the incumbent.The instructions also provide memory-only evaluation, history, architecture, cluster, and budget tools.Source: Appendix H · p. 21
Appendix I: Limitations
The appendix identifies limits in the projection model and tuning search space, including unmodeled network contention, reduced-parallelism fidelity, and hardware and precision scope. It also describes the tool’s potential to reduce tuning GPU-hours and broaden access to configuration exploration.
Appendix I: Limitations
- Network contention is not modeled, so concurrent collectives sharing xGMI or NIC bandwidth may run slower than the dedicated-bandwidth model predicts.The authors have not observed this as a dominant error source in evaluated configurations, but it could add error at extreme scale.Source: Appendix I · p. 22
- Sub-node benchmarking changes the target parallelism configuration; longer analytical reductions can increase residual error, with +6.6% Mixtral error at 2 GPUs as one bound.The framework reports measured base time and analytical deltas for auditing restored dimensions.Source: Appendix I · p. 22
- Nodes are assumed identical in static hardware specifications and runtime behavior; GPU frequency variation and partial-rack failures are outside the model’s scope.Source: Appendix I · p. 22
- The agent searches only optimizations exposed as projection-oracle flags; adding a new optimization requires extending both the oracle and planner.ZeRO-2 is excluded because wherever it would beat ZeRO-1, FSDP2 is available.Source: Appendix I · p. 22
- Precision is evaluated as a throughput knob, not for numerical quality; FP8 recipes require convergence validation before production use.Source: Appendix I · p. 22
- Validating support for new layer types is future work; NVIDIA porting requires a hardware YAML and 8-GPU benchmark, while TPU support requires adapting the collective model to XLA/JAX.Source: Appendix I · p. 22
Appendix J: Broader Impact
- The authors say single-node projection can reduce tuning GPU-hours, while CPU simulation and candidate generation may help smaller groups explore configurations before renting validation capacity.The Mixtral case study used about 30 minutes of single-node agent time plus a small number of validation runs, instead of days of manual multi-node sweeps.Source: Appendix J · p. 22