1000× Faster Backward Pass: Three Sprints That Reshaped the Training Pipeline

Training a state-space model at reasonable scale exposes three hard walls in sequence: CPU overhead in the backward pass, GPU memory fragmentation, and optimiser plateaus that no hyperparameter tune can break. Over the last sprint cycle we hit all three — and fixed them.

Here is how each sprint contributed.

Sprint 1: Making the backward pass nearly free

The backward pass in an autograd engine builds a topological order of operations before executing gradients. On small graphs the cost is invisible. On a multi-layer Mamba model with complex state updates, BuildTopoSafe consumed ~12 ms per step — CPU time that did nothing but bookkeeping.

The fix was a caching layer with structural hash invalidation. BuildTopoSafeWithCache computes a hash of the graph structure once, then reuses the cached order on every subsequent step unless the hash changes. Dynamic graphs are handled correctly: if the structure diverges, the cache rebuilds automatically.

The same sprint introduced a pre-allocated memory pool. Instead of calling cudaMalloc ~2400 times per step, TTensorMemoryPool maintains free-lists by allocation size bucketed to powers of two. After the first step, runtime allocations drop to zero. VRAM pressure spikes disappeared entirely.

cudaMalloc calls per step (before)

2400

cudaMalloc calls per step (after step 1)

0

Combined, these two changes meant the backward pass went from a measurable overhead to a negligible one — leaving room for more expensive optimisers and larger models.

Sprint 2: Breaking the loss plateau with hybrid optimisation

AdamW is reliable, but on the arithmetic reasoning task (Stage 3) it hit a hard floor: loss oscillated around 0.23, and exact-match accuracy stalled at ~82%. No learning rate schedule or weight decay adjustment helped.

The insight was that different parameter types need different preconditioning. Two-dimensional weight matrices benefit from orthogonalisation via the Newton-Schulz iteration — the core idea behind the Muon optimiser. One-dimensional vectors (biases, layer norms) and embedding tables do not, and work better with standard AdamW.

TParamGroupEngine classifies every parameter automatically:

• 2D weight matrices → Muon (Newton-Schulz orthogonalisation)
• 1D vectors (biases, norms) → AdamW
• Embedding tables → AdamW

The THybridOptimizer dispatches a single .Step() across both optimisers transparently. The result broke the plateau:

Adam-only EM

82

Hybrid Muon+AdamW EM

97.2

The ~30 ms epoch overhead from Newton-Schulz iterations was more than compensated by 2.8× faster convergence.

Sprint 3: Full GPU acceleration for Mamba-3

Mamba-3 extends the original Mamba architecture with complex-valued SSM states and multiple input/output streams (MIMO). The Pascal implementation worked, but without GPU kernels it could not scale to Stage 4 (carry arithmetic) and beyond.

Three CUDA kernels closed the gap:

Each kernel has a HAL binding with a CPU fallback, so the same code path works on both platforms. Four new autograd rules (opComplexMul, opComplexAdd, opMixStreams, opBroadcastStreams) integrate complex operations into the cached topological sort from Sprint 1.

The language model now uses TMamba3Layer by default. TMambaLayer is preserved for backward compatibility. Smoke tests confirm correct output shapes, GPU/CPU parity within FP32 tolerance, and gradient flow through all complex operations.

One sprint feeds the next

These three sprints were planned independently, but they compound. Faster topological sorting and zero-runtime allocations (Sprint 1) freed compute budget for the hybrid optimiser (Sprint 2), which converged in fewer steps — steps that now run on a fully GPU-accelerated Mamba-3 (Sprint 3). Each fix unlocked the next bottleneck.

A more detailed breakdown of each subtask is available in the public project overview at lotargo.github.io/public_sonata_ai_landing. The next sprint cycle starts with Stage 4 — carry arithmetic — and the question of whether the same stack can scale further without new surprises.

1000× быстрее backward pass: три спринта, перестроившие пайплайн обучения

При обучении state-space модели в сколько-нибудь разумном масштабе последовательно упираешься в три стены: CPU-overhead в backward pass, фрагментацию GPU-памяти и плато оптимизатора, которое не пробить никакой настройкой гиперпараметров. За последний спринт мы упёрлись во все три — и починили их.

Вот что дал каждый спринт.

Спринт 1: backward pass почти бесплатно

Backward pass в autograd-движке строит топологический порядок операций перед вычислением градиентов. На маленьких графах это незаметно. На многослойной Mamba с комплексными обновлениями состояния BuildTopoSafe потреблял ~12 мс за шаг — чистое CPU-время на бухгалтерию.

Решение — кэш с хэш-инвалидацией структуры графа. BuildTopoSafeWithCache вычисляет хэш один раз, а на следующих шагах переиспользует закэшированный порядок, если хэш не изменился. Динамические графы обрабатываются корректно: при расхождении структуры кэш перестраивается автоматически.

В этом же спринте появился pre-allocated memory pool. Вместо ~2400 вызовов cudaMalloc за шаг, TTensorMemoryPool поддерживает free-lists по размерам аллокаций (степени двойки). После первого шага runtime-аллокации падают до нуля. Скачки давления на VRAM исчезли совсем.

cudaMalloc вызовов за шаг (до)

2400

cudaMalloc вызовов за шаг (после шага 1)

0

Вместе эти два изменения превратили backward pass из измеримого оверхеда в пренебрежимый — освободив вычислительный бюджет для более дорогих оптимизаторов и больших моделей.

Спринт 2: прорыв плато потерь гибридным оптимизатором

AdamW надёжен, но на задаче арифметического рассуждения (Stage 3) он упёрся в твёрдый пол: loss колебался вокруг 0.23, а точность exact-match застряла на ~82%. Ни подбор learning rate, ни weight decay не помогали.

Идея в том, что разным типам параметров нужна разная предобработка. Двумерным матрицам весов полезна ортогонализация через Newton-Schulz — ключевая идея оптимизатора Muon. Одномерным векторам (biases, нормы слоёв) и embedding tables это не нужно — им достаточно стандартного AdamW.

TParamGroupEngine классифицирует каждый параметр автоматически:

• 2D-матрицы весов → Muon (Newton-Schulz ортогонализация)
• 1D-вектора (biases, norms) → AdamW
• Embedding tables → AdamW

THybridOptimizer диспатчит единый .Step() на оба оптимизатора прозрачно. Результат — плато пробито:

Adam-only EM

82

Гибрид Muon+AdamW EM

97.2

Дополнительные ~30 мс на эпоху от Newton-Schulz итераций более чем компенсированы 2.8× ускорением сходимости.

Спринт 3: полное GPU-ускорение Mamba-3

Mamba-3 расширяет исходную архитектуру Mamba комплексными SSM-состояниями и множественными входными/выходными потоками (MIMO). Pascal-реализация работала, но без GPU-ядер не могла масштабироваться на Stage 4 (сложение с переносом) и дальше.

Три CUDA-ядра закрыли разрыв:

Каждое ядро имеет HAL-привязку с CPU-fallback, поэтому один и тот же код работает на обеих платформах. Четыре новых autograd-правила (opComplexMul, opComplexAdd, opMixStreams, opBroadcastStreams) интегрируют комплексные операции в кэшированную топологическую сортировку из Спринта 1.

Языковая модель теперь использует TMamba3Layer по умолчанию. TMambaLayer сохранён для обратной совместимости. Smoke-тесты подтверждают корректные формы выходов, GPU/CPU-паритет в пределах FP32 tolerance и прохождение градиентов через все комплексные операции.

Спринты работают как конвейер

Эти три спринта планировались независимо, но они работают как конвейер: быстрая топологическая сортировка и нулевые runtime-аллокации (Спринт 1) освободили вычислительный бюджет для гибридного оптимизатора (Спринт 2), который сошёлся за меньшее количество шагов — шагов, которые теперь выполняются на полностью GPU-ускоренной Mamba-3 (Спринт 3). Каждый фикс открывал следующий bottleneck.

Более подробная разбивка каждой подзадачи — в публичном обзоре проекта на lotargo.github.io/public_sonata_ai_landing. Следующий цикл спринтов начинается со Stage 4 — арифметики с переносом — и вопроса, сможет ли тот же стек масштабироваться дальше без новых сюрпризов.