jwl50-stanford/loglinear-gpt-hgdn-fla-130m-20k-r3

GPT-style LogLinear H-GDN FLA 130M (20k)

This repo shares the GPT-backed log-linear H-GDN checkpoint from the 20k-step Pile run m05d10y26-loglinear-gdn-fla-130m-vhd64-pile-20k-adamw-r3. It is a raw mlp-mixer Lightning checkpoint, not a standard Transformers save_pretrained bundle.

Code

Use HazyResearch/mlp-mixer main at or after commit c7c36e3bac62af2e6f8ee07337b92989f2be589c. This is the merge commit for PR #7, which adds the GPTLMHeadModel log-linear H-GDN implementation and vendored HAttention/FLA kernel wrapper.

Checkpoint

• Lightning checkpoint: lightning/step_20000.ckpt
• Hydra experiment: mlp_mixer/loglinear_gdn_fla_130m
• W&B run: https://wandb.ai/hazy-research/mlp-mixer/runs/m05d10y26-loglinear-gdn-fla-130m-vhd64-pile-20k-adamw-r3

Metrics

Metric	Value
trainer/global_step	20000
train/ppl_step	6.575155044153354
train/ppl_epoch	9.985622908551177
train/num-tokens_step	10485235712
val/ppl	9.972557950694082
val/loss	2.2998363971710205
val/num-tokens	7283367936

Architecture

• Model class: train.models.gpt.GPTLMHeadModel
• Mixer: LogLinearGatedDeltaNetFLA
• Parameters: 132,191,712
• Layers / hidden / heads: 12 / 768 / 12
• Q/K head dim: 64
• V head dim: 64
• Context length: 2048
• Residual stream: residual_in_fp32=true
• ShortConv: false
• RoPE: disabled
• Fused-kernel sanitizer: sanitize_kernel_output=true

Loading Sketch

import copy
import hydra
from hydra import compose, initialize_config_dir
from omegaconf import OmegaConf

from train.eval import load_checkpoint

config_dir = "train/configs"
with initialize_config_dir(config_dir=config_dir, version_base=None):
    cfg = compose(config_name="config.yaml", overrides=[
        "experiment=mlp_mixer/loglinear_gdn_fla_130m",
        "mode=smoke",
        "print_config=false",
        "logger.wandb.mode=disabled",
    ])
OmegaConf.set_struct(cfg, False)

model_cfg = copy.deepcopy(cfg.model)
config = hydra.utils.instantiate(model_cfg.config, _recursive_=False, _convert_="object")
del model_cfg.config
model = hydra.utils.instantiate(model_cfg, _args_=[config], _recursive_=False)

state = load_checkpoint("lightning/step_20000.ckpt", device="cpu")
model_keys = set(model.state_dict())
state = {k: v for k, v in state.items() if k in model_keys}
model.load_state_dict(state, strict=True)
model.eval()

The original checkpoint includes three Lightning metric counter tensors; the repo eval/export path filters those before strict model loading.