Enable distributed LoRA training

[angeloskath]

The updates to LORA.md are missing but TL;DR we can now do

$ echo "m2-ultra-0 slots=1" >>hostfile
$ echo "m2-ultra-1 slots=1" >>hostfile
$ mpirun --hostfile hostfile -- python -m mlx_lm.lora --train --model mlx-community/Mistral-7B-v0.2-4bit --data /path/to/data --batch-size 16

to train across two nodes (or more really nothing needs to change).

[mzbac]

Is that possible to do distributed inference as well?

[awni]

Is that possible to do distributed inference as well?

Possible yes, but getting a nice speedup is more challenging. That's something we're looking at, but don't have an ETA on right now.

[angeloskath]

@awni feel free to review and then we can merge. I split the launcher to a different branch.

[awni]

Just curious, why do we need this for the test to work?

[awni]

Nvm I see the magicmock thing is messing up the all_sum.

[awni]

Probably tha tshould go under the rank==0 condition as well

[angeloskath]

Yeah ok, it makes sense. I was thinking in general callbacks should always run and if the callback is about reporting then it can choose to only run on rank=0. But our callbacks here are only about reporting so it makes sense to just run them only on node 0.

[awni]

It's a good point actually. It's more flexible that way. I'm on board leaving it to the user to specify the rank.

[awni]

Same there.

[awni]

Looks great!! Let's 🚢

[ivanfioravanti]

This works perfectly! Great job 👏