Kimi K2.5/2.6 LoRA + Megatron-LM backports for older runtime images

[yushengsu-thu]

Summary

This PR is a follow-up to (and superset of) #3 — it brings the Kimi K2.5/K2.6 LoRA support onto bridge-rebase-2026-05 and adds two small backport shims so the same code can run against the older Megatron-LM bundled in radixark/miles:dev without requiring a container refresh.

It contains three commits, applied in this order on top of bridge-rebase-2026-05:

#	Commit	Author	Purpose
1	ff33f80 support kimi 2.5/6 lora(logprob diff exist)	Nan Jiang	Kimi K2.5/K2.6 LoRA — same content as #3
2	c41bfdc backport: try old Megatron-LM import paths when megatron.training.config missing	Yusheng Su	Import fallback for older Megatron-LM training.config layout
3	6326bf5 backport: vendor megatron.core.models.mimo.config.role	Yusheng Su	Vendored fallback for missing megatron.core.models.mimo.config.role

Net diff: 16 files, +1,107 / −18.

---

1. Kimi K2.5/K2.6 LoRA (commit ff33f80)

Ports Nan's Kimi K2.5/K2.6 grouped-expert LoRA work onto this rebase. The K2.5 backbone is an MoE with MLA, and the new adapter machinery is needed so that LoRA can be applied to the grouped-expert linears (TEColumnParallelGroupedLinear / TERowParallelGroupedLinear) while remaining compatible with SGLang's serving contract.

src/megatron/bridge/peft/lora.py — adds experts_shared_outer_loras: bool = False to LoRA. When True, grouped-expert LoRA uses the new SharedOuterGroupedExpertAdapter: gate_up lora_A and down lora_B are shared across local experts (expert_dim=1), matching SGLang PR #21466's experts_shared_outer_loras=True contract. Default False preserves the existing layout selected by share_expert_adapters. Op-fuser is disabled for the shared-outer path.

src/megatron/bridge/peft/lora_layers.py — adds LoRAGroupedExpertLinear (subclass of LoRALinear) that plumbs tokens_per_expert / m_splits from the grouped-MLP forward into the adapter so the per-expert GroupedLinear side can run.

src/megatron/bridge/peft/utils.py (+415 lines) — the main payload:

• SharedGradAllReduce autograd function: identity forward, SUM-allreduce in backward across the intra-PP-stage group (tensor_and_data_parallel_group with CP). Keeps a logically-replicated weight in sync across EP replicas that neither allreduce=True nor allreduce=False covers. SUM (not AVG) is correct: with N EP replicas each holding a partial loss-grad over its (token, expert) subset, total grad is Σ_r g_r; AVG would train at 1/N the intended rate.
• broadcast_across_intra_stage_group(weight): one-shot broadcast from group rank 0 on adapter __init__ so per-rank-initialized weights start bit-identical.
• StageReplicatedColumnParallelLinear / StageReplicatedRowParallelLinear: ColumnParallel/RowParallel linears whose weight gradient is averaged across the intra-PP-stage group in backward (via SharedGradAllReduce). Used with is_expert=True so the underlying TP group is the ETP group; allreduce=False is set automatically and the cross-EP axis is covered by SharedGradAllReduce.
• SharedOuterGroupedExpertAdapter: the new adapter class with mixed shapes — fc1 has 2D shared lora_A (hidden→rank) and 3D per-expert lora_B ([N_local, 2*intermediate, rank]); fc2 has 3D per-expert lora_A and 2D shared lora_B.
• Exported helpers HAVE_TE_COL_GRP_LINEAR / HAVE_TE_ROW_GRP_LINEAR and TE grouped-linear class re-exports so lora.py can dispatch.

src/megatron/bridge/peft/canonical_lora.py — when linear_fc1 is matched but cannot be safely canonicalised (gate/up unfused with grouped-expert), fall through to an unfused ParallelLinearAdapter wrapped in LoRALinear rather than asserting.

src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py (+325 lines) — overrides stream_adapter_weights_megatron_to_hf on KimiK25VLBridge because the base MegatronPeftBridge emission can't handle the mixed 2D/3D shared-outer layout:

• _is_fused_fc1_gate_up: accepts 3D per-expert linear_out ([N_local, 2*intermediate_per_tp, rank]) in addition to the upstream 2D case.
• _gather_expert_adapter_weight / _select_expert_adapter_weight: EP-aware gather/slice that handles both 2D (shared across experts) and 3D (per-expert packed) tensors.
• Mixed-side emission: shared side emits once as [1, …, …] under the .weight0 HF name; per-expert side emits N times under .weight0..weightN-1. Output matches SGLang's expected on-disk LoRA layout for the K2.5 grouped-expert path.

src/megatron/bridge/models/kimi_vl/utils.py — small bug fix: quantize_to_int4 now pins weight_shape to weight.device (previously stayed on CPU and tripped device-mismatch when called from CUDA code).

---

2. Older Megatron-LM training.config backport (commit c41bfdc)

src/megatron/bridge/training/config.py does from megatron.training.config import ... for nine config dataclasses. Upstream Megatron-LM commit 8b00c3ce (2026-03-30) consolidated these into a new megatron.training.config package, but the Megatron-LM bundled inside radixark/miles:dev predates that commit and still has the dataclasses scattered across individual modules. Result: importing megatron.bridge.training.config blows up with ImportError, which transitively blocks every recipe (qwen, gpt-oss, …) on that image.

Fix: wrap the import block in try / except ImportError and fall back to the older locations:

• megatron.training.training_config: CheckpointConfig, LoggerConfig, SchedulerConfig, TrainingConfig, ValidationConfig
• megatron.training.common_config: ProfilingConfig, RNGConfig
• megatron.training.resilience_config: RerunStateMachineConfig, StragglerDetectionConfig

DistributedInitConfig was added in commit 8b00c3ce and has no counterpart on the older layout, so we inline a backport dataclass that mirrors the upstream definition field-for-field (kept in sync with newer Megatron-LM).

Net effect: zero behaviour change on the newer Megatron-LM (the try branch succeeds), and megatron.bridge.training.config becomes importable on the older one.

---

3. MIMO config.role vendored backport (commit 6326bf5)

The same older Megatron-LM has megatron.core.models.mimo but is missing the config.role submodule that newer bridge code imports (MIMO_LANGUAGE_MODULE_KEY). Seven bridge files transitively import it, and even non-MIMO training paths trip on it when importing megatron.bridge.training.config → megatron.bridge.models.megatron_mimo.

Fix: vendor a verbatim copy of role.py from the Megatron-LM submodule at src/megatron/bridge/_compat/mimo_role.py (+167 lines), add the _compat package init explaining its purpose, and wrap each of the seven import sites in try / except ImportError:

try:
    from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
except ImportError:
    from megatron.bridge._compat.mimo_role import MIMO_LANGUAGE_MODULE_KEY

Sites updated:

• src/megatron/bridge/data/megatron_mimo/dp_utils.py
• src/megatron/bridge/models/megatron_mimo/megatron_mimo_config.py
• src/megatron/bridge/models/megatron_mimo/megatron_mimo_ddp.py
• src/megatron/bridge/models/megatron_mimo/megatron_mimo_provider.py
• src/megatron/bridge/training/megatron_mimo_parallel_utils.py
• src/megatron/bridge/training/megatron_mimo_step.py
• src/megatron/bridge/training/train_megatron_mimo.py

Net effect: zero behaviour change when running against a Megatron-LM that already ships mimo.config.role; on older runtimes the vendored copy unblocks the import chain so non-MIMO recipes (qwen, gpt-oss, kimi, …) can proceed without needing a docker image refresh.

---

Relationship to PR #3

PR #3 (nanjiangwill:kimi25 → bridge-rebase-2026-05) is the original Kimi K2.5/K2.6 contribution. This PR carries the same LoRA payload (commit ff33f80) plus the two backport shims (commits c41bfdc, 6326bf5) so the result is runnable on the radixark/miles:dev image.

If #3 lands first, this PR can be re-targeted onto the resulting tip and the diff will shrink to just the two backport commits. If this PR lands first, #3 can be closed or rebased.

---

Compatibility matrix

Megatron-LM commit	training.config package	mimo.config.role	Before this PR	After this PR
Newer (post-8b00c3ce)	✅ present	✅ present	✅ works	✅ works (unchanged)
radixark/miles:dev bundled	❌ scattered modules	❌ missing	❌ ImportError	✅ works via fallback

---

Test plan

• Import sanity: python -c \"import megatron.bridge.training.config\" succeeds on both the upstream NeMo container and radixark/miles:dev.
• Kimi K2.5 LoRA SFT training run with experts_shared_outer_loras=True converges and adapter export produces SGLang-compatible weights.
• Existing non-MoE LoRA recipes (e.g. qwen2.5-0.5B, gpt-oss-20B) are unaffected (experts_shared_outer_loras defaults to False).
• CI on bridge-rebase-2026-05.

Made with Cursor

[Copilot]

Pull request overview

This PR brings Kimi K2.5/K2.6 grouped-expert LoRA support onto bridge-rebase-2026-05 and adds compatibility shims so the bridge can run against older Megatron-LM layouts (notably the radixark/miles:dev image) without requiring a container refresh.

Changes:

• Add “shared-outer” grouped-expert LoRA mode (experts_shared_outer_loras) plus new adapter/linear wrappers and export logic for Kimi K2.5.
• Backport Megatron-LM config imports by falling back to older megatron.training.*_config module locations and inlining a stub DistributedInitConfig when needed.
• Vendor a fallback mimo.config.role implementation and guard MIMO imports to avoid ImportErrors on older Megatron-LM.

Reviewed changes

Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
src/megatron/bridge/peft/lora.py	Adds experts_shared_outer_loras flag and dispatch to shared-outer grouped-expert adapter/wrapper.
src/megatron/bridge/peft/lora_layers.py	Adds LoRAGroupedExpertLinear to plumb tokens_per_expert/splits through to adapters.
src/megatron/bridge/peft/utils.py	Implements shared-outer grouped-expert adapter machinery, stage-replicated gradient sync, and packed per-expert linear.
src/megatron/bridge/peft/canonical_lora.py	Allows linear_fc1 to fall back to an unfused adapter path for unsupported canonicalization cases.
src/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py	Overrides adapter export to HF to support mixed 2D/3D shared-outer grouped-expert LoRA weight layouts.
src/megatron/bridge/models/kimi_vl/utils.py	Fixes quantize_to_int4 device placement for weight_shape.
src/megatron/bridge/training/config.py	Adds try/except import fallback for older Megatron-LM config module layouts; backports DistributedInitConfig.
src/megatron/bridge/_compat/init.py	Introduces _compat package for vendored/compat symbols.
src/megatron/bridge/_compat/mimo_role.py	Vendors missing mimo.config.role (MIMO_LANGUAGE_MODULE_KEY, role/layout dataclasses).
src/megatron/bridge/data/megatron_mimo/dp_utils.py	Guards mimo.config.role import with fallback to vendored _compat module.
src/megatron/bridge/models/megatron_mimo/megatron_mimo_config.py	Guards mimo.config.role import with fallback.
src/megatron/bridge/models/megatron_mimo/megatron_mimo_ddp.py	Guards mimo.config.role import with fallback.
src/megatron/bridge/models/megatron_mimo/megatron_mimo_provider.py	Guards mimo.config.role import with fallback.
src/megatron/bridge/training/megatron_mimo_parallel_utils.py	Guards mimo.config.role import with fallback.
src/megatron/bridge/training/megatron_mimo_step.py	Guards mimo.config.role import with fallback.
src/megatron/bridge/training/train_megatron_mimo.py	Guards mimo.config.role import with fallback.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.