[Bugfix] Add OOT backend option (#32471)
Signed-off-by: Iryna Boiko <iboiko@habana.ai>
This commit is contained in:
@@ -33,6 +33,7 @@ class UnquantizedMoeBackend(Enum):
|
|||||||
CPU = "CPU"
|
CPU = "CPU"
|
||||||
XPU = "XPU"
|
XPU = "XPU"
|
||||||
TPU = "TPU"
|
TPU = "TPU"
|
||||||
|
OOT = "OOT"
|
||||||
|
|
||||||
|
|
||||||
# NOTE(zyongye): Unsupported backend means backend
|
# NOTE(zyongye): Unsupported backend means backend
|
||||||
@@ -42,6 +43,7 @@ UNSUPPORTED_BACKEND = [
|
|||||||
UnquantizedMoeBackend.CPU,
|
UnquantizedMoeBackend.CPU,
|
||||||
UnquantizedMoeBackend.XPU,
|
UnquantizedMoeBackend.XPU,
|
||||||
UnquantizedMoeBackend.TPU,
|
UnquantizedMoeBackend.TPU,
|
||||||
|
UnquantizedMoeBackend.OOT,
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
@@ -95,6 +97,8 @@ def select_unquantized_moe_backend(
|
|||||||
backend = UnquantizedMoeBackend.CPU
|
backend = UnquantizedMoeBackend.CPU
|
||||||
if current_platform.is_tpu():
|
if current_platform.is_tpu():
|
||||||
backend = UnquantizedMoeBackend.TPU
|
backend = UnquantizedMoeBackend.TPU
|
||||||
|
if current_platform.is_out_of_tree():
|
||||||
|
backend = UnquantizedMoeBackend.OOT
|
||||||
|
|
||||||
logger.info_once(_make_log_backend(backend), scope="local")
|
logger.info_once(_make_log_backend(backend), scope="local")
|
||||||
return backend
|
return backend
|
||||||
|
|||||||
Reference in New Issue
Block a user