fix

llumnix/backends/vllm/migration_backend.py

@@ -15,8 +15,12 @@
 import torch
 from func_timeout import func_set_timeout, FunctionTimedOut
 
+import cupy
+from cupy.cuda import nccl
 import ray
 import ray.util.collective as col
+from ray.util.collective.collective_group import nccl_util
+
 from vllm.worker.cache_engine import CacheEngine
 from llumnix.internal_config import MigrationConfig
 from llumnix.backends.migration_backend_interface import MigrationBackendBase, CacheMigrationBackend
@@ -143,9 +147,6 @@ def try_import_gloo():
 class RayColMigrationBackend(CacheMigrationBackend):
     def __init__(self, migration_config: MigrationConfig, cache_engine: CacheEngine, local_rank,
                  scheduling_strategy, is_driver_worker, gpu_cache) -> None:
-        # pylint: disable=C0415
-        import cupy
-
         self.migration_config = migration_config
         self.cache_engine = cache_engine
         self.backend = migration_config.migration_backend
@@ -169,6 +170,7 @@ def __init__(self, migration_config: MigrationConfig, cache_engine: CacheEngine,
             try_import_gloo()
             self.cache_device = "cpu"
         else:
+            nccl_util.TORCH_NCCL_DTYPE_MAP[torch.bfloat16] = nccl.NCCL_FLOAT16
             self.cache_device = torch.device(f"cuda:{self.local_rank}")
 
         pin_memory = (self.backend == 'gloo')

tests/e2e_test/test_e2e.py

@@ -46,7 +46,7 @@ def generate_launch_command(result_filename: str = "", launch_ray_cluster: bool
     disable_init_instance_by_manager, disable_fixed_node_init_instance = parse_launch_mode(launch_mode)
     command = (
         f"RAY_DEDUP_LOGS=0 HEAD_NODE_IP={HEAD_NODE_IP} HEAD_NODE=1 "
-        f"nohup python -m llumnix.entrypoints.vllm.api_server "
+        f"nohup python -u -m llumnix.entrypoints.vllm.api_server "
         f"--host {ip} "
         f"--port {port} "
         f"{'--disable-init-instance-by-manager ' if disable_init_instance_by_manager else ''}"