Kai-2.0-Base / README.md
khtsly's picture
Upload folder using huggingface_hub
32d8112
|
Raw History Blame Contribute Delete
6.27 kB
---
license: apache-2.0
library_name: transformers
language:
- en
tags:
- text-generation
- roblox
- luau
- coder
- on-device
pipeline_tag: text-generation
---
# Kai 2.0 - 0.6B (Base) 🌊
## 1. Model Introduction
Kai 2.0 is an open-weight base model for Luau and Roblox code. It is the mid-sized model in the LuauCoder 2.0 family (Moana 2B, Kai 0.6B, Nalu 0.3B). It is a dense decoder with 28 layers and 643M parameters trained on 63B tokens.
This is a base model with no chat template, so it continues raw text rather than following instructions. Prompt it with the start of the code you want.
## 2. Model Summary
<div align="center">
<table>
<tbody>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Architecture</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">Dense decoder</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Total Parameters</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">643M</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Non-Embedding Parameters</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">596M</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Number of Layers</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">28</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Attention-Layer Composition</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">21 RoPE + 7 NoPE</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Attention Mechanism</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">CCVHA + DiffAttn V2 + XSA</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Hidden Dimension</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">1408</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Number of Query Heads</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">20 (10 computed)</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Number of KV Heads</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">2</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Number of Output Heads</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">10</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Head Dimension</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">128</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Convolution Taps</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">2 depthwise + 2 per head</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Head Postmix Rank</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">2</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Residual Connections</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">Block AttnRes (7 blocks of 4 layers)</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Feed-Forward Hidden Dimension</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">3840</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Activation Function</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">SwiGLU</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Normalization</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">RMSNorm, QK-Norm</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>RoPE Theta</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">500,000</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Vocabulary Size</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">32768</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Embeddings</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">Tied</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Logit Soft Cap</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">±30</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Modality</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">Text</td>
</tr>
<tr>
<td align="center" style="vertical-align: middle; text-align: center"><strong>Context Window</strong></td>
<td align="center" style="vertical-align: middle; text-align: center">65,536</td>
</tr>
</tbody>
</table>
</div>
## 3. Architecture
![Kai 2.0 architecture](assets/architecture.png)
## 4. Usage
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "khtsly/Kai-2.0-Base"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
repo,
trust_remote_code=True,
dtype="bfloat16",
attn_implementation="sdpa", # or "flash_attention_2"
).to("cuda")
prompt = "local function lerp(a, b, t)"
inputs = tok(prompt, return_tensors="pt").to("cuda")
out = model.generate(
**inputs,
max_new_tokens=256,
do_sample=True,
temperature=0.8,
top_p=0.95,
top_k=30,
repetition_penalty=1.05,
)
print(tok.decode(out[0], skip_special_tokens=True))
```
## License
This repository and LuauCoder families model weights are released under the Apache-2.0 License.