Định nghĩa lớp con cho các phương thức của Trainer

Định nghĩa lớp con của các phương thức Trainer để thay đổi hành vi huấn luyện mà không cần viết lại toàn bộ vòng lặp. Việc định nghĩa lớp con sẽ thay đổi vòng lặp huấn luyện, chẳng hạn như bước truyền thẳng hoặc tính toán hàm mất mát.

Trước khi tạo lớp con, hãy cân nhắc xem bạn cần thay đổi điều gì mà Trainer tính toán hay khi nào và liệu có nên thực hiện hành động hay không. Đối với logic về thời gian và điều kiện, hãy sử dụng Callback thay thế. Callback kiểm soát thời điểm các sự kiện xảy ra (ghi log, đánh giá, dừng sớm), trong khi tạo lớp con thay đổi nội dung của các sự kiện đó (tính toán loss, nạp dữ liệu, tối ưu hóa).

Xem tài liệu API của Trainer để có danh sách đầy đủ các phương thức bạn có thể kế thừa. Các phương thức riêng (có tiền tố _) như _save_checkpoint hoặc _evaluate cũng có thể được ghi đè, nhưng chúng có thể thay đổi mà không có thông báo trước.

get_train_dataloader

Phương thức get_train_dataloader() chuẩn sẽ nạp một lô dữ liệu, huấn luyện trên lô đó, loại bỏ nó và nạp lô tiếp theo.

def get_train_dataloader(self):
    return self._get_dataloader(
        batch_size=self._train_batch_size,
        ...
)

GRPO là một thuật toán học tăng cường trực tuyến, tạo ra các chuỗi hoàn thành trước khi huấn luyện trên chúng. Việc tạo chuỗi hoàn thành ở mỗi bước rất tốn kém vì nó mang tính tự hồi quy. Một chuỗi hoàn thành 512 token yêu cầu ~512 lần truyền tiến tuần tự, so với chỉ một lần truyền tiến cho một bước huấn luyện. GRPOTrainer kế thừa lớp con và ghi đè get_train_dataloader() để nhóm các lần tạo chuỗi hoàn thành trên nhiều bước.

trl.GRPOTrainer.get_train_dataloader nạp các batch prompt sinh cho nhiều bước huấn luyện cùng lúc bằng cách nhân kích thước batch với đối số steps_per_generation. Nếu train_batch_size=4 và steps_per_generation=8, dataloader sẽ tạo ra các batch có kích thước 32, giảm chi phí sinh xuống 8 lần.

def get_train_dataloader(self):
    dataloader_params = {
        "batch_size": self._train_batch_size * self.args.steps_per_generation, # this is the only change
        ...
    }

Tính toán loss

compute_loss() trả về tổn thất entropy chéo được tính toán bởi mô hình.

def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    ...
    outputs = model(**inputs)
    ...
    loss = outputs["loss"] # get loss from model

    return (loss, outputs) if return_outputs else loss

DPO đo mức độ mà mô hình chính sách ưu tiên một phản hồi được chọn so với một phản hồi bị từ chối, so với một mô hình tham chiếu. DPOTrainer kế thừa compute_loss() vì việc tính toán loss khác với entropy chéo chuẩn ở một số điểm:

• mô hình không bao giờ nhìn thấy nhãn; nó chỉ trả về logits để DPO tính toán log-probs từ đó
• các phản hồi được chọn và bị từ chối được nối với nhau
• một mô hình tham chiếu tự tính toán log-probs của chính nó
• loss là một hàm của π_chosen, π_rejected, π_ref_chosen, π_ref_rejected

Không có phương án nào ở trên phù hợp với phương thức chuẩn Trainer.compute_loss().

def compute_loss(
    self,
    model: PreTrainedModel | nn.Module,
    inputs: dict[str, torch.Tensor | Any],
    return_outputs=False,
    num_items_in_batch=None,
) -> torch.Tensor | tuple[torch.Tensor, dict[str, float]]:
    ...
    outputs = model(**inputs)
    logits = outputs.logits
    logps = get_logps(logits, inputs)
    chosen_logps, rejected_logps = logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    ref_logits = self.ref_model(**inputs).logits
    ref_logps = get_logps(ref_logits, inputs)
    ref_chosen_logps, ref_rejected_logps = ref_logps.chunk(2, dim=0)  # batch is [chosen, rejected]
    chosen_scores = chosen_logps - ref_chosen_logps
    rejected_scores = rejected_logps - ref_rejected_logps
    per_sequence_loss = -F.logsigmoid(self.beta * chosen_scores - rejected_scores)
    loss = per_sequence_loss.mean()
    return (loss, outputs) if return_outputs else loss

Các bước tiếp theo

• Để xem thêm các ví dụ thực tế, hãy xem cách GRPOTrainer và DPOTrainer mở rộng Trainer trong TRL, hoặc cách Axolotl xây dựng các trainer tùy chỉnh dựa trên nó.
• Xem hướng dẫn Callbacks nếu bạn chỉ cần tùy chỉnh những gì xảy ra trong một sự kiện huấn luyện, chẳng hạn như ghi log các chỉ số vào cuối một bước huấn luyện.