English comments: the published copy is read internationally

modeling_ruqlm.py

@@ -1,16 +1,18 @@
 """
-معمارية Ruq-LM — محوّل صغير مُدرَّب من الصفر.
-
-هذا هو النموذج نفسه: أوزان مهيّأة عشوائياً، لا اشتقاق من أي نموذج جاهز.
-الوصفة حديثة وقياسية: pre-norm + RMSNorm + RoPE + SwiGLU + تضمينات مربوطة.
-
-لماذا هذه الخيارات عند 30M بارامتر تحديداً:
-  - RMSNorm بدل LayerNorm: أقل عمليات، ولا فرق يُذكر في الجودة.
-  - RoPE بدل تضمينات موضعية مُتعلَّمة: لا بارامترات إضافية، وتعميم أفضل
-    على أطوال لم تُرَ أثناء التدريب.
-  - SwiGLU: أفضل من ReLU/GELU عند ثبات عدد البارامترات.
-  - ربط تضمينات الدخل بالخرج: يوفّر 4.2M بارامتر — أي 14% من النموذج
-    عند مفردات 8192. عند هذا الحجم الصغير هذا فرق جوهري لا تحسين هامشي.
+RuqLM architecture — a small transformer trained from scratch.
+
+This is the model itself: randomly initialised weights, not derived from any
+pretrained checkpoint. The recipe is modern and standard: pre-norm, RMSNorm,
+RoPE, SwiGLU, tied embeddings.
+
+Why these choices at 30M parameters specifically:
+  - RMSNorm over LayerNorm: fewer operations, no measurable quality cost.
+  - RoPE over learned positional embeddings: no extra parameters, and better
+    generalisation to lengths not seen during training.
+  - SwiGLU: better than ReLU/GELU at a fixed parameter count.
+  - Tying input and output embeddings: saves 4.2M parameters, 14% of the model
+    at a vocabulary of 8192. At this size that is structural, not a marginal
+    optimisation.
 """
 
 from __future__ import annotations
@@ -29,8 +31,8 @@ class ModelArgs:
     d_model: int = 512
     n_layers: int = 8
     n_heads: int = 8
-    n_kv_heads: int | None = None      # None = انتباه متعدد الرؤوس عادي؛ أقل = GQA
-    ffn_hidden: int | None = None      # None = يُحسب تلقائياً (~8/3 × d مقرّباً لمضاعف 64)
+    n_kv_heads: int | None = None      # None = plain multi-head attention; fewer = GQA
+    ffn_hidden: int | None = None      # None = derived (~8/3 x d, rounded to 64)
     max_seq_len: int = 512
     rope_theta: float = 10000.0
     norm_eps: float = 1e-5
@@ -39,14 +41,14 @@ class ModelArgs:
 
     def __post_init__(self) -> None:
         if self.d_model % self.n_heads:
-            raise ValueError("d_model يجب أن يقبل القسمة على n_heads")
+            raise ValueError("d_model must be divisible by n_heads")
         if self.n_kv_heads is None:
             self.n_kv_heads = self.n_heads
         if self.n_heads % self.n_kv_heads:
-            raise ValueError("n_heads يجب أن يقبل القسمة على n_kv_heads")
+            raise ValueError("n_heads must be divisible by n_kv_heads")
         if self.ffn_hidden is None:
-            # 8/3 × d بدل 4 × d: SwiGLU يستخدم ثلاث مصفوفات لا اثنتين،
-            # فنقلّص العرض للحفاظ على نفس ميزانية البارامترات.
+            # 8/3 x d rather than 4 x d: SwiGLU uses three matrices instead of
+            # two, so the width shrinks to hold the parameter budget constant.
             self.ffn_hidden = 64 * math.ceil((8 * self.d_model / 3) / 64)
 
     @property
@@ -57,7 +59,7 @@ class ModelArgs:
         return asdict(self)
 
 
-# --------------------------------------------------------------------- الطبقات
+# ----------------------------------------------------------------------- layers
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-5):
         super().__init__()
@@ -65,7 +67,8 @@ class RMSNorm(nn.Module):
         self.weight = nn.Parameter(torch.ones(dim))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # يُحسب في float32 دائماً: التطبيع في bf16 يفقد دقة تُهم عند العمق
+        # Always computed in float32: normalising in bf16 loses precision that
+        # matters once the network is deep.
         dtype = x.dtype
         x = x.float()
         x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
@@ -73,7 +76,7 @@ class RMSNorm(nn.Module):
 
 
 def build_rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
-    """يعيد (cos, sin) بشكل (seq_len, head_dim/2)."""
+    """Returns (cos, sin), each of shape (seq_len, head_dim/2)."""
     inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
     pos = torch.arange(seq_len, device=device).float()
     freqs = torch.outer(pos, inv_freq)
@@ -81,7 +84,7 @@ def build_rope_cache(seq_len: int, head_dim: int, theta: float, device, dtype):
 
 
 def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
-    """x بشكل (B, H, S, D) — يدوّر كل زوج إحداثيات بزاوية تتناسب مع الموضع."""
+    """x is (B, H, S, D) — rotates each coordinate pair by an angle set by position."""
     x1, x2 = x.chunk(2, dim=-1)
     cos = cos[None, None, : x.size(-2), :]
     sin = sin[None, None, : x.size(-2), :]
@@ -146,7 +149,7 @@ class Block(nn.Module):
         return x + self.drop(self.ffn(self.ffn_norm(x)))
 
 
-# --------------------------------------------------------------------- النموذج
+# ------------------------------------------------------------------------ model
 class RuqLM(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -161,8 +164,9 @@ class RuqLM(nn.Module):
             self.lm_head.weight = self.tok_emb.weight
 
         self.apply(self._init)
-        # تدرّج المسارات المتبقية ينمو مع العمق؛ نقلّص أوزان الإسقاط الأخير
-        # في كل كتلة بـ 1/sqrt(2L) للحفاظ على تباين ثابت عبر الطبقات (GPT-2).
+        # Variance on the residual stream grows with depth, so the final
+        # projection in each block is scaled down by 1/sqrt(2L) to hold it
+        # roughly constant across layers (GPT-2).
         std = 0.02 / math.sqrt(2 * args.n_layers)
         for block in self.blocks:
             nn.init.normal_(block.attn.wo.weight, mean=0.0, std=std)
@@ -198,7 +202,7 @@ class RuqLM(nn.Module):
 
         loss = None
         if labels is not None:
-            # الإزاحة: الموضع i يتنبأ بالتوكن i+1
+            # Shifted: position i predicts token i+1
             loss = F.cross_entropy(
                 logits[:, :-1].reshape(-1, logits.size(-1)).float(),
                 labels[:, 1:].reshape(-1),
@@ -206,9 +210,9 @@ class RuqLM(nn.Module):
             )
         return logits, loss
 
-    # ------------------------------------------------------------- الإحصاءات
+    # -------------------------------------------------------------------- stats
     def num_params(self, embeddings: bool = True) -> int:
-        """التضمينات المربوطة تُحسب مرة واحدة (lm_head.weight هو نفسه tok_emb.weight)."""
+        """Tied embeddings are counted once (lm_head.weight is tok_emb.weight)."""
         seen, total = set(), 0
         for name, p in self.named_parameters():
             if id(p) in seen:
@@ -222,7 +226,7 @@ class RuqLM(nn.Module):
     @torch.no_grad()
     def generate(self, input_ids, max_new_tokens=128, temperature=0.8,
                  top_k=50, eos_id=None):
-        """توليد بسيط بلا كاش KV — كافٍ للتقييم على تسلسلات قصيرة."""
+        """Plain sampling without a KV cache — adequate for short sequences."""
         self.eval()
         for _ in range(max_new_tokens):
             window = input_ids[:, -self.args.max_seq_len:]