HRM SRAM/DRAM Memory Tiering: Training & Inference Guide

This guide explains how to train, benchmark, and evaluate both the Original (Baseline) and the New (Memory-Tiered) HRM models.

🚀 Setup

1. Prerequisites

Ensure you have a CUDA-capable GPU. The tiered model uses Python Triton for high-speed SRAM kernels.

2. Environment Setup (Recommended)

Create and activate a new virtual environment to keep dependencies isolated:

# Create venv
python -m venv venv

# Activate venv (Mac/Linux)
source venv/bin/activate

# Install dependencies
pip install --upgrade pip
pip install -r requirements.txt
pip install triton matplotlib

🏋️ Training the Models

Training is handled by pretrain.py using Hydra for configuration.

A. Train the Original (Baseline) Model

Uses standard GPU global memory (DRAM) for all layers.

# Example for Sudoku dataset
OMP_NUM_THREADS=8 python pretrain.py \
    arch=hrm_v1 \
    data_path=data/sudoku-extreme-1k-aug-1000 \
    epochs=10000 \
    global_batch_size=384

B. Train the New (Tiered) Model

Optimizes the L-level module by pinning it to simulated SRAM using Triton kernels.

OMP_NUM_THREADS=8 python pretrain.py \
    arch=hrm_tiered \
    data_path=data/sudoku-extreme-1k-aug-1000 \
    epochs=10000 \
    global_batch_size=384

📊 Testing Inference & Benchmarking

The run_benchmark.py script is the primary tool for measuring latencies, throughput, and memory efficiency.

1. Comparison Mode (Recommended)

Automatically runs both models across multiple batch sizes and sequence lengths, then generates performance charts.

python run_benchmark.py \
    --mode compare \
    --batch-sizes 1,8,32 \
    --seq-lens 64,128 \
    --plot \
    --output-dir benchmark_results/

Outputs: benchmark_results/benchmark_comparison.png and results.json.

2. Tiered-Only Benchmark (Deep Dive)

Get detailed metrics for just the tiered model, including memory transfer overhead and SRAM hit rates.

python run_benchmark.py \
    --mode tiered \
    --iterations 50 \
    --batch-sizes 8 \
    --output tiered_stats.json

✅ Evaluation

To verify the accuracy of a trained checkpoint on a test set:

Using the original evaluate script:

python evaluate.py checkpoint=checkpoints/PATH_TO_YOUR_STEP

Using the tiered model class manually:

If you are writing a custom script, you can import the models as follows:

from models.hrm.hrm_act_v1 import HierarchicalReasoningModel_ACTV1 # Original
from models.hrm.hrm_tiered import HRM_Tiered                       # Tiered

🛠 Model Location Reference

Folder/File	Contents
`models/hrm/hrm_act_v1.py`	Original model code (Baseline)
`models/hrm/hrm_tiered.py`	New Tiered model code
`models/triton_kernels.py`	Triton SRAM/DRAM kernels
`models/memory_tier.py`	Memory Tier Manager and CUDA stream logic
`config/arch/hrm_tiered.yaml`	Configuration for the tiered version