github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-26 15:32:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix runtime prepare pass (#4502)

Browse Source 
Co-authored-by: lufang.chen <lufang.chen@nio.com>
This commit is contained in:
Lufang Chen 2023-08-30 17:29:38 +08:00 committed by GitHub
parent 0b00def881
commit 12c95a9fed
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
colossalai/auto_parallel/passes/runtime_preparation_pass.py
View File

@ -144,7 +144,7 @@ def size_value_converting_pass(gm: torch.fx.GraphModule, device_mesh: DeviceMesh
# DeviceMesh information instructs the scaling of the size value
device_mesh_info = {}
for dim, dim_size in enumerate(device_mesh.mesh_shape):
for dim, dim_size in enumerate(device_mesh.shape):
device_mesh_info[dim] = dim_size
def _extract_target_dim(node):