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 <[email protected]>
deepspeedai · Feb 6, 2025 · 301bc28 · 301bc28
1 parent f04649d
commit 301bc28
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