Fix: prevent double accumulation of load balancing loss and z-loss wi…

[thuwzt]

Fix for the issue #1330 : [BUG] MoE load balancing loss is accumulated twice when using activation checkpointing

[taehwakkwon]

Thank you for fixing this issue. I have a question regarding the aux loss.

Although it appears to be doubled, I believe the actual training remains the same since recomputed is turned off. Aux loss is not accumulated to the optimizer. The reason I'm asking is that we previously had recomputed turned on for a period, but now we can turn it off. I'm confused about whether to double the moe-aux-loss-coeff. I don't think it's necessary to double it because loss value is added to aux_losses_tracker.

Megatron-LM/megatron/core/transformer/moe/moe_utils.py

Line 448 in 2d7c521

tracker[name]["values"][layer_number - 1] += loss.detach() # Aggregate the loss for the layer.

.

[thuwzt]

@taehwakkwon Yes, you understand correctly. This is merely a display bug related to the tracker and logger. The compute flow remains unaffected, regardless of whether --moe-layer-recompute is used or not.