Excessive Memory Allocation (56.98GB) on M2 Mac Mini (32GB) Using Qwen2-VL-2B-Instruct-4bit

[WildKratts]

Issue Summary:
When running a single message + image test pair with the specified setup, the process attempts to allocate 56.98GB of memory, far exceeding the available memory on the machine. This leads to a crash with the following error:

libc++abi: terminating due to uncaught exception of type std::runtime_error: Attempting to allocate 56976077312 bytes which is greater than the maximum allowed buffer size of 17179869184 bytes.
zsh: abort python3 -m mlx_vlm.lora --model-path --epochs 2

Environment:
• Machine: M2 Mac Mini (32GB RAM)
• Model: Qwen2-VL-2B-Instruct-4bit
• Framework/Version: Python 3.12

[Blaizzy]

Hey @WildKratts

Could you share a reproducible example?

I know the current trainer has a few bugs to iron out.