Update modeling_mybert.py

modeling_mybert.py

@@ -58,8 +58,7 @@ class MyBertSelfAttention(nn.Module):
         k = k.view(*shp).transpose(1, 2)
         v = v.view(*shp).transpose(1, 2)
         if cos is not None:
-            q, k = _apply_rope(q, k, cos, sin)      # RoPE on Q,K only
-        # attn_mask=None + dropout 0 -> SDPA selects the FlashAttention kernel on sm_89
+            q, k = _apply_rope(q, k, cos, sin)      
         ctx = F.scaled_dot_product_attention(
             q, k, v, attn_mask=attention_mask,
             dropout_p=self.dropout_prob if self.training else 0.0, is_causal=False,
@@ -89,8 +88,6 @@ class MyBertAttention(nn.Module):
 
 
 class MyBertSwiGLU(nn.Module):
-    """SwiGLU FFN: down( silu(gate(x)) * up(x) ). Encoder only -- NOT the MLM head."""
-
     def __init__(self, config):
         super().__init__()
         b = config.use_bias
@@ -131,8 +128,6 @@ class MyBertEncoder(nn.Module):
 
 
 class MyBertPredictionHeadTransform(nn.Module):
-    """MLM head keeps GELU -- SwiGLU deliberately NOT used here."""
-
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
@@ -181,10 +176,6 @@ class MyBertPreTrainedModel(PreTrainedModel):
         elif isinstance(module, nn.LayerNorm):
             module.bias.data.zero_()
             module.weight.data.fill_(1.0)
-        # transformers v5 loads the model on the meta device, then materialises
-        # non-persistent buffers with torch.empty (uninitialised memory) and calls
-        # _init_weights to refill them. Without this branch rope_cos/rope_sin are
-        # garbage after from_pretrained() -> NaN logits in the fill-mask pipeline.
         if isinstance(module, MyBertModel):
             head_dim = self.config.hidden_size // self.config.num_attention_heads
             cos, sin = _build_rope_cache(head_dim, self.config.max_position_embeddings,
@@ -205,7 +196,6 @@ class MyBertModel(MyBertPreTrainedModel):
         self.register_buffer("rope_sin", sin, persistent=False)
         self.post_init()
 
-        # scaled init on residual-output projections (pre-LN depth stability)
         n = config.num_hidden_layers
         for lyr in self.encoder.layer:
             lyr.attention.output.dense.weight.data.normal_(
@@ -232,7 +222,6 @@ class MyBertModel(MyBertPreTrainedModel):
 
 
 class MyBertForMaskedLM(MyBertPreTrainedModel):
-    # modern transformers wants a DICT here (a list raises AttributeError in post_init)
     _tied_weights_keys = {
         "cls.predictions.decoder.weight": "mybert.embeddings.word_embeddings.weight",
     }
@@ -255,9 +244,6 @@ class MyBertForMaskedLM(MyBertPreTrainedModel):
                           return_dict=True).last_hidden_state
 
         if labels is not None and self.config.sparse_prediction:
-            # LM head ONLY on masked positions -- identical loss to dense (diff 0.0),
-            # no B x T x 28996 logit tensor. Fixed-count masking keeps this shape STATIC,
-            # so torch.compile never recompiles.
             flat_h = seq.view(-1, seq.size(-1))
             flat_y = labels.view(-1)
             idx = (flat_y != -100).nonzero(as_tuple=True)[0]
@@ -266,7 +252,7 @@ class MyBertForMaskedLM(MyBertPreTrainedModel):
                                    label_smoothing=label_smoothing)
             return MaskedLMOutput(loss=loss, logits=None)
 
-        logits = self.cls(seq)                      # dense path -> fill-mask pipeline
+        logits = self.cls(seq)                      
         loss = None
         if labels is not None:
             loss = F.cross_entropy(logits.float().view(-1, self.config.vocab_size),