fix(training): count expert parameter norms with TP and EP - #4998
Conversation
yaoyu-33
added
area:training
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
c57cd8f to
482bf93
Compare
|
Review LGTM. Correct, well-scoped bug fix in calc_params_l2_norm. What it fixes Expert (MoE) params with sharded distributed-optimizer main params were appended to sharded_params_data and reduced over the regular DP+CP group (pg_collection.dp_cp). But the distributed optimizer shards expert main params over the expert DP group, so the norm was summed over the wrong process group -- producing an incorrect global norm whenever EP is combined with the distributed optimizer in bf16. The fix routes them into a dedicated sharded_moe_params_data list reduced over pg_collection.expt_dp. Correctness notes
Test coverage Strong. New/updated unit tests assert dense-vs-expert reductions land on matching DP groups, the duplicate filter receives both TP groups, and an end-to-end mixed dense+expert norm (sqrt(41)) validates accumulation. Suggested test cases No perf tests impacted (only train_utils.py and its unit test changed; no scripts/performance/configs/ entries touched). Relevant unit cases:
|
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| if len(sharded_moe_params_data) > 0: | ||
| sharded_moe_norm, _ = multi_tensor_applier( | ||
| multi_tensor_l2norm, | ||
| dummy_overflow_buf, | ||
| [sharded_moe_params_data], | ||
| False, # no per-parameter norm. | ||
| ) | ||
| sharded_moe_norm_2 = sharded_moe_norm * sharded_moe_norm | ||
| else: | ||
| sharded_moe_norm_2 = torch.zeros((1,), dtype=torch.float32, device="cuda") |
There was a problem hiding this comment.
The expert-sharded all_reduce on pg_collection.expt_dp is issued unconditionally (even for empty sharded_moe_params_data), which correctly mirrors the dense sharded_norm_2 reduction and prevents NCCL hangs when ranks have uneven expert-param counts. Nice — just confirming this collective is always reached on every rank in expt_dp, since a conditional collective here would deadlock.
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
yaoyu-33
added
r0.6.0
The existing calc_params_l2_norm tests all mock param_is_not_tensor_parallel_duplicate, so they would pass even against an MCore build that dropped the expert_tp_group kwarg -- silently undercounting expert parameter norms. Add a test that exercises the real MCore function so a pinned MCore regressing below this contract fails loudly. Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
|
/ok to test fcc2c82 |
1 participant
Summary
Root cause and impact
Bridge maintains an independent
calc_params_l2_normutility for parameter-norm reporting. It filtered every parameter with the regular tensor-parallel topology and accumulated all distributed-optimizer main-parameter shards over ordinary DP/CP.For parameters with
allreduce=False, duplicate filtering must use expert TP. Their sharded FP32 main parameters must likewise be reduced over expert DP. The old behavior could omit a logical expert parameter or aggregate its shard over the wrong domain, producing an undercounted parameter L2 norm in training logs.Why the Bridge diff differs from the final MCore diff
MCore #5916 originally implemented the same separate expert-sharded bucket in commit
2629ec4c3ae700d0fd67a4ef4399ad9e2d9c6461. A later GTP refactor absorbed that bucket into MCore's generalized norm/reduction structure, so the final MCore PR diff is smaller.Bridge still owns this independent training-log calculation and uses model-attached process groups (
pg_collection.tp,pg_collection.expt_tp,pg_collection.dp_cp, andpg_collection.expt_dp) instead of MCore's global MPU accessors. Bridge already delegates optimizer gradient norm, clipping, and synchronization to MCore, so this PR does not duplicate any optimizer, TE, DDP, or grouped-linear code.Upstream status
NVIDIA/Megatron-LM#5916 is merged. The current Bridge
mainpin is its exact merge commit:cd4afffa648426a959dc7cb1e24b5ce7d0c3ff54The previous upstream blocker is therefore resolved, and this PR does not change the MCore submodule pointer or dependency metadata.
Test-first validation
Both contracts were executed against current Bridge
mainproduction code before applying this PR's implementation, using the same tests and the pinned MCore merge commit.RED: current main without the Bridge fix
0.02.05.0sqrt(41) = 6.403124...2 failedGREEN: same tests with the Bridge fix
2 passedTestCalcParamsL2Normclass:25 passeduv run pre-commit run --all-files: passedgit diff --check: passedThe sharded numeric test assigns independent remote contributions to DP/CP and expert DP and asserts the final norm, so it validates the observable result rather than only checking mocked call arguments.