github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-10 13:30:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero]added hybrid adam, removed loss scale in adam (#527)

Browse Source 
* [zero]added hybrid adam, removed loss scale of adam

* remove useless code
This commit is contained in:
LuGY
2022-03-25 18:03:54 +08:00
committed by GitHub
parent 8d8c5407c0
commit 105c5301c3
5 changed files with 149 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/nn/optimizer/cpu_adam.py
View File

@@ -17,7 +17,6 @@ class CPUAdam(torch.optim.Optimizer):
eps=1e-8,
weight_decay=0,
adamw_mode=True,
loss_scale=-1,
simd_log=False):
"""
An implementation equivalent to `torch.optim.Adam`.
@@ -29,8 +28,7 @@ class CPUAdam(torch.optim.Optimizer):
super(CPUAdam, self).__init__(model_params, default_args)
self.opt_id = CPUAdam.optimizer_id
CPUAdam.optimizer_id = CPUAdam.optimizer_id + 1
self.adam_w_mode = adamw_mode
self.loss_scale = loss_scale
self.adamw_mode = adamw_mode
try:
import cpu_adam
except ImportError:
@@ -54,12 +52,9 @@ class CPUAdam(torch.optim.Optimizer):
weight_decay,
bias_correction1,
bias_correction2,
loss_scale,
use_adamw=False):
# FIXME(ver217): remove the below line when replace torch adam with fused adam
grad = grad.float()
if loss_scale is not None:
grad.div_(loss_scale)
if weight_decay != 0:
if use_adamw:
@@ -110,7 +105,7 @@ class CPUAdam(torch.optim.Optimizer):
assert state['exp_avg_sq'].device.type == 'cpu', "exp_avg should stay on cpu"
self.cpu_adam_op.adam_update(self.opt_id, state['step'], group['lr'], beta1, beta2, group['eps'],
group['weight_decay'], group['bias_correction'], p.data, p.grad.data,
state['exp_avg'], state['exp_avg_sq'], self.loss_scale)
state['exp_avg'], state['exp_avg_sq'], -1)
elif target_device.type == 'cuda':
assert state['exp_avg'].device.type == 'cuda', "exp_avg should stay on cuda"
assert state['exp_avg_sq'].device.type == 'cuda', "exp_avg should stay on cuda"
@@ -121,7 +116,7 @@ class CPUAdam(torch.optim.Optimizer):
# adam on cuda
self.torch_adam_update(p.data, p.grad.data, state['exp_avg'], state['exp_avg_sq'], group['lr'],
beta1, beta2, group['eps'], group['weight_decay'], bias_correction1,
bias_correction2, self.loss_scale)
bias_correction2, self.adamw_mode)
else:
raise RuntimeError
return loss