NVIDIA-NeMo / NVIDIA-NeMo/Automodel
Support grouped expert bias in checkpoint split/merge adapters
Open
@yuhezhang-ai is already working on this.
Since Aug 26, 2026.
- Dominant language
- Python
- Stars
- 960
- Forks
- 316
- Avg merge
- 3d 20h
- Merged PRs (30d)
- 143
Description
Summary
Add an explicit checkpoint conversion contract for grouped expert MLP bias tensors.
The shared MoE adapter currently converts grouped expert weights, but not gate_up_proj_bias / down_proj_bias. Per-expert HF bias keys can therefore pass through without being regrouped into the native layout. #3610 makes this safe by disabling its low-memory DCP capability for expert_bias=True and rejecting conversion clearly instead of silently loading or exporting incorrect keys.
Scope
- Define the HF per-expert bias keys and native grouped bias layouts for gated and non-gated experts.
- Implement both HF-to-native merge and native-to-HF split conversion.
- Preserve rank-local expert ownership under EP/DTensor configurations.
- Cover ordinary grouped experts and the maintained TE/DeepEP-backed configurations that expose expert bias.
- Keep router correction-bias tensors separate; they already have model-specific rename handling and are not this issue.
Acceptance criteria
- Load, save, and load-save-reload tests cover expert weights and both expert bias tensors.
- Missing or partial bias keys fail loudly.
- Single-device and representative EP layouts produce the same values as a trusted HF-format reference.
- The low-memory DCP capability is enabled only for runtime layouts whose destinations are model-backed or otherwise memory-bounded.
Parent: #3576
Follow-up to #3610.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.