ml-explore / ml-explore/mlx-examples
Huge memory usage when finetuning
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 9k
- Forks
- 1.2k
- PR merge metrics
- No merged PRs in 30d
Description
When fine-tuning Mistral 7B in 4-bit quantization (qlora), I'm seeing huge memory usage (160GB VRAM)
Parameters used:
--batch-size 1--lora-layers 16
The dataset is composed of around 1200 entries. No entry is longer than 7500 tokens, some are significantly shorter, most are right in the middle.
If I reduce --lora-layers to 4, the memory usage peaks at around 30GB, but the end result is very poor quality.
Is this sort of memory usage expected?
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing fine-tuning with Mistral 7B in 4-bit QLoRA using batch size 1 and 16 LoRA layers, then compare memory usage with 4 layers. Trace the fine-tuning entry point and configuration handling to identify what scales with sequence length and LoRA layer count; done means explaining whether the 160GB peak is expected or isolating a reproducible memory problem.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100