split gptbigcode forward (#194)

HabanaAI · Aug 19, 2024 · f7dd91d · f7dd91d
1 parent 8185d76
commit f7dd91d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/vllm/model_executor/models/gpt_bigcode.py b/vllm/model_executor/models/gpt_bigcode.py
@@ -39,6 +39,7 @@
  VocabParallelEmbedding)
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.sampling_metadata import SamplingMetadata
+from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors, SamplerOutput
 
 from .interfaces import SupportsLoRA
@@ -224,9 +225,14 @@ def forward(
  position_embeds = self.wpe(position_ids)
  hidden_states = inputs_embeds + position_embeds
 
+ if current_platform.is_hpu():
+ import habana_frameworks.torch as htorch
+ htorch.core.mark_step()
  for i in range(len(self.h)):
  layer = self.h[i]
  hidden_states = layer(hidden_states, kv_caches[i], attn_metadata)
+ if current_platform.is_hpu():
+ htorch.core.mark_step()
 
  hidden_states = self.ln_f(hidden_states)
  return hidden_states