Text Generation
Transformers
ONNX
Safetensors
English
cagliostro
causal-lm
language-model
base-model
pretrained-from-scratch
small-language-model
custom_code
Instructions to use bench-labs/cagliostro-v3 with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Transformers
How to use bench-labs/cagliostro-v3 with Transformers:
# Use a pipeline as a high-level helper from transformers import pipeline pipe = pipeline("text-generation", model="bench-labs/cagliostro-v3", trust_remote_code=True)# pip install -U transformers accelerate # Load model directly from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("bench-labs/cagliostro-v3", trust_remote_code=True, device_map="auto") - Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- vLLM
How to use bench-labs/cagliostro-v3 with vLLM:
Install from pip and serve model
# Install vLLM from pip: pip install vllm # Start the vLLM server: vllm serve "bench-labs/cagliostro-v3" # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:8000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "bench-labs/cagliostro-v3", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }'Use Docker
docker model run hf.co/bench-labs/cagliostro-v3
- SGLang
How to use bench-labs/cagliostro-v3 with SGLang:
Install from pip and serve model
# Install SGLang from pip: pip install sglang # Start the SGLang server: python3 -m sglang.launch_server \ --model-path "bench-labs/cagliostro-v3" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "bench-labs/cagliostro-v3", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }'Use Docker images
docker run --gpus all \ --shm-size 32g \ -p 30000:30000 \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=<secret>" \ --ipc=host \ lmsysorg/sglang:latest \ python3 -m sglang.launch_server \ --model-path "bench-labs/cagliostro-v3" \ --host 0.0.0.0 \ --port 30000 # Call the server using curl (OpenAI-compatible API): curl -X POST "http://localhost:30000/v1/completions" \ -H "Content-Type: application/json" \ --data '{ "model": "bench-labs/cagliostro-v3", "prompt": "Once upon a time,", "max_tokens": 512, "temperature": 0.5 }' - Docker Model Runner
How to use bench-labs/cagliostro-v3 with Docker Model Runner:
docker model run hf.co/bench-labs/cagliostro-v3
Add HF-format export at step 740,000 (72.745B tokens), model card and training charts
Browse files- .gitattributes +1 -0
- README.md +152 -2
- config.json +30 -0
- configuration_cagliostro.py +43 -0
- generation_config.json +8 -0
- index_curve.png +0 -0
- loss_curve.png +3 -0
- model.safetensors +3 -0
- modeling_cagliostro.py +270 -0
- tokenizer.json +0 -0
- tokenizer_config.json +7 -0
.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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
|
|
|
|
| 3 |
|
|
|
|
| 4 |
|
| 5 |
-
*
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
+

|
| 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 |
+

|
| 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
|
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 |
+
}
|