HRM_sudoku / GUIDE_TRAIN_TEST.md
Code2aum's picture
Upload folder using huggingface_hub
5dc80b3 verified
|
Raw
History Blame Contribute Delete
1.78 kB
# Training Guide: HRM SRAM/DRAM Learning Performance
Use this guide to train the memory-tiered HRM on **real datasets** and verify its reasoning accuracy.
---
### 1. Dataset Generation
The model needs to learn how to solve puzzles. First, generate a training/test set (e.g., Sudoku).
```bash
# Activate environment
source venv/bin/activate
# Build a small Sudoku dataset (1000 examples)
python dataset/build_sudoku_dataset.py \
--output-dir data/sudoku-1k \
--subsample-size 1000 \
--num-aug 10
```
---
### 2. Start Training
Point the trainer to the tiered architecture configuration.
#### Train the Tiered Model (New)
```bash
python pretrain.py \
arch=hrm_tiered \
data_path=data/sudoku-1k \
epochs=1000 \
global_batch_size=384
```
#### Train the Baseline Model (Comparison)
```bash
python pretrain.py \
arch=hrm_v1 \
data_path=data/sudoku-1k \
epochs=1000 \
global_batch_size=384
```
---
### 3. Verify Reasoning Accuracy
Once training is complete (or during training), check the accuracy metrics in your W&B dashboard or via the evaluation script:
```bash
# Replace with the path to your generated checkpoint
python evaluate.py checkpoint=checkpoints/Sudoku_ACT-torch/YOUR_RUN/step_1000
```
**What to look for:**
- **`eval/exact_accuracy`**: Fraction of puzzles solved perfectly.
- **`eval/steps`**: Average number of reasoning steps taken by the ACT (Adaptive Computation Time) module.
- **`eval/q_halt_accuracy`**: How well the model learns when to stop thinking.
---
### 4. Training on ARC (Artificial General Intelligence Benchmark)
To train on the more complex ARC-AGI-2 dataset:
```bash
python dataset/build_arc_dataset.py --output-dir data/arc-2
python pretrain.py arch=hrm_tiered data_path=data/arc-2
```