Train An LLM Using Maximal Update Parameterization

In this guide, you will learn how to configure an LLM in μP (Maximal Update Parametrization) and perform μTransfer in order to transfer optimized hyperparameters from a scaled-down version of the LLM (a “proxy model”) to the target model, thereby reducing the time and cost of hyperparameter sweeps.

If you are already familiar with the μP and μTransfer techniques, feel free to skip ahead to Configuring μP Params For Your Model.

​

Background

Transformer models are commonly trained using Standard Parametrization (SP) where model weights are initialized from normal distributions with constant standard deviation or standard deviation based on the shape of each layer. However, as models scale to many parameters, the parametrization does not account for potential inter-layer interactions, resulting in unstable training. This instability can cause costly restarts. Furthermore, hyperparameters do not transfer well as models scale in size.

μP addresses these challenges by enabling:

• Stable training dynamics at a large scale by controlling the initialization, activations magnitude, and layer-wise adaptive learning rates independent of model width.

• Zero-shot hyperparameter transfer from a smaller model to larger model. Essentially, μP facilitates width invariance to the model’s hyperparameters.

For example, the Cerebras-GPT family of models was trained using both SP and μP parametrizations. For SP configuration, the hyperparameters were chosen from the best results found in the literature. For μP configurations, μTransfer enabled scaling hyperparameters from a 40M parameter model to models up to 2.7B parameters, resulting in a lower loss. These results highlight how μP can improve large model scaling, improving accuracy, and hyperparameter predictability at scale.

​

Transferring hyperparameters

μTransfer is a hyperparameter transfer paradigm that makes zero-shot transfer of near-optimal hyperparameters possible from a small version of the model to a large model via μP. The “small model” is called the proxy-model for which the hyperparameters are tuned, and the “large model” is referred to as the target-model.

The μTransfer procedure can be broken down as follows:

1. Tune nearly optimal hyperparameters of the “proxy-model”. In particular, the optimal learning rate, the initialization standard deviation, and the embedding multiplier will be relevant to derive scaling formulas for the “target-model”.

As an example, the Cerebras-GPT family used a 200 sample random hyperparameter search on a 40M parameter proxy model trained on 600M tokens with a batch size of 131k tokens. You can find results and details on the optimal hyperparameters in Appendix G of the paper.

2. Scale up the hyperparameters based on the size of the target-model. As of release 2.3, much of this scaling is now automated by default when μP params are specified in the config.

For example, let’s use the 40M parameter Cerebras-GPT model as a proxy model and 2.7B parameter Cerebras-GPT model as the target model. The 40M model has a hidden size of 256, and the 2.7B model has a hidden size of 2560. Thus $d_\text{proxy} = 256, \, d_\text{target} = 2560$, and $m_\text{width} = \frac{d_\text{target}}{d_\text{proxy}} = 10$. These equations are found in Table 14 Appendix G of Cerebras-GPT paper.

In a transformer model, it is helpful to identify which hyperparameters can be μTransferred:

• Optimization params: learning rate (LR), momentum, Adam beta, LR schedule, etc.

• Regularization params: Dropout, weight decay, etc.

• Parameter multipliers describe multiplicative constants that apply to weights/biases (embeddings scale, output logits alpha, etc.)

​

Configuring μP Params For Your Model

In this example, you will perform μTransfer on a GPT-3 2.7B model using a 40M proxy model.

You can enable and configure μP through the model and optimizer params. These parameters differ slightly depending on the backbone of the model for which you want to implement muP, but the overall structure stays the same.

The corresponding config for the 2.7B model is:

trainer:
  init:
    ...
    model:
      ...
      hidden_size: 2560
      filter_size: 10240
      num_hidden_layers: 32
      embedding_initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      output_layer_initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      attention_type: scaled_dot_product
      ...
    optimizer:
      AdamW: ...
    precision:
      ...
    schedulers:
    - ...
    callbacks:
      ...
  fit:
    ...

​

Base dimensions

For each model, there is a set of base dimensions that are specifiable. They indicate width dimensions of the proxy model from which you are trying to perform hyperparameter transfer and are the bare minimum params required to enable μP scaling.

In the case of GPT-3, the supported base dimensions are:

• mup_base_hidden_size (Required to enable μP):

  The hidden size of the proxy model.

• mup_base_filter_size (Required to enable μP):

  The filter size of the proxy model.

In the case where GPT-3 40M as the proxy model and 2.7B is the target model, you would set mup_base_hidden_size = 256 and mup_base_filter_size = 1024 in the config of the target model since they correspond to the hidden_size and filter_size of the 40M model config.

​

Initialization Values

μP requires the addition of initialization scaling for certain layers. As of release 2.3, we now handle this scaling automatically based on the specified base dimensions. We use the width ratio of the dimensions and base dimensions to scale the standard deviation of the corresponding layers’ initializers. In addition to this, we also scale the initializers of the encoder/decoder output projection layers by 2 * num_(encoder/decoder)_blocks to aid with transferring across depth.

​

Hyperparameters

While specifying the required base dimensions is the bare minimum requirement to enable μP scaling on the model, it is necessary to additionally specify some μP specific hyperparameters to obtain better performance and stability in the model. These parameters modify the way we scale the embeddings, output logits, and the attention module projections in the transformer to stabilize gradient flow during training.

For GPT-3, the following μP hyperparameters are supported:

• embeddings_scale:

  Scales the embedding hidden states (i.e. the tensor after embeddings & embedding layer norm are applied). Recommended to tune for stabilizing gradient flow during μP training.

• output_logits_alpha:

  Constant applied to the output logits scalar in μP training. The output logits are scaled by output_logits_alpha * mup_base_hidden_size/hidden_size. Recommended to tune for stabilizing output logits in μP training.

• scale_qk_dot_by_d:

  Scales attention QK dot product by d instead of sqrt(d). Must be enabled for muP training.

• attention_logits_alpha:

  Scales the attention QK dot product by the specified value. Recommended to tune for stabilizing attention logits in muP training.

• scale_output_logits_by_d:

  Scales the output logits in μP by mup_base_hidden_size/hidden_size if True and sqrt(mup_base_hidden_size/hidden_size) if False. It is traditionally set to True in the μP implementation of this model.

​

LR Adjustment Groups

μP requires the addition of layer-wise scaling of learning rates to certain layers. As of release 2.3, we now handle this scaling automatically based on the specified base dimensions. The base dimensions are used to calculate the necessary scales on preset LR adjustment groups that apply the scale to layers that correspond to the group.

Following the GPT-3 example, the supported LR adjustment groups are:

• embedding: Targets the embedding weights.

• decoder_attention: Targets the dense layers in the decoder (Q, K, V, Output projections)

• decoder_input_ffn: Targets the first of the two FFN blocks in the decoder.

• decoder_output_ffn: Targets the final FFN block in the decoder.

You also have the option to override these scales in the YAML by specifying the group in the optimizer params. For example, in order to apply some custom scaling to the embedding weights, you will specify the params as such:

optimizer:
  ...
  adjust_learning_rate:
    embedding: <scale-value>

​

Final Config

Putting it all together, the appended params for a GPT-3 2.7B model with a 40M proxy model and custom embedding weights LR scaling would look like this:

model:
  ...
  mup_base_hidden_size: 256
  mup_base_filter_size: 1024
  scale_qk_dot_by_d: True
  scale_output_logits_by_d: True
  embeddings_scale: <custom-value>
  output_logits_alpha: <custom-value>
  attention_logits_alpha: <custom-value>
  ...

optimizer:
  ...
  adjust_learning_rate:
    embedding: <custom-value>

Where <custom-value> is a placeholder for a hyperparameter value you can determine via sweeping or some other means and is different for each parameter.

Appending those params to the final config:

trainer:
  init:
    ...
    model:
      ...
      hidden_size: 2560
      filter_size: 10240
      num_hidden_layers: 32
      embedding_initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      output_layer_initializer:
        a: -0.04
        b: 0.04
        mean: 0.0
        name: truncated_normal
        std: 0.02
      attention_type: scaled_dot_product

      # Maximal Update Parametrization (muP) Settings
      mup_base_hidden_size: 256
      mup_base_filter_size: 1024
      scale_qk_dot_by_d: True
      scale_output_logits_by_d: True
      embeddings_scale: <custom-value>
      output_logits_alpha: <custom-value>
      attention_logits_alpha: <custom-value>
      ...
    optimizer:
      adjust_learning_rate:
        embedding: <custom-value>
      AdamW: ...
    precision:
      ...
    schedulers:
    - ...
    callbacks:
      ...
  fit:
    ...

​

μP-Compatible Models

GPT-2, GPT-3, Bloom, Llama, Falcon, Starcoder, and MPT with decoder-only architecture and autoregressive language modeling objective are supported by μP. As of 2.3, we are also providing beta support for GPTJ, T5, and BERT Pretrain models. μP works with different types of position embeddings such as ALiBi, RoPE, Relative and Fixed position embeddings, with efficient attention architectures like Multi-Query Attention, and activation functions such as SwiGLU.

Please refer to the following guides for a detailed breakdown of model-specific μP params:

• Configure μP for GPT-Style Models
• Configure μP for BERT Pretrain (Beta)
• Configure μP for T5 (Beta)

​

Implementation notes

GPT and BERT type models are implemented based off Table 8 of the Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer paper, while T5’s implementation is based off Table 9 in order to circumvent some conflicts with its pre-existent initialization scheme.

​

Best Practices

When picking the best configuration of the μP hyperparameters from your hyperparameter sweep, it is recommended to pick top-10 configurations from the 40M base model with the best validation loss. Let’s say the goal is to transfer the hyperparameters to 13B model scale, it is advised to first transfer the hyperparameters to an intermediate sized model size such as 256M or 590M.

At the 256M or 590M scale, run training with the top-10 configs for 20 tokens per parameter and evaluate the model using a validation set. To pick the best configuration from validation loss metric, it is recommended to not just look at the lowest loss value but also the training dynamics.

In our experience, for some configurations you may see very low loss values but few instabilities in the training dynamics (which manifest as loss spikes or exploding gradient norm). In that scenario it is best to pick a configuration with second or third best validation loss which exhibit stable training dynamics.

Consult the base μP configuration for the GPT-3 model, available at 40M config. This configuration was employed internally for our hyperparameter optimization processes. Subsequently, we applied μTransfer techniques to upscale to the 256M model. For different model architectures like Llama and Falcon, it’s necessary to develop a new 40M configuration that aligns with their specific structural requirements.

If there is a significant change in the model architecture such as changing position embeddings from Relative to RoPE or ALiBi, changing activation function from GeLU to SwiGLU, it is advised to redo the hyperparameter sweep at 40M scale to pick the best configs. In such cases, using the optimal set of hyperparameters from Cerebras-GPT paper is not recommended.

​

Related research

1. Radford, A. et al. (2019). Language Models are Unsupervised Multitask Learners.

2. Brown, T.B. et al. (2020). Language Models are Few-Shot Learners.

3. Yang, G. et al. (2022). Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer.

4. Dey, N. et al. (2023). Cerebras-GPT: Open Compute-Optimal Language Models Trained on the Cerebras Wafer-Scale Cluster.