Slayer139 base (pretrained, step 90,310)
Browse files- README.md +128 -0
- config.json +19 -0
- model.safetensors +3 -0
- modeling_slayer139.py +174 -0
- tokenizer.json +0 -0
README.md
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
language: [en]
|
| 3 |
+
license: apache-2.0
|
| 4 |
+
library_name: pytorch
|
| 5 |
+
tags: [slayer, small-language-model, glint-tiny-ml, fabryka-ai]
|
| 6 |
+
datasets: [SlayerLab/gollem-v5-arcmix-9b, SlayerLab/minimal-en-corpus-5b, HuggingFaceFW/fineweb-edu, mlfoundations/dclm-baseline-1.0-parquet]
|
| 7 |
+
---
|
| 8 |
+
|
| 9 |
+
# Slayer139 base
|
| 10 |
+
|
| 11 |
+
A small decoder-only English language model (139.3M parameters) trained from scratch by Fabryka AI. This is the **pretrained base** of [Slayer139 1.01](https://huggingface.co/SlayerLab/Slayer139-1.01) and [Slayer139 1.02](https://huggingface.co/SlayerLab/Slayer139-1.02): the end of the 8× RTX 5090 pretraining run (step 90,310), with **no fine-tuning** of any kind. Slayer139 1.01 is this model fine-tuned on ARC-Easy train; Slayer139 1.02 is this model after continued pretraining and a fine-tune. [Slayer139 1.01b](https://huggingface.co/SlayerLab/Slayer139-1.01b) is a separate run with the same recipe on a smaller machine, not this checkpoint.
|
| 12 |
+
|
| 13 |
+
**Author:** Arkadiusz Słota (Fabryka AI / SlayerLab). Training, evaluation and safety gates were run by the Fabryka AI agent team under his direction.
|
| 14 |
+
|
| 15 |
+
## Result
|
| 16 |
+
|
| 17 |
+
lm-evaluation-harness 0.4.13, board protocol (C2), our measurement on 1× RTX 5090; the same checkpoint measured on a second machine (AMD Radeon) gave the same scores.
|
| 18 |
+
|
| 19 |
+
| measure | Slayer139 base | 95% CI |
|
| 20 |
+
|---|---|---|
|
| 21 |
+
| BLiMP (67 tasks) | 79.41 | [79.14; 79.69] |
|
| 22 |
+
| ARC-Easy (test, 2,376) | 57.58 | [55.60; 59.55] |
|
| 23 |
+
| WikiText-2 score | 99.86 | |
|
| 24 |
+
| **Overall** | **78.95** | [78.28; 79.62] |
|
| 25 |
+
|
| 26 |
+
Overall is the mean of the three measures. Confidence intervals: item-level bootstrap (ARC-Easy and each of the 67 BLiMP tasks resampled separately, 10,000 draws). They cover test-set sampling only, not seed-to-seed variation; the run is a single seed. These are the same numbers as the "pretrained base" column in the Slayer139 1.01 and 1.01b cards.
|
| 27 |
+
|
| 28 |
+
Provenance (sha256 prefixes): lm-eval `results.json` RTX 5090 `fabe34ce`, AMD Radeon `719bd11c`; per-item samples RTX 5090 `19f78648`, AMD Radeon `00db3761`.
|
| 29 |
+
|
| 30 |
+
Numbers are our measurements with the board's protocol, not official scores.
|
| 31 |
+
|
| 32 |
+
## Model
|
| 33 |
+
|
| 34 |
+
- Architecture: 17 layers × 768, 12 heads (12 KV heads), SwiGLU, decoder-only; vocabulary 24,576 (BPE), context 1,024.
|
| 35 |
+
- Parameters: 139,279,760 unique parameters (input and output embeddings tied: one shared tensor, counted once; counting the shared tensor twice gives 158,154,128).
|
| 36 |
+
- Checkpoint: the final step 90,310 of the run (sha256 `ba223162…`), the end of the run rather than a selected step (rule written before any test number). `model.safetensors` holds the same weights as that checkpoint (205 of 205 tensors identical; the output head is the tied input embedding).
|
| 37 |
+
- Code: `modeling_slayer139.py` (inference only, plain PyTorch), the same file as in Slayer139 1.01b.
|
| 38 |
+
|
| 39 |
+
## Training
|
| 40 |
+
|
| 41 |
+
- Data: ~23.67B tokens, no repeated epoch (sources and licences below). First the ARC-MIX pool, then a top-up from the same source families in the base corpus's proportions, mixed at document level. Both stages were scanned against WikiText-2 (test, validation) and ARC-Easy/Challenge (validation, test) for shared normalized 13-grams, with short ARC questions matched exactly, and against BLiMP sentences of ≥ 6 words; matching documents were removed. ARC train was not filtered.
|
| 42 |
+
- Optimiser and schedule: Muon (hidden matrices, LR 0.0566) + AdamW (other parameters, LR 1.6971e-3), both ×√8 from the single-node recipe; WSD schedule: 250 warm-up steps, constant to step 72,248, then 1−√ decay to 10% of peak at step 90,310; batch 256 × 1,024 tokens; seed 1337.
|
| 43 |
+
- Same recipe as our single-node final (batch 32, 4× RTX 5090, ~191k tokens/s), scaled to batch 256 with LR ×√8 and the identical data order — 4.4× faster. At equal tokens the 8-GPU run had lower held-out DCLM bits-per-byte at every compared checkpoint (−0.014 at 1.1B tokens, −0.032 at 19.4B tokens).
|
| 44 |
+
- Hardware and time: 8× NVIDIA RTX 5090 (data parallel, bf16 all-reduce), median ~832k tokens/s, about 8 hours of training.
|
| 45 |
+
- Training curves: https://track.fabryka.ai/run/5a4d174e-4a9d-4cc1-9602-838a3ffbb917
|
| 46 |
+
|
| 47 |
+
## Training data
|
| 48 |
+
|
| 49 |
+
Aggregate corpus; each source keeps its upstream terms.
|
| 50 |
+
|
| 51 |
+
- **ARC-MIX pool** (`SlayerLab/gollem-v5-arcmix-9b`, 9.39B tokens before retokenization): the 15-source base `SlayerLab/minimal-en-corpus-5b` (FineWeb-Edu, DCLM-baseline, open-web-math, FineMath, StarCoderData, StackExchange, LoC public-domain books, Wikipedia, Project Gutenberg, scientific papers, UltraChat, WildChat, CC-News, tiny-textbooks, OpenSubtitles) + a FineWeb-Edu expansion + OpenStax textbooks (×4) + extra copies of FineWeb-Edu documents scored as related to ARC science topics (+2 copies for the highest score, +1 for the next). Removed before training: 4,362 documents by the WikiText-2 / ARC validation+test scan and 1,175 documents marked CC BY-NC-SA.
|
| 52 |
+
- **Top-up** (11.4M documents, 57.45 GB of text, build manifest `b73f45f5…`): FineWeb-Edu expansion (37.1%, ODC-By 1.0, plus the same dose of ARC-related copies as the pool), FineWeb-Edu (14.9%, ODC-By 1.0), DCLM-baseline (10.9%, CC BY 4.0), StackExchange via RedPajama (6.1%, CC BY-SA content), LoC public-domain books (5.4%, CC0), StarCoderData (5.1%, The Stack terms), Wikipedia 20231101.en (4.6%, CC BY-SA 3.0 / GFDL), Project Gutenberg (2.9%, US public domain), PMC Open Access commercial-use subset (2.8%, CC0 / CC BY / CC BY-SA per article; ND and NC excluded), open-web-math (2.7%, ODC-By), FineMath (2.6%, ODC-By), WildChat-1M (2.0%, ODC-By; English, non-toxic), CC-News (1.4%, licence unknown), UltraChat 200k (0.7%, MIT), tiny-textbooks (0.7%, Apache 2.0). The top-up scan removed 25,808 documents (0.23%), each together with its ARC-related copies.
|
| 53 |
+
- Licences upstream include ODC-By, CC BY 4.0, CC BY-SA (Wikipedia, StackExchange, some PMC articles), CC0 / public domain, and sources with unknown or restrictive terms. Commercial use: review upstream terms, in particular CC-News, scientific papers, StarCoderData and model-generated chat data (UltraChat, WildChat, tiny-textbooks).
|
| 54 |
+
- OpenStax textbooks (55 titles) are used under CC BY 4.0; attribution: OpenStax, Rice University (see [OpenStax attribution](#openstax-attribution) below).
|
| 55 |
+
|
| 56 |
+
## Limitations
|
| 57 |
+
|
| 58 |
+
- English only; small model; one seed.
|
| 59 |
+
- The pretraining mix deliberately upweights web documents related to ARC science topics (decontaminated against the ARC test set); ARC-Easy is therefore an in-domain benchmark for this model.
|
| 60 |
+
- Not instruction-tuned and not fine-tuned on any benchmark.
|
| 61 |
+
|
| 62 |
+
## OpenStax attribution
|
| 63 |
+
|
| 64 |
+
The training data of this model includes text extracted from the following OpenStax textbooks, each licensed under the Creative Commons Attribution 4.0 International License (CC BY 4.0), © Rice University. Download for free at https://openstax.org. The texts were obtained from the Hugging Face dataset crumb/openstax-text (revision 8f502ca45f9f05cb5673eae445b7a97a4e8c4349). Modified: text extracted, chunked and filtered (chunks matching benchmark test/validation sets were removed).
|
| 65 |
+
|
| 66 |
+
| Title (from source file name) | © year |
|
| 67 |
+
|---|---|
|
| 68 |
+
| APBiology | 2018 |
|
| 69 |
+
| APCollege Physics | 2017 |
|
| 70 |
+
| APMacroeconomics 2e | 2017 |
|
| 71 |
+
| APMicroeconomics 2e | 2017 |
|
| 72 |
+
| Algebra and Trigonometry 2e | 2021 |
|
| 73 |
+
| American Government 3e | 2021 |
|
| 74 |
+
| Anatomy and Physiology 2e | 2022 |
|
| 75 |
+
| Anatomyand Physiology | 2017 |
|
| 76 |
+
| Astronomy 2e | 2022 |
|
| 77 |
+
| Astronomy | 2018 |
|
| 78 |
+
| Biology 2e | 2020 |
|
| 79 |
+
| Business Ethics | 2018 |
|
| 80 |
+
| Chemistry 2e | 2019 |
|
| 81 |
+
| Chemistry Atoms First 2e | 2019 |
|
| 82 |
+
| College Algebra 2e | 2021 |
|
| 83 |
+
| College Algebra Corequisite Support 2e | 2021 |
|
| 84 |
+
| College Physics | 2020 |
|
| 85 |
+
| College Physics 2e | 2022 |
|
| 86 |
+
| College Physics for AP Courses 2e | 2022 |
|
| 87 |
+
| College Success | 2020 |
|
| 88 |
+
| College Success | 2023 |
|
| 89 |
+
| Concepts Biology | 2017 |
|
| 90 |
+
| Contemporary Mathematics | 2023 |
|
| 91 |
+
| Economics 2e | 2018 |
|
| 92 |
+
| Economics 3e | 2022 |
|
| 93 |
+
| Elementary Algebra 2e | 2020 |
|
| 94 |
+
| Entrepreneurship | 2020 |
|
| 95 |
+
| Intermediate Algebra 2e | 2020 |
|
| 96 |
+
| Introduction to Intellectual Property | n/a |
|
| 97 |
+
| Introduction to Philosophy | 2022 |
|
| 98 |
+
| Introduction to Political Science | 2022 |
|
| 99 |
+
| Introductionto Anthropology | 2022 |
|
| 100 |
+
| Introductionto Sociology 3e | 2021 |
|
| 101 |
+
| Introductory Business Statistics | 2018 |
|
| 102 |
+
| Introductory Statistics | 2018 |
|
| 103 |
+
| Macroeconomics 2e | 2018 |
|
| 104 |
+
| Macroeconomics 3e | 2022 |
|
| 105 |
+
| Microbiology | 2021 |
|
| 106 |
+
| Microeconomics 2e | 2018 |
|
| 107 |
+
| Microeconomics 3e | 2022 |
|
| 108 |
+
| Physics | n/a |
|
| 109 |
+
| Prealgebra 2e | 2020 |
|
| 110 |
+
| Precalculus 2e | 2021 |
|
| 111 |
+
| Preparing for College Success | 2023 |
|
| 112 |
+
| Principles Marketing | 2023 |
|
| 113 |
+
| Principlesof Finance | 2022 |
|
| 114 |
+
| Psychology 2e | 2020 |
|
| 115 |
+
| Statistics | n/a |
|
| 116 |
+
| USHistory | 2021 |
|
| 117 |
+
| University Physics Vol 1 | 2021 |
|
| 118 |
+
| University Physics Volume 2 | 2021 |
|
| 119 |
+
| University Physics Volume 3 | 2021 |
|
| 120 |
+
| World History Volume 1 | 2023 |
|
| 121 |
+
| World History Volume 2 | 2022 |
|
| 122 |
+
| Writing Guide | 2021 |
|
| 123 |
+
|
| 124 |
+
Titles licensed CC BY-NC-SA 4.0, non-English titles, and three CC BY 4.0 titles whose text contains elements marked ‘CC BY-NC-SA’ (Introduction to Business, Organizational Behavior, Principles of Management) were not used.
|
| 125 |
+
|
| 126 |
+
## Reproduce
|
| 127 |
+
|
| 128 |
+
Evaluation: lm-evaluation-harness 0.4.13, tasks `blimp`, `arc_easy`, `wikitext` (board protocol C2); per-item samples are kept for paired bootstrap. Method and full run history: paper „Same recipe, 4.4× faster on 8× RTX 5090” (Fabryka AI, in preparation).
|
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": 768,
|
| 7 |
+
"n_head": 12,
|
| 8 |
+
"n_layer": 17,
|
| 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": 24576
|
| 19 |
+
}
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f011c9b842640f23bc8530c29f223ecce9ba224b3c6433083d7dd6005b45e772
|
| 3 |
+
size 557137920
|
modeling_slayer139.py
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Slayer139 1.01 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_slayer139 import load_slayer139, generate
|
| 7 |
+
model, tok = load_slayer139(".") # directory with config.json, model.safetensors, tokenizer.json
|
| 8 |
+
print(generate(model, tok, "The capital of France is", 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_slayer139(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 |
+
state = load_file(os.path.join(path, "model.safetensors"))
|
| 149 |
+
# the output head is tied to the token embedding and stored once (tok.weight); restore the tied key for strict loading
|
| 150 |
+
if "head.weight" in state:
|
| 151 |
+
raise ValueError("model.safetensors must not contain head.weight (tied to tok.weight, stored once)")
|
| 152 |
+
state["head.weight"] = state["tok.weight"]
|
| 153 |
+
model.load_state_dict(state, strict=True)
|
| 154 |
+
return model.to(device).eval(), Tokenizer.from_file(os.path.join(path, "tokenizer.json"))
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
@torch.no_grad()
|
| 158 |
+
def generate(model, tok, prompt, max_new_tokens=64, temperature=0.0, top_k=None, seed=None):
|
| 159 |
+
"""Continue `prompt`. temperature 0 = greedy; otherwise sampling, optionally top-k."""
|
| 160 |
+
g = torch.Generator().manual_seed(seed) if seed is not None else None
|
| 161 |
+
dev = next(model.parameters()).device
|
| 162 |
+
x = torch.tensor([tok.encode(prompt).ids], device=dev)
|
| 163 |
+
for _ in range(max_new_tokens):
|
| 164 |
+
logits = model(x[:, -model.block:])[:, -1].float().cpu()
|
| 165 |
+
if temperature <= 0:
|
| 166 |
+
nxt = logits.argmax(-1, keepdim=True)
|
| 167 |
+
else:
|
| 168 |
+
logits = logits / temperature
|
| 169 |
+
if top_k:
|
| 170 |
+
kth = torch.topk(logits, top_k).values[:, -1:]
|
| 171 |
+
logits = logits.masked_fill(logits < kth, float("-inf"))
|
| 172 |
+
nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=g)
|
| 173 |
+
x = torch.cat([x, nxt.to(dev)], dim=1)
|
| 174 |
+
return tok.decode(x[0].tolist())
|
tokenizer.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|