# 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:

```bash
# 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.

```bash
# 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.

```bash
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.

```bash
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.

```bash
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:
```bash
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:

```python
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 |