VRAM_OOM_NOTES_2026-05-24.md

KoHRM-Text VRAM / OOM Notes

작성일: 2026-05-24

이 문서는 KoHRM-Text-1.4B stage-1 학습 중 VRAM이 시간이 지나며 증가하는 이유, 이전 OOM 원인, 현재 운영 기준을 기록합니다.

현재 관측 상태

현재 stage-1 run은 다음 설정으로 정상 학습 중입니다.

항목	값
GPU	8 x NVIDIA H200
GPU utilization	8장 모두 99%
global batch	180,224 tokens
local token slots/GPU	22,528
context	4,096
VRAM	GPU0 약 129.9GB, 나머지 약 127.6GB / 143.8GB
speed	약 1.02 step/sec
checkpoint interval	5,000 steps

현재 설정은 빠르지만 여유 VRAM이 아주 넓은 편은 아닙니다. H200 장당 약 144GB 중 127-130GB를 사용하므로, NCCL/allocator/compiler/cache/checkpoint 순간 피크가 겹치면 OOM 위험이 다시 생길 수 있습니다.

왜 학습 중 VRAM이 점점 올라가나

VRAM 증가가 곧바로 “메모리 누수”라는 뜻은 아닙니다. 대형 PyTorch/FSDP/compile 학습에서는 다음 요인이 겹치면서 초반보다 뒤에서 VRAM이 더 높아지는 패턴이 흔합니다.

1. torch.compile / CUDA graph / kernel cache

HRM-Text 코드는 여러 forward/backward path를 compile합니다. 초반 몇 step에서는 모든 shape/path가 아직 compile되지 않았고, 학습이 진행되며 추가 graph, Triton kernel, CUDA kernel cache가 만들어집니다.

특히 HRM 구조는 H/L recurrent cycle과 PrefixLM loss가 있어 단순 decoder-only Transformer보다 compile path가 더 복잡합니다. 초반 VRAM만 보고 batch를 크게 잡으면 후속 graph가 생성될 때 추가 메모리를 못 받아 OOM이 날 수 있습니다.

2. final logits buffer 크기

이번 모델은 vocab이 131,072입니다. upstream HRM-Text 논문 설정의 65,536 vocab보다 두 배입니다.

batch token slots가 커질수록 final logits 또는 loss 계산 쪽 임시 버퍼가 매우 커집니다.

예를 들어 local token slots/GPU가 32,768이면 32768 x 131072 bf16 logits 계열 버퍼가 필요할 수 있습니다. 이론상 단일 bf16 dense buffer만 잡아도 약 8GB 이상이고, 실제 backward/temporary/parallel buffer까지 합치면 훨씬 커집니다.

이 때문에 처음에는 global_batch_size=262144 또는 229376이 잠깐 돌아가도, 뒤에서 compile graph와 logits/loss 임시 버퍼가 겹치는 순간 OOM이 날 수 있습니다.

3. FSDP2 / optimizer / EMA 상태

현재 학습은 model weights만 들고 있는 것이 아닙니다.

• model parameters
• gradients
• optimizer state
• Adam-atan2 state
• EMA state
• FSDP shard/all-gather/reduce-scatter buffers
• recurrent carry 관련 state

이 상태들이 step마다 항상 같은 순간에 같은 크기로 보이는 것은 아닙니다. 특정 backward path, optimizer step, checkpoint save 시점에 피크가 올라갈 수 있습니다.

4. NCCL communication buffers

8 GPU 분산 학습에서는 NCCL 통신 버퍼가 필요합니다. all-gather/reduce-scatter 타이밍, bucket 크기, compile된 그래프 실행 순서에 따라 GPU별 피크가 다르게 보일 수 있습니다.

GPU0이 다른 GPU보다 더 높게 보이는 것도 일반적으로 가능합니다. rank0가 로깅, 일부 metadata, checkpoint coordination, dataloader/host interaction을 더 맡는 경우가 있기 때문입니다.

5. CUDA caching allocator

nvidia-smi의 used memory는 “현재 텐서가 실제로 쓰는 메모리”만 뜻하지 않습니다. PyTorch CUDA allocator가 한 번 확보한 블록을 재사용하려고 캐시에 잡고 있으면 nvidia-smi에는 계속 사용 중처럼 보입니다.

따라서 step이 진행될수록 used memory가 올라가고 잘 내려가지 않는 것은 정상일 수 있습니다. 중요한 것은 reserved가 계속 무한 증가하는지, 또는 특정 step 이후 안정 plateau를 만드는지입니다.

6. checkpoint 저장 시 순간 피크

FSDP2 checkpoint 저장 시 .distcp shard, metadata, state_dict materialization, host/device transfer가 겹칩니다. 저장 자체는 주로 CPU/disk 작업이지만, 저장 직전/직후 모델 state 접근 때문에 GPU/CPU 메모리 피크가 생길 수 있습니다.

그래서 너무 잦은 checkpoint 저장은 다음 문제를 만듭니다.

• step 처리 지연
• 디스크 사용량 급증
• HF upload 및 scan 비용 증가
• 저장 시점 피크 메모리 증가

현재 5,000 step마다 약 21GB급 FSDP2 checkpoint가 생깁니다. 500 step마다 저장하면 stage-1 기준으로 체크포인트 수와 저장 부하가 10배 늘어 과합니다.

이전 OOM 원인

이전 OOM은 batch를 크게 잡았을 때 초반 관측 VRAM만 보고 “괜찮다”고 판단한 것이 원인입니다.

핵심은 다음입니다.

1. vocab 131K라 logits/loss 관련 임시 버퍼가 큽니다.
2. HRM recurrent compile path가 초반 몇 step 뒤 추가 메모리를 요구합니다.
3. H200 8장이라 compute는 충분하지만, 1.4B + 131K vocab + EMA + optimizer + FSDP2 조합에서는 batch를 너무 크게 잡으면 후반 피크가 걸립니다.
4. global_batch_size=262144, 229376은 초반에는 가능해 보였지만 안정 마진이 부족했습니다.

현재는 global_batch_size=180224로 내려 안정 진행 중입니다.

운영 기준

현재 stage-1에서는 GPU를 놀리지 않는 것이 우선이지만, OOM으로 run이 죽으면 재시작/검증/체크포인트 정리 비용이 더 큽니다.

권장 기준:

항목	기준
primary batch	global_batch_size=180224
저장 주기	checkpoint_step_interval=5000
로컬 보관	최신 2-3개 checkpoint만 유지
HF main repo	최신 safetensors export 중심
HF raw repo	resume가 필요한 FSDP2 checkpoint만 별도 보관
OOM 재발 시	batch를 5-10% 낮추고 같은 resume checkpoint에서 재시작

500 step checkpoint가 과한 이유

500 step마다 저장하면 다음 문제가 생깁니다.

• 현재 FSDP2 checkpoint 하나가 약 21GB입니다.
• 500 step 간격이면 10,000 step마다 약 20개, 즉 약 420GB가 생깁니다.
• stage-1 전체 88,522 step 기준으로는 단순 계산상 170개 이상이 생겨 수 TB가 됩니다.
• 저장 자체가 학습 루프를 방해하고, HF 업로드/스캔도 커집니다.

따라서 현재처럼 5,000 step 간격으로 저장하고, 로컬은 최신 2-3개만 남기는 편이 맞습니다.

다음 batch 조정 판단

현재 VRAM 사용량은 높지만 학습 속도는 안정적입니다.

다음 stage에서 batch를 올리고 싶으면 한 번에 크게 올리지 말고 다음 순서가 낫습니다.

1. global_batch_size=180224로 안정 완료 확인
2. 다음 dataset stage에서 196608 테스트
3. 2-3천 step 이상 VRAM plateau 확인
4. checkpoint 저장 시점까지 통과하면 유지
5. OOM 또는 피크 불안정 시 즉시 180224 또는 172032로 복귀

논문 설정과 비교하면 H200 8장은 강하지만, 이번 모델은 vocab이 131K라 upstream과 메모리 구조가 다릅니다. 따라서 “H200이니까 무조건 H100 16장 batch를 넘긴다”는 식으로 잡으면 안정성이 떨어집니다.

결론

현재 VRAM 상승은 torch compile/cache, 131K vocab logits buffer, FSDP2/optimizer/EMA/NCCL buffer, checkpoint 순간 피크가 겹친 결과로 보는 것이 맞습니다.

현재 global_batch_size=180224, 5,000 step checkpoint, 최신 2-3개 보관 정책은 빠른 학습과 OOM 회피 사이의 현실적인 균형입니다. 학습이 완전히 안정 plateau를 보이면 다음 stage에서만 소폭 증량을 검토합니다.