[assets] update wechat (#8385)

This commit is contained in:
Yaowei Zheng
2025-06-16 18:23:22 +08:00
committed by GitHub
parent 8e4ac78607
commit 9a2d1dec62
5 changed files with 21 additions and 11 deletions

View File

@@ -188,8 +188,9 @@ class CustomDPOTrainer(DPOTrainer):
batch = nested_detach(batch, clone=True) # avoid error
all_logits: torch.Tensor = model(**batch, return_dict=True, use_cache=False).logits.to(torch.float32)
all_logps, valid_length = get_batch_logps(logits=all_logits, labels=batch["labels"],
ld_alpha=(self.ld_alpha if not is_ref_model else None))
all_logps, valid_length = get_batch_logps(
logits=all_logits, labels=batch["labels"], ld_alpha=(self.ld_alpha if not is_ref_model else None)
)
if self.loss_type in ["ipo", "orpo", "simpo"]:
all_logps = all_logps / valid_length
@@ -219,8 +220,9 @@ class CustomDPOTrainer(DPOTrainer):
ref_context = nullcontext()
with torch.no_grad(), ref_context:
reference_chosen_logps, reference_rejected_logps, *_ = self.concatenated_forward(ref_model, batch,
is_ref_model=True)
reference_chosen_logps, reference_rejected_logps, *_ = self.concatenated_forward(
ref_model, batch, is_ref_model=True
)
return reference_chosen_logps, reference_rejected_logps