From ec8e8e19214aa9c91e4d4a231d76f7bf3a9c63d0 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 --- 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