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...

Prefer an external configuration? See optimizer_from_config.py and custom_policy.json.

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.

BF16 Master Weight Training

BF16 parameters use half the storage of FP32 parameters. Set stochastic_round_writeback=True to help small updates accumulate over training. See bf16_master_weight_training.py for a complete example.

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 0.5–0.9× 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.