83 lines
2.6 KiB
Python
83 lines
2.6 KiB
Python
import fnmatch
|
|
import importlib
|
|
import logging
|
|
from typing import Any, Callable, List, Optional
|
|
|
|
import torch.nn as nn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def register_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None:
|
|
"""
|
|
hook_specs is a list of dicts from server_args.hooks.
|
|
Attaches forward hooks to the matching modules.
|
|
"""
|
|
name_to_module = dict(model.named_modules())
|
|
|
|
for spec in hook_specs:
|
|
spec_name = spec.get("name", "")
|
|
target_patterns = spec.get("target_modules", [])
|
|
if not target_patterns:
|
|
logger.warning(f"Hook spec '{spec_name}' has no 'target_modules', skipping")
|
|
continue
|
|
|
|
hook_factory_path = spec.get("hook_factory")
|
|
if not hook_factory_path:
|
|
logger.warning(f"Hook spec '{spec_name}' has no 'hook_factory', skipping")
|
|
continue
|
|
|
|
config = spec.get("config") or {}
|
|
hook_factory = resolve_callable(hook_factory_path)
|
|
|
|
hook = hook_factory(config) if hook_factory else None
|
|
if hook is None:
|
|
logger.warning(
|
|
f"Hook factory '{hook_factory_path}' for spec '{spec_name}' "
|
|
"returned None, not registering any hook"
|
|
)
|
|
continue
|
|
|
|
# Resolve patterns like "model.layers.*.mlp"
|
|
matched = []
|
|
for name, module in name_to_module.items():
|
|
if any(fnmatch.fnmatch(name, pattern) for pattern in target_patterns):
|
|
matched.append((name, module))
|
|
|
|
if not matched:
|
|
logger.warning(
|
|
f"No modules matched hook spec '{spec_name}' "
|
|
f"patterns={target_patterns}"
|
|
)
|
|
continue
|
|
|
|
for module_name, module in matched:
|
|
_ = module.register_forward_hook(hook)
|
|
logger.info(f"Registered forward hook '{spec_name}' " f"on {module_name}")
|
|
|
|
|
|
def resolve_callable(path: Optional[str]) -> Optional[Callable]:
|
|
if path is None:
|
|
return None
|
|
|
|
if ":" in path:
|
|
module_name, fn_name = path.split(":", 1)
|
|
else:
|
|
parts = path.split(".")
|
|
if len(parts) < 2:
|
|
raise ValueError(
|
|
f"Invalid hook callable path '{path}'. "
|
|
"Expected 'module.submodule:factory' or 'module.submodule.factory'."
|
|
)
|
|
*mod_parts, fn_name = parts
|
|
module_name = ".".join(mod_parts)
|
|
|
|
module = importlib.import_module(module_name)
|
|
try:
|
|
return getattr(module, fn_name)
|
|
except AttributeError as e:
|
|
raise AttributeError(
|
|
f"Module '{module_name}' has no attribute '{fn_name}' "
|
|
f"(from hook path '{path}')"
|
|
) from e
|