v0 param server (using collectives not object store)

ghstack-source-id: 74de8e0ef2fe059390e009332daeb688656a11fa
Pull Request resolved: #2865

Author: mikaylagawarecki

param_server_weight_updater.py

@@ -0,0 +1,263 @@
+import ray
+
+from argparse import ArgumentParser
+from functools import partial
+
+import torch
+from datasets import load_dataset
+from tensordict import TensorDict
+from torch.utils.data import DataLoader
+from torchrl.collectors.weight_update import RayRemoteWeightUpdater
+from transformers import AutoTokenizer, AutoModel
+from vllm import LLM
+
+from vllm.utils import get_ip, get_open_port
+
+from torchrl.collectors.distributed import RayCollector
+from torchrl.envs import LLMEnv
+from torchrl.modules import from_vllm
+
+from torchrl.collectors.vllm_weight_update import vLLMHFLocalWeightUpdater, vLLMRemoteWeightUpdaterBase, WorkerExtension
+
+parser = ArgumentParser()
+parser.add_argument("--dataset", type=str, default="gsm8k")
+parser.add_argument("--batch_size", type=int, default=4)
+parser.add_argument("--epochs", type=int, default=10)
+parser.add_argument("--repeats", type=int, default=10)
+parser.add_argument("--steps_per_batch", type=int, default=16)
+parser.add_argument("--optim_batch_size", type=int, default=4)
+
+
+def make_policy():
+    inference_model = LLM(
+        "facebook/opt-125m",
+        enforce_eager=True,
+        # change to worker_extension_cls when available in stable release
+        worker_cls=WorkerExtension,
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+
+    policy = from_vllm(
+        inference_model, tokenizer=tokenizer, from_text=False, generate=True, return_log_probs=True, generate_kwargs={"temperature": 0.0})
+    return policy
+
+
+def make_env(dataset, batch_size):
+    dataset = load_dataset(dataset, "main")
+    train_dataset = dataset["train"]
+    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.padding_side = "left"
+
+    # Env
+    dataloader = DataLoader(  # noqa: TOR401
+        train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn
+    )
+    env = LLMEnv.from_dataloader(
+        dataloader=dataloader,
+        tokenizer=tokenizer,
+        str2str=True,
+        batch_size=(args.batch_size * args.repeats,),
+        repeats=args.repeats, )
+    return env
+
+
+def collate_fn(batch):
+    batch = torch.stack([TensorDict.from_dict(_batch) for _batch in batch])
+    batch.rename_key_("question", "text")
+    return batch
+
+@ray.remote(num_cpus=1, num_gpus=1)
+class TrainerActor:
+    def __init__(self, model, env_vars):
+        import os
+        import torch
+        import torch.distributed
+        from torch.distributed._composable.fsdp import fully_shard
+
+        torch.cuda.set_device(torch.device('cuda', 0))
+
+        for var in env_vars:
+            os.environ[var] = str(env_vars[var])
+
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(backend="nccl", device_id=torch.device('cuda:0'))
+            print("initialized process group")
+
+        world_size = torch.distributed.get_world_size()
+        rank = torch.distributed.get_rank()
+        print(world_size, rank)
+        self.rank = int(os.environ["RANK"])
+        self.world_size = int(os.environ["WORLD_SIZE"])
+
+
+        # hold back one rank for the parameter server
+        self.fsdp_group = torch.distributed.new_group(ranks=list(range(self.world_size - 1)))
+        self.device_mesh = torch.distributed.device_mesh.DeviceMesh.from_group(self.fsdp_group, device_type="cuda") 
+
+        self.model = AutoModel.from_pretrained(model).cuda()
+
+        fully_shard(self.model, mesh=self.device_mesh)
+    
+    def register_parameter_server(self, param_server):
+        assert self.rank == 0
+        self.param_server = param_server
+    
+    def send_weights_to_param_server(self):
+        if self.rank == 0:
+            ray.get(self.param_server.acquire_state_dict_lock.remote())
+            self.param_server.receive_from_trainer.remote()
+        for k, v in self.model.state_dict().items():
+            replicated_v = v.full_tensor()
+            if self.rank == 0:
+                # dst is global rank, can switch to group_dst arg if not 2.5.1
+                torch.distributed.send(replicated_v, dst=2)
+        if self.rank == 0:
+            ray.get(self.param_server.release_state_dict_lock.remote())
+    
+    def zero_(self):
+        sd = self.model.state_dict()
+        for k, v in sd.items():
+            sd[k] = v.data.zero_()
+    
+    def train(self):
+        import time
+        for _ in range(1):
+            # actually run train loop
+            # ...
+            self.zero_()
+            torch.distributed.barrier(group=self.fsdp_group)
+            self.send_weights_to_param_server()
+            torch.distributed.barrier(group=self.fsdp_group)
+
+
+@ray.remote(num_cpus=1, num_gpus=1)
+class vLLMParameterServer(vLLMRemoteWeightUpdaterBase):
+    def __init__(self, model, vllm_master_address, vllm_master_port, env_vars):
+        super().__init__(model, vllm_master_address, vllm_master_port)
+        import os
+        import torch
+        import torch.distributed
+
+        torch.cuda.set_device(torch.device('cuda', 0))
+
+        for var in env_vars:
+            os.environ[var] = str(env_vars[var])
+
+        if not torch.distributed.is_initialized():
+            torch.distributed.init_process_group(backend="nccl", device_id=torch.device('cuda:0'))
+
+        self.rank = int(os.environ["RANK"])
+        self.world_size = int(os.environ["WORLD_SIZE"])
+        assert self.rank == self.world_size - 1
+
+        self.fsdp_group = torch.distributed.new_group(ranks=list(range(self.world_size - 1)))
+    
+    def receive_from_trainer(self):
+        for k, v in self.state_dict.items():
+            torch.distributed.recv(v, src=0)
+
+    def _skip_update(self, worker_id: int) -> bool:
+        pass
+    
+    def check_weights_changed(self):
+        """
+        Check if the weights are updated to 0.
+        """
+        weights_updated = True
+        for name, p in self.state_dict.items():
+            weights_updated = weights_updated and torch.allclose(
+                p, torch.zeros_like(p))
+        return weights_updated
+
+
+
+def _create_trainer_group(
+    worker_cls,
+    param_server_cls,
+    world_size: int,
+    vllm_master_address,
+    vllm_master_port,
+    model,
+):
+    addr, port = get_ip(), get_open_port()
+    trainer_workers = []
+    fsdp_world_size = world_size - 1
+    for i in range(fsdp_world_size):
+        env_vars = {
+            "RANK": str(i),
+            "WORLD_SIZE": world_size,
+            "MASTER_ADDR": str(addr),
+            "MASTER_PORT": str(port),
+        }
+        worker = worker_cls.remote(model, env_vars)
+        trainer_workers.append(worker)
+    
+    env_vars = {
+        "RANK": str(world_size - 1),
+        "WORLD_SIZE": world_size,
+        "MASTER_ADDR": str(addr),
+        "MASTER_PORT": str(port),
+    }
+    parameter_server = param_server_cls.remote(model, vllm_master_address, vllm_master_port, env_vars)
+    trainer_workers[0].register_parameter_server.remote(parameter_server)
+    return trainer_workers, parameter_server
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+
+    remote_configs = {
+        "num_cpus": 1,
+        "num_gpus": 1,
+        "memory": 2 * 1024**3,
+    }
+
+    model = "facebook/opt-125m"
+
+    ray.init(num_cpus=4, num_gpus=4)
+
+    vllm_master_address, vllm_update_port = get_ip(), get_open_port()
+
+    trainer_workers, parameter_server = _create_trainer_group(
+                                            TrainerActor,
+                                            vLLMParameterServer,
+                                            3,
+                                            vllm_master_address,
+                                            vllm_update_port,
+                                            model,
+                                        )
+
+    handles = []
+    for trainer_worker in trainer_workers:
+        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)
+
+    make_env_parsed = partial(make_env, batch_size=args.batch_size, dataset=args.dataset)
+    collector = RayCollector(
+        [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,
+        },
+        update_after_each_batch=True,
+    )
+    print("done collector init")
+
+    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-125m")
+
+    for i, data in enumerate(collector):
+        print(tokenizer.decode(data["tokens"][0].squeeze()))
+        print(tokenizer.decode(data["tokens_response"][0].squeeze()))
+        if i == 1:
+            break
+    collector.shutdown()

torchrl/collectors/collectors.py

@@ -76,6 +76,15 @@ def cudagraph_mark_step_begin():
         """Placeholder for missing cudagraph_mark_step_begin method."""
         raise NotImplementedError("cudagraph_mark_step_begin not implemented.")
 
+try:
+    import ray
+    from ray.actor import ActorHandle
+
+    _has_ray = True
+except ImportError as err:
+    _has_ray = False
+    RAY_ERR = err
+
 
 _TIMEOUT = 1.0
 INSTANTIATE_TIMEOUT = 20
@@ -174,9 +183,12 @@ def remote_weight_updater(self) -> RemoteWeightUpdaterBase:
     @remote_weight_updater.setter
     def remote_weight_updater(self, value: RemoteWeightUpdaterBase | None):
         if value is not None:
-            value.register_collector(self)
-            if value.collector is not self:
-                raise RuntimeError("Failed to register collector.")
+            if _has_ray and isinstance(value, ray.actor.ActorHandle):
+                value.register_collector.remote(self)
+            else:
+                value.register_collector(self)
+                if value.collector is not self:
+                    raise RuntimeError("Failed to register collector.")
         self._remote_weight_updater = value
 
     def _get_policy_and_device(
@@ -308,7 +320,10 @@ def update_policy_weights_(
         if self.local_weight_updater is not None:
             self.local_weight_updater(policy_weights, **kwargs)
         if self.remote_weight_updater is not None:
-            self.remote_weight_updater(policy_weights, worker_ids=worker_ids, **kwargs)
+            if _has_ray and isinstance(self.remote_weight_updater, ray.actor.ActorHandle):
+                ray.get(self.remote_weight_updater.__call__.remote(policy_weights, worker_ids=worker_ids, **kwargs))
+            else:
+                self.remote_weight_updater(policy_weights, worker_ids=worker_ids, **kwargs)
         elif worker_ids is not None:
             raise TypeError("worker_ids was passed but remote_weight_updater was None.")
 

torchrl/collectors/distributed/ray.py

@@ -759,7 +759,7 @@ def _async_iterator(self) -> Iterator[TensorDictBase]:
             yield out_td
 
             if self.update_after_each_batch or self.max_weight_update_interval > -1:
-                self.update_policy_weights_(worker_ids=collector_index + 1)
+                self.update_policy_weights_(worker_ids=collector_index)
 
             # Schedule a new collection task
             future = collector.next.remote()

torchrl/collectors/vllm_weight_update.py

@@ -0,0 +1,179 @@
+import torch
+import threading
+
+from torchrl.collectors.weight_update import RemoteWeightUpdaterBase
+from torchrl.collectors.weight_update import LocalWeightUpdaterBase
+
+
+VLLM_ERR = None
+try:
+    import vllm
+    from vllm.worker.worker import Worker
+
+    _has_vllm = True
+except ImportError as err:
+    _has_vllm = False
+    VLLM_ERR = err
+
+# These utilities are copied from vLLM's example code.
+def stateless_init_process_group(
+    master_address: str,
+    master_port: int,
+    rank: int,
+    world_size: int,
+    device: torch.device,
+):
+    """
+    vLLM provides `StatelessProcessGroup` to create a process group
+    without considering the global process group in torch.distributed.
+    It is recommended to create `StatelessProcessGroup`, and then initialize
+    the data-plane communication (NCCL) between external (train processes)
+    and vLLM workers.
+    """
+    from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
+    from vllm.distributed.utils import StatelessProcessGroup
+
+    pg = StatelessProcessGroup.create(
+        host=master_address, port=master_port, rank=rank, world_size=world_size
+    )
+    pynccl = PyNcclCommunicator(pg, device=device)
+    return pynccl
+
+
+if _has_vllm:
+    # I should use worker_extension_cls arg and not inherit from worker,
+    # but that is only available on main and not vLLM 0.7.3
+    class WorkerExtension(Worker):
+        """
+        The class for vLLM's worker to inherit from.
+        By defining an extension class, the code can work no matter what is
+        the underlying worker class. This way, the code can be compatible
+        with both vLLM V0 and V1.
+        NOTE: we define this class in a separate module, and the main module
+        should pass the full qualified name as `worker_extension_cls` argument.
+        """
+
+        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
+            self.model_update_group = stateless_init_process_group(
+                master_address,
+                master_port,
+                rank,
+                world_size,
+                self.device,
+            )
+
+        def update_weight(self, name, dtype, shape):
+            weight = torch.empty(shape, dtype=dtype, device="cuda")
+            self.model_update_group.broadcast(weight,
+                                            src=0,
+                                            stream=torch.cuda.current_stream())
+
+            self.model_runner.model.load_weights(weights=[(name, weight)])
+
+            del weight
+        
+        def check_weights_changed(self):
+            """
+            Check if the weights are updated to 0.
+            """
+            weights_updated = True
+            for name, p in self.model_runner.model.named_parameters():
+                weights_updated = weights_updated and torch.allclose(
+                    p, torch.zeros_like(p))
+            return weights_updated
+else:
+    class WorkerExtension:
+        pass
+
+
+class vLLMHFLocalWeightUpdater(LocalWeightUpdaterBase):
+    def __init__(self, master_address, master_port, model_metadata):
+        self.master_address = master_address
+        self.master_port = master_port
+        self.model_metadata = model_metadata
+        self.model_update_group = None
+
+    def _get_server_weights(self):
+        return None
+
+    def _get_local_weights(self):
+        # We don't implement this because we let vLLM's update_weights API handle everything
+        return None
+    
+    def _maybe_map_weights(self, server_weights, local_weights):
+        # vLLM update_weights function handles the mapping from huggingface
+        # so we don't implement this
+        return None
+    
+    def _update_local_weights(self, local_weights, mapped_weights):
+        llm = self.collector.policy["generate"].module
+        if self.model_update_group is None:
+            # FIXME: hardcoded
+            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)
+            )
+        
+        for k, (dtype, shape) in self.model_metadata.items():
+            llm.collective_rpc(
+                "update_weight",
+                args=(k, dtype, shape)
+            )
+
+class vLLMRemoteWeightUpdaterBase(RemoteWeightUpdaterBase):
+    def __init__(self, model, vllm_master_address, vllm_master_port):
+        super().__init__()
+        from transformers import AutoModel
+        self.vllm_master_address = vllm_master_address
+        self.vllm_master_port = vllm_master_port
+        self.state_dict = AutoModel.from_pretrained(model).cuda().eval().state_dict()
+        self.state_dict_lock = threading.Lock()
+        self.vllm_comm_groups = dict()
+        # versioning nyi
+        self.version = 0
+
+    def acquire_state_dict_lock(self):
+        self.state_dict_lock.acquire()
+    
+    def release_state_dict_lock(self):
+        self.state_dict_lock.release()
+    
+    def get_model_metadata(self):
+        return {k: (v.dtype, v.shape) for k, v in self.state_dict.items()}
+    
+    def all_worker_ids(self):
+        return [i for i in range(len(self.collector._remote_collectors))]
+        
+    def _get_server_weights(self):
+        return self.state_dict
+    
+    def _maybe_map_weights(self, server_weights):
+        return server_weights
+    
+    def _init_model_update_group(self, worker_id):
+        # here again, I want to grab the tp size from the vLLM worker... :(
+        # llm.llm_engine.parallel_config.tensor_parallel_size
+        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,
+            0,
+            weight_sync_world_size,
+            torch.device("cuda:0"),
+        )
+        self.vllm_comm_groups[worker_id] = model_update_group
+
+    def _sync_weights_with_worker(
+        self, worker_id: int, server_weights
+    ):
+        self.collector._remote_collectors[worker_id].update_policy_weights_.remote()
+        if worker_id not in self.vllm_comm_groups:
+            self._init_model_update_group(worker_id)
+        with self.state_dict_lock:
+            for i, k in enumerate(server_weights.keys()):
+                self.vllm_comm_groups[worker_id].broadcast(server_weights[k], src=0, stream=torch.cuda.current_stream())