tpu-inference 0.11.1.dev202511150811__py3-none-any.whl → 0.11.1.dev202511270815__py3-none-any.whl

tests/kernels/fused_moe_v1_test.py

@@ -1,7 +1,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from absl.testing import absltest
+from absl.testing import absltest, parameterized
 from jax._src import test_util as jtu
 from jax.sharding import Mesh
 
@@ -10,6 +10,15 @@ from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe, ref_moe
 jax.config.parse_flags_with_absl()
 
 
+def cdiv(a, b):
+    assert b != 0
+    return (a + b - 1) // b
+
+
+def align_to(x, a):
+    return cdiv(x, a) * a
+
+
 def gen_moe_inputs(
     dtype,
     top_k,
@@ -19,11 +28,14 @@ def gen_moe_inputs(
     num_tokens,
     *,
     seed=1234,
+    has_bias=False,
 ):
     key = jax.random.key(seed)
-    k0, k1, k2, k4, k5 = jax.random.split(key, 5)
+    k0, k1, k2, k3, k4, k5, k6 = jax.random.split(key, 7)
+
     a = jax.random.normal(k0, (num_tokens, hidden_size),
                           dtype=jnp.float32).astype(dtype) / 10
+
     w1 = (jax.random.normal(
         k1,
         (num_experts, 2, hidden_size, intermediate_size),
@@ -31,21 +43,54 @@ def gen_moe_inputs(
     ) / 10).astype(dtype)
     w2 = (jax.random.normal(k2, (num_experts, intermediate_size, hidden_size),
                             dtype=jnp.float32) / 10).astype(dtype)
+
+    if has_bias:
+        b1 = (jax.random.normal(k3, (num_experts, 2, intermediate_size),
+                                dtype=jnp.float32) / 10).astype(dtype)
+        b2 = (jax.random.normal(k4, (num_experts, hidden_size),
+                                dtype=jnp.float32) / 10).astype(dtype)
+    else:
+        b1 = b2 = None
+
     gating_output = (
-        jax.random.normal(k4, (num_tokens, num_experts), dtype=jnp.float32) +
+        jax.random.normal(k5, (num_tokens, num_experts), dtype=jnp.float32) +
         jnp.arange(num_tokens * num_experts, dtype=jnp.float32).reshape(
             num_tokens, num_experts) / 100)
+
     # To generate unique top-k!
-    top_k_indices = jax.random.randint(k5, (num_tokens, top_k),
+    top_k_indices = jax.random.randint(k6, (num_tokens, top_k),
                                        minval=0,
                                        maxval=num_experts - 1,
                                        dtype=jnp.int32)
+
     one_hot = (jnp.sum(
         jax.nn.one_hot(top_k_indices, num_experts, dtype=jnp.float32),
         axis=1,
-    ) * 10)
+    ) * 30)
+
     gating_output = (gating_output + one_hot).astype(dtype)
-    return a, w1, w2, gating_output
+
+    return a, w1, w2, b1, b2, gating_output
+
+
+def sub_channel_quantize(x, quant_dtype, wsz=256):
+    """Quantizes x with sub-channel quantization on the 2nd minor."""
+    if jnp.issubdtype(quant_dtype, jnp.floating):
+        dtype_info = jnp.finfo(quant_dtype)
+    else:
+        dtype_info = jnp.iinfo(quant_dtype)
+    dtype_max = float(dtype_info.max)
+    w_lst, scale_lst = [], []
+    assert len(x.shape) >= 2
+    assert x.shape[-2] % wsz == 0
+    for i in range(0, x.shape[-2], wsz):
+        y = x[..., i:i + wsz, :]
+        abs_max = jnp.abs(y).max(axis=-2, keepdims=True)
+        scale = (abs_max / dtype_max).astype(jnp.float32)
+        w = (y / scale).astype(quant_dtype)
+        w_lst.append(w)
+        scale_lst.append(scale)
+    return jnp.concat(w_lst, axis=-2), jnp.concat(scale_lst, axis=-2)
 
 
 @jtu.with_config(jax_numpy_dtype_promotion="standard")
@@ -63,42 +108,266 @@ class MoEKernelTest(jtu.JaxTestCase):
         self.mesh = Mesh(np.array(self.mesh_devices).reshape(1, -1),
                          axis_names=("data", "model"))
 
-    def test_basic(self):
-        dtype = jnp.bfloat16
-        top_k = 2
-        num_experts = 16
-        hidden_size = 256
-        intermediate_size = 256
-        num_tokens = 8 * 2
-
-        a, w1, w2, gating_output = gen_moe_inputs(
+    def _test_moe(
+        self,
+        dtype,
+        top_k,
+        num_experts,
+        hidden_size,
+        intermediate_size,
+        num_tokens,
+        seed,
+        renormalize_topk_logits,
+        bt,
+        bf,
+        bd1,
+        bd2,
+        btc,
+        bfc,
+        bd1c,
+        bd2c,
+        act_fn="silu",
+        w_dtype=None,
+        subc_quant_wsz=None,
+        has_bias=False,
+        atol=2e-1,
+        rtol=2e-1,
+    ):
+        a, w1, w2, b1, b2, gating_output = gen_moe_inputs(
             dtype,
             top_k,
             num_experts,
             hidden_size,
             intermediate_size,
             num_tokens,
+            seed=seed,
+            has_bias=has_bias,
         )
+        w1_scale = None
+        w2_scale = None
+        if w_dtype is not None:
+            if subc_quant_wsz is None:
+                subc_quant_wsz = 256
+            w1, w1_scale = sub_channel_quantize(w1, w_dtype, subc_quant_wsz)
+            w2, w2_scale = sub_channel_quantize(w2, w_dtype, subc_quant_wsz)
 
-        actual = jax.block_until_ready(
-            fused_ep_moe(
-                mesh=self.mesh,
-                tokens=a,
-                w1=w1,
-                w2=w2,
-                gating_output=gating_output,
-                top_k=top_k,
-                bt=32,
-                bf=512,
-                bd1=512,
-                bd2=512,
-                btc=32,
-                bfc=256,
-                bd1c=256,
-                bd2c=256,
-            ))
-        expected = ref_moe(a, w1, w2, gating_output, top_k)
-        self.assertAllClose(expected, actual, atol=2e-2, rtol=2e-2)
+        actual = fused_ep_moe(
+            mesh=self.mesh,
+            tokens=a,
+            w1=w1,
+            w2=w2,
+            gating_output=gating_output,
+            top_k=top_k,
+            renormalize_topk_logits=renormalize_topk_logits,
+            act_fn=act_fn,
+            subc_quant_wsz=subc_quant_wsz,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            b1=b1,
+            b2=b2,
+            bt=bt,
+            bf=bf,
+            bd1=bd1,
+            bd2=bd2,
+            btc=btc,
+            bfc=bfc,
+            bd1c=bd1c,
+            bd2c=bd2c,
+        )
+        expected = ref_moe(
+            a,
+            w1,
+            w2,
+            gating_output,
+            top_k,
+            b1=b1,
+            b2=b2,
+            renormalize_topk_logits=renormalize_topk_logits,
+            activation=act_fn,
+            subc_quant_wsz=subc_quant_wsz,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+        )
+        self.assertAllClose(actual, expected, atol=atol, rtol=rtol)
+
+    @parameterized.product(renormalize_topk_logits=[True, False], )
+    def test_basic(self, renormalize_topk_logits):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=32,
+            bf=1024,
+            bd1=1024,
+            bd2=1024,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    @parameterized.product(act_fn=["silu", "gelu", "swigluoai"], )
+    def test_activation(self, act_fn):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=True,
+            act_fn=act_fn,
+            bt=32,
+            bf=512,
+            bd1=512,
+            bd2=512,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    def test_benchmark_qwen_235(self):
+        num_experts = 128
+        top_k = 8
+        hidden_size = 4096
+        intermediate_size = 1536
+        dtype = jnp.bfloat16
+        num_tokens = 8 * 64
+        seed = 54321
+        renormalize_topk_logits = True
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=seed,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=64,
+            bf=768,
+            bd1=2048,
+            bd2=2048,
+            btc=64,
+            bfc=768,
+            bd1c=2048,
+            bd2c=2048,
+            act_fn="silu",
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    def test_benchmark_qwen_30b_a3b(self):
+        num_experts = 128
+        top_k = 8
+        hidden_size = 2048
+        intermediate_size = 768
+        dtype = jnp.bfloat16
+        num_tokens = 512
+        seed = 54321
+        renormalize_topk_logits = True
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=seed,
+            renormalize_topk_logits=renormalize_topk_logits,
+            bt=16,
+            bf=384,
+            bd1=512,
+            bd2=512,
+            btc=16,
+            bfc=384,
+            bd1c=256,
+            bd2c=256,
+            act_fn="silu",
+            atol=5e-2,
+            rtol=5e-2,
+        )
+
+    @parameterized.product(
+        w_dtype=[jnp.int8, jnp.float8_e5m2, jnp.float4_e2m1fn], )
+    def test_sub_channel_quantization(self, w_dtype):
+        if w_dtype in (
+                jnp.float8_e5m2,
+                jnp.float4_e2m1fn,
+        ) and not jtu.is_device_tpu_at_least(version=7):
+            self.skipTest("Expect TPUv7+")
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=False,
+            w_dtype=w_dtype,
+            subc_quant_wsz=256,
+            bt=32,
+            bf=1024,
+            bd1=1024,
+            bd2=1024,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
+
+    def test_bias(self):
+        dtype = jnp.bfloat16
+        top_k = 8
+        num_experts = 128
+        hidden_size = 1024
+        intermediate_size = 1024
+        num_tokens = 8 * 32
+        self._test_moe(
+            dtype=dtype,
+            top_k=top_k,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_tokens=num_tokens,
+            seed=1234,
+            renormalize_topk_logits=False,
+            has_bias=True,
+            bt=32,
+            bf=512,
+            bd1=512,
+            bd2=512,
+            btc=32,
+            bfc=256,
+            bd1c=256,
+            bd2c=256,
+        )
 
 
 if __name__ == "__main__":

tests/lora/test_layers.py

@@ -91,7 +91,6 @@ def populate_loras(
     index_to_id: list[Optional[int]],
     lora_layer: BaseLayerWithLoRA,
     baselayer_weights: torch.Tensor,
-    generate_embeddings_tensor: int = 0,
     repeats: int = 1,
 ) -> tuple[dict[int, LoRALayerWeights], dict[int, list[LoRALayerWeights]]]:
     """This method populates the lora weights (lora_a and lora_b) in the lora layers (BaseLayerWithLoRA).
@@ -103,8 +102,6 @@ def populate_loras(
         lora_layer: the LoRAlayer to populate.
         baselayer_weights: the PyTorch tensor containing the layer's
             weights.
-        generate_embeddings_tensor: whether to generate an
-            embeddings tensor for each LoRA.
         repeats: must only be set for column parallel packed
             layers. Indicates the number of loras to compose
             together to create a single lora layer.
@@ -131,7 +128,6 @@ def populate_loras(
                     baselayer_weights.device).init_random_lora(
                         module_name=f"fake_{i}",
                         weight=baselayer_weights,
-                        generate_embeddings_tensor=generate_embeddings_tensor,
                     )
                 sublora.lora_b = sublora.lora_b[(sublora_len *
                                                  i):(sublora_len * (i + 1)), :]
@@ -147,7 +143,6 @@ def populate_loras(
                     slot_idx,
                     lora_a=lora.lora_a,
                     lora_b=lora.lora_b,
-                    embeddings_tensor=lora.embeddings_tensor,
                 )
 
             lora_dict[lora_id] = lora
@@ -546,7 +541,6 @@ def _update_punica_wrapper_metadata(punica_wrapper, index_mapping,
             index_to_id,
             lora_config.max_loras,
             vocab_size=512,
-            extra_vocab_size=lora_config.lora_extra_vocab_size,
         )
         assert jax_view(punica_wrapper._lora_indices_per_batch).platform(
         ) == 'tpu', 'punica_wrapper._lora_indices_per_batch should have been moved to TPU.'

tests/lora/utils.py

@@ -24,7 +24,6 @@ class DummyLoRAManager:
         module_name: str,
         weight: torch.Tensor,
         rank: int = 8,
-        generate_embeddings_tensor: int = 0,
     ):
         lora = LoRALayerWeights(
             module_name,
@@ -37,13 +36,6 @@ class DummyLoRAManager:
                               dtype=weight.dtype,
                               device=self._device),
         )
-        if generate_embeddings_tensor:
-            lora.embeddings_tensor = torch.rand(
-                5,
-                generate_embeddings_tensor,
-                dtype=weight.dtype,
-                device=self._device,
-            )
         self.set_module_lora(module_name, lora)
 
         return lora

tpu_inference/__init__.py

@@ -1,21 +1,40 @@
-import os
-
 # The environment variables override should be imported before any other
 # modules to ensure that the environment variables are set before any
 # other modules are imported.
 import tpu_inference.env_override  # noqa: F401
+from tpu_inference import envs
 from tpu_inference import tpu_info as ti
 from tpu_inference.logger import init_logger
 
 logger = init_logger(__name__)
 
-if "proxy" in os.environ.get('JAX_PLATFORMS', '').lower():
+if "proxy" in envs.JAX_PLATFORMS:
     logger.info("Running vLLM on TPU via Pathways proxy.")
     # Must run pathwaysutils.initialize() before any JAX operations
     try:
+        import traceback
+
         import pathwaysutils
+        import vllm
+        from vllm.platforms import (resolve_current_platform_cls_qualname,
+                                    resolve_obj_by_qualname)
         pathwaysutils.initialize()
         logger.info("Module pathwaysutils is imported.")
+
+        # Pathways requires eager resolution of vllm.current_platform instead of
+        # lazy resolution in the normal code path. Since this part involves
+        # global topology discovery across multiple hosts, the platform
+        # resolution must happen before other components are loaded.
+        logger.info("Eagerly resolving vLLM current_platform for Pathways.")
+        platform_cls_qualname = resolve_current_platform_cls_qualname()
+        resolved_platform_instance = resolve_obj_by_qualname(
+            platform_cls_qualname)()
+        vllm.platforms._current_platform = resolved_platform_instance
+        vllm.platforms._init_trace = "".join(traceback.format_stack())
+        logger.info(
+            f"vLLM platform resolved to: {resolved_platform_instance.__class__.__name__}"
+        )
+
     except Exception as e:
         logger.error(
             f"Error occurred while importing pathwaysutils or logging TPU info: {e}"

tpu_inference/core/disagg_utils.py

@@ -1,17 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 
-import os
 from typing import Tuple
 
-PREFILL_SLICES = 'PREFILL_SLICES'
-DECODE_SLICES = 'DECODE_SLICES'
+from tpu_inference import envs
 
 
 def is_disagg_enabled() -> bool:
     # We triggrer our code path as long as prefill slices are set. This
     # allows us to test interleave mode effectively with the code path
     # for comparison purposes.
-    return PREFILL_SLICES in os.environ
+    return bool(envs.PREFILL_SLICES)
 
 
 def _parse_slices(slices_str: str) -> Tuple[int, ...]:
@@ -40,12 +38,12 @@ def _parse_slices(slices_str: str) -> Tuple[int, ...]:
 
 
 def get_prefill_slices() -> Tuple[int, ...]:
-    if PREFILL_SLICES not in os.environ:
+    if not envs.PREFILL_SLICES:
         return ()
-    return _parse_slices(os.environ[PREFILL_SLICES])
+    return _parse_slices(envs.PREFILL_SLICES)
 
 
 def get_decode_slices() -> Tuple[int, ...]:
-    if DECODE_SLICES not in os.environ:
+    if not envs.DECODE_SLICES:
         return ()
-    return _parse_slices(os.environ[DECODE_SLICES])
+    return _parse_slices(envs.DECODE_SLICES)

tpu_inference/distributed/tpu_connector.py

@@ -60,7 +60,6 @@ D workflow:
 
 import copy
 import functools
-import os
 import threading
 import time
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -86,6 +85,7 @@ if TYPE_CHECKING:
     from vllm.v1.core.kv_cache_manager import KVCacheBlocks
     from vllm.v1.request import Request
 
+from tpu_inference import envs
 from tpu_inference.distributed.utils import (get_host_ip, get_kv_ips,
                                              get_kv_ports,
                                              get_kv_transfer_port, get_node_id,
@@ -441,8 +441,7 @@ class TPUConnectorWorker:
 
         self.runner: TPUModelRunner = None
         self.mesh: Mesh = None
-        self.multi_host = os.getenv("TPU_MULTIHOST_BACKEND",
-                                    "").lower() == "ray"
+        self.multi_host = envs.TPU_MULTIHOST_BACKEND == "ray"
         # NOTE(xiang): This can not be the worker rank set in RayDistributedExecutor.
         # The worker rank is assigned with vLLM's sorting logic, which does not work
         # for TPU host topology.

tpu_inference/distributed/utils.py

@@ -2,6 +2,7 @@ import os
 
 from vllm.utils.network_utils import get_ip
 
+from tpu_inference import envs
 from tpu_inference.logger import init_logger
 
 logger = init_logger(__name__)
@@ -17,7 +18,7 @@ def set_node_kv_ip_port(ip_port: tuple[int, str, int]):
 
 
 def get_kv_ips() -> str:
-    if os.getenv("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
+    if envs.TPU_MULTIHOST_BACKEND == "ray":
         num_nodes = len(_NODES_KV_IP_PORT)
         ips = []
         for node_id in range(num_nodes):
@@ -28,7 +29,7 @@ def get_kv_ips() -> str:
 
 
 def get_kv_ports() -> str:
-    if os.getenv("TPU_MULTIHOST_BACKEND", "").lower() == "ray":
+    if envs.TPU_MULTIHOST_BACKEND == "ray":
         num_nodes = len(_NODES_KV_IP_PORT)
         ports = []
         for node_id in range(num_nodes):

tpu_inference/envs.py

@@ -26,7 +26,7 @@ if TYPE_CHECKING:
 environment_variables: dict[str, Callable[[], Any]] = {
     # JAX platform selection (e.g., "tpu", "cpu", "proxy")
     "JAX_PLATFORMS":
-    lambda: os.getenv("JAX_PLATFORMS", ""),
+    lambda: os.getenv("JAX_PLATFORMS", "").lower(),
     # TPU accelerator type (e.g., "v5litepod-16", "v4-8")
     "TPU_ACCELERATOR_TYPE":
     lambda: os.getenv("TPU_ACCELERATOR_TYPE", None),

tpu_inference/executors/ray_distributed_executor.py

@@ -108,6 +108,9 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
             ip_port = self.collective_rpc("get_node_kv_ip_port")
             for item in ip_port:
                 set_node_kv_ip_port(item)
+        self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
+            self.vllm_config.ec_transfer_config is None
+            or not self.vllm_config.ec_transfer_config.is_ec_producer)
 
     def _initialize_ray_cluster(self) -> None:
         """Initialize the distributed cluster with Ray.
@@ -131,10 +134,17 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
                 f"current platform {current_platform.device_name} does not "
                 "support ray.")
 
-        placement_group_specs: List[Dict[str, float]] = [{
-            device_str:
-            node['Resources'][device_str]
-        } for node in ray.nodes()]
+        pp_size = self.parallel_config.pipeline_parallel_size
+        placement_group_specs: List[Dict[str, float]] = []
+        if pp_size == 1:
+            placement_group_specs = [{
+                device_str: node['Resources'][device_str]
+            } for node in ray.nodes()]
+        else:
+            num_devices_per_pp_rank = self.vllm_config.sharding_config.total_devices
+            placement_group_specs = [{
+                device_str: num_devices_per_pp_rank
+            } for _ in range(pp_size)]
 
         # vLLM engine is also a worker to execute model with an accelerator,
         # so it requires to have the device in a current node. Check if
@@ -329,6 +339,8 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
         all_kwargs = []
         for rank, (node_id, _) in enumerate(worker_node_and_tpu_ids):
             local_rank = node_workers[node_id].index(rank)
+            ip = sorted_worker_metadata[rank].ip
+            prev_ip = sorted_worker_metadata[rank - 1].ip if rank > 0 else ""
             kwargs = dict(
                 vllm_config=self.vllm_config,
                 local_rank=local_rank,
@@ -336,22 +348,26 @@ class RayDistributedExecutor(RayDistributedExecutorV1):
                 distributed_init_method=distributed_init_method,
                 is_driver_worker=(not self.parallel_config)
                 or (rank % self.parallel_config.tensor_parallel_size == 0),
+                ip=ip,
+                prev_worker_ip=prev_ip,
             )
             all_kwargs.append(kwargs)
         self.collective_rpc("init_worker", args=(all_kwargs, ))
         self.collective_rpc("init_device")
+        if self.parallel_config.pipeline_parallel_size > 1:
+            self.collective_rpc("initialize_pp_transfer_connect")
         self.collective_rpc("load_model")
 
         if self.use_ray_spmd_worker:
             for pp_rank in range(self.parallel_config.pipeline_parallel_size):
                 self.pp_tp_workers.append([])
-                for tp_rank in range(
-                        int(self.parallel_config.tensor_parallel_size //
-                            num_tpu_per_worker)):
-                    # PP=2, TP=4
-                    # pp_tp_workers = [[0, 1, 2, 3], [4, 5, 6, 7]]
-                    rank = (pp_rank * self.parallel_config.tensor_parallel_size
-                            ) + tp_rank
+                num_tp_workers = int(
+                    self.parallel_config.tensor_parallel_size //
+                    num_tpu_per_worker)
+                for tp_rank in range(num_tp_workers):
+                    # PP=2, TP=4, num_tpu_per_worker=2
+                    # pp_tp_workers = [[0, 1], [2, 3]]
+                    rank = (pp_rank * num_tp_workers) + tp_rank
                     assert len(self.pp_tp_workers[pp_rank]) == tp_rank
                     assert pp_rank < len(self.pp_tp_workers)
                     self.pp_tp_workers[pp_rank].append(self.workers[rank])