github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-06 19:40:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

fixed layout converter caching and updated tester

Browse Source
This commit is contained in:
Edenzzzz
2024-03-26 17:22:27 +08:00
parent a7790a92e8
commit 61da3fbc52
2 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
tests/test_tensor/test_dtensor/test_layout_converter.py
View File

@@ -123,8 +123,15 @@ def check_layout_converting(rank, world_size, port):
assert comm_action_sequence[2].logical_process_axis == 1
# checkout chached_spec_pairs_transform_path
assert layout_converter.cached_solution[("[R, S01, R]", "[S01, R, R]")][0] == transform_path
assert layout_converter.cached_solution[("[R, S01, R]", "[S01, R, R]")][1] == comm_action_sequence
src_shape = source_layout.get_sharded_shape_per_device()
dst_shape = target_layout.get_sharded_shape_per_device()
assert (
layout_converter.cached_solution[(("[R, S01, R]", src_shape), ("[S01, R, R]", dst_shape))][0] == transform_path
)
assert (
layout_converter.cached_solution[(("[R, S01, R]", src_shape), ("[S01, R, R]", dst_shape))][1]
== comm_action_sequence
)
comm_cost = layout_converter.get_total_comm_cost(source_layout, target_layout)