Skip to content

projection.md

Latest commit

 

History

History
212 lines (146 loc) · 12.3 KB

projection.md

File metadata and controls

212 lines (146 loc) · 12.3 KB
Raw

Memory and performance projection

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

Mode Command Role
Memory projection memory Estimates per-GPU memory (parameters, optimizer state, activations) using analytical formulas.
Performance projection performance Benchmarks on a single node (or sub-node), then projects training time to multi-node configurations.
Both projection both Runs 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.

Training backends

Projection reads the framework field of your experiment's pre_trainer module and adapts that backend's own configuration into the shapes and parallel degrees it models. A Llama 3 8B runs the same GEMMs over the same tensors whichever backend trains it, so all three share one profiler tree.

framework Config it reads Benchmark modes Simulate mode
megatron Megatron arguments directly Yes Yes
torchtitan model.flavor, parallelism.* degrees, activation_checkpoint.mode, model converters Yes Yes
maxtext (or jax) model_name, ici_* / dcn_* mesh axes, per_device_batch_size, remat_policy, quantization Yes Yes
torchrec_dlrm DLRM arguments Yes Yes

Benchmark-anchored projection runs the real trainer on one node to measure a representative layer, then scales that measurement analytically. Everything downstream of the measurement — pipeline scheduling, communication modeling, node-count scaling — is identical across backends, so the only thing a backend has to supply is the measurement itself.

Megatron, TorchTitan and DLRM all share one harness, because they share torch: the projection builds the real model, walks to the layer it wants, and times it with CUDA events while reading the caching allocator. A backend joins that harness by saying where its layers live and what shape their forward wants — TorchTitan's blocks take [batch, seq, hidden] and a rope table where Megatron's take [seq, batch, hidden] and a mask, and that is most of the difference between them.

MaxText measures itself instead. Its layers are Flax modules with no autograd to hook and no caching allocator to read, so it times a compiled vjp and takes activation sizes from XLA's own compile-time accounting. Two numbers come out differently as a result, and both are labelled in the saved artifact: the attention/MLP split within a layer is the measured layer total apportioned by the analytical model rather than two separate timings, because a MaxText layer is a single traced graph with no sub-module to time in isolation; and the output head and loss are left to the analytical model. Expert all-to-all is not timed on either non-Megatron backend — TorchTitan puts it in its parallelization plan and MaxText expresses it as sharding, so neither has a dispatch/combine pair to measure — and the projection restores it analytically when it scales to the target topology.

If you have no GPU, or want to skip measurement on a backend that supports it, pass --memory-mode simulate or --profiling-mode simulate.

Where the model architecture comes from

TorchTitan keeps its architectures in a Python flavor table and MaxText keeps its in YAML, so neither is spelled out in the experiment file. Projection resolves them in this order:

  1. Architecture keys set directly in the experiment YAML.
  2. The installed backend — TorchTitan's flavor table, or MaxText's model configs under third_party/maxtext (or PRIMUS_MAXTEXT_PATH).
  3. A transcribed table in primus/core/projection/frameworks/model_specs.py, so sizing a cluster needs neither a GPU nor a training checkout.

tests/unit_tests/core/projection/test_projection_model_specs.py re-checks the transcriptions against each backend whenever one is present, so a spec that drifts from upstream fails a test rather than quietly mis-sizing a cluster.

Adding another backend

Backend selection is registry-driven. Out-of-tree code can add one by writing an adapter that rewrites its trainer namespace in the projection's field names and calling primus.core.projection.frameworks.register_config_adapter(name, adapter). Adapters must normalize in place and be idempotent: the performance driver caps the layer stack, rescales EP onto the bench node and flattens PP on the normalized config, then converts it again.

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

Component Meaning
Parameter memory Model weights assigned to this GPU (respecting parallelism).
Optimizer memory Optimizer state (for example Adam moments), accounting for sharding across data-parallel groups.
Activation memory Activations 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)

Option Description
--config / --exp Path to the Primus YAML configuration (required).
--data_path Data directory (default ./data when included on the parser).
--backend_path Optional Megatron/TorchTitan import path appended to PYTHONPATH.
--export_config Accepted by the shared pretrain parser, but the default core runtime does not currently write a resolved YAML file.

Performance-only options

Option Description
--target-nodes Target number of nodes for scaling projection. Defaults to the minimum nodes required by TP/PP/EP/CP and GPUs per node.
--target-num-nodes Alias-style projection override for target node count.
--target-ep-size Override expert_model_parallel_size for the projection target.
--benchmark-gpus Use fewer than GPUS_PER_NODE GPUs for benchmarking; results are scaled analytically back to a full node.
--hardware-config YAML file with hardware parameters for communication modeling.
--profiling-mode benchmark (default, uses GPU), simulate (analytical / Origami GEMM + SDPA models, no GPU), or both (side-by-side).
--gemm-backend GEMM simulation backend when profiling is simulated (origami).
--gpu-arch Target architecture for simulation (for example mi300x, gfx942, mi355x, gfx950); can use PRIMUS_GPU_ARCH.
--gpu-clock-mhz Override GPU clock in MHz for simulation; can use PRIMUS_GPU_CLOCK_MHZ.
--pipeline-schedule-algorithm Pipeline simulation scheduler (auto, zero-bubble variants, or all for comparison).
--enable-zero-bubble Enable zero-bubble pipeline scheduling for projection.
--enable-deepep Enable DeepEP overlap modeling.
--sync-free-stage Override Sync-Free MoE stage (0 off; stages 1-3 enable additional modeling assumptions).
--num-virtual-stages-per-pipeline-rank Override virtual pipeline stage count for projection.
--micro-batch-size, --global-batch-size Override 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, on any backend.
  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.

Related documentation