diff --git a/modelopt/torch/utils/fsdp2_utils.py b/modelopt/torch/utils/fsdp2_utils.py
index 419abb4ae56..dbc7c60f6e2 100644
--- a/modelopt/torch/utils/fsdp2_utils.py
+++ b/modelopt/torch/utils/fsdp2_utils.py
@@ -91,6 +91,11 @@ def _move_to_fsdp_device(model, params: set[torch.nn.Parameter]) -> None:
 def fsdp2_wrap(model, shard_root=True, mp_policy=None, cpu_offload: bool = False):
     """Auto-detect a HF causal-LM's decoder layers and FSDP2 ``fully_shard`` each one.
 
+    Its callers load causal LMs and text-only VLMs; the expert-parallel dense load, which also takes
+    VLMs, picks its blocks through
+    :func:`~modelopt.torch.utils.plugins.model_load_utils._fsdp2_wrap_modules`, keeping a vision
+    tower replicated.
+
     By default (``shard_root=True``) the root module is wrapped too, so embed/lm_head/norm are
     sharded instead of replicated per rank; pass ``shard_root=False`` to leave the root replicated
     (only decoder layers sharded). Returns the detected decoder layers so callers can reuse the
diff --git a/modelopt/torch/utils/plugins/model_load_utils.py b/modelopt/torch/utils/plugins/model_load_utils.py
index 74ab4a9486d..f67748281ef 100644
--- a/modelopt/torch/utils/plugins/model_load_utils.py
+++ b/modelopt/torch/utils/plugins/model_load_utils.py
@@ -18,19 +18,23 @@
 import contextlib
 import copy
 import dataclasses
+import inspect
 import logging
 import os
 import re
+import sys
 import warnings
 from collections import defaultdict
 from contextlib import contextmanager
 from itertools import chain
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, NamedTuple
 
 import torch
 import torch.distributed.checkpoint as dcp
 import torch.nn as nn
+import transformers
 import transformers.distributed as hf_distributed
+import transformers.modeling_utils as hf_modeling_utils
 from huggingface_hub import snapshot_download
 
 if TYPE_CHECKING:
@@ -38,8 +42,8 @@
 from torch.distributed import device_mesh
 from torch.distributed.checkpoint import HuggingFaceStorageReader
 from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
-from torch.distributed.fsdp import FSDPModule
-from torch.distributed.tensor import DTensor
+from torch.distributed.fsdp import FSDPModule, fully_shard
+from torch.distributed.tensor import DTensor, Replicate
 from transformers import (
     AutoConfig,
     AutoModel,
@@ -52,6 +56,8 @@
 from transformers.conversion_mapping import get_model_conversion_mapping
 from transformers.core_model_loading import WeightConverter, dot_natural_key, rename_source_key
 
+from modelopt.torch.export.model_utils import get_language_model_from_vl
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 from modelopt.torch.utils import distributed as dist
 from modelopt.torch.utils import env, moe
 from modelopt.torch.utils.distributed import (
@@ -62,7 +68,7 @@
     is_initialized,
     set_distribution_plan,
 )
-from modelopt.torch.utils.fsdp2_utils import fsdp2_wrap
+from modelopt.torch.utils.fsdp2_utils import _off_dtype_params, fsdp2_wrap
 from modelopt.torch.utils.logging import print_rank_0, warn_rank_0
 from modelopt.torch.utils.moe import (
     _EXPERTS_STYLES,
@@ -80,6 +86,11 @@
     read_safetensors_subset,
 )
 
+try:  # where newer transformers keeps it; older releases, under integrations (now a partial shim)
+    from transformers.distributed.tensor_parallel import replace_layer_number_by_wildcard
+except ImportError:
+    from transformers.integrations.tensor_parallel import replace_layer_number_by_wildcard
+
 __all__ = [
     "build_meta_causal_lm",
     "fsdp2_wrap",
@@ -889,6 +900,38 @@ def handed() -> DistributionPlan | None:
     return model
 
 
+def _fsdp2_wrap_modules(model: "torch.nn.Module") -> list["torch.nn.Module"]:
+    """The blocks to FSDP2-wrap one by one: the language backbone's ``_no_split_modules`` blocks.
+
+    Falls back to modelopt's decoder-layer detection. A VLM :func:`get_language_model_from_vl`
+    recognizes is searched from its language backbone, so its vision blocks, which the language-only
+    calibration never runs, are not wrapped; any other model is searched whole, so an unrecognized
+    VLM's listed vision blocks are wrapped too. A model whose language backbone shows no blocks is
+    refused.
+    """
+
+    def _boundaries(root: "torch.nn.Module") -> "list[torch.nn.Module] | None":
+        ns = set(getattr(root, "_no_split_modules", None) or [])
+        if ns:
+            mods = [m for m in root.modules() if type(m).__name__ in ns]
+            if mods:
+                return mods
+        layers = LayerActivationCollector.get_decoder_layers(root)
+        return list(layers) if layers is not None else None
+
+    # A recognized VLM -> only its language backbone's blocks (vision/audio unwrapped); else the whole model.
+    lineage = get_language_model_from_vl(model)
+    root = lineage[-1] if lineage else model
+    mods = _boundaries(root)
+    if mods:
+        return mods
+    raise RuntimeError(
+        f"Could not determine FSDP2 wrap boundaries: {type(root).__name__} exposes no "
+        "`_no_split_modules` and modelopt's decoder-layer auto-detection did not match. Provide a "
+        "standard HF causal-LM / VLM checkpoint, or replicate the dense part (MODELOPT_FSDP_SIZE=1)."
+    )
+
+
 def _map_ckpt_keys_to_model(model: "torch.nn.Module", ckpt_keys, model_keys) -> dict:
     """Map each checkpoint key to the model ``state_dict`` key it feeds.
 
@@ -1017,8 +1060,35 @@ def _meta_dcp_load(
 
     Every target param ends up filled -- from the checkpoint, by ``tie_weights`` with a param the
     checkpoint filled, or with the value it already had -- or the load fails: nothing is left as the
-    uninitialized memory ``to_empty`` hands out.
+    uninitialized memory ``to_empty`` hands out. A rank failing before ``dcp.load`` fails every rank.
     """
+    prepared: list[_MetaDcpPrepared] = []
+    # Each rank reads the checkpoint, checks it and materializes its params itself: one failing must
+    # fail every rank here, before the others enter dcp.load's collectives and wait there for it.
+    dist.all_reduce(
+        "the meta-DCP load's preparation",
+        lambda: prepared.append(_meta_dcp_prepare(model, local_dir, device, load_names)),
+        lambda _: None,
+    )
+    _meta_dcp_fill(model, prepared[0])
+
+
+class _MetaDcpPrepared(NamedTuple):
+    """What :func:`_meta_dcp_prepare` hands :func:`_meta_dcp_fill`, by name."""
+
+    reader: HuggingFaceStorageReader
+    load_sd: dict[str, torch.Tensor]
+    full_sd: dict[str, torch.Tensor]
+    want_params: set[str]
+    mapped_names: set[str]
+    emptied: set[str]
+    kept_names: set[str]
+
+
+def _meta_dcp_prepare(
+    model: "torch.nn.Module", local_dir: str, device, load_names: set
+) -> _MetaDcpPrepared:
+    """:func:`_meta_dcp_load`'s checks and materialization, up to the state dict DCP fills."""
     # Read what the checkpoint holds before changing the model, so a refusal leaves it as it was.
     reader = HuggingFaceStorageReader(str(local_dir))
     ckpt_md = reader.read_metadata().state_dict_metadata
@@ -1159,10 +1229,27 @@ def _meta_dcp_load(
 
     full_sd = model.state_dict()  # the materialized tensors: to_empty replaced the meta params
     load_sd = {inv[mk]: full_sd[mk] for mk in full_sd if mk in inv}
-    dcp.load(load_sd, storage_reader=reader)
+    return _MetaDcpPrepared(
+        reader=reader,
+        load_sd=load_sd,
+        full_sd=full_sd,
+        want_params=want_params,
+        mapped_names=mapped_names,
+        emptied=emptied,
+        kept_names=kept_names,
+    )
+
+
+def _meta_dcp_fill(model, prepared: _MetaDcpPrepared) -> None:
+    """:func:`_meta_dcp_load`'s DCP load and re-tie, from :func:`_meta_dcp_prepare`'s state."""
+    dcp.load(prepared.load_sd, storage_reader=prepared.reader)
     # Taken before tying, which can swap a filled param out of its module; in state_dict order, the
     # same on every rank, so every rank keeps the same param of a tied group.
-    filled = {n: model.get_parameter(n) for n in full_sd if n in want_params and n in mapped_names}
+    filled = {
+        n: model.get_parameter(n)
+        for n in prepared.full_sd
+        if n in prepared.want_params and n in prepared.mapped_names
+    }
     if hasattr(model, "tie_weights"):
         model.tie_weights()  # re-establish tied weights the checkpoint omitted
     # Tying keeps one param of each tied group, not necessarily the one the checkpoint filled: a tied
@@ -1179,7 +1266,9 @@ def _meta_dcp_load(
     # initialize it; this loader can't run the model's init on shards, so it fails rather than return
     # the model. Every rank reaches the same verdict.
     unfilled = sorted(
-        n for n in emptied - mapped_names - kept_names if id(model.get_parameter(n)) not in loaded
+        n
+        for n in prepared.emptied - prepared.mapped_names - prepared.kept_names
+        if id(model.get_parameter(n)) not in loaded
     )
     if unfilled:
         raise RuntimeError(
@@ -1490,6 +1579,19 @@ def _decide_distribution(
             f"{knob}={ep_size} must divide the routed experts {hf_config.model_type} shards "
             f"({num_experts})."
         )
+    # Why the expert-parallel load cannot shard the dense part, if it cannot: the export cannot write
+    # it sharded within several expert-data-parallel replicas.
+    unshardable, remedy = None, ""
+    if ep_size > 1 and world_size // ep_size > 1:
+        unshardable = (
+            f"within each of its {world_size // ep_size} expert-data-parallel replicas, a layout "
+            "the export cannot write yet"
+        )
+        remedy = (
+            " Set MODELOPT_EP_DEGREE to the world size if it divides the expert count, "
+            "MODELOPT_EP_DEGREE=1 to FSDP2-shard the whole model without expert parallelism, or "
+            "MODELOPT_FSDP_SIZE=1 if the dense part fits one GPU."
+        )
     if dense == "auto":
         # A rank holds a whole copy of what expert parallelism does not shard, and its 1/ep_size
         # share of the routed experts it does.
@@ -1502,7 +1604,17 @@ def _decide_distribution(
         )
         memory = _gpu_memory()
         too_big = memory is not None and held > _REPLICATED_SHARE * memory
-        dense = "fsdp" if too_big else "replicate"
+        if too_big and unshardable:
+            warn_rank_0(
+                f"{hf_config.model_type}: a rank holding the dense part whole would take more than "
+                f"{_REPLICATED_SHARE:.0%} of a GPU, but it cannot be sharded {unshardable}; "
+                f"replicating it, which may run out of memory.{remedy}"
+            )
+        dense = "fsdp" if too_big and not unshardable else "replicate"
+    elif dense == "fsdp" and unshardable:
+        raise NotImplementedError(
+            f"The dense part of {hf_config.model_type} cannot be sharded {unshardable}.{remedy}"
+        )
     return DistributionPlan(world_size, ep_size, dense)
 
 
@@ -1602,6 +1714,138 @@ def checkpoint_read() -> bool:
     return plan, hf_config
 
 
+def _make_expert_param_matcher(ep_plan: dict[str, Any]):
+    """``is_expert(param_name)``: whether ``ep_plan`` (the plan the EP load shards by) shards it.
+
+    True when the name, layer indices wildcarded and ``.weight``/``.bias`` stripped, contains a
+    non-router plan key as whole dot-separated segments: the key itself, a suffix of the name (plan keys
+    omit the ``model.`` prefix), or a container it nests under. Router styles are replicated.
+    """
+    _router_styles = {"ep_router", "torch_a2a_router"}
+    keys = [k for k, v in (ep_plan or {}).items() if v not in _router_styles]
+
+    def is_expert(pname: str) -> bool:
+        generic = replace_layer_number_by_wildcard(pname)
+        for suffix in (".weight", ".bias"):
+            if generic.endswith(suffix):
+                generic = generic[: -len(suffix)]
+                break
+        padded = "." + generic + "."
+        return any(("." + k + ".") in padded for k in keys)
+
+    return is_expert
+
+
+@contextlib.contextmanager
+def _ep_aware_caching_allocator_warmup(
+    ep_degree: int, ep_plan: dict[str, Any], dense_on_meta: bool = False
+):
+    """Size transformers' load-time allocator warmup to this rank's expert-parallel share.
+
+    transformers credits a param its ``1/world_size`` share only through ``_tp_plan``, which holds no
+    expert param under the EP load, so the warmup would reserve about the whole GPU on every rank and
+    starve the all-to-all's NCCL buffers. While the load runs, an expert param counts ``1/ep_degree`` of
+    itself and, with ``dense_on_meta``, a dense ``.weight``/``.bias`` nothing. It only sizes a
+    reservation: a release without the hook, a hook of another signature, or an error while sizing,
+    keeps the stock sizing, and says so.
+    """
+    if ep_degree <= 1:
+        yield
+        return
+    orig = getattr(hf_modeling_utils, "get_total_byte_count", None)
+    if orig is None:
+        warn_rank_0(
+            f"transformers {transformers.__version__} has no get_total_byte_count in "
+            "modeling_utils to size the load warmup by; using its own sizing, which may reserve "
+            "most of the device."
+        )
+        yield
+        return
+    try:
+        params = list(inspect.signature(orig).parameters)
+    except (TypeError, ValueError):
+        params = []
+    # A release whose hook takes other leading parameters calls it in a way this replacement cannot
+    # know, so it keeps the stock sizing rather than fail the load at transformers' call site. Said,
+    # since that sizing may reserve most of the device and a later NCCL failure would have no cause.
+    if params[:3] != ["model", "accelerator_device_map", "hf_quantizer"]:
+        warn_rank_0(
+            f"transformers {transformers.__version__}'s warmup sizing hook takes {params[:3]}, not "
+            "(model, accelerator_device_map, hf_quantizer); using its full-model sizing, which may "
+            "reserve most of the device."
+        )
+        yield
+        return
+
+    is_expert = _make_expert_param_matcher(ep_plan)
+
+    def _ep_aware(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs):
+        try:
+            try:  # where current transformers keeps it (5.18); older releases, under integrations
+                from transformers.distributed.tensor_parallel import _get_parameter_tp_plan
+            except ImportError:
+                try:
+                    from transformers.integrations.tensor_parallel import _get_parameter_tp_plan
+                except ImportError:
+                    _get_parameter_tp_plan = None
+            ws = (
+                torch.distributed.get_world_size()
+                if torch.distributed.is_available() and torch.distributed.is_initialized()
+                else 1
+            )
+            tp_plan = (getattr(model, "_tp_plan", None) or {}) if ws > 1 else {}
+            tied = set(getattr(model, "all_tied_weights_keys", {}) or {})
+            total: dict = defaultdict(int)
+            for pname, dev in accelerator_device_map.items():
+                if pname in tied:
+                    continue
+                try:
+                    p = model.get_parameter_or_buffer(pname)
+                except Exception:
+                    continue
+                dsize = (
+                    hf_quantizer.param_element_size(model, pname, p)
+                    if hf_quantizer is not None
+                    else p.element_size()
+                )
+                nb = p.numel() * dsize
+                if is_expert(pname):
+                    nb //= ep_degree
+                elif (
+                    tp_plan
+                    and _get_parameter_tp_plan is not None
+                    and _get_parameter_tp_plan(pname, tp_plan, is_weight=True) is not None
+                ):
+                    # Placed at its share over the ep submesh by transformers' TP branch, which ignores
+                    # get_device, so it is on the device even when the rest of the dense part is on meta.
+                    nb //= ep_degree
+                elif dense_on_meta and pname.endswith((".weight", ".bias")):
+                    # On meta through the load (_ep_dense_meta_load); buffers are not, so they count.
+                    continue
+                total[dev] += nb
+            gib = 1024**3
+            sizes = ", ".join(f"{d}: {b / gib:.1f} GiB" for d, b in total.items())
+            print_rank_0(
+                f"[MEM] EP-aware warmup sizing (experts // ep_degree={ep_degree}"
+                f"{', dense on meta -> uncounted' if dense_on_meta else ''}) -> {sizes} "
+                f"(was full-model ~device-max; EP-aware plan-driven sizing)"
+            )
+            return total
+        except Exception as e:
+            # Its fallback reserves most of the device, so say why, or a later NCCL failure has no cause.
+            warn_rank_0(
+                f"EP-aware warmup sizing failed ({type(e).__name__}: {e}); using transformers' "
+                "full-model sizing, which may reserve most of the device."
+            )
+            return orig(model, accelerator_device_map, hf_quantizer, *args, **kwargs)
+
+    try:
+        hf_modeling_utils.get_total_byte_count = _ep_aware
+        yield
+    finally:
+        hf_modeling_utils.get_total_byte_count = orig
+
+
 def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_state) -> None:
     """Seed every module's ``_parallel_state`` for calibration's amax sync.
 
@@ -1617,6 +1861,76 @@ def _seed_parallel_states(model: nn.Module, expert_names, dense_state, expert_st
         )
 
 
+def _fsdp2_shard_dense(model: nn.Module, fsdp_mesh, expert_modules) -> int:
+    """FSDP2-shard the dense part over ``fsdp_mesh``, block by block; return the blocks wrapped.
+
+    Each :func:`_fsdp2_wrap_modules` block is wrapped whole: ``fully_shard`` swaps a wrapped module's
+    class, and quantization and export recognize attention and linear modules by class, so those stay
+    plain inside the block. The routed experts (EP-sharded) and params off the model's dtype are left
+    out; the root (embeddings, LM head, final norm) is not wrapped and stays replicated.
+    """
+    # The experts are EP-sharded already; a block ignores the ones it does not hold.
+    ignored_params = _off_dtype_params(model)
+    for em in expert_modules:
+        ignored_params.update(em.parameters(recurse=True))
+    ignored_param_ids = {id(p) for p in ignored_params}  # id-based: Tensor.__eq__ is elementwise
+
+    # Only the outermost blocks: a model listing a container and its blocks (Kimi's MoonViT) would
+    # wrap a child of a wrapped module, which fully_shard refuses (overlapping meshes).
+    wrap_modules = _fsdp2_wrap_modules(model)
+    wrap_ids = {id(m) for m in wrap_modules}
+    nested_ids: set = set()
+    for m in wrap_modules:
+        for c in m.modules():
+            if id(c) != id(m) and id(c) in wrap_ids:
+                nested_ids.add(id(c))
+    wrap_modules = [m for m in wrap_modules if id(m) not in nested_ids]
+
+    n = 0
+    for layer in wrap_modules:
+        layer_ignored = set(ignored_params)  # experts + any pre-sharded dense (below)
+        has_shardable = False
+        for name, p in list(layer.named_parameters(recurse=True)):
+            if id(p) in ignored_param_ids:
+                continue
+            if isinstance(p.data, DTensor):
+                if all(isinstance(pl, Replicate) for pl in p.data.placements):
+                    # Re-registered as a plain Parameter (a free view): assigning .data would keep the
+                    # DTensor parameter, which fully_shard would place on two meshes.
+                    owner_name, _, leaf = name.rpartition(".")
+                    owner = layer.get_submodule(owner_name) if owner_name else layer
+                    setattr(
+                        owner,
+                        leaf,
+                        nn.Parameter(p.data.to_local(), requires_grad=p.requires_grad),
+                    )
+                else:
+                    layer_ignored.add(
+                        p
+                    )  # already sharded by the loader (TP) -> leave it distributed
+                    continue
+            has_shardable = True
+        if not has_shardable:
+            continue  # nothing left to FSDP-shard in this layer (all expert / already-sharded)
+        fully_shard(
+            layer,
+            mesh=fsdp_mesh,
+            reshard_after_forward=True,
+            ignored_params=layer_ignored,
+        )
+        n += 1
+    if not n:
+        # The dense params were forced to meta for this sharding; with nothing sharded, the load would
+        # materialize the whole dense part on every rank, which the plan shards it to avoid.
+        raise RuntimeError(
+            f"Sharding the dense part found {len(wrap_modules)} block(s) but no dense parameter to "
+            "shard in any of them (each is a routed expert's, already sharded, or off the model's "
+            "dominant dtype), so every rank would hold the whole dense part. Replicate the dense "
+            "part instead (MODELOPT_FSDP_SIZE=1)."
+        )
+    return n
+
+
 # The checkpoint formats load_and_prepare_ep dequantizes as it loads, beside unquantized ones (None).
 # Not MXFP4: transformers dequantizes it into GPT-OSS's transposed expert layout, whatever the model's.
 _EP_SOURCE_FORMATS: frozenset[str | None] = frozenset({None})
@@ -1673,6 +1987,81 @@ def _ep_mesh(*args, **kwargs):
         mixin.initialize_distributed_mesh = original
 
 
+@contextlib.contextmanager
+def _ep_dense_meta_load(is_expert):
+    """Keep the dense ``.weight``/``.bias`` on ``meta`` through ``from_pretrained``.
+
+    transformers places expert params through its TP branch, which never asks ``get_device``, and every
+    other param where ``get_device`` says: ``"meta"`` for a dense key leaves it for
+    :func:`_load_ep_dense_distributed` to shard and load. Buffers load as usual unless named
+    ``.weight``/``.bias``, which no supported MoE has: such a buffer would go to ``meta`` too, where
+    the meta+DCP load refuses it. ``get_device`` is patched in every module the loader looks it up in.
+    """
+    loader = getattr(hf_modeling_utils, "convert_and_load_state_dict_in_model", None)
+    mods = {hf_modeling_utils}
+    loader_module = getattr(loader, "__module__", None)
+    dm = sys.modules.get(loader_module) if loader_module else None
+    if dm is not None:
+        mods.add(dm)
+    targets = [m for m in mods if hasattr(m, "get_device")]
+    if not targets:
+        raise RuntimeError(
+            "Sharding the dense part keeps it off the device during from_pretrained through "
+            f"transformers' get_device, which none of {sorted(m.__name__ for m in mods)} has on "
+            f"transformers {transformers.__version__}. Set MODELOPT_FSDP_SIZE=1 if the dense part "
+            "fits one GPU."
+        )
+
+    def _dense_weight(key) -> bool:
+        k = str(key)
+        return k.endswith((".weight", ".bias")) and not is_expert(k)
+
+    saved = [(m, m.get_device) for m in targets]
+
+    def _make(orig):
+        def patched(device_map, key, *args, **kwargs):
+            if _dense_weight(key):
+                return "meta"
+            return orig(device_map, key, *args, **kwargs)
+
+        return patched
+
+    try:
+        for m, orig in saved:
+            m.get_device = _make(orig)
+        yield
+    finally:
+        for m, orig in saved:
+            m.get_device = orig
+
+
+def _load_ep_dense_distributed(model, fsdp_mesh, local_dir, is_expert, expert_modules):
+    """FSDP2-shard the meta dense part over ``fsdp_mesh`` and DCP-load each rank's shard of it.
+
+    The dense ``.weight``/``.bias`` params arrive on meta (:func:`_ep_dense_meta_load`); the experts,
+    already loaded and EP-sharded, are left as they are. Returns the number of layers sharded.
+    """
+    # The GPU the rank is bound to, as the experts and FSDP2's compute are: not the one its local
+    # rank numbers, which a launch showing each process one GPU does not bind.
+    cuda = fsdp_mesh.device_type == "cuda"
+    device = torch.device(fsdp_mesh.device_type, _own_gpu_index() if cuda else None)
+    # A rank failing to wrap (out of memory, say) fails every rank here, before DCP's collectives.
+    with _refused_together("the dense part's FSDP2 wrapping", alone=False):
+        n_dense = _fsdp2_shard_dense(model, fsdp_mesh, expert_modules)
+        # Every alias of a tied param: the checkpoint may store the one named_parameters() drops.
+        dense_names = {
+            n
+            for n, _ in model.named_parameters(remove_duplicate=False)
+            if n.endswith((".weight", ".bias")) and not is_expert(n)
+        }
+    print_rank_0(
+        f"[ep] dense-distributed load: FSDP2-sharded {n_dense} dense module(s); "
+        f"meta+DCP-loading {len(dense_names)} dense param(s) (experts left EP-sharded)"
+    )
+    _meta_dcp_load(model, local_dir, device, load_names=dense_names)
+    return n_dense
+
+
 def load_and_prepare_ep(
     ckpt_path: str,
     plan: DistributionPlan,
@@ -1696,6 +2085,9 @@ def load_and_prepare_ep(
     the load that rests on ``from_pretrained`` running no collective, as transformers'
     expert-parallel load does not on 5.7 through 5.18.
 
+    With ``plan.dense == "fsdp"`` the dense part is FSDP2-sharded over the ``ep`` group instead,
+    each rank loading only its shard: for a model whose dense part does not fit one GPU.
+
     The model builds with its LM head where an auto class builds one; ``require_lm_head=False``
     builds a VLM without it. Without ``hf_config`` every rank reads the checkpoint's config itself;
     on many ranks, pass it. The caller records the plan on the model
@@ -1704,8 +2096,14 @@ def load_and_prepare_ep(
     """
     if not plan.is_expert_parallel:
         raise ValueError(f"{plan!r} shards no routed experts, which is what this load is for.")
-    if plan.is_fsdp_dense:
-        raise NotImplementedError(f"{plan!r}: this load replicates the dense part on every rank.")
+    if plan.is_fsdp_dense and plan.edp > 1:
+        # Refused before the load: the export, after the whole calibration, is what cannot write it.
+        raise NotImplementedError(
+            f"{plan!r} FSDP2-shards the dense part within each of {plan.edp} expert-data-parallel "
+            "replicas, a layout the export cannot write yet. Set MODELOPT_EP_DEGREE to the world "
+            "size if it divides the expert count, or MODELOPT_FSDP_SIZE=1 if the dense part fits "
+            "one GPU."
+        )
     # Bringing the group up is the launcher's job; this load needs it live and device-bound.
     if not dist.is_initialized():
         raise RuntimeError(
@@ -1889,13 +2287,24 @@ def gpu() -> tuple[bool, bool, str | None]:
     ep_cfg.base_model_ep_plan = swapped_plan
     from_pretrained_kwargs["config"] = hf_config
     print_rank_0(f"[ep] ep_plan={swapped_plan} experts_implementation=torch_a2a")
-    # A rank failing in the load fails every rank here, before the warmup's collectives:
+    is_expert_fn = _make_expert_param_matcher(swapped_plan)
+    # A sharded dense part stays on meta through from_pretrained, to be sharded and loaded below.
+    _dense_ctx = (
+        _ep_dense_meta_load(is_expert_fn) if plan.is_fsdp_dense else contextlib.nullcontext()
+    )
+    # A rank failing in the load fails every rank here, before the collectives below:
     # from_pretrained runs no collective, so the others finish it and meet the failure in the gather.
     with _refused_together("the expert-parallel load", alone=False):
         # The class the plan's decision chose, with its LM head where an auto class builds one.
         loader = _auto_class(hf_config, require_lm_head, trust_remote_code=trust_remote_code)
         print_rank_0(f"[ep] loader={loader.__name__} (require_lm_head={require_lm_head})")
-        with _experts_sharded_over(ep_submesh):
+        with (
+            _ep_aware_caching_allocator_warmup(
+                ep_degree, swapped_plan, dense_on_meta=plan.is_fsdp_dense
+            ),
+            _dense_ctx,
+            _experts_sharded_over(ep_submesh),
+        ):
             model = loader.from_pretrained(ckpt_path, **from_pretrained_kwargs)
         model.eval()
         # The config the model keeps goes back to the checkpoint's plan, as an export writes it out
@@ -1907,6 +2316,22 @@ def gpu() -> tuple[bool, bool, str | None]:
                 vars(kept_ep_config).pop("base_model_ep_plan", None)
             else:
                 kept_ep_config.base_model_ep_plan = declared_plan
+        if plan.is_fsdp_dense:
+            meta = [n for n, p in model.named_parameters() if p.is_meta]
+            dense_meta = [n for n in meta if not is_expert_fn(n)]
+            if not dense_meta:
+                # A transformers that places the dense part without asking get_device loads it whole
+                # on every rank, which is what this layout exists to avoid, and would do so silently.
+                raise RuntimeError(
+                    "Sharding the dense part needs from_pretrained to leave it on meta, but no dense "
+                    f"param came back on meta on transformers {transformers.__version__}: every rank "
+                    "loaded the whole dense part. Set MODELOPT_FSDP_SIZE=1 if the dense part fits "
+                    "one GPU."
+                )
+            print_rank_0(
+                f"[ep] dense-meta load: {len(dense_meta)} dense param(s) on meta, "
+                f"{len(meta) - len(dense_meta)} expert; e.g. {dense_meta[:3]}"
+            )
         # The experts forward reads the implementation from the config when it runs, which
         # from_pretrained left on grouped_mm: torch_a2a on every (sub)config. An export writes none.
         for _m in model.modules():
@@ -1935,6 +2360,19 @@ def gpu() -> tuple[bool, bool, str | None]:
                 "ranks holding different experts. Quantize the model without expert parallelism."
             )
 
+    # The dense part shards over the whole world, the experts' ranks: a forward's dense all-gathers
+    # and expert all-to-alls then run in turn on the same ranks, so they cannot deadlock.
+    if plan.is_fsdp_dense:
+        dense_fsdp_mesh = dense_mesh["dp"]
+        n_dense = _load_ep_dense_distributed(
+            model, dense_fsdp_mesh, ckpt_path, is_expert_fn, expert_modules.values()
+        )
+        print_rank_0(
+            f"[ep] dense sharded: FSDP2-sharded {n_dense} dense module(s) over the whole world, "
+            f"mesh {dense_fsdp_mesh.mesh_dim_names}={tuple(dense_fsdp_mesh.shape)}; experts "
+            "EP-sharded on ep"
+        )
+
     # Each module's ParallelState for calibration's amax sync: dense modules reduce over the whole world;
     # routed experts reduce over their expert-data-parallel replicas, and share one input amax across the
     # layer's experts over the expert-parallel group. A size-1 axis is left unused (-1): no collective.
diff --git a/tests/gpu/torch/utils/test_expert_parallel_load.py b/tests/gpu/torch/utils/test_expert_parallel_load.py
index 849305f4f6a..0e14032e3cf 100644
--- a/tests/gpu/torch/utils/test_expert_parallel_load.py
+++ b/tests/gpu/torch/utils/test_expert_parallel_load.py
@@ -31,8 +31,11 @@
 from torch.distributed.tensor import DTensor
 from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig
 
+import modelopt.torch.quantization as mtq
+from modelopt.torch.quantization.nn import TensorQuantizer
 from modelopt.torch.utils.distributed import (
     DistributionPlan,
+    distributed_forward_loop,
     distribution_plan,
     set_distribution_plan,
 )
@@ -63,10 +66,20 @@ def _local(tensor):
     return tensor.to_local() if isinstance(tensor, DTensor) else tensor
 
 
-def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree):
+# (dense part sharded, share of the GPUs the experts span). A dense part FSDP2-sharded within several
+# expert-parallel groups (sharded with half the GPUs per group) is refused by the loader: the export
+# cannot write it.
+_LAYOUTS = [
+    pytest.param(False, 1, id="dense-replicated-ep-world"),
+    pytest.param(False, 2, id="dense-replicated-ep-half"),
+    pytest.param(True, 1, id="dense-sharded-ep-world"),
+]
+
+
+def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, fsdp_dense, ep_degree):
     device = torch.device(f"cuda:{rank}")
     config = AutoConfig.from_pretrained(ckpt_dir)  # carries qwen3_moe's own base_model_ep_plan
-    plan = DistributionPlan(size, ep_size=ep_degree)
+    plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate")
     model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config).eval()
     # The loader swaps the plan on its own copy: the caller's config keeps the checkpoint's.
     assert config.base_model_ep_plan == AutoConfig.from_pretrained(ckpt_dir).base_model_ep_plan
@@ -101,13 +114,200 @@ def _sharded_model_matches_the_whole_one(rank, size, ckpt_dir, ep_degree):
     torch.testing.assert_close(sharded.float(), reference.float(), atol=2e-2, rtol=2e-2)
 
 
-@pytest.mark.parametrize("ep_share", [1, 2], ids=["ep-world", "ep-half"])
+@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS)
 def test_expert_parallel_model_matches_the_whole_model(
-    need_2_gpus, dist_workers, tmp_path, ep_share
+    need_2_gpus, dist_workers, tmp_path, fsdp_dense, ep_share
+):
+    """With the dense part replicated, or FSDP2-sharded and loaded shard by shard through DCP; and
+    with the experts sharded over every rank, or over half of them and replicated across the halves."""
+    ep_degree = dist_workers.world_size // ep_share
+    if ep_degree < 2 or dist_workers.world_size % ep_share or NUM_EXPERTS % ep_degree:
+        pytest.skip(f"{NUM_EXPERTS} experts do not shard over EP {ep_degree} on these GPUs")
+    _save_tiny_qwen3_moe(tmp_path)
+    dist_workers.run(_sharded_model_matches_the_whole_one, str(tmp_path), fsdp_dense, ep_degree)
+
+
+def _batches(size):
+    """Every rank's batch, the same on all ranks: each trains on its own, the reference on all."""
+    return [
+        torch.randint(0, 128, (2, 16), generator=torch.Generator().manual_seed(r))
+        for r in range(size)
+    ]
+
+
+def _whole_model_grads_per_batch(whole, batches, device):
+    """The whole model's gradient on each batch alone, as ``{parameter name: gradient}``."""
+    grads = []
+    for batch in batches:
+        whole.zero_grad(set_to_none=True)
+        batch = batch.to(device)
+        whole(batch, labels=batch).loss.backward()
+        grads.append({n: p.grad.detach().clone() for n, p in whole.named_parameters()})
+    return grads
+
+
+def _fail_together(problems):
+    """Raise on every rank if any rank found a problem, so that no rank walks on alone into a collective
+    the failed ones never reach (which would hang until the test timeout instead of failing)."""
+    gathered = [None] * torch.distributed.get_world_size()
+    torch.distributed.all_gather_object(gathered, problems)
+    found = [
+        f"rank {r}: {problem}"
+        for r, rank_problems in enumerate(gathered)
+        for problem in rank_problems
+    ]
+    assert not found, "\n".join(found[:20])
+
+
+def _relative_error(actual, expected):
+    # bf16 gradients summed in a different order on each side: compare by relative norm.
+    return (
+        (actual.float() - expected.float()).norm() / expected.float().norm().clamp_min(1e-12)
+    ).item()
+
+
+def _assert_grads_match_the_whole_model(model, per_batch, rank, size, ep_degree):
+    """Routed experts collect the gradient of every token routed to them, their EP group's batches;
+    FSDP2-sharded dense parameters hold the gradient averaged over every rank; replicated dense
+    parameters see only this rank's batch."""
+    ep_rank = rank % ep_degree
+    group = range(rank - ep_rank, rank - ep_rank + ep_degree)
+    per_rank = NUM_EXPERTS // ep_degree
+    owned = slice(ep_rank * per_rank, (ep_rank + 1) * per_rank)
+    # Checked before the comparisons: full_tensor() below is a collective that a rank without the
+    # gradient would skip, leaving the others waiting in it instead of failing.
+    _fail_together(
+        [f"{name} received no gradient" for name, p in model.named_parameters() if p.grad is None]
+    )
+    problems, names = [], set()
+    for name, param in model.named_parameters():
+        names.add(name)
+        if ".mlp.experts." in name:
+            actual = _local(param.grad)
+            expected = sum(per_batch[r][name] for r in group)[owned]
+        elif isinstance(param.grad, DTensor):
+            actual = param.grad.full_tensor()
+            expected = sum(per_batch[r][name] for r in range(size)) / size
+        else:
+            actual = param.grad
+            expected = per_batch[rank][name]
+        error = _relative_error(actual, expected)
+        if not error < 3e-2:
+            problems.append(f"{name}: gradient off by {error:.3g} (relative norm)")
+    if names != set(per_batch[0]):
+        problems.append(
+            f"parameters differ from the whole model's: {sorted(names ^ set(per_batch[0]))[:5]}"
+        )
+    _fail_together(problems)
+
+
+def _sharded_model_backward_matches_the_whole_one(rank, size, ckpt_dir, fsdp_dense, ep_degree):
+    device = torch.device(f"cuda:{rank}")
+    config = AutoConfig.from_pretrained(ckpt_dir)
+    plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate")
+    model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config, freeze=False)
+    set_distribution_plan(
+        model, plan
+    )  # loaded directly, so recorded here, as load_pretrained would
+    whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device)
+
+    batches = _batches(size)
+    mine = batches[rank].to(device)
+    model(mine, labels=mine).loss.backward()
+    _assert_grads_match_the_whole_model(
+        model, _whole_model_grads_per_batch(whole, batches, device), rank, size, ep_degree
+    )
+
+
+@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS)
+def test_expert_parallel_backward_matches_the_whole_model(
+    need_2_gpus, dist_workers, tmp_path, fsdp_dense, ep_share
 ):
-    """With the experts sharded over every rank, or over half of them and replicated across the halves."""
+    """Loaded trainable, the sharded model's gradients are the whole model's: tokens' gradients travel
+    back through the all-to-all to the experts that saw them, and to the router that chose them."""
     ep_degree = dist_workers.world_size // ep_share
     if ep_degree < 2 or dist_workers.world_size % ep_share or NUM_EXPERTS % ep_degree:
         pytest.skip(f"{NUM_EXPERTS} experts do not shard over EP {ep_degree} on these GPUs")
     _save_tiny_qwen3_moe(tmp_path)
-    dist_workers.run(_sharded_model_matches_the_whole_one, str(tmp_path), ep_degree)
+    dist_workers.run(
+        _sharded_model_backward_matches_the_whole_one, str(tmp_path), fsdp_dense, ep_degree
+    )
+
+
+def _calibrated_quantizers(model):
+    return {
+        name: module
+        for name, module in model.named_modules()
+        if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None
+    }
+
+
+def _quantized_sharded_model_trains_like_the_whole_one(
+    rank, size, ckpt_dir, fsdp_dense, fmt, ep_degree
+):
+    cfg = {"fp8": mtq.FP8_DEFAULT_CFG, "nvfp4": mtq.NVFP4_DEFAULT_CFG}[fmt]
+    device = torch.device(f"cuda:{rank}")
+    config = AutoConfig.from_pretrained(ckpt_dir)
+    plan = DistributionPlan(size, ep_size=ep_degree, dense="fsdp" if fsdp_dense else "replicate")
+    model = load_and_prepare_ep(ckpt_dir, plan, hf_config=config, freeze=False)
+    set_distribution_plan(
+        model, plan
+    )  # loaded directly, so recorded here, as load_pretrained would
+    whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device)
+
+    # Calibrate as hf_ptq does: each rank on its own batch, the reference on all of them.
+    batches = _batches(size)
+    mine = batches[rank].to(device)
+    forward_loop = distributed_forward_loop(
+        model, [{"input_ids": mine}], device=device, ep_mesh=expert_parallel_mesh(model)
+    )
+    mtq.quantize(model, cfg, forward_loop)
+    mtq.quantize(whole, cfg, lambda m: [m(b.to(device)) for b in batches])
+
+    # Expert-parallel calibration reaches the whole model's scales: weight scales from this rank's
+    # experts, activation scales synced over the ranks (equal to bf16 rounding of the activations).
+    per_rank = NUM_EXPERTS // ep_degree
+    ep_rank = rank % ep_degree
+    whole_quantizers = _calibrated_quantizers(whole)
+    sharded_quantizers = _calibrated_quantizers(model)
+    problems = [] if sharded_quantizers else ["no quantizer was calibrated"]
+    for name, quantizer in sharded_quantizers.items():
+        if "_weight_quantizers." in name:
+            prefix, index = name.rsplit(".", 1)
+            name = f"{prefix}.{ep_rank * per_rank + int(index)}"
+        reference = whole_quantizers.get(name)
+        if reference is None:
+            problems.append(f"{name} is calibrated here but not in the whole model")
+            continue
+        error = _relative_error(quantizer.amax, reference.amax)
+        if not error < 2e-2:
+            problems.append(f"{name}: amax off by {error:.3g} (relative)")
+        # Train on identical scales, so that the gradients compare the transport and not calibration.
+        quantizer.amax = reference.amax
+    _fail_together(problems)
+
+    model(mine, labels=mine).loss.backward()
+    _assert_grads_match_the_whole_model(
+        model, _whole_model_grads_per_batch(whole, batches, device), rank, size, ep_degree
+    )
+
+
+@pytest.mark.parametrize("fmt", ["fp8", "nvfp4"])
+@pytest.mark.parametrize(("fsdp_dense", "ep_share"), _LAYOUTS)
+def test_quantized_expert_parallel_model_trains_like_the_whole_model(
+    need_2_gpus, dist_workers, tmp_path, fsdp_dense, fmt, ep_share
+):
+    """QAT under expert parallelism: quantized and calibrated expert-parallel -- experts over every
+    rank or half of them, dense replicated or FSDP2-sharded -- the model's scales and its gradients
+    through the fake quantizers match the whole model calibrated on every rank's data."""
+    ep_degree = dist_workers.world_size // ep_share
+    if ep_degree < 2 or dist_workers.world_size % ep_share or NUM_EXPERTS % ep_degree:
+        pytest.skip(f"{NUM_EXPERTS} experts do not shard over EP {ep_degree} on these GPUs")
+    _save_tiny_qwen3_moe(tmp_path)
+    dist_workers.run(
+        _quantized_sharded_model_trains_like_the_whole_one,
+        str(tmp_path),
+        fsdp_dense,
+        fmt,
+        ep_degree,
+    )
diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py
index 4c642b1644c..723c8c576d3 100644
--- a/tests/unit/torch/utils/test_meta_dcp_load.py
+++ b/tests/unit/torch/utils/test_meta_dcp_load.py
@@ -17,16 +17,21 @@
 
 ``_meta_dcp_load`` materializes a model built on ``meta`` and fills it from HF safetensors through
 ``torch.distributed.checkpoint``; ``_map_ckpt_keys_to_model`` bridges checkpoint names to the built
-model's ``state_dict`` names. The loads run for real on CPU, from tiny transformers models into
-hand-written safetensors files. DCP loads in a single process without a process group; the tests that
-need the distributed path start a one-rank gloo group.
+model's ``state_dict`` names; ``_fsdp2_wrap_modules`` picks the blocks FSDP2 shards. The loads run
+for real on CPU, from tiny transformers models into hand-written safetensors files. DCP loads in a
+single process without a process group; the tests that need the distributed path start a one-rank
+gloo group, or run on the pooled two-rank one.
 """
 
+from contextlib import nullcontext
+from unittest import mock
+
 import pytest
 import torch
 import torch.distributed.checkpoint as dcp
 from safetensors.torch import save_file
 from torch import nn
+from torch.distributed.checkpoint import HuggingFaceStorageReader
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import fully_shard
 from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
@@ -34,12 +39,14 @@
 pytest.importorskip("accelerate")
 pytest.importorskip("transformers")
 
+from _test_utils.torch.distributed.utils import bounded_collectives
 from _test_utils.torch.transformers_models import get_tiny_llama
 from accelerate import init_empty_weights
 from transformers import AutoModelForImageTextToText, Qwen2_5_VLConfig
 
 from modelopt.torch.utils.plugins.model_load_utils import (
     _copy_tied_values,
+    _fsdp2_wrap_modules,
     _map_ckpt_keys_to_model,
     _meta_dcp_load,
 )
@@ -170,6 +177,60 @@ def test_map_ckpt_keys_to_model(conversion_mapping, ckpt_keys, model_keys, expec
     assert _map_ckpt_keys_to_model(model, ckpt_keys, model_keys) == expected
 
 
+# -------------------------------------------------------------------------- FSDP2 wrap boundaries
+
+
+def test_wrap_modules_are_the_decoder_layers():
+    """A causal LM is wrapped at its ``_no_split_modules`` blocks: the decoder layers, in order."""
+    model = _meta_twin(get_tiny_llama)
+
+    assert _fsdp2_wrap_modules(model) == list(model.model.layers)
+
+
+def test_vlm_wraps_only_its_language_layers():
+    """Qwen2.5-VL's wrapper also lists the vision block in ``_no_split_modules``, but only the language
+    layers may be sharded: sharded vision weights deadlock calibration, so the tower stays replicated.
+    """
+    model = _meta_twin(_tiny_qwen2_5_vl)
+    assert type(model.model.visual.blocks[0]).__name__ in model._no_split_modules
+
+    assert _fsdp2_wrap_modules(model) == list(model.model.language_model.layers)
+
+
+@pytest.mark.parametrize("no_split_modules", [None, ["NoSuchBlock"]])
+def test_wrap_modules_fall_back_to_decoder_layer_detection(no_split_modules):
+    """Without usable ``_no_split_modules`` (absent, or naming no module present), modelopt's
+    decoder-layer auto-detection still finds the blocks.
+    """
+    model = _meta_twin(get_tiny_llama)
+    model._no_split_modules = no_split_modules
+
+    assert _fsdp2_wrap_modules(model) == list(model.model.layers)
+
+
+def _vlm_whose_language_model_shows_no_blocks():
+    """A VLM whose own ``_no_split_modules`` names its vision block too, as Qwen2.5-VL's does."""
+    model = nn.Module()
+    model._no_split_modules = ["Linear"]
+    model.language_model = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
+    model.visual = nn.Linear(2, 2)
+    return model
+
+
+@pytest.mark.parametrize(
+    "build",
+    [
+        pytest.param(lambda: nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)), id="no-blocks"),
+        pytest.param(_vlm_whose_language_model_shows_no_blocks, id="vlm-language-model"),
+    ],
+)
+def test_wrap_modules_refuse_a_model_without_blocks(build):
+    """With neither source of boundaries matching, fail loudly instead of guessing a layout. For a
+    VLM that is its language model's: the VLM's own list would shard the vision blocks too."""
+    with pytest.raises(RuntimeError, match="Could not determine FSDP2 wrap boundaries"):
+        _fsdp2_wrap_modules(build())
+
+
 # --------------------------------------------------------------------------------- meta + DCP load
 
 
@@ -582,3 +643,45 @@ def test_a_tie_laid_out_differently_is_refused_and_left_as_it_was(gloo_world_of_
         _copy_tied_values(kept, src, "lm_head.weight")
 
     assert torch.equal(local, before)
+
+
+# ---------------------------------------------------------------- one rank failing to prepare
+
+# Where a rank's preparation fails, before the load's collectives: what it raises, and with what.
+_FAILURES = {
+    "metadata": (HuggingFaceStorageReader, "read_metadata", OSError, "cannot read the checkpoint"),
+    "materialize": (nn.Module, "to_empty", MemoryError, "out of device memory"),
+}
+
+
+def _one_rank_fails_preparing_job(rank, size, ckpt_dir, failure):
+    owner, attr, error, message = _FAILURES[failure]
+    model = _meta_twin(get_tiny_llama)
+    entered, refused = [], None
+
+    def dcp_load(state_dict, **kwargs):
+        # As the load's collectives do: a rank in them waits for every peer.
+        entered.append(rank)
+        torch.distributed.barrier()
+
+    with (
+        bounded_collectives(),
+        mock.patch.object(dcp, "load", dcp_load),
+        mock.patch.object(owner, attr, side_effect=error(message)) if rank == 1 else nullcontext(),
+    ):
+        try:  # not pytest.raises: its failure would end the pooled worker instead of the job
+            _meta_dcp_load(model, ckpt_dir, "cpu", load_names=_every_name(model))
+        except Exception as e:
+            refused = str(e)
+    # Rank 1's failure on both ranks, before either entered the load.
+    assert not entered, (rank, entered)
+    assert refused is not None and message in refused, (rank, refused)
+    assert rank == 1 or "rank 1 failed" in refused, refused
+
+
+@pytest.mark.parametrize("failure", list(_FAILURES))
+def test_a_rank_failing_to_prepare_the_load_fails_every_rank(cpu_dist_workers, tmp_path, failure):
+    """Each rank reads the checkpoint and materializes its params before DCP's collective load. One
+    failing there, alone, fails every rank, rather than leave the others waiting in the load."""
+    ckpt_dir = _save(tmp_path, get_tiny_llama().state_dict())
+    cpu_dist_workers.run(_one_rank_fails_preparing_job, ckpt_dir=ckpt_dir, failure=failure)
diff --git a/tests/unit/torch/utils/test_model_load_utils.py b/tests/unit/torch/utils/test_model_load_utils.py
index 55c1dc717a2..15c04ddfadc 100644
--- a/tests/unit/torch/utils/test_model_load_utils.py
+++ b/tests/unit/torch/utils/test_model_load_utils.py
@@ -36,16 +36,23 @@
 import torch
 from _test_utils.torch.distributed.utils import bounded_collectives
 from _test_utils.torch.moe_models import model_with_decoder_at
-from safetensors.torch import save_file
+from safetensors.torch import load_file, save_file
 from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import FSDPModule
+from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
 
 pytest.importorskip("accelerate")
 transformers = pytest.importorskip("transformers")
+import transformers.core_model_loading as core_model_loading
 import transformers.distributed as hf_distributed
+import transformers.modeling_utils as modeling_utils
 from _test_utils.torch.transformers_models import (
     _TinyQwen35Tokenizer,
+    get_tiny_llama,
     get_tiny_qwen3_5_moe_vl,
     get_tiny_qwen3_5_vl,
+    get_tiny_qwen3_moe,
 )
 from packaging.version import Version
 from transformers import (
@@ -82,7 +89,12 @@
     _ckpt_quant_method,
     _conversion_plan,
     _convert_keys,
+    _ep_aware_caching_allocator_warmup,
+    _ep_dense_meta_load,
     _experts_sharded_over,
+    _fsdp2_shard_dense,
+    _load_ep_dense_distributed,
+    _make_expert_param_matcher,
     _resolve_target,
     _seed_parallel_states,
     _transformers_mesh_builder,
@@ -611,13 +623,43 @@ def test_what_a_rank_holds_with_the_dense_part_replicated_must_fit(
 
 
 @needs_qwen3_moe_ep_plan
-def test_the_env_knobs_override_the_automatic_plan(run_as, serve, monkeypatch, tmp_path):
-    monkeypatch.setenv("MODELOPT_EP_DEGREE", "2")
-    monkeypatch.setenv("MODELOPT_FSDP_SIZE", "world")
+def test_a_dense_part_the_expert_parallel_load_cannot_shard_is_replicated(
+    run_as, serve, gpu_memory, tmp_path
+):
+    """The expert-parallel load cannot shard the dense part within several EP groups, which the
+    export cannot write: 2 experts over 4 ranks. Left to decide, the plan replicates it however
+    big, warning with the way out; asked to shard it, it refuses. Without expert parallelism the
+    rule does not apply."""
+    run_as(4)
+    serve(_qwen3_moe(num_experts=2))
+    gpu_memory(1)
+
+    with pytest.warns(UserWarning, match="may run out of memory. Set MODELOPT_EP_DEGREE"):
+        assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2)
+    with pytest.raises(NotImplementedError, match="2 expert-data-parallel replicas"):
+        resolve_distribution_plan(str(tmp_path), dense="fsdp")
+    assert resolve_distribution_plan(str(tmp_path), ep_size=1)[0] == DistributionPlan(
+        4, dense="fsdp"
+    )
+
+
+@needs_qwen3_moe_ep_plan
+@pytest.mark.parametrize(
+    ("knob", "value", "plan"),
+    [
+        ("MODELOPT_EP_DEGREE", "2", DistributionPlan(4, 2)),
+        ("MODELOPT_FSDP_SIZE", "world", DistributionPlan(4, 4, "fsdp")),
+    ],
+    ids=["ep-degree", "fsdp-size"],
+)
+def test_the_env_knobs_override_the_automatic_plan(
+    run_as, serve, monkeypatch, tmp_path, knob, value, plan
+):
+    monkeypatch.setenv(knob, value)
     run_as(4)
     serve(_qwen3_moe(num_experts=8))
 
-    assert resolve_distribution_plan(str(tmp_path))[0] == DistributionPlan(4, 2, "fsdp")
+    assert resolve_distribution_plan(str(tmp_path))[0] == plan
 
 
 @needs_qwen3_moe_ep_plan
@@ -1955,12 +1997,17 @@ def ep_world_of_two(monkeypatch, gloo_world_of_one):
 @pytest.mark.parametrize(
     ("plan", "error", "match"),
     [
-        (DistributionPlan(2), ValueError, "shards no routed experts"),
-        (DistributionPlan(2, ep_size=2, dense="fsdp"), NotImplementedError, "replicates the dense"),
+        pytest.param(DistributionPlan(2), ValueError, "shards no routed experts", id="no-experts"),
+        # FSDP2 within each of several expert-parallel groups: a layout the export cannot write.
+        pytest.param(
+            DistributionPlan(4, 2, "fsdp"),
+            NotImplementedError,
+            "expert-data-parallel",
+            id="dense-hsdp",
+        ),
     ],
-    ids=["no-expert-parallelism", "sharded-dense"],
 )
-def test_the_loader_refuses_a_plan_it_cannot_load_by(plan, error, match):
+def test_the_loader_refuses_a_plan_it_cannot_carry_out(plan, error, match):
     """Before anything else: no process group, transformers release or checkpoint is looked at."""
     with pytest.raises(error, match=match):
         load_and_prepare_ep("unused", plan)
@@ -2595,3 +2642,707 @@ def test_each_rank_builds_the_meshes_on_a_gpu_of_its_own(
     cpu_dist_workers.run(
         _gpu_layout_job, visible=visible, bound=bound, shared=shared, refused=refused
     )
+
+
+# ------------------------------------------------------- the expert-parallel load's dense part
+# ``from_pretrained`` always EP-shards the experts. ``_make_expert_param_matcher`` tells their params
+# from the dense ones, the distribution plan says whether the dense part is replicated or sharded,
+# ``_ep_dense_meta_load`` keeps the dense weights off the device during ``from_pretrained`` when they
+# are sharded, and ``_fsdp2_shard_dense`` / ``_load_ep_dense_distributed`` FSDP2-shard them around
+# the experts. The FSDP2 cases run on a one-rank gloo process group on CPU.
+
+# ------------------------------------------------------------------ _make_expert_param_matcher
+
+# The plan load_and_prepare_ep shards with after the a2a swap: base-relative keys (no ``model.``).
+_A2A_PLAN = {
+    "layers.*.mlp.gate": "torch_a2a_router",
+    "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+    "layers.*.mlp.experts.down_proj": "grouped_gemm",
+    "layers.*.mlp.experts": "torch_a2a_experts",
+}
+
+
+# Only the per-parameter keys: a container key would match the same names on its own.
+_PARAM_KEYS = {k: v for k, v in _A2A_PLAN.items() if v == "grouped_gemm"}
+
+
+@pytest.mark.parametrize(
+    ("plan", "name", "expected"),
+    [
+        pytest.param(_PARAM_KEYS, "layers.0.mlp.experts.gate_up_proj", True, id="key-unprefixed"),
+        pytest.param(
+            _PARAM_KEYS, "model.language_model.layers.3.mlp.experts.down_proj", True, id="vlm"
+        ),
+        pytest.param(
+            _PARAM_KEYS,
+            "language_model.model.layers.3.mlp.experts.gate_up_proj",
+            True,
+            id="vlm-wrapper",
+        ),
+        # NemotronH tags only the container; every param under it is an expert's.
+        pytest.param(
+            {"layers.*.mixer.experts": "torch_a2a_experts"},
+            "backbone.layers.1.mixer.experts.3.down_proj.weight",
+            True,
+            id="under-a-container",
+        ),
+        pytest.param(
+            _A2A_PLAN, "model.layers.0.mlp.shared_experts.gate_proj.weight", False, id="shared"
+        ),
+        pytest.param(
+            {"layers.*.mlp.gate": "ep_router"}, "model.layers.0.mlp.gate.weight", False, id="router"
+        ),
+        # No EP credit then: the load-time warmup falls back to stock sizing.
+        pytest.param(None, "model.layers.0.mlp.experts.gate_up_proj", False, id="no-plan"),
+    ],
+)
+def test_expert_param_matcher(plan, name, expected):
+    """A param is an expert's when the plan says so, under any prefix the base-relative keys omit;
+    a router is replicated whatever its style. The warmup sizing tests classify a whole MoE's
+    experts, router, attention, embedding and head with the same matcher."""
+    assert _make_expert_param_matcher(plan)(name) is expected
+
+
+# ------------------------------------------------------------------ a block-fp8 checkpoint
+
+
+def test_a_block_fp8_checkpoint_is_refused_as_a_format_whatever_its_dense_layout(
+    monkeypatch, gloo_world_of_one
+):
+    """This load does not dequantize block-fp8 checkpoints at all, so that is the refusal a plan
+    sharding their dense part meets, not one advising to replicate it."""
+    monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True)
+    monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2)
+    monkeypatch.setattr(moe, "require_transformers", lambda *args: None)
+    config = SimpleNamespace(quantization_config={"quant_method": "fp8"})
+
+    with pytest.raises(NotImplementedError, match="cannot dequantize a 'fp8' checkpoint"):
+        load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config)
+
+
+# ------------------------------------------------------------------------- _ep_dense_meta_load
+
+
+def _is_expert(name):
+    return ".experts." in name
+
+
+def _get_device_lookups():
+    """The modules whose ``get_device`` global the state-dict loaders call."""
+    return modeling_utils, core_model_loading
+
+
+def test_dense_weights_resolve_to_meta_and_the_rest_to_the_stock_device(monkeypatch):
+    """Only a dense ``.weight`` / ``.bias`` is sent to meta. An expert param or a buffer goes to
+    the stock lookup, with the caller's arguments, in both modules the loaders read it from."""
+    stock_calls = []
+
+    def stock_get_device(device_map, key, *args, **kwargs):
+        stock_calls.append((key, args, kwargs))
+        return device_map[""]
+
+    lookups = _get_device_lookups()
+    for module in lookups:
+        monkeypatch.setattr(module, "get_device", stock_get_device)
+    device_map = {"": "cuda:0"}
+    expert = "model.layers.0.mlp.experts.0.gate_proj.weight"
+    buffer = "model.layers.0.mlp.gate.e_score_correction_bias"
+
+    with _ep_dense_meta_load(_is_expert):
+        for module in lookups:
+            assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.weight") == "meta"
+            assert module.get_device(device_map, "model.layers.0.self_attn.q_proj.bias") == "meta"
+            assert module.get_device(device_map, expert, valid_torch_device=True) == "cuda:0"
+            assert module.get_device(device_map, buffer) == "cuda:0"
+
+    assert stock_calls == [(expert, (), {"valid_torch_device": True}), (buffer, (), {})] * 2
+
+
+def test_the_dense_load_refuses_a_transformers_without_its_get_device_hook(monkeypatch):
+    """Without the hook the dense part would load whole on every rank, the memory sharding it saves."""
+    for module in _get_device_lookups():
+        monkeypatch.delattr(module, "get_device", raising=False)
+
+    with pytest.raises(RuntimeError, match="get_device"), _ep_dense_meta_load(_is_expert):
+        pass
+
+
+def _get_device_patch():
+    return _ep_dense_meta_load(_is_expert), [(m, "get_device") for m in _get_device_lookups()]
+
+
+def _warmup_sizing_patch():
+    if not hasattr(modeling_utils, "get_total_byte_count"):
+        pytest.skip("transformers has no get_total_byte_count to size")
+    sizing = _ep_aware_caching_allocator_warmup(4, _A2A_PLAN)
+    return sizing, [(modeling_utils, "get_total_byte_count")]
+
+
+@pytest.mark.parametrize(
+    "patch", [_get_device_patch, _warmup_sizing_patch], ids=lambda f: f.__name__
+)
+def test_a_load_time_patch_is_undone_even_when_the_load_raises(patch):
+    """``from_pretrained`` reads these transformers globals; a failed load leaves stock ones."""
+    context, patched = patch()
+    stock = [getattr(module, attr) for module, attr in patched]
+
+    with pytest.raises(RuntimeError, match="load failed"), context:
+        assert all(getattr(m, a) is not f for (m, a), f in zip(patched, stock))
+        raise RuntimeError("load failed")
+
+    assert all(getattr(m, a) is f for (m, a), f in zip(patched, stock))
+
+
+def test_from_pretrained_leaves_the_dense_weights_on_meta_and_loads_the_rest(tmp_path):
+    """Through a real CPU ``from_pretrained``: every dense weight stays on meta for the FSDP2 +
+    DCP load to fill, while the MLPs (standing in for experts) and the rotary buffers load."""
+    src = get_tiny_llama()
+    src.save_pretrained(tmp_path)
+
+    with _ep_dense_meta_load(lambda name: ".mlp." in name):
+        model = AutoModelForCausalLM.from_pretrained(tmp_path)
+
+    ref_params, ref_buffers = src.state_dict(), dict(src.named_buffers())
+    for name, param in model.named_parameters():
+        if ".mlp." in name:
+            assert torch.equal(param, ref_params[name]), name
+        else:
+            assert param.is_meta, name
+    for name, buffer in model.named_buffers():
+        assert torch.equal(buffer, ref_buffers[name]), name
+
+
+# ------------------------------------------------------------------------- _fsdp2_shard_dense
+
+
+@pytest.fixture
+def fsdp_mesh(skip_on_windows, tmp_path_factory):
+    """The dense FSDP mesh of a one-rank gloo process group (FSDP2 needs a process group)."""
+    store = tmp_path_factory.mktemp("pg") / "store"
+    torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1)
+    yield init_device_mesh("cpu", (1,), mesh_dim_names=("dp",))
+    torch.distributed.destroy_process_group()
+
+
+class _LayerExperts(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.gate_up_proj = nn.Parameter(torch.randn(4, 8, 4))
+        self.down_proj = nn.Parameter(torch.randn(4, 4, 4))
+
+
+class _DecoderLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = nn.Linear(4, 4)
+        self.mlp = nn.Module()
+        self.mlp.gate = nn.Linear(4, 4, bias=False)  # the router: dense, sharded with its layer
+        self.mlp.experts = _LayerExperts()
+        self.input_layernorm = nn.LayerNorm(4)
+
+
+class _MoeLM(nn.Module):
+    """A causal-LM-shaped MoE whose ``_no_split_modules`` names its decoder layer, as HF's do."""
+
+    _no_split_modules = ["_DecoderLayer"]
+
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.embed_tokens = nn.Embedding(16, 4)
+        self.model.layers = nn.ModuleList(_DecoderLayer() for _ in range(2))
+        self.model.norm = nn.LayerNorm(4)
+        self.lm_head = nn.Linear(4, 16, bias=False)
+
+    def experts(self):
+        return [layer.mlp.experts for layer in self.model.layers]
+
+
+def test_each_decoder_layer_is_sharded_around_its_experts(fsdp_mesh):
+    """A layer's dense params (attention, router, norms) become shards on the FSDP mesh; its
+    experts stay as loaded, for EP. The root's embedding, final norm and head stay replicated, and
+    a layer's submodules keep their classes (modelopt recognizes them by class)."""
+    model = _MoeLM()
+    before = {name: param.detach().clone() for name, param in model.named_parameters()}
+    experts = {name: param for name, param in model.named_parameters() if _is_expert(name)}
+
+    assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2
+
+    assert all(isinstance(layer, FSDPModule) for layer in model.model.layers)
+    assert not isinstance(model, FSDPModule)
+    assert type(model.model.layers[0].self_attn) is nn.Linear
+    for name, param in model.named_parameters():
+        if name in experts:
+            assert param is experts[name] and not isinstance(param, DTensor), name
+        elif name.startswith("model.layers."):
+            assert isinstance(param, DTensor), name
+            assert param.device_mesh == fsdp_mesh and param.placements == (Shard(0),), name
+            assert torch.equal(param.full_tensor(), before[name]), name
+        else:
+            assert not isinstance(param, DTensor), name
+
+
+def test_a_param_off_the_models_dtype_stays_replicated(fsdp_mesh):
+    """FSDP2 shards one dtype per group, and HF keeps a few params in fp32, as a router: such a
+    param stays replicated, as the pure-FSDP2 wrap leaves it, and its layer still unshards."""
+    model = _MoeLM().to(torch.bfloat16)
+    gate = model.model.layers[0].mlp.gate.float()
+
+    assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2
+
+    assert not isinstance(gate.weight, DTensor) and gate.weight.dtype == torch.float32
+    for layer in model.model.layers:
+        layer.unshard()  # FSDP2 checks a group's dtypes here
+
+
+def test_a_layer_holding_only_experts_is_left_unwrapped(fsdp_mesh):
+    model = _MoeLM()
+    bare = model.model.layers[1]
+    bare.self_attn = bare.mlp.gate = bare.input_layernorm = nn.Identity()
+
+    assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 1
+    assert not isinstance(bare, FSDPModule)
+
+
+def test_a_dense_part_left_with_nothing_to_shard_is_refused(fsdp_mesh):
+    """Experts outweighing the dense part in another dtype put every dense param off the model's
+    dominant dtype, out of the shard groups: no layer is left to shard, and the load would hold the
+    whole dense part on every rank."""
+    model = _MoeLM()
+    for experts in model.experts():
+        experts.to(torch.bfloat16)
+
+    with pytest.raises(RuntimeError, match="no dense parameter to shard"):
+        _fsdp2_shard_dense(model, fsdp_mesh, model.experts())
+
+
+def test_a_dense_param_the_loader_already_sharded_stays_as_it_is(fsdp_mesh):
+    """A dense param that arrives as a ``Shard`` DTensor of the loader's own plan is not resharded;
+    the rest of its layer still is."""
+    tp_mesh = init_device_mesh("cpu", (1,), mesh_dim_names=("tp",))
+    model = _MoeLM()
+    attn = model.model.layers[0].self_attn
+    attn.weight = nn.Parameter(distribute_tensor(attn.weight.detach(), tp_mesh, [Shard(1)]))
+    presharded = attn.weight
+
+    assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 2
+
+    assert attn.weight is presharded and attn.weight.placements == (Shard(1),)
+    assert model.model.layers[0].mlp.gate.weight.placements == (Shard(0),)
+
+
+def test_a_replicated_dense_param_is_sharded_over_the_fsdp_mesh_like_a_plain_one(fsdp_mesh):
+    ep_mesh = init_device_mesh("cpu", (1, 1), mesh_dim_names=("edp", "ep"))["ep"]
+    model = _MoeLM()
+    gate = model.model.layers[0].mlp.gate
+    weight = gate.weight.detach().clone()
+    gate.weight = nn.Parameter(distribute_tensor(weight, ep_mesh, [Replicate()]))
+
+    _fsdp2_shard_dense(model, fsdp_mesh, model.experts())
+
+    assert gate.weight.device_mesh == fsdp_mesh and gate.weight.placements == (Shard(0),)
+    assert torch.equal(gate.weight.full_tensor(), weight)
+
+
+class _VisionBlock(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(4, 4)
+
+
+class _VisionTower(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.blocks = nn.ModuleList(_VisionBlock() for _ in range(2))
+
+
+class _MoeVlm(_MoeLM):
+    """Kimi's MoonViT style: ``_no_split_modules`` lists the vision tower *and* its blocks."""
+
+    _no_split_modules = ["_VisionTower", "_VisionBlock", "_DecoderLayer"]
+
+    def __init__(self):
+        super().__init__()
+        self.vision_tower = _VisionTower()
+
+
+def test_a_listed_container_is_wrapped_once_not_again_through_its_blocks(fsdp_mesh):
+    model = _MoeVlm()
+
+    assert _fsdp2_shard_dense(model, fsdp_mesh, model.experts()) == 3
+
+    assert isinstance(model.vision_tower, FSDPModule)
+    assert not any(isinstance(block, FSDPModule) for block in model.vision_tower.blocks)
+
+
+class _ShardMesh:
+    """The meshes, as far as the loader reads them before it loads the dense part."""
+
+    device_type = "cpu"
+
+    def __getitem__(self, name):
+        return self
+
+    def size(self):
+        return 2
+
+    def get_group(self):
+        return -1
+
+
+def test_a_dense_part_from_pretrained_did_not_leave_on_meta_is_refused(
+    monkeypatch, gloo_world_of_one
+):
+    """A transformers that places the dense part without asking ``get_device`` loads it whole on every
+    rank, the memory this layout exists to save: the load refuses that before sharding anything."""
+
+    class _DenseLoadedError(Exception):
+        """Raised by the dense load: the loader got past the check."""
+
+    def dense_load(*args, **kwargs):
+        raise _DenseLoadedError
+
+    loader = SimpleNamespace(__name__="Loader", from_pretrained=lambda path, **kwargs: _MoeLM())
+    monkeypatch.setattr("modelopt.torch.utils.distributed.is_initialized", lambda: True)
+    monkeypatch.setattr("modelopt.torch.utils.distributed.size", lambda group=None: 2)
+    monkeypatch.setattr(moe, "require_transformers", lambda *args: None)
+    monkeypatch.setattr(
+        "torch.distributed.device_mesh.init_device_mesh", lambda *a, **k: _ShardMesh()
+    )
+    monkeypatch.setattr(model_load_utils, "ensure_local_checkpoint", lambda p, **k: (None, p))
+    monkeypatch.setattr(model_load_utils, "_auto_class", lambda *args, **kwargs: loader)
+    monkeypatch.setattr(model_load_utils, "record_unplaced_source_keys", lambda *args: [])
+    monkeypatch.setattr(model_load_utils, "_load_ep_dense_distributed", dense_load)
+    config = SimpleNamespace(
+        base_model_ep_plan={
+            "layers.*.mlp.gate": "ep_router",
+            "layers.*.mlp.experts": "moe_tp_experts",
+            "layers.*.mlp.experts.gate_up_proj": "grouped_gemm",
+            "layers.*.mlp.experts.down_proj": "grouped_gemm",
+        },
+        num_experts=4,
+    )
+
+    with pytest.raises(RuntimeError, match="no dense param came back on meta"):
+        load_and_prepare_ep("unused", DistributionPlan(2, 2, "fsdp"), hf_config=config)
+
+
+# ------------------------------------------------------------------ _load_ep_dense_distributed
+
+
+def test_the_dense_load_shards_first_then_asks_for_exactly_the_meta_dense_params(
+    fsdp_mesh, monkeypatch, tmp_path
+):
+    """As ``_ep_dense_meta_load`` leaves the model: dense weights on meta; experts loaded, here
+    per-expert ``nn.Linear``s, whose ``.weight`` only the expert matcher tells from a dense one;
+    and a Mamba-style ``A_log`` (no ``.weight`` / ``.bias`` suffix) loaded too. The dense layers
+    are FSDP2-sharded before the shared meta+DCP loader runs, and it is asked for exactly the meta
+    dense params, on the mesh's device. A recorder stands in for that load."""
+    with torch.device("meta"):
+        model = _MoeLM()
+    for layer in model.model.layers:
+        layer.mlp.experts = nn.ModuleList(nn.Linear(4, 4, bias=False) for _ in range(4))
+        layer.A_log = nn.Parameter(torch.zeros(4))
+    loads = []
+
+    def record_load(loaded, local_dir, device, *, load_names):
+        sharded = [isinstance(layer.self_attn.weight, DTensor) for layer in loaded.model.layers]
+        loads.append((local_dir, device, load_names, sharded))
+
+    monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load)
+
+    wrapped = _load_ep_dense_distributed(
+        model, fsdp_mesh, str(tmp_path), _is_expert, model.experts()
+    )
+
+    assert wrapped == 2
+    meta_dense = {name for name, param in model.named_parameters() if param.is_meta}
+    assert loads == [(str(tmp_path), torch.device("cpu"), meta_dense, [True, True])]
+    assert not any(_is_expert(name) or name.endswith("A_log") for name in meta_dense)
+
+
+def test_the_dense_load_materializes_on_the_gpu_the_rank_is_bound_to(fsdp_mesh, monkeypatch):
+    """Not on the one its local rank numbers: a launch that shows each process one GPU binds every
+    rank to cuda:0 while the local ranks run 0, 1, ..., and the experts and FSDP2's compute are on
+    the bound one."""
+    loads = []
+    monkeypatch.setenv("LOCAL_RANK", "1")
+    for name, value in [
+        ("is_available", True),
+        ("is_initialized", True),
+        ("current_device", 0),
+        ("device_count", 1),
+    ]:
+        monkeypatch.setattr(torch.cuda, name, lambda value=value: value)
+    monkeypatch.setattr(model_load_utils, "_fsdp2_shard_dense", lambda *args: 0)
+
+    def record_load(model, local_dir, device, *, load_names):
+        loads.append(device)
+
+    monkeypatch.setattr(model_load_utils, "_meta_dcp_load", record_load)
+
+    _load_ep_dense_distributed(
+        nn.Linear(2, 2), SimpleNamespace(device_type="cuda"), "unused", lambda name: False, []
+    )
+
+    assert loads == [torch.device("cuda", 0)]
+
+
+class _PastTheWrappingError(Exception):
+    """Raised in DCP's place: the rank went on past the wrapping."""
+
+
+def _one_rank_fails_the_wrapping_job(rank, size):
+    def wrap(*args):
+        if rank == 1:
+            raise MemoryError("out of memory sharding a dense layer")
+        return 0
+
+    def dcp_load(*args, **kwargs):
+        torch.distributed.barrier()  # DCP's collectives: a rank in them waits for every peer
+        raise _PastTheWrappingError
+
+    refused = None
+    with (
+        bounded_collectives(),
+        mock.patch.object(model_load_utils, "_fsdp2_shard_dense", wrap),
+        mock.patch.object(model_load_utils, "_meta_dcp_load", dcp_load),
+    ):
+        try:  # not pytest.raises: its failure would end the pooled worker instead of the job
+            _load_ep_dense_distributed(
+                nn.Linear(2, 2), SimpleNamespace(device_type="cpu"), "unused", lambda n: False, []
+            )
+        except Exception as error:
+            refused = f"{type(error).__name__}: {error}"
+    assert refused is not None and "out of memory sharding a dense layer" in refused, (
+        rank,
+        refused,
+    )
+    assert rank == 1 or "rank 1 failed" in refused, refused
+
+
+def test_a_rank_failing_to_wrap_the_dense_part_fails_every_rank(cpu_dist_workers):
+    """Wrapping allocates each rank's shards on its own, so it can fail on one rank alone; every
+    rank then raises that failure, rather than leave the others waiting for it in DCP's
+    collectives (stood in by a bounded barrier)."""
+    cpu_dist_workers.run(_one_rank_fails_the_wrapping_job)
+
+
+def test_a_tied_head_the_checkpoint_stores_alone_reaches_both_aliases(fsdp_mesh, tmp_path):
+    """A tied checkpoint can store the pair as ``lm_head.weight`` alone, which ``from_pretrained``
+    loads into both aliases. With a tied MoE as the expert-parallel load leaves it -- the dense part
+    on meta and tied, the experts loaded -- the dense load fills both aliases with it too."""
+    src = get_tiny_qwen3_moe(tie_word_embeddings=True)
+    src.save_pretrained(tmp_path)
+    weights = tmp_path / "model.safetensors"
+    stored = load_file(weights)
+    stored["lm_head.weight"] = stored.pop("model.embed_tokens.weight")
+    save_file(stored, weights, metadata={"format": "pt"})
+    with _ep_dense_meta_load(_is_expert):
+        model = AutoModelForCausalLM.from_pretrained(tmp_path)
+    model.tie_weights()  # as load_and_prepare_ep does once the weights load
+    assert model.lm_head.weight is model.model.embed_tokens.weight
+    assert model.lm_head.weight.is_meta
+
+    experts = [layer.mlp.experts for layer in model.model.layers]
+    _load_ep_dense_distributed(model, fsdp_mesh, str(tmp_path), _is_expert, experts)
+
+    for alias in (model.lm_head.weight, model.model.embed_tokens.weight):
+        assert torch.equal(alias, src.lm_head.weight)
+
+
+# ------------------------------------------------- the expert-parallel load's allocator warmup
+# ``from_pretrained`` reserves the bytes ``get_total_byte_count`` reports for each device in one
+# allocation. Under the EP load it would count every expert at full size (the experts are most of a
+# large MoE), so ``_ep_aware_caching_allocator_warmup`` swaps in a sizing that counts one EP shard of
+# each expert. The sizing is called directly here, and once through ``caching_allocator_warmup`` with
+# the CUDA queries and the allocation faked.
+
+needs_sizing_hook = pytest.mark.skipif(
+    not hasattr(modeling_utils, "get_total_byte_count"),
+    reason="transformers has no get_total_byte_count to size",
+)
+
+_GPU = torch.device("cuda", 0)
+
+# _TinyMoe's fp32 bytes: the fused experts; the dense params (embedding, attention weight + bias,
+# router); the router's persistent ``e_score_correction_bias`` buffer. The tied head adds none.
+_EXPERT_BYTES = 4 * (4 * 16 * 8 + 4 * 8 * 8)
+_DENSE_BYTES = 4 * (16 * 8 + 8 * 8 + 8 + 4 * 8)
+_BUFFER_BYTES = 4 * 4
+_ATTN_BYTES = 4 * (8 * 8 + 8)
+
+
+class _TinyMoe(nn.Module):
+    """One MoE decoder layer, a head tied to the embedding, and a persistent router buffer."""
+
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Module()
+        self.model.embed_tokens = nn.Embedding(16, 8)
+        layer = nn.Module()
+        layer.self_attn = nn.Linear(8, 8)
+        layer.mlp = nn.Module()
+        layer.mlp.gate = nn.Linear(8, 4, bias=False)
+        layer.mlp.gate.register_buffer("e_score_correction_bias", torch.zeros(4))
+        layer.mlp.experts = nn.Module()
+        layer.mlp.experts.gate_up_proj = nn.Parameter(torch.zeros(4, 16, 8))
+        layer.mlp.experts.down_proj = nn.Parameter(torch.zeros(4, 8, 8))
+        self.model.layers = nn.ModuleList([layer])
+        self.lm_head = nn.Linear(8, 16, bias=False)
+        self.lm_head.weight = self.model.embed_tokens.weight
+        self.all_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    def get_parameter_or_buffer(self, name):
+        """``PreTrainedModel``'s lookup, which the sizing goes through."""
+        try:
+            return self.get_parameter(name)
+        except AttributeError:
+            return self.get_buffer(name)
+
+    def device_map(self):
+        """Every param and persistent buffer on one GPU, as ``from_pretrained`` expands it."""
+        names = [name for name, _ in self.named_parameters(remove_duplicate=False)]
+        names += [name for name, _ in self.named_buffers()]
+        return dict.fromkeys(names, _GPU)
+
+
+class _OneBytePerElement:
+    def param_element_size(self, model, name, param):
+        return 1
+
+
+_ALL_BYTES = _EXPERT_BYTES + _DENSE_BYTES + _BUFFER_BYTES
+
+
+@needs_sizing_hook
+@pytest.mark.parametrize(
+    ("dense_on_meta", "quantizer", "stock_bytes", "per_rank_bytes"),
+    [
+        # Experts count one EP shard; the router is replicated, so it counts in full like the rest.
+        pytest.param(
+            False,
+            None,
+            _ALL_BYTES,
+            _EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES,
+            id="dense-replicated",
+        ),
+        # Sharded dense: every dense .weight / .bias stays on meta during from_pretrained, while
+        # persistent buffers still load onto the device.
+        pytest.param(
+            True, None, _ALL_BYTES, _EXPERT_BYTES // 4 + _BUFFER_BYTES, id="dense-on-meta"
+        ),
+        # The quantizer sizes the elements: one byte each where fp32 takes four.
+        pytest.param(
+            False,
+            _OneBytePerElement(),
+            _ALL_BYTES // 4,
+            (_EXPERT_BYTES // 4 + _DENSE_BYTES + _BUFFER_BYTES) // 4,
+            id="quantizer",
+        ),
+    ],
+)
+def test_transformers_reserves_only_the_per_rank_footprint(
+    monkeypatch, dense_on_meta, quantizer, stock_bytes, per_rank_bytes
+):
+    """End to end: the one up-front allocation ``caching_allocator_warmup`` makes shrinks from the
+    whole model (stock sizing, which also checks the byte counts above) to what this EP rank
+    holds."""
+    reserved = []
+
+    def fake_empty(numel, *, dtype, device, **kwargs):
+        reserved.append((numel * torch.tensor([], dtype=dtype).element_size(), device))
+
+    model = _TinyMoe()
+    with monkeypatch.context() as m:
+        m.setattr(torch.cuda, "mem_get_info", lambda index: (80 * 1024**3, 80 * 1024**3))
+        m.setattr(torch.cuda, "memory_reserved", lambda index: 0)
+        m.setattr(torch.cuda, "memory_allocated", lambda index: 0)
+        m.setattr(torch, "empty", fake_empty)
+        modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer)
+        with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=dense_on_meta):
+            modeling_utils.caching_allocator_warmup(model, model.device_map(), quantizer)
+
+    assert reserved == [(stock_bytes, _GPU), (per_rank_bytes, _GPU)]
+
+
+@needs_sizing_hook
+def test_a_tp_planned_dense_param_keeps_its_tp_share_while_the_rest_stays_on_meta(monkeypatch):
+    """transformers places a TP-planned param itself, never asking where the dense part goes, so it
+    is on the device at its TP share even when the rest of the dense part stays on meta. The share
+    is of the ep submesh it is placed over, not of the world, which holds two of them here."""
+    model = _TinyMoe()
+    model._tp_plan = {"model.layers.*.self_attn": "colwise"}
+    monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
+    monkeypatch.setattr(torch.distributed, "get_world_size", lambda group=None: 8)
+    monkeypatch.setattr(torch.distributed, "get_rank", lambda group=None: 0)
+
+    with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN, dense_on_meta=True):
+        total = modeling_utils.get_total_byte_count(model, model.device_map())
+
+    assert total == {_GPU: _EXPERT_BYTES // 4 + _ATTN_BYTES // 4 + _BUFFER_BYTES}
+
+
+@needs_sizing_hook
+def test_an_error_while_sizing_falls_back_to_the_stock_sizing(monkeypatch):
+    """The warmup only sizes an allocator reservation, so whatever the EP-aware sizing cannot
+    handle gets transformers' own number instead of failing the load, with whatever else a later
+    release passes the hook, and a warning: that number can reserve most of the device."""
+    stock_calls = []
+
+    def stock_sizing(model, accelerator_device_map, hf_quantizer=None, *args, **kwargs):
+        stock_calls.append((model, accelerator_device_map, hf_quantizer, args, kwargs))
+        return {_GPU: 123}
+
+    class _UnsupportedQuantizer:
+        def param_element_size(self, model, name, param):
+            raise TypeError("unsupported param")
+
+    monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing)
+    model, quantizer = _TinyMoe(), _UnsupportedQuantizer()
+    device_map = model.device_map()
+
+    with _ep_aware_caching_allocator_warmup(4, _A2A_PLAN):
+        with pytest.warns(UserWarning, match="EP-aware warmup sizing failed"):
+            sizing = modeling_utils.get_total_byte_count(model, device_map, quantizer, 1, flag=True)
+        assert sizing == {_GPU: 123}
+
+    assert stock_calls == [(model, device_map, quantizer, (1,), {"flag": True})]
+
+
+@needs_sizing_hook
+def test_a_sizing_hook_with_other_parameters_keeps_the_stock_sizing(monkeypatch):
+    """A release whose hook takes other leading parameters keeps transformers' own sizing: the
+    replacement would not know how the hook is called. It says so, since that sizing may reserve
+    most of the device."""
+
+    def stock_sizing(model, device_map, quantizer=None):
+        return {_GPU: 123}
+
+    monkeypatch.setattr(modeling_utils, "get_total_byte_count", stock_sizing)
+
+    with (
+        pytest.warns(UserWarning, match="using its full-model sizing"),
+        _ep_aware_caching_allocator_warmup(4, _A2A_PLAN),
+    ):
+        assert modeling_utils.get_total_byte_count is stock_sizing
+
+
+@needs_sizing_hook
+@pytest.mark.parametrize("ep_degree", [1, 0])
+def test_without_expert_parallelism_the_stock_sizing_stays(ep_degree):
+    stock = modeling_utils.get_total_byte_count
+
+    with _ep_aware_caching_allocator_warmup(ep_degree, _A2A_PLAN):
+        assert modeling_utils.get_total_byte_count is stock
+
+
+@needs_sizing_hook
+def test_a_transformers_without_the_sizing_hook_is_left_alone(monkeypatch):
+    monkeypatch.delattr(modeling_utils, "get_total_byte_count")
+
+    with (
+        pytest.warns(UserWarning, match="has no get_total_byte_count"),
+        _ep_aware_caching_allocator_warmup(4, _A2A_PLAN),
+    ):
+        assert not hasattr(modeling_utils, "get_total_byte_count")