scheduler

This commit is contained in:
sayakpaul
2024-12-26 07:55:07 +05:30
parent 28ecfe6b65
commit 0c0a0a5e4a
2 changed files with 22 additions and 11 deletions
+22 -10
View File
@@ -535,19 +535,31 @@ class Trainer:
self.state.train_steps = self.state.train_epochs * num_update_steps_per_epoch self.state.train_steps = self.state.train_epochs * num_update_steps_per_epoch
self.state.overwrote_max_train_steps = True self.state.overwrote_max_train_steps = True
use_deepspeed_lr_schd = ( use_deepspeed_lr_scheduler = (
self.state.accelerator.state.deepspeed_plugin is not None self.state.accelerator.state.deepspeed_plugin is not None
and "scheduler" in self.state.accelerator.state.deepspeed_plugin.deepspeed_config and "scheduler" in self.state.accelerator.state.deepspeed_plugin.deepspeed_config
) )
lr_scheduler = get_scheduler( total_training_steps = self.state.train_steps * self.state.accelerator.num_processes
name=self.args.lr_scheduler, num_warmup_steps = self.args.lr_warmup_steps * self.state.accelerator.num_processes
optimizer=optimizer,
num_warmup_steps=self.args.lr_warmup_steps * self.state.accelerator.num_processes, if use_deepspeed_lr_scheduler:
num_training_steps=self.state.train_steps * self.state.accelerator.num_processes, from accelerate.utils import DummyScheduler
num_cycles=self.args.lr_num_cycles,
power=self.args.lr_power, lr_scheduler = DummyScheduler(
use_deepspeed=use_deepspeed_lr_schd, name=self.args.lr_scheduler,
) optimizer=optimizer,
total_num_steps=total_training_steps,
num_warmup_steps=num_warmup_steps,
)
else:
lr_scheduler = get_scheduler(
name=self.args.lr_scheduler,
optimizer=optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=total_training_steps,
num_cycles=self.args.lr_num_cycles,
power=self.args.lr_power,
)
self.optimizer = optimizer self.optimizer = optimizer
self.lr_scheduler = lr_scheduler self.lr_scheduler = lr_scheduler
-1
View File
@@ -1,5 +1,4 @@
import inspect import inspect
import logging
from accelerate.logging import get_logger from accelerate.logging import get_logger
import torch import torch