Quantized Training

각 자리의 양자화 가능성 —

• weight — BF16 까지는 안전. FP8 weight 는 mma 와 직결. delicate.
• gradient — BF16 표준. FP8 grad 는 delayed scaling 이 필요. 발산의 주요 원인.
• optimizer state (m, v) — 가장 많이 차지. 가장 양자화 안전한 자리. 8-bit Adam 이 표준.
• master weights — FP32 한 벌. 정확도가 직결되니 항상 fp32. 단 stochastic rounding 으로 유지하면 BF16 으로 줄일 수 있다.
• activations — BF16 대부분. recompute 로 줄이는 게 우선, 양자화는 까다로움.

문제 상황을 구체적으로 —

• BF16 master weight: 0.5000 (mantissa 7 비트)
• gradient × lr: 0.0001 (BF16 의 LSB 보다 작음)
• round-to-nearest 결과: 0.5000 (gradient 가 사라짐)
• 여러 step 누적: 0.0001 × 1000 = 0.1 인데 — 매 step 에서 사라지면 영원히 안 더해진다

stochastic rounding 은 — 0.5001 을 받았을 때, 0.5000 으로 99% 확률, 0.5001 의 다음 representable 값으로 1% 확률로 round. 기댓값이 정확하다 (unbiased). 1000 step 누적시 평균적으로 0.1 가 더해진다.

# stochastic rounding — Triton 안에서
@triton.jit
def sr_cast_fp32_to_bf16(x_fp32):
    # 1. fp32 의 lower 16 bit 를 random bit 로 가산
    rand = tl.rand(seed, offset)         # [0, 1) uniform
    rand_bits = (rand * 65536).to(tl.uint32)

    # 2. fp32 비트 직접 manipulation
    bits = x_fp32.to(tl.uint32, bitcast=True)
    bits = bits + rand_bits              # lower 16b 에 random 가산

    # 3. truncate to bf16 (= upper 16 bits)
    return (bits >> 16).to(tl.bfloat16, bitcast=True)

핵심 — round-to-nearest 보다 비용이 거의 같다. random 한 번, addition 한 번. 그런데 학습 양자화의 정확도 손실이 한 자리수 줄어든다.

# 8-bit Adam 의 한 step — Triton (요약)
@triton.jit
def adam_8bit_step(p_ptr, g_ptr,
                    m8_ptr, v8_ptr,
                    sm_ptr, sv_ptr,
                    lr, b1, b2, eps,
                    BLOCK: tl.constexpr):
    pid    = tl.program_id(0)
    offs   = pid * BLOCK + tl.arange(0, BLOCK)

    # 1. shared memory 안에서 dequant
    m8     = tl.load(m8_ptr + offs)             # int8
    v8     = tl.load(v8_ptr + offs)
    sm     = tl.load(sm_ptr + pid)              # scale
    sv     = tl.load(sv_ptr + pid)
    m      = m8.to(tl.float32) * sm
    v      = v8.to(tl.float32) * sv

    # 2. Adam update — fp32 안에서
    g      = tl.load(g_ptr + offs).to(tl.float32)
    m      = b1 * m + (1 - b1) * g
    v      = b2 * v + (1 - b2) * g * g
    p      = tl.load(p_ptr + offs)
    p_new  = p - lr * m / (tl.sqrt(v) + eps)

    # 3. 다시 quantize — 새 scale 산출
    abs_m  = tl.max(tl.abs(m))
    sm_new = abs_m / 127.
    m8_new = (m / sm_new).to(tl.int8)

    tl.store(p_ptr + offs, p_new)
    tl.store(m8_ptr + offs, m8_new)
    tl.store(sm_ptr + pid, sm_new)
    # v 도 같은 패턴 ...

이 커널의 핵심 디자인 결정 —

• dequantized state 를 global memory 에 절대 안 쓴다. 모든 fp32 값이 register/SRAM 에서만 살아있음. 강의 Q&A 의 핵심 질문.
• block-wise scaling. 한 BLOCK 단위(예: 2048 element) 의 amax 로 한 scale. tensor-wise 보다 정확, element-wise 보다 cheap.
• Adam 의 b1, b2 가 1 에 가까우니 m, v 가 매우 안정. 양자화 noise 의 영향이 작다.

강의에서 짚은 미묘한 측정 디테일들 —

• FP8 의 win 은 H100 부터. A100 은 FP8 mma 가 없다. mma 자체는 H100 / H200 / Blackwell.
• 실효 win 은 모델 + batch + seq 에 따라 달라짐. memory-bound 인 자리(작은 batch, KV cache dominant)는 FP8 의 throughput win 이 작다.
• delayed scaling 의 history window. 작으면 (16) 빠르고 stable, 크면 (128) 더 robust 하지만 cold start 가 느림.
• convergence 가 깨지는 자리 — scale 이 sudden 하게 변하는 self-attention 의 logit. 그 자리만 BF16 fallback 하는 패턴.