github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-05 15:44:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

[colotensor] use cpu memory to store state_dict (#1367)

Browse Source
This commit is contained in:
HELSON
2022-07-26 14:13:38 +08:00
committed by GitHub
parent 943a96323e
commit 87775a0682
4 changed files with 26 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_tensor/test_colo_checkpoint_tools.py
View File

@@ -17,7 +17,7 @@ from tests.test_tensor.common_utils import tensor_shard_equal
def run_dist(rank, world_size, port, dp_degree, tp_degree):
colossalai.launch(config={}, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
pg = ProcessGroup(dp_degree=dp_degree, tp_degree=tp_degree)
x = torch.randn(4, 4, device=get_current_device())
x = torch.randn(4, 4)
param = ColoTensor(torch.nn.Parameter(x), spec=ColoTensorSpec(pg))
spec = ShardSpec([-1], [pg.tp_world_size()]), ComputeSpec(ComputePattern.TP1D)
param.set_tensor_spec(*spec)