Gradient accumulation

Các batch lớn tạo ra các kích hoạt lớn làm cạn kiệt bộ nhớ GPU. Gradient accumulation cho phép bạn huấn luyện với kích thước batch hiệu dụng lớn hơn bằng cách phân bổ việc tính toán gradient qua nhiều mini-batch.

Gradient được tích lũy qua n mini-batch trước khi optimizer cập nhật trọng số. Ví dụ, với kích thước batch trên mỗi thiết bị là 8 và 4 bước tích lũy, kích thước batch hiệu dụng là 32.

Step 1: mini-batch 1 → forward → backward → grads = G₁
Step 2: mini-batch 2 → forward → backward → grads = G₁ + G₂
Step 3: mini-batch 3 → forward → backward → grads = G₁ + G₂ + G₃
Step 4: mini-batch 4 → forward → backward → grads = G₁ + G₂ + G₃ + G₄
        → optimizer.step()  ← same update as if batch_size × 4
        → zero_grad()

Chỉ sử dụng gradient accumulation khi batch lớn hơn không vừa với bộ nhớ. Nó không cải thiện throughput so với việc huấn luyện bằng một batch lớn thực sự.

Tích lũy gradient trong gradient_accumulation_steps trên per_device_train_batch_size.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
)

Định dạng mất mát

Đối với hàm loss tùy chỉnh, hãy bao gồm num_items_in_batch để Trainer chia loss cho số lượng mục tiêu dự đoán trên tất cả các mini-batch. Điều này chuẩn hóa theo token thay vì theo số bước cố định với gradient_accumulation_steps. Nếu không, Trainer sẽ chia loss cho gradient_accumulation_steps.

import torch.nn.functional as F

def compute_loss(outputs, labels, num_items_in_batch=None):
    logits = outputs["logits"]
    loss = F.cross_entropy(logits, labels, reduction="sum")
    return loss / num_items_in_batch

Đối với các mô hình LM nhân quả, num_items_in_batch đếm các nhãn đã dịch. Hàm mất mát dịch nhãn sao cho dự đoán tại vị trí i nhắm vào token tại vị trí i + 1, khiến vị trí 0 của mọi chuỗi không có mục tiêu. Trainer loại trừ các vị trí đó và đếm trên labels[..., 1:], do đó mẫu số khớp với số lượng mục tiêu dự đoán mà hàm mất mát sử dụng. Khi một bộ sắp xếp dữ liệu cung cấp trực tiếp shift_labels, chẳng hạn như một bộ sắp xếp không có padding, Trainer sẽ đếm trên tensor đó thay thế. Các loại mất mát khác, như LM có mặt nạ và phân loại, đếm toàn bộ tensor nhãn.

Nếu bạn tạo lớp con của Trainer và ghi đè compute_loss(), hãy đặt loss_is_scaled_for_ga để chỉ ra cách loss của bạn được chuẩn hóa: False cho trung bình trên mỗi batch, khi đó Trainer sẽ chia cho gradient_accumulation_steps, hoặc True nếu bạn đã chia cho num_items_in_batch.

class MyTrainer(Trainer):
    loss_is_scaled_for_ga = True

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        ...
        return loss_sum / num_items_in_batch

Các bước tiếp theo

• Đọc tài liệu GPU memory usage để hiểu điều gì đang gây ra mức sử dụng bộ nhớ trên GPU trong quá trình huấn luyện.
• Xem hướng dẫn Gradient checkpointing để tìm hiểu cách giảm bộ nhớ kích hoạt bằng cách tính lại các kích hoạt thay vì lưu chúng trong bộ nhớ đệm.
• Xem hướng dẫn Đào tạo độ chính xác hỗn hợp để tìm hiểu cách sử dụng các kiểu dữ liệu độ chính xác thấp hơn để giảm bộ nhớ và tăng tốc độ đào tạo.
• Đọc bài đăng trên blog Gradient Accumulation Fix để tìm hiểu cách tính toán tích lũy gradient.