backport: compat shims for older Megatron-LM (miles-main / radixark:miles dev image)

[yushengsu-thu]

Summary

Three backport commits that let megatron.bridge load against the older Megatron-LM bundled in the radixark/miles:dev docker image (the miles-main branch of radixark/Megatron-LM, which forked from NVIDIA upstream before several recent reorgs).

Without these shims, from megatron.bridge.models import ... and from megatron.bridge.training.config import ... fail at module-load on every Bridge consumer (gpt-oss-20b, Qwen, Kimi LoRA training), blocking training entirely on that image.

All three commits follow the same minimal-blast-radius strategy: patch the root-cause module-level import (try newer layout first, fall back to older layout / vendored copy / None sentinel) so non-affected code paths keep working unchanged.

Commits

1. 5efe6efa — backport: try old Megatron-LM import paths when megatron.training.config missing

Wrap the nine `from megatron.training.config import ...` lines in a try/except. If the newer Megatron-LM layout (config dataclasses consolidated into the `megatron.training.config` package, introduced in commit `8b00c3ce` on 2026-03-30) is unavailable, fall back to the individual-file 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 the same migration commit and has no counterpart in the older layout, so we backport an inline stub mirroring the upstream definition.

2. 55beab8f — backport: vendor megatron.core.models.mimo.config.role

Older Megatron-LM (e.g. `radixark/miles:dev` image) has `megatron.core.models.mimo` but is missing the `config.role` submodule. The seven bridge files that import `MIMO_LANGUAGE_MODULE_KEY` from it now wrap the import in try/except and fall back to a vendored copy at `megatron.bridge._compat.mimo_role` (verbatim copy of `role.py` from the `Megatron-Bridge/3rdparty/Megatron-LM` submodule).

This unblocks loading `megatron.bridge.training.config` — and downstream `bridge.models.megatron_mimo` — against the older Megatron-LM, so non-MIMO training paths (qwen, etc.) can proceed without needing a docker image refresh.

3. 9ca0ca8a — backport: try/except parse_hybrid_pattern import when miles-main Megatron-LM lacks it

`from megatron.core.ssm.mamba_hybrid_layer_allocation import parse_hybrid_pattern` in `mamba_provider.py` fails on the `miles-main` branch of `radixark/Megatron-LM` (used in `radixark/miles:dev` docker image), because miles-main branched off NVIDIA upstream before MTP commit `300d1b655`, which is when `parse_hybrid_pattern` was added to `megatron.core.ssm.mamba_hybrid_layer_allocation`.

The failure blocks `from megatron.bridge.models import ...` at module-load on every Bridge consumer (gpt-oss-20b, Qwen, Kimi LoRA training), because `models/init.py:117` imports `nemotron_vl`, which then imports `mamba_provider`, which fails at its module-level import.

`parse_hybrid_pattern` is only used inside method bodies for hybrid/Mamba layer parsing; non-Mamba consumers (gpt-oss MoE, Qwen, Kimi) never hit those code paths. So degrade `parse_hybrid_pattern` to `None` on `ImportError` — module loads, non-Mamba paths keep working, Mamba consumers get a clear `'NoneType' is not callable` at call-site instead.

Same pattern as `c41bfdc7` (`megatron.training.config`) and `6326bf5b` (`mimo.config.role`): patch root-cause module-level imports rather than every downstream consumer.

Files changed

```
src/megatron/bridge/_compat/init.py | 6 +
src/megatron/bridge/_compat/mimo_role.py | 167 +++++++++++++++++++++
src/megatron/bridge/data/megatron_mimo/dp_utils.py | 7 +-
src/megatron/bridge/models/mamba/mamba_provider.py | 12 +-
.../models/megatron_mimo/megatron_mimo_config.py | 7 +-
.../models/megatron_mimo/megatron_mimo_ddp.py | 7 +-
.../models/megatron_mimo/megatron_mimo_provider.py | 7 +-
src/megatron/bridge/training/config.py | 71 +++++++--
.../training/megatron_mimo_parallel_utils.py | 7 +-
src/megatron/bridge/training/megatron_mimo_step.py | 7 +-
.../bridge/training/train_megatron_mimo.py | 7 +-
11 files changed, 288 insertions(+), 17 deletions(-)
```

Test plan

• On radixark/miles:dev image: python -c "from megatron.bridge.models import GPTModelProvider" succeeds (previously failed at module-load).
• On radixark/miles:dev image: python -c "from megatron.bridge.training.config import ConfigContainer" succeeds.
• On a recent NVIDIA Megatron-LM checkout: imports above still resolve to the upstream symbols (try-branch path, not the fallback).
• Run Kimi LoRA training entry point end-to-end on radixark/miles:dev.
• Run gpt-oss-20b / Qwen training entry points end-to-end on radixark/miles:dev.

Made with Cursor