๐ŸŽฏ Primus Auto Benchmarking Tool

December 18, 2025 ยท View on GitHub

Primus Banner

An interactive bash script for automated benchmarking of LLMs on AMD GPUs (MI300X/MI355X) using Megatron or TorchTitan backends supported through Primus.


๐Ÿš€ Quick Start

Step 1: Pull and Launch the Container

docker pull YOUR_IMAGE
docker run -it \
  --device /dev/dri \
  --device /dev/kfd \
  --network host \
  --ipc host \
  --group-add video \
  --cap-add SYS_PTRACE \
  --security-opt seccomp=unconfined \
  --privileged \
  -v $HOME/.ssh:/root/.ssh \
  --name IMAGE_NAME \
  YOUR_IMAGE

Step 2: Navigate to Auto Benchmark Directory

cd /workspace/Primus/tools/auto_benchmark/

Step 3: Run the Benchmarking Tool

bash run_primus_autobenchmark.sh

๐Ÿ“‹ Features

  • โœ… Interactive Menu System - User-friendly CLI with color-coded outputs and ASCII banner
  • โœ… Multi-Backend Support - Compatible with Megatron and TorchTitan with device-specific configs
  • โœ… Batch Processing - Run multiple model configurations sequentially with flexible selection
  • โœ… Configuration Viewing - Preview YAML configs before execution
  • โœ… Configuration Editing - Edit YAML configs individually or in batch before execution
  • โœ… Parameter Overrides - Override specific parameters without editing files permanently
  • โœ… Auto Device Detection - Automatically detects AMD MI300X/MI355X GPUs with intelligent fallback
  • โœ… Device-Specific Paths - Automatically uses device-specific config directories (MI300X/MI355X)
  • โœ… Comprehensive Logging - Timestamped logs saved in organized backend-specific directories
  • โœ… Environment Management - Custom device-specific environment variable support
  • โœ… Automatic Metrics Generation - Backend-specific metrics tables generated after completion
  • โœ… Smart Config Management - Handles edited/override configs properly with automatic cleanup

๐Ÿ“– Complete Walkthrough

1๏ธโƒฃ Backend Selection

When you launch the tool, you'll first choose the backend framework:

โ˜… Choose Backend:
  โ— 1) megatron
  โ— 2) torchtitan

โžœ Enter number or name:

Options:

  • Enter 1 or megatron for Megatron backend
  • Enter 2 or torchtitan for TorchTitan backend

The backend selection determines:

  • Which config directory to use
  • Which metrics script to run after completion
  • Where logs are saved

2๏ธโƒฃ Device Detection

The tool automatically detects your AMD GPU with intelligent fallback:

โ˜… Detecting Device...
  โ— Device found: MI300X
โœ“ GPU Device: MI300X

โœ“ Config directory set to: /workspace/Primus/examples/megatron/configs/MI300X

Auto-detection methods (in order):

  1. Queries rocminfo for "AMD Instinct" devices (direct model name)
  2. Falls back to architecture detection (gfx942 โ†’ MI300X, gfx950 โ†’ MI355X)
  3. Manual selection prompt if both methods fail

Manual Selection (if auto-detection fails):

โœ— Could not detect device automatically
โ˜… Please select Device manually:
  โ— 1) MI300X
  โ— 2) MI355X

โžœ Enter number or name:

Device-Specific Paths: After detection, the config directory is automatically set to:

  • Megatron MI300X: /workspace/Primus/examples/megatron/configs/MI300X
  • Megatron MI355X: /workspace/Primus/examples/megatron/configs/MI355X
  • TorchTitan MI300X: /workspace/Primus/examples/torchtitan/configs/MI300X
  • TorchTitan MI355X: /workspace/Primus/examples/torchtitan/configs/MI355X

3๏ธโƒฃ Model Configuration Selection

The tool scans for available YAML configuration files in the device-specific backend directory:

โ˜… Available Model Configs: (megatron / MI300X)
  โ— 1) llama3_8b.yaml
  โ— 2) llama3_70b.yaml
  โ— 3) qwen2.5_7B-FP8-pretrain.yaml

โžœ Select config number(s) (comma-separated, range, or 'all'):
(Examples: 1,3,5 or 4-8 or all)

Selection Options:

  • Single: 1 - Select one config
  • Multiple: 1,3,5 - Select specific configs (comma-separated)
  • Range: 4-8 - Select a range of configs
  • All: all - Select all available configs

Note: The tool automatically filters out duplicate configs to prevent redundant processing.


4๏ธโƒฃ View Configuration Parameters

Option to preview parameters in your selected configurations:

โ˜… View Configuration Parameters?
โžœ (y/n):

If you choose y, the tool displays the contents of each selected YAML file (excluding comments and empty lines):

Parameters in llama3_8b.yaml:
-----------------------------------
batch_size: 16
learning_rate: 0.0001
max_steps: 1000
seq_length: 2048
-----------------------------------

5๏ธโƒฃ Edit Configuration Files

For Multiple Configs:

โ˜… Edit any configuration files before running?
โžœ (y/n):

If y, you can select which configs to edit:

Selected models:
  โ— 1) llama3_8b.yaml
  โ— 2) llama3_70b.yaml

โ— Enter model numbers to edit (comma-separated, or 'all'):
โžœ

For Single Config:

โ˜… Edit configuration file before running?
โžœ (y/n):

The tool creates a temporary working copy in /tmp and opens it in your default editor (tries nano, vim, vi, code, or $EDITOR).

Important:

  • Edits are applied to the original config location temporarily during benchmark execution
  • Original configs are backed up and restored after each benchmark completes
  • Edited configs are preserved in logs directory for reproducibility

6๏ธโƒฃ Override Parameters

Override specific parameters without editing the entire file:

โ˜… Override any parameters?
  (Format: key=value, e.g., batch_size=32)
โžœ (y/n):

If y, enter overrides one per line:

โžœ Override (or press Enter to finish): batch_size=32
โœ“ Will override: batch_size = 32
โžœ Override (or press Enter to finish): learning_rate=0.001
โœ“ Will override: learning_rate = 0.001
โžœ Override (or press Enter to finish): [Press Enter]

โœ“ 2 parameter(s) will be overridden

Override Behavior:

  • Creates {MODEL}_{BACKEND}_{DEVICE}_{TIMESTAMP}_override.yaml in logs directory
  • Applies overrides using sed for precise YAML modification
  • Can be combined with edited configs
  • Original config is temporarily replaced during execution then restored

7๏ธโƒฃ Device-Specific Environment Variables

Add custom environment variables for your device:

โ˜… Add device-specific environment variables for MI300X?
  (e.g., HSA_OVERRIDE_GFX_VERSION=11.0.0)
โžœ (y/n):

If y, enter variables one per line:

โžœ Variable (or press Enter to finish): HSA_OVERRIDE_GFX_VERSION=11.0.0
โœ“ Will set: HSA_OVERRIDE_GFX_VERSION=11.0.0
โžœ Variable (or press Enter to finish): ROCR_VISIBLE_DEVICES=0,1,2,3
โœ“ Will set: ROCR_VISIBLE_DEVICES=0,1,2,3
โžœ Variable (or press Enter to finish): [Press Enter]

โœ“ 2 environment variable(s) will be set

Format: VAR_NAME=value (allows empty values)


8๏ธโƒฃ Environment Setup

The tool configures the environment:

โ˜… Setting up environment...
โœ“ Set HSA_NO_SCRATCH_RECLAIM=1
โœ“ Set HSA_OVERRIDE_GFX_VERSION=11.0.0
โœ“ Set ROCR_VISIBLE_DEVICES=0,1,2,3
โžœ Enter HuggingFace Token: [hidden input]
โœ“ HuggingFace token set

Automatic settings:

  • HSA_NO_SCRATCH_RECLAIM=1 (always set for AMD GPUs)
  • Any custom environment variables you added
  • HF_TOKEN for HuggingFace authentication (hidden input for security)

9๏ธโƒฃ Benchmark Execution

The tool runs benchmarks for all selected configurations sequentially:

โ„น Total configurations to run: 2
โ„น Configuration list:
   โ— 1. llama3_8b.yaml
   โ— 2. qwen2.5_7B-FP8-pretrain.yaml

โ•”โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•—
โ•‘  LOOP ITERATION: 1/2
โ•‘  CONFIG FILE: llama3_8b.yaml
โ•šโ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•โ•

โ˜… Starting Benchmark 1/2...
   โ— Model: llama3_8b
   โ— Backend: megatron
   โ— Device: MI300X
   โ— Config: /workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45_override.yaml
   โ— Log: /workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45.log

โœ“ Copied edited/overridden config to: /workspace/Primus/examples/megatron/configs/MI300X/llama3_8b.yaml
โœ“ EXP set to: /workspace/Primus/examples/megatron/configs/MI300X/llama3_8b.yaml

  โ— Changing to Primus root directory: /workspace/Primus

[Benchmark output streams here...]

โœ“ Restored original config file

==========================================
 โœ“ Benchmark 1/2 Completed Successfully!
 Log saved at:
   /workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45.log
 Override config saved at:
   /workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45_override.yaml
==========================================

Preparing next benchmark...

โ„น Next: Config 2/2

[Continues with next benchmark...]

For each benchmark:

  1. Displays iteration header with model info
  2. Uses edited config if available, otherwise uses original
  3. Applies parameter overrides to create timestamped override config in logs
  4. Temporarily replaces original config with edited/overridden version
  5. Changes to Primus root directory (/workspace/Primus) for proper path resolution
  6. Exports EXP environment variable pointing to the device-specific config path
  7. Executes /workspace/Primus/examples/run_pretrain.sh
  8. Streams output to both terminal and timestamped log file
  9. Restores original config file after completion
  10. Shows completion status with log file location
  11. Adds 2-second delay between benchmarks for system stability

Error Handling:

  • Script continues even if a benchmark fails
  • Exit codes are captured and displayed
  • Warnings shown for non-zero exit codes

๐Ÿ”Ÿ Metrics Generation

After all benchmarks complete, backend-specific metrics are automatically generated:

=========================================
  All 2 Benchmark(s) Completed!
=========================================

โ˜… Generating Metrics Table...

โœ“ Running: python metrics_megatron.py

[Metrics table displayed here...]

โœ“ Metrics table generated successfully

Automatic Metrics Scripts:

  • Megatron: Runs metrics_megatron.py
  • TorchTitan: Runs metrics_torchtitan.py

The metrics scripts parse the log files and generate formatted summary tables with performance statistics.


๐Ÿ“ Output Files

All output files are organized in the results/ directory within the auto_benchmark tool:

Directory Structure

/workspace/Primus/tools/auto_benchmark/
โ”œโ”€โ”€ run_primus_autobenchmark.sh
โ”œโ”€โ”€ metrics_megatron.py
โ”œโ”€โ”€ metrics_torchtitan.py
โ””โ”€โ”€ results/
    โ”œโ”€โ”€ logs_megatron/
    โ”‚   โ”œโ”€โ”€ llama3_8b_megatron_MI300X_2025-12-17_10-30-45.log
    โ”‚   โ”œโ”€โ”€ llama3_8b_megatron_MI300X_2025-12-17_10-30-45_override.yaml
    โ”‚   โ””โ”€โ”€ llama3_70b_megatron_MI300X_2025-12-17_11-15-20.log
    โ””โ”€โ”€ logs_torchtitan/
        โ”œโ”€โ”€ qwen2.5_7B_torchtitan_MI355X_2025-12-17_14-20-30.log
        โ””โ”€โ”€ qwen2.5_7B_torchtitan_MI355X_2025-12-17_14-20-30_edited.yaml

Log Files

Format: {CONFIG_FILENAME}_{BACKEND}_{DEVICE}_{TIMESTAMP}.log

Examples:

  • /workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45.log
  • /workspace/Primus/tools/auto_benchmark/results/logs_torchtitan/qwen2.5_7B_torchtitan_MI355X_2025-12-17_14-20-30.log

Benefits:

  • Full config filename preserved (not just model name)
  • Backend-specific directories for easy organization
  • Complete benchmark output with metrics
  • Timestamped for version tracking

Override Config Files

Format: {CONFIG_FILENAME}_{BACKEND}_{DEVICE}_{TIMESTAMP}_override.yaml

Example:

/workspace/Primus/tools/auto_benchmark/results/logs_megatron/llama3_8b_megatron_MI300X_2025-12-17_10-30-45_override.yaml

Contents: Copy of original config with parameter overrides applied

Edited Config Files

Format: {CONFIG_FILENAME}_{BACKEND}_{DEVICE}_{TIMESTAMP}_edited.yaml

Example:

/workspace/Primus/tools/auto_benchmark/results/logs_torchtitan/qwen2.5_7B_torchtitan_MI355X_2025-12-17_14-20-30_edited.yaml

Contents: Copy of the manually edited config used for the benchmark


๐Ÿ’ก Tips & Best Practices

  1. Batch Processing: Use all or ranges (e.g., 1-5) to benchmark multiple models efficiently
  2. Device-Specific Configs: Ensure configs exist in the correct device subdirectory (MI300X/MI355X)
  3. Parameter Overrides: Use overrides for quick experiments without modifying config files permanently
  4. Log Organization: Logs are automatically organized by backend in results/logs_{backend}/
  5. Environment Variables: Add device-specific tuning variables (e.g., HSA_OVERRIDE_GFX_VERSION) for optimal performance
  6. Config Editing: Edited configs are applied during execution but originals are preserved
  7. View Before Running: Always preview configs before execution to verify parameters
  8. Metrics Analysis: Backend-specific metrics scripts automatically parse logs after completion
  9. Sequential Runs: 2-second delay between benchmarks ensures system stability
  10. Error Resilience: Script continues running even if individual benchmarks fail
  11. Path Resolution: Script automatically changes to Primus root directory for proper execution
  12. Config Backup: Original configs are automatically backed up and restored after each benchmark

๐Ÿ› ๏ธ Technical Details

Directory Structure

/workspace/Primus/
โ”œโ”€โ”€ examples/
โ”‚   โ”œโ”€โ”€ megatron/
โ”‚   โ”‚   โ”œโ”€โ”€ configs/
โ”‚   โ”‚   โ”‚   โ”œโ”€โ”€ MI300X/           # MI300X-specific Megatron configs
โ”‚   โ”‚   โ”‚   โ””โ”€โ”€ MI355X/           # MI355X-specific Megatron configs
โ”‚   โ”‚   โ””โ”€โ”€ prepare.py
โ”‚   โ”œโ”€โ”€ torchtitan/
โ”‚   โ”‚   โ”œโ”€โ”€ configs/
โ”‚   โ”‚   โ”‚   โ”œโ”€โ”€ MI300X/           # MI300X-specific TorchTitan configs
โ”‚   โ”‚   โ”‚   โ””โ”€โ”€ MI355X/           # MI355X-specific TorchTitan configs
โ”‚   โ”‚   โ””โ”€โ”€ prepare.py
โ”‚   โ””โ”€โ”€ run_pretrain.sh            # Main benchmark execution script
โ””โ”€โ”€ tools/
    โ””โ”€โ”€ auto_benchmark/
        โ”œโ”€โ”€ run_primus_autobenchmark.sh
        โ”œโ”€โ”€ metrics_megatron.py
        โ”œโ”€โ”€ metrics_torchtitan.py
        โ””โ”€โ”€ results/
            โ”œโ”€โ”€ logs_megatron/
            โ””โ”€โ”€ logs_torchtitan/

Environment Variables Set

  • HSA_NO_SCRATCH_RECLAIM=1 (always set for AMD GPUs)
  • HF_TOKEN (user-provided for HuggingFace access)
  • Custom device-specific variables (optional, user-defined)
  • EXP (config path for each benchmark, format: {BACKEND_BASE_DIR}/{DEVICE}/{config}.yaml)

Execution Flow

  1. Backend selection โ†’ Sets BACKEND_BASE_DIR
  2. Device detection โ†’ Sets DEVICE
  3. Config directory construction โ†’ CONFIG_DIR = {BACKEND_BASE_DIR}/{DEVICE}
  4. Config selection โ†’ Scans device-specific directory
  5. Config editing/overrides โ†’ Creates temporary working copies
  6. Environment setup โ†’ Exports required variables
  7. For each benchmark:
    • Backup original config
    • Copy edited/override config to original location
    • Change to /workspace/Primus directory
    • Export EXP variable
    • Execute run_pretrain.sh
    • Restore original config
    • Save logs to results/logs_{backend}/
  8. Generate metrics table using backend-specific script

Supported Editors (Priority Order)

  1. nano - Simple terminal editor
  2. vim - Advanced terminal editor
  3. vi - Classic Unix editor
  4. code - VS Code with --wait flag
  5. $EDITOR - Environment variable fallback

Device Detection Logic

1. Query rocminfo for "AMD Instinct" โ†’ Extract model name (MI300X/MI355X)
2. If empty or invalid โ†’ Query rocminfo for architecture (gfx942 โ†’ MI300X, gfx950 โ†’ MI355X)
3. If still empty โ†’ Prompt for manual selection (1=MI300X, 2=MI355X)

Config Management Strategy

  • Original configs: Never permanently modified
  • Edited configs: Created in /tmp, temporarily replace originals during execution
  • Override configs: Created in logs directory with all changes applied
  • Backup mechanism: .backup_$$ files created before replacement, restored after execution
  • Log preservation: All config versions saved to logs directory for reproducibility

๐Ÿ“ Example Session

# Full example workflow
cd /workspace/Primus/tools/auto_benchmark/
bash run_primus_autobenchmark.sh

# Interactive prompts:
# Backend: megatron (enter '1')
# Device: Auto-detected as MI300X
# Configs: 1,2 (select llama3_8b and qwen2.5_7B)
# View parameters: y (preview configs)
# Edit configs: n (skip editing)
# Override parameters: y
#   batch_size=64
#   learning_rate=0.0005
# Add env vars: y
#   HSA_OVERRIDE_GFX_VERSION=11.0.0
# HF token: hf_xxxxxxxxxxxxx

# Output:
# โœ“ 2 benchmarks run sequentially
# โœ“ Logs saved to results/logs_megatron/
# โœ“ Override configs saved with logs
# โœ“ Metrics table automatically generated
# โœ“ Original configs remain unchanged

๐Ÿ”ง Troubleshooting

Common Issues

Issue: find: 'examples/megatron/configs/MI300X': No such file or directory

  • Cause: Config directory doesn't exist for the detected device
  • Solution: Ensure configs exist in /workspace/Primus/examples/{backend}/configs/{device}/

Issue: Backend prepare script not found

  • Cause: Script not executing from Primus root directory
  • Solution: Script now automatically changes to /workspace/Primus before execution

Issue: Edited config not being used

  • Cause: Config not properly copied to original location
  • Solution: Script now backs up and replaces original configs during execution

Issue: Log files not found

  • Cause: Logs saved in wrong directory
  • Solution: Check results/logs_{backend}/ in auto_benchmark directory

๐Ÿ†˜ Support

For issues or questions:

  1. Check log files in /workspace/Primus/tools/auto_benchmark/results/logs_{backend}/
  2. Verify ROCm installation: rocminfo
  3. Ensure device-specific configs exist in proper directories
  4. Verify script is run from /workspace/Primus/tools/auto_benchmark/
  5. Check that run_pretrain.sh exists in /workspace/Primus/examples/
  6. Review this README for proper usage

๐Ÿ“Š Backend-Specific Metrics

Megatron Metrics

  • TPS (Tokens Per Second) - Throughput metric
  • TFLOPS (Tera FLOPs) - Compute performance
  • Memory (%) - GPU memory utilization
  • Time (ms) - Elapsed execution time

TorchTitan Metrics

  • TPS (Tokens Per Second) - Throughput metric
  • TFLOPS (Tera FLOPs) - Compute performance
  • MFU (Model FLOPs Utilization) - Efficiency metric
  • Memory (%) - GPU memory utilization

Note: Metrics are extracted automatically by running metrics_{backend}.py after all benchmarks complete.


Happy Benchmarking! ๐Ÿš€