Fix dp sync after upstream change #24105 (#179)

- fix behavior of Lazy + `enforce_eager` in which case hpu graph is NOT
used
- disable device group for dp sync when hpu graph is used
- enable DP CI test again

---------

Signed-off-by: Wuxun Zhang <[email protected]>
Co-authored-by: Chendi.Xue <[email protected]>

Author: wuxun-zhang

tests/full_tests/ci_gsm8k_tests.sh

@@ -209,13 +209,12 @@ if [ $? -ne 0 ]; then
 fi
 echo "Embedding-model-support for v1 successful"
 
-# Data Parallel failed with recent upstream changes
-# # DP2
-# echo "Testing data parallel size 2 with vllm-hpu plugin v1"
-# echo HABANA_VISIBLE_DEVICES=all VLLM_SKIP_WARMUP=true PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 python -u vllm-gaudi/examples/data_parallel.py --dp-size 2 --tp-size 2
-# HABANA_VISIBLE_DEVICES=all VLLM_SKIP_WARMUP=true PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 python -u vllm-gaudi/examples/data_parallel.py --dp-size 2 --tp-size 2
-# if [ $? -ne 0 ]; then
-#     echo "Error: Test failed for data parallel size 2" >&2
-#     exit -1
-# fi
-# echo "Test with data parallel size 2 passed"
+# DP2
+echo "Testing data parallel size 2 with vllm-hpu plugin v1"
+echo HABANA_VISIBLE_DEVICES=all VLLM_SKIP_WARMUP=true PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 python -u vllm-gaudi/examples/data_parallel.py --dp-size 2 --tp-size 2
+HABANA_VISIBLE_DEVICES=all VLLM_SKIP_WARMUP=true PT_HPU_LAZY_MODE=1 VLLM_USE_V1=1 python -u vllm-gaudi/examples/data_parallel.py --dp-size 2 --tp-size 2
+if [ $? -ne 0 ]; then
+    echo "Error: Test failed for data parallel size 2" >&2
+    exit -1
+fi
+echo "Test with data parallel size 2 passed"

vllm_gaudi/v1/worker/hpu_model_runner.py

@@ -730,6 +730,13 @@ def __init__(
         self._PAD_BLOCK_ID = -1
         self._tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config)
 
+        if self.vllm_config.parallel_config.data_parallel_size > 1 and htorch.utils.internal.is_lazy(
+        ) and not self.model_config.enforce_eager:
+            from vllm import envs
+            # disable device group for dp synchronization when hpu graph is
+            # turned on since it's not captured and causes issues
+            envs.VLLM_DISABLE_NCCL_FOR_DP_SYNCHRONIZATION = True
+
         # TODO(madamczyk-intel): add a knob for that
         # TODO(madamczyk-intel): debug why increasing it lowers acc
         self.logits_rounding = 1
@@ -2230,7 +2237,7 @@ def _execute_model_generic(self,
         num_blocks = self._num_blocks(attn_metadata)
         self._check_config(batch_size, seq_len, num_blocks, attn_metadata, warmup_mode)
         additional_kwargs = {}
-        if htorch.utils.internal.is_lazy() and not self.model_config.enforce_eager:
+        if htorch.utils.internal.is_lazy():
             use_graphs = self._use_graphs()
             additional_kwargs.update({"bypass_hpu_graphs": not use_graphs})
         else:
@@ -2252,7 +2259,8 @@ def _execute_model_generic(self,
                                                kv_caches=kv_caches,
                                                inputs_embeds=inputs_embeds,
                                                model_mm_kwargs=model_mm_kwargs,
-                                               lora_mask=lora_mask)
+                                               lora_mask=lora_mask,
+                                               **additional_kwargs)
         # NOTE(kzawora): returning hidden_states is required in prompt logprobs
         # scenarios, as they will do logit processing on their own
         if self.use_aux_hidden_state_outputs: