JanusPipe: Efficient Pipeline Parallel Training for Machine Learning Interatomic Potentials

Motivation

• Known: conservative MLIP는 potential energy를 예측하고 이를 atomic position에 대해 미분하여 force를 얻는 방식(F=-∇xE)으로 물리적 보존성(energy conservation)을 만족시키며, 최근 LLM과 유사하게 모델과 데이터를 키우면 정확도와 일반화 성능이 향상되는 scaling-law 경향을 보인다는 것이 알려져 있다.
• Gap: 기존 distributed training system, 특히 1F1B 같은 pipeline parallelism 스케줄은 첫 번째 순서(first-order) 워크로드(LLM, CNN 등)를 가정하여 micro-batch당 forward 한 번, backward 한 번을 전제로 설계되어 있는데, conservative MLIP는 forward 단계 내에서 gradient를 계산하는 double-backward(4단계) 패턴을 가지므로 이를 그대로 적용하면 중복 recomputation, parameter replication으로 인한 메모리 낭비와 phase별 실행시간의 partial order(tFE<tFF<tBE<tBF)로 인한 파이프라인 버블 증가라는 두 가지 성능 문제가 발생한다.
• Why: MLIP는 ab initio 정확도를 가지면서도 O(N)에 가까운 계산 비용으로 대규모 분자동역학(MD) 시뮬레이션을 가능케 하여 배터리, 촉매, 신약 설계 등 과학적 발견에 핵심적인데, scaling law를 확장하려면 효율적인 대규모 분산 학습 시스템이 필수적이며 본 연구는 이 공백을 처음으로 체계적으로 메운다.
• Approach: PP 실행을 device별 instruction list로 추상화한 뒤, SymFold로 first-order PP 스케줄을 conservative MLIP에 맞는 4단계(second-order) instruction list로 변환하여 중복 recomputation과 parameter replication을 제거하고, WaveK로 instruction list를 재정렬해 controllable한 메모리 사용량 하에서 파이프라인 버블을 추가로 줄인다.

Achievement

32개 GPU 환경에서 JanusPipe는 1F1B 대비 평균 1.51배, Hanayo 대비 평균 1.45배의 throughput 향상을 달성했으며, peak GPU memory를 각각 최대 20.56%, 42.70% 절감했다.

How

• PP 실행을 device 단위 instruction list로 추상화하여 4단계(FE/FF/BF/BE) MLIP 학습의 구조적 특성을 명시적으로 표현
• SymFold: first-order PP 스케줄(예: 1F1B)을 second-order instruction list로 변환하여 FF가 FE의 activation과 computation graph를 재사용하도록 하고, FE/FF 간 공유되는 parameter의 gradient를 BF/BE에서 동기화해 중복 recomputation과 parameter replication을 제거
• WaveK: instruction list를 WaveK unit 단위로 재구성/재정렬하여 tFE<tFF<tBE<tBF의 partial order로 인한 파이프라인 버블을 controllable한 메모리 하에서 최소화, 아울러 micro-batch를 repack하여 load imbalance 완화
• JanusPipe를 data parallelism(DP), graph parallelism(GP)과 결합하여 3D-parallel(PP/DP/GP) 학습 시스템으로 확장

Originality

• conservative MLIP의 double-backward(4-phase: FE/FF/BF/BE) 실행 패턴을 최초로 체계적으로 분석하고, 이를 첫 번째 순서 워크로드 기반 PP 스케줄과의 구조적 mismatch로 정식화
• PP 실행을 instruction list로 추상화하는 새로운 표현 방식을 제시
• SymFold라는 스케줄 변환 기법으로 four-phase 워크로드에서의 redundant recomputation/parameter replication 문제를 해결
• WaveK라는 새로운 스케줄링 알고리즘으로 phase별 실행시간의 partial order(tFE<tFF<tBE<tBF)에 특화된 bubble 감소 기법을 제시
• MLIP 학습을 위한 PP/DP/GP 3D-parallel 통합 시스템을 최초로 제안

Limitation & Further Study

• 실험이 32 GPU 규모에서 수행되어 향후 수백~수천 GPU 규모의 초대형 MLIP 학습에서도 동일한 성능 이점이 유지되는지 추가 검증이 필요
• WaveK의 k 파라미터 선택에 따른 memory-throughput trade-off에 대한 이론적 최적화 기준이나 자동 튜닝 방법이 충분히 논의되지 않을 가능성
• 비교 대상이 1F1B, Hanayo 등 기존 PP 스케줄에 한정되어 있어, 최신 non-conservative MLIP나 다른 하이브리드 스케줄과의 비교가 부족
• 특정 GNN 기반 MLIP 아키텍처(UMA, MACE, NequIP 등)에 특화된 설계가 다른 유형의 conservative 모델에 얼마나 일반화되는지 추가 검증 필요

Evaluation

총평: conservative MLIP 학습의 double-backward 구조적 특성을 정확히 짚어내고 이에 특화된 pipeline parallelism 기법을 체계적으로 설계한 실용적이고 임팩트 있는 시스템 논문으로, AI for Science 분야의 대규모 분산학습 인프라 발전에 기여할 것으로 기대된다.