mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 23:53:48 +08:00
8cded575a9
Changes from #6472 caused the no-torch workflow that is an example of
how we build the DeepSpeed release package to fail (so we caught this
before a release, see more in #6402). These changes also copy the style
used to include torch in other accelerator op_builder implementations,
such as npu
[here](https://github.com/microsoft/DeepSpeed/blob/master/op_builder/npu/fused_adam.py#L8)
and hpu
[here](828ddfbbda/op_builder/hpu/fused_adam.py (L15)).
This also updates the no-torch workflow to run on all changes to the
op_builder directory. The test runs quickly and shouldn't add any
additional testing burden there.
Resolves: #6576
44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# Copyright (c) 2024 Cambricon Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from .builder import MLUOpBuilder
|
|
|
|
try:
|
|
import torch
|
|
except ImportError as e:
|
|
pass
|
|
|
|
|
|
class MLUFusedAdam:
|
|
|
|
@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):
|
|
|
|
torch.ops.torch_mlu.fused_adam(noop_flag_buffer, tensor_lists[0], tensor_lists[1], tensor_lists[2],
|
|
tensor_lists[3], lr, beta1, beta2, epsilon, step, adam_w_mode, bias_correction,
|
|
weight_decay)
|
|
|
|
|
|
class FusedAdamBuilder(MLUOpBuilder):
|
|
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 MLUFusedAdam
|