Kỹ Thuật Tối Ưu GPU cho Deep Learning

Mar 25, 2026

Deep Dive

Ngan AI Engineer & HPC

Hướng dẫn thực hành về tối ưu bộ nhớ GPU, kernel fusion và mixed-precision training cho các workload deep learning. Học cách khai thác tối đa từng FLOP trên phần cứng của bạn.

Table of Contents

Giới Thiệu

Tối ưu GPU là yếu tố then chốt để huấn luyện các mô hình deep learning quy mô lớn một cách hiệu quả. Nếu không tinh chỉnh cẩn thận, bạn có thể dễ dàng bỏ phí 60–80% thông lượng lý thuyết của GPU — điều này trực tiếp dẫn đến thời gian huấn luyện dài hơn, chi phí cloud cao hơn, và nỗi lo âu hiện sinh.

Bài viết này bao gồm ba lĩnh vực chính: quản lý bộ nhớ, kernel fusion, và mixed-precision training. Mỗi phần đi kèm code thực tế bạn có thể áp dụng ngay vào dự án của mình.

Chiến Lược Quản Lý Bộ Nhớ

Bộ nhớ GPU là tài nguyên quý giá và có hạn. Hiểu cách quản lý nó là bước đầu tiên để huấn luyện hiệu quả.

Gradient Checkpointing

Gradient checkpointing đánh đổi tính toán lấy bộ nhớ bằng cách tính lại các activation trung gian trong lượt backward thay vì lưu trữ tất cả. Kỹ thuật này có thể giảm bộ nhớ activation lên đến 10× với chi phí tính toán thêm khoảng 30%.

import torch
from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(torch.nn.Module):
    """Bọc một module với gradient checkpointing."""

    def __init__(self, layers: torch.nn.Sequential):
        super().__init__()
        self.layers = layers

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # use_reentrant=False được khuyến nghị cho PyTorch >= 2.0
        return checkpoint(self._forward, x, use_reentrant=False)

    def _forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)

Khi áp dụng checkpointing cho transformer lớn, chỉ bọc các residual block tốn kém — các lớp embedding và projection cuối thường không đáng checkpointing.

Cấu Hình Memory Pool

Cấp phát trước memory pool giảm phân mảnh và overhead cấp phát:

PyTorch: torch.cuda.memory.set_per_process_memory_fraction(0.8)
JAX: XLA_PYTHON_CLIENT_MEM_FRACTION=0.8
TensorFlow: tf.config.experimental.set_memory_growth(gpu, True)

Đo Lường Mức Sử Dụng Bộ Nhớ

Trước khi tối ưu, hãy đo lường. PyTorch đi kèm bộ profiler bộ nhớ tích hợp:

import torch

def profile_memory(model, sample_input):
    torch.cuda.reset_peak_memory_stats()
    with torch.profiler.profile(
        activities=[torch.profiler.ProfilerActivity.CUDA],
        profile_memory=True,
        record_shapes=True,
    ) as prof:
        output = model(sample_input)
        output.sum().backward()

    print(prof.key_averages().table(
        sort_by="self_cuda_memory_usage", row_limit=15
    ))
    peak_mb = torch.cuda.max_memory_allocated() / 1e6
    print(f"Đỉnh đã cấp phát: {peak_mb:.1f} MB")

profile_memory(model, torch.randn(8, 3, 224, 224, device='cuda'))

Kernel Fusion

Kernel fusion kết hợp nhiều thao tác GPU vào một kernel launch duy nhất, giảm overhead băng thông bộ nhớ. Mỗi kernel GPU riêng biệt yêu cầu một vòng lặp đọc-ghi qua global memory; việc fusion giữ dữ liệu trong on-chip cache.

Tại Sao Fusion Quan Trọng

Xét chuỗi thao tác đơn giản: LayerNorm → Dropout → ReLU. Không có fusion:

LayerNorm đọc input → ghi output ra DRAM
Dropout đọc output LayerNorm từ DRAM → ghi kết quả ra DRAM
ReLU đọc output Dropout từ DRAM → ghi kết quả cuối ra DRAM

Đó là 6 vòng lặp DRAM cho ba thao tác đơn giản. Với fusion, chỉ cần 2 (một lần đọc, một lần ghi).

Lợi Ích Của Fusion

Giảm overhead kernel launch
Tính cục bộ bộ nhớ tốt hơn — giá trị trung gian được giữ trong registers / L1 cache
Giảm tiêu thụ băng thông bộ nhớ GPU
Cải thiện cường độ tính toán (FLOPs per byte)

Dùng `torch.compile` Để Tự Động Fusion

Compiler PyTorch 2.x có thể tự động fusion nhiều pattern:

import torch

# Eager mode — các kernel launch riêng biệt
def naive_fused_op(x, weight, bias):
    return torch.nn.functional.layer_norm(x, x.shape[-1:]) @ weight + bias

# torch.compile kích hoạt TorchInductor, fusion + sinh CUDA tối ưu
optimized_op = torch.compile(naive_fused_op, mode="max-autotune")

x = torch.randn(512, 768, device="cuda")
weight = torch.randn(768, 768, device="cuda")
bias = torch.zeros(768, device="cuda")

# Warmup (lần gọi đầu tiên kích hoạt compilation)
for _ in range(3):
    _ = optimized_op(x, weight, bias)

Kernel CUDA Tùy Chỉnh Với Triton

Để kiểm soát tối đa, OpenAI Triton cho phép viết kernel GPU fused bằng Python:

import triton
import triton.language as tl

@triton.jit
def fused_bias_relu_kernel(
    x_ptr, bias_ptr, output_ptr,
    n_elements,
    BLOCK_SIZE: tl.constexpr,
):
    pid = tl.program_id(axis=0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements

    x = tl.load(x_ptr + offsets, mask=mask)
    bias = tl.load(bias_ptr + offsets, mask=mask)
    output = tl.maximum(x + bias, 0.0)  # ReLU(x + bias), fused
    tl.store(output_ptr + offsets, output, mask=mask)

Kernel đơn này cộng vector bias và áp dụng ReLU trong một lần duyệt bộ nhớ — không cần cấp phát tensor trung gian.

Mixed-Precision Training

Mixed-precision sử dụng FP16 (hoặc BF16) cho forward/backward pass trong khi duy trì trọng số master ở FP32. Kết quả: thông lượng tăng gấp ~2×, bộ nhớ giảm một nửa, với độ chính xác gần như không thay đổi.

AMP với PyTorch

API torch.amp xử lý phần lớn sự phức tạp:

import torch
from torch.amp import GradScaler, autocast

model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # scale loss để tránh underflow FP16

for batch in dataloader:
    inputs, targets = batch
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # Scale loss, backward, unscale, clip, step
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

BF16 vs FP16

Định dạng	Bits số mũ	Bits mantissa	Phạm vi động	Ghi chú
FP32	8	23	±3,4×10³⁸	Baseline
FP16	5	10	±65.504	Cần loss scaling
BF16	8	7	±3,4×10³⁸	Không cần scaling

BF16 được ưu tiên trên GPU Ampere (A100) trở lên — cùng phạm vi động với FP32, không cần loss scaling:

with autocast(device_type='cuda', dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

# BF16 không cần GradScaler — backward trực tiếp
loss.backward()
optimizer.step()

Kết Hợp Tất Cả Lại

Vòng lặp huấn luyện production kết hợp cả ba kỹ thuật:

import torch
from torch.amp import GradScaler, autocast
from torch.utils.checkpoint import checkpoint_sequential

def train_epoch(model, dataloader, optimizer, criterion):
    model.train()
    scaler = GradScaler()
    total_loss = 0.0

    for step, (inputs, targets) in enumerate(dataloader):
        inputs = inputs.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)
        optimizer.zero_grad(set_to_none=True)  # nhanh hơn zero_grad()

        with autocast(device_type='cuda', dtype=torch.bfloat16):
            # Gradient checkpointing trên các encoder layer tốn kém
            features = checkpoint_sequential(model.encoder, segments=4, input=inputs)
            outputs = model.head(features)
            loss = criterion(outputs, targets)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()

        total_loss += loss.item()

    return total_loss / len(dataloader)

Các chi tiết đáng chú ý:

non_blocking=True trên .cuda() chồng chéo truyền host→device với tính toán
zero_grad(set_to_none=True) tránh zeroing gradient in-place (tiết kiệm một CUDA kernel)
checkpoint_sequential chia encoder thành 4 segment để checkpointing

Tổng Kết

Tối ưu GPU hiệu quả không phải là một mẹo duy nhất — đó là một kỷ luật. Ba trụ cột được trình bày:

Quản lý bộ nhớ — profile trước, sau đó áp dụng gradient checkpointing và memory pooling khi cần
Kernel fusion — để torch.compile xử lý các pattern phổ biến; dùng Triton cho hot path tùy chỉnh
Mixed-precision — BF16 trên phần cứng hiện đại là lợi thế 2× gần như miễn phí; FP16 với GradScaler hoạt động mọi nơi

Bắt đầu với profiling (bạn không thể tối ưu những gì bạn không đo được), áp dụng mixed-precision trước (ROI cao nhất), sau đó giải quyết bộ nhớ, và cuối cùng tinh chỉnh kernel khi các thắng lợi dễ đã hết.