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

[misc] update pre-commit and run all files (#4752)

Browse Source 
* [misc] update pre-commit

* [misc] run pre-commit

* [misc] remove useless configuration files

* [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
2023-09-19 14:20:26 +08:00
committed by GitHub
parent 3c6b831c26
commit 079bf3cb26
1268 changed files with 50037 additions and 38444 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
colossalai/tensor/d_tensor/utils.py
View File

@@ -7,7 +7,7 @@ from colossalai.tensor.d_tensor.layout import Layout
def get_comm_cost(layout: Layout, comm_spec: CommSpec, forward_only: bool = False) -> Dict[str, float]:
'''
"""
This method is used to compute the communication cost for a given layout and comm_spec.
For all_gather, all2all, and all_reduce operation, the formula provided in DeviceMesh with alpha-beta model is used to
@@ -18,7 +18,7 @@ def get_comm_cost(layout: Layout, comm_spec: CommSpec, forward_only: bool = Fals
comm_spec: the comm_spec to instruct the communication operation.
forward_only: if it is True, we will just count the forward communication cost.
If it is False, we will count both forward and backward communication cost.
'''
"""
comm_size = reduce(operator.mul, layout.get_sharded_shape_per_device(), 1)
device_mesh = layout.device_mesh
comm_pattern = comm_spec.comm_pattern