GPU Optimization

Một GPU hiện đại có hai lớp hệ thống mà bạn phải đồng thời tối ưu: compute fabric (CUDA cores, Tensor Cores, RT cores) và memory hierarchy (HBM → L2 → SMEM → register). Phần lớn performance problem trong ML workload nằm ở lớp thứ hai. Lý do đơn giản: Tensor Core có thể tạo ra 312 TFLOPS, nhưng HBM chỉ cấp được 1555 GB/s. Nếu bạn cần 4 byte (một FP32) cho mỗi FLOP, tức là cần 312 × 10¹² × 4 = 1248 TB/s băng thông — gấp 800 lần những gì HBM có thể cấp. Do đó Tensor Core phải tái sử dụng dữ liệu trong SRAM; nếu không, nó sẽ idle.

Để đo được tất cả những điều trên, bạn dùng profiler. Hai công cụ chính của NVIDIA là:

• Nsight Systems (nsys): xem timeline ở cấp độ process/stream/kernel. Phù hợp để tìm idle gap, data loader bottle‑ neck, kernel launch overhead, sync không cần thiết.
• Nsight Compute (ncu): deep dive vào một kernel cụ thể. Cho bạn occupancy, warp stall reasons, L1/L2 hit rate, Tensor Core utilization.
• PyTorch Profiler: wrapper tiện lợi, export sang Chrome Tracing hoặc TensorBoard. Đủ tốt cho 80% use case training.
• Nvidia-smi / DCGM:thống kê mức độ sử dụng GPU, nhưng chỉ số "GPU util" gây hiểu lầm — nó chỉ nói có kernel chạy hay không, không phản ánh occupancy thật.

Mixed precisionlà đòn bẩy đầu tiên nên kéo. Khái niệm: lưu weight ở FP32 (để cộng gradient không mất độ chính xác), nhưng chạy forward/backward ở FP16 hoặc BF16. Tensor Core xử lý FP16 nhanh gấp 8× FP32 CUDA, và memory giảm một nửa. BF16 (Brain Floating Point) có cùng kích thước FP16 nhưng dynamic range rộng như FP32 → ít gặp overflow hơn, không cần loss scaling. Đây là mặc định cho các model > 1B param hiện nay.

import torch
from torch.cuda.amp import autocast, GradScaler

model = MyTransformer().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # xử lý FP16 overflow tự động

for batch in dataloader:
    optimizer.zero_grad()

    # Autocast: tự động chuyển sang FP16 cho matmul / conv,
    # giữ FP32 cho layernorm / softmax để ổn định.
    with autocast(dtype=torch.float16):
        output = model(batch["input"])
        loss = criterion(output, batch["label"])

    # Loss scaling: nhân loss với số lớn để gradient không underflow ở FP16.
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# Với BF16, không cần GradScaler vì dynamic range đủ rộng:
#     with autocast(dtype=torch.bfloat16):
#         ...
#     loss.backward()
#     optimizer.step()

Gradient checkpointing(còn gọi là activation checkpointing) là đòn bẩy thứ hai. Bình thường, forward pass lưu mọi activation để backward pass có thể dùng khi tính gradient. Với model vài chục tầng và sequence dài, activation chiếm memory nhiều hơn cả weight. Checkpointing chỉ lưu activation ở một số "checkpoint layer" và tính lại phần giữa khi cần. Giá phải trả: forward pass được chạy hai lần (một lần chính thức, một lần recompute) → ~ 30% compute overhead đổi lấy 50-70% giảm memory.

Tensor Coreslà đơn vị phần cứng chuyên biệt cho matmul 4×4 FP16 (hoặc TF32, BF16, FP8 tùy generation). Chúng tăng throughput matmul 8-16× so với CUDA core tiêu chuẩn. Để Tensor Core "bắt" được kernel, shape tensor phải là bội của 8 (cho FP16) hoặc 16 (cho FP8). Đây là lý do nhiều codebase dùng hidden size 4096, 6144, 8192 — đều chia hết cho 8. Dùng 4095 sẽ vô tình đẩy kernel về CUDA core.

import torch

# Shape phải là multiple của 8 cho FP16 Tensor Core.
# 4096 = 8 * 512  → OK
# 4095              → PyTorch sẽ pad hoặc rơi về CUDA core
a = torch.randn(128, 4096, dtype=torch.float16, device="cuda")
b = torch.randn(4096, 4096, dtype=torch.float16, device="cuda")

# Bật TF32 cho matmul FP32 — tăng tốc ~2x, sai số ~1e-3
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# Với Ampere/Hopper, dùng BF16 thay FP16 nếu gặp NaN:
a_bf = a.to(torch.bfloat16)
b_bf = b.to(torch.bfloat16)
out = a_bf @ b_bf  # tự động chạy trên Tensor Core

# Profile để chắc chắn:
#   with torch.profiler.profile(activities=[ProfilerActivity.CUDA]) as p:
#       out = a_bf @ b_bf
#   print(p.key_averages().table(sort_by="cuda_time_total"))
# Tìm kernel 'tensorop' hoặc 'hmma' → Tensor Core đang chạy.

Kernel fusion là bước tối ưu tiếp theo khi mixed precision và checkpointing đã bật. Ý tưởng: thay vì gọi nhiều kernel nhỏ (mỗi kernel đọc/ghi HBM riêng), fuse chúng thành một kernel lớn giữ dữ liệu trong SRAM. Ví dụ nổi tiếng nhất là FlashAttention: attention naive viết ma trận N×N ra HBM rồi đọc lại cho softmax (tốn 4N² bytes memory). FlashAttention chia attention thành tile, tính softmax streaming trong SRAM, chỉ ghi output cuối cùng ra HBM. Kết quả: 2-4× nhanh hơn và dùng ít memory hơn ~20×.

Distributed training mở rộng sang nhiều GPU khi một GPU không đủ. Có ba chiến lược chính:

• Data parallelism (DP/DDP): mỗi GPU có full model, xử lý batch khác nhau, sync gradient qua all-reduce. Đơn giản, scale tốt tới vài chục GPU nếu model đủ nhỏ.
• Tensor parallelism (TP): chia trọng số của mỗi layer ra nhiều GPU. Giao tiếp dày đặc (mỗi layer có 2 all-reduce), yêu cầu NVLink băng thông cao.
• Pipeline parallelism (PP): chia layer thành stage trên các GPU khác nhau, pipeline micro-batch qua các stage để lấp đầy bubble.
• ZeRO / FSDP: shard optimizer states, gradient, và weight qua các GPU. Giữ được mô hình lớn mà vẫn cho cảm giác như data parallel.

Tối ưu GPU là một vòng lặp: profile → tìm bottleneck → áp dụng đúng tool → profile lại. Bạn không thể đoán đúng chỉ bằng đọc paper; mỗi workload có pattern riêng. Công cụ có sẵn (Nsight, PyTorch Profiler, DCGM) đã đủ cho 95% trường hợp. Phần khó là đọc output và biết tìm dấu hiệu gì.

Các chủ đề bạn có thể muốn đọc tiếp: tối ưu inference với KV-cache, model serving ở scale, và tối ưu chi phí GPU cloud.