diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..92f87a9d746 100755 --- a/CHANGELOG.rst +++ b/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(, 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. diff --git a/docs/source/guides/1_quantization.rst b/docs/source/guides/1_quantization.rst index 38ce0956b79..1db915c6975 100644 --- a/docs/source/guides/1_quantization.rst +++ b/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 diff --git a/docs/source/guides/_distributed_loading.rst b/docs/source/guides/_distributed_loading.rst new file mode 100644 index 00000000000..230b0249629 --- /dev/null +++ b/docs/source/guides/_distributed_loading.rst @@ -0,0 +1,119 @@ +======================================== +Loading Hugging Face models across ranks +======================================== + +:func:`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 `. 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 ` +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(, 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. diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index f1601de058d..36b08318a23 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/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) diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index e82257a3132..3f1dd35abeb 100644 --- a/modelopt/torch/utils/distributed.py +++ b/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() diff --git a/modelopt/torch/utils/env.py b/modelopt/torch/utils/env.py index 7e27e7857c3..89f869c3efb 100644 --- a/modelopt/torch/utils/env.py +++ b/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" diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py index 14db7266b5a..419abb4ae56 100644 --- a/modelopt/torch/utils/fsdp2_utils.py +++ b/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 diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py index 1967708c248..094a2d79504 100644 --- a/modelopt/torch/utils/plugins/model_load_utils.py +++ b/modelopt/torch/utils/plugins/model_load_utils.py @@ -13,12 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""HuggingFace model loading: :func:`load_pretrained`, and the plan to distribute a model by.""" +"""HuggingFace model loading: :func:`load_pretrained`, by a plan to distribute the model.""" +import copy import logging import os import re +import warnings from collections import defaultdict +from contextlib import contextmanager from itertools import chain from typing import TYPE_CHECKING, Any @@ -52,6 +55,7 @@ barrier, broadcast_state_dict, is_initialized, + set_distribution_plan, ) from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap from modelopt.torch.utils.logging import print_rank_0, warn_rank_0 @@ -428,74 +432,438 @@ def record_unplaced_source_keys( return keys +_FSDP2 = "FSDP2" +# What each sharded load takes besides the checkpoint. These loads build the model class and place +# it themselves, so from_pretrained's own arguments (device_map, max_memory, ...) do not reach them. +_FSDP2_ONLY = frozenset({"cpu_offload", "mp_policy", "broadcast_chunk_size"}) +_SHARDED_LOAD_OPTIONS = { + _FSDP2: _FSDP2_ONLY | {"trust_remote_code", "attn_implementation", "hf_config"} +} +# The FSDP2-only values that ask nothing of a whole load: the FSDP2 load's defaults, and None where +# that leaves the option unset (a None broadcast_chunk_size asks for all of a source's layers). +_ASKS_NOTHING = {"cpu_offload": (False, None), "mp_policy": (None,), "broadcast_chunk_size": (8,)} + + +def _route(plan: DistributionPlan) -> str | None: + """The sharded load ``plan`` takes, or ``None`` when every rank loads the whole model.""" + if plan.world_size != dist.size(): + raise ValueError( + f"{plan!r} is for {plan.world_size} ranks, but this run has {dist.size()}." + ) + if plan.is_expert_parallel: + raise NotImplementedError( + f"{plan!r} shards the routed experts, which load_pretrained has no load for. Plan " + "ep_size=1 (MODELOPT_EP_DEGREE=1) to give every rank the whole model, or shard it " + "with FSDP2." + ) + return _FSDP2 if plan.is_fsdp_dense else None + + +@contextmanager +def _refused_together(label: str, alone: bool): + """Raise on every rank if the block raises on any, so no rank waits for one that refused.""" + error = None + try: + yield + except BaseException as e: # A SystemExit too: the others would wait for this rank as well. + error = e + + def refuse() -> None: + if error is not None: + raise error + + if alone: # A process of a launch that set no group up has no peers to tell. + refuse() + else: + dist.all_reduce(label, refuse, combine=lambda _: None) + + +def _read_config(path, auto_class, trust_remote_code: bool = False, **hub_options): + """``path``'s config as ``auto_class.from_pretrained`` reads it: config_class, or AutoConfig.""" + if (reader := getattr(auto_class, "config_class", None)) is not None: + return reader.from_pretrained(path, **hub_options) + return AutoConfig.from_pretrained(path, trust_remote_code=trust_remote_code, **hub_options) + + +def _init_quantized_weights() -> bool: + """Whether ``init_quantized_weights`` patched ``from_pretrained`` to quantize and place.""" + return getattr(AutoModelForCausalLM.from_pretrained, "_modelopt_init_quantized_weights", False) + + +def _whole_instead(decided: bool) -> str: + """How to have every rank load the model whole: the knob counts only for a decided plan.""" + plan = f"distribution=DistributionPlan({dist.size()}, dense='replicate')" + return f"MODELOPT_FSDP_SIZE=1 or {plan}" if decided else plan + + +def _refuse_sharded_load_options( + route, auto_class, kwargs, decided: bool, measured: bool = False +) -> None: + """Refuse what the sharded load ``route`` cannot honor; ``measured``: GPU memory chose it.""" + # An option passed as None, from_pretrained's own default (device_map=None, say), asks nothing. + passed = {name for name, value in kwargs.items() if value is not None} + unknown = sorted(passed - _SHARDED_LOAD_OPTIONS[route] - {"dtype", "torch_dtype"}) + # torch_dtype is from_pretrained's older name for dtype, and None means "auto" to it. + dtypes = [name for name in ("dtype", "torch_dtype") if kwargs.get(name) not in (None, "auto")] + if measured and (refused := sorted({*unknown, *dtypes})): + # Say so, rather than refuse the options below as if the caller had asked for this load. + raise TypeError( + f"{refused} do not apply to a sharded load, but no plan or MODELOPT_FSDP_SIZE asked for " + f"a layout, and by rank 0's GPU memory the model does not fit one GPU, so the {route} " + f"load shards it. To load it whole on every rank: {_whole_instead(decided)}." + ) + if _init_quantized_weights(): + raise NotImplementedError( + f"The {route} load builds the model without from_pretrained, so it would skip " + "init_quantized_weights' quantization; load the model whole on every rank: " + f"{_whole_instead(decided)}." + ) + if auto_class not in (None, AutoModelForCausalLM): + name = getattr(auto_class, "__name__", repr(auto_class)) + raise NotImplementedError( + f"The {route} load builds its model class itself, a causal LM's or a text-only VLM's " + f"image-text class, so it takes no auto_class={name}: leave it out to have the load " + f"pick, or load the model whole with it: {_whole_instead(decided)}." + ) + if dtypes: # A sharded load keeps the checkpoint's dtype. + name = dtypes[0] + raise TypeError( + f"The {route} load reads the checkpoint's tensors as stored and keeps their dtype, so " + f"{name}={kwargs[name]!r} does not apply to it; FSDP2's mp_policy sets the dtype it " + "computes in." + ) + if unknown: + raise TypeError( + f"{unknown} do not apply to the {route} load, which places the model itself; it " + f"takes {sorted(_SHARDED_LOAD_OPTIONS[route])}. To place the model yourself, have " + f"every rank load it whole: {_whole_instead(decided)}." + ) + + +def _sharded_load_kwargs( + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided: bool, measured: bool +) -> dict: + """What the sharded load ``route`` names takes, refusing what it cannot honor.""" + from modelopt.torch.models.hf import checkpoint_has_mtp + + _refuse_sharded_load_options(route, auto_class, kwargs, decided, measured) + # Past the refusal, anything else is a dtype the load keeps or an option at its default, None. + kwargs = {k: v for k, v in kwargs.items() if k in _SHARDED_LOAD_OPTIONS[route]} + if checkpoint_has_mtp(model_type, ckpt_path): + # The head is built on a whole model only; carried over, it would export unquantized. + raise NotImplementedError( + f"The {route} load cannot build the MTP head this {model_type} checkpoint carries. To " + f"load it whole on every rank: {_whole_instead(decided)}." + ) + trust_remote_code = kwargs.pop("trust_remote_code", False) + hf_config = kwargs.pop("hf_config", None) + if hf_config is None: + hf_config = AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) + nested = (getattr(hf_config, name, None) for name in ("text_config", "llm_config")) + if any(getattr(c, "quantization_config", None) is not None for c in (hf_config, *nested)): + # Read as stored, quantized weights (a VLM's language model's too) would load as if not. + raise NotImplementedError( + f"The {route} load cannot dequantize this {model_type} checkpoint. To load it whole on " + f"every rank: {_whole_instead(decided)}." + ) + # As build_meta_causal_lm builds: AutoModelForCausalLM refuses most VLMs' configs, and builds + # some as the language model alone, whose other towers the export would carry unquantized. + builds = ( + AutoModelForImageTextToText if is_fsdp2_text_only_vlm(hf_config) else AutoModelForCausalLM + ) + if _auto_class(hf_config, trust_remote_code=trust_remote_code) is not builds: + raise NotImplementedError( + f"The {route} load builds a causal LM, or a text-only VLM it supports, which this " + f"{model_type} checkpoint is not. To load it whole on every rank, pass its model class " + f"as auto_class, with {_whole_instead(decided)}." + ) + if not dist.is_initialized() and "RANK" not in os.environ: + # FSDP2 shards through a process group even on one rank. + raise RuntimeError( + f"The {route} load shards through a process group, even on one rank, which a process " + "started without a launcher cannot set up; launch it with `torchrun --nproc_per_node=1`." + ) + if not _has_own_gpu(): + # Refused with the other ranks, not inside the load, where they would wait for this one. + raise RuntimeError( + f"The {route} load shards onto each rank's own GPU, and this rank has none: it sees " + f"{torch.cuda.device_count()} GPU(s) as local rank {dist.local_rank()}. To load it " + f"whole on every rank: {_whole_instead(decided)}." + ) + return {"trust_remote_code": trust_remote_code, "hf_config": hf_config, **kwargs} + + +def _own_gpu_index() -> int | None: + """This rank's GPU, as torch's DeviceMesh takes it, or ``None`` for a rank with none there. + + The device CUDA is set to once it is set up, else the one the rank's local rank numbers. + """ + if not torch.cuda.is_available(): + return None + if torch.cuda.is_initialized(): + return torch.cuda.current_device() + index = dist.local_rank() + return index if index < torch.cuda.device_count() else None + + +def _has_own_gpu() -> bool: + """Whether this rank has a GPU of its own (:func:`_own_gpu_index`).""" + return _own_gpu_index() is not None + + +def _load_sharded(ckpt_path, load_kwargs) -> nn.Module: + """The FSDP2 load, by the arguments :func:`_sharded_load_kwargs` checked.""" + if not dist.is_initialized(): + # A single process's group of one, which the sharded model needs while it lives. + dist.setup() + return parallel_load_and_prepare_fsdp2( + ckpt_path, + torch.device("cuda", _own_gpu_index()), + dist.rank(), + dist.size(), + **load_kwargs, + ) + + def load_pretrained( ckpt_path: str, auto_class: Any = None, *, model_type: str | None = None, + distribution: DistributionPlan | None = None, **kwargs, ) -> nn.Module: - """Load a Hugging Face model for ModelOpt to optimize and export. - - ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to - ``AutoModelForCausalLM``), plus what ModelOpt needs around the load: + """Load a Hugging Face model for ModelOpt to optimize and export, whole or distributed. + + ``distribution`` is the :class:`~modelopt.torch.utils.distributed.DistributionPlan` to load by: + every rank takes rank 0's or, when rank 0 passes none, :func:`resolve_distribution_plan`'s. + + * ``dense="fsdp"``: :func:`parallel_load_and_prepare_fsdp2` shards a causal LM, or a text-only + VLM it supports, over the process group (one process needs ``torchrun --nproc_per_node=1``). + It builds and places the model itself, so it takes its own options but no ``auto_class``, + placement or ``dtype``, and refuses a quantized checkpoint. + * otherwise ``auto_class.from_pretrained(ckpt_path, **kwargs)`` (``auto_class`` defaults to + ``AutoModelForCausalLM``); on a distributed run every rank loads the whole model, onto its own + GPU unless ``device_map`` is passed or ``init_quantized_weights`` places it. + + No load here shards the routed experts: a plan with ``ep_size > 1`` is refused. A refusal on any + rank raises on every rank, before the load. The plan is recorded on the model + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) for calibration and export, so + :func:`~modelopt.torch.utils.fsdp2_utils.fsdp2_wrap` refuses it: to shard a model yourself, + load it with ``from_pretrained``. :doc:`/guides/_distributed_loading` has the routing in full. + Around the load: * A Hub model ID is made local first (:func:`ensure_local_checkpoint`, which takes the ``revision``, ``token``, ``cache_dir`` and ``local_files_only`` ``kwargs``): the export reads the source checkpoint from disk. Under ``torch.distributed`` every rank must call this. The - config and model loads still take all but ``revision``: an ``auto_map`` entry can name - custom code in another Hub repo, which Transformers fetches with them. A ``subfolder`` is - part of where the checkpoint is: the config, the loading hook and the record below read it, - as the load does. + config read and a whole-model load still take all but ``revision``: an ``auto_map`` entry can + name custom code in another Hub repo, which Transformers fetches with them. A sharded load + builds the model with ``from_config``, which takes none. A ``subfolder`` is part of where the + checkpoint is: the config, the loading hook and the record below read it, as the load does. * The model type's loading hook runs (:func:`~modelopt.torch.models.hf.prepare_model_for_loading`), e.g. to keep the MTP head a Nemotron-H checkpoint carries. ``model_type`` selects it; when omitted, the type of the - ``config`` passed for the load, else of the checkpoint's config. + ``config`` passed for the load, else of the checkpoint's config. A sharded load cannot run + it, so it refuses a checkpoint the hook would build modules for. * The checkpoint tensors the model has no parameter for are recorded (:func:`record_unplaced_source_keys`), so the export carries them over unchanged. The model is returned alone: ``output_loading_info=True`` is refused. So is a ``variant``: the record and the export read the checkpoint's default weight files. """ - if kwargs.pop("output_loading_info", False): - raise ValueError( - "load_pretrained returns the model alone and does not take output_loading_info=True; " - "the checkpoint tensors the model has no parameter for are recorded on it instead " - "(see record_unplaced_source_keys)." - ) - if kwargs.get("variant") is not None: - raise ValueError( - f"load_pretrained does not take variant={kwargs['variant']!r}: the checkpoint tensors " - "the model has no parameter for are recorded from, and carried into the export from, " - "the checkpoint's default weight files." - ) # Importing modelopt.torch.models registers every model spec; only callers of this need them. from modelopt.torch.models.hf import prepare_model_for_loading - auto_class = auto_class or AutoModelForCausalLM + # A launch whose processes set no group up (hf_ptq without --use_fsdp2): each loads alone. + alone = dist.is_distributed_launch() and not dist.is_initialized() + if alone and isinstance(distribution, DistributionPlan) and distribution.world_size > 1: + raise RuntimeError( + f"{distribution!r} is for {distribution.world_size} ranks, whose process group this " + f"{os.environ['WORLD_SIZE']}-rank launch has not set up; load inside distributed_run(), " + "which sets it up." + ) + if alone and isinstance(distribution, DistributionPlan) and _route(distribution) is not None: + # One rank cannot set the launch's group up, and a one-rank plan is no plan for the launch. + launch = os.environ["WORLD_SIZE"] + raise RuntimeError( + f"The {_route(distribution)} load shards through this {launch}-rank launch's process " + "group, which is not set up; load inside distributed_run(), which sets it up, with a " + f"plan for the launch: DistributionPlan({launch}, dense='fsdp')." + ) + if alone and distribution is None: + # Knobs asking to shard over the launch are refused, as is a value that is no EP degree. + ep, world = _ep_degree(None), int(os.environ["WORLD_SIZE"]) + experts = ep != "auto" and ep > 1 + asked = [f"MODELOPT_EP_DEGREE={env.ep_degree()}"] if experts else [] + if _dense_from_env(world) == "fsdp": + asked.append(f"MODELOPT_FSDP_SIZE={env.fsdp_size()}") + if asked: + ep_way = ", with MODELOPT_EP_DEGREE=1: no load here shards the routed experts" + raise RuntimeError( + f"{' and '.join(asked)} ask to shard the model over this " + f"{os.environ['WORLD_SIZE']}-rank launch, whose process group is not set up; load " + f"inside distributed_run(), which sets it up{ep_way if experts else ''}." + ) + # None is these options' own default: passed explicitly, it is no option to refuse. + for name in ("config", "variant", "hf_config"): + if name in kwargs and kwargs[name] is None: + del kwargs[name] + # The config to load with, or, as from_pretrained also takes it, a path to read it from. + config = kwargs.get("config") + config = config if isinstance(config, (PretrainedConfig, str, os.PathLike)) else None + + def handed() -> DistributionPlan | None: + if distribution is None: + return None + if not isinstance(distribution, DistributionPlan): + raise TypeError(f"distribution={distribution!r}: expected a DistributionPlan or None.") + _route(distribution) + return distribution + + # Every rank loads rank 0's distribution, which rank 0 checks before it sends it. + plan = handed() if alone else dist.broadcast_from("distribution", handed) hub_options = { k: kwargs.pop(k) for k in ("token", "cache_dir", "local_files_only") if k in kwargs } revision = {"revision": kwargs.pop("revision")} if "revision" in kwargs else {} - _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) + decided = plan is None and dist.size() > 1 + # Each rank checks its own arguments; a refusal on one fails every rank before the fetch. + with _refused_together("load arguments", alone): + # No load here shards the routed experts yet, so an EP degree left to the decision is 1. + # Parsed in this check, a knob one rank cannot parse fails every rank. + ep_size = 1 if decided and _ep_degree(None) == "auto" else None + if kwargs.pop("output_loading_info", False): + raise ValueError( + "load_pretrained returns the model alone and does not take " + "output_loading_info=True; the checkpoint tensors the model has no parameter for " + "are recorded on it instead (see record_unplaced_source_keys)." + ) + if kwargs.get("variant") is not None: + raise ValueError( + f"load_pretrained does not take variant={kwargs['variant']!r}: the checkpoint " + "tensors the model has no parameter for are recorded from, and carried into the " + "export from, the checkpoint's default weight files." + ) + if config is not None and "hf_config" in kwargs: + raise TypeError("config and hf_config both name the config to load with; pass one.") + if "hf_config" in kwargs and (plan is None or _route(plan) is None): + # Only a sharded load builds from hf_config, so it needs the plan it came with. + raise TypeError( + "hf_config is the config resolve_distribution_plan returned with a plan that " + "shards the model: pass that plan as distribution=, or pass config= to have the " + "plan decided from your config." + ) + if plan is not None and _route(plan) is not None: + # A config passed becomes the load's hf_config; subfolder is where the checkpoint is. + _refuse_sharded_load_options( + _route(plan), + auto_class, + { + k: v + for k, v in kwargs.items() + if k != "subfolder" and not (k == "config" and config is not None) + }, + decided=False, + ) + # Rank 0 fetches the checkpoint; a rank that cannot see it then fails every rank. + with _refused_together("checkpoint fetch", alone): + _, ckpt_path = ensure_local_checkpoint(ckpt_path, **revision, **hub_options) + if config is not None and not isinstance(config, PretrainedConfig): + # Read as the load reads it, so the plan and the loading hook follow the model's config. + trust = kwargs.get("trust_remote_code", False) + config = _read_config(config, auto_class, trust, **hub_options) subfolder = kwargs.pop("subfolder", None) if subfolder: ckpt_path = os.path.join(ckpt_path, subfolder) trust_remote_code = kwargs.get("trust_remote_code", False) - if model_type is None: - # The load uses the config passed for it, if any, so the loading hook follows that one. - model_type = getattr(kwargs.get("config"), "model_type", None) or ( - AutoConfig.from_pretrained( - ckpt_path, trust_remote_code=trust_remote_code, **hub_options - ).model_type + hf_config = kwargs.get("hf_config") + if decided: + plan, hf_config = resolve_distribution_plan( + ckpt_path, + config=config, + auto_class=auto_class, + trust_remote_code=trust_remote_code, + ep_size=ep_size, + **hub_options, ) - with prepare_model_for_loading( - model_type, - ckpt_path, - trust_remote_code, - model_class=None if auto_class in (AutoModel, AutoModelForCausalLM) else auto_class, - ): - model = auto_class.from_pretrained(ckpt_path, **hub_options, **kwargs) - unplaced = record_unplaced_source_keys(model, ckpt_path) + if _route(plan) is not None: + kwargs.setdefault("hf_config", hf_config) + elif plan is None: + # One rank holds the whole model, whatever the model: there is nothing to decide or read. + plan = DistributionPlan() + route = _route(plan) + # Each rank reads the config itself, so a refusal of it fails every rank before the load. + with _refused_together("load checks", alone): + # Options only the FSDP2 load takes would reach from_pretrained, which has no use for them. + if route is None: + for name, nothing in _ASKS_NOTHING.items(): + if name in kwargs and kwargs[name] in nothing: + del kwargs[name] + if route is None and (sharded := sorted(set(kwargs) & _FSDP2_ONLY)): + raise TypeError( + f"{sharded} only apply to the FSDP2 load, but {plan!r} gives every rank the whole " + "model; pass a plan with dense='fsdp' to shard it." + ) + measured = decided and route == _FSDP2 and _dense_from_env(plan.world_size) == "auto" + if route is not None and config is not None: + # The sharded load builds from hf_config, a copy, as building writes to the config. + del kwargs["config"] + if "hf_config" not in kwargs: + kwargs["hf_config"] = copy.deepcopy(config) + if model_type is None: + # The load uses the config passed for it, if any, so the loading hook follows that one. + if route == _FSDP2 and hf_config is None and config is None: + # The FSDP2 load builds from the checkpoint's config too: read it once, for both. + hf_config = AutoConfig.from_pretrained( + ckpt_path, trust_remote_code=trust_remote_code, **hub_options + ) + kwargs["hf_config"] = hf_config + chosen = hf_config if hf_config is not None else config + model_type = getattr(chosen, "model_type", None) or ( + _read_config(ckpt_path, auto_class, trust_remote_code, **hub_options).model_type + ) + load_kwargs = ( + {} + if route is None + else _sharded_load_kwargs( + route, ckpt_path, auto_class, model_type, kwargs, hub_options, decided, measured + ) + ) + if route is not None: + model = _load_sharded(ckpt_path, load_kwargs) + # A sharded load finds the unplaced keys while it places the weights. + keys = model._modelopt_unplaced_source_keys + else: + # init_quantized_weights computes the device map itself, and refuses one it is passed. + spread = plan.world_size > 1 and not _init_quantized_weights() + if spread and (index := _own_gpu_index()) is not None: + # Each rank takes the whole model onto its own GPU: under "auto", every rank's + # accelerate could pick the same idle GPU. + kwargs.setdefault("device_map", {"": f"cuda:{index}"}) + elif spread and "device_map" not in kwargs: + # This rank says so itself: rank 0 may have a GPU of its own, and nothing to say. + warnings.warn( + f"Rank {dist.rank()} has no GPU of its own: it sees {torch.cuda.device_count()} " + f"GPU(s) as local rank {dist.local_rank()}, so it loads the whole model where " + "from_pretrained places it without a device_map, on the CPU.", + stacklevel=2, + ) + auto_class = auto_class or AutoModelForCausalLM + # It runs no collective, so a rank failing in it alone would leave the others in the next. + with ( + _refused_together("whole-model load", alone), + prepare_model_for_loading( + model_type, + ckpt_path, + trust_remote_code, + model_class=None if auto_class in (AutoModel, AutoModelForCausalLM) else auto_class, + ), + ): + model = auto_class.from_pretrained(ckpt_path, **hub_options, **kwargs) + keys = None + set_distribution_plan(model, plan) + unplaced = record_unplaced_source_keys(model, ckpt_path, keys) if unplaced: print_rank_0( f"{len(unplaced)} checkpoint tensor(s) the model has no parameter for " @@ -830,6 +1198,9 @@ def parallel_load_and_prepare_fsdp2( ``broadcast_chunk_size`` sets how many of a source's owned layers are broadcast per collective: a smaller value lowers the peak transient GPU memory at the cost of more collectives (default 8; pass ``None`` to broadcast all of a source's layers at once). + + It records no distribution plan: a caller records its own, as :func:`load_pretrained` does + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`). """ resolved_path = _resolve_checkpoint_dir(ckpt_path, rank) weight_map = indexed_weight_map(resolved_path) @@ -1007,14 +1378,19 @@ def _ep_degree(ep_size) -> int | str: if str(value).lower() in ("auto", "0"): return "auto" try: - return int(value) - except ValueError: + degree = int(value) + except (ValueError, TypeError): # TypeError: an argument that is no number at all, as [2] + degree = 0 + if isinstance(value, float) and value != degree: # int() truncates a fraction + degree = 0 + if degree < 1: name = "MODELOPT_EP_DEGREE" if ep_size is None else "ep_size" - raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") from None + raise ValueError(f"{name}={value!r}: expected auto, 0 (auto) or an EP degree.") + return degree def _decide_distribution( - ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool + ckpt_path: str, hf_config, world_size: int, ep_size, dense, trust_remote_code: bool, auto_class ) -> DistributionPlan: """The plan :func:`resolve_distribution_plan` returns, decided by this rank alone.""" # What the caller set the EP degree with, for a refusal to name it. @@ -1030,6 +1406,25 @@ def _decide_distribution( ) ep_size = _ep_degree(ep_size) dense = _dense_from_env(world_size) if dense is None else dense + if ( + (auto_class not in (None, AutoModelForCausalLM) or _init_quantized_weights()) + and ep_size in ("auto", 1) + and dense in ("auto", "replicate") + ): + # The sharded loads refuse this class and init_quantized_weights: load it whole, unbuilt. + return DistributionPlan(world_size) + if auto_class not in (None, AutoModelForCausalLM): + # Refused here, rather than build the decision's own pick of class, which may not build. + name = getattr(auto_class, "__name__", repr(auto_class)) + raise NotImplementedError( + f"A sharded load builds its model class itself, a causal LM's or a text-only VLM's " + f"image-text class, so it takes no auto_class={name} while ep_size={ep_size!r}, " + f"dense={dense!r} ask to shard. To give every rank the whole model, plan ep_size=1 and " + "dense='replicate' (MODELOPT_EP_DEGREE=1, MODELOPT_FSDP_SIZE=1)." + ) + # Building writes to the config it builds from, and only rank 0 decides: build from a copy, so + # the config returned with the plan, which a sharded load builds from, is alike on every rank. + hf_config = copy.deepcopy(hf_config) try: model = build_meta_causal_lm( ckpt_path, @@ -1097,6 +1492,8 @@ def _decide_distribution( def resolve_distribution_plan( ckpt_path: str, *, + config: PretrainedConfig | None = None, + auto_class: Any = None, trust_remote_code: bool = False, ep_size: int | str | None = None, dense: str | None = None, @@ -1113,7 +1510,7 @@ def resolve_distribution_plan( every later step reads it back from there (:func:`~modelopt.torch.utils.distributed.distribution_plan`). Returns ``(plan, hf_config)``, ``hf_config`` being this rank's own copy of the config (rank 0 decides the plan from its copy), - for a sharded load to build from. + for a sharded load to build from: :func:`load_pretrained` takes it only with a plan that shards. Rank 0 decides, from the model built on the meta device: from its parameters rather than estimates of them, and from its routed experts as an expert-parallel load finds them. Every rank @@ -1140,6 +1537,11 @@ def resolve_distribution_plan( returns no path, so a caller that loads the checkpoint resolves it first, with :func:`ensure_local_checkpoint`, and passes the local directory: the plan is then decided for the snapshot the load reads. + + ``config``, the caller's own config, is decided from instead of the checkpoint's. ``auto_class`` + is the class the load builds with: with one the sharded loads refuse, or inside + ``init_quantized_weights``, and nothing asking to shard, the plan loads the model whole, decided + without a build. """ # Each rank reads the checkpoint for itself. One that cannot must fail every rank here, rather # than leave the others waiting in a broadcast below for a rank that already raised. @@ -1155,14 +1557,18 @@ def checkpoint_read() -> bool: ) if subfolder: local_path = os.path.join(local_path, subfolder) - config = AutoConfig.from_pretrained( + if config is not None: + read.append((local_path, copy.deepcopy(config))) + return True + hf_config = _read_config( local_path, + auto_class, trust_remote_code=trust_remote_code, token=token, cache_dir=cache_dir, local_files_only=local_files_only, ) - read.append((local_path, config)) + read.append((local_path, hf_config)) return True dist.all_reduce("checkpoint read", checkpoint_read, all) @@ -1171,7 +1577,7 @@ def checkpoint_read() -> bool: plan = dist.broadcast_from( "plan", lambda: _decide_distribution( - ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code + ckpt_path, hf_config, world_size, ep_size, dense, trust_remote_code, auto_class ), ) if plan != DistributionPlan(): diff --git a/tests/examples/hf_ptq/test_example_utils.py b/tests/examples/hf_ptq/test_example_utils.py index 392b3d2807e..01e0634c93a 100644 --- a/tests/examples/hf_ptq/test_example_utils.py +++ b/tests/examples/hf_ptq/test_example_utils.py @@ -284,7 +284,7 @@ def test_get_model_uses_expected_dtype_kwarg( torch_dtype=torch.bfloat16, ) - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): calls["eval"] = True @@ -369,7 +369,7 @@ def test_get_model_device_map_for_diffusion_gemma( torch_dtype=torch.bfloat16, ) - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): calls["eval"] = True @@ -498,7 +498,7 @@ def _load_config(*args, **kwargs): raise builtin_config_error return hf_config - class FakeModel: + class FakeModel(torch.nn.Module): # load_pretrained records its plan on every module def eval(self): return None diff --git a/tests/gpu/torch/export/test_export_carry_over.py b/tests/gpu/torch/export/test_export_carry_over.py index fa528f4cdf9..34a54d23f5e 100644 --- a/tests/gpu/torch/export/test_export_carry_over.py +++ b/tests/gpu/torch/export/test_export_carry_over.py @@ -40,7 +40,12 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes -from modelopt.torch.utils.distributed import fsdp2_wrap, is_fsdp2_model +from modelopt.torch.utils.distributed import ( + DistributionPlan, + fsdp2_wrap, + is_fsdp2_model, + set_distribution_plan, +) # Small enough to force the tiny model across several shards, so the multi-file layout -- and the # weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. @@ -108,6 +113,8 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar model.eval() fsdp2_wrap(model) + # Wrapped by hand, the model records how it is distributed, as load_pretrained would. + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # The FSDP2 export path is selected by exactly this predicate, so assert it here: without it # the export would silently fall back to the rank-0 gather and the test would pass while # never touching the code under test. @@ -201,6 +208,7 @@ def _ptq_language_model_and_export(rank, size, *, src_dir, export_dir): ) model.eval() fsdp2_wrap(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) assert is_fsdp2_model(model), "fsdp2_wrap did not shard the VLM" torch.distributed.barrier() diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 3296ccbfc86..54b96461485 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""GPU/distributed tests for the FSDP2 load path and its helpers.""" +"""GPU/distributed tests for the loads: ``load_pretrained`` by its plan, the FSDP2 load, and their +helpers.""" import json import os @@ -24,10 +25,19 @@ from _test_utils.torch.transformers_models import create_tiny_llama_dir from safetensors.torch import load_file from torch.distributed.tensor import DTensor +from transformers import AutoModelForCausalLM, LlamaConfig from modelopt.torch.export.unified_export_hf import export_hf_checkpoint -from modelopt.torch.utils.distributed import broadcast_state_dict -from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 +from modelopt.torch.utils.distributed import ( + DistributionPlan, + broadcast_state_dict, + distribution_plan, + set_distribution_plan, +) +from modelopt.torch.utils.plugins.model_load_utils import ( + load_pretrained, + parallel_load_and_prepare_fsdp2, +) VOCAB_SIZE = 64 @@ -61,6 +71,7 @@ def _test_parallel_load_and_export(rank, size, ckpt_dir, export_dir, cpu_offload size, cpu_offload=cpu_offload, ) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. decoder_params = list(model.model.layers[0].parameters()) @@ -106,6 +117,7 @@ def _test_carry_over_unplaced_weights(rank, size, ckpt_dir, export_dir, orphan_p """Load, export, and require the weights the model could not place to survive.""" device = torch.device(f"cuda:{rank}") model = parallel_load_and_prepare_fsdp2(ckpt_dir, device, rank, size) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) unplaced = getattr(model, "_modelopt_unplaced_source_keys", None) assert unplaced, ( @@ -171,3 +183,68 @@ def test_carry_over_unplaced_weights(dist_workers, tmp_path): orphan_prefix=f"model.layers.{orphan_idx}.", ) ) + + +# ------------------------------------------------------------------ load_pretrained, left to decide + +_TINY = { + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 128, + "max_position_embeddings": 64, +} + + +def _save(config, ckpt_dir): + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(ckpt_dir) + + +def _full(tensor): + return tensor.full_tensor() if isinstance(tensor, DTensor) else tensor + + +def _load_and_check(rank, size, ckpt_dir, knobs, expected): + """Load on every rank of the run with the plan left to decide, and check what each rank holds. + + Left to decide, ``load_pretrained`` FSDP2-shards a dense model when asked to, and otherwise gives + every rank the whole model on its own GPU. Whichever load runs, every rank ends up with the + checkpoint's weights and the load's record of what it could not place. + """ + device = torch.device(f"cuda:{rank}") + os.environ.update(knobs) + try: + model = load_pretrained(ckpt_dir) + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16) + + params = dict(model.named_parameters()) + whole_params = dict(whole.named_parameters()) + assert params.keys() == whole_params.keys() + sharded = {name for name, param in params.items() if isinstance(param, DTensor)} + if expected == "replicated": + assert not sharded + assert all(param.device == device for param in params.values()) + else: + assert {name for name in params if ".layers." in name} <= sharded + for name, param in params.items(): # the same order on every rank: full_tensor gathers + torch.testing.assert_close(_full(param).cpu(), whole_params[name]) + assert model._modelopt_unplaced_source_keys == [] + assert model._modelopt_source_checkpoint == ckpt_dir + dense = "replicate" if expected == "replicated" else "fsdp" + assert distribution_plan(model) == DistributionPlan(size, dense=dense) + finally: + for knob in knobs: + del os.environ[knob] + + +def test_a_dense_model_loads_whole_on_every_rank(need_2_gpus, dist_workers, tmp_path): + _save(LlamaConfig(**_TINY), tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {}, "replicated") + + +def test_a_dense_model_asked_to_shard_loads_fsdp2_sharded(need_2_gpus, dist_workers, tmp_path): + _save(LlamaConfig(**_TINY), tmp_path) + dist_workers.run(_load_and_check, str(tmp_path), {"MODELOPT_FSDP_SIZE": "world"}, "fsdp2") diff --git a/tests/unit/torch/utils/test_fsdp2_utils.py b/tests/unit/torch/utils/test_fsdp2_utils.py index b1bcadc933a..71e6fda5f24 100644 --- a/tests/unit/torch/utils/test_fsdp2_utils.py +++ b/tests/unit/torch/utils/test_fsdp2_utils.py @@ -13,12 +13,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for ``modelopt.torch.utils.fsdp2_utils``: which parameters stay out of FSDP2's shards.""" +"""Tests for ``modelopt.torch.utils.fsdp2_utils``: which parameters stay out of FSDP2's shards, and +the distribution plan a wrap leaves to its caller.""" import pytest import torch import torch.nn as nn +import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils import fsdp2_utils +from modelopt.torch.utils.distributed import ( + DistributionPlan, + distribution_plan, + set_distribution_plan, +) from modelopt.torch.utils.fsdp2_utils import _off_dtype_params @@ -61,3 +70,44 @@ def test_off_dtype_params_dominant_is_by_numel(params, expected): off = _off_dtype_params(model) kept = {n for n, p in model.named_parameters() if p not in off} assert kept == {expected} + + +def test_fsdp2_wrap_leaves_the_plan_to_the_load(monkeypatch): + """load_pretrained records the plan once, at the end of the load; the wrap it calls records none, + so a caller that wraps a model itself must record one before anything reads it.""" + model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + monkeypatch.setattr( + LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + ) + monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda module, **kwargs: None) + monkeypatch.setattr(dist, "size", lambda group=None: 4) + + fsdp2_utils.fsdp2_wrap(model) + + with pytest.raises(RuntimeError, match="no distribution plan recorded"): + distribution_plan(model) + + +@pytest.mark.parametrize( + ("plan", "recorded_on"), + [ + (DistributionPlan(), "model"), + (DistributionPlan(4, dense="fsdp"), "model"), + (DistributionPlan(), "submodule"), + ], + ids=["replicated", "sharded", "in-a-submodule"], +) +def test_fsdp2_wrap_refuses_a_model_with_a_plan_recorded(monkeypatch, plan, recorded_on): + """A plan is recorded once, so a model that records one, as load_pretrained records one on every + model it loads, is refused rather than sharded behind it: a replicated model would read as + replicated, which the export would believe, and a sharded one would be wrapped twice.""" + model = torch.nn.Sequential(torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)) + set_distribution_plan(model if recorded_on == "model" else model[0], plan) + monkeypatch.setattr( + LayerActivationCollector, "get_decoder_layers", staticmethod(lambda model: list(model)) + ) + monkeypatch.setattr(fsdp2_utils, "fully_shard", lambda *a, **k: pytest.fail("sharded")) + + # A ValueError, as set_distribution_plan raises for the same record: one type to catch. + with pytest.raises(ValueError, match="already recorded as distributed"): + fsdp2_utils.fsdp2_wrap(model) diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py index ca9e71b97f5..9266e601180 100644 --- a/tests/unit/torch/utils/test_model_load_utils.py +++ b/tests/unit/torch/utils/test_model_load_utils.py @@ -15,12 +15,14 @@ """Tests for ``modelopt.torch.utils.plugins.model_load_utils``. -Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` answer the process group's rank -and size as a torchrun launch would, and serve tiny real transformers configs from memory, so the -decision builds their real models on the meta device; only the last ones use a process group, two -gloo ranks. +Nothing here needs a GPU. The tests of ``resolve_distribution_plan`` and of the loads +``load_pretrained`` routes a plan to answer the process group's rank and size as a torchrun launch +would, and serve tiny real transformers configs from memory: the decision builds their real models on +the meta device, and stand-ins for the sharded loaders record how they were called. Only the last +tests of each part use a process group, two gloo ranks. """ +import inspect import math import os from contextlib import contextmanager @@ -49,6 +51,7 @@ AutoModelForSeq2SeqLM, AutoModelForSpeechSeq2Seq, LlamaConfig, + LlamaForCausalLM, Phi4MultimodalConfig, Qwen2_5_VLConfig, Qwen3MoeConfig, @@ -56,9 +59,13 @@ WhisperConfig, ) +import modelopt.torch.models.hf as models_hf +import modelopt.torch.quantization as mtq import modelopt.torch.utils.distributed as dist +from modelopt.torch.quantization.plugins import accelerate as accelerate_plugin +from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -from modelopt.torch.utils.distributed import DistributionPlan +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( _auto_class, @@ -66,6 +73,7 @@ _convert_keys, _resolve_target, build_meta_causal_lm, + load_pretrained, record_unplaced_source_keys, resolve_distribution_plan, unplaced_source_keys, @@ -236,17 +244,27 @@ def _gathered_alike(gathered, mine, group=None): def run_as(monkeypatch): """Run as rank ``rank`` of a ``world_size``-rank launch, deciding for itself. - Every rank reads the checkpoint as this one does. How the ranks agree on one plan is tested at - the end of this file. + Every rank acts as its group's first, and reads the checkpoint as this one does. How the ranks + agree on one plan, and on one load, is tested on two gloo ranks at the end of each part. """ monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0) monkeypatch.setattr(torch.distributed, "broadcast_object_list", lambda *a, **k: None) monkeypatch.setattr(torch.distributed, "all_gather_object", _gathered_alike) + broadcast_from = dist.broadcast_from + + def broadcast_as_the_first(label, fn, src=0): + # broadcast_from sends from dist.rank(), which process_group sets for the FSDP2 load. + with monkeypatch.context() as sending: + sending.setattr(dist, "rank", lambda group=None: src) + return broadcast_from(label, fn, src) + + monkeypatch.setattr(dist, "broadcast_from", broadcast_as_the_first) def enter(world_size, rank=0): monkeypatch.setattr(dist, "is_distributed_run", lambda: world_size > 1) monkeypatch.setattr(dist, "size", lambda group=None: world_size) monkeypatch.setattr(dist, "is_master", lambda group=None: rank == 0) + monkeypatch.setattr(dist, "local_rank", lambda: rank) monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: world_size) return enter @@ -486,31 +504,47 @@ def test_an_ep_degree_that_does_not_divide_the_world_names_where_it_came_from( resolve_distribution_plan(str(tmp_path), **arguments) +@pytest.mark.parametrize("value", [[2], 2.9], ids=["a-list", "a-fraction"]) +def test_an_ep_size_that_is_no_whole_number_is_refused(run_as, serve, tmp_path, value): + """An argument int() fails on with a TypeError, or would truncate, is no degree either.""" + run_as(4) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(ValueError, match=r"ep_size=.*EP degree"): + resolve_distribution_plan(str(tmp_path), ep_size=value) + + +@pytest.mark.parametrize("value", ["world", "-2"], ids=["word", "below-1"]) @pytest.mark.parametrize( ("world_size", "source"), [(4, "MODELOPT_EP_DEGREE"), (1, "ep_size")], ids=["knob", "argument-in-a-single-process"], ) def test_an_ep_degree_that_is_not_a_number_names_where_it_came_from( - run_as, serve, monkeypatch, tmp_path, world_size, source + run_as, serve, monkeypatch, tmp_path, world_size, source, value ): - """A single process takes its own argument as given, but parses it as the decision does.""" + """A single process takes its own argument as given, but parses it as the decision does, and a + degree below 1 is no degree either.""" run_as(world_size) serve(LlamaConfig(**_TINY)) arguments = {} if source == "ep_size": - arguments["ep_size"] = "world" + arguments["ep_size"] = value else: - monkeypatch.setenv(source, "world") + monkeypatch.setenv(source, value) - with pytest.raises(ValueError, match=f"{source}='world'"): + with pytest.raises(ValueError, match=f"{source}='{value}'"): resolve_distribution_plan(str(tmp_path), **arguments) @needs_qwen3_moe_ep_plan -@pytest.mark.parametrize("value", ["0", "", "AUTO"], ids=["zero", "exported-empty", "capitals"]) +@pytest.mark.parametrize( + "value", + ["0", "", " ", "AUTO", " auto "], + ids=["zero", "exported-empty", "exported-blank", "capitals", "padded"], +) def test_an_ep_degree_of_0_is_auto(run_as, serve, monkeypatch, tmp_path, value): - """And so is one exported empty, and auto in any case: the plan decides.""" + """And so is one exported empty or blank, and auto in any case or padded: the plan decides.""" monkeypatch.setenv("MODELOPT_EP_DEGREE", value) run_as(4) serve(_qwen3_moe(num_experts=8)) @@ -590,15 +624,17 @@ def test_arguments_override_the_env_knobs(run_as, serve, monkeypatch, tmp_path): ("4", "fsdp"), ("WORLD", "fsdp"), ("", "fsdp"), + (" ", "fsdp"), + (" world ", "fsdp"), ("2", None), ], - ids=["1", "4", "WORLD", "exported-empty", "part-of-the-world"], + ids=["1", "4", "WORLD", "exported-empty", "exported-blank", "padded", "part-of-the-world"], ) def test_the_dense_knob_replicates_or_shards_over_the_whole_world( run_as, serve, gpu_memory, monkeypatch, tmp_path, value, dense ): - """Exported empty, the knob is auto, which shards a model too big for one GPU. Sharding over - part of the world is refused: no layout does that.""" + """Exported empty or blank, the knob is auto, which shards a model too big for one GPU. Sharding + over part of the world is refused: no layout does that.""" monkeypatch.setenv("MODELOPT_FSDP_SIZE", value) run_as(4) serve(LlamaConfig(**_TINY)) @@ -735,6 +771,17 @@ def remote_from_config(config, **kwargs): assert ProbeConfig not in mapping +def test_the_config_returned_with_the_plan_is_as_every_rank_read_it(run_as, serve, tmp_path): + """Rank 0 alone decides, building the model, which writes to the config it builds from; a + sharded load builds from the config returned, so rank 0's must be the one the others read.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + + _, hf_config = resolve_distribution_plan(str(tmp_path)) + + assert vars(hf_config) == vars(LlamaConfig(**_TINY)) + + def test_the_hub_options_locate_the_checkpoint_the_load_reads(tmp_path): """``revision``, ``cache_dir``, ``local_files_only`` and ``subfolder`` find the checkpoint as ``load_pretrained`` takes them, so the plan and its config are those of what the load reads.""" @@ -791,23 +838,995 @@ def test_every_rank_takes_the_plan_rank0_decided(cpu_dist_workers, tmp_path): ) -def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt): +def _one_rank_cannot_see_the_checkpoint_job(rank, size, ckpt, load): isdir = os.path.isdir with ( bounded_collectives(), mock.patch("os.path.isdir", lambda path: (rank == 0 or path != ckpt) and isdir(path)), ): try: - resolve_distribution_plan(ckpt) + if load: + load_pretrained(ckpt, distribution=DistributionPlan(size), device_map="cpu") + else: + resolve_distribution_plan(ckpt) except RuntimeError as error: raised = str(error) else: raised = "no error" - # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the plan's broadcast. + # Rank 1's own error, carried to rank 0, rather than rank 0 timing out in the next collective. assert "which this rank cannot see" in raised, raised -def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path): - """Before the plan's broadcast, which rank 0 would otherwise wait in for a rank that raised.""" +@pytest.mark.parametrize("load", [False, True], ids=["plan", "load"]) +def test_a_rank_that_cannot_see_the_checkpoint_fails_every_rank(cpu_dist_workers, tmp_path, load): + """Deciding the plan, or loading by a handed one: rank 0 would otherwise wait for the rank that + raised in the next collective, the plan's broadcast or the load's checks. The group is the + caller's and the error is caught, so nothing else ends rank 0's wait.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path), load=load) + + +# ------------------------------------------------------------------------------ loading by the plan + + +class _Loaded(torch.nn.Module): + """What a sharded loader returns: a model whose load recorded what it could not place.""" + + def __init__(self): + super().__init__() + self._modelopt_unplaced_source_keys = [] + + +@pytest.fixture +def loaders(monkeypatch): + """Stand-ins for the sharded loaders, each recording its call by parameter name.""" + calls = {} + + def recording(name, loader): + def stub(*args, **kwargs): + calls[name] = inspect.signature(loader).bind(*args, **kwargs).arguments + return _Loaded() + + return stub + + monkeypatch.setattr( + model_load_utils, + "parallel_load_and_prepare_fsdp2", + recording("fsdp2", model_load_utils.parallel_load_and_prepare_fsdp2), + ) + return calls + + +@pytest.fixture +def auto_class(): + """An ``auto_class`` whose from_pretrained records its arguments.""" + + class Auto: + calls: list = [] + + @classmethod + def from_pretrained(cls, ckpt_path, **kwargs): + cls.calls.append(kwargs) + return torch.nn.Linear(1, 1) + + return Auto + + +def _pin_cuda(monkeypatch, gpus, bound=None): + """Pin the CUDA state a rank's GPU is read from: ``gpus`` visible, and the process bound to + cuda:``bound``, or to none yet, so its local rank picks one. Tests then read the case rather + than the host, whose process may have set CUDA up already.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: gpus > 0) + monkeypatch.setattr(torch.cuda, "device_count", lambda: gpus) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: bound is not None) + monkeypatch.setattr(torch.cuda, "current_device", lambda: bound or 0) + + +@pytest.fixture +def a_gpu(monkeypatch): + """This rank has a GPU of its own, which the FSDP2 load shards onto; the stand-in loaders never + use it.""" + _pin_cuda(monkeypatch, 8) + + +@pytest.fixture +def process_group(monkeypatch, a_gpu): + """Rank 1 of a 2-rank process group, as the FSDP2 load reads it.""" + monkeypatch.setattr(dist, "is_initialized", lambda: True) + monkeypatch.setattr(dist, "rank", lambda group=None: 1) + monkeypatch.setattr(dist, "size", lambda group=None: 2) + monkeypatch.setattr(dist, "local_rank", lambda: 1) + + +# ----------------------------------------------------------------------------------- the loads + + +@pytest.mark.parametrize( + ("passed", "gpus", "placed"), + [ + ({}, 4, {"device_map": {"": "cuda:1"}}), + ({"device_map": "cpu"}, 4, {"device_map": "cpu"}), + ({"device_map": None}, 4, {"device_map": None}), + ({}, 1, {}), + ({"mp_policy": None, "cpu_offload": None}, 4, {"device_map": {"": "cuda:1"}}), + # As hf_ptq passes --cpu_offload when it is not set: the FSDP2 load's own default. + ({"cpu_offload": False}, 4, {"device_map": {"": "cuda:1"}}), + ({"broadcast_chunk_size": 8}, 4, {"device_map": {"": "cuda:1"}}), + ], + ids=[ + "defaults_to_the_rank_gpu", + "caller_placement_wins", + "none_loads_on_the_cpu", + "no_gpu_of_its_own", + "fsdp2_options_at_none", + "cpu_offload_at_its_default", + "broadcast_chunk_size_at_its_default", + ], +) +def test_a_replicated_plan_loads_the_whole_model_on_every_rank( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, passed, gpus, placed +): + """Each rank takes the whole model onto its own GPU; "auto" would let every rank's accelerate + pick the same idle GPU. device_map=None is from_pretrained's own: a load on the CPU, which the + caller places, as hf_ptq does. A rank with none of its own, on a host with fewer GPUs than + ranks, is left to from_pretrained's placement rather than handed a device it does not have. The + FSDP2 load's own options left unset or at its defaults ask nothing and do not reach + from_pretrained.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + _pin_cuda(monkeypatch, gpus) + + model = load_pretrained(str(tmp_path), auto_class, **passed) + + assert auto_class.calls == [placed] + assert distribution_plan(model) == DistributionPlan(4) + assert loaders == {} + + +def test_a_rank_bound_to_its_one_visible_gpu_loads_onto_it( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path +): + """A launch that exposes one GPU to each process binds every rank to cuda:0, whatever its local + rank: each takes the whole model onto the GPU it is bound to, rather than read as having + none.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + _pin_cuda(monkeypatch, 1, bound=0) + + load_pretrained(str(tmp_path), auto_class) + + assert auto_class.calls == [{"device_map": {"": "cuda:0"}}] + + +def test_a_rank_without_a_gpu_of_its_own_says_so( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path +): + """On a host with fewer GPUs than ranks, a rank of a replicated load left without one loads the + model where from_pretrained places it while its peers load onto theirs: it warns itself, as rank + 0 may have a GPU of its own and nothing to say.""" + run_as(4, rank=1) + serve(LlamaConfig(**_TINY)) + _pin_cuda(monkeypatch, 1) + + with pytest.warns(UserWarning, match="local rank 1"): + load_pretrained(str(tmp_path), auto_class) + assert auto_class.calls == [{}] + + +@pytest.mark.parametrize( + "defaults", + [ + {"dtype": "auto"}, + {"dtype": None}, + {"torch_dtype": "auto"}, + # As hf_ptq passes them: placement options at from_pretrained's own default. + { + "device_map": None, + "max_memory": None, + "offload_folder": None, + "quantization_config": None, + }, + ], + ids=["auto", "none", "torch_dtype", "placement-none"], +) +def test_an_fsdp_plan_loads_through_the_parallel_fsdp2_loader( + run_as, loaders, process_group, monkeypatch, tmp_path, defaults +): + """FSDP2 shards over the process group it reads its rank and GPU from, by a config read once, + with the Hub options the load was given (custom code an ``auto_map`` names in another repo). + It keeps the checkpoint's dtype, as from_pretrained does for "auto" and for None, under either + name of the argument, and takes a placement option passed as None as not passed.""" + run_as(2, rank=1) + served, reads = LlamaConfig(**_TINY), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(k) or served + ) + plan = DistributionPlan(2, dense="fsdp") + + model = load_pretrained( + str(tmp_path), distribution=plan, cpu_offload=True, local_files_only=True, **defaults + ) + + call = loaders.pop("fsdp2") + assert (call["device"], call["rank"], call["world_size"]) == (torch.device("cuda:1"), 1, 2) + assert call["cpu_offload"] is True + assert call["hf_config"] is served + assert [read.get("local_files_only") for read in reads] == [True] + assert distribution_plan(model) == plan + assert loaders == {} + + +@pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) +def test_a_text_only_vlm_the_fsdp2_load_builds_reaches_it( + run_as, serve, loaders, process_group, tmp_path, model_type +): + """The FSDP2 load builds Qwen3.5's VLMs through their image-text class, as hf_ptq's --use_fsdp2 + does, so a plan that shards one reaches the load rather than being refused as no causal LM.""" + try: + config = serve(AutoConfig.for_model(model_type)) + except ValueError: + pytest.skip("Qwen3.5 requires a recent Transformers") + run_as(2, rank=1) + + load_pretrained(str(tmp_path), distribution=DistributionPlan(2, dense="fsdp")) + + assert loaders.pop("fsdp2")["hf_config"] is config + + +@pytest.mark.parametrize("handed", [True, False], ids=["handed", "decided"]) +def test_a_sharded_load_offers_the_knob_only_for_a_plan_it_decided( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """MODELOPT_FSDP_SIZE counts only for a plan the load decides: to keep a model whole under a + plan of its own, a caller passes a replicated one, which is all the refusal offers it.""" + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with pytest.raises(TypeError, match=r"\['device_map'\] do not apply") as refused: + load_pretrained(str(tmp_path), device_map="auto", **plan) + assert ("MODELOPT_FSDP_SIZE=1" in str(refused.value)) is not handed + assert loaders == {} + + +@pytest.mark.parametrize( + "unset", + [{}, {"config": None}, {"variant": None}, {"hf_config": None}], + ids=["none_passed", "config_none", "variant_none", "hf_config_none"], +) +def test_a_load_it_decides_reads_the_config_once( + run_as, loaders, process_group, monkeypatch, tmp_path, unset +): + """The config the plan was decided from is the one the FSDP2 load gets, not a second read. An + option passed as None, its default, is as if not passed, not one the FSDP2 load refuses.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2, rank=1) + config, reads = LlamaConfig(**_TINY), [] + monkeypatch.setattr( + transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a) or config + ) + + model = load_pretrained(str(tmp_path), **unset) + + assert len(reads) == 1 + assert loaders["fsdp2"]["hf_config"] is config + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + +@pytest.mark.parametrize("handed", [False, True], ids=["decided", "handed"]) +def test_a_config_passed_for_the_load_is_the_one_it_decides_from( + run_as, loaders, process_group, monkeypatch, tmp_path, handed +): + """No config is read: the plan is decided from the one passed for the load, which the FSDP2 + load then builds from, as it does with a plan handed in. Building writes to that config, so the + load works on a copy of it.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2, rank=1) + reads: list = [] + monkeypatch.setattr(transformers.AutoConfig, "from_pretrained", lambda *a, **k: reads.append(a)) + config = LlamaConfig(**{**_TINY, "num_hidden_layers": 2}) + as_passed = config.to_dict() + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + model = load_pretrained(str(tmp_path), config=config, **plan) + + assert reads == [] + assert loaders["fsdp2"]["hf_config"].num_hidden_layers == 2 + assert loaders["fsdp2"]["hf_config"] is not config # building it writes to it + assert config.to_dict() == as_passed + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + +def test_hf_config_is_refused_without_the_plan_it_came_with( + run_as, serve, loaders, process_group, monkeypatch, tmp_path +): + """A sharded load builds from hf_config as given, so the plan must be the one it came with: a + plan decided from the checkpoint's own config would load a model it was not decided for.""" + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + + with pytest.raises(TypeError, match=r"hf_config.*resolve_distribution_plan"): + load_pretrained(str(tmp_path), hf_config=LlamaConfig(**{**_TINY, "num_hidden_layers": 2})) + assert loaders == {} + + +class _OwnConfig(LlamaConfig): + """A config class of the caller's own, which AutoConfig has no mapping for.""" + + model_type = "modelopt_test_own_llama" + + +@pytest.mark.parametrize( + ("fsdp_size", "own_config"), + [(None, False), ("world", False), (None, True)], + ids=["no-knob", "knob-shards", "own-config-class"], +) +def test_a_model_class_of_the_callers_own_loads_whole( + run_as, serve, loaders, auto_class, monkeypatch, tmp_path, fsdp_size, own_config +): + """The sharded loads build their own model class and refuse another, so with nothing asking to + shard, a class of the caller's own loads whole on every rank, and with a knob asking to shard + it is refused. Neither builds the model on the meta device, which the decision's own pick of + class may not build (an encoder-decoder loaded through AutoModel). The decision reads the + checkpoint's config as the class does, through its own config class, which AutoConfig need not + know.""" + run_as(2) + if own_config: + _OwnConfig(**_TINY).save_pretrained(tmp_path) + auto_class.config_class = _OwnConfig + else: + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr( + model_load_utils, "build_meta_causal_lm", lambda *a, **k: pytest.fail("built a model") + ) + if fsdp_size is None: + model = load_pretrained(str(tmp_path), auto_class) + assert distribution_plan(model) == DistributionPlan(2) + assert len(auto_class.calls) == 1 + else: + monkeypatch.setenv("MODELOPT_FSDP_SIZE", fsdp_size) + with pytest.raises(NotImplementedError, match="auto_class=Auto "): + load_pretrained(str(tmp_path), auto_class) + assert loaders == {} + + +@pytest.mark.parametrize("handed", [True, False], ids=["handed", "decided"]) +def test_a_model_class_given_by_name_is_refused_by_name( + run_as, serve, loaders, monkeypatch, tmp_path, handed +): + """A sharded load refuses a model class it does not build, given by name too, rather than fail + on the name having no __name__ of its own.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with pytest.raises(NotImplementedError, match="auto_class='AutoModelForCausalLM'"): + load_pretrained(str(tmp_path), "AutoModelForCausalLM", **plan) + assert loaders == {} + + +@needs_qwen3_moe_ep_plan +@pytest.mark.parametrize("knob", [None, "4"], ids=["left-to-decide", "asked-for"]) +def test_a_load_it_decides_keeps_the_experts_whole( + run_as, serve, loaders, monkeypatch, tmp_path, knob +): + """No load here shards the routed experts yet, so an MoE whose EP degree is left to the decision + is decided without expert parallelism and loads, where an automatic degree would be one the + load refuses; a degree MODELOPT_EP_DEGREE asks for is still refused, not ignored.""" + run_as(4) + serve(_qwen3_moe()) + if knob is not None: + monkeypatch.setenv("MODELOPT_EP_DEGREE", knob) + monkeypatch.setattr( + transformers.AutoModelForCausalLM, + "from_pretrained", + classmethod(lambda cls, *a, **k: torch.nn.Linear(1, 1)), + ) + + if knob is not None: + with pytest.raises(NotImplementedError, match=r"ep_size=4.*MODELOPT_EP_DEGREE=1"): + load_pretrained(str(tmp_path)) + return + assert distribution_plan(load_pretrained(str(tmp_path))) == DistributionPlan(4) + + +def test_an_expert_parallel_plan_is_refused(run_as, serve, loaders, auto_class, tmp_path): + """No load here shards the routed experts: an EP plan fails before anything loads.""" + run_as(2) + serve(Qwen3MoeConfig(num_experts=8, moe_intermediate_size=8, **_TINY)) + + with pytest.raises(NotImplementedError, match=r"ep_size=2.*MODELOPT_EP_DEGREE=1"): + load_pretrained(str(tmp_path), auto_class, distribution=DistributionPlan(2, ep_size=2)) + assert loaders == {} + assert auto_class.calls == [] + + +def _every_rank_loads_rank0s_plan_job(rank, size, ckpt, handed, placed): + model, refused = None, None + # Ranks without GPUs, whatever the host has: the pooled workers keep no CUDA context. + with bounded_collectives(), mock.patch.object(torch.cuda, "is_available", return_value=False): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + model = load_pretrained(ckpt, distribution=handed[rank], **placed) + except Exception as error: + refused = str(error) + if isinstance(handed[0], DistributionPlan): + assert refused is None and distribution_plan(model) == handed[0], (rank, refused) + else: + # Rank 0's refusal, on every rank, rather than a rank left waiting in the broadcast. + assert "expected a DistributionPlan" in str(refused), (rank, refused) + + +@pytest.mark.parametrize( + ("handed", "placed"), + [ + ((DistributionPlan(2), DistributionPlan(2, dense="fsdp")), {"device_map": "cpu"}), + ((DistributionPlan(2), None), {"device_map": "cpu"}), + (("fsdp2", DistributionPlan(2)), {"device_map": "cpu"}), + ((DistributionPlan(2), DistributionPlan(2)), {}), + ], + ids=["another-plan", "none", "not-a-plan", "placed-by-default"], +) +def test_every_rank_loads_the_plan_rank0_was_handed(cpu_dist_workers, tmp_path, handed, placed): + """Over a real broadcast, whatever rank 1 was handed. Ranks loading different plans would build + mismatched meshes, and one handed none would decide a plan through collectives the others never + join: both would hang. A distribution rank 0 refuses fails every rank. Placed by default, each + rank takes the model onto its GPU, and a run without GPUs, as these ranks are, loads it whole.""" + LlamaForCausalLM(LlamaConfig(**_TINY)).save_pretrained(tmp_path) + cpu_dist_workers.run( + _every_rank_loads_rank0s_plan_job, ckpt=str(tmp_path), handed=handed, placed=placed + ) + + +def _fetched_in_step(ckpt_path, **hub_options): + """The checkpoint's fetch, as the collective it is: a rank not in it leaves the others waiting.""" + torch.distributed.barrier() + return None, ckpt_path + + +def _loaded_in_step(*args, **kwargs): + """The FSDP2 load, as the collectives it runs, which a rank without a GPU fails before.""" + if dist.local_rank() >= torch.cuda.device_count(): + raise RuntimeError("no GPU to shard onto") + torch.distributed.barrier() + return _Loaded() + + +_FSDP2_PLAN = DistributionPlan(2, dense="fsdp") + + +def _one_rank_refused_job( + rank, size, ckpt, refused_on_1, gpus_on_1, match, knobs_on_1=None, plan=_FSDP2_PLAN +): + refused = None + gpus = size if rank == 0 else gpus_on_1 # what this rank's host shows it + with ( + bounded_collectives(), + pytest.MonkeyPatch.context() as cuda, + mock.patch.dict(os.environ, (knobs_on_1 or {}) if rank == 1 else {}), + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", _loaded_in_step), + ): + _pin_cuda(cuda, gpus) # the pooled worker's own CUDA state is whatever an earlier job left + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_pretrained(ckpt, distribution=plan, **(refused_on_1 if rank == 1 else {})) + except Exception as error: + refused = str(error) + # Rank 1's refusal on both ranks, rather than rank 0 failing at the group's timeout. + assert refused is not None and match in refused, (rank, refused) + assert rank == 1 or "rank 1 failed" in refused, refused + + +@pytest.mark.parametrize( + ("refused_on_1", "gpus_on_1", "match"), + [ + ( + {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}, + 2, + "config and hf_config", + ), + ({"dtype": torch.float32}, 2, "dtype=torch.float32"), + ( + {"hf_config": LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"})}, + 2, + "dequantize", + ), + ({}, 0, "as local rank 1"), + ({}, 1, "as local rank 1"), + ], + ids=["two-configs", "dtype", "quantized-config", "no-gpu", "fewer-gpus-than-ranks"], +) +def test_a_load_refused_on_one_rank_fails_every_rank( + cpu_dist_workers, tmp_path, refused_on_1, gpus_on_1, match +): + """Rank 1 alone is refused, by its arguments before the checkpoint is fetched, or by the config + it was handed or its lack of a GPU before the FSDP2 load: every rank raises the refusal, rather + than leave rank 0 waiting in the fetch's or the load's collectives, which stand-ins on a bounded + group run here.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run( + _one_rank_refused_job, + ckpt=str(tmp_path), + refused_on_1=refused_on_1, + gpus_on_1=gpus_on_1, + match=match, + ) + + +def test_a_knob_one_rank_cannot_parse_fails_every_rank(cpu_dist_workers, tmp_path): + """With no plan handed, rank 1 alone reads a malformed MODELOPT_EP_DEGREE: every rank raises it, + rather than leave rank 0 waiting in the plan decision's collectives.""" LlamaConfig(**_TINY).save_pretrained(tmp_path) - cpu_dist_workers.run(_one_rank_cannot_see_the_checkpoint_job, ckpt=str(tmp_path)) + cpu_dist_workers.run( + _one_rank_refused_job, + ckpt=str(tmp_path), + refused_on_1={}, + gpus_on_1=2, + match="MODELOPT_EP_DEGREE='eigth'", + knobs_on_1={"MODELOPT_EP_DEGREE": "eigth"}, + plan=None, + ) + + +def _one_visible_gpu_per_process_job(rank, size, ckpt): + devices = [] + + def load(ckpt_path, device, *args, **kwargs): + devices.append(device) + torch.distributed.barrier() + return _Loaded() + + with ( + bounded_collectives(), + pytest.MonkeyPatch.context() as cuda, + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + mock.patch.object(model_load_utils, "parallel_load_and_prepare_fsdp2", load), + ): + # Each rank sees only its own GPU, which its process bound as cuda:0, whatever its local + # rank. + _pin_cuda(cuda, 1, bound=0) + load_pretrained(ckpt, distribution=DistributionPlan(2, dense="fsdp")) + assert devices == [torch.device("cuda", 0)], (rank, devices) + + +def test_one_visible_gpu_per_process_shards_on_every_rank(cpu_dist_workers, tmp_path): + """Two ranks, each shown only its own GPU and bound to it as cuda:0, both shard onto it: rank 1, + whose local rank numbers no GPU it sees, is not refused as having none.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_visible_gpu_per_process_job, ckpt=str(tmp_path)) + + +class _LoadsWholeExceptOnRank1: + """A model class whose whole-model load fails on rank 1 alone, as an out-of-memory GPU would.""" + + @classmethod + def from_pretrained(cls, ckpt_path, **kwargs): + if dist.rank() == 1: + raise OSError("rank 1 could not load the model") + return torch.nn.Linear(1, 1) + + +def _one_rank_fails_to_load_job(rank, size, ckpt): + refused = None + with ( + bounded_collectives(), + mock.patch.object(model_load_utils, "ensure_local_checkpoint", _fetched_in_step), + ): + try: # not pytest.raises: its failure would end the pooled worker instead of the job + load_pretrained(ckpt, _LoadsWholeExceptOnRank1, distribution=DistributionPlan(2)) + except Exception as error: + refused = str(error) + # Rank 1's failure on both ranks, rather than rank 0 going on to the next collective alone. + assert refused is not None and "could not load the model" in refused, (rank, refused) + + +def test_a_whole_load_that_fails_on_one_rank_fails_every_rank(cpu_dist_workers, tmp_path): + """A whole-model load runs no collective, so a rank failing in it alone would leave the others + to wait for it in the next one, calibration's or the export's.""" + LlamaConfig(**_TINY).save_pretrained(tmp_path) + cpu_dist_workers.run(_one_rank_fails_to_load_job, ckpt=str(tmp_path)) + + +@pytest.mark.parametrize( + ("change", "error", "match"), + [ + ({"kwargs": {"device_map": "auto"}}, TypeError, r"\['device_map'\] do not apply"), + ({"kwargs": {"dtype": torch.bfloat16}}, TypeError, r"dtype=torch\.bfloat16"), + ( + {"kwargs": {"torch_dtype": torch.bfloat16}}, + TypeError, + r"torch_dtype=torch\.bfloat16", + ), + ( + {"auto_class": AutoModel}, + NotImplementedError, + r"auto_class=AutoModel: leave it out.*DistributionPlan\(2, dense='replicate'\)", + ), + ({"mtp": True}, NotImplementedError, "MTP head"), + ( + {"kwargs": {"config": LlamaConfig(**_TINY), "hf_config": LlamaConfig(**_TINY)}}, + TypeError, + "config and hf_config", + ), + ], + ids=["placement", "dtype", "torch_dtype", "model_class", "mtp_head", "two_configs"], +) +def test_a_sharded_load_refuses_what_it_cannot_honor( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, change, error, match +): + """Placement, a model class, a head only a whole model builds, two configs to load with: none + reaches a sharded load, which fails before it starts rather than load something other than + what was asked.""" + run_as(2) + serve(LlamaConfig(**_TINY)) + if change.get("mtp"): + monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + + with pytest.raises(error, match=match): + load_pretrained( + str(tmp_path), + change.get("auto_class"), + distribution=DistributionPlan(2, dense="fsdp"), + **change.get("kwargs", {}), + ) + assert loaders == {} + + +@pytest.mark.parametrize( + ("distribution", "checkpoint", "options", "error", "match"), + [ + ( + DistributionPlan(4, dense="fsdp"), + "causal-lm", + None, + ValueError, + "is for 4 ranks", + ), + ( + DistributionPlan(2, dense="fsdp"), + "quantized", + None, + NotImplementedError, + "dequantize", + ), + ( + DistributionPlan(2, dense="fsdp"), + "quantized-language-model", + {"trust_remote_code": True}, + NotImplementedError, + "dequantize", + ), + ( + DistributionPlan(2, dense="fsdp"), + "vlm", + None, + NotImplementedError, + "llava", + ), + ("fsdp2", "causal-lm", None, TypeError, "expected a DistributionPlan"), + ( + DistributionPlan(2), + "causal-lm", + "hf_config", + TypeError, + r"hf_config.*resolve_distribution_plan", + ), + ( + DistributionPlan(2), + "causal-lm", + {"cpu_offload": True}, + TypeError, + r"\['cpu_offload'\] only apply", + ), + ], + ids=[ + "another-world", + "quantized-checkpoint", + "quantized-language-model", + "vlm-checkpoint", + "not-a-plan", + "hf-config-on-a-whole-load", + "fsdp2-option-on-a-whole-load", + ], +) +def test_a_load_refuses_what_it_cannot_load( + run_as, + serve, + loaders, + auto_class, + process_group, + tmp_path, + distribution, + checkpoint, + options, + error, + match, +): + """Before anything loads: a plan for another world; a quantized checkpoint, which FSDP2 would + read as stored, even when only a VLM's language model is quantized and its remote code builds it + as a causal LM; a VLM but a text-only one it supports, which FSDP2 would build through + AutoModelForCausalLM, refused or as its language model alone; anything but a plan; and the + sharded loads' config, or an option only the FSDP2 load takes, which from_pretrained would take + for a model argument and fail far from here, or drop.""" + run_as(2) + config = serve(transformers.LlavaConfig() if checkpoint == "vlm" else LlamaConfig(**_TINY)) + if checkpoint == "quantized": + config.quantization_config = {"quant_method": "fp8"} + if checkpoint == "quantized-language-model": # Kimi-K2.5's layout + config.auto_map = _REMOTE_VLM + config.text_config = LlamaConfig(**_TINY, quantization_config={"quant_method": "fp8"}) + if options == "hf_config": + options = {"auto_class": auto_class, "hf_config": config} + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **(options or {})) + assert loaders == {} + assert auto_class.calls == [] + + +@pytest.mark.parametrize("device_map", ["auto", None], ids=["placed", "none"]) +def test_an_automatic_layout_says_so_when_it_refuses_a_placement( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, device_map +): + """With neither a plan nor MODELOPT_FSDP_SIZE, rank 0's GPU memory decides to shard a model too + big for it, and the sharded load refuses options it does not take. The refusal says the layout + was decided rather than asked for, and how to keep the load whole. device_map=None, as hf_ptq + passes it, places nothing, and the model is sharded.""" + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + + if device_map is None: + load_pretrained(str(tmp_path), AutoModelForCausalLM, device_map=None, dtype="auto") + assert "fsdp2" in loaders + return + with pytest.raises(TypeError, match=r"\['device_map'\] do not apply.*MODELOPT_FSDP_SIZE=1"): + load_pretrained(str(tmp_path), device_map="auto") + assert loaders == {} + + +class _ContextLoadReachedError(Exception): + """Raised where init_quantized_weights' from_pretrained builds the model, past its checks.""" + + +@pytest.mark.parametrize("memory", [1 << 50, 1], ids=["fits", "too-big"]) +def test_a_load_it_decides_inside_init_quantized_weights_is_whole_and_unplaced( + run_as, loaders, process_group, monkeypatch, tmp_path, memory +): + """Inside init_quantized_weights, with nothing asking to shard, every rank gets the whole model, + whatever rank 0's GPU memory, and the load leaves the device map to the context, which computes + it and refuses one passed: the context's from_pretrained gets past that refusal.""" + run_as(2, rank=1) + LlamaConfig(**_TINY).save_pretrained(tmp_path) # read from disk, as the context reads it too + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: memory) + monkeypatch.setattr( + accelerate_plugin, "init_empty_weights", mock.Mock(side_effect=_ContextLoadReachedError) + ) + + with init_quantized_weights(mtq.INT8_DEFAULT_CFG), pytest.raises(_ContextLoadReachedError): + load_pretrained(str(tmp_path)) + assert loaders == {} + + +@pytest.mark.parametrize("handed", [False, True], ids=["knob", "handed"]) +def test_a_sharded_load_is_refused_inside_init_quantized_weights( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, handed +): + """init_quantized_weights quantizes and compresses in the from_pretrained it patches, which the + FSDP2 load never calls: a knob or a plan asking to shard there is refused rather than return the + model unquantized.""" + run_as(2, rank=1) + serve(LlamaConfig(**_TINY)) + monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world") + plan = {"distribution": DistributionPlan(2, dense="fsdp")} if handed else {} + + with ( + init_quantized_weights(mtq.INT8_DEFAULT_CFG), + pytest.raises(NotImplementedError, match=r"init_quantized_weights.*dense='replicate'"), + ): + load_pretrained(str(tmp_path), **plan) + assert loaders == {} + + +@pytest.mark.parametrize( + ("checkpoint", "match"), + [("quantized", "dequantize"), ("mtp", "MTP head")], + ids=["quantized", "mtp"], +) +def test_a_measured_layout_names_the_way_back_when_it_refuses_a_checkpoint( + run_as, serve, loaders, process_group, monkeypatch, tmp_path, checkpoint, match +): + """Nothing asked for the layout rank 0's GPU memory chose, so a checkpoint the FSDP2 load cannot + take is refused with the way to load it whole, as every refusal of that load says.""" + run_as(2, rank=1) + config = serve(LlamaConfig(**_TINY)) + monkeypatch.setattr(model_load_utils, "_gpu_memory", lambda: 1) + if checkpoint == "quantized": + config.quantization_config = {"quant_method": "fp8"} + else: + monkeypatch.setattr(models_hf, "checkpoint_has_mtp", lambda *a: True) + + with pytest.raises(NotImplementedError, match=match) as refused: + load_pretrained(str(tmp_path)) + assert "MODELOPT_FSDP_SIZE=1" in str(refused.value) + assert loaders == {} + + +@pytest.mark.parametrize( + ("distribution", "options", "error", "match"), + [ + (DistributionPlan(4), {}, ValueError, "is for 4 ranks"), + ( + DistributionPlan(2), + {"hf_config": LlamaConfig(**_TINY)}, + TypeError, + r"hf_config.*resolve_distribution_plan", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"device_map": "auto"}, + TypeError, + r"\['device_map'\] do not apply", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"auto_class": AutoModel}, + NotImplementedError, + "auto_class=AutoModel", + ), + ( + DistributionPlan(2, dense="fsdp"), + {"dtype": torch.float32}, + TypeError, + r"dtype=torch\.float32", + ), + ], + ids=[ + "another-world", + "hf-config-on-a-whole-load", + "whole-load-option-on-a-sharded-plan", + "auto-class-on-a-sharded-plan", + "dtype-on-a-sharded-plan", + ], +) +def test_a_handed_plan_is_checked_before_the_checkpoint_is_fetched( + run_as, monkeypatch, tmp_path, distribution, options, error, match +): + """Making the checkpoint local can mean downloading it: a plan the load refuses fails first, as + do hf_config with a plan that loads the model whole and an option a sharded plan's load + refuses.""" + run_as(2) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution, **options) + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_a_single_process_fsdp2_load_sets_up_its_group_of_one( + monkeypatch, serve, loaders, a_gpu, tmp_path, quantized +): + """FSDP2 on one rank still shards through a process group, of one, which the load sets up once + nothing can refuse it: a refused load leaves no group behind.""" + config = serve(LlamaConfig(**_TINY)) + if quantized: + config.quantization_config = {"quant_method": "fp8"} + monkeypatch.setenv("RANK", "0") # torchrun --nproc_per_node=1 + live = [] + monkeypatch.setattr(dist, "is_initialized", lambda: bool(live)) + monkeypatch.setattr(dist, "setup", lambda timeout=None: live.append(0)) + # What a group of one answers once it is up. + monkeypatch.setattr(dist, "rank", lambda group=None: 0) + monkeypatch.setattr(dist, "size", lambda group=None: 1) + monkeypatch.setattr(dist, "local_rank", lambda: 0) + plan = DistributionPlan(dense="fsdp") + + if quantized: + with pytest.raises(NotImplementedError, match="dequantize"): + load_pretrained(str(tmp_path), distribution=plan) + assert live == [] + return + load_pretrained(str(tmp_path), distribution=plan) + + assert live == [0] + call = loaders.pop("fsdp2") + assert (call["rank"], call["world_size"]) == (0, 1) + + +def test_a_single_process_fsdp2_load_needs_a_launcher(monkeypatch, serve, loaders, tmp_path): + """Started without one, the process cannot set its group of one up; it is told to launch with + torchrun rather than fail inside init_process_group.""" + serve(LlamaConfig(**_TINY)) + monkeypatch.delenv("RANK", raising=False) + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) + + with pytest.raises(RuntimeError, match="torchrun --nproc_per_node=1"): + load_pretrained(str(tmp_path), distribution=DistributionPlan(dense="fsdp")) + assert loaders == {} + + +@pytest.fixture +def launch_without_its_group(monkeypatch): + """Rank 0 of an 8-rank launch whose processes set no process group up.""" + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "8") + monkeypatch.setattr(dist, "is_initialized", lambda: False) + monkeypatch.setattr(dist, "setup", lambda timeout=None: pytest.fail("set a group up")) + + +def test_a_launch_that_sets_no_group_up_loads_whole_on_each_process( + launch_without_its_group, serve, loaders, auto_class, tmp_path +): + """A launcher can start several processes that never set a group up (hf_ptq without + --use_fsdp2): each loads the whole model as a single process does, with no collective to wait + in and no plan to agree on.""" + serve(LlamaConfig(**_TINY)) + + model = load_pretrained(str(tmp_path), auto_class) + + assert auto_class.calls == [{}] + assert distribution_plan(model) == DistributionPlan() + assert loaders == {} + + +@pytest.mark.parametrize( + ("distribution", "knobs", "error", "match"), + [ + ( + DistributionPlan(8), + {}, + RuntimeError, + r"is for 8 ranks.*distributed_run", + ), + ( + DistributionPlan(dense="fsdp"), + {}, + RuntimeError, + r"DistributionPlan\(8, dense='fsdp'\)", + ), + ( + None, + {"MODELOPT_FSDP_SIZE": "world"}, + RuntimeError, + "MODELOPT_FSDP_SIZE=world", + ), + ( + None, + {"MODELOPT_EP_DEGREE": "4"}, + RuntimeError, + r"MODELOPT_EP_DEGREE=4 .*distributed_run\(\).*MODELOPT_EP_DEGREE=1", + ), + (None, {"MODELOPT_EP_DEGREE": "eigth"}, ValueError, "MODELOPT_EP_DEGREE='eigth'"), + ], + ids=["several-ranks", "sharded", "fsdp-knob", "ep-knob", "ep-knob-not-a-degree"], +) +def test_what_needs_the_group_is_refused_in_a_launch_that_sets_none_up( + launch_without_its_group, + serve, + loaders, + monkeypatch, + tmp_path, + distribution, + knobs, + error, + match, +): + """A plan for several ranks, a sharded load, and knobs asking to shard over the launch need its + group, which one rank cannot set up for itself: refused before anything is fetched, rather than + leave each process the whole model the knobs did not ask for. A knob that is not one is refused + as the plan decision refuses it, not dropped.""" + serve(LlamaConfig(**_TINY)) + for name, value in knobs.items(): + monkeypatch.setenv(name, value) + monkeypatch.setattr( + model_load_utils, "ensure_local_checkpoint", lambda *a, **k: pytest.fail("fetched it") + ) + + with pytest.raises(error, match=match): + load_pretrained(str(tmp_path), distribution=distribution) + assert loaders == {} diff --git a/tests/unit/torch/utils/test_unplaced_source_keys.py b/tests/unit/torch/utils/test_unplaced_source_keys.py index 31920ea6799..46a5e7ef634 100644 --- a/tests/unit/torch/utils/test_unplaced_source_keys.py +++ b/tests/unit/torch/utils/test_unplaced_source_keys.py @@ -41,6 +41,7 @@ Mxfp4Config, ) +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.plugins import model_load_utils from modelopt.torch.utils.plugins.model_load_utils import ( load_pretrained, @@ -105,6 +106,7 @@ def test_records_the_stray_tensors_and_the_checkpoint(tmp_path, extra, subfolder assert model._modelopt_unplaced_source_keys == sorted(extra), "sorted, for a stable export" assert model._modelopt_source_checkpoint == str(checkpoint) assert model.dtype == torch.bfloat16, "from_pretrained kwargs must pass through" + assert distribution_plan(model) == DistributionPlan(), "a single process holds it whole" class _Recorder: @@ -192,17 +194,65 @@ def test_an_option_it_cannot_honor_is_refused_before_anything_loads(option, refu load_pretrained("no/such/checkpoint", **option) -def test_a_config_passed_for_the_load_needs_none_on_disk(tmp_path): - """The load uses the config it is passed, as ``from_pretrained`` does, and so does the choice of - its loading hook: a directory without a config.json loads.""" +class _OwnConfig(LlamaConfig): + """A config class of the caller's own, which AutoConfig has no mapping for.""" + + model_type = "modelopt_test_own_llama" + + +class _OwnModel(LlamaForCausalLM): + config_class = _OwnConfig + + +@pytest.mark.parametrize( + ("as_path", "model_type", "model_class"), + [ + (False, None, None), + (True, None, LlamaForCausalLM), + (True, "llama", LlamaForCausalLM), + (True, None, _OwnModel), + ], + ids=["config", "config-path", "config-path-and-model-type", "own-config-class-path"], +) +def test_a_config_passed_for_the_load_needs_none_on_disk( + tmp_path, as_path, model_type, model_class +): + """The load uses the config it is passed, or reads it from the directory it is passed, as + ``from_pretrained`` does, through the model class's own config class, and so do its plan and the + choice of its loading hook: a checkpoint directory without a config.json loads, through a model + class of the caller's own too.""" config = _tiny_llama_config() + if model_class is _OwnModel: + config = _OwnConfig(**config.to_diff_dict()) + weights = tmp_path / "weights" + weights.mkdir() + _save(weights, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) + if as_path: + (tmp_path / "config").mkdir() + (weights / "config.json").rename(tmp_path / "config" / "config.json") + passed = {"auto_class": model_class, "config": str(tmp_path / "config")} + else: + (weights / "config.json").unlink() + passed = {"config": config} + (weights / "generation_config.json").write_text("{}") # from_pretrained reads it instead + + model = _load(weights, model_type=model_type, **passed) + + assert model._modelopt_unplaced_source_keys == [STRAY] + + +@pytest.mark.parametrize("model_type", [_OwnConfig.model_type, None], ids=["given", "read"]) +def test_a_model_class_of_the_callers_own_reads_its_config_itself(tmp_path, model_type): + """A single process plans nothing for a model class of the caller's own, and a model type left + to read is read as the class reads its config, through its own config class, which AutoConfig + need not know: the plan of a single process is the same whatever the model.""" + config = _OwnConfig(**_tiny_llama_config().to_diff_dict()) _save(tmp_path, config, {**LlamaForCausalLM(config).state_dict(), STRAY: torch.ones(32)}) - (tmp_path / "config.json").unlink() - (tmp_path / "generation_config.json").write_text("{}") # from_pretrained reads it instead - model = _load(tmp_path, config=config) + model = _load(tmp_path, _OwnModel, model_type=model_type) assert model._modelopt_unplaced_source_keys == [STRAY] + assert distribution_plan(model) == DistributionPlan() def test_an_ignore_rule_does_not_hide_an_unplaced_tensor(tmp_path):