From 82ca0ea9cd4d4482f861ec6e6f27801e9450f5a0 Mon Sep 17 00:00:00 2001 From: Omar Elayan <142979319+oelayan7@users.noreply.github.com> Date: Fri, 7 Feb 2025 00:05:31 +0200 Subject: [PATCH] import triton files when triton is supported and installed (#6989) Those files have code that gets run when importing them, so in systems that doesn't support triton but have triton installed this causes issues. In general, I think it is better to import triton when it is installed and supported. Signed-off-by: Omar Elayan Signed-off-by: siqi --- 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 360113b78a3db..7e3c81b714c06 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