Gradient checkpointing

Quá trình truyền thẳng thường lưu đệm tất cả các kích hoạt trung gian để quá trình truyền ngược tái sử dụng. Tuy nhiên, kích hoạt tỷ lệ thuận với kích thước lô và độ dài chuỗi. Gradient checkpointing chỉ lưu một số kích hoạt nhất định và loại bỏ phần còn lại. Điều này buộc quá trình truyền ngược phải tính toán lại một số kích hoạt theo thời gian thực khi chúng được cần đến.

Normal training:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save ALL activations)
  Backward:  ←uses cached activations everywhere

Gradient checkpointing:
  Forward:   [L1]→[L2]→[L3]→[L4]   (save only at checkpoints, discard the rest)
  Backward:  ←reaches L2, recomputes L2→L3 from scratch, uses it, discards it

Quá trình huấn luyện thường chậm hơn vì bước truyền ngược phải tính toán lại các kích hoạt đã bị loại bỏ, nhưng việc checkpoint giúp giảm bộ nhớ kích hoạt.

Đặt gradient_checkpointing=True để bật.

Sử dụng cùng với gradient accumulation để giảm thêm mức sử dụng bộ nhớ.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
)

Lưu checkpoint một phần

Gradient checkpointing đầy đủ sẽ tính toán lại mọi lớp có thể checkpoint. Nếu lượt chạy của bạn còn một chút dung lượng bộ nhớ dự phòng, hãy checkpoint ít lớp hơn để đánh đổi một phần tiết kiệm bộ nhớ lấy tốc độ.

Truyền every_n_layers vào gradient_checkpointing_enable() để chọn khoảng cách checkpoint.

every_n_layers=2

Forward:   input  -> [L1] -> [L2] -> [L3] -> [L4] -> [L5] -> [L6]
                      CP     keep     CP     keep     CP     keep

Backward:  output <- [L6] <- [L5] <- [L4] <- [L3] <- [L2] <- [L1]
                     keep   rerun    keep   rerun    keep   rerun

Với every_n_layers=2, lớp đầu tiên và mỗi lớp thứ hai sau đó sẽ được checkpoint. Các lớp được checkpoint sẽ loại bỏ activation của chúng trong quá trình forward và tính lại chúng trong quá trình backward, trong khi các lớp khác giữ activation của chúng trong bộ nhớ.

model.gradient_checkpointing_enable(every_n_layers=2)

Mặc định, every_n_layers=1, sẽ checkpoint mọi layer. Các giá trị lớn hơn sẽ checkpoint layer đầu tiên và sau đó là mỗi n layer tiếp theo, để các activation của các layer khác trong bộ nhớ. Ví dụ, every_n_layers=2 sẽ checkpoint các layer 1, 3, 5, và cứ thế. Chỉ các module kế thừa từ GradientCheckpointingLayer mới được tính. Các module khác hỗ trợ gradient checkpointing vẫn được bật.

Để sử dụng kiểm tra điểm gradient một phần với Trainer, hãy đặt every_n_layers trong gradient_checkpointing_kwargs.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"every_n_layers": 4},
)

Chuyển các kích hoạt đã lưu

Gradient checkpointing giữ một activation cho mỗi layer được checkpoint trên GPU. Với độ dài sequence lớn, điều này có thể tiêu tốn đáng kể bộ nhớ (xấp xỉ layers x sequence x hidden x bytes_per_element). Đặt offload để giữ các activation đó trong bộ nhớ host được ghim (pinned host memory) thay thế, với chi phí là một thao tác copy từ device sang host trong quá trình forward và một thao tác copy từ host sang device trong quá trình backward.

from transformers import TrainingArguments

args = TrainingArguments(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"offload": True},
)

Cả hai bản sao đều chạy trên luồng tính toán, do đó cách này đánh đổi một bước chậm hơn để lấy bộ nhớ. Hãy dùng nó khi một lần chạy không vừa bộ nhớ theo cách khác, chứ không phải để tăng tốc.

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 Đà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 nhằm giảm thêm bộ nhớ và tăng tốc độ đào tạo.
• Xem hướng dẫn Kernels để tìm hiểu cách tăng tốc huấn luyện bằng các kernel hợp nhất tùy chỉnh.