학습 처리량 3.7배 향상과 메모리 32% 절감

복잡한 거대언어모델을 튜닝할 때마다 메모리 부족으로 학습이 중단되거나, 병목 현상으로 인해 GPU가 제 성능을 내지 못하는 상황은 흔히 겪는 불편이다. 이제는 단순히 모델을 불러오는 단계를 넘어, 하드웨어 자원을 얼마나 효율적으로 점유하고 연산 처리량을 극대화할 수 있는지가 파인튜닝의 성패를 가른다. NVIDIA가 공개한 NeMo AutoModel은 기존 Transformers v5 라이브러리 기반의 MoE(Mixture of Experts, 전문가 혼합 모델) 학습 환경에서 코드 수정 없이도 학습 처리량을 3.4배에서 최대 3.7배까지 높이고, GPU 메모리 사용량을 29%에서 32%까지 절감하는 성능을 달성했다.

NeMo AutoModel은 HuggingFace의 표준 API인 from_pretrained()와 호환되도록 설계되어 기존 코드를 유지하면서도 즉각적인 성능 개선이 가능하다. 개발자는 복잡한 인프라 설정 없이 임포트 문을 변경하는 것만으로 최적화된 학습 환경을 구축할 수 있다. 다음은 기존 HuggingFace 모델 로드 방식을 NeMo AutoModel로 전환하는 예시 코드다.

python
from nemo.collections.nlp.models.automodel import NeMoAutoModelForCausalLM

이 라이브러리는 Expert Parallelism(전문가 병렬 처리)을 통해 모델의 전문가 가중치를 여러 GPU에 분산 배치함으로써 메모리 점유율을 획기적으로 낮춘다. 예를 들어 Qwen3 모델의 경우 기존 대비 피크 메모리를 68.2GiB에서 48.1GiB로 줄였으며, Nemotron 3 Nano 30B 모델은 62.1GiB에서 42.5GiB까지 메모리 사용량을 최적화했다. 이러한 메모리 여유는 더 큰 배치 사이즈를 설정하거나 더 긴 시퀀스를 학습할 수 있는 환경을 제공한다.

기술적 핵심은 DeepEP(Deep Expert Parallelism)와 TransformerEngine 커널의 결합에 있다. DeepEP는 토큰을 전문가에게 전달하는 디스패치(Dispatch) 과정과 결과를 합치는 과정을 GPU 커널로 통합하여, 통신과 전문가 연산을 동시에 수행하는 중첩(Overlap) 구조를 구현했다. 또한 TransformerEngine의 퓨즈드 어텐션(Fused Attention)과 선형 레이어 커널을 활용해 연산 효율을 높였다. NeMo AutoModel은 Qwen3, NVIDIA Nemotron, GPT-OSS, DeepSeek V3와 같은 주요 모델에 대해 수동으로 튜닝된 최적화 구현체를 제공하며, 지원되지 않는 모델의 경우 Liger 커널 패칭 등을 적용한 표준 방식으로 자동 전환되어 작동한다.

Expert Parallelism과 DeepEP의 최적화 기제

학습 효율은 모델의 크기가 아니라 데이터를 어떻게 배치하느냐에 따라 결정된다. 대규모 MoE 모델을 튜닝할 때 개발자가 가장 먼저 부딪히는 벽은 GPU 메모리 부족으로 인한 실행 불가 상태다. NeMo AutoModel은 전문가 병렬 처리(Expert Parallelism, EP)를 통해 이 문제를 해결한다. EP는 모델 내의 전문가 가중치를 단일 GPU가 아닌 여러 GPU에 분산해 배치하는 방식이다. Qwen3 모델의 경우 이 방식을 적용했을 때 피크 메모리가 68.2GiB에서 48.1GiB로 낮아졌다. 메모리 점유율이 약 29% 감소하면서 더 큰 배치 크기를 설정하거나 더 긴 문장 시퀀스를 처리할 수 있는 물리적 여유 공간이 확보된다.

단순히 메모리를 줄이는 것을 넘어 연산 속도를 높이려면 데이터가 이동하는 통신 시간을 없애야 한다. MoE 구조에서는 토큰을 적절한 전문가에게 보내는 전송(Dispatch)과 계산된 결과를 다시 합치는 결합(Combine) 과정에서 막대한 통신 비용이 발생한다. DeepEP는 이 전송과 결합 과정을 별도의 단계로 나누지 않고 하나의 GPU 커널로 통합해 처리한다. 이를 통해 데이터가 GPU 사이를 이동하는 통신 시간과 전문가가 실제 계산을 수행하는 연산 시간을 겹쳐서 처리하는 중첩(Overlap)을 구현했다. 통신이 연산의 발목을 잡는 병목 현상을 제거해 전체적인 학습 처리량을 높인다.

기초 연산 단계의 효율은 TransformerEngine(TE) 커널이 담당한다. TE는 퓨즈드 어텐션(Fused Attention, 여러 연산을 하나로 묶어 메모리 접근 횟수를 줄인 어텐션)과 선형 레이어, RMSNorm(평균과 분산 대신 제곱평균제곱근을 사용해 정규화하는 방식)을 최적화된 형태로 구현했다. 이는 MoE 레이어뿐만 아니라 모델 전체의 레이어 타입에서 PyTorch나 Flash Attention 같은 기존 구현체보다 일관된 속도 향상을 제공한다. 개별 연산의 실행 시간을 줄이고 메모리 읽기 쓰기 횟수를 최소화함으로써 H100과 같은 고성능 GPU의 하드웨어 자원 활용도를 극대화한다.

사용자는 `BackendConfig`를 통해 전문가 백엔드 설정을 제어하며 모델의 최적화 경로를 지정할 수 있다. 특히 다중 GPU 환경에서 학습을 수행하려면 장치 간의 논리적 연결 구조를 정의하는 `device_mesh` 설정이 필수적이다. 8개의 GPU를 사용하는 단일 노드 환경에서는 다음과 같이 설정하여 전문가 병렬 처리를 활성화하고 메모리 효율을 높인다.

python
device_mesh = DeviceMesh("cuda", [[0, 1, 2, 3, 4, 5, 6, 7]])

이 설정을 통해 모델은 FSDP2(완전 샤딩 데이터 병렬 처리 2세대)와 EP, DeepEP의 이점을 동시에 누리며 대규모 파라미터를 효율적으로 처리한다. 결과적으로 기존의 단순 데이터 병렬 처리 방식에서 겪었던 메모리 압박과 통신 지연 문제를 동시에 해결하는 구조를 갖춘다.

550B 모델 학습 가능 여부를 가른 라이브러리 격차

거대 모델을 학습시키다 보면 코드의 논리보다 하드웨어의 물리적 한계에 먼저 부딪히는 경우가 많다. 특히 파라미터 규모가 커질수록 메모리 부족으로 인해 프로세스가 강제 종료되는 현상은 엔지니어에게 매우 흔한 경험이다. Nemotron 3 Ultra 550B A55B 모델의 풀 파인튜닝 과정이 정확히 이 지점에 있었다. 풀 파인튜닝은 모델의 모든 파라미터를 업데이트하고 Adam 옵티마이저 상태를 물리적으로 생성하는 작업으로, 모델 크기에 비례해 메모리 소모가 극심하게 늘어난다. 이 모델은 Mamba2, LatentMoE, Multi-Token Prediction(MTP, 다음 토큰을 여러 개 동시에 예측하는 기술)이 결합된 550B 규모의 하이브리드 모델이다. 엔비디아는 16개 노드에 총 128개의 H100 GPU를 투입해 이 초거대 모델의 학습 가능 여부를 검증하며 라이브러리별 한계를 측정했다.

표준 라이브러리인 Transformers v4는 학습 시작 단계에서 데드락(두 프로세스가 서로의 자원을 기다리며 무한히 대기하는 상태) 현상으로 인해 실행에 실패했다. Qwen3 MoE 모델의 전문가들을 `ModuleList`라는 모듈 리스트 형태로 저장하고, 각각을 FSDP(Fully Sharded Data Parallel, 모델 가중치를 여러 GPU에 분산 저장하는 기술)로 개별 래핑했기 때문이다. 순전파 과정에서 입력 데이터에 따라 호출되는 전문가가 달라지자, 각 GPU 랭크마다 수행하는 FSDP 집합 통신인 AllGather와 ReduceScatter의 호출 순서가 불일치하며 시스템이 무한 대기 상태에 빠졌다. Transformers v5는 전문가를 퓨즈드 3D 파라미터 텐서(가중치를 하나의 3차원 배열로 통합 저장하는 방식)로 저장해 개별 전문가 단위의 FSDP 통신 과정을 완전히 제거함으로써 이 데드락 문제를 해결했다.

하지만 v5에서도 550B 규모의 모델은 결국 메모리 부족(OOM, Out of Memory)으로 인해 실행이 불가능했다. 텐서 구조를 개선해 통신 오류는 잡았지만, 단일 GPU가 감당해야 할 가중치의 절대량이 하드웨어 메모리 용량을 초과했기 때문이다. NeMo AutoModel은 여기서 Expert Parallelism(EP, 전문가 가중치를 여러 GPU에 나누어 배치하는 병렬 처리 방식)을 통해 돌파구를 찾았다. 전문가들을 GPU 전체에 샤딩(분산 저장)하여 개별 GPU가 보유해야 할 메모리 풋프린트를 가용 예산 범위 내로 낮춘 것이다. 결과적으로 표준 라이브러리가 메모리 한계로 인해 벤치마크 수치조차 기록하지 못한 550B 모델의 풀 파인튜닝을 완수하며, 단순한 연산 가속을 넘어 모델 규모의 한계를 결정짓는 라이브러리 설계의 격차를 증명했다.

HuggingFace 생태계 유지와 실무 도입 가치

새로운 프레임워크를 도입할 때 가장 큰 부담은 기존에 짜놓은 코드를 전부 갈아엎어야 한다는 시간적 손실이다. 수개월의 개발 공수가 들어간 파이프라인을 포기하고 인프라를 옮기는 결정은 매우 무겁다. NeMo AutoModel은 `AutoModelForCausalLM`을 상속받아 구현되었다. 상속이란 기존 클래스의 기능을 그대로 물려받는 프로그래밍 방식으로, 덕분에 실무자는 HuggingFace 모델용으로 작성한 기존 코드를 수정 없이 그대로 사용할 수 있다. 라이브러리 임포트 문 하나만 바꾸면 기존의 모델 로드 및 학습 흐름이 유지된다.

특정 모델군에는 하드웨어 수준에서 최적화된 전용 구현체를 제공한다. Qwen3, NVIDIA Nemotron, GPT-OSS, DeepSeek V3 모델이 이에 해당한다. 이 모델들은 TransformerEngine(트랜스포머 엔진, 연산 가속 라이브러리)의 어텐션과 퓨즈드 선형 레이어(여러 연산을 하나로 합쳐 처리하는 층), 그리고 커스텀 전문가 커널을 통해 구동된다. 모델마다 다른 체크포인트 구조를 일일이 맞추는 번거로운 작업인 체크포인트 플러밍(Checkpoint Plumbing) 과정을 생략하고, 재사용 가능한 핵심 연산 최적화에만 집중한 결과다.

지원 목록에 포함되지 않은 모델의 경우 vanilla HF 방식으로 폴백(Fallback, 기본 설정으로 되돌아감) 작동한다. 단순히 기본 상태로 돌아가는 것이 아니라 Liger 커널 패칭(메모리 효율을 높이는 커널 수정) 같은 최적화 기법을 적용해 최소한의 성능 이득을 보장한다. 어떤 경로를 통해 모델이 로드되든 장치 메시(Device Mesh, GPU 간 통신 구조) 설정만 추가하면 추가적인 코드 수정 없이 멀티 GPU 학습 환경으로 확장할 수 있다.

학습 종료 후 가중치를 저장하는 단계에서도 생태계 호환성을 유지한다. `save_pretrained()` 함수를 호출하면 표준 HuggingFace 체크포인트 형식으로 저장된다. 이렇게 저장된 결과물은 vLLM(고속 LLM 추론 엔진)이나 SGLang(구조화된 언어 생성 프레임워크) 같은 외부 추론 도구에서 별도의 변환 과정 없이 즉시 로드할 수 있다. 인프라 교체 비용을 들이지 않고도 최신 MoE 최적화 기술을 실무에 즉시 적용할 수 있는 구조다. 개발자는 모델 아키텍처나 추론 엔진을 바꾸는 리스크 없이 학습 효율과 처리량이라는 실질적인 이득만 취하면 된다.

Transformers v5 대비 최대 3.7배의 처리량 향상과 32%의 메모리 절감은 하드웨어 한계로 포기했던 550B 규모 모델의 풀 파인튜닝을 실무 영역으로 끌어온 결과다. 개발자는 기존 HuggingFace 기반 코드를 유지한 채 임포트 문을 변경하는 것만으로 대규모 MoE 모델의 학습 효율을 즉시 높이는 선택을 할 수 있다. 이제 라이브러리의 선택은 단순한 편의를 넘어 학습 가능한 모델의 물리적 상한선을 결정하는 핵심 변수가 된다.