[OMNIML-6081] [1/2] Pack sharded weights in FSDP2 weight-update windows, and NVFP4 DTensor shards

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,86 @@ 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.
+
+    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.
+        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:
+        # 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)
+    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 +130,12 @@ 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 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)
+        return
     _export_quantized_weight(module, ctx.dtype, weight_name)
 
 

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)

modelopt/torch/export/registry.py

@@ -28,10 +28,15 @@
 
 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.export.hf_export_handlers import _HeldWindow
+    from modelopt.torch.quantization.utils.core_utils import ModuleNames
+
 __all__ = [
     "ExportContext",
     "ExportHandler",
@@ -59,6 +64,22 @@ class ExportContext:
     model's spec in ``modelopt.torch.models``. ``None`` means unknown: spec lookups then
     fail loudly instead of guessing."""
 
+    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
+    (: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: "_HeldWindow | None" = None
+    """The FSDP2 unit's weight-update window the handlers hold open across its modules."""
+
 
 ExportHandler = Callable[[str, nn.Module, ExportContext], None]
 

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,12 @@ 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)
+    # 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)
 
 
 def _prepare_model_for_export(model, dtype, is_modelopt_qlora):

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.
 

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)

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
 

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})