Potential memory leak during Llama 3 8b model fine-tuning with LoRA about mlx-examples HOT 9 CLOSED

Seems to be related to gradient checkpointing but that's all I know so far.. turning off gradient checkpointing and the peak memory is constant for a fixed sequence length as expected.

from mlx-examples.

@mzbac, a bit off topic and please let me know if I'm breaking some rules so I'll delete the comment.

But can you share an example of your dataset used for finetuning llama3 instruct? I am a bit confued by its template.

Thanks

@Satyam7166-tech You can check the data here at https://huggingface.co/datasets/mzbac/function-calling-llama-3-format-v1.1/viewer and feel free to create a discussion there if you have any questions about the dataset.

from mlx-examples.

@awni I can confirm that the issue is related to gradient checkpointing. With gradient checkpointing disabled, I trained the model for 8000 iterations and the peak memory is consistent and works as expected.

from mlx-examples.

Just FYI, the leak should be fixed now with gradient checkpointing enabled.

from mlx-examples.

Seems to be related to gradient checkpointing but that's all I know so far.. turning off gradient checkpointing and the peak memory is constant for a fixed sequence length as expected.

After disabling gradient checkpointing, the peak memory has stabilized. I will continue training for a few thousand iterations to see if everything is running as expected. Thanks @awni 🚀

from mlx-examples.

Thanks for the fix!

from mlx-examples.

Very strange .. seems like there is a leak of some sort.

from mlx-examples.

and the crash logs:

Iter 7320: Train loss 0.216, Learning Rate 1.000e-06, It/sec 0.116, Tokens/sec 40.214, Trained Tokens 3893259, Peak mem 342.467 GB
libc++abi: terminating due to uncaught exception of type std::runtime_error: [malloc_or_wait] Unable to allocate 554319936 bytes.
zsh: abort      mlx_lm.lora --config lora_config.yaml

from mlx-examples.

@mzbac, a bit off topic and please let me know if I'm breaking some rules so I'll delete the comment.

But can you share an example of your dataset used for finetuning llama3 instruct? I am a bit confued by its template.

Thanks

from mlx-examples.