sabari2005/cyberslm-base
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: sabari2005/cyberslm-base
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-08-29 19:27:18 +08:00
commit 8035c0135e
19 changed files with 281707 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
.gitattributes vendored Normal file
View File

@@ -0,0 +1,35 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text

146
README.md Normal file
View File

@@ -0,0 +1,146 @@
---
license: apache-2.0
language:
- en
library_name: pytorch
pipeline_tag: text-generation
tags:
- cybersecurity
- security
- small-language-model
- from-scratch
- causal-lm
- pretrained
---
# CyberSLM-base — 33.5M-parameter cybersecurity language model
A decoder-only transformer pretrained from scratch on a cybersecurity corpus.
This is the **base** model: it continues text. It has not been instruction-tuned
and will not answer questions.
For question answering use
[**sabari2005/cyberslm-instruct**](https://huggingface.co/sabari2005/cyberslm-instruct).
**Code:** [github.com/Sabari2005/cyberslm](https://github.com/Sabari2005/cyberslm)
## What this model is
Give it the start of a sentence and it continues it:
```
Prompt: "SQL injection is"
Output: "SQL injection is a common issue in the web interface of Cisco IOS and
IOS XE Software. It has been declared as critical for its security,
integrity, and availability. The vulnerability exists because the
affected software does not properly validate user-supplied input..."
```
Ask it a question and it will continue the *question*, not answer it.
## Model details
| | |
|---|---|
| parameters | 33,531,264 |
| layers | 12 |
| d_model | 384 |
| heads / head_dim | 6 / 64 |
| FFN (SwiGLU) | 1024 |
| context | 2048 |
| vocab | 32,000 (SentencePiece BPE, byte-fallback) |
| positional encoding | RoPE, base 10000 |
| normalisation | RMSNorm, pre-norm |
| LM head | tied to embedding |
| precision | trained in bf16 |
**Training.** 786,432,000 tokens = 4.04 epochs over a 194.8M-token corpus
(~60% cybersecurity across 16 subdomains, ~20% general English and reasoning,
~15% programming, ~5% CS fundamentals). 6,000 steps at 131,072 tokens/step,
AdamW, lr 3e-4 → 3e-5, 600 warmup, cosine decay, grad clip 1.0.
Single A100-40GB, 71 minutes, 184,084 tokens/sec.
## Evaluation
409,600 held-out tokens. Compared against an earlier checkpoint of the same
architecture, both scored by one process on **identical windows at identical
context** (a longer conditioning window lowers loss on its own, so scoring each
at its own maximum would not be a fair comparison):
| metric | this model | earlier checkpoint |
|---|---|---|
| validation loss | **2.3627** | 2.6255 |
| perplexity | **10.62** | 13.81 |
| bits / token | **3.4086** | 3.7878 |
| top-1 accuracy | **57.21%** | 54.38% |
| top-5 accuracy | **72.64%** | 69.55% |
| 8-gram repetition | **23.7%** | 34.0% |
Training-time validation loss at step 6,000 was 2.0247, measured on a different
subset; only the columns above are like-for-like.
No benchmark accuracy is claimed — there is no contamination-checked security
question bank, so nothing beyond next-token metrics is asserted. Single seed.
## Usage
```bash
pip install torch sentencepiece
git clone https://huggingface.co/sabari2005/cyberslm-base
cd cyberslm-base
python infer_base.py --prompt "SQL injection is"
```
Options:
```bash
python infer_base.py \
--prompt "A buffer overflow occurs when" \
--max-new-tokens 120 \
--temperature 0.8 \ # 0 = greedy/deterministic
--top-k 50 --top-p 0.95 \
--repetition-penalty 1.15
```
### Loading directly
```python
import torch, sentencepiece as spm
from cyberslm.model.config import CyberSLMConfig
from cyberslm.model.model import build_model
payload = torch.load("models/base.pt", map_location="cpu", weights_only=False)
model = build_model(CyberSLMConfig(**payload["config"]), device=torch.device("cpu"))
model.load_state_dict(payload["model_state"])
model.eval()
sp = spm.SentencePieceProcessor(); sp.load("tokenizer/tokenizer.model")
ids = [sp.bos_id()] + sp.encode("SQL injection is", out_type=int)
out = model.generate(torch.tensor([ids]), max_new_tokens=60,
temperature=0.0, eos_id=sp.eos_id())
print(sp.decode(out[0].tolist()))
```
`generate()` uses a KV cache, so decoding is O(n) — roughly 50–70 tok/s on CPU.
## Limitations
A 33.5M-parameter model trained on 786M tokens. It produces fluent,
domain-flavoured security prose and is **not factually reliable**. Output drifts
into CVE-advisory boilerplate because that pattern is common in the corpus, and
longer generations repeat (23.7% 8-gram repetition measured).
Intended for research into small language models and as a base for further
scaling or fine-tuning. Not intended for security advice or any use where being
wrong matters.
## Training data
Not published. Curated from public cybersecurity, programming and
general-English sources; not redistributed with the model.
## License
Apache-2.0 for the code and weights. Verify licensing for downstream use
against the sources the corpus was curated from.

32
config.json Normal file
View File

@@ -0,0 +1,32 @@
{
"architectures": [
"LlamaForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": 2,
"dtype": "float32",
"eos_token_id": 3,
"head_dim": 64,
"hidden_act": "silu",
"hidden_size": 384,
"initializer_range": 0.02,
"intermediate_size": 1024,
"max_position_embeddings": 2048,
"mlp_bias": false,
"model_type": "llama",
"num_attention_heads": 6,
"num_hidden_layers": 12,
"num_key_value_heads": 6,
"pad_token_id": 0,
"pretraining_tp": 1,
"rms_norm_eps": 1e-06,
"rope_parameters": {
"rope_theta": 10000.0,
"rope_type": "default"
},
"tie_word_embeddings": true,
"transformers_version": "5.15.1",
"use_cache": true,
"vocab_size": 32000
}

35
cyberslm/model/__init__.py Normal file
View File

@@ -0,0 +1,35 @@
"""
CyberSLM model package.
Public API (grows as phases are added):
Phase 1: CyberSLMConfig, default_config, RMSNorm, RotaryPositionEmbedding, apply_rope
Phase 2: MultiHeadSelfAttention, SwiGLUFeedForward
Phase 3: DecoderBlock, CyberSLM, build_model
Phase 4: (training engine lives in cyberslm.training)
"""
from cyberslm.model.config import CyberSLMConfig, default_config
from cyberslm.model.norm import RMSNorm
from cyberslm.model.rope import RotaryPositionEmbedding, apply_rope
from cyberslm.model.attention import MultiHeadSelfAttention
from cyberslm.model.ffn import SwiGLUFeedForward
from cyberslm.model.block import DecoderBlock
from cyberslm.model.model import CyberSLM, build_model, count_parameters, model_summary
__all__ = [
# Phase 1
"CyberSLMConfig",
"default_config",
"RMSNorm",
"RotaryPositionEmbedding",
"apply_rope",
# Phase 2
"MultiHeadSelfAttention",
"SwiGLUFeedForward",
# Phase 3
"DecoderBlock",
"CyberSLM",
"build_model",
"count_parameters",
"model_summary",
]

302
cyberslm/model/attention.py Normal file
View File

@@ -0,0 +1,302 @@
"""
Multi-Head Self Attention (MHSA)
=================================
Standard scaled dot-product multi-head self attention with:
- Rotary Position Embedding (RoPE) on queries and keys
- Causal (auto-regressive) masking
- No bias on projection layers
- Pre-norm placement handled by the enclosing DecoderBlock
Mathematical definition
-----------------------
Given input X ∈ ℝ^{B×T×d}:
Q = X Wq, K = X Wk, V = X Wv (projections, no bias)
Split into H heads, each of dimension d_h = d / H:
Qₕ, Kₕ = RoPE(Qₕ), RoPE(Kₕ) (apply rotary embeddings)
Scaled dot-product attention per head:
Aₕ = softmax( (Qₕ Kₕᵀ) / √d_h + mask ) Vₕ
where mask[i,j] = 0 if j ≤ i else −∞ (causal constraint).
Concatenate and project:
output = concat(A₁, ..., A_H) Wo
Complexity
----------
Time : O(T² · d) — quadratic in sequence length (standard attention)
Space: O(T² · H) — attention weight matrix per head
Numerical stability
-------------------
- Scaling by 1/√d_h keeps the pre-softmax logits in a well-conditioned
range, preventing vanishing gradients from very peaked softmax outputs.
- Softmax is computed by PyTorch's numerically stable implementation
(subtract max before exp).
- RoPE is applied in float32 (see rope.py).
- Causal mask adds −∞ (not a large negative number) so masked positions
become exactly 0 after softmax — no gradient leakage.
FlashAttention compatibility
-----------------------------
The forward pass is written in a way that is structurally compatible with
a future drop-in replacement by ``torch.nn.functional.scaled_dot_product_attention``
(PyTorch 2.0+) or the ``flash-attn`` library. To migrate:
1. Replace the manual QKᵀ/softmax/V block with:
F.scaled_dot_product_attention(q, k, v, attn_mask=None,
dropout_p=0.0, is_causal=True)
2. Remove the manual mask addition (is_causal=True handles it).
"""
from __future__ import annotations
import math
from typing import Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from cyberslm.model.config import CyberSLMConfig
from cyberslm.model.rope import RotaryPositionEmbedding, apply_rope
def _causal_bias(
q_len: int,
k_len: int,
past_len: int,
dtype: torch.dtype,
device: torch.device,
) -> Tensor:
"""
Additive causal mask of shape ``(q_len, k_len)`` for a query block that
starts at absolute position ``past_len``.
Query row ``i`` represents absolute position ``past_len + i`` and may attend
to key columns ``0 .. past_len + i`` inclusive; everything after is -inf.
With ``past_len == 0`` this reduces to the usual upper-triangular mask.
"""
q_pos = torch.arange(q_len, device=device).unsqueeze(1) + past_len # (q,1)
k_pos = torch.arange(k_len, device=device).unsqueeze(0) # (1,k)
return torch.where(
k_pos <= q_pos,
torch.zeros((), dtype=dtype, device=device),
torch.full((), float("-inf"), dtype=dtype, device=device),
)
class MultiHeadSelfAttention(nn.Module):
"""
Multi-Head Self Attention with RoPE and causal masking.
This module owns the four projection matrices (Wq, Wk, Wv, Wo),
the RoPE cache, and the causal mask buffer.
Parameters
----------
config : CyberSLMConfig
Validated model configuration.
Attributes
----------
q_proj : nn.Linear ``(hidden_dim, hidden_dim)``, no bias
k_proj : nn.Linear ``(hidden_dim, hidden_dim)``, no bias
v_proj : nn.Linear ``(hidden_dim, hidden_dim)``, no bias
o_proj : nn.Linear ``(hidden_dim, hidden_dim)``, no bias
rope : RotaryPositionEmbedding
Shape
-----
Input : ``(batch, seq_len, hidden_dim)``
Output : ``(batch, seq_len, hidden_dim)``
"""
def __init__(
self,
config: CyberSLMConfig,
rope: Optional[RotaryPositionEmbedding] = None,
) -> None:
super().__init__()
self.hidden_dim = config.hidden_dim
self.num_heads = config.num_heads
self.head_dim = config.head_dim
self.scale = 1.0 / math.sqrt(self.head_dim)
self.attn_dropout_p = config.attn_dropout
# ------------------------------------------------------------------ #
# Projection layers — no bias (modern practice, saves ~4×384 params) #
# ------------------------------------------------------------------ #
self.q_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
self.o_proj = nn.Linear(config.hidden_dim, config.hidden_dim, bias=False)
# ------------------------------------------------------------------ #
# RoPE cache #
# ------------------------------------------------------------------ #
# The cos/sin tables depend only on (head_dim, max_seq_len, base), so
# every layer's would be byte-identical. CyberSLM builds ONE and passes
# it in; previously each of the 12 layers constructed its own, costing
# ~12 MB of duplicated buffers. Falls back to building its own so the
# module stays usable standalone (tests, ablations).
self.rope = rope if rope is not None else RotaryPositionEmbedding(
head_dim=config.head_dim,
max_seq_len=config.max_seq_len,
base=config.rope_base,
)
# Causality is enforced by scaled_dot_product_attention(is_causal=...)
# rather than a materialised (max_seq_len × max_seq_len) mask buffer,
# which previously cost ~67 MB per layer.
def forward(
self,
x: Tensor,
attention_mask: Optional[Tensor] = None,
return_attn_weights: bool = False,
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
use_cache: bool = False,
) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor, Tensor]]]:
"""
Compute multi-head self attention.
Parameters
----------
x : Tensor
Input of shape ``(batch, seq_len, hidden_dim)``.
attention_mask : Optional[Tensor]
Key-padding mask of shape ``(batch, seq_len)`` with 1 for real
tokens and 0 for padding. When provided, padded keys are excluded
from every query's attention (in addition to the causal mask).
``None`` means no padding (the common training case with packed
sequences).
return_attn_weights : bool
If True, also return the attention weight matrix for inspection.
This forces the slower explicit-softmax path; leave False for
training so the fused kernel is used.
Returns
-------
output : Tensor
Shape ``(batch, seq_len, hidden_dim)``.
attn_weights : Optional[Tensor]
Shape ``(batch, num_heads, seq_len, seq_len)`` if
``return_attn_weights=True``, else ``None``.
present : Optional[Tuple[Tensor, Tensor]]
The concatenated ``(k, v)`` for this layer when ``use_cache=True``,
to be fed back on the next decoding step. ``None`` otherwise.
"""
B, T, _ = x.shape
# Number of tokens already in the cache == absolute position of x[0].
past_len = kv_cache[0].size(2) if kv_cache is not None else 0
# ------------------------------------------------------------------ #
# 1. Linear projections #
# ------------------------------------------------------------------ #
q = self.q_proj(x) # (B, T, hidden_dim)
k = self.k_proj(x)
v = self.v_proj(x)
# ------------------------------------------------------------------ #
# 2. Reshape to (B, H, T, head_dim) for multi-head computation #
# ------------------------------------------------------------------ #
q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2) # (B,H,T,D)
k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
v = v.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
# ------------------------------------------------------------------ #
# 3. Apply Rotary Position Embeddings to Q and K #
# ------------------------------------------------------------------ #
# offset=past_len so a cached decode step rotates the new token by its
# TRUE absolute position rather than position 0.
q, k = apply_rope(q, k, self.rope, offset=past_len)
# ------------------------------------------------------------------ #
# 3b. Prepend the cache. RoPE is applied to the new k BEFORE the
# concat, and cached keys were already rotated when they were first
# computed -- so each key keeps the rotation for its own position.
# ------------------------------------------------------------------ #
if kv_cache is not None:
k = torch.cat([kv_cache[0], k], dim=2)
v = torch.cat([kv_cache[1], v], dim=2)
present = (k, v) if use_cache else None
S = k.size(2) # total key length (past + current)
# ------------------------------------------------------------------ #
# 4. Build the additive key-padding bias (if any). #
# Shape broadcasts over heads and query positions: (B, 1, 1, T). #
# ------------------------------------------------------------------ #
pad_bias: Optional[Tensor] = None
if attention_mask is not None:
# 0 where padding → -inf added to those key columns.
pad = (attention_mask == 0)[:, None, None, :] # (B,1,1,S) bool
pad_bias = torch.zeros(
(B, 1, 1, pad.size(-1)), dtype=q.dtype, device=q.device
).masked_fill(pad, float("-inf"))
if not return_attn_weights:
# Fused, memory-efficient path (FlashAttention when available).
# is_causal=True applies the causal mask without materialising it.
if pad_bias is None and past_len == 0:
context = F.scaled_dot_product_attention(
q, k, v,
is_causal=True,
dropout_p=self.attn_dropout_p if self.training else 0.0,
)
elif pad_bias is None and T == 1:
# Single-token decode: every cached key is in the past, so the
# causal constraint is already satisfied and no mask is needed.
context = F.scaled_dot_product_attention(
q, k, v, dropout_p=0.0,
)
else:
# Combine causal + padding into one additive float mask.
# Query i sits at absolute position past_len + i and may attend
# to keys 0..past_len+i, so the triangle is offset by past_len.
causal = _causal_bias(T, S, past_len, q.dtype, q.device)
attn_bias = causal[None, None, :, :]
if pad_bias is not None:
attn_bias = attn_bias + pad_bias # (B,1,T,S)
context = F.scaled_dot_product_attention(
q, k, v,
attn_mask=attn_bias,
dropout_p=self.attn_dropout_p if self.training else 0.0,
)
attn_weights = None
else:
# Explicit path — needed only when the caller wants the weights.
scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale # (B,H,T,S)
causal = _causal_bias(T, S, past_len, scores.dtype, scores.device)
scores = scores + causal[None, None, :, :]
if pad_bias is not None:
scores = scores + pad_bias
attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32)
if self.attn_dropout_p > 0.0 and self.training:
attn_weights = F.dropout(attn_weights, p=self.attn_dropout_p)
attn_weights = attn_weights.to(v.dtype)
context = torch.matmul(attn_weights, v) # (B,H,T,head_dim)
# ------------------------------------------------------------------ #
# 8. Merge heads: (B, H, T, D) → (B, T, H*D) = (B, T, hidden_dim) #
# ------------------------------------------------------------------ #
context = context.transpose(1, 2).contiguous().view(B, T, self.hidden_dim)
# ------------------------------------------------------------------ #
# 9. Output projection #
# ------------------------------------------------------------------ #
output = self.o_proj(context)
return output, attn_weights, present
def extra_repr(self) -> str:
return (
f"hidden_dim={self.hidden_dim}, "
f"num_heads={self.num_heads}, "
f"head_dim={self.head_dim}"
)

144
cyberslm/model/block.py Normal file
View File

@@ -0,0 +1,144 @@
"""
Transformer Decoder Block
==========================
A single Pre-Norm residual decoder block consisting of:
1. RMSNorm → Multi-Head Self Attention → residual add
2. RMSNorm → SwiGLU FFN → residual add
Pre-Norm architecture
---------------------
Post-Norm (original Transformer):
x = LayerNorm(x + SubLayer(x))
Pre-Norm (modern: GPT-2 onward, LLaMA, etc.):
x = x + SubLayer(LayerNorm(x))
Pre-Norm is strongly preferred for deep networks because:
- Gradients flow through the residual connection bypassing the
normalisation, preventing vanishing gradients in very deep stacks.
- Training is more stable without learning-rate warmup tricks.
- Final RMSNorm on the output is added at the model level (not here)
to normalise the final residual stream before the LM head.
Residual stream
---------------
The residual stream x ∈ ℝ^{B×T×d} is the backbone of the model. Each
sub-layer reads from it, computes a delta, and adds back:
Δ_attn = MHSA( RMSNorm(x) )
x = x + Δ_attn
Δ_ffn = FFN( RMSNorm(x) )
x = x + Δ_ffn
This additive structure means the gradient of the loss with respect to
early layers contains a direct path through the identity (residual),
enabling reliable training of 12+ layer models.
"""
from __future__ import annotations
from typing import Optional, Tuple
import torch
import torch.nn as nn
from torch import Tensor
from cyberslm.model.config import CyberSLMConfig
from cyberslm.model.norm import RMSNorm
from cyberslm.model.attention import MultiHeadSelfAttention
from cyberslm.model.ffn import SwiGLUFeedForward
class DecoderBlock(nn.Module):
"""
Pre-Norm Transformer Decoder Block.
Parameters
----------
config : CyberSLMConfig
Validated model configuration.
layer_idx : int
Zero-based index of this block in the stack (used for display only).
Sub-modules
-----------
attn_norm : RMSNorm
Normalises the residual stream before attention.
attn : MultiHeadSelfAttention
Self attention with RoPE and causal masking.
ffn_norm : RMSNorm
Normalises the residual stream before the FFN.
ffn : SwiGLUFeedForward
SwiGLU position-wise feed-forward network.
Shape
-----
Input : ``(batch, seq_len, hidden_dim)``
Output : ``(batch, seq_len, hidden_dim)``
"""
def __init__(
self,
config: CyberSLMConfig,
layer_idx: int = 0,
rope=None,
) -> None:
super().__init__()
self.layer_idx = layer_idx
# Pre-norm before attention.
self.attn_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
# Multi-head self attention (owns RoPE + causal mask buffers).
self.attn = MultiHeadSelfAttention(config, rope=rope)
# Pre-norm before FFN.
self.ffn_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
# SwiGLU feed-forward.
self.ffn = SwiGLUFeedForward(config)
def forward(
self,
x: Tensor,
attention_mask: Optional[Tensor] = None,
return_attn_weights: bool = False,
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
use_cache: bool = False,
) -> Tuple[Tensor, Optional[Tensor], Optional[Tuple[Tensor, Tensor]]]:
"""
Apply one Pre-Norm residual decoder block.
Parameters
----------
x : Tensor
Residual stream of shape ``(batch, seq_len, hidden_dim)``.
attention_mask : Optional[Tensor]
Key-padding mask ``(batch, seq_len)`` (1=keep, 0=pad), forwarded
to the attention sub-layer. ``None`` for packed/unpadded batches.
return_attn_weights : bool
Propagated to the attention sub-layer.
Returns
-------
x : Tensor
Updated residual stream, same shape as input.
attn_weights : Optional[Tensor]
Attention weights if requested, else None.
"""
# ---- Attention sub-layer ----------------------------------------- #
attn_out, attn_weights, present = self.attn(
self.attn_norm(x),
attention_mask=attention_mask,
return_attn_weights=return_attn_weights,
kv_cache=kv_cache,
use_cache=use_cache,
)
x = x + attn_out
# ---- FFN sub-layer ------------------------------------------------ #
x = x + self.ffn(self.ffn_norm(x))
return x, attn_weights, present
def extra_repr(self) -> str:
return f"layer_idx={self.layer_idx}"

234
cyberslm/model/config.py Normal file
View File

@@ -0,0 +1,234 @@
"""
CyberSLM Model Configuration
=============================
Defines the complete, validated configuration for the CyberSLM decoder-only
transformer. All hyperparameters are frozen at construction time and validated
for mathematical consistency before any model component is instantiated.
Architecture summary
--------------------
- Hidden dim : 384
- Decoder layers : 12
- Attention heads : 6
- Head dim : 64 (hidden_dim / num_heads = 384 / 6 = 64)
- FFN inner dim : 1024 (SwiGLU gate + value = 2 × 1024 → projects back to 384)
- Vocab size : 32 000
- Max context : 4 096
- Approx params : 33.53 M
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
@dataclass(frozen=True)
class CyberSLMConfig:
"""
Immutable configuration for the CyberSLM decoder-only transformer.
All fields are set once at construction; the frozen dataclass guarantees
no accidental mutation during training. Call :meth:`validate` immediately
after construction or use the convenience constructor
:func:`default_config`.
Attributes
----------
vocab_size : int
Number of tokens in the SentencePiece BPE vocabulary.
max_seq_len : int
Maximum token sequence length (context window).
hidden_dim : int
Embedding dimension ``d_model``.
num_layers : int
Number of stacked transformer decoder blocks.
num_heads : int
Number of attention heads. Must evenly divide ``hidden_dim``.
head_dim : int
Dimension of each attention head. Must equal ``hidden_dim // num_heads``.
ffn_hidden_dim : int
Inner dimension of the SwiGLU feed-forward network.
The gate and value projections each map hidden_dim → ffn_hidden_dim,
and the output projection maps ffn_hidden_dim → hidden_dim.
rope_base : int
Base for Rotary Position Embedding frequency computation (θ = 10 000).
norm_eps : float
Epsilon added inside RMSNorm to prevent division by zero.
tie_weights : bool
When True the output projection shares weights with the token embedding.
bias : bool
When True linear layers include a bias term (False = modern practice).
dropout : float
Residual / feed-forward dropout probability (0.0 = disabled).
attn_dropout : float
Attention weight dropout probability (0.0 = disabled).
pad_token_id : Optional[int]
Token ID used for padding; None if the dataset never pads.
bos_token_id : Optional[int]
Beginning-of-sequence token ID.
eos_token_id : Optional[int]
End-of-sequence token ID.
"""
# ------------------------------------------------------------------ #
# Vocabulary & sequence #
# ------------------------------------------------------------------ #
vocab_size: int = 32_000
max_seq_len: int = 4_096
# ------------------------------------------------------------------ #
# Transformer dimensions #
# ------------------------------------------------------------------ #
hidden_dim: int = 384
num_layers: int = 12
num_heads: int = 6
head_dim: int = 64 # must equal hidden_dim // num_heads
ffn_hidden_dim: int = 1_024
# ------------------------------------------------------------------ #
# Positional encoding #
# ------------------------------------------------------------------ #
rope_base: int = 10_000
# ------------------------------------------------------------------ #
# Normalization #
# ------------------------------------------------------------------ #
norm_eps: float = 1e-6
# ------------------------------------------------------------------ #
# Architecture flags #
# ------------------------------------------------------------------ #
tie_weights: bool = True
bias: bool = False
dropout: float = 0.0
attn_dropout: float = 0.0
# ------------------------------------------------------------------ #
# Special token IDs (set by tokenizer integration layer) #
# ------------------------------------------------------------------ #
pad_token_id: Optional[int] = None
bos_token_id: Optional[int] = 2 # real SentencePiece BOS id
eos_token_id: Optional[int] = 3 # real SentencePiece EOS id
# ------------------------------------------------------------------ #
# Validation #
# ------------------------------------------------------------------ #
def validate(self) -> "CyberSLMConfig":
"""
Assert mathematical consistency of every hyperparameter.
Raises
------
ValueError
If any hyperparameter violates an architectural constraint.
Returns
-------
CyberSLMConfig
Self, to allow chaining: ``cfg = CyberSLMConfig().validate()``.
"""
errors: list[str] = []
# Positivity checks
for name, value in [
("vocab_size", self.vocab_size),
("max_seq_len", self.max_seq_len),
("hidden_dim", self.hidden_dim),
("num_layers", self.num_layers),
("num_heads", self.num_heads),
("head_dim", self.head_dim),
("ffn_hidden_dim", self.ffn_hidden_dim),
("rope_base", self.rope_base),
]:
if value <= 0:
errors.append(f"{name} must be positive, got {value}")
# Attention head consistency
if self.hidden_dim % self.num_heads != 0:
errors.append(
f"hidden_dim ({self.hidden_dim}) must be divisible by "
f"num_heads ({self.num_heads})"
)
expected_head_dim = self.hidden_dim // self.num_heads
if self.head_dim != expected_head_dim:
errors.append(
f"head_dim ({self.head_dim}) must equal "
f"hidden_dim // num_heads = {expected_head_dim}"
)
# RoPE requires even head_dim (pairs of sin/cos)
if self.head_dim % 2 != 0:
errors.append(
f"head_dim ({self.head_dim}) must be even for RoPE"
)
# Dropout bounds
for name, value in [("dropout", self.dropout), ("attn_dropout", self.attn_dropout)]:
if not (0.0 <= value < 1.0):
errors.append(f"{name} must be in [0, 1), got {value}")
# norm_eps positivity
if self.norm_eps <= 0.0:
errors.append(f"norm_eps must be positive, got {self.norm_eps}")
if errors:
raise ValueError(
"CyberSLMConfig validation failed:\n"
+ "\n".join(f" • {e}" for e in errors)
)
return self
# ------------------------------------------------------------------ #
# Derived properties #
# ------------------------------------------------------------------ #
@property
def total_attention_dim(self) -> int:
"""``num_heads × head_dim`` — equals ``hidden_dim`` by construction."""
return self.num_heads * self.head_dim
@property
def rope_half_dim(self) -> int:
"""Number of frequency pairs in RoPE (``head_dim // 2``)."""
return self.head_dim // 2
# ------------------------------------------------------------------ #
# Display #
# ------------------------------------------------------------------ #
def __str__(self) -> str:
lines = [
"CyberSLMConfig",
"=" * 40,
f" vocab_size : {self.vocab_size:,}",
f" max_seq_len : {self.max_seq_len:,}",
f" hidden_dim : {self.hidden_dim}",
f" num_layers : {self.num_layers}",
f" num_heads : {self.num_heads}",
f" head_dim : {self.head_dim}",
f" ffn_hidden_dim : {self.ffn_hidden_dim}",
f" rope_base : {self.rope_base}",
f" norm_eps : {self.norm_eps}",
f" tie_weights : {self.tie_weights}",
f" bias : {self.bias}",
f" dropout : {self.dropout}",
f" attn_dropout : {self.attn_dropout}",
]
return "\n".join(lines)
def default_config() -> CyberSLMConfig:
"""
Return the validated default CyberSLM configuration.
This is the single source of truth for all training runs.
All hyperparameters match the finalized architecture specification.
Returns
-------
CyberSLMConfig
A validated, immutable configuration object.
"""
cfg = CyberSLMConfig()
cfg.validate()
return cfg

127
cyberslm/model/ffn.py Normal file
View File

@@ -0,0 +1,127 @@
"""
SwiGLU Feed-Forward Network (FFN)
===================================
Reference: "GLU Variants Improve Transformer" (Noam Shazeer, 2020)
https://arxiv.org/abs/2002.05202
Mathematical definition
-----------------------
Standard FFN (for contrast):
FFN(x) = activation(x W₁) W₂
SwiGLU FFN:
SwiGLU(x) = (x W_gate ⊙ swish(x W_gate)) W₂ ← WRONG shorthand
Correct form with separate gate and value projections:
gate(x) = x W_gate ∈ ℝ^{B×T×ffn_dim}
val(x) = x W_val ∈ ℝ^{B×T×ffn_dim}
hidden = swish(gate(x)) ⊙ val(x)
out = hidden W_out ∈ ℝ^{B×T×hidden_dim}
where swish(z) = z · sigmoid(z) = z · σ(z).
Why SwiGLU
----------
- The gating mechanism (⊙) gives the network a multiplicative path to
control information flow — intuitively "how much" of each feature passes
through at each position.
- swish is smooth and non-monotonic, empirically outperforming ReLU and
GELU in large-scale experiments (PaLM, LLaMA, etc.).
- The factored W_gate / W_val structure adds one extra matrix but improves
quality relative to a standard 2-layer FFN of the same parameter budget.
Parameter count
---------------
Three matrices: W_gate, W_val, W_out
W_gate : hidden_dim × ffn_hidden_dim (384 × 1024 = 393 216)
W_val : hidden_dim × ffn_hidden_dim (384 × 1024 = 393 216)
W_out : ffn_hidden_dim × hidden_dim (1024 × 384 = 393 216)
Total per layer: 1 179 648 ≈ 1.18 M
No bias on any projection (consistent with modern practice).
Numerical stability
-------------------
- swish(z) = z · σ(z) is numerically safe for all real z.
- σ(z) = 1/(1+exp(−z)) in PyTorch uses a numerically stable implementation.
- The element-wise product of two bounded (after sigmoid) quantities does
not amplify values explosively.
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from cyberslm.model.config import CyberSLMConfig
class SwiGLUFeedForward(nn.Module):
"""
SwiGLU Feed-Forward Network.
Consists of three bias-free linear projections:
``gate_proj``, ``val_proj``, and ``out_proj``.
Parameters
----------
config : CyberSLMConfig
Validated model configuration.
Shape
-----
Input : ``(batch, seq_len, hidden_dim)``
Output : ``(batch, seq_len, hidden_dim)``
"""
def __init__(self, config: CyberSLMConfig) -> None:
super().__init__()
self.hidden_dim = config.hidden_dim
self.ffn_hidden_dim = config.ffn_hidden_dim
# Gate projection: produces the gating signal fed through swish.
self.gate_proj = nn.Linear(
config.hidden_dim, config.ffn_hidden_dim, bias=False
)
# Value projection: produces the value signal gated element-wise.
self.val_proj = nn.Linear(
config.hidden_dim, config.ffn_hidden_dim, bias=False
)
# Output projection: maps back to residual stream dimension.
self.out_proj = nn.Linear(
config.ffn_hidden_dim, config.hidden_dim, bias=False
)
def forward(self, x: Tensor) -> Tensor:
"""
Apply SwiGLU feed-forward transformation.
Parameters
----------
x : Tensor
Input of shape ``(batch, seq_len, hidden_dim)``.
Returns
-------
Tensor
Output of shape ``(batch, seq_len, hidden_dim)``.
"""
# gate: (B, T, ffn_hidden_dim)
# val: (B, T, ffn_hidden_dim)
gate = self.gate_proj(x)
val = self.val_proj(x)
# SwiGLU: swish(gate) ⊙ val
# F.silu is swish: silu(z) = z * sigmoid(z)
hidden = F.silu(gate) * val # (B, T, ffn_hidden_dim)
# Project back to hidden_dim
return self.out_proj(hidden) # (B, T, hidden_dim)
def extra_repr(self) -> str:
return (
f"hidden_dim={self.hidden_dim}, "
f"ffn_hidden_dim={self.ffn_hidden_dim}"
)

514
cyberslm/model/model.py Normal file
View File

@@ -0,0 +1,514 @@
"""
CyberSLM — Complete Decoder-Only Language Model
================================================
Assembles all components into the full model:
Token Embedding → Decoder Stack (×12) → Final RMSNorm → LM Head
Weight tying
------------
The LM head (output projection that maps hidden_dim → vocab_size) shares its
weight matrix with the token embedding (vocab_size × hidden_dim).
Mathematical justification: both the embedding matrix E and the unembedding
matrix U operate in the same semantic space. Tying U = Eᵀ forces consistency
("a token's output representation should be similar to its input representation"),
reduces parameters by vocab_size × hidden_dim = 32 000 × 384 ≈ 12.3 M, and
empirically improves perplexity on small models.
Parameter count breakdown
--------------------------
Component Params
---------------------------------------- ------
Token embedding (vocab × hidden) 12 288 000
↳ shared with LM head (no extra cost) 0
Decoder blocks × 12:
attn_norm (RMSNorm) 384 ×12 = 4 608
attn Q/K/V/O proj 589 824 ×12 = 7 077 888
ffn_norm (RMSNorm) 384 ×12 = 4 608
ffn gate/val/out 1 179 648 ×12 = 14 155 776
Final RMSNorm 384
---------------------------------------- ------
Total ≈ 33 531 264 (≈33.53 M)
Note: RoPE buffers and causal mask buffers are NOT parameters.
Initialisation
--------------
- Embeddings: N(0, 0.02) — small but non-zero, standard practice.
- All linear weights: N(0, 0.02)
- All RMSNorm γ: 1.0 (already set by RMSNorm.__init__)
- Output projection weights = Embedding weights (weight tying).
- Scaled output projections: attention o_proj and FFN out_proj are
scaled by 1/√(2·num_layers) to prevent residual stream variance
from growing with depth (following GPT-2 / LLaMA init practice).
"""
from __future__ import annotations
import math
from typing import Dict, List, Optional, Tuple
import torch
import torch.nn as nn
from torch import Tensor
from cyberslm.model.config import CyberSLMConfig, default_config
from cyberslm.model.block import DecoderBlock
from cyberslm.model.norm import RMSNorm
from cyberslm.model.rope import RotaryPositionEmbedding
class CyberSLM(nn.Module):
"""
CyberSLM Decoder-Only Transformer.
Parameters
----------
config : CyberSLMConfig
Validated model configuration.
Attributes
----------
config : CyberSLMConfig
embedding : nn.Embedding
Token embedding table, shape ``(vocab_size, hidden_dim)``.
layers : nn.ModuleList[DecoderBlock]
Stack of ``num_layers`` decoder blocks.
final_norm : RMSNorm
Applied to the residual stream after the last block.
lm_head : nn.Linear
Projects hidden_dim → vocab_size. Weight tied to ``embedding``.
"""
def __init__(self, config: CyberSLMConfig) -> None:
super().__init__()
config.validate()
self.config = config
# ------------------------------------------------------------------ #
# Token embedding #
# ------------------------------------------------------------------ #
self.embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
# ------------------------------------------------------------------ #
# Decoder stack #
# ------------------------------------------------------------------ #
# One RoPE table shared by every layer (identical by construction).
self.rope = RotaryPositionEmbedding(
head_dim=config.head_dim,
max_seq_len=config.max_seq_len,
base=config.rope_base,
)
self.layers = nn.ModuleList(
[
DecoderBlock(config, layer_idx=i, rope=self.rope)
for i in range(config.num_layers)
]
)
# ------------------------------------------------------------------ #
# Final normalisation #
# ------------------------------------------------------------------ #
self.final_norm = RMSNorm(config.hidden_dim, eps=config.norm_eps)
# ------------------------------------------------------------------ #
# Language model head #
# ------------------------------------------------------------------ #
# bias=False: unembedding never needs a bias term.
self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
# ------------------------------------------------------------------ #
# Weight tying: lm_head.weight ≡ embedding.weight #
# ------------------------------------------------------------------ #
if config.tie_weights:
self.lm_head.weight = self.embedding.weight
# ------------------------------------------------------------------ #
# Parameter initialisation #
# ------------------------------------------------------------------ #
self._init_weights()
# ---------------------------------------------------------------------- #
# Initialisation #
# ---------------------------------------------------------------------- #
def _init_weights(self) -> None:
"""
Initialise all parameters with production-quality values.
Strategy
--------
- Embedding : N(0, 0.02)
- All Linear weights : N(0, 0.02)
- o_proj and out_proj: scaled down by 1/√(2·L) where L = num_layers
to stabilise the residual stream variance at initialisation.
- All RMSNorm γ : 1.0 (already set in RMSNorm.__init__)
- All biases : 0.0 (none exist in this config)
"""
std = 0.02
scaled_std = std / math.sqrt(2.0 * self.config.num_layers)
for name, module in self.named_modules():
if isinstance(module, nn.Embedding):
nn.init.normal_(module.weight, mean=0.0, std=std)
elif isinstance(module, nn.Linear):
# Scaled init for residual output projections.
if name.endswith("o_proj") or name.endswith("out_proj"):
nn.init.normal_(module.weight, mean=0.0, std=scaled_std)
else:
nn.init.normal_(module.weight, mean=0.0, std=std)
if module.bias is not None:
nn.init.zeros_(module.bias)
# Weight tying must be re-applied after init because _init_weights
# initialised embedding.weight; lm_head.weight already points to the
# same tensor (Python object reference), so no extra step needed.
# Verify it is still tied.
if self.config.tie_weights:
assert self.lm_head.weight is self.embedding.weight, (
"Weight tying broken after _init_weights"
)
# ---------------------------------------------------------------------- #
# Forward pass #
# ---------------------------------------------------------------------- #
def forward(
self,
input_ids: Tensor,
attention_mask: Optional[Tensor] = None,
return_all_attn_weights: bool = False,
) -> Tuple[Tensor, List[Optional[Tensor]]]:
"""
Run the full forward pass.
Parameters
----------
input_ids : Tensor
Long tensor of shape ``(batch, seq_len)`` with token IDs in
``[0, vocab_size)``.
attention_mask : Optional[Tensor]
Optional key-padding mask ``(batch, seq_len)`` (1=keep, 0=pad).
Pass this when batches contain right-padded sequences so padded
positions do not corrupt real tokens; leave ``None`` for packed,
unpadded training batches.
return_all_attn_weights : bool
If True, collect and return attention weights from every layer.
Disabled by default for training efficiency.
Returns
-------
logits : Tensor
Shape ``(batch, seq_len, vocab_size)``. Raw (pre-softmax) scores.
all_attn_weights : List[Optional[Tensor]]
One entry per decoder block; each is either the attention weight
tensor or ``None``.
Notes
-----
For language model training the standard loss is:
loss = cross_entropy(logits[:, :-1].reshape(-1, V),
input_ids[:, 1:].reshape(-1))
where we predict the next token at every position.
"""
# ------------------------------------------------------------------ #
# 1. Token embedding #
# ------------------------------------------------------------------ #
x = self.embedding(input_ids) # (B, T, hidden_dim)
# ------------------------------------------------------------------ #
# 2. Decoder stack #
# ------------------------------------------------------------------ #
all_attn_weights: List[Optional[Tensor]] = []
for block in self.layers:
# use_cache=False -> `present` is None; forward() deliberately keeps
# its (logits, attn_weights) return signature unchanged. Cached
# decoding lives in generate() instead of overloading this method.
x, attn_w, _ = block(
x,
attention_mask=attention_mask,
return_attn_weights=return_all_attn_weights,
)
all_attn_weights.append(attn_w)
# ------------------------------------------------------------------ #
# 3. Final normalisation #
# ------------------------------------------------------------------ #
x = self.final_norm(x) # (B, T, hidden_dim)
# ------------------------------------------------------------------ #
# 4. LM head (weight-tied unembedding) #
# ------------------------------------------------------------------ #
logits = self.lm_head(x) # (B, T, vocab_size)
return logits, all_attn_weights
# ---------------------------------------------------------------------- #
# Cached autoregressive generation #
# ---------------------------------------------------------------------- #
@torch.no_grad()
def generate(
self,
input_ids: Tensor,
max_new_tokens: int = 256,
temperature: float = 0.0,
top_k: int = 0,
top_p: float = 1.0,
repetition_penalty: float = 1.0,
eos_id: Optional[int] = None,
) -> Tensor:
"""
Generate continuations using a KV cache.
Why this exists
---------------
The previous generation loop re-ran the whole 12-layer stack over the
entire prefix for every single token, making decoding O(n^2) in
sequence length. With a cache each step attends over the cached keys and
only computes the new token, which is O(n) overall.
Sampling is applied per row, so batched prompts are supported. Rows that
have emitted ``eos_id`` are frozen (further tokens are forced to
``eos_id``) and generation stops once every row is finished.
Parameters
----------
input_ids : Tensor ``(batch, prompt_len)`` of token ids.
temperature : 0.0 selects greedy argmax; >0 samples.
top_k / top_p : 0 and 1.0 respectively disable the filter.
repetition_penalty : >1.0 divides logits of already-present tokens.
eos_id : stop token; ``None`` means never stop early.
Returns
-------
Tensor ``(batch, prompt_len + generated)`` including the prompt.
"""
self.eval()
device = input_ids.device
B = input_ids.size(0)
max_ctx = self.config.max_seq_len
if input_ids.size(1) >= max_ctx:
input_ids = input_ids[:, -(max_ctx - 1):]
caches: List[Optional[tuple]] = [None] * len(self.layers)
finished = torch.zeros(B, dtype=torch.bool, device=device)
out = input_ids
cur = input_ids
for _ in range(max_new_tokens):
if out.size(1) >= max_ctx:
break
h = self.embedding(cur)
new_caches = []
for block, layer_cache in zip(self.layers, caches):
h, _, present = block(h, kv_cache=layer_cache, use_cache=True)
new_caches.append(present)
caches = new_caches
logits = self.lm_head(self.final_norm(h))[:, -1, :].float()
if repetition_penalty != 1.0:
for b in range(B):
seen = torch.unique(out[b])
lg = logits[b, seen]
# Divide positives, multiply negatives, so the penalty always
# pushes a token DOWN regardless of its logit's sign.
logits[b, seen] = torch.where(
lg > 0, lg / repetition_penalty, lg * repetition_penalty
)
if temperature == 0.0:
nxt = logits.argmax(dim=-1)
else:
logits = logits / temperature
if top_k > 0:
k = min(top_k, logits.size(-1))
thresh = torch.topk(logits, k, dim=-1).values[:, -1, None]
logits = logits.masked_fill(logits < thresh, float("-inf"))
if top_p < 1.0:
srt, idx = torch.sort(logits, descending=True, dim=-1)
probs = torch.softmax(srt, dim=-1)
cum = probs.cumsum(dim=-1) - probs # prob mass strictly before this token
srt = srt.masked_fill(cum > top_p, float("-inf"))
logits = torch.full_like(logits, float("-inf")).scatter(1, idx, srt)
nxt = torch.multinomial(torch.softmax(logits, dim=-1), 1).squeeze(-1)
if eos_id is not None:
nxt = torch.where(finished, torch.full_like(nxt, eos_id), nxt)
finished = finished | (nxt == eos_id)
cur = nxt.unsqueeze(1)
out = torch.cat([out, cur], dim=1)
if eos_id is not None and bool(finished.all()):
break
return out
# ---------------------------------------------------------------------- #
# Convenience: next-token logits #
# ---------------------------------------------------------------------- #
def get_next_token_logits(self, input_ids: Tensor) -> Tensor:
"""
Return logits for the next token after the last input position.
Parameters
----------
input_ids : Tensor
Shape ``(batch, seq_len)``.
Returns
-------
Tensor
Shape ``(batch, vocab_size)``.
"""
logits, _ = self.forward(input_ids)
return logits[:, -1, :] # (B, vocab_size)
# --------------------------------------------------------------------------- #
# Parameter counting #
# --------------------------------------------------------------------------- #
def count_parameters(model: nn.Module) -> Dict[str, int]:
"""
Count trainable and total parameters.
Because of weight tying, lm_head.weight is counted only once
(it shares storage with embedding.weight).
Parameters
----------
model : nn.Module
Returns
-------
dict with keys:
``total`` — total parameter elements (no double-counting)
``trainable`` — trainable parameter elements
"""
seen: set = set()
total = 0
trainable = 0
for param in model.parameters():
# data_ptr() is unique per underlying storage tensor.
if param.data_ptr() in seen:
continue
seen.add(param.data_ptr())
n = param.numel()
total += n
if param.requires_grad:
trainable += n
return {"total": total, "trainable": trainable}
# --------------------------------------------------------------------------- #
# Model summary #
# --------------------------------------------------------------------------- #
def model_summary(model: CyberSLM) -> str:
"""
Return a human-readable model summary string.
Parameters
----------
model : CyberSLM
Returns
-------
str
Formatted multi-line summary including per-component parameter counts.
"""
cfg = model.config
param_info = count_parameters(model)
lines = [
"=" * 60,
f" CyberSLM Model Summary",
"=" * 60,
f" Architecture : Decoder-only Transformer",
f" Hidden dim : {cfg.hidden_dim}",
f" Num layers : {cfg.num_layers}",
f" Num heads : {cfg.num_heads}",
f" Head dim : {cfg.head_dim}",
f" FFN hidden dim : {cfg.ffn_hidden_dim}",
f" Vocab size : {cfg.vocab_size:,}",
f" Max seq len : {cfg.max_seq_len:,}",
f" Weight tied : {cfg.tie_weights}",
f" RoPE base : {cfg.rope_base}",
"-" * 60,
f" Total params : {param_info['total']:>14,}",
f" Trainable : {param_info['trainable']:>14,}",
"-" * 60,
" Per-component:",
]
# Embedding
emb_p = model.embedding.weight.numel()
lines.append(f" Embedding : {emb_p:>12,}")
# Per-block breakdown (just the first block, all identical)
block = model.layers[0]
attn_norm_p = sum(p.numel() for p in block.attn_norm.parameters())
attn_p = sum(p.numel() for p in block.attn.parameters())
ffn_norm_p = sum(p.numel() for p in block.ffn_norm.parameters())
ffn_p = sum(p.numel() for p in block.ffn.parameters())
block_total = attn_norm_p + attn_p + ffn_norm_p + ffn_p
lines.append(f" Decoder block (×{cfg.num_layers:2d}) : {block_total:>12,} per block")
lines.append(f" attn_norm : {attn_norm_p:>12,}")
lines.append(f" attention (Q/K/V/O) : {attn_p:>12,}")
lines.append(f" ffn_norm : {ffn_norm_p:>12,}")
lines.append(f" ffn (gate/val/out) : {ffn_p:>12,}")
lines.append(f" Decoder stack total : {block_total * cfg.num_layers:>12,}")
# Final norm
final_norm_p = sum(p.numel() for p in model.final_norm.parameters())
lines.append(f" Final RMSNorm : {final_norm_p:>12,}")
# LM head — note: weight tied, so 0 additional params
lm_tied_note = " (weight-tied, no extra params)" if cfg.tie_weights else ""
lines.append(f" LM head : {0:>12,}{lm_tied_note}")
lines.append("=" * 60)
return "\n".join(lines)
# --------------------------------------------------------------------------- #
# Model builder #
# --------------------------------------------------------------------------- #
def build_model(
config: Optional[CyberSLMConfig] = None,
device: Optional[torch.device] = None,
) -> CyberSLM:
"""
Build, initialise, and optionally place the CyberSLM model.
Parameters
----------
config : CyberSLMConfig, optional
Validated config. Uses :func:`default_config` if None.
device : torch.device, optional
Target device. Stays on CPU if None.
Returns
-------
CyberSLM
Fully initialised model ready for training.
"""
if config is None:
config = default_config()
else:
config.validate()
model = CyberSLM(config)
if device is not None:
model = model.to(device)
return model

114
cyberslm/model/norm.py Normal file
View File

@@ -0,0 +1,114 @@
"""
RMSNorm — Root Mean Square Layer Normalisation
===============================================
Reference: "Root Mean Square Layer Normalization" (Zhang & Sennrich, 2019)
https://arxiv.org/abs/1910.07467
Mathematical definition
-----------------------
Given an input vector **x** ∈ ℝ^d:
RMS(x) = sqrt( (1/d) * Σ xᵢ² + ε )
RMSNorm(x) = (x / RMS(x)) * γ
where γ ∈ ℝ^d is a learned per-channel scale (initialised to 1.0) and
ε > 0 is a small constant for numerical stability.
Key differences from LayerNorm
-------------------------------
- No mean subtraction (no re-centering step).
- No learned bias β (the bias-free variant).
- ~30 % fewer operations than LayerNorm, which matters across 12 layers.
- Empirically matches or exceeds LayerNorm in modern transformer training.
Numerical stability
-------------------
- The RMS is computed in float32 regardless of input dtype. This prevents
underflow/overflow when activations are in bfloat16 or float16. For our
FP32 training runs this cast is a no-op but it is correct and future-proof.
- ε = 1e-6 (default) prevents division by zero even for near-zero inputs.
- The scale γ is cast back to the input dtype before multiplication.
"""
from __future__ import annotations
import torch
import torch.nn as nn
from torch import Tensor
class RMSNorm(nn.Module):
"""
Root Mean Square Layer Normalisation without mean-centering or bias.
Parameters
----------
dim : int
Feature dimension to normalise over (last dimension of the input).
eps : float
Small constant added to the RMS denominator for numerical stability.
Defaults to 1e-6.
Shape
-----
Input : ``(*, dim)`` — any leading batch / sequence dimensions.
Output : ``(*, dim)`` — same shape as input.
Examples
--------
>>> norm = RMSNorm(384, eps=1e-6)
>>> x = torch.randn(2, 512, 384)
>>> y = norm(x)
>>> y.shape
torch.Size([2, 512, 384])
"""
def __init__(self, dim: int, eps: float = 1e-6) -> None:
super().__init__()
if dim <= 0:
raise ValueError(f"dim must be positive, got {dim}")
if eps <= 0.0:
raise ValueError(f"eps must be positive, got {eps}")
self.dim = dim
self.eps = eps
# Learned per-channel scale, initialised to 1 (identity transform).
self.weight = nn.Parameter(torch.ones(dim))
def _compute_rms(self, x: Tensor) -> Tensor:
"""
Compute the RMS over the last dimension.
Always promotes to float32 to prevent numerical issues with
reduced-precision dtypes. For FP32 training this is a no-op.
Returns
-------
Tensor
Shape ``(*, 1)`` — one RMS value per token position.
"""
return x.float().pow(2).mean(dim=-1, keepdim=True).add(self.eps).sqrt()
def forward(self, x: Tensor) -> Tensor:
"""
Normalise ``x`` by its per-token root mean square.
Parameters
----------
x : Tensor
Input of shape ``(*, dim)``.
Returns
-------
Tensor
Normalised output of the same shape and dtype as ``x``.
"""
rms = self._compute_rms(x)
# Normalise in float32, then cast back to original dtype.
x_normed = x.float() / rms
# Scale by learned weights (cast to input dtype for type safety).
return (x_normed * self.weight.float()).to(x.dtype)
def extra_repr(self) -> str:
return f"dim={self.dim}, eps={self.eps}"

291
cyberslm/model/rope.py Normal file
View File

@@ -0,0 +1,291 @@
"""
Rotary Position Embedding (RoPE)
=================================
Reference: "RoFormer: Enhanced Transformer with Rotary Position Embedding"
(Su et al., 2021) — https://arxiv.org/abs/2104.09864
Mathematical definition
-----------------------
For a query (or key) vector **q** at position ``m`` with head dimension ``d``:
1. Partition **q** into pairs: (q₁, q₂), (q₃, q₄), ..., (q_{d-1}, q_d).
2. For each pair index ``i ∈ {0, 1, ..., d/2 - 1}`` define the frequency:
θᵢ = 1 / base^(2i / d) (base = 10 000)
3. Apply a 2-D rotation to each pair at position ``m``:
R(m, θᵢ) · (q_{2i}, q_{2i+1}) =
(q_{2i} cos(m·θᵢ) − q_{2i+1} sin(m·θᵢ),
q_{2i} sin(m·θᵢ) + q_{2i+1} cos(m·θᵢ))
This is equivalent to multiplying **q** (viewed as complex numbers) by
``exp(i · m · θ)``, which preserves the inner product of relative positions:
⟨R(m)q, R(n)k⟩ depends only on (m − n),
giving translation-equivariant attention without absolute position tokens.
Efficient implementation
------------------------
The rotation can be expressed without complex arithmetic:
q_rot = [q_even · cos − q_odd · sin,
q_even · sin + q_odd · cos]
where ``q_even = q[..., 0::2]``, ``q_odd = q[..., 1::2]``.
Numerically interleaved form (even/odd) vs. split-half form (first/second half)
are equivalent; we use the interleaved form for clarity.
Precomputation
--------------
``cos`` and ``sin`` tensors of shape ``(max_seq_len, head_dim // 2)`` are
computed once and registered as non-parameter buffers so they move with the
module (CPU ↔ GPU) and are not included in ``state_dict`` checkpoints.
Stability notes
---------------
- Frequencies are computed in float64 then cast to float32 to minimise
floating-point error in ``pow`` and ``arange``.
- All rotations are executed in float32 to prevent loss of precision.
"""
from __future__ import annotations
import math
from typing import Tuple
import torch
import torch.nn as nn
from torch import Tensor
class RotaryPositionEmbedding(nn.Module):
"""
Precomputed Rotary Position Embedding cache.
Registers ``cos`` and ``sin`` buffers of shape
``(max_seq_len, head_dim // 2)`` at construction time. Applying RoPE
to a query or key tensor costs only element-wise multiplications and
additions — no matrix multiplications.
Parameters
----------
head_dim : int
Dimension of each attention head. Must be even.
max_seq_len : int
Maximum sequence length to pre-compute. Sequences longer than this
will raise an error at runtime.
base : int
RoPE base frequency (10 000 in the original paper).
Shape of ``apply``
------------------
Input : ``(batch, num_heads, seq_len, head_dim)``
Output : ``(batch, num_heads, seq_len, head_dim)``
"""
def __init__(
self,
head_dim: int,
max_seq_len: int,
base: int = 10_000,
) -> None:
super().__init__()
if head_dim <= 0 or head_dim % 2 != 0:
raise ValueError(
f"head_dim must be a positive even integer, got {head_dim}"
)
if max_seq_len <= 0:
raise ValueError(f"max_seq_len must be positive, got {max_seq_len}")
if base <= 0:
raise ValueError(f"base must be positive, got {base}")
self.head_dim = head_dim
self.max_seq_len = max_seq_len
self.base = base
# Pre-compute and register buffers (not model parameters).
cos_cache, sin_cache = self._build_cache(head_dim, max_seq_len, base)
self.register_buffer("cos_cache", cos_cache, persistent=False)
self.register_buffer("sin_cache", sin_cache, persistent=False)
@staticmethod
def _build_cache(
head_dim: int,
max_seq_len: int,
base: int,
) -> Tuple[Tensor, Tensor]:
"""
Build ``(cos, sin)`` caches of shape ``(max_seq_len, head_dim // 2)``.
Computation is performed in float64 for precision, then cast to
float32 for storage.
Returns
-------
Tuple[Tensor, Tensor]
``cos_cache`` and ``sin_cache``, each of shape
``(max_seq_len, head_dim // 2)``.
"""
half_dim = head_dim // 2
# θᵢ = 1 / base^(2i / head_dim) for i ∈ {0, ..., half_dim - 1}
# Computed in float64 to avoid precision loss in the exponent.
inv_freq = 1.0 / (
base ** (torch.arange(0, head_dim, 2, dtype=torch.float64) / head_dim)
)
# Shape: (half_dim,)
# Position indices m ∈ {0, 1, ..., max_seq_len - 1}
positions = torch.arange(max_seq_len, dtype=torch.float64)
# Shape: (max_seq_len,)
# Outer product: angles[m, i] = m * θᵢ
angles = torch.outer(positions, inv_freq)
# Shape: (max_seq_len, half_dim)
cos_cache = angles.cos().to(torch.float32)
sin_cache = angles.sin().to(torch.float32)
return cos_cache, sin_cache
@staticmethod
def _rotate_half(x: Tensor) -> Tensor:
"""
Rotate the last dimension by interleaving even/odd pairs.
For input ``x`` of shape ``(..., head_dim)``:
x_even = x[..., 0::2] (positions 0, 2, 4, ...)
x_odd = x[..., 1::2] (positions 1, 3, 5, ...)
Returns ``[-x_odd, x_even]`` interleaved back into ``(..., head_dim)``.
This is the standard rotation that implements the complex-number trick.
Parameters
----------
x : Tensor
Shape ``(..., head_dim)`` where ``head_dim`` is even.
Returns
-------
Tensor
Same shape as ``x``.
"""
x_even = x[..., 0::2] # (..., head_dim // 2)
x_odd = x[..., 1::2] # (..., head_dim // 2)
# Interleave: stack along new dim then flatten.
rotated = torch.stack([-x_odd, x_even], dim=-1)
# Shape: (..., head_dim // 2, 2) → (..., head_dim)
return rotated.flatten(start_dim=-2)
def apply(self, x: Tensor, offset: int = 0) -> Tensor:
"""
Apply Rotary Position Embeddings to ``x``.
Parameters
----------
x : Tensor
Query or key tensor of shape
``(batch, num_heads, seq_len, head_dim)``.
offset : int
Absolute position of ``x[..., 0, :]`` in the full sequence.
This is what makes incremental decoding correct. With a KV cache the
model feeds one token at a time, so ``seq_len == 1`` and the naive
``cos_cache[:1]`` would rotate every generated token as if it were at
position 0 -- destroying all positional information after the prompt.
Passing ``offset=len(cache)`` selects the true absolute position.
Returns
-------
Tensor
Rotated tensor with the same shape and dtype as ``x``.
Raises
------
ValueError
If ``seq_len`` exceeds ``max_seq_len``.
"""
seq_len = x.size(2)
if offset < 0:
raise ValueError(f"offset must be >= 0, got {offset}")
if offset + seq_len > self.max_seq_len:
raise ValueError(
f"Positions [{offset}, {offset + seq_len}) exceed RoPE cache size "
f"{self.max_seq_len}. Re-instantiate with a larger max_seq_len."
)
# Retrieve cached values for this absolute position span.
# cos_cache: (seq_len, head_dim // 2)
# sin_cache: (seq_len, head_dim // 2)
cos = self.cos_cache[offset : offset + seq_len] # type: ignore[index]
sin = self.sin_cache[offset : offset + seq_len] # type: ignore[index]
# Expand to broadcast over batch and head dimensions:
# (1, 1, seq_len, head_dim // 2) → broadcasts with (B, H, T, D/2)
cos = cos.unsqueeze(0).unsqueeze(0)
sin = sin.unsqueeze(0).unsqueeze(0)
# Interleave cos and sin to match full head_dim.
# Each of (cos, sin) has shape (1, 1, T, D/2).
# We need (1, 1, T, D) by interleaving even positions with cos,
# odd positions with sin. The _rotate_half trick handles this:
#
# x_rot = x * cos_full + rotate_half(x) * sin_full
#
# where cos_full[..., 0::2] = cos and cos_full[..., 1::2] = cos,
# i.e. each cos value applies to both the even AND its paired odd slot.
# Achieved by repeating each half-dim value into both slots.
cos_full = cos.repeat_interleave(2, dim=-1) # (1, 1, T, D)
sin_full = sin.repeat_interleave(2, dim=-1) # (1, 1, T, D)
# Work in float32 for stability, then restore original dtype.
x_fp32 = x.float()
x_rot = x_fp32 * cos_full + self._rotate_half(x_fp32) * sin_full
return x_rot.to(x.dtype)
def forward(self, x: Tensor, offset: int = 0) -> Tensor:
"""Alias for :meth:`apply` to support ``nn.Sequential`` usage."""
return self.apply(x, offset=offset)
def extra_repr(self) -> str:
return (
f"head_dim={self.head_dim}, "
f"max_seq_len={self.max_seq_len}, "
f"base={self.base}"
)
# ---------------------------------------------------------------------------
# Functional helper
# ---------------------------------------------------------------------------
def apply_rope(
q: Tensor,
k: Tensor,
rope: RotaryPositionEmbedding,
offset: int = 0,
) -> Tuple[Tensor, Tensor]:
"""
Apply the same RoPE instance to both query and key tensors.
Parameters
----------
q : Tensor
Query tensor of shape ``(batch, num_heads, seq_len, head_dim)``.
k : Tensor
Key tensor of shape ``(batch, num_heads, seq_len, head_dim)``.
rope : RotaryPositionEmbedding
Pre-built RoPE module (carries the cos/sin cache on the correct device).
Returns
-------
Tuple[Tensor, Tensor]
``(q_rot, k_rot)`` — rotated queries and keys.
"""
return rope.apply(q, offset=offset), rope.apply(k, offset=offset)

10
generation_config.json Normal file
View File

@@ -0,0 +1,10 @@
{
"_from_model_config": true,
"bos_token_id": 2,
"eos_token_id": 3,
"output_attentions": false,
"output_hidden_states": false,
"pad_token_id": 0,
"transformers_version": "5.15.1",
"use_cache": true
}

124
infer_base.py Normal file
View File

@@ -0,0 +1,124 @@
"""
Base model inference — raw text continuation.
python Final/infer_base.py --prompt "SQL injection is"
python Final/infer_base.py --interactive
The base model is a *continuation* model, not a chat model. Give it the start of
a sentence and it continues. Asking it a question will not get an answer; it
will continue the question. Use infer_chat.py for question answering.
"""
from __future__ import annotations
import argparse
import sys
import time
from pathlib import Path
import torch
_HERE = Path(__file__).resolve().parent
if str(_HERE) not in sys.path:
sys.path.insert(0, str(_HERE))
from cyberslm.model.config import CyberSLMConfig, default_config # noqa: E402
from cyberslm.model.model import build_model, count_parameters # noqa: E402
def load(ckpt: Path, device: torch.device):
payload = torch.load(ckpt, map_location=device, weights_only=False)
if isinstance(payload, dict) and "config" in payload:
cfg, state = CyberSLMConfig(**payload["config"]), payload["model_state"]
elif isinstance(payload, dict) and "model_state" in payload:
cfg, state = default_config(), payload["model_state"]
else:
cfg, state = default_config(), payload
model = build_model(cfg, device=device)
model.load_state_dict(state)
model.eval()
return model, cfg, payload
def main() -> int:
ap = argparse.ArgumentParser(description="CyberSLM base model (text continuation)")
ap.add_argument("--prompt", "-p", default="SQL injection is")
ap.add_argument("--interactive", "-i", action="store_true")
ap.add_argument("--checkpoint", "-c", default=str(_HERE / "models" / "base.pt"))
ap.add_argument("--tokenizer", default=str(_HERE / "tokenizer" / "tokenizer.model"))
ap.add_argument("--max-new-tokens", "-m", type=int, default=120)
ap.add_argument("--temperature", "-t", type=float, default=0.8,
help="0 = greedy/deterministic")
ap.add_argument("--top-k", type=int, default=50)
ap.add_argument("--top-p", type=float, default=0.95)
ap.add_argument("--repetition-penalty", type=float, default=1.15)
ap.add_argument("--device", default=None)
ap.add_argument("--model-info", action="store_true")
args = ap.parse_args()
import sentencepiece as spm
device = torch.device(args.device) if args.device else torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
ckpt = Path(args.checkpoint)
if not ckpt.exists():
print(f"Checkpoint not found: {ckpt}", file=sys.stderr)
return 1
sp = spm.SentencePieceProcessor()
sp.load(args.tokenizer)
model, cfg, payload = load(ckpt, device)
print(f"model : {ckpt.name} ({count_parameters(model)['total']:,} params)")
print(f"context: {cfg.max_seq_len} vocab: {cfg.vocab_size:,} device: {device}")
if isinstance(payload, dict) and payload.get("step"):
print(f"trained: step {payload['step']} val_loss {payload.get('val_loss'):.4f}")
if args.model_info:
return 0
def run(text: str) -> None:
ids = sp.encode(text, out_type=int)
if sp.bos_id() >= 0:
ids = [sp.bos_id()] + ids
x = torch.tensor([ids], dtype=torch.long, device=device)
t0 = time.perf_counter()
out = model.generate(
x, max_new_tokens=args.max_new_tokens, temperature=args.temperature,
top_k=args.top_k, top_p=args.top_p,
repetition_penalty=args.repetition_penalty,
eos_id=sp.eos_id() if sp.eos_id() >= 0 else None,
)
dt = time.perf_counter() - t0
new = out[0, len(ids):].tolist()
if sp.eos_id() in new:
new = new[: new.index(sp.eos_id())]
print("\n" + "-" * 66)
# Decode prompt+continuation TOGETHER. Decoding the continuation
# alone and concatenating drops the word-boundary marker on its
# first token, gluing the halves ("...is" + "a common" -> "isa").
prompt_ids = [t for t in ids if t != sp.bos_id()]
print(sp.decode(prompt_ids + new))
print("-" * 66)
print(f"{len(new)} tokens in {dt:.2f}s ({len(new)/dt if dt else 0:.1f} tok/s)\n")
if args.interactive:
print("\nBase model - type the START of a sentence, it continues it.")
print("('exit' to quit)")
while True:
try:
q = input("\n> ").strip()
except (EOFError, KeyboardInterrupt):
print("\nBye."); break
if not q:
continue
if q.lower() in {"exit", "quit", "q"}:
print("Bye."); break
run(q)
return 0
run(args.prompt)
return 0
if __name__ == "__main__":
sys.exit(main())

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:53a36d26deb9c2c05c367a78ec10c42abaf43a6591a39969ca212b8b293f8d95
size 134137152

3
models/base.pt Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:743556ae1b31a19579249832cc9334441e618546810a0788614af020bcfa8e6e
size 134159111

6
special_tokens_map.json Normal file
View File

@@ -0,0 +1,6 @@
{
"bos_token": "<bos>",
"eos_token": "<eos>",
"unk_token": "<unk>",
"pad_token": "<pad>"
}

279575
tokenizer.json Normal file
View File

File diff suppressed because it is too large Load Diff

3
tokenizer/tokenizer.model Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:0c1a2ad32fd7d576e42bad69c99312896a077913a543a25af73831e3266a68e6
size 735456

9
tokenizer_config.json Normal file
View File

@@ -0,0 +1,9 @@
{
"tokenizer_class": "PreTrainedTokenizerFast",
"model_max_length": 2048,
"bos_token": "<bos>",
"eos_token": "<eos>",
"unk_token": "<unk>",
"pad_token": "<pad>",
"clean_up_tokenization_spaces": false
}