METHODOLOGY_ARCHITECTURE_NOTES_2026-05-24.md

KoHRM-Text Methodology and Architecture Notes

작성일: 2026-05-24

이 문서는 KoHRM-Text-1.4B가 HRM-Text 논문 방식과 어떤 점에서 같고, 어떤 점에서 운영상 다른지 정리합니다.

참고 문서:

• HRM-Text paper: https://arxiv.org/html/2605.20613
• Upstream code: https://github.com/sapientinc/HRM-Text
• KoHRM-Text code: https://github.com/LLM-OS-Models/KoHRM-text

결론

우리의 현재 학습은 방법론 관점에서는 HRM-Text 논문식 single-stage instruction pretraining에 맞춰져 있습니다.

다만 실행 운영 관점에서는 논문과 완전히 같지 않습니다. 논문은 40B unique tokens를 단일 연속 run으로 학습했고, 중간 checkpointing/crash recovery를 쓰지 않았다고 설명합니다. 우리는 데이터 준비와 HF 업로드, OOM 회피, 체크포인트 보존 때문에 stage-0, stage0b, stage-1처럼 나누어 실행하고 있습니다.

핵심 차이는 다음입니다.

항목	HRM-Text 논문	KoHRM-Text 현재 방식
학습 목적	instruction-response task-completion objective	동일
loss	response-only NLL	동일
attention	PrefixLM, instruction bidirectional + response causal	동일 코드 경로 사용
raw LM pretraining 후 SFT	하지 않음	하지 않음
SFT 후보 데이터	instruction pretraining에 포함	포함
실행 형태	단일 연속 run	staged resume run
checkpoint	논문은 중간 checkpointing 없음	운영상 5,000 step마다 저장
tokenizer	65,536 BPE	131,072 Korean/terminal BPE
hardware	16 x H100	8 x H200

따라서 “논문처럼 single-stage 지시문 사전학습인가?”에 대한 답은 다음처럼 정리하는 것이 정확합니다.

학습 objective와 데이터 포맷은 single-stage instruction pretraining입니다.
그러나 실행은 한 프로세스의 단일 연속 run이 아니라, 같은 objective를 유지하면서 checkpoint resume으로 이어가는 staged pretraining입니다.

논문 방법론 요약

HRM-Text 논문은 기존 LLM의 대규모 raw-text causal LM 사전학습 후 mid-training/SFT로 가는 다단계 recipe를 비효율적이라고 보고, 처음부터 instruction-response pair만으로 학습합니다.

논문 핵심:

1. raw text 전체 토큰을 예측하지 않습니다.
2. instruction tokens에는 loss를 주지 않습니다.
3. response tokens에만 negative log-likelihood loss를 줍니다.
4. instruction segment는 PrefixLM mask로 bidirectional attention을 허용합니다.
5. response segment는 autoregressive causal attention을 유지합니다.
6. 데이터는 direct, cot, synth, noisy 같은 condition tag를 instruction 앞에 붙여 응답 스타일을 제어합니다.
7. <think>...</think> 같은 explicit long-CoT trace는 제거해 내부 recurrent computation이 문제 해결을 맡게 합니다.

논문은 1B HRM-Text를 scratch로 학습했고, 약 40B unique tokens 및 16 x H100에서 약 46시간을 사용했다고 설명합니다. 공개 평가에는 EMA checkpoint를 사용합니다.

현재 KoHRM-Text 학습 방식

현재 KoHRM-Text-1.4B도 raw causal LM이 아니라 HRM-Text V1Dataset 포맷으로 학습합니다.

각 샘플은 기본적으로 다음 구조를 갖습니다.

instruction span -> response span

토큰 레벨에서는 다음처럼 구성됩니다.

<|im_start|> condition_tokens instruction <|im_end|> response <|box_end|>

condition token은 다음 mapping을 사용합니다.

condition	token
direct	`<
cot	`<
noisy	`<
synth	`<

dataset_new.py 기준으로 instruction span은 inputs에는 들어가지만 target_only=True일 때 label은 IGNORE_LABEL_ID가 됩니다. response span만 실제 cross entropy loss를 받습니다.

즉 현재 학습은 “문서를 처음부터 끝까지 다 맞히는 raw LM”이 아니라 “주어진 지시/문맥을 보고 응답을 완성하는 task-completion pretraining”입니다.

PrefixLM 구현 확인

현재 코드에서 PrefixLM 경로는 다음 파일들로 확인됩니다.

파일	역할
dataset_new.py	instruction/response span 분리, response-only label 구성
models/flash_attention_prefixlm_v2.py	prefix bidirectional attention + response causal attention 구현
models/layers.py	attention type prefixlm 사용
models/lm_head.py	IGNORE_LABEL_ID를 제외하고 response label에만 CE loss 계산

dataset_new.py에서는 각 샘플에 대해 instruction 길이를 prefix_lens, response 길이를 causal_lens로 넘깁니다.

flash_attention_prefixlm_v2.py는 attention을 두 부분으로 처리합니다.

1. prefix 구간: instruction tokens끼리 bidirectional attention
2. causal 구간: response tokens가 prefix 전체와 이전 response tokens를 보는 causal attention

이 구조가 논문의 PrefixLM과 맞는 핵심입니다.

모델 아키텍처

현재 표준 모델명은 KoHRM-Text-1.4B이고, arch/size@arch=XL에 해당하는 설정을 사용합니다.

항목	값
params	1,384,120,320
hidden size	1,536
total configured layers	32
half layers	true
H module layers	16
L module layers	16
heads	12
head dim	128
expansion	4
intermediate size	4,096
context	4,096
RoPE theta	10,000
norm	RMSNorm-style parameterless norm
init	LeCun normal
dtype	bfloat16
tokenizer vocab	131,072

HRM recurrent schedule:

항목	값
H cycles	2
L cycles per H cycle	3
effective H/L recurrence	H2L3
bp min steps	2
bp max steps	5
bp warmup ratio	0.2

코드상 흐름은 다음과 같습니다.

1. token embedding에서 시작한 hidden state를 z_H로 둡니다.
2. learned/fixed low-level initial state z_L_init에서 z_L을 시작합니다.
3. 각 H cycle마다 L module을 3번 반복 업데이트합니다.
4. 그 뒤 H module을 1번 업데이트합니다.
5. 총 2번의 H cycle을 수행합니다.
6. 최종 z_H에 LM head를 붙여 vocabulary logits를 냅니다.

논문 표현으로는 slow-evolving strategic layer인 H module과 fast-evolving execution layer인 L module의 dual-timescale recurrent design입니다.

MagicNorm / 안정화

논문은 recurrent depth가 깊어지면 activation variance와 gradient instability가 커지므로 MagicNorm과 warmup deep credit assignment를 사용합니다.

현재 코드에서는 norm_type: pre를 쓰되, Transformer module 끝에 final RMSNorm을 적용합니다. 즉 내부 block은 PreNorm 스타일이고 module output은 norm으로 capped됩니다. 이것이 논문에서 말하는 MagicNorm 계열 안정화와 대응됩니다.

Backward는 처음부터 모든 recurrent step을 다 통과하지 않고, bp_steps를 warmup합니다.

현재 설정:

bp_warmup_ratio: 0.2
bp_min_steps: 2
bp_max_steps: 5

초기에는 마지막 2 recurrent steps 위주로 gradient를 흘리고, 학습이 진행되며 최대 5 steps까지 늘립니다. 이 점도 논문 방식과 맞습니다.

Optimizer / 스케줄

현재 pretraining config:

항목	값
optimizer	Adam-atan2
beta1	0.9
beta2	0.95
weight decay	0.1
lr	2.2e-4
lr warmup	2,000 steps
lr min ratio	1.0
EMA	0.9999
gradient clipping	없음

논문도 Adam-atan2, warmup 2,000 steps, weight decay 0.1, EMA 0.9999, bf16을 사용합니다. 공개/평가는 EMA checkpoint를 기준으로 합니다.

현재 staged run이 논문과 다른 이유

논문은 “single continuous run”이라고 설명합니다. 우리는 다음 현실적 이유 때문에 staged run으로 운영합니다.

1. 한국어/터미널/tokenizer 전처리가 순차적으로 끝나고 있습니다.
2. GPU를 놀리지 않기 위해 준비된 데이터부터 학습을 시작했습니다.
3. H200 8장 환경에서 131K vocab 때문에 batch OOM 안정 마진을 실측해야 했습니다.
4. HF 업로드와 raw checkpoint 보존이 필요합니다.
5. full HRM 328G no-cap retokenization이 아직 진행 중입니다.

하지만 각 stage가 다른 목적함수로 바뀌는 것은 아닙니다.

Stage	Objective	성격
stage-0	PrefixLM response-only	준비 완료된 711.3M mix
stage0b	PrefixLM response-only	같은 mix 추가 pass
stage-1	PrefixLM response-only	HRM fast-cap 14.55B
later stage	PrefixLM response-only	full HRM 328G retokenized + 한국어/터미널/툴콜 mix
final SFT	PrefixLM response-only 또는 SFT용 response-only	품질 높은 subset으로 후처리

중요한 점은 stage-0에서 stage-1로 넘어갈 때 model/optimizer/EMA/carry를 이어받고, resume_step_offset과 total_steps_override로 global step/LR schedule을 이어가도록 수정했다는 것입니다.

즉 “학습 방법론”은 single-stage instruction pretraining이고, “운영 방식”은 staged continuation입니다.

현재 데이터가 single-stage 원칙에 맞는지

현재 prepared dataset들은 전부 가능한 한 instruction-response 형태로 변환했습니다.

데이터	single-stage 원칙 적합성
HRM cleaned data	원래 HRM instruction/response/condition 구조라 적합
ToolBench	tool instruction -> tool-call/answer response 구조라 적합
SWE-ZERO/local terminal	terminal context -> next action/answer 구조라 적합
GLM/Claude reasoning	final answer 중심으로 정리하면 적합
한국어 법률/위키 원문	chunked instruction/response task로 바꿔 투입하므로 적합

주의할 점:

• 한국어 위키/법률 raw chunk는 “그냥 다음 텍스트 예측”처럼 넣으면 논문식 task-completion에서 멀어집니다.
• 따라서 title/context를 instruction으로 두고 chunk/summary/extraction을 response로 두는 식이 더 맞습니다.
• local terminal dataset은 objective에 잘 맞지만 전체 비중이 너무 커지면 일반 지식/한국어 균형이 무너질 수 있습니다.

현재 방식 평가

현재 방식은 논문 핵심과 잘 맞습니다.

맞는 부분:

• scratch training
• HRM H2L3 recurrent architecture
• PrefixLM attention
• response-only loss
• condition token 사용
• Adam-atan2
• bf16
• EMA 0.9999
• 4,096 context
• LeCun normal init

다른 부분:

• vocab 65,536이 아니라 131,072입니다.
• 16 x H100이 아니라 8 x H200입니다.
• 논문은 단일 연속 run, 우리는 staged resume run입니다.
• 논문은 40B unique tokens를 보고했고, 현재 public checkpoint는 stage-1 fast-cap 중간 산출물입니다.
• 한국어/터미널/툴콜 비중이 논문보다 훨씬 큽니다.

위 차이는 의도된 변경입니다. 특히 한국어/터미널/툴콜을 목표로 하므로 tokenizer와 데이터 분포를 바꾼 것은 맞습니다. 다만 논문의 효율성을 재현하려면 최종적으로 full HRM cleaned data와 balanced Korean/terminal/tool mix를 합쳐 40B+ token 수준으로 이어 학습해야 합니다.

운영 결론

현재는 다음 기준으로 계속 가는 것이 맞습니다.

1. 현재 stage-1은 계속 유지합니다.
2. 전처리 완료 데이터는 HF dataset repo에 올려 다른 머신에서도 재현 가능하게 둡니다.
3. full HRM 328G no-cap retokenization이 끝나면 next stage로 이어 학습합니다.
4. SFT 후보 데이터도 pretraining에 먼저 포함합니다.
5. 별도 final SFT는 마지막에 품질 높은 subset으로 다시 수행합니다.
6. model repo는 최신 safetensors 중심, raw checkpoint repo는 resume용으로 분리합니다.