[WIP] run any model from huggingface

.circleci/config.yml

@@ -178,6 +178,50 @@ jobs:
           inference_engine: mlx
           model_id: llama-3.2-1b
 
+  chatgpt_api_integration_test_torch_linux_cpu:
+    machine:
+      image: ubuntu-2404:2024.08.1
+    resource_class: large
+    steps:
+      - checkout
+      - run:
+          name: Set up Python
+          command: |
+            brew install [email protected]
+            python3.12 -m venv env
+            source env/bin/activate
+      - run:
+          name: Install dependencies
+          command: |
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install .
+      - run_chatgpt_api_test:
+          inference_engine: torch
+          model_id: llama-3.2-1b
+
+  chatgpt_api_integration_test_torch_mac:
+    macos:
+      xcode: "15.4.0"
+    resource_class: macos.m1.large.gen1
+    steps:
+      - checkout
+      - run:
+          name: Set up Python
+          command: |
+            brew install [email protected]
+            python3.12 -m venv env
+            source env/bin/activate
+      - run:
+          name: Install dependencies
+          command: |
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install .
+      - run_chatgpt_api_test:
+          inference_engine: torch
+          model_id: llama-3.2-1b
+
   test_macos_m1:
     macos:
       xcode: "16.0.0"
@@ -211,9 +255,72 @@ jobs:
 workflows:
   version: 2
   build_and_test:
+    jobs:
+      - approve_run:
+          type: approval
+          requires: []
+          filters:
+            branches:
+              ignore: main
+      - unit_test:
+          requires:
+            - approve_run
+      - discovery_integration_test:
+          requires:
+            - approve_run
+      - chatgpt_api_integration_test_mlx:
+          requires:
+            - approve_run
+      - test_macos_m1:
+          requires:
+            - approve_run
+      - chatgpt_api_integration_test_torch_linux_cpu:
+          requires:
+            - approve_run
+      - chatgpt_api_integration_test_torch_mac:
+          requires:
+            - approve_run
+
+  # Workflow for forked PRs without approval
+  forked_pr_workflow:
     jobs:
       - unit_test
       - discovery_integration_test
       - chatgpt_api_integration_test_mlx
       - test_macos_m1
-      # - chatgpt_api_integration_test_tinygrad
+      - chatgpt_api_integration_test_torch_linux_cpu
+      - chatgpt_api_integration_test_torch_mac
+    # The trigger condition ensures this workflow runs for forked PRs
+    triggers:
+      - type: pull_request
+        filters:
+          branches:
+            ignore: main
+
+  # Existing workflow for main branch
+  main_branch_workflow:
+    jobs:
+      - unit_test:
+          filters:
+            branches:
+              only: main
+      - discovery_integration_test:
+          filters:
+            branches:
+              only: main
+      - chatgpt_api_integration_test_mlx:
+          filters:
+            branches:
+              only: main
+      - test_macos_m1:
+          filters:
+            branches:
+              only: main
+      - chatgpt_api_integration_test_torch_linux_cpu:
+          filters:
+            branches:
+              only: main
+      - chatgpt_api_integration_test_torch_mac:
+          filters:
+            branches:
+              only: main

.gitignore

@@ -170,3 +170,9 @@ cython_debug/
 #.idea/
 
 **/*.xcodeproj/*
+
+# PyTorch interface
+.offload
+
+# neovim/vim settings 
+.vimrc

exo/api/chatgpt_api.py

@@ -17,7 +17,6 @@
 from exo.models import model_base_shards
 from typing import Callable
 
-
 class Message:
   def __init__(self, role: str, content: Union[str, List[Dict[str, Union[str, Dict[str, str]]]]]):
     self.role = role
@@ -60,6 +59,9 @@ def generate_completion(
       "finish_reason": finish_reason,
     }],
   }
+
+  if DEBUG >= 3:
+    print(f"completion: {completion}")
 
   if not stream:
     completion["usage"] = {

exo/helpers.py

@@ -33,7 +33,6 @@ def get_system_info():
     return "Linux"
   return "Non-Mac, non-Linux system"
 
-
 def find_available_port(host: str = "", min_port: int = 49152, max_port: int = 65535) -> int:
   used_ports_file = os.path.join(tempfile.gettempdir(), "exo_used_ports")
 

exo/inference/inference_engine.py

@@ -8,7 +8,7 @@
 
 class InferenceEngine(ABC):
   @abstractmethod
-  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, image_str: Optional[str] = None, inference_state: Optional[str] = None) -> (np.ndarray, str, bool):
+  async def infer_prompt(self, request_id: str, shard: Shard, prompt: str, image_str: Optional[str] = None, inference_state: Optional[str] = None) -> Tuple[np.ndarray, str, bool]:
     pass
 
   @abstractmethod
@@ -27,5 +27,8 @@ def get_inference_engine(inference_engine_name: str, shard_downloader: 'ShardDow
     tinygrad.helpers.DEBUG.value = int(os.getenv("TINYGRAD_DEBUG", default="0"))
 
     return TinygradDynamicShardInferenceEngine(shard_downloader)
+  elif inference_engine_name == "torch":
+    from exo.inference.torch.inference import TorchDynamicShardInferenceEngine
+    return TorchDynamicShardInferenceEngine(shard_downloader)
   else:
     raise ValueError(f"Inference engine {inference_engine_name} not supported")

exo/inference/torch/.gitignore

@@ -0,0 +1,2 @@
+data/
+model/archive/

exo/inference/torch/README.md

@@ -0,0 +1,44 @@
+# PyTorch & HuggingFace inference engine
+
+## Tech
+
+Tested on
+
+```bash
+# Laptop/PC
+Distributor ID: Pop
+Description:    Pop!_OS 22.04 LTS
+Release:        22.04
+Codename:       jammy
+CUDA Version: 12.4 
+Nvidia Driver Version: 550.107.02
+
+GPU 1: Nvidia GeForce RTX 3060 6GB Laptop
+```
+```bash
+# Server
+Distributor ID: Pop
+Description:    Pop!_OS 22.04 LTS
+Release:        22.04
+Codename:       jammy
+CUDA Version:   12.4
+Nvidia Driver Version: 550.90.07
+
+GPU 1: NVIDIA T1000 8GB
+GPU 2: NVIDIA Quadro M2000 4GB
+GPU 3: NVIDIA Quadro M2000 4GB
+GPU 4: NVIDIA Quadro P400 2GB
+GPU 5: NVIDIA Quadro P400 2GB 
+```
+
+
+## Notes/Issues
+### 10/10/2024
+- To select a pytorch device via environment variables, set the variable TORCH_DEVICE
+  - XLA is currently not installed and will need to be added to inference.py, looking into doing this on a TPU VM
+  - With pytorch, CUDA and ROCm are the same so specifying CUDA also enables ROCm support. See this [post](https://github.com/pytorch/pytorch/issues/55223#issuecomment-812587373)
+  - Looking into adding mobile device support properly
+- If device is not CPU the data type defaults to float32 else float16.
+
+### 10/13/2024
+Still working on split model development (see test_split_model.py). Right now, it seems to do it but still transformers is loading more in the RAM and GPU as it loads up a larger models (causing an OOM). Will research and add to next update. Right now, tests are added and are in development.