Tuning agent

August 12, 2026 · View on GitHub

The Tuning Agent is an LLM-driven search for a near-optimal Primus training configuration—the full parallelism strategy plus the coupled batching, pipeline-schedule, memory, MoE-communication, and precision knobs — on a target GPU cluster, without running the workload at scale. It drives the Primus Projection tool as an evaluation oracle. Projection provides two estimates—memory and performance—each of which runs benchmark-anchored by default (measuring what fits on a sub-node run and scaling the rest analytically) with a fully analytical no-GPU simulate fallback. The agent returns the configuration that maximizes tokens/s/GPU subject to a per-GPU memory safety margin. See Knobs Searched for the full set of levers it tunes.

Package: primus/agents/tuning_agent/
Entry point: python -m primus.agents.tuning_agent

This document is both the user/operator guide (installation, configuration, execution modes, CLI reference, worked example, troubleshooting) and the design write-up (paper-ready problem statement and deferred future features) for the agent.

Why a tuning agent
How it works
Knobs searched
Installation
LLM setup
Quickstart
Execution modes
CLI reference
Target-cluster YAML
The search loop
Evaluator and projection modes
Output artefacts
Worked example
Troubleshooting
Limitations
Design notes and future features

Why a tuning agent

Choosing a training configuration for a large training (or inference) workload is a combinatorial problem. The configuration is the joint choice of the parallelism dimensions:

Data Parallel (DP)—derived from world size and the other axes,
Tensor Parallel (TP),
Expert Parallel (EP)—for MoE models,
Context Parallel (CP),
Pipeline Parallel (PP), with virtual pipeline (VPP) and the pipeline schedule (1F1B / interleaved / zero-bubble / ZBV-* / Megatron-ILP),

together with the strongly coupled training knobs that decide how those dimensions translate into in-flight work and memory pressure: global batch size (GBS), micro batch size (MBS), activation recomputation (recompute_granularity, recompute_num_layers), the overlap flags (overlap_grad_reduce, overlap_param_gather), and a set of higher-impact levers—FP8 precision, MoE DeepEP / sync-free communication, fused cross-entropy, and optimizer-state sharding (distributed optimizer / FSDP2). The full set is enumerated in Knobs Searched.

The legal Cartesian product is large, the objective surface is non-convex and architecture-specific, and exhaustive evaluation on real hardware is infeasible. The agent replaces exhaustive sweeps with a small number of informed analytical evaluations: an LLM proposes candidate configurations conditioned on the resolved model architecture, the target cluster, and the history of prior trials; the projection evaluator scores them for memory feasibility and throughput. The loop keeps an incumbent and returns the best legal configuration within a user-specified budget.

How it works

┌────────────────────────────────────────────────────────────────────────┐
│  workload YAML ──► resolve architecture record (layers, hidden, MoE …)   │
│  target-cluster YAML ──► cluster + budget + LLM config                   │
│                                                                          │
│   derive per-axis legal sets (divisibility + cluster size, in code)      │
│                              │                                           │
│            ┌─────────────────┴─────────────────┐                        │
│            ▼                                     ▼                        │
│   systematic seed grid                 DSPy planner + RLM rounds          │
│   (no LLM, coarse TP/PP/EP/CP)         (LLM proposes, tools score)        │
│            │                                     │                        │
│            └──────────────► Evaluator ◄──────────┘                       │
│                  projection memory / performance                         │
│                  (simulate · no GPU | benchmark · GPU)                    │
│                              │                                           │
│                              ▼                                           │
│        trials.jsonl · trials.png · scratchpad.txt · summary.json          │
│              best config (PRIMUS_* exports + re-runnable YAML)            │
└────────────────────────────────────────────────────────────────────────┘

The key idea is a separation of concerns: legality is computed in code (the LLM can never spend budget proposing an obviously illegal config), while strategy is delegated to the LLM (which axis to push next, when to trade recompute for MBS, whether CP helps a given MoE shape).

Knobs searched

The agent sweeps far more than the five parallelism dimensions. Its trial configuration (TrialConfig in legality.py) carries the full set of knobs below; every knob is either inherited from the workload YAML (when the agent leaves it unset) or overridden for a trial and translated into the corresponding projection flags by the evaluator. Each is legality-checked in code before it ever reaches the projection tool.

Parallelism and batching

Knob	Legal values	What it controls
`tp`	divisors of `num_attention_heads` and `hidden_size`, ≤ `gpus_per_node`	Tensor parallelism
`pp`	divisors of `num_layers` (plus layout-aware depths 2/4/8/16 for layout workloads, and the workload's own PP)	Pipeline parallelism
`ep`	divisors of `num_experts` (MoE only, else 1)	Expert parallelism
`cp`	dense: divisors of `seq_length` ≤ `gpus_per_node`; MoE: `CP ≤ EP` and `EP % CP == 0` (parallel folding)	Context parallelism
`vpp`	divisors of `num_layers / PP` (≤ 8), plus layout-aware VPP	Virtual pipeline (interleaving)
`mbs`	powers of two ≤ 16	Micro batch size
`gbs`	multiple of `MBS × DP`	Global batch size
`overlap_grad_reduce`	`true` / `false`	Overlap the DP gradient all-reduce with backward

Pipeline schedule

Knob	Legal values	What it controls
`pp_schedule`	VPP=1: `auto`, `zerobubble`, `zerobubble-heuristic`, `seaailab-ilp`; VPP=2: `auto`, `zbv-formatted`, `zbv-greedy-half`, `zbv-greedy-min`; other VPP: `auto`	Schedule algorithm (paired with VPP)
`enable_zero_bubble`	`true` / `false` / inherit	Split backward into B + W to fill pipeline bubbles

Memory levers

Knob	Legal values	What it controls
`recompute_granularity`	`none` / `selective` / `full`	Activation recomputation strategy
`recompute_num_layers`	int (≤ layers per VPP stage)	Layers recomputed per stage under `full`
`cross_entropy_loss_fusion`	`true` / `false` / inherit	Fused cross-entropy—large-vocab memory + compute win
`use_distributed_optimizer`	`true` / `false` / inherit	ZeRO-1 optimizer-state sharding across DP
`use_torch_fsdp2`	`true` / `false` / inherit	FSDP2 sharding (mutually exclusive with `use_distributed_optimizer`)

MoE communication—MoE only, high impact

Knob	Legal values	What it controls
`use_turbo_deepep`	`true` / `false` / inherit	DeepEP dispatch/combine kernels for MoE All-to-All (large MoE-comm win)
`sync_free_stage`	`0` / `1` / `2` / `3`	Sync-free MoE pipelining; stage ≥ 2 auto-enables DeepEP
`target_ep_size`	positive int / inherit	EP override used for All-to-All modeling

Precision—high impact

Knob	Legal values	What it controls
`fp8`	`none` / `hybrid` (also `e4m3`, `delayed`)	FP8 on linear layers—roughly 2× compute on GEMMs

Coupling rules enforced in code

Some knobs interact; the validator rejects incoherent combinations before any projection call, so the LLM never wastes budget on them:

use_torch_fsdp2 and use_distributed_optimizer are mutually exclusive (FSDP2 already shards the optimizer state).
sync_free_stage ≥ 2 auto-enables DeepEP, so use_turbo_deepep=false alongside it is contradictory.
use_turbo_deepep, sync_free_stage, and target_ep_size are MoE-only and rejected on dense workloads.
fp8 combined with MLA attention is steered to a tensorwise recipe for container compatibility.

The optimization.axes toggles in the target-cluster YAML gate the core parallelism / batching / schedule / recompute / overlap axes (set one to false to pin it to the workload baseline). The remaining Tier-A knobs (DeepEP, sync-free, FP8, loss fusion, distributed-optimizer / FSDP2) are explored by the LLM proposer whenever they apply to the workload.

Installation

python3 -m venv .venv-agent
source .venv-agent/bin/activate
pip install -r primus/agents/tuning_agent/requirements.txt

# Optional: only needed for the `simulate` projection backend (no GPU).
pip install git+https://github.com/ROCm/rocm-libraries.git#subdirectory=shared/origami/python

The agent itself is lightweight; the heavy dependencies (torch, Megatron, Origami) are only needed by the evaluator paths you actually use:

You want to run…	You need…
`--mode dry` (loop sanity check)	nothing beyond the agent requirements
`--mode memory-real`	Primus + torch (CPU is fine)
`--profiling-mode simulate`	+ Origami
`--profiling-mode benchmark`	+ a ROCm GPU on the local host

LLM setup

The agent uses DSPy, which routes LLM calls through LiteLLM internally—no separate proxy process is required. Set credentials for whichever provider you use:

# OpenAI
export OPENAI_API_KEY=sk-...
export LLM_MODEL=openai/gpt-4o            # default if unset

# Anthropic
export ANTHROPIC_API_KEY=sk-ant-...
export LLM_MODEL=anthropic/claude-opus-4-5

# Any OpenAI-compatible endpoint (Ollama, vLLM, local LiteLLM proxy, …)
export OPENAI_API_KEY=dummy              # or the real key if required
export OPENAI_API_BASE=http://localhost:11434/v1
export LLM_MODEL=openai/llama3

The model string follows LiteLLM's provider-prefixed convention (<provider>/<model-name>). Credentials can also live in a .env file; the agent searches $CWD/.env, <repo-root>/.env, then ~/.env. You may also set them in the target-cluster YAML under agent.llm (see below); the environment takes precedence over the YAML for LLM_MODEL.

The LLM stage is optional. --seed-only runs the systematic seed grid against the real projection tool and reports the best seed config with no LLM credentials at all.

Quickstart

# 1) Dry-run: no primus-cli, no LLM. Exercises the loop end-to-end with
#    synthesised metrics — verifies the install on a CPU-only host.
python -m primus.agents.tuning_agent \
    --workload examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml \
    --target-cluster examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml \
    --out-dir tuning_runs/dry-run \
    --dry-run --seed-only

# 2) Seed-only with the real projection tool (Origami needed for simulate):
python -m primus.agents.tuning_agent \
    --workload examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml \
    --target-cluster examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml \
    --out-dir tuning_runs/mixtral-22b-seed \
    --seed-only

# 3) Full agent (planner + DSPy.RLM rounds):
python -m primus.agents.tuning_agent \
    --workload examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml \
    --target-cluster examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml \
    --out-dir tuning_runs/mixtral-22b-full

Execution modes

The agent has two orthogonal mode switches: what the evaluator does (--mode) and whether the LLM stage runs (--seed-only / --agent-only).

Evaluator mode (`--mode`)

Mode	GPU	What it does
`dry`	No	Synthesises metrics; never calls `primus-cli`. For testing the agent loop on a host with no Primus/torch/Origami install. (`--dry-run` is shorthand.)
`memory-real`	No	Calls `projection memory` for real and synthesises `tokens/s` from a memory + axes heuristic. Use when `projection performance` cannot run (no Origami / no GPU).
`full`	No*	Calls `projection memory` and `projection performance`. Default. `simulate` profiling needs Origami (no GPU); `benchmark` profiling needs a GPU.

* full + --profiling-mode simulate runs entirely on CPU.

LLM stage

Flag	Effect
(default)	Run the systematic seed grid, then the DSPy planner + RLM rounds.
`--seed-only`	Evaluate the systematic seed plan only, then exit. No LLM.
`--no-agent`	Alias for `--seed-only`.
`--agent-only`	Skip the seed grid; run the LLM agent on the existing `trials.jsonl`.
`--resume`	Reuse an existing `trials.jsonl` in `--out-dir` (suppresses the overwrite warning).

CLI reference

python -m primus.agents.tuning_agent \
    --workload <primus_pretrain.yaml> \
    --target-cluster <target_cluster.yaml> \
    [--out-dir <dir>] \
    [--mode {dry,memory-real,full}] \
    [--profiling-mode {simulate,benchmark}] \
    [--dry-run] [--seed-only] [--no-agent] [--agent-only] [--resume] \
    [--seed-budget N]

Flag	Default	Description
`--workload`	(required)	Path to a Primus pretrain YAML. The agent follows the `extends:` chain to flatten the architecture.
`--target-cluster`	(required)	Path to a target-cluster YAML (see below).
`--out-dir`	`./tuning_runs/<cluster-name>`	Output directory for trials, plot, scratchpad, summary.
`--mode`	`full`	Evaluator mode (`dry` / `memory-real` / `full`).
`--profiling-mode`	`simulate`	Profiling backend for `projection performance` in `full` mode. `simulate` = Origami, no GPU; `benchmark` = real GPUs via `torch.distributed.run` on the local node, projected to `target_cluster.num_nodes` (requires `has_gpu: true`).
`--dry-run`	off	Shorthand for `--mode dry`.
`--seed-only`	off	Evaluate the systematic seed plan only; skip the LLM.
`--no-agent`	off	Alias for `--seed-only`.
`--agent-only`	off	Skip seed evaluation; run the LLM agent on existing history.
`--resume`	off	Reuse an existing `trials.jsonl` in `--out-dir`.
`--seed-budget`	`12`	Max seed candidates evaluated before the LLM takes over.

Target-cluster YAML

A thin wrapper around the existing Primus hardware_config convention, so no new networking format has to be invented; topology, bandwidths, and latencies are consumed by the analytical communication model (see projection.md → Assumptions (performance projection)).

A complete example ships at examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml:

target_cluster:
  name: mi355x_4nodes
  num_nodes: 4
  gpus_per_node: 8
  gpu_arch: mi355x                          # selects the Origami / SDPA hardware profile
  hardware_config: examples/hardware_configs/mi355x.yaml
  gpu_clock_mhz: null                       # optional Origami clock override

available_for_benchmark:
  has_gpu: true                             # true → --profiling-mode benchmark enabled
  benchmark_gpus: 8                         # GPUs the agent may use locally
  benchmark_arch: mi355x                    # arch of the local host (may differ from target)

optimization:
  objective: tokens_per_s_per_gpu
  memory_safety_margin: 0.10                # hard cap = HBM × (1 − margin)
  hbm_capacity_gb: 288.0                    # MI300X/MI325X/MI355X = 192 / 256 / 288
  budget:
    max_proposals: 30                       # LLM-proposed candidates
    max_perf_calls: 30                      # simulate calls
    max_benchmark_calls: 12                 # GPU benchmark calls (ignored if has_gpu=false)
    max_rounds: 8                           # outer agent rounds
    max_rlm_iterations: 30                  # inner RLM iterations per round
  axes:                                     # toggle individual axes off to fix them
    tp: true
    pp: true
    ep: true
    cp: true
    mbs: true
    gbs: false
    vpp: true
    pp_schedule: true
    recompute: true
    overlap_grad_reduce: false

agent:
  llm:
    model: anthropic/claude-opus-4-5        # optional; env LLM_MODEL overrides
    api_key: sk-ant-...                     # optional if set in env
    base_url: ...                           # optional; override for custom endpoints
  prompt_extras:                            # optional free-form context for the LLM
    - "Prefer keeping EP within a single node when possible."

Field reference

Block / field	Default	Meaning
`target_cluster.num_nodes` / `gpus_per_node`	`1` / `8`	Target world size = `num_nodes × gpus_per_node`.
`target_cluster.gpu_arch`	`mi355x`	Selects the Origami / SDPA hardware profile for `simulate`.
`target_cluster.hardware_config`	`null`	Path to a Primus hardware YAML for communication modeling.
`available_for_benchmark.has_gpu`	`false`	When `true`, the `benchmark` profiling path is enabled.
`available_for_benchmark.benchmark_gpus`	`0`	GPUs the agent may use for sub-node benchmarking.
`optimization.objective`	`tokens_per_s_per_gpu`	The scalar maximized by the search.
`optimization.memory_safety_margin`	`0.10`	Feasibility cap: `projected_mem ≤ HBM × (1 − margin)`.
`optimization.hbm_capacity_gb`	`192.0`	Per-GPU HBM capacity used for the cap.
`optimization.budget.*`	see above	Stops the search when any budget is exhausted.
`optimization.axes.*`	most `true`	Set an axis `false` to pin it to the workload baseline.
`agent.llm.*`	env-resolved	Per-cluster LLM overrides (model / key / base_url).
`agent.prompt_extras`	`[]`	Free-form hints appended to the LLM prompt.

The search loop

Resolve the workload. Load the workload YAML, follow modules.pre_trainer.model into primus/configs/models/megatron/<model>.yaml, chase the extends: chain, and flatten an architecture record (num_layers, hidden_size, ffn_hidden_size, num_attention_heads, num_query_groups, num_experts, moe_router_topk, moe_ffn_hidden_size, attention type, seq_length, …) plus a baseline overrides record (the current TP/PP/EP/CP/MBS/GBS/recompute) used as the search starting point.

Derive legal axes in code (not delegated to the LLM):

Axis	Legality
TP	divides `num_attention_heads` and `hidden_size`; ≤ `gpus_per_node`
PP	divides `num_layers` (or any if explicit `pipeline_model_parallel_layout`)
EP	divides `num_experts` (MoE only), else 1
CP	dense: free; MoE: `EP % CP == 0` and `CP ≤ EP` (parallel folding)
DP	derived: `world_size / (TP × PP × EP)` (MoE) or `world_size / (TP × PP × CP)` (dense)
MBS	positive integer with `GBS % (MBS × DP) == 0`
GBS	multiple of `MBS × DP`
VPP	divides `num_layers / PP`
pipeline schedule	depends on VPP: 1F1B / Megatron-ILP / Zero-Bubble need VPP=1; Interleaved 1F1B needs VPP>1; ZBV-* need VPP=2
recompute	`full` / `selective` / off; `recompute_num_layers ≤ layers_per_vpp_stage`

Seed (no LLM): the workload baseline plus a coarse legal grid over (TP, PP, EP, CP), capped at --seed-budget candidates, each filtered by projection memory and scored by projection performance.
LLM rounds (DSPy ChainOfThought planner + RLM driver): each round sees the architecture, cluster, axis legality, the full compressed trial history, and a durable scratchpad, and proposes k new candidates. The agent validates legality in code, filters with memory, scores survivors, appends (config, result) to history, and updates the incumbent. It early-stops when no improvement occurs within budget.
Emit the best config (as PRIMUS_* env exports and a re-runnable workload overlay YAML), the trial log, and the plot.

Evaluator and projection modes

The evaluator wraps the Primus Projection CLI behind a uniform interface, so the agent does not need to know which mode produced a number:

evaluate(config) -> {
  legal: bool,
  reason: str | None,                 # only when legal=False
  memory_per_gpu_gb: float,
  param_optimizer_gb: float,
  activation_gb: float,
  iteration_ms: float | None,
  tokens_per_s_per_gpu: float | None,
  tflops_per_s_per_gpu: float | None,
  bubble_ratio: float | None,
  comm_breakdown: dict[str, float],   # ms by collective
  source: "memory_only" | "simulate" | "benchmark" | "both",
}

The wrapper always picks the cheapest sufficient mode: a memory-only pre-filter to reject infeasible configs before paying for a performance call, then simulate (or benchmark for promising candidates when a GPU is available). The tool belt exposed to the LLM mirrors this:

evaluate_memory_only(config_json)—cheap pre-filter
evaluate_simulate(config_json)—primary scoring path
evaluate_with_benchmark(config_json)—only if has_gpu: true
get_history, get_best, get_legal_axes, get_architecture, get_cluster, get_budget_status
note_to_scratchpad, read_scratchpad
query_llm(prompt, system?)—one-shot "LLM-inside-LLM" consultation

For the underlying projection math—memory components, the simulate vs. benchmark trade-off, and the benchmark-based memory projection the agent uses for OOM-accurate feasibility—see projection.md.

Output artefacts

Everything lands in --out-dir:

trials.jsonl       all attempted configs and their results
trials.png         incumbent objective vs. trial number
scratchpad.txt     durable LLM notes carried across rounds
summary.json       agent-summarised winner
trials/*.yaml      one re-runnable workload-overlay YAML per trial

The run also prints the best configuration and ready-to-paste exports:

=== BEST CONFIGURATION ===
  trial #17  source=simulate
  TP=1 PP=4 EP=8 CP=1
  MBS=2 GBS=128 VPP=2 schedule=interleaved_1f1b
  recompute=selective/None
  → tokens/s/GPU = 3,260
  → memory/GPU   = 182.8 GB

=== EXPORTS for primus-cli ===
  export PRIMUS_TP=1
  export PRIMUS_PP=4
  export PRIMUS_EP=8
  # then pass --micro-batch-size 2 --global-batch-size 128
``$

---

## \text{Worked} \text{example}

\text{Search} \text{for} \text{the} \text{best} \text{Mixtral} 8 \times 22\text{B} \text{configuration} \text{on} \text{a} 4-\text{node} \text{MI355X} \text{pod}, \text{with} \text{a}
\text{single} \text{idle} 8-\text{GPU} \text{node} \text{available} \text{for} \text{benchmarking}:

$``bash
# .env or exported: OPENAI_API_KEY / ANTHROPIC_API_KEY + LLM_MODEL

python -m primus.agents.tuning_agent \
    --workload examples/megatron/configs/MI355X/mixtral_8x22B_v0.1-BF16-pretrain.yaml \
    --target-cluster examples/agents/tuning_agent/target_cluster_mi355x_4nodes.yaml \
    --out-dir tuning_runs/mixtral-22b-mi355x \
    --mode full --profiling-mode benchmark

What happens:

The agent resolves Mixtral 8×22B (56 MoE layers, hidden 6144, 8 experts, topk 2) and derives legal axes given 32 GPUs.
The seed grid + LLM rounds propose configs; each is memory-filtered, then scored. With --profiling-mode benchmark, promising candidates are benchmarked on the local 8-GPU node and projected to 4 nodes.
The best legal config (highest tokens/s/GPU under the 10% memory margin) is printed with PRIMUS_* exports, and all trials are written to tuning_runs/mixtral-22b-mi355x/.

For a pure capacity-planning run with no GPU, swap in --profiling-mode simulate and set has_gpu: false in the cluster YAML.

Troubleshooting

Symptom	Likely cause / fix
`WARNING: cannot import agent (...); skipping LLM stage`	`dspy` / `python-dotenv` not installed. Install `primus/agents/tuning_agent/requirements.txt`, or run with `--seed-only`.
`--profiling-mode benchmark but available_for_benchmark.has_gpu=false`	Benchmark trials are marked illegal. Set `has_gpu: true` (and `benchmark_gpus`) in the target-cluster YAML, or use `--profiling-mode simulate`.
`simulate` errors about Origami	Install Origami (see Installation) or fall back to `--mode memory-real`.
`<file> exists; re-using` warning	A prior `trials.jsonl` is in `--out-dir`. Pass `--resume` to reuse it intentionally, or point `--out-dir` elsewhere.
`no legal trial found`	Every candidate failed memory/legality. Loosen `memory_safety_margin`, raise `hbm_capacity_gb`, or widen the legal axes.
LLM stage produces no improvement	The seed incumbent may already be near-optimal; check `trials.png`. Increase `budget.max_rounds` / `max_proposals` to explore further.

Limitations

These are honest caveats, not future features:

Simulator-vs-reality gap. --profiling-mode simulate is analytical; projection.md explains why benchmarking is more accurate. In no-GPU mode the agent's numbers are analytical, not anchored. When a GPU is available, run a few benchmark trials as a sanity check.
Memory-projection blind spots in simulate: A2A buffers, allocator fragmentation, and communication scratch are not fully modeled. The memory_safety_margin compensates conservatively; the benchmark-based memory path (see projection.md) closes most of this gap.
Search-space explosion with all axes on. The agent mitigates with an impact-ordered deterministic seed plan—high-leverage levers first (recompute, MoE DeepEP / sync-free, FP8, then schedule and VPP/MBS neighbors), with the broad TP×PP×EP×CP grid evaluated last and capped by --seed-budget—so the LLM starts from an informed incumbent and spends its budget on polish.
Cluster-description lossiness: averaged bandwidth/latency cannot capture contention or per-rail asymmetry.

Design notes and future features

This section captures the design rationale and the paper-ready problem statement, plus the list of features deliberately deferred from v1.

Problem statement (paper-ready)

We address the problem of automatically selecting a near-optimal training configuration—the parallelism strategy plus the coupled batching, schedule, memory, MoE-communication, and precision knobs—for a large-scale distributed training workload on a target GPU cluster, without executing the workload at scale. The configuration space is combinatorial: for each axis only a small set of values is legal (constrained by divisibility against num_attention_heads, num_experts, num_layers, GBS / (MBS × DP); by the minimum-GPU requirement TP × PP × CP for dense models or TP × PP × EP for MoE under parallel folding; by per-GPU memory; and by axis interactions such as EP % CP == 0 or schedule-specific VPP), but their legal Cartesian product is large enough that exhaustive evaluation on real hardware is infeasible.

The objective surface is non-convex and architecture-specific: increasing TP reduces per-GPU parameters and activations but introduces AllReduce on the critical path; increasing PP reduces per-GPU parameters but introduces bubbles and requires GA ≥ PP; increasing EP shards expert weights but introduces inter-node All-to-All when EP spans nodes; increasing MBS improves GEMM efficiency but linearly inflates activation memory; recomputation trades ~33% compute for one-to-two orders of magnitude less activation memory on MoE layers.

We formulate the search as LLM-as-policy over a hybrid analytical / benchmark-driven evaluator. The evaluator is the Primus Projection tool, which provides three calls of increasing fidelity and cost: (i) a memory projection that runs either analytically (no GPU) or, by default, benchmark-anchored—measuring the real per-rank peak on a sub-node run and extrapolating it for an OOM-accurate estimate; (ii) a fully analytical performance projection built on the Origami GEMM model and an SDPA simulator; and (iii) a hybrid benchmark that measures per-layer compute on as few as one GPU and analytically scales PP/EP/DP to the target cluster. The LLM proposes candidate configurations conditioned on the resolved architecture, the cluster description, and the history of prior trials; the evaluator returns memory feasibility, throughput, and a structured breakdown. The loop returns, within a user-specified budget, the configuration that maximizes tokens/s/GPU subject to a configurable per-GPU memory safety margin.

The contribution is a configuration-search methodology that exploits an LLM's ability to reason over architectural priors (topk dominance of MoE activations, the impact of MQA on attention activation, the inter-node/intra-node boundary for All-to-All) to direct an analytical oracle—replacing exhaustive sweeps on real hardware with a small number of informed analytical evaluations. The system runs either entirely on a CPU-only host (simulate backend only) or in a mixed mode where a small number of real-hardware benchmark runs calibrate the analytical predictions.

Scope

A standalone agent, distributed alongside Primus, that drives the Projection tool in a closed loop.
A target-cluster YAML wrapping the existing examples/hardware_configs/*.yaml convention, so no new networking format has to be invented.
A single objective: maximize tokens/s/GPU subject to projected_memory ≤ HBM_capacity × (1 − safety_margin).
Agent search over all parallelism and coupled axes (TP, PP, EP, CP, MBS, GBS, VPP, pipeline schedule, recompute, overlap flags) plus the higher-impact levers (FP8, MoE DeepEP / sync-free, fused cross-entropy, distributed-optimizer / FSDP2)—restricted by per-architecture legality. See Knobs Searched for the complete list.
Two evaluator paths: a no-GPU path (projection memory --memory-mode simulate + projection performance --profiling-mode simulate), always available; and a with-GPU path that uses the default benchmark-anchored memory and performance projections for OOM-accurate feasibility and higher-fidelity throughput on promising candidates.
LLM driven via DSPy/LiteLLM (no separate proxy process); provider, model, and budget are user-configurable.
Output: the best legal configuration (YAML overlay + PRIMUS_* env overrides), the projected metrics, and a trial log.

Deferred future features

These are recorded so they are not lost; they are deliberately left out to keep the agent small.

F1. Multi-objective Pareto—replace the scalar tokens/s/GPU objective with a Pareto frontier over (throughput, MFU, memory headroom, projected $/token, projected energy/token).
F2. Online cluster-spec retrieval—pull the cluster description from an internal registry or known-archs catalog (MI300X / MI325X / MI355X reference pods) and fall back to user overrides; optionally infer topology from a small DCGM / ROCm-SMI dump.
F3. Persistent memory / configuration cache—cache (model_signature, cluster_signature) → best_known_configs across runs; invalidate when ROCm / hipBLASLt / framework versions change.
F4. Agent-proposed scale-downs and microbenchmarks—let the agent propose reduced-model proxies and targeted microbenchmarks to reduce the uncertainty of its current top-k (reusing the moe_proxy_single_node.yaml pattern).
F5. Telemetry plug-ins (rocprofiler / TraceLens / Magpie)—after a benchmark run, optionally extract per-kernel time, GEMM efficiency, A2A bytes, NIC utilization to calibrate the analytical models and explain underperformance back to the agent. Exposed as a SKILL.md-described plug-in. References: TraceLens, Magpie.
F6. Calibration learning—under --profiling-mode both, record per-(model, arch, dim) residuals between simulate and benchmark, fit a small correction model, and report a confidence band on subsequent simulate runs.
F7. Robustness / sensitivity report—for the winning config, sweep ±1 step on each axis and report whether the optimum is a sharp peak or a broad basin (cheap; only simulate calls).
F8. Cross-axis priors as DSPy modules—a library of tunable "rules of thumb" that DSPy's optimizer can refine over time using the trial logs.
F9. Sub-agent "test-proposer"—delegate targeted experiments (e.g. a 2-layer scale-down forward+backward microbenchmark, or a stand-alone A2A probe at the proposed EP × hidden_size × topk), profiling the test with explicit synchronisation rather than the sandbox. A run_proposed_experiment(plan, code) tool can be added to tools.py without restructuring the loop.
F10. Sub-LLM expert router—extend the existing query_llm tool into a named expert router (query_llm(expert='moe', …)).

Known design holes

Simulator-vs-reality gap—see Limitations above; in no-GPU mode the agent reports a confidence caveat.
Memory-projection blind spots—A2A buffers, allocator fragmentation, and comm scratch are not modeled analytically; the benchmark-based memory projection closes most of this gap by anchoring on a measured peak.
Search-space explosion—mitigated by an impact-ordered deterministic seed plan (high-leverage levers first, broad parallelism grid last) plus LLM-guided search from the seed incumbent.
Cluster-description lossiness—averaged bandwidth/latency cannot capture contention or per-rail asymmetry.

Projection—memory + performance projection internals, including the benchmark-based memory projection the agent relies on.
Tuning Agent package README — quickstart reference inside the source tree.