diff --git a/docs/advanced_features/forward_hooks.md b/docs/advanced_features/forward_hooks.md index 33f5f6ce4..4dc170b33 100644 --- a/docs/advanced_features/forward_hooks.md +++ b/docs/advanced_features/forward_hooks.md @@ -20,14 +20,14 @@ Hooks are configured via a `ServerArgs` field: class ServerArgs: ... # For forward hooks - hooks: Optional[List[dict[str, Any]]] = None + forward_hooks: Optional[List[dict[str, Any]]] = None ```` In JSON form, a minimal configuration looks like: ```jsonc { - "hooks": [ + "forward_hooks": [ { "name": "outer_linear_hooks", "target_modules": ["outer.0", "outer.1"], @@ -42,7 +42,7 @@ In JSON form, a minimal configuration looks like: #### Top-level fields -* `hooks` (optional list of objects) +* `forward_hooks` (optional list of objects) Each element is a hook spec describing: * Which modules to target @@ -53,7 +53,7 @@ In JSON form, a minimal configuration looks like: ### Hook spec schema -Each entry in `hooks` is a JSON object with the following shape: +Each entry in `forward_hooks` is a JSON object with the following shape: ```jsonc { @@ -149,16 +149,16 @@ The first two cause initialization to fail fast with a descriptive error; the la Hooks are registered in `ModelRunner.initialize()`: ```python -if server_args.hooks: - register_hooks(self.model, server_args.hooks) +if server_args.forward_hooks: + register_forward_hooks(self.model, server_args.forward_hooks) ``` -The actual registration logic is implemented by `register_hooks`: +The actual registration logic is implemented by `register_forward_hooks`: ```python -def register_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None: +def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None: """ - hook_specs is a list of dicts from server_args.hooks. + hook_specs is a list of dicts from server_args.forward_hooks. Attaches forward hooks to the matching modules. """ name_to_module = dict(model.named_modules()) @@ -255,7 +255,7 @@ In JSON: ```jsonc { - "hooks": [ + "forward_hooks": [ { "name": "capture_outer", "target_modules": ["outer.0", "outer.1"], @@ -279,7 +279,7 @@ This will: ### Summary -* Define `hooks` as a list of specs in `ServerArgs` to turn on the feature. +* Define `forward_hooks` as a list of specs in `ServerArgs` to turn on the feature. * Each spec: diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 33583cf1f..ed68bf965 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -403,7 +403,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s ## Forward hooks | Argument | Description | Defaults | Options | | --- | --- | --- | --- | -| `--hooks` | JSON-formatted list of hook specifications. Each element must include `target_modules` (list of glob patterns matched against `model.named_modules()` names) and `hook_factory` (Python import path to a factory, e.g. `my_package.hooks:make_hook`). An optional `name` field is used for logging, and an optional `config` object is passed as a `dict` to the factory. | `None` | Type: JSON list | +| `--forward-hooks` | JSON-formatted list of forward hook specifications. Each element must include `target_modules` (list of glob patterns matched against `model.named_modules()` names) and `hook_factory` (Python import path to a factory, e.g. `my_package.hooks:make_hook`). An optional `name` field is used for logging, and an optional `config` object is passed as a `dict` to the factory. | `None` | Type: JSON list | ## Debug tensor dumps | Argument | Description | Defaults | Options | diff --git a/python/sglang/srt/model_executor/hook_manager.py b/python/sglang/srt/model_executor/hook_manager.py index 42c2828af..5f6d78a51 100644 --- a/python/sglang/srt/model_executor/hook_manager.py +++ b/python/sglang/srt/model_executor/hook_manager.py @@ -8,9 +8,9 @@ import torch.nn as nn logger = logging.getLogger(__name__) -def register_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None: +def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None: """ - hook_specs is a list of dicts from server_args.hooks. + hook_specs is a list of dicts from server_args.forward_hooks. Attaches forward hooks to the matching modules. """ name_to_module = dict(model.named_modules()) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e62f07746..7e4ab70e9 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -122,7 +122,7 @@ from sglang.srt.mem_cache.memory_pool import ( from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_executor.hook_manager import register_hooks +from sglang.srt.model_executor.hook_manager import register_forward_hooks from sglang.srt.model_executor.npu_graph_runner import NPUGraphRunner from sglang.srt.model_executor.piecewise_cuda_graph_runner import ( PiecewiseCudaGraphRunner, @@ -535,8 +535,8 @@ class ModelRunner: self.graph_mem_usage = 0 self.init_attention_backend() - if server_args.hooks: - register_hooks(self.model, server_args.hooks) + if server_args.forward_hooks: + register_forward_hooks(self.model, server_args.forward_hooks) # auxiliary hidden capture mode. TODO: expose this to server args? if self.spec_algorithm.is_eagle3() and not self.is_draft_worker: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7a4f53202..3768057e8 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -588,7 +588,7 @@ class ServerArgs: mm_enable_dp_encoder: bool = False # For forward hooks - hooks: Optional[List[dict[str, Any]]] = None + forward_hooks: Optional[List[dict[str, Any]]] = None def __post_init__(self): """ @@ -3870,10 +3870,10 @@ class ServerArgs: # For registering hooks parser.add_argument( - "--hooks", + "--forward-hooks", type=json_list_type, - default=None, - help="The hooks to be attached.", + default=ServerArgs.forward_hooks, + help="JSON-formatted forward hook specifications to attach to the model.", ) @classmethod diff --git a/test/srt/test_model_hooks.py b/test/srt/test_model_hooks.py index 656b124d0..473aafb3e 100644 --- a/test/srt/test_model_hooks.py +++ b/test/srt/test_model_hooks.py @@ -4,7 +4,7 @@ import json import torch import torch.nn as nn -from sglang.srt.model_executor.hook_manager import register_hooks +from sglang.srt.model_executor.hook_manager import register_forward_hooks from sglang.srt.server_args import ServerArgs from sglang.test.test_utils import CustomTestCase @@ -46,7 +46,7 @@ class TinyModel(nn.Module): class TestAttachHooks(CustomTestCase): - """Tests for ModelRunner.register_hooks / resolve_callable integration.""" + """Tests for register_forward_hooks / resolve_callable integration.""" def setUp(self): HOOK_CALLS.clear() @@ -67,7 +67,7 @@ class TestAttachHooks(CustomTestCase): ] model = TinyModel() - register_hooks(model, hook_specs) + register_forward_hooks(model, hook_specs) x = torch.randn(3, 4) _ = model(x) @@ -92,7 +92,7 @@ class TestAttachHooks(CustomTestCase): } ] - register_hooks(model, hook_specs) + register_forward_hooks(model, hook_specs) x = torch.randn(3, 4) _ = model(x) @@ -103,7 +103,7 @@ class TestAttachHooks(CustomTestCase): def test_cli_hooks_reach_model(self): """ Ensure that when hooks are provided via CLI, they are parsed into - ServerArgs, passed to ModelRunner.register_hooks, and actually + ServerArgs, passed to register_forward_hooks, and actually run during a forward pass. """ parser = argparse.ArgumentParser() @@ -121,17 +121,17 @@ class TestAttachHooks(CustomTestCase): cli_args = [ "--model-path", "Qwen/Qwen2-7B-Instruct", # Dummy value; not used in this test - "--hooks", + "--forward-hooks", json.dumps(hooks_spec), ] args = parser.parse_args(cli_args) server_args = ServerArgs.from_cli_args(args) - self.assertEqual(server_args.hooks, hooks_spec) + self.assertEqual(server_args.forward_hooks, hooks_spec) model = TinyModel() - register_hooks(model, server_args.hooks) + register_forward_hooks(model, server_args.forward_hooks) x = torch.randn(3, 4) _ = model(x)