[BugFix] Fix incorrect GPU device mapping in multi-replica stages (#4132)

Signed-off-by: ZhengWG <zwg0606@gmail.com>

tests/entrypoints/test_stage_utils.py

@@ -135,22 +135,6 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte
 # ---- _map_device_list unit tests ----
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency():
-    """Device IDs already in the visible set are returned as-is (idempotency)."""
-    result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"])
-    assert result == ["0", "1"]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency_single():
-    """Single device ID in visible set passes through."""
-    result = _map_device_list(1, ["1"], ["0", "1"])
-    assert result == ["1"]
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_index_mapping():
@@ -211,6 +195,18 @@ def test_map_device_list_partial_mapping():
     assert result == ["5", "6"]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_multi_replica_offset_cvd():
+    """Logical indices from split_devices_for_replicas must map through an offset
+    CUDA_VISIBLE_DEVICES, not be treated as physical IDs.
+    """
+    visible = ["2", "3", "4", "5"]
+    assert _map_device_list(1, ["1"], visible) == ["3"]
+    assert _map_device_list(1, ["2"], visible) == ["4"]
+    assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():