import triton files when triton is supported and installed (deepspeed…

…ai#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 <[email protected]> Signed-off-by: siqi <[email protected]>
siqi654321 · Feb 7, 2025 · 82ca0ea · 82ca0ea
1 parent f9c8c1d
commit 82ca0ea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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