github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-10-11 18:01:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[zero] update zero context init with the updated test utils (#327)

Browse Source
This commit is contained in:
Jiarui Fang
2022-03-08 14:45:01 +08:00
committed by Frank Lee
parent 6268446b81
commit 11bddb6e55
10 changed files with 96 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_zero_data_parallel/test_shard_param.py
View File

@@ -46,6 +46,8 @@ def _run_shard_param_v2(rank, world_size, port):
sparam = ShardedParamV2(param=param, process_group=None)
allclose(sparam.data, param_ref.data)
sparam.remove_torch_payload()
assert (param.data.numel() == 1)