423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2025-12-15 03:10:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Cleaner integration.

Browse Source
This commit is contained in:
Jonery
2024-06-19 12:29:40 +08:00
parent 97c5235160
commit 5c2ff1b749
8 changed files with 24 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/llamafactory/train/rm/trainer.py
View File

@@ -72,15 +72,9 @@ class PairwiseTrainer(Trainer):
self.processor = processor
self.can_return_loss = True # override property to return eval_loss
if finetuning_args.use_badam:
from badam import clip_grad_norm_for_sparse_tensor
self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_for_sparse_tensor, self.accelerator)
if (self.args.deepspeed_plugin is not None
and self.args.deepspeed_plugin.zero_stage == 3
):
from badam.utils import BAdamZeRO3Callback
self.callback_handler.add_callback(BAdamZeRO3Callback)
from badam import clip_grad_norm_old_version, BAdamCallback
self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, self.accelerator)
self.callback_handler.add_callback(BAdamCallback)
def create_optimizer(self) -> "torch.optim.Optimizer":
if self.optimizer is None: