No module named 'triton' - on Mac

[smriyaz]

Hey there,
All of my installation steps were successful. However, unfortunately, when I run the 'python ./examples/run_music_generation.py --model_path=./ckpt --version="3B" ' command to test the creation, I'm getting the following error:
import error: No module named 'triton'
Model components will be loaded to devices as specified:
mula: cuda
codec: cuda
You have set lazy_load = False. Loading HeartMuLa and HeartCodec onto device...

and it goes on with various line errors in the .py files. I think I don't have 'triton' installed on my Mac. Is there any workaround to get this work without Triton? I see it's almost impossible to get the triton installed on the Mac.

Any assistance would be highly appreciated.

Regards,
Riyaz

[derwaldgeist]

Same here. (I also got an error because the setup was not compatible with the current version of hugging face hub.)

[anthonypero]

+1 for macOS/Apple Silicon support. M2 Ultra here, hit the same triton/bitsandbytes wall. Would love an MPS fallback path — happy to beta test.

[anthonypero]

Here is a summary of the changes required to get HeartMuLa running on macOS (Apple Silicon) using the MPs backend.

We successfully generated 10s of audio on a Mac Studio (M2 Ultra) with these patches.

1. Verification

• Hardware: Mac Studio (M2 Ultra)
• OS: macOS 14+
• Python: 3.10
• PyTorch: 2.4.1 (MPS backend)

2. Changes Required

A. Remove bitsandbytes dependency

In pyproject.toml, remove bitsandbytes. It is listed but not actually used in the codebase, and it prevents installation on non-CUDA platforms.

-    "bitsandbytes==0.49.0",

B. Fix Autocast for MPS

In src/heartlib/pipelines/music_generation.py, torch.autocast throws a RuntimeError on MPS. We patched it to use nullcontext when the device is MPS.

Don't forget to import contextlib at the top of the file.

# src/heartlib/pipelines/music_generation.py

# ... inside _forward method ...
        self.mula.setup_caches(bs_size)
        
+       # specific patch for MPS: autocast not supported
+       if self.mula_device.type == "mps":
+           autocast_ctx = contextlib.nullcontext()
+       else:
+           autocast_ctx = torch.autocast(device_type=self.mula_device.type, dtype=self.mula_dtype)

-       with torch.autocast(device_type=self.mula_device.type, dtype=self.mula_dtype):
+       with autocast_ctx:
            curr_token = self.mula.generate_frame(...)

(Same patch applied to the second generation loop in the file)

C. Fix Tensor Type Conversion

In src/heartlib/heartcodec/models/transformer.py, the legacy .type() method with string arguments causes a ValueError on MPS tensors. Replaced with .to().

# src/heartlib/heartcodec/models/transformer.py

    def timestep_embedding(self, timesteps, max_period=10000, scale=1000):
        # ...
        freqs = torch.exp(
            -math.log(max_period)
            * torch.arange(start=0, end=half, device=timesteps.device)
            / half
-       ).type(timesteps.type())
+       ).to(dtype=timesteps.dtype)

3. Usage

Run with mps device and float16 (bfloat16 support on MPS can be spotty depending on PyTorch version):

python examples/run_music_generation.py \
    --model_path ./ckpt \
    --version "3B" \
    --mula_device mps \
    --codec_device mps \
    --mula_dtype float16

[surak]

Here is a summary of the changes required to get HeartMuLa running on macOS (Apple Silicon) using the MPs backend.

@anthonypero would this break functionality on other platforms? There's this already: #23

[qwertie]

Worth noting that people on Windows also encounter this error, since the instructions don't say to install triton, or how, and then there's an error importing triton (import error: No module named 'triton'), so I installed triton-2.1.0-cp311-cp311-win_amd64.whl, and it didn't work:

G:\AI stuff\Music\HeartMuLa\heartlib>py -3.11 ./examples/run_music_generation.py --model_path=./ckpt --version="3B"
Traceback (most recent call last):
  File "G:\AI stuff\Music\HeartMuLa\heartlib\examples\run_music_generation.py", line 1, in <module>
    from heartlib import HeartMuLaGenPipeline
  File "G:\AI stuff\Music\HeartMuLa\heartlib\src\heartlib\__init__.py", line 1, in <module>
    from .pipelines.music_generation import HeartMuLaGenPipeline
  File "G:\AI stuff\Music\HeartMuLa\heartlib\src\heartlib\pipelines\music_generation.py", line 2, in <module>
    from ..heartmula.modeling_heartmula import HeartMuLa
  File "G:\AI stuff\Music\HeartMuLa\heartlib\src\heartlib\heartmula\modeling_heartmula.py", line 4, in <module>
    from transformers.modeling_utils import PreTrainedModel
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\transformers\modeling_utils.py", line 57, in <module>
    from .integrations.flex_attention import flex_attention_forward
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\transformers\integrations\flex_attention.py", line 46, in <module>
    class WrappedFlexAttention:
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\transformers\integrations\flex_attention.py", line 61, in WrappedFlexAttention
    @torch.compiler.disable(recursive=False)
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\compiler\__init__.py", line 256, in disable
    import torch._dynamo
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\__init__.py", line 13, in <module>
    from . import (
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\aot_compile.py", line 14, in <module>
    from torch._dynamo.convert_frame import GraphRuntimeEnv
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\convert_frame.py", line 60, in <module>
    from torch._dynamo.symbolic_convert import TensorifyState
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\symbolic_convert.py", line 53, in <module>
    from torch._dynamo.exc import ObservedException, TensorifyScalarRestartAnalysis
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\exc.py", line 44, in <module>
    from .utils import counters
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\_dynamo\utils.py", line 2500, in <module>
    if has_triton_package():
       ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\torch\utils\_triton.py", line 9, in has_triton_package
    import triton  # noqa: F401
    ^^^^^^^^^^^^^
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\__init__.py", line 8, in <module>
    from .runtime import (
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\runtime\__init__.py", line 1, in <module>
    from .autotuner import (Autotuner, Config, Heuristics, OutOfResources, autotune, heuristics)
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\runtime\autotuner.py", line 7, in <module>
    from ..testing import do_bench
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\testing.py", line 7, in <module>
    from . import language as tl
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\language\__init__.py", line 6, in <module>
    from .standard import (
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\language\standard.py", line 3, in <module>
    from ..runtime.jit import jit
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\runtime\jit.py", line 10, in <module>
    from ..runtime.driver import driver
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\runtime\driver.py", line 1, in <module>
    from ..backends import backends
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\backends\__init__.py", line 50, in <module>
    backends = _discover_backends()
               ^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\backends\__init__.py", line 44, in _discover_backends
    driver = _load_module(name, os.path.join(root, name, 'driver.py'))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\backends\__init__.py", line 12, in _load_module
    spec.loader.exec_module(module)
  File "C:\Users\Family\AppData\Roaming\Python\Python311\site-packages\triton\backends\nvidia\driver.py", line 18, in <module>
    library_dir += [os.path.join(os.environ.get("CUDA_PATH"), "lib", "x64")]
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen ntpath>", line 108, in join
TypeError: expected str, bytes or os.PathLike object, not NoneType

(Note: Rumor has it that this project is incompatible with python 3.12+, so I've been using py -3.11 instead of python)

My AI suggested uninstalling that triton (py -3.11 -m pip uninstall -y triton) and getting a CUDA version specifically (NVIDIA GPU required):

py -3.11 -m pip uninstall -y torch torchvision torchaudio
py -3.11 -m pip install --index-url https://download.pytorch.org/whl/cu128 torch torchvision torchaudio

After doing that, the import error: No module named 'triton' still appears, but then it actually starts. And it looks like it will take two or three hours to generate a song. I suppose there is some expected amount of VRAM that I don't have.