Choosing the right optimization algorithm is a significant factor when training large-scale models, particularly in distributed environments. While standard optimizers like Stochastic Gradient Descent (SGD) with momentum or Adam form the basis, scaling up training often necessitates optimizers designed to handle very large batch sizes or specific regularization needs effectively. The techniques discussed earlier in this chapter, such as data parallelism with pmap and gradient accumulation, directly lead to larger effective batch sizes, influencing optimizer behavior.

Common Optimizers for Large Models

AdamW

Adam remains a widely used optimizer due to its adaptive learning rates. However, the standard implementation of L2 regularization in Adam is often suboptimal. AdamW modifies Adam by decoupling the weight decay calculation from the gradient update associated with the adaptive learning rates.

In standard Adam with L2 regularization, the decay term interacts with the adaptive moments ( $m_t$ and $v_t$ ), potentially leading to weights with large historical gradients decaying less than those with smaller gradients. AdamW applies the weight decay directly to the weights after the Adam step, behaving more like the weight decay used with SGD.

The conceptual update for a parameter $\theta$ looks like:

Compute gradient $g_t = \nabla L(\theta_{t-1})$ .
Update biased first and second moment estimates $m_t, v_t$ .
Compute bias-corrected estimates $\hat{m}_t, \hat{v}_t$ .
Update parameters using Adam logic: $\theta'_{t} = \theta_{t-1} - \eta \frac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}$ .
Apply weight decay: $\theta_t = \theta'_{t} - \eta \lambda \theta_{t-1}$ . (Note: The decay is proportional to the learning rate $\eta$ and decay factor $\lambda$ ).

This decoupling often results in better generalization performance and more stable training for large models like Transformers compared to the original Adam implementation with L2 regularization. It has become a standard choice for many large language model training recipes.

LAMB (Layer-wise Adaptive Moments optimizer for Batching training)

When using massive batch sizes, often achieved through distributed training across many accelerators using pmap, standard optimizers like Adam can sometimes become unstable or require careful learning rate tuning (especially extensive warm-up). LAMB was developed specifically to enable stable training with extremely large batch sizes (tens of thousands or more).

The main idea behind LAMB is to apply layer-wise normalization to the parameter updates. It computes the Adam update step similarly to AdamW but then normalizes the update for each layer based on the ratio of the norms of the weights and the Adam update for that layer.

Conceptually, for each layer $l$ :

Compute the Adam(W) update step $\Delta \theta_l$ .
Compute the ratio $r_l = \frac{\|\theta_l\|}{\|\Delta \theta_l\|}$ .
Scale the update for that layer: $\theta_{l, t} = \theta_{l, t-1} - \eta \cdot r_l \cdot \Delta \theta_l$ . (This is a simplified view; the actual implementation involves trust ratio calculations).

This layer-specific trust ratio scaling helps prevent updates from becoming excessively large or small for certain layers, which can happen with large batches where gradient variance is reduced, potentially leading Adam to take overly aggressive steps. LAMB has been shown effective in training models like BERT with very large batch sizes significantly faster than previously possible.

Other Important Considerations

Beyond the choice of the core algorithm (AdamW, LAMB, etc.), other factors are essential for successful large-scale optimization:

Learning Rate Scheduling: Almost all large model training relies heavily on carefully designed learning rate schedules. Common strategies include:
- Warm-up: Starting with a small learning rate and linearly increasing it over the first few thousand steps stabilizes training early on, especially when weights are randomly initialized.
- Decay: After the warm-up phase, the learning rate is typically decayed over the rest of training. Common decay schedules include linear decay, cosine annealing, or inverse square root decay. The choice often depends on the specific model architecture and dataset.
Optimizer State Management: Adaptive optimizers like Adam(W) and LAMB maintain state (e.g., momentum and variance estimates) for each parameter. For large models, this state can consume a significant amount of memory, sometimes rivaling the model parameters themselves. In distributed settings using pmap, this optimizer state must also be distributed (sharded) across the devices along with the parameters and gradients. Libraries like optax (a popular gradient processing and optimization library for JAX) typically handle this automatically when used correctly within a pmap-decorated function. Ensure your training setup correctly partitions the optimizer state across devices.

Optimization in Practice with JAX

Libraries like optax provide implementations for many common and advanced optimizers, integrating smoothly with JAX's functional programming model and transformations. Using optax typically involves:

Initialization: Creating the optimizer state, often based on the model parameters.
Update Step: A function that takes gradients and the current optimizer state to produce parameter updates and the new optimizer state.

Here's a conceptual example of an update step using optax:

import jax
import jax.numpy as jnp
import optax

# Assume 'params' are model parameters (e.g., a PyTree)
# Assume 'grads' are gradients computed by jax.grad
# Assume 'opt_state' is the current optimizer state

# Define the optimizer (e.g., AdamW with cosine decay schedule)
learning_rate_schedule = optax.warmup_cosine_decay_schedule(...)
optimizer = optax.adamw(learning_rate=learning_rate_schedule, weight_decay=0.01)

# Initialize optimizer state (typically done once outside the training loop)
# opt_state = optimizer.init(params)

# Inside the training step (potentially pmap'd)
@jax.jit
def update_step(params, grads, opt_state):
  updates, opt_state = optimizer.update(grads, opt_state, params)
  params = optax.apply_updates(params, updates)
  return params, opt_state

# Apply the update
# params, opt_state = update_step(params, grads, opt_state)

When this update_step function is used within a pmap, optax helps ensure that the gradient calculations, updates, and state management are handled correctly across devices, including necessary gradient aggregation (e.g., using lax.pmean) before the optimizer update is computed.

Choosing and tuning the optimizer, along with its associated learning rate schedule and weight decay, is an iterative process. While optimizers like AdamW provide a strong starting point, and LAMB offers advantages for extremely large batches, experimentation is often needed to find the best combination for a specific large-scale training task, considering the model architecture, dataset characteristics, and available computational resources.