Add KV replicate feature when kv_head<tp_size

[kaixuanliu]

[kaixuanliu]

@Cyrilvallez @3outeille This PR tries to support 1 scenario: when the model need to be loaded with tp mode, but the tp_size is larger than kv_head, we need to do KV replicate like the feature in VLLM. I have verified the correctness of this PR using bigcode/starcoder2-3b model w/ kv_head=2 and tp_size=4. pls help review, thx!

[3outeille]

Hey super cool ! Just a note that we are undergoing a refactor of our Tensor Parallel cf #47579. If that's not too much, is it possible to rebase your work on top of the new branch ?

[kaixuanliu]

Here is the example code:

import os

import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh

from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig


MODEL_ID = "bigcode/starcoder2-3b"
PROMPT = "Explain tensor parallelism in one concise sentence."
MAX_NEW_TOKENS = 128


def main():
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    dist.init_process_group(backend="nccl")

    try:
        world_size = dist.get_world_size()
        tp_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=("tp",))

        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            device_mesh=tp_mesh,
            distributed_config=DistributedConfig(tp_size=world_size),
            dtype=torch.bfloat16,
            attn_implementation="sdpa",
        ).eval()
    
        inputs = tokenizer(PROMPT, return_tensors="pt").to(device)
        output_ids = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            use_cache=True,
            cache_implementation="static",
        )

        # TP ranks must produce the same complete sequence. This catches missing all-reduces as well as incorrect
        # KV head replication/repeat_kv handling during cached decoding.
        gathered_output_ids = [torch.empty_like(output_ids) for _ in range(world_size)]
        dist.all_gather(gathered_output_ids, output_ids)
        if not all(torch.equal(output_ids, other) for other in gathered_output_ids[1:]):
            raise AssertionError("TP ranks generated different token sequences.")

        if dist.get_rank() == 0:
            completion = tokenizer.decode(output_ids[0, inputs.input_ids.shape[-1] :], skip_special_tokens=True)
            print(f"PASS: {MODEL_ID} generated consistently with TP={world_size} and KV head replication.")
            print(f"Completion: {completion}")
    finally:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

[kaixuanliu]

Hey super cool ! Just a note that we are undergoing a refactor of our Tensor Parallel cf #47579. If that's not too much, is it possible to rebase your work on top of the new branch ?

Sure, will do the rebase work after #47579 is merged.

[3outeille]

Hey @kaixuanliu, PR has been merged ! Feel free to rebase, i'll review it

[kaixuanliu]

@3outeille ,rebase work is done, pls help review, thx!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[3outeille]

does that still holds for 2D mesh ?

[kaixuanliu]

pls help review again

[3outeille]

double check because iirc id(mesh) as a key doesnt work. Had issue back in the day

import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh

dist.init_process_group("gloo")
mesh = init_device_mesh("cpu", (2,), mesh_dim_names=["tp"])
a, b = mesh["tp"], mesh["tp"]

print(id(a) is id(b)) # False

[kaixuanliu]

Same as above

[3outeille]

to double check it works, pick a model with tp_size > num_kv_heads (i.e: "Qwen/Qwen2.5-VL-3B-Instruct") and add it to TP_DISTRIBUTED_TEST_MODEL_TYPES in test_tensor_parallel_mixin.py and test its forward + backward to see if everything pass

I think we can leave the model in the list so that we can catch KV_replication regression later

[kaixuanliu]

Added test_tp_kv_head_replication to TensorParallelTesterMixin: it forces num_key_value_heads=1 on the tiny config so num_kv_heads < tp_size, asserts the layers actually went through ReplicateKVHeadsParallel, and then runs the full forward + backward (incl. per-parameter grad comparison). This runs for every model already in TP_DISTRIBUTED_TEST_MODEL_TYPES (qwen2/qwen3/qwen3_moe/...), so we get regression coverage without having to wire the TP mixin into a VLM test class (Qwen2_5_VLModelTest doesn't use CausalLMModelTester, so it would be skipped anyway).

[3outeille]

correct for 1D, but not for 2D mesh

[kaixuanliu]

Since apply_tensor_parallelism is always handed device_mesh["tp"] and TP/FSDP/PP are mutually exclusive, this is unreachable today: L194-L206, I have added a guard in __init__ func

[3outeille]

if enable_expert_parallel=True, it will swap out model.tp_plan to ep_plan. Double check if replication is still applied

[kaixuanliu]

Thx for advice!! Have updated.

[ArthurZucker]

Really not aligned as this IMO should be much much simpler to do

[ArthurZucker]

I really don't think it makes sense to put this attention, kv specific argument into a very general shading scheme.

[kaixuanliu]

Well, I agree. It is not a good design indeed..., have fixed it.

[ArthurZucker]

this is highly bloated, and that makes me think it just does not follow the design.
We are probably missing ShardingOps, that you would want for that, otherwise we can just properly resolve, based on the dim, what to do with the weights, regardless of attention or not. The same happens for WP sharding: see #48237

[kaixuanliu]

Well the diff is a little big. But I don't think a purely dim-based rule can cover this case. Also take bigcode/starcoder2-3b as example, num_key_value_heads=2, head_dim=128, so k_proj.out_features=256. On tp=4, 256 % 4 == 0 — a dim-based rule shards it happily, and attention then dies on k_proj(x).view(*input_shape, -1, self.head_dim) with 64 features per rank. Whenever kv_heads < tp_size you get less than one head per rank, so this always happens. It can't be fixed in attention either: rotate_half pairs dim i with i + head_dim/2, and qk/softmax/FA kernels all need a whole head. head_dim is a hard granularity floor, so replication is the only option.

[kaixuanliu]

@ArthurZucker I have refactored the code to align with your comments: the sharding machinery is based on head_dim.

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 33601862252
Result: success | Grafana metrics are not available yet.