Maggio33 commited on
Commit
98eae28
·
verified ·
1 Parent(s): d71e3ee

TokenizerTest-32m-B: final + every-10k weights, card, bpb curve

Browse files
README.md ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: cc-by-sa-4.0
3
+ language:
4
+ - pl
5
+ library_name: pytorch
6
+ tags:
7
+ - polish
8
+ - language-model
9
+ - research
10
+ - tokenizer-comparison
11
+ pipeline_tag: text-generation
12
+ ---
13
+
14
+ # TokenizerTest-32m-B (Polish research base model, tokenizer comparison — arm B)
15
+
16
+ A small Polish **base** language model trained as arm **B** of a controlled tokenizer experiment. It is a research
17
+ checkpoint for comparing tokenizers, not a model for use.
18
+
19
+ ## Companion model (tokenizer A/B experiment)
20
+
21
+ This model is one half of a **controlled tokenizer comparison**. Two models were trained with the **same architecture, training recipe, seed and the same Polish text**. The **only** difference is the tokenizer:
22
+
23
+ | arm | tokenizer | repo | final step |
24
+ |---|---|---|---|
25
+ | **A** | Fabryka AI BPE, vocab 32 000 (`60e23148`) | SlayerLab/TokenizerTest-32m-A | 89 000 |
26
+ | **B** | tokenizer by stubbornGuy, vocab 32 000 (`77eeebed`) | SlayerLab/TokenizerTest-32m-B | 95 000 |
27
+
28
+ This model: **B**.
29
+
30
+ **Shared recipe:** ~38.9M parameters (tied input/output embeddings); "32m" in the repo name is the recipe name ("Glint 32M r6"), not the parameter count. Architecture: 15 layers, d_model 384, 6 heads, context 1024, Muon optimizer, batch 32, seed 1337, trainer `train_gpt_ref_r6.py` (`0223c083`).
31
+
32
+ Both arms saw the same ~12.4 GB of Polish text (public dataset [`SlayerLab/slayer-pl-8x3b`](https://huggingface.co/datasets/SlayerLab/slayer-pl-8x3b) @ `d50df04a`, files `pack-00`, `pack-01` and `shared/{core,core_sa,legal,legal_sa}`; validation = `pack-07` after 13-gram decontamination); B splits it more finely (on the training text 3.96 vs 4.25 bytes per token incl. EOS; on validation 4.14 vs 4.41), so it needed more steps for the same bytes. For the same reason B has lower bits per token (val 4.395 vs 4.668), which says nothing about model quality.
33
+
34
+ **How to compare:** loss per token is **not** comparable across tokenizers. We compare **bits per byte** on the same bytes (lower is better):
35
+
36
+ | | A | B | B − A |
37
+ |---|---|---|---|
38
+ | validation (7 232 820 B) | **1.0583** | 1.0622 | +0.0040 (+0.37 %) |
39
+ | held-out W (1 430 217 B) | **1.0781** | 1.0883 | +0.0101 (+0.94 %) |
40
+
41
+ From ~4 B bytes onward the gap is roughly constant (val +0.003…+0.004, W +0.008…+0.010) and does not close by the end of training. Curve every 10k steps (bits per byte; note: at the same **step** B has seen fewer bytes than A):
42
+
43
+ | step | A val | B val | A W | B W |
44
+ |---|---|---|---|---|
45
+ | 10 000 | 1.1910 | 1.2009 | 1.2327 | 1.2463 |
46
+ | 20 000 | 1.1448 | 1.1523 | 1.1799 | 1.1896 |
47
+ | 30 000 | 1.1200 | 1.1278 | 1.1499 | 1.1634 |
48
+ | 40 000 | 1.1032 | 1.1109 | 1.1304 | 1.1435 |
49
+ | 50 000 | 1.0888 | 1.0969 | 1.1141 | 1.1291 |
50
+ | 60 000 | 1.0773 | 1.0850 | 1.1000 | 1.1145 |
51
+ | 70 000 | 1.0675 | 1.0752 | 1.0885 | 1.1025 |
52
+ | 80 000 | 1.0610 | 1.0677 | 1.0818 | 1.0944 |
53
+ | 89 000 | 1.0583 | — | 1.0781 | — |
54
+ | 90 000 | — | 1.0633 | — | 1.0892 |
55
+ | 95 000 | — | 1.0622 | — | 1.0883 |
56
+
57
+ The same numbers, with checkpoint and data hashes, are in `bpb_curve.json`.
58
+
59
+ **Caveat:** one seed per arm, no noise estimate. The honest reading is *"at this scale the B tokenizer brings no gain in bits per byte"*, not *"B is worse"*.
60
+
61
+ ## Training data
62
+
63
+ Public dataset [`SlayerLab/slayer-pl-8x3b`](https://huggingface.co/datasets/SlayerLab/slayer-pl-8x3b) at revision `d50df04a`; the same files for both arms:
64
+
65
+ | file | use | sha256 |
66
+ |---|---|---|
67
+ | `pack-00/web.parquet` | training (web) | `5cbb6dfa3a276a30db3c2034105ffffa75cc29b767456e7d75c305f9bd7ec562` |
68
+ | `pack-01/web.parquet` | training (web) | `802f4829f41ef83e23979d2460500907bcfd44278e300eccb5e39ac666a04c2c` |
69
+ | `shared/core.parquet` | training (core) | `ea0b4b901ace3780501b02a3972f014f77ce722315476e6f93681616958e0a99` |
70
+ | `shared/core_sa.parquet` | training (core, CC BY-SA) | `910b18eeebbc26efbbd76554d7484d016c6df692a3b6a18492b5c9ceb32549cb` |
71
+ | `shared/legal.parquet` | training (legal) | `6d4cc2f604d76452cd7e7ea52d1a04826f5e44f09bb035ba0853232fe373e74a` |
72
+ | `shared/legal_sa.parquet` | training (legal, CC BY-SA) | `a1609cd31e12e0f6fdfdb030edb538bc230fc5f5ab376d23136180338e9dbd1e` |
73
+ | `pack-07/web.parquet` | validation (13-gram decontaminated against training) | `3d93331061595949f2e2cab187cfe688a117785c048e94688808b399b3f53b67` |
74
+
75
+ ## Usage
76
+
77
+ The model is a custom PyTorch architecture (same code as `SlayerLab/GoLLeM-v6-250M`), not a `transformers` class.
78
+
79
+ | file | content |
80
+ |---|---|
81
+ | `model.safetensors` | final weights, step 95 000 (fp32, 38,866,958 parameters; output head stored as a copy of the tied embedding) |
82
+ | `config.json` | architecture |
83
+ | `tokenizer.json` | arm B tokenizer by stubbornGuy, 32 000 tokens (`tokenizers`) |
84
+ | `modeling_gollem_v6.py` | model code (inference) and a small `generate` helper |
85
+ | `step_10000/` … `step_90000/` | intermediate checkpoints every 10k steps (same three files each; weights only, no optimizer state) |
86
+ | `bpb_curve.json` | bits-per-byte curve (step, checkpoint sha256, val, W) |
87
+ | `SHA256SUMS` | sha256 of every file |
88
+
89
+ ```python
90
+ # pip install torch safetensors tokenizers huggingface_hub
91
+ from huggingface_hub import snapshot_download
92
+ import sys
93
+
94
+ path = snapshot_download("SlayerLab/TokenizerTest-32m-B") # pin revision="<commit>" for reproducibility
95
+ sys.path.insert(0, path)
96
+ from modeling_gollem_v6 import load_gollem_v6, generate
97
+
98
+ model, tok = load_gollem_v6(path) # final checkpoint
99
+ print(generate(model, tok, "Kraków jest miastem", max_new_tokens=40, temperature=0.8, top_k=50, seed=1))
100
+ model_10k, _ = load_gollem_v6(f"{path}/step_10000") # an intermediate checkpoint
101
+ ```
102
+
103
+ Loading any `model.safetensors` of this repository into `modeling_gollem_v6.py` reproduces the logits of the
104
+ corresponding training checkpoint exactly (checked with `torch.equal` for every step).
105
+
106
+ ## Limitations
107
+
108
+ - Small **base** model (~38.9M parameters): it continues Polish text; it does not answer questions or follow
109
+ instructions. Limited knowledge; may reproduce errors and biases of web text. Research use only.
110
+ - Single run, single seed per arm; see the caveat above.
111
+
112
+ ## License
113
+
114
+ **Weights: CC BY-SA 4.0** (the training text includes CC BY-SA parts of `SlayerLab/slayer-pl-8x3b`). Attribution:
115
+ TokenizerTest-32m-B, Fabryka AI (SlayerLab), link to this repository; derivative weights under the same licence.
116
+ The tokenizer (`tokenizer.json`) was created by stubbornGuy and is released here under CC BY-SA 4.0 with his consent.
117
+ Training data keep their upstream licences; see the dataset card of `SlayerLab/slayer-pl-8x3b`.
SHA256SUMS ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ a0e944099340c90223f7e02b166d317e87f699b7cf811a279c9e6bbb9d39f475 README.md
2
+ 5a52c0867ea28be8ef8506f830d803ce98fcdea59fb654c8f4fe6ffa75c95e90 bpb_curve.json
3
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 config.json
4
+ 1a2fd327fc4069fe739e4152810acbfb66d23cf873e674d55dc85c839149d354 model.safetensors
5
+ ee2ea8c4e957413f3a7c1e8fe93a0a0b4431a2bcabaa6cc5ceeaea103fcd58ee modeling_gollem_v6.py
6
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_10000/config.json
7
+ 4fa3ae6e45f848b2c386f16c9c73a12bfcd52d5cf93eb348550a5a23dcfd5456 step_10000/model.safetensors
8
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_10000/tokenizer.json
9
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_20000/config.json
10
+ b5c10b7a8a4c3ce36aefa06de0279b07201233360304607d4eace1cd0e9be8c2 step_20000/model.safetensors
11
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_20000/tokenizer.json
12
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_30000/config.json
13
+ ebd312799b49497d878f55e7aaeb99cfb10e7ead11b56997f51617ff82500767 step_30000/model.safetensors
14
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_30000/tokenizer.json
15
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_40000/config.json
16
+ d7e74c53c9f62d453dfe43966811cad432e389b8c7b1115d749d16b4dfc68fb7 step_40000/model.safetensors
17
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_40000/tokenizer.json
18
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_50000/config.json
19
+ e663da4de9c9c353a1bfb1fd2b94f41eb52c0928b89be5765c8bf4b70627a75e step_50000/model.safetensors
20
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_50000/tokenizer.json
21
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_60000/config.json
22
+ 84b3c8517269a5f1168e58299d4695a60082eb7212a2e3b1bdcfb3828554ec28 step_60000/model.safetensors
23
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_60000/tokenizer.json
24
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_70000/config.json
25
+ e4870e7e196d98047efe559f6e3bca7ebd88a81f50ccf2994b89b256ff4ddfee step_70000/model.safetensors
26
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_70000/tokenizer.json
27
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_80000/config.json
28
+ d63c02503612f59a2a039c304028664403d32b446729418584c4e955aa73bde4 step_80000/model.safetensors
29
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_80000/tokenizer.json
30
+ 525c8d131963e6dc92f798bc602e4367e5a26fc274016ddb08b8a87b54f7fa51 step_90000/config.json
31
+ 7f56f3bdca911cf3be9958107160aa49e6a48b0fd83db463630a2ba700738371 step_90000/model.safetensors
32
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 step_90000/tokenizer.json
33
+ 77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7 tokenizer.json
bpb_curve.json ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "definition": "bits/byte = sum of NLL over text tokens / ln 2 / UTF-8 bytes; non-overlapping 1024-token windows; EOS targets excluded",
3
+ "trainer_sha256": "0223c0830ed4ae323f6424038a2de7bfb6b425653f3591dd55844f95706f5d17",
4
+ "tokenizer_sha256": "77eeebed509827f18b859ff159f29b2405d24d5b84381d6226f22e4fb6ca98d7",
5
+ "sets": {
6
+ "W": {
7
+ "sha256": "234e7ec21bf2236bc74425591531f2c4772f8c49ff853d57f6fa16f33ee90ac7",
8
+ "docs": 463,
9
+ "text_bytes": 1430217,
10
+ "tokens_text": 351268
11
+ },
12
+ "val": {
13
+ "sha256": "748d13898e736666d96f1c9307222c52414a9b66b06bf212e7a07e765babc668",
14
+ "docs": 2723,
15
+ "text_bytes": 7232820,
16
+ "tokens_text": 1748065
17
+ }
18
+ },
19
+ "device": "cpu",
20
+ "ckpts": [
21
+ {
22
+ "step": 10000,
23
+ "ckpt_sha256": "13a9130e6ac322c3a342c6bbcdf0838bdaa82e7e6962c0b4ab4ee1bdb677d2c8",
24
+ "W": 1.246255139078163,
25
+ "val": 1.2008817462071877
26
+ },
27
+ {
28
+ "step": 20000,
29
+ "ckpt_sha256": "78d637de184f406d22621b4512fc79da9d2b47f9ba532525aba44199505bdd5f",
30
+ "W": 1.189575099180923,
31
+ "val": 1.152323463429622
32
+ },
33
+ {
34
+ "step": 30000,
35
+ "ckpt_sha256": "2ec241ef98aaf7d54a15398bbb3d64dd2f0d3d6e0d562cd5e2aa4b7f7a468ac3",
36
+ "W": 1.163382329460513,
37
+ "val": 1.1278191393784163
38
+ },
39
+ {
40
+ "step": 40000,
41
+ "ckpt_sha256": "0bf61a0a7595c56f22af420dcd7a956a087c49100be258660fbc5ab109934b1b",
42
+ "W": 1.1434992106638442,
43
+ "val": 1.1109230609739182
44
+ },
45
+ {
46
+ "step": 50000,
47
+ "ckpt_sha256": "2c8189ed98f6bd27064205725ac1430c815cb4a718c764a57178632ff676b590",
48
+ "W": 1.1290988704930747,
49
+ "val": 1.096905167732812
50
+ },
51
+ {
52
+ "step": 60000,
53
+ "ckpt_sha256": "a24e66893894405f763a83ea28eee3432e8ccb6911f15e59715119ae1705089f",
54
+ "W": 1.1144814554280478,
55
+ "val": 1.0850306176283422
56
+ },
57
+ {
58
+ "step": 70000,
59
+ "ckpt_sha256": "2cedd2505161e340b2ef8ad79199d431ae54aec6cecac6d131f7d76dddfcea34",
60
+ "W": 1.1024713768041312,
61
+ "val": 1.0752361156563053
62
+ },
63
+ {
64
+ "step": 80000,
65
+ "ckpt_sha256": "f074891643fe2efa2bcf2e8545cd6ade5e024baf29fe47f161fcc43ac092c76c",
66
+ "W": 1.0943567744308031,
67
+ "val": 1.0677255240917687
68
+ },
69
+ {
70
+ "step": 90000,
71
+ "ckpt_sha256": "16f2089bd49d108383660a206a93551edacaa69ebe76686b65b2e32f9df5104a",
72
+ "W": 1.0892155925274738,
73
+ "val": 1.0633125873779654
74
+ },
75
+ {
76
+ "step": 95000,
77
+ "ckpt_sha256": "538505bc15e4332231b3e2958b02b7567348b3d43b8974e57f48526266879ab9",
78
+ "W": 1.0882527260062658,
79
+ "val": 1.062229520611975
80
+ }
81
+ ]
82
+ }
config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1a2fd327fc4069fe739e4152810acbfb66d23cf873e674d55dc85c839149d354
3
+ size 204636280
modeling_gollem_v6.py ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """GoLLeM v6 model code (inference only), extracted from the training script used for this run.
2
+
3
+ The module definitions are those of the training script; only the weight initialisation and the training loss were
4
+ removed and comments translated. Weights load with strict=True. Requirements: torch, safetensors, tokenizers.
5
+
6
+ from modeling_gollem_v6 import load_gollem_v6, generate
7
+ model, tok = load_gollem_v6(".") # directory with config.json, model.safetensors, tokenizer.json
8
+ print(generate(model, tok, "Najwyższym szczytem Polski jest", max_new_tokens=40))
9
+
10
+ Notes: RoPE uses the interleaved convention (pairs of adjacent channels), not rotate_half; the output head is tied to
11
+ the token embedding; context length is 1,024 tokens.
12
+ """
13
+ from __future__ import annotations
14
+
15
+ import json
16
+ import os
17
+ import types
18
+
19
+ import torch
20
+ import torch.nn as nn
21
+ import torch.nn.functional as F
22
+
23
+
24
+ class RMSNorm(nn.Module):
25
+ """RMSNorm with fp32 compute for stability."""
26
+ def __init__(self, d, eps=1e-6):
27
+ super().__init__()
28
+ self.weight = nn.Parameter(torch.ones(d))
29
+ self.eps = eps
30
+
31
+ def forward(self, x):
32
+ return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
33
+
34
+
35
+ def make_norm(d, cfg):
36
+ return RMSNorm(d, cfg.norm_eps) if cfg.norm == "rmsnorm" else nn.LayerNorm(d)
37
+
38
+
39
+ def apply_rope(x, base=100000.0):
40
+ """Parameter-free RoPE on [B, H, T, D], interleaved convention (even/odd channel pairs)."""
41
+ _, _, T, dim = x.shape
42
+ pos = torch.arange(T, device=x.device, dtype=torch.float32)
43
+ freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=x.device, dtype=torch.float32) / dim))
44
+ ang = torch.outer(pos, freq)
45
+ cos, sin = ang.cos().to(x.dtype)[None, None], ang.sin().to(x.dtype)[None, None]
46
+ even, odd = x[..., ::2], x[..., 1::2]
47
+ return torch.stack((even * cos - odd * sin, even * sin + odd * cos), dim=-1).flatten(-2)
48
+
49
+
50
+ class SwiGLU(nn.Module):
51
+ """Gated MLP: down(silu(gate(x)) * up(x))."""
52
+ def __init__(self, d, hidden):
53
+ super().__init__()
54
+ self.gate = nn.Linear(d, hidden, bias=False)
55
+ self.up = nn.Linear(d, hidden, bias=False)
56
+ self.down = nn.Linear(hidden, d, bias=False)
57
+
58
+ def forward(self, x):
59
+ return self.down(F.silu(self.gate(x)) * self.up(x))
60
+
61
+
62
+ class Block(nn.Module):
63
+ def __init__(self, d, nh, block, cfg, is_first=False):
64
+ super().__init__()
65
+ self.ln1 = make_norm(d, cfg)
66
+ self.ln2 = make_norm(d, cfg)
67
+ self.qkv = nn.Linear(d, 3 * d)
68
+ self.proj = nn.Linear(d, d)
69
+ if cfg.ffn == "swiglu":
70
+ self.mlp = SwiGLU(d, int(round(cfg.ffn_mult * d)))
71
+ else:
72
+ self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
73
+ self.nh = nh
74
+ self.d = d
75
+ self.cfg = cfg
76
+ self.is_first = is_first
77
+ if cfg.value_residual and not is_first:
78
+ self.vr_lambda = nn.Parameter(torch.zeros(1))
79
+ if cfg.qk_norm:
80
+ hd = d // nh
81
+ self.q_norm = RMSNorm(hd, cfg.norm_eps)
82
+ self.k_norm = RMSNorm(hd, cfg.norm_eps)
83
+
84
+ def forward(self, x, v0=None):
85
+ B, T, D = x.size()
86
+ h = self.ln1(x)
87
+ q, k, v = self.qkv(h).split(self.d, dim=2)
88
+ hd = D // self.nh
89
+ q = q.view(B, T, self.nh, hd).transpose(1, 2)
90
+ k = k.view(B, T, self.nh, hd).transpose(1, 2)
91
+ v = v.view(B, T, self.nh, hd).transpose(1, 2)
92
+ if self.cfg.qk_norm:
93
+ q = self.q_norm(q)
94
+ k = self.k_norm(k)
95
+ if self.cfg.pos == "rope":
96
+ q = apply_rope(q, self.cfg.rope_theta)
97
+ k = apply_rope(k, self.cfg.rope_theta)
98
+ if self.cfg.value_residual:
99
+ if self.is_first:
100
+ v0 = v
101
+ else:
102
+ v = v + self.vr_lambda * v0
103
+ y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
104
+ y = y.transpose(1, 2).contiguous().view(B, T, D)
105
+ x = x + self.proj(y)
106
+ x = x + self.mlp(self.ln2(x))
107
+ return x, v0
108
+
109
+
110
+ class GPT(nn.Module):
111
+ def __init__(self, vocab, n_layer, n_embd, n_head, block, cfg):
112
+ super().__init__()
113
+ self.cfg = cfg
114
+ self.tok = nn.Embedding(vocab, n_embd)
115
+ self.use_rope = cfg.pos == "rope"
116
+ if not self.use_rope:
117
+ self.pos = nn.Embedding(block, n_embd)
118
+ self.blocks = nn.ModuleList([Block(n_embd, n_head, block, cfg, is_first=(i == 0)) for i in range(n_layer)])
119
+ self.lnf = make_norm(n_embd, cfg)
120
+ self.head = nn.Linear(n_embd, vocab, bias=False)
121
+ self.head.weight = self.tok.weight # tied
122
+ self.block = block
123
+
124
+ def forward(self, idx):
125
+ B, T = idx.size()
126
+ x = self.tok(idx)
127
+ if not self.use_rope:
128
+ pos = torch.arange(T, device=idx.device)
129
+ x = x + self.pos(pos)[None]
130
+ v0 = None
131
+ for b in self.blocks:
132
+ x, v0 = b(x, v0)
133
+ logits = self.head(self.lnf(x))
134
+ cap = getattr(self.cfg, "logit_cap", 0.0)
135
+ if cap and cap > 0:
136
+ logits = cap * torch.tanh(logits / cap)
137
+ return logits
138
+
139
+
140
+ def load_gollem_v6(path, device="cpu"):
141
+ """Model (eval mode) and tokenizer from a directory with config.json, model.safetensors and tokenizer.json."""
142
+ from safetensors.torch import load_file
143
+ from tokenizers import Tokenizer
144
+
145
+ with open(os.path.join(path, "config.json"), encoding="utf-8") as f:
146
+ cfg = types.SimpleNamespace(**json.load(f))
147
+ model = GPT(cfg.vocab, cfg.n_layer, cfg.n_embd, cfg.n_head, cfg.block, cfg)
148
+ model.load_state_dict(load_file(os.path.join(path, "model.safetensors")), strict=True)
149
+ return model.to(device).eval(), Tokenizer.from_file(os.path.join(path, "tokenizer.json"))
150
+
151
+
152
+ @torch.no_grad()
153
+ def generate(model, tok, prompt, max_new_tokens=64, temperature=0.0, top_k=None, seed=None):
154
+ """Continue `prompt`. temperature 0 = greedy; otherwise sampling, optionally top-k."""
155
+ g = torch.Generator().manual_seed(seed) if seed is not None else None
156
+ dev = next(model.parameters()).device
157
+ x = torch.tensor([tok.encode(prompt).ids], device=dev)
158
+ for _ in range(max_new_tokens):
159
+ logits = model(x[:, -model.block:])[:, -1].float().cpu()
160
+ if temperature <= 0:
161
+ nxt = logits.argmax(-1, keepdim=True)
162
+ else:
163
+ logits = logits / temperature
164
+ if top_k:
165
+ kth = torch.topk(logits, top_k).values[:, -1:]
166
+ logits = logits.masked_fill(logits < kth, float("-inf"))
167
+ nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=g)
168
+ x = torch.cat([x, nxt.to(dev)], dim=1)
169
+ return tok.decode(x[0].tolist())
step_10000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_10000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4fa3ae6e45f848b2c386f16c9c73a12bfcd52d5cf93eb348550a5a23dcfd5456
3
+ size 204636280
step_10000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_20000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_20000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b5c10b7a8a4c3ce36aefa06de0279b07201233360304607d4eace1cd0e9be8c2
3
+ size 204636280
step_20000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_30000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_30000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ebd312799b49497d878f55e7aaeb99cfb10e7ead11b56997f51617ff82500767
3
+ size 204636280
step_30000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_40000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_40000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d7e74c53c9f62d453dfe43966811cad432e389b8c7b1115d749d16b4dfc68fb7
3
+ size 204636280
step_40000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_50000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_50000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e663da4de9c9c353a1bfb1fd2b94f41eb52c0928b89be5765c8bf4b70627a75e
3
+ size 204636280
step_50000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_60000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_60000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:84b3c8517269a5f1168e58299d4695a60082eb7212a2e3b1bdcfb3828554ec28
3
+ size 204636280
step_60000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_70000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_70000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e4870e7e196d98047efe559f6e3bca7ebd88a81f50ccf2994b89b256ff4ddfee
3
+ size 204636280
step_70000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_80000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_80000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d63c02503612f59a2a039c304028664403d32b446729418584c4e955aa73bde4
3
+ size 204636280
step_80000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
step_90000/config.json ADDED
@@ -0,0 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "block": 1024,
3
+ "ffn": "swiglu",
4
+ "ffn_mult": 2.667,
5
+ "logit_cap": 0.0,
6
+ "n_embd": 384,
7
+ "n_head": 6,
8
+ "n_layer": 15,
9
+ "norm": "rmsnorm",
10
+ "norm_eps": 1e-06,
11
+ "pos": "rope",
12
+ "qk_norm": true,
13
+ "rope_theta": 100000.0,
14
+ "tied_weights": {
15
+ "tok.weight": "head.weight"
16
+ },
17
+ "value_residual": true,
18
+ "vocab": 32000
19
+ }
step_90000/model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7f56f3bdca911cf3be9958107160aa49e6a48b0fd83db463630a2ba700738371
3
+ size 204636280
step_90000/tokenizer.json ADDED
The diff for this file is too large to render. See raw diff
 
tokenizer.json ADDED
The diff for this file is too large to render. See raw diff