{"id":97917,"date":"2025-07-07T08:00:00","date_gmt":"2025-07-07T15:00:00","guid":{"rendered":""},"modified":"2026-03-05T09:59:59","modified_gmt":"2026-03-05T17:59:59","slug":"optimizing-memory-usage-in-large-language-models-fine-tuning-with-kaito-best-practices-from-phi-3","status":"publish","type":"post","link":"https:\/\/opensource.microsoft.com\/blog\/2025\/07\/07\/optimizing-memory-usage-in-large-language-models-fine-tuning-with-kaito-best-practices-from-phi-3\/","title":{"rendered":"Optimizing memory usage in large language models fine-tuning with KAITO: Best practices from Phi-3"},"content":{"rendered":"\n

Large Language Models (LLMs) have revolutionized AI, but fine-tuning these massive models remains a significant challenge\u2014especially for organizations with limited computing resources. To address this, the Cloud Native team at Azure is working to make AI on Kubernetes more cost-effective and approachable for a broader range of users. <\/p>\n\n\n\n

Fine-tuning is the process of adapting pre-trained LLMs to perform better on specific tasks by training them on specialized data\u2014whether that\u2019s domain-specific text (legal or medical), task-specific examples (summaries or chat transcripts), or business-specific content (internal documentation or customer interactions). This critical step allows organizations to significantly improve model accuracy, tailor outputs to their needs, and leverage powerful foundation models without the enormous cost of training from scratch. <\/p>\n\n\n\n

In this post, we share best practices based on insights from our experiments fine-tuning Microsoft’s Phi-3-mini-128k model using KAITO<\/strong><\/a>, a Cloud Native Computing Foundation (CNCF)-governed open source project that simplifies running AI workloads on Kubernetes and the associated managed add-on for Azure Kubernetes Service (AKS). You’ll learn strategies to help you fine-tune powerful LLMs even within reasonable hardware constraints, making advanced AI a realistic option for your organization. <\/p>\n\n\n\n

\n
Learn more about KAITO<\/a><\/div>\n<\/div>\n\n\n\n

Understanding the model <\/h2>\n\n\n\n

Understanding our model\u2019s characteristics is crucial for developing effective fine-tuning strategies, especially when working with limited computational resources.  <\/p>\n\n\n\n

Phi-3-mini-128k-instruct is a 3.8 billion parameter model (~14.6 GB)\u2014compact enough to run efficiently on most single- or dual-GPU setups, yet powerful enough to support a wide range of tasks. <\/p>\n\n\n\n

Microsoft provides both 4k and 128k context length variants of this model on Hugging Face<\/a>. These numbers refer to the maximum number of tokens the model can process at once\u2014essentially the amount of text the AI can “remember” in a single conversation. The 4k variant handles about 6-7 pages of text (~3,000 words), while the 128k version processes an entire novel’s worth of content simultaneously (~90,000 words).  <\/p>\n\n\n\n

However, longer context windows require significantly more memory and computational resources during fine-tuning, which is one of the key challenges we’ll address in this article.  <\/p>\n\n\n\n

The memory challenge in fine-tuning <\/h2>\n\n\n\n

During the fine-tuning process, memory becomes the primary bottleneck\u2014especially when training on longer sequences of tokens (larger chunks of text) or when working with limited hardware. <\/p>\n\n\n\n

Unlike inference\u2014which simply generates text\u2014fine-tuning must store additional data like computational graphs and gradients. As a result, inference memory grows roughly linearly <\/strong>with sequence length, while fine-tuning memory grows non-linearly and can quickly exhaust GPU resources. <\/p>\n\n\n\n

In practice, while inference might handle 32,000 tokens on a single GPU, fine-tuning the same model might be limited to just 2,000-4,000 tokens before running out of memory. This creates a fundamental challenge: even though Phi-3-mini-128k-instruct can process novel-length text during use, memory constraints during fine-tuning often force developers to work with much shorter sequences, limiting the model’s full potential. <\/p>\n\n\n\n

Key memory optimizations <\/h2>\n\n\n\n

To identify effective memory optimizations, we leveraged KAITO’s existing inference<\/a> and fine-tuning<\/a> services with added memory logging to conduct our tests. All experiments were conducted on an NVIDIA A100 (80GB) GPU using standard KAITO deployments (code available in our GitHub repository<\/a>). <\/p>\n\n\n\n

For those looking to reproduce these results, both our fine-tuning and inference services now expose a \/metrics<\/em> endpoint that allows you to track memory usage with these optimizations in your own environment. <\/p>\n\n\n\n

Through this approach, we identified several effective strategies for optimizing memory usage during fine-tuning: <\/p>\n\n\n\n

1. Precision format selection <\/h2>\n\n\n\n

Here\u2019s how memory usage compares across formats: <\/p>\n\n\n\n

Precision format<\/strong> <\/td>Memory usage<\/strong> <\/td><\/tr>
Float32 <\/td>15.73 GB <\/td><\/tr>
Float16\/BFloat16 <\/td>8.09 GB <\/td><\/tr>
8-bit Quantized (INT8) <\/td>4.64 GB <\/td><\/tr>
4-bit Quantized (INT4) <\/td>3.00 GB <\/td><\/tr><\/tbody><\/table><\/figure>\n\n\n\n

As shown, selecting the right precision format can be the difference between being able to fine-tune a model or not. The 80% memory savings from 4-bit quantization allows you to work with models that would otherwise require high-end GPUs or distributed setups, significantly reducing your cloud compute costs. <\/p>\n\n\n\n

In KAITO, you can configure precision using your tuning parameters ConfigMap like so: <\/p>\n\n\n

\nModelConfig:\u00a0\n\u00a0\u00a0torch_dtype: \"bfloat16\" # Options: float32, float16, bfloat16\u00a0\n\u00a0\nQuantizationConfig:\u00a0\n\u00a0 load_in_4bit: true\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0 # Enable 4-bit quantization\u00a0\n\u00a0 bnb_4bit_quant_type: \"nf4\"\u00a0\n\u00a0 bnb_4bit_compute_dtype: \"bfloat16\"\u00a0\n\u00a0 bnb_4bit_use_double_quant: true\u00a0\n<\/pre><\/div>\n\n\n
For more details on precision formats and parameters, see the KAITO QLoRA parameters template<\/a>. For full examples, visit the fine-tuning examples directory<\/a> to see how these settings are used in practice. <\/p>\n\n\n\n
2. LoRA vs. QLoRA <\/h2>\n\n\n\n
LoRA (<\/strong>Low-Rank Adaptation<\/strong>)<\/strong> fine-tunes models by freezing the original weights and inserting small trainable adapter layers. This drastically reduces compute by training only a tiny fraction of the model parameters. <\/p>\n\n\n\n
QLoRA (Quantized LoRA)<\/strong> takes this further by storing the frozen weights in 4-bit precision instead of 16-bit, reducing memory usage without sacrificing LoRA\u2019s core benefits. <\/p>\n\n\n\n
Memory efficiency<\/strong> 
Standard LoRA (16-bit) showed steep memory growth\u2014exceeding 80GB at ~3,500 tokens. QLoRA, in contrast, maintained a stable profile and reduced memory usage by ~75%, enabling fine-tuning with much longer sequences on the same hardware. <\/p>\n\n\n\n
This comes with a modest tradeoff in processing speed due to quantization overhead, but the benefits are clear\u2014especially for domains like legal or medical, where preserving long context improves model comprehension and accuracy. <\/p>\n\n\n\n
In short:<\/strong> LoRA is faster but memory-heavy. QLoRA is slower but far more memory-efficient. <\/p>\n\n\n\n
In KAITO, you can enable QLoRA using your tuning parameters ConfigMap like so: <\/p>\n\n\n
\nQuantizationConfig:\u00a0\n\u00a0 load_in_4bit: true\u00a0\n\u00a0 bnb_4bit_quant_type: \"nf4\"\u00a0\n\u00a0 bnb_4bit_compute_dtype: \"bfloat16\"\u00a0\n\u00a0 bnb_4bit_use_double_quant: true\u00a0\n<\/pre><\/div>\n\n\n
3. Batch size optimization <\/h2>\n\n\n\n
Contrary to intuition, increasing batch size can deliver several benefits: <\/p>\n\n\n\n