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

[Tensor] fix optimizer for CPU parallel (#1069)

Browse Source
This commit is contained in:
Ziyue Jiang
2022-06-06 17:36:11 +08:00
committed by GitHub
parent 49832b2344
commit 4fc748f69b
2 changed files with 16 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/test_tensor/test_hybrid_device.py
View File

@@ -9,6 +9,7 @@ from colossalai.context import ParallelMode
from colossalai.nn.parallel.layers import init_colo_module
from colossalai.nn.parallel.data_parallel import ColoDDP
from colossalai.nn.optimizer import ColoOptimizer
import colossalai
import torch
@@ -56,10 +57,11 @@ def run_hybrid_device(use_ddp):
print(f'embedding weight size: {real_model.embed.weight.size()} | new device: {real_model.embed.weight.device}')
#print(f'linear weight size: {real_model.proj.weight.size()} | new device: {real_model.proj.weight.device}')
optimizer = ColoOptimizer(dict(model.named_parameters()), torch.optim.SGD, lr=0.1)
data = torch.randint(low=0, high=20, size=(16,), device=get_current_device())
out = model(data)
out.sum().backward()
optimizer.step()
def run_dist(rank, world_size, port, use_ddp):
if use_ddp and world_size == 1:
@@ -81,4 +83,4 @@ def _test_hybrid_device(world_size, use_ddp):
if __name__ == '__main__':
_test_hybrid_device(1, False)
_test_hybrid_device(4, True)