I've been experimenting with different optimizer setups for training diffusion models over the last few days.
The original motivation was pretty simple: reduce optimizer memory usage enough so I could fit larger diffusion models on a single GPU.
The Problem with Existing Optimizers
While working on this, I ran into specific limitations with the usual suspects:
AdamW8bit: Works great, but the optimizer state memory is still a huge bottleneck for larger models.
Lion8bit: Efficient, but not enough, and in my experience with small-batch setups, the training often becomes unstable and just blows up the model after a while.
Standard Adafactor: Very memory efficient, but existing implementations (PyTorch / HuggingFace) have rigid scheduling behaviors. PyTorch couples the internal LR schedule to training steps with no option to disable it. HuggingFace allows LR decoupling, but the second-moment decay still cannot be turned off. In long continuous training or fine-tuning, this leads to a "blunting" of the second-order moment estimate, weakening the optimizer's ability to adapt to new gradient distributions.
The Solution
To explore these issues further, I built an 8-bit Adafactor variant with:
Fully configurable scheduling and optional fixed beta2 behavior.
8-bit optimizer state representation.
CUDA-fused kernel implementation for performance.
Optional APOLLO-style low-rank subspace projection with Fira limiter to speed up convergence and stabilize updates.
Current Status & Experience
At the moment, this has largely become my default optimizer for diffusion model training, since it's much lighter than AdamW8bit and more stable than Lion8bit in my case.
A few interesting observations from my recent runs:
Faster Convergence: Enabling APOLLO-style projection seems to improve convergence speed and generalization compared to the standard Adafactor path.
Occasionally, changing projection subspaces can introduce small gradient spikes, but the built-in Fira Limiter for the APOLLO path handles it well, and I can relax the external gradient clipping.
About the Memory Usage
Here are the optimizer checkpoint sizes from my setup (not a full benchmark yet, but gives a good idea):
~1.7B Diffusion UNet:
AdamW8bit (bitsandbytes): ~3.22 GB
Adafactor (FP32): ~734 MB
Adafactor8Bit: ~188 MB (approx. 17× optimizer state reduction vs AdamW8bit.)
Text Encoder (Custom CLIP, mostly transformer weights):
AdamW8bit: ~1.04 GB
Adafactor (FP32): ~2.48 MB
Adafactor8Bit: ~1.75 MB (Standard Adafactor is already extremely memory efficient here, so quantization gains are smaller).
Getting Started:
pip install -U adafactor8bitRecommended Usage (v0.2.1): Hybrid Routing & Fira Limiter
For most diffusion training setups, I currently recommend a hybrid routing strategy:
Embeddings → Momentum-free Adam style (
factored=False,scale_parameter=False,d=1e9)
Using element-wise variance scaling in log-space provides fine-grained, per-token updates. This avoids the “cold-start” over-scaling issue of standard Adafactor when an embedding row is activated for the first time in a while, without the overhead of APOLLO projection. Pair this with an Adam-style learning rate (e.g., 1e-4) for best results.2D weights (Linear layers) → APOLLO
In my experience, APOLLO generally converged faster and showed better generalization than the standard Adafactor path, while keeping memory usage similarly low.Convolutions and other >2D tensors → Full-Rank (
factored=False)
For finer gradient scaling, we can disable row/column factorization to keep the native spatial structure intact, maintaining independent variance for each spatial position in the convolution kernel.Norms and biases → FP32 Adafactor, no weight decay
The standard recipe for stable training.*New in v0.2.1: Optional 4-bit Packed Momentum
v0.2.1 adds optional 4-bit packed first-moment (beta1) support, allowing momentum to be enabled selectively for parameter groups where it provides the most benefit, while keeping the additional optimizer memory very small. In my current configuration, I enablebeta1only for dense weight matrices (Linear / Conv), while embeddings, norms, and biases remain momentum-free.
You can adjust apollo_rank based on memory budget:
0— Disable APOLLO and use the standard Adafactor path.16— Default used by LLaMA-Factory.256— Recommended by the official APOLLO repository for 1B–7B models.
Enabling the Fira Limiter on the Adafactor paths suppresses gradient spikes, often making external clip_grad_norm_() unnecessary.
# Define learning rates
lr = 1e-3
lr_emb = 1e-4 # For Embedding layers, we use an Adam-style learning rate
def get_param_groups(model, lr_emb, weight_decay, apollo_rank=256):
group_1d, group_embed, group_2d, group_nd = [], [], [], []
for name, param in model.named_parameters():
if not param.requires_grad: continue
is_1d = param.ndim <= 1 or "bias" in name or "norm" in name
# Match true Token Embeddings, excluding Position and Time Embeddings
is_embedding = ("embed" in name.lower()
and "position" not in name.lower()
and "pos_embed" not in name.lower()
and "time" not in name.lower())
if is_1d:
group_1d.append(param)
elif is_embedding:
group_embed.append(param)
elif param.ndim == 2:
group_2d.append(param)
else:
group_nd.append(param)
return [
# 1. 1D / Sensitive: FP32, No Weight Decay
{"params": group_1d, "weight_decay": 0.0, "quantize": False, "apollo_rank": 0},
# 2. Embeddings: Recreating a momentum-free Adam
{
"params": group_embed,
"weight_decay": 0.0,
"quantize": False,
"apollo_rank": 0,
"factored": False, # Enable element-wise variance
"scale_parameter": False, # Disable internal RMS scaling
"d": 1e9, # Disable global Trust-Region clipping
"lr": lr_emb # Override global learning rate
},
# 3. 2D Weights: 8-bit quantization, Weight Decay, APOLLO low-rank projection
{
"params": group_2d,
"weight_decay": weight_decay,
"quantize": True,
"apollo_rank": apollo_rank,
"beta1":0.9, # Remove if minimizing optimizer memory is the priority.
},
# 4. >2D Weights: 8-bit quantization, Weight Decay, Full-Rank
{
"params": group_nd,
"weight_decay": weight_decay,
"quantize": True,
"apollo_rank": 0,
"beta1":0.9, # Remove if minimizing optimizer memory is the priority.
"factored": False # Disables factorization to preserve spatial structures, enabling finer gradient scaling.
# Note: This increases state memory for >2D weights, depending on your model architecture.
# If VRAM is constrained, reverting to factored=True is a safe alternative.
},
]
model = MyModel().cuda()
optimizer = Adafactor8Bit(
get_param_groups(model, lr_emb = lr_emb, weight_decay=1e-2, apollo_rank=256),
lr=lr,
# For continual learning or when using an external LR scheduler
relative_step=False, # Disable internal LR scheduling
beta2=0.999, # Lock EMA window to prevent "blunting" over steps
enable_fira_for_adafactor=True # Enable Fira Limiter globally; external grad clipping can be safely removed
)
# With Fira enabled, torch.nn.utils.clip_grad_norm_() in your training loop can usually be removed
This is the configuration I currently use for most of my own diffusion training. If you're already using adafactor8bit>=0.2.1, feel free to use it as a starting point and adapt the routing to your own model architecture and memory budget.
GitHub [Documentation and source code]:
8-bit Adafactor Optimizer with Fused CUDA Kernels
Feel free to try it out and let me know how it behaves in your training! If this helps your training, a GitHub Star would be hugely appreciated! :)