MPS: Cannot add LoRA to Unet (LoftQ)

[bghira]

System Info

peft v0.9.0
accelerate v0.26.1

apple m3 max

Who can help?

@sayakpaul

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder
• My own task or dataset (give details below)

Reproduction

@torch.no_grad()
def loftq_init(weight: Union[torch.Tensor, torch.nn.Parameter], num_bits: int, reduced_rank: int, num_iter=1):
    if num_bits not in [2, 4, 8]:
        raise ValueError("Only support 2, 4, 8 bits quantization")
    if num_iter <= 0:
        raise ValueError("Number of iterations must be greater than 0")

    out_feature, in_feature = weight.size()
    device = weight.device
    dtype = weight.dtype

    logging.info(
        f"Weight: ({out_feature}, {in_feature}) | Rank: {reduced_rank} "
        f"| Num Iter: {num_iter} | Num Bits: {num_bits}"
    )
    if not is_bnb_4bit_available() or num_bits in [2, 8]:
        quantizer = NFQuantizer(num_bits=num_bits, device=device, method="normal", block_size=64)
        compute_device = device
    else:
        compute_device = "cuda"

Expected behavior

The above code defaults to "cuda", but it should consider whether mps is available first, instead.

[bghira]

actually i'm not sure even if it works with MPS, once that's set up due to bitsandbytes-foundation/bitsandbytes#252

maybe just blocking the usage on mps systems is needed until the plumbing is there for it

[sayakpaul]

Hmm bitsandbytes-foundation/bitsandbytes#252 is probably the reason why we default to cuda here. @younesbelkada can share more.

[BenjaminBossan]

I don't have a MB to test this but it looks indeed like it wouldn't work. Note that you can download some LoftQ-initialized models from the hub if you don't want to run this process yourself, but of course using them still requires bnb to work.

In general, I think it's a good reminder for us to be more careful about hard-coding the device in our code, as the number of viable hardware accelerators continues to climb.

[bghira]

for simpletuner my workaround is to use gaussian init instead of loftq on mps boxes. obviously it would be nice to have feature parity across the board, but we just ain't there yet

[younesbelkada]

Hi @bghira
Thanks for the issue! Indeed bnb is not compatible with MPS devices yet, LoftQ requires bnb so this won't work.
You can subscribe to: bitsandbytes-foundation/bitsandbytes#997 for more details about the ongoing progress with respect to bnb + apple sillicon integration

[bghira]

kinda sad to put a dependency on such a non-portable library, makes the feature nvidia-only.