HRM-MoE: Efficient Sparse Pretraining with Hierarchical Reasoning

Model Structure Comparison

HRM-MoE is a sparse Mixture-of-Experts extension of sapientinc/HRM-Text. It keeps the original HRM-Text recipe for hierarchical recurrent modeling, PrefixLM sequence packing, FlashAttention 3, PyTorch FSDP2 training, checkpointing, evaluation, and conversion, while replacing the dense FFN path with a routed MoE FFN.

Config	Layers	Hidden	Heads	FFN / experts	Active FFN width per token	Total parameters	Status
HRM-Text XL dense	32	1536	12	dense SwiGLU, intermediate 4096	4096	~1.18B	dense baseline
HRM-MoE XL_moe64x8_grouped_triton	32	1536	12	64 routed SwiGLU experts, top-k 8, expert width 512	8 x 512 = 4096	~5.41B	released checkpoint
HRM-MoE XL_moe128x8	32	1536	12	128 routed SwiGLU experts, top-k 8, expert width 512	8 x 512 = 4096	~10.25B	experimental comparison

Both MoE variants keep the per-token activated FFN width aligned with the dense XL HRM-Text FFN: only 8 experts are active for each token, so the activated FFN width remains 4096. The 128x8 run doubles the inactive expert pool compared with 64x8 while keeping the same per-token activation. In our completed pretraining comparison, that extra sparse capacity did not consistently improve quality over 64x8; the Hugging Face release therefore uses the 64x8 checkpoint.

The released 64x8 path uses fp32 router softmax, normalized top-k routing, auxiliary load-balancing loss, and grouped Triton GEMMs for expert compute.

32-GPU Pretraining Results

The table below compares the 32-GPU dense XL run, the released 64x8 HRM-MoE run, and the completed 128x8 same-activation experiment on the same sampled HRM pretraining data and global_batch_size=196608. All three columns report completed epoch-4 checkpoint evaluation.

Benchmark	Metric	Dense XL epoch 4	HRM-MoE 64x8 epoch 4	HRM-MoE 128x8 epoch 4
GSM8k	acc	83.93	84.99	86.20
MATH	acc	54.96	60.08	59.56
DROP	em	79.45	80.86	80.25
DROP	f1	83.06	84.53	83.86
MMLU	acc	61.38	61.18	62.80
ARC	acc	83.02	87.80	86.95
HellaSwag	acc	61.96	73.89	70.93
Winogrande	acc	71.98	73.88	73.95
BoolQ	acc	87.25	88.75	88.44
MMLU-Pro	acc	32.72	37.57	34.44
AIME25	maj_pass@1	13.33	16.67	20.00
AIME25	maj_pass@10	36.67	36.67	40.00
AIME25	maj_pass@100	53.33	56.67	46.67

64x8 is the stronger default at epoch 4: it wins on MATH, DROP, ARC, HellaSwag, BoolQ, MMLU-Pro, and AIME25 maj_pass@100. The 128x8 experiment improves GSM8k, MMLU, Winogrande, and AIME25 maj_pass@1/10, but it also shows a much higher MMLU-Pro invalid rate (9.94% vs 1.57% for 64x8), so we do not treat it as the primary release model.

Launch the MoE Pretraining

Required Resources

The intended training target is Hopper-class GPUs because the attention path depends on FlashAttention 3 and the final MoE expert path uses Triton kernels.

The final MoE preset is designed for multi-GPU pretraining:

Model	Preset	GPUs	Notes
HRM-MoE XL 64x8	XL_moe64x8_grouped_triton	8+ H100/H200	use shared storage for data and checkpoints

1. Prepare Data

HRM-MoE trains from the same sampled, tokenized data layout as HRM-Text. The training data.path directory must contain:

metadata.json
tokens.npy
epoch_<n>/inst_start.npy
epoch_<n>/inst_len.npy
epoch_<n>/resp_start.npy
epoch_<n>/resp_len.npy

Prepare sampled data with the companion sapientinc/data_io pipeline:

cd <DATA_IO_PATH>
python sample_tokenized.py epochs=4 output_path=/path/to/sampled > show_analytics.md

For reusable experiments, keep sampled data on shared storage and pass it with data.path=/path/to/sampled. Only use /dev/shm for short same-container debugging.

2. Start the Environment

The upstream tested Docker image is:

docker run --gpus all --ipc=host --network=host -it \
  -v "$PWD":/workspace \
  sapientai/hrm-text:latest

If you install from source, follow the tested CUDA, PyTorch, and FlashAttention versions in docker/Dockerfile, then run:

pip install -r requirements.txt

For multi-node training, mount the same workspace and checkpoint path on every node. Verify NCCL before starting a long job.

3. Launch Pretraining

Single-node example:

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
HRM_MOE_TRITON_AUTOTUNE=1 \
HRM_MOE_TRITON_SM_MARGIN=16 \
torchrun --nproc_per_node=8 pretrain.py \
  arch/size@arch=XL_moe64x8_grouped_triton \
  data.path=/path/to/sampled \
  global_batch_size=196608

Multi-node example:

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
HRM_MOE_TRITON_AUTOTUNE=1 \
HRM_MOE_TRITON_SM_MARGIN=16 \
torchrun \
  --nproc_per_node=8 \
  --nnodes=<NUM_NODES> \
  --node_rank=<NODE_RANK> \
  --master_addr=<MASTER_ADDR> \
  --master_port=<MASTER_PORT> \
  pretrain.py \
  arch/size@arch=XL_moe64x8_grouped_triton \
  data.path=/path/to/sampled

Useful MoE training switches:

• allow_compile_moe=false is the default; sparse routing currently runs in eager mode for stability.
• fsdp_wrap_moe_experts=true wraps packed expert parameters separately when possible.
• HRM_MOE_PROFILE=1 prints MoE phase timings for router, dispatch, grouped GEMMs, activation, combine, and auxiliary metrics.

On the rjob cluster, the same final preset can be launched through the wrapper:

num_gpus=8 \
arch_size=XL_moe64x8_grouped_triton \
data_path=/path/to/sampled \
bash scripts/rjob_hrm_pretrain.sh

Use dry_run=true before submitting a real rjob launch.

4. Evaluate

Evaluation loads the latest checkpoint epoch automatically when ckpt_epoch is not provided:

python -m evaluation.main ckpt_path=/path/to/checkpoint_dir

To run a benchmark subset and lower memory use:

python -m evaluation.main \
  ckpt_path=/path/to/checkpoint_dir \
  run_only='[GSM8k,MATH]' \
  generation_config.batch_size=16

For 8-GPU fanout evaluation on rjob:

ckpt_path=/path/to/checkpoint_dir \
num_gpus=8 \
batch_size=16 \
entrypoint=scripts/hrm_eval_fanout_entrypoint.sh \
bash scripts/rjob_hrm_eval.sh

Fine-Tuning (SFT)

Continue from a pretrain checkpoint on instruction data. Full-parameter SFT uses the same pretrain.py entrypoint with --config-name cfg_sft.

Input is a JSONL file with one object per line:

{"instruction": "<full prompt>", "response": "<expected output>", "condition": "direct"}

Prepare SFT data:

python scripts/prepare_sft_data.py \
  --train input.jsonl \
  --tokenizer /path/to/tokenizer.json \
  --output /path/to/sft_data \
  --epochs 5

Launch SFT:

OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 \
HRM_MOE_TRITON_AUTOTUNE=1 \
HRM_MOE_TRITON_SM_MARGIN=16 \
torchrun --nproc_per_node=8 pretrain.py \
  --config-name cfg_sft \
  arch/size@arch=XL_moe64x8_grouped_triton \
  data.path=/path/to/sft_data \
  resume_from=/path/to/pretrain_ckpt \
  +checkpoint_path=/path/to/sft_out

--epochs for data preparation must match the SFT training config. Add weights_only_resume_from_ema=true when fine-tuning from pretrain EMA weights with a fresh optimizer.

Verify the Final MoE Path

Run a local equivalence smoke test:

python scripts/test_moe_shard_equivalence.py

Run the CUDA/rjob MoE gate before trusting a kernel or routing change:

bash scripts/rjob_hrm_moe_equiv.sh

Repository Layout

HRM-MoE/
|-- config/                       # Hydra configs for model, data, and training
|-- config/arch/size/XL_moe64x8_grouped_triton.yaml
|-- conversion/convert_to_hf.py    # FSDP2 checkpoint -> HF-style export
|-- evaluation/                    # Evaluation engines, benchmark wrappers, configs
|-- models/layers.py               # Attention, dense FFN, and final sparse MoE layer
|-- models/moe_triton_grouped_gemm.py
|-- models/moe_profile.py
|-- docker/                        # Tested CUDA/PyTorch/FlashAttention environment
|-- scripts/                       # Data prep, rjob launchers, eval, validation
|-- dataset_new.py                 # PrefixLM packed dataset loader
|-- multipack_sampler.py           # Distributed multipack batch sampler
|-- pretrain.py                    # FSDP2 pretraining/SFT entrypoint
`-- simple_inference_engine.py     # Native checkpoint inference helper

Technical Notes

• models/layers.py contains the final sparse MoE FFN path: fp32 router softmax, top-k 8 routing, grouped Triton expert compute, weighted combine, and auxiliary load-balancing loss.
• models/moe_triton_grouped_gemm.py implements the grouped Triton expert GEMM path used by the final preset.
• models/moe_profile.py records optional CUDA event timings when HRM_MOE_PROFILE=1.
• dataset_new.py loads PrefixLM packed samples and emits FlashAttention sequence metadata.
• pretrain.py handles Hydra config, FSDP2 wrapping, optimizer creation, LR schedule, W&B logging, code/config snapshots, and checkpointing.

Upstream

This repository is derived from HRM-Text:

• Paper: HRM-Text: Efficient Pretraining Beyond Scaling
• Upstream code: sapientinc/HRM-Text
• Upstream model: sapientinc/HRM-Text-1B

If you use the original HRM-Text work, please cite:

@misc{wang2026hrmtextefficientpretrainingscaling,
      title={HRM-Text: Efficient Pretraining Beyond Scaling},
      author={Guan Wang and Changling Liu and Chenyu Wang and Cai Zhou and Yuhao Sun and Yifei Wu and Shuai Zhen and Luca Scimeca and Yasin Abbasi Yadkori},
      year={2026},
      eprint={2605.20613},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2605.20613},
}

License

Apache License 2.0