Gather tensor-parallel sharded parameters on read in a trace

[khaiwang]

Problem

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 the vocab), so the steered output diverged from the single-GPU run without an error.

Change

Rebuilt on 0.8 inside the TP envoy structures (the earlier version of this PR was on dev and hooked Batcher / base Envoy.__getattr__; none of that remains).

ParallelEnvoy.__getattr__ (modeling/vllm/envoys.py) and TPEnvoy.__getattr__ (modeling/tp/envoys.py) 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 (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.

Semantics stated in docs/models/tensor-parallel.md (this replaces the "parameters are not gathered" paragraph):

• One all-gather per read, on every rank, allocating the full tensor per rank. Read once and reuse rather than indexing layer.weight in a loop.
• Every rank must perform the read; a read under rank-dependent control flow deadlocks, like any collective in a block.
• The gathered tensor is a copy: layer.weight[i] = v does not reach the model.
• 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.
• Outside a trace layer.weight is still the slice.

Verification (two A100s, tp=2)

• tests/vllm/test_tensor_parallel.py (vllm 0.19.1, Qwen2.5-0.5B): 25/25, 19 existing plus new TestShardedParameters (row-parallel weights torch.equal to the tp=1 reference, fused column weights equal as a row multiset, lm_head.weight full vocab with weight[vocab-5] equal to the reference row, weight outside a trace still the slice).
• tests/tp (transformers 5.15.0, tiny-random-Llama): 42 passed, including four new recorded values (gate_proj_weight, down_proj_weight, lm_head_weight, and a lens computed as norm(hidden) @ lm_head.weight.T), bit-exact across ranks and within drift of the single-GPU run.

Note for the tests: saving a raw vLLM nn.Parameter (envoy.weight.save()) on a tp=1 engine fails at collect_nnsight pickling (cannot pickle 'weakref.ReferenceType', vLLM's weight_loader stamp). Independent of this change; the tests save .weight.data.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HEt5ijqEYH9RN64FNzPtpz

[JadenFiotto-Kaufman]

@khaiwang Hey can you make this off of the 0.8 branch? Theres a TPEnvoy and TPInterleaver so I'd put them there. Theres a similar thing for modeling/tp/ for HF Transformers tensor parallelism