V0.7.2 ulysses v1 debug (vllm-project#6)

* print numforward

* print forward context

* print forward context

* print forward context

* print forward context

* synchronize after Ulysses in graph capture

* synchronize after Ulysses in graph capture

* synchronize NCCL operations (all-to-all) for graph capture

Author: sfc-gh-mhidayetoglu

vllm/model_executor/models/llama.py

@@ -539,6 +539,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors)
 
+        self.numforward = 0
+
     def _init_model(self, vllm_config: VllmConfig, prefix: str = ""):
         return LlamaModel(vllm_config=vllm_config, prefix=prefix)
 
@@ -570,6 +572,12 @@ def forward(
         #     print(f"input_ids: {input_ids.shape}")
         #     print(f"positions: {positions.shape}")
         #     print(f"N {N}, SP {SP}, N_ranks {N_ranks} sum {sum(N_ranks)}")
+        # from vllm.forward_context import get_forward_context
+        # if torch.distributed.get_rank() == 0:
+        #     print(f"numforward {self.numforward} N {N} "
+        #           f"N_ranks {N_ranks}")
+        # if get_forward_context().attn_metadata is not None:
+        #     self.numforward += 1
 
         input_ids[0:N_ulysses] = input_ids[N_start:N_start + N_ulysses]
         positions[0:N_ulysses] = positions[N_start:N_start + N_ulysses]

vllm/v1/attention/backends/flash_attn.py

@@ -316,6 +316,9 @@ def forward(
                                             group=self.device_group)
         c = torch.transpose(c, 0, 1).reshape(
             N_ulysses, self.num_heads * self.SP * self.head_size)
+        # synchronization of NCCL operations is necessary for graph capture
+        if attn_metadata is None:
+            torch.cuda.synchronize()
         return c