more changes

Author: luccafong

debug_logging.sh

@@ -0,0 +1,77 @@
+# Intermediate Tensor Logging
+
+This document provides guidance on using the intermediate tensor logging feature in vLLM, which allows you to capture and save intermediate tensors during model execution.
+
+## Overview
+
+The intermediate tensor logging feature enables you to:
+
+- Log input and output tensors from a configured set of filters
+- Filter modules by name using regex patterns
+- Filter module fwd call index (e.g. dump 2nd call of forward pass on same module)
+- Filter tensors by device
+- Filter whole model fwd step id
+
+
+This is manily useful for debugging model accucacy gaps with 2 runs
+
+
+## Usage
+
+### Enabling via Configuration File
+
+The easiest way to enable intermediate logging is by providing a configuration file:
+
+```bash
+python -m vllm.entrypoints.openai.api_server \
+  --model <your_model> \
+  --il-config-path /path/to/config.json
+```
+
+### Configuration Options
+
+The configuration file should be a JSON file with the following structure:
+
+```json
+{
+  "output_dir": "/tmp/vllm_intermediates",
+  "module_call_match": ["layers\\.0\\.(?!.*rotary_emb).*", "rotary_emb:0", "embed_tokens", "model\\.norm"],
+  "log_step_ids": [0, 1],
+  "device_names": ["cuda:0"]
+}
+```
+
+#### Configuration Parameters
+
+| Parameter | Type | Description | Default |
+|-----------|------|-------------|---------|
+| `output_dir` | string | Directory where to save the intermediate tensors | `/tmp/vllm_intermediates` |
+| `module_call_match` | array | Regex patterns to filter module names, if limti to ith call only, add `:i` | `null` (log all modules) |
+| `log_step_ids` | array | List of step IDs to log | `[0]` |
+| `max_tensor_size` | integer | Maximum number of elements in tensors to log | `null` (no limit) |
+| `device_names` | array | List of device names to log | `[]` (log all devices) |
+
+### Output Directory Structure
+
+When you enable intermediate logging, the system creates a timestamped directory under your specified `output_dir`. This helps organize multiple logging sessions:
+
+```
+/tmp/vllm_intermediates/010fed05-4a36-4c19-ab44-7cd67e3f63ce/
+└── step_0
+    ├── model.embed_tokens
+    │   ├── inputs_0_cuda_0.pt
+    │   ├── inputs.json
+    │   ├── outputs_cuda_0.pt
+    │   └── outputs.json
+    ├── model.layers.0.input_layernorm
+    │   ├── inputs_0_cuda_0.pt
+    │   ├── inputs.json
+    │   ├── outputs_cuda_0.pt
+    │   └── outputs.json
+    └── step_1/
+        └── ...
+```
+
+Each tensor is saved in two formats:
+1. `.json` files containing metadata and small tensor values
+2. `.pt` files containing the full PyTorch tensors (can be loaded with `torch.load()`)

docs/contributing/intermediate_logging.md

@@ -17,6 +17,14 @@ def create_test_prompts() -> list[tuple[str, SamplingParams]]:
         (
             "A robot may not injure a human being",
             SamplingParams(temperature=0.0, logprobs=1, prompt_logprobs=1),
+        ),
+        (
+            "To be or not to be,",
+            SamplingParams(temperature=0.8, top_k=5, presence_penalty=0.2),
+        ),
+        (
+            "What is the meaning of life?",
+            SamplingParams(n=2, temperature=0.8, top_p=0.95, frequency_penalty=0.1),
         )
     ]
 

examples/offline_inference/llm_engine_example.py

@@ -17,7 +17,8 @@
 from functools import cached_property
 from importlib.util import find_spec
 from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Literal, Optional,
-                    Protocol, TypeVar, Union, cast, get_args)
+                    Protocol, TypeVar, Union, cast, get_args, List, Set)
+from re import Pattern
 
 import regex as re
 import torch
@@ -3952,6 +3953,119 @@ class KVEventsConfig:
     """
 
 
+@config
+@dataclass
+class IntermediateLoggingConfig:
+    """Configuration for intermediate tensor logging."""
+    
+    output_dir: str = "/tmp/vllm_intermediates"
+    """Directory where to save the intermediate tensors."""
+    
+    module_call_match: Optional[List[str]] = None
+    """Match modules by name regex and call index (
+    a module can be called multiple times in a step)
+    List of regex:call_idx, call_idx is -1  for default for all calls """
+    
+    log_step_ids: List[int] = field(default_factory=lambda: [0])
+    """List of step IDs to log (empty list means log all steps)."""
+    
+    log_post_fwd_inputs: bool = False
+    """Whether logging inputs after forwards for each module"""
+
+    max_tensor_size: Optional[int] = None
+    """Maximum number of elements in tensors to log (None = no limit)."""
+    
+    enabled: bool = True
+    """Whether logging is enabled."""
+    device_names: List[str] = field(default_factory=list)
+    """List of device names to log (empty list means log all devices)."""
+    
+    _compiled_module_calls: dict[Pattern,int] = field(default_factory=dict, init=False)
+    """Compiled regex patterns for module filtering."""
+    
+    _module_call: dict[str, int] = field(default_factory=dict, init=False)
+    _step_id_set: Set[int] = field(default_factory=set, init=False)
+    """Set of step IDs for faster lookup."""
+    _output_run_dir: str = "/tmp/vllm_intermediates"
+    """Unique directory to save single run/serve logging result."""
+    
+    def __post_init__(self):
+        """Initialize derived fields after instance creation."""
+        self._compile_regex_patterns()
+        self._output_run_dir = self.output_dir + "/" + str(uuid.uuid4())
+        self._step_id_set = set(self.log_step_ids)
+        
+    def _compile_regex_patterns(self):
+        """Compile regex patterns for module name filtering."""
+        from vllm.logger import init_logger
+        logger = init_logger(__name__)
+        
+        self._compiled_module_matches = []
+        
+        if self.module_call_match is None:
+            logger.info("No module name regex patterns provided, will log all modules")
+            return
+        
+        # Compile all patterns
+        for regex_pattern_call_idx in self.module_call_match:
+            try:
+                splits = regex_pattern_call_idx.split(":", 2)
+                regex_pattern = splits[0]
+                call_idx = -1
+                if len(splits) > 1:
+                    call_idx = int(splits[1])
+                compiled_pattern: Pattern[str] = re.compile(regex_pattern)
+                self._compiled_module_calls[compiled_pattern] = call_idx
+                logger.info(f"Successfully compiled regex pattern: '{regex_pattern}'")
+            except re.error as e:
+                logger.error(f"Invalid regex pattern from '{regex_pattern_call_idx}': {e}")
+                raise ValueError(f"Invalid regex pattern '{regex_pattern_call_idx}': {e}")
+            except Exception as e:
+                logger.error(f"Failed to parse module_call_match")
+
+        
+        logger.info(f"Compiled {len(self._compiled_module_calls)} regex patterns")
+
+    def to_dict(self) -> dict:
+        """Convert the config to a dictionary for serialization."""
+        return {
+            "output_run_dir": self.output_run_dir,
+            "module_call_match": self.module_call_match,
+            "log_step_ids": self.log_step_ids,
+            "max_tensor_size": self.max_tensor_size,
+            "enabled": self.enabled,
+            "device_names": self.device_names
+        }
+    
+    @classmethod
+    def from_dict(cls, dict_value: dict) -> "IntermediateLoggingConfig":
+        """Parse the CLI value for the speculative config."""
+        return cls(**dict_value)
+    
+    @property
+    def output_run_dir(self) -> str:
+        return self._output_run_dir
+
+    def compute_hash(self) -> str:
+        """
+        WARNING: Whenever a new field is added to this config,
+        ensure that it is included in the factors list if
+        it affects the computation graph.
+
+        Provide a hash that uniquely identifies all the configs
+        that affect the structure of the computation
+        graph from input ids/embeddings to the final hidden states,
+        excluding anything before input ids/embeddings and after
+        the final hidden states.
+        """
+        # Intermediate logging doesn't affect the computation graph
+        factors: list[Any] = []
+        hash_str = hashlib.md5(str(factors).encode(),
+                               usedforsecurity=False).hexdigest()
+        return hash_str
+    
+
+
 class CompilationLevel:
     # constants for the levels of the compilation process
     NO_COMPILATION = 0
@@ -4409,6 +4523,10 @@ class VllmConfig:
     """The configurations for distributed KV cache transfer."""
     kv_events_config: Optional[KVEventsConfig] = None
     """The configurations for event publishing."""
+    il_config: Optional[IntermediateLoggingConfig] = None
+    """Configuration for intermediate tensor logging."""
+    il_config_path: Optional[str] = None
+    """Path to a JSON file containing intermediate logging configuration."""
     # some opaque config, only used to provide additional information
     # for the hash computation, mainly used for testing, debugging or out of
     # tree config registration.
@@ -4497,6 +4615,10 @@ def compute_hash(self) -> str:
             vllm_factors.append(self.kv_transfer_config.compute_hash())
         else:
             vllm_factors.append("None")
+        if self.il_config:
+            vllm_factors.append(self.il_config.compute_hash())
+        else:
+            vllm_factors.append("None")
         if self.additional_config:
             if isinstance(additional_config := self.additional_config, dict):
                 additional_config_hash = hashlib.md5(
@@ -4705,6 +4827,18 @@ def __post_init__(self):
             if self.kv_events_config is not None:
                 # Hybrid KV cache manager is not compatible with KV events.
                 self.scheduler_config.disable_hybrid_kv_cache_manager = True
+        # Load intermediate logging config from file if provided
+        if self.il_config_path is not None:
+            try:
+                logger.info(f"Loading intermediate logging config from {self.il_config_path}")
+                with open(self.il_config_path, 'r') as f:
+                    il_config_dict = json.load(f)
+                    self.il_config = IntermediateLoggingConfig.from_dict(il_config_dict)
+                    logger.info(f"Successfully loaded intermediate logging config: {self.il_config.to_dict()}")
+            except Exception as e:
+                logger.error(f"Failed to load intermediate logging config from {self.il_config_path}: {e}")
+                raise ValueError(f"Failed to load intermediate logging config from {self.il_config_path}: {e}") from e
+
 
     def update_sizes_for_sequence_parallelism(self,
                                               possible_sizes: list) -> list:

vllm/config.py

@@ -440,6 +440,7 @@ class EngineArgs:
         ParallelConfig.enable_multimodal_encoder_data_parallel
 
     async_scheduling: bool = SchedulerConfig.async_scheduling
+    il_config_path: Optional[str] = None
 
     def __post_init__(self):
         # support `EngineArgs(compilation_config={...})`
@@ -858,6 +859,8 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
                                 **vllm_kwargs["compilation_config"])
         vllm_group.add_argument("--additional-config",
                                 **vllm_kwargs["additional_config"])
+        vllm_group.add_argument("--il-config-path",
+                                **vllm_kwargs["il_config_path"])
 
         # Other arguments
         parser.add_argument('--use-v2-block-manager',
@@ -1276,7 +1279,6 @@ def create_engine_config(
             otlp_traces_endpoint=self.otlp_traces_endpoint,
             collect_detailed_traces=self.collect_detailed_traces,
         )
-
         config = VllmConfig(
             model_config=model_config,
             cache_config=cache_config,
@@ -1292,6 +1294,7 @@ def create_engine_config(
             compilation_config=self.compilation_config,
             kv_transfer_config=self.kv_transfer_config,
             kv_events_config=self.kv_events_config,
+            il_config_path=self.il_config_path,
             additional_config=self.additional_config,
         )
 

vllm/engine/arg_utils.py

@@ -43,6 +43,9 @@
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.version import __version__ as VLLM_VERSION
+from vllm.v1.worker.intermediates_logging import (
+    intermediate_logging
+)
 
 logger = init_logger(__name__)
 
@@ -73,33 +76,7 @@ def __init__(self,
 
         # Setup Model.
         self.model_executor = executor_class(vllm_config)
-
-        
-        # Hook the dumping logic
-        from vllm.v1.worker.il_config import IntermediateLoggingConfig
-        
-        # Create a configuration for intermediate logging
-        logger.info("Setting up intermediate tensor logging")
-        
-        # Define the regex patterns to match module names
-        # These patterns will match any module with "layers.0" or "embed_tokens" in its name
-        module_patterns = ["layers\\.0", "embed_tokens"]
-        logger.info(f"Using module name regex patterns: {module_patterns}")
-        
-        il_config = IntermediateLoggingConfig(
-            output_dir="/tmp/vllm_intermediates",  # Directory to save intermediates
-            module_name_regex=module_patterns,     # Log layer 0 and embedding modules
-            log_step_ids=[0, 1, 2, 3, 4, 5],       # Log steps 0-5
-            max_tensor_size=1000000,               # Limit to 1M elements
-            enabled=True                           # Enable logging
-        )
-        
-        logger.info(f"Intermediate logging config: {il_config.to_dict()}")
-        
-        # Register hooks for intermediate tensor logging
-        logger.info("Calling register_intermediate_hooks via collective_rpc")
-        self.collective_rpc("register_intermediate_hooks", args=(il_config,))
-        logger.info("Finished setting up intermediate tensor logging")
+        self.collective_rpc("register_intermediate_hooks", args=(vllm_config.il_config,))
 
         if executor_fail_callback is not None:
             self.model_executor.register_failure_callback(
@@ -243,19 +220,9 @@ def abort_requests(self, request_ids: list[str]):
 
     def execute_model(self, scheduler_output: SchedulerOutput):
         try:
-            # Increment the step counter for intermediate logging
-            try:
-                from vllm.v1.worker.intermediates_logging import increment_step
-                logger.info("Incrementing intermediate logging step counter before model execution")
-                increment_step()
-            except Exception as e:
-                logger.warning(f"Failed to increment intermediate logging step counter: {e}")
-            
             # Execute the model
-            result = self.model_executor.execute_model(scheduler_output)
-            
-            logger.info(f"Model execution completed for step with {scheduler_output.total_num_scheduled_tokens} tokens")
-            return result
+            with intermediate_logging(self.vllm_config.il_config):
+                return self.model_executor.execute_model(scheduler_output)
         except Exception as err:
             # We do not want to catch BaseException here since we're only
             # interested in dumping info when the exception is due to an