Fix the jax device ordering. (#915)

Signed-off-by: Lance Wang <[email protected]>

Author: wang2yn84

tpu_inference/runner/tpu_jax_runner.py

@@ -154,6 +154,15 @@ def _init_mesh(self) -> None:
         except KeyError:
             sharding_strategy = {"tensor_parallelism": len(self.devices)}
 
+        try:
+            enforce_device_order = self.vllm_config.additional_config[
+                "sharding"]["sharding_strategy"]["device_indexes"] is not None
+
+        except KeyError:
+            enforce_device_order = False
+
+        logger.info(f"Device sequence enforced: {enforce_device_order}")
+
         if os.getenv("NEW_MODEL_DESIGN", False):
             self.mesh = build_mesh(self.devices, sharding_strategy)
         else:
@@ -169,9 +178,14 @@ def _init_mesh(self) -> None:
             axis_names = ("data", "model")
             mesh_shape = (dp, tp)
 
-            self.mesh = make_optimized_mesh(mesh_shape,
-                                            axis_names,
-                                            devices=self.devices)
+            if enforce_device_order:
+                self.mesh = jax.make_mesh(mesh_shape,
+                                          axis_names,
+                                          devices=self.devices)
+            else:
+                self.mesh = make_optimized_mesh(mesh_shape,
+                                                axis_names,
+                                                devices=self.devices)
         logger.info(f"Init mesh | mesh={self.mesh}")
 
     def _init_phased_profiling(self) -> None:

tpu_inference/worker/tpu_worker_jax.py

@@ -121,14 +121,29 @@ def initialize_cache(self, num_gpu_blocks: int,
 
     def init_device(self):
         if not self.devices:
+            device_indexes = []
+            tp = self.parallel_config.tensor_parallel_size
             try:
                 device_indexes = self.vllm_config.additional_config[
                     "sharding"]["sharding_strategy"]["device_indexes"]
-                self.devices = [jax.devices()[i] for i in device_indexes]
             except KeyError:
-                tp = self.parallel_config.tensor_parallel_size
                 self.devices = jax.devices()[:tp]
 
+            # Enforcing the devices sequence to be consistent with the specified device indexes
+            if not self.devices:
+                all_devices = jax.devices()
+                device_dict = {device.id: device for device in all_devices}
+                self.devices = []
+                for device_index in device_indexes:
+                    device = device_dict[device_index]
+                    if device is None:
+                        raise KeyError(
+                            f"Device index {device_index} not found in "
+                            f"jax.devices() with IDs {list(device_dict.keys())}!"
+                        )
+                    self.devices.append(device)
+                self.devices = self.devices[:tp]
+
         # Initialize the vLLM distribution layer as a single chip environment,
         # we'll swap the model's parallel modules with TPU SPMD equivalents.
         with set_current_vllm_config(self.vllm_config):