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 <[email protected]>
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.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.

0 commit comments

Comments
 (0)