Skip to content

tuning_agent.md

Latest commit

 

History

History
716 lines (587 loc) · 35.2 KB

tuning_agent.md

File metadata and controls

716 lines (587 loc) · 35.2 KB
Raw

Tuning Agent

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.

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.

Table of Contents

  1. Why a Tuning Agent
  2. How It Works
  3. Knobs Searched
  4. Installation
  5. LLM Setup
  6. Quickstart
  7. Execution Modes
  8. CLI Reference
  9. Target-Cluster YAML
  10. The Search Loop
  11. Evaluator and Projection Modes
  12. Output Artefacts
  13. Worked Example
  14. Troubleshooting
  15. Limitations
  16. Design Notes & 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 & 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 → Communication Modeling).

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

  1. 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.

  2. 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

  3. 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.

  4. 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.

  5. 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

Worked Example

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

# .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:

  1. The agent resolves Mixtral 8×22B (56 MoE layers, hidden 6144, 8 experts, topk 2) and derives legal axes given 32 GPUs.
  2. 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.
  3. 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:

  1. 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.
  2. 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.
  3. 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.
  4. Cluster-description lossiness: averaged bandwidth/latency cannot capture contention or per-rail asymmetry.

Design Notes & 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

  1. A standalone agent, distributed alongside Primus, that drives the Projection tool in a closed loop.
  2. A target-cluster YAML wrapping the existing examples/hardware_configs/*.yaml convention, so no new networking format has to be invented.
  3. A single objective: maximize tokens/s/GPU subject to projected_memory ≤ HBM_capacity × (1 − safety_margin).
  4. 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.
  5. 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.
  6. LLM driven via DSPy/LiteLLM (no separate proxy process); provider, model, and budget are user-configurable.
  7. 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

  1. Simulator-vs-reality gap — see Limitations above; in no-GPU mode the agent reports a confidence caveat.
  2. 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.
  3. 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.
  4. Cluster-description lossiness — averaged bandwidth/latency cannot capture contention or per-rail asymmetry.

Related Documentation

  • 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.