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

[hotfix] add shard dim to aviod backward communication error (#2954)

Browse Source
This commit is contained in:
YuliangLiu0306
2023-03-01 11:41:53 +08:00
committed by GitHub
parent 090f14fd6b
commit 47fb214b3b
2 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/auto_parallel/tensor_shard/node_handler/strategy/reshape_generator.py
View File

@@ -343,6 +343,7 @@ class DefaultReshapeGenerator(ReshapeGenerator):
comm_type=CommType.BEFORE,
arg_index=0)
input_comm_action.comm_spec.gather_dim = total_mesh_dim_list
input_comm_action.comm_spec.shard_dim = total_mesh_dim_list
elif len(total_mesh_dim_list) >= 2:
source_spec = sharding_spec_mapping["input"]