AMD has released the Primus Tuning Agent, a tool designed to automatically search for optimal configurations for large-scale AI training jobs. The tool uses the Primus projection engine as a fast scoring oracle to explore the joint configuration space for high-throughput, memory-legal recipes. In a case study involving the Mixtral 8×22B model, the agent identified a configuration that delivered a 27% increase in throughput compared to AMD’s published 4-node BF16 reference, without requiring full-cluster profiling or hand-written Megatron configurations. The tool aims to reduce the time and cost associated with trial-and-error approaches to configuring distributed training runs. Source: amd
The tuning agent operates in two phases: a deterministic seed phase and an LLM investigation phase. The deterministic seed phase evaluates a structured set of legal configurations based on the user's baseline, ensuring that any improvements are directly attributable to specific configuration changes. This phase guarantees coverage of high-leverage single-axis optimizations without relying on an LLM. The LLM investigation phase uses a DSPy planner and reasoning LM to propose cross-axis combinations that a single-axis sweep cannot reach. Both phases score candidates using the same projection oracle, which reduces the per-candidate cost from tens of minutes of multi-node profiling to seconds of single-node analysis. Source: amd
Configuring a distributed LLM training run involves selecting five interacting parallelism dimensions — Tensor (TP), Pipeline (PP), Expert (EP), Context (CP), and Data (DP) parallelism — alongside micro-batch size, gradient accumulation, recomputation granularity, FSDP/optimizer sharding, MoE communication backend, pipeline schedule, and precision. Each combination trades memory, compute utilization, and communication overhead differently, creating an enormous and complex search space. The boundary between feasible and infeasible configurations is razor-thin, with a 4-byte miscount of activation memory per element potentially pushing a configuration from legal to out-of-memory (OOM). Source: amd