wehub-resource-sync
75 lines
2.1 KiB
Python
75 lines
2.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from .builder import NPUOpBuilder
|
|
|
|
try:
|
|
import torch_npu
|
|
except ImportError as e:
|
|
pass
|
|
|
|
|
|
class NPUFusedAdam:
|
|
|
|
@staticmethod
|
|
def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, adam_w_mode,
|
|
bias_correction, weight_decay, *args):
|
|
bias_correction1 = beta1**(step - 1)
|
|
bias_correction2 = beta2**(step - 1)
|
|
|
|
# iteration group['params']
|
|
for i in range(len(tensor_lists[0])):
|
|
grad_flat = tensor_lists[0][i]
|
|
param_flat = tensor_lists[1][i]
|
|
m_flat = tensor_lists[2][i]
|
|
v_flat = tensor_lists[3][i]
|
|
|
|
if adam_w_mode:
|
|
param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam_w(
|
|
bias_correction1,
|
|
bias_correction2,
|
|
lr,
|
|
weight_decay,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
grad_flat,
|
|
None, # max_grad_norm
|
|
False, # amsgrad
|
|
False, # maximize
|
|
out=(param_flat.data, m_flat, v_flat))
|
|
else:
|
|
param_flat.data, m_flat, v_flat = torch_npu.npu_apply_adam(
|
|
bias_correction1,
|
|
bias_correction2,
|
|
lr,
|
|
beta1,
|
|
beta2,
|
|
epsilon,
|
|
grad_flat,
|
|
False, # use_locking
|
|
False, # use_nesterov
|
|
out=(param_flat.data, m_flat, v_flat))
|
|
|
|
|
|
class FusedAdamBuilder(NPUOpBuilder):
|
|
BUILD_VAR = "DS_BUILD_FUSED_ADAM"
|
|
NAME = "fused_adam"
|
|
|
|
def __init__(self):
|
|
super().__init__(name=self.NAME)
|
|
|
|
def absolute_name(self):
|
|
return f'deepspeed.ops.adam.{self.NAME}_op'
|
|
|
|
def sources(self):
|
|
return []
|
|
|
|
def include_paths(self):
|
|
return []
|
|
|
|
def load(self, verbose=True):
|
|
return NPUFusedAdam
|