서론

수천 개의 GPU로 구성된 클러스터에서 거대 언어 모델(LLM)을 학습시키는 현장을 상상해 보십시오. 막대한 전력이 소모되고, 팬이 윙윙거리는 소음 속에서 학습이 진행됩니다. 이때 단 하나의 네트워크 스위치 결함이나, 일부 GPU의 메모리 오류로 인해 전체 학습 작업이 중단된다면 어떨까요? 기존의 분산 학습 방식인 데이터 병렬 처리(Data Parallelism)는 모든 워커(Worker)가 매 스텝마다 그래디언트(Gradient)를 동기화해야 하므로, 네트워크 지연이나 노드 장애에 극도로 취약합니다. 대규모 클러스터에서 “장애는 필연"이라는 관점에서 볼 때, 이러한 강력한 결합(Tight Coupling) 구조는 운영상 재앙과도 같습니다.

DeepMind가 제안하는 **Decoupled DiLoCo(Distributed Low-Communication Training)**는 이러한 통신 병목과 결합도 문제를 해결하기 위한 정교한 접근법입니다. 기존 DiLoCo가 통신 빈도를 줄이는 데 초점을 맞췄다면, Decoupled DiLoCo는 시스템의 **탄력성(Resilience)**에 집중합니다. 전체 클러스터가 동시에 완벽하게 동기화될 필요 없이, 각 워커가 독립적으로 학습을 진행하고 느슨하게(Loosely) 상태를 공유함으로써 특정 노드의 장애가 전체 학습을 멈추게 하지 않는 환경을 구현합니다. 이는 특히 네트워크 대역폭이 제한적이거나, 스팟 인스턴스(Spot Instance)와 같은 중단 가능한 리소스를 활용해야 하는 상용 환경에서 게임 체인저가 될 수 있습니다.

본론

기술적 원리: 동기화의 탈결합 (Decoupling)

기존의 분산 학습에서는 AllReduce 알고리즘을 통해 매 스텝 혹은 미니배치마다 모든 워커의 파라미터를 평균화합니다. 이는 통신 오버헤드가 매우 큽니다. 반면, DiLoCo 계열의 알고리즘은 로컬 SGD(Local SGD) 방식을 기반으로 합니다. 각 워커는 일정 횟수($H$)의 스텝 동안 전혀 통신 없이 자신만의 데이터로 학습(Local Training)을 진행합니다. 그 후, 주기적으로 전역 동기화(Global Synchronization) 단계에서 모델 파라미터가 아닌 **옵티마이저 상태(Optimizer State)**를 공유합니다.

Decoupled DiLoCo의 핵심 혁신은 이 동기화 과정을 “탈결합"했다는 점입니다. 특정 워커가 네트워크 오류로 연결이 끊기거나 늦어지더라도, 글로벌 코디네이터는 기다리지 않고 사용 가능한 워커들의 상태만으로 모델을 업데이트합니다. 장애가 발생한 워커는 복구 시 이전 체크포인트부터 다시 시작하는 것이 아니라, 현재의 글로벌 상태를 동기화받아 학습에 재합류할 수 있습니다. 이는 **옵티마이저 상태 평균화(Optimizer State Averaging)**를 통해 가능하며, 단순한 파라미터 평균화보다 학습 안정성이 훨씬 높습니다.

다음은 기존 방식과 Decoupled DiLoCo의 학습 흐름을 비교한 간단한 다이어그램입니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

graph LR
    subgraph Standard_DDP
        W1[Worker 1] -->|매 스텝 동기화| Sync[Barrier Sync]
        W2[Worker 2] -->|매 스텝 동기화| Sync
        W3[Worker 3] -->|매 스텝 동기화| Sync
    end

    subgraph Decoupled_DiLoCo
        A[Worker 1] -->|로컬 스텝 N회 수행| B[비동기 상태 공유]
        C[Worker 2] -->|로컬 스텝 N회 수행| B
        D[Worker 3] -->|장애 발생 및 재연결| B
        B -->|옵티마이저 상태 평균| A
        B -->|옵티마이저 상태 평균| C
        B -->|옵티마이저 상태 평균| D
    end

성능 비교 및 통신 효율성

Decoupled DiLoCo를 도입했을 때 얻을 수 있는 이점은 명확합니다. 아래 표는 전통적인 DDP(Data Distributed Parallel) 방식과 DiLoCo, 그리고 Decoupled DiLoCo의 특징을 비교한 것입니다.

| 비교 항목 | Standard DDP | DiLoCo | Decoupled DiLoCo | | :— | :— | :— | :— | | 동기화 주기 | 매 스텝 (Step) | $H$ 스텝마다 (약간 느슨함) | 비동기적 (완전히 탈결합) | | 통신 대상 | 그래디언트 (Gradient) | 옵티마이저 상태 (Adam Moments) | 옵티마이저 상태 (Adam Moments) | | 네트워크 대역폭 | 매우 높음 (병목 발생) | 낮음 (효율적) | 매우 낮음 (탄력적 운용 가능) | | 결함 허용 (Fault Tolerance) | 낮음 (하나의 노드 장애로 전체 중단) | 중간 (동기화 시점까지 영향 없음) | 높음 (장애 노드 무시 및 자동 재합류) | | 주요 사용 사례 | 단일 클러스터 내 고속 학습 | 대역폭이 제한된 분산 환경 | 공용 클라우드, 스팟 인스턴스, 지리적 분산 학습 |

PyTorch를 활용한 개념적 구현

Decoupled DiLoCo는 아직 PyTorch의 torch.distributed 코어에 기본 내장된 기능은 아니지만, 커스텀 로직을 통해 구현할 수 있습니다. 아래 코드는 간단한 Local SGD와 옵티마이저 상태 공유를 시뮬레이션한 개념적 예시입니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def train_diloco_step(model, optimizer, dataloader, steps_to_communicate=100):
    """
    DiLoCo 스타일의 학습 루프 예시
    :param steps_to_communicate: 로컬 학습 스텝 수 (H)
    """
    model.train()
    local_step_count = 0
    
    for inputs, targets in dataloader:

출처: https://deepmind.google/blog/decoupled-diloco/