github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-17 07:00:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into kto

Browse Source
This commit is contained in:
Tong Li
2024-07-19 15:23:31 +08:00
committed by GitHub
parent 544b7a38a1 f585d4e38e
commit d08c99be0d
14 changed files with 89 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
applications/ColossalChat/coati/trainer/sft.py
View File

@@ -102,7 +102,6 @@ class SFTTrainer(SLTrainer):
batch_size = batch["input_ids"].size(0)
outputs = self.model(batch["input_ids"], attention_mask=batch["attention_mask"], labels=batch["labels"])
loss = outputs.loss
step_bar.set_description(f"Epoch {epoch + 1}/{self.max_epochs} Loss: {loss.detach().cpu().item():.4f}")
self.booster.backward(loss=loss, optimizer=self.optimizer)
@@ -115,6 +114,7 @@ class SFTTrainer(SLTrainer):
self.optimizer.zero_grad()
self.scheduler.step()
step_bar.set_postfix({"train/loss": self.accumulative_meter.get("loss")})
if self.writer:
self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), self.num_train_step)
self.writer.add_scalar("train/lr", self.scheduler.get_last_lr()[0], self.num_train_step)