[OMNIML-6063] [18/21] Decompress compressed-tensors MoE experts on the fly as they load

modelopt/torch/models/_remote_code.py

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Loading the modeling code a checkpoint ships with itself."""
+
+import glob
+import os
+import shutil
+import tempfile
+import warnings
+
+
+def stage_full_remote_closure() -> None:
+    """Make transformers stage every module of a local checkpoint's remote code, not just one hop.
+
+    transformers' ``get_cached_module_file`` copies a remote module and its *direct* relative imports
+    into ``transformers_modules/<repo>/<hash>/``, but does not recurse. Remote code more than one import
+    deep -- Kimi-K2.5's ``modeling_kimi_k25 -> modeling_deepseek -> configuration_deepseek`` -- is
+    then staged incomplete, and the later recursive walk in ``get_relative_import_files`` fails with
+    ``FileNotFoundError``, from ``get_class_from_dynamic_module`` and ``from_pretrained`` alike.
+
+    Wrap ``get_cached_module_file`` so that, loading from a local checkpoint dir, it copies every
+    sibling ``*.py`` into the stage dir it just populated. Idempotent.
+    """
+    import transformers.dynamic_module_utils as dmu
+    from transformers.utils import HF_MODULES_CACHE
+
+    if getattr(dmu.get_cached_module_file, "_modelopt_full_closure", False):
+        return
+    orig = dmu.get_cached_module_file
+
+    def _staged(pretrained_model_name_or_path, module_file, *args, **kwargs):
+        resolved = orig(pretrained_model_name_or_path, module_file, *args, **kwargs)
+        src = str(pretrained_model_name_or_path)
+        if os.path.isdir(src):
+            staged = str(resolved)
+            if not os.path.isabs(staged):
+                staged = os.path.join(HF_MODULES_CACHE, staged)
+            stage_dir = os.path.dirname(staged)
+            try:
+                for py in glob.glob(os.path.join(src, "*.py")):
+                    dst = os.path.join(stage_dir, os.path.basename(py))
+                    if not os.path.exists(dst):
+                        # Every rank loading the checkpoint stages into this dir and imports from it:
+                        # renamed into place (atomic), a file appears there whole, never half written.
+                        fd, tmp = tempfile.mkstemp(dir=stage_dir)
+                        os.close(fd)
+                        shutil.copy(py, tmp)
+                        os.replace(tmp, dst)
+            except OSError as e:  # a failed copy only matters if a deeper import needs the file
+                warnings.warn(f"Could not stage the remote code of {src}: {e}")
+        return resolved
+
+    setattr(_staged, "_modelopt_full_closure", True)
+    dmu.get_cached_module_file = _staged

modelopt/torch/utils/plugins/model_load_utils.py

@@ -2100,6 +2100,13 @@ def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int:
 _EP_SOURCE_FORMATS = frozenset({None, "fp8", "mxfp4"})
 
 
+def _decompressed_as_loaded(cfg) -> bool:
+    """Whether a model spec installed the on-the-fly decompressor for this config's checkpoint."""
+    from modelopt.torch.models import _onfly_decompress
+
+    return _onfly_decompress.is_marked(cfg)
+
+
 def _ckpt_quant_method(cfg):
     """Return the checkpoint's native quantization method, or ``None``.
 
@@ -2367,7 +2374,7 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A
             f"{plan!r}: the dense part of a block-fp8 checkpoint cannot be sharded: Distributed "
             "Checkpoint reads it as stored, with no dequantization. Plan dense='replicate'."
         )
-    if quant_method not in _EP_SOURCE_FORMATS:
+    if quant_method not in _EP_SOURCE_FORMATS and not _decompressed_as_loaded(hf_config):
         # transformers' own quantizer would load it as stored, for PTQ to treat as plain weights.
         dequantized = ", ".join(sorted(f for f in _EP_SOURCE_FORMATS if f))
         raise NotImplementedError(
@@ -2377,6 +2384,20 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A
         )
     # The swapped plan is this load's own: the caller's config keeps the plan its checkpoint declares.
     hf_config = copy.deepcopy(hf_config)
+    if _decompressed_as_loaded(hf_config):
+        # The decompressor hands the experts over dequantized. Left on the config (Kimi-K2.5's config
+        # class copies its text config's to the top level), the checkpoint's quantization_config
+        # would build transformers' compressed-tensors quantizer, which from 5.13 claims the same
+        # expert converters first. Deleted rather than set to None: before 5.11 transformers takes
+        # any quantization_config attribute, None too, for a quantized checkpoint. Every level
+        # _ckpt_quant_method reads is cleared.
+        for cfg in (
+            hf_config,
+            getattr(hf_config, "text_config", None),
+            getattr(hf_config, "llm_config", None),
+        ):
+            if cfg is not None and hasattr(cfg, "quantization_config"):
+                del cfg.quantization_config
     ep_cfg = _ep_config(hf_config)  # for VLMs the plan lives on .text_config
     ep_plan = dict(getattr(ep_cfg, "base_model_ep_plan", None) or {})
     if not ep_plan:
@@ -2400,6 +2421,12 @@ class builds one; ``require_lm_head=False`` builds a VLM without it, through ``A
     # from, take the local snapshot: a Hub ID recorded there would leave the export nothing to read.
     # Resolved once the plan and the config pass, so a refused load downloads nothing.
     ckpt_path = _resolve_checkpoint_dir(ckpt_path, torch.distributed.get_rank())
+    if _decompressed_as_loaded(hf_config):
+        # Only the experts decompress as they load, and the config cannot say what else the
+        # checkpoint compressed, or in which dtypes: what it stores is checked before anything loads.
+        from modelopt.torch.models import _onfly_decompress
+
+        _onfly_decompress.check_compressed_keys(hf_config, ckpt_path)
 
     # Two INDEPENDENT meshes (NeMo-AutoModel style), named by the model part they govern -- two
     # factorizations of the same world:

noxfile.py

@@ -68,7 +68,12 @@ def _cov_args():
 @nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS])
 def unit(session, torch_ver, tf_ver):
     """Unit tests — parametrized over torch and transformers versions."""
-    session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]")
+    # compressed-tensors runs the real-library tests of on-the-fly expert decompression. It needs
+    # torch>=2.10, so the older legs skip those tests; it goes in with the torchvision pin so the
+    # resolver keeps the leg's torch. Not in dev-test: the transformers-free partial_unit(torch)
+    # lane installs that extra, and compressed-tensors pulls in transformers.
+    extra = [] if torch_ver in ("torch_28", "torch_29") else ["compressed-tensors"]
+    session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]", *extra)
     tf_pins = TRANSFORMERS_VERSIONS[tf_ver]
     if tf_pins:
         session.install(*tf_pins)