From 59d09d44c98700f870a631dd6c69421e1e8fdc10 Mon Sep 17 00:00:00 2001 From: khaiwang Date: Tue, 25 Aug 2026 10:49:51 -0400 Subject: [PATCH] Gather tensor-parallel sharded parameters on read in a trace Under tensor parallelism a parameter read inside a trace returned this rank's slice. A steering cell reading `lm_head.weight[token_id]` on the rank that does not own the token indexed a different token's row (Qwen2.5-0.5B, tp=2: `lm_head.weight.shape[0]` was 75968, half vocab), so the steered output diverged from the single-GPU run without an error. `ParallelEnvoy.__getattr__` (vLLM) and `TPEnvoy.__getattr__` (transformers) now all-gather a tensor attribute read while interleaving, gated on the interleaver's `fragments.enabled`, the same check their ad-hoc `__call__` uses, so a one-rank engine is untouched. vLLM: the sharded dim is the `output_dim` / `input_dim` stamp vLLM puts on each parameter it shards (row-parallel splits the input dim, every other parallel layer the output dim); a tensor without the stamp (a row-parallel bias, a scale) is replicated and passes through. The vocab-parallel head's padding rows are dropped to `org_vocab_size`. transformers: the dim comes from transformers' own `ALL_PARALLEL_STYLES.plan_to_weight_dim` / `plan_to_bias_dim` table and the gather is its `gather_full_tensor`. Fused weights (`qkv_proj`, `gate_up_proj`, `packed_colwise`) come back with rows grouped by rank, the layout their gathered `.output` already has, so `x @ weight.T` and `.output` agree. The gathered tensor is a copy: an in-place edit to it does not reach the model. Outside a trace `layer.weight` is still the slice. Verified on two A100s: tests/vllm/test_tensor_parallel.py 25/25 at tp=2 (vllm 0.19.1), tests/tp 42 passed at tp=2 (transformers 5.15.0). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01HEt5ijqEYH9RN64FNzPtpz --- docs/models/tensor-parallel.md | 27 +++++++++---- src/nnsight/modeling/tp/envoys.py | 53 ++++++++++++++++++++++++- src/nnsight/modeling/vllm/envoys.py | 54 ++++++++++++++++++++++++- tests/tp/test_sharded_tracing.py | 6 ++- tests/tp/worker.py | 16 ++++++++ tests/vllm/test_tensor_parallel.py | 61 +++++++++++++++++++++++++++++ 6 files changed, 205 insertions(+), 12 deletions(-) diff --git a/docs/models/tensor-parallel.md b/docs/models/tensor-parallel.md index a9e0f794..fee8dd17 100644 --- a/docs/models/tensor-parallel.md +++ b/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 diff --git a/src/nnsight/modeling/tp/envoys.py b/src/nnsight/modeling/tp/envoys.py index ac34d737..62f14f80 100644 --- a/src/nnsight/modeling/tp/envoys.py +++ b/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`. diff --git a/src/nnsight/modeling/vllm/envoys.py b/src/nnsight/modeling/vllm/envoys.py index 7f11ca33..30501541 100644 --- a/src/nnsight/modeling/vllm/envoys.py +++ b/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 diff --git a/tests/tp/test_sharded_tracing.py b/tests/tp/test_sharded_tracing.py index 7dcf7a39..7a9f320a 100644 --- a/tests/tp/test_sharded_tracing.py +++ b/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", ] diff --git a/tests/tp/worker.py b/tests/tp/worker.py index 944f9e49..923bd03a 100644 --- a/tests/tp/worker.py +++ b/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() diff --git a/tests/vllm/test_tensor_parallel.py b/tests/vllm/test_tensor_parallel.py index f8e9bbb7..c6e0faa8 100644 --- a/tests/vllm/test_tensor_parallel.py +++ b/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]