nvidia-modelopt 0.35.0__py3-none-any.whl

modelopt/torch/sparsity/mode.py

@@ -0,0 +1,203 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Sparsity mode descriptor."""
+
+from torch import nn
+
+from modelopt.torch.opt.config import ModeloptBaseConfig
+from modelopt.torch.opt.conversion import ApplyModeError
+from modelopt.torch.opt.dynamic import DynamicSpace
+from modelopt.torch.opt.mode import (
+    ConvertEntrypoint,
+    ConvertReturnType,
+    MetadataDict,
+    ModeDescriptor,
+    RestoreEntrypoint,
+    UpdateEntrypoint,
+    _ModeRegistryCls,
+)
+from modelopt.torch.opt.searcher import BaseSearcher
+from modelopt.torch.utils import compare_dict, unwrap_model
+
+from .config import ExportSparseConfig, SparseGPTConfig, SparseMagnitudeConfig
+from .magnitude import MagnitudeSearcher
+from .module import SpDMRegistry
+from .sparsegpt import SparseGPTSearcher
+
+SparsityModeRegistry = _ModeRegistryCls("sparsity")
+
+
+def convert_sparse_model(model: nn.Module, config: ModeloptBaseConfig) -> ConvertReturnType:
+    """Function for converting a model to a sparsity meta-model."""
+    # we use the search space utility here with a custom registry to convert the model
+    dynamic_space = DynamicSpace(model)
+    dynamic_space.convert_to_dynamic(config.model_dump(), SpDMRegistry)
+
+    return dynamic_space.model, {"subnet_config": DynamicSpace(model).config()}
+
+
+def restore_sparse_model(
+    model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict
+) -> nn.Module:
+    """Function for restoring a previously convert model to a sparsity meta-model."""
+    model, _ = convert_sparse_model(model, config)
+
+    if "subnet_config" in metadata:
+        DynamicSpace(model).select(metadata["subnet_config"])
+
+    return model
+
+
+def update_sparse_metadata(
+    model: nn.Module, config: ModeloptBaseConfig, metadata: MetadataDict
+) -> None:
+    """Update subnet config to current subnet config of model."""
+    metadata["subnet_config"] = DynamicSpace(model).config()
+
+
+def export_sparse(model: nn.Module, config: ExportSparseConfig) -> ConvertReturnType:
+    """Export a sparse model to a regular model."""
+    # sanity check to avoid DP/DDP here in the entrypoint
+    model = unwrap_model(model, raise_error=True)
+
+    # store config from model if we can find it for a future convert/restore process
+    metadata = {"subnet_config": DynamicSpace(model).config()}
+
+    # export model in-place
+    model = DynamicSpace(model).export(SpDMRegistry)
+
+    return model, metadata
+
+
+def restore_export_sparse(
+    model: nn.Module, config: ExportSparseConfig, metadata: MetadataDict
+) -> nn.Module:
+    """Restore & export a sparse model to a regular model."""
+    # select activated/deactivated sparse modules
+    DynamicSpace(model).select(metadata["subnet_config"])
+
+    # run export
+    model, metadata_new = export_sparse(model, config)
+
+    # double check metadata
+    unmatched_keys = compare_dict(metadata, metadata_new)
+    if unmatched_keys:
+        raise ApplyModeError(f"Unmatched metadata={unmatched_keys}!")
+
+    return model
+
+
+@SparsityModeRegistry.register_mode
+class SparseMagnitudeModeDescriptor(ModeDescriptor):
+    """Class to define and describe magnitude-based sparsification."""
+
+    @property
+    def name(self) -> str:
+        """Returns the name of the mode."""
+        return "sparse_magnitude"
+
+    @property
+    def config_class(self) -> type[ModeloptBaseConfig]:
+        """Specifies the config class for the mode."""
+        return SparseMagnitudeConfig
+
+    @property
+    def next_modes(self) -> set[str] | None:
+        """Specifies the next modes for the mode."""
+        return {"export_sparse", "kd_loss", "quantize"}
+
+    @property
+    def export_mode(self) -> str | None:
+        """The mode that corresponds to the export mode of this mode."""
+        return "export_sparse"
+
+    @property
+    def search_algorithm(self) -> type[BaseSearcher]:
+        """Specifies the search algorithm for the mode."""
+        return MagnitudeSearcher
+
+    @property
+    def convert(self) -> ConvertEntrypoint:
+        """The mode's entrypoint for converting a model."""
+        return convert_sparse_model
+
+    @property
+    def restore(self) -> RestoreEntrypoint:
+        """The mode's entrypoint for restoring a model."""
+        return restore_sparse_model
+
+    @property
+    def update_for_save(self) -> UpdateEntrypoint:
+        """The mode's entrypoint for updating the models metadata."""
+        return update_sparse_metadata
+
+    @property
+    def update_for_new_mode(self) -> UpdateEntrypoint:
+        """The mode's entrypoint for updating the models metadata."""
+        return update_sparse_metadata
+
+
+@SparsityModeRegistry.register_mode
+class SparseGPTModeDescriptor(SparseMagnitudeModeDescriptor):
+    """Class to define and describe sparsification based on SparseGPT."""
+
+    @property
+    def name(self) -> str:
+        """Returns the name of the mode."""
+        return "sparsegpt"
+
+    @property
+    def config_class(self) -> type[ModeloptBaseConfig]:
+        """Specifies the config class for the mode."""
+        return SparseGPTConfig
+
+    @property
+    def search_algorithm(self) -> type[BaseSearcher]:
+        """Specifies the search algorithm for the mode."""
+        return SparseGPTSearcher
+
+
+@SparsityModeRegistry.register_mode
+class ExportSparseModeDescriptor(ModeDescriptor):
+    """Class to describe the ``"export_sparse"`` mode.
+
+    The properties of this mode can be inspected via the source code.
+    """
+
+    @property
+    def name(self) -> str:
+        """Returns the value (str representation) of the mode."""
+        return "export_sparse"
+
+    @property
+    def config_class(self) -> type[ModeloptBaseConfig]:
+        """Specifies the config class for the mode."""
+        return ExportSparseConfig
+
+    @property
+    def is_export_mode(self) -> bool:
+        """Specifies if this mode is an export mode."""
+        return True
+
+    @property
+    def convert(self) -> ConvertEntrypoint:
+        """The mode's entrypoint for converting a model."""
+        return export_sparse
+
+    @property
+    def restore(self) -> RestoreEntrypoint:
+        """The mode's entrypoint for restoring a model."""
+        return restore_export_sparse

modelopt/torch/sparsity/module.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Dynamic class for all sparse modules."""
+
+import torch
+from torch import nn
+
+from modelopt.torch.opt.dynamic import DynamicModule, _DMRegistryCls
+from modelopt.torch.opt.hparam import Hparam
+
+__all__ = ["SpDMRegistry", "SparseModule"]
+
+
+SpDMRegistry = _DMRegistryCls(prefix="Sparse")  # global instance for the sparsity registry
+
+
+@SpDMRegistry.register({nn.Linear: "nn.Linear", nn.Conv2d: "nn.Conv2d"})
+class SparseModule(DynamicModule):
+    """Base dynamic class for all sparse modules."""
+
+    @staticmethod
+    def _get_weight(mod: "SparseModule", weight: torch.Tensor) -> torch.Tensor:
+        if mod.is_sparse and mod._weight_mask is not None:
+            masked_weight = weight * mod._weight_mask
+            # Quick workaround for the custom attribute for Megatron.
+            # TODO: maybe we need a more general way for customized attributes
+            if hasattr(weight, "main_grad"):
+                masked_weight.main_grad = weight.main_grad
+            return masked_weight
+        return weight
+
+    def _setup(self):
+        # define hparam to check if sparsity is activated
+        hp = Hparam([0, -1], original=0)
+        hp.active = 0
+        self._register_hparam("is_sparse", hp)
+
+        # define the sparse mask here (don't pre-allocate memory to maximize memory savings)
+        self._register_temp_attribute("_weight_mask", None, lambda m, n, v: m.register_buffer(n, v))
+
+        # register dynamic attributes of the class
+        self._register_dynamic_attribute("weight", self._get_weight)
+
+    def modify(self, *args, **kwargs):
+        """Initialize the sparsity mask when this is called.
+
+        Note that for any module that is not frozen via ``None`` in the rules, this function will be
+        called. Hence, we use this function to initialize the sparsity mask only when necessary.
+        """
+        hp = self.get_hparam("is_sparse")
+        if -1 in hp.choices and self._weight_mask is None:
+            hp.active = -1
+            self._weight_mask = torch.ones_like(self.weight, dtype=torch.bool)
+
+    def set_mask(self, value: torch.BoolTensor | None):
+        """Set the active sparse mask of the module weights."""
+        if value is None:
+            self._weight_mask = None
+            return
+
+        # sanity checks on the mask
+        w_shape = self.weight.shape
+        assert value is not None, "Mask cannot be None."
+        assert value.shape == w_shape, f"Mask must have shape {w_shape}, got {value.shape} instead."
+        assert value.dtype == torch.bool, f"Mask must be of type torch.bool, but got {value.dtype}."
+
+        # assign mask
+        with torch.no_grad():
+            if torch.all(value):
+                self._weight_mask = None
+            elif self._weight_mask is None:
+                self._weight_mask = value.detach().clone().to(self.weight.device)
+            else:
+                self._weight_mask.copy_(value.to(self._weight_mask.device))

modelopt/torch/sparsity/plugins/__init__.py

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Handles sparsity plugins for third-party modules.
+
+Currently, we support plugins for
+
+- :meth:`megatron<modelopt.torch.sparsity.plugins.megatron>`
+
+"""
+
+from modelopt.torch.utils import import_plugin
+
+with import_plugin("megatron"):
+    from .megatron import *

modelopt/torch/sparsity/plugins/megatron.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Support sparsify and save/resore for Megatron."""
+
+import megatron.core.transformer.mlp as megatron_mlp
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
+from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint
+
+from modelopt.torch.opt.plugins.megatron import _MegatronMLP
+
+from ..config import SparseGPTConfig, SparseMagnitudeConfig
+from ..module import SparseModule, SpDMRegistry
+
+
+class _MegatronParallelLinear(SparseModule):
+    def _get_shard_axis_dict(self, state_dict):
+        raise NotImplementedError
+
+    def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None):
+        sharded_state_dict = super().sharded_state_dict(prefix, sharded_offsets)
+
+        sparse_state_dict = {
+            k: v
+            for k, v in self.state_dict(prefix="", keep_vars=True).items()
+            if k == "_weight_mask"
+        }
+
+        sharded_axis_dict = self._get_shard_axis_dict(sparse_state_dict)
+
+        if sparse_state_dict:
+            sharded_state_dict.update(
+                **make_sharded_tensors_for_checkpoint(
+                    sparse_state_dict, prefix, sharded_axis_dict, sharded_offsets
+                )
+            )
+
+        return sharded_state_dict
+
+
+@SpDMRegistry.register(
+    {ColumnParallelLinear: "megatron.core.tensor_parallel.layers.ColumnParallelLinear"}
+)
+class _MegatronColumnParallelLinear(_MegatronParallelLinear):
+    def _get_shard_axis_dict(self, state_dict):
+        return {"_weight_mask": 0}
+
+
+@SpDMRegistry.register(
+    {RowParallelLinear: "megatron.core.tensor_parallel.layers.RowParallelLinear"}
+)
+class _MegatronRowParallelLinear(_MegatronParallelLinear):
+    def _get_shard_axis_dict(self, state_dict):
+        return {"_weight_mask": 1}
+
+
+@SpDMRegistry.register({megatron_mlp.MLP: "megatron.core.transformer.mlp.MLP"})
+class _SparseMegatronMLP(_MegatronMLP):
+    """Module to support special handling of `linear_fc1` in `sharded_state_dict()` of MCore `MLP`."""
+
+    _modelopt_state_keys = [r"\._weight_mask$"]
+
+
+def _get_extra_rules():
+    """Get the extra rules for megatron."""
+    return {
+        "megatron.core.tensor_parallel.layers.ColumnParallelLinear": {
+            "*": {},
+            "*output_layer*": None,
+        },
+        "megatron.core.tensor_parallel.layers.RowParallelLinear": {
+            "*": {},
+            "*output_layer*": None,
+        },
+        "megatron.core.transformer.mlp.MLP": {},
+    }
+
+
+# Update the default rules
+SparseMagnitudeConfig.register_default(_get_extra_rules())
+SparseGPTConfig.register_default(_get_extra_rules())

modelopt/torch/sparsity/searcher.py

@@ -0,0 +1,84 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+"""Searcher interface for sparsity algorithms."""
+
+from abc import abstractmethod
+from collections.abc import Iterator
+
+import torch
+import torch.nn as nn
+
+from modelopt.torch.opt.searcher import BaseSearcher, SearchConfig, SearchStateDict
+from modelopt.torch.utils import print_rank_0
+
+from . import magnitude as asp
+from .module import SparseModule
+
+
+class BaseSparseSearcher(BaseSearcher):
+    """A generic sparse mask searching algorithm."""
+
+    _pattern_2_4 = "2:4 sparsity"
+
+    @property
+    def default_search_config(self) -> SearchConfig:
+        """Get the default config for the searcher."""
+        return {**super().default_search_config, "pattern": self._pattern_2_4}
+
+    @property
+    def default_state_dict(self) -> SearchStateDict:
+        """Return default state dict."""
+        return {}
+
+    def sanitize_search_config(self, config: SearchConfig | None) -> SearchStateDict:
+        """Sanitize the search config dict."""
+        config_sanitized = super().sanitize_search_config(config)
+
+        # sanity check of sparsity format
+        is_nm_prune, n, m = asp.get_nmprune_info(config_sanitized["pattern"])
+        assert is_nm_prune and n == 2 and m == 4, (
+            f"Unsupported pattern {self.config['pattern']} for sparsity"
+        )
+
+        return config_sanitized
+
+    @abstractmethod
+    def _check_weight_size(self, weight, mod_name) -> bool:
+        """Check if the weight size is supported by the algorithm."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _compute_mask(self, module: SparseModule) -> torch.BoolTensor:
+        """Compute the mask and update weight for a given module."""
+        raise NotImplementedError
+
+    def _named_sparsifiable_modules(self) -> Iterator[tuple[str, nn.Module]]:
+        """Get the named sparsifiable modules."""
+        for name, module in self.model.named_modules():
+            if (
+                isinstance(module, SparseModule)
+                and module.is_sparse
+                and self._check_weight_size(module.weight, name)
+            ):
+                yield name, module
+
+    def run_search(self):
+        """Search for sparse mask."""
+        for name, module in self._named_sparsifiable_modules():
+            # compute the mask (and potentially weight update inside compute_mask)
+            print_rank_0(f"Searching for sparse mask and weight update for module {name}.")
+            with torch.no_grad():
+                module.set_mask(self._compute_mask(module))