fix(deps): AUT-1177 finalize merged MCore dev runtime

.dev.commit

@@ -1 +1 @@
-af91b6b5cc5aec01458b8e6d24ccb0f4717e49e3
+de1ebd63c7935886150364f9fc1a4150c7bdb37c

docker/Dockerfile.ci

@@ -117,13 +117,17 @@ ARG UV_CACHE_PRUNE_ARGS
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     --mount=type=cache,target=/root/.cache/uv \
-    export MAMBA_FORCE_BUILD=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE && \
+    export MAMBA_FORCE_BUILD=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE \
+        FAST_HADAMARD_TRANSFORM_FORCE_BUILD=TRUE && \
     # flash-mla (Megatron-LM's no_pypi_wheels group) has no PyPI wheel and is built from source by
     # the uv sync below. Point the compilers at the CCCL/libcu++ headers (under cccl/ in this base
     # image, not the default CUDA include path) and scope the build to the target archs.
     export FLASH_MLA_DISABLE_SM90=1 NVCC_THREADS=16 \
         CFLAGS="-I/usr/local/cuda/include/cccl${CFLAGS:+ $CFLAGS} -DNDEBUG" \
         CXXFLAGS="-I/usr/local/cuda/include/cccl${CXXFLAGS:+ $CXXFLAGS} -DNDEBUG" && \
+    # Do not reuse a wheel previously downloaded by the dependency's setup.py. The force-build
+    # variable is not part of uv's cache key, so evict the package before rebuilding locked source.
+    uv cache clean fast-hadamard-transform && \
     # Reinstall nvidia-cutlass-dsl in system site packages
     # uv install can run into non-deterministic install issues
     # Keep this pin compatible with cuDNN Frontend; mismatched versions can break

docker/common/install.sh

@@ -82,6 +82,17 @@ main() {
     unset PIP_CONSTRAINT
 
     if [[ "$USE_UV" == "true" ]]; then
+        # MCore dev requires TransformerEngine as part of its dev extra. TE's bundled NCCL-EP
+        # extension needs newer NCCL device headers than the NGC 25.05 installation-test image
+        # provides. Disable only that optional extension when the gitlink is the recorded dev ref.
+        submodule_sha="$(git -C 3rdparty/Megatron-LM rev-parse HEAD 2>/dev/null || true)"
+        dev_sha="$(tr -d '[:space:]' <.dev.commit)"
+        main_sha="$(tr -d '[:space:]' <.main.commit)"
+        if [[ -n "$submodule_sha" && "$submodule_sha" == "$dev_sha" && "$dev_sha" != "$main_sha" ]]; then
+            export NVTE_WITH_NCCL_EP="${NVTE_WITH_NCCL_EP:-0}"
+            echo "🔧 MCore dev compatibility: NVTE_WITH_NCCL_EP=$NVTE_WITH_NCCL_EP"
+        fi
+
         if [[ "$BASE_IMAGE" == "pytorch" ]]; then
             UV_ARGS=(
                 "--no-install-package" "torch"

pyproject.toml

@@ -252,10 +252,10 @@ build = [
 # them at test time with scripts/install_diffusion_deps.sh. This group is kept for
 # documentation only; `uv sync --group diffusion` will NOT install them on Linux.
 diffusion = ["imageio", "imageio-ffmpeg", "av"]
-# flash-mla has no PyPI wheel and is built from source by uv (no-build-isolation) during the
-# image's `uv sync --all-groups`. Megatron-LM declares it in the same group, but uv does not
-# install a path dependency's groups, so it is declared here too.
-no_pypi_wheels = ["flash_mla"]
+# flash-mla and fast-hadamard-transform have no usable PyPI wheels and are built from source by uv
+# (no-build-isolation) during the image's `uv sync --all-groups`. Megatron-LM declares them in the
+# same group, but uv does not install a path dependency's groups, so they are declared here too.
+no_pypi_wheels = ["flash_mla", "fast-hadamard-transform"]
 
 [project.entry-points."nemo_run.cli"]
 lm = "megatron.bridge"

src/megatron/bridge/models/distillation_provider.py

@@ -13,13 +13,16 @@
 # limitations under the License.
 
 import logging
+from copy import deepcopy
 from dataclasses import dataclass, fields
 from typing import TYPE_CHECKING, Any, Optional
 
 import modelopt.torch.distill as mtd
 import modelopt.torch.distill.plugins.megatron as mtd_mcore
 from megatron.core.models.common.language_module.language_module import LanguageModule
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.utils import unwrap_model
+from megatron.training.models.base import ModelConfig
 
 from megatron.bridge.models.gpt_provider import GPTModelProvider
 from megatron.bridge.models.hybrid.hybrid_provider import HybridModelProvider
@@ -33,6 +36,110 @@
 logger = logging.getLogger(__name__)
 
 
+def _clone_model_config(config: ModelConfig) -> ModelConfig:
+    """Clone a model config without resolving proxied ``__deepcopy__`` methods."""
+    config_copy = type(config).__new__(type(config))
+    memo = {id(config): config_copy}
+    for name, value in config.__dict__.items():
+        object.__setattr__(config_copy, name, deepcopy(value, memo))
+    return config_copy
+
+
+class DistillationModelConfig:
+    """Builder-backed model config for knowledge distillation.
+
+    Instances retain the student's concrete config class and builder while
+    registering a pre-wrap hook that builds the teacher from its own config and
+    applies ModelOpt distillation after both models are constructed.
+    """
+
+    _SHARED_ATTRIBUTES = (
+        "tensor_model_parallel_size",
+        "pipeline_model_parallel_size",
+        "context_parallel_size",
+        "seq_length",
+        "pipeline_dtype",
+    )
+
+    def _initialize_distillation(
+        self,
+        teacher: ModelConfig,
+        kd_config: Optional["ModelOptDistillConfig"],
+        distill_submodule: Optional[str],
+    ) -> None:
+        """Attach teacher state and validate the shared distributed settings."""
+        object.__setattr__(self, "teacher", teacher)
+        object.__setattr__(self, "kd_config", kd_config)
+        object.__setattr__(self, "distill_submodule", distill_submodule)
+        object.__setattr__(self, "_mirror_to_teacher", True)
+
+        self.cross_entropy_fusion_impl = "native"
+
+    def as_dict(self) -> dict[str, Any]:
+        """Serialize as the original student config for checkpoint restoration."""
+        result = super().as_dict()
+        student_class = self._student_config_class
+        result["_target_"] = f"{student_class.__module__}.{student_class.__qualname__}"
+        return result
+
+    def finalize(self) -> None:
+        """Finalize student and teacher configs without mirroring derived values."""
+        object.__setattr__(self, "_mirror_to_teacher", False)
+        try:
+            super().finalize()
+            if hasattr(self.teacher, "finalize"):
+                self.teacher.finalize()
+        finally:
+            object.__setattr__(self, "_mirror_to_teacher", True)
+        for attribute_name in self._SHARED_ATTRIBUTES:
+            if getattr(self, attribute_name) != getattr(self.teacher, attribute_name):
+                raise ValueError(f"Student and teacher configs must have the same {attribute_name}.")
+
+    def _convert_hook(self, model_chunks: list) -> list:
+        """Build the teacher and convert the student after model construction."""
+        assert len(model_chunks) == 1, "ModelOpt KD does not support virtual pipeline (>1 model chunk)."
+        student_model = unwrap_model(model_chunks[0])
+
+        teacher_builder = self.teacher.get_builder_cls()(self.teacher)
+        teacher_chunks = teacher_builder.build_distributed_models(
+            pg_collection=ProcessGroupCollection.use_mpu_process_groups(),
+            wrap_with_ddp=False,
+            mixed_precision_wrapper=None,
+        )
+        assert len(teacher_chunks) == 1, "ModelOpt KD does not support virtual pipeline (>1 model chunk)."
+        teacher_model = unwrap_model(teacher_chunks[0])
+
+        if self.distill_submodule is not None:
+            self.full_model = student_model
+            student_model = getattr(student_model, self.distill_submodule)
+            teacher_model = getattr(teacher_model, self.distill_submodule)
+
+        kd_cfg = mtd_mcore.setup_distillation_config(self.kd_config, student_model.config, teacher_model.config)
+        modelopt_cfg = {
+            "teacher_model": teacher_model,
+            "criterion": kd_cfg.criterion,
+            "loss_balancer": kd_cfg.loss_balancer,
+        }
+        kd_model = mtd.convert(student_model, mode=[("kd_loss", modelopt_cfg)])
+        if self.distill_submodule is not None:
+            assert getattr(self.full_model, self.distill_submodule) is kd_model
+            model_chunks[0] = self.full_model
+        else:
+            model_chunks[0] = kd_model
+        mtd_mcore.adjust_distillation_model_for_mcore(kd_model, kd_cfg)
+        return model_chunks
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name in {"teacher", "kd_config", "distill_submodule", "full_model", "_mirror_to_teacher"}:
+            object.__setattr__(self, name, value)
+            return
+
+        super().__setattr__(name, value)
+        teacher = getattr(self, "teacher", None)
+        if self._mirror_to_teacher and teacher is not None and hasattr(teacher, name):
+            setattr(teacher, name, value)
+
+
 @dataclass
 class DistillationProvider(TransformerConfig):
     """Provider for Bridge language models in distillation mode.
@@ -160,12 +267,12 @@ def __setattr__(self, name, value):
 
 
 def convert_to_distillation_provider(
-    student_provider: GPTModelProvider | HybridModelProvider,
-    teacher_provider: GPTModelProvider | HybridModelProvider,
+    student_provider: GPTModelProvider | HybridModelProvider | ModelConfig,
+    teacher_provider: GPTModelProvider | HybridModelProvider | ModelConfig,
     kd_config: Optional["ModelOptDistillConfig"] = None,
     *,
     distill_submodule: Optional[str] = None,
-) -> "DistillationProvider":
+) -> "DistillationProvider | DistillationModelConfig":
     """Convert a given model provider to a DistillationProvider.
 
     The KD conversion runs in a pre-wrap hook (after the student's weights are loaded), not in
@@ -182,8 +289,24 @@ def convert_to_distillation_provider(
             submodule is trained).
     """
 
+    if isinstance(student_provider, ModelConfig):
+        assert isinstance(teacher_provider, ModelConfig), "Teacher config must be a ModelConfig."
+        student_class = type(student_provider)
+        student_copy = _clone_model_config(student_provider)
+        distillation_class = type(
+            f"Distillation{student_class.__name__}",
+            (DistillationModelConfig, student_class),
+            {"__module__": __name__},
+        )
+        student_copy.__class__ = distillation_class
+        object.__setattr__(student_copy, "_student_config_class", student_class)
+        student_copy._initialize_distillation(teacher_provider, kd_config, distill_submodule)
+        hooks = [*student_copy.pre_wrap_hooks, student_copy._convert_hook]
+        object.__setattr__(student_copy, "pre_wrap_hooks", hooks)
+        return student_copy
+
     assert isinstance(student_provider, (GPTModelProvider, HybridModelProvider)), (
-        "Student provider must be a subclass of GPTModelProvider or HybridModelProvider."
+        "Student provider must be a subclass of GPTModelProvider, HybridModelProvider, or ModelConfig."
     )
     assert isinstance(teacher_provider, (GPTModelProvider, HybridModelProvider)), (
         "Teacher provider must be a subclass of GPTModelProvider or HybridModelProvider."

src/megatron/bridge/training/checkpointing.py

@@ -346,6 +346,16 @@ def read_metadata(tracker_filename: str) -> tuple[int, bool]:
     return max_iter, release
 
 
+def _get_model_parallel_sizes(model_config: dict[str, Any]) -> tuple[int, int]:
+    """Read tensor and pipeline parallel sizes from provider or builder configs."""
+    transformer_config = model_config.get("transformer")
+    parallel_config = transformer_config if isinstance(transformer_config, dict) else model_config
+    return (
+        parallel_config["tensor_model_parallel_size"],
+        parallel_config["pipeline_model_parallel_size"],
+    )
+
+
 def _extract_megatron_lm_args_from_state_dict(state_dict: dict[str, Any]) -> dict[str, Any]:
     """Extract and convert legacy Megatron-LM args from checkpoint state_dict to Megatron-Bridge config format.
 
@@ -2749,10 +2759,7 @@ def _load_checkpoint_from_path(
             tp_pp_match = True
             mismatch_msg = ""
         else:
-            ckpt_tp_pp = (
-                run_config["model"]["tensor_model_parallel_size"],
-                run_config["model"]["pipeline_model_parallel_size"],
-            )
+            ckpt_tp_pp = _get_model_parallel_sizes(run_config["model"])
             run_tp_pp = (
                 cfg.model.tensor_model_parallel_size,
                 cfg.model.pipeline_model_parallel_size,
@@ -2869,10 +2876,7 @@ def _load_checkpoint_from_path(
             run_config_filename = get_checkpoint_run_config_filename(checkpoint_name)
             if file_exists(run_config_filename):
                 run_config = read_run_config(run_config_filename)
-                ckpt_tp_pp = (
-                    run_config["model"]["tensor_model_parallel_size"],
-                    run_config["model"]["pipeline_model_parallel_size"],
-                )
+                ckpt_tp_pp = _get_model_parallel_sizes(run_config["model"])
                 run_tp_pp = (
                     cfg.model.tensor_model_parallel_size,
                     cfg.model.pipeline_model_parallel_size,

src/megatron/bridge/training/distill.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from megatron.bridge.models.distillation_provider import DistillationProvider
+from megatron.bridge.models.distillation_provider import DistillationModelConfig, DistillationProvider
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.gpt_step import forward_step_modelopt
 from megatron.bridge.training.pretrain import pretrain
@@ -32,6 +32,8 @@ def distill(
         This is an experimental API and is subject to change in backwards
         incompatible ways without notice.
     """
-    assert isinstance(config.model, DistillationProvider), "Distillation requires a DistillationProvider"
+    assert isinstance(config.model, (DistillationProvider, DistillationModelConfig)), (
+        "Distillation requires a DistillationProvider or DistillationModelConfig"
+    )
 
     return pretrain(config, forward_step_modelopt)

tests/functional_tests/test_groups/recipes/test_llama_recipes_distill_3b-1b.py

@@ -101,7 +101,6 @@ def run_distill_recipe_test(
         config.checkpoint.load = str(checkpoint_dir)
         config.logger.tensorboard_dir = str(tensorboard_dir)
 
-        # Combine into a distillation provider
         config.model = convert_to_distillation_provider(config.model, teacher_config.model)
 
         # Set default distillation configuration