Memory and performance projection

August 12, 2026 · View on GitHub

Primus projection tools estimate per-GPU memory and training throughput for large-scale distributed jobs without requiring the full target cluster. Two modes are available: analytical memory projection and performance projection that combines profiling with simulation.

Implementation: primus/cli/subcommands/projection.py

ModeCommandRole
Memoryprojection memoryEstimates per-GPU memory (parameters, optimizer state, activations) using analytical formulas.
Performanceprojection performanceBenchmarks on a single node (or sub-node), then projects training time to multi-node configurations.
Bothprojection bothRuns a single benchmark and produces both the performance and (benchmark-anchored) memory projections from it. Recommended for cluster-sizing workflows.

Core logic

  • Memory: primus/core/projection/memory_projection/
  • Performance: primus/core/projection/performance_projection/

Related: Micro-benchmarking suite, Preflight diagnostics, Megatron parameters.


Memory projection

Quick start

export NNODES=1
export HSA_NO_SCRATCH_RECLAIM=1

./runner/primus-cli direct --script primus/cli/main.py -- \
  projection memory \
  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml

Adjust --config to your experiment YAML. Memory estimation is analytical; the CLI still expects a normal Primus launch path (including distributed initialization where applicable).

What it estimates

ComponentMeaning
Parameter memoryModel weights assigned to this GPU (respecting parallelism).
Optimizer memoryOptimizer state (for example Adam moments), accounting for sharding across data-parallel groups.
Activation memoryActivations retained for the backward pass for a given microbatch and sequence length.

The tool walks a hierarchical profiler structure aligned with the model (embeddings, dense and MoE layers, output head, loss) and aggregates per-component contributions.

How to interpret results

Console output includes per-component breakdowns and a summary such as parameter count, param+optimizer memory, activation memory for the configured batch size and sequence length, and a projected total. Use these to answer whether a configuration fits in HBM before you allocate large clusters.


Performance projection

Quick start

Minimum required nodes (derived from parallelism):

export NNODES=1
export HSA_NO_SCRATCH_RECLAIM=1

./runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml

How it works

  1. Profile layer-level behavior on one node (or a subset of GPUs with automatic scaling rules).
  2. Simulate pipeline scheduling, data parallelism, and communication using analytical models.
  3. Project iteration time and tokens/s to a target node count when you specify one.

Projecting to a specific node count

./runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \
  --target-nodes 4

If --target-nodes is omitted, the tool defaults to the minimum number of nodes implied by your parallelism configuration (TP, PP, EP, CP, GPUs per node).

Parallelism overrides (environment)

You can override parallelism for what-if analysis:

export PRIMUS_TP=1
export PRIMUS_PP=3
export PRIMUS_EP=8

./runner/primus-cli direct --script primus/cli/main.py -- \
  projection performance \
  --config examples/megatron/configs/MI300X/deepseek_v2_lite-BF16-pretrain.yaml \
  --target-nodes 6

Command reference

Syntax

primus-cli [global-options] <mode> [mode-args] -- projection {memory,performance,both} [options]

Shared options (both modes)

OptionDescription
--config / --expPath to the Primus YAML configuration (required).
--data_pathData directory (default ./data when included on the parser).
--backend_pathOptional Megatron/TorchTitan import path appended to PYTHONPATH.
--export_configAccepted by the shared pretrain parser, but the default core runtime does not currently write a resolved YAML file.

Performance-only options

OptionDescription
--target-nodesTarget number of nodes for scaling projection. Defaults to the minimum nodes required by TP/PP/EP/CP and GPUs per node.
--target-num-nodesAlias-style projection override for target node count.
--target-ep-sizeOverride expert_model_parallel_size for the projection target.
--benchmark-gpusUse fewer than GPUS_PER_NODE GPUs for benchmarking; results are scaled analytically back to a full node.
--hardware-configYAML file with hardware parameters for communication modeling.
--profiling-modebenchmark (default, uses GPU), simulate (analytical / Origami GEMM + SDPA models, no GPU), or both (side-by-side).
--gemm-backendGEMM simulation backend when profiling is simulated (origami).
--gpu-archTarget architecture for simulation (for example mi300x, gfx942, mi355x, gfx950); can use PRIMUS_GPU_ARCH.
--gpu-clock-mhzOverride GPU clock in MHz for simulation; can use PRIMUS_GPU_CLOCK_MHZ.
--pipeline-schedule-algorithmPipeline simulation scheduler (auto, zero-bubble variants, or all for comparison).
--enable-zero-bubbleEnable zero-bubble pipeline scheduling for projection.
--enable-deepepEnable DeepEP overlap modeling.
--sync-free-stageOverride Sync-Free MoE stage (0 off; stages 1-3 enable additional modeling assumptions).
--num-virtual-stages-per-pipeline-rankOverride virtual pipeline stage count for projection.
--micro-batch-size, --global-batch-sizeOverride batch sizes for projection without editing the YAML.

Assumptions and limitations

Assumptions (performance projection)

  1. Data-parallel scaling—Compute time scales with ideal weak-scaling assumptions versus data-parallel width.
  2. Communication model—Uses simplified bandwidth and latency models (defaults such as efficiency factors may apply).
  3. Pipeline scheduling—Bubble and overlap behavior is modeled with fixed splits; real frameworks may differ.
  4. Gradients and MoE—Gradient all-reduce overlap and MoE all-to-all behavior follow the implemented model (for example overlap flags, EP scaling).

Limitations

  1. Single-node benchmark accuracy—Reduced PP/EP on the benchmark GPU count may not capture every production behavior.
  2. Contention—Network contention between jobs is not modeled.
  3. Memory vs speed—Activation recomputation reduces memory but adds compute; performance projection may not fully reflect that trade-off unless modeled.
  4. Heterogeneity—Assumes homogeneous nodes; GPU frequency drift across nodes is not modeled.

Tips

  1. Run projection memory first to confirm a configuration is feasible in HBM before spending time on performance projection.
  2. Always establish a single-node baseline before interpreting multi-node projections.
  3. Data-parallel scaling is bounded by batching: if you run out of microbatches (global_batch_size / micro_batch_size), adding nodes may not increase throughput.
  4. If the YAML requires multiple nodes (for example large PP), the performance path may automatically reduce parallelism for benchmarking and restore it analytically—read the console summary carefully.
  5. No GPU available: use --profiling-mode simulate for CPU-side analytical timing.
  6. Validate models: use --profiling-mode both to compare GPU benchmark timing with simulation on the same config.
  7. For MoE models, activation memory from MoE layers often dominates; memory projection highlights when recomputation is worth considering.