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

[moe] init moe plugin comm setting with sp

Browse Source
This commit is contained in:
hxwang
2024-07-18 08:37:06 +00:00
committed by Hongxin Liu
parent 09d6280d3e
commit 877d94bb8c
7 changed files with 101 additions and 95 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/test_moe/test_moe_checkpoint.py
View File

@@ -26,9 +26,7 @@ top_k = 2
def check_model_equal(model1, model2):
assert set(model1.state_dict().keys()) == set(model2.state_dict().keys())
for i, ((name, p1), p2) in enumerate(zip(model1.named_parameters(), model2.parameters())):
if loose_close(p1, p2, p1.dtype):
print(f"Model parameter {name} is not equal. is_moe_tensor: {is_moe_tensor(p1)}")
raise AssertionError(f"Model parameter {name} is not equal")
loose_close(p1, p2, p1.dtype)
def get_optimizer_snapshot(optim):