SupraGDN-5M • GatedDeltaNet • Tiny-SOTA

We are introducing SupraGDN-5M, a GatedDeltaNet-architecture-based model with 5 million parameters, pretrained from scratch as a base model on a single RTX Pro 4500 SE on 5B Fineweb-Edu tokens in about ~2.5 hours.
This is the base work for future models, as the upcoming Supra3-family with our most capable tiny SOTA SLM models.
Please note, that this is an undertrained, experimental base model with no high capabilities.

Pretraining

The pretraining ran for exactly 1 epoch on the first 5B tokens of Fineweb-Edu sample-10BT on a single RTX Pro 4500 SE rented from Runpod.

• Batch Size: 1024
• Context: 256 tokens
• Pretraining tokens: 5B
• Vocab size of custom tokenizer: 6000 tokens
• GPU: RTX Pro 4500 SE on Runpod
• Total time: ~2.5 hours
• Total cost: ~$2

Final Loss

The Val Loss dropped from ~7.1 to ~3.5251 and a Val-PPL of 33.96. Same goes for the Train Loss.

Samples

PROMPT: The history of
OUTPUT: The history of Native Americans during the Civil War was much worse during the Civil War. Since then, and the culture of North America has changed, the history of Native Americans from the beginning, the story of the revolutionary colonies of America, and the history of America has changed dramatically. The history of North America in 1876 marks the twentieth-morning of America’
--------------------------------------------------------------------------------
PROMPT: In order to understand how
OUTPUT: In order to understand how the problem may interfere, most importantly, in a way that can be applied to other things in a way that is useful to society.

--------------------------------------------------------------------------------
PROMPT: The most important thing about science is
OUTPUT: The most important thing about science is that the first step in the scientific profession is to study more and more science projects on the planet. Some people may believe that science projects are a way of life but others do not realize what is going on in their entire lives.

--------------------------------------------------------------------------------

Benchmarks

How this model compares to other SOTA mini ~5M parameter models

Verdict: We can see clearly, that our GDN model is highly competitive in hard tasks like ARC-Easy, ARC-Challenge and HellaSwag, while in PIQA it's not very good. Future work - to improve.

How to run the model

You need modeling_gatedflow.py in the same folder as this code!

import json
import torch
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
from safetensors.torch import load_file
from modeling_gatedflow import GatedFlowConfig, GatedFlowForCausalLM

REPO_ID = "SupraLabs/SupraGDN-5M"

PROMPT = "The history of"
MAX_NEW_TOKENS = 50
TEMPERATURE = 0.8
TOP_K = 40

CONFIG_PATH = hf_hub_download(REPO_ID, "config.json")
WEIGHTS_PATH = hf_hub_download(REPO_ID, "model.safetensors")
TOKENIZER_PATH = hf_hub_download(REPO_ID, "tokenizer.json")

config_dict = json.load(open(CONFIG_PATH))
config = GatedFlowConfig(
    vocab_size=config_dict["vocab_size"],
    d_model=config_dict["d_model"],
    n_layers=config_dict["n_layers"],
    n_heads=config_dict["n_heads"],
    ffn_mult=config_dict["ffn_mult"],
    window=config_dict["window"],
    max_seq_len=config_dict["max_seq_len"],
)

model = GatedFlowForCausalLM(config)
state_dict = load_file(WEIGHTS_PATH)
model.load_state_dict({f"model.{k}": v for k, v in state_dict.items()}, strict=False)
model.eval()

tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
eos_id = tokenizer.token_to_id("<eos>")

input_ids = torch.tensor([tokenizer.encode(PROMPT).ids])

output_ids = model.generate(
    input_ids,
    max_new_tokens=MAX_NEW_TOKENS,
    temperature=TEMPERATURE,
    top_k=TOP_K,
    eos_token_id=eos_id,
)

print(tokenizer.decode(output_ids[0].tolist()))