Add pruning checkpoints for the compress algorithm

modelopt/torch/_compress/activation_scoring/score_pruning_activations.py

@@ -18,13 +18,13 @@
 import hydra
 import torch
 from omegaconf import DictConfig
-from modelopt.torch._compress.utils.parsing import format_global_config
 
 from modelopt.torch._compress.tools.hydra_utils import register_hydra_resolvers
 from modelopt.torch._compress.tools.logger import mprint
 from modelopt.torch._compress.tools.runtime import BaseRuntime, NativeDdpRuntime
 from modelopt.torch._compress.tools.validate_model import validate_model
 from modelopt.torch._compress.utils.dist_utils import is_distributed
+from modelopt.torch._compress.utils.parsing import format_global_config
 
 
 def has_checkpoint_support(activation_hooks_kwargs: dict) -> bool:

modelopt/torch/_compress/compress.py

@@ -23,12 +23,12 @@
 import build_library_and_stats
 import mip_and_realize_models
 import pruning_ckpts
-import modelopt.torch._compress.activation_scoring.score_pruning_activations as score_pruning_activations
 import scoring
 from omegaconf import DictConfig
-from modelopt.torch._compress.tools.runtime import IRuntime
 
+import modelopt.torch._compress.activation_scoring.score_pruning_activations as score_pruning_activations
 from modelopt.torch._compress.tools.hydra_utils import initialize_hydra_config_for_dir
+from modelopt.torch._compress.tools.runtime import IRuntime
 
 
 def compress(

modelopt/torch/_compress/nas/plugins/compress_nas_plugin.py

@@ -25,11 +25,11 @@
 
 import build_library_and_stats
 import mip_and_realize_models
-import pruning_ckpts
 import scoring
 import torch
 from torch import nn
 
+import modelopt.torch._compress.pruning.pruning_ckpts as pruning_ckpts
 from modelopt.torch._compress.activation_scoring import score_pruning_activations
 from modelopt.torch._compress.decilm.converters.convert_llama3_to_decilm import (
     convert_llama3_to_decilm,

modelopt/torch/_compress/sewing_kit/__init__.py

@@ -13,22 +13,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # mypy: ignore-errors
+
 from .core import (
-    Needle,
+    CantResolveNodeDependenciesException,
+    ConstantTarget,
+    ExternalTarget,
+    FunctionTarget,
+    InputsLoopFoundException,
     KnotException,
     LoopFoundException,
-    InputsLoopFoundException,
+    ModuleTarget,
     MultipleExternalNodesException,
+    Needle,
     OnlyInternalNodesException,
     OutputsLoopFoundException,
-    ExternalTarget,
-    ModuleTarget,
-    ConstantTarget,
-    FunctionTarget,
     RemoteTarget,
     StitchedModule,
     StitchedModuleException,
-    CantResolveNodeDependenciesException,
     StitchedModuleOutput,
 )
-from .passage import always_false_predicate, always_true_predicate, InputArgs
+from .passage import InputArgs, always_false_predicate, always_true_predicate

modelopt/torch/_compress/sewing_kit/core.py

@@ -14,11 +14,14 @@
 # limitations under the License.
 
 # mypy: ignore-errors
+
 from __future__ import annotations
+
 from abc import ABC
 from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Any, Callable, Iterable, Literal, Optional, Sequence, Union
+
 from typing_extensions import override
 
 try:
@@ -30,19 +33,18 @@
 import torch.distributed
 import torch.nn as nn
 
-from .utils import distributed_isend_obj, distributed_recv_obj, dynamo_skip
 from .passage import (
-    Passage,
     InputArgs,
     OutputValue,
-    Predicate,
-    always_false_predicate,
+    Passage,
     PassageInputAdapter,
-    PassageOutputAdapter,
     PassageInputOverrides,
+    PassageOutputAdapter,
     PassageOutputOverrides,
+    Predicate,
+    always_false_predicate,
 )
-
+from .utils import distributed_isend_obj, distributed_recv_obj, dynamo_skip
 
 InputAdapter = Callable[[InputArgs], InputArgs]
 OutputAdapter = Callable[..., OutputValue]

modelopt/torch/_compress/sewing_kit/passage/__init__.py

@@ -13,16 +13,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+
 from .core import (
-    Passage,
-    PassageOutput,
     InputArgs,
     OutputValue,
-    Predicate,
+    Passage,
     PassageInputAdapter,
-    PassageOutputAdapter,
     PassageInputOverrides,
+    PassageOutput,
+    PassageOutputAdapter,
     PassageOutputOverrides,
-    always_true_predicate,
+    Predicate,
     always_false_predicate,
+    always_true_predicate,
 )

modelopt/torch/_compress/sewing_kit/passage/core.py

@@ -15,10 +15,9 @@
 
 # mypy: ignore-errors
 from __future__ import annotations
-import sys
-
-from collections.abc import Sequence, Callable
 
+import sys
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from typing import Any, ContextManager, Iterable, Mapping, Optional, Union
 
@@ -27,19 +26,19 @@
 except ImportError:
     from typing_extensions import Self
 
+import torch.nn as nn
 from typing_extensions import override
 
-import torch.nn as nn
+from ..common import logger
 from ..utils import (
     ActivityContext,
-    has_fake_tensor,
+    dynamo_skip,
     fake_tensors,
+    has_fake_tensor,
     is_submodule_of,
     is_submodule_or_same,
     real_tensors,
-    dynamo_skip,
 )
-from ..common import logger
 
 
 @dataclass

modelopt/torch/_compress/sewing_kit/utils.py

@@ -16,7 +16,7 @@
 from __future__ import annotations
 
 import inspect
-from collections.abc import Sequence, Mapping
+from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from typing import (
     Any,
@@ -31,17 +31,17 @@
     cast,
     overload,
 )
-from typing_extensions import override
+
 import torch
-import torch.distributed
-import torch._dynamo
 import torch._C
-from torch import Tensor
-import torch.utils._pytree as pytree
+import torch._dynamo
+import torch.distributed
 import torch.nn as nn
 import torch.nn.functional as F
+import torch.utils._pytree as pytree
+from torch import Tensor
 from torch._subclasses import FakeTensor, FakeTensorMode
-
+from typing_extensions import override
 
 Fn = TypeVar("Fn", bound=Callable)