refactor(export): declare the module's dependencies at the top

layerwise_export.py had 17 function-local imports, none of them cycle-driven:
nothing under export/ imports this module, and export/__init__ does not either,
so the deferral bought nothing that CONTRIBUTING's module-top rule asks us to
justify. Verified by importing the module from both entry orders and running the
export suites.

The one that stays local is plugins.huggingface, which pulls in transformers --
unified_export_hf, unified_export_hf_streaming and moe_utils all defer that same
import, so this now matches its neighbours.

Separately, _reconcile_export_with_resume moves to utils/layerwise_calib.py,
next to the manifest reader and detect_resume_point whose contract it encodes,
keeping model_calib.py to the algorithms.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

modelopt/torch/export/layerwise_export.py

@@ -22,10 +22,34 @@
 
 import torch
 import torch.nn as nn
+from safetensors import safe_open
 from safetensors.torch import save_file
 
+from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.utils.core_utils import (
+    enable_weight_access_and_writeback,
+    requires_weight_materialization,
+)
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+from modelopt.torch.utils import distributed as dist
+
+from .layer_utils import is_moe, sync_moe_gate_up_amax
 from .model_config import FUSION_FREE_FORMATS, QUANTIZATION_NVFP4
-from .quant_utils import get_quant_config, get_quantization_format
+from .model_utils import TiedWeightMap
+from .quant_aware_conversion import build_reverse_name_mapper, revert_quant_config_names
+from .quant_utils import _postprocess_single_tensor, get_quant_config, get_quantization_format
+from .registry import ExportContext, PrepareMoEInputsRegistry
+from .unified_export_hf import (
+    _add_mtp_exclusions,
+    _dispatch_export_handler,
+    _fuse_shared_input_modules,
+    _prepare_moe_inputs,
+    _resolve_export_dtype,
+    _write_hf_export_config,
+    collect_shared_input_modules,
+    save_non_weight_artifacts,
+)
+from .unified_export_hf_streaming import _assert_no_split_rules
 
 # Fusable per layer because the groups (q/k/v, gate/up) never cross a decoder boundary.
 # AWQ and SVDQuant also need pre-quant-scale steps, which are still whole-model.
@@ -44,8 +68,6 @@ def layer_shard_name(layer_idx: int) -> str:
 
 def _is_quantized_module(module: nn.Module) -> bool:
     """By type, not name: fused experts name theirs ``gate_up_proj_weight_quantizer``."""
-    from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer
-
     return any(
         isinstance(child, (TensorQuantizer, SequentialQuantizer)) for child in module.children()
     )
@@ -67,8 +89,6 @@ def _tied_quantized_modules(model: nn.Module) -> list[str]:
     weights are on meta and would pass vacuously. Falls back to ``data_ptr`` when the model
     publishes no map (transformers < 5).
     """
-    from .model_utils import TiedWeightMap
-
     tied_map = TiedWeightMap(model)
     groups: dict[str, list[str]] = {}
     by_ptr: dict[int, list[str]] = {}
@@ -105,8 +125,6 @@ def assert_formats_supported(module: nn.Module, scope: str) -> None:
 
 def assert_layerwise_export_supported(model: nn.Module) -> None:
     """Raise unless per-layer export is valid for this model."""
-    from modelopt.torch.utils import distributed as dist
-
     assert_formats_supported(model, "before calibration")
 
     tied = _tied_quantized_modules(model)
@@ -145,14 +163,6 @@ def __init__(
 
         Runs before calibration, so nothing amax-dependent exists yet.
         """
-        from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
-
-        from .layer_utils import is_moe
-        from .quant_aware_conversion import build_reverse_name_mapper
-        from .registry import ExportContext, PrepareMoEInputsRegistry
-        from .unified_export_hf import _resolve_export_dtype
-        from .unified_export_hf_streaming import _assert_no_split_rules
-
         assert_layerwise_export_supported(model)
         # Splits regroup tensors across the whole state dict; no per-layer pass reverses that.
         _assert_no_split_rules(model)
@@ -214,10 +224,9 @@ def export_layer(
         a fusing format can rediscover which modules share an input; omit them only when
         nothing fuses.
         """
+        # Local, as in every other export module: the plugin imports transformers.
         from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear
 
-        from .unified_export_hf import _dispatch_export_handler, _prepare_moe_inputs
-
         assert not self._finalized, "export_layer() called after finalize()"
         if layer_module is not self._layers[layer_idx]:
             # Not an assert: -O would strip it, and the failure is silent -- layer N's
@@ -257,8 +266,6 @@ def _unify_shared_quantization_params(
         shared-input group, one weight_scale_2 per expert gate/up pair. Its pre-quant-scale
         steps are AWQ/SVDQuant-only and refused.
         """
-        from .layer_utils import sync_moe_gate_up_amax
-
         # A set, not get_quantization_format: that stops at the first hit, so a mixed
         # FP8-attention/NVFP4-expert layer would report fp8 and skip fusing entirely.
         if _module_formats(layer_module) - FUSION_FREE_FORMATS:
@@ -267,9 +274,6 @@ def _unify_shared_quantization_params(
 
     def _fuse_shared_input_scales(self, layer_module: nn.Module, layer_inputs: list | None) -> None:
         """Rediscover the groups that share an input, on real activations, and fuse them."""
-        from .quant_utils import get_quantization_format
-        from .unified_export_hf import _fuse_shared_input_modules, collect_shared_input_modules
-
         layer_format = get_quantization_format(layer_module)
         if not layer_inputs:
             raise RuntimeError(
@@ -290,19 +294,6 @@ def finalize(self) -> dict:
 
         Leaves ``export_dir`` a complete checkpoint; no ``export_hf_checkpoint()`` needed.
         """
-        from modelopt.torch.quantization.utils.core_utils import (
-            enable_weight_access_and_writeback,
-            requires_weight_materialization,
-        )
-
-        from .quant_aware_conversion import revert_quant_config_names
-        from .unified_export_hf import (
-            _add_mtp_exclusions,
-            _dispatch_export_handler,
-            _write_hf_export_config,
-            save_non_weight_artifacts,
-        )
-
         assert not self._finalized, "finalize() called twice"
         self._finalized = True
 
@@ -407,8 +398,6 @@ def assert_shards_present(self, upto: int) -> None:
 
     def _collect(self, out: dict[str, torch.Tensor], full_key: str, tensor: torch.Tensor) -> None:
         """Apply per-tensor export postprocessing and hub-name reversal, or drop the tensor."""
-        from .quant_utils import _postprocess_single_tensor
-
         if tensor is None or tensor.is_meta:
             return
         new_key, new_value = _postprocess_single_tensor(
@@ -426,8 +415,6 @@ def _write_index(self) -> None:
         From disk because a resumed run never saw the earlier shards in memory; by layer
         count rather than a glob, so a longer previous run's leftovers cannot leak in.
         """
-        from safetensors import safe_open
-
         # Out of the index already; delete them so the directory *is* the checkpoint.
         for stale in self._export_dir.glob("model-layer-*.safetensors"):
             if int(stale.stem.rsplit("-", 1)[1]) >= len(self._layers):

modelopt/torch/quantization/model_calib.py

@@ -34,7 +34,7 @@
 from modelopt.torch.quantization.utils.layerwise_calib import (
     LayerActivationCollector,
     _CheckpointState,
-    _read_manifest,
+    _reconcile_export_with_resume,
 )
 from modelopt.torch.utils import print_rank_0, warn_rank_0
 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master
@@ -2048,37 +2048,6 @@ def postprocess(module, name):
     max_calibrate(model, forward_loop)
 
 
-def _reconcile_export_with_resume(
-    exporter, checkpoint_dir: str | None, start_layer: int, num_layers: int
-) -> bool:
-    """Reconcile the shards on disk with the layer calibration will start from.
-
-    Returns True when every layer already has a shard, so only ``finalize()`` is owed.
-    """
-    manifest = _read_manifest(checkpoint_dir) if checkpoint_dir is not None else None
-    last = (manifest or {}).get("last_completed_layer")
-    total = (manifest or {}).get("num_layers")
-    if total is not None and total != num_layers:
-        raise ValueError(
-            f"Layerwise checkpoint at {checkpoint_dir} was written for {total} layers "
-            f"but this model has {num_layers}. Use a fresh checkpoint_dir."
-        )
-    # detect_resume_point returns None once the manifest is complete, which puts start_layer
-    # back at 0 and would recalibrate everything the shards already hold.
-    if last is not None and last + 1 >= num_layers:
-        exporter.assert_shards_present(num_layers)
-        return True
-
-    if start_layer > 0:
-        exporter.assert_shards_present(start_layer)
-    elif checkpoint_dir is not None:
-        # Starting at 0 with a checkpoint_dir means no usable resume record, so calibration
-        # would silently overwrite finished shards. Without one there is no resume to lose,
-        # and re-exporting is the documented behaviour.
-        exporter.assert_no_orphan_shards()
-    return False
-
-
 @torch.no_grad()
 def layerwise_calibrate(
     model: nn.Module,

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -566,6 +566,37 @@ def detect_resume_point(checkpoint_dir: str) -> tuple[int, dict] | None:
     return (last + 1, manifest)
 
 
+def _reconcile_export_with_resume(
+    exporter, checkpoint_dir: str | None, start_layer: int, num_layers: int
+) -> bool:
+    """Reconcile the shards on disk with the layer calibration will start from.
+
+    Returns True when every layer already has a shard, so only ``finalize()`` is owed.
+    """
+    manifest = _read_manifest(checkpoint_dir) if checkpoint_dir is not None else None
+    last = (manifest or {}).get("last_completed_layer")
+    total = (manifest or {}).get("num_layers")
+    if total is not None and total != num_layers:
+        raise ValueError(
+            f"Layerwise checkpoint at {checkpoint_dir} was written for {total} layers "
+            f"but this model has {num_layers}. Use a fresh checkpoint_dir."
+        )
+    # detect_resume_point returns None once the manifest is complete, which puts start_layer
+    # back at 0 and would recalibrate everything the shards already hold.
+    if last is not None and last + 1 >= num_layers:
+        exporter.assert_shards_present(num_layers)
+        return True
+
+    if start_layer > 0:
+        exporter.assert_shards_present(start_layer)
+    elif checkpoint_dir is not None:
+        # Starting at 0 with a checkpoint_dir means no usable resume record, so calibration
+        # would silently overwrite finished shards. Without one there is no resume to lose,
+        # and re-exporting is the documented behaviour.
+        exporter.assert_no_orphan_shards()
+    return False
+
+
 class _CheckpointState:
     """Manages checkpoint save and restore for layerwise calibration.