[li-lizhe]

Summary

NpuGroupedLoraLinear registers one weight{i} (and optional bias{i}) per local expert, but those parameters never received the topology metadata Megatron uses to tell expert parameters apart from ordinary dense ones. Megatron's parameter/gradient buffer classifies an expert parameter with

is_expert_param = not getattr(param, "allreduce", True)

so without allreduce = False these per-expert adapters were indistinguishable from dense parameters and could be reduced over the dense TP/DP groups instead of the expert groups.

This PR sets allreduce = False on the per-expert weights and biases when is_expert is true. Dense adapters (is_expert=False) keep the default so their reduction is unchanged, and expert data parallel still synchronizes replicas of the same expert.

Verification

Environment: Ascend 910B4 (device 0), torch 2.15.0.dev20260917+cpu, torch_npu 2.15.0.dev20260917+gitec69335, Python 3.12. npu-smi reported 910B4 OK during the run.

megatron-core / transformer_engine / peft are not installed on that host, so the real NpuGroupedLoraLinear.__init__ was executed with the file taken verbatim from this branch and only those third-party imports stubbed, with real torch_npu, and the resulting parameters inspected the same way Megatron's buffer inspects them.

Upstream main (before), 4 experts, is_expert=True, use_cpu_initialization=True:

weight0  shape=(64, 32) dtype=torch.bfloat16 device=cpu allreduce=<absent> -> classified_as_expert=False
bias0    shape=(64,) dtype=torch.bfloat16 device=cpu allreduce=<absent> -> classified_as_expert=False
=> 0/8 params classified as expert params

This branch (after), 4 experts, is_expert=True:

weight0  shape=(64, 32) dtype=torch.bfloat16 device=npu:0 allreduce=False -> classified_as_expert=True
bias0    shape=(64,) dtype=torch.bfloat16 device=npu:0 allreduce=False -> classified_as_expert=True
=> 8/8 params classified as expert params

This branch (after), dense adapter, is_expert=False:

weight0  shape=(64, 32) dtype=torch.bfloat16 device=npu:0 allreduce=<absent> -> classified_as_expert=False
=> 0/8 params classified as expert params

A regression test is added in tests/test_npu_grouped_lora_expert_params.py.

Not tested

• No end-to-end MoE training run was performed, so the actual collective group chosen at runtime is not observed directly; what is verified is the classification input Megatron reads.
• No multi-rank torchrun run (HCCL/EDP group assignment is therefore not observed).

Refs #229

[addsubmuldiv]

Could you please run pre-commit run --all-files and fix the lint issues it reports? Thanks.

[li-lizhe]

Thanks! I ran the pinned hooks' tools locally (flake8, isort 8.0.1, yapf 0.43.0, plus the pre-commit-hooks checks incl. trailing whitespace / EOF / line endings). The only failure was isort on tests/test_npu_grouped_lora_expert_params.py: it wants a section separator (blank line) between the third-party (megatron) and first-party (mcore_bridge) import blocks. Fixed in dc87cd18 and pushed to the branch; isort/flake8/yapf are clean on the changed files now. The lint workflow should re-run on the new head.