[wip] [Feature] [V1] intermediate logging #21215

luccafong · 2025-07-19T03:11:45Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Test Plan

python3 /data/users/fanglu/gitrepos/vllm/examples/o
ffline_inference/llm_engine_example.py --model "meta-llama/Llama-3.1-8B-Instruct"  --enforce-eager --il-config-path ~/logging.json

Test Result

/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
...

(Optional) Documentation Update

github-actions · 2025-07-19T03:11:52Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

luccafong · 2025-07-19T03:12:35Z

Initial prototype, still need to adjust structure for intermediates_logging.py (add tests) and add a script for comparison report generation cc @simon-mo @yeqcharlotte @houseroad

gemini-code-assist

Code Review

This pull request introduces a new intermediate tensor logging feature, which is a great addition for debugging model behavior. The implementation is comprehensive, covering configuration, argument parsing, engine integration, and the core logging logic with PyTorch hooks.

I've identified a few critical and high-severity issues that should be addressed to ensure the feature is robust and doesn't introduce unintended side effects. These include a potential crash in the hook removal logic, silent exception swallowing, and debug artifacts that could impact production environments.

Once these points are addressed, this will be a solid contribution to the project.

gemini-code-assist · 2025-07-19T03:13:19Z

vllm/v1/worker/intermediates_logging.py

+        step_hook = model.register_forward_hook(step_fwd)
+        self.hooks.append(("", None, step_hook))


This line appends a 3-element tuple to self.hooks for the step_hook. However, the remove_hooks method expects 4-element tuples when unpacking. This will cause a ValueError: not enough values to unpack when remove_hooks is called.

To fix this, you should ensure all tuples in self.hooks have the same structure. Since step_hook is a forward hook (a post-hook), it should be the 4th element in the tuple, with the pre-hook being None.

Suggested change

step_hook = model.register_forward_hook(step_fwd)

self.hooks.append(("", None, step_hook))

self.hooks.append(('', None, None, step_hook))

gemini-code-assist · 2025-07-19T03:13:19Z

vllm/config.py

+        from vllm.logger import init_logger
+        logger = init_logger(__name__)
+
+        self._compiled_module_matches = []


This line initializes self._compiled_module_matches, but this attribute is never used elsewhere in the class. The code seems to intend to use self._compiled_module_calls, which is already initialized as a field. This line appears to be a typo and adds confusion. It should be removed.

vllm/config.py

vllm/v1/worker/intermediates_logging.py

gemini-code-assist · 2025-07-19T03:13:20Z

vllm/v1/worker/intermediates_logging.py

+    if log_call_idx >= 0 and current_call_idx != log_call_idx:
+        should_log = False
+
+    print(f"{should_log=} for {module_name}, {current_call_idx=}, {log_call_idx=}")


This print() statement appears to be for debugging. It will write to stdout on every forward pass of a logged module, which can be very noisy in a production environment and clutter the logs. Please remove it before merging.

gemini-code-assist · 2025-07-19T03:13:20Z

vllm/v1/worker/intermediates_logging.py

+        self.hooks = []
+
+        path = Path(config.output_run_dir)
+        path.mkdir(parents=True)


path.mkdir(parents=True) will raise a FileExistsError if the directory already exists. Since this code might be run multiple times or the directory might be created by another process, you should add exist_ok=True to prevent this error.

Suggested change

path.mkdir(parents=True)

path.mkdir(parents=True, exist_ok=True)

vllm/v1/worker/worker_base.py

mergify · 2025-07-19T03:50:13Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @luccafong.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

houseroad · 2025-07-21T20:42:57Z

vllm/v1/worker/intermediates_logging.py

@@ -0,0 +1,545 @@
+"""


more changes Apply suggestions from code review Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> fix tp Signed-off-by: Lu Fang <fanglu@fb.com> add comparison tool tmp add unit test and fix format Signed-off-by: Lu Fang <fanglu@fb.com>

Signed-off-by: Lu Fang <fanglu@fb.com>

yeqcharlotte · 2025-07-22T04:08:31Z

tools/compare_il.py

+    # Filter steps
+    if steps:
+        step_names = [f"step_{step}" for step in steps]
+        steps_to_include = {step: True for step in step_names}


this can be a set?

yeqcharlotte · 2025-07-22T04:14:33Z

vllm/config.py

@@ -3952,6 +3953,117 @@ class KVEventsConfig:
    """


+@config
+@dataclass
+class IntermediateLoggingConfig:


possible for us to provide some default logging config. so users don't have to set every configuration in this dataclass to dump tensors?

yeqcharlotte · 2025-07-22T04:17:39Z

vllm/config.py

@@ -4409,6 +4521,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


nit: il for intermediate logging is very unintuitive 😅 maybe just spell out the whole thing

yeqcharlotte · 2025-07-22T04:19:55Z

vllm/v1/worker/intermediates_logging.py

@@ -0,0 +1,550 @@
+# SPDX-License-Identifier: Apache-2.0


i think this file belongs to utils not worker tho

yeqcharlotte · 2025-07-22T04:24:42Z

vllm/v1/engine/core.py

@@ -73,6 +73,8 @@ def __init__(self,

        # Setup Model.
        self.model_executor = executor_class(vllm_config)
+        self.collective_rpc("register_intermediate_hooks", args=(vllm_config.il_config,))


would we register intermediate hooks by invoking kwargs even when config is None?

mergify bot added documentation Improvements or additions to documentation v1 labels Jul 19, 2025

gemini-code-assist bot reviewed Jul 19, 2025

View reviewed changes

luccafong changed the title ~~[wip] intermediate logging~~ [wip] [Feature] [V1] intermediate logging Jul 19, 2025

luccafong force-pushed the il_tool branch 3 times, most recently from ab68fd3 to 9252ae5 Compare July 19, 2025 03:44

mergify bot added the needs-rebase label Jul 19, 2025

houseroad reviewed Jul 21, 2025

View reviewed changes

vllm/v1/worker/intermediates_logging.py

@@ -0,0 +1,545 @@

"""

Copy link

Collaborator

houseroad Jul 21, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

copyright?

luccafong force-pushed the il_tool branch from 12ec7d1 to 826726c Compare July 22, 2025 01:38

add comparison script and documentation

8c53027

Signed-off-by: Lu Fang <fanglu@fb.com>

yeqcharlotte reviewed Jul 22, 2025

View reviewed changes

		step_hook = model.register_forward_hook(step_fwd)
		self.hooks.append(("", None, step_hook))

	step_hook = model.register_forward_hook(step_fwd)
	self.hooks.append(("", None, step_hook))
	self.hooks.append(('', None, None, step_hook))

	path.mkdir(parents=True)
	path.mkdir(parents=True, exist_ok=True)

Uh oh!

[wip] [Feature] [V1] intermediate logging #21215

Are you sure you want to change the base?

[wip] [Feature] [V1] intermediate logging #21215

Conversation

luccafong commented Jul 19, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 19, 2025

Uh oh!

luccafong commented Jul 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jul 19, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist bot Jul 19, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jul 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mergify bot commented Jul 19, 2025

Uh oh!

houseroad Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

yeqcharlotte Jul 22, 2025

Choose a reason for hiding this comment

Uh oh!

yeqcharlotte Jul 22, 2025

Choose a reason for hiding this comment

Uh oh!

yeqcharlotte Jul 22, 2025

Choose a reason for hiding this comment

Uh oh!

yeqcharlotte Jul 22, 2025

Choose a reason for hiding this comment

Uh oh!

yeqcharlotte Jul 22, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

luccafong commented Jul 19, 2025 •

edited by github-actions bot

Loading

luccafong commented Jul 19, 2025 •

edited

Loading