github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-24 18:51:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'coati/support-pp' of github.com:TongLi3701/ColossalAI into coati/support-pp

Browse Source
This commit is contained in:
Tong Li 2024-08-12 11:35:40 +00:00
commit 641867018f
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
applications/ColossalChat/coati/trainer/rm.py
View File

@ -48,7 +48,7 @@ class RewardModelTrainer(SLTrainer):
model: Any, model: Any,
booster: Booster, booster: Booster,
optimizer: Optimizer, optimizer: Optimizer,
plugin: Plugin, plugin: Plugin,
lr_scheduler: _LRScheduler, lr_scheduler: _LRScheduler,
tokenizer: PreTrainedTokenizerBase, tokenizer: PreTrainedTokenizerBase,
loss_fn: Optional[Callable] = None, loss_fn: Optional[Callable] = None,
@ -60,7 +60,9 @@ class RewardModelTrainer(SLTrainer):
save_dir: str = None, save_dir: str = None,
coordinator: DistCoordinator = None, coordinator: DistCoordinator = None,
) -> None: ) -> None:
super().__init__(booster, max_epochs=max_epochs, model=model, optimizer=optimizer, plugin=plugin, start_epoch=start_epoch) super().__init__(
booster, max_epochs=max_epochs, model=model, optimizer=optimizer, plugin=plugin, start_epoch=start_epoch
)
self.actor_scheduler = lr_scheduler self.actor_scheduler = lr_scheduler
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.loss_fn = loss_fn if loss_fn is not None else LogSigLoss(beta=beta) self.loss_fn = loss_fn if loss_fn is not None else LogSigLoss(beta=beta)