[WIP] Add fake HPU mode to Habana components.

.github/workflows/cpu-test.yml

@@ -0,0 +1,34 @@
+name: cpu-test
+
+on:
+  # Trigger the workflow on push or pull request,
+  # but only for the habana_main branch
+  push:
+    branches:
+      - habana_main
+  pull_request:
+    branches:
+      - habana_main
+
+
+jobs:
+  cputest:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install -r requirements-hpu.txt
+        VLLM_TARGET_DEVICE=hpu python setup.py develop
+    - name: cpu-test
+      run: |
+        VLLM_SKIP_WARMUP=true VLLM_PROMPT_SEQ_BUCKET_MAX=128 VLLM_USE_FAKE_HPU=1  python examples/offline_inference_fakehpu.py --fake_hpu

examples/offline_inference_fakehpu.py

@@ -0,0 +1,41 @@
+from os import environ
+
+if environ.get('VLLM_USE_FAKE_HPU', '0') == '1':
+    from vllm.utils import migrate_to_cpu
+    print("CHECK1")
+    migrate_to_cpu()
+    print("CHECK2")
+
+# Sample prompts.
+prompts = [
+    "Berlin is the capital city of ",
+    "Louvre is located in the city called ",
+    "Barack Obama was the 44th president of ",
+    "Warsaw is the capital city of ",
+    "Gniezno is a city in ",
+    "Hebrew is an official state language of ",
+    "San Francisco is located in the state of ",
+    "Llanfairpwllgwyngyll is located in country of ",
+]
+ref_answers = [
+    "Germany", "Paris", "United States", "Poland", "Poland", "Israel",
+    "California", "Wales"
+]
+
+from vllm import LLM, SamplingParams
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0, n=1, use_beam_search=False)
+
+# Create an LLM.
+llm = LLM(model="facebook/opt-125m", max_model_len=32, max_num_seqs=4)
+# Generate texts from the prompts. The output is a list of RequestOutput objects
+# that contain the prompt, generated text, and other information.
+outputs = llm.generate(prompts, sampling_params)
+# Print the outputs.
+for output, answer in zip(outputs, ref_answers):
+    prompt = output.prompt
+    generated_text = output.outputs[0].text
+    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
+    #assert answer in generated_text, (
+     #   f"The generated text does not contain the correct answer: {answer}")
+print('PASSED')

vllm/__init__.py

@@ -1,5 +1,10 @@
 """vLLM: a high-throughput and memory-efficient inference engine for LLMs"""
 
+from vllm.utils import is_fake_hpu, migrate_to_cpu
+
+if is_fake_hpu():
+    migrate_to_cpu()
+
 from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
 from vllm.engine.llm_engine import LLMEngine

vllm/distributed/device_communicators/hpu_communicator.py

@@ -4,9 +4,8 @@
 
 from vllm.platforms import current_platform
 
-if current_platform.is_hpu():
-    import habana_frameworks.torch as htorch  # noqa: F401
 
+import habana_frameworks.torch as htorch  # noqa: F401)
 
 class HpuCommunicator:
 

vllm/executor/ray_habana_executor.py

@@ -13,7 +13,7 @@
 from vllm.utils import (_run_task_with_lock,
                         error_on_invalid_device_count_status,
                         get_distributed_init_method, get_ip, get_open_port,
-                        get_vllm_instance_id, make_async)
+                        get_vllm_instance_id, is_fake_hpu, make_async)
 
 if ray is not None:
     from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
@@ -87,18 +87,20 @@ def _init_workers_ray(self, placement_group: "PlacementGroup",
         driver_ip = get_ip()
         worker_wrapper_kwargs = self._get_worker_wrapper_args()
         for bundle_id, bundle in enumerate(placement_group.bundle_specs):
-            if not bundle.get("HPU", 0):
+            resource_name = "HPU" if not is_fake_hpu() else "CPU"
+            if not bundle.get(resource_name, 0):
                 continue
             scheduling_strategy = PlacementGroupSchedulingStrategy(
                 placement_group=placement_group,
                 placement_group_capture_child_tasks=True,
                 placement_group_bundle_index=bundle_id,
             )
-
+            resources = {'HPU': num_gpus} if not is_fake_hpu() else {}
+            num_cpus = 0 if not is_fake_hpu() else num_gpus
             worker = ray.remote(
-                num_cpus=0,
+                num_cpus=num_cpus,
                 num_gpus=0,
-                resources={'HPU': num_gpus},
+                resources=resources,
                 scheduling_strategy=scheduling_strategy,
                 **ray_remote_kwargs,
             )(RayWorkerWrapper).remote(**worker_wrapper_kwargs)

vllm/executor/ray_utils.py

@@ -3,7 +3,7 @@
 from vllm.config import ParallelConfig
 from vllm.logger import init_logger
 from vllm.sequence import ExecuteModelRequest
-from vllm.utils import get_ip, is_hip, is_hpu, is_tpu, is_xpu
+from vllm.utils import get_ip, is_fake_hpu, is_hip, is_hpu, is_tpu, is_xpu
 from vllm.worker.worker_base import WorkerWrapperBase
 
 logger = init_logger(__name__)
@@ -97,7 +97,7 @@ def initialize_ray_cluster(
     if is_tpu():
         device_str = "TPU"
     elif is_hpu():
-        device_str = "HPU"
+        device_str = "HPU" if not is_fake_hpu() else 'CPU'
     # Create placement group for worker processes
     current_placement_group = ray.util.get_current_placement_group()
     if current_placement_group:

vllm/hpu/ops.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 import habana_frameworks.torch as htorch
+
 import torch
 import torch.nn.functional as F
 
@@ -265,4 +266,4 @@ def dispatch_bgmv_embedding(
     x = x.unsqueeze(1)
     out = x @ wa
     out = out.squeeze(1)
-    y += out * scale
+    y += out * scale

vllm/hpu/utils.py

@@ -8,6 +8,7 @@
 from functools import wraps
 
 import habana_frameworks.torch as htorch
+
 import torch
 
 from vllm.hpu.cache_ops import insert_or_update_cache

vllm/model_executor/models/opt.py

@@ -254,11 +254,9 @@ def forward(
         if self.project_in is not None:
             inputs_embeds, _ = self.project_in(inputs_embeds)
         hidden_states = inputs_embeds + pos_embeds
-
         for i in range(len(self.layers)):
             layer = self.layers[i]
             hidden_states = layer(hidden_states, kv_caches[i], attn_metadata)
-
         if self.final_layer_norm is not None:
             hidden_states = self.final_layer_norm(hidden_states)
         if self.project_out is not None:

vllm/utils.py

@@ -208,9 +208,30 @@ def is_neuron() -> bool:
 
 @lru_cache(maxsize=None)
 def is_hpu() -> bool:
+    return _is_habana_frameworks_installed() or _is_built_for_hpu()
+
+
+@lru_cache(maxsize=None)
+def is_fake_hpu() -> bool:
+    return os.environ.get('VLLM_USE_FAKE_HPU', '0') != '0' or (
+        not _is_habana_frameworks_installed() and _is_built_for_hpu())
+
+
+@lru_cache(maxsize=None)
+def _is_habana_frameworks_installed() -> bool:
     from importlib import util
+    if os.environ.get('VLLM_USE_FAKE_HPU', '0') == '1':
+        return False
     return util.find_spec('habana_frameworks') is not None
 
+@lru_cache(maxsize=None)
+def _is_built_for_hpu() -> bool:
+    from importlib.metadata import PackageNotFoundError, version
+    try:
+        return "gaudi" in version("vllm")
+    except PackageNotFoundError:
+        return False
+
 
 @lru_cache(maxsize=None)
 def is_tpu() -> bool:
@@ -624,18 +645,24 @@ def __init__(self, device=None):
 
     @staticmethod
     def current_device_memory_usage() -> float:
+        if is_fake_hpu():
+            return 0
         # Return the device memory usage in bytes.
         free_hpu_memory, total_hpu_memory = torch.hpu.mem_get_info()
         return total_hpu_memory - free_hpu_memory
 
     @staticmethod
     def current_free_device_memory() -> float:
+        if is_fake_hpu():
+            return 0
         # Return the device memory usage in bytes.
         free_hpu_memory, _ = torch.hpu.mem_get_info()
         return free_hpu_memory
 
     @staticmethod
     def total_device_memory() -> float:
+        if is_fake_hpu():
+            return 0
         # Return the device memory usage in bytes.
         _, total_hpu_memory = torch.hpu.mem_get_info()
         return total_hpu_memory
@@ -1088,3 +1115,44 @@ async def _run_task_with_lock(task: Callable, lock: asyncio.Lock, *args,
     """Utility function to run async task in a lock"""
     async with lock:
         return await task(*args, **kwargs)
+
+def _create_dummy_modules():
+    import types
+    import importlib
+
+    habana_frameworks = types.ModuleType('habana_frameworks')
+    spec = importlib.util.spec_from_loader('habana_frameworks', loader=None)
+    habana_frameworks.__spec__ = spec
+    sys.modules['habana_frameworks'] = habana_frameworks
+    sys.modules['habana_frameworks.torch'] = habana_frameworks.torch = types.ModuleType('habana_frameworks.torch')
+    sys.modules['habana_frameworks.torch.core'] = habana_frameworks.torch.core = types.ModuleType('habana_frameworks.torch.core')
+
+    sys.modules['habana_frameworks.torch.utils'] = habana_frameworks.torch.utils = types.ModuleType('habana_frameworks.torch.utils')
+    sys.modules['habana_frameworks.torch.utils.internal'] = habana_frameworks.torch.utils.internal = types.ModuleType('habana_frameworks.torch.utils.internal')
+
+    sys.modules['torch.hpu'] = torch.hpu = types.ModuleType('torch.hpu')
+
+    habana_frameworks.torch.core.mark_step = lambda: print('calling mark_step')
+    habana_frameworks.torch.utils.internal.is_lazy = lambda: print('calling is_lazy') 
+    torch.hpu.synchronize = lambda: print('calling synchronize')
+
+def _do_nothing():
+    pass
+
+def _return_false():
+    return False
+
+def _migrate_to_cpu():
+    import habana_frameworks.torch as htorch
+
+    print(sys.modules['torch'].__spec__)
+    print(sys.modules['habana_frameworks'].__spec__)
+
+    htorch.core.mark_step = _do_nothing
+    htorch.utils.internal.is_lazy = _return_false
+    torch.hpu.synchronize = _do_nothing
+
+def migrate_to_cpu():
+    _create_dummy_modules()
+    _migrate_to_cpu()
+

vllm/worker/cache_engine.py

@@ -6,7 +6,7 @@
 from vllm.attention import get_attn_backend
 from vllm.config import CacheConfig, DeviceConfig, ModelConfig, ParallelConfig
 from vllm.logger import init_logger
-from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, get_dtype_size,
+from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, get_dtype_size, is_fake_hpu,
                         is_pin_memory_available)
 
 logger = init_logger(__name__)
@@ -78,7 +78,7 @@ def _allocate_kv_cache(
         pin_memory = is_pin_memory_available() if device == "cpu" else False
         kv_cache: List[torch.Tensor] = []
         for _ in range(self.num_attention_layers):
-            if device == 'hpu':
+            if device == 'hpu' or is_fake_hpu():
                 key_cache = torch.zeros(kv_cache_shape,
                                         dtype=self.dtype,
                                         device=device)

vllm/worker/habana_model_runner.py

@@ -14,7 +14,11 @@
 from typing import (TYPE_CHECKING, Any, Callable, Dict, List, NamedTuple,
                     Optional, Set, Tuple, Type, TypeVar, Union)
 
+from vllm.utils import (HabanaMemoryProfiler, format_bytes, is_fake_hpu,
+                        is_pin_memory_available, make_tensor_with_pad)
+
 import habana_frameworks.torch as htorch
+
 import torch
 
 from vllm.attention import AttentionMetadata, get_attn_backend
@@ -31,8 +35,6 @@
 from vllm.sampling_params import SamplingParams
 from vllm.sequence import (IntermediateTensors, SamplerOutput, SequenceData,
                            SequenceGroupMetadata)
-from vllm.utils import (HabanaMemoryProfiler, format_bytes,
-                        is_pin_memory_available, make_tensor_with_pad)
 from vllm.worker.model_runner_base import (
     ModelRunnerBase, ModelRunnerInputBase,
     _add_attn_metadata_broadcastable_dict,
@@ -189,8 +191,9 @@ def __init__(self, model, enforce_eager):
         self.model = model
         self.prefill_use_fusedsdpa = os.getenv('VLLM_PROMPT_USE_FUSEDSDPA',
                                                '0').lower() in ['1', 'true']
-
-        if not htorch.utils.internal.is_lazy() and not enforce_eager:
+
+        if not is_fake_hpu()  and not htorch.utils.internal.is_lazy(
+        ) and not enforce_eager:
             self.model = torch.compile(self.model,
                                        backend='hpu_backend',
                                        dynamic=False)
@@ -419,7 +422,9 @@ def __init__(
                                if model_config is not None else None)
         self.device_config = (device_config
                               if device_config is not None else DeviceConfig())
-
+        if is_fake_hpu():
+            device_config.device = torch.device('cpu')
+            device_config.device_type = 'cpu'
         self.device = self.device_config.device
         self.enforce_eager = self.model_config.enforce_eager
         self.max_num_seqs = self.scheduler_config.max_num_seqs
@@ -1406,6 +1411,8 @@ def mem_margin(self, value):
 
 
 def _maybe_wrap_in_hpu_graph(*args, **kwargs):
+    if is_fake_hpu():
+        return HpuModelAdapter(*args, **kwargs)
     return htorch.hpu.wrap_in_hpu_graph(HpuModelAdapter(
         *args, **
         kwargs)) if htorch.utils.internal.is_lazy() else HpuModelAdapter(