Performance tuning guide

Performance tuning guide#

This guide covers AMD-focused performance work in Primus: HipBLASLt autotuning for GEMMs, Primus-Turbo optional kernels, mixed precision, activation recomputation, communication overlap, memory settings, and MoE-specific flags. It references Primus examples and Megatron module YAMLs.

1. HipBLASLt autotuning#

Transformer Engine and GEMM-heavy training benefit from HipBLASLt kernel selection. Primus integrates a three-stage workflow controlled by PRIMUS_HIPBLASLT_TUNING_STAGE (see examples/README.md and examples/run_pretrain.sh).

Activate tuning first. The stage variable is only honored when the master switch PRIMUS_HIPBLASLT_TUNING=1 is set (and PRIMUS_DETERMINISTIC is not 1). Without PRIMUS_HIPBLASLT_TUNING=1, both run_pretrain.sh and the CLI hook runner/helpers/hooks/train/pretrain/prepare_experiment.sh skip tuning entirely and force TE_HIPBLASLT_TUNING_RUN_COUNT=0 / TE_HIPBLASLT_TUNING_ALGO_COUNT=0. Export PRIMUS_HIPBLASLT_TUNING=1 alongside the stage in every command below.

Stage 0 (default)#

No tuning:

export PRIMUS_HIPBLASLT_TUNING_STAGE=0  # default

Stage 1: Dump GEMM shapes#

Run a short training job so shapes are collected during real forward/backward passes. Reduce train_iters (or equivalent) for faster shape collection.

export PRIMUS_HIPBLASLT_TUNING=1
export PRIMUS_HIPBLASLT_TUNING_STAGE=1
./runner/primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

Output layout (from examples/README.md):

  • ./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_shape

Stage 2: Offline tuning#

Runs offline tuning from dumped shapes (often 10–30 minutes depending on model and shapes):

export PRIMUS_HIPBLASLT_TUNING=1
export PRIMUS_HIPBLASLT_TUNING_STAGE=2
./runner/primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

Expected output:

  • ./output/tune_hipblaslt/${PRIMUS_MODEL}/gemm_tune/tune_hipblas_gemm_results.txt

Stage 3: Train with tuned kernels#

Point the runtime at the tuned override file:

export PRIMUS_HIPBLASLT_TUNING=1
export PRIMUS_HIPBLASLT_TUNING_STAGE=3
export HIPBLASLT_TUNING_OVERRIDE_FILE=/path/to/tune_hipblas_gemm_results.txt
./runner/primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

Standalone offline tool#

For manual HipBLASLt bench workflows, see examples/offline_tune/offline_tune_gemm.py and examples/offline_tune/README.md (hipblaslt-bench integration and HIPBLASLT_TUNING_OVERRIDE_FILE usage).

2. Primus-Turbo optimization#

Primus-Turbo is a separate package of optimized AMD GPU kernels used by Primus Megatron and TorchTitan integrations. It is controlled by the master flag enable_primus_turbo in Megatron configs (primus/configs/modules/megatron/primus_turbo.yaml extends into trainer/model as needed). You must install the external primus_turbo package for these paths to be available.

Master flag (Megatron)#

enable_primus_turbo: true

Feature flags (Megatron)#

Defaults in primus/configs/modules/megatron/primus_turbo.yaml are mostly false until enabled.

Flag

Purpose

use_turbo_attention

Optimized attention kernels.

use_turbo_parallel_linear

Optimized tensor-parallel linear layers.

use_turbo_grouped_gemm

Optimized grouped GEMM for MoE.

use_turbo_grouped_mlp

Removed—use use_turbo_grouped_gemm (passing this key now raises an error).

use_turbo_rms_norm

Optimized RMSNorm.

moe_use_fused_router_with_aux_score

Fused MoE router (requires Primus-Turbo backend; see Backend Patch Notes).

use_turbo_deepep

DeepEP token dispatcher; set with enable_primus_turbo: true.

turbo_deepep_num_cu

Compute units for DeepEP (patch notes suggest practices such as 64 or 80 for EP8, 32 for EP16–64).

turbo_sync_free_moe_stage

Sync-free MoE stages (0–3; 0 disables, stage 2 recommended for performance per patch notes). See MoE training deep-dive.

use_turbo_fused_act_with_probs

Fused activation with probabilities to reduce redundant work.

Feature flags (TorchTitan)#

TorchTitan presets include primus_turbo in primus/configs/modules/torchtitan/pre_trainer.yaml

Example keys:

primus_turbo:
  enable_primus_turbo: true
  use_turbo_attention: true
  use_turbo_async_tp: true
  use_turbo_float8_linear: true
  use_turbo_grouped_mm: false

Documentation#

Extended Megatron arguments and Turbo-related behavior are summarized in Backend Patch Notes.

3. Mixed precision training#

Megatron (trainer_base.yaml patterns)#

Setting

Description

bf16: true

BFloat16 training (default true in trainer_base.yaml).

fp16: false

FP16 training (optional).

fp8

FP8 recipe control (null / recipes such as delayed scaling in upstream Megatron).

fp8_recipe, fp8_margin, fp8_interval

FP8 scaling behavior.

fp4, fp4_recipe

Experimental FP4 paths.

first_last_layers_bf16: true

Keep first/last layers in BF16 for stability (num_layers_at_start_in_bf16, num_layers_at_end_in_bf16 fine-tune).

TorchTitan#

Setting

Location

training.mixed_precision_param: bfloat16

pre_trainer.yaml default

training.mixed_precision_reduce: float32

Reduce precision

FP8 / quantization

quantize.linear.float8.* in primus/configs/modules/torchtitan/quantize.yaml

Loss fusion (Megatron model)#

From primus/configs/models/megatron/language_model.yaml:

  • Default: cross_entropy_loss_fusion: false with cross_entropy_fusion_impl: "native".

  • To use Transformer Engine fused cross entropy where supported, enable the fusion explicitly and set cross_entropy_fusion_impl: "te".

4. Activation recomputation#

Megatron#

Parameter

Typical values

Notes

recompute_granularity

full, selective

full recomputes more; max memory savings.

recompute_method

uniform, block

How recomputation is distributed.

recompute_num_layers

integer

Layers to recompute when using selective or uniform strategies.

recompute_layer_ids

list or null

Primus extension: global layer indices from 0 to num_layers - 1 (the patch resolves block-local indices to global ids via layer_offset). Use with recompute_granularity: full and supported recompute methods.

TorchTitan#

Parameter

Location

activation_checkpoint.mode

none in default pre_trainer.yaml

activation_checkpoint.selective_ac_option

Selective AC options

5. Communication overlap#

Megatron#

From trainer_base.yaml and model settings:

Setting

Purpose

overlap_grad_reduce

Overlap gradient reduction with backward.

overlap_param_gather

Overlap parameter gather with forward.

overlap_p2p_comm

Pipeline P2P overlap.

async_tensor_model_parallel_allreduce

Async TP all-reduce (model config).

TorchTitan#

Setting

Purpose

parallelism.enable_async_tensor_parallel: true

Async tensor parallelism.

Environment#

CUDA_DEVICE_MAX_CONNECTIONS=1 is commonly required for correct overlap behavior in TP/PP stacks (see docs/03-configuration-reference/environment-variables.md and Megatron tests). Primus launch scripts or your cluster setup might set this.

6. Memory optimization#

Megatron#

Parameter

Purpose

optimizer_cpu_offload: true

Offload optimizer state to CPU.

optimizer_offload_fraction

Fraction to offload (1.0 in trainer_base.yaml).

use_distributed_optimizer: true

Shards optimizer state across DP ranks (when enabled).

empty_unused_memory_level

Aggressive emptying of unused memory (0 default).

global_batch_size + micro_batch_size

Increase global batch via gradient accumulation without increasing per-step activation memory.

TorchTitan#

Parameter

Purpose

training.enable_cpu_offload

CPU offload path in pre_trainer.yaml

7. MoE-specific optimization#

Megatron (model + turbo)#

Setting

Purpose

moe_permute_fusion: true

Fuse permutation / unpermutation (patch-notes.md).

moe_use_fused_router_with_aux_score: true

Fused router + aux loss (Primus-Turbo).

use_turbo_deepep: true

DeepEP dispatcher (enable_primus_turbo must be true).

turbo_sync_free_moe_stage: 2

Recommended stage for sync-free MoE (per patch notes).

overlap_moe_expert_parallel_comm: true

Overlap expert parallel communication (trainer_base.yaml).

Quick checklist#

  1. GEMMs: run HipBLASLt stages 1–3 or use offline_tune_gemm.py for custom workflows.

  2. Kernels: enable primus_turbo after installing primus_turbo; turn on attention/MoE flags as needed.

  3. Precision: BF16 by default; add FP8/FP4 only with recipe testing.

  4. Memory: recomputation + distributed optimizer + CPU offload + accumulation before buying more GPUs.

  5. MoE: fusion + DeepEP + sync-free stages + EP comm overlap when supported.

Contents