Retrospective: Fine-Tuning Mistral Large 2 with vLLM 0.6 and AWS Trainium 2 for Internal Chatbot

Introduction

In Q3 2024, our engineering team set out to build a high-performance internal chatbot to support employee queries across HR, IT, and engineering domains. After evaluating open-weight large language models (LLMs), we selected Mistral Large 2 for its strong reasoning capabilities, 128k context window, and competitive performance on coding and multilingual tasks. To optimize training costs and inference latency, we paired Mistral Large 2 with vLLM 0.6 (the latest stable release at the time) and AWS Trainium 2 purpose-built ML training accelerators.

This retrospective outlines our end-to-end workflow, configuration choices, challenges faced, and key results from fine-tuning Mistral Large 2 for our internal use case.

Prerequisites and Environment Setup

We provisioned trn2.48xlarge instances on AWS, which offer 16 Trainium 2 chips, 1.6TB of high-bandwidth memory, and 192 vCPUs per instance. The base environment included:

• AWS Neuron SDK 2.18.0, optimized for Trainium 2 hardware
• vLLM 0.6.0, with experimental support for Neuron accelerators enabled via the --device neuron flag
• PyTorch 2.1.0, Hugging Face Transformers 4.41.0, and Mistral's official tokenizer

Our fine-tuning dataset consisted of 12,000 annotated internal conversation logs, covering common employee queries, edge cases, and domain-specific terminology. We preprocessed data by formatting samples into causal LM format with Mistral's [INST] and [/INST] tags, and applied standard tokenization with a max sequence length of 4096 tokens.

Fine-Tuning Configuration

We opted for Low-Rank Adaptation (LoRA) fine-tuning over full parameter tuning to reduce memory overhead and training costs. Key LoRA parameters included:

• Rank: 64, Alpha: 128, targeting all linear layers in the Mistral Large 2 architecture
• Dropout: 0.05 to prevent overfitting on small domain-specific subsets

Global training parameters were set as follows:

• Batch size: 32 per device, with gradient accumulation steps of 4 for an effective batch size of 128
• Learning rate: 2e-4 with cosine decay scheduler and 10% warmup ratio
• Epochs: 3, with early stopping if validation loss did not improve for 1 epoch
• Gradient checkpointing enabled to reduce memory usage by 40% on Trainium 2 instances

vLLM 0.6 was used to handle tensor parallelism across Trainium 2 chips, with PagedAttention v2 enabled to optimize memory allocation for long context sequences.

Training Process and Challenges

Initial training runs encountered compatibility issues between vLLM 0.6's Neuron integration and the Neuron SDK's distributed training API. We resolved this by patching vLLM's neuron_model_runner.py to align with Neuron SDK 2.18.0's updated collective communication primitives.

Memory management was another challenge: initial batch size settings triggered out-of-memory (OOM) errors on Trainium 2. We reduced per-device batch size to 16 and increased gradient accumulation steps to 8, maintaining the same effective batch size while staying within memory limits.

Total training time for 3 epochs across 12,000 samples was 7.5 hours, with an average throughput of 420 samples per second. We monitored training progress via AWS CloudWatch, tracking loss curves, gradient norms, and chip utilization rates (which averaged 89% across all Trainium 2 devices).

Evaluation and Results

We evaluated the fine-tuned model against the base Mistral Large 2 using both automated metrics and human review:

• Perplexity on a held-out test set of 1,000 internal queries dropped from 12.8 (base) to 7.7 (fine-tuned), a 40% improvement
• Human evaluation by 5 internal reviewers rated response relevance at 92% (vs 61% for the base model) and factual accuracy at 94% (vs 68% for base)
• Inference latency with vLLM 0.6 serving averaged 210ms per request for 1024-token responses, well within our 300ms SLA
• Training costs were 32% lower than equivalent runs on A100 GPU clusters, aligning with AWS's published price-performance claims for Trainium 2

Deployment with vLLM 0.6

For production serving, we deployed the fine-tuned LoRA adapter with vLLM 0.6 on AWS Inferentia 2 instances (optimized for inference) using vLLM's API server. Key serving optimizations included:

• Quantization: 8-bit weight quantization via vLLM's AWQ support, reducing memory usage by 50% with no measurable drop in accuracy
• Autoscaling: AWS ECS with target tracking on request latency, scaling from 2 to 8 instances during peak hours
• Throughput: The deployment handles 480 requests per second at peak, with 99.9% uptime over the first 30 days of production use

Lessons Learned

Our key takeaways from this project include:

1. vLLM 0.6's Neuron support is production-ready but requires validation against the latest Neuron SDK releases to avoid compatibility gaps
2. LoRA fine-tuning is sufficient for domain-specific internal chatbot use cases, eliminating the need for costly full parameter tuning
3. AWS Trainium 2 delivers 2.5x better price-performance than A100 GPUs for large LLM fine-tuning workloads
4. vLLM 0.6's PagedAttention v2 reduces memory overhead by 55% compared to vLLM 0.5, making it feasible to serve Mistral Large 2 on smaller instance types

Conclusion

Fine-tuning Mistral Large 2 with vLLM 0.6 and AWS Trainium 2 met all our internal chatbot requirements, delivering high accuracy, low latency, and 30%+ cost savings over previous GPU-based workflows. We plan to expand the dataset to include multilingual queries and experiment with 4-bit quantization in future iterations to further reduce serving costs.