Optimizing Optimizers

개별 optimizer step 을 본다면 — Adam 한 줄은 w -= lr · m / (sqrt(v) + eps) 같이 가벼운 산술이다. 그런데 학습 루프 전체에서 봤을 때, 큰 모델에서 optimizer 가 한 step 의 10–30% 까지 차지하는 일이 흔하다. 왜? 모델의 모든 파라미터 텐서마다 같은 update 를 따로 launch 하기 때문.

강의가 답하는 질문 셋.

1. 왜 optimizer 가 무거운가 — 산술이 무거운 게 아니라 launch 가 많은 것이 문제. (§02–§03 의 정량.)
2. PyTorch 가 이 문제를 어떻게 풀어왔는가 — foreach → fused → multi_tensor_apply → torch.compile 의 사다리.
3. 큰 모델에서 optimizer state 자체가 메모리 한계인 자리 — CPU offload, unified memory 같은 후속 영역의 첫 그림.

Adam 의 update 식. 한 파라미터 텐서 p 에 대해 첫째 모멘트 m, 둘째 모멘트 v 를 들고 있고, gradient g 가 들어왔을 때.

m = β₁ * m + (1 - β₁) * g                # 1 mul + 1 mul + 1 add = ~3 launch
v = β₂ * v + (1 - β₂) * g²               # pow + 1 mul + 1 add = ~3-4 launch
m̂ = m / (1 - β₁^t)
v̂ = v / (1 - β₂^t)
p = p - lr * m̂ / (sqrt(v̂) + ε)            # 1 sqrt + 1 div + 1 mul + 1 sub = ~4 launch

이 한 줄짜리 식이 PyTorch eager 에서는 한 파라미터 텐서당 약 10개 안팎의 op로 풀린다 — pointwise mul, add, sub, sqrt, div 의 시퀀스. Inductor 또는 fused kernel 이 없으면 각자가 따로 launch.

강의에서 Jane 이 처음 보여주는 정량적 사실 — 모델이 LLaMA-7B 같은 규모면 파라미터 텐서가 수천 개다. 한 step 의 launch 수가 급격히 늘어난다.

이 표가 강의의 motivation. 실제 산술 시간은 거의 안 늘어났는데 (단순한 elementwise op), launch overhead 가 step 의 30~50% 까지 차지한다. fusion 또는 묶음으로 launch 수를 줄이면 그대로 해결.

강의의 정량. 한 launch 의 fixed cost ~10 μs (host driver 기준 추정), 작은 elementwise 커널의 실제 GPU 시간 ~5 μs. 즉 커널 자체보다 launch 가 더 길다. 이 비율이 작은 op + 많은 tensor 인 경우의 본질.

같은 일을 푸는 두 방향.

• tensor 차원의 fusion — N 개의 같은 op 를 한 launch 로 묶는다. foreach, multi_tensor_apply, fused 가 다 이 방향.
• op 차원의 fusion — 같은 tensor 에 대한 여러 op 를 한 커널로 묶는다. torch.compile 이 이 방향.

두 방향이 결합하면 — 한 step 의 모든 optimizer launch 를 한 두 개의 커널 launch 로. 이게 PyTorch 가 도달한 현재 상태.

PyTorch 의 첫 해결책. torch._foreach_add_(tensors, alpha=…) 같은 함수가 tensor 의 list 를 받아 같은 op 를 일괄 적용. 내부적으로는 한 launch (또는 매우 적은 수의 launch) 로 처리.

# Adam 의 한 step (foreach 스타일, 약식)
def adam_step_foreach(params, grads, m_list, v_list, lr, b1, b2, eps, t):
    torch._foreach_mul_(m_list, b1)
    torch._foreach_add_(m_list, grads, alpha=1 - b1)
    torch._foreach_mul_(v_list, b2)
    torch._foreach_addcmul_(v_list, grads, grads, value=1 - b2)

    bias1 = 1 - b1**t
    bias2 = 1 - b2**t
    step_size = lr / bias1

    denom = torch._foreach_sqrt(v_list)
    torch._foreach_div_(denom, math.sqrt(bias2))
    torch._foreach_add_(denom, eps)

    torch._foreach_addcdiv_(params, m_list, denom, value=-step_size)

foreach 가 “한 op 를 여러 tensor 에” 묶었다면, fused 는 한 단 더 — Adam 의 모든 산술 (mul, add, sqrt, div, addcdiv) 을 한 CUDA 커널로. 한 step 의 모든 launch 가 사실상 한 두 개로 줄어든다.

# Python 에서는 한 줄
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4, fused=True)

# 내부적으로는 — 약간 단순화한 fused Adam kernel signature
__global__ void adam_kernel(
    TensorListMetadata<4> meta,    // params, grads, m, v 의 pointer list
    float lr, float b1, float b2, float eps, int t)
{
    // 한 thread 가 한 element 의 모든 산술을 끝낸다
    int tensor_idx = blockIdx.y;
    float* p = meta.tensors[0][tensor_idx];
    float* g = meta.tensors[1][tensor_idx];
    float* m = meta.tensors[2][tensor_idx];
    float* v = meta.tensors[3][tensor_idx];
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < meta.sizes[tensor_idx]) {
        m[i] = b1 * m[i] + (1 - b1) * g[i];
        v[i] = b2 * v[i] + (1 - b2) * g[i] * g[i];
        float mhat = m[i] / (1 - powf(b1, t));
        float vhat = v[i] / (1 - powf(b2, t));
        p[i] -= lr * mhat / (sqrtf(vhat) + eps);
    }
}

강의에서 Jane 이 가장 흥미롭게 깐 자리. 한 커널에 “여러 tensor” 를 전달하려면 — 각 tensor 의 device pointer 를 launch 시점에 다 넘겨야 한다. 그런데 CUDA 의 kernel 인자는 4 KB 한도 (cudaLaunchKernel arg buffer). pointer 한 개가 8 byte 면, ~500 개 tensor 까지밖에 못 넘긴다. 큰 모델은 그 이상.

해결책 — pointer 를 struct 안에 묶고, struct 들을 chunk 로 나눠 여러 launch. 강의에서 보여준 패턴.

// PyTorch 의 multi_tensor_apply 의 본질 (단순화)
template<int N>
struct TensorListMetadata {
    void* tensors[N][CHUNK_SIZE];   // 각 tensor list 의 pointer 들
    int   sizes[CHUNK_SIZE];        // 각 tensor 의 element 수
};

void multi_tensor_apply(
    std::vector<std::vector<Tensor>>& tensor_lists,
    Op op)
{
    int n_tensors = tensor_lists[0].size();
    for (int chunk_start = 0; chunk_start < n_tensors; chunk_start += CHUNK_SIZE) {
        TensorListMetadata<N> meta;
        // fill meta with chunk_start..chunk_start+CHUNK_SIZE pointers
        op<<<blocks, threads>>>(meta, ...);   // 한 chunk 당 한 launch
    }
}

이 한 디테일이 시리즈 안에서 자주 다시 등장 — “한 launch 에 어디까지 묶을 수 있는가” 가 항상 인자 한도와 같이. CUDA Graphs 가 또 다른 답 (§07 의 직전).

PyTorch 의 가장 최근 사다리. torch.compile 이 optimizer 의 Python 코드를 통째로 trace 해서 Triton fused 커널로 만들어 준다. 사용자가 직접 fused 버전을 짜지 않아도 — Inductor 가 같은 fusion 을 자동으로.

# 한 줄 추가로 모든 optimizer 가 “자동 fused”
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
optimizer.step = torch.compile(optimizer.step)

# custom optimizer 도 동일하게
class MyOptimizer(torch.optim.Optimizer):
    def step(self):
        # 평범한 Python 산술
        ...

opt = MyOptimizer(...)
opt.step = torch.compile(opt.step)   # 같은 fusion 을 받는다

이 4단 사다리 (eager → foreach → fused → compile) 가 같은 알고리즘에 대한 PyTorch 의 진화 패턴이다. matmul, attention, normalization 모두 비슷한 사다리를 따라간다 — 이 강의가 그 패턴의 가장 명확한 한 사례.

강의 후반에 Jane 이 짚는 다음 영역. Adam 은 파라미터당 m, v 의 두 state 를 들고 있어 → 모델 메모리의 3배. 큰 모델 (LLaMA-7B FP32 = 28 GB · Adam state = 56 GB · 총 84 GB) 면 GPU 한 장에 안 들어간다.

해결책 셋.

• FSDP — Fully Sharded Data Parallel. optimizer state 를 GPU 들 사이에 sharding. 각 GPU 가 자기 shard 만.
• CPU offload — Adam state 를 CPU 에 두고 step 시점에만 GPU 로. torch.optim.Adam(..., capturable=True) 또는 fused-cpu Adam.
• unified memory — CUDA 의 unified address space. CPU/GPU 가 같은 pointer 를 보지만 OS 가 page migration. 강의에서 Jane 이 “현재 PyTorch 에서 직접 쓰기는 어렵다” 고 명시 — “C++ 단에서 직접 cudaMallocManaged 해야 함.”

강의의 마지막 데모. 같은 모델 (작은 transformer block) 의 한 Adam step 을 네 가지 모드로 측정 — eager, foreach, fused, compile. PyTorch profiler 의 chrome trace 위에서 직접 본다.

이 표가 PyTorch optimizer team 이 한 일의 압축. 같은 알고리즘이 ~80 ms → ~5 ms 로 16배 가속. 산술 자체는 한 글자도 안 바뀜.

PyTorch optimizer 의 4단 사다리와 그 motivation 을 다시 펼치기 전 기억할 사실들.

손에 새기기 — 실습 시퀀스

1. 같은 모델로 4모드 비교 — 작은 transformer block 1 step 의 시간을 eager, foreach, fused, compile 로 각각 측정. 자기 GPU 에서 강의의 16배 가속이 재현되는가.
2. chrome trace 로 launch 수 세기 — 4 모드의 trace 를 export 해서 optimizer step 안의 막대 수를 직접 센다. ~수만 → ~수십 의 차이.
3. 파라미터 수 vs 시간 — 모델 크기를 4단 (small/base/large/xl) 로 변화시키면서 각 모드의 시간 추세. eager 는 거의 linear, fused 는 거의 flat.
4. custom optimizer 에 compile 적용 — 자기만의 optimizer (예: Lion, Tiger) 를 만들고 opt.step = torch.compile(opt.step) 로 wrap. 가속 비율 측정.
5. multi_tensor_apply 의 chunk 직접 보기 — PyTorch source 의 aten/src/ATen/native/cuda/MultiTensorApply.cuh 를 읽어 CHUNK_SIZE 와 metadata struct 의 정의 확인.
6. foreach 의 in-place 검증 — torch._foreach_add_ 호출 전후 tensor 의 data_ptr 가 같은지 확인. 새 alloc 안 일어남.
7. fused 와 numerical 차이 — eager 와 fused 가 정확히 같은 결과를 주는가? torch.allclose(rtol=…). fp32 면 거의 동일, fp16 면 미세한 차이.
8. 한 페이지 plan — 자기 학습 코드의 한 step 에서 “optimizer 가 차지하는 비율” 을 측정하고, 어느 단계 (foreach/fused/compile) 가 가장 큰 효과인지 결정.

tensor list fusion · launch overhead 의 정량 · multi_tensor_apply 패턴이 시리즈 여러 강의의 핵심 어휘.

강의 안에서 흐릿하게 지나간 자리들과, 자기 환경에서 직접 측정해야 손에 박히는 사실들.

• 16배 가속의 정확한 수치 — 강의에서 정확한 수치는 모델/GPU 에 따라 다름. 자기 환경에서 직접 측정 필요. (확인 필요)
• CUDA Graphs 와의 관계 — Q&A 에서 Mark 가 “언제 cudaGraphs 를 쓸 것인가” 를 별도 강의 주제로 언급. fused optimizer 와 graphs 의 결합은 큰 모델에서 핵심.
• fused vs compile 의 정확한 차이 — 같은 결과를 주지만, 어느 한 쪽이 특정 케이스에서 빠른가? compile 의 첫 step 컴파일 비용은? graph break 의 조건? L053 의 영역.
• sparse Adam 의 fused 여부 — 강의 시점에 fused 없음. 현재 (2026) 상태는 별도 확인.
• capturable=True 의 의미 — Adam 의 옵션 중 capturable 이 등장하지만 강의에서 깊이 다루지 않음. CUDA Graph 와의 호환을 위한 변형.
• unified memory 의 PyTorch 지원 — Jane 이 “현재 PyTorch 에서 직접 쓰기 어렵다” 고 명시. 그 이후 (2024–2026) 변화가 있었는지 확인.
• fp16 / bf16 fused Adam 의 numerical 차이 — 산술 순서가 약간 다른 fused 가 fp16 에서 model accuracy 에 영향? 강의에서 깊이 안 다룸. L007 와 L028 에서 일부.