From e26f08fa82920cdd681b59a055314878edcde6c6 Mon Sep 17 00:00:00 2001 From: Omar Elayan Date: Sat, 1 Feb 2025 13:38:53 +0200 Subject: [PATCH] import triton files when triton is supported and installed Signed-off-by: Omar Elayan --- deepspeed/model_implementations/transformers/ds_transformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/model_implementations/transformers/ds_transformer.py b/deepspeed/model_implementations/transformers/ds_transformer.py index 360113b78a3d..7e3c81b714c0 100644 --- a/deepspeed/model_implementations/transformers/ds_transformer.py +++ b/deepspeed/model_implementations/transformers/ds_transformer.py @@ -14,7 +14,7 @@ from deepspeed.ops.transformer.inference.op_binding.workspace import WorkspaceOp from deepspeed.accelerator import get_accelerator import deepspeed -if deepspeed.HAS_TRITON: +if deepspeed.HAS_TRITON and get_accelerator().is_triton_supported(): from deepspeed.ops.transformer.inference.triton.mlp import TritonMLP from deepspeed.ops.transformer.inference.triton.attention import TritonSelfAttention