GaLore: Advancing Large Model Training on Consumer-grade Hardware

GaLore enables the training of large language models (LLMs) with up to 7 billion parameters on consumer-grade hardware, such as the NVIDIA RTX 4090, by significantly reducing the memory footprint of optimizer states. This democratization of AI research allows practitioners to train large-scale models without requiring high-end industrial computational resources.

Memory Efficiency via Low-Rank Gradient Projection

GaLore reduces memory consumption by projecting gradients into a lower-dimensional subspace before they are processed by the optimizer. This approach leverages the inherent low-rank structure of gradients in deep neural networks to minimize the amount of data that must be stored and manipulated during training.

For adaptive optimization algorithms like Adam, the optimizer state typically represents a major portion of the memory footprint. By applying this projection, GaLore achieves a reported reduction in memory for storing optimizer states of more than 82.5% during training.

Dynamic Subspace Switching

To maintain the capacity for full-parameter learning and avoid confining the model to a limited portion of the parameter space, GaLore utilizes a dynamic subspace switching mechanism. This mechanism allows the model to navigate through different low-rank subspaces throughout the training process.

The frequency of these switches is balanced to ensure a consistent optimization trajectory while adapting to the evolving low-rank structure of the gradients. This allows for nuanced control over the trade-off between memory efficiency and optimization performance.

Integration with 8-bit Optimizers

Combining GaLore with 8-bit precision optimizers further maximizes memory efficiency by quantizing optimizer states. This synergy allows for the training of larger models or the use of larger batch sizes within the same hardware constraints without compromising model accuracy or convergence speed.

Algorithmic Process of 8-bit Optimization with GaLore

  1. Gradient Projection: Full-precision gradients are projected into a low-rank subspace using projection matrices, then quantized to 8-bit format.
  2. Quantization: Projected gradients, model weights, and optimizer states (such as Adam's moving averages) are quantized from 32-bit floating-point to 8-bit integer representations.
  3. Optimizer Update: 8-bit quantized gradients are updated; this involves de-quantizing gradients to floating-point, applying the update rule, and re-quantizing the updated optimizer states to 8-bit.
  4. De-quantization and Weight Update: Weights are de-quantized to floating-point for processing. GaLore then employs a final projection to map the de-quantized low-rank updates back into the original parameter space before the weight update is applied.

Implementation with Hugging Face Transformers

GaLore is integrated into the Hugging Face transformers library (version 4.39.0 or higher) and the galore-torch library. Users can implement GaLore by specifying the optimizer in TrainingArguments using options such as galore_adamw, galore_adamw_8bit, or galore_adafactor and defining the target modules via optim_target_modules.

Layer-wise Updates

To further reduce the memory footprint, GaLore supports layer-wise weight updates. Instead of performing a single update for all layers after backpropagation, the optimizer updates weights one layer at a time using PyTorch post-accumulation hooks. This feature is activated by appending _layerwise to the optimizer name (e.g., galore_adamw_layerwise).

Sources