diff --git a/colossalai/moe/_operation.py b/colossalai/moe/_operation.py index bfbf00b5a..ba087a03b 100644 --- a/colossalai/moe/_operation.py +++ b/colossalai/moe/_operation.py @@ -452,4 +452,4 @@ def all_to_all_uneven( assert ( inputs.requires_grad ), "Input must require grad to assure that backward is executed, otherwise it might hang the program." - return AllToAllUneven.apply(inputs, input_split_sizes, output_split_sizes, group, overlap, fp8_communicatio) + return AllToAllUneven.apply(inputs, input_split_sizes, output_split_sizes, group, overlap, fp8_communication)