tachiwin/tiyat_alpha_arch1

Tiyat Alpha (arch1)

Custom 184M-param nnx transformer — not a standard transformers architecture, so AutoModel/FlaxAutoModelForCausalLM will not load this. Reconstruct the MiniGPT class (see the training/eval notebooks) and load weights like this:

import numpy as np, jax, flax.nnx as nnx
from flax.nnx import traversals
import json

config = json.load(open("config.json"))
model = MiniGPT(config, rngs=nnx.Rngs(0))

npz = np.load("flax_model.weights.npz")
flat = {tuple(k.split("/")): jax.numpy.array(v) for k, v in npz.items()}
pure_dict = traversals.unflatten_mapping(flat)

state = nnx.state(model)
nnx.replace_by_pure_dict(state, pure_dict)
nnx.update(model, state)

• Trained on: tachiwin/tiyat-ground-pretrain-m1024
• Tokenizer: tachiwin/tokenizer_64k
• Steps: 141,672