github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2026-01-29 21:49:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

add moe context, moe utilities and refactor gradient handler (#455)

Browse Source
This commit is contained in:
HELSON
2022-03-18 16:38:32 +08:00
committed by GitHub
parent af185b5519
commit 84fd7c1d4d
11 changed files with 255 additions and 125 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/context/random/_helper.py
View File

@@ -147,15 +147,10 @@ def with_seed(func, parallel_mode: ParallelMode):
def moe_set_seed(seed):
if torch.cuda.is_available():
from colossalai.core import global_context as gpc
moe_mp_rank = gpc.get_local_rank(ParallelMode.MOE_MODEL)
moe_mp_seed = seed + moe_mp_rank
add_seed(ParallelMode.MOE_MODEL, moe_mp_seed)
global_rank = gpc.get_global_rank()
add_seed(ParallelMode.TENSOR, global_rank, True)
print(f"moe seed condition: {global_rank} with moe seed {moe_mp_seed}, ",
f"tensor seed {global_rank}",
flush=True)
diff_seed = seed + global_rank
add_seed(ParallelMode.TENSOR, diff_seed, True)
print(f"moe seed condition: {global_rank} with tensor seed {diff_seed}", flush=True)
def reset_seeds():