🤖 machine learning

Model Parallelism With Subnetwork Data Parallelism

이 논문은 모델을 구조화된 서브네트워크로 분할하여 활성화 값을 교환하지 않고도 여러 워커에 걸쳐 학습함으로써, 다양한 아키텍처와 규모에 걸쳐 성능을 유지하거나 향상시키면서도 상당한 메모리 절감(28%-60%)을 달성하는 분산 학습 프레임워크인 Subnetwork Data Parallelism (SDP)을 소개한다.

원저자: Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

게시일 2026-06-02

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

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

당신은 거대하고 똑똑하지만, 엄청나게 배가 고픈 로봇(대규모 AI 모델)에게 말하는 법이나 이미지를 인식하는 법을 가르치려 한다고 상상해 보세요. 문제는 이 로봇을 위해 너무 크고 비싼 주방이 필요해서, 오직 소수의 사람들만이 그 주방을 지을 여유가 있다는 점입니다. AI의 세계에서 이 "주방"은 강력한 그래픽 카드(GPU)에 있는 컴퓨터 메모리(RAM)를 의미합니다.

이 논문은 이 거대한 로봇들을 훈련시키는 새로운 방법인 **서브네트워크 데이터 병렬성(Subnetwork Data Parallelism, SDP)**을 소개합니다. 여기서는 이해를 돕기 위해 간단한 비유를 통해 이 방식이 어떻게 작동하는지 설명합니다.

기존 방식: "전체 복제"의 문제

전통적으로 거대한 AI를 훈련시키기 위해 연구자들은 **데이터 병렬 처리(Data Parallelism, DDD)**라는 방법을 사용합니다.

비유: 당신에게 거대한 연회를 준비하려는 8명의 요리사(GPU) 팀이 있다고 상상해 보세요. 기존 방식에서는 모든 요리사가 전체 레시피 북과 모든 식재료의 완전한 복사본을 하나씩 다 가집니다. 그들은 모두 전체 요리를 만들고, 맛을 본 다음, 다음 라운드에서 레시피를 어떻게 개선할지 서로 의견을 나누기 위해 소리를 높여 대화합니다.
문제점: 이것은 매우 낭비적입니다. 각 요리사는 자신의 레시피와 재료를 담아둘 수 있는 아주 큰 주방을 가져야 합니다. 만약 레시피가 너무 크다면, 주방이 폭발해 버리거나(메모리 부족) 요리사들이 아예 요리를 할 수 없게 됩니다.

새로운 방식: "특화된 팀" (SDP)

저자들은 **서브네트워크 데이터 병렬성(SDP)**을 제안합니다. 모든 요리사에게 전체 레시피를 주는 대신, 레시피를 특정 섹션으로 나누어 서로 다른 요리사들에게 할당하는 방식입니다.

비유: 이제 요리사 A는 전채 요리만을 담당하고, 요리사 B는 스프를, 요리사 C는 메인 코스 요리를 담당합니다.
- 전체 요리를 공유하지 않음: 결정적으로, 그들은 전체 음식을 주고받을 필요가 없습니다. 그들은 자신들이 작업 중인 특정 식재료에 대해서만 이야기합니다.
- 주방이 작아짐: 요리사 A는 전채 요리 레시피만 가지고 있으면 되기 때문에, 그들의 주방은 훨씬 작아질 수 있습니다. 이를 통해 훨씬 작고 저렴한 주방에서도 거대한 연회를 준비할 수 있게 됩니다.
- 결과: 훨씬 더 큰 로봇을 훈련시키거나(또는 더 빠르게 훈련하거나), 훨씬 더 비싼 주방 없이도 가능해집니다.

작업을 나누는 두 가지 방법

논문은 이 "서브 레시피"를 요리사들에게 할당하는 두 가지 다른 방식을 테스트합니다.

포워드 마스킹 (Forward Masking, "재료를 잘라내는" 방식):
- 작동 방식: 요리사는 요리를 시작하기 전에 자신이 사용하지 않기로 한 재료들을 말 그대로 버립니다. 그들은 자신에게 할당된 요리의 부분만을 요리합니다.
- 장점: 사용하지 않는 재료를 저장조차 하지 않기 때문에 공간을 가장 많이 절약합니다. 이는 마치 수프를 끓이면서 당근과 양파만 사고, 나머지 식재료 목록은 완전히 무시하는 것과 같습니다.
- 단점: 레시피의 일부를 무시하기 때문에, 전체적인 그림을 학습하기 위해 몇 번의 배치를 더 거쳐야 합니다.
백워드 마스킹 (Backward Masking, "책 전체를 읽되, 노트는 자기 것만 쓰는" 방식):
- 작동 방식: 요리사는 맥락을 이해하기 위해 전체 레시피 북(전체 모델)을 읽지만, 요리를 어떻게 개선할지에 대한 노트를 적을 때는 자신이 책임지는 섹션에 대해서만 기록합니다.
- 장점: 요리사가 전체 식사를 이해하고 있으므로 수학적으로 조금 더 "정직한" 방식이며, 공간을 절약하면서도 더 안전하고 안정적인 학습이 가능합니다.

무엇을 발견했는가?

연구진은 이 두 가지 매우 다른 유형의 AI를 대상으로 테스트했습니다:

언어 모델 (LLMs): 텍스트를 작성하는 "LLaMA" 모델 같은 것들입니다. 이들은 5억 개와 10억 개의 "뉴런"(파라미터)을 가진 모델을 훈련시켰습니다.
- 결과: 메모리 요구량을 28%에서 60%까지 줄였습니다. 어떤 경우에는 이전에는 거대한 주방이 필요했던 모델을 훨씬 작은 주방에 넣을 수 있었으며, AI의 언어 능력이 떨어지는 일도 없었습니다.
- 보너스: SDP가 다른 메모리 절약 기술들("액티베이션 체크포인팅" 및 "FSDP")과 완벽하게 호환된다는 것을 발견했습니다. 이는 마치 레고 블록을 쌓는 것과 같아서, SDP를 다른 방법들과 결합하여 메모리 사용량을 무려 **85%**까지 획기적으로 줄일 수 있습니다.
이미지 분류기: 사진을 인식하는 모델(ResNet 및 Swin Transformer 등)입니다.
- 결과: 표준 이미지 데이터셋(CIFAR)을 사용하여 이 모델들을 훈련시켰습니다. 메모리를 훨씬 적게 사용함에도 불구하고(때로는 원래의 40% 수준까지), 이 AI는 전체 크기 버전만큼이나 이미지를 잘 인식했습니다. 어떤 경우에는 "특화된 팀" 방식이 오히려 AI가 더 잘 학습하도록 도와주는 "정규화(Regularization, 과하게 생각하는 것을 방지하는 방법)" 역할을 하기도 했습니다.

핵심 요약

이 논문은 새로운 종류의 AI를 발명하거나 병원 또는 자율주행차에서 사용하는 새로운 AI 방식을 제안하는 것이 아닙니다. 대신, 이 논문은 물류 문제를 해결합니다.

이것을 건설 현장의 인력을 조직하는 새로운 방법이라고 생각해보세요. 모든 노동자에게 초고층 빌딩의 전체 설계도를 주어 주머니를 너무 무겁게 만드는 대신, 각 노동자에게 그들이 맡은 특정 층의 설계도만 주는 것입니다. 그들은 여전히 같은 초고층 빌딩을 짓지만, 훨씬 적은 공간을 차지하며, 동일한 자원으로 더 높은 건물을 지을 수 있습니다.

핵심 요점:

SDP는 거대한 AI 모델을 여러 컴퓨터에 걸쳐 작은 조각으로 나눕니다.
모든 컴퓨터가 전체 모델을 보유할 필요를 없애, 메모리를 28%~60% 절약합니다.
사용 방식의 변경 없이 기존 AI 모델(LLaMA, ResNet 등)과 함께 사용할 수 있습니다.
다른 메모리 절약 기술과 결합하여 훨씬 더 많은 공간을 절약할 수 있습니다(최대 85%).
그럼에도 불구하고, AI는 전통적인 방식만큼(또는 때로는 그보다 더 잘) 성능을 유지합니다.

기술 요약: 서브네트워크 데이터 병렬성 (Subnetwork Data Parallelism, SDP)

문제 정의

딥 뉴럴 네트워크의 급격한 규모 확장은 가속기(accelerator)에 전례 없는 메모리 요구 사항을 초래하였으며, 이는 종종 분산 학습 전략을 필요로 합니다. **데이터 병렬성 (Data Parallelism, DDP)**은 널리 사용되지만, 모든 GPU에 전체 모델을 복제하므로 높은 메모리 오버헤드와 값비싼 그래디언트 동기화 문제를 야기합니다. 반대로, 모델 병렬성 (Model Parallelism) (예: 파이프라인 병렬성, 텐서 병렬성)은 더 큰 아키텍처를 수용하기 위해 장치 간에 모델을 분할하지만, 활성화 함수(activation)에 대한 상당한 통신 비용을 발생시키며 종종 파이프라인 버블(pipeline bubble) 및 부하 불균형 문제를 일으킵니다.

FSDP(Fully Sharded Data Parallelism) 및 ZeRO와 같은 기존 솔루션은 파라미터, 그래디언트, 옵티마이저 상태를 샤딩하여 디바이스당 메모리를 줄여주지만, 특히 그래디언트 동기화 과정에서 상당한 통신 오버헤드를 여전히 발생시킵니다. 활성화 대역폭과 노드당 메모리 사용량을 줄이면서도 수렴 품질을 희생하지 않거나 고대역폭 인터커넥트를 요구하지 않는 해결책을 찾는 것은 여전히 남아 있는 과제입니다.

방법론: 서브네트워크 데이터 병렬성 (SDP)

저자들은 모델을 구조화된 서브네트워크로 분할하여 활성화 값을 교환하지 않고도 워커(worker) 간에 학습시키는 분산 학습 프레임워크인 **서브네트워크 데이터 병렬성 (SDP)**을 제안합니다. 계산을 순차적으로 나누는 파이프라인 방식과 달리, SDP는 각 워커에 입력부터 손실(loss)까지의 전체 경로를 보존하는 구조적으로 완전한 부분(서브네트워크)을 할당하여 독립적인 그래디언트 계산을 가능하게 합니다.

핵심 메커니즘

서브네트워크 구축:
- 블록 수준 (B-SDP): 전체 아키텍처 블록(예: 트랜스포머 블록 또는 ResNet 베이직 블록)을 제거합니다.
- 뉴런/채널 수준 (N-SDP): 인접한 레이어 간의 일관성을 보장하면서, 완전 연결 계층(fully connected layers)의 특정 뉴런이나 컨볼루션 계층의 채널을 제거합니다.
- 워커들은 단계적 평균화(stepwise averaging)를 통해 중첩되는 파라미터를 동기화합니다.
마스킹 체제 (Masking Regimes):
본 논문은 학습 중 서브네트워크가 어떻게 활용되는지를 정의하기 위해 두 가지 상호 보완적인 마스킹 전략을 조사합니다.
- 포워드 마스킹 (Forward-Masking): 모델이 마스킹된 파라미터( $m_{fwd} = m$ )를 사용하여 평가됩니다. 이는 마스킹된 구성 요소에 대한 파라미터, 활성화 및 그래디언트를 제거하여 상당한 메모리와 연산 절감 효과를 가져옵니다. 그러나 그래디언트는 마스킹된 순전파(forward pass)를 반영합니다.
- 백워드 마스킹 (Backward-Masking): 순전파 시에는 전체 모델( $m_{fwd} = m_{uni}$ )을 사용하고, 역전파 및 집계 시에만 희소성(sparsity)을 적용합니다( $m_{bwd} = m$ ). 이는 편향되지 않은 그래디언트 추정치를 유지하며(이론적으로 입증됨) 원칙적인 베이스라인을 제공하지만, 활성화보다는 주로 그래디언트와 옵티마이저 상태에 대한 메모리 절감 효과를 제공합니다.
이론적 기초:
저자들은 $L$ -smooth 환경에서의 백워드 마스킹 체제에 대한 수렴 분석을 제공합니다. 이들은 수렴 속도가 유니폼 마스킹과 특정 마스크 분포 사이의 거리로 정의되는 "마스킹 오차(masking error)"에 의해 결정됨을 입증합니다. 또한, 주기적 평균화를 통한 중첩된 파라미터 할당이 부분적인 동기화를 유지함을 보여줌으로써 마스크 그래프의 스펙트럼 갭(spectral gap)과 수렴 품질 사이의 관계를 연결합니다.

주요 기여

새로운 패러다임: 모델을 노드 간에 분산시켜 디바이스당 메모리를 28%~60% 감소시키면서 성능을 유지하거나 향상시키는 메모리 효율적인 분산 학습 패러다임으로서 SDP를 도입했습니다.
이론적 보장: 백워드 마스킹된 SDP의 수렴을 마스크 그래프의 스펙트럼 갭 조건과 연결하여, 특정 마스킹 조건 하에서 편향되지 않은 그래디언트가 유지됨을 증명했습니다.
결합 가능성 (Composability): SDP가 기존의 메모리 절감 기술과 직교(orthogonal)함을 입증했습니다. SDP는 물리적으로 모델 구성 요소를 제거함으로써 파라미터, 그래디언트, 옵티마이저 상태 및 활성화를 동시에 축소합니다. 이를 통해 SDP는 FSDP (ZeRO-3) 및 **활성화 체크포인팅 (Activation Checkpointing, AC)**과 깔끔하게 결합될 수 있습니다.
일반화 가능성: LLaMA(대규모 언어 모델) 및 이미지 분류 모델(ResNet, Swin Transformer)을 포함한 다양한 아키텍처에 대해 SDP를 검증하여, 이것이 표준 데이터 병렬성의 "드롭인(drop-in)" 대안임을 보여주었습니다.

실험 결과

대규모 언어 모델 (LLaMA)

FineWeb 데이터셋을 사용하여 134M, 500M, 1B 파라미터 LLaMA 모델에 대한 실험을 수행했습니다.

메모리 효율성: 1B 규모에서, 69%의 활성 구성 요소( $C=0.69$ )를 가진 B-SDP는 DDP 대비 피크 GPU 메모리를 28% 감소시키면서도 동일하거나 더 낮은 검증 손실을 달성했습니다. 500M 규모에서는 50%의 활성 구성 요소를 가진 B-SDP가 메모리를 40% 감소시켰습니다.
다운스트림 성능: 5가지 다운스트림 벤치마크(ARC-E, BoolQ, HellaSwag, OBQA, SciQ)에서 모든 SDP 변형은 매크로 평균 기준 약 1포인트 이내에서 DDP 성능과 일치했으며, 일부 변형(B-SDP)은 특정 작업에서 약간의 개선을 보였습니다.
결합 가능성: B-SDP를 FSDP 및 활성화 체크포인팅과 결합했을 때, 1B 모델에 대해 디바이스당 피크 메모리 사용량은 9.7 GB를 기록했습니다. 이는 DDP(70.9 GB) 대비 85% 감소한 수치이며, FSDP+AC 단독 사용 시(11.0 GB)보다 엄격히 낮습니다.
통신: SDP는 스텝당 통신 볼륨을 약 28% 감소시켰습니다 (DDP의 3,143 MB 대비 2,253 MB 이동). 노드 간 설정에서 B-SDP와 FSDP의 조합은 DDP와 비교하여 실제 학습 시간(wall-clock time)을 19.9% 단축했습니다.

이미지 분류

CIFAR-10 및 CIFAR-100을 사용하여 ResNet-18, WideResNet-18, Swin-Tiny에 대한 실험을 수행했습니다.

ResNet-18: SDP 설정은 DDP 메모리 예산의 **40%**만을 사용하여 경쟁력 있거나 더 높은 정확도를 달로했습니다. N-SDP(뉴런 수준)는 64% 메모리 사용량에서 DDP를 능가하는 강력한 정규화 효과를 보여주었습니다.
Swin Transformer: B-SDP는 CIFAR-10에서 DDP 정확도와 일치하면서도 68%의 메모리를 사용했고, CIFAR-100에서는 32%의 메모리 감소와 함께 정확도를 2% 향상시켰습니다.
희소성 한계: 중간 정도의 희소성에서는 포워드 마스킹(B-SDP/N-SDP)이 잘 작동했지만, 극단적인 희소성 수준(예: $C=3/8$ )에서는 **백워드 마스킹(Bb-SDP)**이 다른 변형들이 실패하는 모델 붕괴(collapse)를 방지하며 우월함을 입증했습니다.

의의 및 주장

본 논문은 SDP가 데이터 병렬성과 모델 병렬성을 보완하는 실질적인 "세 번째 축"을 제공한다고 주장합니다. SDP의 주요 의의는 다음과 같습니다.

메모리와 대역폭의 분리: 워커 간에 활성화 값을 교환할 필요를 없앰으로써, SDP는 파이프라인 및 텐서 병렬성을 제한하는 흔한 병목 현상인 대역폭 문제를 완화합니다.
FLOP 매칭 스케일링 가능: SDP를 통해 동일한 하드웨어 예산 내에서 더 큰 모델이나 더 긴 시퀀스를 학습할 수 있습니다. 스텝당 연산 및 통신량이 감소함에 따라 발생하는 약간의 수렴 속도 저하는 더 많은 반복 횟수(FLOP-matched)를 수행할 수 있는 능력으로 상쇄됩니다.
드롭인 솔루션 제공: SDP는 마스킹 로직 외에는 기존의 분산 인프라를 변경할 필요 없이 표준 데이터 병렬의 드롭인 대체제로 기능합니다.

저자들은 SDP가 특히 메모리와 대역폭이 주요 제약 조건인 시나리오에서 매우 가치 있다고 강조하며, 복잡한 파이프라인 스케줄링이나 전체 샤딩의 통신 오버헤드 없이 메모리가 제한된 가속기에서 대규모 모델을 학습할 수 있는 확장 가능한 경로를 제공한다고 설명합니다.