Upload folder using huggingface_hub

amplify_te.py

@@ -34,6 +34,7 @@ class AMPLIFYConfig(PretrainedConfig):
         layer_norm_after_embedding: bool = False,
         layer_norm_before_last_layer: bool = True,
         vocab_size: int = 27,
+        padded_vocab_size: int = 32,
         ffn_bias: bool = False,
         att_bias: bool = False,
         pad_token_id: int = 0,
@@ -56,6 +57,7 @@ class AMPLIFYConfig(PretrainedConfig):
             layer_norm_after_embedding (bool): Whether to use layer normalization after the embedding.
             layer_norm_before_last_layer (bool): Whether to use layer normalization before the last layer.
             vocab_size (int): The vocabulary size of the model.
+            padded_vocab_size (int): The padded vocabulary size of the model to support fp8.
             ffn_bias (bool): Whether to use bias in the feedforward network.
             att_bias (bool): Whether to use bias in the attention.
             pad_token_id (int): The padding token id.
@@ -77,11 +79,16 @@ class AMPLIFYConfig(PretrainedConfig):
         self.layer_norm_after_embedding = layer_norm_after_embedding
         self.layer_norm_before_last_layer = layer_norm_before_last_layer
         self.vocab_size = vocab_size
+        self.padded_vocab_size = padded_vocab_size
         self.ffn_bias = ffn_bias
         self.att_bias = att_bias
         self.pad_token_id = pad_token_id
         self.max_length = max_length
 
+        assert self.padded_vocab_size >= self.vocab_size, (
+            "padded_vocab_size must be greater than or equal to vocab_size"
+        )
+
 
 class AMPLIFYPreTrainedModel(PreTrainedModel):
     """AMPLIFY pre-trained model."""
@@ -114,7 +121,7 @@ class AMPLIFY(AMPLIFYPreTrainedModel):
         self.config = config
 
         self.encoder = nn.Embedding(
-            config.vocab_size,
+            config.padded_vocab_size,
             config.hidden_size,
             padding_idx=config.pad_token_id,
             dtype=config.torch_dtype,
@@ -245,7 +252,7 @@ class AMPLIFYForMaskedLM(AMPLIFYPreTrainedModel):
         if config.layer_norm_before_last_layer:
             self.decoder = transformer_engine.pytorch.LayerNormLinear(
                 config.hidden_size,
-                config.vocab_size,
+                config.padded_vocab_size,
                 config.norm_eps,
                 params_dtype=config.torch_dtype,
                 normalization="RMSNorm" if config.rms_norm else "LayerNorm",
@@ -292,6 +299,8 @@ class AMPLIFYForMaskedLM(AMPLIFYPreTrainedModel):
 
         # Classification head with layer norm
         logits = self.decoder(outputs.last_hidden_state)
+        if self.config.padded_vocab_size != self.config.vocab_size:
+            logits = logits[:, :, : self.config.vocab_size]
 
         if labels is not None:
             loss = nn.functional.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1))

config.json

@@ -28,10 +28,11 @@
   "num_hidden_layers": 32,
   "other_special_token_ids": null,
   "pad_token_id": 0,
+  "padded_vocab_size": 32,
   "pre_activation_layer_norm": true,
   "rms_norm": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.54.0.dev0",
+  "transformers_version": "4.53.2",
   "unk_token_id": 1,
   "vocab_path": "conf/tokenizer/amplify_vocab.txt",
   "vocab_size": 27

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2df10849de933a3d9917ae2bf543154ffe77091d9c84f4cbb384d666f48dfd03
-size 1416078404
+oid sha256:821e89362bf8b393963d70dfd70a4b60f0bb6a83f485785526531f8bb8f26060
+size 1416116824