wehub-resource-sync
108 lines
4.5 KiB
Python
108 lines
4.5 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import math
|
|
|
|
from deepspeed.utils import logger
|
|
# from deepspeed.runtime.lr_schedules import WarmupLR
|
|
from ..constants import *
|
|
|
|
#####based on the paper random-ltd: https://arxiv.org/abs/2211.11586
|
|
|
|
|
|
class BaseScheduler(object):
|
|
|
|
def __init__(self):
|
|
self.state = {}
|
|
|
|
def __fixed_root_get_value(self, global_steps, root_degree=None):
|
|
s_state = self.state[RANDOM_LTD_SCHEDULE_CONFIG]
|
|
if root_degree is None:
|
|
root_degree = s_state['root_degree']
|
|
next_seq = (float(global_steps) / s_state[RANDOM_LTD_REQUIRE_STEP])**(1.0 / root_degree)
|
|
next_seq = math.floor(next_seq * (self.state[RANDOM_LTD_MAX_VALUE] - self.state[RANDOM_LTD_MIN_VALUE]) +
|
|
self.state[RANDOM_LTD_MIN_VALUE])
|
|
next_seq -= (next_seq % s_state[RANDOM_LTD_INCREASE_STEP])
|
|
next_seq = min(next_seq, self.state[RANDOM_LTD_MAX_VALUE])
|
|
return next_seq
|
|
|
|
def get_value(self, global_steps):
|
|
if self.state[RANDOM_LTD_SCHEDULER_TYPE] == 'fixed_linear':
|
|
return self.__fixed_root_get_value(global_steps, 1)
|
|
else:
|
|
raise RuntimeError('Unsupported random LTD schedule type')
|
|
|
|
|
|
class RandomLTDScheduler(BaseScheduler):
|
|
|
|
def __init__(self, config):
|
|
super().__init__()
|
|
self.model_layer_num = config[RANDOM_LTD_TOTAL_LAYER_NUM]
|
|
self.random_ltd_layer_num = config[RANDOM_LTD_LAYER_NUM]
|
|
self.config_schedule = config[RANDOM_LTD_SCHEDULER]
|
|
self.global_batch_size = config[RANDOM_LTD_GLOBAL_BATCH_SIZE]
|
|
self.reset_to_init()
|
|
|
|
if config[RANDOM_LTD_LAYER_TOKEN_LR_SCHEDULE][RANDOM_LTD_LAYER_TOKEN_LR_ENABLED]:
|
|
logger.warning("**********Work In Progress************")
|
|
raise NotImplementedError
|
|
|
|
self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0
|
|
|
|
# self.first_step = True
|
|
def get_total_layer_tokens(self, train_iters):
|
|
for step in range(train_iters):
|
|
self.update_seq(step)
|
|
return self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS]
|
|
|
|
def reset_to_init(self):
|
|
if self.config_schedule is not None:
|
|
self.state[RANDOM_LTD_MIN_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE]
|
|
self.state[RANDOM_LTD_MAX_VALUE] = self.config_schedule[RANDOM_LTD_MAX_VALUE]
|
|
self.state[RANDOM_LTD_CURRENT_VALUE] = self.config_schedule[RANDOM_LTD_MIN_VALUE]
|
|
self.state[RANDOM_LTD_SCHEDULE_CONFIG] = self.config_schedule[RANDOM_LTD_SCHEDULE_CONFIG]
|
|
self.state[RANDOM_LTD_SCHEDULER_TYPE] = self.config_schedule[RANDOM_LTD_SCHEDULER_TYPE]
|
|
self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = 0
|
|
self.state[RANDOM_LTD_CURR_STEP] = -1
|
|
|
|
def get_current_seq(self):
|
|
return self.state[RANDOM_LTD_CURRENT_VALUE]
|
|
|
|
def set_current_seq(self, seq_length):
|
|
self.state[RANDOM_LTD_CURRENT_VALUE] = seq_length
|
|
|
|
def get_random_ltd_layer_num(self):
|
|
return self.random_ltd_layer_num
|
|
|
|
def get_state(self):
|
|
return self.state
|
|
|
|
def set_state(self, state):
|
|
self.state = state
|
|
|
|
def update_seq(self, global_steps):
|
|
if self.state[RANDOM_LTD_CURRENT_VALUE] < self.state[RANDOM_LTD_MAX_VALUE]:
|
|
self.state[RANDOM_LTD_CURRENT_VALUE] = self.get_value(global_steps)
|
|
if global_steps != self.state[RANDOM_LTD_CURR_STEP]:
|
|
self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] += self.global_batch_size*(self.state[RANDOM_LTD_CURRENT_VALUE] * self.random_ltd_layer_num \
|
|
+ self.state[RANDOM_LTD_MAX_VALUE] * (self.model_layer_num - self.random_ltd_layer_num))
|
|
self.state[RANDOM_LTD_CURR_STEP] = global_steps
|
|
|
|
def state_dict(self):
|
|
return {
|
|
RANDOM_LTD_CONSUMED_LAYER_TOKENS: self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS],
|
|
RANDOM_LTD_CURR_STEP: self.state[RANDOM_LTD_CURR_STEP],
|
|
RANDOM_LTD_CURRENT_VALUE: self.state[RANDOM_LTD_CURRENT_VALUE],
|
|
RANDOM_LTD_MIN_VALUE: self.state[RANDOM_LTD_MIN_VALUE],
|
|
RANDOM_LTD_MAX_VALUE: self.state[RANDOM_LTD_MAX_VALUE],
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
self.state[RANDOM_LTD_CONSUMED_LAYER_TOKENS] = state_dict[RANDOM_LTD_CONSUMED_LAYER_TOKENS]
|
|
self.state[RANDOM_LTD_CURR_STEP] = state_dict[RANDOM_LTD_CURR_STEP]
|
|
self.state[RANDOM_LTD_CURRENT_VALUE] = state_dict[RANDOM_LTD_CURRENT_VALUE]
|
|
self.state[RANDOM_LTD_MIN_VALUE] = state_dict[RANDOM_LTD_MIN_VALUE]
|
|
self.state[RANDOM_LTD_MAX_VALUE] = state_dict[RANDOM_LTD_MAX_VALUE]
|