github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-23 10:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

[pre-commit.ci] auto fixes from pre-commit.com hooks

Browse Source 
for more information, see https://pre-commit.ci
This commit is contained in:
pre-commit-ci[bot] 2024-08-12 11:29:21 +00:00 committed by Tong Li
parent ba80449d62
commit e6245485fc
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)