Skip to content

[NPU] Adapt LoRA and expert handling to Megatron Core 0.18 and TENPU - #189

Draft
addsubmuldiv wants to merge 1 commit into
modelscope:mainfrom
addsubmuldiv:adapt_018_main
Draft

[NPU] Adapt LoRA and expert handling to Megatron Core 0.18 and TENPU#189
addsubmuldiv wants to merge 1 commit into
modelscope:mainfrom
addsubmuldiv:adapt_018_main

Conversation

@addsubmuldiv

@addsubmuldiv addsubmuldiv commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adapt the NPU model configuration and LoRA integration to Megatron Core 0.18 with MegatronAdaptor and TransformerEngineNPU (TENPU), removing the dependency on MindSpeed-specific APIs.

Companion integration: modelscope/ms-swift#10087

Changes

  • Remove MindSpeed argument parsing, cached defaults, and dynamic configuration augmentation from ModelConfig.
  • Use TELinear for NPU LoRA adapters instead of selecting an nn.Linear fallback by MindSpeed version.
  • For NPU TELayerNormColumnParallelLinear, obtain the linear output and post-normalization activation from one base-layer forward using return_layernorm_output. Preserve the bias contract and restore the original flag in finally, rather than reconstructing normalization through MindSpeed private methods.
  • Correct column-parallel LoRA output sizing at the local-to-global dimension boundary and use _tp_group as the TP-group fallback.
  • Identify NPU expert layers through MCore TEGroupedLinear and expert configuration, rather than a MindSpeed module name. Pass expert tensor/data-parallel groups when constructing grouped LoRA sharded checkpoint tensors.

Validation

Tested together with the companion Swift change on Ascend 910B3, Python 3.12, CANN 9.1.0, PyTorch 2.10.0, torch_npu 2.10.0.post2, and Transformers 5.12.1. Backend snapshots: MCore ba7b5ebce12a (0.18.0), MegatronAdaptor 2a164d4548e2, and TENPU 0a389a677586.

  • Dense Qwen3-0.6B LoRA training and checkpoint loading with TP2/PP2.
  • Four-layer Qwen3.5 GDN full/LoRA SFT with TP2/PP2/CP2 and packing.
  • Four-layer Qwen3.5 MoE full/LoRA SFT with TP2/PP2/EP2. The full smoke used no_save_optim; it does not establish MoE optimizer-resume correctness.
  • Dense HF export/reload: maximum logits difference 7.82e-5.
  • Dense FSDP LoRA save/resume with merge_lora=false: matching logged loss across continuous training and two successive resume stages.

These are bounded 2–4-step integration checks, not long-run or full-size-model qualification. GPU/FP8 paths were not tested.

Scope and known limitation

USE_MCORE_GDN and LoRA merge/export behavior are unchanged. BF16 in-place merge followed by subtraction-based unmerge can change training weights through rounding; exact restoration/non-mutating export is a separate follow-up, not fixed here. Use merge_lora=false when validating uninterrupted-versus-resumed LoRA training equivalence.

@addsubmuldiv addsubmuldiv changed the title adapt_megatron_018 [NPU] Adapt LoRA and expert handling to Megatron Core 0.18 and TENPU Sep 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant