DeepSeek Engram과 OLMo-core를 활용한 분산 학습 구현

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

기존의 피드포워드 네트워크(FFN)나 혼합 전문가(MoE) 구조를 대체하기 위해 DeepSeek의 Engram 메모리 레이어를 OLMo-core 아키텍처에 이식했습니다. 이 구조는 토큰 라우팅을 통해 메모리 행렬 전체가 아닌 특정 Engram 블록만 활성화하며, GPU 간 통신 오버헤드를 최소화하도록 설계되었습니다. AI2의 내부 클러스터 환경에 최적화된 OLMo-core를 일반 클라우드인 RunPod에서 실행하기 위해 환경 구축 파이프라인을 재설계하고 독자적인 텔레메트리 우회 기법을 적용했습니다. 4개의 NVIDIA A40 GPU 클러스터에서 3.5B 파라미터 모델로 테스트한 결과, 장치당 초당 36개 토큰의 처리량과 안정적인 손실값 하락을 기록하며 분산 학습의 안정성을 확인했습니다. 이번 연구는 고성능 프레임워크의 이식성을 높이고 새로운 메모리 아키텍처의 실무 적용 가능성을 제시했다는 점에서 의미가 있습니다.

의미 / 영향

특정 기업의 슈퍼컴퓨터 환경에 종속된 고성능 AI 프레임워크를 일반 클라우드 인프라로 이식하는 구체적인 방법론을 제시했으며, 새로운 메모리 아키텍처인 Engram의 실전 배치 가능성을 확인했습니다.

빠른 이해

요약 브리프

DeepSeek의 Engram 메모리 구조를 OLMo-core 프레임워크에 이식하여 분산 학습에 성공한 사례입니다. 일반 클라우드 환경에서 발생할 수 있는 I/O 및 종속성 문제를 해결하고 3.5B 모델의 안정적인 학습 성능을 수치로 입증했습니다.

새로운 점

AI2의 폐쇄적인 OLMo-core 환경을 RunPod으로 이식하고 DeepSeek의 최신 Engram 레이어를 결합한 최초의 오픈소스 PoC입니다.

핵심 메커니즘

입력 토큰 → Engram 라우팅 알고리즘 → 특정 메모리 블록 활성화 → HSDP 기반 분산 처리 → 안정적인 그래디언트 업데이트 및 손실값 하락

핵심 수치

Throughput: ~36 Tokens / Second / Device- 4x A40 GPU 환경
Memory Reserved: 28.6 GiB- A40 48GB 중 효율적 사용
Training Loss: 12.21 → 12.05- 196 steps 기준

섹션별 상세

Engram 아키텍처 설계 및 통합

표준 FFN 대신 Engram 메모리 레이어를 도입하여 토큰이 동적으로 특정 메모리 블록에 라우팅되도록 설계했습니다. 이 방식은 전체 메모리 행렬을 활성화하지 않아 연산 효율성을 높이며, 1D 및 2D 블록 병렬 처리를 통해 다중 GPU 환경에서 조회 단계의 통신 비용을 줄입니다. OLMo-core의 HSDP(Fully Sharded Data Parallel) 및 torchrun 분산 토폴로지와 기본적으로 통합되어 대규모 확장이 가능합니다.

RunPod 환경에서의 MLOps 최적화

AI2 내부 시스템인 Beaker와 Gantry에 종속된 OLMo-core를 독립적인 클라우드 환경에서 실행하기 위한 최적화 전략을 수립했습니다. I/O 병목을 피하기 위해 네트워크 볼륨 대신 100GB NVMe 컨테이너 디스크를 할당하고, 환경 구축 스크립트를 통해 flash-attn 등 CUDA 커널 컴파일 시간을 단축했습니다. 또한 sed 명령어를 사용하여 소스 코드 내의 독자적인 모니터링 데코레이터를 제거함으로써 외부 클러스터에서의 실행 차단 요소를 해결했습니다.

분산 학습 검증 및 성능 지표

4개의 NVIDIA A40 GPU로 구성된 클러스터에서 OLMo 3B 모델에 Engram을 추가한 3.5B 파라미터 설정으로 학습 안정성을 검증했습니다. 장치당 약 36 Tokens/sec의 처리량을 유지하며 48GB 메모리 중 약 28.6 GiB를 예약하여 효율적인 자원 사용률을 보였습니다. 196단계의 학습 과정에서 교차 엔트로피 손실이 12.21에서 12.05로 안정적으로 하락했으며, 수치적 불안정성이나 데드락 현상은 발생하지 않았습니다.

실무 Takeaway

DeepSeek Engram 레이어를 OLMo-core에 통합하여 FFN 대비 연산 효율성을 개선하고 분산 환경에서의 확장성을 확보함
RunPod과 같은 범용 클라우드에서 OLMo-core를 실행하기 위해 NVMe 기반의 빠른 환경 구축 파이프라인과 텔레메트리 우회 스크립트를 적용함
4x A40 GPU 환경에서 3.5B 모델 학습 시 장치당 36 Tokens/sec의 처리량과 안정적인 그래디언트 흐름을 수치로 증명함

언급된 리소스

GitHubOLMo-core Engram PoC GitHub Repository

문서원문 링크