fromziro/test-new-arch

Overview of our New Architecture

We're switching over from Qwen to a custum implementation:

1. Per-head gating - Applies a softmax gate to each attention head, allowing the model to suppress individual heads on a per-token basis.
2. Attention Res - Based on the Kimi K3 implementation. See the paper here.
3. SwiGlue - Obvious choice.
4. GQA - Same as SwiGlue.
5. RoPe - Same thing here.

The Model

2581 params, dillionv2 tokenizer, random train data.

It was trained on a .txt file filled with one singular word ('the') repeated.

We used it to test our pipeline, from training to inference.

Test inference

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

MODEL_DIR = r"fromziro/test-new-arch"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_DIR,
    dtype=torch.float16 if device == "cuda" else torch.float32,
    trust_remote_code=True,
    low_cpu_mem_usage=True,
).to(device)

model.eval()

def generate(prompt: str, max_new_tokens: int = 20) -> str:
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        output_ids = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,          # greedy decoding
            temperature=None,
            top_p=None,
            top_k=None,
            num_beams=1,
            use_cache=False,          # your model does not implement KV cache
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

print(generate("The", max_new_tokens=20))