Intel Gpu support initialization (#4340)

This commit is contained in:
Abhilash Majumder 2023-10-27 08:09:51 +05:30 committed by GitHub
parent 317e2c857e
commit 778a010df8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 106 additions and 42 deletions

View file

@ -9,6 +9,7 @@ from pathlib import Path
import numpy as np
from tokenizers import Tokenizer
from transformers import is_torch_xpu_available
import modules.shared as shared
from modules.callbacks import Iteratorize
@ -27,7 +28,7 @@ class RWKVModel:
pass
@classmethod
def from_pretrained(self, path, dtype="fp16", device="cuda"):
def from_pretrained(self, path, dtype="bf16" if is_torch_xpu_available() else "fp16", device="xpu" if is_torch_xpu_available() else "cuda"):
tokenizer_path = Path(f"{path.parent}/20B_tokenizer.json")
if shared.args.rwkv_strategy is None:
model = RWKV(model=str(path), strategy=f'{device} {dtype}')