veScale-FSDP: Flexible and High-Performance FSDP at Scale

이 논문은 기존 FSDP 시스템의 구조 인식 훈련 및 비 요소 단위 옵티마이저 지원 한계를 해결하고, 유연한 'RaggedShard' 포맷과 구조 인식 계획 알고리즘을 도입하여 처리량을 5~66% 향상시키고 메모리 사용량을 16~30% 절감하며 대규모 GPU 환경에서 효율적으로 확장 가능한 'veScale-FSDP'를 제안합니다.

핵심

거대 AI 모델을 훈련시키는 새로운 비법: veScale-FSDP

안녕하세요! 오늘 소개해 드릴 논문은 **"veScale-FSDP"**라는 이름의 새로운 기술입니다. 이 기술은 거대한 인공지능 (AI) 모델을 가르치는 (훈련시키는) 방식을 완전히 바꾼 혁신적인 시스템입니다.

이걸 이해하기 쉽게, 거대한 레고 성을 짓는 상황에 비유해 설명해 드릴게요.

---

1. 문제: 왜 기존 방식은 힘들었을까? (레고 조각이 너무 작아서)

거대 AI 모델을 훈련시킨다는 건, 수조 개의 레고 조각 (데이터와 지식) 을 수천 개의 작업자 (GPU) 가 나눠서 조립하는 것과 같습니다.

• 기존 방식 (FSDP): 과거에는 이 레고 조각들을 무작위로 잘게 부숴서 각 작업자에게 나누어 주었습니다.
  • 문제 1 (잘못된 자르기): 어떤 레고 조각은 '벽돌' 단위로 잘려야 하는데, 기존 방식은 '벽돌의 반'이나 '조각' 단위로 잘라버렸습니다. 그래서 작업자가 조립할 때, "아! 이 조각은 내 것이 아니야, 옆 친구에게서 가져와야 해!"라고 계속 주고받으며 시간이 낭비되었습니다.
  • 문제 2 (비효율적인 공간): 작업대 (메모리) 가 너무 좁아서, 필요한 레고 조각만 가져오려고 하면 쓸데없는 빈 공간까지 차지하게 되어, 작업자가 "여기 공간이 부족해!"라고 외치며 멈추는 경우가 많았습니다.
  • 문제 3 (새로운 도구 사용 불가): 최근에는 '벽돌 전체'를 한 번에 다루는 새로운 도구 (특수한 최적화 알고리즘) 가 나왔는데, 기존 방식은 조각이 너무 잘게 나뉘어 있어 이 도구를 쓸 수 없었습니다.

2. 해결책: veScale-FSDP 의 등장 (똑똑한 레고 정리사)

이 문제를 해결하기 위해 등장한 것이 veScale-FSDP입니다. 이 시스템은 두 가지 핵심 아이디어로 작동합니다.

① 'RaggedShard(불규칙한 조각)' 방식: 상황에 맞춰 잘라라!

기존에는 무조건 똑같은 크기로만 잘랐다면, veScale-FSDP 는 레고 조각의 모양에 맞춰서 잘라줍니다.

• 비유: 만약 어떤 레고 블록이 '벽돌 10 개' 단위라면, 10 개 단위로 딱 잘라 작업자에게 줍니다. '벽돌 3 개' 단위라면 3 개 단위로 줍니다.
• 효과: 작업자는 조각을 다시 조립할 필요가 없으니, 조립 속도 (훈련 속도) 가 훨씬 빨라집니다. 또한, 새로운 도구 (Muon 옵티마이저 등) 를 쓰더라도 조각이 깨지지 않아 바로 사용할 수 있습니다.

② 'DBuffer(공유 창고)'와 '스마트 계획': 공간과 시간을 아껴라!

• DBuffer (공유 창고): 각 작업자가 따로따로 창고를 쓰는 대신, 거대한 공유 창고를 만들어서 레고를 바로 꺼내 쓸 수 있게 했습니다. 이렇게 하면 레고를 옮기는 (복사하는) 시간이 거의 사라집니다.
• 스마트 계획 알고리즘: "누가 어떤 레고를 몇 개씩 가져가야 가장 빨리, 그리고 가장 적게 공간만 차지할까?"를 계산하는 초지능 계획사가 있습니다. 이 계획사는 불필요한 빈 공간 (패딩) 을 최소화하고, 작업자들 간의 이동 경로를 최적화합니다.

3. 실제 성과: 얼마나 빨라졌을까?

이 새로운 시스템을 도입한 결과, 놀라운 변화가 일어났습니다.

• 속도: 기존 시스템보다 5% 에서 최대 66% 까지 더 빠릅니다. (레고 성을 짓는 속도가 거의 두 배에 가까워진 셈입니다.)
• 공간: 필요한 메모리 (작업대 공간) 가 16% 에서 30% 까지 줄었습니다. 같은 공간으로 더 많은 작업을 할 수 있게 된 것입니다.
• 확장성: 이제 **수만 개의 GPU(작업자)**가 함께 일해도 효율이 떨어지지 않습니다. 마치 10,000 명의 레고 장인들이 함께 일해도 서로 부딪히지 않고 질서 있게 일하는 것과 같습니다.

4. 왜 이것이 중요한가요?

이 기술은 단순히 "더 빠르다"는 것을 넘어, AI 연구자들이 더 자유롭게 상상할 수 있게 해줍니다.

• 자유로운 설계: 연구자들은 "내 모델은 이렇게 생겼으니, 조각을 이렇게 잘라줘"라고 말하면 시스템이 알아서 맞춰줍니다. 기존처럼 시스템에 맞춰 모델을 고칠 필요가 없습니다.
• 새로운 기술 수용: 앞으로 나올 더 복잡한 AI 모델이나 새로운 학습 방법도 이 시스템이 자연스럽게 받아들여줍니다.

요약

veScale-FSDP는 거대 AI 모델을 훈련시킬 때, 레고 조각을 무작위로 자르지 않고 모양에 맞춰 똑똑하게 잘라주고, 공유 창고를 효율적으로 관리하며, 최적의 이동 경로를 계획해주는 시스템입니다. 덕분에 우리는 더 빠르고, 더 저렴하게, 더 큰 AI 모델을 만들 수 있게 되었습니다.

이 기술은 이미 실제 산업 현장에서 수만 개의 GPU 를 활용해 검증되었으며, 앞으로 우리가 만나게 될 더 똑똑한 AI 들의 기반이 될 것입니다.

기술 요약

1. 문제 정의 (Problem)

기존의 FSDP (또는 ZeRO) 시스템은 다음과 같은 주요 문제점으로 인해 최신 모델 학습에 한계를 겪고 있습니다:

• 구조 인식 학습과의 비호환성:
  • 최신 모델 (Gemini, Kimi K2 등) 은 Shampoo, Muon 과 같은 행렬 기반 옵티마이저와 DeepSeek-V3 와 같은 블록 단위 양자화 (Block-wise Quantization) 기술을 사용합니다.
  • 기존 FSDP 는 파라미터를 요소 단위 (element-wise) 또는 행 단위 (row-wise) 로 균등하게 분할 (sharding) 합니다. 이로 인해 학습에 필요한 원자적 블록 (atomic tensor blocks) 경계가 분할 경계와 일치하지 않아, 모델 개발자가 코드를 침입적으로 수정하거나 시스템이 복잡한 패딩 및 통신 로직을 처리해야 합니다.
• 성능 및 메모리 비효율성:
  • 메모리: GPU 메모리가 주요 병목 현상입니다. 기존 시스템은 메모리 파편화 (fragmentation) 와 비효율적인 할당으로 인해 OOM(Out-Of-Memory) 이 발생하거나, 불필요한 오버프로비저닝으로 자원이 낭비됩니다.
  • 통신: FSDP1 은 AllGather 효율성이 낮고, FSDP2 는 파라미터 단위 분할로 인해 AllGather 후 Copy-Out 및 ReduceScatter 전 Copy-In 과정에서의 복사 오버헤드가 큽니다. 또한 정렬되지 않은 통신 버퍼로 인해 통신 지연이 발생합니다.
  • 확장성: 수만 개의 GPU 로 확장할 때 효율성을 유지하는 시스템이 부족합니다.

---

2. 방법론 (Methodology)

veScale-FSDP 는 유연성과 성능을 동시에 달성하기 위해 세 가지 핵심 기술을 도입했습니다.

2.1 RaggedShard (유연한 분할 포맷)

• 개념: 기존 DTensor 의 Shard 포맷을 확장하여, **임의의 분할 입자성 (granularity)**과 임의의 분포를 지원하는 새로운 포맷입니다.
• 특징:
  • 블록 단위 분할 지원: 텐서를 임의 크기의 2D 블록으로 분할할 수 있어, 블록 단위 양자화 (Block-wise Quantization) 와 행렬 옵티마이저 (Muon 등) 를 자연스럽게 지원합니다.
  • 비균등 분할: 각 디바이스에 할당되는 블록 수가 다를 수 있어 (JaggedTensor 개념), 행렬 옵티마이저가 필요한 비균등 데이터 배치를 처리할 수 있습니다.
  • 기존 포맷과의 호환성: TP(텐서 병렬) 나 EP(전문가 병렬) 와 같은 기존 DTensor 배치 전략과 결합하여 사용할 수 있습니다.

2.2 구조 인식 계획 알고리즘 (Structure-Aware Planning Algorithm)

• 목표: RaggedShard 텐서들을 통신 버퍼에 배치할 때, 패딩 (padding) 을 최소화하고 통신 효율을 극대화하는 최적의 레이아웃을 찾습니다.
• 문제 정의: 이는 NP-hard 최적화 문제 (Partition problem 으로 환원 가능) 로 정의됩니다.
• 해결책:
  • 휴리스틱 기반 DP 알고리즘: ILP 솔버는 대규모 환경에서 비현실적이므로, 트랜스포머 모델의 규칙성을 활용한 다항 시간 (polynomial-time) 동적 계획법 (DP) 휴리스틱을 사용합니다.
  • 배치 전략: 텐서를 재배열 (permute) 하고, 텐서 사이에만 패딩을 추가하여 텐서 내부의 메모리 연속성을 유지하고 로드 불균형을 해소합니다.
  • 결과: 통신 버퍼 내에서 블록 경계가 정렬되도록 하여, 불필요한 통신과 메모리 낭비를 방지합니다.

2.3 분산 버퍼 (Distributed Buffer, DBuffer)

• 기능: RaggedShard 텐서를 위한 고성능 원시 연산자 (primitive) 입니다.
• 핵심 이점:
  • Zero-Copy: 전역 버퍼의 슬라이스를 직접 참조하여 복사 (copy) 오버헤드를 제거합니다.
  • 그룹 단위 연산: 개별 텐서별 커널 실행 대신 그룹 단위로 연산을 융합 (fuse) 하여 커널 발동 오버헤드를 줄입니다.
  • 메모리 관리: 배치 할당 (batched allocation) 을 통해 메모리 파편화를 줄이고, 스트림 종속성을 명시적으로 관리하여 메모리 해제를 최적화합니다.

---

3. 주요 기여 (Key Contributions)

1. RaggedShard 포맷 도입: 구조 인식 학습 (블록 단위 양자화, 행렬 옵티마이저) 을 가능하게 하는 유연한 분할 메커니즘을 제안했습니다.
2. 고성능 계획 알고리즘: NP-hard 문제를 해결하는 효율적인 휴리스틱 알고리즘을 통해 패딩 오버헤드를 최소화하고 통신 효율을 극대화했습니다.
3. DBuffer 원시 연산자: 복사 없는 (zero-copy) 통신과 그룹 단위 커널 융합을 통해 메모리 및 계산 오버헤드를 대폭 줄였습니다.
4. PyTorch 네이티브 통합: fully_shard API 와 호환되어 모델 코드를 수정하지 않고도 기존 FSDP2 를 대체할 수 있는 플러그 앤 플레이 모듈로 구현되었습니다.

---

4. 실험 결과 (Results)

1024 개 이상의 GPU 클러스터에서 다양한 모델 (LLaMA-3-70B, GPT-OSS-120B, 내부 MoE 모델) 을 대상으로 평가했습니다.

• 처리량 (Throughput):
  • 기존 FSDP 시스템 (DeepSpeed, PyTorch FSDP1/2, Megatron-FSDP) 대비 5% ~ 66% 높은 처리량을 달성했습니다.
  • 특히 MoE 모델에서는 11~66% 의 향상을 보였으며, 이는 패딩 오버헤드 제거와 DBuffer 기반 제로-복사 통신 덕분입니다.
• 메모리 사용량:
  • 피크 메모리 사용량이 16% ~ 30% 감소했습니다.
  • 이는 DBuffer 의 배치 할당 정책과 메모리 파편화 감소 기법 때문입니다.
• 확장성 (Scalability):
  • 약한 확장성 (Weak Scaling): 10,000 개 GPU 까지 거의 선형적인 확장성을 보였습니다.
  • 강한 확장성 (Strong Scaling): 128M 토크의 글로벌 배치 크기에서 10,000 개 GPU 까지 선형 확장성을 유지했습니다.
  • 모델 크기 확장: 1,000 개 GPU 에서 2.4 조 (2.4T) 파라미터 모델 학습이 가능하며, MFU(모델 FLOPS 활용률) 가 모델 크기가 커질수록 오히려 향상되었습니다.
• 특수 옵티마이저 지원:
  • 8-bit Adam: RaggedShard 를 통해 블록 단위 양자화를 자연스럽게 지원하며, 기존 시스템 대비 복잡한 커스텀 통신 로직 없이 구현 가능했습니다.
  • Muon 옵티마이저: 행렬 단위 연산이 필요한 Muon 을 분산 환경에서 효율적으로 실행할 수 있게 하여, 256 GPU 에서 47.3% 의 MFU 를 달성했습니다.

---

5. 의의 및 결론 (Significance)

• 차세대 대규모 모델 학습의 표준: veScale-FSDP 는 구조 인식 학습과 비 요소 단위 옵티마이저를 지원하는 최초의 고성능 FSDP 시스템으로, 차세대 초대규모 모델 (Gemini, Kimi 등) 의 학습을 가능하게 합니다.
• 개발 생산성 향상: 모델 개발자가 복잡한 분산 로직이나 패딩 처리를 신경 쓰지 않고도 최첨단 옵티마이저와 양자화 기법을 적용할 수 있어, 연구 및 개발 속도가 가속화됩니다.
• 실무 검증: 10,000 개 이상의 GPU 를 사용하는 실제 산업용 워크로드에서 성공적으로 배포 및 검증되었으며, PyTorch 생태계 (TorchTitan 등) 와의 긴밀한 협력을 통해 오픈소스 생태계에 기여하고 있습니다.

결론적으로, veScale-FSDP 는 **유연성 (RaggedShard)**과 **성능 (Planning Algorithm + DBuffer)**을 결합하여 대규모 GPU 클러스터에서의 학습 효율성을 혁신적으로 개선한 시스템입니다.