Gather tensor-parallel sharded parameters on read in a trace

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

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.
 

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}")

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

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)