Quantized and Memory-Efficient Training
How to reduce LLM fine-tuning memory with GraphOptimizer, QLoRA, mixed precision, FP8, Adam8bit, gradient accumulation, and gradient checkpointing.
Large model fine-tuning is usually memory-bound before it is compute-bound. Deeplearning4j 1.0.0-rewrite adds several SameDiff tools for reducing memory pressure while keeping the standard training loop: graph optimization, mixed precision, FP8 configuration, 8-bit optimizer state, gradient accumulation, gradient checkpointing, QLoRA, and post-training quantization.
The main runnable example is GraphOptimizerQuantizedTrainingExample.java. It demonstrates 11 scenarios from post-training quantization through a full GraphOptimizer -> Adam8bit -> DSP training workflow.
Choose a Memory Strategy
Base model weights do not fit comfortably
QLoRA
Quantized base plus trainable LoRA adapters
Optimizer state is too large
Adam8bit
Adam moment buffers use INT8 block quantization
Batch size is too small
Gradient accumulation
Preserves effective batch size across micro-batches
Activations dominate memory
Gradient checkpointing
Recomputes selected activations during backward pass
GPU supports FP8 tensor cores
FP8TrainingConfig
Lowers matmul activation/weight precision where eligible
Inference deployment needs smaller weights
Post-training quantization
Casts trained weights for deployment after training
Training graph has unnecessary ops
GraphOptimizer
Removes dead code and applies graph-level simplifications
These methods compose. A practical large-model recipe is QLoRA or LoRA adapters, BF16 compute, Adam8bit optimizer state, gradient accumulation, checkpointing, and graph optimization before steady-state training.
Graph Optimization Before Training
GraphOptimizer can simplify a SameDiff graph before training. When optimizing a training graph, keep both the model output and the loss as required outputs. Otherwise dead-code elimination can remove the loss branch.
List<OptimizerSet> passes = GraphOptimizer.defaultOptimizations();
SameDiff optimized = GraphOptimizer.optimize(rawSd, "output", "loss");
optimized.setLossVariables("loss");
optimized.setTrainingConfig(TrainingConfig.builder()
.updater(Adam8bit.builder()
.learningRate(1e-3)
.blockSize(2048)
.build())
.dataSetFeatureMapping("input")
.dataSetLabelMapping("labels")
.build());For LLM fine-tuning, required outputs usually include logits and the scalar loss. Keep the loss reachable from the optimized graph before calling fit().
QLoRA for Base-Weight Memory
QLoRA is the first choice when the base model is large but you only need adapter updates. It keeps the base in a quantized representation and trains LoRA matrices:
Use QLoRA for instruction tuning and preference alignment when full fine-tuning does not fit. Use full fine-tuning only when you need to update all weights and have enough memory for weights, gradients, activations, and optimizer state.
Adam8bit for Optimizer State
Adam stores two running buffers per parameter. For large models, optimizer state can exceed the size of the model weights. Adam8bit stores those buffers with block-wise INT8 quantization and per-block scale factors.
Use block sizes large enough for efficient metadata overhead but small enough to isolate outliers. The examples use 2048 elements per block.
Mixed Precision and Loss Scaling
Mixed precision keeps numerically sensitive state in FP32 while allowing forward/backward compute to use lower precision where safe. For LLM fine-tuning, the usual pattern is:
Keep master weights in FP32.
Use BF16 or FP16 compute where supported.
Keep layer norm, softmax, and loss reductions in higher precision.
Use loss scaling for FP16 to avoid gradient underflow.
Use BF16 when available because it has a wider exponent range than FP16 and usually needs less manual loss-scale tuning.
FP8 Training Config
FP8 is useful for eligible matmul/linear operations on supported NVIDIA hardware. It should not be applied blindly to every op.
The example checks op eligibility. Matrix multiplication and linear projections are the main beneficiaries. Softmax, normalization, GELU, and loss reductions generally need higher precision.
Gradient Accumulation
Gradient accumulation lets you use a small micro-batch while preserving a larger effective batch size:
Use it when the model fits but the desired batch size does not. Accumulate in FP32 when using low-precision compute so repeated additions do not lose too much precision.
For SFT, SFTConfig exposes this directly:
Gradient Checkpointing
Gradient checkpointing trades compute for activation memory. Instead of storing every activation from the forward pass, selected activations are recomputed during backpropagation.
The examples cover these strategies:
Every-N checkpointing
You want explicit control over layer intervals
Sqrt-N checkpointing
You want a balanced default as depth grows
Async offload
You want to move checkpoints through host memory with prefetching
Checkpointing composes well with Adam8bit and QLoRA: adapters reduce trainable parameter memory, Adam8bit reduces optimizer state, and checkpointing reduces activations.
Full Optimized Training Path
The complete production-style path in the example is:
Build or import the SameDiff graph.
Optimize the graph with required outputs including loss.
Set the loss variables explicitly.
Configure Adam8bit or another memory-efficient updater.
Train through
SameDiff.fit().Let DSP compile the steady-state training graph.
Apply deployment quantization only after training if the graph needs lower-precision inference weights.
This matters because training and deployment quantization are different. Some FP16 or INT8 weight casts are deployment optimizations; applying them before fit() can create dtype combinations the backward pass should not use.
Post-Training Quantization
Post-training quantization is for deployment. Train normally, validate the model, then quantize weights for inference. It is the simplest path when accuracy loss is acceptable and retraining is expensive.
Use quantization-aware training when post-training quantization causes too much accuracy loss. QAT inserts fake-quantization behavior during training so the model learns to tolerate deployment precision.
Memory Checklist
Before changing model architecture, check these in order:
Are you training all weights when LoRA or QLoRA would be enough?
Is optimizer state larger than the weights? Use
Adam8bit.Is activation memory the peak? Use gradient checkpointing.
Is micro-batch too small? Add gradient accumulation.
Are long examples wasting padding? Use bucketing or packing.
Are teacher logits dense? Store FP16 or top-K targets.
Are you optimizing a training graph with the loss retained?
Are deployment quantization passes being applied only after training?
Run the Example
Related Pages
Last updated
Was this helpful?