Mixed precision FP16 BF16 và gradient checkpointing

Trong systems engineering có một câu hỏi cổ điển: bạn sẽ chi 10K cho RAM thêm, hay tốn 2 tuần optimize code để giảm memory? Câu trả lời phụ thuộc vào tỷ lệ cost: nếu hardware rẻ hơn dev time, mua hardware. Nếu ngược lại, optimize.

Trong LLM training, tỷ lệ luôn nghiêng về optimize. Tại sao? Vì hardware không scale tuyến tính. Một A100 80GB giá $15K, nhưng một GPU 160GB không tồn tại. Khi model > 80GB, không có giải pháp “mua thêm RAM cho 1 GPU”. Phải dùng nhiều GPU, hoặc optimize memory trên GPU hiện có.

Bài này nói về hai kỹ thuật quan trọng nhất để cắt memory trong training: mixed precision (FP16/BF16) và gradient checkpointing. Áp dụng cả hai, một model 7B vốn cần 80GB có thể train được trên 24GB VRAM. Đây là khác biệt giữa “phải thuê A100 80GB $1.5/h” và “dùng RTX 3090 ở nhà”.

Mental model: memory ăn từ đâu

Trong training, memory chia làm 4 phần:

Total Memory
    |
    |-- Model weights (params)
    |-- Optimizer state (Adam: 2x params)
    |-- Gradients (1x params)
    |-- Activations (depends on batch size, seq_len, layers)

Với Llama-3-8B FP32:

• Weights: 8B * 4 bytes = 32GB
• Optimizer (AdamW): 8B * 8 bytes (2x FP32) = 64GB
• Gradients: 8B * 4 bytes = 32GB
• Activations: ~20-40GB tùy seq_len, batch size

Tổng: ~150GB. Một A100 80GB không đủ. Hai A100 cũng vừa khít.

Mixed precision cắt weights, gradients, activations xuống một nửa (FP32, FP16, hoặc BF16). Optimizer state vẫn giữ FP32 (gọi là master weights).

Sau mixed precision:

• Weights FP16: 16GB
• Optimizer FP32: 64GB (vẫn vậy)
• Gradients FP16: 16GB
• Activations FP16: 10-20GB

Tổng: ~110GB. Vẫn không đủ 1 A100 80GB.

Gradient checkpointing cắt activations xuống 30-40%. Trade-off: thêm 30% compute (recompute activation lúc backward).

Sau cả hai:

• Weights FP16: 16GB
• Optimizer FP32: 64GB
• Gradients FP16: 16GB
• Activations FP16 + checkpointing: 4-6GB

Tổng: ~100GB. Vẫn còn ZeRO/FSDP cắt optimizer state qua nhiều GPU (bài 17).

Phần 1: FP32, FP16, BF16, sự khác biệt

Floating point trong CPU/GPU là cách lưu số thực dùng số bit cố định.

FP32 (single precision): 32 bit. 1 bit sign, 8 bit exponent, 23 bit mantissa.

sign | exponent (8 bit) | mantissa (23 bit)
 1   |     8            |       23

Range: ~1e-38 đến ~3.4e38. Precision: 7 chữ số thập phân.

FP16 (half precision): 16 bit. 1 bit sign, 5 bit exponent, 10 bit mantissa.

sign | exponent (5 bit) | mantissa (10 bit)
 1   |     5            |     10

Range: ~6e-5 đến 65504. Precision: 3-4 chữ số thập phân.

BF16 (brain float 16): 16 bit. 1 bit sign, 8 bit exponent, 7 bit mantissa.

sign | exponent (8 bit) | mantissa (7 bit)
 1   |     8            |    7

Range: ~1e-38 đến ~3.4e38 (giống FP32). Precision: 2-3 chữ số thập phân.

So sánh visual:

Format	Total bits	Range	Precision
FP32	32	1e-38 ~ 3.4e38	7 digits
FP16	16	6e-5 ~ 65504	3-4 digits
BF16	16	1e-38 ~ 3.4e38	2-3 digits
FP8 (E4M3)	8	0.002 ~ 448	1-2 digits
INT8	8	-128 ~ 127	integer

Đặc tính quan trọng: BF16 có cùng range với FP32 nhưng ít precision. FP16 có range hẹp hơn FP32 rất nhiều.

Phần 2: Tại sao FP16 không đủ, và BF16 cứu

Đầu 2018, mixed precision đầu tiên dùng FP16. Mau chóng phát hiện vấn đề:

Overflow. Gradient ban đầu rất lớn (tới 1e5), FP16 max ở 65504. Gradient overflow thành infinity, training NaN.

Underflow. Gradient cuối training nhỏ (1e-7), FP16 min ở 6e-5. Gradient rounded về 0, model không learn nữa.

Giải pháp NVIDIA đưa ra: loss scaling. Nhân loss với một hằng số S (ví dụ 1024) trước backward, làm gradient lớn lên S lần để fit vào FP16 range. Lúc optimizer step, chia gradient cho S trở lại.

Code:

scaler = torch.cuda.amp.GradScaler()

for batch in dataloader:
    with torch.cuda.amp.autocast(dtype=torch.float16):
        output = model(batch)
        loss = loss_fn(output, target)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

GradScaler tự động tăng/giảm S theo gradient overflow detection. Khá phức tạp, đôi khi vẫn fail.

Năm 2019, Google publish BF16 (đã tồn tại từ trước trên TPU). BF16 giải quyết overflow/underflow vì range giống FP32. Precision thấp hơn FP16 nhưng đủ cho neural network.

Từ 2020 trở đi, A100 và H100 đều support BF16 native. Mọi LLM training hiện đại dùng BF16, không dùng FP16:

for batch in dataloader:
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        output = model(batch)
        loss = loss_fn(output, target)

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

Không cần GradScaler. Đơn giản hơn nhiều, ít fail hơn nhiều.

Khi nào dùng	FP16	BF16
GPU support	V100, T4, RTX 20/30 series	A100, H100, RTX 30/40 series (newer)
Cần loss scaling	Có	Không
Overflow risk	Cao	Rất thấp
Inference	OK	OK
Training	Tránh nếu có thể	Default modern

Phần 3: Mixed precision flow đầy đủ

“Mixed” precision nghĩa là dùng nhiều precision khác nhau trong cùng một training step:

1. Forward:
   - Weights: FP16/BF16
   - Activations: FP16/BF16
   - Output (loss): FP32

2. Backward:
   - Gradient compute: FP16/BF16
   - Gradient accumulate (master): FP32

3. Optimizer:
   - Master weights: FP32
   - Optimizer state (m, v): FP32
   - Update: FP32 -> cast xuống FP16/BF16 cho forward step kế tiếp

Lưu hai version weights: FP32 (master) và FP16/BF16 (compute). Forward/backward dùng compute version cho nhanh. Optimizer step dùng master version để giữ precision.

Memory:

• Master weights: N * 4 bytes
• Compute weights: N * 2 bytes
• Gradients: N * 2 bytes
• Optimizer state: N * 8 bytes (Adam: m, v đều FP32)

Total: 16 bytes / param.

Llama-3-8B: 8B * 16 bytes = 128GB. Vẫn cần multi-GPU. Nhưng nếu FP32 thuần thì sẽ là 8B * 16 = 128GB weight+optim+grad cộng 32GB activation, total 160GB+.

Mixed precision tiết kiệm khoảng 30% memory total. Đáng kể.

Phần 4: Gradient checkpointing

Activations là phần memory cuối cần optimize. Khi forward pass đi qua layer 1, output activation của layer 1 được lưu lại để dùng cho backward. Layer 2, 3, … N đều phải lưu activation.

Memory activation = batch_size * seq_len * hidden_dim * num_layers * 2 (bytes for BF16).

Với GPT-3 175B (96 layer, hidden 12288, batch_size=2048, seq_len=2048):

2048 * 2048 * 12288 * 96 * 2 = ~9.9 TB

Không có GPU nào fit được activation này. Đây là lý do gradient checkpointing.

Ý tưởng: thay vì lưu activation của mọi layer, chỉ lưu activation của một số layer “checkpoint” (ví dụ mỗi 4 layer). Khi backward đến layer giữa hai checkpoint, recompute activation bằng cách forward lại từ checkpoint gần nhất.

Bình thường:
  forward:  L1 -> L2 -> L3 -> L4 -> ... -> Ln (lưu tất cả)
  backward: cần activation tất cả -> OK

Checkpointing:
  forward:  L1 -> [L2 -> L3 -> L4] -> L5 -> [L6 -> L7 -> L8] -> ...
            lưu activation L1, L5, L9, ...
            không lưu L2, L3, L4, L6, L7, L8

  backward đến L7:
    1. Lấy L5 activation từ memory
    2. Forward L5 -> L6 -> L7 (recompute)
    3. Dùng activation L7 vừa recompute để backward

Trade-off:

• Memory activation giảm khoảng 1 / sqrt(N) cho checkpoint mỗi sqrt(N) layer
• Compute tăng ~30-40% (vì forward 1.3x)

PyTorch implement:

from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        x = self.attn(x)
        x = self.ffn(x)
        return x

class TransformerModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.blocks = nn.ModuleList([TransformerBlock() for _ in range(24)])

    def forward(self, x):
        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False)
        return x

Wrap mỗi block bằng checkpoint(). PyTorch tự động:

1. Forward không lưu activation trong block
2. Backward tự forward lại block để có activation, rồi compute gradient

Với HuggingFace Transformers, bật bằng:

model.gradient_checkpointing_enable()

Phần 5: Đo lường thực tế

Train một transformer 350M params, batch_size=8, seq_len=1024 trên RTX 3090 24GB.

Setting	Memory peak	Step time	Tokens/sec
FP32, no checkpoint	OOM	-	-
BF16, no checkpoint	18.2 GB	0.85s	9650
BF16, checkpoint	11.4 GB	1.12s	7300

BF16 cắt memory 50%. Checkpointing cắt thêm 37% nhưng chậm 32%. Tradeoff hợp lý nếu OOM với non-checkpoint setting.

Train cùng model trên A100 40GB, batch_size=32:

Setting	Memory peak	Step time
BF16 no checkpoint	31.5 GB	1.4s
BF16 checkpoint	19.8 GB	1.85s

Pattern lặp lại: cắt memory ~37%, tăng time ~32%.

Quyết định dùng checkpoint hay không phụ thuộc: memory hay throughput quan trọng hơn?

• Memory bound (OOM nếu không checkpoint): bật.
• Compute bound (memory dư): tắt, nhanh hơn.

Một nguyên tắc: nếu GPU memory utilization > 90% và OOM dễ xảy ra, bật checkpoint. Nếu < 70%, tắt.

Pitfall: ZeRO không thay thế mixed precision

Có một lần một dev hỏi: “tôi dùng DeepSpeed ZeRO-3 rồi, có cần BF16 không?”. Câu trả lời: có, vẫn cần.

ZeRO chia optimizer state qua nhiều GPU, nhưng compute weights và activations vẫn được lưu local trên mỗi GPU. Mixed precision cắt cả 4 thành phần memory (weights, optim state, gradients, activations). Hai kỹ thuật bổ sung nhau, không thay thế.

Trong production, đúng setup là:

1. Mixed precision BF16 (bật ở mọi training)
2. Gradient checkpointing (bật nếu activation tốn lớn)
3. ZeRO / FSDP (bật nếu model > VRAM 1 GPU)

Bỏ bất kỳ cái nào, memory tăng đáng kể.

Cheatsheet

Tính năng	Code PyTorch	Cắt memory
Mixed BF16	with torch.cuda.amp.autocast(dtype=torch.bfloat16):	~30% total
Mixed FP16 + scaler	scaler = torch.cuda.amp.GradScaler()	~30% total
Gradient checkpoint	torch.utils.checkpoint.checkpoint(block, x)	~37% activation
HF auto checkpoint	model.gradient_checkpointing_enable()	~37% activation
AdamW 8-bit	bitsandbytes.optim.AdamW8bit(...)	~50% optimizer

Memory formula (mixed precision)	Bytes / param
Master weight FP32	4
Compute weight FP16/BF16	2
Gradient FP16/BF16	2
Adam m FP32	4
Adam v FP32	4
Total per param	16 bytes

Quick estimate: model N params cần 16N bytes cho weights + optim + gradients, cộng activation phụ thuộc batch size.

GPU	VRAM	Max model size FP32	Max model size BF16	Với checkpoint + ZeRO-3
RTX 3090	24GB	0.8B	1.5B	7B
RTX 4090	24GB	0.8B	1.5B	7B
A100 40GB	40GB	1.5B	2.5B	13B
A100 80GB	80GB	3B	5B	30B
H100 80GB	80GB	3B	5B	70B (FSDP)

Lời kết

Mixed precision và gradient checkpointing là hai kỹ thuật mọi LLM trainer phải có trong toolbox. Bật cả hai gần như miễn phí (1 dòng code, vài % throughput), nhưng cắt memory 50-70%. Khác biệt giữa “không train được” và “train được” trên hardware tiêu dùng.

Hands-on song song:

1. Mở Colab free tier, T4 16GB VRAM. Lấy nanoGPT, train với FP32 và đo memory. Sau đó bật BF16 (nếu T4 support, T4 chỉ có FP16, dùng FP16 + scaler). Đo lại memory và throughput.
2. Bật gradient_checkpointing_enable() trên một HuggingFace model nhỏ (gpt2-medium 355M params). Đo memory trước và sau. Verify thấy giảm activation memory rõ rệt.
3. Đọc paper “Training Deep Nets with Sublinear Memory Cost” (Chen 2016), paper gốc của gradient checkpointing. Phần thuật toán không khó, chỉ cần hiểu graph traversal.
4. Nếu muốn ngoài 24GB tiêu dùng, RunPod có A100 40GB từ $0.79/giờ, Vast.ai có 3090 từ $0.20/giờ. Một experiment train 1B model với BF16 + checkpoint mất 3-5 giờ, chi phí $1-4. Đủ để hands-on.

Bài 17 sẽ vào distributed training: DP, DDP, FSDP, pipeline parallel. Khi model vẫn lớn quá sau khi optimize hết memory, phải chia model qua nhiều GPU. Đây là phần khó nhất của LLM training, nhưng cũng là phần cần biết nếu muốn train model > 7B params.