reduce cpu host overhead when using moe #5578
Conversation
| @@ -366,7 +366,7 @@ def top2gating(logits: Tensor, | |||
| combine_weights = combine1_sec + combine2_sec | |||
| dispatch_mask = combine_weights.bool() | |||
|
|
|||
| return l_aux, combine_weights, dispatch_mask, exp_counts.detach().to('cpu') | |||
| return l_aux, combine_weights, dispatch_mask, exp_counts | |||
There was a problem hiding this comment.
Currently exp_counts is unused at forward() of any of MoE classes, right?
There was a problem hiding this comment.
Currently
exp_countsis unused atforward()of any of MoE classes, right?
Yes, I have test it in Megatron-deepspeed and find that exp_counts is unused at forward() of any of MoE classes.
There was a problem hiding this comment.
I have test these changes on my testing GPU platform, and work fine, no error and the loss keeps same to the original way.
There was a problem hiding this comment.
@ranzhejiang Thank you for your contribution! I have a few questions about your changes. Can you clarify them?
deepspeed/moe/sharded_moe.py
Outdated
| @@ -322,7 +322,7 @@ def top2gating(logits: Tensor, | |||
| l_aux = torch.mean(me * ce) * num_experts * num_experts | |||
|
|
|||
| # gating decisions | |||
| exp_counts = torch.sum(mask1 + mask2, dim=0) | |||
| exp_counts = torch.sum(mask1 + mask2, dim=0).detach().to(logits.device) | |||
There was a problem hiding this comment.
Can the device of mask1 and mask1 be different from logits?
There was a problem hiding this comment.
Can the device of
mask1andmask1be different fromlogits?
From line 296 to 301, we can find that , the calculation of mask1 depends on logits, and all torch operation will keep the original device, so the device of mask1 and logits be the same one. The same to mask1 and mask2 line 309 to 311
ranzhejiang
force-pushed
the
zhejiang/reduce_host_overhead_moe
branch
from
e9e32f4
to
d860d2c
Compare
|
Hi, @tohtana I have clarified the modifications you mentioned and retest this PR with Megatron-Deepspeed on GPU platform(8xA800). It runs well and loss remains consistent with the original method, Could you please help review it again? Thanks! |
The operation
.to('cpu')is not necessary for exp_counts, and it will cause device to host synchronization which damage performance.