Đào tạo trên TPU

TPU (Tensor Processing Unit) là một loại phần cứng được thiết kế đặc biệt để huấn luyện mô hình một cách hiệu quả. Accelerate hỗ trợ huấn luyện trên TPU, nhưng có một vài điều bạn cần lưu ý, cụ thể là biên dịch đồ thị. Bài hướng dẫn này thảo luận ngắn gọn về việc biên dịch, và để biết thêm chi tiết, hãy xem hướng dẫn Training on TPUs with Accelerate.

Biên dịch

TPU tạo ra một đồ thị chứa tất cả các phép toán trong bước huấn luyện, chẳng hạn như bước truyền thuận, bước truyền ngược và bước tối ưu hóa. Đây là lý do tại sao bước huấn luyện đầu tiên luôn mất một khoảng thời gian nhất định, vì việc xây dựng và biên dịch đồ thị này tốn thời gian. Tuy nhiên, một khi quá trình biên dịch hoàn tất, nó sẽ được lưu vào bộ nhớ đệm và tất cả các bước tiếp theo sẽ nhanh hơn đáng kể.

Chìa khóa là tránh việc biên dịch lại mã của bạn, nếu không quá trình huấn luyện sẽ cực kỳ chậm. Điều này có nghĩa là tất cả các thao tác của bạn phải hoàn toàn giống nhau:

• tất cả các tensor trong các batch của bạn phải có cùng độ dài (ví dụ: không có dynamic padding cho các tác vụ NLP)
• mã của bạn phải là tĩnh (ví dụ: không có các lớp với vòng lặp for có độ dài khác nhau tùy thuộc vào đầu vào như LSTM)

Chia sẻ trọng số

Một thiết kế mô hình ngôn ngữ phổ biến là gán chung trọng số giữa lớp nhúng và lớp softmax. Tuy nhiên, khi chuyển mô hình sang TPU (dù bạn tự thực hiện hay truyền nó vào phương thức prepare()), việc gán chung trọng số sẽ bị phá vỡ và bạn sẽ cần gán lại các trọng số.

Để thêm hành vi đặc biệt (như weight tying) trong script của bạn cho TPU, hãy đặt distributed_type thành DistributedType.TPU trước. Sau đó, bạn có thể sử dụng phương thức tie_weights để liên kết các trọng số.

if accelerator.distributed_type == DistributedType.TPU:
    model.tie_weights()