Rename: --hooks to --forward-hooks (#13994)
This commit is contained in:
@@ -20,14 +20,14 @@ Hooks are configured via a `ServerArgs` field:
|
|||||||
class ServerArgs:
|
class ServerArgs:
|
||||||
...
|
...
|
||||||
# For forward hooks
|
# 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:
|
In JSON form, a minimal configuration looks like:
|
||||||
|
|
||||||
```jsonc
|
```jsonc
|
||||||
{
|
{
|
||||||
"hooks": [
|
"forward_hooks": [
|
||||||
{
|
{
|
||||||
"name": "outer_linear_hooks",
|
"name": "outer_linear_hooks",
|
||||||
"target_modules": ["outer.0", "outer.1"],
|
"target_modules": ["outer.0", "outer.1"],
|
||||||
@@ -42,7 +42,7 @@ In JSON form, a minimal configuration looks like:
|
|||||||
|
|
||||||
#### Top-level fields
|
#### Top-level fields
|
||||||
|
|
||||||
* `hooks` (optional list of objects)
|
* `forward_hooks` (optional list of objects)
|
||||||
Each element is a hook spec describing:
|
Each element is a hook spec describing:
|
||||||
|
|
||||||
* Which modules to target
|
* Which modules to target
|
||||||
@@ -53,7 +53,7 @@ In JSON form, a minimal configuration looks like:
|
|||||||
|
|
||||||
### Hook spec schema
|
### 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
|
```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()`:
|
Hooks are registered in `ModelRunner.initialize()`:
|
||||||
|
|
||||||
```python
|
```python
|
||||||
if server_args.hooks:
|
if server_args.forward_hooks:
|
||||||
register_hooks(self.model, server_args.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
|
```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.
|
Attaches forward hooks to the matching modules.
|
||||||
"""
|
"""
|
||||||
name_to_module = dict(model.named_modules())
|
name_to_module = dict(model.named_modules())
|
||||||
@@ -255,7 +255,7 @@ In JSON:
|
|||||||
|
|
||||||
```jsonc
|
```jsonc
|
||||||
{
|
{
|
||||||
"hooks": [
|
"forward_hooks": [
|
||||||
{
|
{
|
||||||
"name": "capture_outer",
|
"name": "capture_outer",
|
||||||
"target_modules": ["outer.0", "outer.1"],
|
"target_modules": ["outer.0", "outer.1"],
|
||||||
@@ -279,7 +279,7 @@ This will:
|
|||||||
|
|
||||||
### Summary
|
### 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:
|
* Each spec:
|
||||||
|
|
||||||
|
|||||||
@@ -403,7 +403,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
|||||||
## Forward hooks
|
## Forward hooks
|
||||||
| Argument | Description | Defaults | Options |
|
| 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
|
## Debug tensor dumps
|
||||||
| Argument | Description | Defaults | Options |
|
| Argument | Description | Defaults | Options |
|
||||||
|
|||||||
@@ -8,9 +8,9 @@ import torch.nn as nn
|
|||||||
logger = logging.getLogger(__name__)
|
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.
|
Attaches forward hooks to the matching modules.
|
||||||
"""
|
"""
|
||||||
name_to_module = dict(model.named_modules())
|
name_to_module = dict(model.named_modules())
|
||||||
|
|||||||
@@ -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.cpu_graph_runner import CPUGraphRunner
|
||||||
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
|
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.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.npu_graph_runner import NPUGraphRunner
|
||||||
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
|
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
|
||||||
PiecewiseCudaGraphRunner,
|
PiecewiseCudaGraphRunner,
|
||||||
@@ -535,8 +535,8 @@ class ModelRunner:
|
|||||||
self.graph_mem_usage = 0
|
self.graph_mem_usage = 0
|
||||||
self.init_attention_backend()
|
self.init_attention_backend()
|
||||||
|
|
||||||
if server_args.hooks:
|
if server_args.forward_hooks:
|
||||||
register_hooks(self.model, server_args.hooks)
|
register_forward_hooks(self.model, server_args.forward_hooks)
|
||||||
|
|
||||||
# auxiliary hidden capture mode. TODO: expose this to server args?
|
# auxiliary hidden capture mode. TODO: expose this to server args?
|
||||||
if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
|
if self.spec_algorithm.is_eagle3() and not self.is_draft_worker:
|
||||||
|
|||||||
@@ -588,7 +588,7 @@ class ServerArgs:
|
|||||||
mm_enable_dp_encoder: bool = False
|
mm_enable_dp_encoder: bool = False
|
||||||
|
|
||||||
# For forward hooks
|
# For forward hooks
|
||||||
hooks: Optional[List[dict[str, Any]]] = None
|
forward_hooks: Optional[List[dict[str, Any]]] = None
|
||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
"""
|
"""
|
||||||
@@ -3870,10 +3870,10 @@ class ServerArgs:
|
|||||||
|
|
||||||
# For registering hooks
|
# For registering hooks
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--hooks",
|
"--forward-hooks",
|
||||||
type=json_list_type,
|
type=json_list_type,
|
||||||
default=None,
|
default=ServerArgs.forward_hooks,
|
||||||
help="The hooks to be attached.",
|
help="JSON-formatted forward hook specifications to attach to the model.",
|
||||||
)
|
)
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
|
|||||||
@@ -4,7 +4,7 @@ import json
|
|||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
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.srt.server_args import ServerArgs
|
||||||
from sglang.test.test_utils import CustomTestCase
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
@@ -46,7 +46,7 @@ class TinyModel(nn.Module):
|
|||||||
|
|
||||||
|
|
||||||
class TestAttachHooks(CustomTestCase):
|
class TestAttachHooks(CustomTestCase):
|
||||||
"""Tests for ModelRunner.register_hooks / resolve_callable integration."""
|
"""Tests for register_forward_hooks / resolve_callable integration."""
|
||||||
|
|
||||||
def setUp(self):
|
def setUp(self):
|
||||||
HOOK_CALLS.clear()
|
HOOK_CALLS.clear()
|
||||||
@@ -67,7 +67,7 @@ class TestAttachHooks(CustomTestCase):
|
|||||||
]
|
]
|
||||||
|
|
||||||
model = TinyModel()
|
model = TinyModel()
|
||||||
register_hooks(model, hook_specs)
|
register_forward_hooks(model, hook_specs)
|
||||||
|
|
||||||
x = torch.randn(3, 4)
|
x = torch.randn(3, 4)
|
||||||
_ = model(x)
|
_ = 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)
|
x = torch.randn(3, 4)
|
||||||
_ = model(x)
|
_ = model(x)
|
||||||
@@ -103,7 +103,7 @@ class TestAttachHooks(CustomTestCase):
|
|||||||
def test_cli_hooks_reach_model(self):
|
def test_cli_hooks_reach_model(self):
|
||||||
"""
|
"""
|
||||||
Ensure that when hooks are provided via CLI, they are parsed into
|
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.
|
run during a forward pass.
|
||||||
"""
|
"""
|
||||||
parser = argparse.ArgumentParser()
|
parser = argparse.ArgumentParser()
|
||||||
@@ -121,17 +121,17 @@ class TestAttachHooks(CustomTestCase):
|
|||||||
cli_args = [
|
cli_args = [
|
||||||
"--model-path",
|
"--model-path",
|
||||||
"Qwen/Qwen2-7B-Instruct", # Dummy value; not used in this test
|
"Qwen/Qwen2-7B-Instruct", # Dummy value; not used in this test
|
||||||
"--hooks",
|
"--forward-hooks",
|
||||||
json.dumps(hooks_spec),
|
json.dumps(hooks_spec),
|
||||||
]
|
]
|
||||||
|
|
||||||
args = parser.parse_args(cli_args)
|
args = parser.parse_args(cli_args)
|
||||||
server_args = ServerArgs.from_cli_args(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()
|
model = TinyModel()
|
||||||
register_hooks(model, server_args.hooks)
|
register_forward_hooks(model, server_args.forward_hooks)
|
||||||
|
|
||||||
x = torch.randn(3, 4)
|
x = torch.randn(3, 4)
|
||||||
_ = model(x)
|
_ = model(x)
|
||||||
|
|||||||
Reference in New Issue
Block a user