change mup param names

config.json

@@ -15,7 +15,7 @@
   },
   "bos_token_id": 50256,
   "embd_pdrop": 0.0,
-  "embeddings_scale": 14.6,
+  "mup_embeddings_scale": 14.6,
   "eos_token_id": 50256,
   "initializer_range": 0.073,
   "layer_norm_epsilon": 1e-05,
@@ -25,16 +25,16 @@
   "n_inner": 6826,
   "n_layer": 32,
   "n_positions": 8192,
-  "output_logits_scale": 0.22200000000000003,
+  "mup_output_alpha": 2.2200000000000003,
   "position_embedding_type": "alibi",
   "reorder_and_upcast_attn": false,
   "resid_pdrop": 0.0,
   "scale_attn_by_inverse_layer_idx": false,
   "scale_attn_weights": true,
-  "scale_qk_dot_by_d": true,
+  "mup_scale_qk_dot_by_d": true,
   "torch_dtype": "bfloat16",
   "transformers_version": "4.30.0",
   "use_cache": true,
   "vocab_size": 50257,
-  "width_scale": 0.1
+  "mup_width_scale": 0.1
 }

configuration_btlm.py

@@ -23,7 +23,7 @@ from transformers.utils import logging
 logger = logging.get_logger(__name__)
 
 BTLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
-    "cerebras/BTLM-3B": "https://huggingface.co/cerebras/BTLM-3B/resolve/main/config.json",
+    "cerebras/btlm-3b-8k-base": "https://huggingface.co/cerebras/btlm-3b-8k-base/resolve/main/config.json",
 }
 
 
@@ -74,14 +74,14 @@ class BTLMConfig(PretrainedConfig):
             dot-product/softmax to float() when training with mixed precision.
         position_embedding_type (`str`, *optional*, defaults to `"learned"`):
             Positional embedding can be either `"alibi"` or `"learned"`.
-        width_scale (`float`, *optional*, defaults to 1.0):
+        mup_width_scale (`float`, *optional*, defaults to 1.0):
             muP parameter to scale learning rate and initializers. Calculated as (`d_model,0 / d_model`), where
             `d_model` is the model's width and `d_model,0` is the proxy model's width.
-        embeddings_scale (`float`, *optional*, defaults to 1.0):
+        mup_embeddings_scale (`float`, *optional*, defaults to 1.0):
             muP parameter to scale token and position embeddings.
-        output_logits_scale (`float`, *optional*, defaults to 1.0):
-            muP parameter to scale output logits. Calculated as (`output_alpha * width_scale`)
-        scale_qk_dot_by_d (`bool`, *optional*, defaults to `False`):
+        mup_output_alpha (`float`, *optional*, defaults to 1.0):
+            muP parameter to scale output logits (`output_logits_scale = mup_output_alpha * mup_width_scale`).
+        mup_scale_qk_dot_by_d (`bool`, *optional*, defaults to `False`):
             Scale attention weights by dividing by hidden_size instead of sqrt(hidden_size). Need to set
             scale_attn_weights to `True` as well.
 
@@ -130,10 +130,10 @@ class BTLMConfig(PretrainedConfig):
         scale_attn_by_inverse_layer_idx=False,
         reorder_and_upcast_attn=False,
         position_embedding_type="learned",
-        width_scale=1.0,
-        embeddings_scale=1.0,
-        output_logits_scale=1.0,
-        scale_qk_dot_by_d=False,
+        mup_width_scale=1.0,
+        mup_embeddings_scale=1.0,
+        mup_output_alpha=1.0,
+        mup_scale_qk_dot_by_d=False,
         **kwargs,
     ):
         self.vocab_size = vocab_size
@@ -157,9 +157,9 @@ class BTLMConfig(PretrainedConfig):
         self.eos_token_id = eos_token_id
 
         self.position_embedding_type = position_embedding_type
-        self.width_scale = width_scale
-        self.embeddings_scale = embeddings_scale
-        self.output_logits_scale = output_logits_scale
-        self.scale_qk_dot_by_d = scale_qk_dot_by_d
+        self.mup_width_scale = mup_width_scale
+        self.mup_embeddings_scale = mup_embeddings_scale
+        self.mup_output_alpha = mup_output_alpha
+        self.mup_scale_qk_dot_by_d = mup_scale_qk_dot_by_d
 
         super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

modeling_btlm.py

@@ -48,11 +48,11 @@ from .configuration_btlm import BTLMConfig
 
 logger = logging.get_logger(__name__)
 
-_CHECKPOINT_FOR_DOC = "cerebras/BTLM-3B"
+_CHECKPOINT_FOR_DOC = "cerebras/btlm-3b-8k-base"
 _CONFIG_FOR_DOC = "BTLMConfig"
 
 BTLM_PRETRAINED_MODEL_ARCHIVE_LIST = [
-    "cerebras/BTLM-3B",
+    "cerebras/btlm-3b-8k-base",
     # See all BTLM models at https://huggingface.co/models?filter=btlm
 ]
 
@@ -204,7 +204,7 @@ class BTLMAttention(nn.Module):
 
         self.pruned_heads = set()
 
-        self.attn_scale_power = 1.0 if config.scale_qk_dot_by_d else 0.5
+        self.attn_scale_power = 1.0 if config.mup_scale_qk_dot_by_d else 0.5
 
     def prune_heads(self, heads):
         if len(heads) == 0:
@@ -511,7 +511,7 @@ class BTLMPreTrainedModel(PreTrainedModel):
 
     def _init_weights(self, module):
         """Initialize the weights."""
-        mup_init_scale = math.sqrt(self.config.width_scale)
+        mup_init_scale = math.sqrt(self.config.mup_width_scale)
         if isinstance(module, (nn.Linear, Conv1D)):
             # Slightly different from the TF version which uses truncated_normal for initialization
             # cf https://github.com/pytorch/pytorch/pull/5617
@@ -576,7 +576,7 @@ class BTLMPreTrainedModel(PreTrainedModel):
                         return 1
             return 0
 
-        width_scale = self.config.width_scale
+        width_scale = self.config.mup_width_scale
         new_param_groups = []
         new_param_groups.append({"params": [], "lr": lr * width_scale, "weight_decay": weight_decay})
         if not decoupled_wd:
@@ -754,7 +754,7 @@ class BTLMModel(BTLMPreTrainedModel):
             if config.position_embedding_type != "alibi"
             else None
         )
-        self.embeddings_scale = config.embeddings_scale
+        self.embeddings_scale = config.mup_embeddings_scale
 
         self.drop = nn.Dropout(config.embd_pdrop)
         self.h = nn.ModuleList([BTLMBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])
@@ -1062,7 +1062,7 @@ class BTLMLMHeadModel(BTLMPreTrainedModel):
         super().__init__(config)
         self.transformer = BTLMModel(config)
         self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        self.output_logits_scale = config.output_logits_scale
+        self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale
 
         # Model parallel
         self.model_parallel = False
@@ -1264,7 +1264,7 @@ class BTLMForSequenceClassification(BTLMPreTrainedModel):
         self.num_labels = config.num_labels
         self.transformer = BTLMModel(config)
         self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
-        self.output_logits_scale = config.output_logits_scale
+        self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale
 
         # Model parallel
         self.model_parallel = False
@@ -1397,7 +1397,7 @@ class BTLMForTokenClassification(BTLMPreTrainedModel):
             classifier_dropout = 0.1
         self.dropout = nn.Dropout(classifier_dropout)
         self.classifier = nn.Linear(config.hidden_size, config.num_labels)
-        self.output_logits_scale = config.output_logits_scale
+        self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale
 
         # Model parallel
         self.model_parallel = False
@@ -1492,7 +1492,7 @@ class BTLMForQuestionAnswering(BTLMPreTrainedModel):
         self.num_labels = config.num_labels
         self.transformer = BTLMModel(config)
         self.qa_outputs = nn.Linear(config.hidden_size, 2)
-        self.output_logits_scale = config.output_logits_scale
+        self.output_logits_scale = config.mup_output_alpha * config.mup_width_scale
 
         # Model parallel
         self.model_parallel = False