[Feat] Support T5 Tensor Parallelism (#1881)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>

tests/diffusion/models/t5_encoder/test_t5_encoder_tp.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import pytest
+import torch
+from transformers import T5Config
+from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+from vllm_omni.diffusion.models.t5_encoder.t5_encoder import (
+    T5EncoderModel,
+    T5SelfAttention,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder"
+
+SMALL_T5_CONFIG = dict(
+    d_model=64,
+    d_kv=8,
+    d_ff=128,
+    num_heads=8,
+    num_layers=2,
+    vocab_size=256,
+    relative_attention_num_buckets=32,
+    relative_attention_max_distance=128,
+    is_gated_act=True,
+    dense_act_fn="gelu_new",
+    layer_norm_epsilon=1e-6,
+    feed_forward_proj="gated-gelu",
+)
+
+
+@pytest.fixture(scope="module")
+def t5_config() -> T5Config:
+    return T5Config(**SMALL_T5_CONFIG)
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_tp_group(monkeypatch, mocker):
+    """Set up TP=2, rank=0, VllmConfig, and mock activation for all tests."""
+    device_config = DeviceConfig(device="cpu")
+
+    # TP world size
+    monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size",
+        lambda: 2,
+    )
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0)
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank",
+        lambda: 0,
+    )
+
+    # TP group
+    mock_tp_group = mocker.MagicMock()
+    mock_tp_group.world_size = 2
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group)
+
+    monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU())
+
+    with set_current_vllm_config(VllmConfig(device_config=device_config)):
+        yield
+
+
+class TestRelativePositionBiasTPSlicing:
+    """Verify compute_bias slices heads correctly per TP rank."""
+
+    def test_compute_bias_shape(self, t5_config):
+        attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
+
+        seq_len = 6
+        bias = attn.compute_bias(seq_len, seq_len, device=torch.device("cpu"))
+
+        local_heads = t5_config.num_heads // 2
+        assert bias.shape == (1, local_heads, seq_len, seq_len)
+
+    def test_all_ranks_cover_all_heads(self, t5_config, monkeypatch):
+        seq_len = 4
+
+        biases = []
+        ref_weight = None
+        for rank in range(2):
+            monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda r=rank: r)
+            attn = T5SelfAttention(t5_config, has_relative_attention_bias=True)
+            if rank > 0:
+                attn.relative_attention_bias.weight.data.copy_(ref_weight)
+            else:
+                ref_weight = attn.relative_attention_bias.weight.data.clone()
+            biases.append(attn.compute_bias(seq_len, seq_len, device=torch.device("cpu")))
+
+        full_bias = torch.cat(biases, dim=1)
+        assert full_bias.shape == (1, t5_config.num_heads, seq_len, seq_len)
+
+
+class TestT5EncoderModelWeightLoading:
+    """Test weight loading at the top-level T5EncoderModel."""
+
+    def test_model_instantiation(self, t5_config):
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        assert model.config is t5_config
+        assert model.encoder is not None
+        assert len(model.encoder.block) == t5_config.num_layers
+
+    def test_embedding_shape(self, t5_config):
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        assert model.shared.embedding_dim == t5_config.d_model
+
+    def test_embed_input_ids(self, t5_config, monkeypatch):
+        # Verify method and output shape
+        model = T5EncoderModel(t5_config, prefix="text_encoder")
+
+        # Mock all-reduce to be identity (no actual TP communication)
+        monkeypatch.setattr(
+            "vllm.model_executor.layers.vocab_parallel_embedding.tensor_model_parallel_all_reduce",
+            lambda x: x,
+        )
+
+        input_ids = torch.randint(0, t5_config.vocab_size, (2, 8))
+        embeddings = model.embed_input_ids(input_ids)
+
+        assert embeddings.shape == (2, 8, t5_config.d_model)
+
+    def test_qkv_weights_loaded_through_blocks(self):
+        # Verify that HF-style separate Q/K/V weights can be loaded through
+        # the block hierarchy
+        config = T5Config(**{**SMALL_T5_CONFIG, "num_layers": 1})
+        model = T5EncoderModel(config, prefix="text_encoder")
+
+        inner_dim = config.num_heads * config.d_kv
+        prefix = "encoder.block.0.layer.0.SelfAttention."
+
+        loaded = model.load_weights(
+            [
+                (prefix + "q.weight", torch.randn(inner_dim, config.d_model)),
+                (prefix + "k.weight", torch.randn(inner_dim, config.d_model)),
+                (prefix + "v.weight", torch.randn(inner_dim, config.d_model)),
+            ]
+        )
+
+        assert len(loaded) > 0
+        attn = model.encoder.block[0].layer[0].SelfAttention
+        expected_qkv_dim = 3 * (config.num_heads // 2) * config.d_kv
+        assert attn.qkv_proj.weight.shape == (expected_qkv_dim, config.d_model)
+
+
+class TestTPConstraints:
+    """Verify that invalid TP configurations raise clear errors."""
+
+    def test_num_heads_not_divisible_by_tp(self):
+        config = T5Config(**{**SMALL_T5_CONFIG, "num_heads": 7})
+        with pytest.raises(AssertionError, match=r"num_heads.*must be divisible by tp_size"):
+            T5SelfAttention(config)
+
+    def test_num_heads_divisible_by_tp(self, t5_config):
+        attn = T5SelfAttention(t5_config)
+        assert attn.n_heads_per_partition == 4