modern-transformer-gqa-370M-15B-slimpajama

Modern-Transformer-GQA-370M-15B-SlimPajama

This model is for research purposes only and is not intended for production use.

A modern transformer language model (373.6M parameters) with Qwen3-Next-inspired architectural features, pretrained from scratch on 15B tokens from SlimPajama.

This model incorporates several design choices from recent efficient transformer architectures: grouped-query attention (GQA), partial rotary position embeddings, output gating, QK-normalization, and zero-centered RMSNorm. It serves as a modernized attention baseline for comparison with linear attention and state-space models in the Sequence Modeling Baselines collection.

Trained with flash-linear-attention and Flame.

Usage

Requires: pip install flash-linear-attention

import torch
import fla.models  # registers the Transformer architecture with HuggingFace Auto classes

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "puigde/modern-transformer-gqa-370M-15B-slimpajama",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).cuda()
tokenizer = AutoTokenizer.from_pretrained(
    "puigde/modern-transformer-gqa-370M-15B-slimpajama"
)

inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")
output = model.generate(**inputs, max_new_tokens=50, do_sample=False)
print(tokenizer.decode(output[0], skip_special_tokens=True))

The import fla.models line registers the Transformer architecture with HuggingFace's Auto classes. Without it, from_pretrained will fail with an unknown model type error.

Architecture

Training

Tokenizer: LlamaTokenizer (from fla-hub/gla-1.3B-100B), vocab 32,000.

Evaluation

Zero-shot, lm-evaluation-harness:

RULER (needle-in-a-haystack):

Note: RULER at 4K is beyond the 2K training context and collapses for GQA. The MHA variant (modern-transformer-mha-370M-15B-slimpajama) retains partial 4K performance (S1=0.68, MK1=0.22).

Recall:

Citation

@article{yang2024fla,
  title={Gated Linear Attention Transformers with Hardware-Efficient Training},
  author={Yang, Songlin and Wang, Bailin and Shen, Yikang and Panda, Rameswar and Kim, Yoon},
  journal={arXiv preprint arXiv:2312.06635},
  year={2024}
}