Add MPS support

[maximegmd]

Description

Supports basic training on Mac M series.

Motivation and Context

It partially solves Mac support.

How has this been tested?

Ran a train job with lora-mps.yml from start to finish.

[ehartford]

This is awesome! @maximegmd

[winglian]

I'm not sure what's up with packaging and CI. I've run into this issue with our CI before and I think I just gave up on using it at one point.

[maximegmd]

I can always "parse" version numbers with a quick regex but it's not ideal (thanks ChatGPT)

[winglian]

Thanks @maximegmd 🚀

[fakerybakery]

Thanks @maximegmd!! @winglian is MLX support still planned?

[maximegmd]

Thanks @maximegmd!! @winglian is MLX support still planned?

The issue with supporting MLX is that we will pretty much have to rewrite axolotl entirely to replace all torch specific code (all of it) to MLX, so no bitsandbytes, no xformers, no flash-attn, no transformers... At this point I think we are better off waiting for the transformers support of MPS to reach an acceptable state.

[flexchar]

Thanks for this, @maximegmd !!
May I know how long it took to fine-tune and what are your M specs?

[maximegmd]

Thanks for this, @maximegmd !! May I know how long it took to fine-tune and what are your M specs?

Will run it again tomorrow, because I was today years old when I discovered the performance power option 😅

M3 Max 16/40 64GB

[fakerybakery]

@maximegmd For installation, since flash-attn & deepspeed aren't supported on macOS, should I use pip install -e . instead of pip install -e '.[flash-attn,deepspeed]'? (And are you using accelerate?)

[maximegmd]

Right just install the base without any options.

I use accelerate:

accelerate launch -m axolotl.cli.train examples/tiny-llama/lora-mps.yml

This training takes around 56 minutes on my machine.

[fakerybakery]

Nice, thx! I'll try it out