Long Training Times and GPU Utilization Problem

[meirm]

Issue: Long Training Times and GPU Utilization Problem

Description:

Hi everyone,

I am currently working on fine-tuning an LLM using the MLX library on my MacBook Pro M1 with 16GB of RAM. While the training process works, it is taking an exceptionally long time to complete. From what I observe, the training is not utilizing the GPU and is instead maxing out the Performance-CPUs at 100%.

Here are the key details of my setup and issue:

• Machine: MacBook Pro M1 with 16GB RAM
• Library: MLX
• Model: TinyLlama-1.1B-Chat-v1.0-4bit
• Training Command:

  python -m mlx_lm.lora --model mlx-community/TinyLlama-1.1B-Chat-v1.0-4bit \
    --data data --train --batch-size 2  --lora-layers -1 --iters 4000 \
    --learning-rate 0.00001 --grad-checkpoint  --seed 42 --save-every 200 \
    --steps-per-eval 20

Symptoms:

• Training process is very slow.
• Performance-CPUs are at 100% utilization.
• GPU does not appear to be utilized at all.

Questions:

1. Is there a way to ensure that the GPU is used for the training process?
2. Are there specific configurations or settings in the MLX library that I might be missing to enable GPU utilization?
3. Any tips or recommendations to speed up the training process on an M1 MacBook?

Additional Information:
I have verified that the necessary packages and dependencies are installed correctly. Any insights or guidance on how to resolve this issue would be greatly appreciated.

Thanks in advance for your help!

Best regards,
Meir

[awni]

It looks like you are using the GPU but the utilization is not great. You can double check with:

python -c "import mlx.core as mx; print(mx.default_device())"

Now why it's under-utilized is a harder question to answer. One thing that is very odd is the way your toks/sec are periodic: high, then low. Could you share some details about the dataset? Are there highly variable sequence lengths?