Giảm mức sử dụng bộ nhớ

Quy trình huấn luyện thường có thể được tối ưu hóa để giảm mức tiêu thụ bộ nhớ, và TRL cung cấp một số tính năng tích hợp sẵn để hỗ trợ đạt được điều này.

Dưới đây, chúng tôi phác thảo các kỹ thuật này và khuyến nghị thử nghiệm với các tổ hợp khác nhau để xác định cấu hình nào hoạt động tốt nhất cho thiết lập cụ thể của bạn.

Mỗi phương pháp bao gồm các ví dụ cho các trainer được hỗ trợ. Nếu bạn không chắc một kỹ thuật có tương thích với trainer của bạn hay không, vui lòng xem tài liệu của trainer tương ứng.

Để biết thêm về các chiến lược khác, chẳng hạn như gradient checkpointing, được hỗ trợ trên tất cả các trainer, hãy xem hướng dẫn hiệu suất của transformers.

Cắt ngắn

Chiều dài chuỗi trong tập dữ liệu có thể dao động rất lớn. Khi dữ liệu được nhóm thành các lô (batch), các chuỗi sẽ được đệm (padding) để khớp với chuỗi dài nhất trong lô, điều này có thể gây ra mức sử dụng bộ nhớ cao, ngay cả khi phần lớn các chuỗi tương đối ngắn.

Để giảm mức sử dụng bộ nhớ, việc cắt ngắn các chuỗi xuống độ dài hợp lý là rất quan trọng. Mặc dù các trình huấn luyện của TRL cắt ngắn chuỗi theo mặc định, bạn có thể muốn điều chỉnh độ dài cắt ngắn mặc định để phù hợp hơn với trường hợp sử dụng cụ thể của mình.

from trl import DPOConfig

training_args = DPOConfig(..., max_length=...)

from trl import SFTConfig

training_args = SFTConfig(..., max_length=...)

Cách chọn giá trị max_length?

Nếu max_length quá nhỏ, một phần đáng kể token của bạn sẽ bị loại bỏ và không đóng góp vào quá trình huấn luyện. Nếu quá lớn, mức sử dụng bộ nhớ có thể tăng đột biến, có thể dẫn đến lỗi hết bộ nhớ (OOM). Nếu không có packing hoặc padding-free, một max_length lớn cũng có thể dẫn đến huấn luyện kém hiệu quả, vì nhiều token sẽ là padding.

Để giúp bạn chọn một giá trị phù hợp, chúng tôi cung cấp một tiện ích để trực quan hóa phân phối độ dài chuỗi trong bộ dữ liệu của bạn.

Đóng gói

Kỹ thuật này chỉ khả dụng cho huấn luyện SFT và các cấu hình sử dụng FlashAttention (hoặc các biến thể của nó).

Cắt ngắn có một số nhược điểm:

1. Mất thông tin: Các token quan trọng ở cuối chuỗi có thể bị loại bỏ.
2. Chọn độ dài cắt ngắn: Quá ngắn sẽ mất dữ liệu; quá dài sẽ giảm hiệu quả.

Packing khắc phục các vấn đề này bằng cách nhóm nhiều chuỗi vào cùng một hàng huấn luyện, lấp đầy mỗi hàng đến max_length.

TRL triển khai packing bằng cách sử dụng bin packing Best-Fit Decreasing (BFD), giúp nhóm các chuỗi một cách hiệu quả đồng thời giảm thiểu padding. Khi một chuỗi vượt quá max_length, các chiến lược khác nhau sẽ quyết định cách xử lý các token tràn ra.

TRL hỗ trợ ba chiến lược:

• "bfd" (mặc định): Sử dụng đóng gói Best-Fit Decreasing. Nếu một chuỗi vượt quá max_length, các token vượt mức sẽ bị loại bỏ.

• "bfd_split": Sử dụng đóng gói Best-Fit Decreasing, nhưng các chuỗi dài sẽ được chia thành các khối ≤ max_length trước khi đóng gói. Cách này giữ nguyên tất cả các token và tuân theo phương pháp được đề xuất trong Fewer Truncations Improve Language Modeling.

• "wrapped": Tất cả các token được nối thành một luồng và chia thành các khối có độ dài cố định. Chiến lược này giảm thiểu padding nhưng có thể trộn lẫn các ví dụ không liên quan. Chiến lược này tương ứng với tiền xử lý concatenate-then-split được mô tả trong tài liệu (ví dụ, Fewer Truncations Improve Language Modeling). Nó có nhược điểm là phá vỡ tính liên tục của chuỗi đối với một phần lớn bộ dữ liệu, làm giảm hiệu suất, như đã thảo luận trong Qwen3-Coder-Next Technical Report.

Nếu tất cả các chuỗi đều ngắn hơn max_length, bfd và bfd_split hoạt động giống hệt nhau, vì không cần cắt ngắn hay tách chuỗi.

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    packing=True,
    packing_strategy="bfd",
    max_length=512,
)

PEFT để tinh chỉnh hiệu quả tham số

Các phương pháp Fine-Tuning hiệu quả theo tham số (PEFT) như LoRA là một trong những kỹ thuật hiệu quả nhất để giảm mức sử dụng bộ nhớ trong quá trình huấn luyện. Thay vì huấn luyện tất cả các tham số của mô hình, các phương pháp PEFT chỉ huấn luyện một số lượng nhỏ các tham số adapter, từ đó giảm đáng kể yêu cầu về bộ nhớ và cho phép fine-tuning các mô hình lớn hơn trên phần cứng hạn chế.

Để biết chi tiết đầy đủ về việc sử dụng PEFT với TRL, bao gồm các phương pháp adapter khác nhau, các tùy chọn lượng tử hóa và các cấu hình nâng cao, hãy xem Tích hợp PEFT.

Để sử dụng PEFT nhằm giảm mức sử dụng bộ nhớ:

from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer

dataset = load_dataset("trl-lib/Capybara", split="train")

peft_config = LoraConfig()

trainer = SFTTrainer(
    model="Qwen/Qwen2.5-0.5B",
    train_dataset=dataset,
    peft_config=peft_config,
)

PEFT có thể kết hợp với các kỹ thuật giảm bộ nhớ khác như lượng tử hóa (4-bit hoặc 8-bit) để tiết kiệm bộ nhớ nhiều hơn nữa. Xem Tích hợp PEFT để biết các ví dụ về lượng tử hóa.

Không có padding

Batching không padding là một phương pháp thay thế để giảm mức sử dụng bộ nhớ. Trong phương pháp này, một batch được lấy mẫu trước, sau đó được làm phẳng thành một chuỗi duy nhất, tránh việc padding. Khác với packing, vốn có thể dẫn đến các chuỗi không hoàn chỉnh do kết hợp các phần từ những mẫu khác nhau, batching không padding đảm bảo rằng tất cả các chuỗi vẫn hoàn chỉnh và nguyên vẹn.

from trl import DPOConfig

training_args = DPOConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

from trl import SFTConfig

training_args = SFTConfig(..., padding_free=True, model_init_kwargs={"attn_implementation": "kernels-community/flash-attn2"})

Tải xuống kích hoạt

Activation offloading là một kỹ thuật hiệu quả về bộ nhớ, giúp giảm mức sử dụng VRAM của GPU bằng cách tạm thời di chuyển các tensor kích hoạt sang RAM của CPU trong quá trình truyền thẳng (forward pass) và chỉ đưa chúng trở lại khi cần thiết cho quá trình truyền ngược (backward pass). Điều này làm giảm đáng kể mức sử dụng bộ nhớ đỉnh, với cái giá là thời gian huấn luyện tăng lên một chút.

Để bật tính năng offloading activation trong cấu hình huấn luyện SFT của bạn:

from trl import SFTConfig

training_args = SFTConfig(..., activation_offloading=True)

Bên trong, activation offloading triển khai saved_tensors_hooks của PyTorch để chặn các activation trong quá trình forward pass. Nó quản lý thông minh các tensor cần offload dựa trên kích thước và ngữ cảnh, tránh việc offload các tensor output vì sẽ không hiệu quả. Để tối ưu hiệu năng, thông qua một cờ (mặc định là true), nó có thể sử dụng CUDA streams để chồng chéo tính toán với các thao tác chuyển dữ liệu CPU-GPU.

Đệm các chuỗi lên bội

Kỹ thuật này hiện được hỗ trợ cho các trình huấn luyện SFT và Reward.

Khi được bật, tùy chọn này đảm bảo rằng tất cả các chuỗi đều được lấp đầy đến bội số của giá trị được chỉ định.
Điều này có thể cải thiện hiệu quả tính toán trên một số phần cứng bằng cách căn chỉnh độ dài chuỗi với các ranh giới thân thiện với bộ nhớ.

from trl import SFTConfig

training_args = SFTConfig(..., pad_to_multiple_of=2048)

from trl import RewardConfig

training_args = RewardConfig(..., pad_to_multiple_of=2048)

Trình cấp phát bộ nhớ đệm PyTorch

Trong các phiên chạy dài, đặc biệt là RL trực tuyến (GRPO, RLOO, Online DPO), bộ nhớ GPU có thể bị phân mảnh. Việc đặt expandable_segments:True cho phép bộ cấp phát bộ nhớ có bộ nhớ đệm của PyTorch mở rộng các phân đoạn hiện có, từ đó giảm khoảng cách giữa bộ nhớ đã cấp phát và bộ nhớ đã đặt chỗ:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export PYTORCH_ALLOC_CONF=expandable_segments:True  # canonical name since PyTorch 2.10

Với chế độ ngủ của vLLM (vllm_enable_sleep_mode=True), vLLM chỉ đọc PYTORCH_CUDA_ALLOC_CONF để khắc phục cài đặt này, vì vậy đừng chỉ đặt PYTORCH_ALLOC_CONF.

Tắt việc thu thập mô hình cho quá trình tạo trong các phương pháp trực tuyến

Khi sử dụng DeepSpeed ZeRO-3, trọng số mô hình được chia nhỏ trên nhiều GPU. Các phương pháp trực tuyến bao gồm việc tạo ra các phần hoàn thành từ mô hình như một phần của quá trình huấn luyện. Trong bước này, trọng số mô hình được tập hợp tạm thời trên một GPU duy nhất để tạo ra nội dung. Đối với các mô hình rất lớn, việc tập hợp này có thể dẫn đến lỗi OOM, như được mô tả trong vấn đề này: #2250.

Nếu bạn gặp sự cố này, bạn có thể tắt việc thu thập trọng số mô hình cho quá trình tạo bằng cách đặt tham số sau:

from trl import GRPOConfig

training_args = GRPOConfig(..., ds3_gather_for_generation=False)

from trl.experimental.online_dpo import OnlineDPOConfig

training_args = OnlineDPOConfig(..., ds3_gather_for_generation=False)

from trl import RLOOConfig

training_args = RLOOConfig(..., ds3_gather_for_generation=False)

Điều chỉnh này ngăn việc thu thập trọng số mô hình, tránh các lỗi OOM, nhưng có thể dẫn đến tốc độ sinh chậm hơn.

Chế độ ngủ của vLLM

Khi sử dụng vLLM làm backend sinh cho các phương pháp huấn luyện trực tuyến, bạn có thể bật sleep mode để chuyển các tham số và cache của vLLM sang RAM CPU trong bước tối ưu hóa, sau đó tải chúng trở lại VRAM GPU khi cần đồng bộ trọng số và sinh.

from trl import GRPOConfig

training_args = GRPOConfig(..., vllm_enable_sleep_mode=True)

from trl import RLOOConfig

training_args = RLOOConfig(..., vllm_enable_sleep_mode=True)

Việc offload trọng số và cache của vLLM giúp giữ mức sử dụng bộ nhớ GPU ở mức thấp, điều này đặc biệt hữu ích khi huấn luyện các mô hình lớn hoặc khi sử dụng tài nguyên GPU hạn chế. Tuy nhiên, việc đánh thức engine vLLM từ chế độ ngủ sẽ gây ra một số độ trễ trong quá trình truyền dữ liệu giữa host và thiết bị, có thể ảnh hưởng nhẹ đến tốc độ huấn luyện.

Gradient checkpointing

Gradient checkpointing đánh đổi tính toán lấy bộ nhớ bằng cách không lưu trữ tất cả các kích hoạt trung gian trong quá trình truyền thuận, thay vào đó là tính lại chúng trong quá trình truyền ngược.

from trl import SFTConfig

training_args = SFTConfig(..., gradient_checkpointing=True)

Gradient checkpointing được bật mặc định trong tất cả các trainer để tối ưu hóa việc sử dụng bộ nhớ. Bạn có thể tắt tính năng này bằng cách đặt gradient_checkpointing=False nếu cần.

Kiểm điểm kích hoạt có chọn lọc

Với SFTTrainer, bạn có thể lưu trữ đầu ra của attention trong quá trình forward thay vì tính lại trong quá trình backward. Điều này khôi phục phần lớn hiệu suất bị giảm do checkpointing ở ngữ cảnh dài, với chi phí thêm một tensor có kích thước bằng hidden-state cho mỗi lớp. Tính năng này buộc sử dụng checkpointing không tái nhập (non-reentrant).

from trl import SFTConfig

training_args = SFTConfig(
    ...,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"selective": True},
)

Để biết thêm các kỹ thuật tối ưu hóa bộ nhớ, hãy xem Transformers Performance Guide.