Skip to content

Commit 9232e2f

Browse files
authored
[1/n] Reorganize sparsity module to separate weight and attention sparsity (#517)
## What does this PR do? **Type of change:** ? <!-- Use one of the following: Bug fix, new feature, new example, new tests, documentation. --> **Overview:** ? Reorganizes the `modelopt/torch/sparsity` module to prepare for adding attention sparsity support. All existing weight sparsity functionality (source code and tests) is moved into a dedicated `weight_sparsity` subdirectory while maintaining full backward compatibility. ``` # These imports still work exactly as before from modelopt.torch.sparsity import sparsify, magnitude_prune from modelopt.torch.sparsity.mode import SparsityMode ``` ## Usage <!-- You can potentially add a usage example below. --> ```python # Add a code snippet demonstrating how to use this ``` ## Testing <!-- Mention how have you tested your change if applicable. --> ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No <!--- If No, explain why. --> - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> Signed-off-by: Kai Xu <kaix@nvidia.com>
1 parent ca94c96 commit 9232e2f

File tree

15 files changed

+34
-6
lines changed

15 files changed

+34
-6
lines changed

modelopt/torch/sparsity/__init__.py

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,12 @@
1313
# See the License for the specific language governing permissions and
1414
# limitations under the License.
1515

16-
"""API for sparsification algorithms."""
16+
"""API for sparsification algorithms.
1717
18-
from . import mode, module, plugins
19-
from .sparsification import *
18+
This module provides access to both weight sparsity and attention sparsity algorithms.
19+
For backward compatibility, weight sparsity APIs are re-exported at the module level.
20+
"""
21+
22+
# Import weight sparsity for backward compatibility
23+
from .weight_sparsity import mode, module, plugins
24+
from .weight_sparsity.sparsification import *
Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2+
# SPDX-License-Identifier: Apache-2.0
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
"""API for weight sparsification algorithms."""
17+
18+
from . import mode, module, plugins
19+
20+
# Explicitly expose commonly used items
21+
from .mode import SparsityModeRegistry
22+
from .module import SparseModule, SpDMRegistry
23+
from .sparsification import *
File renamed without changes.

0 commit comments

Comments
 (0)