Missing triton import for sam_fast

[mdejaegher]

sam_fast model in https://github.com/pytorch/benchmark/tree/main/torchbenchmark/models/sam_fast is failing to run in both eager and compile mode because it's missing the triton import requirement.

The full error is (after I installed both https://github.com/pytorch/benchmark/blob/main/torchbenchmark/models/sam_fast/requirements.txt and https://github.com/pytorch/benchmark/blob/main/requirements.txt):

(torchbench) user@user:~/work$ python benchmark/run.py sam_fast -t eval --metrics cpu_peak_mem,ttfb
Traceback (most recent call last):
  File "/home/user/work/benchmark/run.py", line 599, in <module>
    main()  # pragma: no cover
  File "/home/user/work/benchmark/run.py", line 512, in main
    m = load_model(config)
  File "/home/user/work/benchmark/torchbenchmark/util/experiment/instantiator.py", line 85, in load_model
    Model = load_model_by_name(config.name)
  File "/home/user/work/benchmark/torchbenchmark/__init__.py", line 667, in load_model_by_name
    module = importlib.import_module(module_path, package=__name__)
  File "/usr/lib/python3.10/importlib/__init__.py", line 126, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1050, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1027, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1006, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 688, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 883, in exec_module
  File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed
  File "/home/user/work/benchmark/torchbenchmark/models/sam_fast/__init__.py", line 5, in <module>
    from segment_anything_fast.build_sam import sam_model_fast_registry
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/__init__.py", line 7, in <module>
    from .build_sam import (
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/build_sam.py", line 11, in <module>
    from .modeling import ImageEncoderViT, MaskDecoder, PromptEncoder, Sam, TwoWayTransformer
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/modeling/__init__.py", line 7, in <module>
    from .sam import Sam
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/modeling/sam.py", line 13, in <module>
    from .image_encoder import ImageEncoderViT
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/modeling/image_encoder.py", line 15, in <module>
    from segment_anything_fast.flash_4 import _attention_rel_h_rel_w
  File "/home/user/work/torchbench/lib/python3.10/site-packages/segment_anything_fast/flash_4.py", line 23, in <module>
    import triton
ModuleNotFoundError: No module named 'triton'

Basically what is happening is that the model is using segment_anything_fast which itself needs triton and tries to import it at some point, but triton is not currently in the requirements file of the sam_fast model.

[FindHao]

I think triton is one of the dependency of torch. Since torchbench is benchmarking torch, if you have installed troch, the triton(pytorch-triton) should be installed automatically, right?

[mdejaegher]

Took me a little while to figure out what was happening but your question was very helpful. Turns out my environment wasn't 100% clean and I used the 2.4.1+cpu version of torch before rather than vanilla 2.4.1. As the name suggests, that doesn't come with triton. When I use vanilla 2.4.1 from a clean environment then, as you said, it does include triton and works as expected.

Will close this as it works as expected. Thanks for the help!