feat:add device indexes (#470)

Co-authored-by: pathfinder-fp <[email protected]>

Author: pathfinder-pf

.github/workflows/pr-test.yml

@@ -90,19 +90,25 @@ jobs:
   unit-test-4-tpu:
     needs: [check-changes]
     if: github.event.pull_request.draft == false && needs.check-changes.outputs.main_package == 'true'
-    #runs-on: arc-runner-v6e-1
-    runs-on: ubuntu-slim
+    runs-on: arc-runner-v6e-4
     strategy:
       fail-fast: false
       matrix:
         part: [0, 1]
     steps:
-      # TODO: please rename when implement unit test
-      - name: Mock
+      - name: Checkout code
+        uses: actions/checkout@v5
+      - name: Run unit test
         timeout-minutes: 30
         env:
           SGLANG_JAX_IS_IN_CI: true
-        run: echo ${{ matrix.part }}
+        run: |
+          python3.12 -m venv .venv
+          source .venv/bin/activate
+          pip install uv
+          uv pip install -e "python[all]"
+          bash scripts/killall_sglang.sh
+          python test/srt/run_suite.py --suite unit-test-tpu-v6e-4 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2
 
   # =============================================== e2e-test ===============================================
   e2e-test-1-tpu:

python/sgl_jax/srt/managers/scheduler.py

@@ -218,7 +218,11 @@ def __init__(
         # init distribution
         if self.nnodes > 1:
             jax.distributed.initialize(server_args.dist_init_addr, self.nnodes, self.node_rank)
-        self.mesh = create_device_mesh(ici_parallelism=[-1, self.tp_size], dcn_parallelism=[1, 1])
+        self.mesh = create_device_mesh(
+            ici_parallelism=[-1, self.tp_size],
+            dcn_parallelism=[1, 1],
+            device_indexes=server_args.device_indexes,
+        )
 
         TpWorkerClass = ModelWorkerClient if self.enable_overlap else ModelWorker
 

python/sgl_jax/srt/server_args.py

@@ -66,6 +66,7 @@ class ServerArgs:
 
     # Runtime options
     device: str | None = None
+    device_indexes: list[int] | None = None
     tp_size: int = 1
     ep_size: int = 1
     stream_interval: int = 1
@@ -216,6 +217,10 @@ def __post_init__(self):
             "1" if self.enable_deterministic_sampling else "0"
         )
 
+        if self.nnodes > 1 and self.device_indexes is not None:
+            logger.warning("In a multi-machine scenario, device_indexes will be set to None.")
+            self.device_indexes = None
+
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
         # Model and tokenizer
@@ -499,6 +504,14 @@ def add_cli_args(parser: argparse.ArgumentParser):
             default=ServerArgs.device,
             help="The device to use ('cuda', 'xpu', 'hpu', 'npu', 'cpu'). Defaults to auto-detection if not specified.",
         )
+
+        parser.add_argument(
+            "--device-indexes",
+            type=int,
+            nargs="+",
+            help="The device indexes to use build mesh. Defaults is all if not specified.",
+        )
+
         parser.add_argument(
             "--tensor-parallel-size",
             "--tp-size",

python/sgl_jax/srt/utils/mesh_utils.py

@@ -9,6 +9,7 @@ def create_device_mesh(
     ici_parallelism: Sequence[int],
     dcn_parallelism: Sequence[int],
     devices=None,
+    device_indexes: list[int] = None,
     num_slices: int = 1,
     allow_split_physical_axes: bool = True,
     use_explicit_sharding: bool = True,
@@ -17,6 +18,13 @@ def create_device_mesh(
     if devices is None:
         devices = jax.devices()
 
+    if device_indexes is not None:
+        max_index = max(device_indexes)
+        if max_index >= len(devices):
+            raise RuntimeError("Device index out of range")
+        devices_dict = {device.id: device for device in devices}
+        devices = [devices_dict.get(i) for i in list(set(device_indexes))]
+
     ici_parallelism = fill_unspecified_parallelism(ici_parallelism, len(devices))
     if num_slices > 1:
         dcn_parallelism = fill_unspecified_parallelism(dcn_parallelism, num_slices)

python/sgl_jax/test/test_mesh.py

@@ -0,0 +1,39 @@
+import unittest
+
+from sgl_jax.srt.utils.mesh_utils import create_device_mesh
+from sgl_jax.test.test_utils import CustomTestCase
+
+
+class TestMesh(CustomTestCase):
+    def test_mesh_with_no_device_indexes(self):
+        mesh = create_device_mesh(ici_parallelism=[1, -1], dcn_parallelism=[1, 1])
+        self.assertEqual(mesh.shape.get("data"), 1, "dp should be 1")
+        self.assertEqual(mesh.shape.get("tensor"), 4, "tp should be 1")
+
+    def test_mesh_with_device_indexes(self):
+        mesh = create_device_mesh(
+            ici_parallelism=[1, -1], dcn_parallelism=[1, 1], device_indexes=[0, 1]
+        )
+        self.assertEqual(mesh.shape.get("data"), 1, "dp should be 1")
+        self.assertEqual(mesh.shape.get("tensor"), 2, "tp should be 1")
+
+    def test_mesh_with_duplicated_device_indexes(self):
+        mesh = create_device_mesh(
+            ici_parallelism=[1, -1], dcn_parallelism=[1, 1], device_indexes=[0, 1, 0, 1]
+        )
+        self.assertEqual(mesh.shape.get("data"), 1, "dp should be 1")
+        self.assertEqual(mesh.shape.get("tensor"), 2, "tp should be 1")
+
+    def test_mesh_with_large_device_indexes(self):
+        try:
+            _ = create_device_mesh(
+                ici_parallelism=[1, -1], dcn_parallelism=[1, 1], device_indexes=[0, 4]
+            )
+        except Exception as e:
+            self.assertTrue(
+                isinstance(e, RuntimeError), "the device indexes have exceeded the len of devices"
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

python/sgl_jax/test/test_utils.py

@@ -60,6 +60,7 @@ def create_device_mesh(
     ici_parallelism: Sequence[int],
     dcn_parallelism: Sequence[int],
     devices=None,
+    device_indexes: list[int] = None,
     num_slices: int = 1,
     allow_split_physical_axes: bool = True,
     use_explicit_sharding: bool = True,
@@ -68,6 +69,13 @@ def create_device_mesh(
     if devices is None:
         devices = jax.devices()
 
+    if device_indexes is not None:
+        max_index = max(device_indexes)
+        if max_index >= len(devices):
+            raise RuntimeError("Device index out of range")
+        devices_dict = {device.id: device for device in devices}
+        devices = [devices_dict.get(i) for i in list(set(device_indexes))]
+
     ici_parallelism = fill_unspecified_parallelism(ici_parallelism, len(devices))
     if num_slices > 1:
         dcn_parallelism = fill_unspecified_parallelism(dcn_parallelism, num_slices)

test/srt/run_suite.py

@@ -168,6 +168,9 @@ def run_one_file(filename):
         TestFile("python/sgl_jax/test/speculative/test_eagle_tree_build.py", 1),
         TestFile("python/sgl_jax/test/speculative/test_eagle_utils.py", 1),
     ],
+    "unit-test-tpu-v6e-4": [
+        TestFile("python/sgl_jax/test/test_mesh.py", 1),
+    ],
     "kernel-performance-test-tpu-v6e-1": [
         TestFile("benchmark/kernels/flash_attention/bench_flashattention.py", 5),
         TestFile("benchmark/kernels/megablox_gmm/bench_megablox_gmm.py", 2),