Gather tensor-parallel sharded parameters on read in a trace

docs/models/tensor-parallel.md

@@ -80,18 +80,31 @@ final norm all arrive complete. Only two kinds of value are really a slice:
 | Whole modules | no | `model.layers[i].output`, `mlp.output`, `norm.output` |
 | The LM head | no — gathered by transformers | `lm_head.output` |
 | Embeddings | no — all-reduced | `embed_tokens.output` |
-| **Parameters** | **yes — not gathered** | `q_proj.weight`, `down_proj.weight` |
+| **Parameters** | yes, gathered for you | `q_proj.weight`, `down_proj.weight`, `lm_head.weight` |
 
 Calling a sharded module **ad hoc** — the logit lens, `model.lm_head(hidden)` —
 is corrected the same way its activations are, so it returns the full-width
 result it would on one GPU.
 
-Parameters are the exception, and the one place a trace does not read as it would
-on one GPU: `layer.weight` is this rank's slice, at `1/tp_size` of the real
-width. Weights are what tensor parallelism exists to split, so nnsight does not
-quietly reassemble one — that would allocate the whole tensor on every rank, in
-the situation where memory was tight enough to reach for TP. Gather it yourself
-if you need it whole, and remember every rank must do so.
+A **parameter** read inside a trace is gathered too: `lm_head.weight[token_id]`
+is the real row on every rank, and `layer.weight.shape` is the full shape. Four
+things to know about it:
+
+- It costs one all-gather per read and allocates the whole tensor on every rank,
+  in the situation where memory was tight enough to reach for TP. Read it once
+  and reuse the result rather than indexing `layer.weight` in a loop.
+- Every rank must perform the read. A read under rank-dependent control flow
+  deadlocks, as any collective in a block would.
+- The gathered tensor is a copy. `layer.weight[i] = v` edits the copy and never
+  reaches the model; edit activations through `.output` instead. Assigning
+  `layer.weight = t` sets an attribute on the envoy and shadows the module's
+  parameter from then on.
+- A fused weight (`qkv_proj`, `gate_up_proj`) comes back with its rows grouped
+  by rank, the same layout its gathered `.output` has, so `x @ weight.T` and
+  `.output` agree.
+
+Outside a trace `layer.weight` is still this rank's slice, at `1/tp_size` of
+the real width.
 
 The gather only fires when an intervention is actually parked on that location,
 so reading a handful of locations does not pay for the hundreds you ignored. A

src/nnsight/modeling/tp/envoys.py

@@ -27,8 +27,11 @@
 input is conditional: see
 [`SPLITS_ITS_OWN_INPUT`][nnsight.modeling.tp.envoys.SPLITS_ITS_OWN_INPUT].
 
-Parameters are left alone. ``layer.weight`` is this rank's real slice, as it is
-anywhere else under transformers tensor parallelism.
+A **parameter read** is the other thing outside the bracket. ``layer.weight``
+is this rank's slice, so ``lm_head.weight[token_id]`` on the wrong rank indexes
+a different token's row. Read inside a trace it is all-gathered to its full
+shape (`TPEnvoy.__getattr__`), off transformers' own style-to-dim table;
+outside a trace it is still the slice.
 """
 
 from __future__ import annotations
@@ -101,6 +104,52 @@ def __call__(self, *args: Any, hook: bool = False, **kwargs: Any) -> Any:
             result = _gather(result, module._hf_device_mesh)
         return result
 
+    def __getattr__(self, name: str) -> Any:
+        """A tensor attribute read inside a trace, gathered to its full shape.
+
+        Which dim a style splits, for a weight and for a bias, is transformers'
+        own table (``ALL_PARALLEL_STYLES.plan_to_weight_dim`` /
+        ``plan_to_bias_dim``, the one its checkpoint saver uses); a style that
+        replicates the parameter maps to ``None`` and passes through.
+
+        One all-gather per read, on every rank: a read under rank-dependent
+        control flow deadlocks, the same condition every collective here
+        carries. A ``packed_colwise`` weight comes back with its rows grouped
+        by rank, the same layout its gathered ``.output`` has. The result is a
+        copy: ``layer.weight[i] = v`` does not reach the model.
+        """
+        value = super().__getattr__(name)
+        fragments = self.interleaver.fragments
+        if (
+            not isinstance(value, torch.Tensor)
+            or not self.interleaver.interleaving
+            or fragments is None
+            or not fragments.enabled
+        ):
+            return value
+
+        module = self._module
+        # Stamped by transformers only on a module it sharded.
+        style = getattr(module, "_hf_tp_plan", None)
+        if style is None or module._hf_device_mesh.size() < 2:
+            return value
+        mesh = module._hf_device_mesh
+
+        from transformers.integrations.tensor_parallel import (
+            ALL_PARALLEL_STYLES,
+            gather_full_tensor,
+        )
+
+        dims = (
+            ALL_PARALLEL_STYLES.plan_to_bias_dim
+            if name == "bias"
+            else ALL_PARALLEL_STYLES.plan_to_weight_dim
+        )
+        dim = dims.get(style)
+        if dim is None or value.dim() == 0:
+            return value
+        return gather_full_tensor(value.data, dim, mesh)
+
 
 def tp_envoys() -> dict:
     """The ``envoys`` map pairing shardable module types with `TPEnvoy`.

src/nnsight/modeling/vllm/envoys.py

@@ -15,8 +15,10 @@
 [`VLLMFragments`][nnsight.modeling.vllm.fragments.VLLMFragments] already recorded
 for exactly this envoy's two locations.
 
-Parameters are left alone. ``layer.weight`` is this rank's real slice here, as it
-is anywhere else in vLLM, and gathering one is the caller's business.
+A **parameter read** is the other thing outside the bracket. ``layer.weight``
+is this rank's slice, so ``lm_head.weight[token_id]`` on the wrong rank indexes
+a different token's row. Read inside a trace it is all-gathered to its full
+shape (`ParallelEnvoy.__getattr__`); outside a trace it is still the slice.
 """
 
 from __future__ import annotations
@@ -92,3 +94,51 @@ def __call__(self, *args: Any, hook: bool = False, **kwargs: Any) -> Any:
             result = fragments.whole(outof, result)
 
         return result
+
+    def __getattr__(self, name: str) -> Any:
+        """A tensor attribute read inside a trace, gathered to its full shape.
+
+        vLLM stamps every parameter it shards with the dim it split
+        (``output_dim`` / ``input_dim``); a tensor without that stamp (a
+        row-parallel bias, a scale) is replicated and passes through. A
+        row-parallel layer splits its input dim, every other parallel layer its
+        output dim. The vocab-parallel head is padded to a TP-divisible size, so
+        the padding rows are dropped to give the true ``[vocab, hidden]``.
+
+        One all-gather per read, on every rank: a read under rank-dependent
+        control flow deadlocks, the same condition every collective here
+        carries. A fused layer (``qkv_proj``, ``gate_up_proj``) comes back with
+        its rows grouped by rank, the same layout its gathered ``.output`` has.
+        The result is a copy: ``layer.weight[i] = v`` does not reach the model.
+        """
+        value = super().__getattr__(name)
+        fragments = self.interleaver.fragments
+        if (
+            not isinstance(value, torch.Tensor)
+            or not self.interleaver.interleaving
+            or fragments is None
+            or not fragments.enabled
+        ):
+            return value
+
+        module = self._module
+        if module.tp_size < 2:
+            # Built with `disable_tp=True`: replicated on every rank.
+            return value
+
+        from vllm.distributed.communication_op import tensor_model_parallel_all_gather
+        from vllm.model_executor.layers.linear import RowParallelLinear
+        from vllm.model_executor.layers.vocab_parallel_embedding import (
+            VocabParallelEmbedding,
+        )
+
+        axis = "input_dim" if isinstance(module, RowParallelLinear) else "output_dim"
+        # Genuinely optional: vLLM sets it only on the parameters it sharded.
+        dim = getattr(value, axis, None)
+        if dim is None or value.dim() <= dim:
+            return value
+
+        whole = tensor_model_parallel_all_gather(value.data, dim=dim)
+        if isinstance(module, VocabParallelEmbedding) and dim == 0:
+            whole = whole[: module.org_vocab_size]
+        return whole

tests/tp/test_sharded_tracing.py

@@ -126,7 +126,7 @@ def runs(tmp_path_factory) -> tuple[dict, list[dict]]:
 
 
 # Every value the worker records: a sharded read, a boundary-straddling edit, a
-# cached read, and a generation.
+# cached read, a generation, and a parameter read.
 VALUES = [
     "gate_proj_out",
     "down_proj_in",
@@ -141,6 +141,10 @@ def runs(tmp_path_factory) -> tuple[dict, list[dict]]:
     "adhoc_colwise",
     "adhoc_rowwise",
     "adhoc_lens",
+    "gate_proj_weight",
+    "down_proj_weight",
+    "lm_head_weight",
+    "weight_lens",
 ]
 
 

tests/tp/worker.py

@@ -107,6 +107,22 @@ def record(name: str, tensor: torch.Tensor) -> None:
         record("adhoc_rowwise", layer.mlp.down_proj(down_in).save())
         record("adhoc_lens", model.lm_head(model.model.norm(hidden)).save())
 
+    # A parameter read inside a trace is the full weight: a column-parallel one
+    # split on its output dim, a row-parallel one on its input dim, and the
+    # gathered head. Reading it in place of the head's forward reproduces the
+    # lens, so the layout is the single-GPU one and not merely the right shape.
+    with model.trace(PROMPT):
+        record("gate_proj_weight", layer.mlp.gate_proj.weight.save())
+        record("down_proj_weight", layer.mlp.down_proj.weight.save())
+        record("lm_head_weight", model.lm_head.weight.save())
+        hidden = model.model.norm(layer.output[0])
+        record("weight_lens", (hidden @ model.lm_head.weight.T).save())
+    if args.tp > 1:
+        local_rows = layer.mlp.gate_proj._module.weight.shape[0]
+        assert results["gate_proj_weight"].shape[0] == local_rows * args.tp, (
+            "gate_proj.weight read in a trace is still this rank's slice"
+        )
+
     # A cache must record whole tensors too, and only for what it selects.
     with model.trace(PROMPT) as tracer:
         cache = tracer.cache(modules=[layer.mlp.gate_proj], include_inputs=True).save()

tests/vllm/test_tensor_parallel.py

@@ -262,3 +262,64 @@ def test_an_installed_block_leaves_nothing_either(self, vllm_qwen_tp, ET_prompt)
             assert counts == [0] * len(counts), f"workers left behind: {counts}"
         finally:
             edit.clear()
+
+
+class TestShardedParameters:
+    """A parameter read inside a trace is the full tensor, not this rank's shard."""
+
+    @pytest.mark.parametrize("path", ROW_PARALLEL)
+    @torch.no_grad()
+    def test_row_weight_matches_reference(
+        self, vllm_qwen_ref, vllm_qwen_tp, path, ET_prompt
+    ):
+        with vllm_qwen_ref.trace(ET_prompt, temperature=0.0, top_p=1):
+            ref = _submodule(vllm_qwen_ref, path).weight.data.save()
+        with vllm_qwen_tp.trace(ET_prompt, temperature=0.0, top_p=1):
+            tp = _submodule(vllm_qwen_tp, path).weight.data.save()
+
+        # A row-parallel weight splits the input dim; gathered in rank order it
+        # is laid out exactly as the single-rank weight.
+        assert tp.shape == ref.shape
+        assert torch.equal(tp.cpu(), ref.cpu())
+
+    @pytest.mark.parametrize("path", COLUMN_PARALLEL)
+    @torch.no_grad()
+    def test_column_weight_gathers_every_row(
+        self, vllm_qwen_ref, vllm_qwen_tp, path, ET_prompt
+    ):
+        with vllm_qwen_ref.trace(ET_prompt, temperature=0.0, top_p=1):
+            ref = _submodule(vllm_qwen_ref, path).weight.data.save()
+        with vllm_qwen_tp.trace(ET_prompt, temperature=0.0, top_p=1):
+            tp = _submodule(vllm_qwen_tp, path).weight.data.save()
+
+        # Fused column-parallel weights come back with rows grouped by rank, the
+        # layout the gathered output has, so compare as a multiset of rows.
+        assert tp.shape == ref.shape
+        ref_rows = ref.float().cpu().sum(dim=-1).sort().values
+        tp_rows = tp.float().cpu().sum(dim=-1).sort().values
+        assert torch.allclose(tp_rows, ref_rows)
+
+    @torch.no_grad()
+    def test_lm_head_weight_full_vocab(self, vllm_qwen_ref, vllm_qwen_tp, ET_prompt):
+        with vllm_qwen_ref.trace(ET_prompt, temperature=0.0, top_p=1):
+            ref = vllm_qwen_ref.lm_head.weight.data.save()
+        vocab = ref.shape[0]
+        tid = vocab - 5  # a row the last rank owns
+
+        with vllm_qwen_tp.trace(ET_prompt, temperature=0.0, top_p=1):
+            tp = vllm_qwen_tp.lm_head.weight.data.save()
+            tp_row = vllm_qwen_tp.lm_head.weight[tid].save()
+
+        # Full vocab with vLLM's TP padding rows dropped, and the upper-shard
+        # row is the real token's row, not an index into this rank's slice.
+        assert tp.shape == ref.shape
+        assert torch.equal(tp_row.cpu(), ref[tid].cpu())
+
+    @torch.no_grad()
+    def test_weight_outside_trace_is_slice(self, vllm_qwen_ref, vllm_qwen_tp):
+        # The gather is a trace-time correction; outside one the module holds
+        # this rank's slice, as anywhere else in vLLM.
+        tp = vllm_qwen_tp.model.layers[LAYER].mlp.down_proj.weight
+        ref = vllm_qwen_ref.model.layers[LAYER].mlp.down_proj.weight
+        ranks = vllm_qwen_tp.model.layers[LAYER].mlp.down_proj._module.tp_size
+        assert tp.shape[1] * ranks == ref.shape[1]