스케일링 — 분산 학습, FSDP, DeepSpeed (Scaling: Distributed Training, FSDP, DeepSpeed)

124M 모델은 GPU 하나에서 학습했습니다. 이제 70억 파라미터(parameter)를 시도해봅니다. 모델은 메모리에 들어가지 않습니다. 데이터는 단일 머신에서 몇 주가 걸립니다. 규모가 커지면 분산 학습(Distributed Training)은 선택이 아닙니다. 앞으로 나아갈 유일한 길입니다.

유형: Build 언어: Python 선수 지식: Phase 10, Lesson 04 (미니 GPT 사전학습) 예상 시간: 약 120분

학습 목표

• 세 가지 병렬화(Parallelism)인 데이터 병렬화(Data Parallelism), 텐서 병렬화(Tensor Parallelism), 파이프라인 병렬화(Pipeline Parallelism)를 설명하고, 모델 크기와 클러스터 크기에 따라 각각이 언제 필요한지 판단합니다.
• 여러 GPU에 걸쳐 그래디언트 동기화(Gradient Synchronization)를 수행하는 PyTorch DDP 기반 데이터 병렬 학습을 구현합니다.
• 주어진 모델 크기에 대해 가중치(Weights), 옵티마이저 상태(Optimizer States), 그래디언트(Gradients), 활성값(Activations)을 포함한 메모리 예산(Memory Budget)을 계산하고 최소 하드웨어를 산정합니다.
• FSDP 또는 DeepSpeed ZeRO 단계를 설정해 모델 상태(Model States)를 여러 GPU에 샤딩(Sharding)하고, 단일 GPU 메모리를 넘는 모델을 적재합니다.

문제

7B 파라미터 모델은 FP16에서 가중치만 14GB가 필요합니다. Adam 옵티마이저는 모든 파라미터마다 두 개의 추가 사본, 즉 1차/2차 모멘트 추정치(First and Second Moment Estimates)를 저장합니다. 이것이 28GB를 더 차지합니다. 역전파 중에 그래디언트가 14GB를 더 추가합니다. 활성값 하나도 저장하기 전에 이미 56GB입니다.

NVIDIA A100에는 80GB 메모리가 있습니다.

80GB 중 56GB가 이미 소비되었습니다. 남은 24GB는 활성값, 즉 순전파 중에 계산되어 역전파를 위해 살아 있어야 하는 중간값(Intermediate Values)을 담는 공간입니다. 2048토큰 시퀀스와 4096차원 모델에서 한 층의 활성값은 약 64MB를 사용합니다. 32개 층이면 샘플 하나당 2GB가 필요합니다. 배치 크기 8이면 16GB입니다. 남은 공간은 24GB뿐입니다. 배치 크기를 12로 늘리면 메모리가 터집니다.

이제 70B 파라미터를 시도해봅니다. 가중치만 FP16 기준 140GB입니다. 단일 GPU에는 들어가지 않습니다. 가중치만 담으려 해도 최소 2개의 A100이 필요합니다(2 x 80GB = 160GB). 옵티마이저 상태와 그래디언트까지 더하면 훨씬 더 많은 GPU가 필요합니다. 최소 3개 이상, 현실적으로는 샤딩 전략에 따라 8~16개가 필요합니다.

Llama 3 405B는 NVIDIA H100 16,384개에서 학습되었습니다. 한 번의 학습 실행(training run)에 들어간 연산(compute) 비용은 약 1억 달러로 추정됩니다. DeepSeek V3는 아키텍처(Mixture of Experts는 토큰마다 일부 파라미터만 활성화합니다)와 학습 효율을 영리하게 설계해 비슷한 규모의 모델을 약 560만 달러에 학습했습니다.

이번 lesson은 대규모 학습을 가능하게 하는 네 가지 전략을 다룹니다. 데이터 병렬화, 텐서 병렬화, 파이프라인 병렬화, 완전 샤딩 데이터 병렬화(Fully Sharded Data Parallelism)입니다. 분산 학습 프레임워크를 다루기 전에 순수 Python으로 각 메커니즘을 시뮬레이션해 원리를 이해합니다.

개념

왜 분산이 필요한가 (Why Distribution is Required)

실제 모델에 대한 메모리 계산입니다. 모든 숫자는 추정이 아니라 계산된 값입니다.

"Adam States" 열이 진짜 골치입니다. Adam은 모든 파라미터마다 1차 모멘트의 이동 평균(running mean) m과 2차 모멘트의 이동 분산(running variance) v를 저장하며, 둘 다 FP32입니다. 70B 모델에서는 70B x 4 bytes x 2 = 560GB입니다. 옵티마이저만으로 A100 일곱 대가 필요합니다.

H100 하나는 80GB입니다. Llama 3 405B는 가중치, 옵티마이저 상태, 그래디언트를 담기 위해 최소 61개의 H100이 필요합니다. 활성값까지 더하면 숫자는 더 커집니다. Meta가 16,384개의 GPU를 사용한 것은 원해서가 아니라 그래야만 했기 때문입니다.

데이터 병렬화 (Data Parallelism)

가장 단순한 분산 전략입니다. 전체 모델을 N개의 GPU에 복사합니다. 각 학습 배치를 N개의 같은 크기 조각으로 나눕니다. 각 GPU는 자기 데이터 조각에서 순전파와 역전파를 수행합니다. 역전파가 끝나면 모든 GPU의 그래디언트를 평균냅니다. 모든 GPU는 같은 평균 그래디언트로 자기 가중치 사본을 업데이트하므로 모든 사본이 동기화됩니다.

장점: 처리량(throughput)이 선형으로 확장됩니다. N개의 GPU가 한 스텝(step)마다 N배 많은 데이터를 처리합니다. 통신은 그래디언트 평균 한 번으로 제한되며 계산과 겹칠 수 있습니다.

단점: 모든 GPU가 모델, 옵티마이저 상태, 그래디언트의 완전한 사본을 들고 있어야 합니다. 70B 모델에서는 GPU마다 840GB가 필요합니다. 데이터 병렬화는 GPU당 메모리를 줄여주지 않습니다. 학습 시간만 줄여줍니다.

수식: 유효 배치 크기(effective batch size) = GPU당 배치 크기 x N. N=64 GPU이고 GPU당 배치가 16이면 유효 배치는 1,024입니다. Llama 3는 스텝마다 1,600만 토큰의 유효 배치 크기를 사용했습니다.

graph TD
    subgraph DataParallel["Data Parallelism (N=4 GPUs)"]
        B["전체 배치\n(샘플 1024개)"] --> S["분할(Split)"]
        S --> G1["GPU 1\n전체 모델 사본\n샘플 256개"]
        S --> G2["GPU 2\n전체 모델 사본\n샘플 256개"]
        S --> G3["GPU 3\n전체 모델 사본\n샘플 256개"]
        S --> G4["GPU 4\n전체 모델 사본\n샘플 256개"]
        G1 --> AR["AllReduce\n그래디언트 평균"]
        G2 --> AR
        G3 --> AR
        G4 --> AR
        AR --> U["업데이트\n(모든 GPU 동일)"]
    end

    style B fill:#1a1a2e,stroke:#e94560,color:#fff
    style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G3 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G4 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style AR fill:#1a1a2e,stroke:#51cf66,color:#fff
    style U fill:#1a1a2e,stroke:#51cf66,color:#fff

텐서 병렬화 (Tensor Parallelism)

개별 층(Layer)을 여러 GPU에 나눕니다. 하나의 행렬 곱셈(Matrix Multiplication)을 GPU들에 분할하고, 각 GPU가 결과의 일부를 계산합니다.

피드포워드 층의 가중치 행렬이 (8192, 8192) 형태라고 합시다. 4-way 텐서 병렬화에서는 각 GPU가 (8192, 2048) 샤드(shard)를 가집니다. 각 GPU는 입력에 자기 샤드를 곱해 부분 결과(partial result)를 만듭니다. 부분 결과는 all-reduce 또는 all-gather로 결합되어 전체 출력을 만듭니다.

장점: GPU당 모델 가중치 메모리를 줄여줍니다. 70B 모델을 8개 GPU에 나누면 각 GPU는 약 8.75B 파라미터에 해당하는 가중치만 들고 있으면 됩니다.

단점: 매 층 뒤에 빠른 GPU 간 통신이 필요합니다. 각 matmul 뒤의 all-reduce가 지연 시간(latency)을 추가합니다. NVLink에서는 잘 동작합니다(같은 노드(node) 안의 GPU 사이 900 GB/s). 하지만 InfiniBand로 연결된 노드 사이에서는 좋지 않습니다(400 Gb/s, 약 50 GB/s). 텐서 병렬화는 거의 항상 단일 노드 안(GPU 8개)으로 제한됩니다.

실제 사용: Megatron-LM이 텐서 병렬화를 널리 알렸습니다. Llama 3 405B는 각 노드 안에서 8-way 텐서 병렬화를 사용합니다.

파이프라인 병렬화 (Pipeline Parallelism)

모델을 층 기준으로 나눕니다. GPU 1은 18층, GPU 2는 916층, GPU 3은 1724층, GPU 4는 2532층을 실행합니다. 데이터는 파이프라인을 따라 흐릅니다. GPU 1이 자기 층을 계산하고 활성값(activation)을 GPU 2로 보내면, GPU 2가 자기 층을 계산해 GPU 3으로 보내는 식입니다.

장점: GPU 사이의 통신이 적습니다. 층 경계의 활성값만 보내면 되며, 이는 그래디언트나 가중치에 비해 작은 양입니다. 대역폭 요구가 낮기 때문에 노드 사이에서도 잘 동작합니다.

단점: 파이프라인 버블(pipeline bubbles)이 생깁니다. GPU 4가 마이크로 배치(micro-batch) 1의 순전파를 계산하는 동안 GPU 1, 2, 3은 이미 자기 부분을 보냈으므로 놀고 있는(idle) 상태입니다. 역전파에서는 패턴이 반대로 나타납니다. 순진한(naive) 파이프라이닝에서는 N개 단계(stage)에서 GPU 활용률이 1/N에 불과합니다.

GPipe와 PipeDream은 배치를 마이크로 배치로 나누어 버블 문제를 해결합니다. GPU 1은 마이크로 배치 1의 순전파를 마치자마자 마이크로 배치 2를 시작합니다. 이렇게 파이프라인 단계 전반에 계산을 겹칩니다. 마이크로 배치가 M개이고 단계가 N개라면 버블 비율은 (N-1)/M으로 줄어듭니다. N=4, M=16이면 버블은 3/16 = 18.75% 유휴 시간입니다.

FSDP: 완전 샤딩 데이터 병렬 (Fully Sharded Data Parallel)

FSDP는 데이터 병렬화의 확장성과 샤딩의 메모리 효율을 결합합니다. 각 GPU가 전체 모델 사본을 들고 있는 대신, 각 GPU는 파라미터, 그래디언트, 옵티마이저 상태의 1/N만 보유합니다.

어떤 층의 순전파 전에 FSDP는 all-gather를 실행해 모든 GPU에서 전체 파라미터를 모아 각 GPU 메모리에 재구성합니다. 순전파 후에는 로컬이 아닌 파라미터를 버립니다. 역전파 중에는 그래디언트 계산을 위해 all-gather가 다시 실행됩니다. 역전파 후에는 reduce-scatter가 그래디언트 샤드를 분산해 각 GPU가 그래디언트의 1/N만 저장하도록 합니다.

8개 GPU에서 70B 모델을 학습하는 경우의 계산:

FSDP가 없으면 70B 모델은 80GB GPU 하나에 들어갈 수 없습니다. 8개 GPU에서 FSDP를 써도 GPU당 105GB입니다. 아직도 들어가지 않습니다. GPU당 80GB 아래로 내리려면 최소 16개 GPU가 필요하거나, 활성값 체크포인팅(activation checkpointing)과 결합해 역전파 때 활성값을 저장하지 않고 재계산해야 합니다.

통신 비용은 일반(vanilla) 데이터 병렬화보다 큽니다. 각 층 전에 all-gather가 필요하기 때문입니다. 하지만 메모리 절약 덕분에 이전에는 불가능했던 학습 실행이 가능해집니다.

graph TD
    subgraph FSDP["FSDP: Fully Sharded Data Parallel (4 GPUs)"]
        direction TB
        S["모델: 4개 층, 샤딩됨"]

        subgraph GPU1["GPU 1"]
            G1S["Shard: 파라미터 1/4\n옵티마이저 1/4\n그래디언트 1/4"]
        end
        subgraph GPU2["GPU 2"]
            G2S["Shard: 파라미터 1/4\n옵티마이저 1/4\n그래디언트 1/4"]
        end
        subgraph GPU3["GPU 3"]
            G3S["Shard: 파라미터 1/4\n옵티마이저 1/4\n그래디언트 1/4"]
        end
        subgraph GPU4["GPU 4"]
            G4S["Shard: 파라미터 1/4\n옵티마이저 1/4\n그래디언트 1/4"]
        end

        AG["All-Gather\n(각 층 전에 전체 파라미터 재구성)"]
        FW["순전파\n(전체 파라미터 일시 보유)"]
        RS["Reduce-Scatter\n(역전파 후 그래디언트 샤드 분배)"]

        S --> GPU1
        S --> GPU2
        S --> GPU3
        S --> GPU4
        GPU1 --> AG
        GPU2 --> AG
        GPU3 --> AG
        GPU4 --> AG
        AG --> FW
        FW --> RS
    end

    style G1S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G3S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G4S fill:#1a1a2e,stroke:#0f3460,color:#fff
    style AG fill:#1a1a2e,stroke:#e94560,color:#fff
    style FW fill:#1a1a2e,stroke:#51cf66,color:#fff
    style RS fill:#1a1a2e,stroke:#e94560,color:#fff

DeepSpeed ZeRO

DeepSpeed의 ZeRO(Zero Redundancy Optimizer)는 개념적으로 FSDP와 동일하지만 Microsoft가 독립적으로 개발했습니다. ZeRO는 세 단계를 정의하며, 각 단계는 더 공격적으로 샤딩합니다.

ZeRO-3는 FSDP와 동등합니다. 이름은 다르지만 메커니즘은 같습니다. PyTorch는 DeepSpeed가 이 개념을 입증한 뒤 자체(native) 구현으로 FSDP를 추가했습니다.

DeepSpeed는 ZeRO-Offload(옵티마이저 상태를 더 싸고 큰 CPU RAM으로 오프로드(offload))와 ZeRO-Infinity(NVMe SSD로 오프로드)도 도입했습니다. 이들은 학습 속도를 메모리 용량과 맞바꿉니다. 오프로드된 연산은 느리지만 GPU 메모리를 비워줍니다.

혼합 정밀도 학습 (Mixed Precision Training)

현대 학습은 여러 부동소수점 형식(floating-point format)을 동시에 사용합니다.

• 순전파(Forward pass): FP16 또는 BF16(16비트). FP32의 절반 메모리입니다. 텐서 코어(tensor core)에서 matmul이 2배 빠르게 실행됩니다.
• 마스터 가중치(Master weights): FP32(32비트). 옵티마이저가 가중치 업데이트 중 수치 정밀도(numerical precision)를 유지하기 위해 보관합니다.
• 손실 스케일링(Loss scaling): FP16 그래디언트가 0으로 언더플로(underflow)되는 것을 막기 위해 역전파 전에 손실(loss)에 큰 상수를 곱합니다. 옵티마이저 스텝 전에 같은 상수로 나눕니다.

BF16(Brain Float 16)은 FP32와 같은 지수 범위(exponent range, 지수 비트 8개)를 가지지만 정밀도는 낮습니다(가수 비트(mantissa bit) 7개로, FP32의 23개에 비해 적습니다). 같은 범위의 값을 표현할 수 있기 때문에 손실 스케일링이 거의 필요 없습니다. FP16은 지수 비트 5개와 가수 비트 10개를 가집니다. 세밀한 값을 표현할 수 있지만 극단적인 크기에서는 오버플로(overflow)/언더플로가 발생합니다.

Google TPU는 BF16을 자체적으로 사용합니다. NVIDIA A100과 H100은 FP16과 BF16을 모두 지원합니다. 업계는 대체로 BF16으로 이동했습니다. 손실 스케일링 문제를 줄여주기 때문입니다.

7B 모델의 메모리 비교:

혼합 정밀도는 이 모델에서 28GB를 절약합니다. 옵티마이저 상태는 여전히 FP32로 남습니다. 메모리 대부분은 결국 여기에서 잡아먹습니다.

Megatron-LM과 3D 병렬화 (Megatron-LM and 3D Parallelism)

실제 대규모 학습은 세 가지 병렬화를 모두 결합합니다.

• 노드 그룹 전체에 걸친 데이터 병렬화(data parallelism): 배치 크기를 키웁니다.
• 노드 내부의 텐서 병렬화(tensor parallelism): 층을 GPU 8개에 나눕니다.
• 노드 간 파이프라인 병렬화(pipeline parallelism): 층 그룹을 여러 머신에 나눕니다.

H100 16,384개에서의 Llama 3 405B 학습:

• 각 노드 내부의 8-way 텐서 병렬화 (노드당 GPU 8개)
• 노드 간 16-way 파이프라인 병렬화 (16개 파이프라인 단계)
• 남은 차원의 128-way 데이터 병렬화 (16,384 / 8 / 16 = 128)

이 3D 분해(8 x 16 x 128 = 16,384)가 수천 개 GPU로 확장하는 방법입니다. 각 GPU는 서로 다른 데이터 샤드를 보고(데이터 병렬), 각 층의 한 조각을 보유하며(텐서 병렬), 서로 다른 층 집합을 계산합니다(파이프라인 병렬).

DeepSeek V3는 다른 접근을 취했습니다. Mixture of Experts 아키텍처는 토큰마다 671B 파라미터 중 37B만 활성화합니다. 따라서 각 GPU는 활성 파라미터만 계산하고 그에 대한 활성값만 저장하면 됩니다. 이들은 2,048개의 H800 GPU에서 학습했습니다. Meta GPU 수의 1/8도 안 되며, 비용도 Meta의 추정 1억 달러 대비 560만 달러 수준이었습니다.

graph TD
    subgraph ThreeD["3D Parallelism (Llama 3 405B)"]
        direction TB
        subgraph DP["Data Parallel (128-way)\n128개 그룹에 배치 분할"]
            subgraph PP["Pipeline Parallel (16-way)\n16개 단계에 층 분할"]
                subgraph TP["Tensor Parallel (8-way)\n각 층을 GPU 8개에 분할"]
                    G1["GPU 1\n1~N층의 slice"]
                    G2["GPU 2\n1~N층의 slice"]
                    G8["GPU 8\n1~N층의 slice"]
                end
            end
        end
    end

    N1["총합: 8 x 16 x 128 = 16,384 GPU"]

    style G1 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G2 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style G8 fill:#1a1a2e,stroke:#0f3460,color:#fff
    style N1 fill:#1a1a2e,stroke:#e94560,color:#fff

직접 만들기

Step 1: 데이터 병렬화 시뮬레이션 (Simulate Data Parallelism)

배치를 시뮬레이션된 GPU들에 나눕니다. 각 GPU는 자기 샤드에서 순전파를 계산합니다. "그래디언트"는 손실 값으로 시뮬레이션하고 평균을 냅니다.

import numpy as np

def simulate_data_parallelism(data, num_gpus, model_fn):
    batch_size = len(data)
    shard_size = batch_size // num_gpus
    remainder = batch_size % num_gpus

    gpu_losses = []
    gpu_gradients = []

    offset = 0
    for gpu_id in range(num_gpus):
        extra = 1 if gpu_id < remainder else 0
        shard = data[offset:offset + shard_size + extra]
        offset += shard_size + extra

        loss, grad = model_fn(shard)
        gpu_losses.append(loss)
        gpu_gradients.append(grad)

    avg_loss = np.mean(gpu_losses)
    avg_gradient = np.mean(gpu_gradients, axis=0)

    return avg_loss, avg_gradient

all-reduce 연산, 즉 그래디언트 평균이 데이터 병렬화의 유일한 통신입니다. 실제로는 NVIDIA GPU에서 NCCL 라이브러리를 사용하며, 링 all-reduce(ring all-reduce)를 구현합니다. 각 GPU는 자기 그래디언트의 1/N을 한쪽 이웃에게 보내고 반대쪽 이웃에게서 1/N을 받습니다. N-1 스텝 뒤에는 모든 GPU가 완전한 평균을 갖게 됩니다. 전체 통신량은 2 x gradient_size x (N-1)/N이며, N이 커질수록 그래디언트 크기의 약 2배에 가까워집니다.

Step 2: 텐서 병렬화 시뮬레이션 (Simulate Tensor Parallelism)

가중치 행렬을 여러 GPU에 나눕니다. 각 GPU가 부분 행렬 곱셈을 계산합니다. 그 결과를 결합합니다.

def simulate_tensor_parallelism(input_data, weight_matrix, num_gpus):
    d_in, d_out = weight_matrix.shape
    assert d_out % num_gpus == 0, f"d_out {d_out} not divisible by num_gpus {num_gpus}"
    shard_size = d_out // num_gpus

    partial_results = []
    for gpu_id in range(num_gpus):
        start = gpu_id * shard_size
        end = start + shard_size
        weight_shard = weight_matrix[:, start:end]

        partial = input_data @ weight_shard
        partial_results.append(partial)

    full_output = np.concatenate(partial_results, axis=-1)

    direct_output = input_data @ weight_matrix
    error = np.abs(full_output - direct_output).max()

    return full_output, error

오차는 정확히 0이거나 머신 엡실론(machine epsilon) 수준이어야 합니다. 텐서 병렬화는 수학적으로 정확합니다. 단일 GPU에서 전체 matmul을 계산한 것과 같은 결과를 냅니다. 출력 차원 기준으로 나누기 때문에 각 GPU는 서로 다른 열 묶음을 만들고, 이를 이어 붙이면(concatenate) 전체 결과가 재구성됩니다.

열 병렬(column-parallel) 선형 층(출력 차원 분할)에서는 이어 붙입니다. 행 병렬(row-parallel, 입력 차원 분할)에서는 합산합니다. 트랜스포머 FFN에서는 첫 번째 선형 층(확장, expand)이 열 병렬을 사용하고 두 번째 선형 층(축소, contract)이 행 병렬을 사용합니다. 이렇게 하면 두 층 사이에 일어날 all-reduce 한 번을 아낄 수 있습니다.

Step 3: 파이프라인 병렬화 시뮬레이션 (Simulate Pipeline Parallelism)

모델 층을 가상 GPU들에 나눕니다. 앞쪽 단계가 놀고 있는 버블 문제를 보여줍니다.

def simulate_pipeline_parallelism(num_layers, num_stages, num_microbatches):
    layers_per_stage = num_layers // num_stages

    timeline = {}
    clock = 0

    for mb in range(num_microbatches):
        for stage in range(num_stages):
            start_time = max(
                timeline.get((stage, mb - 1, "fwd"), (0, 0))[1] if mb > 0 else 0,
                timeline.get((stage - 1, mb, "fwd"), (0, 0))[1] if stage > 0 else 0,
            )
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "fwd")] = (start_time, end_time)

    last_fwd_end = max(v[1] for v in timeline.values())

    for mb in range(num_microbatches - 1, -1, -1):
        for stage in range(num_stages - 1, -1, -1):
            deps = [last_fwd_end]
            if mb < num_microbatches - 1 and (stage, mb + 1, "bwd") in timeline:
                deps.append(timeline[(stage, mb + 1, "bwd")][1])
            if stage < num_stages - 1 and (stage + 1, mb, "bwd") in timeline:
                deps.append(timeline[(stage + 1, mb, "bwd")][1])
            start_time = max(deps)
            end_time = start_time + layers_per_stage
            timeline[(stage, mb, "bwd")] = (start_time, end_time)

    total_time = max(v[1] for v in timeline.values())
    compute_time = num_microbatches * num_stages * layers_per_stage * 2
    bubble_fraction = 1.0 - compute_time / (total_time * num_stages)

    return timeline, total_time, bubble_fraction

4개 단계와 1개 마이크로 배치에서는 버블 비율이 75%입니다. 항상 네 GPU 중 세 개가 놀고 있습니다. 16개 마이크로 배치를 쓰면 약 19%로 떨어집니다. 버블을 줄이는 비용은 메모리입니다. 동시에 진행 중인 모든 마이크로 배치의 활성값을 저장해야 하기 때문입니다.

Step 4: 메모리 계산기 (Memory Calculator)

어떤 모델 크기에 대해서도 정확한 학습 메모리 요구량을 계산합니다.

def memory_calculator(
    params_billions,
    precision_bytes=2,
    optimizer="adam",
    num_gpus=1,
    sharding="none",
    sequence_length=2048,
    batch_size_per_gpu=1,
    hidden_dim=None,
    num_layers=None,
):
    params = params_billions * 1e9

    weight_memory = params * precision_bytes

    if optimizer == "adam":
        optimizer_memory = params * 4 * 2
    elif optimizer == "sgd":
        optimizer_memory = params * 4
    else:
        optimizer_memory = 0

    gradient_memory = params * precision_bytes

    total_no_activation = weight_memory + optimizer_memory + gradient_memory

    if hidden_dim and num_layers:
        activation_per_layer = (
            sequence_length * batch_size_per_gpu * hidden_dim * precision_bytes * 4
        )
        activation_memory = activation_per_layer * num_layers
    else:
        activation_memory = params * precision_bytes * 0.5

    if sharding == "fsdp" or sharding == "zero3":
        weight_memory /= num_gpus
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero2":
        optimizer_memory /= num_gpus
        gradient_memory /= num_gpus
    elif sharding == "zero1":
        optimizer_memory /= num_gpus

    per_gpu_total = weight_memory + optimizer_memory + gradient_memory + activation_memory

    return {
        "params_billions": params_billions,
        "weights_gb": weight_memory / 1e9,
        "optimizer_gb": optimizer_memory / 1e9,
        "gradients_gb": gradient_memory / 1e9,
        "activations_gb": activation_memory / 1e9,
        "per_gpu_total_gb": per_gpu_total / 1e9,
        "total_across_gpus_gb": per_gpu_total * num_gpus / 1e9,
        "fits_on_80gb": per_gpu_total / 1e9 <= 80,
        "num_gpus": num_gpus,
        "sharding": sharding,
    }

이 계산기는 모든 ML 엔지니어가 묻는 질문에 답합니다. "GPU가 몇 개 필요할까?" 모델 크기를 넣어 보고 메모리에 들어가는지 확인합니다. GPU당 총량이 80GB 아래로 내려갈 때까지 샤딩 전략을 조정합니다.

Step 5: 혼합 정밀도 시뮬레이션 (Mixed Precision Simulation)

FP32, FP16, 혼합 정밀도 학습의 메모리 사용량을 비교합니다.

def mixed_precision_comparison(params_billions):
    params = params_billions * 1e9

    fp32_weights = params * 4
    fp32_optimizer = params * 4 * 2
    fp32_gradients = params * 4
    fp32_total = fp32_weights + fp32_optimizer + fp32_gradients

    fp16_weights = params * 2
    fp16_master = params * 4
    fp16_optimizer = params * 4 * 2
    fp16_gradients = params * 2
    fp16_total = fp16_weights + fp16_master + fp16_optimizer + fp16_gradients

    mixed_weights = params * 2
    mixed_optimizer = params * 4 * 2
    mixed_gradients = params * 2
    mixed_total = mixed_weights + mixed_optimizer + mixed_gradients

    return {
        "fp32_total_gb": fp32_total / 1e9,
        "fp16_with_master_gb": fp16_total / 1e9,
        "mixed_bf16_gb": mixed_total / 1e9,
        "savings_vs_fp32": 1 - mixed_total / fp32_total,
    }

대부분의 사람이 가장 놀라는 점은 이것입니다. 혼합 정밀도는 메모리를 절반으로 줄여주지 않습니다. 옵티마이저 상태, 즉 Adam의 m과 v는 정밀도와 무관하게 FP32로 남습니다. 7B 모델에서 FP32 학습은 112GB를 사용하고, 혼합 정밀도는 84GB를 사용합니다. 50%가 아니라 25% 감소입니다. 옵티마이저가 메모리를 지배합니다.

사용해보기

모든 시뮬레이션 실행 (Run All Simulations)

def run_all_demos():
    print("=" * 70)
    print("DATA PARALLELISM SIMULATION")
    print("=" * 70)

    np.random.seed(42)
    data = np.random.randn(64, 32)
    weight = np.random.randn(32, 16)

    def model_fn(batch):
        output = batch @ weight
        loss = np.mean(output ** 2)
        grad = 2 * batch.T @ (batch @ weight) / len(batch)
        return loss, grad

    for n_gpus in [1, 2, 4, 8]:
        loss, grad = simulate_data_parallelism(data, n_gpus, model_fn)
        print(f"  {n_gpus} GPUs: loss={loss:.4f}, grad_norm={np.linalg.norm(grad):.4f}")

    print()
    print("=" * 70)
    print("TENSOR PARALLELISM SIMULATION")
    print("=" * 70)

    x = np.random.randn(4, 8192)
    W = np.random.randn(8192, 8192)

    for n_gpus in [1, 2, 4, 8]:
        output, error = simulate_tensor_parallelism(x, W, n_gpus)
        print(f"  {n_gpus} GPUs: output_shape={output.shape}, max_error={error:.2e}")

    print()
    print("=" * 70)
    print("PIPELINE PARALLELISM SIMULATION")
    print("=" * 70)

    for n_mb in [1, 4, 8, 16, 32]:
        _, total_t, bubble = simulate_pipeline_parallelism(32, 4, n_mb)
        print(f"  {n_mb:2d} micro-batches: total_time={total_t:4d}, bubble={bubble:.1%}")

    print()
    print("=" * 70)
    print("MEMORY CALCULATOR")
    print("=" * 70)

    configs = [
        (7, "none", 1),
        (7, "fsdp", 8),
        (70, "none", 1),
        (70, "fsdp", 8),
        (70, "fsdp", 16),
        (405, "fsdp", 64),
        (405, "fsdp", 128),
    ]

    print(f"  {'Model':>8} {'Sharding':>8} {'GPUs':>5} {'Per-GPU':>10} {'Fits 80GB':>10}")
    print("  " + "-" * 50)
    for params, shard, gpus in configs:
        result = memory_calculator(params, num_gpus=gpus, sharding=shard)
        fits = "Yes" if result["fits_on_80gb"] else "No"
        print(f"  {params:>6}B {shard:>8} {gpus:>5} {result['per_gpu_total_gb']:>8.1f}GB {fits:>10}")

    print()
    print("=" * 70)
    print("MIXED PRECISION COMPARISON")
    print("=" * 70)

    for params_b in [7, 13, 70, 405]:
        result = mixed_precision_comparison(params_b)
        print(f"  {params_b}B: FP32={result['fp32_total_gb']:.0f}GB, "
              f"Mixed BF16={result['mixed_bf16_gb']:.0f}GB, "
              f"Savings={result['savings_vs_fp32']:.0%}")

산출물 만들기

이 lesson은 outputs/prompt-distributed-training-planner.md를 산출합니다. 모델 크기와 사용 가능한 하드웨어를 입력으로 받아 완전한 분산 학습 계획을 만들어냅니다. 병렬화 전략, 메모리 예산, 통신 오버헤드, 예상 처리량을 포함합니다.

연습문제

1. (쉬움) 메모리 계산기에 활성값 체크포인팅(activation checkpointing)을 추가합니다. 체크포인팅을 사용하면 K번째 층마다 활성값만 저장합니다. 보통 K=1은 모두 재계산한다는 뜻입니다. 메모리-연산 트레이드오프를 보여줍니다. 체크포인팅은 메모리를 얼마나 절약하고, 학습을 얼마나 느리게 만들까요? 전체 체크포인팅은 대략 33% 더 많은 연산을 사용합니다.

2. (중간) 파이프라인 병렬화 시뮬레이션을 확장해 PipeDream이 사용하는 1F1B(one forward, one backward) 스케줄을 구현합니다. 4개 단계와 8개 마이크로 배치에서 순진한(naive) 스케줄과 버블 비율을 비교합니다. 1F1B는 역전파를 더 일찍 시작하므로 최대 메모리(peak memory)가 더 작아야 합니다.

3. (중간) 그래디언트 누적 시뮬레이터(gradient accumulation simulator)를 구현합니다. 마이크로 배치마다 all-reduce를 수행하는 대신 K 스텝 동안 로컬에서 그래디언트를 누적한 뒤 all-reduce를 수행합니다. 이렇게 하면 통신이 K배 줄지만 최종 그래디언트는 동일하다는 것을 보여줍니다.

4. (어려움) 비용 추정기(cost estimator)를 만듭니다. 모델 크기, 목표 토큰 수, GPU 유형(A100은 $2/hr, H100은 $3.50/hr), 병렬화 전략이 주어졌을 때 전체 학습 비용을 달러로 추정합니다. 알려진 비용과 비교해 검증합니다. Llama 3 405B는 약 1억 달러, DeepSeek V3는 약 560만 달러로 보고되었습니다.

5. (어려움) 메모리 계산기에 ZeRO-Offload를 추가합니다. 노드당 CPU RAM 512GB, NVMe 2TB가 있다고 가정합니다. 옵티마이저 상태를 CPU로 오프로드하면 70B 모델을 16개가 아닌 4개 GPU에서 학습할 수 있음을 보여줍니다. 단, 옵티마이저 스텝은 30~50% 느려집니다.

핵심 용어

더 읽을거리

• Rajbhandari et al., 2020 -- "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models" - 세 가지 샤딩 단계를 정의한 DeepSpeed ZeRO 논문입니다.
• Shoeybi et al., 2020 -- "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism" - 트랜스포머를 위한 NVIDIA의 텐서 병렬화입니다.
• Narayanan et al., 2021 -- "Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM" - 데이터, 텐서, 파이프라인을 결합한 3D 병렬화를 다룹니다.
• Zhao et al., 2023 -- "PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel" - PyTorch의 자체 FSDP 구현 경험을 다룹니다.
• Llama 3 Technical Report - 16,384개 GPU 학습과 3D 병렬화 세부 내용을 다룹니다.
• DeepSeek-V3 Technical Report - MoE 아키텍처가 학습 비용을 한 자릿수 규모로 줄이는 방법을 설명합니다.