💻 computer science

MatrixFSDP: communication-free matrix optimizers under ZeRO-3 parameter sharding

MatrixFSDP는 각 2D 가중치 행렬이 단일 랭크에 완전히 존재하도록 파라미터 배치를 재구성함으로써, 메모리 효율성을 유지하면서도 최적화 단계 중 비용이 많이 드는 행렬 재구성을 제거하고 상당한 지연 시간 감소를 달성하여 ZeRO-3 샤딩 하에서 Muon과 같은 행렬 옵티마이저를 통한 통신 프리 대규모 학습을 가능하게 합니다.

원저자: Ming Gao, Yanwu Xu, Hao Zhang

게시일 2026-07-08✓ Author reviewed ⓘ

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ming Gao, Yanwu Xu, Hao Zhang

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대한 로봇 팀(컴퓨터 클러스터)에게 소설을 쓰는 법을 가르치려 한다고 상상해 보세요. 로봇들은 방대한 텍스트가 담긴 책 한 권으로부터 함께 학습하고 있습니다. 이를 효율적으로 수행하기 위해, 그들은 Muon이라고 불리는 특별한 "학습 규칙"을 사용합니다.

문제점: "전체 그림" vs "퍼즐 조각"

보통 로봇들이 학습할 때는 거대한 책을 아주 작은 퍼즐 조각들로 나눕니다. 각 로봇은 책의 몇 페이지(샤드, shard)만을 가집니다. 이는 단 하나의 로봇도 책 전체를 다 들고 있을 필요가 없게 만들어 메모리를 절약하는 데 매우 효과적입니다. 이 방법을 ZeRO-3라고 부릅니다.

하지만 Muon 학습 규칙은 조금 까다롭습니다. 이 규칙은 한 번에 몇 페이지씩 나누어 학습하는 것을 원하지 않습니다. 완벽하게 제 역할을 수행하기 위해서는, 단어들 사이의 관계를 이해할 수 있도록 전체 2D 페이지(전체 행렬)를 한꺼번에 봐야만 합니다.

충돌 지점:

ZeRO-3는 말합니다: "우리에게는 퍼즐 조각들뿐이야."
Muon은 말합니다: "나는 학습을 위해 전체 페이지가 필요해."

기존의 해결책들 (나쁜 옵션들):

"재구성(Reconstruction)" 방식: 로봇들이 학습할 때마다 모든 작업을 멈추고, 모든 로봇으로부터 퍼즐 조각들을 모아 전체 페이지를 다시 붙이고, Muon이 학습하게 한 뒤, 즉시 페이지를 다시 찢어서 분해합니다.
- 단점: 이것은 마치 사람들이 작업을 하다가 계속 멈춰 서서, 거대한 직소 퍼즐을 조립했다가 다시 해체하는 과정을 반복하는 것과 같습니다. 매 단계마다 엄청난 양의 시간과 에너지(통신 비용)를 낭비하게 됩니다.
"전체 복사(Full Copy)" 방식: 조각을 나누는 대신, 모든 로봇이 책 전체의 복사본을 하나씩 가집니다. Muon은 즉시 학습할 수 있습니다. 모두가 전체 그림을 가지고 있기 때문입니다.
- 단점: 이 방식은 너무 많은 메모리를 요구하며, 책이 너무 커지면 로봇들의 두뇌(GPU)가 폭발해 버립니다. 즉, 공간이 부족해지는 것입니다.

새로운 해결책: MatrixFSDP

이 논문의 저자들은 MatrixFSDP라는 영리한 제3의 길을 찾아냈습니다. 그들은 학습 규칙(Muon)을 바꾸지도 않았고, 모두에게 전체 책을 들고 있으라고 강요하지도 않았습니다. 대신, 누가 책을 들고 있을지를 바꿨습니다.

비유: "전담 사서"

도서관에서 책들이 보통 조각조각 나뉘어 여러 사서들에게 분배되어 있다고 상상해 보세요.

MatrixFSDP의 아이디어: 모든 "페이지"(행렬)에 대해, 그 페이지를 담당할 특정 사서 한 명을 **"소유자(Owner)"**로 임명합니다.
- 이 소유자는 완전하고 전체적인 페이지를 통째로 가집니다.
- 그 외의 다른 사서들은 해당 페이지에 대해서는 아무것도 가지지 않습니다(빈 공간).
- 특별한 Muon 규칙이 필요 없는 나머지 부분들에 대해서는 기존의 "퍼즐 조각" 방식을 그대로 유지합니다.

실제 작동 방식:

학습 중 (Optimizer 단계): "소유자"는 이미 전체 페이지를 가지고 있기 때문에, Muon이 즉시 학습할 수 있습니다. 누구도 조각을 모으거나 무엇인가를 붙일 필요가 없습니다. 마치 사서가 자신의 책상 위에 놓인 책을 바로 읽는 것과 같습니다. 통신이 전혀 필요하지 않습니다.
읽기/쓰기 중 (Forward/Backward Pass): 로봇들이 책을 읽거나 써야 할 때, 일시적으로 조각들을 다시 모아 작업을 수행한 뒤, 즉시 다시 "소유자"의 슬롯으로 되돌려 놓습니다.

이것이 왜 중요한가

이 논문은 이 접근 방식이 거대 AI 모델을 훈련할 때 발생하는 가장 큰 병목 현상을 해결한다고 주장합니다.

속도: "붙이고 찢는" 과정을 멈췄기 때문에 학습 단계가 믿을 수 없을 정도로 빨라졌습니다. 단일 컴퓨터 노드에서는 4.2배 더 빨라졌습니다. 8개의 노드로 구성된 대규모 클러스터에서는 54.6배 더 빨라졌는데, 이는 기존 방식이 컴퓨터 간에 데이터를 주고받는 네트워크 과정에서 시간을 낭비했던 반면, MatrixFSDP는 데이터를 로컬에 유지했기 때문입니다.
메모리: "전체 복사" 방식과 달리, MatrixFSDP는 여전히 "퍼즐 조각" 방식의 메모리만 사용합니다. 덕분에 "전체 복사" 방식으로는 감당할 수 없는 거대한 모델도 훈련할 수 있게 해줍니다.
정확도: 그들은 "소유자"가 전체를 모았을 때와 동일한 데이터를 얻기 때문에, 학습 결과가 완벽하고 느린 방식과 동일하다는 것을 증명했습니다.

요약

MatrixFSDP는 팀을 재편성하는 것과 같습니다. 즉, 특정 업무를 수행하기 위해 전체 문서가 필요한 사람에게 실제로 그 전체 문서를 들려주는 것입니다. 그 외의 사람들은 해당 작업에 대해 아무것도 들고 있지 않습니다. 이를 통해 데이터를 끊임없이 주고받는 수고를 덜어줌으로써, 더 큰 책상(더 많은 메모리)을 필요로 하지 않으면서도 팀이 훨씬 더 빠르게 일할 수 있게 만듭니다.

기술 요약: MatrixFSDP

문제 정의

대규모 언어 모델 학습은 점차 Muon, Shampoo, SOAP와 같은 행렬 구조 최적화 도구(matrix-structured optimizers)에 의존하고 있습니다. 이러한 방식은 좌표별(coordinate-wise) 최적화 도구(예: AdamW)와 달리, Newton–Schulz 직교화(orthogonalization)나 행렬 전처리(preconditioning)와 같은 연산을 수행하기 위해 완전한 2D 그래디언트 행렬을 입력값으로 필요로 합니다. 이는 파라미터, 그래디언트, 옵티마이저 상태를 여러 랭크(rank)에 분할하여 메모리 효율성을 극대화하는 FSDP (Fully Sharded Data Parallel) 및 ZeRO-3와의 근본적인 시스템 불일치를 야기합니다.

표준 ZeRO-3 환경에서 각 랭크는 모델의 일부 슬라이스만을 보유합니다. 행렬 최적화 도구를 적용하려면 시스템은 매 옵티마이저 스텝마다 전체 2D 그래디언트 행렬을 재구성해야 합니다. 이는 백워드 패스(backward pass) 이후 가중치 크기만큼의 all-gather 통신을 필요로 하며, 이 통신은 연산과 중첩(overlap)될 수 없는 임계 경로 병목 현상이 되어 노드 수와 모델 크기에 따라 성능이 급격히 저하됩니다.

기존 솔루션들은 불완전합니다:

재구성 (Stock FSDP2-Muon): ZeRO-3의 메모리 절감 효과는 유지하지만, 매 옵티마이저 스텝마다 막대한 통신 오버헤드가 발생합니다.
ZeRO-1 소유자 배치 (ZeRO-1 Owner Placement): 특정 랭크(소유자)에 전체 행렬을 할당하여 옵티마이저 스텝을 로컬에서 수행하게 하지만, 모든 랭크에 전체 파라미터가 상주해야 하므로 ZeRO-3의 메모리 효율성을 포기해야 합니다.

목표는 ZeRO-3 수준의 파라미터 메모리를 희생하지 않으면서 전체 행렬 옵티마이저 입력을 확보하는 것입니다.

방법론: MatrixFSDP

MatrixFSDP는 제3의 경로인 **소유자 형태의 ZeRO-3 샤드(Owner-shaped ZeRO-3 shards)**를 도입합니다. 이는 최적화 알고리즘을 변경하거나 샤딩을 포기하는 대신, 데이터 병렬 그룹 내에서 샤드가 존재하는 위치를 재정의합니다.

1. 소유자 형태 배치 (Owner-Shaped Placement)

모델의 모든 2D 가중치 행렬에 대해, MatrixFSDP는 단일 데이터 병렬 랭크를 **소유자(owner)**로 지정합니다.

소유자 랭크 (Owner Rank): 완전한 2D 행렬을 보유합니다.
비소유자 랭크 (Non-Owner Ranks): 해당 특정 파라미터에 대해 빈(empty) 샤드를 보유합니다.
비행렬 텐서 (Non-Matrix Tensors): 1D 벡터 및 기타 비행렬 텐서들은 "테일(tail)"로 묶여 별도의 소유자에 할당되며, 이들은 기존의 AdamW 최적화 도구를 계속 사용합니다.

이를 통해 백워드 패스와 그래디언트 리덕션(reduction)이 완료된 후, Muon의 입력(완전한 2D 그래디언트)이 자연스럽게 소유자 랭크에 위치하게 됩니다. 결과적으로 Newton–Schulz 업데이트는 제로(zero) 옵티마이저 스텝 행렬 통신으로 소유자 랭크에서 로컬하게 실행됩니다.

2. 런타임 과제 및 해결책

이러한 불균형한 레이아웃(일부 랭크는 전체 행렬을 보유하고 다른 랭크는 빈 샤드를 보유함)은 표준 FSDP2의 동일 크기 집합(equal-size collective) 가정을 위반합니다. MatrixFSDP는 다음 네 가지 핵심 런타임 메커니즘을 통해 이를 해결합니다.

MatrixShard 메타데이터 및 글로벌 플래너: 글로벌 플래너는 클러스터 전체에서 상주 바이트(resident bytes)와 옵티마이저 워크로드를 균형 있게 배분하도록 소유자를 할당합니다. 이 메타데이터는 어떤 랭크가 어떤 텐서 역할을 소유하는지 추적하여, 시스템이 특수 사례들의 집합이 아닌 하나의 일관된 ZeRO-3 레이아웃으로 유지되도록 보장합니다.
소유자-세그먼트 P2P 집합 연산 (Owner-Segment P2P Collectives): 표준 all-gather는 결정론적인 소유자-세그먼트 집합 연산으로 대체됩니다. 포워드/백워드 실체화(materialization) 과정에서, 소유자로부터 필요한 랭크로 비어 있지 않은 세그먼트만 전송됩니다. 그래디언트 리덕션 과정에서는 결과가 비어 있지 않은 소유자 세그먼트에만 기록됩니다. 이는 빈 랭크에 패딩을 넣는 것을 방지하고, 전체 행렬 all-gather를 제거합니다.
소유자-버퍼 피닝 (Owner-Buffer Pinning): FSDP2 라이프사이클(unshard $\to$ compute $\to$ reshard)을 처리하기 위해, MatrixFSDP는 "피닝(pinned)" 상태를 도입합니다. 포워드 패스 이후 버퍼가 축소될 때, 스토리지는 공유 풀로 반환되지 않고 피닝됩니다. 이를 통해 백워드 패스에서 전체 행렬을 동일한 메모리 위치에서 "회생(revive)"시킬 수 있어, autograd 뷰 에러를 방지하고 소유자가 버퍼를 해제하기 전에 업데이트를 수행할 수 있도록 합니다.
체크포인팅 및 라우팅: 체크포인팅(DCP)과 옵티마이저 상태 라우팅은 MatrixShard 메타데이터를 사용합니다. 옵티마이저 상태는 올바른 소유자(행렬의 경우 Muon, 테일의 경우 AdamW)로 라우팅되며, 체크포인트는 서로 다른 월드 사이즈(world size) 간의 리샤딩을 지원하기 위해 특정 배치 레이아웃을 저장합니다.

주요 기여

소유자 형태의 ZeRO-3: 2D 행렬은 단일 소유자에게 전체로 존재하고 다른 곳은 비어 있게 만드는 새로운 배치 전략으로, 샤드된 메모리 체계 내에서 로컬 행렬 최적화를 가능하게 합니다.
통신이 없는 옵티마이저 스텝: 그래디언트 리덕션 목적지를 옵티마이저의 입력 요구사항과 일치시킴으로써, MatrixFSDP는 기존 ZeRO-3에서 행렬 최적화 도구에 내재되어 있던 가중치 크기의 all-gather 병목 현상을 제거합니다.
런타임 인프라스트럭처: 불균형한 소유자 레이아웃을 FSDP2 프레임워크 내에서 효율적이고 정확하게 작동하게 하는 플래너, 커스텀 집합 연산, 버퍼 피닝 등의 일련의 런타임 구성 요소들을 제공합니다.
범용성: 이 접근 방식은 옵티마이저에 구애받지 않으며, 사전 조건 도구(preconditioner)의 분산 버전 없이도 Muon, Shampoo, SOAP와 함께 작동함을 입증했습니다.

실험 결과

다양한 모델 크기(최대 32B 파라미터)에 대해 64개의 NVIDIA A100 GPU(8개 노드)에서 평가되었습니다.

옵티마이저 스텝 지연 시간: MatrixFSDP는 단일 노드에서 stock FSDP2-Muon 대비 4.2배, 8개 노드에서 54.6배의 옵티마이저 스텝 지연 시간 감소를 달이트합니다. 베이스라인의 통신은 노드 간 패브릭을 가로지르는 반면, MatrixFSDP의 스텝은 로컬에 머물기 때문에 규모가 커질수록 속도 향상이 증가합니다.
엔드 투 엔드(End-to-End) 속도 향상: 최대 2.15배의 엔드 투 엔드 훈련 속도 향상을 달성했습니다.
메모리 효율성: MatrixFSDP는 ZeRO-3 수준의 메모리를 유지합니다. 32B 모델의 경우 랭크당 10 GB를 사용하는 반면, ZeRO-1 소유자 배치 방식은 GPU 한도인 80 GB를 초과합니다(14B 모델 기준 랭크당 >54 GB 예상). 이는 "gather-once" 방식의 FSDP2-Muon보다 3~7배 더 메모리 효율적입니다.
정확성: 훈련 궤적은 높은 정밀도로 전체 행렬 DDP Muon 레퍼런스와 일치하며(출력된 손실 차이 $\leq 0$ ), 이는 소유자가 업데이트에 필요한 정확한 그래디언트 행렬을 수신했음을 확인시켜 줍니다.

의의 및 주장

본 논문은 MatrixFSDP가 행렬 최적화 도구와 파라미터 샤딩 사이의 "시스템 불일치"를 해결한다고 주장합니다. 주요 의의는 행렬 최적화 도구(더 나은 토큰 효율성과 수렴성을 제공함)를 사용하면서도, 이전 방식들이 가졌던 과도한 통신 비용이나 메모리 오버헤드 없이 대규모 학습을 가능하게 한다는 점에 있습니다.

저자들은 MatrixFSDP가 특히 대규모 샤드 스팬(shard span)을 가진 멀티 노드 ZeRO-3 실행에서 효과적이라고 강조합니다. 이는 매 스텝마다 행렬을 재구성하는 비용이 훈련 시간을 지배하게 될 상황에서 매우 중요합니다. 또한, 단일 노드 실행이나 고정된 모델 강한 스케일링(strong scaling) 환경에서는 이득이 작을 수 있으며, 이 경우 소유자 팬아웃(fanout)이 새로운 병목이 될 수 있음을 겸허히 언급합니다. 이 연구는 비행렬 파라미터를 위한 좌표별 최적화 도구를 대체하려는 것이 아니라, 현대적 모델의 2D 행렬 구성 요소를 처리하기 위한 통합되고 효율적인 경로를 제공하는 것을 목적으로 합니다.