Huấn luyện trên GPU

Chưa cần GPU để kiểm mạng biết học

Overfit một batch nghĩa là cố ý học lại một nhóm rất nhỏ nhiều lần. Với dữ liệu không có nhãn mâu thuẫn và mạng đủ capacity, train loss thường có thể giảm rất sâu. Nếu chưa làm được, đừng vội tăng dataset: có thể loss ghép sai, nhãn lệch mẫu, optimizer thiếu tham số hoặc graph bị cắt. Đây là kiểm sức khỏe pipeline, không phải điểm validation.

Ta dùng bốn góc của một hình vuông, nhãn theo dấu tọa độ thứ nhất. Tắt dropout, không augmentation để mục tiêu ổn định. Mạng nhỏ và seed cố định chạy CPU:

import torch
from torch import nn

torch.set_num_threads(1)
torch.manual_seed(31)
xb = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]])
yb = torch.tensor([0, 0, 1, 1], dtype=torch.long)
model = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))
optimizer = torch.optim.Adam(model.parameters(), lr=.05)
criterion = nn.CrossEntropyLoss()
initial_loss = criterion(model(xb), yb).item()
for step in range(100):
    model.train()
    optimizer.zero_grad(set_to_none=True)
    logits = model(xb)
    assert logits.shape == (4, 2)
    loss = criterion(logits, yb)
    assert torch.isfinite(loss)
    loss.backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all()
               for p in model.parameters())
    optimizer.step()
model.eval()
with torch.no_grad():
    final_logits = model(xb)
    final_loss = criterion(final_logits, yb).item()
    accuracy = (final_logits.argmax(dim=1) == yb).float().mean().item()
print('loss:', round(initial_loss, 6), '->', round(final_loss, 6))
print('one-batch accuracy:', accuracy)

Kết quảloss: 0.776305 -> 5.4e-05
one-batch accuracy: 1.0

Đọc loss đầu/cuối và accuracy trên cùng bốn mẫu đã học. Accuracy 1 ở đây chỉ xác nhận model có thể khớp nhóm nhỏ; nó không chứng minh khả năng dự đoán mẫu chưa thấy. Kiểm grad is not None thích hợp với mạng này vì mọi tham số tham gia forward. Model có nhánh điều kiện hoặc tham số đóng băng cần kiểm theo thiết kế, không bắt mọi tham số đều có gradient.

Tìm lỗi theo điểm xuất hiện đầu tiên

Khi pipeline không học hoặc báo lỗi, ta lần ngược tới chỗ sai đầu tiên thay vì vá chỗ triệu chứng lộ ra:

• Shape: in shape của input, logits và target, gọi tên từng chiều (batch, lớp, đặc trưng) chứ không chỉ đếm số chiều.
• Dtype: target của cross-entropy là chỉ số lớp kiểu long, còn target của BCE là số thực.
• Device: model và mọi tensor đi vào cùng một phép phải nằm cùng nơi.
• NaN: kiểm lần lượt input, activation, logits, loss rồi gradient để tìm phép đầu tiên tạo giá trị không hữu hạn; torch.isfinite phát hiện cả NaN lẫn vô cùng. Một nan ở cuối có thể bắt đầu từ một phép chia cho 0 rất sớm.

Có thể bật torch.autograd.detect_anomaly() trên vài bước kiểm để thư viện hỗ trợ truy dấu backward; nó chậm nên không dùng trong mọi lượt đo.

Nếu loss hữu hạn nhưng gradient nổ, giảm learning rate chưa chắc chữa được phép đạo hàm không ổn định: kiểm log, chia và căn bậc hai tại miền không hợp lệ. torch.where vẫn lấy đạo hàm cả nhánh không được chọn: torch.where(x > 0, x.sqrt(), 0) với $x = -1$ có sqrt(-1) là NaN, và NaN nhân 0 trong backward vẫn là NaN; kẹp đầu vào trước khi lấy căn. Nếu gradient bằng None, kiểm detach/no_grad hoặc tham số chưa được dùng.

Để minh họa cách khoanh vùng mà không gây training nổ, ta kiểm một đầu vào lỗi:

bad_input = torch.tensor([[1., float('nan')], [2., 3.]])
finite_rows = torch.isfinite(bad_input).all(dim=1)
print('finite rows:', finite_rows.tolist())
print('bad row indices:', (~finite_rows).nonzero(as_tuple=True)[0].tolist())

Kết quảfinite rows: [False, True]
bad row indices: [0]

Hàng đầu sai ngay từ input; đổi optimizer không sửa được dữ liệu ấy. Tùy bài, ta cần điều tra nguồn, điền thiếu bằng thống kê train hoặc loại dòng theo quy định dữ liệu; không tùy tiện thay mọi NaN thành 0 rồi coi pipeline đã đúng.

Lỗi device thường hiện ra ở dòng khác với chỗ gây ra nó. Model đã chuyển lên GPU, input cũng đã chuyển, nhưng một tensor khác tạo trong vòng lặp (target, mask, trọng số lớp) vẫn nằm ở CPU; phép đầu tiên chạm tới nó mới báo lỗi.

Bộ nhớ có nhiều phần

Một tham số FP32 chiếm 4 byte. Training còn có gradient, và Adam giữ hai moment cùng thang FP32. Chỉ các phần này đã khoảng $4P+4P+8P=16P$ byte cho $P$ tham số, chưa tính activation, batch input và vùng nhớ tạm.

Đổi ra con số: 1 triệu tham số FP32 là 4 MB. Bảng embedding 100 000 từ, mỗi vector 300 chiều FP32, là 120 MB.

Activation phải giữ lại vì backward cần chúng (bài 3.2); suy luận không có backward nên dùng xong là bỏ. Mỗi ảnh tạo 1 triệu giá trị ở mỗi lớp, 50 lớp, batch 64, FP32 đã là $64 \cdot 50 \cdot 10^6 \cdot 4$ byte $\approx 12.8$ GB.

Gradient checkpointing chỉ lưu đầu ra của vài lớp rồi tính lại phần giữa khi backward, đổi thêm một lượt forward lấy bộ nhớ: 100 lớp chia 10 đoạn thì chỉ giữ cỡ 10 điểm lưu cộng 10 lớp của đoạn đang tính lại, thay vì 100.

parameter_count = sum(p.numel() for p in model.parameters())
weight_bytes = sum(p.numel() * p.element_size() for p in model.parameters())
print('parameters:', parameter_count)
print('weight bytes:', weight_bytes)
print('FP32 Adam parameter-state estimate:', 16 * parameter_count, 'bytes')
B, heads, L = 16, 8, 1024
score_bytes = B * heads * L * L * 4
print('materialized attention scores:', score_bytes // 2**20, 'MiB')

Kết quảparameters: 42
weight bytes: 168
FP32 Adam parameter-state estimate: 672 bytes
materialized attention scores: 512 MiB

Mạng 42 tham số rất nhỏ, nhưng một tensor attention score (bài 3.9) giả định (16,8,1024,1024) FP32 đã là 512 MiB. Dòng đó chỉ tính kích thước, không cấp phát tensor ấy. Bình phương độ dài chính là $O(L^2)$ quen thuộc của thuật toán duyệt mọi cặp.

Kernel attention tối ưu (như FlashAttention) tính theo từng khối và không dựng toàn bộ bảng score, nên 512 MiB là chi phí của cách cài thẳng, không phải của mọi kernel.

Giảm batch thì không phải phần nào cũng giảm theo. Với MLP, weight không đổi khi giảm batch 64 xuống 32; activation dạng (B,h) thường giảm khoảng một nửa. Gradient và optimizer state cũng không giảm vì số tham số giữ nguyên. Vì vậy đỉnh bộ nhớ GPU (VRAM) không giảm đúng một nửa.

Với attention dựng toàn bộ bảng score, giảm độ dài 1024 xuống 512 làm riêng score giảm bốn lần: tensor (16,8,512,512) FP32 chỉ còn $16\cdot8\cdot512^2\cdot4$ byte $= 128$ MiB thay vì 512 MiB, trong khi trọng số giữ nguyên.

Hết bộ nhớ: microbatch và gradient accumulation

Lỗi CUDA out of memory (OOM) nghĩa là không xin được thêm vùng nhớ cho tensor kế tiếp. Cách sửa trực tiếp là giảm phần lớn nhất: batch, độ dài chuỗi, kích thước ảnh hoặc độ rộng mạng. torch.cuda.empty_cache() chỉ trả lại phần cache chưa được tensor nào dùng; nó không giải phóng tensor còn được tham chiếu.

Nguồn rò hay gặp là giữ tensor loss qua các step, như history.append(loss) hay total = total + loss: mỗi loss kéo theo graph của step đó. Nếu vòng validation vừa quên torch.no_grad() vừa giữ loss như vậy, backward không bao giờ chạy nên mỗi graph còn nguyên các activation đã lưu, và bộ nhớ tăng theo từng batch. Metric nên lưu bằng loss.item(). Sau khi sửa OOM, khởi động lại tiến trình để bỏ trạng thái cũ.

Giảm batch làm gradient nhiễu hơn và đổi động học huấn luyện. Muốn giữ batch hiệu dụng (effective batch) cũ, ta chia batch thành $k$ microbatch, gọi backward() cho từng microbatch và chỉ gọi optimizer.step() sau $k$ lần. PyTorch cộng gradient vào .grad qua các lần backward; loss mặc định lại là trung bình trong batch, nên mỗi loss microbatch phải chia cho $k$ thì tổng mới bằng gradient trung bình của cả batch. Thí nghiệm dưới so gradient của một batch 32 mẫu với bốn microbatch 8 mẫu, có và không chia loss:

torch.manual_seed(7)
X_acc = torch.randn(32, 2)
y_acc = (X_acc[:, 0] > 0).long()
net = nn.Sequential(nn.Linear(2, 8), nn.ReLU(), nn.Linear(8, 2))

def gradient_vector(module):
    return torch.cat([p.grad.flatten() for p in module.parameters()])

net.zero_grad(set_to_none=True)
criterion(net(X_acc), y_acc).backward()            # cả batch 32 mẫu một lần
full_grad = gradient_vector(net).clone()

accumulation_steps = 4
for divide_loss in (True, False):
    net.zero_grad(set_to_none=True)
    x_parts = X_acc.chunk(accumulation_steps)
    y_parts = y_acc.chunk(accumulation_steps)
    for x_part, y_part in zip(x_parts, y_parts):
        micro_loss = criterion(net(x_part), y_part)  # mỗi microbatch 8 mẫu
        if divide_loss:
            micro_loss = micro_loss / accumulation_steps
        micro_loss.backward()                         # .grad cộng dồn qua các lần backward
    accumulated = gradient_vector(net)
    ratio = (accumulated.norm() / full_grad.norm()).item()
    print('divide by 4:', divide_loss, '| norm ratio:', round(ratio, 4),
          '| matches full batch:', torch.allclose(accumulated, full_grad, atol=1e-6))

Kết quảdivide by 4: True | norm ratio: 1.0 | matches full batch: True
divide by 4: False | norm ratio: 4.0 | matches full batch: False

Khi chia loss cho 4, gradient cộng dồn trùng gradient của batch 32 tới sai số làm tròn. Không chia, gradient lớn gấp đúng 4 lần: với SGD, điều đó tương đương lặng lẽ nhân learning rate với 4. Phép trùng cần hai điều kiện:

• Các microbatch phải bằng nhau; microbatch cuối thiếu mẫu thì phải chia theo số mẫu thật.
• Mạng không có lớp tính thống kê theo batch: BatchNorm chuẩn hóa theo từng microbatch 8 mẫu, nên accumulation không tái tạo được thống kê của batch 32.

Bộ nhớ activation chỉ còn cỡ một microbatch, còn weight, gradient và trạng thái optimizer giữ nguyên. Mỗi optimizer step tốn $k$ lượt forward/backward nối tiếp, nên accumulation tiết kiệm bộ nhớ chứ không tăng tốc. Số lần cập nhật chia cho $k$: 50 000 mẫu, batch 128 có 390 batch đủ mỗi epoch, 10 epoch với $k = 4$ cho 975 lần step.

Mixed precision là lựa chọn số học

Mixed precision dùng kiểu số ít bit hơn (FP16, 2 byte) cho những phép chịu được và giữ những phép nhạy ở FP32. autocast chọn dtype theo từng phép: nhân ma trận và convolution chạy FP16, còn softmax, log, exp, phép tổng và phần lớn hàm loss giữ FP32. Model không bị ép toàn bộ thành FP16, và bản trọng số mà optimizer cập nhật vẫn là FP32.

Lợi ích là nhiều activation chỉ còn 2 byte thay vì 4, và Tensor Core của GPU chạy phép FP16 nhanh hơn nhiều. Cái giá là FP16 chỉ biểu diễn được một khoảng số hẹp.

Gradient ở nhiều lớp rất nhỏ. FP16 không biểu diễn được số dương nhỏ hơn khoảng $6\cdot10^{-8}$, nên một gradient cỡ $10^{-8}$ bị làm tròn về 0. GradScaler nhân loss với một hệ số lớn (mặc định khởi đầu $2^{16} = 65\,536$) trước backward; theo quy tắc chuỗi, mọi gradient cũng được nhân theo, rồi được chia lại trước khi optimizer cập nhật. Phép thử trên CPU cho thấy cơ chế:

tiny_grad = torch.tensor(1e-8)
scale = 65536.0
print('FP16 of 1e-8:', tiny_grad.half().item())
scaled_half = (tiny_grad * scale).half()
print('FP16 of 1e-8 * 65536:', round(scaled_half.item(), 8))
print('unscaled in FP32:', f'{scaled_half.float().item() / scale:.4e}')

Kết quảFP16 of 1e-8: 0.0
FP16 of 1e-8 * 65536: 0.00065517
unscaled in FP32: 9.9972e-09

Đưa thẳng $10^{-8}$ vào FP16 cho 0: gradient ấy mất hẳn. Nhân 65536 trước thì được $6.55\cdot10^{-4}$, nằm gọn trong khoảng của FP16; chia lại ở FP32 cho $9.9972\cdot10^{-9}$, lệch khoảng 0.03% vì FP16 chỉ giữ được khoảng ba chữ số có nghĩa. Nếu hệ số lớn quá làm gradient tràn thành inf hoặc NaN, scaler.step bỏ qua bước đó và scaler.update() giảm hệ số; sau một chuỗi bước ổn định, hệ số được tăng lại. GradScaler không chữa được nhãn sai hay learning rate quá lớn.

Bản trọng số FP32 cũng cần vì một lý do tương tự: cập nhật thường nhỏ hơn trọng số rất nhiều, và trong FP16 phép cộng $1 + 10^{-4}$ làm tròn về đúng 1, nên cập nhật sẽ mất nếu cộng thẳng vào bản 16 bit.

BF16 có cùng khoảng số mũ với FP32 nên thường không cần scaler, nhưng chỉ chạy nhanh trên GPU từ thế hệ Ampere trở đi; với T4, mixed precision nghĩa là FP16 kèm GradScaler.

Thứ tự một step trong script: autocast cho forward và loss → scaler.scale(loss).backward() → scaler.unscale_(optimizer) → clip nếu cần → scaler.step(optimizer) → scaler.update(). Clip phải đứng sau unscale_: trước đó norm gradient còn lớn gấp hệ số scale, nên ngưỡng clip 1.0 gần như luôn bị kích hoạt và sau khi chia lại, gradient bị thu nhỏ quá mức. Script dùng API torch.amp theo tài liệu PyTorch 2.9; mã cũ hơn viết torch.cuda.amp.

Đo trên máy có GPU

Phần đo tốc độ nằm trong script gpu_probe.py, chạy riêng trên máy có GPU. Script tự tạo một MLP 512→256→10 với AdamW và một batch ngẫu nhiên ngay trên GPU, nên không cần tải dữ liệu hay model. Lệnh đầu chỉ kiểm môi trường: in phiên bản PyTorch, bản CUDA mà PyTorch được build kèm và trạng thái CUDA_READY hoặc CPU_ONLY. Hai lệnh sau chạy benchmark ở FP32 và ở mixed precision:

python content/phase-3/recipes/gpu_probe.py
python content/phase-3/recipes/gpu_probe.py --run --precision fp32
python content/phase-3/recipes/gpu_probe.py --run --precision amp

Benchmark cần GPU CUDA hoạt động và bản PyTorch tương thích. Trên máy chỉ có CPU, script dừng với CUDA_NOT_AVAILABLE; khi hết bộ nhớ, nó dừng với CUDA_OOM và gợi ý giảm --batch. Mỗi lượt chạy 5 step warmup, đồng bộ, đặt lại bộ đếm peak memory, đo 50 step (đổi bằng --warmup, --steps) rồi đồng bộ lần nữa.

Kết quả là một dòng JSON gồm:

• device: tên GPU;
• milliseconds_per_step, samples_per_second (batch × số step chia thời gian);
• peak_allocated_MiB: bộ nhớ tensor dùng lúc cao nhất;
• peak_reserved_MiB: phần PyTorch giữ, thường lớn hơn vì có cache.

So hai dòng FP32 và AMP cùng batch để thấy mixed precision tiết kiệm bao nhiêu thời gian và bộ nhớ trên đúng card đó. Nếu môi trường thi cấp T4, chạy chính phép đo ấy trên T4: số giây không suy ra được từ tên card. Phần đo chỉ gồm loop training trên batch đã ở GPU (includes_data_loading: false); nếu đọc ảnh hay chép CPU→GPU mới là nút thắt, phải đo riêng pipeline đầy đủ.

CUDA chạy bất đồng bộ: CPU xếp lệnh vào hàng đợi rồi đi tiếp ngay, nên đồng hồ CPU đọc lúc ấy chỉ đo việc gửi lệnh dù GPU chưa làm xong. Đồng bộ trước và sau đoạn đo, hoặc dùng CUDA events, mới đo được công việc thực. Không gọi .item() ở mỗi batch trong benchmark nếu không cần, vì lấy số về CPU buộc phải đồng bộ và làm thay đổi chính thứ đang đo.

GPU nhanh nhờ chạy song song hàng nghìn phép trên tensor lớn. Mạng 42 tham số với batch 4 ở đầu bài thì ngược lại: mỗi step chỉ vài trăm phép nhân, còn chi phí khởi chạy kernel, chép dữ liệu và đồng bộ là cố định, nên trên GPU nó có thể chậm hơn trên CPU.

Ước lượng thời gian chạy

Gọi $t$ là median step time đo trên những batch đại diện. Với $N$ mẫu, batch $B$ và $E$ epoch, phần train xấp xỉ

Lấy median sau warmup, vì vài step đầu chậm do khởi tạo kernel và allocator. Chẳng hạn $N = 50\,000$, $B = 128$ cho $\lceil 390.6 \rceil = 391$ step mỗi epoch; nếu đo được $t = 0.12$ giây thì 20 epoch tốn $20 \cdot 391 \cdot 0.12 \approx 938$ giây, khoảng 15.6 phút. Đó mới là phần train: còn validation, đọc dữ liệu, lưu checkpoint và thời gian khởi động.

Batch tăng thông lượng chứ không giảm độ trễ: một batch 32 ảnh mất 0.4 giây là 80 ảnh mỗi giây, gấp 4 lần chạy từng ảnh 0.05 giây, nhưng ảnh nào cũng phải chờ cả batch xong. $t$ phải đo với shape giống dữ liệu thật trên đúng phần cứng được cấp; số đo của mạng rất nhỏ không dùng để suy ra cho model pretrained lớn.