GigaAM on MLX: the transcripts matched, but the model was broken

GigaAM v3 is Sber’s open-source Russian speech recognition model: a 16-layer Conformer with four head variants — CTC, RNNT, and two end-to-end versions that output punctuated text directly. Officially it runs on PyTorch and ONNX. I needed local dictation on my Mac, so in the spring I ported the model to MLX, Apple’s framework for Apple Silicon. The port “matched PyTorch”: I ran phrases through it, and the transcripts came out the same.

In the autumn I put together a proper eval — 200 recordings, comparing tensors rather than transcripts — and it turned out that the port had been carrying a bug the whole time, one that shifted the encoder output by almost the magnitude of the output itself. Yet the transcripts matched on 195–199 of the 200 recordings, and WER didn’t change at all. Below is that bug, a few other places where GigaAM differs from a textbook Conformer, and honest speed numbers. If you’re porting GigaAM anywhere — CoreML, C++, your own runtime — these same spots are waiting for you.

What we’re porting

• 16 kHz mono → log-mel, 64 bins (window 320, hop 160, no centering);
• two Conv1d layers with stride 2 — 4× downsampling in time;
• 16 Conformer layers, d_model 768, 16 heads, rotary positions;
• a CTC or RNNT head (a single-layer LSTM of size 320 + joint).

The character-level models work with 33 letters, the end-to-end ones with SentencePiece: 256 pieces for E2E CTC, 1024 for E2E RNNT. The fp16 weights are about 421–425 MB.

How to verify the port

The corpus is 200 short phrases from the test split of Golos Crowd (823.85 s). It’s the same sample that garrrikkotua/gigaam-vs-whisper-russian used to compare GigaAM against Whisper: his repository includes the SHA-256 of every clip, so I downloaded the audio from the HF mirror of Golos and checked all 200 hashes. A bonus is a ready-made reference point: he has archived hypotheses from PyTorch GigaAM RNNT.

For each clip, the eval runs upstream PyTorch (CPU, fp32 — the reference) and MLX, and records:

• whether the transcript matches;
• max |PT − MLX| on the encoder output;
• for CTC, whether the argmax matches on every frame;
• WER and CER against the references.

The second item is the one that matters. Why transcripts are a poor criterion becomes clear in the next section.

Trap 1: the RoPE base is 5000, not 10000

In the very first run (5 clips per model), the transcripts matched in 19 cases out of 20, but the encoder output diverged by 0.4–1.8 while the output itself had a scale of 1.7–3.5. A layer-by-layer comparison showed that the mel matches to within 3e-5, the subsampling to within 1.4e-4, and after the very first Conformer layer the difference is 0.39. The culprit turned out to be the cos/sin tables for RoPE: they diverged by almost 2.

Here’s how upstream creates the positional encoding:

self.pos_enc = RotaryPositionalEmbedding(d_model // n_heads, pos_emb_max_len)

And here’s the base class signature:

class PositionalEncoding(nn.Module):
    def __init__(self, dim: int, base: int): ...

The second positional argument is base. What goes into it is pos_emb_max_len, i.e. 5000. In “textbook” RoPE, and in my port, the base was 10000. The model was trained with 5000, so 5000 is correct. After the fix, every one of the 16 layers matches PyTorch to within 2e-5.

The most interesting part is what this bug actually cost. I ran all 200 clips with the old base and compared against PyTorch:

For the character-level models the bug is invisible in both transcripts and WER. The discrepancies are rare and look like ordinary recognition errors: «забей мне место» → «займей мне место» (“save me a seat” → a non-word), «муз тиви» → «муз тив» (“Muz TV” → a truncated “Muz T”). For E2E RNNT, the buggy version’s WER is even one error lower. So “the transcripts match” and “WER is the same” are not proof that the port matches. The proof is the tensor difference.

Trap 2: RoPE before the projections

In the usual rotary attention, the input is first projected into q and k, and it’s those that get rotated. GigaAM does it the other way around — here’s RotaryPositionMultiHeadAttention.forward, abridged:

query = query.transpose(0, 1).view(t, b, self.h, self.d_k)
key   = key.transpose(0, 1).view(t, b, self.h, self.d_k)
value = value.transpose(0, 1).view(t, b, self.h, self.d_k)

query, key = apply_rotary_pos_emb(query, key, cos, sin, offset=0)

q, k, v = self.forward_qkv(query..., key..., value...)  # linear_q/k/v happen here

What gets rotated is the raw input, split into heads, and only afterwards does it go through linear_q and linear_k. Value isn’t rotated. In MLX:

x_heads = x.reshape(B, T, H, d_k)
q_rot = x_heads * cos + rotate_half(x_heads) * sin
k_rot = x_heads * cos + rotate_half(x_heads) * sin   # same x — self-attention
q = linear_q(q_rot.reshape(B, T, D))
k = linear_k(k_rot.reshape(B, T, D))
v = linear_v(x)                                       # no rotation

I found this one back in the spring by printing the source of forward from a debug script. The takeaway for other ports: off-the-shelf fused RoPE kernels assume q and k after projection, so for GigaAM you can’t just drop them in as-is.

Trap 3: mel filters — carry them over, don’t recompute them

My own HTK mel filter bank “like torchaudio’s” didn’t match the original. It turned out simpler not to recompute it: the filter bank and the STFT window already live in the PyTorch model as buffers (mel_scale.fb, spectrogram.window). The converter puts them into model.safetensors, and the runtime loads them from there. The log-mel now matches PyTorch to within 3e-5.

Advice that goes beyond GigaAM: if a model’s front end is a fixed matrix, port the matrix.

Trap 4: weight layout and precision

• Conv1d: PyTorch stores [out, in, K], MLX stores [out, K, in]. In GigaAM, all three-dimensional weights are convolutions. After transposing, the array has to be made contiguous (np.ascontiguousarray), otherwise safetensors won’t save it.
• LSTM: PyTorch has two biases, mlx.nn.LSTM has one — add them together; weight_ih → Wx, weight_hh → Wh.
• fp16 or fp32. GigaAM checkpoints are already stored in fp16: 548 of the 552 tensors in the CTC model. Converting the encoder to fp32 is pointless — the file is twice as large, the numbers are the same.

There’s a catch with fp16 in MLX itself. The model input is log-mel in fp32, so with fp16 weights the computation still runs in fp32, and every operation pays for a type cast. If you upcast the weights to fp32 once at load time, CTC speeds up from 167× to 201× real time, and RNNT from 73× to 91× (+0.45 GB of memory). load_model now does this by default.

End-to-end: SentencePiece without SentencePiece

The E2E models output normalized, punctuated text, but their output is indices of SentencePiece pieces. During conversion the whole piece vocabulary is dumped into config.json, and the decoder fits in a dozen lines:

• control tokens are skipped;
• <unk> → « ⁇ »;
• <0xNN> is byte fallback, written out as a byte;
• ▁ → space, and the bytes are assembled into UTF-8.

A subtlety I tripped over: SentencePiece only removes the leading ▁ of the first piece; it doesn’t trim whitespace at both ends. The first version called strip() and diverged from spm.decode on edge cases. The check was fuzzing: 100,000 random sequences per model, and after the fix, zero discrepancies.

Along the way I found a converter bug: for E2E CTC it wrote 33 letters into the vocabulary instead of 256 pieces, so that model didn’t work in MLX at all. The vocabulary is now taken from the tokenizer of the loaded PyTorch model.

Memory: the MLX cache grows with input variety

The first benchmark run ate 17.2 GB on a 24 GB Mac and pushed the system into swap. There was no leak: the model’s working memory stayed at 0.41 GB. What grew was MLX’s buffer cache. Every new audio length means new tensor shapes, and MLX keeps freed buffers of every shape — by default, almost up to the GPU memory limit.

For speech recognition, where every clip has its own length, this matters, and for streaming with a growing buffer even more so: there, every step is a new shape. The fix is one line:

mx.set_cache_limit(1 << 30)  # 1 GB

On 100 short clips it costs about 4 % in speed (2.29 → 2.39 s).

Speed

Apple M4 Pro, 24 GB. Audio in memory, a warm-up, then three passes, taking the median. “Short” is 200 clips one at a time (median 3.9 s), “long” is 38 segments of ≈20 s. The figure is how many times faster than real time.

The honest conclusion: for the CTC models, MLX and PyTorch on MPS are neck and neck; MLX wins on RNNT, by 2.3–2.5×. Judging by the numbers (not by profiling), in RNNT most of the cost is the greedy decoding loop, step by step, and in MLX it is cheaper. Against PyTorch on CPU the gap is 7 to 20×. The model loads in MLX in 0.1–0.2 s versus ~1 s for PyTorch.

Quality

MLX transcripts match PyTorch on 200 of 200 clips for all four models, so the WER is identical:

RNNT made 15 errors on 1001 words — exactly as in garrrikkotua’s archive, and all 200 hypotheses match his archived ones byte for byte.

The E2E models write “HD”, “4K”, “15-й” (Russian for “15th”) and phone numbers in digits, while Golos spells all of these out in Cyrillic words. So their WER on the full corpus mostly measures formatting, and I give a second column — the clips where no such formatting occurs. Golos Crowd consists of short voice commands, so these numbers can’t be compared with the GigaAM authors’ figures on their own datasets.

Streaming: honestly, it’s pseudo-streaming

GigaAM is not a streaming model. Streaming in the port works through a growing buffer: every step re-recognizes all the audio, but no more than the last 30 seconds. The cost of a step grows with the buffer. The E2E models have one more quirk: punctuation and capitalization at the start of a phrase change from step to step, so the “new chunk of text” almost always turns out to be the whole text over again.

Where the code is

• The converter, the MLX runtime and the eval: mlx_convert/ in the misteral/GigaAM fork; the upstream PR is salute-developers/GigaAM#62.
• The eval can be reproduced with a single sequence of commands: mlx_convert/eval/README.md.