TobiasLogic commited on
Commit
ad9fb89
·
verified ·
1 Parent(s): 71b956f

Add HF-format export at step 740,000 (72.745B tokens), model card and training charts

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ loss_curve.png filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,5 +1,155 @@
1
- this is just for transfering checkpoints between instances, why are you here
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
 
3
 
 
4
 
5
- *PLEASE STOP LIKING THIS ONE*
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: apache-2.0
3
+ language:
4
+ - en
5
+ pipeline_tag: text-generation
6
+ library_name: transformers
7
+ tags:
8
+ - causal-lm
9
+ - language-model
10
+ - base-model
11
+ - pretrained-from-scratch
12
+ - small-language-model
13
+ datasets:
14
+ - HuggingFaceTB/smollm-corpus
15
+ - mlfoundations/dclm-baseline-1.0
16
+ - HuggingFaceTB/finemath
17
+ - nvidia/OpenMathInstruct-2
18
+ - HuggingFaceTB/smoltalk
19
+ ---
20
 
21
+ # cagliostro-v3
22
 
23
+ A 146M parameter decoder-only language model pretrained from scratch on 75B tokens of open web, synthetic textbook and mathematics data. It is the third model in the cagliostro line and the first to clear an Index of 26 on the [Open SLM Leaderboard](https://huggingface.co/spaces/AxiomicLabs/Open_SLM_Leaderboard) metric.
24
 
25
+ > **This is a mid-training checkpoint.** The run is at **72.745B of 75B tokens** and **80% of the way through** its learning rate cooldown. Roughly 2.3B tokens remain. The numbers below are measured on the weights in this repository as they stand today, not on a finished model. They should improve slightly when the run completes, and this card will be updated when it does.
26
+
27
+ ## Results
28
+
29
+ Zero-shot, measured with lm-evaluation-harness and the leaderboard's own ArithMark-3 script, on the exported float32 weights in this repository.
30
+
31
+ | Benchmark | Metric | Score |
32
+ |---|---|---:|
33
+ | HellaSwag | acc_norm | 42.42 |
34
+ | ARC-Easy | acc_norm | 54.17 |
35
+ | ARC-Challenge | acc_norm | 27.99 |
36
+ | PIQA | acc_norm | 67.46 |
37
+ | ArithMark-3 | acc_norm | 43.80 |
38
+ | **Open SLM Index** | | **26.27** |
39
+
40
+ The Index is the leaderboard's own formula, `(N(HellaSwag,25) + N(CombinedARC,25) + N(PIQA,50) + 0.65*N(ArithMark,25)) / 3.65` where `N(v,c) = 100(v-c)/(100-c)` and CombinedARC is the mean of ARC-Easy and ARC-Challenge.
41
+
42
+ For context against other models at this scale, using the leaderboard's published figures:
43
+
44
+ | Model | Params | Tokens | Index |
45
+ |---|---:|---:|---:|
46
+ | SmolLM2-135M | 135M | 2T | 27.13 |
47
+ | **cagliostro-v3** | **146M** | **72.7B** | **26.27** |
48
+ | SmolLM-135M | 135M | 600B | 25.74 |
49
+ | GPT-X2.5-135M | 135M | n/a | 25.17 |
50
+ | Haidass1.5-143M | 143M | n/a | 25.07 |
51
+ | BananaMind-2-Pro | 139M | 100B | 24.96 |
52
+
53
+ ![Index during training](index_curve.png)
54
+
55
+ The flat stretch from 30B to 64B is not a stall. A warmup-stable-decay schedule holds the learning rate at its peak for the first 85% of the run, and constant peak learning rate depresses zero-shot multiple-choice accuracy even while validation loss keeps improving. The recovery from 22.0 to 26.3 is what the cooldown unlocks, not new knowledge appearing in the last 9B tokens.
56
+
57
+ ![Training loss](loss_curve.png)
58
+
59
+ The sharp drop in training loss at 63.75B is the data mixture changing, not the model improving. Validation loss on a fixed held-out set is the honest line, and it moves smoothly.
60
+
61
+ ## Model details
62
+
63
+ | Field | Value |
64
+ |---|---:|
65
+ | Parameters | 146,352,000 |
66
+ | Non-embedding parameters | 85.7% |
67
+ | Layers | 30 |
68
+ | Hidden size | 640 |
69
+ | Intermediate size | 1,536 |
70
+ | Attention heads | 10 |
71
+ | Key/value heads | 5 |
72
+ | Attention | Grouped query attention with cross-head subspace attenuation |
73
+ | Activation | SwiGLU |
74
+ | Normalization | RMSNorm, eps 1e-6 |
75
+ | Positional encoding | RoPE, theta 100,000 |
76
+ | Context length | 2,048 |
77
+ | Vocabulary | 32,768 BPE |
78
+ | Embeddings | Tied input and output |
79
+ | Logit cap | 15.0 |
80
+ | Weights | float32 safetensors |
81
+
82
+ The architecture is defined in this repository. `trust_remote_code=True` is required because `CagliostroForCausalLM` is not part of transformers.
83
+
84
+ ## Training data
85
+
86
+ Two mixtures. The first covers the stable phase, the second takes over when the cooldown begins at 85% of the run.
87
+
88
+ | Source | Stable phase | Cooldown |
89
+ |---|---:|---:|
90
+ | FineWeb-Edu (deduplicated) | 43.7% | 37.0% |
91
+ | DCLM-Baseline | 28.3% | 5.0% |
92
+ | Cosmopedia v2 | 16.0% | 25.0% |
93
+ | FineMath 3+ | 5.0% | 15.0% |
94
+ | OpenMathInstruct-2 | 3.0% | 13.0% |
95
+ | InfiWebMath 3+ | 2.0% | 0.0% |
96
+ | SmolTalk | 2.0% | 5.0% |
97
+
98
+ Mathematics rises from 10% to 28% during the cooldown. No source exceeds 0.4 epochs across the full run, so nothing is repeated enough to memorize.
99
+
100
+ ## Training setup
101
+
102
+ | Field | Value |
103
+ |---|---:|
104
+ | Optimizer | AdamW, weight decay 0.01 |
105
+ | Schedule | Warmup-stable-decay |
106
+ | Warmup | 2,000 steps |
107
+ | Stable phase | steps 2,000 to 648,498 at peak learning rate |
108
+ | Cooldown | 114,441 steps, cosine to zero |
109
+ | Tokens per step | 98,304 |
110
+ | Total steps | 762,939 |
111
+ | Precision | bfloat16 with float32 master weights |
112
+ | Hardware | one RTX 5090 |
113
+ | Throughput | 90,000 to 103,000 tokens per second |
114
+ | Wall clock | about 9 days |
115
+
116
+ ## Usage
117
+
118
+ ```python
119
+ import torch
120
+ from transformers import AutoModelForCausalLM, AutoTokenizer
121
+
122
+ model_id = "bench-labs/cagliostro-v3"
123
+ tok = AutoTokenizer.from_pretrained(model_id)
124
+ model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, dtype=torch.float32)
125
+
126
+ ids = tok("The capital of France is", return_tensors="pt")
127
+ out = model.generate(**ids, max_new_tokens=32, do_sample=False)
128
+ print(tok.decode(out[0], skip_special_tokens=True))
129
+ ```
130
+
131
+ This is a base model with no instruction tuning and no chat template. It completes text.
132
+
133
+ ## Reproducing the evaluation
134
+
135
+ ```
136
+ pip install lm-eval
137
+ python -m lm_eval --model hf \
138
+ --model_args pretrained=bench-labs/cagliostro-v3,dtype=float32,trust_remote_code=True \
139
+ --tasks hellaswag,arc_easy,arc_challenge,piqa \
140
+ --num_fewshot 0 --batch_size 8 --device cuda:0
141
+ ```
142
+
143
+ ArithMark-3 uses the script linked from the leaderboard, pointed at the same model id. Evaluate in float32. A bfloat16 round trip moves logits by about 1.3e-1 at this logit cap, which is enough to change borderline multiple-choice answers, while float32 agrees with the training weights to 3.2e-05.
144
+
145
+ ## Provenance
146
+
147
+ This repository holds the full training history. A checkpoint was pushed every 30 minutes from the first step, giving 363 commits between 11 and 20 September 2026. Any intermediate checkpoint can be retrieved by revision, and the benchmark trajectory in the chart above can be reproduced from them.
148
+
149
+ ## Limitations
150
+
151
+ English only. 2,048 token context. No instruction tuning, no safety tuning, no RLHF. At 146M parameters it confabulates freely and should not be relied on for factual questions. The mathematics ability measured by ArithMark is arithmetic and short symbolic work, not general mathematical reasoning.
152
+
153
+ ## License
154
+
155
+ Apache-2.0. The training data is drawn from FineWeb-Edu and FineMath (ODC-By), DCLM-Baseline and OpenMathInstruct-2 (CC-BY-4.0), Cosmopedia v2 and SmolTalk (Apache-2.0).
config.json ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "CagliostroForCausalLM"
4
+ ],
5
+ "bos_token_id": 0,
6
+ "canon": false,
7
+ "dtype": "float32",
8
+ "eos_token_id": 0,
9
+ "hidden_size": 640,
10
+ "intermediate_size": 1536,
11
+ "layer_repeat": 1,
12
+ "logit_cap": 15.0,
13
+ "max_position_embeddings": 2048,
14
+ "mlp_type": "swiglu",
15
+ "model_type": "cagliostro",
16
+ "num_attention_heads": 10,
17
+ "num_hidden_layers": 30,
18
+ "num_key_value_heads": 5,
19
+ "rms_norm_eps": 1e-06,
20
+ "rope_theta": 100000.0,
21
+ "tie_word_embeddings": true,
22
+ "transformers_version": "5.17.0",
23
+ "value_residual": false,
24
+ "vocab_size": 32768,
25
+ "xsa": true,
26
+ "auto_map": {
27
+ "AutoConfig": "configuration_cagliostro.CagliostroConfig",
28
+ "AutoModelForCausalLM": "modeling_cagliostro.CagliostroForCausalLM"
29
+ }
30
+ }
configuration_cagliostro.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import PretrainedConfig
2
+
3
+
4
+ class CagliostroConfig(PretrainedConfig):
5
+ model_type = "cagliostro"
6
+ keys_to_ignore_at_inference = ["past_key_values"]
7
+
8
+ def __init__(
9
+ self,
10
+ vocab_size=32768,
11
+ hidden_size=640,
12
+ num_hidden_layers=30,
13
+ num_attention_heads=10,
14
+ num_key_value_heads=5,
15
+ intermediate_size=1600,
16
+ max_position_embeddings=4096,
17
+ rope_theta=100000.0,
18
+ rms_norm_eps=1e-6,
19
+ tie_word_embeddings=True,
20
+ mlp_type="swiglu",
21
+ xsa=False,
22
+ canon=False,
23
+ value_residual=False,
24
+ layer_repeat=1,
25
+ logit_cap=0.0,
26
+ **kwargs,
27
+ ):
28
+ self.vocab_size = vocab_size
29
+ self.hidden_size = hidden_size
30
+ self.num_hidden_layers = num_hidden_layers
31
+ self.num_attention_heads = num_attention_heads
32
+ self.num_key_value_heads = num_key_value_heads
33
+ self.intermediate_size = intermediate_size
34
+ self.max_position_embeddings = max_position_embeddings
35
+ self.rope_theta = rope_theta
36
+ self.rms_norm_eps = rms_norm_eps
37
+ self.mlp_type = mlp_type
38
+ self.xsa = xsa
39
+ self.canon = canon
40
+ self.value_residual = value_residual
41
+ self.layer_repeat = layer_repeat
42
+ self.logit_cap = logit_cap
43
+ super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
generation_config.json ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_from_model_config": true,
3
+ "bos_token_id": 0,
4
+ "eos_token_id": 0,
5
+ "output_attentions": false,
6
+ "output_hidden_states": false,
7
+ "transformers_version": "5.17.0"
8
+ }
index_curve.png ADDED
loss_curve.png ADDED

Git LFS Details

  • SHA256: 8e3182e05ff39bc59e0af5294facf9864a3bb9c19844e45b5cfb777267cb2909
  • Pointer size: 131 Bytes
  • Size of remote file: 112 kB
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:11e31616df91cb5b0a4a9a1c352f769f25cb4a2bf6edb4ff221c6295bb5e82a1
3
+ size 585441264
modeling_cagliostro.py ADDED
@@ -0,0 +1,270 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ from transformers import PreTrainedModel
7
+ from transformers.generation import GenerationMixin
8
+ from transformers.modeling_outputs import CausalLMOutputWithPast
9
+
10
+ try:
11
+ from .configuration_cagliostro import CagliostroConfig
12
+ except ImportError:
13
+ from configuration_cagliostro import CagliostroConfig
14
+
15
+
16
+ class RMSNorm(nn.Module):
17
+ def __init__(self, dim, eps=1e-6):
18
+ super().__init__()
19
+ self.eps = eps
20
+ self.weight = nn.Parameter(torch.ones(dim))
21
+
22
+ def forward(self, x):
23
+ dtype = x.dtype
24
+ x = x.float()
25
+ x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
26
+ return x.to(dtype) * self.weight
27
+
28
+
29
+ def rope_freqs(seq_len, head_dim, theta, device=None):
30
+ inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
31
+ t = torch.arange(seq_len, device=device).float()
32
+ freqs = torch.outer(t, inv_freq)
33
+ return torch.cat([freqs, freqs], dim=-1)
34
+
35
+
36
+ def rotate_half(x):
37
+ x1, x2 = x.chunk(2, dim=-1)
38
+ return torch.cat([-x2, x1], dim=-1)
39
+
40
+
41
+ def apply_rope(x, cos, sin):
42
+ return x * cos + rotate_half(x) * sin
43
+
44
+
45
+ class Attention(nn.Module):
46
+ def __init__(self, dim, n_heads, n_kv_heads, head_dim, norm_eps=1e-6,
47
+ xsa=False, value_residual=False):
48
+ super().__init__()
49
+ if n_heads % n_kv_heads != 0:
50
+ raise ValueError(
51
+ f"num_attention_heads ({n_heads}) must be divisible by "
52
+ f"num_key_value_heads ({n_kv_heads})")
53
+ self.n_heads = n_heads
54
+ self.n_kv_heads = n_kv_heads
55
+ self.head_dim = head_dim
56
+ self.n_rep = n_heads // n_kv_heads
57
+ self.xsa = xsa
58
+
59
+ self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
60
+ self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
61
+ self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
62
+ self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
63
+
64
+ self.q_norm = RMSNorm(head_dim, eps=norm_eps)
65
+ self.k_norm = RMSNorm(head_dim, eps=norm_eps)
66
+
67
+ self.value_residual = value_residual
68
+ if value_residual:
69
+ self.v_mix = nn.Parameter(torch.full((1,), -4.0))
70
+
71
+ def forward(self, x, cos, sin, attn_mask=None, v_first=None):
72
+ B, T, _ = x.shape
73
+ q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim)
74
+ k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
75
+ v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim)
76
+
77
+ q = self.q_norm(q)
78
+ k = self.k_norm(k)
79
+ q = apply_rope(q, cos, sin)
80
+ k = apply_rope(k, cos, sin)
81
+
82
+ v_raw = v
83
+ if self.value_residual and v_first is not None:
84
+ lam = torch.sigmoid(self.v_mix)
85
+ v = (1 - lam) * v + lam * v_first
86
+
87
+ q = q.transpose(1, 2)
88
+ k = k.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
89
+ v = v.transpose(1, 2).repeat_interleave(self.n_rep, dim=1)
90
+
91
+ if attn_mask is None:
92
+ out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
93
+ else:
94
+ out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
95
+
96
+ if self.xsa:
97
+ num = (out * v).sum(dim=-1, keepdim=True)
98
+ den = (v * v).sum(dim=-1, keepdim=True).clamp_min(1e-6)
99
+ out = out - (num / den) * v
100
+
101
+ out = out.transpose(1, 2).contiguous().view(B, T, -1)
102
+ return self.o_proj(out), v_raw
103
+
104
+
105
+ class SwiGLU(nn.Module):
106
+ def __init__(self, dim, hidden):
107
+ super().__init__()
108
+ self.gate_proj = nn.Linear(dim, hidden, bias=False)
109
+ self.up_proj = nn.Linear(dim, hidden, bias=False)
110
+ self.down_proj = nn.Linear(hidden, dim, bias=False)
111
+
112
+ def forward(self, x):
113
+ return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
114
+
115
+
116
+ class ReLU2MLP(nn.Module):
117
+ def __init__(self, dim, hidden):
118
+ super().__init__()
119
+ self.up_proj = nn.Linear(dim, hidden, bias=False)
120
+ self.down_proj = nn.Linear(hidden, dim, bias=False)
121
+
122
+ def forward(self, x):
123
+ return self.down_proj(F.relu(self.up_proj(x)).square())
124
+
125
+
126
+ class Canon(nn.Module):
127
+ def __init__(self, dim, kernel=4):
128
+ super().__init__()
129
+ self.kernel = kernel
130
+ self.conv = nn.Conv1d(dim, dim, kernel, groups=dim, bias=False)
131
+
132
+ def forward(self, x):
133
+ y = F.pad(x.transpose(1, 2), (self.kernel - 1, 0))
134
+ return x + self.conv(y).transpose(1, 2)
135
+
136
+
137
+ class DecoderLayer(nn.Module):
138
+ def __init__(self, dim, n_heads, n_kv_heads, head_dim, mlp_hidden, norm_eps,
139
+ xsa, mlp_type, canon, value_residual):
140
+ super().__init__()
141
+ self.attn_norm = RMSNorm(dim, eps=norm_eps)
142
+ self.attn = Attention(dim, n_heads, n_kv_heads, head_dim, norm_eps,
143
+ xsa=xsa, value_residual=value_residual)
144
+ self.mlp_norm = RMSNorm(dim, eps=norm_eps)
145
+ self.mlp = ReLU2MLP(dim, mlp_hidden) if mlp_type == "relu2" else SwiGLU(dim, mlp_hidden)
146
+ self.canon = Canon(dim) if canon else None
147
+
148
+ def forward(self, x, cos, sin, attn_mask=None, v_first=None):
149
+ h = self.attn_norm(x)
150
+ if self.canon is not None:
151
+ h = self.canon(h)
152
+ a, v = self.attn(h, cos, sin, attn_mask, v_first)
153
+ x = x + a
154
+ x = x + self.mlp(self.mlp_norm(x))
155
+ return x, v
156
+
157
+
158
+ class CagliostroPreTrainedModel(PreTrainedModel):
159
+ config_class = CagliostroConfig
160
+ base_model_prefix = "model"
161
+ supports_gradient_checkpointing = False
162
+ _no_split_modules = ["DecoderLayer"]
163
+ _tied_weights_keys = {"lm_head.weight": "embed_tokens.weight"}
164
+
165
+ def _init_weights(self, module):
166
+ if isinstance(module, nn.Linear):
167
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
168
+ elif isinstance(module, nn.Embedding):
169
+ nn.init.normal_(module.weight, mean=0.0, std=0.02)
170
+
171
+
172
+ class CagliostroForCausalLM(CagliostroPreTrainedModel, GenerationMixin):
173
+ def __init__(self, config):
174
+ super().__init__(config)
175
+ dim = config.hidden_size
176
+ self.head_dim = dim // config.num_attention_heads
177
+ self.layer_repeat = max(1, config.layer_repeat)
178
+ self.logit_cap = config.logit_cap
179
+
180
+ mlp_hidden = config.intermediate_size
181
+ if config.mlp_type == "relu2":
182
+ mlp_hidden = int(round(mlp_hidden * 1.5 / 64) * 64)
183
+
184
+ self.embed_tokens = nn.Embedding(config.vocab_size, dim)
185
+ self.layers = nn.ModuleList([
186
+ DecoderLayer(
187
+ dim, config.num_attention_heads, config.num_key_value_heads,
188
+ self.head_dim, mlp_hidden, config.rms_norm_eps,
189
+ config.xsa, config.mlp_type, config.canon,
190
+ config.value_residual and i > 0,
191
+ )
192
+ for i in range(config.num_hidden_layers)
193
+ ])
194
+ self.norm_out = RMSNorm(dim, eps=config.rms_norm_eps)
195
+ self.lm_head = nn.Linear(dim, config.vocab_size, bias=False)
196
+
197
+ self.rope_theta = config.rope_theta
198
+ self._rope_cache = None
199
+
200
+ self.post_init()
201
+
202
+ def _rope(self, seq_len, device, dtype):
203
+ key = (seq_len, str(device), str(dtype))
204
+ if self._rope_cache is None or self._rope_cache[0] != key:
205
+ freqs = rope_freqs(seq_len, self.head_dim, self.rope_theta, device=device)
206
+ self._rope_cache = (key,
207
+ freqs.cos()[None, :, None, :].to(dtype),
208
+ freqs.sin()[None, :, None, :].to(dtype))
209
+ return self._rope_cache[1], self._rope_cache[2]
210
+
211
+ def get_input_embeddings(self):
212
+ return self.embed_tokens
213
+
214
+ def set_input_embeddings(self, value):
215
+ self.embed_tokens = value
216
+
217
+ def get_output_embeddings(self):
218
+ return self.lm_head
219
+
220
+ def set_output_embeddings(self, value):
221
+ self.lm_head = value
222
+
223
+ def _causal_mask(self, attention_mask, B, T, dtype, device):
224
+ causal = torch.ones(T, T, dtype=torch.bool, device=device).tril()
225
+ mask = causal[None, None, :, :].expand(B, 1, T, T)
226
+ if attention_mask is not None:
227
+ pad = attention_mask.to(torch.bool)[:, None, None, :]
228
+ mask = mask & pad
229
+ out = torch.zeros(mask.shape, dtype=dtype, device=device)
230
+ return out.masked_fill(~mask, torch.finfo(dtype).min)
231
+
232
+ def forward(self, input_ids=None, attention_mask=None, labels=None,
233
+ inputs_embeds=None, use_cache=None, output_attentions=None,
234
+ output_hidden_states=None, return_dict=None, **kwargs):
235
+ if inputs_embeds is None:
236
+ x = self.embed_tokens(input_ids)
237
+ else:
238
+ x = inputs_embeds
239
+ B, T, _ = x.shape
240
+
241
+ cos, sin = self._rope(T, x.device, x.dtype)
242
+
243
+ attn_mask = None
244
+ if attention_mask is not None and attention_mask.dim() == 2 and not bool(attention_mask.all()):
245
+ attn_mask = self._causal_mask(attention_mask, B, T, x.dtype, x.device)
246
+
247
+ v_first = None
248
+ for layer in self.layers:
249
+ for _ in range(self.layer_repeat):
250
+ x, v = layer(x, cos, sin, attn_mask, v_first)
251
+ if v_first is None:
252
+ v_first = v
253
+
254
+ x = self.norm_out(x)
255
+ logits = self.lm_head(x)
256
+ if self.logit_cap:
257
+ logits = self.logit_cap * torch.tanh(logits / self.logit_cap)
258
+
259
+ loss = None
260
+ if labels is not None:
261
+ shift_logits = logits[:, :-1].float().reshape(-1, logits.size(-1))
262
+ shift_labels = labels[:, 1:].reshape(-1)
263
+ loss = F.cross_entropy(shift_logits, shift_labels, ignore_index=-100)
264
+
265
+ if return_dict is False:
266
+ return (loss, logits) if loss is not None else (logits,)
267
+ return CausalLMOutputWithPast(loss=loss, logits=logits)
268
+
269
+ def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
270
+ return {"input_ids": input_ids, "attention_mask": attention_mask}
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer_config.json ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ {
2
+ "tokenizer_class": "PreTrainedTokenizerFast",
3
+ "bos_token": "<|endoftext|>",
4
+ "eos_token": "<|endoftext|>",
5
+ "unk_token": "<|endoftext|>",
6
+ "model_max_length": 2048
7
+ }