github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-08 20:40:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

[kernel] move all symlinks of kernel to colossalai._C (#1971)

Browse Source
This commit is contained in:
ver217
2022-11-17 13:42:33 +08:00
committed by GitHub
parent 7e24b9b9ee
commit f8a7148dec
27 changed files with 463 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/nn/optimizer/fused_sgd.py
View File

@@ -20,7 +20,7 @@ class FusedSGD(Optimizer):
:class:`colossalai.nn.optimizer.FusedSGD` may be used as a drop-in replacement for ``torch.optim.SGD``
:class:`colossalai.nn.optimizer.FusedSGD` may be used with or without Amp.
:class:`colossalai.nn.optimizer.FusedSGD` may be used with or without Amp.
Nesterov momentum is based on the formula from
`On the importance of initialization and momentum in deep learning`__.
@@ -80,12 +80,13 @@ class FusedSGD(Optimizer):
self.wd_after_momentum = wd_after_momentum
if multi_tensor_applier.available:
import colossal_C
import colossalai._C.fused_optim
# Skip buffer
self._dummy_overflow_buf = torch.tensor([0],
dtype=torch.int,
device=self.param_groups[0]["params"][0].device)
self.multi_tensor_sgd = colossal_C.multi_tensor_sgd
self.multi_tensor_sgd = colossalai._C.fused_optim.multi_tensor_sgd
else:
raise RuntimeError('FusedSGD requires cuda extensions')