Skip to content

Commit

Permalink
import triton files when triton is supported and installed (#6989)
Browse files Browse the repository at this point in the history
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 <oelayan@habana.ai>
  • Loading branch information
oelayan7 authored and traincheck-team committed Feb 9, 2025
1 parent 397b654 commit efeec37
Showing 1 changed file with 1 addition and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down

0 comments on commit efeec37

Please # to comment.