Cleaner integration.
Former-commit-id: 26d4b05d424bd71f570195dd433258caf6465d92
This commit is contained in:
@@ -96,15 +96,9 @@ class CustomDPOTrainer(DPOTrainer):
|
||||
self.save_model(os.path.join(self.args.output_dir, "pissa_init"))
|
||||
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user