JuliaFluxGPT

JuliaFluxGPT

A ~23M parameter LLaMA-style decoder-only model with Grouped Query Attention (GQA), trained on classical philosophy and mathematics texts, implemented in Julia with Flux.jl.

Model Family Context

JuliaFluxGPT is the largest model in the Julia SLM collection, using a different framework (Flux.jl vs Lux.jl) and a more modern attention design (GQA):

Architecture

GPT (LLaMA-style)
+-- wte: Embedding(2000 -> 512)         [weight-tied with output projection]
+-- blocks x 8:
|   +-- ln1: RMSNorm(512)
|   +-- attn: CausalSelfAttention
|   |   +-- wq: Dense(512 -> 512)       [8 query heads, 64 dim each]
|   |   +-- wkv: Dense(512 -> 256)      [2 KV heads, 64 dim each, fused K+V]
|   |   +-- proj: Dense(512 -> 512)
|   +-- ln2: RMSNorm(512)
|   +-- ffwd: SwiGLUFFN
|       +-- w_gate: Dense(512 -> 1344)  [gate path]
|       +-- w_up: Dense(512 -> 1344)    [value path]
|       +-- w_down: Dense(1344 -> 512)
+-- ln_f: RMSNorm(512)
+-- [output: weight-tied with wte]

Grouped Query Attention (GQA)

GQA (Ainslie et al., 2023) uses fewer key-value heads than query heads, reducing KV-cache memory during inference while maintaining quality:

• 8 query heads (64 dim each) = full expressiveness in queries
• 2 KV heads (64 dim each) = 4x KV memory reduction
• 4 query heads per KV group = each KV head is shared by 4 query heads
• KV heads are repeated (expanded) to match query head count before attention computation

Attention parameter savings:

• Standard MHA: Q(512x512) + K(512x512) + V(512x512) + O(512x512) = 1,048,576
• GQA 8Q/2KV: Q(512x512) + KV(512x256) + O(512x512) = 655,360 (37% reduction)

RoPE (Rotary Position Embeddings)

Applied to Q and K after projection, before attention scores:

cos_cache, sin_cache = precompute_rope_freqs(head_dim=64, max_seq_len=256)
q_rotated = apply_rope(q, cos, sin, T)
k_rotated = apply_rope(k, cos, sin, T)

SwiGLU FFN

hidden = max(64, round_to_64(4 * 512 * 2/3)) = 1344
gate = swish(w_gate(x))
value = w_up(x)
output = w_down(gate * value)

Model Details

Training

Implementation Notes

Flux.jl vs Lux.jl

JuliaFluxGPT uses Flux.jl (implicit parameters, @layer macro) rather than Lux.jl (explicit parameters). Key differences:

Weight Tying Implementation

Weight tying is implemented in the forward pass rather than through a separate tied layer:

function (m::GPT)(idx)
    # ... forward through blocks ...
    x = m.ln_f(x)
    W = m.wte.weight           # reuse embedding weights
    out = W' * reshape(x, C, T*B)  # transpose matmul
    reshape(out, vocab_size, T, B)
end

This avoids complications with Flux.loadmodel! when loading checkpoints.

Usage

OpenAI-Compatible API

Served via JuliaFluxGPT Space:

curl -X POST https://lisamegawatts-juliafluxgpt.hf.space/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "messages": [{"role": "user", "content": "the nature of"}],
    "max_tokens": 200,
    "temperature": 0.8,
    "top_k": 40
  }'

Streaming supported with "stream": true.

Files

Checkpoint contains:

• model_state — Flux model weights
• hyperparams — Dict with vocab_size, n_embd, block_size, n_layer, n_head, n_kv_head
• step — Training step at checkpoint
• best_val_loss — Best validation loss achieved

Provenance

• Author: LisaMegaWatts
• Source: DavinciDreams/symbiogenesis
• Training notebook: juliaflux_v2.ipynb
• Training date: February 2026
• Architecture reference: LLaMA (Touvron et al., 2023) with GQA (Ainslie et al., 2023)

References

• Touvron, H., et al. (2023). LLaMA: Open and Efficient Foundation Language Models.
• Ainslie, J., et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.
• Karpathy, A. (2023). nanoGPT. GitHub repository.

Citation

@misc{juliafluxgpt2026,
  title={JuliaFluxGPT: A LLaMA-style GQA Model in Julia/Flux.jl},
  author={LisaMegaWatts},
  year={2026},
  url={https://huggingface.co/LisaMegaWatts/JuliaFluxGPT}
}

License

MIT