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

[rename] convert_to_dist -> redistribute (#1243)

Browse Source
This commit is contained in:
Jiarui Fang
2022-07-11 13:05:44 +08:00
committed by GitHub
parent f6add9b720
commit 2699dfbbfd
7 changed files with 18 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/test_tensor/test_loss_func.py
View File

@@ -22,7 +22,7 @@ def check_cross_entropy():
world_size = torch.distributed.get_world_size()
pg = ProcessGroup(tp_degree=world_size)
input_t_colo = ColoTensor.from_torch_tensor(tensor=input_ct, spec=ColoTensorSpec(pg))
input_shard = input_t_colo.convert_to_dist_spec(distspec.shard([-1], [pg.tp_world_size()]))
input_shard = input_t_colo.redistribute(distspec.shard([-1], [pg.tp_world_size()]))
input_shard.set_tensor_spec(dist_spec=None, compute_spec=ComputeSpec(ComputePattern.TP1D))
output = F.cross_entropy(input_t, target)