Use sync recv and modify sampler (#109)

* Modify sampler

* Use sync send recv

* revert benchmark serving

Author: gty111

gllm/dist_utils.py

@@ -16,12 +16,12 @@ def recv_pp_data(src, shape, has_residual):
     hidden_states = torch.zeros(torch.Size(shape))
     if has_residual:
         residual = hidden_states.clone().detach()
-        hidden_states_future = dist.irecv(hidden_states, src)
-        residual_future = dist.irecv(residual, src)
-        return hidden_states_future, residual_future, hidden_states, residual
+        dist.recv(hidden_states, src)
+        dist.recv(residual, src)
+        return hidden_states, residual
     else:
-        hidden_states_future = dist.irecv(hidden_states, src)
-        return hidden_states_future, hidden_states
+        dist.recv(hidden_states, src)
+        return hidden_states
 
 def send_obj_list(obj_list, dst):
     dist.send_object_list(obj_list, dst=dst)

gllm/layers/sampler.py

@@ -13,10 +13,11 @@ def forward(self, logits: torch.Tensor, input_data: InputData):
         # top_p top_k
         logits = self._apply_top_k_top_p(logits, input_data.top_p, input_data.top_k)
         probs = torch.softmax(logits, dim=1)
-        # q = torch.empty_like(probs)
-        # q.exponential_()
-        # return probs.div_(q).argmax(dim=1).cpu().numpy().tolist()
-        return torch.multinomial(probs, 1).squeeze(1).cpu().numpy().tolist()
+        
+        q = torch.empty_like(probs)
+        q.exponential_()
+        return probs.div_(q).argmax(dim=1).cpu().numpy().tolist()
+        # return torch.multinomial(probs, 1).squeeze(1).cpu().numpy().tolist()
 
     def _apply_top_k_top_p(
         self,

gllm/worker.py

@@ -108,17 +108,10 @@ def forward_pp(self):
             hidden_states = None
             residual = None
             if self.ret_residual:
-                input_data, (hidden_states_future, residual_future,
-                             hidden_states, residual) = self.run_queue[0]
-                if not (hidden_states_future.is_completed() and residual_future.is_completed()):
-                    return
+                input_data, (hidden_states, residual) = self.run_queue.popleft()
             else:
-                input_data, (hidden_states_future,
-                             hidden_states) = self.run_queue[0]
-                if not hidden_states_future.is_completed():
-                    return
-
-            self.run_queue.popleft()
+                input_data, hidden_states = self.run_queue.popleft()
+            
             output = self.model_runner.step_once(
                 input_data, hidden_states, residual)
             if is_output_rank():