LisaMegaWatts/MonarchSLM
0
1#=2model.jl — Self-contained inference engine for MonarchSLM3 4Implements the Monarch Mixer architecture (sub-quadratic sequence mixing using5structured matrices) with RMSNorm, SwiGLU, and weight-tied output.6No Lux dependency — parameters loaded directly from JLD2. CPU-only inference.7 8Architecture per block:9 MonarchSequenceMixer (8 heads × MonarchMatrix + CausalDepthwiseConv + LearnedGate)10 → RMSNorm pre-norm + residual11 SwiGLU FFN → RMSNorm pre-norm + residual12 13References:14 Monarch Mixer (Dao et al., 2023): Sub-quadratic GEMM-based architecture15=#16 17using NNlib18using NNlib: batched_mul19using Statistics20using Random21using JSON322using TOML23 24# ═══════════════════════════════════════════════════════════════════25# Model configuration26# ═══════════════════════════════════════════════════════════════════27 28struct ModelConfig29 arch::String30 embed_dim::Int31 n_layers::Int32 n_monarch_heads::Int33 conv_kernel_size::Int34 context_length::Int35 vocab_size::Int36 weight_tying::Bool37 bias::Bool38end39 40function load_config_toml(path::String; vocab_size::Int=0)41 raw = TOML.parsefile(path)42 m = get(raw, "model", Dict())43 return ModelConfig(44 get(m, "arch", "monarch"),45 get(m, "embed_dim", 256),46 get(m, "n_layers", 8),47 get(m, "n_monarch_heads", 8),48 get(m, "conv_kernel_size", 4),49 get(m, "context_length", 256),50 vocab_size,51 get(m, "weight_tying", true),52 get(m, "bias", false),53 )54end55 56# ═══════════════════════════════════════════════════════════════════57# Character-level tokenizer58# ═══════════════════════════════════════════════════════════════════59 60struct CharTokenizer61 char_to_idx::Dict{Char, Int}62 idx_to_char::Vector{Char}63 vocab_size::Int64end65 66function load_char_vocab_json(path::String)67 raw = JSON3.read(read(path, String))68 chars = Char[only(String(s)) for s in raw]69 char_to_idx = Dict(c => i for (i, c) in enumerate(chars))70 return CharTokenizer(char_to_idx, chars, length(chars))71end72 73function encode(t::CharTokenizer, text::String)74 indices = Int[]75 sizehint!(indices, length(text))76 for c in text77 idx = get(t.char_to_idx, c, nothing)78 idx !== nothing && push!(indices, idx)79 end80 return indices81end82 83function decode(t::CharTokenizer, indices::AbstractVector{<:Integer})84 buf = IOBuffer()85 for idx in indices86 if 1 <= idx <= t.vocab_size87 write(buf, t.idx_to_char[idx])88 end89 end90 return String(take!(buf))91end92 93# ═══════════════════════════════════════════════════════════════════94# BPE Tokenizer (GPT-2 style)95# ═══════════════════════════════════════════════════════════════════96 97struct BPETokenizer98 encoder::Dict{String, Int}99 decoder::Dict{Int, String}100 merges::Vector{Tuple{String, String}}101 merge_ranks::Dict{Tuple{String, String}, Int}102 byte_to_unicode::Dict{UInt8, Char}103 unicode_to_byte::Dict{Char, UInt8}104 vocab_size::Int105 pat::Regex106end107 108function load_bpe_tokenizer(vocab_path::String, merges_path::String)109 encoder = JSON3.read(read(vocab_path, String), Dict{String, Int})110 decoder = Dict{Int, String}(v => k for (k, v) in encoder)111 112 merge_lines = readlines(merges_path)113 start = startswith(first(merge_lines), "#") ? 2 : 1114 merges = Tuple{String, String}[]115 for line in merge_lines[start:end]116 parts = split(strip(line))117 length(parts) == 2 && push!(merges, (String(parts[1]), String(parts[2])))118 end119 merge_ranks = Dict{Tuple{String, String}, Int}(m => i for (i, m) in enumerate(merges))120 121 b2u = _build_byte_to_unicode()122 u2b = Dict{Char, UInt8}(v => k for (k, v) in b2u)123 124 pat = r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"125 126 return BPETokenizer(encoder, decoder, merges, merge_ranks, b2u, u2b,127 length(encoder), pat)128end129 130function encode(t::BPETokenizer, text::String)131 tokens = Int[]132 for m in eachmatch(t.pat, text)133 word = m.match134 encoded_chars = [string(t.byte_to_unicode[b]) for b in Vector{UInt8}(word)]135 bpe_tokens = _bpe_encode_word(encoded_chars, t.merge_ranks)136 for tok in bpe_tokens137 id = get(t.encoder, tok, nothing)138 id !== nothing && push!(tokens, id + 1)139 end140 end141 return tokens142end143 144function decode(t::BPETokenizer, ids::AbstractVector{<:Integer})145 token_strs = [get(t.decoder, id - 1, "") for id in ids]146 joined = join(token_strs)147 out = UInt8[]148 sizehint!(out, length(joined))149 for c in joined150 b = get(t.unicode_to_byte, c, nothing)151 if b !== nothing152 push!(out, b)153 else154 append!(out, codeunits(string(c)))155 end156 end157 return String(out)158end159 160function _bpe_encode_word(symbols::Vector{String}, merge_ranks::Dict{Tuple{String, String}, Int})161 while length(symbols) > 1162 best_pair = nothing163 best_rank = typemax(Int)164 for i in 1:length(symbols)-1165 pair = (symbols[i], symbols[i+1])166 rank = get(merge_ranks, pair, typemax(Int))167 if rank < best_rank168 best_rank = rank169 best_pair = pair170 end171 end172 best_rank == typemax(Int) && break173 174 new_symbols = String[]175 i = 1176 while i <= length(symbols)177 if i < length(symbols) && symbols[i] == best_pair[1] && symbols[i+1] == best_pair[2]178 push!(new_symbols, best_pair[1] * best_pair[2])179 i += 2180 else181 push!(new_symbols, symbols[i])182 i += 1183 end184 end185 symbols = new_symbols186 end187 return symbols188end189 190function _build_byte_to_unicode()191 bs = UInt8[]192 append!(bs, UInt8('!'):UInt8('~'))193 append!(bs, UInt8('¡'):UInt8('¬'))194 append!(bs, UInt8('®'):UInt8('ÿ'))195 cs = Int[Int(b) for b in bs]196 n = 0197 for b in 0x00:0xff198 if !(b in bs)199 push!(bs, b)200 push!(cs, 256 + n)201 n += 1202 end203 end204 return Dict{UInt8, Char}(b => Char(c) for (b, c) in zip(bs, cs))205end206 207# ═══════════════════════════════════════════════════════════════════208# Unified tokenizer interface209# ═══════════════════════════════════════════════════════════════════210 211const Tokenizer = Union{CharTokenizer, BPETokenizer}212 213tokenizer_vocab_size(t::CharTokenizer) = t.vocab_size214tokenizer_vocab_size(t::BPETokenizer) = t.vocab_size215 216# ═══════════════════════════════════════════════════════════════════217# Causal mask (multiplicative 0/1 for Monarch)218# ═══════════════════════════════════════════════════════════════════219 220function make_causal_mask(seq_len::Int)221 return Float32[j <= i ? 1.0f0 : 0.0f0 for i in 1:seq_len, j in 1:seq_len]222end223 224# ═══════════════════════════════════════════════════════════════════225# Layer primitives (shared with transformer)226# ═══════════════════════════════════════════════════════════════════227 228function rmsnorm_forward(x, weight; eps=1.0f-6)229 rms = sqrt.(mean(x .^ 2; dims=1) .+ eps)230 return weight .* (x ./ rms)231end232 233function swiglu_forward(x, ps)234 D = size(x, 1)235 x_flat = reshape(x, D, :)236 gate = ps.w1 * x_flat237 val = ps.v * x_flat238 hidden = NNlib.swish.(gate) .* val239 out = ps.w2 * hidden240 return reshape(out, D, size(x)[2:end]...)241end242 243# ═══════════════════════════════════════════════════════════════════244# Monarch Matrix realization245# ═══════════════════════════════════════════════════════════════════246 247"""248 monarch_realize(L1, L2, p) -> Matrix{Float32}249 250Materialize T×T Monarch matrix: M = Pᵀ · BlockDiag(L1) · P · BlockDiag(L2)251where L1, L2 are (p, p, p) block-diagonal factors and T = p².252"""253function monarch_realize(L1, L2, p::Int)254 T = p * p255 256 # Start with identity matrix257 I_T = Float32[i == j ? 1.0f0 : 0.0f0 for i in 1:T, j in 1:T]258 259 # Reshape columns: (T, T) → (p, p, T)260 x = reshape(I_T, p, p, T)261 262 # Apply L2 block-diagonal: for each block k, multiply L2[:,:,k] @ x[:,k,:]263 x = permutedims(x, (1, 3, 2)) # (p, T, p)264 x = batched_mul(L2, x) # (p, p, p) × (p, T, p) → (p, T, p)265 x = permutedims(x, (1, 3, 2)) # (p, p, T)266 267 # Permutation P: transpose the p×p grid268 x = permutedims(x, (2, 1, 3))269 270 # Apply L1 block-diagonal271 x = permutedims(x, (1, 3, 2)) # (p, T, p)272 x = batched_mul(L1, x) # (p, T, p)273 x = permutedims(x, (1, 3, 2)) # (p, p, T)274 275 # Undo permutation276 x = permutedims(x, (2, 1, 3))277 278 return reshape(x, T, T)279end280 281# ═══════════════════════════════════════════════════════════════════282# Causal Depthwise Conv1d283# ═══════════════════════════════════════════════════════════════════284 285"""286 causal_depthwise_conv1d(x, kernel) -> Array287 288x: (D, T, B), kernel: (K, D)289Causal convolution: pad K-1 zeros on the left, sum over kernel taps.290"""291function causal_depthwise_conv1d(x, kernel)292 D, T, B = size(x)293 K = size(kernel, 1)294 295 # Causal pad: K-1 zeros on the left296 pad = zeros(Float32, D, K - 1, B)297 x_padded = cat(pad, x; dims=2) # (D, T+K-1, B)298 299 # Sum over kernel taps300 out = sum(1:K) do k301 reshape(kernel[k:k, :], D, 1, 1) .* x_padded[:, k:k+T-1, :]302 end303 304 return out305end306 307# ═══════════════════════════════════════════════════════════════════308# Pre-compute inference caches (Monarch matrices + causal mask)309# ═══════════════════════════════════════════════════════════════════310 311"""312 precompute_inference_caches(config, ps) -> NamedTuple313 314Pre-realize all Monarch matrices and apply causal mask once at startup.315Avoids recomputing them on every forward pass during generation.316"""317function precompute_inference_caches(config::ModelConfig, ps)318 p = isqrt(config.context_length)319 mask = make_causal_mask(config.context_length)320 321 # Pre-realize all Monarch matrices: monarch_ms[layer][head] = masked T×T matrix322 monarch_ms = Vector{Vector{Matrix{Float32}}}(undef, config.n_layers)323 for i in 1:config.n_layers324 name = Symbol("block_$i")325 bp = getproperty(ps.blocks, name)326 layer_ms = Vector{Matrix{Float32}}(undef, config.n_monarch_heads)327 for j in 1:config.n_monarch_heads328 head_name = Symbol("head_$j")329 ps_m = getproperty(bp.seq_mixer.monarchs, head_name)330 M = monarch_realize(ps_m.L1, ps_m.L2, p) .* mask331 layer_ms[j] = M332 end333 monarch_ms[i] = layer_ms334 end335 336 return (; mask, monarch_ms)337end338 339# ═══════════════════════════════════════════════════════════════════340# Monarch Sequence Mixer forward pass (uses cached matrices)341# ═══════════════════════════════════════════════════════════════════342 343function monarch_sequence_mixer_forward(x, ps, n_heads::Int, monarch_ms_layer)344 D, T, B = size(x)345 H = n_heads346 HD = D ÷ H347 348 # 1. Causal depthwise conv for local context349 conv_out = causal_depthwise_conv1d(x, ps.conv.kernel)350 351 # 2. Multi-head Monarch mixing (pre-realized matrices)352 monarch_slices = map(1:H) do i353 # Slice cached matrix to actual sequence length354 M_t = monarch_ms_layer[i][1:T, 1:T]355 356 # Extract this head's channel slice: (HD, T, B)357 ch_start = (i - 1) * HD + 1358 ch_end = i * HD359 x_slice = x[ch_start:ch_end, :, :]360 361 # Matmul: (T, T) × (T, HD*B) → (T, HD*B)362 x_flat = reshape(permutedims(x_slice, (2, 1, 3)), T, HD * B)363 y_flat = M_t * x_flat364 365 # Reshape back: (T, HD*B) → (T, HD, B) → (HD, T, B)366 permutedims(reshape(y_flat, T, HD, B), (2, 1, 3))367 end368 369 # Concatenate heads along channel dimension370 monarch_out = cat(monarch_slices...; dims=1)371 372 # 3. Combine conv (local) + Monarch (global), then gate373 combined = conv_out .+ monarch_out374 gate = NNlib.sigmoid_fast.(ps.gate.weight)375 gated = gate .* combined376 377 return gated378end379 380# ═══════════════════════════════════════════════════════════════════381# Full model forward pass (uses cached data)382# ═══════════════════════════════════════════════════════════════════383 384function model_forward(config::ModelConfig, ps, x, caches)385 T = size(x, 1) # x: (seq_len, batch) of integer token IDs386 387 # Token embedding: (seq_len, batch) → (embed_dim, seq_len, batch)388 h = ps.tok_emb.weight[:, x]389 390 # Monarch blocks391 for i in 1:config.n_layers392 name = Symbol("block_$i")393 bp = getproperty(ps.blocks, name)394 395 # Pre-norm sequence mixing + residual396 normed = rmsnorm_forward(h, bp.ln1.weight)397 mixed = monarch_sequence_mixer_forward(normed, bp.seq_mixer,398 config.n_monarch_heads,399 caches.monarch_ms[i])400 h = h .+ mixed401 402 # Pre-norm FFN + residual403 normed2 = rmsnorm_forward(h, bp.ln2.weight)404 ffn_out = swiglu_forward(normed2, bp.ffn)405 h = h .+ ffn_out406 end407 408 # Final norm409 h = rmsnorm_forward(h, ps.ln_f.weight)410 411 # Output projection412 D, T_out, B = size(h)413 h_flat = reshape(h, D, T_out * B)414 if config.weight_tying415 logits = ps.tok_emb.weight' * h_flat416 else417 logits = ps.head.weight * h_flat418 end419 return reshape(logits, :, T_out, B)420end421 422# ═══════════════════════════════════════════════════════════════════423# Sampling helpers424# ═══════════════════════════════════════════════════════════════════425 426function top_k_filter(logits::AbstractVector, k::Int)427 k = min(k, length(logits))428 sorted = sort(Array(logits); rev=true)429 threshold = sorted[k]430 return map(l -> l >= threshold ? l : typemin(eltype(logits)), logits)431end432 433function top_p_filter(logits::AbstractVector, p::Float64)434 sorted_indices = sortperm(Array(logits); rev=true)435 sorted_logits = logits[sorted_indices]436 probs = NNlib.softmax(sorted_logits)437 cumprobs = cumsum(Array(probs))438 cutoff = something(findfirst(>=(p), cumprobs), length(probs))439 result = fill(typemin(eltype(logits)), length(logits))440 for i in 1:cutoff441 result[sorted_indices[i]] = logits[sorted_indices[i]]442 end443 return result444end445 446function sample_categorical(probs::AbstractVector)447 r = rand(Float32)448 cumulative = 0.0f0449 for i in eachindex(probs)450 cumulative += probs[i]451 r <= cumulative && return i452 end453 return length(probs)454end455 456# ═══════════════════════════════════════════════════════════════════457# Text generation with streaming callback458# ═══════════════════════════════════════════════════════════════════459 460function generate_streaming(config::ModelConfig, ps,461 tokenizer::Tokenizer, prompt::String;462 max_tokens::Int=200,463 temperature::Float64=0.8,464 top_k::Int=0,465 top_p::Float64=1.0,466 on_token=nothing,467 caches=nothing)468 tokens = encode(tokenizer, prompt)469 if isempty(tokens)470 tokens = [rand(1:tokenizer_vocab_size(tokenizer))]471 end472 473 # Use provided caches or compute them once474 if caches === nothing475 caches = precompute_inference_caches(config, ps)476 end477 478 generated = String[]479 480 for _ in 1:max_tokens481 ctx = if length(tokens) > config.context_length482 tokens[end-config.context_length+1:end]483 else484 copy(tokens)485 end486 487 x = reshape(ctx, :, 1)488 logits = model_forward(config, ps, x, caches)489 next_logits = Vector{Float32}(logits[:, end, 1])490 491 if temperature != 1.0492 next_logits ./= Float32(temperature)493 end494 495 if top_k > 0496 next_logits = top_k_filter(next_logits, top_k)497 end498 499 if top_p < 1.0500 next_logits = top_p_filter(next_logits, top_p)501 end502 503 probs = NNlib.softmax(next_logits)504 next_token = sample_categorical(probs)505 506 push!(tokens, next_token)507 token_str = decode(tokenizer, [next_token])508 push!(generated, token_str)509 510 on_token !== nothing && on_token(token_str)511 end512 513 return join(generated)514end515