github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-13 21:22:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix style, add kto data sample

Browse Source
This commit is contained in:
YeAnbang
2024-07-18 08:38:56 +00:00
parent 845ea7214e
commit 544b7a38a1
11 changed files with 68 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
applications/ColossalChat/coati/trainer/kto.py
View File

@@ -180,7 +180,7 @@ class KTOTrainer(SLTrainer):
self.optimizer.zero_grad()
self.actor_scheduler.step()
# # sync
# sync
loss_mean = all_reduce_mean(tensor=loss)
chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards.mean())
rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards.mean())
@@ -297,7 +297,7 @@ class KTOTrainer(SLTrainer):
chosen_logprob, rejected_logprob, kl_logprob, ref_chosen_logprob, ref_rejected_logprob, ref_kl_logprob
)
# # sync
# sync
loss_mean = all_reduce_mean(tensor=loss)
chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards.mean())
rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards.mean())