[Paper] MemFine: 메모리 인식 세밀한 스케줄링 for MoE Training

Source: arXiv - 2511.21431v1

Overview

대규모 Mixture‑of‑Experts (MoE) 모델 학습이 심각한 장벽에 부딪히고 있습니다: 토큰의 동적 라우팅이 큰 로드 불균형을 초래해 GPU 메모리 사용량이 급증하기 때문입니다. MemFine은 토큰 스트림과 전문가 작업량을 더 작은 “청크”로 나누는 메모리 인식 세밀 스케줄링 시스템을 도입하여, 트레이너가 필요한 부분만 재계산하면서 일반적인 GPU의 메모리 한도 내에 머물 수 있게 합니다. 그 결과, 이전에는 메모리 부족으로 실행이 불가능했던 하드웨어에서도 실행 가능한 보다 확장 가능한 MoE 학습 파이프라인을 제공합니다.

Key Contributions

• Chunk‑Based Decomposition: 토큰 분배와 전문가 연산을 관리 가능한 조각으로 분해하여 선택적 재계산을 가능하게 합니다.
• Theoretical Memory Model: 메모리 소비를 예측하고 스케줄러의 실시간 의사결정을 안내하는 폐쇄형 식을 제공합니다.
• Dynamic Scheduling Algorithm: 수동 튜닝 없이 메모리 절감과 연산 처리량 사이의 트레이드‑오프를 즉시 최적화합니다.
• Empirical Gains: 기존 전체 재계산 방식 대비 **48 %**의 활성화 메모리 감소와 **4.4 %**의 처리량 향상을 보여줍니다.
• Hardware‑Friendly Design: 메모리가 제한된 일반 GPU에서도 동작하도록 설계되어, 특수 하드웨어나 대규모 배치 사이즈 감소가 필요 없습니다.

Methodology

1. Token & Expert Chunking – 전체 배치를 하나의 덩어리로 다루는 대신, MemFine은 들어오는 토큰 스트림과 전문가 집합을 더 작은 청크로 분할합니다. 각 청크는 독립적으로 처리될 수 있어 단일 연산에 필요한 피크 메모리를 제한합니다.
2. Chunked Recomputation – 역전파를 위해 활성화를 저장해야 하는 레이어에서, MemFine은 활성화를 선택적으로 버리고 역전파 시 필요한 청크만 재계산합니다. 이는 메모리 모델에 의해 가이드되어 재계산 비용이 메모리 절감 효과를 초과하지 않도록 합니다.
3. Memory Model‑Driven Scheduler – 가벼운 분석 모델이 각 청크 구성의 메모리 발자국을 추정합니다. 스케줄러는 GPU 메모리 예산을 만족하면서 처리량을 최대화하는 구성을 선택합니다. 이 결정은 매 학습 단계마다 업데이트되어 MoE의 끊임없이 변하는 토큰 라우팅 패턴에 적응합니다.
4. Implementation Details – 인기 있는 딥러닝 프레임워크(예: PyTorch)에 기존 MoE 디스패처를 대체하는 형태로 통합되어, 사용자 측 코드 변경이 최소화됩니다.

Results & Findings

• Memory Savings: 활성화 메모리를 거의 절반으로 줄여 16 GB GPU에서 이전에 충돌하던 모델도 안정적으로 학습할 수 있게 되었습니다.
• Throughput: 세밀한 스케줄링이 추가하는 재계산 오버헤드가 매우 작아, 메모리로 인한 정체를 감소시켜 전체 처리량이 향상됩니다.
• Accuracy: MemFine은 버려진 활성화에 대해 정확히 동일한 순전파를 재계산하기 때문에 최종 모델 품질에 거의 영향을 주지 않습니다.

Practical Implications

• Cost‑Effective Scaling: 기업은 비싼 다중 GPU 서버를 구매하지 않고도 수십억 파라미터 규모의 MoE 모델을 확장할 수 있습니다; 16‑32 GB GPU 하나만으로도 많은 작업을 수행할 수 있습니다.
• Simpler DevOps: 메모리에 맞추기 위해 전문가 용량이나 배치 사이즈를 수동으로 조정할 필요가 없으며, MemFine의 스케줄러가 자동으로 처리해 엔지니어링 부담을 줄여줍니다.
• Broader Accessibility: 제한된 하드웨어 예산을 가진 연구자와 스타트업도 이전에 접근하기 어려웠던 최첨단 MoE 아키텍처를 실험할 수 있습니다.
• Integration Path: MemFine은 기존 MoE 라이브러리에 플러그인 형태로 삽입되므로, 몇 줄의 코드만 추가하면 즉시 메모리 이점을 얻을 수 있으며 모델 로직을 재작성할 필요가 없습니다.

Limitations & Future Work

• Recomputation Overhead on Very Large Batches: 스케줄러가 이를 완화하지만, 매우 큰 배치 사이즈에서는 여전히 눈에 띄는 재계산 지연이 발생할 수 있습니다.
• Hardware Diversity: 현재 평가가 NVIDIA GPU에 초점을 맞추고 있어, TPU나 AMD GPU로 확장하려면 추가 튜닝이 필요할 수 있습니다.
• Dynamic Expert Count: MemFine은 레이어당 전문가 수가 고정돼 있다고 가정하므로, 학습 중 전문가를 추가/제거하는 모델을 다루는 것은 아직 해결되지 않은 과제입니다.
• Future Directions: 저자들은 런타임 프로파일링에 기반한 적응형 청크 크기 탐색, MemFine을 혼합 정밀도 학습과 결합, 그리고 보다 일반화된 스케줄러 API를 오픈소스로 제공해 커뮤니티 채택을 확대할 계획입니다.

Authors

• Lu Zhao
• Rong Shi
• Shaoqing Zhang
• Yueqiang Chen
• Baoguo He
• Hongfeng Sun
• Ziqing Yin
• Shangchao Su
• Zhiyan Cui
• Liang Dong
• Xiyuan Li
• Lingbin Wang
• Jianwei He
• Jiesong Ma
• Weikang Huang
• Jianglei Tong
• Dongdong Gao
• Jian Zhang
• Hong Tian

Paper Information

• arXiv ID: 2511.21431v1
• Categories: cs.DC
• Published: November 26, 2025
• PDF: Download PDF