From 944c8056382477662f153a0fcc52487d81355281 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Fri, 5 Jun 2026 01:00:59 -0400 Subject: [PATCH 1/2] docs(dev): document tracer.barrier() vLLM bug (Barrier object not shared across invokes) barrier() is broken on the vLLM path (reproduces at tp1/pp1, non-PP): each invoke is serialized into its own globals, so each gets a private copy of the Barrier with its own participants set. The count never reaches n, both invokes take the no-op send(BARRIER, None) branch, the workers block at the barrier and are abandoned, and all post-barrier code is silently dropped. Diagnosed (instrumentation reverted); fix is a design choice (interleaver-owned barrier registry keyed by a serialization-stable id, preferred; or graft the Barrier into canonical globals). Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/developing/barrier-vllm-not-shared.md | 109 +++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 docs/developing/barrier-vllm-not-shared.md diff --git a/docs/developing/barrier-vllm-not-shared.md b/docs/developing/barrier-vllm-not-shared.md new file mode 100644 index 00000000..710ba766 --- /dev/null +++ b/docs/developing/barrier-vllm-not-shared.md @@ -0,0 +1,109 @@ +--- +title: tracer.barrier() is broken on the vLLM path — the Barrier object isn't shared across invokes +one_liner: On vLLM, each invoke is serialized into its own globals, so each gets a private copy of the Barrier; the participant count never completes, both invokes fall through to the no-op branch, and all post-barrier code is silently dropped. +tags: [internals, dev, vllm, barrier, cross-invoke, bug] +status: DIAGNOSED — not fixed (fix is a design choice; see below). Non-PP; reproduces at tp=1/pp=1. +related: [docs/usage/barrier.md, docs/concepts/batching-and-invokers.md, docs/concepts/threading-and-mediators.md] +sources: [src/nnsight/intervention/tracing/tracer.py, src/nnsight/intervention/interleaver.py, src/nnsight/modeling/vllm/model_runners/GPUModelRunner.py] +--- + +# `tracer.barrier()` is broken on the vLLM path + +> **Scope.** `barrier` is an advanced cross-invoke primitive (capture an activation in one invoke, +> use it in another). This bug is **vLLM-path-specific** and **not** related to pipeline parallelism — +> it reproduces at `tensor_parallel_size=1, pipeline_parallel_size=1`. It was surfaced while +> stress-testing PP but is orthogonal to that work. + +## Symptom + +A trace that uses `tracer.barrier(n)` to coordinate two invokes silently drops everything after the +`barrier()` call. The clearest form — a value saved/appended **after** the barrier comes back empty: + +```python +with model.trace(temperature=0.0, top_p=1) as tracer: + barrier = tracer.barrier(2) + captured = list().save() + out = list().save() + with tracer.invoke(clean_prompt): + captured.append(model.model.layers[8].mlp.output.clone()) + barrier() + with tracer.invoke(corrupt_prompt): + barrier() + model.model.layers[8].mlp.output = captured[0] # patch + out.append(model.logits) # <-- never lands +# out == [] -> IndexError on out[0] +``` + +Characterization (Qwen2.5-0.5B, tp1/pp1): + +| case | result | +|---|---| +| no barrier, save in 2nd invoke | works (saved-globals union, commit `8c08897`) | +| barrier, append **before** `barrier()` | survives | +| barrier, append **after** `barrier()` | **dropped** | +| barrier, append in both (before+after) | only the **before**-barrier one survives | + +No exception, no hang. + +## Root cause: the `Barrier` object isn't shared across invokes on vLLM + +`tracer.barrier(n)` returns a `Barrier` holding a `participants` set; calling it does +(`tracing/tracer.py`): + +```python +def __call__(self): + mediator = self.model.interleaver.current + self.participants.add(mediator.name) + if len(self.participants) == self.n_participants: + mediator.send(Events.BARRIER, self.participants) # real barrier + else: + mediator.send(Events.BARRIER, None) # not all here yet +``` + +On the **vLLM path each `tracer.invoke` is serialized and deserialized into its own per-invoke +`__globals__` on the worker** (`GPUModelRunner.process_new_reqs`). The `Barrier` object is referenced by +every invoke, but it is **not a `.save()`d name**, so the canonical-globals *union* added in `8c08897` +(which reconciles only saved names across invokes) never reconciles it. Each invoke therefore gets its +**own copy** of the `Barrier`, each with its **own** `participants` set. + +Consequence: each `barrier()` call adds only its own mediator → `len(participants)` is always 1, never +reaches `n` → **every** invoke takes the `else: send(BARRIER, None)` branch. The real barrier walk in +`Interleaver.handle_barrier_event` only runs when `participants is not None`; with `None` it does nothing +and never `respond()`s to the worker. So both worker threads **block at `barrier()` forever** and are +abandoned when the trace exits — every statement after `barrier()` is silently skipped. + +**Evidence** (instrumented `handle_barrier_event`, since reverted): the walk fired with +`participants=None` on *both* invokes' calls — the real participant set was never assembled. Body prints +placed after `barrier()` never executed; the (correctly grafted) saved list was collected with length 0. + +## Why it works on `LanguageModel` + +`LanguageModel` runs invokes in-process with shared globals — there is no per-invoke serialization, so +`b` is genuinely the same object across invokes and the participant count completes normally. The barrier +docs/examples are written against that path. The break is specific to vLLM's serialize-per-invoke model. + +## Fix direction (a design choice — not yet implemented) + +The barrier's participant/sync state must be **shared across invokes** on the worker, the same way saved +vars are. Two options: + +1. **Interleaver-owned barrier registry, keyed by a serialization-stable barrier id.** `barrier()` counts + participants in `interleaver.barriers[barrier_id]` instead of the per-copy `Barrier.participants`. + The id is assigned at `tracer.barrier()` creation and travels with the serialized object. Most robust; + mirrors how cross-invoke saved state is centralized. Preferred. +2. **Graft the `Barrier` object into canonical globals** (like `8c08897` does for saved names) so every + invoke references the one canonical `Barrier`. Smaller, but requires detecting barrier objects in + globals (they are not in `nnsight_saved_names`). + +Either way, also make the no-participants path not strand the worker (it currently never `respond()`s). + +## Workaround (today) + +Don't rely on a `barrier()` to gate state you collect afterward on vLLM: do the save **before** the +barrier, or collect it in the **first** invoke. (Cross-invoke value *sharing* for an in-place patch may +still be affected, since the patch also runs after `barrier()`.) + +## Reproduce + +`/tmp/pp_stress/cp_isolate.py` (variants v1–v7) and `/tmp/pp_stress/cp_dbg.py` isolate the trigger to the +barrier and show the `participants=None` smoking gun under `CP_DBG=1`. From ed981b8e4dc75d22a137ff74d6d8bad4fedec30d Mon Sep 17 00:00:00 2001 From: khaiwang Date: Sat, 20 Jun 2026 13:41:43 -0400 Subject: [PATCH 2/2] feat(vllm): gather tensor-parallel sharded parameters on read in a trace Under tensor parallelism vLLM shards parameters across ranks (lm_head/embeddings are vocab-sharded; attention/MLP projections are output- or input-sharded). nnsight handed intervention code the local shard, so reading a parameter inside a trace -- e.g. `lm_head.weight[token_id]` to build a steering direction -- returned the wrong row on a rank that does not own that token, silently diverging from single-GPU. This is the parameter analogue of the existing activation gather (VLLMBatcher gathers RowParallelLinear/ColumnParallelLinear I/O). Parameter reads now route through the batcher: Envoy.__getattr__ delegates a tensor attribute to interleaver.batcher.gather_param while interleaving (so the collective fires on every rank); the base Batcher returns it unchanged (non-vLLM and tp=1 untouched), and VLLMBatcher all-gathers the shard to its full logical shape. The sharded dim comes from the module class (RowParallelLinear -> input dim; ColumnParallelLinear/VocabParallelEmbedding -> output/vocab dim), because vLLM sets BOTH output_dim and input_dim on every linear weight (they label the dims, not which is sharded); vocab padding is stripped to org_vocab_size. Verified on Qwen2.5-0.5B tp=2: lm_head, qkv_proj, gate_up_proj, o_proj and down_proj all gather to the full tp=1 shape with matching norms; a steering write reading lm_head[token_id] goes from divergent (maxabs 26, wrong global token in the top-5) to equivalent. Regression tests in tests/test_tp_param_gather.py. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_015y5Sy9vzzc9YJZXCtewSdQ --- src/nnsight/intervention/batching.py | 15 +++++ src/nnsight/intervention/envoy.py | 9 +++ src/nnsight/modeling/vllm/batching.py | 50 ++++++++++++++ tests/test_tp_param_gather.py | 95 +++++++++++++++++++++++++++ 4 files changed, 169 insertions(+) create mode 100644 tests/test_tp_param_gather.py diff --git a/src/nnsight/intervention/batching.py b/src/nnsight/intervention/batching.py index c3fcf088..f7fdc510 100755 --- a/src/nnsight/intervention/batching.py +++ b/src/nnsight/intervention/batching.py @@ -195,6 +195,21 @@ def batch( return (args, kwargs), None + def gather_param(self, module, value): + """Transform a module parameter read inside a trace before it reaches intervention code. + + The base batcher models no parallelism, so the parameter is returned unchanged. A + parallelism-aware subclass (``VLLMBatcher``) overrides this to all-gather a tensor-parallel + sharded parameter back to its full logical shape, so intervention code reading e.g. + ``module.weight[token_id]`` sees the complete weight on every rank rather than a local shard. + + Args: + module: The module the parameter was read from (its TP/shard metadata lives here). + value: The raw attribute value (e.g. the local-shard ``weight`` tensor). + """ + + return value + def narrow(self, batch_group: Optional[List[int]]) -> Any: """Extract an invoke's slice from the current activation value. diff --git a/src/nnsight/intervention/envoy.py b/src/nnsight/intervention/envoy.py index 8c6663c6..79c45a8b 100755 --- a/src/nnsight/intervention/envoy.py +++ b/src/nnsight/intervention/envoy.py @@ -1031,6 +1031,15 @@ def trace(*args, trace: bool = True, **kwargs): return util.fetch_attr(self, value.__path__[len(self.path) :]) return self._add_envoy(value, name) else: + # Parameter/buffer reads inside a trace are routed through the batcher so a + # parallelism-aware batcher (vLLM tensor parallelism) can gather a sharded weight + # back to its full logical shape — e.g. a vocab-sharded `lm_head.weight`, so that + # `head.weight[token_id]` indexes the true row on every rank, not a local shard. + # `interleaving` guarantees `self.interleaver` is set AND that we are inside a trace + # on every rank (so the collective gather can't desync); the base batcher's + # gather_param is a no-op, leaving non-parallel models untouched. + if self.interleaving and isinstance(value, torch.Tensor): + value = self.interleaver.batcher.gather_param(self._module, value) return value else: raise AttributeError(f"{self} has no attribute {name}") diff --git a/src/nnsight/modeling/vllm/batching.py b/src/nnsight/modeling/vllm/batching.py index 0b6855be..a358b597 100644 --- a/src/nnsight/modeling/vllm/batching.py +++ b/src/nnsight/modeling/vllm/batching.py @@ -10,6 +10,7 @@ split_tensor_along_last_dim, tensor_model_parallel_all_reduce, ) +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding class VLLMBatcher(Batcher): @@ -29,9 +30,14 @@ def __init__(self, *args, **kwargs): self.parallel = False self.gathered = False self.type = None + # True once wrap() has run, i.e. only when tensor_parallel_size > 1 (wrap is gated on that in + # GPUModelRunner). gather_param keys off this so tp=1 is byte-identical (no gather, no unpad). + self.tp_active = False def wrap(self, model: Envoy): + self.tp_active = True + def pre_input_hook(module: torch.nn.Module, args: Any, kwargs: Any): self.current_module = module self.type = "input" @@ -166,3 +172,47 @@ def swap(self, batch_group: Union[int, None], swap_value: Any): self.check_gathered() return super().swap(batch_group, swap_value) + + def gather_param(self, module, value): + """All-gather a tensor-parallel sharded parameter back to its full logical shape. + + Called from ``Envoy.__getattr__`` when intervention code reads a module attribute (e.g. + ``lm_head.weight``) inside a trace. vLLM shards parameters across TP ranks; without this the + trace sees only the local shard, so a vocab-indexed read like ``head.weight[token_id]`` lands + on the wrong global row on a rank that does not own that token. + + The sharded dim is keyed off the MODULE CLASS — vLLM's own parallelism abstraction — NOT the + parameter's ``output_dim``/``input_dim`` attrs (vLLM sets BOTH on every linear weight; they + only label the dims, they don't say which is sharded): + - ``RowParallelLinear`` shards the INPUT dim -> dim 1 (weight only; bias is replicated) + - ``ColumnParallelLinear`` shards the OUTPUT dim -> dim 0 (weight and bias) + - ``VocabParallelEmbedding`` shards the VOCAB dim -> dim 0 (then strip vocab padding) + Subclasses are covered by isinstance (QKV/MergedColumn -> Column; ParallelLMHead -> Vocab). + Anything else (e.g. ReplicatedLinear, norms) is returned unchanged. Only fires when TP is + active (``tp_active``), so tp=1 is byte-identical. + """ + + if not self.tp_active: + return value + + if isinstance(module, RowParallelLinear): + # Only the 2-D weight is sharded (on the input dim); the bias is replicated, leave it. + if not (isinstance(value, torch.Tensor) and value.ndim >= 2): + return value + shard_dim = 1 + elif isinstance(module, (ColumnParallelLinear, VocabParallelEmbedding)): + shard_dim = 0 + else: + return value # not a TP-sharded module we gather + + full = tensor_model_parallel_all_gather(value.data, dim=shard_dim) + + # VocabParallelEmbedding/ParallelLMHead pad the vocab to a TP-divisible size; for standard + # (non-LoRA) models the real rows are [0:org_vocab_size] with padding appended at the end, so + # strip it to recover the true [vocab_size, hidden]. (Interleaved per-partition padding from + # added-vocab/LoRA would need shard_indices-based reconstruction — not handled here.) + org_vocab_size = getattr(module, "org_vocab_size", None) + if org_vocab_size is not None and shard_dim == 0: + full = full[:org_vocab_size] + + return full diff --git a/tests/test_tp_param_gather.py b/tests/test_tp_param_gather.py new file mode 100644 index 00000000..95187a1a --- /dev/null +++ b/tests/test_tp_param_gather.py @@ -0,0 +1,95 @@ +"""Tests for tensor-parallel PARAMETER gathering. + +vLLM shards parameters across TP ranks (e.g. ``lm_head``/``embed_tokens`` are vocab-sharded by +``VocabParallelEmbedding``). Without gathering, intervention code that reads a parameter inside a +trace sees only the local shard, so a vocab-indexed read like ``lm_head.weight[token_id]`` lands on +the WRONG global row on a rank that does not own that token (it indexes into that rank's shard). This +is the parameter analogue of the activation gather in ``VLLMBatcher`` — see +``Envoy.__getattr__``/``Batcher.gather_param``. + +Run with: pytest tests/test_tp_param_gather.py --tp 2 -v +""" +import pytest +import torch + +try: + from nnsight.modeling.vllm import VLLM +except Exception as e: + pytest.skip(f"Skipping VLLM tests: \n{e}", allow_module_level=True) + +import nnsight + +MODEL = "Qwen/Qwen2.5-0.5B-Instruct" +PROMPT = "The Eiffel Tower is in the city of" + + +@pytest.fixture(scope="module") +def tp(request): + tp = request.config.getoption("--tp") + if tp > torch.cuda.device_count() or tp < 1: + pytest.exit("--tp can't be higher than the number of available GPUs.") + return tp + + +@pytest.fixture(scope="module") +def model(tp: int): + return VLLM( + MODEL, + tensor_parallel_size=tp, + gpu_memory_utilization=0.3, + dispatch=True, + dtype=torch.float16, + ) + + +class TestTPParamGather: + @torch.no_grad() + def test_lm_head_weight_is_full_vocab(self, tp, model): + """Reading ``lm_head.weight`` inside a trace yields the full ``[vocab, hidden]`` tensor on + every rank. At tp=1 it always did; under TP>1 it must be gathered, not the per-rank shard.""" + vocab = model.model.config.vocab_size + with model.trace(temperature=0.0, max_tokens=1) as tracer: + with tracer.invoke(PROMPT): + rows = nnsight.save(model.lm_head.weight.shape[0]) + assert int(rows) == vocab, ( + f"lm_head.weight has {int(rows)} rows under tp={tp}, expected full vocab {vocab}; " + f"a sharded ~{vocab // max(tp, 1)} means the TP parameter gather did not fire." + ) + + @torch.no_grad() + def test_upper_shard_token_row_accessible(self, tp, model): + """A token id in the UPPER vocab shard (>= vocab/tp) must be indexable and return a real + unembed row. On the raw per-rank shard that index is out of range / a different token, so this + only passes once the parameter is gathered to its full logical shape.""" + if tp < 2: + pytest.skip("Upper-shard indexing is only meaningful with TP>1") + vocab = model.model.config.vocab_size + tid = vocab - 5 # firmly inside the last rank's shard + with model.trace(temperature=0.0, max_tokens=1) as tracer: + with tracer.invoke(PROMPT): + row = nnsight.save(model.lm_head.weight[tid].float()) + assert torch.isfinite(row).all() + assert row.norm() > 0 + + @torch.no_grad() + def test_row_and_column_parallel_weights_gathered(self, tp, model): + """EVERY TP-sharded weight type gathers to its full [output_size, input_size], not just the + vocab-parallel lm_head. RowParallelLinear (o_proj) shards the INPUT dim (1); ColumnParallelLinear + (qkv_proj) shards the OUTPUT dim (0). Regression for keying the gather dim off the parameter's + output_dim/input_dim attrs (vLLM sets BOTH on every linear weight), which gathered row-parallel + weights on the wrong dim — e.g. an o_proj shard [896,448] became [1792,448] instead of [896,896].""" + if tp < 2: + pytest.skip("Sharding only occurs with TP>1") + L = model.model.config.num_hidden_layers // 2 + o = model.model.layers[L].self_attn.o_proj._module # RowParallelLinear + qkv = model.model.layers[L].self_attn.qkv_proj._module # ColumnParallelLinear (QKV) + exp_o = (o.output_size, o.input_size) # full logical shape (LinearBase stores both) + exp_qkv = (qkv.output_size, qkv.input_size) + with model.trace(temperature=0.0, max_tokens=1) as tracer: + with tracer.invoke(PROMPT): + ow = model.model.layers[L].self_attn.o_proj.weight + o_shape = nnsight.save([ow.shape[0], ow.shape[1]]) + qw = model.model.layers[L].self_attn.qkv_proj.weight + qkv_shape = nnsight.save([qw.shape[0], qw.shape[1]]) + assert tuple(int(x) for x in o_shape) == exp_o, (tuple(int(x) for x in o_shape), exp_o) + assert tuple(int(x) for x in qkv_shape) == exp_qkv, (tuple(int(x) for x in qkv_shape), exp_qkv)