Huấn luyện hiệu quả bộ nhớ

🤗 PEFT giúp việc tinh chỉnh tham số trở nên hiệu quả, nhưng không tự động tối ưu bộ nhớ. Tổng quan này tổng hợp các mẹo để giảm bộ nhớ huấn luyện và liên kết đến các hướng dẫn chi tiết.

Luôn cân nhắc các phương pháp cơ bản như chọn mô hình cơ sở nhỏ hơn, kích thước lô nhỏ hơn hoặc độ dài chuỗi ngắn hơn để giảm mức sử dụng bộ nhớ.

Tổng quan bộ nhớ huấn luyện

Hãy phân tích phân bố bộ nhớ huấn luyện để chúng ta có thể suy luận về các biện pháp đối phó tiềm năng. Chúng ta sẽ sử dụng một mô hình ngôn ngữ lớn được huấn luyện với bộ tối ưu Adam làm ví dụ. Khi thực hiện tinh chỉnh đầy đủ, các thành phần sau sẽ chiếm dụng bộ nhớ:

1. tham số của mô hình cơ sở: mức tiêu thụ bộ nhớ phụ thuộc rất nhiều vào dtype được chọn. Số bit trên mỗi tham số càng ít (16 cho float16), bộ nhớ chiếm dụng càng nhỏ. Một mô hình 1B ở float16 (16 bit/2 byte cho mỗi tham số) sẽ chiếm xấp xỉ 1e9 × 2 byte = 1.863 GiB bộ nhớ.
2. tất cả các tham số của mô hình cơ sở có thể huấn luyện ×3 cho gradient (1×) và trạng thái của bộ tối ưu Adam (2×), do đó cần 5,59GB bộ nhớ
3. bộ nhớ cho các kích hoạt trung gian giữa các lớp, những giá trị này khó dự đoán nhưng chủ yếu phụ thuộc vào dtype tính toán được sử dụng và độ dài chuỗi / kích thước lô.

Việc sử dụng một mô hình cơ sở nhỏ hơn hoặc một kiểu dữ liệu tính toán nhỏ hơn sẽ làm giảm tất cả các thành phần, trong khi việc sử dụng các chuỗi ngắn hơn hoặc các lô nhỏ hơn chủ yếu ảnh hưởng đến bộ nhớ cho gradient và kích hoạt. Việc áp dụng các phương pháp PEFT sẽ làm giảm số lượng tham số có thể huấn luyện và do đó giảm đáng kể cả gradient lẫn trạng thái của bộ tối ưu, tiết kiệm được một lượng lớn bộ nhớ.

Chọn phương pháp phù hợp

Không phải mọi phương pháp PEFT đều được xây dựng như nhau và một số công thức dễ dàng được xây dựng theo cách tiết kiệm bộ nhớ hơn. Nếu bạn đang bị giới hạn về bộ nhớ, việc tham khảo bộ so sánh phương pháp PEFT và lọc theo mức sử dụng bộ nhớ accelerator tối đa là hợp lý. Mức sử dụng bộ nhớ accelerator trung bình có thể khá tương đương giữa các phương pháp, nhưng không phải phương pháp nào cũng mở rộng theo kích thước kích hoạt và độ dài chuỗi một cách đồng đều; một số phương pháp dễ gặp hiện tượng tăng vọt bộ nhớ hơn những phương pháp khác.

Hãy cân nhắc sử dụng các token có thể huấn luyện khi nhắm vào các lớp lớn như đầu mô hình ngôn ngữ hoặc lớp nhúng để tinh chỉnh các token cụ thể.

Lượng tử hóa

Lượng tử hóa là một trong những cách tốt nhất để giảm mức tiêu thụ bộ nhớ của mô hình cơ sở và, tùy thuộc vào phương pháp lượng tử hóa được sử dụng, cũng sẽ giảm bộ nhớ kích hoạt. Vì các phương pháp PEFT chỉ chiếm một phần nhỏ trong tổng số tham số, PEFT mặc định sử dụng độ chính xác cao hơn so với mô hình cơ sở. Điều này cũng có thể có tác dụng giúp các adapter giảm bớt một phần mất chất lượng do các phương pháp lượng tử hóa gây ra. Đọc hướng dẫn lượng tử hóa PEFT.

Biên dịch

Các mô hình chúng ta huấn luyện được cấu thành từ các phép toán như nhân ma trận, cộng và gán, trong đó mỗi phép toán tạo ra một kết quả mới và, do đó, cần chiếm dụng bộ nhớ. Nếu các kết quả trung gian đó không còn cần thiết, chúng ta có thể hợp nhất (fuse) các phép toán này để tiết kiệm bộ nhớ. Đây chỉ là một trong nhiều tối ưu hóa mà torch.compile có thể thực hiện cho bạn, vì vậy hãy xem hướng dẫn PEFT torch.compile.

Gradient Checkpointing

Bạn có thể đánh đổi bộ nhớ với phép tính bằng cách chỉ lưu lại mỗi gradient thứ n giữa các lớp và tính phần còn lại theo thời gian thực. Xem tài liệu gradient checkpointing của Transformers để biết thêm.

Khi không sử dụng Diffusers hoặc Transformers, bạn có thể cần tự triển khai logic gradient checkpointing của riêng mình, tùy thuộc vào framework huấn luyện mà bạn đang dùng.

Mất NLL phân mảnh

Việc sử dụng NLLLoss rất phổ biến khi huấn luyện các mô hình ngôn ngữ (hoặc các tác vụ phân loại). Bạn cấp phát một ma trận có kích thước batch × sequence × vocabulary. Với các chuỗi hoặc từ vựng đặc biệt dài, điều này có thể trở nên tốn kém rất nhanh.

Khi sử dụng TRL, bạn có thể dùng Liger kernel integration hoặc dùng Chunked NLLLoss. Phương án sau sẽ chia chuỗi thành các khối có kích thước 256 để giữ mức tiêu thụ bộ nhớ tối đa ở mức không đổi.

Nếu kích thước chunk mặc định không tối ưu cho cài đặt của bạn, hãy xem PR TRL gốc để biết thêm thông tin về cách điều chỉnh kích thước chunk.