Adafactor8Bit: Memory-Efficient Optimizer with Adaptive Log-Space Quantization

Here we demonstrate a hybrid grouping strategy for complex hybrid architectures (e.g., Vision-Language Models, Diffusion UNets) to achieve stable and efficient training.

📌 The following strategies are applied:

Implementation:

from adafactor8bit import Adafactor8Bit
# Define learning rates
lr = 1e-3
lr_adam = 1e-4 # For Embedding and N-D layers, we use an Adam-style learning rate
apollo_targets = ["attn", "mlp"]
def get_param_groups(model, lr_adam, weight_decay, apollo_rank=256):
 group_1d, group_embed, group_lm_head, group_apollo, 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())
 is_lm_head = param.ndim == 2 and "lm_head" in name.lower()
 is_apollo_target = param.ndim == 2 and any(t in name for t in apollo_targets)
 
 if is_1d:
 group_1d.append(param)
 elif is_embedding:
 group_embed.append(param)
 elif is_lm_head:
 group_lm_head.append(param)
 elif is_apollo_target:
 group_apollo.append(param)
 elif param.ndim == 2:
 group_2d.append(param)
 else:
 group_nd.append(param)
 # Common configuration for Adam-style groups (element-wise, no RMS scaling, no clipping)
 adam_style = {
 "factored": False, # full-rank V
 "scale_parameter": False, # no parameter RMS scaling
 "d": 0.0, # no RMS clipping
 "beta3": None, # CAME requires factored=True
 "apollo_rank": 0,
 "lr": lr_adam, 
 }
 return [
 # 1. 1D / Sensitive: FP32, No Weight Decay
 {"params": group_1d, "weight_decay": 0.0, "quantize": False, "apollo_rank": 0},
 
 # 2. Embeddings: Momentum-free Adam-style scaling
 {
 "params": group_embed, 
 "weight_decay": 0.0, 
 "quantize": False,
 "beta1": None, # Momentum-free
 **adam_style,
 },
 # 3. LM Head: Adam-style with momentum (avoid factored distortion)
 {
 "params": group_lm_head, 
 "weight_decay": 0.0, 
 "quantize": True,
 "beta1": 0.9,
 **adam_style,
 },
 
 # 4. 2D Weights (APOLLO targets): 8-bit quantization, APOLLO low-rank projection
 {
 "params": group_apollo, 
 "weight_decay": weight_decay, 
 "quantize": True, 
 "apollo_rank": apollo_rank,
 "beta1": 0.9, # Remove if minimizing optimizer memory is the priority.
 },
 # 5. 2D Weights (Others): Default Adafactor path (factored), 8-bit quantization
 {
 "params": group_2d, 
 "weight_decay": weight_decay, 
 "quantize": True, 
 },
 
 # 6. >2D Weights: 8-bit quantization, Weight Decay, Full-Rank
 {
 "params": group_nd, 
 "weight_decay": weight_decay, 
 "quantize": True, 
 "beta1": 0.9,
 **adam_style,
 },
 ]
model = MyModel().cuda()
optimizer = Adafactor8Bit(
 get_param_groups(model, lr_adam=lr_adam, 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 for all non-APOLLO paths
)
# Training loop...

For more complete examples, please refer to the examples folder.

Continual Learning (beta2 & relative_step)

By default, Adafactor's second-moment decay rate dynamically decays with the training step, and the internal learning rate schedule (relative_step) scales the learning rate accordingly.

For endless fine-tuning or lifelong learning, this often leads to overly small learning rates and "blunted" second-moment estimates. To avoid these issues and keep the optimizer responsive:

• Set relative_step=False to disable the built-in LR schedule (allowing you to use an external scheduler).
• Set beta2=0.999 to lock the EMA window (similar to Adam).

Decoupled Weight Decay (scale_weight_decay=False)

By default, Adafactor's weight decay is coupled with the parameter's RMS scale.

• If you prefer the AdamW-style decoupled weight decay, set scale_weight_decay=False.

Fira Limiter

The Norm-Growth Limiter limits the relative increase of update norms to suppress destructive gradient spikes.

• enable_fira_for_apollo: Defaults to True. The APOLLO path applies the limiter by default to guard against sudden gradient rises in the low-rank subspace.
• enable_fira_for_adafactor: Defaults to False. Enables the limiter for the non-APOLLO paths.
• fira_margin: Defaults to 0.01. The tolerance margin for norm growth, shared by both switches. The limiter activates only when the current update norm grows by more than this margin (e.g., 0.01 = 1%) compared to the previous step.

Full-Rank and Factorized Variance (factored)

By default, Adafactor factorizes the second moment of $\ge$ 2D tensors into row and column statistics (factored=True) to minimize state memory. Setting factored=False switches to full element-wise variance (similar to RMSProp), while still retaining Adafactor's update RMS clipping mechanism (controlled by the d parameter). This configuration can be useful in the following scenarios:

• Convolution Weights (>2D): Setting factored=False maintains independent variance for each spatial position in the convolution kernel, enabling finer per-element gradient scaling.
• Sparse Embeddings: Combining factored=False, scale_parameter=False, d=0.0 and a lower learning rate creates a momentum-free adaptive optimizer. This allows for fine-grained, per-token updates while avoiding cold-token interference.

No-Compiler Environments (use_cuda_kernel=False)

If you are in an environment without a CUDA compiler and want to bypass JIT compilation entirely:

• Set use_cuda_kernel=False to fall back to the pure PyTorch implementation.

Enable the APOLLO path to compute gradient scaling factors in a memory-efficient low-rank subspace. Compared to Adafactor's standard row/column factorization (which assumes spatial independence), APOLLO uses random subspace projection to capture cross-dimensional covariance information while keeping memory overhead low.

• apollo_rank: The target rank for the projection subspace. The default is 0 (disabled).

  • The official APOLLO GitHub repository recommends a rank of 256 for 1B and 7B models.
  • The LLaMA-Factory default is 16.
  • Setting this to 1 (APOLLO-Mini style) minimizes the low-rank state (saves even more VRAM than the Adafactor path). The original APOLLO-Mini relies on the first-moment (beta1) to smooth out projection noise. To replicate this, set beta1=0.9 alongside apollo_rank=1.

• apollo_scale: Heuristic scale parameter used to compensate for approximation error introduced by low-rank gradient scaling. The implementation applies sqrt(apollo_scale) to the resulting scaling factor. Defaults to 1.0.

• apollo_scale_type: Determines how the scaling factor is applied. 'channel' applies it per channel (Standard APOLLO), while 'tensor' applies it globally (APOLLO-Mini).

• apollo_update_proj_gap: Steps between projection matrix refreshes. Defaults to 200.

• apollo_factorize (Experimental): Applies Adafactor-style row/column factorization within the low-rank subspace to further reduce state memory overhead.

Enable the CAME (Confidence-guided Adaptive Memory Efficient Optimization) path to add a confidence estimation stage after momentum accumulation:

Adaptive Scaling ($V$) → Momentum Accumulation ($M$) → Confidence Weighting ($C$)

Key Parameters & Tuning

The confidence stage measures the consistency between the current update direction and historical momentum, adaptively suppressing highly oscillatory updates.

• beta3: EMA decay coefficient for the confidence matrix. Requires beta1 (momentum) and factored=True. Defaults to None (disabled).
• Learning Rate: The official CAME implementation recommends ×ばつ the AdamW learning rate (see official tuning guide). To align with the original CAME behavior in this library, disable Adafactor's RMS scaling (scale_parameter=False) and set d=1.0 (which corresponds to CAME's clip_threshold).
• Warmup: Since the confidence matrix is zero-initialized without bias correction, a learning rate warmup is recommended to safely establish the confidence baseline.
• Choosing beta3: beta3 should generally be larger than beta2 so the confidence estimate evolves more slowly than the variance estimate. A practical starting range is 0.9995–0.99995 when beta2=0.999.

Configuration Example

To replicate "vanilla" CAME (stripping Adafactor's native modifications), you can use the following configuration:

{
 "params": param_group,
 "lr": lr, # Original CAME recommends 0.5-0.9x AdamW LR
 "beta1": 0.9,
 "beta2": 0.999,
 "beta3": 0.9999, # Enable CAME confidence guidance
 "apollo_rank": 0, # Set to 0 for vanilla CAME. (Set >0 to enable APOLLO+CAME fusion)
 "weight_decay": weight_decay,
 "scale_weight_decay": False,
 "scale_parameter": False, # Disable Adafactor RMS scaling to align with vanilla CAME
 "d": 1.0,
 "relative_step": False,
},

If you are migrating from optimizers like AdamW, Adafactor's learning rate behavior might feel a bit different. This is mainly due to the scale_parameter option.

• scale_parameter=True (default) Because of RMS scaling, a very small lr (e.g., 1e-5) often leads to extremely slow progress. Start with lr=1e-3 and adjust in the range 1e-4–5e-3 if needed.

• scale_parameter=False Disables RMS scaling, making the update scale more similar to AdamW. Use the learning rates you're familiar with for AdamW and tune as usual. (Note: the second moment is still factorized, so behavior is not identical.)

These are safe starting points. Always validate on your own task and batch size.