Skip to content

[megatron] fix: return 3-tuple under calculate_per_token_loss to fix MoE aux/z-loss grad blowup at CP>1#6836

Open
EricMarcus-ai wants to merge 4 commits into
verl-project:mainfrom
kaiko-ai:fix/megatron-moe-zloss-per-token-cp
Open

[megatron] fix: return 3-tuple under calculate_per_token_loss to fix MoE aux/z-loss grad blowup at CP>1#6836
EricMarcus-ai wants to merge 4 commits into
verl-project:mainfrom
kaiko-ai:fix/megatron-moe-zloss-per-token-cp