wehub-resource-sync
88 lines
3.4 KiB
Python
88 lines
3.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# DeepSpeed Team
|
|
|
|
import math
|
|
|
|
try:
|
|
import torch
|
|
import torch_supa_ext.deepspeed # noqa: F401 — registers torch.ops.deepspeed
|
|
except ImportError:
|
|
pass
|
|
|
|
try:
|
|
_has_kernel = hasattr(torch.ops, 'deepspeed') and hasattr(torch.ops.deepspeed, 'multi_tensor_adam')
|
|
except Exception:
|
|
_has_kernel = False
|
|
|
|
from .builder import SUPAOpBuilder
|
|
|
|
|
|
class SUPAFusedAdam:
|
|
"""
|
|
Fused Adam for Biren SUPA GPUs.
|
|
|
|
Calls torch.ops.deepspeed.multi_tensor_adam (registered by torch_supa_ext.deepspeed)
|
|
when the compiled extension is available; falls back to a numerically equivalent
|
|
pure-PyTorch loop for cmodel / functional testing.
|
|
"""
|
|
|
|
@staticmethod
|
|
def multi_tensor_adam(chunk_size, noop_flag_buffer, tensor_lists, lr, beta1, beta2, epsilon, step, mode,
|
|
bias_correction, weight_decay):
|
|
import torch # ensure torch is available at runtime
|
|
|
|
# noop_flag guard (kernel also checks internally, but short-circuit here is cheap)
|
|
if noop_flag_buffer.item() == 1:
|
|
return
|
|
|
|
if _has_kernel:
|
|
# MR #96 API: four separate Tensor-list arguments (not a nested list)
|
|
grads, params, exp_avgs, exp_avg_sqs = tensor_lists
|
|
torch.ops.deepspeed.multi_tensor_adam(chunk_size, noop_flag_buffer, grads, params, exp_avgs, exp_avg_sqs,
|
|
lr, beta1, beta2, epsilon, step, mode, bias_correction, weight_decay)
|
|
return
|
|
|
|
# Pure-PyTorch fallback (cmodel / no compiled backend)
|
|
bias_correction1 = 1.0 - beta1**step if bias_correction else 1.0
|
|
bias_correction2 = 1.0 - beta2**step if bias_correction else 1.0
|
|
for i in range(len(tensor_lists[0])):
|
|
g = tensor_lists[0][i].float()
|
|
p = tensor_lists[1][i]
|
|
m = tensor_lists[2][i]
|
|
v = tensor_lists[3][i]
|
|
if mode == 1: # AdamW: decoupled weight decay
|
|
m.mul_(beta1).add_(g, alpha=1.0 - beta1)
|
|
v.mul_(beta2).addcmul_(g, g, value=1.0 - beta2)
|
|
denom = (v.sqrt() / math.sqrt(bias_correction2)).add_(epsilon)
|
|
# Decouple weight decay on the old param before the Adam step so the
|
|
# result matches the kernel's p_old*(1 - lr*wd) - lr*adam_update.
|
|
p.data.add_(p.data, alpha=-lr * weight_decay)
|
|
p.data.addcdiv_(m, denom, value=-(lr / bias_correction1))
|
|
else: # Adam: L2 regularization
|
|
g_wd = g.add(p.float(), alpha=weight_decay)
|
|
m.mul_(beta1).add_(g_wd, alpha=1.0 - beta1)
|
|
v.mul_(beta2).addcmul_(g_wd, g_wd, value=1.0 - beta2)
|
|
denom = (v.sqrt() / math.sqrt(bias_correction2)).add_(epsilon)
|
|
p.data.addcdiv_(m, denom, value=-(lr / bias_correction1))
|
|
|
|
|
|
class FusedAdamBuilder(SUPAOpBuilder):
|
|
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 load(self, verbose=True):
|
|
return SUPAFusedAdam
|
|
|
|
def is_compatible(self, verbose=False):
|
|
import torch # ensure torch is available at runtime
|
|
return hasattr(torch.ops, 'deepspeed') and hasattr(torch.ops.deepspeed, 'multi_tensor_adam')
|