amd-agi/primus/primus-projection

Primus Projection Skill

Purpose

Provide a quick, opinionated guide for using Primus Projection to:

• Choose parallelism and pipeline schedules
• Validate memory fit on target nodes
• Understand communication collectives/algorithms
• Explore optimization trade-offs with minimal compute
• Run the right commands fast

Scope

• Docs: docs/02-user-guide/projection.md
• Projection code: primus/core/projection/
• Projection performance entrypoints: primus/core/projection/performance_projection/
• Projection memory entrypoints: primus/core/projection/memory_projection/
• Pipeline schedulers: primus/core/pipeline_parallel/scheduler/
• Backend config adapters: primus/core/projection/frameworks/
• Example configs: examples/megatron/configs/, examples/torchtitan/configs/, examples/maxtext/configs/
• Hardware configs: examples/hardware_configs/

Training Backends

Projection reads the framework field of the experiment's pre_trainer module and adapts that backend's own config into the shapes and parallel degrees it models. All language-model backends share one profiler tree, so every recommendation below applies whichever backend trains the model.

TorchTitan benchmarks through the same torch harness as Megatron. MaxText measures itself: it times a compiled vjp and takes activation sizes from XLA's compile-time accounting, so its attention/MLP split is apportioned analytically and its output head and loss are analytical. Expert all-to-all is not timed on either and is restored analytically. Architecture comes from the installed backend when present (TorchTitan's flavor table, MaxText's model YAMLs) and otherwise from the transcribed table in primus/core/projection/frameworks/model_specs.py, so no GPU or training checkout is needed to size a cluster in simulate mode.

Note the vocabulary differences when reading a config: TorchTitan states parallelism as parallelism.*_parallel_degree (with -1 meaning "the leftover ranks") and MaxText as orthogonal ici_* / dcn_* mesh axes (with -1 meaning the same). MaxText's per_device_batch_size counts every device, not every data-parallel rank.

Projection Advisor (What users want to know)

1) Best parallelism strategy

• Evaluate TP/PP/EP/CP/DP for the target cluster using performance projection.
• Guidance:
• Prefer keeping TP intra-node (for faster AR); scale EP and DP across nodes.
• For MoE, use EP that avoids inter-node A2A when possible; otherwise enable DeepEP.
• Use CP only when sequence length forces it; in MoE, CP may be folded into EP.
• Increase DP last (weak scaling) once a minimal viable PP×EP×TP is established.

2) Best pipeline schedule

• Compare (for the chosen PP/VPP):
• 1F1B
• Interleaved 1F1B (VPP > 1)
• Zero-Bubble (B/W split, VPP = 1)
• ZBV Formatted (VPP = 2)
• ZBV Greedy (VPP = 2, memory modes: min/half)
• Megatron ILP (VPP = 1, where available)
• Recommendation pattern:
• Dense models: interleaved or ZB depending on VPP feasibility and memory.
• MoE (e.g., Mixtral): VPP=2 with ZBV Formatted often wins, especially with DeepEP ON.

3) Memory fit at target nodes

• Run memory projection first. Review:
• Parameters + optimizer (sharded by DP)
• Activations (attention vs MoE MLP; MoE often dominates)
• Pipeline scaling: peak activations scale with PP and VPP schedule
• If not fitting:
• Enable recomputation (full or partial)
• Adjust microbatch size and PP/VPP
• Reduce top-k or expert sizes for MoE if acceptable

4) Communication collectives and algorithms

• In use (typical):
• AllReduce: TP and DP gradients (best-of Ring/Hypercube/Bruck/Single-shot)
• All-to-All: EP token dispatch/combine (pairwise, topology-aware)
• Reduce-Scatter / All-Gather: optimizer/activation sharding (if configured)
• P2P Send/Recv: PP stage activations/gradients
• Distinguish intra-node (xGMI/NVLink/UALink) vs inter-node (IB/RoCE). Provide --hardware-config to model your network accurately.

5) Sample command lines

# Memory projection (check fit)
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection memory \
  --config <model_config.yaml>
 
# Performance projection (benchmark mode)
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --target-nodes <N>
 
# Performance projection (simulation-only, no GPU)
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --profiling-mode simulate \
  --gpu-arch <mi300x|mi325x|mi355x> \
  --target-nodes <N>
 
# Sub-node benchmarking (e.g., on 1 GPU), project to many nodes
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --benchmark-gpus 1 \
  --target-nodes <N>
 
# Compare benchmark vs simulation (accuracy check)
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --profiling-mode both \
  --target-nodes <N>
 
# Override parallelism from environment (optional)
export PRIMUS_TP=1
export PRIMUS_PP=3
export PRIMUS_EP=8
 
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --target-nodes <N>
 
# Provide cluster network topology
bash runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config <model_config.yaml> \
  --hardware-config examples/hardware_configs/mi355x.yml \
  --target-nodes <N>

Optimization Space Exploration

Primus can project throughput and memory across dozens of optimization combinations from a single-node benchmark (8 GPUs), eliminating the need for expensive multi-node runs. This was validated on DeepSeek V2 (236B MoE) across 17 configurations on 4-node MI355X — the projection tracked measured results across all optimization combinations and correctly identified the highest-impact optimizations.

Category 1: MoE Communication Optimizations (CLI flags)

Category 2: Pipeline Schedule & Layout (CLI flags + config)

Category 3: Memory Optimizations (config)

Category 4: Parallelism Configuration (config + env)

Category 5: Precision & Attention (config)