Profiling and observability

Profiling and observability#

This guide covers how to capture and analyze performance data in Primus: the PyTorch/Kineto profiler, GPU memory snapshots, AMD’s TraceLens trace analysis, ROCm memory sampling, memory/performance projection, and pipeline-schedule visualization. Parameters are grounded in primus/configs/modules/megatron/, primus/configs/modules/torchtitan/pre_trainer.yaml, and primus/configs/modules/maxtext/pre_trainer.yaml.

1. Torch profiler (Megatron)#

Megatron training integrates the PyTorch profiler. Defaults live in primus/configs/modules/megatron/trainer_base.yaml and primus_megatron_module.yaml.

Parameter	Default	Purpose
`profile`	`false`	Master switch for profiling.
`use_pytorch_profiler`	`false`	Use the PyTorch (Kineto) profiler path.
`profile_ranks`	`[0]`	Which global ranks to profile.
`profile_step_start`	`10`	First step to capture.
`profile_step_end`	`12`	Last step to capture.
`disable_profiler_activity_cpu`	`false`	Drop CPU-side activity to shrink traces (GPU-only trace).
`torch_profiler_record_shapes`	`true`	Record tensor shapes per op.
`torch_profiler_with_stack`	`true`	Record Python/C++ stacks (larger traces).
`torch_profiler_use_gzip`	`false`	Gzip the exported trace.

Enable via CLI overrides on train pretrain:

./runner/primus-cli direct -- train pretrain \
  --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml \
  --profile True \
  --use_pytorch_profiler True \
  --profile_step_start 5 \
  --profile_step_end 6 \
  --disable_profiler_activity_cpu False

Keep the capture window short (a few steps after warmup)—traces grow quickly, especially with with_stack and CPU activity enabled. Load the resulting trace in Perfetto to inspect CPU/GPU overlap, kernel launch delays, and idle gaps.

2. GPU memory profiling (Megatron)#

Two complementary mechanisms:

Memory history snapshot (trainer_base.yaml):

Parameter	Default	Purpose
`record_memory_history`	`false`	Record the CUDA/HIP allocator history for snapshot analysis.
`memory_snapshot_path`	`snapshot.pickle`	Output path for the allocator snapshot.

Load the pickle with PyTorch’s memory visualizer to find fragmentation and peak allocations.

ROCm memory sampling (primus_megatron_module.yaml):

Parameter	Default	Purpose
`use_rocm_mem_info`	`false`	When `true`, collect ROCm memory info via `rocm-smi` every iteration.
`use_rocm_mem_info_iters`	`[1, 2]`	When `use_rocm_mem_info=false`, only sample at these iterations.

Also relevant: log_memory_to_tensorboard (trainer_base.yaml) writes memory metrics to TensorBoard.

3. TraceLens automated trace analysis (Megatron)#

TraceLens turns raw profiler traces into hierarchical breakdowns (roofline/efficiency, compute-vs-memory bound kernels, communication-vs-sync separation, trace diffing). Primus can generate and optionally upload these reports. Configured in primus/configs/modules/megatron/primus_megatron_module.yaml:

Parameter	Default	Purpose
`generate_tracelens_report`	`false`	Generate TraceLens reports locally (auto-enabled when upload is on).
`mlflow_upload_tracelens_report`	`false`	Upload reports to MLflow (auto-enables generation, profiling, tensorboard).
`mlflow_tracelens_ranks`	`null`	Ranks to analyze (`null` = all; e.g. `[0, 8]` for one rank/node).
`mlflow_tracelens_output_format`	`xlsx`	`xlsx` (fastest), `csv`, or `all`.
`mlflow_tracelens_cleanup_after_upload`	`false`	Delete local reports after upload to save disk.
`mlflow_tracelens_auto_install`	`true`	Auto-install TraceLens if missing (set `false` to disable).

Related profiler/log uploads: mlflow_upload_traces (upload raw trace files) and mlflow_upload_logs (upload training logs). See Logging & experiment tracking for MLflow setup.

4. Performance metrics to MLflow (Megatron)#

mlflow_upload_performance_metrics: false (primus_megatron_module.yaml) enables a comprehensive scaling-test metric set when turned on (implicitly enabling throughput calculation):

perf/throughput_tflops_per_gpu, perf/tps_tokens_per_sec_per_gpu, perf/iteration_time_ms
perf/{rocm,hip}_current_mem_gb, perf/{rocm,hip}_mem_utilization_pct
perf/gpu_utilization_pct_rank{N}, perf/gpu_utilization_pct_avg

GPU utilization collection uses an all_gather every log_interval, which synchronizes ranks—keep this in mind for throughput-sensitive runs.

5. Profiling (TorchTitan)#

Configured under profiling: in primus/configs/modules/torchtitan/pre_trainer.yaml:

Parameter	Default	Purpose
`enable_profiling`	`false`	Enable the Torch profiler.
`profile_freq`	`10`	Capture every N steps.
`save_traces_folder`	`profile_traces`	Output folder for traces.
`enable_memory_snapshot`	`false`	Capture GPU memory snapshots.
`save_memory_snapshot_folder`	`memory_snapshot`	Output folder for snapshots.

Communication tracing is configured under comm: (trace_buf_size, save_traces_folder: comm_traces).

6. Profiling (MaxText)#

Configured in primus/configs/modules/maxtext/pre_trainer.yaml:

Parameter	Default	Purpose
`profiler`	`xplane`	Profiler backend (XPlane traces).
`skip_first_n_steps_for_profiler`	`3`	Warmup steps to skip before capture.
`profiler_steps`	`1`	Number of steps to capture.

7. Memory and performance projection#

Project resource usage before launching, without consuming a full cluster. Exposed through the Primus CLI projection subcommand (primus/cli/subcommands/projection.py).

# Memory projection from an experiment config
./primus-cli direct -- projection memory --config examples/megatron/configs/MI300X/llama2_7B-BF16-pretrain.yaml

# Performance projection (single-node benchmarking)
./primus-cli direct -- projection performance --config <exp>.yaml

# Performance projection scaled to N nodes (simulation)
./primus-cli direct -- projection performance --config <exp>.yaml --target-nodes 4

Memory projection breaks VRAM down across parameters, gradients, activations, optimizer states, and mixed-precision overhead—invaluable for MoE and ultra-large models where activations dominate. See Projection for the full reference.

8. Pipeline schedule visualization#

Diagnose pipeline bubbles and stage imbalance with the built-in tool tools/visualization/pp_vis/.

Dump per-rank schedule data during training with --dump_pp_data true (Megatron flag dump_pp_data, primus_megatron_module.yaml). Output lands under output/pp_data/ (config.json, pp_rank_*.json).
Install the tool deps and run the viewer:

pip install -r tools/visualization/pp_vis/requirements.txt
python tools/visualization/pp_vis/vis.py   # open http://127.0.0.1:8988

Configure task_list in vis.py to point at your dumped log_path and the iterations to render. The tool can also visualize the PP simulator output (see tools/visualization/pp_vis/README.md).

9. Recommended workflow#

Project first—run projection memory to confirm the config fits before booking GPUs.
Capture a short trace—a few steps after warmup with profile + use_pytorch_profiler.
Inspect—Perfetto for the timeline; TraceLens for automated hierarchical analysis and trace diffs.
Check memory—record_memory_history / ROCm sampling for fragmentation and peaks.
Check pipelines—dump_pp_data + pp_vis for bubbles and stage imbalance.