Upload custom model with source code and tokenizer

config.json

@@ -10,6 +10,7 @@
   },
   "cls_token_id": 1,
   "context_length": 1024,
+  "dropout_p": 0.0,
   "dropout_rate": 0.0,
   "embedding_dim": 384,
   "head_dim": 128,

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:393214c1dd86d89df226806c609500210a77fc89e3a0f3da3bbcf78506365afe
+oid sha256:2ff9098ed08946c6a31592eceabe3adaeb14c53f03f1882534dc8f295e55a7e5
 size 274585131

qwen.py

@@ -49,7 +49,7 @@ class FlexQwenConfig(PretrainedConfig):
         cls_token_id: int = 1,
         pad_token_id: int = 3,
         tie_word_embeddings: bool = False,
-        dropout_rate: float = 0.0,
+        dropout_p: float = 0.0,
         **kwargs,
     ):
         super().__init__(
@@ -87,7 +87,7 @@ class FlexQwenConfig(PretrainedConfig):
         # Standard HF Config params
         self.tie_word_embeddings = tie_word_embeddings
 
-        self.dropout_rate = dropout_rate
+        self.dropout_p = dropout_p
 
 
 class FlexQwenPreTrainedModel(PreTrainedModel):
@@ -117,6 +117,7 @@ class GroupedQueryAttention(nn.Module):
         head_dim: int | None = None,
         qk_norm: bool = False,
         rms_norm_eps: float = 1e-6,
+        dropout_p: float = 0.0,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ):
@@ -129,6 +130,7 @@ class GroupedQueryAttention(nn.Module):
         self.num_heads = num_heads
         self.num_kv_groups = num_kv_groups
         self.group_size = num_heads // num_kv_groups
+        self.dropout_p = dropout_p
 
         if head_dim is None:
             assert in_features % num_heads == 0, (
@@ -210,7 +212,7 @@ class GroupedQueryAttention(nn.Module):
             key,
             value,
             attn_mask=attention_mask,
-            dropout_p=0.0,
+            dropout_p=self.dropout_p,
             enable_gqa=True,
         )
         out = self.out_proj(
@@ -234,6 +236,7 @@ class Transformer(nn.Module):
         moe_num_experts: int = 0,
         moe_hidden_dim: int = 128,
         rms_norm_eps: float = 1e-6,
+        dropout_p: float = 0.0,
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ):
@@ -245,6 +248,8 @@ class Transformer(nn.Module):
             head_dim=head_dim,
             num_kv_groups=num_kv_groups,
             qk_norm=qk_norm,
+            rms_norm_eps=rms_norm_eps,
+            dropout_p=dropout_p,
             **factory_kwargs,
         )
 
@@ -331,6 +336,7 @@ class FlexQwen(FlexQwenPreTrainedModel):
                     moe_num_experts=config.moe_num_experts,
                     moe_hidden_dim=config.moe_hidden_dim,
                     rms_norm_eps=config.rms_norm_eps,
+                    dropout_p=config.dropout_p,
                     device=device,
                     dtype=dtype,
                 )
@@ -499,7 +505,7 @@ class FlexQwenForSequenceClassification(FlexQwenPreTrainedModel):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = FlexQwen(config, device=device, dtype=dtype)
-        self.dropout = nn.Dropout(p=config.dropout_rate)
+        self.dropout = nn.Dropout(p=config.dropout_p)
         self.score = CastedLinear(
             config.embedding_dim,
             self.num_labels,