Support synchronizing shard weights from training instance to inferencing instance

[hxdtest]

This is an issue to request for comments (RFC) from veRL maintainers on an implementation that would fundamentally support synchronizing shard weights from training instance to inferencing instance.

• Background

If the training instance and inference instance (VerlEngine) is launched separately, i.e. they are not in the same process, model weights needed to transported from training instance to inference engine. It is a pretty large overhead, especially for models like DeepSeek V3 671B with bf16 datatype. If tp size is 32 for inferencing, every inference instance only needs 671/32=21B parameter.

• Use-case:

Inference Side

Build sync weights process group with training instance

        if self.tp_rank == 0:
            self.inference_engine._engine.init_weights_update_group(
                master_address=ip,
                master_port=port,
                rank_offset= self.dp_rank * self.tp_size + inference_world_size,   
                world_size=self.sync_weight_world_size,  # 1 + 8
                group_name="test_parameter_update_group",
                backend="nccl",
            )

Receive shard weights from training instance and update weights in place

            sync_sharded_weights = torch.tensor([0], dtype=torch.int64).cuda()
            torch.distributed.recv(sync_sharded_weights, src=src, group=self._model_update_group)
            sync_sharded_weights = sync_sharded_weights.item()
            num_layers = torch.tensor([0], dtype=torch.int64, device=self.gpu_id)
            torch.distributed.recv(num_layers, src=src, group=self._model_update_group)
            num_layers = num_layers.item()
 

            for i in range(num_layers + 1):
                tensor_size = torch.tensor([0], dtype=torch.long, device=self.gpu_id)
                torch.distributed.recv(tensor_size, src=src, group=self._model_update_group)
                tensor = torch.empty(tensor_size.item(), dtype=torch.uint8,device=self.gpu_id)
                torch.distributed.recv(tensor, src=src, group=self._model_update_group)
                import pickle
                serialized_data = bytes(tensor.cpu().numpy())
                state_dict_shape = pickle.loads(serialized_data)
                state_dict = {}
                ops = []
   
                if sync_sharded_weights:
                    params_dict = dict(self.model.named_parameters())
                    for k, v in state_dict_shape.items():
                        param = params_dict[k]
                        recv_op = dist.P2POp(dist.irecv, param.data, src, group=self._model_update_group)
                        ops.append(recv_op)
                    reqs = dist.batch_isend_irecv(ops)
                    for req in reqs:
                        req.wait()
                dist.barrier(group=self._model_update_group, device_ids=[self.gpu_id])

            self.model.post_process()

Training Side
Build sync weights process group with inference instance

            group = init_custom_process_group(
                backend="nccl",
                init_method=f"tcp://{ip}:{port}",
                world_size=self.sync_weight_world_size,
                rank=self.rank,
                group_name="test_parameter_update_group",
            )

             sglang_state_dict = convert_hf_to_sglang(state_dict, sglang_tp_rank, sglang_tp_size, 
                                                                                        first_k_dense_layer, num_experts, num_layers)
            state_dict = sglang_state_dict
            state_dict_size = {}
            for k, v in state_dict.items():
                state_dict_size[k] = list(v.shape)

            serialized_data = pickle.dumps(state_dict_size)
            size = len(serialized_data)
            tensor_data = torch.ByteTensor(
                np.frombuffer(serialized_data, dtype=np.uint8)
            ) .cuda()
            tensor_size = torch.tensor([size], dtype=torch.long).cuda()

            if self.participate_sync_weight_group():  
                torch.distributed.send(tensor_size, dst=dst, group =  self.weight_sync_group)
                torch.distributed.send(tensor_data, dst=dst, group =  self.weight_sync_group)
                ops = []
                for k, v in state_dict.items():
                    # torch.distributed.send(v, dst=dst, group = self.weight_sync_group)      
                    send_op = dist.P2POp(dist.isend, v.contiguous(), dst, group = self.weight_sync_group)      
                    ops.append(send_op)
            
                reqs = dist.batch_isend_irecv(ops)
                for req in reqs:
                    req.wait()
                torch.distributed.barrier(group=self.weight_sync_group, device_ids=[self.device])