From 8674b73d7ceebe696ae040d97385cff6f6526665 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 28 Oct 2026 23:37:15 +0000 Subject: [PATCH 1/9] 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 2/9] 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 3/9] 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 4/9] 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 5/9] 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 6/9] 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 d568dacad398bf2566c598dc5d3a92047ef00644 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 17:43:45 +0000 Subject: [PATCH 7/9] 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 8/9] 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 46742ffd24edf602cc6ac56544590f9b268e4358 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 18:06:23 +0000 Subject: [PATCH 9/9] 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)