From ef4d4b9e0a8e80dc20872e06de1d518d89b8ad35 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:37:15 +0000 Subject: [PATCH 01/67] Pack NVFP4 weights that are DTensor shards The no-gather export packs weights while the model is still sharded. NVFP4's packing broke there in three ways: the block scales' zero fill used a boolean-mask assignment, which lowers to index_put_, which DTensor cannot dispatch; the fp4 cast mixed the DTensor with plain constants; and the nibble pack's << is a silent no-op on a DTensor. The scales use torch.where, the cast runs on the local shard and rewraps it with its global shape and stride, so an uneven shard keeps its real size, and the pack uses torch.bitwise_left_shift. Plain tensors pack exactly as before. The one-rank gloo world the meta-DCP load tests used moves to the shared conftest, for these tests too. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- .../quantization/qtensor/nvfp4_tensor.py | 56 ++++++++++++-- tests/unit/torch/conftest.py | 15 ++++ .../export/test_sharded_weight_packing.py | 74 +++++++++++++++++++ .../torch/quantization/test_nvfp4_tensor.py | 26 +++++++ tests/unit/torch/utils/test_meta_dcp_load.py | 9 --- 5 files changed, 165 insertions(+), 15 deletions(-) create mode 100644 tests/unit/torch/export/test_sharded_weight_packing.py diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 8726dc5c719..3ac450c2908 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -163,7 +163,9 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - per_block_scale[per_block_scale == 0] = 1.0 + # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, + # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure num_blocks_per_row = weight.shape[-1] // block_size @@ -213,8 +215,11 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0 - per_block_scale[per_block_scale == 0] = 1.0 + # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: + # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device + # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is + # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: per_block_scale = _cast_per_block_scale_to_fp8(per_block_scale) @@ -248,6 +253,33 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" + # The distributed export path packs a SHARDED weight. Every step below is elementwise over the + # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound + # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and + # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and + # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time + # instead just moves the failure to the next one. + _dtensor_meta = None + try: + from torch.distributed.tensor import DTensor as _DTensor + + if isinstance(weight, _DTensor): + # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as + # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and + # gives each rank a different global shape. Downstream that desynchronises the + # exporter's per-rank tensor classification and surfaces as a KeyError deep in the + # writer, long after the real mistake. + _dtensor_meta = ( + _DTensor, + weight.device_mesh, + weight.placements, + weight.shape, + weight.stride(), + ) + weight = weight.to_local() + except ImportError: + pass + device = weight.device # Extract sign and compute absolute values in one pass @@ -266,7 +298,18 @@ def _cast_fp4(cls, weight: torch.Tensor): ) # Combine sign, ordinal, and rounding adjustment - return (sign_bit << 3) + ord + equals_odd_bounds + packed = (sign_bit << 3) + ord + equals_odd_bounds + if _dtensor_meta is not None: + _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta + packed = _dt_cls.from_local( + packed, + _mesh, + _placements, + run_check=False, + shape=_gshape, + stride=_gstride, + ) + return packed @classmethod def quantize( @@ -347,8 +390,9 @@ def quantize( return scaled_weight # Cast weights to fp4 q_weight = cls._cast_fp4(scaled_weight) - # Pack weights - packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] + # Pack weights. torch.bitwise_left_shift, not ``<< 4``: on a DTensor the operator is a silent + # no-op (torch 2.11), which would pack unrelated codes on the distributed export path. + packed_weight = torch.bitwise_left_shift(q_weight[..., 1::2], 4) | q_weight[..., 0::2] return ( cls(input_shape, input_dtype, packed_weight), weights_scaling_factor, diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index 852412d2565..de576f629c4 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -16,6 +16,7 @@ from collections import defaultdict import pytest +import torch from modelopt.torch.opt.mode import _ModeRegistryCls @@ -52,3 +53,17 @@ def use_channel_div_4(): for name, kv_for_all in lookup.items(): for regex, (k, v) in kv_for_all.items(): config_class.model_fields[name].default[regex][k] = v + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group, so the code under test takes its distributed path, and its + CPU device mesh.""" + from torch.distributed.device_mesh import init_device_mesh + + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + try: + yield init_device_mesh("cpu", (1,)) + finally: + torch.distributed.destroy_process_group() diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py new file mode 100644 index 00000000000..a1c9fb19e58 --- /dev/null +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" + +import torch +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from modelopt.torch.quantization.qtensor import NVFP4QTensor + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): + """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. + + The block scales come out sharded like the weight, an all-zero block included: filling it by + boolean-mask assignment lowers to ``index_put_``, which DTensor cannot dispatch. + """ + torch.manual_seed(0) + weight = torch.randn(4, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 # an all-zero block + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + + sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert isinstance(sharded_scale, DTensor) and sharded_scale.placements == (Shard(0),) + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale)) + assert torch.equal(packed.full_tensor(), want) + # The second-level scale read off the sharded weight itself, as without one handed in. + assert torch.equal( + NVFP4QTensor.get_weights_scaling_factor(sharded, 16)[1], + NVFP4QTensor.get_weights_scaling_factor(weight, 16)[1], + ) + + +def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): + """``from_local`` would infer local x world rows; an uneven shard must keep the real total. + + Two rows held of a three-row tensor is what one rank holds when three rows split over two. + """ + local = torch.randn(2, 32) * 4 + sharded = DTensor.from_local( + local.clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 32]), + stride=(32, 1), + ) + + packed = NVFP4QTensor._cast_fp4(sharded) + + assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) + assert packed.shape == (3, 32) + assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index a7b64515e3a..ce1159711ba 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -106,6 +106,32 @@ def test_helper_clamps_underflow_to_min(self): out = _cast_per_block_scale_to_fp8(tiny).float() assert (out > 0).all(), f"FP8 cast produced zero scales: {out.tolist()}" + @staticmethod + def _dynamic_block_scales(): + weight = torch.zeros(1, 48) + weight[0, 16] = 6.0 # amax 6 at a second-level scale of 1: block scale exactly 1 + weight[0, 32] = 3.0 # amax 3: block scale 0.5 + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) + return scale.float()[0] + + @staticmethod + def _static_block_scales(): + q = SimpleNamespace( + block_sizes={-1: 16}, + global_amax=torch.tensor(6.0 * 448.0), # a second-level scale of 1 + _amax=torch.tensor([0.0, 6.0, 3.0]), + ) + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(q, torch.randn(1, 48)) + return scale.float()[0] + + @pytest.mark.parametrize("path", ["dynamic", "static"]) + def test_an_all_zero_block_gets_a_scale_of_one(self, path): + """Blocks of amax 0, 6 and 3: the zero block gets 1, as the amax-6 block does, not the clamp.""" + zero, unit, half = getattr(self, f"_{path}_block_scales")().tolist() + + assert zero == unit == 1.0 + assert half == 0.5 + def test_static_path_no_nan_when_block_amax_zero(self): """Static path: zero-amax block + small global_amax must clamp to 448, not cast to NaN.""" block_size = 16 diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 764bce545c6..4c642b1644c 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -474,15 +474,6 @@ def recording_load(state_dict, **kwargs): # ------------------------------------------------------------------- under a process group -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group, so the loader takes its distributed path.""" - store = tmp_path_factory.mktemp("pg") / "store" - torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield - torch.distributed.destroy_process_group() - - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the From 34d5415aa54df4cb85c6b4e5c62662423b1839ff Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:38:33 +0000 Subject: [PATCH 02/67] Pack an FSDP2-sharded model's weights in held weight-update windows The no-gather export packs weights while the model is still sharded, with no weight-access window open. ExportContext.keep_fused_experts marks that path. On it, a quantized linear whose dense part FSDP2 shards is packed inside a weight-update window, which hands FSDP2 the packed weights; packed outside one, the unit's next reshard would put its unquantized shard back. Under held_fsdp2_weight_windows, which the export's dispatch loop now enters, a unit's window stays open from one module to the next, so a decoder layer is gathered once rather than once per linear; any other loop gets a window per module, closed by its handler. Where expert parallelism leaves the root unwrapped, the window opens on the layer's own unit, as on a wrapped root it would do nothing. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 98 +++++++++++++++++++ modelopt/torch/export/registry.py | 20 ++++ modelopt/torch/export/unified_export_hf.py | 14 ++- .../torch/quantization/utils/core_utils.py | 4 +- .../export/test_sharded_weight_packing.py | 63 ++++++++++++ 5 files changed, 195 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 200ec6e5dae..d28e7be022d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -16,9 +16,16 @@ """Built-in module handlers for unified Hugging Face export.""" import collections.abc +import contextlib import warnings +from dataclasses import dataclass import torch.nn as nn +from torch.distributed.fsdp import FSDPModule + +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps +from modelopt.torch.quantization.utils.core_utils import ModuleNames, _get_enclosing_fsdp_module +from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .moe_utils import _export_fused_experts @@ -34,6 +41,87 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") +@dataclass +class _HeldWindow: + """An FSDP2 unit's weight-update window, and the modules packed in it so far.""" + + unit: nn.Module + modules: list[nn.Module] + stack: contextlib.ExitStack + + +def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: + """The innermost FSDP2 unit enclosing ``module``, or the model if none does.""" + if ctx.module_names is None: + ctx.module_names = module_name_maps(ctx.model) + return _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + + +def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: + """The root a weight-update window on ``unit`` is opened on, and its module names. + + The window opens on the innermost unit enclosing the module, but only when the root it is handed + is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the + dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which + packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated + codes next to correct scales). So there that unit is handed in as the root, with names relative + to it, which the window resolves the modules' parameters against. + """ + if isinstance(ctx.model, FSDPModule) or unit is ctx.model: + # The model's own names, which _enclosing_unit has built. + return ctx.model, ctx.module_names + return unit, module_name_maps(unit) + + +def _fsdp2_weight_window( + ctx: ExportContext, module: nn.Module +) -> contextlib.AbstractContextManager: + """The ``fsdp2_aware_weight_update`` window to pack ``module``'s weights in. + + Inside :func:`held_fsdp2_weight_windows`, the window of the module's FSDP2 unit, held open + across the unit's modules. Anywhere else, a window of the module's own, closed when the handler + is done, so a dispatch loop that knows nothing of held windows cannot leave a unit unsharded. + """ + unit = _enclosing_unit(ctx, module) + if not ctx.hold_fsdp2_windows: + root, names = _window_root(ctx, unit) + return fsdp2_aware_weight_update(root, module, names=names) + held = ctx.fsdp2_window + if held is not None and held.unit is unit: + held.modules.append(module) + return contextlib.nullcontext() + _close_fsdp2_weight_window(ctx) + root, names = _window_root(ctx, unit) + held = _HeldWindow(unit, [module], contextlib.ExitStack()) + held.stack.enter_context(fsdp2_aware_weight_update(root, held.modules, names=names)) + ctx.fsdp2_window = held + return contextlib.nullcontext() + + +@contextlib.contextmanager +def held_fsdp2_weight_windows(ctx: ExportContext): + """Within the block, pack each FSDP2 unit's modules in one weight-update window. + + The window unshards the whole unit and reshards it when it closes, so a window per module + gathers a decoder layer once per quantized linear in it. Here the unit's window stays open from + one handler to the next until the export reaches another unit, each module joining the list + whose weights it re-registers with FSDP2, and leaving the block closes the last. + """ + ctx.hold_fsdp2_windows = True + try: + yield + finally: + ctx.hold_fsdp2_windows = False + _close_fsdp2_weight_window(ctx) + + +def _close_fsdp2_weight_window(ctx: ExportContext) -> None: + """Close the weight-update window the handlers hold open, if any, resharding its unit.""" + held, ctx.fsdp2_window = ctx.fsdp2_window, None + if held is not None: + held.stack.close() + + def _export_weight( module: nn.Module, ctx: ExportContext, @@ -43,6 +131,16 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight + # The no-gather writer packs while the model is still sharded, so unlike the + # collect_export_tensors path there is no open weight-access window here and the weight is a + # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an + # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, + # which reshards the unit when it closes: a layer left unsharded holds all of its weights on + # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: + with _fsdp2_weight_window(ctx, module): + _export_quantized_weight(module, ctx.dtype, weight_name) + return _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 1aa575d92de..808210a4f68 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,10 +28,14 @@ from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn +if TYPE_CHECKING: + from modelopt.torch.quantization.utils.core_utils import ModuleNames + __all__ = [ "ExportContext", "ExportHandler", @@ -59,6 +63,22 @@ class ExportContext: model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" + # Distributed (no-gather) export: the handlers pack while the model is still sharded, an + # FSDP2-owned weight inside its unit's weight-update window. + keep_fused_experts: bool = False + + module_names: "ModuleNames | None" = None + """``model``'s name maps, built once per export by the first handler that looks a module up + (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + + hold_fsdp2_windows: bool = False + """Whether the handlers hold an FSDP2 unit's weight-update window open across its modules: set + only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which + closes the last window. Elsewhere each module is packed in a window of its own.""" + + fsdp2_window: Any = None + """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 14f3c82fd36..3e849506ebb 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -104,7 +104,7 @@ export_sparse_attention_config = None # Importing the built-in handlers installs their entries in the two registries. -from . import hf_export_handlers as _hf_export_handlers # noqa: F401 +from . import hf_export_handlers as _hf_export_handlers from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import get_experts_list, is_layernorm, is_quantlinear, sync_moe_gate_up_amax from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model @@ -636,6 +636,12 @@ def _export_quantized_weight( The export includes converting weight tensor to correct quantized values and quantized dtype, and registering scaling factors. + An export handler packs through ``hf_export_handlers._export_weight`` rather than calling this + directly: on the keep-fused export an FSDP2-owned weight must be packed inside its unit's + weight-update window, which that helper opens. + Packed outside it, the unit's next reshard puts FSDP2's unquantized shard back, and the + checkpoint is written unquantized. + Tied-weight dedup is not handled here: both sides of a tie are packed independently (identically, once ``sync_tied_input_amax`` has equalized their scales), and the duplicate is dropped by name in :func:`postprocess_state_dict`. Deduping per-module at @@ -996,8 +1002,10 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - for name, sub_module in model.named_modules(): - _dispatch_export_handler(name, sub_module, ctx) + # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. + with _hf_export_handlers.held_fsdp2_weight_windows(ctx): + for name, sub_module in model.named_modules(): + _dispatch_export_handler(name, sub_module, ctx) def _prepare_model_for_export(model, dtype, is_modelopt_qlora): diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b4120720b5c..0915a9d1720 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1054,7 +1054,9 @@ def fsdp2_aware_weight_update( Args: root_model (nn.Module): The root model of the FSDPModule. modules_to_update (list): The list of modules to update which should be a list of modules that are - direct children of the FSDPModule. + direct children of the FSDPModule. A list is read again when the context exits, so a caller + may append modules of the same FSDP unit to it inside the context, and their weights are + re-registered with FSDP2 too (the export's held weight windows rely on this). reshard (bool): Whether to reshard the FSDP root module after the weight update. Returns: diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index a1c9fb19e58..0e666784169 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -15,10 +15,21 @@ """Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +import contextlib +import copy + +import pytest import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def _raw_bytes(tensor): @@ -72,3 +83,55 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) assert packed.shape == (3, 32) assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) + + +@pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) +@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( + monkeypatch, gloo_world_of_one, root_sharded, held +): + """On the no-gather path the model is still sharded when a linear is packed. + + Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the + DTensor's local view is refused under autograd), so the handlers pack inside a weight-update + window, on the layer's own unit where expert parallelism leaves the root unsharded. The window + gathers the whole unit, so within held windows it opens once for all of the unit's linears; + any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is + sharded again afterwards, and the sharded state dict carries the packed weights, as an + unsharded pack has them. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + refs = copy.deepcopy(model)[0] + for i, ref in enumerate(refs): + ExportModuleRegistry.match(ref)( + f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) + ) + fully_shard(layer, mesh=gloo_world_of_one) + if root_sharded: + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + unshards = [] + unshard = FSDPModule.unshard + monkeypatch.setattr( + FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) + ) + + with handlers.held_fsdp2_weight_windows(ctx) if held else contextlib.nullcontext(): + for i, linear in enumerate(layer): + ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) + + assert unshards == ([layer] if held else [layer, layer]) + assert all(isinstance(linear.weight, DTensor) for linear in layer) + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for i, ref in enumerate(refs): + packed = sd[f"0.{i}.weight"] + assert isinstance(packed, DTensor) + assert packed.dtype == torch.uint8 + assert torch.equal(packed.full_tensor(), ref.weight) + for name in ("weight_scale", "weight_scale_2", "input_scale"): + want = _raw_bytes(getattr(ref, name)) + assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name From f774bed9a29bfec39e7070ba47c15129b76c42aa Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:09:46 +0000 Subject: [PATCH 03/67] Fold fused experts' quantizers into their weights shard-locally The fused-experts export splits each 3-D expert weight into full per-expert tensors before quantizing them. Under FSDP2 that first all-gathers the weight, so every rank ends up holding every expert: a 235B MoE OOMed at ~157 GB per rank regardless of node count. With keep_fused_experts set on the export context, the export handler instead folds each expert's quantizer into the fused weight in place. On a sharded weight it quantizes only this rank's experts, taking each one's quantizer by its global expert number, and keeps the result sharded with the original expert count, so an uneven split (3 experts over 2 ranks) keeps its numbering. split_fused_experts_state_dict later cuts the folded weight and its stacked scales into exactly the per-expert keys the single-process export writes. Under FSDP2 the fold, and every dense weight's packing on this path, runs in an FSDP2 weight window opened on the unit that owns the weight -- the model when it is wrapped, else the enclosing layer, which is all expert parallelism wraps. The window hands FSDP2 the packed weights and reshards on the way out; without it, the next reshard puts FSDP2's own unquantized shards back. NVFP4 can pack a sharded weight: the per-block scale avoids index_put_, which DTensor cannot dispatch, _cast_fp4 packs the local shard and keeps its global shape, and the nibble pack uses torch.bitwise_left_shift, since on a DTensor `x << 4` returns x unchanged (torch 2.11). Nothing sets keep_fused_experts yet; the writer that consumes the fold follows. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 58 +- modelopt/torch/export/moe_utils.py | 386 +++++++++++ modelopt/torch/export/registry.py | 5 + .../quantization/qtensor/nvfp4_tensor.py | 60 +- .../export/test_fused_experts_keep_fused.py | 641 ++++++++++++++++++ 5 files changed, 1142 insertions(+), 8 deletions(-) create mode 100644 tests/unit/torch/export/test_fused_experts_keep_fused.py diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 200ec6e5dae..3b9f4023509 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -20,8 +20,11 @@ import torch.nn as nn +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update +from modelopt.torch.utils.distributed import is_fsdp2_model + from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax -from .moe_utils import _export_fused_experts +from .moe_utils import _export_fused_experts, _export_fused_experts_keep_fused from .quant_format import QUANTIZATION_NONE from .quant_utils import get_quantization_format from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -34,6 +37,30 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") +def _fsdp2_weight_window(model: nn.Module, module: nn.Module): + """``fsdp2_aware_weight_update`` on the FSDP2 unit that owns ``module``'s weights. + + That is ``model`` when the model itself is wrapped. Expert parallelism FSDP2-shards only the dense + layers and leaves the root unwrapped, and a window opened on an unwrapped root does nothing: the + weight stays a sharded DTensor, which packing mishandles (``DTensor << int`` is a silent no-op in + torch 2.11, so NVFP4 packed unrelated codes next to correct scales). So open it on the enclosing + unit instead. + """ + from torch.distributed.fsdp import FSDPModule + + root = model + if not isinstance(model, FSDPModule): + root = next( + ( + unit + for unit in model.modules() + if isinstance(unit, FSDPModule) and any(m is module for m in unit.modules()) + ), + model, + ) + return fsdp2_aware_weight_update(root, module) + + def _export_weight( module: nn.Module, ctx: ExportContext, @@ -43,6 +70,17 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight + # The no-gather writer packs while the model is still sharded, so unlike the + # collect_export_tensors path there is no open weight-access window here and the weight is a + # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an + # index_put_, which DTensor cannot dispatch. Open a window per module for that path only, and + # reshard on the way out: a layer left unsharded holds all of its weights on every rank, and + # FSDP2 later puts its own shards back over anything packed in it outside a window. + if ctx.keep_fused_experts and is_fsdp2_model(ctx.model): + with _fsdp2_weight_window(ctx.model, module): + _export_quantized_weight(module, ctx.dtype, weight_name) + return + _export_quantized_weight(module, ctx.dtype, weight_name) @@ -138,8 +176,24 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex Tied experts are packed independently and their duplicate keys are dropped by name in postprocess_state_dict; no per-module dedup cache is used. + + Distributed no-gather export (``ctx.keep_fused_experts``) folds the quantizers into the fused + weight IN PLACE, quantizing only this rank's local experts and keeping the result ``Shard(0)`` -- + the per-expert split happens later on the sharded weight (``split_fused_experts_state_dict``). + Otherwise (single process, or gather export) it splits into full per-expert submodules here. """ - _export_fused_experts(module, ctx.dtype) + if ctx.keep_fused_experts and is_fsdp2_model(ctx.model): + # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands + # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its + # original shards the next time the layer reshards, and exported unquantized. + with _fsdp2_weight_window(ctx.model, module): + _export_fused_experts_keep_fused(module, ctx.dtype) + elif ctx.keep_fused_experts: + # No-gather distributed export: fold each rank's LOCAL experts and leave the weight + # fused + sharded; the per-expert split happens after consolidation. + _export_fused_experts_keep_fused(module, ctx.dtype) + else: + _export_fused_experts(module, ctx.dtype) @ExportModuleRegistry.register(predicate=is_quantlinear) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 02660142736..79db95916dd 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -110,6 +110,152 @@ def _fallback_amax(weight: torch.Tensor, quantizer) -> torch.Tensor: return torch.maximum(groups.amax(dim=-1), -groups.amin(dim=-1)).float().reshape(like.shape) +def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int): + """Wrap this rank's experts as a DTensor spanning ``num_experts`` rows on the expert axis.""" + from torch.distributed.tensor import DTensor + + shape = torch.Size((num_experts, *local.shape[1:])) + stride = torch.empty(shape, device="meta").stride() + return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) + + +def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> None: + """Fold per-expert quantizers into the FUSED 3-D expert weights *in place*, staying sharded. + + Distributed (DCP) counterpart to :func:`_export_fused_experts`. That function slices the fused + weight into full per-expert tensors -- under FSDP2 it must first all-gather the ``Shard(0)`` + fused weight (via the export handler's ``unshard``) and then materializes every expert on every + rank, putting the whole model on each GPU (OOMs at 235B/480B). This keeps each projection a + single fused param: on a ``Shard(0)`` DTensor it quantizes ONLY this rank's local experts + (``to_local()`` + the DTensor's dim-0 global offset for the matching per-expert weight quantizer) + and rewraps the result ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. + On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model + path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to + the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` + reproduce :func:`_export_fused_experts`'s output exactly. Registers, per projection ``P``: + ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), + ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). + """ + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + try: + from torch.distributed.tensor import DTensor, Replicate, Shard + except Exception: # pragma: no cover - non-distributed install + DTensor = None # noqa: N806 (the optional-import fallback for a class) + try: + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + except Exception: # pragma: no cover - older torch: fall back to even-split offset + compute_local_shape_and_global_offset = None + + first = getattr(module, "_first_proj_attr", "gate_up_proj") + n = module.num_experts + for proj in (first, "down_proj"): + weight = getattr(module, proj) + weight_quantizers = getattr(module, f"{proj}_weight_quantizers") + input_quantizer = getattr(module, f"{proj}_input_quantizer", None) + + # Under FSDP2 the fused 3-D weight is a DTensor sharded on dim 0 (the expert axis). Quantize + # ONLY this rank's local experts and rewrap the result sharded identically -- indexing the + # DTensor (weight[i]) would all-gather the whole fused weight onto every rank, de-sharding + # the model and OOMing the GPU at scale. Plain tensor (single process): all experts. + wdata = weight.data + is_dt = DTensor is not None and isinstance(wdata, DTensor) + if is_dt: + mesh, placements = wdata.device_mesh, wdata.placements + local_w = wdata.to_local() + offset = None + if compute_local_shape_and_global_offset is not None: + try: + _, global_offset = compute_local_shape_and_global_offset( + tuple(wdata.shape), mesh, placements + ) + offset = int(global_offset[0]) + except Exception: + offset = None + if offset is None: # even-split fallback (dim-0 Shard on a 1-D mesh) + offset, nshards = 0, 1 + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + nshards = mesh.size(mdim) + offset = mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) + break + scale_placements = [ + p if (isinstance(p, Shard) and p.dim == 0) else Replicate() for p in placements + ] + local_indices = list(range(offset, offset + local_w.shape[0])) + else: + local_w = wdata + local_indices = list(range(n)) + + fp8_weights: list[torch.Tensor] = [] + weight_scales: list[torch.Tensor] = [] + weight_scale_2s: list[torch.Tensor] = [] + input_scale: torch.Tensor | None = None + for local_idx, global_idx in enumerate(local_indices): + w_quantizer = weight_quantizers[global_idx] + w_slice = local_w[local_idx] + # Uncalibrated-expert fallback: derive amax from this expert's local weight slice + # (matches _export_fused_experts), so a never-routed expert still exports sane scales. + if ( + hasattr(w_quantizer, "is_enabled") + and w_quantizer.is_enabled + and ( + not hasattr(w_quantizer, "_amax") + or w_quantizer._amax is None + or torch.all(w_quantizer._amax == 0) + ) + ): + w_quantizer.amax = w_slice.abs().amax().to(torch.float32) + warnings.warn( + f"Expert {global_idx} {proj} weight quantizer was not calibrated (amax missing " + f"or zero); using weight-derived amax. Increase calibration size to activate all " + f"experts.", + stacklevel=2, + ) + wrapper = nn.Module() + wrapper.weight = nn.Parameter(w_slice.contiguous(), requires_grad=False) + wrapper.weight_quantizer = w_quantizer + if input_quantizer is not None: + wrapper.input_quantizer = input_quantizer + _export_quantized_weight(wrapper, dtype) + fp8_weights.append(wrapper.weight.data) + weight_scales.append(wrapper.weight_scale) + # NVFP4 carries a SECOND per-tensor weight scale (weight_scale_2) that dequantizes the + # per-block weight_scale; FP8 has none. Keep it so the fused buffer (and the per-expert + # split) stay dequantizable. + if hasattr(wrapper, "weight_scale_2"): + weight_scale_2s.append(wrapper.weight_scale_2) + if hasattr(wrapper, "input_scale"): + input_scale = wrapper.input_scale + local_fp8 = torch.stack(fp8_weights, dim=0) + local_ws = torch.stack(weight_scales, dim=0) + if is_dt: + # Keep the global expert count: from_local would infer it as local x world, which is + # wrong for an uneven shard (3 experts on 2 ranks) and misnumbers the experts. + local_fp8 = _rewrap_expert_rows(local_fp8, mesh, placements, wdata.shape[0]) + local_ws = _rewrap_expert_rows(local_ws, mesh, scale_placements, wdata.shape[0]) + setattr(module, proj, nn.Parameter(local_fp8, requires_grad=False)) + module.register_buffer(f"{proj}_weight_scale", local_ws) + if weight_scale_2s: + # Per-expert scalar (E,) -- shards on the expert axis exactly like the weight scale. + local_ws2 = torch.stack(weight_scale_2s, dim=0) + if is_dt: + local_ws2 = _rewrap_expert_rows(local_ws2, mesh, scale_placements, wdata.shape[0]) + module.register_buffer(f"{proj}_weight_scale_2", local_ws2) + if input_scale is not None: + module.register_buffer(f"{proj}_input_scale", input_scale) + + # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers. + for attr in ( + f"{first}_weight_quantizers", + f"{first}_input_quantizer", + "down_proj_weight_quantizers", + "down_proj_input_quantizer", + ): + if hasattr(module, attr): + delattr(module, attr) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, @@ -347,3 +493,243 @@ def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | Non output_path = Path(output_dir) / ".moe.html" output_path.write_text(html_content, encoding="utf-8") print(f"\033[1mExpert token count table saved to {output_path}\033[0m") + + +# ----- Distributed-export helpers (fused-experts split for the no-gather FSDP2 write) ----- +_FUSED_PROJ = ("gate_up_proj", "up_proj", "down_proj") + + +def _fused_experts_prefixes(state_dict: dict) -> list[str]: + """Return the prefixes of the modules holding a fused 3-D expert weight. + + Return the sorted set of ``...experts`` prefixes that carry a fused 3-D expert weight + (``.gate_up_proj`` / ``up_proj`` / ``down_proj``) in ``state_dict``. + """ + prefixes: set[str] = set() + for key, tensor in state_dict.items(): + for proj in _FUSED_PROJ: + if key.endswith(f".experts.{proj}") and hasattr(tensor, "dim") and tensor.dim() == 3: + prefixes.add(key[: -(len(proj) + 1)]) + break + return sorted(prefixes) + + +def split_fused_experts_state_dict(state_dict: dict) -> dict: + """Rewrite fused-experts entries in a (consolidated) HF state dict into per-expert entries. + + The distributed export writes experts *fused* (3-D ``experts.`` weights + fused scale + buffers, kept sharded for DCP). This post-pass splits them into the same per-expert layout the + in-model :func:`_export_fused_experts` produces, so the two export paths yield identical + checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no + quantization. Non-fused-experts keys pass through unchanged. + + Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): + ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` + ``P.down_proj`` ``(E, O, I)`` + ``P.down_proj_weight_scale`` + ``P.down_proj_input_scale`` + Outputs, for each expert ``i``: + ``P.{i}.gate_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) [gated only] + ``P.{i}.up_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) + ``P.{i}.down_proj.weight`` ``(O, I)`` (+ ``.weight_scale``/``.input_scale``) + """ + prefixes = _fused_experts_prefixes(state_dict) + if not prefixes: + return state_dict + + # Keys consumed (and thus removed) while emitting per-expert keys. + consumed: set[str] = set() + out: dict = {} + + for prefix in prefixes: + gated = f"{prefix}.gate_up_proj" in state_dict + first = "gate_up_proj" if gated else "up_proj" + first_w = state_dict[f"{prefix}.{first}"] + down_w = state_dict[f"{prefix}.down_proj"] + n_experts = first_w.shape[0] + + def _scales(proj_name): + ws = state_dict.get(f"{prefix}.{proj_name}_weight_scale") + ins = state_dict.get(f"{prefix}.{proj_name}_input_scale") + # NVFP4 second-level (per-tensor) weight scale; None for FP8. (E,) one scalar per expert. + ws2 = state_dict.get(f"{prefix}.{proj_name}_weight_scale_2") + return ws, ins, ws2 + + first_ws, first_in, first_ws2 = _scales(first) + down_ws, down_in, down_ws2 = _scales("down_proj") + + def _slice_weight_scale(ws, row_slice, fused_rows): + # per-tensor-per-expert scale -> scalar; per-output-channel -> slice the rows. + if ws is None: + return None + if ws.dim() <= 1: # (E,) one scalar per expert + return ws[expert_idx] + sub = ws[expert_idx] # (fused_rows,) or (fused_rows, ...) + return sub if row_slice is None else sub[row_slice] + + for expert_idx in range(n_experts): + if gated: + inter = first_w.shape[1] // 2 + # gate and up are halves of the fused gate_up: they share its single per-tensor + # weight_scale_2 (first_ws2), matching the in-model export + vLLM's W1/W3 fusion. + projections = [ + ( + "gate_proj", + first_w[expert_idx, :inter, :], + first_ws, + first_in, + first_ws2, + slice(0, inter), + ), + ( + "up_proj", + first_w[expert_idx, inter:, :], + first_ws, + first_in, + first_ws2, + slice(inter, None), + ), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + else: + projections = [ + ("up_proj", first_w[expert_idx], first_ws, first_in, first_ws2, None), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + for proj_name, weight, ws, ins, ws2, row_slice in projections: + base = f"{prefix}.{expert_idx}.{proj_name}" + # .clone() every emitted tensor so no two keys share storage. gate/up share the + # per-tensor scale object (and slices alias the fused parent); without cloning, + # safetensors/save_pretrained shared-tensor dedup silently drops the duplicate + # (e.g. up_proj.weight_scale/input_scale would go missing). + out[f"{base}.weight"] = weight.detach().clone().contiguous() + sliced_ws = _slice_weight_scale(ws, row_slice, weight.shape[0]) + if sliced_ws is not None: + out[f"{base}.weight_scale"] = sliced_ws.detach().clone().contiguous() + # weight_scale_2 is a per-expert scalar (shared by gate/up): _slice_weight_scale's + # dim<=1 branch returns ws2[expert_idx] regardless of row_slice. NVFP4 only; None for FP8. + sliced_ws2 = _slice_weight_scale(ws2, row_slice, weight.shape[0]) + if sliced_ws2 is not None: + out[f"{base}.weight_scale_2"] = sliced_ws2.detach().clone().contiguous() + if ins is not None: + out[f"{base}.input_scale"] = ins.detach().clone() + + # Mark the fused tensors + their scale buffers consumed. + for proj in (first, "down_proj"): + for suffix in ("", "_weight_scale", "_weight_scale_2", "_input_scale"): + consumed.add(f"{prefix}.{proj}{suffix}") + + out.update({key: tensor for key, tensor in state_dict.items() if key not in consumed}) + return out + + +def _dtensor_dim0_offset(dt) -> int: + """Global dim-0 start index of this rank's local shard of a ``Shard(0)`` DTensor.""" + from torch.distributed.tensor import Shard + + mesh, placements = dt.device_mesh, dt.placements + try: + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + + _, global_offset = compute_local_shape_and_global_offset(tuple(dt.shape), mesh, placements) + return int(global_offset[0]) + except Exception: + n = dt.shape[0] + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + nshards = mesh.size(mdim) + return mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) + return 0 + + +def _split_local_fused_module( + prefix: str, sharded_sd: dict, ep_rank: int = 0, total_experts: int = 0 +) -> dict: + """Split this rank's local shard of a fused-experts module into per-expert keys. + + Split THIS rank's *local* shard of one fused-experts module into per-expert keys with GLOBAL + expert indices (plain, local tensors -- no all-gather). + + FSDP2/EP shard the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank + owns is whole and local. We reuse :func:`split_fused_experts_state_dict` on the ``to_local()`` + shard (which numbers experts ``0..E_local-1``) and then shift each expert index by the rank's + global dim-0 offset, so the resulting per-expert keys match the single-process export and stay + on this rank for a direct distributed write (no consolidation). + + The fused expert tensor can arrive in two EP layouts, distinguished by its dim-0 size vs + ``total_experts`` (the config expert count): + * **per-EP-group** (size < total): transformers pre-sliced the param to this rank's EP group + (FSDP x EP -> a dp-sharded DTensor of the group; or a plain ep-shard). ``_dtensor_dim0_offset`` + only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``. + * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts + (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base. + ``total_experts == 0`` keeps the legacy per-EP-group behavior. + """ + from torch.distributed.tensor import DTensor + + gated = f"{prefix}.gate_up_proj" in sharded_sd + first = "gate_up_proj" if gated else "up_proj" + fw = sharded_sd[f"{prefix}.{first}"] + # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns. + dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0 + n_local = ( + fw.to_local().shape[0] + if isinstance(fw, DTensor) + else (fw.shape[0] if fw is not None else 0) + ) + group_size = int(fw.shape[0]) if fw is not None else 0 + # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or + # total is unknown (legacy). When it spans all experts, dp_offset is already global. + add_ep_base = ( + bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts) + ) + offset = dp_offset + (ep_rank * group_size if add_ep_base else 0) + + def _loc(key): + v = sharded_sd.get(key) + if v is None: + return None + return v.to_local() if isinstance(v, DTensor) else v + + def _loc_scale(key): + # Weight scales are dim-0 (per-expert) tensors. A DTensor scale shards with the weight -> + # to_local() aligns. But under EP the keep-fused fold emits the scale as a PLAIN buffer + # spanning the whole EP group (not dp-sharded like the weight param); slice it by the + # weight's within-group dp_offset so each scale pairs with its own dp-local expert. + v = sharded_sd.get(key) + if v is None: + return None + if isinstance(v, DTensor): + return v.to_local() + if v.dim() >= 1 and v.shape[0] > n_local and dp_offset + n_local <= v.shape[0]: + return v[dp_offset : dp_offset + n_local] + return v + + # Build a fused sub-dict of this rank's local experts (plain), then split + reindex. + local_fused: dict = { + f"{prefix}.{first}": _loc(f"{prefix}.{first}"), + f"{prefix}.down_proj": _loc(f"{prefix}.down_proj"), + } + for proj in (first, "down_proj"): + ws = _loc_scale(f"{prefix}.{proj}_weight_scale") + if ws is not None: + local_fused[f"{prefix}.{proj}_weight_scale"] = ws + # NVFP4 second-level per-tensor weight scale: also a per-expert (dim-0) tensor, so the same + # dp-offset slicing pairs each scalar with its dp-local expert. None for FP8. + ws2 = _loc_scale(f"{prefix}.{proj}_weight_scale_2") + if ws2 is not None: + local_fused[f"{prefix}.{proj}_weight_scale_2"] = ws2 + ins = sharded_sd.get(f"{prefix}.{proj}_input_scale") # shared scalar (replicated) + if ins is not None: + local_fused[f"{prefix}.{proj}_input_scale"] = ins + + split_local = split_fused_experts_state_dict( + local_fused + ) # keys {prefix}.{i}.* (i in 0..E_local-1) + if offset == 0: + return split_local + + plen = len(prefix) + 1 + out: dict = {} + for key, val in split_local.items(): + i_str, tail = key[plen:].split(".", 1) + out[f"{prefix}.{int(i_str) + offset}.{tail}"] = val + return out diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 1aa575d92de..d0187c59321 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -59,6 +59,11 @@ class ExportContext: model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" + # Distributed (no-gather) export: keep MoE experts FUSED + sharded instead of splitting them + # into full per-expert tensors, so each rank materializes only its local experts (the split + # happens post-consolidation on the sharded fused weight). See _export_fused_experts_keep_fused. + keep_fused_experts: bool = False + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 8726dc5c719..4d46e203341 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -163,7 +163,11 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - per_block_scale[per_block_scale == 0] = 1.0 + # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, + # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + per_block_scale = torch.where( + per_block_scale == 0, torch.ones_like(per_block_scale), per_block_scale + ) # Reshape per_block_scale to match weight's block structure num_blocks_per_row = weight.shape[-1] // block_size @@ -213,8 +217,13 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0 - per_block_scale[per_block_scale == 0] = 1.0 + # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: + # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device + # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is + # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + per_block_scale = torch.where( + per_block_scale == 0, torch.ones_like(per_block_scale), per_block_scale + ) if not keep_high_precision: per_block_scale = _cast_per_block_scale_to_fp8(per_block_scale) @@ -248,6 +257,33 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" + # The distributed export path packs a SHARDED weight. Every step below is elementwise over the + # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound + # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and + # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and + # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time + # instead just moves the failure to the next one. + _dtensor_meta = None + try: + from torch.distributed.tensor import DTensor as _DTensor + + if isinstance(weight, _DTensor): + # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as + # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and + # gives each rank a different global shape. Downstream that desynchronises the + # exporter's per-rank tensor classification and surfaces as a KeyError deep in the + # writer, long after the real mistake. + _dtensor_meta = ( + _DTensor, + weight.device_mesh, + weight.placements, + weight.shape, + weight.stride(), + ) + weight = weight.to_local() + except ImportError: + pass + device = weight.device # Extract sign and compute absolute values in one pass @@ -266,7 +302,18 @@ def _cast_fp4(cls, weight: torch.Tensor): ) # Combine sign, ordinal, and rounding adjustment - return (sign_bit << 3) + ord + equals_odd_bounds + packed = (sign_bit << 3) + ord + equals_odd_bounds + if _dtensor_meta is not None: + _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta + packed = _dt_cls.from_local( + packed, + _mesh, + _placements, + run_check=False, + shape=_gshape, + stride=_gstride, + ) + return packed @classmethod def quantize( @@ -347,8 +394,9 @@ def quantize( return scaled_weight # Cast weights to fp4 q_weight = cls._cast_fp4(scaled_weight) - # Pack weights - packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] + # Pack weights. torch.bitwise_left_shift, not ``<< 4``: on a DTensor the operator is a silent + # no-op (torch 2.11), which would pack unrelated codes on the distributed export path. + packed_weight = torch.bitwise_left_shift(q_weight[..., 1::2], 4) | q_weight[..., 0::2] return ( cls(input_shape, input_dtype, packed_weight), weights_scaling_factor, diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py new file mode 100644 index 00000000000..88f97c32a8d --- /dev/null +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -0,0 +1,641 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Keep-fused MoE expert export used by the no-gather FSDP2 writer (CPU, no GPU needed). + +The no-gather writer never splits a fused 3-D expert weight while the model is sharded: +``_export_fused_experts_keep_fused`` folds every expert's quantizer into the fused weight in place, +and ``split_fused_experts_state_dict`` / ``_split_local_fused_module`` cut the result into +per-expert keys later. That pipeline must write exactly what the single-process +``_export_fused_experts`` writes, so most tests compare the two byte for byte on a tiny synthetic +fused-experts module quantized with FP8 and NVFP4. Sharded cases run on a gloo process group: one +rank in process, and two spawned ranks where a non-zero expert offset matters. +""" + +import copy +import math +import types +import warnings + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from _test_utils.torch.distributed.utils import spawn_multiprocess_job +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +pytest.importorskip("transformers") +pytest.importorskip("safetensors") +from safetensors.torch import load_file, save_file + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import ( + _dtensor_dim0_offset, + _export_fused_experts, + _export_fused_experts_keep_fused, + _fused_experts_prefixes, + _split_local_fused_module, + split_fused_experts_state_dict, +) +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.qtensor import NVFP4QTensor + +NUM_EXPERTS = 4 +HIDDEN_DIM = 32 +INTERMEDIATE_DIM = 16 +PREFIX = "model.layers.0.mlp.experts" + +FP8_WEIGHT_ONLY_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG) +FP8_WEIGHT_ONLY_CFG["quant_cfg"].append({"quantizer_name": "*input_quantizer", "enable": False}) + +QUANT_CFGS = pytest.mark.parametrize( + "quant_cfg", + [mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG], + ids=["fp8", "fp8_weight_only", "nvfp4"], +) + + +class _KeepFusedGatedExperts(nn.Module): + """Fused gate+up experts in the transformers 5.x layout (Qwen3-MoE / Mixtral).""" + + def __init__(self, num_experts=NUM_EXPERTS): + super().__init__() + self.num_experts = num_experts + self.gate_up_proj = nn.Parameter( + torch.randn(num_experts, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 + ) + self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def forward(self, hidden_states, top_k_index, top_k_weights): + out = torch.zeros_like(hidden_states) + mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: + top_k_pos, token_idx = torch.where(mask[expert_idx]) + gate, up = F.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]).chunk( + 2, dim=-1 + ) + current = F.linear(self.act_fn(gate) * up, self.down_proj[expert_idx]) + out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) + return out + + +class _KeepFusedUngatedExperts(nn.Module): + """Ungated fused experts (NemotronH): a single ``up_proj`` and ``down_proj``.""" + + def __init__(self, num_experts=NUM_EXPERTS): + super().__init__() + self.num_experts = num_experts + self.up_proj = nn.Parameter(torch.randn(num_experts, INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02) + self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def forward(self, hidden_states, top_k_index, top_k_weights): + out = torch.zeros_like(hidden_states) + mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: + top_k_pos, token_idx = torch.where(mask[expert_idx]) + up = self.act_fn(F.linear(hidden_states[token_idx], self.up_proj[expert_idx])) + current = F.linear(up, self.down_proj[expert_idx]) + out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) + return out + + +LAYOUTS = pytest.mark.parametrize( + "experts_cls", + [_KeepFusedGatedExperts, _KeepFusedUngatedExperts], + ids=["gated", "ungated"], +) + + +@pytest.fixture(autouse=True) +def _unregister_synthetic_experts(): + """``mtq.quantize`` registers the synthetic classes on the fly; keep the registry clean.""" + yield + for cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts): + if QuantModuleRegistry.get(cls) is not None: + QuantModuleRegistry.unregister(cls) + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group and its CPU device mesh.""" + store = tmp_path_factory.mktemp("pg") / "store" + dist.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + yield init_device_mesh("cpu", (1,)) + dist.destroy_process_group() + + +def _quantized_experts(experts_cls, quant_cfg, num_experts=NUM_EXPERTS): + """Build, quantize and calibrate a fused-experts module; every expert sees one token.""" + torch.manual_seed(0) + model = nn.Module() + model.experts = experts_cls(num_experts) + + def forward_loop(m): + hidden = torch.randn(num_experts, HIDDEN_DIM) + m.experts(hidden, torch.arange(num_experts).unsqueeze(1), torch.ones(num_experts, 1)) + + mtq.quantize(model, quant_cfg, forward_loop) + return model.experts + + +def _prefixed(state_dict): + return {f"{PREFIX}.{key}": value for key, value in state_dict.items()} + + +def _per_expert_export(experts): + """What the single-process export writes: ``_export_fused_experts`` on a copy.""" + module = copy.deepcopy(experts) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + _export_fused_experts(module, torch.bfloat16) + return _prefixed(module.state_dict()) + + +def _keep_fused_export(experts): + """The fused state ``_export_fused_experts_keep_fused`` leaves on a copy.""" + module = copy.deepcopy(experts) + _export_fused_experts_keep_fused(module, torch.bfloat16) + return _prefixed(module.state_dict()) + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def _assert_byte_identical(got, want): + assert set(got) == set(want), (sorted(set(want) - set(got)), sorted(set(got) - set(want))) + for key, ref in want.items(): + assert got[key].dtype == ref.dtype, key + assert got[key].shape == ref.shape, key + assert torch.equal(_raw_bytes(got[key]), _raw_bytes(ref)), key + + +def _assert_quantized(state_dict): + """Guard against a config that quantized nothing, which would make any comparison trivial.""" + weights = [v for k, v in state_dict.items() if k.endswith(".weight")] + assert weights, "no per-expert weights" + assert all(v.dtype in (torch.float8_e4m3fn, torch.uint8) for v in weights) + + +# ------------------------------------------------------------------ keep-fused + split + + +@LAYOUTS +@QUANT_CFGS +def test_keep_fused_then_split_writes_exactly_the_per_expert_export(experts_cls, quant_cfg): + experts = _quantized_experts(experts_cls, quant_cfg) + want = _per_expert_export(experts) + _assert_quantized(want) + + _assert_byte_identical(split_fused_experts_state_dict(_keep_fused_export(experts)), want) + + +@LAYOUTS +@QUANT_CFGS +def test_keep_fused_leaves_one_fused_weight_and_stacked_scales_per_projection( + experts_cls, quant_cfg +): + experts = _quantized_experts(experts_cls, quant_cfg) + first = experts._first_proj_attr + nvfp4 = quant_cfg is mtq.NVFP4_DEFAULT_CFG + has_input_scale = quant_cfg is not FP8_WEIGHT_ONLY_CFG + + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + suffixes = {"", "_weight_scale"} + suffixes |= {"_weight_scale_2"} if nvfp4 else set() + suffixes |= {"_input_scale"} if has_input_scale else set() + assert set(experts.state_dict()) == {f"{p}{s}" for p in (first, "down_proj") for s in suffixes} + for proj in (first, "down_proj"): + weight = getattr(experts, proj) + assert weight.dim() == 3 and weight.shape[0] == NUM_EXPERTS + assert weight.dtype == (torch.uint8 if nvfp4 else torch.float8_e4m3fn) + assert getattr(experts, f"{proj}_weight_scale").shape[0] == NUM_EXPERTS + if nvfp4: + assert getattr(experts, f"{proj}_weight_scale_2").shape == (NUM_EXPERTS,) + if has_input_scale: + assert getattr(experts, f"{proj}_input_scale").dim() == 0 + # The quantizers' information now lives in the buffers above. + assert not [n for n, _ in experts.named_children() if "quantizer" in n] + + +@LAYOUTS +def test_uncalibrated_expert_takes_the_same_weight_derived_amax_as_the_per_expert_export( + experts_cls, +): + """An expert calibration never reached falls back to an amax read off its own weight. + + For the gated layout that amax must come from the whole fused ``gate_up`` slice, so gate and up + keep one shared ``weight_scale_2`` -- the per-expert export guarantees that, and vLLM's W1/W3 + fusion depends on it. + """ + experts = _quantized_experts(experts_cls, mtq.NVFP4_DEFAULT_CFG) + first = experts._first_proj_attr + for proj in (first, "down_proj"): + getattr(experts, f"{proj}_weight_quantizers")[1]._amax.zero_() + want = _per_expert_export(experts) + + with pytest.warns(UserWarning, match=r"Expert 1 .* not calibrated"): + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + + _assert_byte_identical(got, want) + if first == "gate_up_proj": + gate_ws2 = got[f"{PREFIX}.1.gate_proj.weight_scale_2"] + assert torch.equal(gate_ws2, got[f"{PREFIX}.1.up_proj.weight_scale_2"]) + + +# ------------------------------------------------------------------ split_fused_experts_state_dict + + +LOCAL_EXPERTS = 2 # experts in each hand-built fused entry below + + +def _fused_state_dict(*, gated, rows=4, cols=6): + """A fused-experts entry with position-coded values, so every slice is identifiable.""" + first = "gate_up_proj" if gated else "up_proj" + fused_rows = 2 * rows if gated else rows + n = LOCAL_EXPERTS + return { + f"{PREFIX}.{first}": torch.arange(n * fused_rows * cols, dtype=torch.float32) + .reshape(n, fused_rows, cols) + .to(torch.float8_e4m3fn), + # Per-row (per-channel) weight scale: its rows are split along with the weight's. + f"{PREFIX}.{first}_weight_scale": torch.arange(n * fused_rows).view(n, fused_rows) + 100.0, + # Per-expert scalars. + f"{PREFIX}.{first}_weight_scale_2": torch.arange(n) + 7.0, + f"{PREFIX}.{first}_input_scale": torch.tensor(0.5), + f"{PREFIX}.down_proj": torch.arange(n * cols * rows, dtype=torch.float32) + .reshape(n, cols, rows) + .to(torch.float8_e4m3fn), + f"{PREFIX}.down_proj_weight_scale": torch.arange(n) + 3.0, + f"{PREFIX}.down_proj_input_scale": torch.tensor(0.25), + } + + +def test_split_cuts_gate_and_up_rows_with_their_scales(): + sd = _fused_state_dict(gated=True) + fused = dict(sd) + untouched = torch.ones(3) + sd["model.layers.0.mlp.gate.weight"] = untouched + + out = split_fused_experts_state_dict(sd) + + gate_up, ws = fused[f"{PREFIX}.gate_up_proj"], fused[f"{PREFIX}.gate_up_proj_weight_scale"] + for e in range(LOCAL_EXPERTS): + base = f"{PREFIX}.{e}" + assert torch.equal(_raw_bytes(out[f"{base}.gate_proj.weight"]), _raw_bytes(gate_up[e, :4])) + assert torch.equal(_raw_bytes(out[f"{base}.up_proj.weight"]), _raw_bytes(gate_up[e, 4:])) + assert torch.equal(out[f"{base}.gate_proj.weight_scale"], ws[e, :4]) + assert torch.equal(out[f"{base}.up_proj.weight_scale"], ws[e, 4:]) + # gate and up are halves of one projection: they share its per-tensor second scale. + for proj in ("gate_proj", "up_proj"): + assert out[f"{base}.{proj}.weight_scale_2"].item() == 7.0 + e + assert out[f"{base}.{proj}.input_scale"].item() == 0.5 + down = fused[f"{PREFIX}.down_proj"][e] + assert torch.equal(_raw_bytes(out[f"{base}.down_proj.weight"]), _raw_bytes(down)) + assert out[f"{base}.down_proj.weight_scale"].item() == 3.0 + e + assert out[f"{base}.down_proj.input_scale"].item() == 0.25 + assert not any(k in out for k in fused), "fused entries must be consumed" + assert out["model.layers.0.mlp.gate.weight"] is untouched + + +def test_split_keeps_an_ungated_up_proj_whole(): + sd = _fused_state_dict(gated=False) + up = sd[f"{PREFIX}.up_proj"] + + out = split_fused_experts_state_dict(sd) + + assert not any(".gate_proj." in k for k in out) + for e in range(LOCAL_EXPERTS): + assert torch.equal(_raw_bytes(out[f"{PREFIX}.{e}.up_proj.weight"]), _raw_bytes(up[e])) + assert torch.equal( + out[f"{PREFIX}.{e}.up_proj.weight_scale"], sd[f"{PREFIX}.up_proj_weight_scale"][e] + ) + assert f"{PREFIX}.up_proj" not in out + + +def test_split_gives_every_tensor_its_own_storage(tmp_path): + """Slices alias the fused parent, and gate/up share one scale per expert. + + Aliasing keeps the whole fused tensor alive behind every per-expert key, and a shared scale is + refused by safetensors (``save_pretrained`` dedups it away instead). + """ + out = split_fused_experts_state_dict(_fused_state_dict(gated=True)) + + storages = [t.untyped_storage().data_ptr() for t in out.values()] + assert len(set(storages)) == len(storages) + path = tmp_path / "experts.safetensors" + save_file(out, str(path)) + reloaded = load_file(str(path)) + assert set(reloaded) == set(out) + for key, value in out.items(): + assert torch.equal(_raw_bytes(reloaded[key]), _raw_bytes(value)), key + + +def test_split_leaves_a_state_dict_without_fused_experts_unchanged(): + sd = { + "model.embed_tokens.weight": torch.ones(4, 2), + f"{PREFIX}.0.up_proj.weight": torch.ones(4, 2), # already per-expert + } + + out = split_fused_experts_state_dict(sd) + + assert out.keys() == sd.keys() + assert all(out[k] is sd[k] for k in sd) + + +def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): + fused = torch.zeros(2, 4, 4) + sd = { + "model.layers.1.mlp.experts.gate_up_proj": fused, + "model.layers.1.mlp.experts.down_proj": fused, + "model.layers.1.mlp.experts.gate_up_proj_weight_scale": fused, + "model.layers.0.mlp.experts.up_proj": fused, + "model.layers.0.mlp.experts.down_proj": fused, + # Not a fused-experts weight: already split, 2-D, outside an ``experts`` module, no tensor. + "model.layers.2.mlp.experts.0.down_proj": fused, + "model.layers.3.mlp.experts.down_proj": torch.zeros(4, 4), + "model.layers.4.mlp.gate_up_proj": fused, + "model.layers.5.mlp.experts.down_proj": "not a tensor", + } + assert _fused_experts_prefixes(sd) == [ + "model.layers.0.mlp.experts", + "model.layers.1.mlp.experts", + ] + + +@pytest.mark.parametrize( + ("ep_rank", "total_experts", "first_index"), + [ + (0, 0, 0), + # This EP group's slice of a larger expert set: shift by the group base. + (1, 2 * LOCAL_EXPERTS, LOCAL_EXPERTS), + (3, 4 * LOCAL_EXPERTS, 3 * LOCAL_EXPERTS), + # The tensor already spans every expert, so its indices are global. + (1, LOCAL_EXPERTS, 0), + # Total unknown: the legacy per-group shift. + (1, 0, LOCAL_EXPERTS), + ], +) +def test_split_local_numbers_experts_globally(ep_rank, total_experts, first_index): + sd = _fused_state_dict(gated=True) + unshifted = split_fused_experts_state_dict(dict(sd)) + + out = _split_local_fused_module(PREFIX, sd, ep_rank=ep_rank, total_experts=total_experts) + + want = {} + for key, value in unshifted.items(): + expert, tail = key[len(PREFIX) + 1 :].split(".", 1) + want[f"{PREFIX}.{int(expert) + first_index}.{tail}"] = value + _assert_byte_identical(out, want) + + +# ------------------------------------------------------------------ sharded (DTensor) experts + + +def _shard_fused_weights(experts, mesh): + for proj in (experts._first_proj_attr, "down_proj"): + weight = getattr(experts, proj).data + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + setattr(experts, proj, nn.Parameter(sharded, requires_grad=False)) + return experts + + +@LAYOUTS +@QUANT_CFGS +def test_keep_fused_on_a_sharded_weight_stays_sharded_and_splits_like_the_per_expert_export( + gloo_world_of_one, experts_cls, quant_cfg +): + experts = _quantized_experts(experts_cls, quant_cfg) + plain = _keep_fused_export(experts) + sharded = _shard_fused_weights(copy.deepcopy(experts), gloo_world_of_one) + + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + for key, value in fused.items(): + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + else: # weights and per-expert scales stay sharded on the expert axis + assert isinstance(value, DTensor), key + assert value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + value = value.to_local() + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + assert _dtensor_dim0_offset(fused[f"{PREFIX}.{sharded._first_proj_attr}"]) == 0 + + split = _split_local_fused_module(PREFIX, fused) + + assert not any(isinstance(v, DTensor) for v in split.values()) + _assert_byte_identical(split, _per_expert_export(experts)) + + +def test_keep_fused_keeps_the_expert_count_of_an_uneven_shard(gloo_world_of_one): + """Two experts held of three is what rank 0 holds when three experts split over two ranks.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=3) + for proj in ("gate_up_proj", "down_proj"): + full = getattr(experts, proj).data + held = DTensor.from_local( + full[:2].clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=full.shape, + stride=full.stride(), + ) + setattr(experts, proj, nn.Parameter(held, requires_grad=False)) + + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + assert experts.gate_up_proj.shape[0] == 3 + assert experts.gate_up_proj_weight_scale.shape[0] == 3 + + +def _two_rank_split_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + for experts_cls, quant_cfg in [ + (_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG), + (_KeepFusedUngatedExperts, mtq.NVFP4_DEFAULT_CFG), + ]: + # Every rank calibrates the same seeded module, so all ranks hold identical quantizers. + experts = _quantized_experts(experts_cls, quant_cfg) + want = _per_expert_export(experts) + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + first = fused[f"{PREFIX}.{sharded._first_proj_attr}"] + n_local = first.to_local().shape[0] + offset = rank * math.ceil(NUM_EXPERTS / world_size) + assert _dtensor_dim0_offset(first) == offset + mine = { + k: v + for k, v in want.items() + if offset <= int(k[len(PREFIX) + 1 :].split(".")[0]) < offset + n_local + } + assert mine, "this rank owns no experts" + + _assert_byte_identical(_split_local_fused_module(PREFIX, fused), mine) + + # A per-expert scale held as a plain tensor over ALL experts (as keep-fused emits it for an + # unsharded weight) must be cut to this rank's experts before pairing. + unsharded = _keep_fused_export(experts) + plain_scales = { + k: unsharded[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v + for k, v in fused.items() + } + _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) + + +def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(skip_on_windows): + """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index.""" + spawn_multiprocess_job(2, _two_rank_split_job, backend="gloo") + + +# ------------------------------------------------------------------ NVFP4 on sharded weights + + +def test_nvfp4_block_scale_maps_an_all_zero_block_to_one(): + weight = torch.randn(2, 32) + weight[0, :16] = 0.0 # an all-zero block + weight[1, :16] = 0.0 + weight[1, 0] = 6.0 # amax 6 at a second-level scale of 1 -> block scale exactly 1 + + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) + + assert scale.float()[0, 0].item() == 1.0 + assert scale.float()[1, 0].item() == 1.0 + + +def test_nvfp4_static_block_scale_maps_a_zero_amax_block_to_one(): + # Two rows of two blocks: a zero-amax block next to ones whose amax maps to a scale of 1. + quantizer = types.SimpleNamespace( + block_sizes={-1: 16}, + global_amax=torch.tensor(6.0 * 448.0), + _amax=torch.tensor([0.0, 6.0, 6.0, 3.0]), + ) + + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(quantizer, torch.randn(2, 32)) + + assert scale.float()[0, 0] == scale.float()[0, 1] == scale.float()[1, 0] + assert scale.float()[1, 1] != scale.float()[1, 0] + + +def test_nvfp4_block_scale_accepts_a_sharded_weight(gloo_world_of_one): + weight = torch.randn(4, 32) + weight[1, :16] = 0.0 + sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) + + want, want_2 = NVFP4QTensor.get_weights_scaling_factor(weight, 16) + got, got_2 = NVFP4QTensor.get_weights_scaling_factor(sharded, 16) + + assert isinstance(got, DTensor) and got.placements == (Shard(0),) + assert torch.equal(_raw_bytes(got.to_local()), _raw_bytes(want)) + assert torch.equal(got_2, want_2) + + +def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): + """``from_local`` would infer local x world rows; an uneven shard must keep the real total. + + Two rows held of a three-row tensor is what one rank holds when three rows split over two. + """ + local = torch.randn(2, 32) * 4 + sharded = DTensor.from_local( + local.clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 32]), + stride=(32, 1), + ) + + packed = NVFP4QTensor._cast_fp4(sharded) + + assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) + assert packed.shape == (3, 32) + assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) + + +# ------------------------------------------------------------------ export handlers + + +@pytest.mark.parametrize("keep_fused", [True, False], ids=["keep_fused", "split"]) +def test_fused_experts_handler_keeps_experts_fused_only_when_asked(keep_fused): + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.experts = experts + assert ExportContext(model=container, dtype=torch.bfloat16).keep_fused_experts is False + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=keep_fused) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + + names = set(experts.state_dict()) + if keep_fused: + assert {"gate_up_proj", "gate_up_proj_weight_scale", "down_proj"} <= names + assert not any(n.startswith("0.") for n in names) + else: + assert "0.gate_proj.weight" in names + assert "gate_up_proj" not in names + + +def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_one): + """On the no-gather path the model is still sharded when a linear is packed. + + Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the + DTensor's local view is refused under autograd), so the handler packs inside a weight-update + window; the sharded state dict must then carry the packed weight, as an unsharded pack has it. + """ + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + ref = copy.deepcopy(model) + ExportModuleRegistry.match(ref[0])("0", ref[0], ExportContext(model=ref, dtype=torch.bfloat16)) + fully_shard(model[0], mesh=gloo_world_of_one) + fully_shard(model, mesh=gloo_world_of_one) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + ExportModuleRegistry.match(model[0])("0", model[0], ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + packed = sd["0.weight"] + assert isinstance(packed, DTensor) + assert packed.dtype == torch.uint8 + assert torch.equal(packed.full_tensor(), ref[0].weight) + for name in ("weight_scale", "weight_scale_2", "input_scale"): + assert torch.equal(_raw_bytes(sd[f"0.{name}"]), _raw_bytes(getattr(ref[0], name))), name + + +def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): + """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator.""" + torch.manual_seed(0) + weight = torch.randn(4, 32, dtype=torch.bfloat16) + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + + sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert torch.equal(packed.full_tensor(), want) From cc0384cca6bdac6deed6892edd616008b8b3b725 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:21:09 +0000 Subject: [PATCH 04/67] Open the FSDP2 weight window by the model's distribution plan The export handlers pack inside a weight-update window when the model's recorded plan FSDP2-shards its dense part, rather than when they find an FSDP2 module in it. A test now covers the fused experts' window, which had none. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 6 ++-- .../export/test_fused_experts_keep_fused.py | 34 +++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 3b9f4023509..5a1a80b6bf2 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -21,7 +21,7 @@ import torch.nn as nn from modelopt.torch.quantization.utils import fsdp2_aware_weight_update -from modelopt.torch.utils.distributed import is_fsdp2_model +from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .moe_utils import _export_fused_experts, _export_fused_experts_keep_fused @@ -76,7 +76,7 @@ def _export_weight( # index_put_, which DTensor cannot dispatch. Open a window per module for that path only, and # reshard on the way out: a layer left unsharded holds all of its weights on every rank, and # FSDP2 later puts its own shards back over anything packed in it outside a window. - if ctx.keep_fused_experts and is_fsdp2_model(ctx.model): + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: with _fsdp2_weight_window(ctx.model, module): _export_quantized_weight(module, ctx.dtype, weight_name) return @@ -182,7 +182,7 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex the per-expert split happens later on the sharded weight (``split_fused_experts_state_dict``). Otherwise (single process, or gather export) it splits into full per-expert submodules here. """ - if ctx.keep_fused_experts and is_fsdp2_model(ctx.model): + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its # original shards the next time the layer reshards, and exported unquantized. diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 88f97c32a8d..30bdeefccfd 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -24,6 +24,7 @@ rank in process, and two spawned ranks where a non-zero expert offset matters. """ +import contextlib import copy import math import types @@ -56,6 +57,7 @@ from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan NUM_EXPERTS = 4 HIDDEN_DIM = 32 @@ -599,6 +601,37 @@ def test_fused_experts_handler_keeps_experts_fused_only_when_asked(keep_fused): assert "gate_up_proj" not in names +@pytest.mark.parametrize("dense", ["fsdp", "replicate"]) +def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_part( + monkeypatch, dense +): + """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a + weight-update window, they would be overwritten with its original shards at the next reshard.""" + import modelopt.torch.export.hf_export_handlers as handlers + + events = [] + + @contextlib.contextmanager + def window(root, module): + events.append("open") + yield + events.append("close") + + monkeypatch.setattr(handlers, "_fsdp2_weight_window", window) + monkeypatch.setattr( + handlers, "_export_fused_experts_keep_fused", lambda module, dtype: events.append("fold") + ) + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.experts = experts + set_distribution_plan(container, DistributionPlan(2, ep_size=2, dense=dense)) + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) + + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + + assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) + + def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_one): """On the no-gather path the model is still sharded when a linear is packed. @@ -613,6 +646,7 @@ def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_ ExportModuleRegistry.match(ref[0])("0", ref[0], ExportContext(model=ref, dtype=torch.bfloat16)) fully_shard(model[0], mesh=gloo_world_of_one) fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) ExportModuleRegistry.match(model[0])("0", model[0], ctx) From d69333e389ef51f95d8338d54ded8bbd37d0e2d2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 22:09:13 +0000 Subject: [PATCH 05/67] Run the keep-fused two-rank test on the shared gloo workers It spawned two ranks, about 14 s of the linux unit job. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_fused_experts_keep_fused.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 30bdeefccfd..a645555c99d 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -35,7 +35,6 @@ import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F -from _test_utils.torch.distributed.utils import spawn_multiprocess_job from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import fully_shard @@ -508,9 +507,9 @@ def _two_rank_split_job(rank, world_size): _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) -def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(skip_on_windows): +def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index.""" - spawn_multiprocess_job(2, _two_rank_split_job, backend="gloo") + cpu_dist_workers.run(_two_rank_split_job) # ------------------------------------------------------------------ NVFP4 on sharded weights From 4d2d40a5b4c35bea469faf7282ee983f8f5e007e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 20:45:32 +0000 Subject: [PATCH 06/67] Export a fused-experts projection with no quantized expert as is The keep-fused export read every expert's weight_scale, so a recipe that leaves an MoE layer's experts unquantized crashed the no-gather export with an AttributeError after calibration. A projection whose experts are all unquantized now stays as it is, which is what the per-expert export writes for it. One whose experts mix quantized and unquantized cannot be one fused tensor and is refused before either projection is folded. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 17 ++++++++- .../export/test_fused_experts_keep_fused.py | 37 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 79db95916dd..d16697cf659 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -132,7 +132,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` - reproduce :func:`_export_fused_experts`'s output exactly. Registers, per projection ``P``: + reproduce :func:`_export_fused_experts`'s output exactly. Registers, per quantized projection ``P``: ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). """ @@ -149,7 +149,22 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N first = getattr(module, "_first_proj_attr", "gate_up_proj") n = module.num_experts + # One fused tensor cannot hold quantized and unquantized experts side by side, so a projection + # is folded only when all its experts are quantized; one with none quantized stays as it is, + # which is what the per-expert export writes for it. Both are decided before either is folded, + # so a refusal leaves the module untouched. + folded = [] for proj in (first, "down_proj"): + enabled = [q.is_enabled for q in getattr(module, f"{proj}_weight_quantizers")] + if all(enabled): + folded.append(proj) + elif any(enabled): + raise NotImplementedError( + f"{type(module).__name__}.{proj}: weight quantizers " + f"{[i for i, e in enumerate(enabled) if not e]} of {len(enabled)} are disabled and " + "the rest are not; a fused expert weight cannot mix quantized and unquantized experts." + ) + for proj in folded: weight = getattr(module, proj) weight_quantizers = getattr(module, f"{proj}_weight_quantizers") input_quantizer = getattr(module, f"{proj}_input_quantizer", None) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index a645555c99d..129cd253318 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -264,6 +264,43 @@ def test_uncalibrated_expert_takes_the_same_weight_derived_amax_as_the_per_exper assert torch.equal(gate_ws2, got[f"{PREFIX}.1.up_proj.weight_scale_2"]) +def _disable_weight_quantizers(experts, proj, indices): + for i in indices: + getattr(experts, f"{proj}_weight_quantizers")[i].disable() + + +@LAYOUTS +@QUANT_CFGS +@pytest.mark.parametrize("unquantized", ["first", "down_proj", "both"]) +def test_a_projection_with_no_quantized_expert_is_written_as_the_per_expert_export_writes_it( + experts_cls, quant_cfg, unquantized +): + """A recipe can leave an MoE layer's experts unquantized: such a projection has nothing to fold.""" + experts = _quantized_experts(experts_cls, quant_cfg) + first = experts._first_proj_attr + for proj in {"first": [first], "down_proj": ["down_proj"], "both": [first, "down_proj"]}[ + unquantized + ]: + _disable_weight_quantizers(experts, proj, range(NUM_EXPERTS)) + want = _per_expert_export(experts) + + _assert_byte_identical(split_fused_experts_state_dict(_keep_fused_export(experts)), want) + + +@LAYOUTS +def test_a_projection_mixing_quantized_and_unquantized_experts_is_refused_untouched(experts_cls): + """One fused tensor cannot hold both, and the refusal comes before the other projection folds.""" + experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) + _disable_weight_quantizers(experts, "down_proj", [1]) + before = {k: v.clone() for k, v in experts.state_dict().items()} + + with pytest.raises(NotImplementedError, match=r"down_proj: weight quantizers \[1\] of 4"): + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + _assert_byte_identical(experts.state_dict(), before) + assert experts.down_proj_weight_quantizers[0].is_enabled + + # ------------------------------------------------------------------ split_fused_experts_state_dict From 87a638611f89a16b6562aebaf7699421b978f42f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 17:49:24 +0000 Subject: [PATCH 07/67] Trim the keep-fused fold's tests: one export table, real two-rank shards The per-expert export comparison, the uncalibrated-expert test and the 18-case unquantized-projection test are one table now: the six layout x format cases, an uncalibrated expert, and each projection left unquantized. Byte identity with _export_fused_experts subsumes the dropped layout test's keys, shapes and dtypes, the uncalibrated test's warning text and its gate/up weight_scale_2 check. The gate/up cut and the ungated up_proj split tests go: the table fails on each regression they guarded. The one-rank sharded test and the faked uneven shard fold into the two-rank job, which now runs all six cases with the placement, global shape and local bytes checks, plus three experts over two ranks. The two NVFP4 zero-block tests are one parametrized test beside main's scale-clamping tests, and the sharded-scale test folds into the pack test. The conftest's gloo_world_of_one yields the one-rank mesh, so the local copy goes. Signed-off-by: Shengliang Xu --- tests/unit/torch/conftest.py | 13 + .../export/test_fused_experts_keep_fused.py | 376 ++++++------------ .../torch/quantization/test_nvfp4_tensor.py | 26 ++ 3 files changed, 159 insertions(+), 256 deletions(-) diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index 852412d2565..ec11dd90999 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -16,6 +16,7 @@ from collections import defaultdict import pytest +import torch from modelopt.torch.opt.mode import _ModeRegistryCls @@ -52,3 +53,15 @@ def use_channel_div_4(): for name, kv_for_all in lookup.items(): for regex, (k, v) in kv_for_all.items(): config_class.model_fields[name].default[regex][k] = v + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group, so the code under test takes its distributed path, and its + CPU device mesh.""" + from torch.distributed.device_mesh import init_device_mesh + + 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,)) + torch.distributed.destroy_process_group() diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 129cd253318..7c01074921f 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -27,12 +27,10 @@ import contextlib import copy import math -import types import warnings import pytest import torch -import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict @@ -66,12 +64,6 @@ FP8_WEIGHT_ONLY_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG) FP8_WEIGHT_ONLY_CFG["quant_cfg"].append({"quantizer_name": "*input_quantizer", "enable": False}) -QUANT_CFGS = pytest.mark.parametrize( - "quant_cfg", - [mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG], - ids=["fp8", "fp8_weight_only", "nvfp4"], -) - class _KeepFusedGatedExperts(nn.Module): """Fused gate+up experts in the transformers 5.x layout (Qwen3-MoE / Mixtral).""" @@ -135,15 +127,6 @@ def _unregister_synthetic_experts(): QuantModuleRegistry.unregister(cls) -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group and its CPU device mesh.""" - store = tmp_path_factory.mktemp("pg") / "store" - dist.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield init_device_mesh("cpu", (1,)) - dist.destroy_process_group() - - def _quantized_experts(experts_cls, quant_cfg, num_experts=NUM_EXPERTS): """Build, quantize and calibrate a fused-experts module; every expert sees one token.""" torch.manual_seed(0) @@ -200,98 +183,85 @@ def _assert_quantized(state_dict): # ------------------------------------------------------------------ keep-fused + split -@LAYOUTS -@QUANT_CFGS -def test_keep_fused_then_split_writes_exactly_the_per_expert_export(experts_cls, quant_cfg): - experts = _quantized_experts(experts_cls, quant_cfg) - want = _per_expert_export(experts) - _assert_quantized(want) +def _uncalibrated(experts, projections): + """Expert 1's calibration never ran: zero its weight amax.""" + for proj in projections: + getattr(experts, f"{proj}_weight_quantizers")[1]._amax.zero_() - _assert_byte_identical(split_fused_experts_state_dict(_keep_fused_export(experts)), want) +def _unquantized(experts, projections): + """A recipe left these projections' experts unquantized.""" + for proj in projections: + for quantizer in getattr(experts, f"{proj}_weight_quantizers"): + quantizer.disable() -@LAYOUTS -@QUANT_CFGS -def test_keep_fused_leaves_one_fused_weight_and_stacked_scales_per_projection( - experts_cls, quant_cfg -): - experts = _quantized_experts(experts_cls, quant_cfg) - first = experts._first_proj_attr - nvfp4 = quant_cfg is mtq.NVFP4_DEFAULT_CFG - has_input_scale = quant_cfg is not FP8_WEIGHT_ONLY_CFG - - _export_fused_experts_keep_fused(experts, torch.bfloat16) - - suffixes = {"", "_weight_scale"} - suffixes |= {"_weight_scale_2"} if nvfp4 else set() - suffixes |= {"_input_scale"} if has_input_scale else set() - assert set(experts.state_dict()) == {f"{p}{s}" for p in (first, "down_proj") for s in suffixes} - for proj in (first, "down_proj"): - weight = getattr(experts, proj) - assert weight.dim() == 3 and weight.shape[0] == NUM_EXPERTS - assert weight.dtype == (torch.uint8 if nvfp4 else torch.float8_e4m3fn) - assert getattr(experts, f"{proj}_weight_scale").shape[0] == NUM_EXPERTS - if nvfp4: - assert getattr(experts, f"{proj}_weight_scale_2").shape == (NUM_EXPERTS,) - if has_input_scale: - assert getattr(experts, f"{proj}_input_scale").dim() == 0 - # The quantizers' information now lives in the buffers above. - assert not [n for n, _ in experts.named_children() if "quantizer" in n] +def _export_cases(): + layouts = [(_KeepFusedGatedExperts, "gated"), (_KeepFusedUngatedExperts, "ungated")] + cfgs = [ + (mtq.FP8_DEFAULT_CFG, "fp8"), + (FP8_WEIGHT_ONLY_CFG, "fp8_weight_only"), + (mtq.NVFP4_DEFAULT_CFG, "nvfp4"), + ] + cases = [ + pytest.param(cls, cfg, None, (), id=f"{name}-{cfg_name}") + for cls, name in layouts + for cfg, cfg_name in cfgs + ] + cases += [ + pytest.param( + cls, + mtq.NVFP4_DEFAULT_CFG, + _uncalibrated, + ("first", "down_proj"), + id=f"{name}-uncalibrated", + ) + for cls, name in layouts + ] + cases += [ + pytest.param( + cls, mtq.FP8_DEFAULT_CFG, _unquantized, projs, id=f"{name}-unquantized-{which}" + ) + for cls, name in layouts + for which, projs in ( + ("first", ("first",)), + ("down", ("down_proj",)), + ("both", ("first", "down_proj")), + ) + ] + return cases -@LAYOUTS -def test_uncalibrated_expert_takes_the_same_weight_derived_amax_as_the_per_expert_export( - experts_cls, + +@pytest.mark.parametrize(("experts_cls", "quant_cfg", "change", "projections"), _export_cases()) +def test_keep_fused_then_split_writes_exactly_the_per_expert_export( + experts_cls, quant_cfg, change, projections ): - """An expert calibration never reached falls back to an amax read off its own weight. + """Also when an expert calibration never reached falls back to an amax read off its own weight. For the gated layout that amax must come from the whole fused ``gate_up`` slice, so gate and up - keep one shared ``weight_scale_2`` -- the per-expert export guarantees that, and vLLM's W1/W3 - fusion depends on it. + keep one shared ``weight_scale_2``, as the per-expert export guarantees and vLLM's W1/W3 fusion + needs. And when a recipe leaves a projection's experts unquantized: it has nothing to fold. """ - experts = _quantized_experts(experts_cls, mtq.NVFP4_DEFAULT_CFG) - first = experts._first_proj_attr - for proj in (first, "down_proj"): - getattr(experts, f"{proj}_weight_quantizers")[1]._amax.zero_() + experts = _quantized_experts(experts_cls, quant_cfg) + if change is not None: + first = experts._first_proj_attr + change(experts, [first if proj == "first" else proj for proj in projections]) want = _per_expert_export(experts) + if change is None: + _assert_quantized(want) - with pytest.warns(UserWarning, match=r"Expert 1 .* not calibrated"): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # an uncalibrated expert warns got = split_fused_experts_state_dict(_keep_fused_export(experts)) _assert_byte_identical(got, want) - if first == "gate_up_proj": - gate_ws2 = got[f"{PREFIX}.1.gate_proj.weight_scale_2"] - assert torch.equal(gate_ws2, got[f"{PREFIX}.1.up_proj.weight_scale_2"]) - - -def _disable_weight_quantizers(experts, proj, indices): - for i in indices: - getattr(experts, f"{proj}_weight_quantizers")[i].disable() - - -@LAYOUTS -@QUANT_CFGS -@pytest.mark.parametrize("unquantized", ["first", "down_proj", "both"]) -def test_a_projection_with_no_quantized_expert_is_written_as_the_per_expert_export_writes_it( - experts_cls, quant_cfg, unquantized -): - """A recipe can leave an MoE layer's experts unquantized: such a projection has nothing to fold.""" - experts = _quantized_experts(experts_cls, quant_cfg) - first = experts._first_proj_attr - for proj in {"first": [first], "down_proj": ["down_proj"], "both": [first, "down_proj"]}[ - unquantized - ]: - _disable_weight_quantizers(experts, proj, range(NUM_EXPERTS)) - want = _per_expert_export(experts) - - _assert_byte_identical(split_fused_experts_state_dict(_keep_fused_export(experts)), want) @LAYOUTS def test_a_projection_mixing_quantized_and_unquantized_experts_is_refused_untouched(experts_cls): """One fused tensor cannot hold both, and the refusal comes before the other projection folds.""" experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) - _disable_weight_quantizers(experts, "down_proj", [1]) + experts.down_proj_weight_quantizers[1].disable() before = {k: v.clone() for k, v in experts.state_dict().items()} with pytest.raises(NotImplementedError, match=r"down_proj: weight quantizers \[1\] of 4"): @@ -329,48 +299,6 @@ def _fused_state_dict(*, gated, rows=4, cols=6): } -def test_split_cuts_gate_and_up_rows_with_their_scales(): - sd = _fused_state_dict(gated=True) - fused = dict(sd) - untouched = torch.ones(3) - sd["model.layers.0.mlp.gate.weight"] = untouched - - out = split_fused_experts_state_dict(sd) - - gate_up, ws = fused[f"{PREFIX}.gate_up_proj"], fused[f"{PREFIX}.gate_up_proj_weight_scale"] - for e in range(LOCAL_EXPERTS): - base = f"{PREFIX}.{e}" - assert torch.equal(_raw_bytes(out[f"{base}.gate_proj.weight"]), _raw_bytes(gate_up[e, :4])) - assert torch.equal(_raw_bytes(out[f"{base}.up_proj.weight"]), _raw_bytes(gate_up[e, 4:])) - assert torch.equal(out[f"{base}.gate_proj.weight_scale"], ws[e, :4]) - assert torch.equal(out[f"{base}.up_proj.weight_scale"], ws[e, 4:]) - # gate and up are halves of one projection: they share its per-tensor second scale. - for proj in ("gate_proj", "up_proj"): - assert out[f"{base}.{proj}.weight_scale_2"].item() == 7.0 + e - assert out[f"{base}.{proj}.input_scale"].item() == 0.5 - down = fused[f"{PREFIX}.down_proj"][e] - assert torch.equal(_raw_bytes(out[f"{base}.down_proj.weight"]), _raw_bytes(down)) - assert out[f"{base}.down_proj.weight_scale"].item() == 3.0 + e - assert out[f"{base}.down_proj.input_scale"].item() == 0.25 - assert not any(k in out for k in fused), "fused entries must be consumed" - assert out["model.layers.0.mlp.gate.weight"] is untouched - - -def test_split_keeps_an_ungated_up_proj_whole(): - sd = _fused_state_dict(gated=False) - up = sd[f"{PREFIX}.up_proj"] - - out = split_fused_experts_state_dict(sd) - - assert not any(".gate_proj." in k for k in out) - for e in range(LOCAL_EXPERTS): - assert torch.equal(_raw_bytes(out[f"{PREFIX}.{e}.up_proj.weight"]), _raw_bytes(up[e])) - assert torch.equal( - out[f"{PREFIX}.{e}.up_proj.weight_scale"], sd[f"{PREFIX}.up_proj_weight_scale"][e] - ) - assert f"{PREFIX}.up_proj" not in out - - def test_split_gives_every_tensor_its_own_storage(tmp_path): """Slices alias the fused parent, and gate/up share one scale per expert. @@ -458,73 +386,39 @@ def _shard_fused_weights(experts, mesh): return experts -@LAYOUTS -@QUANT_CFGS -def test_keep_fused_on_a_sharded_weight_stays_sharded_and_splits_like_the_per_expert_export( - gloo_world_of_one, experts_cls, quant_cfg -): - experts = _quantized_experts(experts_cls, quant_cfg) - plain = _keep_fused_export(experts) - sharded = _shard_fused_weights(copy.deepcopy(experts), gloo_world_of_one) - - _export_fused_experts_keep_fused(sharded, torch.bfloat16) - - fused = _prefixed(sharded.state_dict()) - assert set(fused) == set(plain) - for key, value in fused.items(): - if key.endswith("_input_scale"): # one scalar shared by every expert - assert not isinstance(value, DTensor), key - else: # weights and per-expert scales stay sharded on the expert axis - assert isinstance(value, DTensor), key - assert value.placements == (Shard(0),), key - assert value.shape == plain[key].shape, key - value = value.to_local() - assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key - assert _dtensor_dim0_offset(fused[f"{PREFIX}.{sharded._first_proj_attr}"]) == 0 - - split = _split_local_fused_module(PREFIX, fused) - - assert not any(isinstance(v, DTensor) for v in split.values()) - _assert_byte_identical(split, _per_expert_export(experts)) - - -def test_keep_fused_keeps_the_expert_count_of_an_uneven_shard(gloo_world_of_one): - """Two experts held of three is what rank 0 holds when three experts split over two ranks.""" - experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=3) - for proj in ("gate_up_proj", "down_proj"): - full = getattr(experts, proj).data - held = DTensor.from_local( - full[:2].clone(), - gloo_world_of_one, - [Shard(0)], - run_check=False, - shape=full.shape, - stride=full.stride(), - ) - setattr(experts, proj, nn.Parameter(held, requires_grad=False)) - - _export_fused_experts_keep_fused(experts, torch.bfloat16) - - assert experts.gate_up_proj.shape[0] == 3 - assert experts.gate_up_proj_weight_scale.shape[0] == 3 - - def _two_rank_split_job(rank, world_size): mesh = init_device_mesh("cpu", (world_size,)) - for experts_cls, quant_cfg in [ - (_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG), - (_KeepFusedUngatedExperts, mtq.NVFP4_DEFAULT_CFG), - ]: + cases = [ + (experts_cls, quant_cfg, NUM_EXPERTS) + for experts_cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) + for quant_cfg in (mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG) + ] + # Three experts over two ranks: rank 0 holds two, rank 1 one. The folded tensors must keep the + # global expert count, not infer local x world rows. + cases.append((_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG, 3)) + for experts_cls, quant_cfg, num_experts in cases: # Every rank calibrates the same seeded module, so all ranks hold identical quantizers. - experts = _quantized_experts(experts_cls, quant_cfg) + experts = _quantized_experts(experts_cls, quant_cfg, num_experts) want = _per_expert_export(experts) + plain = _keep_fused_export(experts) sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) _export_fused_experts_keep_fused(sharded, torch.bfloat16) fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) first = fused[f"{PREFIX}.{sharded._first_proj_attr}"] n_local = first.to_local().shape[0] - offset = rank * math.ceil(NUM_EXPERTS / world_size) + offset = rank * math.ceil(num_experts / world_size) assert _dtensor_dim0_offset(first) == offset + for key, value in fused.items(): + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + else: # weights and per-expert scales stay sharded on the expert axis + assert isinstance(value, DTensor), key + assert value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + local = plain[key][offset : offset + n_local] + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(local)), key mine = { k: v for k, v in want.items() @@ -536,83 +430,19 @@ def _two_rank_split_job(rank, world_size): # A per-expert scale held as a plain tensor over ALL experts (as keep-fused emits it for an # unsharded weight) must be cut to this rank's experts before pairing. - unsharded = _keep_fused_export(experts) plain_scales = { - k: unsharded[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v + k: plain[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v for k, v in fused.items() } _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): - """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index.""" + """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index, and the + fused tensors stay sharded on the expert axis at their global shape, an uneven shard included.""" cpu_dist_workers.run(_two_rank_split_job) -# ------------------------------------------------------------------ NVFP4 on sharded weights - - -def test_nvfp4_block_scale_maps_an_all_zero_block_to_one(): - weight = torch.randn(2, 32) - weight[0, :16] = 0.0 # an all-zero block - weight[1, :16] = 0.0 - weight[1, 0] = 6.0 # amax 6 at a second-level scale of 1 -> block scale exactly 1 - - scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) - - assert scale.float()[0, 0].item() == 1.0 - assert scale.float()[1, 0].item() == 1.0 - - -def test_nvfp4_static_block_scale_maps_a_zero_amax_block_to_one(): - # Two rows of two blocks: a zero-amax block next to ones whose amax maps to a scale of 1. - quantizer = types.SimpleNamespace( - block_sizes={-1: 16}, - global_amax=torch.tensor(6.0 * 448.0), - _amax=torch.tensor([0.0, 6.0, 6.0, 3.0]), - ) - - scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(quantizer, torch.randn(2, 32)) - - assert scale.float()[0, 0] == scale.float()[0, 1] == scale.float()[1, 0] - assert scale.float()[1, 1] != scale.float()[1, 0] - - -def test_nvfp4_block_scale_accepts_a_sharded_weight(gloo_world_of_one): - weight = torch.randn(4, 32) - weight[1, :16] = 0.0 - sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) - - want, want_2 = NVFP4QTensor.get_weights_scaling_factor(weight, 16) - got, got_2 = NVFP4QTensor.get_weights_scaling_factor(sharded, 16) - - assert isinstance(got, DTensor) and got.placements == (Shard(0),) - assert torch.equal(_raw_bytes(got.to_local()), _raw_bytes(want)) - assert torch.equal(got_2, want_2) - - -def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): - """``from_local`` would infer local x world rows; an uneven shard must keep the real total. - - Two rows held of a three-row tensor is what one rank holds when three rows split over two. - """ - local = torch.randn(2, 32) * 4 - sharded = DTensor.from_local( - local.clone(), - gloo_world_of_one, - [Shard(0)], - run_check=False, - shape=torch.Size([3, 32]), - stride=(32, 1), - ) - - packed = NVFP4QTensor._cast_fp4(sharded) - - assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) - assert packed.shape == (3, 32) - assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) - - # ------------------------------------------------------------------ export handlers @@ -697,9 +527,14 @@ def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_ def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): - """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator.""" + """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. + + The block scales come out sharded like the weight, an all-zero block included: filling it by + boolean-mask assignment lowers to ``index_put_``, which DTensor cannot dispatch. + """ torch.manual_seed(0) weight = torch.randn(4, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 # an all-zero block scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data @@ -708,4 +543,33 @@ def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + assert isinstance(sharded_scale, DTensor) and sharded_scale.placements == (Shard(0),) + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale)) assert torch.equal(packed.full_tensor(), want) + # The second-level scale read off the sharded weight itself, as without one handed in. + assert torch.equal( + NVFP4QTensor.get_weights_scaling_factor(sharded, 16)[1], + NVFP4QTensor.get_weights_scaling_factor(weight, 16)[1], + ) + + +def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): + """``from_local`` would infer local x world rows; an uneven shard must keep the real total. + + Two rows held of a three-row tensor is what one rank holds when three rows split over two. + """ + local = torch.randn(2, 32) * 4 + sharded = DTensor.from_local( + local.clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 32]), + stride=(32, 1), + ) + + packed = NVFP4QTensor._cast_fp4(sharded) + + assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) + assert packed.shape == (3, 32) + assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index a7b64515e3a..ce1159711ba 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -106,6 +106,32 @@ def test_helper_clamps_underflow_to_min(self): out = _cast_per_block_scale_to_fp8(tiny).float() assert (out > 0).all(), f"FP8 cast produced zero scales: {out.tolist()}" + @staticmethod + def _dynamic_block_scales(): + weight = torch.zeros(1, 48) + weight[0, 16] = 6.0 # amax 6 at a second-level scale of 1: block scale exactly 1 + weight[0, 32] = 3.0 # amax 3: block scale 0.5 + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) + return scale.float()[0] + + @staticmethod + def _static_block_scales(): + q = SimpleNamespace( + block_sizes={-1: 16}, + global_amax=torch.tensor(6.0 * 448.0), # a second-level scale of 1 + _amax=torch.tensor([0.0, 6.0, 3.0]), + ) + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(q, torch.randn(1, 48)) + return scale.float()[0] + + @pytest.mark.parametrize("path", ["dynamic", "static"]) + def test_an_all_zero_block_gets_a_scale_of_one(self, path): + """Blocks of amax 0, 6 and 3: the zero block gets 1, as the amax-6 block does, not the clamp.""" + zero, unit, half = getattr(self, f"_{path}_block_scales")().tolist() + + assert zero == unit == 1.0 + assert half == 0.5 + def test_static_path_no_nan_when_block_amax_zero(self): """Static path: zero-amax block + small global_amax must clamp to 448, not cast to NaN.""" block_size = 16 From ec602efbf7ceb1d4a0496aeeb354830958b1a604 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:50:10 +0000 Subject: [PATCH 08/67] Import the keep-fused fold test's export handlers at module scope The FSDP2 weight-window test imported modelopt.torch.export.hf_export_handlers inside the test, though it is neither optional nor circular: the export package already imports it. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_fused_experts_keep_fused.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 7c01074921f..2dc9347cd78 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -42,6 +42,7 @@ pytest.importorskip("safetensors") from safetensors.torch import load_file, save_file +import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( _dtensor_dim0_offset, @@ -473,8 +474,6 @@ def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_pa ): """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a weight-update window, they would be overwritten with its original shards at the next reshard.""" - import modelopt.torch.export.hf_export_handlers as handlers - events = [] @contextlib.contextmanager From f8657d573e273d83a44a93a8c9840bbf7356a3f7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 23:32:32 +0000 Subject: [PATCH 09/67] Fall back per quantizer stage in the keep-fused fold, as the split export does Main's _export_fused_experts (#2807) derives an uncalibrated expert's amax per TensorQuantizer stage and in the shape each stage expects: W4A8 AWQ's SequentialQuantizer has no amax of its own, and its INT4 block stage rejects a scalar. The keep-fused fold, which must match that export byte for byte, still set one scalar amax on the whole quantizer. Use the same _tensor_quantizers, _is_uncalibrated and _fallback_amax helpers; a per-tensor stage gets the same abs-max value as before. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index d16697cf659..8c1ae9eca71 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -211,16 +211,11 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N w_slice = local_w[local_idx] # Uncalibrated-expert fallback: derive amax from this expert's local weight slice # (matches _export_fused_experts), so a never-routed expert still exports sane scales. - if ( - hasattr(w_quantizer, "is_enabled") - and w_quantizer.is_enabled - and ( - not hasattr(w_quantizer, "_amax") - or w_quantizer._amax is None - or torch.all(w_quantizer._amax == 0) - ) - ): - w_quantizer.amax = w_slice.abs().amax().to(torch.float32) + # Per stage, in the shape each expects (W4A8's SequentialQuantizer, block-wise INT4). + uncalibrated = [q for q in _tensor_quantizers(w_quantizer) if _is_uncalibrated(q)] + for q in uncalibrated: + q.amax = _fallback_amax(w_slice, q) + if uncalibrated: warnings.warn( f"Expert {global_idx} {proj} weight quantizer was not calibrated (amax missing " f"or zero); using weight-derived amax. Increase calibration size to activate all " From eafde05bed99c252382234ab23cbe165cee815a7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 01:08:24 +0000 Subject: [PATCH 10/67] Test the keep-fused fold's per-stage amax fallback with a W4A8 expert The keep-fused fold derives an uncalibrated expert's amax per TensorQuantizer stage, through main's _tensor_quantizers / _is_uncalibrated / _fallback_amax, but no test held a W4A8 SequentialQuantizer to it. Both cases fail with the fold's old single scalar amax. Ported from the epdp stack (bfca073416). Signed-off-by: Shengliang Xu --- .../export/test_fused_experts_keep_fused.py | 42 ++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 2dc9347cd78..2e8c1aa4603 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -53,7 +53,8 @@ split_fused_experts_state_dict, ) from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry -from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import NVFP4QTensor from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan @@ -258,6 +259,45 @@ def test_keep_fused_then_split_writes_exactly_the_per_expert_export( _assert_byte_identical(got, want) +@pytest.mark.parametrize("amax_state", ["calibrated", "zero"]) +def test_the_fold_derives_a_w4a8_experts_amax_per_quantizer_stage(monkeypatch, amax_state): + """W4A8's SequentialQuantizer has no amax of its own, and its INT4 block stage rejects a scalar: + an expert calibration never reached gets a per-block amax there and a scalar on the FP8 stage.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + block = 16 + rows = {"gate_up_proj": 2 * INTERMEDIATE_DIM, "down_proj": HIDDEN_DIM} + blocks = {"gate_up_proj": HIDDEN_DIM // block, "down_proj": INTERMEDIATE_DIM // block} + for proj in rows: + for i in range(NUM_EXPERTS): + int4 = TensorQuantizer( + QuantizerAttributeConfig(num_bits=4, block_sizes={-1: block, "type": "static"}) + ) + fp8 = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + fill = 0.0 if amax_state == "zero" else 1.0 + int4.amax = torch.full( + (rows[proj] * blocks[proj], 1), fill + ) # as block calibration records + fp8.amax = torch.tensor(fill) + getattr(experts, f"{proj}_weight_quantizers")[i] = SequentialQuantizer(int4, fp8) + + seen = [] + + def spy(wrapper, dtype, **_kwargs): + int4, fp8 = wrapper.weight_quantizer + seen.append((int4._amax.clone(), fp8._amax.clone())) + wrapper.weight_scale = torch.ones(1) + + monkeypatch.setattr("modelopt.torch.export.unified_export_hf._export_quantized_weight", spy) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + assert len(seen) == 2 * NUM_EXPERTS + assert {amax.shape for amax, _ in seen} == {(64, 1), (32, 1)} + assert all(torch.all(int4_amax > 0) and torch.all(fp8_amax > 0) for int4_amax, fp8_amax in seen) + assert any("was not calibrated" in str(w.message) for w in caught) == (amax_state == "zero") + + @LAYOUTS def test_a_projection_mixing_quantized_and_unquantized_experts_is_refused_untouched(experts_cls): """One fused tensor cannot hold both, and the refusal comes before the other projection folds.""" From 703f299bd47077378e17cebea04b4cc4390e4cce Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 19:30:24 +0000 Subject: [PATCH 11/67] Drop the source checkpoint's expert scales in the keep-fused fold too The per-expert export removes a fused-experts module's source block and activation scales (_delete_fused_moe_source_attrs) once it has replaced the weights they described. The keep-fused fold removed only the quantizers, so a native-FP8 module would export its stale *_scale_inv and *_activation_scale beside the new scales, which the per-expert export does not write. No load hands the fold such a module today (the expert-parallel load dequantizes block-FP8, the FSDP2 load refuses it), but the fold now removes the same attributes. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 8 +++++++- .../export/test_fused_experts_keep_fused.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 8c1ae9eca71..54e8ed4e5f9 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -255,12 +255,18 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if input_scale is not None: module.register_buffer(f"{proj}_input_scale", input_scale) - # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers. + # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers -- and + # the source checkpoint's scales, which described the weights the fold replaced, as the + # per-expert export does (_delete_fused_moe_source_attrs). for attr in ( f"{first}_weight_quantizers", f"{first}_input_quantizer", "down_proj_weight_quantizers", "down_proj_input_quantizer", + f"{first}_scale_inv", + "down_proj_scale_inv", + f"{first}_activation_scale", + "down_proj_activation_scale", ): if hasattr(module, attr): delattr(module, attr) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 2e8c1aa4603..531fd83fe68 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -198,6 +198,13 @@ def _unquantized(experts, projections): quantizer.disable() +def _with_source_scales(experts, projections): + """A native-FP8 checkpoint's experts keep their source block and activation scales.""" + for proj in projections: + experts.register_buffer(f"{proj}_scale_inv", torch.ones(NUM_EXPERTS, 1, 1)) + experts.register_buffer(f"{proj}_activation_scale", torch.ones(())) + + def _export_cases(): layouts = [(_KeepFusedGatedExperts, "gated"), (_KeepFusedUngatedExperts, "ungated")] cfgs = [ @@ -220,6 +227,16 @@ def _export_cases(): ) for cls, name in layouts ] + cases += [ + pytest.param( + cls, + mtq.FP8_DEFAULT_CFG, + _with_source_scales, + ("first", "down_proj"), + id=f"{name}-source-scales", + ) + for cls, name in layouts + ] cases += [ pytest.param( cls, mtq.FP8_DEFAULT_CFG, _unquantized, projs, id=f"{name}-unquantized-{which}" From 158cc7fba49ea61d432cd5e12471e9b0d034e790 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:47:38 +0000 Subject: [PATCH 12/67] Split block and packed expert scales by their rows, and refuse a rank with no experts split_fused_experts_state_dict cut gate's and up's weight scales at the weight's row split, which holds only when a scale has one row per weight row. An INT4 weight packs two rows a byte while its scales stay per row, and an FP8 2-D block scale covers 128 rows, so for INT4 blockwise, W4A8 AWQ and FP8 2-D blockwise experts gate and up got the wrong scales. Each now takes the scale rows in proportion to the weight rows it takes; the per-expert export comparison covers the three formats. The keep-fused fold refuses, on every rank, a layout that leaves a rank no experts: that rank would stack nothing and register no input scale, so the ranks' state dicts would differ. It reads the shard's expert offset through _dtensor_dim0_offset, as the split does. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 50 +++++++++---------- .../export/test_fused_experts_keep_fused.py | 32 +++++++++++- 2 files changed, 55 insertions(+), 27 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 54e8ed4e5f9..1d6f3ea6bce 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -142,10 +142,6 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N from torch.distributed.tensor import DTensor, Replicate, Shard except Exception: # pragma: no cover - non-distributed install DTensor = None # noqa: N806 (the optional-import fallback for a class) - try: - from torch.distributed.tensor._utils import compute_local_shape_and_global_offset - except Exception: # pragma: no cover - older torch: fall back to even-split offset - compute_local_shape_and_global_offset = None first = getattr(module, "_first_proj_attr", "gate_up_proj") n = module.num_experts @@ -177,23 +173,20 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N is_dt = DTensor is not None and isinstance(wdata, DTensor) if is_dt: mesh, placements = wdata.device_mesh, wdata.placements + # A rank holding no experts would stack none and register no input scale, so the + # ranks' state dicts would differ for the DCP write. Every rank chunks the experts as + # torch does, so every rank refuses together. + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + chunk = -(-wdata.shape[0] // mesh.size(mdim)) + if not chunk or -(-wdata.shape[0] // chunk) < mesh.size(mdim): + raise NotImplementedError( + f"{type(module).__name__}.{proj}: {wdata.shape[0]} experts over " + f"{mesh.size(mdim)} shards leaves a rank none; the keep-fused fold " + "needs every rank to hold an expert." + ) local_w = wdata.to_local() - offset = None - if compute_local_shape_and_global_offset is not None: - try: - _, global_offset = compute_local_shape_and_global_offset( - tuple(wdata.shape), mesh, placements - ) - offset = int(global_offset[0]) - except Exception: - offset = None - if offset is None: # even-split fallback (dim-0 Shard on a 1-D mesh) - offset, nshards = 0, 1 - for mdim, p in enumerate(placements): - if isinstance(p, Shard) and p.dim == 0: - nshards = mesh.size(mdim) - offset = mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) - break + offset = _dtensor_dim0_offset(wdata) scale_placements = [ p if (isinstance(p, Shard) and p.dim == 0) else Replicate() for p in placements ] @@ -572,14 +565,19 @@ def _scales(proj_name): first_ws, first_in, first_ws2 = _scales(first) down_ws, down_in, down_ws2 = _scales("down_proj") - def _slice_weight_scale(ws, row_slice, fused_rows): - # per-tensor-per-expert scale -> scalar; per-output-channel -> slice the rows. + def _slice_weight_scale(ws, expert_idx, row_slice, rows): + # A per-expert scalar as is; else the expert's scale rows covering ``row_slice`` of its + # ``rows`` weight rows, in proportion: an INT4 weight packs two rows a byte while its + # scales stay per row, and an FP8 2-D block scale covers 128 rows. if ws is None: return None if ws.dim() <= 1: # (E,) one scalar per expert return ws[expert_idx] - sub = ws[expert_idx] # (fused_rows,) or (fused_rows, ...) - return sub if row_slice is None else sub[row_slice] + sub = ws[expert_idx] + if row_slice is None: + return sub + start, stop, _ = row_slice.indices(rows) + return sub[start * sub.shape[0] // rows : stop * sub.shape[0] // rows] for expert_idx in range(n_experts): if gated: @@ -617,12 +615,12 @@ def _slice_weight_scale(ws, row_slice, fused_rows): # safetensors/save_pretrained shared-tensor dedup silently drops the duplicate # (e.g. up_proj.weight_scale/input_scale would go missing). out[f"{base}.weight"] = weight.detach().clone().contiguous() - sliced_ws = _slice_weight_scale(ws, row_slice, weight.shape[0]) + sliced_ws = _slice_weight_scale(ws, expert_idx, row_slice, first_w.shape[1]) if sliced_ws is not None: out[f"{base}.weight_scale"] = sliced_ws.detach().clone().contiguous() # weight_scale_2 is a per-expert scalar (shared by gate/up): _slice_weight_scale's # dim<=1 branch returns ws2[expert_idx] regardless of row_slice. NVFP4 only; None for FP8. - sliced_ws2 = _slice_weight_scale(ws2, row_slice, weight.shape[0]) + sliced_ws2 = _slice_weight_scale(ws2, expert_idx, row_slice, first_w.shape[1]) if sliced_ws2 is not None: out[f"{base}.weight_scale_2"] = sliced_ws2.detach().clone().contiguous() if ins is not None: diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 531fd83fe68..e3f7d23efb2 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -276,6 +276,29 @@ def test_keep_fused_then_split_writes_exactly_the_per_expert_export( _assert_byte_identical(got, want) +@pytest.mark.parametrize( + "quant_cfg", + [ + mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG, + mtq.W4A8_AWQ_BETA_CFG, + mtq.FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG, + ], + ids=["int4-blockwise", "w4a8-awq", "fp8-2d-blockwise"], +) +def test_the_split_cuts_packed_and_block_scales_with_their_rows(monkeypatch, quant_cfg): + """An INT4 weight packs two rows a byte while its scales stay per row, and an FP8 2-D block + scale covers 128 rows: gate and up each take the scale rows of their own weight rows. Their + blocks need dimensions of whole 128-blocks.""" + monkeypatch.setitem(globals(), "HIDDEN_DIM", 128) + monkeypatch.setitem(globals(), "INTERMEDIATE_DIM", 128) + experts = _quantized_experts(_KeepFusedGatedExperts, quant_cfg) + want = _per_expert_export(experts) + + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + + _assert_byte_identical(got, want) + + @pytest.mark.parametrize("amax_state", ["calibrated", "zero"]) def test_the_fold_derives_a_w4a8_experts_amax_per_quantizer_stage(monkeypatch, amax_state): """W4A8's SequentialQuantizer has no amax of its own, and its INT4 block stage rejects a scalar: @@ -494,10 +517,17 @@ def _two_rank_split_job(rank, world_size): } _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) + # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone + # failing while rank 0 goes on to write. + lone = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=1) + with pytest.raises(NotImplementedError, match="leaves a rank none"): + _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) + def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index, and the - fused tensors stay sharded on the expert axis at their global shape, an uneven shard included.""" + fused tensors stay sharded on the expert axis at their global shape, an uneven shard included. + A layout that leaves a rank no expert is refused on every rank.""" cpu_dist_workers.run(_two_rank_split_job) From 8502776d351edf785b5fb15f6b143d50321b4beb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:41:05 +0000 Subject: [PATCH 13/67] Pair shard-local expert quantizers and refuse what the fold can't write Under expert parallelism a fused-experts module keeps only its own experts' weight quantizers, in shard order, while its fused weight's DTensor spans every expert. The keep-fused fold indexed them by global expert, so every rank but the first would raise IndexError. It now pairs each local expert with the quantizer at its position in the shard when the list is shorter than the global extent, and refuses on every rank a quantizer count that is neither every expert's nor the shard's. The fold also refuses native FP8 expert weights, which only the scale_inv it drops dequantizes, and the split refuses a block scale whose rows cannot be cut at the gate|up boundary (a 128-row block across gpt-oss's 2880-row gate, say) instead of flooring the proportion. Every refusal of the fold runs before either projection is folded. The FSDP2 weight window found its unit by scanning the model's modules once per exported module and rebuilt the model's name maps each time. It now walks up the module's name through maps built once per export (ExportContext.module_names) and opens on the innermost enclosing unit. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 33 +++---- modelopt/torch/export/moe_utils.py | 79 +++++++++++---- modelopt/torch/export/registry.py | 8 ++ .../export/test_fused_experts_keep_fused.py | 97 ++++++++++++++++--- 4 files changed, 164 insertions(+), 53 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 5a1a80b6bf2..4983eadfc3c 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -20,7 +20,8 @@ import torch.nn as nn -from modelopt.torch.quantization.utils import fsdp2_aware_weight_update +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps +from modelopt.torch.quantization.utils.core_utils import _get_enclosing_fsdp_module from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax @@ -37,28 +38,24 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") -def _fsdp2_weight_window(model: nn.Module, module: nn.Module): +def _fsdp2_weight_window(ctx: ExportContext, module: nn.Module): """``fsdp2_aware_weight_update`` on the FSDP2 unit that owns ``module``'s weights. - That is ``model`` when the model itself is wrapped. Expert parallelism FSDP2-shards only the dense + That is the model when the model itself is wrapped. Expert parallelism FSDP2-shards only the dense layers and leaves the root unwrapped, and a window opened on an unwrapped root does nothing: the weight stays a sharded DTensor, which packing mishandles (``DTensor << int`` is a silent no-op in - torch 2.11, so NVFP4 packed unrelated codes next to correct scales). So open it on the enclosing - unit instead. + torch 2.11, so NVFP4 packed unrelated codes next to correct scales). So open it on the innermost + enclosing unit instead, found by walking up the module's name: the export calls this per module. """ from torch.distributed.fsdp import FSDPModule - root = model - if not isinstance(model, FSDPModule): - root = next( - ( - unit - for unit in model.modules() - if isinstance(unit, FSDPModule) and any(m is module for m in unit.modules()) - ), - model, - ) - return fsdp2_aware_weight_update(root, module) + if ctx.module_names is None: + ctx.module_names = module_name_maps(ctx.model) + if isinstance(ctx.model, FSDPModule): + return fsdp2_aware_weight_update(ctx.model, module, names=ctx.module_names) + unit = _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + # Names relative to the unit, which the window resolves the module's parameters against. + return fsdp2_aware_weight_update(unit, module) def _export_weight( @@ -77,7 +74,7 @@ def _export_weight( # reshard on the way out: a layer left unsharded holds all of its weights on every rank, and # FSDP2 later puts its own shards back over anything packed in it outside a window. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: - with _fsdp2_weight_window(ctx.model, module): + with _fsdp2_weight_window(ctx, module): _export_quantized_weight(module, ctx.dtype, weight_name) return @@ -186,7 +183,7 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its # original shards the next time the layer reshards, and exported unquantized. - with _fsdp2_weight_window(ctx.model, module): + with _fsdp2_weight_window(ctx, module): _export_fused_experts_keep_fused(module, ctx.dtype) elif ctx.keep_fused_experts: # No-gather distributed export: fold each rank's LOCAL experts and leave the weight diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 1d6f3ea6bce..a05ffb1ce12 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -119,6 +119,36 @@ def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int) return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) +def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: + """Refuse a sharded fused expert weight the keep-fused fold cannot fold, on every rank alike. + + Reads only the global shape, the mesh and the quantizer count, which every rank agrees on. + """ + from torch.distributed.tensor import Shard + + name = f"{type(module).__name__}.{proj}" + shards = 1 + for mdim, placement in enumerate(weight.placements): + if isinstance(placement, Shard) and placement.dim == 0: + size = weight.device_mesh.size(mdim) + # A rank holding no experts would stack none and register no input scale, so the ranks' + # state dicts would differ for the DCP write. The experts are chunked as torch does. + chunk = -(-weight.shape[0] // size) + if not chunk or -(-weight.shape[0] // chunk) < size: + raise NotImplementedError( + f"{name}: {weight.shape[0]} experts over {size} shards leaves a rank none; the " + "keep-fused fold needs every rank to hold an expert." + ) + shards *= size + # FSDP2 keeps every expert's quantizer on every rank; expert parallelism keeps only this rank's, + # in the order of its shard, as _QuantFusedExperts._local_experts reads them. + if num_quantizers != weight.shape[0] and num_quantizers * shards != weight.shape[0]: + raise NotImplementedError( + f"{name}: {num_quantizers} weight quantizers for {weight.shape[0]} experts in {shards} " + "shards; the keep-fused fold needs one per expert, or one per expert of this rank's shard." + ) + + def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> None: """Fold per-expert quantizers into the FUSED 3-D expert weights *in place*, staying sharded. @@ -127,8 +157,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N fused weight (via the export handler's ``unshard``) and then materializes every expert on every rank, putting the whole model on each GPU (OOMs at 235B/480B). This keeps each projection a single fused param: on a ``Shard(0)`` DTensor it quantizes ONLY this rank's local experts - (``to_local()`` + the DTensor's dim-0 global offset for the matching per-expert weight quantizer) - and rewraps the result ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. + (``to_local()``, each with its own weight quantizer: by global index where every rank keeps every + expert's, by position in the shard where a rank keeps only its own) and rewraps the result + ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` @@ -147,12 +178,22 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N n = module.num_experts # One fused tensor cannot hold quantized and unquantized experts side by side, so a projection # is folded only when all its experts are quantized; one with none quantized stays as it is, - # which is what the per-expert export writes for it. Both are decided before either is folded, - # so a refusal leaves the module untouched. + # which is what the per-expert export writes for it. Every refusal comes before either is + # folded, so it leaves the module untouched. folded = [] for proj in (first, "down_proj"): - enabled = [q.is_enabled for q in getattr(module, f"{proj}_weight_quantizers")] + weight_quantizers = getattr(module, f"{proj}_weight_quantizers") + enabled = [q.is_enabled for q in weight_quantizers] if all(enabled): + wdata = getattr(module, proj).data + if wdata.element_size() == 1: + # get_expert_weight dequantizes these with their scale_inv, which the fold drops. + raise NotImplementedError( + f"{type(module).__name__}.{proj}: native FP8 expert weights cannot be folded " + "as stored; dequantize them first (unpack_weight())." + ) + if DTensor is not None and isinstance(wdata, DTensor): + _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) folded.append(proj) elif any(enabled): raise NotImplementedError( @@ -171,26 +212,18 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N # the model and OOMing the GPU at scale. Plain tensor (single process): all experts. wdata = weight.data is_dt = DTensor is not None and isinstance(wdata, DTensor) + quantizers_are_global = True if is_dt: mesh, placements = wdata.device_mesh, wdata.placements - # A rank holding no experts would stack none and register no input scale, so the - # ranks' state dicts would differ for the DCP write. Every rank chunks the experts as - # torch does, so every rank refuses together. - for mdim, p in enumerate(placements): - if isinstance(p, Shard) and p.dim == 0: - chunk = -(-wdata.shape[0] // mesh.size(mdim)) - if not chunk or -(-wdata.shape[0] // chunk) < mesh.size(mdim): - raise NotImplementedError( - f"{type(module).__name__}.{proj}: {wdata.shape[0]} experts over " - f"{mesh.size(mdim)} shards leaves a rank none; the keep-fused fold " - "needs every rank to hold an expert." - ) local_w = wdata.to_local() offset = _dtensor_dim0_offset(wdata) scale_placements = [ p if (isinstance(p, Shard) and p.dim == 0) else Replicate() for p in placements ] local_indices = list(range(offset, offset + local_w.shape[0])) + # Every expert's quantizers (FSDP2), or only this rank's, in shard order (expert + # parallelism): _refuse_unfoldable_shards allows no other count. + quantizers_are_global = len(weight_quantizers) == wdata.shape[0] else: local_w = wdata local_indices = list(range(n)) @@ -200,7 +233,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N weight_scale_2s: list[torch.Tensor] = [] input_scale: torch.Tensor | None = None for local_idx, global_idx in enumerate(local_indices): - w_quantizer = weight_quantizers[global_idx] + w_quantizer = weight_quantizers[global_idx if quantizers_are_global else local_idx] w_slice = local_w[local_idx] # Uncalibrated-expert fallback: derive amax from this expert's local weight slice # (matches _export_fused_experts), so a never-routed expert still exports sane scales. @@ -577,7 +610,15 @@ def _slice_weight_scale(ws, expert_idx, row_slice, rows): if row_slice is None: return sub start, stop, _ = row_slice.indices(rows) - return sub[start * sub.shape[0] // rows : stop * sub.shape[0] // rows] + n = sub.shape[0] + # Exact only for scale rows spread evenly over the weight rows, with the cut between + # two of them; otherwise a block of rows sharing one scale straddles gate and up. + if (rows % n and n % rows) or start * n % rows or stop * n % rows: + raise NotImplementedError( + f"A {n}-row weight scale over {rows} weight rows cannot be split at the gate|up " + f"boundary, row {start or stop}: a block of rows sharing one scale crosses it." + ) + return sub[start * n // rows : stop * n // rows] for expert_idx in range(n_experts): if gated: diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index d0187c59321..6640b5ceb30 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,10 +28,14 @@ from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING import torch import torch.nn as nn +if TYPE_CHECKING: + from modelopt.torch.quantization.utils.core_utils import ModuleNames + __all__ = [ "ExportContext", "ExportHandler", @@ -64,6 +68,10 @@ class ExportContext: # happens post-consolidation on the sharded fused weight). See _export_fused_experts_keep_fused. keep_fused_experts: bool = False + module_names: "ModuleNames | None" = None + """``model``'s name maps, built once per export by the first handler that looks a module up + (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index e3f7d23efb2..8618950fc1d 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -338,14 +338,30 @@ def spy(wrapper, dtype, **_kwargs): assert any("was not calibrated" in str(w.message) for w in caught) == (amax_state == "zero") +def _mixed(experts): + experts.down_proj_weight_quantizers[1].disable() + + +def _native_fp8(experts): + experts.down_proj = nn.Parameter( + experts.down_proj.data.to(torch.float8_e4m3fn), requires_grad=False + ) + + @LAYOUTS -def test_a_projection_mixing_quantized_and_unquantized_experts_is_refused_untouched(experts_cls): - """One fused tensor cannot hold both, and the refusal comes before the other projection folds.""" +@pytest.mark.parametrize( + ("change", "match"), + [(_mixed, r"down_proj: weight quantizers \[1\] of 4"), (_native_fp8, "native FP8")], + ids=["mixed", "native-fp8"], +) +def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, change, match): + """One fused tensor cannot mix quantized and unquantized experts, and native FP8 weights need the + scale_inv the fold drops. The refusal comes before the other projection folds.""" experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) - experts.down_proj_weight_quantizers[1].disable() + change(experts) before = {k: v.clone() for k, v in experts.state_dict().items()} - with pytest.raises(NotImplementedError, match=r"down_proj: weight quantizers \[1\] of 4"): + with pytest.raises(NotImplementedError, match=match): _export_fused_experts_keep_fused(experts, torch.bfloat16) _assert_byte_identical(experts.state_dict(), before) @@ -410,6 +426,23 @@ def test_split_leaves_a_state_dict_without_fused_experts_unchanged(): assert all(out[k] is sd[k] for k in sd) +@pytest.mark.parametrize( + ("rows", "scale_rows"), [(6, 3), (10, 4)], ids=["block-across-the-cut", "uneven-blocks"] +) +def test_split_refuses_a_scale_block_gate_and_up_would_share(rows, scale_rows): + """Two-row blocks over 6 rows put gate's last row and up's first under one scale; three-row + blocks over 10 rows end in a partial block, which no proportion places. Neither half can be given + scales of its own, so the split refuses rather than round.""" + sd = { + f"{PREFIX}.gate_up_proj": torch.zeros(1, rows, 4).to(torch.float8_e4m3fn), + f"{PREFIX}.gate_up_proj_weight_scale": torch.ones(1, scale_rows, 1), + f"{PREFIX}.down_proj": torch.zeros(1, 4, rows // 2).to(torch.float8_e4m3fn), + } + + with pytest.raises(NotImplementedError, match=r"cannot be split at the gate\|up boundary"): + split_fused_experts_state_dict(sd) + + def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): fused = torch.zeros(2, 4, 4) sd = { @@ -523,11 +556,39 @@ def _two_rank_split_job(rank, world_size): with pytest.raises(NotImplementedError, match="leaves a rank none"): _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) + # Expert parallelism leaves a rank only its own experts' quantizers, in shard order, under a + # fused weight spanning every expert. Kept last: without that pairing rank 1 alone fails, and + # no collective follows for rank 0 to wait in. + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG) + plain = _keep_fused_export(experts) + n_local = NUM_EXPERTS // world_size + mine = slice(rank * n_local, (rank + 1) * n_local) + + def sharded_keeping_quantizers(kept): + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + for proj in (sharded._first_proj_attr, "down_proj"): + name = f"{proj}_weight_quantizers" + setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[kept])) + return sharded + + with pytest.raises(NotImplementedError, match="3 weight quantizers for 4 experts in 2 shards"): + _export_fused_experts_keep_fused(sharded_keeping_quantizers(slice(0, 3)), torch.bfloat16) + sharded = sharded_keeping_quantizers(mine) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + for key, value in fused.items(): + if isinstance(value, DTensor): + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(plain[key][mine])), key + else: + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index, and the fused tensors stay sharded on the expert axis at their global shape, an uneven shard included. - A layout that leaves a rank no expert is refused on every rank.""" + Its quantizers are every expert's (FSDP2) or, under expert parallelism, its own two. A layout + that leaves a rank no expert, or quantizers that are neither, is refused on every rank.""" cpu_dist_workers.run(_two_rank_split_job) @@ -584,32 +645,36 @@ def window(root, module): assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) -def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_one): +@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_one, root_sharded): """On the no-gather path the model is still sharded when a linear is packed. Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the DTensor's local view is refused under autograd), so the handler packs inside a weight-update - window; the sharded state dict must then carry the packed weight, as an unsharded pack has it. + window, on the layer's own unit where expert parallelism leaves the root unsharded; the sharded + state dict must then carry the packed weight, as an unsharded pack has it. """ torch.manual_seed(0) - model = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + model = nn.Sequential(layer) mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) - ref = copy.deepcopy(model) - ExportModuleRegistry.match(ref[0])("0", ref[0], ExportContext(model=ref, dtype=torch.bfloat16)) - fully_shard(model[0], mesh=gloo_world_of_one) - fully_shard(model, mesh=gloo_world_of_one) + ref = copy.deepcopy(model)[0][0] + ExportModuleRegistry.match(ref)("0.0", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + fully_shard(layer, mesh=gloo_world_of_one) + if root_sharded: + fully_shard(model, mesh=gloo_world_of_one) set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) - ExportModuleRegistry.match(model[0])("0", model[0], ctx) + ExportModuleRegistry.match(layer[0])("0.0", layer[0], ctx) sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) - packed = sd["0.weight"] + packed = sd["0.0.weight"] assert isinstance(packed, DTensor) assert packed.dtype == torch.uint8 - assert torch.equal(packed.full_tensor(), ref[0].weight) + assert torch.equal(packed.full_tensor(), ref.weight) for name in ("weight_scale", "weight_scale_2", "input_scale"): - assert torch.equal(_raw_bytes(sd[f"0.{name}"]), _raw_bytes(getattr(ref[0], name))), name + assert torch.equal(_raw_bytes(sd[f"0.0.{name}"]), _raw_bytes(getattr(ref, name))), name def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): From 8492e0c561389532092a19a5d3212f74c2cc5fe9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:19:46 +0000 Subject: [PATCH 14/67] Refuse native FP8 experts the fold would leave without their scale_inv The keep-fused fold refused native FP8 expert weights only in a projection whose experts are all quantized, but it deletes _scale_inv and _activation_scale for both projections. A native FP8 projection whose experts a recipe leaves unquantized was therefore kept as raw FP8 bytes with its scale_inv dropped, and split_fused_experts_state_dict then cut them into per-expert weights nothing can dequantize, where the per-expert export writes them dequantized. The check now runs for every projection, before either is folded, so that case is refused with the module untouched. Also: _fsdp2_weight_window's docstring says why an unwrapped root hands the window the unit itself (fsdp2_aware_weight_update opens it on the innermost enclosing unit, but only under a wrapped root), and _split_local_fused_module's comment no longer says the fold emits plain EP-group scales, which it emits as DTensors. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 11 +++++----- modelopt/torch/export/moe_utils.py | 22 ++++++++++--------- .../export/test_fused_experts_keep_fused.py | 18 +++++++++++---- 3 files changed, 32 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 4983eadfc3c..9d06fb52a6e 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -41,11 +41,12 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: def _fsdp2_weight_window(ctx: ExportContext, module: nn.Module): """``fsdp2_aware_weight_update`` on the FSDP2 unit that owns ``module``'s weights. - That is the model when the model itself is wrapped. Expert parallelism FSDP2-shards only the dense - layers and leaves the root unwrapped, and a window opened on an unwrapped root does nothing: the - weight stays a sharded DTensor, which packing mishandles (``DTensor << int`` is a silent no-op in - torch 2.11, so NVFP4 packed unrelated codes next to correct scales). So open it on the innermost - enclosing unit instead, found by walking up the module's name: the export calls this per module. + The window opens on the innermost unit enclosing the module, but only when the root it is handed + is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the + dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which + packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated + codes next to correct scales). So there that unit is handed in as the root, found by walking up + the module's name: the export calls this per module. """ from torch.distributed.fsdp import FSDPModule diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index a05ffb1ce12..52df40be6c7 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -184,14 +184,16 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N for proj in (first, "down_proj"): weight_quantizers = getattr(module, f"{proj}_weight_quantizers") enabled = [q.is_enabled for q in weight_quantizers] + wdata = getattr(module, proj).data + if wdata.element_size() == 1: + # get_expert_weight dequantizes these with their scale_inv, which the fold drops below + # whether it folds the projection or leaves it as it is. + raise NotImplementedError( + f"{type(module).__name__}.{proj}: native FP8 expert weights cannot be exported as " + "stored, since the fold drops their scale_inv; dequantize them first " + "(unpack_weight())." + ) if all(enabled): - wdata = getattr(module, proj).data - if wdata.element_size() == 1: - # get_expert_weight dequantizes these with their scale_inv, which the fold drops. - raise NotImplementedError( - f"{type(module).__name__}.{proj}: native FP8 expert weights cannot be folded " - "as stored; dequantize them first (unpack_weight())." - ) if DTensor is not None and isinstance(wdata, DTensor): _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) folded.append(proj) @@ -746,9 +748,9 @@ def _loc(key): def _loc_scale(key): # Weight scales are dim-0 (per-expert) tensors. A DTensor scale shards with the weight -> - # to_local() aligns. But under EP the keep-fused fold emits the scale as a PLAIN buffer - # spanning the whole EP group (not dp-sharded like the weight param); slice it by the - # weight's within-group dp_offset so each scale pairs with its own dp-local expert. + # to_local() aligns. The keep-fused fold emits DTensor scales for a sharded weight, but a + # scale handed in as a PLAIN tensor over every expert of the group is accepted too: slice + # it by the weight's within-group dp_offset so each scale pairs with its own dp-local expert. v = sharded_sd.get(key) if v is None: return None diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 8618950fc1d..ac6b9e155c7 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -348,15 +348,25 @@ def _native_fp8(experts): ) +def _native_fp8_unquantized(experts): + _native_fp8(experts) + _unquantized(experts, ["down_proj"]) + + @LAYOUTS @pytest.mark.parametrize( ("change", "match"), - [(_mixed, r"down_proj: weight quantizers \[1\] of 4"), (_native_fp8, "native FP8")], - ids=["mixed", "native-fp8"], + [ + (_mixed, r"down_proj: weight quantizers \[1\] of 4"), + (_native_fp8, "native FP8"), + (_native_fp8_unquantized, "native FP8"), + ], + ids=["mixed", "native-fp8", "native-fp8-unquantized"], ) def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, change, match): """One fused tensor cannot mix quantized and unquantized experts, and native FP8 weights need the - scale_inv the fold drops. The refusal comes before the other projection folds.""" + scale_inv the fold drops, also where a recipe leaves them unquantized. The refusal comes before + the other projection folds.""" experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) change(experts) before = {k: v.clone() for k, v in experts.state_dict().items()} @@ -365,7 +375,7 @@ def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, ch _export_fused_experts_keep_fused(experts, torch.bfloat16) _assert_byte_identical(experts.state_dict(), before) - assert experts.down_proj_weight_quantizers[0].is_enabled + assert getattr(experts, f"{experts._first_proj_attr}_weight_quantizers")[0].is_enabled # ------------------------------------------------------------------ split_fused_experts_state_dict From 73570e583f35129e69260caf9c5a2cf5be06a730 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:55:58 +0000 Subject: [PATCH 15/67] Decide a fused expert projection's fold on every rank alike Under expert parallelism each rank keeps only its own experts' weight quantizers, so the keep-fused fold decided on a rank-local mask: one rank could fold a projection that another left as it was (all of its experts disabled) or refused (some of them), and the ranks' state dicts and the writer's collectives then diverged. The fold now gathers the projection's quantizer states over the mesh dim that shards the experts, joined by global offset, and folds, leaves or refuses it on the gathered mask, so every rank decides alike and a refusal names the global experts. The shard checks now run first for every sharded projection, so the ranks also agree on whether to gather. The shard checks also refuse a fused weight sharded on another dim or over more than one mesh dim: the fold shards the scales on the expert axis alone and gathers over one mesh dim. An unsharded weight is refused when num_experts, its experts and its quantizers disagree, instead of folding only the first experts. split_fused_experts_state_dict's docstring now asks for one module's shard at a time, since its input keeps the fused tensors while the per-expert copies are made. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 45 ++++++++++++-- .../export/test_fused_experts_keep_fused.py | 58 ++++++++++++++++--- 2 files changed, 90 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 52df40be6c7..cb8b0931413 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -122,11 +122,21 @@ def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int) def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: """Refuse a sharded fused expert weight the keep-fused fold cannot fold, on every rank alike. - Reads only the global shape, the mesh and the quantizer count, which every rank agrees on. + Every rank reaches the same verdict: the global shape and the mesh are the same on all of them, + and their quantizer counts differ only where each keeps its own experts' quantizers of an uneven + shard, which the checks below refuse on every rank. """ from torch.distributed.tensor import Shard name = f"{type(module).__name__}.{proj}" + sharded = [placement for placement in weight.placements if isinstance(placement, Shard)] + if len(sharded) > 1 or any(placement.dim != 0 for placement in sharded): + # The fold shards the scales on the expert axis alone, and the ranks agree on their experts' + # quantization over the one mesh dim that shards it. + raise NotImplementedError( + f"{name}: placements {weight.placements}; the keep-fused fold needs the experts (dim 0) " + "sharded over one mesh dim and no other dim sharded." + ) shards = 1 for mdim, placement in enumerate(weight.placements): if isinstance(placement, Shard) and placement.dim == 0: @@ -159,7 +169,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N single fused param: on a ``Shard(0)`` DTensor it quantizes ONLY this rank's local experts (``to_local()``, each with its own weight quantizer: by global index where every rank keeps every expert's, by position in the shard where a rank keeps only its own) and rewraps the result - ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. + ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. Where a rank keeps only + its own experts' quantizers, the ranks first gather each projection's quantizer states + (collective over the expert mesh dim), so every rank folds, leaves or refuses it alike. On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` @@ -193,9 +205,28 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N "stored, since the fold drops their scale_inv; dequantize them first " "(unpack_weight())." ) + is_dt = DTensor is not None and isinstance(wdata, DTensor) + if is_dt: + _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) + if len(weight_quantizers) != wdata.shape[0]: + # Expert parallelism leaves each rank only its own experts' quantizers, so the ranks + # decide on the projection together: alone, one rank could fold it while another + # leaves it as it is or refuses it. Joined by offset rather than by group rank. + mdim = next(i for i, p in enumerate(wdata.placements) if isinstance(p, Shard)) + gathered: list = [None] * wdata.device_mesh.size(mdim) + torch.distributed.all_gather_object( + gathered, + (_dtensor_dim0_offset(wdata), enabled), + group=wdata.device_mesh.get_group(mdim), + ) + enabled = [e for _, shard in sorted(gathered) for e in shard] if all(enabled): - if DTensor is not None and isinstance(wdata, DTensor): - _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) + if not is_dt and not len(weight_quantizers) == wdata.shape[0] == n: + raise NotImplementedError( + f"{type(module).__name__}.{proj}: num_experts is {n}, but the unsharded weight " + f"holds {wdata.shape[0]} experts and there are {len(weight_quantizers)} weight " + "quantizers; the keep-fused fold needs the three equal." + ) folded.append(proj) elif any(enabled): raise NotImplementedError( @@ -559,13 +590,15 @@ def _fused_experts_prefixes(state_dict: dict) -> list[str]: def split_fused_experts_state_dict(state_dict: dict) -> dict: - """Rewrite fused-experts entries in a (consolidated) HF state dict into per-expert entries. + """Rewrite the fused-experts entries of a state dict into per-expert entries. The distributed export writes experts *fused* (3-D ``experts.`` weights + fused scale buffers, kept sharded for DCP). This post-pass splits them into the same per-expert layout the in-model :func:`_export_fused_experts` produces, so the two export paths yield identical checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no - quantization. Non-fused-experts keys pass through unchanged. + quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and + the input keeps its fused tensors, so the experts are held twice until it returns: pass one + module's shard at a time, as :func:`_split_local_fused_module` does, not a whole model's. Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index ac6b9e155c7..427e212ae22 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -27,6 +27,7 @@ import contextlib import copy import math +import re import warnings import pytest @@ -353,6 +354,10 @@ def _native_fp8_unquantized(experts): _unquantized(experts, ["down_proj"]) +def _fewer_experts_counted(experts): + experts.num_experts = NUM_EXPERTS - 1 + + @LAYOUTS @pytest.mark.parametrize( ("change", "match"), @@ -360,13 +365,15 @@ def _native_fp8_unquantized(experts): (_mixed, r"down_proj: weight quantizers \[1\] of 4"), (_native_fp8, "native FP8"), (_native_fp8_unquantized, "native FP8"), + (_fewer_experts_counted, "num_experts is 3, but the unsharded weight holds 4 experts"), ], - ids=["mixed", "native-fp8", "native-fp8-unquantized"], + ids=["mixed", "native-fp8", "native-fp8-unquantized", "num-experts-short"], ) def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, change, match): """One fused tensor cannot mix quantized and unquantized experts, and native FP8 weights need the - scale_inv the fold drops, also where a recipe leaves them unquantized. The refusal comes before - the other projection folds.""" + scale_inv the fold drops, also where a recipe leaves them unquantized. A ``num_experts`` short of + the weight's rows would fold only the first experts. The refusal comes before the other + projection folds.""" experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) change(experts) before = {k: v.clone() for k, v in experts.state_dict().items()} @@ -567,8 +574,9 @@ def _two_rank_split_job(rank, world_size): _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) # Expert parallelism leaves a rank only its own experts' quantizers, in shard order, under a - # fused weight spanning every expert. Kept last: without that pairing rank 1 alone fails, and - # no collective follows for rank 0 to wait in. + # fused weight spanning every expert. Kept last, with the joint decision below: without that + # pairing rank 1 alone fails, and rank 0 then waits only in the collectives below, which the + # worker pool ends once a rank has failed. experts = _quantized_experts(_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG) plain = _keep_fused_export(experts) n_local = NUM_EXPERTS // world_size @@ -593,15 +601,51 @@ def sharded_keeping_quantizers(kept): else: assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + # Each rank sees only its own experts' quantizers there, so the ranks decide on a projection + # together: rank 1 disabling some or all of its experts' is refused on both ranks, rather than + # folded on rank 0 alone, and every expert disabled leaves the projection as it is on both. + for disabled in ([3], [2, 3], [0, 1, 2, 3]): + sharded = sharded_keeping_quantizers(mine) + for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, start=mine.start): + if expert in disabled: + quantizer.disable() + if len(disabled) < NUM_EXPERTS: + with pytest.raises(NotImplementedError, match=re.escape(f"quantizers {disabled} of 4")): + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + else: + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert hasattr(sharded, "gate_up_proj_weight_scale") + assert not hasattr(sharded, "down_proj_weight_scale") + def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index, and the fused tensors stay sharded on the expert axis at their global shape, an uneven shard included. - Its quantizers are every expert's (FSDP2) or, under expert parallelism, its own two. A layout - that leaves a rank no expert, or quantizers that are neither, is refused on every rank.""" + Its quantizers are every expert's (FSDP2) or, under expert parallelism, its own two, and then + the ranks decide on each projection's quantization together. A layout that leaves a rank no + expert, or quantizers that are neither, is refused on every rank.""" cpu_dist_workers.run(_two_rank_split_job) +@pytest.mark.parametrize( + ("mesh_shape", "placements"), + [((1,), [Shard(2)]), ((1, 1), [Shard(0), Shard(0)])], + ids=["off-the-expert-axis", "over-two-mesh-dims"], +) +def test_a_fused_weight_sharded_other_than_by_expert_is_refused( + gloo_world_of_one, mesh_shape, placements +): + """The fold shards the scales on the expert axis alone, and the ranks agree on their experts + over one mesh dim: the scales of another layout would claim one rank's slice as the whole.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + mesh = init_device_mesh("cpu", mesh_shape) + weight = distribute_tensor(experts.down_proj.data, mesh, placements) + experts.down_proj = nn.Parameter(weight, requires_grad=False) + + with pytest.raises(NotImplementedError, match="sharded over one mesh dim"): + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + # ------------------------------------------------------------------ export handlers From 390346a254717162e15a6055cd6f1d466ee27db4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:18:37 +0000 Subject: [PATCH 16/67] Fold fused experts in the weight domain the per-expert export packs The per-expert export reads each expert through get_expert_weight, which casts native FP8 experts' unpacked weights, held at the default dtype, to the export dtype. The fold read the raw slices, so it derived NVFP4 block scales from fp32 and left a projection with no enabled quantizer in fp32. It now reads each expert through that accessor, from this rank's rows, and writes a projection it leaves as it is in the same domain. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 27 ++++++++++++- .../export/test_fused_experts_keep_fused.py | 38 +++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index cb8b0931413..b0509f46b71 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -18,6 +18,7 @@ import copy import warnings from pathlib import Path +from types import SimpleNamespace import torch import torch.nn as nn @@ -119,6 +120,16 @@ def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int) return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) +def _local_expert_weight(module: nn.Module, proj: str, local: torch.Tensor, idx: int, dtype): + """``module.get_expert_weight(proj, idx, dtype)``, read from ``local`` rather than the module. + + The per-expert export packs each expert as that accessor returns it, which casts unpacked native + FP8 experts, held at the default dtype, to the export dtype. The accessor indexes the fused + parameter itself, a sharded DTensor here, so it is handed this rank's rows instead. + """ + return type(module).get_expert_weight(SimpleNamespace(**{proj: local}), proj, idx, dtype) + + def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: """Refuse a sharded fused expert weight the keep-fused fold cannot fold, on every rank alike. @@ -267,7 +278,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N input_scale: torch.Tensor | None = None for local_idx, global_idx in enumerate(local_indices): w_quantizer = weight_quantizers[global_idx if quantizers_are_global else local_idx] - w_slice = local_w[local_idx] + w_slice = _local_expert_weight(module, proj, local_w, local_idx, dtype) # Uncalibrated-expert fallback: derive amax from this expert's local weight slice # (matches _export_fused_experts), so a never-routed expert still exports sane scales. # Per stage, in the shape each expects (W4A8's SequentialQuantizer, block-wise INT4). @@ -314,6 +325,20 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if input_scale is not None: module.register_buffer(f"{proj}_input_scale", input_scale) + # A projection left as it is goes out in the weight domain the per-expert export writes it in. + for proj in (first, "down_proj"): + if proj in folded: + continue + wdata = getattr(module, proj).data + is_dt = DTensor is not None and isinstance(wdata, DTensor) + local = wdata.to_local() if is_dt else wdata + domain = _local_expert_weight(module, proj, local, 0, dtype).dtype + if domain != local.dtype: + local = local.to(domain) + if is_dt: + local = _rewrap_expert_rows(local, wdata.device_mesh, wdata.placements, wdata.shape[0]) + setattr(module, proj, nn.Parameter(local, requires_grad=False)) + # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers -- and # the source checkpoint's scales, which described the weights the fold replaced, as the # per-expert export does (_delete_fused_moe_source_attrs). diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 427e212ae22..ccd9a59ef83 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -277,6 +277,44 @@ def test_keep_fused_then_split_writes_exactly_the_per_expert_export( _assert_byte_identical(got, want) +@pytest.mark.parametrize("unquantized", [(), ("down_proj",)], ids=["quantized", "down-unquantized"]) +def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_expert_export( + unquantized, +): + """Native FP8 experts unpack at the default dtype, and the per-expert export reads each expert in + the export dtype (``get_expert_weight``): the fold reads them alike, for the scales it derives + and for a projection it leaves unquantized.""" + from transformers import PretrainedConfig + from transformers.integrations.finegrained_fp8 import FP8Experts + + torch.manual_seed(0) + config = PretrainedConfig( + hidden_size=HIDDEN_DIM, + intermediate_size=INTERMEDIATE_DIM, + num_experts=NUM_EXPERTS, + hidden_act="silu", + ) + experts = FP8Experts(config, block_size=None) + for name in ("gate_up_proj", "down_proj"): + weight = getattr(experts, name) + with torch.no_grad(): + weight.copy_((torch.randn(weight.shape) * 4).to(weight.dtype)) + getattr(experts, f"{name}_scale_inv").fill_(0.01) + hidden = torch.randn(NUM_EXPERTS, HIDDEN_DIM) + route = (torch.arange(NUM_EXPERTS).unsqueeze(1), torch.ones(NUM_EXPERTS, 1)) + registered_before = QuantModuleRegistry.get(FP8Experts) + try: + mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG, lambda m: m(hidden, *route)) + experts.unpack_weight() + _unquantized(experts, unquantized) + want = _per_expert_export(experts) + got = split_fused_experts_state_dict(_keep_fused_export(experts)) + finally: + if registered_before is None and QuantModuleRegistry.get(FP8Experts) is not None: + QuantModuleRegistry.unregister(FP8Experts) + _assert_byte_identical(got, want) + + @pytest.mark.parametrize( "quant_cfg", [ From d62f873c0d9078ee1054d3c5801842482f57e21a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:19:40 +0000 Subject: [PATCH 17/67] Refuse a keep-fused export of fused experts not named experts The fold claims a fused-experts module by its structure, but the split into per-expert keys, and the writer, find its fused weights by their experts. key. A container named otherwise was folded and then written fused, to fail only at load. Refuse it before anything is folded. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 8 ++++++++ .../export/test_fused_experts_keep_fused.py | 16 ++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 9d06fb52a6e..6d99d1e7071 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -180,6 +180,14 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex the per-expert split happens later on the sharded weight (``split_fused_experts_state_dict``). Otherwise (single process, or gather export) it splits into full per-expert submodules here. """ + attr = name.rpartition(".")[2] + if ctx.keep_fused_experts and attr != "experts": + # The split into per-expert keys finds a fused weight by its ``experts.`` key, as the + # writer does, so one named otherwise would be written fused and fail only at load. + raise NotImplementedError( + f"{name}: the keep-fused export splits fused experts into per-expert keys by their " + f"`experts` attribute name, and this module is named {attr!r}." + ) if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index ccd9a59ef83..a1e819d68fd 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -708,6 +708,22 @@ def test_fused_experts_handler_keeps_experts_fused_only_when_asked(keep_fused): assert "gate_up_proj" not in names + +def test_keep_fused_refuses_fused_experts_not_named_experts(): + """The split finds a fused weight by its ``experts.`` key, as the writer does: a container + named otherwise is refused before anything is folded, not written fused to fail at load.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + container = nn.Module() + container.moe = experts + before = copy.deepcopy(experts.state_dict()) + ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) + + with pytest.raises(NotImplementedError, match="named 'moe'"): + ExportModuleRegistry.match(experts)("model.layers.0.mlp.moe", experts, ctx) + + torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) + + @pytest.mark.parametrize("dense", ["fsdp", "replicate"]) def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_part( monkeypatch, dense From 89504e4e33310d155d56cb085b315c5f8037716f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:23:32 +0000 Subject: [PATCH 18/67] Pack each FSDP2 unit of a keep-fused export in one weight window fsdp2_aware_weight_update unshards a module's whole FSDP2 unit and reshards it when it closes, so a window per module gathered a decoder layer once per quantized linear in it. The handlers now hold the unit's window open from one module to the next, each joining the modules it re-registers with FSDP2, and the export driver closes the last. The window also resolves the unit's names once rather than walking the unit per parameter. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 62 ++++++++++++++----- modelopt/torch/export/registry.py | 6 +- modelopt/torch/export/unified_export_hf.py | 11 +++- .../export/test_fused_experts_keep_fused.py | 46 +++++++++----- 4 files changed, 88 insertions(+), 37 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 6d99d1e7071..da409e30bed 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -16,7 +16,9 @@ """Built-in module handlers for unified Hugging Face export.""" import collections.abc +import contextlib import warnings +from dataclasses import dataclass import torch.nn as nn @@ -38,25 +40,55 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") -def _fsdp2_weight_window(ctx: ExportContext, module: nn.Module): - """``fsdp2_aware_weight_update`` on the FSDP2 unit that owns ``module``'s weights. +@dataclass +class _HeldWindow: + """An FSDP2 unit's weight-update window, and the modules packed in it so far.""" + + unit: nn.Module + modules: list[nn.Module] + stack: contextlib.ExitStack + + +def _hold_fsdp2_weight_window(ctx: ExportContext, module: nn.Module) -> None: + """Hold open the ``fsdp2_aware_weight_update`` window of the FSDP2 unit owning ``module``'s weights. + + The window unshards the whole unit and reshards it when it closes, so the export packs every + module of a unit in one window: it stays open from one handler to the next until the export + reaches another unit, and :func:`close_fsdp2_weight_window` closes the last. Each module joins + the list the window was opened with, whose weights it re-registers with FSDP2 when it closes. The window opens on the innermost unit enclosing the module, but only when the root it is handed is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated codes next to correct scales). So there that unit is handed in as the root, found by walking up - the module's name: the export calls this per module. + the module's name. """ from torch.distributed.fsdp import FSDPModule if ctx.module_names is None: ctx.module_names = module_name_maps(ctx.model) - if isinstance(ctx.model, FSDPModule): - return fsdp2_aware_weight_update(ctx.model, module, names=ctx.module_names) unit = _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model - # Names relative to the unit, which the window resolves the module's parameters against. - return fsdp2_aware_weight_update(unit, module) + held = ctx.fsdp2_window + if held is not None and held.unit is unit: + held.modules.append(module) + return + close_fsdp2_weight_window(ctx) + if isinstance(ctx.model, FSDPModule): + root, names = ctx.model, ctx.module_names + else: + # Names relative to the unit, which the window resolves the modules' parameters against. + root, names = unit, module_name_maps(unit) + held = _HeldWindow(unit, [module], contextlib.ExitStack()) + held.stack.enter_context(fsdp2_aware_weight_update(root, held.modules, names=names)) + ctx.fsdp2_window = held + + +def close_fsdp2_weight_window(ctx: ExportContext) -> None: + """Close the weight-update window the handlers hold open, if any, resharding its unit.""" + held, ctx.fsdp2_window = ctx.fsdp2_window, None + if held is not None: + held.stack.close() def _export_weight( @@ -71,14 +103,11 @@ def _export_weight( # The no-gather writer packs while the model is still sharded, so unlike the # collect_export_tensors path there is no open weight-access window here and the weight is a # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an - # index_put_, which DTensor cannot dispatch. Open a window per module for that path only, and - # reshard on the way out: a layer left unsharded holds all of its weights on every rank, and - # FSDP2 later puts its own shards back over anything packed in it outside a window. + # index_put_, which DTensor cannot dispatch. Pack in the unit's window for that path only, which + # reshards it once the export moves on: a layer left unsharded holds all of its weights on every + # rank, and FSDP2 later puts its own shards back over anything packed in it outside a window. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: - with _fsdp2_weight_window(ctx, module): - _export_quantized_weight(module, ctx.dtype, weight_name) - return - + _hold_fsdp2_weight_window(ctx, module) _export_quantized_weight(module, ctx.dtype, weight_name) @@ -192,9 +221,8 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its # original shards the next time the layer reshards, and exported unquantized. - with _fsdp2_weight_window(ctx, module): - _export_fused_experts_keep_fused(module, ctx.dtype) - elif ctx.keep_fused_experts: + _hold_fsdp2_weight_window(ctx, module) + if ctx.keep_fused_experts: # No-gather distributed export: fold each rank's LOCAL experts and leave the weight # fused + sharded; the per-expert split happens after consolidation. _export_fused_experts_keep_fused(module, ctx.dtype) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 6640b5ceb30..90efb594e69 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,7 +28,7 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -72,6 +72,10 @@ class ExportContext: """``model``'s name maps, built once per export by the first handler that looks a module up (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + fsdp2_window: Any = None + """The FSDP2 unit's weight-update window the handlers hold open across its modules, closed by + :func:`~modelopt.torch.export.hf_export_handlers.close_fsdp2_weight_window`.""" + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 14f3c82fd36..7c5d3c5088d 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -104,7 +104,7 @@ export_sparse_attention_config = None # Importing the built-in handlers installs their entries in the two registries. -from . import hf_export_handlers as _hf_export_handlers # noqa: F401 +from . import hf_export_handlers as _hf_export_handlers from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import get_experts_list, is_layernorm, is_quantlinear, sync_moe_gate_up_amax from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model @@ -996,8 +996,13 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - for name, sub_module in model.named_modules(): - _dispatch_export_handler(name, sub_module, ctx) + try: + for name, sub_module in model.named_modules(): + _dispatch_export_handler(name, sub_module, ctx) + finally: + # A keep-fused export's handlers hold an FSDP2 unit's weight window open across its + # modules; close the last one. + _hf_export_handlers.close_fsdp2_weight_window(ctx) def _prepare_model_for_export(model, dtype, is_modelopt_qlora): diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index a1e819d68fd..f2f9b651532 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -36,7 +36,7 @@ import torch.nn.functional as F from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.fsdp import fully_shard +from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor pytest.importorskip("transformers") @@ -733,12 +733,12 @@ def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_pa events = [] @contextlib.contextmanager - def window(root, module): + def window(root, modules, names=None): events.append("open") yield events.append("close") - monkeypatch.setattr(handlers, "_fsdp2_weight_window", window) + monkeypatch.setattr(handlers, "fsdp2_aware_weight_update", window) monkeypatch.setattr( handlers, "_export_fused_experts_keep_fused", lambda module, dtype: events.append("fold") ) @@ -749,40 +749,54 @@ def window(root, module): ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) + handlers.close_fsdp2_weight_window(ctx) assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) @pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) -def test_quant_linear_handler_packs_an_fsdp2_sharded_nvfp4_weight(gloo_world_of_one, root_sharded): +def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( + monkeypatch, gloo_world_of_one, root_sharded +): """On the no-gather path the model is still sharded when a linear is packed. Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the - DTensor's local view is refused under autograd), so the handler packs inside a weight-update - window, on the layer's own unit where expert parallelism leaves the root unsharded; the sharded - state dict must then carry the packed weight, as an unsharded pack has it. + DTensor's local view is refused under autograd), so the handlers pack inside a weight-update + window, on the layer's own unit where expert parallelism leaves the root unsharded. The window + gathers the whole unit, so it opens once for all of the unit's linears; the sharded state dict + must then carry the packed weights, as an unsharded pack has them. """ torch.manual_seed(0) layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) model = nn.Sequential(layer) mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) - ref = copy.deepcopy(model)[0][0] - ExportModuleRegistry.match(ref)("0.0", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + refs = copy.deepcopy(model)[0] + for i, ref in enumerate(refs): + ExportModuleRegistry.match(ref)(f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16)) fully_shard(layer, mesh=gloo_world_of_one) if root_sharded: fully_shard(model, mesh=gloo_world_of_one) set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + unshards = [] + unshard = FSDPModule.unshard + monkeypatch.setattr( + FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) + ) - ExportModuleRegistry.match(layer[0])("0.0", layer[0], ctx) + for i, linear in enumerate(layer): + ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) + handlers.close_fsdp2_weight_window(ctx) + assert unshards == [layer] sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) - packed = sd["0.0.weight"] - assert isinstance(packed, DTensor) - assert packed.dtype == torch.uint8 - assert torch.equal(packed.full_tensor(), ref.weight) - for name in ("weight_scale", "weight_scale_2", "input_scale"): - assert torch.equal(_raw_bytes(sd[f"0.0.{name}"]), _raw_bytes(getattr(ref, name))), name + for i, ref in enumerate(refs): + packed = sd[f"0.{i}.weight"] + assert isinstance(packed, DTensor) + assert packed.dtype == torch.uint8 + assert torch.equal(packed.full_tensor(), ref.weight) + for name in ("weight_scale", "weight_scale_2", "input_scale"): + assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), _raw_bytes(getattr(ref, name))), name def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): From efca07b994dfc3a225beb19c68edfb84ccbd86a7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:24:08 +0000 Subject: [PATCH 19/67] Wrap the keep-fused fixes' long lines Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 4 +++- tests/unit/torch/export/test_fused_experts_keep_fused.py | 8 +++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index b0509f46b71..ce821b7cb48 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -336,7 +336,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if domain != local.dtype: local = local.to(domain) if is_dt: - local = _rewrap_expert_rows(local, wdata.device_mesh, wdata.placements, wdata.shape[0]) + local = _rewrap_expert_rows( + local, wdata.device_mesh, wdata.placements, wdata.shape[0] + ) setattr(module, proj, nn.Parameter(local, requires_grad=False)) # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers -- and diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index f2f9b651532..febefc06dfe 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -708,7 +708,6 @@ def test_fused_experts_handler_keeps_experts_fused_only_when_asked(keep_fused): assert "gate_up_proj" not in names - def test_keep_fused_refuses_fused_experts_not_named_experts(): """The split finds a fused weight by its ``experts.`` key, as the writer does: a container named otherwise is refused before anything is folded, not written fused to fail at load.""" @@ -772,7 +771,9 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) refs = copy.deepcopy(model)[0] for i, ref in enumerate(refs): - ExportModuleRegistry.match(ref)(f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + ExportModuleRegistry.match(ref)( + f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) + ) fully_shard(layer, mesh=gloo_world_of_one) if root_sharded: fully_shard(model, mesh=gloo_world_of_one) @@ -796,7 +797,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( assert packed.dtype == torch.uint8 assert torch.equal(packed.full_tensor(), ref.weight) for name in ("weight_scale", "weight_scale_2", "input_scale"): - assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), _raw_bytes(getattr(ref, name))), name + want = _raw_bytes(getattr(ref, name)) + assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): From 64b9c940666fafb10870df2c4e985f77456b2b4c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:48:06 +0000 Subject: [PATCH 20/67] Hold an FSDP2 unit's weight window open only where the export opts in A held window leaves its unit unsharded until something closes it, and only the _process_quantized_modules loop did. The layerwise and streaming exports dispatch the same handlers from their own loops, and would have left a unit gathered on every rank. The handlers now hold a unit's window only inside held_fsdp2_weight_windows, which that loop enters and which closes the last window on the way out; any other loop packs each module in a window of its own, closed by its handler. Closing a window is no longer a caller's job. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 86 +++++++++++++------ modelopt/torch/export/registry.py | 8 +- modelopt/torch/export/unified_export_hf.py | 7 +- .../export/test_fused_experts_keep_fused.py | 22 +++-- 4 files changed, 81 insertions(+), 42 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index da409e30bed..1e40ab6b7ee 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -23,7 +23,7 @@ import torch.nn as nn from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps -from modelopt.torch.quantization.utils.core_utils import _get_enclosing_fsdp_module +from modelopt.torch.quantization.utils.core_utils import ModuleNames, _get_enclosing_fsdp_module from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax @@ -49,42 +49,73 @@ class _HeldWindow: stack: contextlib.ExitStack -def _hold_fsdp2_weight_window(ctx: ExportContext, module: nn.Module) -> None: - """Hold open the ``fsdp2_aware_weight_update`` window of the FSDP2 unit owning ``module``'s weights. +def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: + """The innermost FSDP2 unit enclosing ``module``, or the model if none does.""" + if ctx.module_names is None: + ctx.module_names = module_name_maps(ctx.model) + return _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + - The window unshards the whole unit and reshards it when it closes, so the export packs every - module of a unit in one window: it stays open from one handler to the next until the export - reaches another unit, and :func:`close_fsdp2_weight_window` closes the last. Each module joins - the list the window was opened with, whose weights it re-registers with FSDP2 when it closes. +def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: + """The root a weight-update window on ``unit`` is opened on, and its module names. The window opens on the innermost unit enclosing the module, but only when the root it is handed is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated - codes next to correct scales). So there that unit is handed in as the root, found by walking up - the module's name. + codes next to correct scales). So there that unit is handed in as the root, with names relative + to it, which the window resolves the modules' parameters against. """ from torch.distributed.fsdp import FSDPModule - if ctx.module_names is None: - ctx.module_names = module_name_maps(ctx.model) - unit = _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + if isinstance(ctx.model, FSDPModule): + return ctx.model, ctx.module_names + return unit, module_name_maps(unit) + + +def _fsdp2_weight_window( + ctx: ExportContext, module: nn.Module +) -> contextlib.AbstractContextManager: + """The ``fsdp2_aware_weight_update`` window to pack ``module``'s weights in. + + Inside :func:`held_fsdp2_weight_windows`, the window of the module's FSDP2 unit, held open + across the unit's modules. Anywhere else, a window of the module's own, closed when the handler + is done, so a dispatch loop that knows nothing of held windows cannot leave a unit unsharded. + """ + unit = _enclosing_unit(ctx, module) + if not ctx.hold_fsdp2_windows: + root, names = _window_root(ctx, unit) + return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: held.modules.append(module) - return - close_fsdp2_weight_window(ctx) - if isinstance(ctx.model, FSDPModule): - root, names = ctx.model, ctx.module_names - else: - # Names relative to the unit, which the window resolves the modules' parameters against. - root, names = unit, module_name_maps(unit) + return contextlib.nullcontext() + _close_fsdp2_weight_window(ctx) + root, names = _window_root(ctx, unit) held = _HeldWindow(unit, [module], contextlib.ExitStack()) held.stack.enter_context(fsdp2_aware_weight_update(root, held.modules, names=names)) ctx.fsdp2_window = held + return contextlib.nullcontext() -def close_fsdp2_weight_window(ctx: ExportContext) -> None: +@contextlib.contextmanager +def held_fsdp2_weight_windows(ctx: ExportContext): + """Within the block, pack each FSDP2 unit's modules in one weight-update window. + + The window unshards the whole unit and reshards it when it closes, so a window per module + gathers a decoder layer once per quantized linear in it. Here the unit's window stays open from + one handler to the next until the export reaches another unit, each module joining the list + whose weights it re-registers with FSDP2, and leaving the block closes the last. + """ + ctx.hold_fsdp2_windows = True + try: + yield + finally: + ctx.hold_fsdp2_windows = False + _close_fsdp2_weight_window(ctx) + + +def _close_fsdp2_weight_window(ctx: ExportContext) -> None: """Close the weight-update window the handlers hold open, if any, resharding its unit.""" held, ctx.fsdp2_window = ctx.fsdp2_window, None if held is not None: @@ -103,11 +134,13 @@ def _export_weight( # The no-gather writer packs while the model is still sharded, so unlike the # collect_export_tensors path there is no open weight-access window here and the weight is a # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an - # index_put_, which DTensor cannot dispatch. Pack in the unit's window for that path only, which - # reshards it once the export moves on: a layer left unsharded holds all of its weights on every - # rank, and FSDP2 later puts its own shards back over anything packed in it outside a window. + # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, + # which reshards the unit when it closes: a layer left unsharded holds all of its weights on + # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: - _hold_fsdp2_weight_window(ctx, module) + with _fsdp2_weight_window(ctx, module): + _export_quantized_weight(module, ctx.dtype, weight_name) + return _export_quantized_weight(module, ctx.dtype, weight_name) @@ -221,8 +254,9 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its # original shards the next time the layer reshards, and exported unquantized. - _hold_fsdp2_weight_window(ctx, module) - if ctx.keep_fused_experts: + with _fsdp2_weight_window(ctx, module): + _export_fused_experts_keep_fused(module, ctx.dtype) + elif ctx.keep_fused_experts: # No-gather distributed export: fold each rank's LOCAL experts and leave the weight # fused + sharded; the per-expert split happens after consolidation. _export_fused_experts_keep_fused(module, ctx.dtype) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 90efb594e69..ca222a5cc97 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -72,9 +72,13 @@ class ExportContext: """``model``'s name maps, built once per export by the first handler that looks a module up (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + hold_fsdp2_windows: bool = False + """Whether the handlers hold an FSDP2 unit's weight-update window open across its modules: set + only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which + closes the last window. Elsewhere each module is packed in a window of its own.""" + fsdp2_window: Any = None - """The FSDP2 unit's weight-update window the handlers hold open across its modules, closed by - :func:`~modelopt.torch.export.hf_export_handlers.close_fsdp2_weight_window`.""" + """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 7c5d3c5088d..63a104fbd43 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -996,13 +996,10 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - try: + # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. + with _hf_export_handlers.held_fsdp2_weight_windows(ctx): for name, sub_module in model.named_modules(): _dispatch_export_handler(name, sub_module, ctx) - finally: - # A keep-fused export's handlers hold an FSDP2 unit's weight window open across its - # modules; close the last one. - _hf_export_handlers.close_fsdp2_weight_window(ctx) def _prepare_model_for_export(model, dtype, is_modelopt_qlora): diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index febefc06dfe..6131d9dd201 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -728,7 +728,8 @@ def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_pa monkeypatch, dense ): """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a - weight-update window, they would be overwritten with its original shards at the next reshard.""" + weight-update window, they would be overwritten with its original shards at the next reshard. + Outside held windows the handler closes the window itself.""" events = [] @contextlib.contextmanager @@ -748,22 +749,24 @@ def window(root, modules, names=None): ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) - handlers.close_fsdp2_weight_window(ctx) assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) +@pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) @pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( - monkeypatch, gloo_world_of_one, root_sharded + monkeypatch, gloo_world_of_one, root_sharded, held ): """On the no-gather path the model is still sharded when a linear is packed. Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the DTensor's local view is refused under autograd), so the handlers pack inside a weight-update window, on the layer's own unit where expert parallelism leaves the root unsharded. The window - gathers the whole unit, so it opens once for all of the unit's linears; the sharded state dict - must then carry the packed weights, as an unsharded pack has them. + gathers the whole unit, so within held windows it opens once for all of the unit's linears; + any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is + sharded again afterwards, and the sharded state dict carries the packed weights, as an + unsharded pack has them. """ torch.manual_seed(0) layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) @@ -785,11 +788,12 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) ) - for i, linear in enumerate(layer): - ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - handlers.close_fsdp2_weight_window(ctx) + with handlers.held_fsdp2_weight_windows(ctx) if held else contextlib.nullcontext(): + for i, linear in enumerate(layer): + ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - assert unshards == [layer] + assert unshards == ([layer] if held else [layer, layer]) + assert all(isinstance(linear.weight, DTensor) for linear in layer) sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) for i, ref in enumerate(refs): packed = sd[f"0.{i}.weight"] From faa81c63bc62fc3818dcc08d8c668ff325927a6b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:48:45 +0000 Subject: [PATCH 21/67] Add a module to the held weight window once The BMM experts handler packs gate_up_proj and then down_proj of the same module, which joined the unit's held window once per weight. Closing the window then re-registered the module's weights with FSDP2 twice: the second pass built an FSDPParam from the DTensor shard the first pass had left, which FSDP2 takes for a tensor-parallel param and refuses. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 5 ++- .../export/test_fused_experts_keep_fused.py | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 1e40ab6b7ee..1d2a7a5073b 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -88,7 +88,10 @@ def _fsdp2_weight_window( return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: - held.modules.append(module) + # A module with several quantized weights (fused BMM experts) packs each in turn: it joins + # the window once, or the window would re-register its weights with FSDP2 twice. + if not any(m is module for m in held.modules): + held.modules.append(module) return contextlib.nullcontext() _close_fsdp2_weight_window(ctx) root, names = _window_root(ctx, unit) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 6131d9dd201..98427f143e8 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -805,6 +805,40 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name +class GptOssExperts(nn.Module): + """Named as transformers' fused BMM experts, whose handler packs each of its weights in turn.""" + + +def test_bmm_experts_pack_both_weights_in_their_unit_s_held_window(gloo_world_of_one): + """``gate_up_proj`` and then ``down_proj`` of one module pack in the unit's held window, which + the module joins once: listed twice, the window would re-register its weights with FSDP2 a + second time, off the shards the first pass left.""" + torch.manual_seed(0) + linears = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 32, bias=False)) + mtq.quantize(linears, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + experts = GptOssExperts() + for proj, linear in zip(("gate_up_proj", "down_proj"), linears): + experts.register_parameter(proj, nn.Parameter(linear.weight.detach().clone())) + setattr(experts, f"{proj}_weight_quantizer", linear.weight_quantizer) + setattr(experts, f"{proj}_input_quantizer", linear.input_quantizer) + ref = copy.deepcopy(experts) + ExportModuleRegistry.match(ref)(PREFIX, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + layer = nn.Sequential(experts) + model = nn.Sequential(layer) + fully_shard(layer, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + ExportModuleRegistry.match(experts)("0.0", experts, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for key, want in ref.state_dict().items(): + got = sd[f"0.0.{key}"] + got = got.full_tensor() if isinstance(got, DTensor) else got + assert torch.equal(_raw_bytes(got), _raw_bytes(want)), key + + def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. From 044460ae0d8742657ad3482f995bc3f47d19057e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 06:49:25 +0000 Subject: [PATCH 22/67] Put the keep-fused fold's stacked scales on the weight's device Per-tensor FP8 makes each expert's weight scale from a Python float, so on the CPU wherever the weight lives, and a sharded fold rewrapped the stacked CPU scales as a DTensor on the weight's mesh: a CUDA mesh over a CPU shard, which no collective can serve. The fold now moves the stacked scales to the weight's device first. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 6 ++++-- .../export/test_fused_experts_keep_fused.py | 20 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index ce821b7cb48..9dbd1377848 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -308,7 +308,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if hasattr(wrapper, "input_scale"): input_scale = wrapper.input_scale local_fp8 = torch.stack(fp8_weights, dim=0) - local_ws = torch.stack(weight_scales, dim=0) + # On the weight's device: per-tensor FP8 builds each expert's scale from a Python float, on + # the CPU, and a DTensor rewrap would label a CPU shard with the weight's device mesh. + local_ws = torch.stack(weight_scales, dim=0).to(local_w.device) if is_dt: # Keep the global expert count: from_local would infer it as local x world, which is # wrong for an uneven shard (3 experts on 2 ranks) and misnumbers the experts. @@ -318,7 +320,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N module.register_buffer(f"{proj}_weight_scale", local_ws) if weight_scale_2s: # Per-expert scalar (E,) -- shards on the expert axis exactly like the weight scale. - local_ws2 = torch.stack(weight_scale_2s, dim=0) + local_ws2 = torch.stack(weight_scale_2s, dim=0).to(local_w.device) if is_dt: local_ws2 = _rewrap_expert_rows(local_ws2, mesh, scale_placements, wdata.shape[0]) module.register_buffer(f"{proj}_weight_scale_2", local_ws2) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 98427f143e8..bb308ff1c27 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -377,6 +377,26 @@ def spy(wrapper, dtype, **_kwargs): assert any("was not calibrated" in str(w.message) for w in caught) == (amax_state == "zero") +def test_the_fold_puts_its_scales_on_the_weights_device(monkeypatch): + """Per-tensor FP8 makes each expert's scale from a Python float, on the CPU wherever the weight + is, and a sharded fold rewraps the scales on the weight's mesh. The weights here sit on meta, + the one other device a CPU runner has, so a stand-in packs them, its scales on the CPU.""" + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) + for proj in ("gate_up_proj", "down_proj"): + setattr(experts, proj, nn.Parameter(getattr(experts, proj).detach().to("meta"))) + + def pack(wrapper, dtype, **_kwargs): + wrapper.weight_scale = torch.tensor(1.0) + wrapper.weight_scale_2 = torch.tensor(1.0) + + monkeypatch.setattr("modelopt.torch.export.unified_export_hf._export_quantized_weight", pack) + _export_fused_experts_keep_fused(experts, torch.bfloat16) + + for proj in ("gate_up_proj", "down_proj"): + for scale in ("weight_scale", "weight_scale_2"): + assert getattr(experts, f"{proj}_{scale}").is_meta, (proj, scale) + + def _mixed(experts): experts.down_proj_weight_quantizers[1].disable() From 833d95733bdc71ba3cddeb784e41ae058d25e9e0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:18:13 +0000 Subject: [PATCH 23/67] Hand the fold's expert accessor the module, with only the projection swapped The keep-fused fold reads each expert through the module's own get_expert_weight, and handed it this rank's rows of the fused weight on a bare namespace. The fused-experts PR (#2745) has that accessor pick this rank's experts with the module's _local_experts, which the namespace lacks: with both in, every fold raised AttributeError, failing 26 keep-fused and fused-experts unit tests. The accessor now reads a view of the module whose only difference is the projection. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 9dbd1377848..1886e009fb2 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -18,7 +18,6 @@ import copy import warnings from pathlib import Path -from types import SimpleNamespace import torch import torch.nn as nn @@ -120,14 +119,25 @@ def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int) return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) +class _LocalRows: + """``module``, with its fused ``proj`` replaced by this rank's rows of it.""" + + def __init__(self, module: nn.Module, proj: str, local: torch.Tensor): + self._module, self._proj, self._local = module, proj, local + + def __getattr__(self, name: str): + return self._local if name == self._proj else getattr(self._module, name) + + def _local_expert_weight(module: nn.Module, proj: str, local: torch.Tensor, idx: int, dtype): """``module.get_expert_weight(proj, idx, dtype)``, read from ``local`` rather than the module. The per-expert export packs each expert as that accessor returns it, which casts unpacked native FP8 experts, held at the default dtype, to the export dtype. The accessor indexes the fused - parameter itself, a sharded DTensor here, so it is handed this rank's rows instead. + parameter itself, a sharded DTensor here, so it is handed this rank's rows instead. Whatever + else it reads, such as how it picks this rank's experts out of a weight, is the module's own. """ - return type(module).get_expert_weight(SimpleNamespace(**{proj: local}), proj, idx, dtype) + return type(module).get_expert_weight(_LocalRows(module, proj, local), proj, idx, dtype) def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: From 6c93a25991669cfba89f36757b9050d13956c225 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:57:23 +0000 Subject: [PATCH 24/67] Stack the keep-fused fold's packed experts in place The fold kept each expert's packed weight and scale in a list and stacked them at the end, so for a moment it held the packed projection twice, beside the high-precision fused weight. It now writes each expert into one stack made from the first expert's packed shape, so the packed projection is held once. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 35 +++++++++++++++++++++++------- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 1886e009fb2..24e3f959639 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -140,6 +140,23 @@ def _local_expert_weight(module: nn.Module, proj: str, local: torch.Tensor, idx: return type(module).get_expert_weight(_LocalRows(module, proj, local), proj, idx, dtype) +def _stack_row(stack: torch.Tensor | None, idx: int, row: torch.Tensor, rows: int, device): + """Write ``row`` as row ``idx`` of a ``rows``-row stack, made on ``device`` for the first row. + + The fold packs one expert at a time: copied into one stack as each is packed, the projection is + held once, where a list of the packed experts and their ``torch.stack`` would hold it twice. + """ + if stack is None: + stack = torch.empty((rows, *row.shape), dtype=row.dtype, device=device) + elif row.shape != stack.shape[1:] or row.dtype != stack.dtype: + raise ValueError( + f"Expert {idx} packs to {row.dtype} {tuple(row.shape)}, the experts before it to " + f"{stack.dtype} {tuple(stack.shape[1:])}; they cannot be stacked." + ) + stack[idx].copy_(row) + return stack + + def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantizers: int) -> None: """Refuse a sharded fused expert weight the keep-fused fold cannot fold, on every rank alike. @@ -282,8 +299,8 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N local_w = wdata local_indices = list(range(n)) - fp8_weights: list[torch.Tensor] = [] - weight_scales: list[torch.Tensor] = [] + local_fp8: torch.Tensor | None = None + local_ws: torch.Tensor | None = None weight_scale_2s: list[torch.Tensor] = [] input_scale: torch.Tensor | None = None for local_idx, global_idx in enumerate(local_indices): @@ -308,8 +325,12 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if input_quantizer is not None: wrapper.input_quantizer = input_quantizer _export_quantized_weight(wrapper, dtype) - fp8_weights.append(wrapper.weight.data) - weight_scales.append(wrapper.weight_scale) + rows = len(local_indices) + packed = wrapper.weight.data + local_fp8 = _stack_row(local_fp8, local_idx, packed, rows, packed.device) + # On the weight's device: per-tensor FP8 builds each expert's scale from a Python float, + # on the CPU, and a DTensor rewrap would label a CPU shard with the weight's device mesh. + local_ws = _stack_row(local_ws, local_idx, wrapper.weight_scale, rows, local_w.device) # NVFP4 carries a SECOND per-tensor weight scale (weight_scale_2) that dequantizes the # per-block weight_scale; FP8 has none. Keep it so the fused buffer (and the per-expert # split) stay dequantizable. @@ -317,10 +338,8 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N weight_scale_2s.append(wrapper.weight_scale_2) if hasattr(wrapper, "input_scale"): input_scale = wrapper.input_scale - local_fp8 = torch.stack(fp8_weights, dim=0) - # On the weight's device: per-tensor FP8 builds each expert's scale from a Python float, on - # the CPU, and a DTensor rewrap would label a CPU shard with the weight's device mesh. - local_ws = torch.stack(weight_scales, dim=0).to(local_w.device) + # Every rank holds an expert (_refuse_unfoldable_shards), so both stacks have a row. + assert local_fp8 is not None and local_ws is not None if is_dt: # Keep the global expert count: from_local would infer it as local x world, which is # wrong for an uneven shard (3 experts on 2 ranks) and misnumbers the experts. From 455fa15e44ad1136068f3c0e6b0865d44ec4a4af Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:58:00 +0000 Subject: [PATCH 25/67] Say what orders an expert-parallel rank's quantizers, without a symbol this tree lacks The comment cited a helper that comes with another PR of the series. It now states the invariant itself: transformers sets num_experts to the local count, so the quantized module makes one quantizer per expert of its local weight, in the shard's order. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 24e3f959639..46b76ddd680 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -189,7 +189,8 @@ def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantize ) shards *= size # FSDP2 keeps every expert's quantizer on every rank; expert parallelism keeps only this rank's, - # in the order of its shard, as _QuantFusedExperts._local_experts reads them. + # in the order of its shard: transformers sets num_experts to the local count, so the quantized + # module makes one quantizer per expert of its local weight, which its forward picks by offset. if num_quantizers != weight.shape[0] and num_quantizers * shards != weight.shape[0]: raise NotImplementedError( f"{name}: {num_quantizers} weight quantizers for {weight.shape[0]} experts in {shards} " From aaf9c1e9a122e6d23503e8fd49587ea560eaacfb Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:58:20 +0000 Subject: [PATCH 26/67] Say which experts a rank folds inside an FSDP2 weight window The window unshards FSDP2's params before the fold reads them, so under FSDP2 alone the fused experts arrive as the unit's whole unsharded weight and each rank folds every expert of the unit, one unit at a time; expert weights FSDP2 does not shard, as expert parallelism keeps them, stay sharded over the experts and each rank folds its own. The fold's comment no longer says an FSDP2 fused weight is always a sharded DTensor there. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 5 ++++- modelopt/torch/export/moe_utils.py | 9 +++++---- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 1d2a7a5073b..18ace3bbb87 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -256,7 +256,10 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: # FSDP2 owns the fused weights here: fold them inside a weight-update window, which hands # FSDP2 the packed weights. Replaced behind its back, they would be overwritten with its - # original shards the next time the layer reshards, and exported unquantized. + # original shards the next time the layer reshards, and exported unquantized. The window + # unshards FSDP2's params, so under FSDP2 alone each rank folds every expert of the unit, one + # unit at a time as a forward holds it; expert weights FSDP2 does not shard, as expert + # parallelism keeps them, stay sharded over the experts, and each rank folds its own. with _fsdp2_weight_window(ctx, module): _export_fused_experts_keep_fused(module, ctx.dtype) elif ctx.keep_fused_experts: diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 46b76ddd680..a7157237c9b 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -278,10 +278,11 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N weight_quantizers = getattr(module, f"{proj}_weight_quantizers") input_quantizer = getattr(module, f"{proj}_input_quantizer", None) - # Under FSDP2 the fused 3-D weight is a DTensor sharded on dim 0 (the expert axis). Quantize - # ONLY this rank's local experts and rewrap the result sharded identically -- indexing the - # DTensor (weight[i]) would all-gather the whole fused weight onto every rank, de-sharding - # the model and OOMing the GPU at scale. Plain tensor (single process): all experts. + # A fused 3-D weight sharded on dim 0 (the expert axis) is a DTensor: quantize ONLY this + # rank's local experts and rewrap the result sharded identically -- indexing the DTensor + # (weight[i]) would all-gather the whole fused weight onto every rank, de-sharding the model + # and OOMing the GPU at scale. A plain tensor (a single process, or FSDP2's unsharded weight + # inside a weight window): all experts. wdata = weight.data is_dt = DTensor is not None and isinstance(wdata, DTensor) quantizers_are_global = True From da2c12f3aeb12322d7165cdbf3025ea9b7fd7e8f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 07:59:32 +0000 Subject: [PATCH 27/67] Replace zero NVFP4 block scales with a scalar, not a tensor of ones torch.where takes the 1.0 as a Python scalar, which DTensor dispatches as it does the tensor form (checked on torch 2.8 and 2.14), so the zero-scale replacement no longer makes a full-size tensor of ones for every packed weight. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/qtensor/nvfp4_tensor.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 4d46e203341..3ac450c2908 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -165,9 +165,7 @@ def get_weights_scaling_factor_from_quantizer( per_block_scale = per_block_amax / E2M1_MAX # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. - per_block_scale = torch.where( - per_block_scale == 0, torch.ones_like(per_block_scale), per_block_scale - ) + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure num_blocks_per_row = weight.shape[-1] // block_size @@ -221,9 +219,7 @@ def get_weights_scaling_factor( # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is # pointwise, so DTensor handles it, and the result is identical for a plain tensor. - per_block_scale = torch.where( - per_block_scale == 0, torch.ones_like(per_block_scale), per_block_scale - ) + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: per_block_scale = _cast_per_block_scale_to_fp8(per_block_scale) From 3ed0377edcd18efb110f348f62d87305f478be50 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 08:00:17 +0000 Subject: [PATCH 28/67] Reuse the model's module names for a window outside every FSDP2 unit A module no FSDP2 unit encloses opens its window on the model itself, whose names the export context already holds; the window no longer builds the model-wide map again for it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 18ace3bbb87..c01424637ca 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -68,7 +68,8 @@ def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, Module """ from torch.distributed.fsdp import FSDPModule - if isinstance(ctx.model, FSDPModule): + if isinstance(ctx.model, FSDPModule) or unit is ctx.model: + # The model's own names, which _enclosing_unit has built. return ctx.model, ctx.module_names return unit, module_name_maps(unit) From b0e1038131037d69337d116d5a94b82099ddb304 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:53:04 +0000 Subject: [PATCH 29/67] Say that a weight-update window reads its module list again on exit The export's held windows append each module of an FSDP2 unit to the list they passed to fsdp2_aware_weight_update after entering it, and rely on that list being read again when the window closes. A copy taken at entry would leave every appended module's packed weights unregistered with FSDP2, to be overwritten by its original shards at the next reshard. The docstring now states the contract. Signed-off-by: Shengliang Xu --- modelopt/torch/quantization/utils/core_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b4120720b5c..0915a9d1720 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1054,7 +1054,9 @@ def fsdp2_aware_weight_update( Args: root_model (nn.Module): The root model of the FSDPModule. modules_to_update (list): The list of modules to update which should be a list of modules that are - direct children of the FSDPModule. + direct children of the FSDPModule. A list is read again when the context exits, so a caller + may append modules of the same FSDP unit to it inside the context, and their weights are + re-registered with FSDP2 too (the export's held weight windows rely on this). reshard (bool): Whether to reshard the FSDP root module after the weight update. Returns: From 62a0893464f95d2f40e58fd48add7ea344c797c5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:53:13 +0000 Subject: [PATCH 30/67] Refuse GGML formats in the keep-fused fold before folding anything A GGML format packs each expert's weight into its blocks and registers no weight_scale, which the fold stacks per expert: the fold failed with an AttributeError midway, possibly after replacing the first projection. It is refused with the fold's other unsupported layouts now, before any projection is folded and on every rank alike; the per-expert export still writes it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 12 ++++++++++++ .../export/test_fused_experts_keep_fused.py | 19 ++++++++++++++++--- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index a7157237c9b..e9262314aa4 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -218,6 +218,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). """ + from modelopt.torch.export.quant_format import GGML_FORMATS from modelopt.torch.export.unified_export_hf import _export_quantized_weight try: @@ -259,6 +260,17 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N group=wdata.device_mesh.get_group(mdim), ) enabled = [e for _, shard in sorted(gathered) for e in shard] + if any(enabled) and any( + getattr(q, "num_bits", None) in GGML_FORMATS for q in weight_quantizers + ): + # A GGML format packs each expert's weight into its blocks and registers no weight_scale + # for the fold to stack. Decided on the gathered states and the recipe's format, so every + # rank refuses alike. + raise NotImplementedError( + f"{type(module).__name__}.{proj}: GGML formats pack each expert's weight with no " + "weight_scale, which the keep-fused fold cannot stack; only the per-expert export " + "writes them." + ) if all(enabled): if not is_dt and not len(weight_quantizers) == wdata.shape[0] == n: raise NotImplementedError( diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index bb308ff1c27..233afec3f03 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -55,6 +55,7 @@ ) from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.qtensor import NVFP4QTensor from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan @@ -416,6 +417,17 @@ def _fewer_experts_counted(experts): experts.num_experts = NUM_EXPERTS - 1 +def _ggml(experts): + for i in range(len(experts.down_proj_weight_quantizers)): + experts.down_proj_weight_quantizers[i] = TensorQuantizer( + QuantizerAttributeConfig( + num_bits="q8_0", + block_sizes={-1: GGML_FORMAT_REGISTRY["q8_0"].block_size}, + backend="ggml", + ) + ) + + @LAYOUTS @pytest.mark.parametrize( ("change", "match"), @@ -424,14 +436,15 @@ def _fewer_experts_counted(experts): (_native_fp8, "native FP8"), (_native_fp8_unquantized, "native FP8"), (_fewer_experts_counted, "num_experts is 3, but the unsharded weight holds 4 experts"), + (_ggml, "GGML formats"), ], - ids=["mixed", "native-fp8", "native-fp8-unquantized", "num-experts-short"], + ids=["mixed", "native-fp8", "native-fp8-unquantized", "num-experts-short", "ggml"], ) def test_a_projection_the_fold_cannot_write_is_refused_untouched(experts_cls, change, match): """One fused tensor cannot mix quantized and unquantized experts, and native FP8 weights need the scale_inv the fold drops, also where a recipe leaves them unquantized. A ``num_experts`` short of - the weight's rows would fold only the first experts. The refusal comes before the other - projection folds.""" + the weight's rows would fold only the first experts, and a GGML format packs each expert's weight + with no scale to stack. The refusal comes before the other projection folds.""" experts = _quantized_experts(experts_cls, mtq.FP8_DEFAULT_CFG) change(experts) before = {k: v.clone() for k, v in experts.state_dict().items()} From 71b2564c4f17ec468ee93ffc52a9f30c15a1adb4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:52:09 +0000 Subject: [PATCH 31/67] Say which reads of the projection the local rows stand in for The rows replace self. where get_expert_weight reads it itself. A method it calls is bound to the module and would read the whole fused DTensor, so an accessor written that way would gather it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index e9262314aa4..52b4d71a685 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -136,6 +136,8 @@ def _local_expert_weight(module: nn.Module, proj: str, local: torch.Tensor, idx: FP8 experts, held at the default dtype, to the export dtype. The accessor indexes the fused parameter itself, a sharded DTensor here, so it is handed this rank's rows instead. Whatever else it reads, such as how it picks this rank's experts out of a weight, is the module's own. + The rows stand in for ``self.`` only where the accessor reads it itself: a method it calls + is bound to the module, and would read the whole fused DTensor. """ return type(module).get_expert_weight(_LocalRows(module, proj, local), proj, idx, dtype) From 906c355b1e043b2ba48073513f1235c860d19d34 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 20:29:57 +0000 Subject: [PATCH 32/67] Run the keep-fused tests' ungated case on the real NemotronHExperts The ungated stand-in relied on the generic plugin detecting any up_proj-only container. The real class converts through whichever wrapper quantizes Nemotron-H's experts: that detection today, a registration of their own once Nemotron-H's experts get one. The fixture leaves a class the registry covers through its base class alone, since unregistering it would fail. Signed-off-by: Shengliang Xu --- .../export/test_fused_experts_keep_fused.py | 47 +++++++++++-------- 1 file changed, 28 insertions(+), 19 deletions(-) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 233afec3f03..ed8d76377cb 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -42,6 +42,8 @@ pytest.importorskip("transformers") pytest.importorskip("safetensors") from safetensors.torch import load_file, save_file +from transformers import NemotronHConfig +from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq @@ -94,25 +96,24 @@ def forward(self, hidden_states, top_k_index, top_k_weights): return out -class _KeepFusedUngatedExperts(nn.Module): - """Ungated fused experts (NemotronH): a single ``up_proj`` and ``down_proj``.""" +class _KeepFusedUngatedExperts(NemotronHExperts): + """Nemotron-H's ungated fused experts: a single ``up_proj`` and ``down_proj``. - def __init__(self, num_experts=NUM_EXPERTS): - super().__init__() - self.num_experts = num_experts - self.up_proj = nn.Parameter(torch.randn(num_experts, INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02) - self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) - self.act_fn = nn.SiLU() + The real class, on its eager forward, so that whichever wrapper ModelOpt quantizes Nemotron-H's + experts with converts it. + """ - def forward(self, hidden_states, top_k_index, top_k_weights): - out = torch.zeros_like(hidden_states) - mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) - for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: - top_k_pos, token_idx = torch.where(mask[expert_idx]) - up = self.act_fn(F.linear(hidden_states[token_idx], self.up_proj[expert_idx])) - current = F.linear(up, self.down_proj[expert_idx]) - out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) - return out + def __init__(self, num_experts=NUM_EXPERTS): + config = NemotronHConfig( + n_routed_experts=num_experts, + hidden_size=HIDDEN_DIM, + moe_intermediate_size=INTERMEDIATE_DIM, + ) + config._experts_implementation = "eager" + super().__init__(config) + with torch.no_grad(): + self.up_proj.normal_(std=0.02) + self.down_proj.normal_(std=0.02) LAYOUTS = pytest.mark.parametrize( @@ -124,9 +125,17 @@ def forward(self, hidden_states, top_k_index, top_k_weights): @pytest.fixture(autouse=True) def _unregister_synthetic_experts(): - """``mtq.quantize`` registers the synthetic classes on the fly; keep the registry clean.""" + """``mtq.quantize`` registers the synthetic classes on the fly; keep the registry clean. + + A class the registry already covers, through its base class's registration, stays as it is. + """ + on_the_fly = [ + cls + for cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) + if cls not in QuantModuleRegistry + ] yield - for cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts): + for cls in on_the_fly: if QuantModuleRegistry.get(cls) is not None: QuantModuleRegistry.unregister(cls) From eebbcbe3a0e8ea279de25e838f4bf7932d9b0de3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:06:14 +0000 Subject: [PATCH 33/67] Refuse BMM experts in the keep-fused export, whose split would cut them wrong The writer splits a keep-fused export's experts. weights into per-expert keys as (num_experts, out_dim, in_dim). GPT-OSS's and Llama-4's BMM experts are named experts too but store theirs transposed (GPT-OSS also interleaves gate and up and keeps biases), so the split would write every expert wrong, quantized or not. Their handler now refuses the keep-fused export before packing anything; the gather export writes them fused, as their checkpoints are. Nothing else packs several weights of one module into the held window, so the window's once-per-module check goes too. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 15 ++++++-- .../export/test_fused_experts_keep_fused.py | 37 ++++++------------- 2 files changed, 22 insertions(+), 30 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index c01424637ca..2498b59aa55 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -89,10 +89,7 @@ def _fsdp2_weight_window( return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: - # A module with several quantized weights (fused BMM experts) packs each in turn: it joins - # the window once, or the window would re-register its weights with FSDP2 twice. - if not any(m is module for m in held.modules): - held.modules.append(module) + held.modules.append(module) return contextlib.nullcontext() _close_fsdp2_weight_window(ctx) root, names = _window_root(ctx, unit) @@ -325,6 +322,16 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non Its weight quantizer has one amax covering all experts, so under FSDP2 each rank can pack just the experts it owns and produce identical bytes -- no gather or unshard needed. """ + if ctx.keep_fused_experts: + # The keep-fused export leaves fused experts for the writer to split by their + # ``experts.`` key, as ``(num_experts, out_dim, in_dim)`` weights. These store theirs + # transposed (GPT-OSS also interleaves gate and up, and keeps biases), so the split would cut + # every expert wrong, quantized or not. + raise NotImplementedError( + f"{name}: the keep-fused export splits fused experts stored (num_experts, out_dim, " + f"in_dim), and {type(module).__name__} stores its experts transposed; export it through " + "the gather export instead." + ) if get_quantization_format(module) == QUANTIZATION_NONE: return # TODO: consolidate uncalibrated experts handling logic diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index ed8d76377cb..fcfa4e1f35c 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -848,37 +848,22 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( class GptOssExperts(nn.Module): - """Named as transformers' fused BMM experts, whose handler packs each of its weights in turn.""" + """Named as transformers' fused BMM experts, which store their weights transposed.""" -def test_bmm_experts_pack_both_weights_in_their_unit_s_held_window(gloo_world_of_one): - """``gate_up_proj`` and then ``down_proj`` of one module pack in the unit's held window, which - the module joins once: listed twice, the window would re-register its weights with FSDP2 a - second time, off the shards the first pass left.""" - torch.manual_seed(0) - linears = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 32, bias=False)) - mtq.quantize(linears, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) +def test_keep_fused_refuses_bmm_experts(): + """The writer splits a keep-fused export's ``experts.`` weights as ``(num_experts, out_dim, + in_dim)``; BMM experts store theirs transposed, so the split would cut every expert wrong, and + the export refuses them before packing anything.""" experts = GptOssExperts() - for proj, linear in zip(("gate_up_proj", "down_proj"), linears): - experts.register_parameter(proj, nn.Parameter(linear.weight.detach().clone())) - setattr(experts, f"{proj}_weight_quantizer", linear.weight_quantizer) - setattr(experts, f"{proj}_input_quantizer", linear.input_quantizer) - ref = copy.deepcopy(experts) - ExportModuleRegistry.match(ref)(PREFIX, ref, ExportContext(model=ref, dtype=torch.bfloat16)) - layer = nn.Sequential(experts) - model = nn.Sequential(layer) - fully_shard(layer, mesh=gloo_world_of_one) - set_distribution_plan(model, DistributionPlan(dense="fsdp")) - ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + experts.gate_up_proj = nn.Parameter(torch.randn(NUM_EXPERTS, HIDDEN_DIM, 2 * INTERMEDIATE_DIM)) + before = copy.deepcopy(experts.state_dict()) + ctx = ExportContext(model=experts, dtype=torch.bfloat16, keep_fused_experts=True) - with handlers.held_fsdp2_weight_windows(ctx): - ExportModuleRegistry.match(experts)("0.0", experts, ctx) + with pytest.raises(NotImplementedError, match="stores its experts transposed"): + ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) - sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) - for key, want in ref.state_dict().items(): - got = sd[f"0.0.{key}"] - got = got.full_tensor() if isinstance(got, DTensor) else got - assert torch.equal(_raw_bytes(got), _raw_bytes(want)), key + torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): From 00804c5e964c5dd0999f3f8a8b00bf8d04edf783 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:06:30 +0000 Subject: [PATCH 34/67] Say what an unknown expert count means to the local split, not that it is legacy Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 52b4d71a685..7794c933c58 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -825,7 +825,8 @@ def _split_local_fused_module( only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``. * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base. - ``total_experts == 0`` keeps the legacy per-EP-group behavior. + ``total_experts == 0``, a count the caller could not read off the model's config, takes the + tensor for a per-EP-group slice. """ from torch.distributed.tensor import DTensor @@ -841,7 +842,7 @@ def _split_local_fused_module( ) group_size = int(fw.shape[0]) if fw is not None else 0 # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or - # total is unknown (legacy). When it spans all experts, dp_offset is already global. + # total is unknown (0). When it spans all experts, dp_offset is already global. add_ep_base = ( bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts) ) From d242e3ca072bd27a4893e2e5c16a8eb09c9bdb53 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:06:47 +0000 Subject: [PATCH 35/67] Say export handlers pack through _export_weight, inside the held window Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 63a104fbd43..7b4dbb5f596 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -636,6 +636,12 @@ def _export_quantized_weight( The export includes converting weight tensor to correct quantized values and quantized dtype, and registering scaling factors. + An export handler packs through ``hf_export_handlers._export_weight`` rather than calling this + directly: on the keep-fused export an FSDP2-owned weight must be packed inside its unit's + weight-update window, which that helper (or the fused-experts handler, around its fold) opens. + Packed outside it, the unit's next reshard puts FSDP2's unquantized shard back, and the + checkpoint is written unquantized. + Tied-weight dedup is not handled here: both sides of a tie are packed independently (identically, once ``sync_tied_input_amax`` has equalized their scales), and the duplicate is dropped by name in :func:`postprocess_state_dict`. Deduping per-module at From d8f9a0d198fd9db58e4d62f9ef95d457a30eef39 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:45:12 +0000 Subject: [PATCH 36/67] Say why the fold's GGML refusal needs no gathered format The format comes from a recipe that sets it by module name, so it is the same for every expert of a projection on every rank, and the ranks refuse alike without gathering it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 7794c933c58..74888cffb26 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -266,8 +266,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N getattr(q, "num_bits", None) in GGML_FORMATS for q in weight_quantizers ): # A GGML format packs each expert's weight into its blocks and registers no weight_scale - # for the fold to stack. Decided on the gathered states and the recipe's format, so every - # rank refuses alike. + # for the fold to stack. Decided on the gathered states and on the format, which a recipe + # sets by module name and so is the same for every expert on every rank: every rank + # refuses alike. raise NotImplementedError( f"{type(module).__name__}.{proj}: GGML formats pack each expert's weight with no " "weight_scale, which the keep-fused fold cannot stack; only the per-expert export " From 9b063357de6f390e12f1846561990067aa55bb22 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:45:26 +0000 Subject: [PATCH 37/67] Say why the fold's replicated input_scale agrees across ranks It is the one value the fold writes replicated rather than sharded; calibration syncs the shared input quantizer's amax over the data-parallel ranks and, under expert parallelism, over the expert-parallel group, so every rank registers the same value. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 74888cffb26..7951a430c24 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -371,6 +371,9 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N local_ws2 = _rewrap_expert_rows(local_ws2, mesh, scale_placements, wdata.shape[0]) module.register_buffer(f"{proj}_weight_scale_2", local_ws2) if input_scale is not None: + # Replicated, not sharded like the rest: calibration syncs the shared input quantizer's + # amax over the data-parallel ranks and, under expert parallelism, over the + # expert-parallel group (max_calibrate), so every rank registers the same value. module.register_buffer(f"{proj}_input_scale", input_scale) # A projection left as it is goes out in the weight domain the per-expert export writes it in. From 60c3d81d88f40a5d5040bfefa787df5665abe653 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 21:45:40 +0000 Subject: [PATCH 38/67] Drop the split's unreachable checks for a missing fused weight fw comes from a direct subscript of the state dict, which raises on a missing key, so it is always a tensor and the None arms are dead. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 7951a430c24..7962307a248 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -839,12 +839,8 @@ def _split_local_fused_module( fw = sharded_sd[f"{prefix}.{first}"] # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns. dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0 - n_local = ( - fw.to_local().shape[0] - if isinstance(fw, DTensor) - else (fw.shape[0] if fw is not None else 0) - ) - group_size = int(fw.shape[0]) if fw is not None else 0 + n_local = fw.to_local().shape[0] if isinstance(fw, DTensor) else fw.shape[0] + group_size = int(fw.shape[0]) # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or # total is unknown (0). When it spans all experts, dp_offset is already global. add_ep_base = ( From 912e953645557cf2abae8a7f2108f13870fb3a73 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 23:04:23 +0000 Subject: [PATCH 39/67] Destroy the gloo fixture's group even when its mesh fails to build Signed-off-by: Shengliang Xu --- tests/unit/torch/conftest.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index ec11dd90999..de576f629c4 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -63,5 +63,7 @@ def gloo_world_of_one(skip_on_windows, tmp_path_factory): 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,)) - torch.distributed.destroy_process_group() + try: + yield init_device_mesh("cpu", (1,)) + finally: + torch.distributed.destroy_process_group() From 0a40404f83b1414a03d924009a2ea1670fb87589 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 00:10:33 +0000 Subject: [PATCH 40/67] Let the meta-DCP load tests use the shared gloo_world_of_one fixture Their module-local copy shadowed the one this PR adds to the torch unit conftest. Deleting it is the same change #2741 makes, so whichever lands first, the other merges cleanly. Signed-off-by: Shengliang Xu --- tests/unit/torch/utils/test_meta_dcp_load.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 764bce545c6..4c642b1644c 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -474,15 +474,6 @@ def recording_load(state_dict, **kwargs): # ------------------------------------------------------------------- under a process group -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group, so the loader takes its distributed path.""" - store = tmp_path_factory.mktemp("pg") / "store" - torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield - torch.distributed.destroy_process_group() - - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the From b8be3139a31efc949700a9be070faaaefcb0bf91 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:27:46 +0000 Subject: [PATCH 41/67] Import the keep-fused fold's torch.distributed and GGML names at module scope moe_utils imports DTensor, Replicate, Shard and GGML_FORMATS at the top, and hf_export_handlers FSDPModule: none breaks a cycle, and with torch 2.10 the DTensor = None fallback cannot happen, so its is-not-None checks go too. _export_quantized_weight stays local, as it breaks one, and so does the private compute_local_shape_and_global_offset fallback. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 3 +-- modelopt/torch/export/moe_utils.py | 23 ++++++--------------- 2 files changed, 7 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 2498b59aa55..fbd1fa5475e 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -21,6 +21,7 @@ from dataclasses import dataclass import torch.nn as nn +from torch.distributed.fsdp import FSDPModule from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps from modelopt.torch.quantization.utils.core_utils import ModuleNames, _get_enclosing_fsdp_module @@ -66,8 +67,6 @@ def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, Module codes next to correct scales). So there that unit is handed in as the root, with names relative to it, which the window resolves the modules' parameters against. """ - from torch.distributed.fsdp import FSDPModule - if isinstance(ctx.model, FSDPModule) or unit is ctx.model: # The model's own names, which _enclosing_unit has built. return ctx.model, ctx.module_names diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 7962307a248..780d0cec023 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -21,6 +21,9 @@ import torch import torch.nn as nn +from torch.distributed.tensor import DTensor, Replicate, Shard + +from modelopt.torch.export.quant_format import GGML_FORMATS def _delete_fused_moe_source_attrs(module: nn.Module) -> None: @@ -112,8 +115,6 @@ def _fallback_amax(weight: torch.Tensor, quantizer) -> torch.Tensor: def _rewrap_expert_rows(local: torch.Tensor, mesh, placements, num_experts: int): """Wrap this rank's experts as a DTensor spanning ``num_experts`` rows on the expert axis.""" - from torch.distributed.tensor import DTensor - shape = torch.Size((num_experts, *local.shape[1:])) stride = torch.empty(shape, device="meta").stride() return DTensor.from_local(local, mesh, placements, run_check=False, shape=shape, stride=stride) @@ -166,8 +167,6 @@ def _refuse_unfoldable_shards(module: nn.Module, proj: str, weight, num_quantize and their quantizer counts differ only where each keeps its own experts' quantizers of an uneven shard, which the checks below refuse on every rank. """ - from torch.distributed.tensor import Shard - name = f"{type(module).__name__}.{proj}" sharded = [placement for placement in weight.placements if isinstance(placement, Shard)] if len(sharded) > 1 or any(placement.dim != 0 for placement in sharded): @@ -220,14 +219,8 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). """ - from modelopt.torch.export.quant_format import GGML_FORMATS from modelopt.torch.export.unified_export_hf import _export_quantized_weight - try: - from torch.distributed.tensor import DTensor, Replicate, Shard - except Exception: # pragma: no cover - non-distributed install - DTensor = None # noqa: N806 (the optional-import fallback for a class) - first = getattr(module, "_first_proj_attr", "gate_up_proj") n = module.num_experts # One fused tensor cannot hold quantized and unquantized experts side by side, so a projection @@ -247,7 +240,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N "stored, since the fold drops their scale_inv; dequantize them first " "(unpack_weight())." ) - is_dt = DTensor is not None and isinstance(wdata, DTensor) + is_dt = isinstance(wdata, DTensor) if is_dt: _refuse_unfoldable_shards(module, proj, wdata, len(weight_quantizers)) if len(weight_quantizers) != wdata.shape[0]: @@ -299,7 +292,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N # and OOMing the GPU at scale. A plain tensor (a single process, or FSDP2's unsharded weight # inside a weight window): all experts. wdata = weight.data - is_dt = DTensor is not None and isinstance(wdata, DTensor) + is_dt = isinstance(wdata, DTensor) quantizers_are_global = True if is_dt: mesh, placements = wdata.device_mesh, wdata.placements @@ -381,7 +374,7 @@ def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> N if proj in folded: continue wdata = getattr(module, proj).data - is_dt = DTensor is not None and isinstance(wdata, DTensor) + is_dt = isinstance(wdata, DTensor) local = wdata.to_local() if is_dt else wdata domain = _local_expert_weight(module, proj, local, 0, dtype).dtype if domain != local.dtype: @@ -791,8 +784,6 @@ def _slice_weight_scale(ws, expert_idx, row_slice, rows): def _dtensor_dim0_offset(dt) -> int: """Global dim-0 start index of this rank's local shard of a ``Shard(0)`` DTensor.""" - from torch.distributed.tensor import Shard - mesh, placements = dt.device_mesh, dt.placements try: from torch.distributed.tensor._utils import compute_local_shape_and_global_offset @@ -832,8 +823,6 @@ def _split_local_fused_module( ``total_experts == 0``, a count the caller could not read off the model's config, takes the tensor for a per-EP-group slice. """ - from torch.distributed.tensor import DTensor - gated = f"{prefix}.gate_up_proj" in sharded_sd first = "gate_up_proj" if gated else "up_proj" fw = sharded_sd[f"{prefix}.{first}"] From 298223055224203c29344c4a1d8aaeaa8dacbad0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:40:58 +0000 Subject: [PATCH 42/67] Leave the sharded-packing tests to [1/2] and the local split to #2752 The held-window and NVFP4 DTensor tests are in [1/2]'s test_sharded_weight_packing.py now, so their copies here go. _split_local_fused_module, which renumbers one rank's split experts globally, moves to #2752 with the writer that calls it, as the review asked, and its two tests go with it. The fold keeps its single-rank tests and the refusal of a weight sharded other than by expert. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 92 +----- .../export/test_fused_experts_keep_fused.py | 263 +----------------- 2 files changed, 6 insertions(+), 349 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 780d0cec023..da5122a59ea 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -669,7 +669,7 @@ def split_fused_experts_state_dict(state_dict: dict) -> dict: checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and the input keeps its fused tensors, so the experts are held twice until it returns: pass one - module's shard at a time, as :func:`_split_local_fused_module` does, not a whole model's. + module's shard at a time, not a whole model's. Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` @@ -797,93 +797,3 @@ def _dtensor_dim0_offset(dt) -> int: nshards = mesh.size(mdim) return mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) return 0 - - -def _split_local_fused_module( - prefix: str, sharded_sd: dict, ep_rank: int = 0, total_experts: int = 0 -) -> dict: - """Split this rank's local shard of a fused-experts module into per-expert keys. - - Split THIS rank's *local* shard of one fused-experts module into per-expert keys with GLOBAL - expert indices (plain, local tensors -- no all-gather). - - FSDP2/EP shard the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank - owns is whole and local. We reuse :func:`split_fused_experts_state_dict` on the ``to_local()`` - shard (which numbers experts ``0..E_local-1``) and then shift each expert index by the rank's - global dim-0 offset, so the resulting per-expert keys match the single-process export and stay - on this rank for a direct distributed write (no consolidation). - - The fused expert tensor can arrive in two EP layouts, distinguished by its dim-0 size vs - ``total_experts`` (the config expert count): - * **per-EP-group** (size < total): transformers pre-sliced the param to this rank's EP group - (FSDP x EP -> a dp-sharded DTensor of the group; or a plain ep-shard). ``_dtensor_dim0_offset`` - only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``. - * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts - (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base. - ``total_experts == 0``, a count the caller could not read off the model's config, takes the - tensor for a per-EP-group slice. - """ - gated = f"{prefix}.gate_up_proj" in sharded_sd - first = "gate_up_proj" if gated else "up_proj" - fw = sharded_sd[f"{prefix}.{first}"] - # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns. - dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0 - n_local = fw.to_local().shape[0] if isinstance(fw, DTensor) else fw.shape[0] - group_size = int(fw.shape[0]) - # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or - # total is unknown (0). When it spans all experts, dp_offset is already global. - add_ep_base = ( - bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts) - ) - offset = dp_offset + (ep_rank * group_size if add_ep_base else 0) - - def _loc(key): - v = sharded_sd.get(key) - if v is None: - return None - return v.to_local() if isinstance(v, DTensor) else v - - def _loc_scale(key): - # Weight scales are dim-0 (per-expert) tensors. A DTensor scale shards with the weight -> - # to_local() aligns. The keep-fused fold emits DTensor scales for a sharded weight, but a - # scale handed in as a PLAIN tensor over every expert of the group is accepted too: slice - # it by the weight's within-group dp_offset so each scale pairs with its own dp-local expert. - v = sharded_sd.get(key) - if v is None: - return None - if isinstance(v, DTensor): - return v.to_local() - if v.dim() >= 1 and v.shape[0] > n_local and dp_offset + n_local <= v.shape[0]: - return v[dp_offset : dp_offset + n_local] - return v - - # Build a fused sub-dict of this rank's local experts (plain), then split + reindex. - local_fused: dict = { - f"{prefix}.{first}": _loc(f"{prefix}.{first}"), - f"{prefix}.down_proj": _loc(f"{prefix}.down_proj"), - } - for proj in (first, "down_proj"): - ws = _loc_scale(f"{prefix}.{proj}_weight_scale") - if ws is not None: - local_fused[f"{prefix}.{proj}_weight_scale"] = ws - # NVFP4 second-level per-tensor weight scale: also a per-expert (dim-0) tensor, so the same - # dp-offset slicing pairs each scalar with its dp-local expert. None for FP8. - ws2 = _loc_scale(f"{prefix}.{proj}_weight_scale_2") - if ws2 is not None: - local_fused[f"{prefix}.{proj}_weight_scale_2"] = ws2 - ins = sharded_sd.get(f"{prefix}.{proj}_input_scale") # shared scalar (replicated) - if ins is not None: - local_fused[f"{prefix}.{proj}_input_scale"] = ins - - split_local = split_fused_experts_state_dict( - local_fused - ) # keys {prefix}.{i}.* (i in 0..E_local-1) - if offset == 0: - return split_local - - plen = len(prefix) + 1 - out: dict = {} - for key, val in split_local.items(): - i_str, tail = key[plen:].split(".", 1) - out[f"{prefix}.{int(i_str) + offset}.{tail}"] = val - return out diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index fcfa4e1f35c..54687a436d2 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -17,27 +17,22 @@ The no-gather writer never splits a fused 3-D expert weight while the model is sharded: ``_export_fused_experts_keep_fused`` folds every expert's quantizer into the fused weight in place, -and ``split_fused_experts_state_dict`` / ``_split_local_fused_module`` cut the result into -per-expert keys later. That pipeline must write exactly what the single-process -``_export_fused_experts`` writes, so most tests compare the two byte for byte on a tiny synthetic -fused-experts module quantized with FP8 and NVFP4. Sharded cases run on a gloo process group: one -rank in process, and two spawned ranks where a non-zero expert offset matters. +and ``split_fused_experts_state_dict`` cuts the result into per-expert keys later. That pipeline +must write exactly what the single-process ``_export_fused_experts`` writes, so most tests compare +the two byte for byte on a tiny synthetic fused-experts module quantized with FP8 and NVFP4. +Sharded cases run on a one-rank gloo process group. """ import contextlib import copy -import math -import re import warnings import pytest import torch import torch.nn as nn import torch.nn.functional as F -from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.fsdp import FSDPModule, fully_shard -from torch.distributed.tensor import DTensor, Shard, distribute_tensor +from torch.distributed.tensor import Shard, distribute_tensor pytest.importorskip("transformers") pytest.importorskip("safetensors") @@ -48,18 +43,15 @@ import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( - _dtensor_dim0_offset, _export_fused_experts, _export_fused_experts_keep_fused, _fused_experts_prefixes, - _split_local_fused_module, split_fused_experts_state_dict, ) from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer -from modelopt.torch.quantization.qtensor import NVFP4QTensor from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan NUM_EXPERTS = 4 @@ -560,153 +552,9 @@ def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): ] -@pytest.mark.parametrize( - ("ep_rank", "total_experts", "first_index"), - [ - (0, 0, 0), - # This EP group's slice of a larger expert set: shift by the group base. - (1, 2 * LOCAL_EXPERTS, LOCAL_EXPERTS), - (3, 4 * LOCAL_EXPERTS, 3 * LOCAL_EXPERTS), - # The tensor already spans every expert, so its indices are global. - (1, LOCAL_EXPERTS, 0), - # Total unknown: the legacy per-group shift. - (1, 0, LOCAL_EXPERTS), - ], -) -def test_split_local_numbers_experts_globally(ep_rank, total_experts, first_index): - sd = _fused_state_dict(gated=True) - unshifted = split_fused_experts_state_dict(dict(sd)) - - out = _split_local_fused_module(PREFIX, sd, ep_rank=ep_rank, total_experts=total_experts) - - want = {} - for key, value in unshifted.items(): - expert, tail = key[len(PREFIX) + 1 :].split(".", 1) - want[f"{PREFIX}.{int(expert) + first_index}.{tail}"] = value - _assert_byte_identical(out, want) - - # ------------------------------------------------------------------ sharded (DTensor) experts -def _shard_fused_weights(experts, mesh): - for proj in (experts._first_proj_attr, "down_proj"): - weight = getattr(experts, proj).data - sharded = distribute_tensor(weight, mesh, [Shard(0)]) - setattr(experts, proj, nn.Parameter(sharded, requires_grad=False)) - return experts - - -def _two_rank_split_job(rank, world_size): - mesh = init_device_mesh("cpu", (world_size,)) - cases = [ - (experts_cls, quant_cfg, NUM_EXPERTS) - for experts_cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) - for quant_cfg in (mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG) - ] - # Three experts over two ranks: rank 0 holds two, rank 1 one. The folded tensors must keep the - # global expert count, not infer local x world rows. - cases.append((_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG, 3)) - for experts_cls, quant_cfg, num_experts in cases: - # Every rank calibrates the same seeded module, so all ranks hold identical quantizers. - experts = _quantized_experts(experts_cls, quant_cfg, num_experts) - want = _per_expert_export(experts) - plain = _keep_fused_export(experts) - sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) - _export_fused_experts_keep_fused(sharded, torch.bfloat16) - fused = _prefixed(sharded.state_dict()) - assert set(fused) == set(plain) - first = fused[f"{PREFIX}.{sharded._first_proj_attr}"] - n_local = first.to_local().shape[0] - offset = rank * math.ceil(num_experts / world_size) - assert _dtensor_dim0_offset(first) == offset - for key, value in fused.items(): - if key.endswith("_input_scale"): # one scalar shared by every expert - assert not isinstance(value, DTensor), key - assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key - else: # weights and per-expert scales stay sharded on the expert axis - assert isinstance(value, DTensor), key - assert value.placements == (Shard(0),), key - assert value.shape == plain[key].shape, key - local = plain[key][offset : offset + n_local] - assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(local)), key - mine = { - k: v - for k, v in want.items() - if offset <= int(k[len(PREFIX) + 1 :].split(".")[0]) < offset + n_local - } - assert mine, "this rank owns no experts" - - _assert_byte_identical(_split_local_fused_module(PREFIX, fused), mine) - - # A per-expert scale held as a plain tensor over ALL experts (as keep-fused emits it for an - # unsharded weight) must be cut to this rank's experts before pairing. - plain_scales = { - k: plain[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v - for k, v in fused.items() - } - _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) - - # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone - # failing while rank 0 goes on to write. - lone = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=1) - with pytest.raises(NotImplementedError, match="leaves a rank none"): - _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) - - # Expert parallelism leaves a rank only its own experts' quantizers, in shard order, under a - # fused weight spanning every expert. Kept last, with the joint decision below: without that - # pairing rank 1 alone fails, and rank 0 then waits only in the collectives below, which the - # worker pool ends once a rank has failed. - experts = _quantized_experts(_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG) - plain = _keep_fused_export(experts) - n_local = NUM_EXPERTS // world_size - mine = slice(rank * n_local, (rank + 1) * n_local) - - def sharded_keeping_quantizers(kept): - sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) - for proj in (sharded._first_proj_attr, "down_proj"): - name = f"{proj}_weight_quantizers" - setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[kept])) - return sharded - - with pytest.raises(NotImplementedError, match="3 weight quantizers for 4 experts in 2 shards"): - _export_fused_experts_keep_fused(sharded_keeping_quantizers(slice(0, 3)), torch.bfloat16) - sharded = sharded_keeping_quantizers(mine) - _export_fused_experts_keep_fused(sharded, torch.bfloat16) - fused = _prefixed(sharded.state_dict()) - assert set(fused) == set(plain) - for key, value in fused.items(): - if isinstance(value, DTensor): - assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(plain[key][mine])), key - else: - assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key - - # Each rank sees only its own experts' quantizers there, so the ranks decide on a projection - # together: rank 1 disabling some or all of its experts' is refused on both ranks, rather than - # folded on rank 0 alone, and every expert disabled leaves the projection as it is on both. - for disabled in ([3], [2, 3], [0, 1, 2, 3]): - sharded = sharded_keeping_quantizers(mine) - for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, start=mine.start): - if expert in disabled: - quantizer.disable() - if len(disabled) < NUM_EXPERTS: - with pytest.raises(NotImplementedError, match=re.escape(f"quantizers {disabled} of 4")): - _export_fused_experts_keep_fused(sharded, torch.bfloat16) - else: - _export_fused_experts_keep_fused(sharded, torch.bfloat16) - assert hasattr(sharded, "gate_up_proj_weight_scale") - assert not hasattr(sharded, "down_proj_weight_scale") - - -def test_two_ranks_fold_and_split_their_own_experts_under_global_indices(cpu_dist_workers): - """Rank 1 holds experts 2-3: each must be folded with ITS quantizer and keep its index, and the - fused tensors stay sharded on the expert axis at their global shape, an uneven shard included. - Its quantizers are every expert's (FSDP2) or, under expert parallelism, its own two, and then - the ranks decide on each projection's quantization together. A layout that leaves a rank no - expert, or quantizers that are neither, is refused on every rank.""" - cpu_dist_workers.run(_two_rank_split_job) - - @pytest.mark.parametrize( ("mesh_shape", "placements"), [((1,), [Shard(2)]), ((1, 1), [Shard(0), Shard(0)])], @@ -795,58 +643,6 @@ def window(root, modules, names=None): assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) -@pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) -@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) -def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( - monkeypatch, gloo_world_of_one, root_sharded, held -): - """On the no-gather path the model is still sharded when a linear is packed. - - Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the - DTensor's local view is refused under autograd), so the handlers pack inside a weight-update - window, on the layer's own unit where expert parallelism leaves the root unsharded. The window - gathers the whole unit, so within held windows it opens once for all of the unit's linears; - any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is - sharded again afterwards, and the sharded state dict carries the packed weights, as an - unsharded pack has them. - """ - torch.manual_seed(0) - layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) - model = nn.Sequential(layer) - mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) - refs = copy.deepcopy(model)[0] - for i, ref in enumerate(refs): - ExportModuleRegistry.match(ref)( - f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) - ) - fully_shard(layer, mesh=gloo_world_of_one) - if root_sharded: - fully_shard(model, mesh=gloo_world_of_one) - set_distribution_plan(model, DistributionPlan(dense="fsdp")) - ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) - unshards = [] - unshard = FSDPModule.unshard - monkeypatch.setattr( - FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) - ) - - with handlers.held_fsdp2_weight_windows(ctx) if held else contextlib.nullcontext(): - for i, linear in enumerate(layer): - ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - - assert unshards == ([layer] if held else [layer, layer]) - assert all(isinstance(linear.weight, DTensor) for linear in layer) - sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) - for i, ref in enumerate(refs): - packed = sd[f"0.{i}.weight"] - assert isinstance(packed, DTensor) - assert packed.dtype == torch.uint8 - assert torch.equal(packed.full_tensor(), ref.weight) - for name in ("weight_scale", "weight_scale_2", "input_scale"): - want = _raw_bytes(getattr(ref, name)) - assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name - - class GptOssExperts(nn.Module): """Named as transformers' fused BMM experts, which store their weights transposed.""" @@ -864,52 +660,3 @@ def test_keep_fused_refuses_bmm_experts(): ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) - - -def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): - """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. - - The block scales come out sharded like the weight, an all-zero block included: filling it by - boolean-mask assignment lowers to ``index_put_``, which DTensor cannot dispatch. - """ - torch.manual_seed(0) - weight = torch.randn(4, 32, dtype=torch.bfloat16) - weight[1, :16] = 0.0 # an all-zero block - scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) - scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] - want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data - - sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) - sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] - packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data - - assert isinstance(sharded_scale, DTensor) and sharded_scale.placements == (Shard(0),) - assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale)) - assert torch.equal(packed.full_tensor(), want) - # The second-level scale read off the sharded weight itself, as without one handed in. - assert torch.equal( - NVFP4QTensor.get_weights_scaling_factor(sharded, 16)[1], - NVFP4QTensor.get_weights_scaling_factor(weight, 16)[1], - ) - - -def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): - """``from_local`` would infer local x world rows; an uneven shard must keep the real total. - - Two rows held of a three-row tensor is what one rank holds when three rows split over two. - """ - local = torch.randn(2, 32) * 4 - sharded = DTensor.from_local( - local.clone(), - gloo_world_of_one, - [Shard(0)], - run_check=False, - shape=torch.Size([3, 32]), - stride=(32, 1), - ) - - packed = NVFP4QTensor._cast_fp4(sharded) - - assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) - assert packed.shape == (3, 32) - assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) From 136c76618b8cd1d634e1e59fc6b764c5309ca54f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:36:07 +0000 Subject: [PATCH 43/67] List a module in its unit's held weight window once, however many weights it packs A handler that packs several weights of one module asks for the window once per weight. Listed twice, the module was re-registered with FSDP2 a second time when the window closed, from the shard the first registration had put back, which FSDP2 rejects. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 5 +++- .../export/test_sharded_weight_packing.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index d28e7be022d..eaf0c5607db 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -88,7 +88,10 @@ def _fsdp2_weight_window( return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: - held.modules.append(module) + # Once per module, though a handler packing several weights asks once per weight: the + # window re-registers each listed module's weights with FSDP2 when it closes. + if not any(m is module for m in held.modules): + held.modules.append(module) return contextlib.nullcontext() _close_fsdp2_weight_window(ctx) root, names = _window_root(ctx, unit) diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 0e666784169..5a5a1856048 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -135,3 +135,29 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for name in ("weight_scale", "weight_scale_2", "input_scale"): want = _raw_bytes(getattr(ref, name)) assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name + + +def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): + """A handler packing several weights of one module asks for the window once per weight. + + The module must join its unit's held window once: listed twice, it would be re-registered with + FSDP2 a second time when the window closes, from the shard the first registration put back. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + ref = copy.deepcopy(layer[0]) + ExportModuleRegistry.match(ref)("0.0", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + fully_shard(layer, mesh=gloo_world_of_one) + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + ExportModuleRegistry.match(layer[0])("0.0", layer[0], ctx) + with handlers._fsdp2_weight_window(ctx, layer[0]): # as for the module's next weight + pass + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) From ebfb6a7b05026efd408edadfc01b8fa8dc71c99d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:37:31 +0000 Subject: [PATCH 44/67] Keep the sharded packing's comments to their constraint, and its imports at module scope The comments said what to avoid and then retold how each torch failure surfaced; that account is in the PR description. The static block-scale path says why it matches the dynamic one: its amax buffer is never a DTensor. DTensor and init_device_mesh are imported at module scope (torch>=2.10 has both), and ExportContext.fsdp2_window is typed as the held window it holds. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 16 ++--- modelopt/torch/export/registry.py | 5 +- .../quantization/qtensor/nvfp4_tensor.py | 65 ++++++------------- tests/unit/torch/conftest.py | 3 +- 4 files changed, 29 insertions(+), 60 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index eaf0c5607db..27d44bd072d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -60,12 +60,8 @@ def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: """The root a weight-update window on ``unit`` is opened on, and its module names. - The window opens on the innermost unit enclosing the module, but only when the root it is handed - is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the - dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which - packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated - codes next to correct scales). So there that unit is handed in as the root, with names relative - to it, which the window resolves the modules' parameters against. + A window does nothing on an unwrapped root, which expert parallelism leaves (it FSDP2-shards + only the dense layers), so there the unit itself is the root, with names relative to it. """ if isinstance(ctx.model, FSDPModule) or unit is ctx.model: # The model's own names, which _enclosing_unit has built. @@ -134,12 +130,8 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight - # The no-gather writer packs while the model is still sharded, so unlike the - # collect_export_tensors path there is no open weight-access window here and the weight is a - # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an - # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, - # which reshards the unit when it closes: a layer left unsharded holds all of its weights on - # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + # The no-gather export packs a still-sharded model: an FSDP2-owned weight goes in its unit's + # weight-update window, or the unit's next reshard puts FSDP2's unquantized shard back. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: with _fsdp2_weight_window(ctx, module): _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 808210a4f68..ca99532dfae 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,12 +28,13 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import torch import torch.nn as nn if TYPE_CHECKING: + from modelopt.torch.export.hf_export_handlers import _HeldWindow from modelopt.torch.quantization.utils.core_utils import ModuleNames __all__ = [ @@ -76,7 +77,7 @@ class ExportContext: only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which closes the last window. Elsewhere each module is packed in a window of its own.""" - fsdp2_window: Any = None + fsdp2_window: "_HeldWindow | None" = None """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 3ac450c2908..ab394b6667f 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -16,6 +16,7 @@ """Implements NVFP4 quantization for efficient tensor storage and computation.""" import torch +from torch.distributed.tensor import DTensor from ..backends.utils import fp4_compatible from ..qtensor.base_qtensor import BaseQuantizedTensor @@ -163,8 +164,7 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, - # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + # torch.where, to match the dynamic path below, which must avoid index_put_ on a DTensor. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure @@ -215,10 +215,8 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: - # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device - # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is - # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + # Set all zero values in scale to 1.0. torch.where, not a boolean-mask assignment: that lowers + # to index_put_, which DTensor (a sharded weight in the distributed export) cannot dispatch. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: @@ -253,33 +251,11 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" - # The distributed export path packs a SHARDED weight. Every step below is elementwise over the - # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound - # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and - # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and - # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time - # instead just moves the failure to the next one. - _dtensor_meta = None - try: - from torch.distributed.tensor import DTensor as _DTensor - - if isinstance(weight, _DTensor): - # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as - # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and - # gives each rank a different global shape. Downstream that desynchronises the - # exporter's per-rank tensor classification and surfaces as a KeyError deep in the - # writer, long after the real mistake. - _dtensor_meta = ( - _DTensor, - weight.device_mesh, - weight.placements, - weight.shape, - weight.stride(), - ) - weight = weight.to_local() - except ImportError: - pass - + # A DTensor is cast on its local shard: the steps below mix it with plain tensors, which + # DTensor refuses. The cast is elementwise, so rewrapping the result gives the same tensor. + sharded = weight if isinstance(weight, DTensor) else None + if sharded is not None: + weight = sharded.to_local() device = weight.device # Extract sign and compute absolute values in one pass @@ -299,17 +275,18 @@ def _cast_fp4(cls, weight: torch.Tensor): # Combine sign, ordinal, and rounding adjustment packed = (sign_bit << 3) + ord + equals_odd_bounds - if _dtensor_meta is not None: - _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta - packed = _dt_cls.from_local( - packed, - _mesh, - _placements, - run_check=False, - shape=_gshape, - stride=_gstride, - ) - return packed + if sharded is None: + return packed + # The global shape and stride as given: from_local would infer local x world, wrong for an + # uneven shard. + return DTensor.from_local( + packed, + sharded.device_mesh, + sharded.placements, + run_check=False, + shape=sharded.shape, + stride=sharded.stride(), + ) @classmethod def quantize( diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index de576f629c4..d751079bfd2 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -17,6 +17,7 @@ import pytest import torch +from torch.distributed.device_mesh import init_device_mesh from modelopt.torch.opt.mode import _ModeRegistryCls @@ -59,8 +60,6 @@ def use_channel_div_4(): def gloo_world_of_one(skip_on_windows, tmp_path_factory): """A one-rank gloo process group, so the code under test takes its distributed path, and its CPU device mesh.""" - from torch.distributed.device_mesh import init_device_mesh - store = tmp_path_factory.mktemp("pg") / "store" torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) try: From d34142120fddaee54ef7ceb6c37a2c51774fe926 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:57:57 +0000 Subject: [PATCH 45/67] Drop the BMM experts handler's claim that a rank packs its experts without a gather Every path that reaches this handler now packs whole weights: the gather export inside its gather window, and the keep-fused export refuses these experts. The docstring says that instead. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index a96bfa426f1..1418e8f0efb 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -313,8 +313,8 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) -> def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> None: """Export fused BMM-style expert weights (Llama4 / GPT-OSS). - Its weight quantizer has one amax covering all experts, so under FSDP2 each rank can pack - just the experts it owns and produce identical bytes -- no gather or unshard needed. + Their weight quantizer has one amax covering all experts. The keep-fused export refuses them, + since they store their experts transposed. """ if ctx.keep_fused_experts: # The keep-fused export leaves fused experts for the writer to split by their From 3a86980a421bb08d3d38845e0fbc87d4b13b48f2 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 01:36:38 +0000 Subject: [PATCH 46/67] Test the sharded fold on two ranks again, without the per-rank split The split took the two-rank test to #2752 with _split_local_fused_module, which left the shard-local fold without a check of its output. The test is back, checking each rank's fold against the unsharded one at its offset (_dtensor_dim0_offset), an uneven 3-over-2 shard included, and, gathered whole with full_tensor(), against the per-expert export. It keeps the expert-parallel quantizer pairing, the ranks' joint decision on a projection, the mismatched-count refusal and the refusal of a layout that leaves a rank no expert. gloo cannot gather FP8, so one-byte values are gathered as bytes. The native-FP8 test imports PretrainedConfig at module scope; FP8Experts stays behind importorskip, since transformers 5.5 imports triton in its module. Signed-off-by: Shengliang Xu --- .../export/test_fused_experts_keep_fused.py | 146 ++++++++++++++++-- 1 file changed, 137 insertions(+), 9 deletions(-) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 54687a436d2..5a7df9003ad 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -20,11 +20,14 @@ and ``split_fused_experts_state_dict`` cuts the result into per-expert keys later. That pipeline must write exactly what the single-process ``_export_fused_experts`` writes, so most tests compare the two byte for byte on a tiny synthetic fused-experts module quantized with FP8 and NVFP4. -Sharded cases run on a one-rank gloo process group. +Sharded cases run on a gloo process group: one rank in process, and two spawned ranks where a +non-zero expert offset matters. """ import contextlib import copy +import math +import re import warnings import pytest @@ -32,17 +35,18 @@ import torch.nn as nn import torch.nn.functional as F from torch.distributed.device_mesh import init_device_mesh -from torch.distributed.tensor import Shard, distribute_tensor +from torch.distributed.tensor import DTensor, Shard, distribute_tensor pytest.importorskip("transformers") pytest.importorskip("safetensors") from safetensors.torch import load_file, save_file -from transformers import NemotronHConfig +from transformers import NemotronHConfig, PretrainedConfig from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( + _dtensor_dim0_offset, _export_fused_experts, _export_fused_experts_keep_fused, _fused_experts_prefixes, @@ -286,8 +290,8 @@ def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_exp """Native FP8 experts unpack at the default dtype, and the per-expert export reads each expert in the export dtype (``get_expert_weight``): the fold reads them alike, for the scales it derives and for a projection it leaves unquantized.""" - from transformers import PretrainedConfig - from transformers.integrations.finegrained_fp8 import FP8Experts + # transformers 5.5 imports triton in this module, which not every platform has. + fp8_experts_cls = pytest.importorskip("transformers.integrations.finegrained_fp8").FP8Experts torch.manual_seed(0) config = PretrainedConfig( @@ -296,7 +300,7 @@ def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_exp num_experts=NUM_EXPERTS, hidden_act="silu", ) - experts = FP8Experts(config, block_size=None) + experts = fp8_experts_cls(config, block_size=None) for name in ("gate_up_proj", "down_proj"): weight = getattr(experts, name) with torch.no_grad(): @@ -304,7 +308,7 @@ def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_exp getattr(experts, f"{name}_scale_inv").fill_(0.01) hidden = torch.randn(NUM_EXPERTS, HIDDEN_DIM) route = (torch.arange(NUM_EXPERTS).unsqueeze(1), torch.ones(NUM_EXPERTS, 1)) - registered_before = QuantModuleRegistry.get(FP8Experts) + registered_before = QuantModuleRegistry.get(fp8_experts_cls) try: mtq.quantize(experts, mtq.NVFP4_DEFAULT_CFG, lambda m: m(hidden, *route)) experts.unpack_weight() @@ -312,8 +316,8 @@ def test_keep_fused_then_split_writes_unpacked_native_fp8_experts_as_the_per_exp want = _per_expert_export(experts) got = split_fused_experts_state_dict(_keep_fused_export(experts)) finally: - if registered_before is None and QuantModuleRegistry.get(FP8Experts) is not None: - QuantModuleRegistry.unregister(FP8Experts) + if registered_before is None and QuantModuleRegistry.get(fp8_experts_cls) is not None: + QuantModuleRegistry.unregister(fp8_experts_cls) _assert_byte_identical(got, want) @@ -555,6 +559,130 @@ def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): # ------------------------------------------------------------------ sharded (DTensor) experts +def _shard_fused_weights(experts, mesh): + for proj in (experts._first_proj_attr, "down_proj"): + weight = getattr(experts, proj).data + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + setattr(experts, proj, nn.Parameter(sharded, requires_grad=False)) + return experts + + +def _whole(value): + """A sharded value gathered whole. gloo cannot gather FP8, so one-byte values move as bytes.""" + if not isinstance(value, DTensor): + return value + if value.element_size() != 1: + return value.full_tensor() + raw = value.to_local().view(torch.uint8) + raw = DTensor.from_local( + raw, + value.device_mesh, + value.placements, + run_check=False, + shape=value.shape, + stride=value.stride(), + ) + return raw.full_tensor().view(value.dtype) + + +def _two_rank_fold_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + cases = [ + (experts_cls, quant_cfg, NUM_EXPERTS) + for experts_cls in (_KeepFusedGatedExperts, _KeepFusedUngatedExperts) + for quant_cfg in (mtq.FP8_DEFAULT_CFG, FP8_WEIGHT_ONLY_CFG, mtq.NVFP4_DEFAULT_CFG) + ] + # Three experts over two ranks: rank 0 holds two, rank 1 one. The folded tensors must keep the + # global expert count, not infer local x world rows. + cases.append((_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG, 3)) + for experts_cls, quant_cfg, num_experts in cases: + # Every rank calibrates the same seeded module, so all ranks hold identical quantizers. + experts = _quantized_experts(experts_cls, quant_cfg, num_experts) + want = _per_expert_export(experts) + plain = _keep_fused_export(experts) + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + first = fused[f"{PREFIX}.{sharded._first_proj_attr}"] + n_local = first.to_local().shape[0] + offset = rank * math.ceil(num_experts / world_size) + assert _dtensor_dim0_offset(first) == offset + for key, value in fused.items(): + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + else: # weights and per-expert scales stay sharded on the expert axis + assert isinstance(value, DTensor), key + assert value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + local = plain[key][offset : offset + n_local] + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(local)), key + # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. + whole = {key: _whole(value) for key, value in fused.items()} + _assert_byte_identical(split_fused_experts_state_dict(whole), want) + + # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone + # failing while rank 0 goes on to write. + lone = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG, num_experts=1) + with pytest.raises(NotImplementedError, match="leaves a rank none"): + _export_fused_experts_keep_fused(_shard_fused_weights(lone, mesh), torch.bfloat16) + + # Expert parallelism leaves a rank only its own experts' quantizers, in shard order, under a + # fused weight spanning every expert. Kept last, with the joint decision below: without that + # pairing rank 1 alone fails, and rank 0 then waits only in the collectives below, which the + # worker pool ends once a rank has failed. + experts = _quantized_experts(_KeepFusedGatedExperts, mtq.NVFP4_DEFAULT_CFG) + plain = _keep_fused_export(experts) + n_local = NUM_EXPERTS // world_size + mine = slice(rank * n_local, (rank + 1) * n_local) + + def sharded_keeping_quantizers(kept): + sharded = _shard_fused_weights(copy.deepcopy(experts), mesh) + for proj in (sharded._first_proj_attr, "down_proj"): + name = f"{proj}_weight_quantizers" + setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[kept])) + return sharded + + with pytest.raises(NotImplementedError, match="3 weight quantizers for 4 experts in 2 shards"): + _export_fused_experts_keep_fused(sharded_keeping_quantizers(slice(0, 3)), torch.bfloat16) + sharded = sharded_keeping_quantizers(mine) + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + fused = _prefixed(sharded.state_dict()) + assert set(fused) == set(plain) + for key, value in fused.items(): + if isinstance(value, DTensor): + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(plain[key][mine])), key + else: + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + + # Each rank sees only its own experts' quantizers there, so the ranks decide on a projection + # together: rank 1 disabling some or all of its experts' is refused on both ranks, rather than + # folded on rank 0 alone, and every expert disabled leaves the projection as it is on both. + for disabled in ([3], [2, 3], [0, 1, 2, 3]): + sharded = sharded_keeping_quantizers(mine) + for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, start=mine.start): + if expert in disabled: + quantizer.disable() + if len(disabled) < NUM_EXPERTS: + with pytest.raises(NotImplementedError, match=re.escape(f"quantizers {disabled} of 4")): + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + else: + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert hasattr(sharded, "gate_up_proj_weight_scale") + assert not hasattr(sharded, "down_proj_weight_scale") + + +def test_two_ranks_fold_their_own_experts_under_global_indices(cpu_dist_workers): + """Rank 1 holds experts 2-3: each must be folded with ITS quantizer at its global offset, and the + fused tensors stay sharded on the expert axis at their global shape, an uneven shard included; + gathered, they split into the per-expert export. Its quantizers are every expert's (FSDP2) or, + under expert parallelism, its own two, and then the ranks decide on each projection's + quantization together. A layout that leaves a rank no expert, or quantizers that are neither, + is refused on every rank.""" + cpu_dist_workers.run(_two_rank_fold_job) + + @pytest.mark.parametrize( ("mesh_shape", "placements"), [((1,), [Shard(2)]), ((1, 1), [Shard(0), Shard(0)])], From 54f8632c78bf8704b72e8550d9313c25c077faef Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:20:36 +0000 Subject: [PATCH 47/67] Hand FSDP2 the packed weight of a unit that owns it directly, and test the packing on two ranks Under an unwrapped root, a module that is an FSDP2 unit of its own is the window's root and owns its weight directly. get_prefixed_param_names named that weight's module after the weight itself, so the window's FSDPParam mapping held 'weight.weight' while its exit looked up '.weight': it skipped the packed weight, and the reshard put the float shard back. The owner of a parameter the root holds is now named '', and both sides join names alike. The handler test gains that layout, which fails without the fix. A two-rank job packs an uneven NVFP4 shard, three rows over two ranks, into the plain pack's codes and scales, and runs the handlers over two FSDP2 units, a 17-row weight among them, against an unsharded export. Signed-off-by: Shengliang Xu --- .../torch/quantization/utils/core_utils.py | 16 ++-- .../export/test_sharded_weight_packing.py | 96 +++++++++++++++++-- .../torch/quantization/test_param_index.py | 2 +- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 0915a9d1720..ca7168ac1da 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -942,7 +942,13 @@ def get_prefixed_param_names(parent_model, target_module, param_index=None): target_ids = {id(p) for p in target_module.parameters()} # Lowest position == first in named_parameters() order, matching a linear scan's result. match = min((param_index[pid] for pid in target_ids if pid in param_index), default=None) - return match[1].rsplit(".", 1)[0] if match is not None else None + # "" for a parameter the parent owns itself. + return match[1].rpartition(".")[0] if match is not None else None + + +def _param_key(module_name: str, param_name: str) -> str: + """The full name of ``module_name``'s parameter ``param_name``; ``""`` names the root module.""" + return f"{module_name}.{param_name}" if module_name else param_name def create_fsdp_param_mapping(fsdp_param_list, model): @@ -966,8 +972,7 @@ def create_fsdp_param_mapping(fsdp_param_list, model): if module_name is not None: # Get the parameter name from _module_info and construct full param name param_name = param._module_info.param_name - full_param_name = f"{module_name}.{param_name}" - mapping[full_param_name] = param + mapping[_param_key(module_name, param_name)] = param return mapping @@ -1096,7 +1101,7 @@ def fsdp2_aware_weight_update( module_name = _get_module_name(module, root_model, names) # Check if any parameter from this module is in the mapping module_params_in_mapping = any( - f"{module_name}.{n}" in fsdp_param_mapping + _param_key(module_name, n) in fsdp_param_mapping for n, _ in module.named_parameters() ) assert module_params_in_mapping, ( @@ -1112,8 +1117,7 @@ def fsdp2_aware_weight_update( # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): - name = _get_module_name(module, root_model, names) - name = f"{name}.{param_name}" + name = _param_key(_get_module_name(module, root_model, names), param_name) if name not in fsdp_param_mapping: continue diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 5a5a1856048..b61b11b8740 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +"""Packing weights that are still sharded, as the no-gather export does (CPU, gloo, one and two ranks).""" import contextlib import copy @@ -21,7 +21,9 @@ import pytest import torch import torch.nn as nn +from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor @@ -86,9 +88,9 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) @pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) -@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +@pytest.mark.parametrize("layout", ["fsdp2", "expert-parallel", "linear-units"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( - monkeypatch, gloo_world_of_one, root_sharded, held + monkeypatch, gloo_world_of_one, layout, held ): """On the no-gather path the model is still sharded when a linear is packed. @@ -98,7 +100,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( gathers the whole unit, so within held windows it opens once for all of the unit's linears; any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is sharded again afterwards, and the sharded state dict carries the packed weights, as an - unsharded pack has them. + unsharded pack has them. A linear that is a unit of its own under an unwrapped root owns its + weight as the window's root: its packed weight must be handed to FSDP2 all the same. """ torch.manual_seed(0) layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) @@ -109,8 +112,10 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( ExportModuleRegistry.match(ref)( f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) ) - fully_shard(layer, mesh=gloo_world_of_one) - if root_sharded: + units = list(layer) if layout == "linear-units" else [layer] + for unit in units: + fully_shard(unit, mesh=gloo_world_of_one) + if layout == "fsdp2": fully_shard(model, mesh=gloo_world_of_one) set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) @@ -124,7 +129,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for i, linear in enumerate(layer): ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - assert unshards == ([layer] if held else [layer, layer]) + per_linear = [linear if layout == "linear-units" else layer for linear in layer] + assert unshards == (list(dict.fromkeys(per_linear)) if held else per_linear) assert all(isinstance(linear.weight, DTensor) for linear in layer) sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) for i, ref in enumerate(refs): @@ -161,3 +167,79 @@ def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _two_rank_packing_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + torch.manual_seed(0) + # Three rows over two ranks, an all-zero block included: rank 0 packs two rows, rank 1 one. + weight = torch.randn(3, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(0, 2) if rank == 0 else slice(2, 3) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor(), want) + + # The handlers over two units in held windows, a 17-row weight sharded 9 + 8 among them. + model = _TwoUnits() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with bounded_collectives(), handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor(), ref.weight), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_two_ranks_pack_their_shards_into_the_unsharded_bytes(cpu_dist_workers): + """On two ranks a shard holds other rows than the whole weight: an uneven NVFP4 pack gives the + plain pack's codes and scales, and the handlers, run over two units, leave the sharded state + dict holding what an unsharded export packs.""" + cpu_dist_workers.run(_two_rank_packing_job) diff --git a/tests/unit/torch/quantization/test_param_index.py b/tests/unit/torch/quantization/test_param_index.py index a78dd7b8cff..54ade0a317a 100644 --- a/tests/unit/torch/quantization/test_param_index.py +++ b/tests/unit/torch/quantization/test_param_index.py @@ -28,7 +28,7 @@ def _linear_scan(parent_model, target_module): target_ids = {id(p) for p in target_module.parameters()} return next( ( - name.rsplit(".", 1)[0] + name.rpartition(".")[0] for name, param in parent_model.named_parameters() if id(param) in target_ids ), From 51f187fb0294431ea85697264439c3272a4c3a5b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:50:23 +0000 Subject: [PATCH 48/67] Test the sharded packing over NCCL on every visible GPU The two-rank gloo job, on CUDA: an uneven NVFP4 shard, two rows a rank and one on the last, packs into the plain pack's codes and scales, and the handlers, run over two FSDP2 units in held windows with keep_fused_experts, leave the sharded state dict holding what an unsharded export packs. It needs two GPUs, so a single-GPU runner skips it. Signed-off-by: Shengliang Xu --- .../export/test_sharded_weight_packing_gpu.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/gpu/torch/export/test_sharded_weight_packing_gpu.py diff --git a/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py new file mode 100644 index 00000000000..9284f755e48 --- /dev/null +++ b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does, over NCCL on every GPU.""" + +import copy + +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import Shard, distribute_tensor + +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu() + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _packing_job(rank, world_size): + mesh = init_device_mesh("cuda", (world_size,)) + torch.manual_seed(0) + # Two rows a rank, the last rank one, an all-zero block included. + rows = 2 * world_size - 1 + weight = torch.randn(rows, 32, dtype=torch.bfloat16, device="cuda") + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(2 * rank, min(2 * rank + 2, rows)) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor().cpu(), want.cpu()) + + # The handlers over two units in held windows, a 17-row weight sharded unevenly among them. + model = _TwoUnits().cuda() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32, device="cuda"))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor().cpu(), ref.weight.cpu()), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_sharded_packing_matches_the_unsharded_bytes_over_nccl(need_2_gpus, dist_workers): + """An uneven NVFP4 shard packs into the plain pack's codes and scales, and the handlers, run over + two FSDP2 units in held windows, leave the sharded state dict holding what an unsharded export + packs.""" + dist_workers.run(_packing_job) From af4c213ef640d785d2e7f73441d1e4ff5b2f01e8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:37:15 +0000 Subject: [PATCH 49/67] Pack NVFP4 weights that are DTensor shards The no-gather export packs weights while the model is still sharded. NVFP4's packing broke there in three ways: the block scales' zero fill used a boolean-mask assignment, which lowers to index_put_, which DTensor cannot dispatch; the fp4 cast mixed the DTensor with plain constants; and the nibble pack's << is a silent no-op on a DTensor. The scales use torch.where, the cast runs on the local shard and rewraps it with its global shape and stride, so an uneven shard keeps its real size, and the pack uses torch.bitwise_left_shift. Plain tensors pack exactly as before. The one-rank gloo world the meta-DCP load tests used moves to the shared conftest, for these tests too. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- .../quantization/qtensor/nvfp4_tensor.py | 56 ++++++++++++-- tests/unit/torch/conftest.py | 15 ++++ .../export/test_sharded_weight_packing.py | 74 +++++++++++++++++++ .../torch/quantization/test_nvfp4_tensor.py | 26 +++++++ tests/unit/torch/utils/test_meta_dcp_load.py | 9 --- 5 files changed, 165 insertions(+), 15 deletions(-) create mode 100644 tests/unit/torch/export/test_sharded_weight_packing.py diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 8726dc5c719..3ac450c2908 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -163,7 +163,9 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - per_block_scale[per_block_scale == 0] = 1.0 + # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, + # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure num_blocks_per_row = weight.shape[-1] // block_size @@ -213,8 +215,11 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0 - per_block_scale[per_block_scale == 0] = 1.0 + # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: + # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device + # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is + # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: per_block_scale = _cast_per_block_scale_to_fp8(per_block_scale) @@ -248,6 +253,33 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" + # The distributed export path packs a SHARDED weight. Every step below is elementwise over the + # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound + # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and + # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and + # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time + # instead just moves the failure to the next one. + _dtensor_meta = None + try: + from torch.distributed.tensor import DTensor as _DTensor + + if isinstance(weight, _DTensor): + # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as + # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and + # gives each rank a different global shape. Downstream that desynchronises the + # exporter's per-rank tensor classification and surfaces as a KeyError deep in the + # writer, long after the real mistake. + _dtensor_meta = ( + _DTensor, + weight.device_mesh, + weight.placements, + weight.shape, + weight.stride(), + ) + weight = weight.to_local() + except ImportError: + pass + device = weight.device # Extract sign and compute absolute values in one pass @@ -266,7 +298,18 @@ def _cast_fp4(cls, weight: torch.Tensor): ) # Combine sign, ordinal, and rounding adjustment - return (sign_bit << 3) + ord + equals_odd_bounds + packed = (sign_bit << 3) + ord + equals_odd_bounds + if _dtensor_meta is not None: + _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta + packed = _dt_cls.from_local( + packed, + _mesh, + _placements, + run_check=False, + shape=_gshape, + stride=_gstride, + ) + return packed @classmethod def quantize( @@ -347,8 +390,9 @@ def quantize( return scaled_weight # Cast weights to fp4 q_weight = cls._cast_fp4(scaled_weight) - # Pack weights - packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] + # Pack weights. torch.bitwise_left_shift, not ``<< 4``: on a DTensor the operator is a silent + # no-op (torch 2.11), which would pack unrelated codes on the distributed export path. + packed_weight = torch.bitwise_left_shift(q_weight[..., 1::2], 4) | q_weight[..., 0::2] return ( cls(input_shape, input_dtype, packed_weight), weights_scaling_factor, diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index 852412d2565..de576f629c4 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -16,6 +16,7 @@ from collections import defaultdict import pytest +import torch from modelopt.torch.opt.mode import _ModeRegistryCls @@ -52,3 +53,17 @@ def use_channel_div_4(): for name, kv_for_all in lookup.items(): for regex, (k, v) in kv_for_all.items(): config_class.model_fields[name].default[regex][k] = v + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group, so the code under test takes its distributed path, and its + CPU device mesh.""" + from torch.distributed.device_mesh import init_device_mesh + + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + try: + yield init_device_mesh("cpu", (1,)) + finally: + torch.distributed.destroy_process_group() diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py new file mode 100644 index 00000000000..a1c9fb19e58 --- /dev/null +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" + +import torch +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from modelopt.torch.quantization.qtensor import NVFP4QTensor + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): + """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. + + The block scales come out sharded like the weight, an all-zero block included: filling it by + boolean-mask assignment lowers to ``index_put_``, which DTensor cannot dispatch. + """ + torch.manual_seed(0) + weight = torch.randn(4, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 # an all-zero block + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + + sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert isinstance(sharded_scale, DTensor) and sharded_scale.placements == (Shard(0),) + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale)) + assert torch.equal(packed.full_tensor(), want) + # The second-level scale read off the sharded weight itself, as without one handed in. + assert torch.equal( + NVFP4QTensor.get_weights_scaling_factor(sharded, 16)[1], + NVFP4QTensor.get_weights_scaling_factor(weight, 16)[1], + ) + + +def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): + """``from_local`` would infer local x world rows; an uneven shard must keep the real total. + + Two rows held of a three-row tensor is what one rank holds when three rows split over two. + """ + local = torch.randn(2, 32) * 4 + sharded = DTensor.from_local( + local.clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 32]), + stride=(32, 1), + ) + + packed = NVFP4QTensor._cast_fp4(sharded) + + assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) + assert packed.shape == (3, 32) + assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index a7b64515e3a..ce1159711ba 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -106,6 +106,32 @@ def test_helper_clamps_underflow_to_min(self): out = _cast_per_block_scale_to_fp8(tiny).float() assert (out > 0).all(), f"FP8 cast produced zero scales: {out.tolist()}" + @staticmethod + def _dynamic_block_scales(): + weight = torch.zeros(1, 48) + weight[0, 16] = 6.0 # amax 6 at a second-level scale of 1: block scale exactly 1 + weight[0, 32] = 3.0 # amax 3: block scale 0.5 + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) + return scale.float()[0] + + @staticmethod + def _static_block_scales(): + q = SimpleNamespace( + block_sizes={-1: 16}, + global_amax=torch.tensor(6.0 * 448.0), # a second-level scale of 1 + _amax=torch.tensor([0.0, 6.0, 3.0]), + ) + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(q, torch.randn(1, 48)) + return scale.float()[0] + + @pytest.mark.parametrize("path", ["dynamic", "static"]) + def test_an_all_zero_block_gets_a_scale_of_one(self, path): + """Blocks of amax 0, 6 and 3: the zero block gets 1, as the amax-6 block does, not the clamp.""" + zero, unit, half = getattr(self, f"_{path}_block_scales")().tolist() + + assert zero == unit == 1.0 + assert half == 0.5 + def test_static_path_no_nan_when_block_amax_zero(self): """Static path: zero-amax block + small global_amax must clamp to 448, not cast to NaN.""" block_size = 16 diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 764bce545c6..4c642b1644c 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -474,15 +474,6 @@ def recording_load(state_dict, **kwargs): # ------------------------------------------------------------------- under a process group -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group, so the loader takes its distributed path.""" - store = tmp_path_factory.mktemp("pg") / "store" - torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield - torch.distributed.destroy_process_group() - - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the From 889d63073838928ca52bc9eacf5415c783fa2e56 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:38:33 +0000 Subject: [PATCH 50/67] Pack an FSDP2-sharded model's weights in held weight-update windows The no-gather export packs weights while the model is still sharded, with no weight-access window open. ExportContext.keep_fused_experts marks that path. On it, a quantized linear whose dense part FSDP2 shards is packed inside a weight-update window, which hands FSDP2 the packed weights; packed outside one, the unit's next reshard would put its unquantized shard back. Under held_fsdp2_weight_windows, which the export's dispatch loop now enters, a unit's window stays open from one module to the next, so a decoder layer is gathered once rather than once per linear; any other loop gets a window per module, closed by its handler. Where expert parallelism leaves the root unwrapped, the window opens on the layer's own unit, as on a wrapped root it would do nothing. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 98 +++++++++++++++++++ modelopt/torch/export/registry.py | 20 ++++ modelopt/torch/export/unified_export_hf.py | 14 ++- .../torch/quantization/utils/core_utils.py | 4 +- .../export/test_sharded_weight_packing.py | 63 ++++++++++++ 5 files changed, 195 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 200ec6e5dae..d28e7be022d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -16,9 +16,16 @@ """Built-in module handlers for unified Hugging Face export.""" import collections.abc +import contextlib import warnings +from dataclasses import dataclass import torch.nn as nn +from torch.distributed.fsdp import FSDPModule + +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps +from modelopt.torch.quantization.utils.core_utils import ModuleNames, _get_enclosing_fsdp_module +from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .moe_utils import _export_fused_experts @@ -34,6 +41,87 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") +@dataclass +class _HeldWindow: + """An FSDP2 unit's weight-update window, and the modules packed in it so far.""" + + unit: nn.Module + modules: list[nn.Module] + stack: contextlib.ExitStack + + +def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: + """The innermost FSDP2 unit enclosing ``module``, or the model if none does.""" + if ctx.module_names is None: + ctx.module_names = module_name_maps(ctx.model) + return _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + + +def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: + """The root a weight-update window on ``unit`` is opened on, and its module names. + + The window opens on the innermost unit enclosing the module, but only when the root it is handed + is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the + dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which + packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated + codes next to correct scales). So there that unit is handed in as the root, with names relative + to it, which the window resolves the modules' parameters against. + """ + if isinstance(ctx.model, FSDPModule) or unit is ctx.model: + # The model's own names, which _enclosing_unit has built. + return ctx.model, ctx.module_names + return unit, module_name_maps(unit) + + +def _fsdp2_weight_window( + ctx: ExportContext, module: nn.Module +) -> contextlib.AbstractContextManager: + """The ``fsdp2_aware_weight_update`` window to pack ``module``'s weights in. + + Inside :func:`held_fsdp2_weight_windows`, the window of the module's FSDP2 unit, held open + across the unit's modules. Anywhere else, a window of the module's own, closed when the handler + is done, so a dispatch loop that knows nothing of held windows cannot leave a unit unsharded. + """ + unit = _enclosing_unit(ctx, module) + if not ctx.hold_fsdp2_windows: + root, names = _window_root(ctx, unit) + return fsdp2_aware_weight_update(root, module, names=names) + held = ctx.fsdp2_window + if held is not None and held.unit is unit: + held.modules.append(module) + return contextlib.nullcontext() + _close_fsdp2_weight_window(ctx) + root, names = _window_root(ctx, unit) + held = _HeldWindow(unit, [module], contextlib.ExitStack()) + held.stack.enter_context(fsdp2_aware_weight_update(root, held.modules, names=names)) + ctx.fsdp2_window = held + return contextlib.nullcontext() + + +@contextlib.contextmanager +def held_fsdp2_weight_windows(ctx: ExportContext): + """Within the block, pack each FSDP2 unit's modules in one weight-update window. + + The window unshards the whole unit and reshards it when it closes, so a window per module + gathers a decoder layer once per quantized linear in it. Here the unit's window stays open from + one handler to the next until the export reaches another unit, each module joining the list + whose weights it re-registers with FSDP2, and leaving the block closes the last. + """ + ctx.hold_fsdp2_windows = True + try: + yield + finally: + ctx.hold_fsdp2_windows = False + _close_fsdp2_weight_window(ctx) + + +def _close_fsdp2_weight_window(ctx: ExportContext) -> None: + """Close the weight-update window the handlers hold open, if any, resharding its unit.""" + held, ctx.fsdp2_window = ctx.fsdp2_window, None + if held is not None: + held.stack.close() + + def _export_weight( module: nn.Module, ctx: ExportContext, @@ -43,6 +131,16 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight + # The no-gather writer packs while the model is still sharded, so unlike the + # collect_export_tensors path there is no open weight-access window here and the weight is a + # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an + # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, + # which reshards the unit when it closes: a layer left unsharded holds all of its weights on + # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: + with _fsdp2_weight_window(ctx, module): + _export_quantized_weight(module, ctx.dtype, weight_name) + return _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 1aa575d92de..808210a4f68 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,10 +28,14 @@ from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn +if TYPE_CHECKING: + from modelopt.torch.quantization.utils.core_utils import ModuleNames + __all__ = [ "ExportContext", "ExportHandler", @@ -59,6 +63,22 @@ class ExportContext: model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" + # Distributed (no-gather) export: the handlers pack while the model is still sharded, an + # FSDP2-owned weight inside its unit's weight-update window. + keep_fused_experts: bool = False + + module_names: "ModuleNames | None" = None + """``model``'s name maps, built once per export by the first handler that looks a module up + (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + + hold_fsdp2_windows: bool = False + """Whether the handlers hold an FSDP2 unit's weight-update window open across its modules: set + only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which + closes the last window. Elsewhere each module is packed in a window of its own.""" + + fsdp2_window: Any = None + """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 14f3c82fd36..3e849506ebb 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -104,7 +104,7 @@ export_sparse_attention_config = None # Importing the built-in handlers installs their entries in the two registries. -from . import hf_export_handlers as _hf_export_handlers # noqa: F401 +from . import hf_export_handlers as _hf_export_handlers from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import get_experts_list, is_layernorm, is_quantlinear, sync_moe_gate_up_amax from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model @@ -636,6 +636,12 @@ def _export_quantized_weight( The export includes converting weight tensor to correct quantized values and quantized dtype, and registering scaling factors. + An export handler packs through ``hf_export_handlers._export_weight`` rather than calling this + directly: on the keep-fused export an FSDP2-owned weight must be packed inside its unit's + weight-update window, which that helper opens. + Packed outside it, the unit's next reshard puts FSDP2's unquantized shard back, and the + checkpoint is written unquantized. + Tied-weight dedup is not handled here: both sides of a tie are packed independently (identically, once ``sync_tied_input_amax`` has equalized their scales), and the duplicate is dropped by name in :func:`postprocess_state_dict`. Deduping per-module at @@ -996,8 +1002,10 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - for name, sub_module in model.named_modules(): - _dispatch_export_handler(name, sub_module, ctx) + # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. + with _hf_export_handlers.held_fsdp2_weight_windows(ctx): + for name, sub_module in model.named_modules(): + _dispatch_export_handler(name, sub_module, ctx) def _prepare_model_for_export(model, dtype, is_modelopt_qlora): diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b4120720b5c..0915a9d1720 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1054,7 +1054,9 @@ def fsdp2_aware_weight_update( Args: root_model (nn.Module): The root model of the FSDPModule. modules_to_update (list): The list of modules to update which should be a list of modules that are - direct children of the FSDPModule. + direct children of the FSDPModule. A list is read again when the context exits, so a caller + may append modules of the same FSDP unit to it inside the context, and their weights are + re-registered with FSDP2 too (the export's held weight windows rely on this). reshard (bool): Whether to reshard the FSDP root module after the weight update. Returns: diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index a1c9fb19e58..0e666784169 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -15,10 +15,21 @@ """Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +import contextlib +import copy + +import pytest import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def _raw_bytes(tensor): @@ -72,3 +83,55 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) assert packed.shape == (3, 32) assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) + + +@pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) +@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( + monkeypatch, gloo_world_of_one, root_sharded, held +): + """On the no-gather path the model is still sharded when a linear is packed. + + Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the + DTensor's local view is refused under autograd), so the handlers pack inside a weight-update + window, on the layer's own unit where expert parallelism leaves the root unsharded. The window + gathers the whole unit, so within held windows it opens once for all of the unit's linears; + any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is + sharded again afterwards, and the sharded state dict carries the packed weights, as an + unsharded pack has them. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + refs = copy.deepcopy(model)[0] + for i, ref in enumerate(refs): + ExportModuleRegistry.match(ref)( + f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) + ) + fully_shard(layer, mesh=gloo_world_of_one) + if root_sharded: + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + unshards = [] + unshard = FSDPModule.unshard + monkeypatch.setattr( + FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) + ) + + with handlers.held_fsdp2_weight_windows(ctx) if held else contextlib.nullcontext(): + for i, linear in enumerate(layer): + ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) + + assert unshards == ([layer] if held else [layer, layer]) + assert all(isinstance(linear.weight, DTensor) for linear in layer) + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for i, ref in enumerate(refs): + packed = sd[f"0.{i}.weight"] + assert isinstance(packed, DTensor) + assert packed.dtype == torch.uint8 + assert torch.equal(packed.full_tensor(), ref.weight) + for name in ("weight_scale", "weight_scale_2", "input_scale"): + want = _raw_bytes(getattr(ref, name)) + assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name From 8dd716985196566376e86375774f58846cfa3838 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:36:07 +0000 Subject: [PATCH 51/67] List a module in its unit's held weight window once, however many weights it packs A handler that packs several weights of one module asks for the window once per weight. Listed twice, the module was re-registered with FSDP2 a second time when the window closed, from the shard the first registration had put back, which FSDP2 rejects. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 5 +++- .../export/test_sharded_weight_packing.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index d28e7be022d..eaf0c5607db 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -88,7 +88,10 @@ def _fsdp2_weight_window( return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: - held.modules.append(module) + # Once per module, though a handler packing several weights asks once per weight: the + # window re-registers each listed module's weights with FSDP2 when it closes. + if not any(m is module for m in held.modules): + held.modules.append(module) return contextlib.nullcontext() _close_fsdp2_weight_window(ctx) root, names = _window_root(ctx, unit) diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 0e666784169..5a5a1856048 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -135,3 +135,29 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for name in ("weight_scale", "weight_scale_2", "input_scale"): want = _raw_bytes(getattr(ref, name)) assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name + + +def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): + """A handler packing several weights of one module asks for the window once per weight. + + The module must join its unit's held window once: listed twice, it would be re-registered with + FSDP2 a second time when the window closes, from the shard the first registration put back. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + ref = copy.deepcopy(layer[0]) + ExportModuleRegistry.match(ref)("0.0", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + fully_shard(layer, mesh=gloo_world_of_one) + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + ExportModuleRegistry.match(layer[0])("0.0", layer[0], ctx) + with handlers._fsdp2_weight_window(ctx, layer[0]): # as for the module's next weight + pass + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) From 34b4f7c16bc04cbdd785acb7196f5b91b7c83873 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:37:31 +0000 Subject: [PATCH 52/67] Keep the sharded packing's comments to their constraint, and its imports at module scope The comments said what to avoid and then retold how each torch failure surfaced; that account is in the PR description. The static block-scale path says why it matches the dynamic one: its amax buffer is never a DTensor. DTensor and init_device_mesh are imported at module scope (torch>=2.10 has both), and ExportContext.fsdp2_window is typed as the held window it holds. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 16 ++--- modelopt/torch/export/registry.py | 5 +- .../quantization/qtensor/nvfp4_tensor.py | 65 ++++++------------- tests/unit/torch/conftest.py | 3 +- 4 files changed, 29 insertions(+), 60 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index eaf0c5607db..27d44bd072d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -60,12 +60,8 @@ def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: """The root a weight-update window on ``unit`` is opened on, and its module names. - The window opens on the innermost unit enclosing the module, but only when the root it is handed - is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the - dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which - packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated - codes next to correct scales). So there that unit is handed in as the root, with names relative - to it, which the window resolves the modules' parameters against. + A window does nothing on an unwrapped root, which expert parallelism leaves (it FSDP2-shards + only the dense layers), so there the unit itself is the root, with names relative to it. """ if isinstance(ctx.model, FSDPModule) or unit is ctx.model: # The model's own names, which _enclosing_unit has built. @@ -134,12 +130,8 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight - # The no-gather writer packs while the model is still sharded, so unlike the - # collect_export_tensors path there is no open weight-access window here and the weight is a - # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an - # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, - # which reshards the unit when it closes: a layer left unsharded holds all of its weights on - # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + # The no-gather export packs a still-sharded model: an FSDP2-owned weight goes in its unit's + # weight-update window, or the unit's next reshard puts FSDP2's unquantized shard back. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: with _fsdp2_weight_window(ctx, module): _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 808210a4f68..ca99532dfae 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,12 +28,13 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import torch import torch.nn as nn if TYPE_CHECKING: + from modelopt.torch.export.hf_export_handlers import _HeldWindow from modelopt.torch.quantization.utils.core_utils import ModuleNames __all__ = [ @@ -76,7 +77,7 @@ class ExportContext: only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which closes the last window. Elsewhere each module is packed in a window of its own.""" - fsdp2_window: Any = None + fsdp2_window: "_HeldWindow | None" = None """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 3ac450c2908..ab394b6667f 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -16,6 +16,7 @@ """Implements NVFP4 quantization for efficient tensor storage and computation.""" import torch +from torch.distributed.tensor import DTensor from ..backends.utils import fp4_compatible from ..qtensor.base_qtensor import BaseQuantizedTensor @@ -163,8 +164,7 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, - # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + # torch.where, to match the dynamic path below, which must avoid index_put_ on a DTensor. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure @@ -215,10 +215,8 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: - # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device - # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is - # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + # Set all zero values in scale to 1.0. torch.where, not a boolean-mask assignment: that lowers + # to index_put_, which DTensor (a sharded weight in the distributed export) cannot dispatch. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: @@ -253,33 +251,11 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" - # The distributed export path packs a SHARDED weight. Every step below is elementwise over the - # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound - # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and - # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and - # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time - # instead just moves the failure to the next one. - _dtensor_meta = None - try: - from torch.distributed.tensor import DTensor as _DTensor - - if isinstance(weight, _DTensor): - # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as - # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and - # gives each rank a different global shape. Downstream that desynchronises the - # exporter's per-rank tensor classification and surfaces as a KeyError deep in the - # writer, long after the real mistake. - _dtensor_meta = ( - _DTensor, - weight.device_mesh, - weight.placements, - weight.shape, - weight.stride(), - ) - weight = weight.to_local() - except ImportError: - pass - + # A DTensor is cast on its local shard: the steps below mix it with plain tensors, which + # DTensor refuses. The cast is elementwise, so rewrapping the result gives the same tensor. + sharded = weight if isinstance(weight, DTensor) else None + if sharded is not None: + weight = sharded.to_local() device = weight.device # Extract sign and compute absolute values in one pass @@ -299,17 +275,18 @@ def _cast_fp4(cls, weight: torch.Tensor): # Combine sign, ordinal, and rounding adjustment packed = (sign_bit << 3) + ord + equals_odd_bounds - if _dtensor_meta is not None: - _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta - packed = _dt_cls.from_local( - packed, - _mesh, - _placements, - run_check=False, - shape=_gshape, - stride=_gstride, - ) - return packed + if sharded is None: + return packed + # The global shape and stride as given: from_local would infer local x world, wrong for an + # uneven shard. + return DTensor.from_local( + packed, + sharded.device_mesh, + sharded.placements, + run_check=False, + shape=sharded.shape, + stride=sharded.stride(), + ) @classmethod def quantize( diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index de576f629c4..d751079bfd2 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -17,6 +17,7 @@ import pytest import torch +from torch.distributed.device_mesh import init_device_mesh from modelopt.torch.opt.mode import _ModeRegistryCls @@ -59,8 +60,6 @@ def use_channel_div_4(): def gloo_world_of_one(skip_on_windows, tmp_path_factory): """A one-rank gloo process group, so the code under test takes its distributed path, and its CPU device mesh.""" - from torch.distributed.device_mesh import init_device_mesh - store = tmp_path_factory.mktemp("pg") / "store" torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) try: From 2da1b3f9446808aac7da59f7d04a5b69ef266375 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:20:36 +0000 Subject: [PATCH 53/67] Hand FSDP2 the packed weight of a unit that owns it directly, and test the packing on two ranks Under an unwrapped root, a module that is an FSDP2 unit of its own is the window's root and owns its weight directly. get_prefixed_param_names named that weight's module after the weight itself, so the window's FSDPParam mapping held 'weight.weight' while its exit looked up '.weight': it skipped the packed weight, and the reshard put the float shard back. The owner of a parameter the root holds is now named '', and both sides join names alike. The handler test gains that layout, which fails without the fix. A two-rank job packs an uneven NVFP4 shard, three rows over two ranks, into the plain pack's codes and scales, and runs the handlers over two FSDP2 units, a 17-row weight among them, against an unsharded export. Signed-off-by: Shengliang Xu --- .../torch/quantization/utils/core_utils.py | 16 ++-- .../export/test_sharded_weight_packing.py | 96 +++++++++++++++++-- .../torch/quantization/test_param_index.py | 2 +- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 0915a9d1720..ca7168ac1da 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -942,7 +942,13 @@ def get_prefixed_param_names(parent_model, target_module, param_index=None): target_ids = {id(p) for p in target_module.parameters()} # Lowest position == first in named_parameters() order, matching a linear scan's result. match = min((param_index[pid] for pid in target_ids if pid in param_index), default=None) - return match[1].rsplit(".", 1)[0] if match is not None else None + # "" for a parameter the parent owns itself. + return match[1].rpartition(".")[0] if match is not None else None + + +def _param_key(module_name: str, param_name: str) -> str: + """The full name of ``module_name``'s parameter ``param_name``; ``""`` names the root module.""" + return f"{module_name}.{param_name}" if module_name else param_name def create_fsdp_param_mapping(fsdp_param_list, model): @@ -966,8 +972,7 @@ def create_fsdp_param_mapping(fsdp_param_list, model): if module_name is not None: # Get the parameter name from _module_info and construct full param name param_name = param._module_info.param_name - full_param_name = f"{module_name}.{param_name}" - mapping[full_param_name] = param + mapping[_param_key(module_name, param_name)] = param return mapping @@ -1096,7 +1101,7 @@ def fsdp2_aware_weight_update( module_name = _get_module_name(module, root_model, names) # Check if any parameter from this module is in the mapping module_params_in_mapping = any( - f"{module_name}.{n}" in fsdp_param_mapping + _param_key(module_name, n) in fsdp_param_mapping for n, _ in module.named_parameters() ) assert module_params_in_mapping, ( @@ -1112,8 +1117,7 @@ def fsdp2_aware_weight_update( # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): - name = _get_module_name(module, root_model, names) - name = f"{name}.{param_name}" + name = _param_key(_get_module_name(module, root_model, names), param_name) if name not in fsdp_param_mapping: continue diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 5a5a1856048..b61b11b8740 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +"""Packing weights that are still sharded, as the no-gather export does (CPU, gloo, one and two ranks).""" import contextlib import copy @@ -21,7 +21,9 @@ import pytest import torch import torch.nn as nn +from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor @@ -86,9 +88,9 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) @pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) -@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +@pytest.mark.parametrize("layout", ["fsdp2", "expert-parallel", "linear-units"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( - monkeypatch, gloo_world_of_one, root_sharded, held + monkeypatch, gloo_world_of_one, layout, held ): """On the no-gather path the model is still sharded when a linear is packed. @@ -98,7 +100,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( gathers the whole unit, so within held windows it opens once for all of the unit's linears; any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is sharded again afterwards, and the sharded state dict carries the packed weights, as an - unsharded pack has them. + unsharded pack has them. A linear that is a unit of its own under an unwrapped root owns its + weight as the window's root: its packed weight must be handed to FSDP2 all the same. """ torch.manual_seed(0) layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) @@ -109,8 +112,10 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( ExportModuleRegistry.match(ref)( f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) ) - fully_shard(layer, mesh=gloo_world_of_one) - if root_sharded: + units = list(layer) if layout == "linear-units" else [layer] + for unit in units: + fully_shard(unit, mesh=gloo_world_of_one) + if layout == "fsdp2": fully_shard(model, mesh=gloo_world_of_one) set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) @@ -124,7 +129,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for i, linear in enumerate(layer): ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - assert unshards == ([layer] if held else [layer, layer]) + per_linear = [linear if layout == "linear-units" else layer for linear in layer] + assert unshards == (list(dict.fromkeys(per_linear)) if held else per_linear) assert all(isinstance(linear.weight, DTensor) for linear in layer) sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) for i, ref in enumerate(refs): @@ -161,3 +167,79 @@ def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _two_rank_packing_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + torch.manual_seed(0) + # Three rows over two ranks, an all-zero block included: rank 0 packs two rows, rank 1 one. + weight = torch.randn(3, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(0, 2) if rank == 0 else slice(2, 3) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor(), want) + + # The handlers over two units in held windows, a 17-row weight sharded 9 + 8 among them. + model = _TwoUnits() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with bounded_collectives(), handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor(), ref.weight), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_two_ranks_pack_their_shards_into_the_unsharded_bytes(cpu_dist_workers): + """On two ranks a shard holds other rows than the whole weight: an uneven NVFP4 pack gives the + plain pack's codes and scales, and the handlers, run over two units, leave the sharded state + dict holding what an unsharded export packs.""" + cpu_dist_workers.run(_two_rank_packing_job) diff --git a/tests/unit/torch/quantization/test_param_index.py b/tests/unit/torch/quantization/test_param_index.py index a78dd7b8cff..54ade0a317a 100644 --- a/tests/unit/torch/quantization/test_param_index.py +++ b/tests/unit/torch/quantization/test_param_index.py @@ -28,7 +28,7 @@ def _linear_scan(parent_model, target_module): target_ids = {id(p) for p in target_module.parameters()} return next( ( - name.rsplit(".", 1)[0] + name.rpartition(".")[0] for name, param in parent_model.named_parameters() if id(param) in target_ids ), From d0cef073454724aa9a4c685f5657a2658da2bd17 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:50:23 +0000 Subject: [PATCH 54/67] Test the sharded packing over NCCL on every visible GPU The two-rank gloo job, on CUDA: an uneven NVFP4 shard, two rows a rank and one on the last, packs into the plain pack's codes and scales, and the handlers, run over two FSDP2 units in held windows with keep_fused_experts, leave the sharded state dict holding what an unsharded export packs. It needs two GPUs, so a single-GPU runner skips it. Signed-off-by: Shengliang Xu --- .../export/test_sharded_weight_packing_gpu.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/gpu/torch/export/test_sharded_weight_packing_gpu.py diff --git a/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py new file mode 100644 index 00000000000..9284f755e48 --- /dev/null +++ b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does, over NCCL on every GPU.""" + +import copy + +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import Shard, distribute_tensor + +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu() + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _packing_job(rank, world_size): + mesh = init_device_mesh("cuda", (world_size,)) + torch.manual_seed(0) + # Two rows a rank, the last rank one, an all-zero block included. + rows = 2 * world_size - 1 + weight = torch.randn(rows, 32, dtype=torch.bfloat16, device="cuda") + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(2 * rank, min(2 * rank + 2, rows)) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor().cpu(), want.cpu()) + + # The handlers over two units in held windows, a 17-row weight sharded unevenly among them. + model = _TwoUnits().cuda() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32, device="cuda"))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor().cpu(), ref.weight.cpu()), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_sharded_packing_matches_the_unsharded_bytes_over_nccl(need_2_gpus, dist_workers): + """An uneven NVFP4 shard packs into the plain pack's codes and scales, and the handlers, run over + two FSDP2 units in held windows, leave the sharded state dict holding what an unsharded export + packs.""" + dist_workers.run(_packing_job) From 8674b73d7ceebe696ae040d97385cff6f6526665 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:37:15 +0000 Subject: [PATCH 55/67] Pack NVFP4 weights that are DTensor shards The no-gather export packs weights while the model is still sharded. NVFP4's packing broke there in three ways: the block scales' zero fill used a boolean-mask assignment, which lowers to index_put_, which DTensor cannot dispatch; the fp4 cast mixed the DTensor with plain constants; and the nibble pack's << is a silent no-op on a DTensor. The scales use torch.where, the cast runs on the local shard and rewraps it with its global shape and stride, so an uneven shard keeps its real size, and the pack uses torch.bitwise_left_shift. Plain tensors pack exactly as before. The one-rank gloo world the meta-DCP load tests used moves to the shared conftest, for these tests too. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- .../quantization/qtensor/nvfp4_tensor.py | 56 ++++++++++++-- tests/unit/torch/conftest.py | 15 ++++ .../export/test_sharded_weight_packing.py | 74 +++++++++++++++++++ .../torch/quantization/test_nvfp4_tensor.py | 26 +++++++ tests/unit/torch/utils/test_meta_dcp_load.py | 9 --- 5 files changed, 165 insertions(+), 15 deletions(-) create mode 100644 tests/unit/torch/export/test_sharded_weight_packing.py diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 8726dc5c719..3ac450c2908 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -163,7 +163,9 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - per_block_scale[per_block_scale == 0] = 1.0 + # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, + # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure num_blocks_per_row = weight.shape[-1] // block_size @@ -213,8 +215,11 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0 - per_block_scale[per_block_scale == 0] = 1.0 + # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: + # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device + # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is + # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: per_block_scale = _cast_per_block_scale_to_fp8(per_block_scale) @@ -248,6 +253,33 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" + # The distributed export path packs a SHARDED weight. Every step below is elementwise over the + # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound + # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and + # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and + # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time + # instead just moves the failure to the next one. + _dtensor_meta = None + try: + from torch.distributed.tensor import DTensor as _DTensor + + if isinstance(weight, _DTensor): + # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as + # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and + # gives each rank a different global shape. Downstream that desynchronises the + # exporter's per-rank tensor classification and surfaces as a KeyError deep in the + # writer, long after the real mistake. + _dtensor_meta = ( + _DTensor, + weight.device_mesh, + weight.placements, + weight.shape, + weight.stride(), + ) + weight = weight.to_local() + except ImportError: + pass + device = weight.device # Extract sign and compute absolute values in one pass @@ -266,7 +298,18 @@ def _cast_fp4(cls, weight: torch.Tensor): ) # Combine sign, ordinal, and rounding adjustment - return (sign_bit << 3) + ord + equals_odd_bounds + packed = (sign_bit << 3) + ord + equals_odd_bounds + if _dtensor_meta is not None: + _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta + packed = _dt_cls.from_local( + packed, + _mesh, + _placements, + run_check=False, + shape=_gshape, + stride=_gstride, + ) + return packed @classmethod def quantize( @@ -347,8 +390,9 @@ def quantize( return scaled_weight # Cast weights to fp4 q_weight = cls._cast_fp4(scaled_weight) - # Pack weights - packed_weight = (q_weight[..., 1::2] << 4) | q_weight[..., 0::2] + # Pack weights. torch.bitwise_left_shift, not ``<< 4``: on a DTensor the operator is a silent + # no-op (torch 2.11), which would pack unrelated codes on the distributed export path. + packed_weight = torch.bitwise_left_shift(q_weight[..., 1::2], 4) | q_weight[..., 0::2] return ( cls(input_shape, input_dtype, packed_weight), weights_scaling_factor, diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index 852412d2565..de576f629c4 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -16,6 +16,7 @@ from collections import defaultdict import pytest +import torch from modelopt.torch.opt.mode import _ModeRegistryCls @@ -52,3 +53,17 @@ def use_channel_div_4(): for name, kv_for_all in lookup.items(): for regex, (k, v) in kv_for_all.items(): config_class.model_fields[name].default[regex][k] = v + + +@pytest.fixture +def gloo_world_of_one(skip_on_windows, tmp_path_factory): + """A one-rank gloo process group, so the code under test takes its distributed path, and its + CPU device mesh.""" + from torch.distributed.device_mesh import init_device_mesh + + store = tmp_path_factory.mktemp("pg") / "store" + torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) + try: + yield init_device_mesh("cpu", (1,)) + finally: + torch.distributed.destroy_process_group() diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py new file mode 100644 index 00000000000..a1c9fb19e58 --- /dev/null +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" + +import torch +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +from modelopt.torch.quantization.qtensor import NVFP4QTensor + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8) + + +def test_nvfp4_packs_a_sharded_weight_into_the_same_codes(gloo_world_of_one): + """``DTensor << int`` is a silent no-op, so the nibble pack must not rely on the operator. + + The block scales come out sharded like the weight, an all-zero block included: filling it by + boolean-mask assignment lowers to ``index_put_``, which DTensor cannot dispatch. + """ + torch.manual_seed(0) + weight = torch.randn(4, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 # an all-zero block + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + + sharded = distribute_tensor(weight, gloo_world_of_one, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert isinstance(sharded_scale, DTensor) and sharded_scale.placements == (Shard(0),) + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale)) + assert torch.equal(packed.full_tensor(), want) + # The second-level scale read off the sharded weight itself, as without one handed in. + assert torch.equal( + NVFP4QTensor.get_weights_scaling_factor(sharded, 16)[1], + NVFP4QTensor.get_weights_scaling_factor(weight, 16)[1], + ) + + +def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one): + """``from_local`` would infer local x world rows; an uneven shard must keep the real total. + + Two rows held of a three-row tensor is what one rank holds when three rows split over two. + """ + local = torch.randn(2, 32) * 4 + sharded = DTensor.from_local( + local.clone(), + gloo_world_of_one, + [Shard(0)], + run_check=False, + shape=torch.Size([3, 32]), + stride=(32, 1), + ) + + packed = NVFP4QTensor._cast_fp4(sharded) + + assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) + assert packed.shape == (3, 32) + assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) diff --git a/tests/unit/torch/quantization/test_nvfp4_tensor.py b/tests/unit/torch/quantization/test_nvfp4_tensor.py index a7b64515e3a..ce1159711ba 100644 --- a/tests/unit/torch/quantization/test_nvfp4_tensor.py +++ b/tests/unit/torch/quantization/test_nvfp4_tensor.py @@ -106,6 +106,32 @@ def test_helper_clamps_underflow_to_min(self): out = _cast_per_block_scale_to_fp8(tiny).float() assert (out > 0).all(), f"FP8 cast produced zero scales: {out.tolist()}" + @staticmethod + def _dynamic_block_scales(): + weight = torch.zeros(1, 48) + weight[0, 16] = 6.0 # amax 6 at a second-level scale of 1: block scale exactly 1 + weight[0, 32] = 3.0 # amax 3: block scale 0.5 + scale, _ = NVFP4QTensor.get_weights_scaling_factor(weight, 16, torch.tensor(1.0)) + return scale.float()[0] + + @staticmethod + def _static_block_scales(): + q = SimpleNamespace( + block_sizes={-1: 16}, + global_amax=torch.tensor(6.0 * 448.0), # a second-level scale of 1 + _amax=torch.tensor([0.0, 6.0, 3.0]), + ) + scale, _ = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(q, torch.randn(1, 48)) + return scale.float()[0] + + @pytest.mark.parametrize("path", ["dynamic", "static"]) + def test_an_all_zero_block_gets_a_scale_of_one(self, path): + """Blocks of amax 0, 6 and 3: the zero block gets 1, as the amax-6 block does, not the clamp.""" + zero, unit, half = getattr(self, f"_{path}_block_scales")().tolist() + + assert zero == unit == 1.0 + assert half == 0.5 + def test_static_path_no_nan_when_block_amax_zero(self): """Static path: zero-amax block + small global_amax must clamp to 448, not cast to NaN.""" block_size = 16 diff --git a/tests/unit/torch/utils/test_meta_dcp_load.py b/tests/unit/torch/utils/test_meta_dcp_load.py index 764bce545c6..4c642b1644c 100644 --- a/tests/unit/torch/utils/test_meta_dcp_load.py +++ b/tests/unit/torch/utils/test_meta_dcp_load.py @@ -474,15 +474,6 @@ def recording_load(state_dict, **kwargs): # ------------------------------------------------------------------- under a process group -@pytest.fixture -def gloo_world_of_one(skip_on_windows, tmp_path_factory): - """A one-rank gloo process group, so the loader takes its distributed path.""" - store = tmp_path_factory.mktemp("pg") / "store" - torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) - yield - torch.distributed.destroy_process_group() - - def test_fsdp2_sharded_params_are_filled_in_place(tmp_path, gloo_world_of_one): """The EP hybrid's layout: the decoder layers are ``fully_shard``-ed on meta before the load, the root is not. The layers' params must stay sharded ``DTensor``s holding the checkpoint values, the From c18632238d722f1a63ed236e25df1a6939a9d7c0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:38:33 +0000 Subject: [PATCH 56/67] Pack an FSDP2-sharded model's weights in held weight-update windows The no-gather export packs weights while the model is still sharded, with no weight-access window open. ExportContext.keep_fused_experts marks that path. On it, a quantized linear whose dense part FSDP2 shards is packed inside a weight-update window, which hands FSDP2 the packed weights; packed outside one, the unit's next reshard would put its unquantized shard back. Under held_fsdp2_weight_windows, which the export's dispatch loop now enters, a unit's window stays open from one module to the next, so a decoder layer is gathered once rather than once per linear; any other loop gets a window per module, closed by its handler. Where expert parallelism leaves the root unwrapped, the window opens on the layer's own unit, as on a wrapped root it would do nothing. Split out of #2750 as [1/2], at its review's request. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 98 +++++++++++++++++++ modelopt/torch/export/registry.py | 20 ++++ modelopt/torch/export/unified_export_hf.py | 14 ++- .../torch/quantization/utils/core_utils.py | 4 +- .../export/test_sharded_weight_packing.py | 63 ++++++++++++ 5 files changed, 195 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 200ec6e5dae..d28e7be022d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -16,9 +16,16 @@ """Built-in module handlers for unified Hugging Face export.""" import collections.abc +import contextlib import warnings +from dataclasses import dataclass import torch.nn as nn +from torch.distributed.fsdp import FSDPModule + +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps +from modelopt.torch.quantization.utils.core_utils import ModuleNames, _get_enclosing_fsdp_module +from modelopt.torch.utils.distributed import distribution_plan from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .moe_utils import _export_fused_experts @@ -34,6 +41,87 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool: return hasattr(module, f"{first_proj_attr}_weight_quantizers") +@dataclass +class _HeldWindow: + """An FSDP2 unit's weight-update window, and the modules packed in it so far.""" + + unit: nn.Module + modules: list[nn.Module] + stack: contextlib.ExitStack + + +def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: + """The innermost FSDP2 unit enclosing ``module``, or the model if none does.""" + if ctx.module_names is None: + ctx.module_names = module_name_maps(ctx.model) + return _get_enclosing_fsdp_module(module, ctx.model, ctx.module_names) or ctx.model + + +def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: + """The root a weight-update window on ``unit`` is opened on, and its module names. + + The window opens on the innermost unit enclosing the module, but only when the root it is handed + is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the + dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which + packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated + codes next to correct scales). So there that unit is handed in as the root, with names relative + to it, which the window resolves the modules' parameters against. + """ + if isinstance(ctx.model, FSDPModule) or unit is ctx.model: + # The model's own names, which _enclosing_unit has built. + return ctx.model, ctx.module_names + return unit, module_name_maps(unit) + + +def _fsdp2_weight_window( + ctx: ExportContext, module: nn.Module +) -> contextlib.AbstractContextManager: + """The ``fsdp2_aware_weight_update`` window to pack ``module``'s weights in. + + Inside :func:`held_fsdp2_weight_windows`, the window of the module's FSDP2 unit, held open + across the unit's modules. Anywhere else, a window of the module's own, closed when the handler + is done, so a dispatch loop that knows nothing of held windows cannot leave a unit unsharded. + """ + unit = _enclosing_unit(ctx, module) + if not ctx.hold_fsdp2_windows: + root, names = _window_root(ctx, unit) + return fsdp2_aware_weight_update(root, module, names=names) + held = ctx.fsdp2_window + if held is not None and held.unit is unit: + held.modules.append(module) + return contextlib.nullcontext() + _close_fsdp2_weight_window(ctx) + root, names = _window_root(ctx, unit) + held = _HeldWindow(unit, [module], contextlib.ExitStack()) + held.stack.enter_context(fsdp2_aware_weight_update(root, held.modules, names=names)) + ctx.fsdp2_window = held + return contextlib.nullcontext() + + +@contextlib.contextmanager +def held_fsdp2_weight_windows(ctx: ExportContext): + """Within the block, pack each FSDP2 unit's modules in one weight-update window. + + The window unshards the whole unit and reshards it when it closes, so a window per module + gathers a decoder layer once per quantized linear in it. Here the unit's window stays open from + one handler to the next until the export reaches another unit, each module joining the list + whose weights it re-registers with FSDP2, and leaving the block closes the last. + """ + ctx.hold_fsdp2_windows = True + try: + yield + finally: + ctx.hold_fsdp2_windows = False + _close_fsdp2_weight_window(ctx) + + +def _close_fsdp2_weight_window(ctx: ExportContext) -> None: + """Close the weight-update window the handlers hold open, if any, resharding its unit.""" + held, ctx.fsdp2_window = ctx.fsdp2_window, None + if held is not None: + held.stack.close() + + def _export_weight( module: nn.Module, ctx: ExportContext, @@ -43,6 +131,16 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight + # The no-gather writer packs while the model is still sharded, so unlike the + # collect_export_tensors path there is no open weight-access window here and the weight is a + # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an + # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, + # which reshards the unit when it closes: a layer left unsharded holds all of its weights on + # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: + with _fsdp2_weight_window(ctx, module): + _export_quantized_weight(module, ctx.dtype, weight_name) + return _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 1aa575d92de..808210a4f68 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,10 +28,14 @@ from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn +if TYPE_CHECKING: + from modelopt.torch.quantization.utils.core_utils import ModuleNames + __all__ = [ "ExportContext", "ExportHandler", @@ -59,6 +63,22 @@ class ExportContext: model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" + # Distributed (no-gather) export: the handlers pack while the model is still sharded, an + # FSDP2-owned weight inside its unit's weight-update window. + keep_fused_experts: bool = False + + module_names: "ModuleNames | None" = None + """``model``'s name maps, built once per export by the first handler that looks a module up + (:func:`~modelopt.torch.quantization.utils.module_name_maps`).""" + + hold_fsdp2_windows: bool = False + """Whether the handlers hold an FSDP2 unit's weight-update window open across its modules: set + only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which + closes the last window. Elsewhere each module is packed in a window of its own.""" + + fsdp2_window: Any = None + """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" + ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ffbd62ddfba..6ec8f7b714d 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -104,7 +104,7 @@ export_sparse_attention_config = None # Importing the built-in handlers installs their entries in the two registries. -from . import hf_export_handlers as _hf_export_handlers # noqa: F401 +from . import hf_export_handlers as _hf_export_handlers from .convert_hf_config import convert_hf_quant_config_format from .layer_utils import get_experts_list, is_layernorm, is_quantlinear, sync_moe_gate_up_amax from .model_utils import TiedWeightMap, get_language_model_from_vl, is_multimodal_model @@ -636,6 +636,12 @@ def _export_quantized_weight( The export includes converting weight tensor to correct quantized values and quantized dtype, and registering scaling factors. + An export handler packs through ``hf_export_handlers._export_weight`` rather than calling this + directly: on the keep-fused export an FSDP2-owned weight must be packed inside its unit's + weight-update window, which that helper opens. + Packed outside it, the unit's next reshard puts FSDP2's unquantized shard back, and the + checkpoint is written unquantized. + Tied-weight dedup is not handled here: both sides of a tie are packed independently (identically, once ``sync_tied_input_amax`` has equalized their scales), and the duplicate is dropped by name in :func:`postprocess_state_dict`. Deduping per-module at @@ -996,8 +1002,10 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - for name, sub_module in model.named_modules(): - _dispatch_export_handler(name, sub_module, ctx) + # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. + with _hf_export_handlers.held_fsdp2_weight_windows(ctx): + for name, sub_module in model.named_modules(): + _dispatch_export_handler(name, sub_module, ctx) def _prepare_model_for_export(model, dtype, is_modelopt_qlora): diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index b4120720b5c..0915a9d1720 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -1054,7 +1054,9 @@ def fsdp2_aware_weight_update( Args: root_model (nn.Module): The root model of the FSDPModule. modules_to_update (list): The list of modules to update which should be a list of modules that are - direct children of the FSDPModule. + direct children of the FSDPModule. A list is read again when the context exits, so a caller + may append modules of the same FSDP unit to it inside the context, and their weights are + re-registered with FSDP2 too (the export's held weight windows rely on this). reshard (bool): Whether to reshard the FSDP root module after the weight update. Returns: diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index a1c9fb19e58..0e666784169 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -15,10 +15,21 @@ """Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +import contextlib +import copy + +import pytest import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def _raw_bytes(tensor): @@ -72,3 +83,55 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) assert isinstance(packed, DTensor) and packed.placements == (Shard(0),) assert packed.shape == (3, 32) assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) + + +@pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) +@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( + monkeypatch, gloo_world_of_one, root_sharded, held +): + """On the no-gather path the model is still sharded when a linear is packed. + + Packing NVFP4 straight off the sharded weight is not possible (the in-place cast on the + DTensor's local view is refused under autograd), so the handlers pack inside a weight-update + window, on the layer's own unit where expert parallelism leaves the root unsharded. The window + gathers the whole unit, so within held windows it opens once for all of the unit's linears; + any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is + sharded again afterwards, and the sharded state dict carries the packed weights, as an + unsharded pack has them. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + refs = copy.deepcopy(model)[0] + for i, ref in enumerate(refs): + ExportModuleRegistry.match(ref)( + f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) + ) + fully_shard(layer, mesh=gloo_world_of_one) + if root_sharded: + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + unshards = [] + unshard = FSDPModule.unshard + monkeypatch.setattr( + FSDPModule, "unshard", lambda unit, *a, **k: unshards.append(unit) or unshard(unit, *a, **k) + ) + + with handlers.held_fsdp2_weight_windows(ctx) if held else contextlib.nullcontext(): + for i, linear in enumerate(layer): + ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) + + assert unshards == ([layer] if held else [layer, layer]) + assert all(isinstance(linear.weight, DTensor) for linear in layer) + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for i, ref in enumerate(refs): + packed = sd[f"0.{i}.weight"] + assert isinstance(packed, DTensor) + assert packed.dtype == torch.uint8 + assert torch.equal(packed.full_tensor(), ref.weight) + for name in ("weight_scale", "weight_scale_2", "input_scale"): + want = _raw_bytes(getattr(ref, name)) + assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name From 8d1caf3b993524978d96beacfc9dbd65dad1dfe8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:36:07 +0000 Subject: [PATCH 57/67] List a module in its unit's held weight window once, however many weights it packs A handler that packs several weights of one module asks for the window once per weight. Listed twice, the module was re-registered with FSDP2 a second time when the window closed, from the shard the first registration had put back, which FSDP2 rejects. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 5 +++- .../export/test_sharded_weight_packing.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index d28e7be022d..eaf0c5607db 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -88,7 +88,10 @@ def _fsdp2_weight_window( return fsdp2_aware_weight_update(root, module, names=names) held = ctx.fsdp2_window if held is not None and held.unit is unit: - held.modules.append(module) + # Once per module, though a handler packing several weights asks once per weight: the + # window re-registers each listed module's weights with FSDP2 when it closes. + if not any(m is module for m in held.modules): + held.modules.append(module) return contextlib.nullcontext() _close_fsdp2_weight_window(ctx) root, names = _window_root(ctx, unit) diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 0e666784169..5a5a1856048 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -135,3 +135,29 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for name in ("weight_scale", "weight_scale_2", "input_scale"): want = _raw_bytes(getattr(ref, name)) assert torch.equal(_raw_bytes(sd[f"0.{i}.{name}"]), want), name + + +def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): + """A handler packing several weights of one module asks for the window once per weight. + + The module must join its unit's held window once: listed twice, it would be re-registered with + FSDP2 a second time when the window closes, from the shard the first registration put back. + """ + torch.manual_seed(0) + layer = nn.Sequential(nn.Linear(32, 16, bias=False)) + model = nn.Sequential(layer) + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + ref = copy.deepcopy(layer[0]) + ExportModuleRegistry.match(ref)("0.0", ref, ExportContext(model=ref, dtype=torch.bfloat16)) + fully_shard(layer, mesh=gloo_world_of_one) + fully_shard(model, mesh=gloo_world_of_one) + set_distribution_plan(model, DistributionPlan(dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + ExportModuleRegistry.match(layer[0])("0.0", layer[0], ctx) + with handlers._fsdp2_weight_window(ctx, layer[0]): # as for the module's next weight + pass + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) From 9f30d031fa692c13b463dac4ec6b4a8e0ed3534b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 00:37:31 +0000 Subject: [PATCH 58/67] Keep the sharded packing's comments to their constraint, and its imports at module scope The comments said what to avoid and then retold how each torch failure surfaced; that account is in the PR description. The static block-scale path says why it matches the dynamic one: its amax buffer is never a DTensor. DTensor and init_device_mesh are imported at module scope (torch>=2.10 has both), and ExportContext.fsdp2_window is typed as the held window it holds. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 16 ++--- modelopt/torch/export/registry.py | 5 +- .../quantization/qtensor/nvfp4_tensor.py | 65 ++++++------------- tests/unit/torch/conftest.py | 3 +- 4 files changed, 29 insertions(+), 60 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index eaf0c5607db..27d44bd072d 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -60,12 +60,8 @@ def _enclosing_unit(ctx: ExportContext, module: nn.Module) -> nn.Module: def _window_root(ctx: ExportContext, unit: nn.Module) -> tuple[nn.Module, ModuleNames | None]: """The root a weight-update window on ``unit`` is opened on, and its module names. - The window opens on the innermost unit enclosing the module, but only when the root it is handed - is wrapped itself: on an unwrapped root it does nothing. Expert parallelism FSDP2-shards only the - dense layers and leaves the root unwrapped, where the weight would stay a sharded DTensor, which - packing mishandles (``DTensor << int`` is a silent no-op in torch 2.11, so NVFP4 packed unrelated - codes next to correct scales). So there that unit is handed in as the root, with names relative - to it, which the window resolves the modules' parameters against. + A window does nothing on an unwrapped root, which expert parallelism leaves (it FSDP2-shards + only the dense layers), so there the unit itself is the root, with names relative to it. """ if isinstance(ctx.model, FSDPModule) or unit is ctx.model: # The model's own names, which _enclosing_unit has built. @@ -134,12 +130,8 @@ def _export_weight( # install the built-in handlers while retaining this legacy helper's import path. from .unified_export_hf import _export_quantized_weight - # The no-gather writer packs while the model is still sharded, so unlike the - # collect_export_tensors path there is no open weight-access window here and the weight is a - # DTensor. FP8 tolerates that; NVFP4 does not -- get_weights_scaling_factor does an - # index_put_, which DTensor cannot dispatch. Pack in a weight-update window for that path only, - # which reshards the unit when it closes: a layer left unsharded holds all of its weights on - # every rank, and FSDP2 later puts its own shards back over anything packed outside a window. + # The no-gather export packs a still-sharded model: an FSDP2-owned weight goes in its unit's + # weight-update window, or the unit's next reshard puts FSDP2's unquantized shard back. if ctx.keep_fused_experts and distribution_plan(ctx.model).is_fsdp_dense: with _fsdp2_weight_window(ctx, module): _export_quantized_weight(module, ctx.dtype, weight_name) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 808210a4f68..ca99532dfae 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -28,12 +28,13 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import torch import torch.nn as nn if TYPE_CHECKING: + from modelopt.torch.export.hf_export_handlers import _HeldWindow from modelopt.torch.quantization.utils.core_utils import ModuleNames __all__ = [ @@ -76,7 +77,7 @@ class ExportContext: only inside :func:`~modelopt.torch.export.hf_export_handlers.held_fsdp2_weight_windows`, which closes the last window. Elsewhere each module is packed in a window of its own.""" - fsdp2_window: Any = None + fsdp2_window: "_HeldWindow | None" = None """The FSDP2 unit's weight-update window the handlers hold open across its modules.""" diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index 3ac450c2908..ab394b6667f 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -16,6 +16,7 @@ """Implements NVFP4 quantization for efficient tensor storage and computation.""" import torch +from torch.distributed.tensor import DTensor from ..backends.utils import fp4_compatible from ..qtensor.base_qtensor import BaseQuantizedTensor @@ -163,8 +164,7 @@ def get_weights_scaling_factor_from_quantizer( # Compute scales in float per_block_scale_max = global_amax / E2M1_MAX per_block_scale = per_block_amax / E2M1_MAX - # torch.where, not a boolean-mask assignment: the latter lowers to aten::index_put_, - # which DTensor cannot dispatch, and this runs on DTensor weights during distributed export. + # torch.where, to match the dynamic path below, which must avoid index_put_ on a DTensor. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) # Reshape per_block_scale to match weight's block structure @@ -215,10 +215,8 @@ def get_weights_scaling_factor( per_block_scale = per_block_amax / ( E2M1_MAX * weights_scaling_factor_2.to(per_block_amax.device) ) - # Set all zero values in scale to 1.0. Use torch.where rather than a boolean-mask assignment: - # the latter lowers to aten::index_put_, which DTensor cannot dispatch ("Cannot find device - # mesh from args"), and on the distributed export path this weight IS a DTensor. torch.where is - # pointwise, so DTensor handles it, and the result is identical for a plain tensor. + # Set all zero values in scale to 1.0. torch.where, not a boolean-mask assignment: that lowers + # to index_put_, which DTensor (a sharded weight in the distributed export) cannot dispatch. per_block_scale = torch.where(per_block_scale == 0, 1.0, per_block_scale) if not keep_high_precision: @@ -253,33 +251,11 @@ def get_activation_scaling_factor(cls, quantizer): @classmethod def _cast_fp4(cls, weight: torch.Tensor): """Converts tensor to uint4.""" - # The distributed export path packs a SHARDED weight. Every step below is elementwise over the - # weight, but two of them mix it with plain constants -- searchsorted against the E2M1 bound - # table, and the odd-bound comparison -- which DTensor refuses ("got mixed torch.Tensor and - # DTensor"). Run on this rank's local shard and rewrap: the result is identical elementwise, and - # the no-gather writer emits exactly this rank's shard anyway. Patching the ops one at a time - # instead just moves the failure to the next one. - _dtensor_meta = None - try: - from torch.distributed.tensor import DTensor as _DTensor - - if isinstance(weight, _DTensor): - # Capture the GLOBAL shape/stride: from_local() would otherwise re-infer them as - # local x world, which is wrong for an UNEVEN shard (e.g. 96 rows over 64 ranks) and - # gives each rank a different global shape. Downstream that desynchronises the - # exporter's per-rank tensor classification and surfaces as a KeyError deep in the - # writer, long after the real mistake. - _dtensor_meta = ( - _DTensor, - weight.device_mesh, - weight.placements, - weight.shape, - weight.stride(), - ) - weight = weight.to_local() - except ImportError: - pass - + # A DTensor is cast on its local shard: the steps below mix it with plain tensors, which + # DTensor refuses. The cast is elementwise, so rewrapping the result gives the same tensor. + sharded = weight if isinstance(weight, DTensor) else None + if sharded is not None: + weight = sharded.to_local() device = weight.device # Extract sign and compute absolute values in one pass @@ -299,17 +275,18 @@ def _cast_fp4(cls, weight: torch.Tensor): # Combine sign, ordinal, and rounding adjustment packed = (sign_bit << 3) + ord + equals_odd_bounds - if _dtensor_meta is not None: - _dt_cls, _mesh, _placements, _gshape, _gstride = _dtensor_meta - packed = _dt_cls.from_local( - packed, - _mesh, - _placements, - run_check=False, - shape=_gshape, - stride=_gstride, - ) - return packed + if sharded is None: + return packed + # The global shape and stride as given: from_local would infer local x world, wrong for an + # uneven shard. + return DTensor.from_local( + packed, + sharded.device_mesh, + sharded.placements, + run_check=False, + shape=sharded.shape, + stride=sharded.stride(), + ) @classmethod def quantize( diff --git a/tests/unit/torch/conftest.py b/tests/unit/torch/conftest.py index de576f629c4..d751079bfd2 100644 --- a/tests/unit/torch/conftest.py +++ b/tests/unit/torch/conftest.py @@ -17,6 +17,7 @@ import pytest import torch +from torch.distributed.device_mesh import init_device_mesh from modelopt.torch.opt.mode import _ModeRegistryCls @@ -59,8 +60,6 @@ def use_channel_div_4(): def gloo_world_of_one(skip_on_windows, tmp_path_factory): """A one-rank gloo process group, so the code under test takes its distributed path, and its CPU device mesh.""" - from torch.distributed.device_mesh import init_device_mesh - store = tmp_path_factory.mktemp("pg") / "store" torch.distributed.init_process_group("gloo", init_method=store.as_uri(), rank=0, world_size=1) try: From 3c7c5d9c3ef652c5c48f2cd037fa6d68bee72a80 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:20:36 +0000 Subject: [PATCH 59/67] Hand FSDP2 the packed weight of a unit that owns it directly, and test the packing on two ranks Under an unwrapped root, a module that is an FSDP2 unit of its own is the window's root and owns its weight directly. get_prefixed_param_names named that weight's module after the weight itself, so the window's FSDPParam mapping held 'weight.weight' while its exit looked up '.weight': it skipped the packed weight, and the reshard put the float shard back. The owner of a parameter the root holds is now named '', and both sides join names alike. The handler test gains that layout, which fails without the fix. A two-rank job packs an uneven NVFP4 shard, three rows over two ranks, into the plain pack's codes and scales, and runs the handlers over two FSDP2 units, a 17-row weight among them, against an unsharded export. Signed-off-by: Shengliang Xu --- .../torch/quantization/utils/core_utils.py | 16 ++-- .../export/test_sharded_weight_packing.py | 96 +++++++++++++++++-- .../torch/quantization/test_param_index.py | 2 +- 3 files changed, 100 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 0915a9d1720..ca7168ac1da 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -942,7 +942,13 @@ def get_prefixed_param_names(parent_model, target_module, param_index=None): target_ids = {id(p) for p in target_module.parameters()} # Lowest position == first in named_parameters() order, matching a linear scan's result. match = min((param_index[pid] for pid in target_ids if pid in param_index), default=None) - return match[1].rsplit(".", 1)[0] if match is not None else None + # "" for a parameter the parent owns itself. + return match[1].rpartition(".")[0] if match is not None else None + + +def _param_key(module_name: str, param_name: str) -> str: + """The full name of ``module_name``'s parameter ``param_name``; ``""`` names the root module.""" + return f"{module_name}.{param_name}" if module_name else param_name def create_fsdp_param_mapping(fsdp_param_list, model): @@ -966,8 +972,7 @@ def create_fsdp_param_mapping(fsdp_param_list, model): if module_name is not None: # Get the parameter name from _module_info and construct full param name param_name = param._module_info.param_name - full_param_name = f"{module_name}.{param_name}" - mapping[full_param_name] = param + mapping[_param_key(module_name, param_name)] = param return mapping @@ -1096,7 +1101,7 @@ def fsdp2_aware_weight_update( module_name = _get_module_name(module, root_model, names) # Check if any parameter from this module is in the mapping module_params_in_mapping = any( - f"{module_name}.{n}" in fsdp_param_mapping + _param_key(module_name, n) in fsdp_param_mapping for n, _ in module.named_parameters() ) assert module_params_in_mapping, ( @@ -1112,8 +1117,7 @@ def fsdp2_aware_weight_update( # Update FSDPParam list for module in modules_to_update: for param_name, param in module.named_parameters(): - name = _get_module_name(module, root_model, names) - name = f"{name}.{param_name}" + name = _param_key(_get_module_name(module, root_model, names), param_name) if name not in fsdp_param_mapping: continue diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index 5a5a1856048..b61b11b8740 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Packing weights that are still sharded, as the no-gather export does (CPU, one gloo rank).""" +"""Packing weights that are still sharded, as the no-gather export does (CPU, gloo, one and two ranks).""" import contextlib import copy @@ -21,7 +21,9 @@ import pytest import torch import torch.nn as nn +from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import FSDPModule, fully_shard from torch.distributed.tensor import DTensor, Shard, distribute_tensor @@ -86,9 +88,9 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) @pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) -@pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) +@pytest.mark.parametrize("layout", ["fsdp2", "expert-parallel", "linear-units"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( - monkeypatch, gloo_world_of_one, root_sharded, held + monkeypatch, gloo_world_of_one, layout, held ): """On the no-gather path the model is still sharded when a linear is packed. @@ -98,7 +100,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( gathers the whole unit, so within held windows it opens once for all of the unit's linears; any other dispatch loop gets a window per linear, closed by its handler. Either way the unit is sharded again afterwards, and the sharded state dict carries the packed weights, as an - unsharded pack has them. + unsharded pack has them. A linear that is a unit of its own under an unwrapped root owns its + weight as the window's root: its packed weight must be handed to FSDP2 all the same. """ torch.manual_seed(0) layer = nn.Sequential(nn.Linear(32, 16, bias=False), nn.Linear(16, 8, bias=False)) @@ -109,8 +112,10 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( ExportModuleRegistry.match(ref)( f"0.{i}", ref, ExportContext(model=ref, dtype=torch.bfloat16) ) - fully_shard(layer, mesh=gloo_world_of_one) - if root_sharded: + units = list(layer) if layout == "linear-units" else [layer] + for unit in units: + fully_shard(unit, mesh=gloo_world_of_one) + if layout == "fsdp2": fully_shard(model, mesh=gloo_world_of_one) set_distribution_plan(model, DistributionPlan(dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) @@ -124,7 +129,8 @@ def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( for i, linear in enumerate(layer): ExportModuleRegistry.match(linear)(f"0.{i}", linear, ctx) - assert unshards == ([layer] if held else [layer, layer]) + per_linear = [linear if layout == "linear-units" else layer for linear in layer] + assert unshards == (list(dict.fromkeys(per_linear)) if held else per_linear) assert all(isinstance(linear.weight, DTensor) for linear in layer) sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) for i, ref in enumerate(refs): @@ -161,3 +167,79 @@ def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _two_rank_packing_job(rank, world_size): + mesh = init_device_mesh("cpu", (world_size,)) + torch.manual_seed(0) + # Three rows over two ranks, an all-zero block included: rank 0 packs two rows, rank 1 one. + weight = torch.randn(3, 32, dtype=torch.bfloat16) + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(0, 2) if rank == 0 else slice(2, 3) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor(), want) + + # The handlers over two units in held windows, a 17-row weight sharded 9 + 8 among them. + model = _TwoUnits() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with bounded_collectives(), handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor(), ref.weight), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_two_ranks_pack_their_shards_into_the_unsharded_bytes(cpu_dist_workers): + """On two ranks a shard holds other rows than the whole weight: an uneven NVFP4 pack gives the + plain pack's codes and scales, and the handlers, run over two units, leave the sharded state + dict holding what an unsharded export packs.""" + cpu_dist_workers.run(_two_rank_packing_job) diff --git a/tests/unit/torch/quantization/test_param_index.py b/tests/unit/torch/quantization/test_param_index.py index a78dd7b8cff..54ade0a317a 100644 --- a/tests/unit/torch/quantization/test_param_index.py +++ b/tests/unit/torch/quantization/test_param_index.py @@ -28,7 +28,7 @@ def _linear_scan(parent_model, target_module): target_ids = {id(p) for p in target_module.parameters()} return next( ( - name.rsplit(".", 1)[0] + name.rpartition(".")[0] for name, param in parent_model.named_parameters() if id(param) in target_ids ), From a48ce7aa55a293718973c2866cb19331b4e1d891 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 06:50:23 +0000 Subject: [PATCH 60/67] Test the sharded packing over NCCL on every visible GPU The two-rank gloo job, on CUDA: an uneven NVFP4 shard, two rows a rank and one on the last, packs into the plain pack's codes and scales, and the handlers, run over two FSDP2 units in held windows with keep_fused_experts, leave the sharded state dict holding what an unsharded export packs. It needs two GPUs, so a single-GPU runner skips it. Signed-off-by: Shengliang Xu --- .../export/test_sharded_weight_packing_gpu.py | 112 ++++++++++++++++++ 1 file changed, 112 insertions(+) create mode 100644 tests/gpu/torch/export/test_sharded_weight_packing_gpu.py diff --git a/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py new file mode 100644 index 00000000000..9284f755e48 --- /dev/null +++ b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py @@ -0,0 +1,112 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Packing weights that are still sharded, as the no-gather export does, over NCCL on every GPU.""" + +import copy + +import torch +import torch.nn as nn +from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict +from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard +from torch.distributed.tensor import Shard, distribute_tensor + +import modelopt.torch.export.hf_export_handlers as handlers +import modelopt.torch.quantization as mtq +from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry +from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu() + + +class _Unit(nn.Module): + """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" + + def __init__(self): + super().__init__() + self.a = nn.Linear(32, 17, bias=False) + self.b = nn.Linear(32, 16, bias=False) + + def forward(self, x): + return self.a(x), self.b(x) + + +class _TwoUnits(nn.Module): + def __init__(self): + super().__init__() + self.units = nn.ModuleList([_Unit(), _Unit()]) + + def forward(self, x): + return [unit(x) for unit in self.units] + + +def _packing_job(rank, world_size): + mesh = init_device_mesh("cuda", (world_size,)) + torch.manual_seed(0) + # Two rows a rank, the last rank one, an all-zero block included. + rows = 2 * world_size - 1 + weight = torch.randn(rows, 32, dtype=torch.bfloat16, device="cuda") + weight[1, :16] = 0.0 + scale_2 = (weight.abs().amax().float() / (6.0 * 448.0)).reshape(()) + scale = NVFP4QTensor.get_weights_scaling_factor(weight, 16, scale_2)[0] + want = NVFP4QTensor.quantize(weight, 16, scale, scale_2)[0]._quantized_data + mine = slice(2 * rank, min(2 * rank + 2, rows)) + + sharded = distribute_tensor(weight, mesh, [Shard(0)]) + sharded_scale = NVFP4QTensor.get_weights_scaling_factor(sharded, 16, scale_2)[0] + packed = NVFP4QTensor.quantize(sharded, 16, sharded_scale, scale_2)[0]._quantized_data + + assert sharded_scale.shape == scale.shape and packed.shape == want.shape + assert torch.equal(_raw_bytes(sharded_scale.to_local()), _raw_bytes(scale[mine])) + assert torch.equal(packed.full_tensor().cpu(), want.cpu()) + + # The handlers over two units in held windows, a 17-row weight sharded unevenly among them. + model = _TwoUnits().cuda() + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, lambda m: m(torch.randn(4, 32, device="cuda"))) + linears = [ + (f"units.{u}.{p}", getattr(unit, p)) for u, unit in enumerate(model.units) for p in "ab" + ] + refs = copy.deepcopy(model) + for name, _ in linears: + ref = refs.get_submodule(name) + ExportModuleRegistry.match(ref)(name, ref, ExportContext(model=ref, dtype=torch.bfloat16)) + for unit in model.units: + fully_shard(unit, mesh=mesh) + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + + with handlers.held_fsdp2_weight_windows(ctx): + for name, linear in linears: + ExportModuleRegistry.match(linear)(name, linear, ctx) + + sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + for name, _ in linears: + ref = refs.get_submodule(name) + assert torch.equal(sd[f"{name}.weight"].full_tensor().cpu(), ref.weight.cpu()), name + for leaf in ("weight_scale", "weight_scale_2", "input_scale"): + got = _raw_bytes(sd[f"{name}.{leaf}"]) + assert torch.equal(got, _raw_bytes(getattr(ref, leaf))), (name, leaf) + + +def test_sharded_packing_matches_the_unsharded_bytes_over_nccl(need_2_gpus, dist_workers): + """An uneven NVFP4 shard packs into the plain pack's codes and scales, and the handlers, run over + two FSDP2 units in held windows, leave the sharded state dict holding what an unsharded export + packs.""" + dist_workers.run(_packing_job) From 26e9863b87db99282f2a9b791afdfa068f9fe8c9 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:14:17 +0000 Subject: [PATCH 61/67] Drop the fused-experts handler's claim that FSDP2 ranks pack their experts fused The paragraph predates the keep-fused fold. Without that fold the handler splits every expert into per-expert weights, under FSDP2 too, and the docstring's last paragraph says what the fold does instead. Signed-off-by: Shengliang Xu --- modelopt/torch/export/hf_export_handlers.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 1418e8f0efb..f99d9020f74 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -226,9 +226,6 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None: """Split and quantize a fused-experts module with plural weight quantizers. - Under FSDP2 each rank holds only some experts, so it packs just those (kept fused) and the - split into per-expert keys is deferred until the gather brings all experts together. - Tied experts are packed independently and their duplicate keys are dropped by name in postprocess_state_dict; no per-module dedup cache is used. From d568dacad398bf2566c598dc5d3a92047ef00644 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:43:45 +0000 Subject: [PATCH 62/67] Refuse a DTensor shard in the 4-bit packers that cannot pack one `<<` is a silent no-op on a DTensor, and the INT4, NF4 and MXFP4 packers, and the export's INT4 packer, still shifted with it. They now use torch.bitwise_left_shift, which gives the same bytes on a plain tensor. None of them packs a shard either way: each combines its input with plain tensors, which DTensor refuses deep inside, or, on INT4's fallback, gathers it whole on every rank. So each now refuses a DTensor up front, and NVFP4 refuses one that comes with a plain per-block scale, such as a static quantizer's, which covers the whole weight. Signed-off-by: Shengliang Xu --- modelopt/torch/export/quant_utils.py | 6 ++- .../quantization/qtensor/base_qtensor.py | 12 +++++ .../torch/quantization/qtensor/int4_tensor.py | 5 ++- .../quantization/qtensor/mxfp4_tensor.py | 6 ++- .../torch/quantization/qtensor/nf4_tensor.py | 8 +++- .../quantization/qtensor/nvfp4_tensor.py | 12 ++++- .../export/test_sharded_weight_packing.py | 44 ++++++++++++++++++- 7 files changed, 83 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/export/quant_utils.py b/modelopt/torch/export/quant_utils.py index bf4db7f7d55..c77f651324a 100755 --- a/modelopt/torch/export/quant_utils.py +++ b/modelopt/torch/export/quant_utils.py @@ -41,6 +41,7 @@ NVFP4QTensor, QTensorWrapper, ) +from modelopt.torch.quantization.qtensor.base_qtensor import _refuse_dtensor from modelopt.torch.quantization.utils import ( quantizer_attr_names, representative_weight_quantizer, @@ -828,6 +829,7 @@ def _validate_int4_block_size(in_dim, block_size): def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): """Packs the INT4 weights into uint8 tensor.""" + _refuse_dtensor(weight, "INT4 packing") out_dim = weight.shape[-2] assert out_dim % 2 == 0, f"Cannot pack weight. Out dimension {out_dim} is not an even number." in_dim = weight.shape[-1] @@ -851,7 +853,7 @@ def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): # Pack two 4-bit values (val0,val1) into a single byte: val0 = transpose[..., 0] & 0x0F val1 = transpose[..., 1] & 0x0F - packed_byte = val0 | (val1 << 4) + packed_byte = val0 | torch.bitwise_left_shift(val1, 4) # Transpose back to the shape (experts, out_dim // 2, in_dim) return packed_byte.permute(0, 2, 1).contiguous().view(torch.uint8) @@ -864,7 +866,7 @@ def pack_int4_in_uint8(weight, weights_scaling_factor, block_size): # Pack two 4-bit values into one byte val0 = reshaped[..., 0] & 0x0F val1 = reshaped[..., 1] & 0x0F - packed_byte = val0 | (val1 << 4) + packed_byte = val0 | torch.bitwise_left_shift(val1, 4) # Return shape (out_dim // 2, in_dim) return packed_byte.T.contiguous().view(torch.uint8) diff --git a/modelopt/torch/quantization/qtensor/base_qtensor.py b/modelopt/torch/quantization/qtensor/base_qtensor.py index 86d4eb493c3..ba4404b15c8 100644 --- a/modelopt/torch/quantization/qtensor/base_qtensor.py +++ b/modelopt/torch/quantization/qtensor/base_qtensor.py @@ -37,6 +37,18 @@ class QTensorType(enum.Enum): __all__ = ["BaseQuantizedTensor", "QTensorWrapper", "pack_real_quantize_weight"] +def _refuse_dtensor(tensor: torch.Tensor, packer: str) -> None: + """Refuse a DTensor shard in a packer that only packs a whole tensor. + + DTensor refuses the plain tensors such a packer combines with its input, or the packer gathers + the input whole on every rank; of the 4-bit packers, only NVFP4's packs a shard. + """ + if isinstance(tensor, DTensor): + raise NotImplementedError( + f"{packer} packs a whole tensor, not a DTensor shard; pack the unsharded tensor." + ) + + class BaseQuantizedTensor: """Base class for quantized tensors, providing methods for quantization and dequantization. diff --git a/modelopt/torch/quantization/qtensor/int4_tensor.py b/modelopt/torch/quantization/qtensor/int4_tensor.py index f78d12aad75..15ae1518cc2 100644 --- a/modelopt/torch/quantization/qtensor/int4_tensor.py +++ b/modelopt/torch/quantization/qtensor/int4_tensor.py @@ -19,7 +19,7 @@ import torch from ..extensions import get_cuda_ext -from ..qtensor.base_qtensor import BaseQuantizedTensor +from ..qtensor.base_qtensor import BaseQuantizedTensor, _refuse_dtensor from ..utils import reduce_amax, reduce_block_padding __all__ = ["INT4QTensor"] @@ -47,6 +47,7 @@ def quantize(cls, input: torch.Tensor, block_size: int) -> tuple: Returns: tuple: Contains quantized data, input quantization config, and scale quantization config. """ + _refuse_dtensor(input, cls.__name__) cuda_ext = get_cuda_ext() scale_quant_maxbound = cls._get_quant_maxbound(num_bits=4) @@ -81,7 +82,7 @@ def quantize(cls, input: torch.Tensor, block_size: int) -> tuple: # pack the int4 weights into a uint8 tensor # packing format: w0, w1, w2, w3, w4, w5, ... # | byte | byte | byte | - packed_output_uint8 = flattened[::2] << 4 | flattened[1::2] + packed_output_uint8 = torch.bitwise_left_shift(flattened[::2], 4) | flattened[1::2] # reshape the packed_output_uint8 to the original dimensions for sharding packed_output_uint8 = packed_output_uint8.reshape(*original_input.shape[:-1], -1) diff --git a/modelopt/torch/quantization/qtensor/mxfp4_tensor.py b/modelopt/torch/quantization/qtensor/mxfp4_tensor.py index 022825e4062..102c00e6c71 100644 --- a/modelopt/torch/quantization/qtensor/mxfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/mxfp4_tensor.py @@ -17,7 +17,7 @@ import torch -from ..qtensor.base_qtensor import BaseQuantizedTensor +from ..qtensor.base_qtensor import BaseQuantizedTensor, _refuse_dtensor __all__ = ["MXFP4QTensor"] @@ -57,10 +57,12 @@ def fuse_uint4_to_uint8(x): # If this behavior is not desired, please modify the code accordingly left_side = x[..., 0::2] # Even indices (0, 2, 4...) right_side = x[..., 1::2] # Odd indices (1, 3, 5...) - new_data = right_side.clone() << 4 # Put odd indices (higher addresses) in high bits + # Put odd indices (higher addresses) in high bits + new_data = torch.bitwise_left_shift(right_side, 4) new_data[..., : left_side.shape[-1]] += left_side # Put even indices in low bits return new_data + _refuse_dtensor(input, cls.__name__) if block_size is None: block_size = 32 diff --git a/modelopt/torch/quantization/qtensor/nf4_tensor.py b/modelopt/torch/quantization/qtensor/nf4_tensor.py index 5b647d4dabc..bac74bc533c 100644 --- a/modelopt/torch/quantization/qtensor/nf4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nf4_tensor.py @@ -19,7 +19,7 @@ import torch from ..extensions import get_cuda_ext -from ..qtensor.base_qtensor import BaseQuantizedTensor +from ..qtensor.base_qtensor import BaseQuantizedTensor, _refuse_dtensor from ..utils import reduce_amax, reduce_block_padding nf4_table = torch.tensor( @@ -87,6 +87,7 @@ def quantize( Returns: tuple: Contains quantized data, input quantization config, and scale quantization config. """ + _refuse_dtensor(input, cls.__name__) cuda_ext = get_cuda_ext() # pad the input if needed @@ -116,7 +117,10 @@ def quantize( # pack the int4 weights into a uint8 tensor # packing format: w0, w1, w2, w3, w4, w5, ... # | byte | byte | byte | - packed_output_uint8 = quantized_output_uint8[::2] << 4 | quantized_output_uint8[1::2] + packed_output_uint8 = ( + torch.bitwise_left_shift(quantized_output_uint8[::2], 4) + | quantized_output_uint8[1::2] + ) # pad the scales if needed scales = reduce_block_padding(scales.view(-1), block_sizes={-1: scale_block_size}) diff --git a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py index ab394b6667f..89cf9066a80 100644 --- a/modelopt/torch/quantization/qtensor/nvfp4_tensor.py +++ b/modelopt/torch/quantization/qtensor/nvfp4_tensor.py @@ -310,6 +310,16 @@ def quantize( Returns: tuple: Contains quantized data, quantized per block scaling factor, and per tensor scaling factor. """ + if ( + isinstance(input, DTensor) + and weights_scaling_factor is not None + and not isinstance(weights_scaling_factor, DTensor) + ): + # A plain per-block scale, as a static quantizer's, covers the whole weight. + raise NotImplementedError( + "NVFP4 packs a DTensor shard with the per-block scale computed off the shard " + "(get_weights_scaling_factor), not a plain one such as a static quantizer's." + ) # Get original input shape input_shape = input.shape input_dtype = input.dtype @@ -368,7 +378,7 @@ def quantize( # Cast weights to fp4 q_weight = cls._cast_fp4(scaled_weight) # Pack weights. torch.bitwise_left_shift, not ``<< 4``: on a DTensor the operator is a silent - # no-op (torch 2.11), which would pack unrelated codes on the distributed export path. + # no-op, which would pack unrelated codes on the distributed export path. packed_weight = torch.bitwise_left_shift(q_weight[..., 1::2], 4) | q_weight[..., 0::2] return ( cls(input_shape, input_dtype, packed_weight), diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index b61b11b8740..a98fe08b4c4 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -29,8 +29,11 @@ import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq +from modelopt.torch.export.quant_utils import pack_int4_in_uint8 from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry -from modelopt.torch.quantization.qtensor import NVFP4QTensor +from modelopt.torch.quantization.config import QuantizerAttributeConfig +from modelopt.torch.quantization.nn import NVFP4StaticQuantizer +from modelopt.torch.quantization.qtensor import INT4QTensor, MXFP4QTensor, NF4QTensor, NVFP4QTensor from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan @@ -87,6 +90,45 @@ def test_nvfp4_cast_keeps_the_global_shape_of_an_uneven_shard(gloo_world_of_one) assert torch.equal(packed.to_local(), NVFP4QTensor._cast_fp4(local.clone())) +def _nvfp4_through_a_static_quantizer(weight): + whole = weight.full_tensor() if isinstance(weight, DTensor) else weight + quantizer = NVFP4StaticQuantizer( + quant_attribute_cfg=QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)} + ) + ) + quantizer.amax = whole.reshape(-1, 16).abs().amax(-1) + quantizer.global_amax = whole.abs().amax() + scale, scale_2 = NVFP4QTensor.get_weights_scaling_factor_from_quantizer(quantizer, weight) + return NVFP4QTensor.quantize(weight, 16, scale, scale_2) + + +@pytest.mark.parametrize( + "pack", + [ + lambda weight: INT4QTensor.quantize(weight, 16), + lambda weight: NF4QTensor.quantize(weight, 16, 4), + lambda weight: MXFP4QTensor.quantize(weight, 32), + lambda weight: pack_int4_in_uint8(weight, torch.ones(4, 2), 16), + _nvfp4_through_a_static_quantizer, + ], + ids=["int4", "nf4", "mxfp4", "int4-export", "nvfp4-static"], +) +def test_a_packer_that_packs_only_whole_tensors_refuses_a_shard(gloo_world_of_one, pack): + """Only NVFP4, with the per-block scale it computes off the shard, packs a DTensor shard. + + The other 4-bit packers combine their input with plain tensors, which DTensor refuses deep + inside, or gather it whole (INT4's fallback, whose ``<<`` also dropped the shift on a DTensor). + They refuse a shard up front, as NVFP4 does a plain per-block scale, a static quantizer's. + """ + torch.manual_seed(0) + weight = torch.randn(4, 32) + pack(weight) + + with pytest.raises(NotImplementedError): + pack(distribute_tensor(weight, gloo_world_of_one, [Shard(0)])) + + @pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) @pytest.mark.parametrize("layout", ["fsdp2", "expert-parallel", "linear-units"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( From 6033bb8e2a8b3c29dbaf98b1ba7317d117548660 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:44:09 +0000 Subject: [PATCH 63/67] Keep a module the weight window cannot name off the root's own key The window keys a parameter its root owns by its bare name, and _param_key took any falsy module name for the root. A module the window's name map does not know (None) then took the root's key, and the window would have handed FSDP2 that module's weight as the root's own. Only "" names the root now; an unknown module keeps its old None. key, which no parameter has, so the window skips it as before. keep_fused_experts is documented in a docstring, like the fields after it, as the marker of the no-gather export. Signed-off-by: Shengliang Xu --- modelopt/torch/export/registry.py | 4 ++-- .../torch/quantization/utils/core_utils.py | 10 +++++++--- .../export/test_sharded_weight_packing.py | 19 +++++++++++++++++++ 3 files changed, 28 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index ca99532dfae..a0ad2c1deeb 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -64,9 +64,9 @@ class ExportContext: model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then fail loudly instead of guessing.""" - # Distributed (no-gather) export: the handlers pack while the model is still sharded, an - # FSDP2-owned weight inside its unit's weight-update window. keep_fused_experts: bool = False + """Marks the distributed no-gather export, whose handlers pack while the model is still sharded: + an FSDP2-owned weight is packed inside its unit's weight-update window.""" module_names: "ModuleNames | None" = None """``model``'s name maps, built once per export by the first handler that looks a module up diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index ca7168ac1da..8bb3ac69288 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -946,9 +946,13 @@ def get_prefixed_param_names(parent_model, target_module, param_index=None): return match[1].rpartition(".")[0] if match is not None else None -def _param_key(module_name: str, param_name: str) -> str: - """The full name of ``module_name``'s parameter ``param_name``; ``""`` names the root module.""" - return f"{module_name}.{param_name}" if module_name else param_name +def _param_key(module_name: str | None, param_name: str) -> str: + """The full name of ``module_name``'s parameter ``param_name``; ``""`` names the root module. + + A module the name map does not know (``None``) keeps the ``None.`` key no parameter has, so it is + skipped rather than taken for the root. + """ + return param_name if module_name == "" else f"{module_name}.{param_name}" def create_fsdp_param_mapping(fsdp_param_list, model): diff --git a/tests/unit/torch/export/test_sharded_weight_packing.py b/tests/unit/torch/export/test_sharded_weight_packing.py index a98fe08b4c4..4893fefcfc8 100644 --- a/tests/unit/torch/export/test_sharded_weight_packing.py +++ b/tests/unit/torch/export/test_sharded_weight_packing.py @@ -34,6 +34,7 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn import NVFP4StaticQuantizer from modelopt.torch.quantization.qtensor import INT4QTensor, MXFP4QTensor, NF4QTensor, NVFP4QTensor +from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, module_name_maps from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan @@ -211,6 +212,24 @@ def test_a_module_joins_its_units_held_window_once(gloo_world_of_one): assert torch.equal(sd["0.0.weight"].full_tensor(), ref.weight) +def test_a_module_the_window_cannot_name_never_takes_the_roots_own_weight(gloo_world_of_one): + """A unit that owns its weight directly keys it by its bare name in the window. + + A module the window's name map does not know must not take that key when the window closes: it + would be handed to FSDP2 as the unit's own weight. + """ + torch.manual_seed(0) + unit, stranger = nn.Linear(32, 16, bias=False), nn.Linear(32, 16, bias=False) + fully_shard(unit, mesh=gloo_world_of_one) + want = unit.weight.full_tensor().clone() + modules = [unit] + + with fsdp2_aware_weight_update(unit, modules, names=module_name_maps(unit)): + modules.append(stranger) # as a held window takes a module in + + assert torch.equal(unit.weight.full_tensor(), want) + + class _Unit(nn.Module): """An FSDP2 unit of two linears fed the same input; 17 rows shard unevenly.""" From ec2741cdef3965080541fadca510d5d4c5514ce6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:56:32 +0000 Subject: [PATCH 64/67] Test the sharded fold over NCCL on every visible GPU Each rank folds the experts it holds on its GPU, at their global offsets, and the folds, gathered over NCCL, split into the per-expert export, an uneven shard included. Under expert parallelism the ranks decide on each projection together: their quantizer states go over Gloo in the cpu:gloo,cuda:nccl group the export runs on, and over NCCL in a group of NCCL alone. Every collective runs on a group that gives up after a minute. The test skips below two GPUs. Signed-off-by: Shengliang Xu --- .../test_fused_experts_keep_fused_gpu.py | 211 ++++++++++++++++++ 1 file changed, 211 insertions(+) create mode 100644 tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py diff --git a/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py new file mode 100644 index 00000000000..b2407d0f63a --- /dev/null +++ b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py @@ -0,0 +1,211 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The keep-fused fold of sharded fused experts on every visible GPU, over NCCL.""" + +import copy +import re +from datetime import timedelta + +import pytest +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor, Shard, distribute_tensor + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.moe_utils import ( + _dtensor_dim0_offset, + _export_fused_experts, + _export_fused_experts_keep_fused, + split_fused_experts_state_dict, +) + +HIDDEN_DIM = 32 +INTERMEDIATE_DIM = 16 +PREFIX = "model.layers.0.mlp.experts" + + +class _FusedExperts(nn.Module): + """Fused gate+up experts in the transformers 5.x layout (Qwen3-MoE / Mixtral).""" + + def __init__(self, num_experts): + super().__init__() + self.num_experts = num_experts + self.gate_up_proj = nn.Parameter( + torch.randn(num_experts, 2 * INTERMEDIATE_DIM, HIDDEN_DIM) * 0.02 + ) + self.down_proj = nn.Parameter(torch.randn(num_experts, HIDDEN_DIM, INTERMEDIATE_DIM) * 0.02) + self.act_fn = nn.SiLU() + + def forward(self, hidden_states, top_k_index, top_k_weights): + out = torch.zeros_like(hidden_states) + mask = F.one_hot(top_k_index, num_classes=self.num_experts).permute(2, 1, 0) + for expert_idx in torch.greater(mask.sum(dim=(-1, -2)), 0).nonzero()[:, 0]: + top_k_pos, token_idx = torch.where(mask[expert_idx]) + gate, up = F.linear(hidden_states[token_idx], self.gate_up_proj[expert_idx]).chunk( + 2, dim=-1 + ) + current = F.linear(self.act_fn(gate) * up, self.down_proj[expert_idx]) + out.index_add_(0, token_idx, current * top_k_weights[token_idx, top_k_pos, None]) + return out + + +def _quantized_experts(quant_cfg, num_experts): + """Fused experts calibrated on one token each, moved to this rank's GPU to be folded. + + Every rank builds and calibrates the same seeded module, so all hold identical quantizers. + """ + torch.manual_seed(0) + model = nn.Module() + model.experts = _FusedExperts(num_experts) + hidden = torch.randn(num_experts, HIDDEN_DIM) + tokens = (torch.arange(num_experts).unsqueeze(1), torch.ones(num_experts, 1)) + mtq.quantize(model, quant_cfg, lambda m: m.experts(hidden, *tokens)) + return model.experts.cuda() + + +def _state(experts): + return {f"{PREFIX}.{key}": value for key, value in experts.state_dict().items()} + + +def _raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).cpu() + + +def _whole(value): + """A sharded value gathered whole, one-byte values moved as bytes.""" + if not isinstance(value, DTensor): + return value + if value.element_size() != 1: + return value.full_tensor() + raw = DTensor.from_local( + value.to_local().view(torch.uint8), + value.device_mesh, + value.placements, + run_check=False, + shape=value.shape, + stride=value.stride(), + ) + return raw.full_tensor().view(value.dtype) + + +def _sharded(experts, mesh): + experts = copy.deepcopy(experts) + for proj in ("gate_up_proj", "down_proj"): + weight = distribute_tensor(getattr(experts, proj).data, mesh, [Shard(0)]) + setattr(experts, proj, nn.Parameter(weight, requires_grad=False)) + return experts + + +def _fold_job(rank, world_size, backend): + # Every collective below runs on this group, which gives up after a minute. + group = dist.new_group(backend=backend, timeout=timedelta(seconds=60)) + try: + # A mesh compares equal to any other over the same ranks and names, whatever its group, so + # DTensor's caches could hand a later job a mesh over this group once it is destroyed. + mesh = DeviceMesh.from_group(group, "cuda", mesh_dim_names=(f"ep{group.group_name}",)) + # Every rank keeps every expert's quantizers, as under FSDP2. Two experts a rank, then one + # expert fewer: the last rank holds one, and the folds must keep the global expert count. + cases = [(cfg, 2 * world_size) for cfg in (mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG)] + cases.append((mtq.NVFP4_DEFAULT_CFG, 2 * world_size - 1)) + for quant_cfg, num_experts in cases: + experts = _quantized_experts(quant_cfg, num_experts) + reference = copy.deepcopy(experts) + _export_fused_experts(reference, torch.bfloat16) + plain = copy.deepcopy(experts) + _export_fused_experts_keep_fused(plain, torch.bfloat16) + plain = _state(plain) + sharded = _sharded(experts, mesh) + + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + + fused = _state(sharded) + assert set(fused) == set(plain) + offset = _dtensor_dim0_offset(fused[f"{PREFIX}.gate_up_proj"]) + for key, value in fused.items(): + assert value.device.type == "cuda", key + if key.endswith("_input_scale"): # one scalar shared by every expert + assert not isinstance(value, DTensor), key + assert torch.equal(_raw_bytes(value), _raw_bytes(plain[key])), key + continue + assert isinstance(value, DTensor) and value.placements == (Shard(0),), key + assert value.shape == plain[key].shape, key + rows = plain[key][offset : offset + value.to_local().shape[0]] + assert torch.equal(_raw_bytes(value.to_local()), _raw_bytes(rows)), key + # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. + split = split_fused_experts_state_dict({k: _whole(v) for k, v in fused.items()}) + want = _state(reference) + assert set(split) == set(want) + for key, value in want.items(): + assert split[key].dtype == value.dtype and split[key].shape == value.shape, key + assert torch.equal(_raw_bytes(split[key]), _raw_bytes(value)), key + + # A layout that leaves a rank no expert is refused on every rank. + lone = _sharded(_quantized_experts(mtq.FP8_DEFAULT_CFG, world_size - 1), mesh) + with pytest.raises(NotImplementedError, match="leaves a rank none"): + _export_fused_experts_keep_fused(lone, torch.bfloat16) + + # Under expert parallelism a rank keeps only its own experts' quantizers, so the ranks gather + # each projection's quantizer states over the group before they fold it. + experts = _quantized_experts(mtq.NVFP4_DEFAULT_CFG, 2 * world_size) + plain = copy.deepcopy(experts) + _export_fused_experts_keep_fused(plain, torch.bfloat16) + plain = _state(plain) + mine = slice(2 * rank, 2 * rank + 2) + + def keeping_own_quantizers(): + sharded = _sharded(experts, mesh) + for proj in ("gate_up_proj", "down_proj"): + name = f"{proj}_weight_quantizers" + setattr(sharded, name, nn.ModuleList(list(getattr(sharded, name))[mine])) + return sharded + + sharded = keeping_own_quantizers() + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + for key, value in _state(sharded).items(): + local = value.to_local() if isinstance(value, DTensor) else value + want = plain[key][mine] if isinstance(value, DTensor) else plain[key] + assert torch.equal(_raw_bytes(local), _raw_bytes(want)), key + + # The last rank disabling one or both of its experts' quantizers is refused on every rank, + # naming the global experts; every expert disabled leaves the projection as it is on all. + last = 2 * world_size - 1 + for disabled in ([last], [last - 1, last], list(range(2 * world_size))): + sharded = keeping_own_quantizers() + for expert, quantizer in enumerate(sharded.down_proj_weight_quantizers, mine.start): + if expert in disabled: + quantizer.disable() + if len(disabled) < 2 * world_size: + match = re.escape(f"quantizers {disabled} of {2 * world_size}") + with pytest.raises(NotImplementedError, match=match): + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + else: + _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert hasattr(sharded, "gate_up_proj_weight_scale") + assert not hasattr(sharded, "down_proj_weight_scale") + finally: + dist.destroy_process_group(group) + + +@pytest.mark.parametrize("backend", [None, "nccl"], ids=["objects-over-gloo", "objects-over-nccl"]) +def test_ranks_fold_their_own_experts_over_nccl(need_2_gpus, dist_workers, backend): + """Each rank folds the experts it holds on its GPU, at their global offsets, and the folds, + gathered over NCCL, split into the per-expert export, an uneven shard included. Under expert + parallelism the ranks decide on each projection together: in the export's ``cpu:gloo,cuda:nccl`` + group their quantizer states go over Gloo, in a group of NCCL alone over NCCL.""" + dist_workers.run(_fold_job, backend=backend) From 46742ffd24edf602cc6ac56544590f9b268e4358 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:06:23 +0000 Subject: [PATCH 65/67] Say where the held weight windows take effect, and what ends a hung GPU test _process_quantized_modules enters held_fsdp2_weight_windows around its dispatch loop, but its context leaves keep_fused_experts unset and its assert keeps sharded models out, so no window opens there yet. Its comment now says so instead of describing a live property. The GPU test's comment names what ends a rank left waiting in a window: pytest's per-test timeout, after which the worker pool kills its ranks. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 4 +++- tests/gpu/torch/export/test_sharded_weight_packing_gpu.py | 2 ++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 6ec8f7b714d..e6d75612869 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1002,7 +1002,9 @@ def _process_quantized_modules( model_type=hf_model_type(model), ) - # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. + # For the no-gather export, which is to set keep_fused_experts on this context and let a sharded + # model past the assert above: each FSDP2 unit's modules are then packed in one weight window, + # held across them. Until then no window opens here. with _hf_export_handlers.held_fsdp2_weight_windows(ctx): for name, sub_module in model.named_modules(): _dispatch_export_handler(name, sub_module, ctx) diff --git a/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py index 9284f755e48..74b615ba8f6 100644 --- a/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py +++ b/tests/gpu/torch/export/test_sharded_weight_packing_gpu.py @@ -92,6 +92,8 @@ def _packing_job(rank, world_size): set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) ctx = ExportContext(model=model, dtype=torch.bfloat16, keep_fused_experts=True) + # A rank left waiting in a window's unshard fails the test at pytest's per-test timeout, after + # which the worker pool kills its ranks. with handlers.held_fsdp2_weight_windows(ctx): for name, linear in linears: ExportModuleRegistry.match(linear)(name, linear, ctx) From fc6d99829d808ba2f042233dcebc8bfcc6ba8567 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:21:17 +0000 Subject: [PATCH 66/67] Say how the per-expert split numbers the experts it is given It numbers each module's experts from 0 by their row in the tensors it receives, so whole tensors get their global numbers, and a caller that splits a rank's local shard shifts each index by the shard's offset. "One module's shard at a time" meant one module's tensors, which read as a rank's shard. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index da5122a59ea..37d44465f40 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -669,7 +669,9 @@ def split_fused_experts_state_dict(state_dict: dict) -> dict: checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and the input keeps its fused tensors, so the experts are held twice until it returns: pass one - module's shard at a time, not a whole model's. + module's tensors at a time, not a whole model's. Each module's experts are numbered from 0 by + their row in the tensors given, which must be plain: whole ones get their global numbers, and a + caller splitting a rank's local shard (``to_local()``) shifts each by the shard's global offset. Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` From 945271deb04d379db8017165e2650516be5e0006 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:48:34 +0000 Subject: [PATCH 67/67] Refuse fused projections whose expert counts differ, and check the GPU test's keys The split numbers each module's experts by the first projection, so a down_proj holding more experts lost the extra ones unwritten, and one holding fewer failed deep in the loop. It is refused up front, naming the module. The scale slicer's comment says that only the gate|up cut reads its row count. The GPU test's expert-parallel comparison checks the folded keys too, as the two-rank unit test does. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 8 +++++--- .../torch/export/test_fused_experts_keep_fused_gpu.py | 1 + .../unit/torch/export/test_fused_experts_keep_fused.py | 10 ++++++++++ 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 37d44465f40..07e0ddd4129 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -695,6 +695,8 @@ def split_fused_experts_state_dict(state_dict: dict) -> dict: first_w = state_dict[f"{prefix}.{first}"] down_w = state_dict[f"{prefix}.down_proj"] n_experts = first_w.shape[0] + if down_w.shape[0] != n_experts: + raise ValueError(f"{prefix}: {n_experts} {first} experts, {down_w.shape[0]} down_proj.") def _scales(proj_name): ws = state_dict.get(f"{prefix}.{proj_name}_weight_scale") @@ -707,9 +709,9 @@ def _scales(proj_name): down_ws, down_in, down_ws2 = _scales("down_proj") def _slice_weight_scale(ws, expert_idx, row_slice, rows): - # A per-expert scalar as is; else the expert's scale rows covering ``row_slice`` of its - # ``rows`` weight rows, in proportion: an INT4 weight packs two rows a byte while its - # scales stay per row, and an FP8 2-D block scale covers 128 rows. + # A per-expert scalar as is; else the expert's scale rows covering ``row_slice``, a cut of + # the fused gate|up's ``rows`` rows, in proportion (only gate and up are cut): an INT4 + # weight packs two rows a byte while its scales stay per row, an FP8 2-D block 128 rows. if ws is None: return None if ws.dim() <= 1: # (E,) one scalar per expert diff --git a/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py index b2407d0f63a..83e2e974f3f 100644 --- a/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py +++ b/tests/gpu/torch/export/test_fused_experts_keep_fused_gpu.py @@ -177,6 +177,7 @@ def keeping_own_quantizers(): sharded = keeping_own_quantizers() _export_fused_experts_keep_fused(sharded, torch.bfloat16) + assert set(_state(sharded)) == set(plain) for key, value in _state(sharded).items(): local = value.to_local() if isinstance(value, DTensor) else value want = plain[key][mine] if isinstance(value, DTensor) else plain[key] diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index 5a7df9003ad..0a516db55aa 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -536,6 +536,16 @@ def test_split_refuses_a_scale_block_gate_and_up_would_share(rows, scale_rows): split_fused_experts_state_dict(sd) +def test_split_refuses_projections_holding_different_expert_counts(): + """The split numbers experts by the first projection, so a down_proj expert beyond its count + would be dropped unwritten.""" + sd = _fused_state_dict(gated=True) + sd[f"{PREFIX}.down_proj"] = torch.zeros(LOCAL_EXPERTS + 1, 6, 4).to(torch.float8_e4m3fn) + + with pytest.raises(ValueError, match="down_proj"): + split_fused_experts_state_dict(sd) + + def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): fused = torch.zeros(2, 4, 4) sd = {