Update on "v0 param server (using collectives not object store)"

[ghstack-poisoned]

Author: mikaylagawarecki

param_server_weight_updater.py

@@ -218,16 +218,18 @@ def _create_trainer_group(
 
     model = "facebook/opt-125m"
 
-    ray.init(num_cpus=4, num_gpus=4)
+    ray.init(num_cpus=5, num_gpus=5)
 
-    vllm_master_address, vllm_update_port = get_ip(), get_open_port()
+    vllm_addresses = [get_ip()] * 2
+    vllm_ports = [get_open_port() for i in range(2)]
+    print(vllm_ports)
 
     trainer_workers, parameter_server = _create_trainer_group(
                                             TrainerActor,
                                             vLLMParameterServer,
                                             3,
-                                            vllm_master_address,
-                                            vllm_update_port,
+                                            vllm_addresses,
+                                            vllm_ports,
                                             model,
                                         )
 
@@ -236,19 +238,28 @@ def _create_trainer_group(
         handles.append(trainer_worker.train.remote())
 
     model_metadata = ray.get(parameter_server.get_model_metadata.remote())
-    local_weight_updater = vLLMHFLocalWeightUpdater(vllm_master_address, vllm_update_port, model_metadata)
+    local_weight_updaters = [
+        vLLMHFLocalWeightUpdater(vllm_master_address, vllm_update_port, model_metadata) for
+        vllm_master_address, vllm_update_port in zip(vllm_addresses, vllm_ports)
+    ]
 
     make_env_parsed = partial(make_env, batch_size=args.batch_size, dataset=args.dataset)
     collector = RayCollector(
-        [make_env_parsed],
+        [make_env_parsed, make_env_parsed],
         policy_factory=make_policy,
         frames_per_batch=40,
         total_frames=200,
         remote_configs=remote_configs,
         remote_weight_updater=parameter_server,
-        collector_kwargs={
-            "local_weight_updater": local_weight_updater,
-        },
+        num_collectors=2,
+        collector_kwargs=[
+            {
+                "local_weight_updater": local_weight_updaters[0],
+            },
+            {
+                "local_weight_updater": local_weight_updaters[1],
+            }
+        ],
         update_after_each_batch=True,
     )
     print("done collector init")
@@ -258,6 +269,6 @@ def _create_trainer_group(
     for i, data in enumerate(collector):
         print(tokenizer.decode(data["tokens"][0].squeeze()))
         print(tokenizer.decode(data["tokens_response"][0].squeeze()))
-        if i == 1:
+        if i == 3:
             break
     collector.shutdown()

torchrl/collectors/vllm_weight_update.py

@@ -56,7 +56,8 @@ class WorkerExtension(Worker):
         def init_weight_update_group(self, master_address, master_port,
                                     rank_offset, world_size):
             from vllm.distributed.parallel_state import get_world_group
-            rank = get_world_group().rank + rank_offset
+            # rank = get_world_group().rank + rank_offset
+            rank = rank_offset
             self.model_update_group = stateless_init_process_group(
                 master_address,
                 master_port,
@@ -91,10 +92,11 @@ class WorkerExtension:
 
 class vLLMHFLocalWeightUpdater(LocalWeightUpdaterBase):
     def __init__(self, master_address, master_port, model_metadata):
+        print(f"{master_address=}, {master_port=}")
         self.master_address = master_address
         self.master_port = master_port
         self.model_metadata = model_metadata
-        self.model_update_group = None
+        self.initialized_group = None
 
     def _get_server_weights(self):
         return None
@@ -110,13 +112,13 @@ def _maybe_map_weights(self, server_weights, local_weights):
 
     def _update_local_weights(self, local_weights, mapped_weights):
         llm = self.collector.policy["generate"].module
-        if self.model_update_group is None:
-            # FIXME: hardcoded
+        if self.initialized_group is None:
             weight_sync_world_size = llm.llm_engine.parallel_config.tensor_parallel_size + 1
             llm.collective_rpc(
                 "init_weight_update_group",
                 args=(self.master_address, self.master_port, 1, weight_sync_world_size)
             )
+            self.initialized_group = True
 
         for k, (dtype, shape) in self.model_metadata.items():
             llm.collective_rpc(
@@ -125,11 +127,11 @@ def _update_local_weights(self, local_weights, mapped_weights):
             )
 
 class vLLMRemoteWeightUpdaterBase(RemoteWeightUpdaterBase):
-    def __init__(self, model, vllm_master_address, vllm_master_port):
+    def __init__(self, model, vllm_master_addresses, vllm_master_ports):
         super().__init__()
         from transformers import AutoModel
-        self.vllm_master_address = vllm_master_address
-        self.vllm_master_port = vllm_master_port
+        self.vllm_master_addresses = vllm_master_addresses
+        self.vllm_master_ports = vllm_master_ports
         self.state_dict = AutoModel.from_pretrained(model).cuda().eval().state_dict()
         self.state_dict_lock = threading.Lock()
         self.vllm_comm_groups = dict()
@@ -160,8 +162,8 @@ def _init_model_update_group(self, worker_id):
         vllm_tp_size = 1
         weight_sync_world_size = vllm_tp_size + 1
         model_update_group = stateless_init_process_group(
-            self.vllm_master_address,
-            self.vllm_master_port,
+            self.vllm_master_addresses[worker_id],
+            self.vllm_master_ports[worker_id],
             0,
             weight_sync_world_size,
             torch.device("cuda:0"),