Mixed Precision (FP16/BF16) — Nửa precision, gấp đôi tốc độ

Dùng 16-bit thay vì 32-bit để train model lớn nhanh gấp 2 lần, tiết kiệm nửa bộ nhớ — bí kíp giúp training GPT-4 không tốn 200 triệu USD

Training GPT-4 được ước tính tốn hơn 100 triệu USD tiền điện và GPU. Một phần lớn chi phí đến từ việc phải thuê cluster khổng lồ chạy liên tục hàng tháng trời. Mixed precision là lý do tại sao điều này khả thi về mặt kinh tế — nó cho phép chúng ta dùng nửa bộ nhớ, tính toán nhanh gấp đôi trên Tensor Cores, mà vẫn giữ được độ chính xác cần thiết để model học sâu. Nếu không có kỹ thuật này, bạn sẽ cần gấp đôi số lượng GPU để train Llama 3 hoặc GPT-4, biến nó thành bài toán tài chính không thể.

Vấn đề

Trước khi có mixed precision, training neural network dùng FP32 (32-bit floating point) — mỗi parameter chiếm 4 bytes. Nghe có vẻ nhỏ nhặt, nhưng hãy làm phép tính với quy mô LLM hiện đại:

Một model 175B parameters (như GPT-3) cần 700GB chỉ để lưu weights. Optimizer Adam lại cần thêm 2 bản sao nữa (momentum và variance) cho mỗi parameter — thêm 1.4TB. Gradients thêm 700GB nữa. Tổng cộng hơn 2.8TB bộ nhớ chỉ cho model states. GPU mạnh nhất hiện nay (H100) chỉ có 80GB VRAM.

Không có mixed precision, bạn không thể train model này trên bất kỳ GPU nào tồn tại. Bạn sẽ buộc phải chia nhỏ model quá mức (tensor parallelism) hoặc dùng gradient checkpointing với overhead tính toán khổng lồ.

Vấn đề thứ hai là bottleneck băng thông nhớ (memory bandwidth). Mỗi lần forward/backward pass, GPU phải load tất cả weights từ VRAM vào SRAM. Dữ liệu càng lớn, thời gian chờ càng lâu — dù GPU có hàng trăm TFLOPS compute capacity, nó vẫn ngồi chơi vì không đọc kịp dữ liệu.

Ý tưởng cốt lõi

Insight cốt lõi: Neural networks không cần 32-bit precision cho mọi phép tính. Giống như bạn không cần cân phân tích để đong gạo nấu cơm — độ chính xác "vừa đủ" là được. Mixed precision training dùng FP16 (16-bit) hoặc BF16 (Brain Floating Point 16-bit) cho phần lớn tính toán, chỉ giữ FP32 cho những chỗ thực sự cần.

Đây là "bí kíp" có 3 thành phần:

1. Master Weights (FP32) + Compute (FP16/BF16) Giữ một bản sao weights chính xác cao (FP32) trên CPU hoặc VRAM. Trong mỗi iteration, copy xuống dạng 16-bit để tính forward và backward (nhanh, nhẹ), rồi dùng kết quả gradient để cập nhật bản FP32. Như vậy accumulation qua nhiều step vẫn chính xác, nhưng tính toán nặng nhọc lại dùng 16-bit.

2. Dynamic Loss Scaling FP16 có vấn đề "underflow" — gradient quá nhỏ (< 2⁻²⁴) sẽ bị làm tròn thành 0, model ngừng học. Giải pháp: nhân loss bằng hằng số lớn (thường là 1024 hoặc 65536) trước khi backward, rồi chia gradient ngược lại cùng hệ số đó. Điều này đẩy gradient ra khỏi vùng underflow mà không làm sai lệch toán học.

3. BF16 thay vì FP16 BF16 là "biến thể" 16-bit của Google với 8-bit exponent (như FP32) và 7-bit mantissa. Nó có range rộng như FP32 nhưng precision thấp hơn. Với LLM, range quan trọng hơn precision (tránh overflow trong attention scores), nên BF16 an toàn hơn FP16 mà vẫn nhanh như nhau.

That's it. Chúng ta vừa cắt đôi bộ nhớ cần dùng, tận dụng được Tensor Cores chuyên dụng của NVIDIA (chỉ hoạt động với 16-bit), mà model vẫn converge như thường.

Misconception buster: Nhiều người sợ "16-bit sẽ làm mất thông tin, model kém chính xác". Thực nghiệm cho thấy mixed precision training cho kết quả gần như tương đương FP32 purist — vì gradient descent vốn đã chịu được nhiễu (stochasticity), và master weights FP32 đảm bảo không có accumulation error lớn theo thời gian.

Tại sao nó hoạt động

Tại sao việc đơn giản "cắt đôi bit-width" lại cho speedup gần 2x? Vì nó giải quyết đúng bottleneck:

Memory Bandwidth: A100 có băng thông nhớ ~2 TB/s nhưng compute 312 TFLOPS. Nếu mỗi op cần 1 byte dữ liệu, bạn cần 312 TB/s bandwidth để GPU không nghỉ — điều impossible. Dùng 16-bit giảm lượng data cần fetch đi một nửa, hiệu quả giống như tăng gấp đôi bandwidth.

Tensor Cores: Phần cứng NVIDIA Ampere/HOpper thiết kế riêng cho FP16/BF16 matrix multiply (GEMM), cho throughput gấp 2-8x so với FP32 trên cùng chip.

Loss Scaling Chi Tiết:

# Pseudocode của dynamic loss scaling
loss = model.forward(input) * current_scale  # scale up
scaled_grads = backward(loss)
if any(grad.is_inf_or_nan() for grad in scaled_grads):
    current_scale /= 2  # overflow, giảm scale
    skip_step()
else:
    grads = [grad / current_scale for grad in scaled_grads]  # unscale
    optimizer.step(grads)
    current_scale *= 2  # tăng scale dần để tận dụng dynamic range

Cơ chế này tự động điều chỉnh để gradient luôn nằm trong "sweet spot" của FP16 mà không cần manual tuning.

BF16 vs FP16 Trade-off:

FP16: 1 bit sign, 5 bits exponent, 10 bits mantissa. Precision cao nhưng range nhỏ (±65,504 max). Dễ bị overflow khi attention scores lớn.
BF16: 1 bit sign, 8 bits exponent, 7 bits mantissa. Range giống FP32 (±3.4×10³⁸) nhưng precision thấp. An toàn hơn cho training LLM, nhưng đôi khi cần gradient clipping cho stability.

Ý nghĩa thực tế

Impact thực tế trên quy mô lớn:

GPT-4, Llama 3, PaLM: Tất cả đều dùng BF16 làm định dạng mặc định. Không có mixed precision, chi phí training GPT-4 sẽ tăng từ ~100M USD lên ~200-250M USD — biến nó thành bài toán không khả thi cho OpenAI lúc đó.
Memory Footprint: Giảm 50% model states, cho phép tăng batch size gấp đôi hoặc train model gấp đôi kích thước trên cùng phần cứng.
Throughput: Training throughput tăng 2-3x trên A100/H100 khi chuyển từ FP32 sang mixed precision.

Limitations:

FP16 Underflow: Một số layer như LayerNorm hoặc activation functions nhạy cảm với precision, thường phải giữ ở FP32 (hybrid precision).
BF16 Precision: Một số bài toán như training GANs hoặc reinforcement learning cần precision cao hơn, BF16 đôi khi không đủ.
Hardware Lock-in: Lợi ích lớn nhất đến từ Tensor Cores của NVIDIA; trên CPU hoặc GPU cũ hơn (không có mixed precision hardware support), speedup không rõ rệt.

Who uses it: Mọi framework hiện đại — PyTorch AMP (Automatic Mixed Precision), JAX, TensorFlow, DeepSpeed — đều bật mixed precision mặc định. NVIDIA Megatron-LM và Zero Optimizer kết hợp chặt chẽ với mixed precision để train model nghìn tỷ parameters.

Đào sâu hơn

Paper gốc:

"Mixed Precision Training" (Micikevicius et al., 2018, ICLR) — bài báo kinh điển từ NVIDIA giới thiệu dynamic loss scaling.
"BFloat16: The secret to high performance on Cloud TPUs" (Google, 2019) — lý do BF16 trở thành chuẩn cho TPU và sau đó là GPU.

Bài liên quan TroiSinh:

Cùng cụm (training-at-scale):

Data Parallelism — Chia batch cho nhiều GPU, đơn giản nhất: kết hợp với mixed precision để tận dụng tối đa VRAM.
Gradient Checkpointing — Đánh đổi compute lấy memory: khi mixed precision vẫn chưa đủ, checkpointing giúp bạn train model lớn hơn nữa.
Scaling Laws (Chinchilla) — Quy luật giữa model size, data và compute: mixed precision giúp bạn đạt compute-optimal point với chi phí thấp hơn.

Đọc tiếp (Level 1 - training-efficiency):

Zero Optimizer — Shard optimizer state, train trillion-param models: giải quyết phần còn lại của bài toán memory (optimizer states) sau khi đã dùng mixed precision cho weights.
Gradient Accumulation — Giả lập batch lớn khi GPU nhỏ: kết hợp với mixed precision để train hiệu quả trên consumer hardware.