HRM-Text (GitHub 저장소)

HRM-Text는 HRM 아키텍처 기반의 1B 텍스트 생성 모델로, task completion과 latent space reasoning을 강화해 사전학습 효율을 높였다. 저장소는 PrefixLM sequence packing, FlashAttention 3, PyTorch FSDP2 학습, 평가, 체크포인트 변환까지 포함한 전체 사전학습 프레임워크를 제공한다. 아키텍처는 HRM, Transformer, TRM, RINS, UT baseline을 지원한다.

공개된 기준에서는 foundation model 사전학습 비용을 약 1000달러 수준까지 낮추는 것을 목표로 한다. 기존 대비 130~600배 적은 연산량과 150~900배 적은 데이터를 내세우며, 참고 실행은 L(0.6B) 이 8 H100 / 50시간 / 약 800달러, XL(1B) 이 16 H100 / 46시간 / 약 1,472달러다.

벤치마크에서는 L이 GSM8k 77.6%, MATH 51.2%, MMLU 56.6%를 기록했고, XL은 GSM8k 84.7%, MATH 56.5%, MMLU 60.7%, ARC-C 81.9%, BoolQ 86.2%까지 올랐다. Hopper-class GPU가 권장되며, attention 경로가 FlashAttention 3에 의존한다.

• data_io로 코퍼스를 정제·토큰화한 뒤 stratified sampling을 수행한다. 기본 학습은 4 epoch 기준이므로 샘플링도 epochs=4에 맞춘다.
• 단일 노드는 /dev/shm/sampled에, 멀티노드는 공유 스토리지에 샘플을 두고 각 노드가 동일한 데이터를 재생성한다.
• Docker 이미지 또는 소스 설치를 지원하고, 멀티노드는 NCCL 통신을 먼저 검증한 뒤 wandb login으로 실험 추적을 붙인다.
• torchrun으로 학습하고 에폭마다 checkpoints/에 저장한다. evaluation.main은 최신 체크포인트를 자동으로 불러오며, conversion.convert_to_hf로 Transformers 포맷으로 export하고 평가와 export는 EMA weights를 기본으로 쓴다.

현재 저장소에는 학습, 체크포인트, 평가가 구현돼 있고, Transformers 포맷 export도 제공된다. Native Transformers 지원은 병합됐고 다음 릴리스에 포함될 예정이며, native vLLM 지원은 진행 중이다.