JohnEnev/modern-llm-v2-base

modern-llm V2 (base, pretrained)

A ~315M-parameter decoder-only transformer, modernized from V1 and trained from scratch in raw PyTorch (no Hugging Face Trainer). This is the base pretrained model (EMA weights): fluent English completion, not instruction-tuned. It completes text; it does not follow instructions.

📖 Full writeup: github.com/JohnEnev/modern-llm

Architecture

Parameters	315,758,848
d_model	1024
Layers	24
Query heads	16
KV heads	4 (GQA, 4:1)
Attention	Differential Attention
QK-Norm	Yes
Position encoding	RoPE
Normalization	RMSNorm, Pre-LN
Feed-forward	SwiGLU
Context length	1024
Vocab size	50,304 (tiktoken GPT-2 BPE, padded)
Optimizer	Muon (2D matrices) + AdamW (rest)
Tied embeddings	Yes

How to load

from huggingface_hub import snapshot_download
from safetensors.torch import load_file
import sys
 
local_dir = snapshot_download("JohnEnev/modern-llm-v2-base")
sys.path.insert(0, local_dir)
 
from modeling.gpt import GPT, GPTConfig
 
# GPTConfig() defaults may drift over time — set the V2 shape explicitly.
config = GPTConfig(
    vocab_size=50304, d_model=1024, n_layers=24, n_heads=16, n_kv_heads=4,
    max_seq_len=1024, use_flash=True, tie_weights=True,
    use_qk_norm=True, use_diff_attn=True, use_mhc=False,
)
model = GPT(config)
 
state_dict = load_file(f"{local_dir}/model.safetensors")
model.load_state_dict(state_dict, strict=False)  # lm_head re-tied below
model.lm_head.weight = model.token_embeddings.weight
model.eval()

How to sample from it

import torch, tiktoken
 
enc = tiktoken.get_encoding("gpt2")
input_ids = torch.tensor([enc.encode("The meaning of life is")])
output = model.generate(input_ids, max_new_tokens=50, temperature=0.8, top_k=50)
print(enc.decode(output[0].tolist()))

Limitations

English only, 1024-token context. Small model, expect factual errors and limited reasoning.