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