minor

Signed-off-by: Jingyu Xin <[email protected]>

Author: jingyu-ml

modelopt/torch/peft/conversion.py

@@ -15,13 +15,11 @@
 
 """PEFT conversion and restore utilities for LoRA modules."""
 
-import fnmatch
-from collections.abc import Callable, Iterable
-
 import torch.nn as nn
 
 from modelopt.torch.opt.conversion import ModelLikeModule, ModeloptStateManager
 from modelopt.torch.opt.mode import ConvertReturnType, MetadataDict
+from modelopt.torch.utils.network import matches_pattern
 
 from .config import PEFTConfig
 from .lora.layer import LoRAModule, LoRAModuleRegistry
@@ -117,7 +115,7 @@ def add_adapter(model, config: PEFTConfig):
             # Later patterns override earlier ones
             merged_setting = None
             for wildcard_or_filter_func, adapter_setting in adapter_cfg.items():
-                if _matches(name, wildcard_or_filter_func):
+                if matches_pattern(name, wildcard_or_filter_func):
                     if merged_setting is None:
                         merged_setting = adapter_setting.copy()
                     else:
@@ -134,44 +132,9 @@ def add_adapter(model, config: PEFTConfig):
     return model
 
 
-def _matches(
-    name: str,
-    patterns: str | Callable[[str], bool] | Iterable[str | Callable[[str], bool]] | None,
-    *,
-    allow_callable: bool = True,
-) -> bool:
-    if patterns is None:
-        return True
-
-    if isinstance(patterns, (str, bytes)):
-        patterns_iter: Iterable[str | Callable[[str], bool]] = (patterns,)
-    elif callable(patterns):
-        if not allow_callable:
-            raise TypeError("Callable patterns are not supported in this context.")
-        patterns_iter = (patterns,)
-    elif isinstance(patterns, Iterable):
-        patterns_iter = tuple(patterns)
-    else:
-        raise TypeError(f"Unsupported pattern type: {type(patterns)}")
-
-    for pattern in patterns_iter:
-        if isinstance(pattern, (str, bytes)):
-            if fnmatch.fnmatch(name, pattern):
-                return True
-        elif callable(pattern):
-            if not allow_callable:
-                raise TypeError("Callable patterns are not supported in this context.")
-            if pattern(name):
-                return True
-        else:
-            raise TypeError(f"Unsupported pattern type: {type(pattern)}")
-
-    return False
-
-
 def _iter_lora_modules(model, layer_patterns=None):
     for module_name, module in model.named_modules():
-        if isinstance(module, LoRAModule) and _matches(module_name, layer_patterns):
+        if isinstance(module, LoRAModule) and matches_pattern(module_name, layer_patterns):
             yield module_name, module
 
 
@@ -192,14 +155,14 @@ def _set_base_requires_grad(model, *, requires_grad: bool, layer_patterns=None):
         # If layer_patterns is specified, only affect matching layers
         if layer_patterns is not None:
             module_name = ".".join(name.split(".")[:-1])  # Get module name without param name
-            if not _matches(module_name, layer_patterns):
+            if not matches_pattern(module_name, layer_patterns):
                 continue
         param.requires_grad = requires_grad
 
 
 def _iter_adapter_names(module, adapter_patterns=None):
     for adapter_name in module._lora_adapters:
-        if _matches(adapter_name, adapter_patterns, allow_callable=False):
+        if matches_pattern(adapter_name, adapter_patterns, allow_callable=False):
             yield adapter_name
 
 

modelopt/torch/peft/convert.py

@@ -15,14 +15,14 @@
 
 """User-facing PEFT API for LoRA module conversion and adapter management."""
 
-import fnmatch
 from typing import Any
 
 import torch.nn as nn
 
 from modelopt.torch.opt import apply_mode
 from modelopt.torch.peft.config import PEFTConfig
 from modelopt.torch.peft.conversion import add_adapter
+from modelopt.torch.utils.network import matches_pattern
 
 try:
     from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
@@ -101,30 +101,15 @@ def _set_adapter_state(model, enable_state, layer_patterns=None, adapter_pattern
     if not is_peft_model(model):
         raise ValueError("Model must be a PEFT model to set adapter states.")
 
-    def matches_any_pattern(name, patterns, allow_callable=True):
-        for pattern in patterns:
-            if isinstance(pattern, str):
-                if fnmatch.fnmatch(name, pattern):
-                    return True
-            elif allow_callable and callable(pattern):
-                if pattern(name):
-                    return True
-            else:
-                pattern_type = "pattern" if allow_callable else "adapter pattern"
-                raise TypeError(f"Unsupported {pattern_type} type: {type(pattern)}")
-        return False
-
     for module_name, module in model.named_modules():
         if isinstance(module, LoRAModule):
             if layer_patterns is not None:
-                if not matches_any_pattern(module_name, layer_patterns, allow_callable=True):
+                if not matches_pattern(module_name, layer_patterns, allow_callable=True):
                     continue
 
             for adapter_name, adapter_dict in module._lora_adapters.items():
                 if adapter_patterns is not None:
-                    if not matches_any_pattern(
-                        adapter_name, adapter_patterns, allow_callable=False
-                    ):
+                    if not matches_pattern(adapter_name, adapter_patterns, allow_callable=False):
                         continue
 
                 adapter_dict["enable"] = enable_state

modelopt/torch/utils/network.py

@@ -15,6 +15,7 @@
 
 """Utility functions for PyTorch models."""
 
+import fnmatch
 import inspect
 import types
 import warnings
@@ -55,6 +56,7 @@ def _convert_to_wrapped_module_name(name: str) -> str:
     "is_channels_last",
     "is_parallel",
     "make_divisible",
+    "matches_pattern",
     "model_to",
     "param_num",
     "param_num_from_forward",
@@ -87,6 +89,73 @@ def _convert_to_wrapped_module_name(name: str) -> str:
 ConstructorLike = Callable | tuple
 
 
+def matches_pattern(
+    name: str,
+    patterns: str | Callable[[str], bool] | Iterable[str | Callable[[str], bool]] | None,
+    *,
+    allow_callable: bool = True,
+) -> bool:
+    """Check if a name matches any of the given patterns.
+
+    This utility function checks if a given name (e.g., module name, layer name)
+    matches any pattern in a collection of patterns. Patterns can be:
+    - String wildcards (using fnmatch syntax, e.g., "*.attention.*")
+    - Callable predicates that take a name and return bool
+    - None (matches everything)
+
+    Args:
+        name: The name to check (e.g., "model.layer1.attention.weight")
+        patterns: A single pattern, iterable of patterns, or None.
+                 If None, returns True (matches everything).
+        allow_callable: If False, raises TypeError when encountering callable patterns.
+                       Useful for contexts where only string patterns are allowed.
+
+    Returns:
+        True if the name matches any pattern, False otherwise.
+
+    Raises:
+        TypeError: If pattern type is unsupported or if callable patterns are
+                  provided when allow_callable=False.
+
+    Examples:
+        >>> matches_pattern("model.attention.query", "*.attention.*")
+        True
+        >>> matches_pattern("model.mlp.linear", ["*.attention.*", "*.mlp.*"])
+        True
+        >>> matches_pattern("model.layer1", lambda x: "layer" in x)
+        True
+        >>> matches_pattern("anything", None)
+        True
+    """
+    if patterns is None:
+        return True
+
+    if isinstance(patterns, (str, bytes)):
+        patterns_iter: Iterable[str | Callable[[str], bool]] = (patterns,)
+    elif callable(patterns):
+        if not allow_callable:
+            raise TypeError("Callable patterns are not supported in this context.")
+        patterns_iter = (patterns,)
+    elif isinstance(patterns, Iterable):
+        patterns_iter = tuple(patterns)
+    else:
+        raise TypeError(f"Unsupported pattern type: {type(patterns)}")
+
+    for pattern in patterns_iter:
+        if isinstance(pattern, (str, bytes)):
+            if fnmatch.fnmatch(name, pattern):
+                return True
+        elif callable(pattern):
+            if not allow_callable:
+                raise TypeError("Callable patterns are not supported in this context.")
+            if pattern(name):
+                return True
+        else:
+            raise TypeError(f"Unsupported pattern type: {type(pattern)}")
+
+    return False
+
+
 def is_parallel(model: nn.Module) -> bool:
     """Check if a PyTorch model is parallelized."""
     return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel))