File size: 3,081 Bytes
5dc80b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | # 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 |
|