TRL documentation
Tùy chỉnh huấn luyện
Tùy chỉnh huấn luyện
TRL được thiết kế với tính mô-đun hóa để người dùng có thể tùy chỉnh vòng lặp huấn luyện một cách hiệu quả theo nhu cầu của mình. Dưới đây là các ví dụ về cách bạn có thể áp dụng và kiểm tra các kỹ thuật khác nhau.
Mặc dù các ví dụ này sử dụng DPOTrainer, các phương pháp tùy chỉnh này áp dụng cho hầu hết (nếu không muốn nói là tất cả) các trainer trong TRL.
Sử dụng các bộ tối ưu hóa và trình lập lịch khác nhau
Mặc định, DPOTrainer tạo một optimizer torch.optim.AdamW. Bạn có thể tạo và định nghĩa một optimizer khác và truyền nó vào DPOTrainer như sau:
from datasets import load_dataset
from torch import optim
from transformers import AutoModelForCausalLM
from trl import DPOTrainer
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
optimizer = optim.SGD(model.parameters(), lr=1e-6)
trainer = DPOTrainer(
model=model,
train_dataset=dataset,
optimizers=(optimizer, None),
)
trainer.train()Thêm trình lập lịch tốc độ học
Bạn cũng có thể thêm các trình lập lịch tốc độ học bằng cách truyền cả optimizer và scheduler:
from torch import optim
optimizer = optim.AdamW(model.parameters(), lr=1e-6)
lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)
trainer = DPOTrainer(..., optimizers=(optimizer, lr_scheduler))Các mô hình tham chiếu 8-bit
Vì trl hỗ trợ tất cả các đối số từ khóa khi tải mô hình từ transformers bằng cách sử dụng from_pretrained, bạn cũng có thể tận dụng load_in_8bit từ transformers để tinh chỉnh hiệu quả hơn về bộ nhớ.
Đọc thêm về việc nạp mô hình 8-bit trong transformers Nạp ở 8bit hoặc 4bit.
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
ref_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", quantization_config=quantization_config)
trainer = DPOTrainer(..., ref_model=ref_model)Thêm callback tùy chỉnh
Bạn có thể tùy chỉnh vòng lặp huấn luyện bằng cách thêm các callback để ghi log, giám sát hoặc dừng sớm. Callback cho phép bạn thực thi mã tùy chỉnh tại các thời điểm cụ thể trong quá trình huấn luyện.
from transformers import TrainerCallback
class CustomLoggingCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if logs is not None:
print(f"Step {state.global_step}: {logs}")
trainer = DPOTrainer(..., callbacks=[CustomLoggingCallback()])Thêm các chỉ số đánh giá tùy chỉnh
Bạn có thể định nghĩa các chỉ số đánh giá tùy chỉnh để theo dõi trong quá trình huấn luyện. Điều này hữu ích để giám sát hiệu suất của mô hình trên các tác vụ cụ thể.
def compute_metrics(eval_preds):
logits, labels = eval_preds
# Add your metric computation here
return {"custom_metric": 0.0}
training_args = DPOConfig(..., eval_strategy="steps", eval_steps=100)
trainer = DPOTrainer(..., eval_dataset=eval_dataset, compute_metrics=compute_metrics)Sử dụng huấn luyện độ chính xác hỗn hợp
Huấn luyện độ chính xác hỗn hợp có thể tăng tốc đáng kể quá trình huấn luyện và giảm mức sử dụng bộ nhớ. Bạn có thể bật tính năng này bằng cách đặt bf16=True hoặc fp16=True trong cấu hình huấn luyện.
# Use bfloat16 precision (recommended for modern GPUs)
training_args = DPOConfig(..., bf16=True)Lưu ý: Sử dụng bf16=True cho các GPU Ampere (A100, RTX 30xx) hoặc mới hơn, và fp16=True cho các GPU đời cũ.
Sử dụng tích lũy gradient
Khi huấn luyện với bộ nhớ GPU hạn chế, tích lũy gradient cho phép bạn mô phỏng kích thước lô lớn hơn bằng cách tích lũy gradient qua nhiều bước trước khi cập nhật trọng số.
# Simulate a batch size of 32 with per_device_train_batch_size=4 and gradient_accumulation_steps=8
training_args = DPOConfig(
...,
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
)Thay đổi mục tiêu huấn luyện
Việc kế thừa một trainer là cách đơn giản nhất để tùy chỉnh quá trình huấn luyện. Loss được định nghĩa trong phương thức compute_loss, vì vậy để huấn luyện với một mục tiêu khác, hãy kế thừa trainer và ghi đè phương thức này. Việc chuẩn bị dữ liệu, sinh dữ liệu, ghi log và lưu checkpoint được kế thừa, nên lớp con chỉ chứa những gì thực sự thay đổi.
import torch
from trl import DPOTrainer
from trl.trainer.utils import selective_log_softmax
class MyDPOTrainer(DPOTrainer):
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
# In this example: the hinge loss from https://huggingface.co/papers/2309.06657
input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
shift_labels, shift_completion_mask = input_ids[:, 1:], inputs["completion_mask"][:, 1:]
logits = model(input_ids=input_ids, attention_mask=attention_mask).logits
with torch.no_grad():
ref_logits = self.ref_model(input_ids=input_ids, attention_mask=attention_mask).logits
# Sum the log-probs over the completion tokens, for the policy and for the reference model
logps = (selective_log_softmax(logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
ref_logps = (selective_log_softmax(ref_logits[:, :-1], shift_labels) * shift_completion_mask).sum(dim=1)
chosen_logps, rejected_logps = logps.chunk(2, dim=0) # batch is [chosen, rejected]
ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)
delta = (chosen_logps - ref_chosen_logps) - (rejected_logps - ref_rejected_logps)
return torch.relu(1 - self.beta * delta).mean()Thêm các tham số vào cấu hình
Khai báo các tham số mà hàm loss của bạn cần bằng cách tạo lớp con của config:
from dataclasses import dataclass, field
from trl import DPOConfig
@dataclass
class MyDPOConfig(DPOConfig):
my_coef: float = field(default=0.1, metadata={"help": "Coefficient of the custom term."})Thay đổi định dạng lô
Khi loss cần các đầu vào mà collator mặc định không tạo ra, hãy truyền collator của riêng bạn và ghi đè _prepare_dataset để giữ nguyên dataset khi nó đã ở định dạng mong đợi:
class MyDPOTrainer(DPOTrainer):
def _prepare_dataset(self, dataset, *args, **kwargs):
return dataset
trainer = MyDPOTrainer(..., data_collator=my_collator)Ví dụ hoàn chỉnh
Ví dụ SFT block-diffusion examples/sft_diffusion_gemma/sft_diffusion_gemma.py kết hợp cả ba: DiffusionGemmaSFTConfig kế thừa SFTConfig với các tham số canvas và corruption, còn DiffusionGemmaSFTTrainer kế thừa SFTTrainer và thay thế cross-entropy tự hồi quy bằng mục tiêu khử nhiễu block-diffusion. Cả hai đều không yêu cầu thay đổi thư viện.
Antidoom, một công cụ mã nguồn mở từ Liquid AI, mở rộng DPOTrainer bằng Final Token Preference Optimization (Antislop), một hàm mất mát ưu tiên trên một vị trí token duy nhất nhằm giảm các vòng lặp lặp lại trong các mô hình suy luận.
Update on GitHub