NVIDIA-NeMo / NVIDIA-NeMo/Automodel

Support grouped expert bias in checkpoint split/merge adapters

Open
#3,688 0 comments 0 reactions 1 assignee View on GitHub

@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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.