Author(s): Chew Loong Nian – AI ENGINEER

Originally published on Towards AI.

Unsloth vs Axolotl vs TRL: 87% of Your Fine-Tuning VRAM Goes to a Tensor You Never Wrote

I spent an afternoon pricing every byte of a LoRA fine-tuning step from first principles, and the result was not what I expected. For Llama 3.1 8B, 87.3% of the memory that scales with your sequence length is not the model, not the optimizer, and not the attention activations. It is the cross-entropy loss head — one tensor that exists for a few microseconds, produces a single scalar, and never appears in your training script.

After introducing the surprising memory culprit (the cross-entropy loss head), the article explains how that tensor scales with vocabulary size, showing that for common models the loss head can consume ~87–97% of the marginal per-token VRAM budget. It provides quick arithmetic and a 120-line calculator to estimate loss-head VRAM, then validates the approach against published benchmarks—finding close agreement and isolating why framework comparisons often mislead (speed gains are smaller than structural memory savings). The author then maps how Unsloth, Axolotl, TRL, and LLaMA-Factory reduce or avoid that loss-head cost via techniques like fused/streamed loss computation, chunking, gradient checkpointing activation offload, expert quantization, and kernel integrations. Finally, it gives practical guidance on which framework to choose based on GPU count and training type, concluding that most “framework wars” boil down to how each handles this one hidden tensor.

Read the full blog for free on Medium.

Join thousands of data leaders on the AI newsletter. Join over 80,000 subscribers and keep up to date with the latest developments in AI. From research to projects and ideas. If you are building an AI startup, an AI-related product, or a service, we invite you to consider becoming a sponsor.

Published via Towards AI