github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-02 09:38:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add interface for colo tesnor dp size (#3227)

Browse Source
This commit is contained in:
YH
2023-03-27 10:42:21 +09:00
committed by GitHub
parent 1653063fce
commit 1a229045af
2 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
colossalai/tensor/colo_tensor.py
View File

@@ -138,6 +138,15 @@ class ColoTensor(torch.Tensor):
def get_tp_world_size(self) -> int:
return self.process_group.tp_world_size()
def get_dp_world_size(self) -> int:
"""get_dp_world_size
get the dp world size of the tensor.
Returns:
int: dp world size
"""
return self.process_group.dp_world_size()
def set_dist_spec(self, dist_spec: _DistSpec):
"""set_dist_spec
set dist spec and change the payloads.