> For the complete documentation index, see [llms.txt](https://deeplearning4j.konduit.ai/llms.txt). Markdown versions of documentation pages are available by appending `.md` to page URLs; this page is available as [Markdown](https://deeplearning4j.konduit.ai/en-1.0.0-rewrite/deeplearning4j/overview-4/quantized-training.md).

# 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

| Constraint                                 | Use                        | Why                                                       |
| ------------------------------------------ | -------------------------- | --------------------------------------------------------- |
| 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.

```java
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:

```java
QLoraConfig qloraConfig = QLoraConfig.builder()
    .r(16)
    .loraAlpha(32)
    .loraDropout(0.05)
    .bits(4)
    .quantType("nf4")
    .doubleQuant(true)
    .computeDataType(DataType.BFLOAT16)
    .loraDataType(DataType.BFLOAT16)
    .blockSize(64)
    .targetModules(targetModules)
    .taskType(TaskType.CAUSAL_LM)
    .build();

PeftModel model = PeftModel.fromPretrained(baseModel, qloraConfig);
```

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.

```java
Adam8bit optimizer = Adam8bit.builder()
    .learningRate(1e-3)
    .beta1(0.9)
    .beta2(0.999)
    .epsilon(1e-8)
    .blockSize(2048)
    .pagedOptimizer(false)
    .build();

sd.setTrainingConfig(TrainingConfig.builder()
    .updater(optimizer)
    .dataSetFeatureMapping("input")
    .dataSetLabelMapping("labels")
    .build());
```

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.

```java
FP8TrainingConfig fp8Config = FP8TrainingConfig.builder()
    .useE4M3ForForward(true)
    .perTensorScaling(true)
    .amaxHistoryLength(16)
    .build();

fp8Config.validate();
FP8ScaleManager scaleManager = new FP8ScaleManager(fp8Config);
```

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:

```
effective batch = micro batch * gradientAccumulationSteps
```

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:

```java
SFTConfig config = SFTConfig.builder()
    .gradientAccumulationSteps(8)
    .computeDataType(DataType.BFLOAT16)
    .build();
```

## 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:

| Strategy              | Use when                                                          |
| --------------------- | ----------------------------------------------------------------- |
| 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:

1. Build or import the SameDiff graph.
2. Optimize the graph with required outputs including loss.
3. Set the loss variables explicitly.
4. Configure Adam8bit or another memory-efficient updater.
5. Train through `SameDiff.fit()`.
6. Let DSP compile the steady-state training graph.
7. 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:

1. Are you training all weights when LoRA or QLoRA would be enough?
2. Is optimizer state larger than the weights? Use `Adam8bit`.
3. Is activation memory the peak? Use gradient checkpointing.
4. Is micro-batch too small? Add gradient accumulation.
5. Are long examples wasting padding? Use bucketing or packing.
6. Are teacher logits dense? Store FP16 or top-K targets.
7. Are you optimizing a training graph with the loss retained?
8. Are deployment quantization passes being applied only after training?

## Run the Example

```bash
cd ../deeplearning4j-examples/samediff-examples
mvn exec:java -Dexec.mainClass="org.nd4j.examples.samediff.quickstart.training.GraphOptimizerQuantizedTrainingExample"
```

## Related Pages

* [LLM Fine-Tuning Overview](/en-1.0.0-rewrite/deeplearning4j/overview-4/fine-tuning.md)
* [Instruction SFT, LoRA, and QLoRA](/en-1.0.0-rewrite/deeplearning4j/overview-4/instruction-sft-lora.md)
* [Knowledge Distillation](/en-1.0.0-rewrite/deeplearning4j/overview-4/knowledge-distillation.md)
* [Preference and RL Alignment](/en-1.0.0-rewrite/deeplearning4j/overview-4/preference-rl-alignment.md)
* [PEFT and RL Alignment Training](/en-1.0.0-rewrite/deeplearning4j/peft-and-rl.md)