[OMNIML-6054] [2/2] Route load_pretrained by the distribution plan

CHANGELOG.rst

@@ -43,6 +43,7 @@ Changelog
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.
+- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, FSDP2-sharded or whole on every rank, by the plan ``resolve_distribution_plan`` decides or one passed as its new ``distribution`` argument.
 
 **Backward Breaking Changes**
 
@@ -51,6 +52,9 @@ Changelog
 - Bump minimum required PyTorch version to 2.10.
 - Bump minimum required ``diffusers`` version to 0.37, ``peft`` to 0.18.1, ``accelerate`` to 1.6 and ``nltk`` to 3.9 for the ``[hf]`` extra.
 - Drop support for Transformer Engine 1.x. Quantizing Transformer Engine modules now requires Transformer Engine 2.0 or later.
+- ``load_pretrained`` called in a process group of several ranks now decides how to distribute the model rather than load it whole on every rank: a model too big for one GPU is FSDP2-sharded, and the sharded load refuses placement arguments, a ``dtype``, a quantized checkpoint, one carrying an MTP head, a model class it does not build, and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, which ``MODELOPT_EP_DEGREE`` cannot ask for until the expert-parallel load lands. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(<ranks>, dense="replicate")``, to keep loading it whole on every rank; the guide "Loading Hugging Face models across ranks" has the details.
+- ``load_pretrained`` loading a model whole on several ranks now places it on each rank's own GPU unless ``device_map`` is passed.
+- ``load_pretrained`` now records how it distributed the model on the model it returns, even in a single process, so ``fsdp2_wrap`` refuses that model; to shard a model yourself, wrap one loaded with ``from_pretrained``.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
 - ``hf_quant_config.json`` of an INT4 or W4A8 AWQ checkpoint now sets ``pre_quant_scale`` per layer, only for layers that have the tensor. A model whose fused-MoE experts have none (AWQ never applies to them) is now exported as ``MIXED_PRECISION`` with per-layer entries instead of a uniform algorithm.
 - The Windows FVD example rejects I3D checkpoints requiring unrestricted pickle loading. Custom TorchScript archives now require ``--weights-sha256`` with a trusted checksum; tensor state dictionaries continue to load directly.

docs/source/guides/1_quantization.rst

@@ -19,6 +19,7 @@ Below, you can find the documentation for the quantization toolkit in ModelOpt:
     ./_basic_quantization.rst
     ./_choosing_quant_methods.rst
     ./_pytorch_quantization.rst
+    ./_distributed_loading.rst
     ./_quant_cfg.rst
     ./_customized_model_quantization.rst
     ./_compress_quantized_models.rst

docs/source/guides/_distributed_loading.rst

@@ -0,0 +1,119 @@
+========================================
+Loading Hugging Face models across ranks
+========================================
+
+:func:`load_pretrained <modelopt.torch.utils.plugins.model_load_utils.load_pretrained>` loads a
+Hugging Face model for ModelOpt to optimize and export, whole on every rank or distributed. It loads
+by a :class:`DistributionPlan <modelopt.torch.utils.distributed.DistributionPlan>`. Every rank takes
+rank 0's plan: the one rank 0 passes as ``distribution``, or, when it passes none, the one
+:func:`resolve_distribution_plan <modelopt.torch.utils.plugins.model_load_utils.resolve_distribution_plan>`
+decides from ``MODELOPT_FSDP_SIZE``, ``MODELOPT_EP_DEGREE`` and rank 0's GPU memory. The plan is
+recorded on the model, where calibration and export read it.
+
+.. code-block:: python
+
+    # torchrun --nproc_per_node 8 quantize.py
+    from modelopt.torch.utils.distributed import distributed_run
+    from modelopt.torch.utils.plugins.model_load_utils import load_pretrained
+
+    with distributed_run():
+        model = load_pretrained("meta-llama/Llama-3.1-8B")  # a whole copy on every rank's GPU
+        # With MODELOPT_FSDP_SIZE=world: FSDP2-sharded over the 8 ranks instead.
+
+Routes
+======
+
+* A plan with ``dense="fsdp"`` takes the FSDP2 load (``parallel_load_and_prepare_fsdp2``): a causal
+  LM, or a text-only VLM it supports (Qwen3.5), is FSDP2-sharded over the process group, each rank
+  reading its own layers.
+* Any other plan gives every rank the whole model, through ``auto_class.from_pretrained``
+  (``auto_class`` defaults to ``AutoModelForCausalLM``).
+
+No load shards the routed experts yet. A plan the load decides keeps the expert-parallel degree at
+1, so an MoE loads without expert parallelism; an EP degree ``MODELOPT_EP_DEGREE`` asks for, or a
+plan with ``ep_size > 1``, is refused.
+
+A refusal on any rank raises on every rank, before the load, and says how to load the model whole:
+``MODELOPT_FSDP_SIZE=1`` for a plan the load decided, or
+``distribution=DistributionPlan(<ranks>, dense="replicate")``.
+
+The plan is recorded once, so ``fsdp2_wrap`` refuses a model ``load_pretrained`` returned: to shard
+a model yourself, load it with ``from_pretrained``.
+
+Whole-model loads
+=================
+
+* On a distributed run, each rank loads the model onto its own GPU when it has one, no
+  ``device_map`` is passed and ``init_quantized_weights`` does not place it: the GPU its process is
+  bound to once CUDA is set up, else the one its local rank numbers. ``device_map=None``,
+  ``from_pretrained``'s own default, loads it on the CPU for the caller to place. A rank without a
+  GPU of its own warns and loads where ``from_pretrained`` places it.
+* The options only the FSDP2 load takes are refused, except at the FSDP2 load's own defaults
+  (``cpu_offload=False``, ``mp_policy=None``, ``broadcast_chunk_size=8``) or
+  ``cpu_offload=None``, which ask nothing of a whole-model load, so a caller can forward them
+  whatever the plan. ``broadcast_chunk_size=None`` asks for all of a source's layers at once.
+
+Sharded loads
+=============
+
+* The FSDP2 load builds the model class and places the model itself: a text-only VLM through its
+  image-text class, anything else through ``AutoModelForCausalLM``. Leave ``auto_class`` out.
+* Besides its own options (``cpu_offload``, ``mp_policy``, ``broadcast_chunk_size``) it takes
+  ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other keyword
+  rather than ignore it. An option passed as ``None``, ``from_pretrained``'s own default, asks
+  nothing and is not refused.
+* It reads the checkpoint's tensors as stored and keeps their dtype, so it refuses a ``dtype`` and
+  a quantized checkpoint. It cannot run the model type's loading hook, so it refuses a checkpoint
+  the hook would build modules for, such as a Nemotron-H MTP head.
+* A single process, or a group of one, reads no knob: it takes the FSDP2 load only from a plan
+  passed as ``distribution``, and sets a group of one up for it, which needs
+  ``torchrun --nproc_per_node=1``.
+
+The automatic layout
+====================
+
+With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory decides whether the dense part
+is sharded. The model is weighed in the checkpoint's own dtype (its config's, else bfloat16), not
+one the load asks for: hand a load that widens it, ``dtype=torch.float32`` say, its plan.
+
+A call that places the model (``device_map``, ``quantization_config``, a ``dtype``) loads only
+where the model fits one GPU; where it does not, the refusal says the layout was decided rather
+than asked for. Pass ``distribution=`` or set ``MODELOPT_FSDP_SIZE=1`` to keep the load whole. An
+option at ``from_pretrained``'s own default, ``device_map=None`` included, places nothing and does
+not keep the load whole.
+
+Configs
+=======
+
+``hf_config`` is the config ``resolve_distribution_plan`` returned with the plan handed in as
+``distribution``: the load builds from it as given, so without that plan it is refused. A
+``config`` passed for the load, or a path to one, is the one the plan is decided from, and a sharded
+load takes it as its ``hf_config``.
+
+Hub options
+===========
+
+A Hub model ID is made local first, so the export can read the source checkpoint from disk; the
+``revision``, ``token``, ``cache_dir`` and ``local_files_only`` options go to that fetch. The config
+read and a whole-model load also take all but ``revision``: an ``auto_map`` entry can name custom
+code in another Hub repo, which Transformers fetches with them. The FSDP2 load builds the model
+from the config with ``from_config``, which takes none of them, so such code comes with the
+environment's Hub settings.
+
+Launches without a process group
+================================
+
+A launch whose processes set no process group up (a launcher that only exports ``WORLD_SIZE``, as
+``hf_ptq.py`` without ``--use_fsdp2``) has each process load as a single process does: the whole
+model, with no plan to agree on and no ``device_map`` default. A plan for several ranks, a sharded
+load, and a knob asking to shard need the group, so they are refused there: set it up first, with
+``distributed_run()``.
+
+Inside ``init_quantized_weights``
+=================================
+
+``init_quantized_weights`` (see :doc:`_compress_quantized_models`) quantizes and compresses the
+model in the ``from_pretrained`` it patches, which the FSDP2 load never calls. Inside it, a plan
+the load decides gives every rank the whole model, with nothing to measure, and the load leaves
+the device map to ``init_quantized_weights``, which computes it and refuses one it is passed. A
+knob or a plan that asks to shard is refused there.

modelopt/torch/quantization/plugins/accelerate.py

@@ -358,6 +358,8 @@ def patched_from_pretrained(cls, /, pretrained_model_name_or_path, *args, **kwar
             **dispatch_kwargs,
         )
 
+    # A sharded load_pretrained builds the model without from_pretrained, so it reads this mark.
+    patched_from_pretrained._modelopt_init_quantized_weights = True  # type: ignore[attr-defined]
     try:
         # Apply our patches
         transformers.AutoModelForCausalLM.from_pretrained = classmethod(patched_from_pretrained)

modelopt/torch/utils/distributed.py

@@ -351,6 +351,16 @@ def is_fsdp_dense(self) -> bool:
 _PLAN_ATTR = "_modelopt_distribution_plan"
 
 
+def _recorded_plan(model: torch.nn.Module) -> tuple[str, DistributionPlan] | None:
+    """Which of ``model`` and its submodules records a plan, named, and that plan; else ``None``."""
+    for name, module in model.named_modules():
+        recorded = getattr(module, _PLAN_ATTR, None)
+        if recorded is not None:
+            where = f"{type(module).__name__} {name!r}" if name else type(module).__name__
+            return where, recorded
+    return None
+
+
 def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> None:
     """Record that ``model`` is distributed as ``plan``, on it and on every submodule it has now.
 
@@ -374,14 +384,12 @@ def set_distribution_plan(model: torch.nn.Module, plan: DistributionPlan) -> Non
             f"{plan!r} is for {plan.world_size} ranks, but the run it is recorded in has "
             f"{size()}; a plan's world_size is the ranks that hold the model."
         )
-    for name, module in model.named_modules():
-        recorded = getattr(module, _PLAN_ATTR, None)
-        if recorded is not None:
-            where = f"{type(module).__name__} {name!r}" if name else type(module).__name__
-            raise ValueError(
-                f"{where} is already recorded as distributed by {recorded!r}; a model is "
-                f"distributed and its plan recorded once, so {plan!r} cannot be recorded over it."
-            )
+    if (found := _recorded_plan(model)) is not None:
+        where, recorded = found
+        raise ValueError(
+            f"{where} is already recorded as distributed by {recorded!r}; a model is distributed "
+            f"and its plan recorded once, so {plan!r} cannot be recorded over it."
+        )
     for module in model.modules():
         setattr(module, _PLAN_ATTR, plan)
 
@@ -403,9 +411,10 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan:
     if sharded or size() > 1:
         how = "is FSDP2-sharded" if sharded else f"is held in a run of {size()} ranks"
         raise RuntimeError(
-            f"{type(model).__name__} {how}, but has no distribution plan recorded. Record how it "
-            "is distributed with set_distribution_plan(); a module added to a model after its plan "
-            "was recorded has none, so read the plan off that model."
+            f"{type(model).__name__} {how}, but has no distribution plan recorded. Load it with "
+            "load_pretrained, or record how it is distributed with set_distribution_plan(); a "
+            "module added to a model after its plan was recorded has none, so read the plan off "
+            "that model."
         )
     return DistributionPlan()
 

modelopt/torch/utils/env.py

@@ -35,10 +35,10 @@
 
 
 def ep_degree() -> str:
-    """Expert-parallel degree knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``."""
-    return os.environ.get("MODELOPT_EP_DEGREE") or "auto"
+    """Expert-parallel degree knob for ``resolve_distribution_plan``; blank/unset -> ``"auto"``."""
+    return (os.environ.get("MODELOPT_EP_DEGREE") or "").strip() or "auto"
 
 
 def fsdp_size() -> str:
-    """Dense-sharding knob for ``resolve_distribution_plan``; ``""``/unset -> ``"auto"``."""
-    return os.environ.get("MODELOPT_FSDP_SIZE") or "auto"
+    """Dense-sharding knob for ``resolve_distribution_plan``; blank/unset -> ``"auto"``."""
+    return (os.environ.get("MODELOPT_FSDP_SIZE") or "").strip() or "auto"

modelopt/torch/utils/fsdp2_utils.py

@@ -26,6 +26,7 @@
 from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard
 from torch.distributed.tensor import DTensor
 
+from modelopt.torch.utils.distributed import _recorded_plan
 from modelopt.torch.utils.logging import warn_rank_0
 
 __all__ = ["fsdp2_wrap"]
@@ -99,7 +100,20 @@ def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False
     :func:`_off_dtype_params`), since FSDP2 rejects a shard group that mixes dtypes. They stay
     replicated and are moved onto the shards' device, which ``fully_shard`` does not do for the
     params it ignores.
+
+    It records no distribution plan, and refuses a model that, or any of whose submodules, records
+    one, as a model ``load_pretrained`` returned does: wrap one loaded with ``from_pretrained``.
     """
+    if (found := _recorded_plan(model)) is not None:
+        # A model recorded as replicated would read so while sharded, misleading the export, and one
+        # recorded as sharded is wrapped already. A ValueError, as set_distribution_plan raises.
+        where, recorded = found
+        raise ValueError(
+            f"{where} is already recorded as distributed by {recorded!r}, and a plan is recorded "
+            "once. To shard a model yourself, load it with from_pretrained rather than "
+            "load_pretrained, which records the plan it loaded by, then record yours with "
+            "set_distribution_plan() once it is wrapped."
+        )
     # Importing the quantization package registers every mode and plugin; only a wrap needs it.
     from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector