Skip to content

profiling-and-observability.md

Latest commit

 

History

History
170 lines (115 loc) · 8.26 KB

profiling-and-observability.md

File metadata and controls

170 lines (115 loc) · 8.26 KB
Raw

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/.

  1. 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).
  2. 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

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

Related documentation