[OMNIML-6060] Load an expert-parallel plan through load_pretrained

CHANGELOG.rst

@@ -43,7 +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.
+- ``modelopt.torch.utils.plugins.model_load_utils.load_pretrained`` now loads a model distributed under ``torchrun``, with an MoE's routed experts sharded across the ranks, 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**
 
@@ -52,7 +52,7 @@ 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, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE loads without expert parallelism, and is refused when ``MODELOPT_EP_DEGREE`` asks for it, until the expert-parallel load lands; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(<ranks>, dense="replicate")``, to load it whole on every rank.
+- ``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, a load that refuses placement arguments, quantized checkpoints, ones carrying an MTP head and a call inside ``init_quantized_weights``; an MoE the expert-parallel load takes has its routed experts sharded across the ranks, or the whole model FSDP2-sharded when a rank's share with the rest of it replicated would not fit one GPU, and is refused where it would otherwise load whole on every rank; a model class other than a causal LM is refused when ``MODELOPT_FSDP_SIZE`` or ``MODELOPT_EP_DEGREE`` asks to shard; and a model loaded whole goes onto each rank's own GPU unless ``device_map`` places it. Set ``MODELOPT_FSDP_SIZE=1`` and ``MODELOPT_EP_DEGREE=1``, or pass ``distribution=DistributionPlan(<ranks>, dense="replicate")``, to load it whole on every rank.
 - ``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.

docs/source/guides/_distributed_loading.rst

@@ -19,23 +19,36 @@ recorded on the model, where calibration and export read it.
     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.
+        moe = load_pretrained("Qwen/Qwen3-30B-A3B")  # its routed experts sharded over the 8 ranks
 
 Routes
 ======
 
+* A plan with ``ep_size > 1`` takes the expert-parallel load (``load_and_prepare_ep``): each rank
+  holds ``num_experts / ep_size`` of an MoE's routed experts and a replica of the dense part, and
+  tokens reach their experts' ranks by all-to-all. It replicates the dense part, so a plan that
+  FSDP2-shards it as well is refused.
 * 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.
+Left to decide, the EP degree of an MoE is the largest that divides both its routed experts and the
+world size. The decision plans without expert parallelism, and says why, an MoE the expert-parallel
+load would refuse: one its model spec rules out (GPT-OSS's, Llama-4's), one below the transformers
+release it needs, a quantized checkpoint, or a plan that shards no expert weights. It FSDP2-shards
+the whole model instead when a rank's share of the experts with a whole copy of the dense part would
+not fit one GPU. An EP degree a knob or argument asks for stands, and the load refuses what it
+cannot take. A plan that leaves an expert-parallel MoE whole on every rank, which no knob asked
+for, is refused, naming the ways out; ``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1`` asks
+for it.
 
 A refusal on any rank raises on every rank, before the load, and says how to load the model whole:
 ``MODELOPT_FSDP_SIZE=1`` for a plan the load decided, or
-``distribution=DistributionPlan(<ranks>, dense="replicate")``.
+``distribution=DistributionPlan(<ranks>, dense="replicate")``; for an expert-parallel plan,
+``MODELOPT_EP_DEGREE=1`` with ``MODELOPT_FSDP_SIZE=1``, or
+``distribution=DistributionPlan(<ranks>, 1, "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``.
@@ -56,15 +69,23 @@ Whole-model loads
 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 sharded load builds the model class and places the model itself, so leave ``auto_class`` out.
+  The FSDP2 load builds a text-only VLM through its image-text class and anything else through
+  ``AutoModelForCausalLM``; the expert-parallel load builds the checkpoint's own class, a VLM's
+  image-text class included.
+* Besides its own options (the FSDP2 load's ``cpu_offload``, ``mp_policy`` and
+  ``broadcast_chunk_size``; the expert-parallel load's ``freeze`` and ``require_lm_head``) each
+  takes ``trust_remote_code``, ``attn_implementation`` and ``hf_config``, and refuses any other
+  keyword rather than ignore it, the other sharded load's options as that load's own. An option
+  passed as ``None``, ``from_pretrained``'s own default, asks nothing and is not refused, nor is an
+  FSDP2 option at that load's own default. The expert-parallel load leaves the model unfrozen, as
+  the other loads do, unless ``freeze=True``.
+* Each keeps the checkpoint's dtype, so it refuses a ``dtype``. The FSDP2 load reads the tensors as
+  stored, so it refuses a quantized checkpoint; the expert-parallel load loads unquantized ones
+  only. Neither can run the model type's loading hook, so each refuses a checkpoint the hook would
+  build modules for, such as a Nemotron-H MTP head.
+* Each shards onto every rank's own GPU, and a rank without one is refused with the others. The
+  expert-parallel load runs with one GPU per process, each bound to ``cuda:0``, too.
 * 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``.
@@ -76,11 +97,12 @@ With neither a plan nor ``MODELOPT_FSDP_SIZE``, rank 0's GPU memory decides whet
 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.
+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. So does the expert-parallel load's, for an MoE whose experts the decision shards with neither a
+plan nor ``MODELOPT_EP_DEGREE``. 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
 =======
@@ -95,10 +117,10 @@ 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.
+read, a whole-model load and the expert-parallel 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
 ================================
@@ -113,7 +135,7 @@ 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
+model in the ``from_pretrained`` it patches, which the sharded loads never call. 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.