LM Engine
Table of contents
  1. 1 Introduction
  2. 2 Architecture
  3. MoE
  4. Mamba-2
  5. Exclusive Self Attention (XSA)
  6. 3 Pretraining Data
  7. 4 Training Cluster/Infrastructure
  8. Cluster
  9. Kernels
  10. Context parallelism for attention
  11. Context parallelism for Mamba-2
  12. 5 Model Training
  13. Long-context extension
  14. 6 Evaluation
  15. 7 Conclusion
  16. Citation
  17. References

Rıgel Base: Reaching Llama-3.2 Quality with <1% of its Compute

Mayank Mishra1, Bharat Runwal2, Ion Stoica1, Tri Dao3,4*, Joseph E. Gonzalez1*

1UC Berkeley, 2MIT-IBM Watson Lab, 3Princeton University, 4Together AI

* Equal advising

Abstract

We release Rigel, a 2.3B-total, 360M-active-parameter hybrid mixture-of-experts base model. Three of every four sequence-mixer layers are Mamba-2 state-space layers; the fourth is grouped-query attention with exclusive self attention and a 4,096-token sliding window, so the inference state stays bounded regardless of context length. The feed-forward blocks are a fine-grained 128-expert, top-2 MoE. Rigel is pretrained at a 4,096-token context over a six-phase data curriculum and then extended to a 294,912-token context. With no dedicated cluster, the run moved between H100, A100, V100 and TPU v5p/v6e chips on a single codebase, with kernels written or adapted for each. Rigel used 5.4×10215.4 \times 10^{21} non-embedding pretraining FLOPs, about 117×117\times fewer than Llama-3.2-1B/3B, 53×53\times fewer than Granite-4.2-3B and 35×35\times fewer than SmolLM3-3B, while landing within a few points of Granite-4.2-3B and Llama-3.2-3B on zero-shot averages and above Llama-3.2-1B. We describe the architecture, the data mixture, the multi-chip infrastructure, the optimization recipe and the evaluations.

Introduction

We release Rigel: a 2.3B-total, 360M-active-parameter (260M active non-embedding parameter) hybrid Mamba-2 (Dao and Gu, 2024) + attention (Vaswani et al., 2017) mixture-of-experts (Fedus et al., 2022) base model. Three of every four sequence-mixer layers are Mamba-2 state-space layers, carrying a fixed-size recurrent state regardless of context length; the fourth is grouped-query attention (Ainslie et al., 2023) with a group size of 4 (query to key-value head ratio) and exclusive self attention (XSA) (Zhai, 2026). The model is pretrained at a 4,096-token context length, then extended to a 294,912 token context with 4,096-token sliding window attention towards the end of training.

Rigel was trained on five different kinds of accelerators: H100s, A100s (both 40GB and 80GB), V100s (32GB), and TPU v5p and v6e. We had no dedicated cluster and hence, we picked chips up as they became available and moved the run between them.

Rigel also used 5.4×10215.4 \times 10^{21} non-embedding pretraining FLOPs, about 117×117\times fewer than the Llama-3.2-1B and Llama-3.2-3B models (we assume Llama-3.1-8B’s training compute and assume distillation and pruning are free and cost no FLOPs), 53×53\times fewer than Granite-4.2-3B, and 35×35\times fewer than SmolLM3-3B, while landing within a few points of Granite-4.2-3B and Llama-3.2-3B on zero-shot averages and above Llama-3.2-1B.

Rigel therefore targets a different point on the scaling curve than trillion-parameter frontier modelsNotepartly because we couldn’t get compute to train a trillion parameter model: a base model small enough to pretrain, ablate, and iterate on with a modest and changing pool of hardware, at a total-parameter budget measured in billions and an active-parameter budget within reach of a single accelerator. This blog gives a brief overview of Rigel: the architecture, the six-phase pretraining mixture, the multi-chip clusters and the kernels and context-parallel machinery it needed, the optimization recipe, and how the result compares to models trained with far more compute.

The following principle guided this project from the start:

Work with the compute you have. With no dedicated cluster, we treated whatever hardware became available as an opportunity rather than a compromise.

Training across different hardware, whether from different vendors or from different generations of a single vendor, brought three challenges:

  1. Uneven kernel support. Kernels that are well optimized for one chip may exist but not for another. This limits how fast we can run our experiments.
  2. Architecture followed kernel availability. We wrote our own MoE kernels for V100s, but implementing Gated DeltaNet (GDN) (Yang et al., 2025) on V100s and TPUs looked like a much bigger undertaking, so we adopted Mamba-2 (Dao and Gu, 2024), which was far simpler to implement. We hope to implement GDN for our future training runs.
  3. Numerics differ across hardware. V100s support FP16 but not BF16. This did not hurt our current model, but it could affect a larger training run.

Architecture

For the feed-forward blocks, we use MoE layers since they disentangle training FLOPs from the total parameter count of the model. For the sequence mixer blocks, we adopted a hybrid design: GQA (Ainslie et al., 2023) with sliding window length of 4,096 tokens with Mamba-2 (Dao and Gu, 2024) layers. While we wanted to use Gated DeltaNet (Yang et al., 2025) or Kimi Delta Attention (Kimi Team et al., 2025), writing forward and backward kernels for V100s and TPUs for GDN layers seemed like a much bigger undertaking. We hope to adopt GDN (Yang et al., 2025) (or KDA (Kimi Team et al., 2025)) in our next models.

MoE

Two axes have driven recent language-model efficiency work. The first decouples a model’s capacity from the FLOPs required to train the model: sparse mixture-of-experts routing (Fedus et al., 2022) lets total parameters (and consequently model capacity) grow far faster than the active parameters. This allows training much larger models with a FLOPs budget that is governed by only the active parameters of the model rather than the total model parameters. Recent frontier LLMs like DeepSeek (DeepSeek-AI et al., 2026; Dai et al., 2024), Kimi (Kimi Team et al., 2026), Qwen (Yang et al., 2025) have further pushed towards increasingly fine-grained sparser MoE layers and towards hybrid recurrent/attention based sequence mixing blocks.

Both the Mamba-2 (Dao and Gu, 2024) and the attention (Vaswani et al., 2017) layers share the same feed-forward block: a fine-grained MoE with 128 SwiGLU (Shazeer, 2020) experts, an intermediate size of just 128 per expert, and top-2 routing. We use an auxiliary load-balancing loss for training the model. We don’t use any shared experts purely to save FLOPs. Expert weights are stored in an interleaved layout for the fused SonicMoE (Guo et al., 2026)NoteMore about infra later 🙂.

We use tied word-embeddings for our model to keep the parameter count low. Every non-MoE component is dense; while we experimented with MoA (Mixture of Attention) (Zhang et al., 2022), we found that our kernel was not optimized enough due to a lack of enough fusion at the time of starting model training (and simply running full dense GQA (Ainslie et al., 2023) was faster on both H100s and A100s) and hence we chose to stick to using dense attention for our model. An efficient MoA kernel could have enabled longer training than our current run 🫠. We aim to make MoA kernels faster for our future runs and make the model bigger for the same FLOPs by making the attention (Vaswani et al., 2017) layers (and possibly Mamba-2 (Dao and Gu, 2024) layers) sparse.

Mamba-2

Left: the residual stream alternates three Mamba-2 layers for every one attention layer, each interleaved with a MoE. Bottom right: inside the Mamba-2 layer, qt,kt,vtq_t, k_t, v_t each get their own depthwise causal convolution and activation before feeding the Mamba-2 recurrence, whose output is gated by gtg_t, normalized and projected out. Top right: inside the attention layer, qt,kt,vtq_t, k_t, v_t feed causal GQA attention to produce ztz_t, which is then corrected by projecting out its own vtv_t component before the output projection — the XSA (Zhai, 2026) step detailed below.

Mamba-2 (Dao and Gu, 2024) has been recently proposed as a reformulation of the selective state-space model introduced by Mamba (Gu and Dao, 2024), expressed as an instance of structured state-space duality (SSD). It casts the SSM recurrence as a form of linear attention with a data-dependent, scalar state-transition gate. Unlike softmax attention (Vaswani et al., 2017), whose KV-cache and per-token compute both grow with context length, a Mamba-2 (Dao and Gu, 2024) layer carries a fixed-size recurrent state regardless of how many tokens came before. While this makes Mamba-2 (or any other linear attention layers) efficient at inference time, it comes at the cost of the precise per-token retrieval where attention still does best. Rigel uses Mamba-2 (Dao and Gu, 2024) for 30 of its 40 sequence-mixing layers, interleaved with grouped-query attention (Ainslie et al., 2023) every fourth layer to recover that retrieval capability where it matters most.

Mamba-2 (Dao and Gu, 2024) admits the following recurrent form:

Ht=ft Ht−1+ktvt⊤yt=Ht⊤qt\begin{aligned} \mathbf{H}_t &= f_t\,\mathbf{H}_{t-1} + k_t v_t^\top \\ y_t &= \mathbf{H}_t^\top q_t \end{aligned}

where ft∈R,kt∈RK×1,vt∈RV×1,Ht∈RK×Vf_t \in \mathbb{R}, k_t \in \mathbb{R}^{K \times 1}, v_t \in \mathbb{R}^{V \times 1}, \textbf{H}_t \in \mathbb{R}^{K \times V}NoteNote that ktk_t here has been assumed to be discretized; we omit the discretization term here for simpler notation..

While the recurrent form is extremely useful for O(1)O(1) inference, it cannot use matrix multiply units: tensor cores on NVIDIA GPUs or MXUs on TPUs (Norrie et al., 2020; Jouppi et al., 2023). Each step is a sequential rank-one update ktvt⊤k_t v_t^\top on Ht\mathbf{H}_t, not the large matrix multiply that AI accelerators need to reach good utilization. So training with the recurrent form directly is inefficient.

Mamba-2 (Dao and Gu, 2024) like Linear Attention (Katharopoulos et al., 2020) admits a chunkwise-parallel form (Yang et al., 2024) that can make efficient use of these matrix multiply units. The chunkwise-parallel form can be derived by running the recurrence over blocks of CC contiguous tokens: the inter-chunk term stays a genuine sequential recurrence over the chunk-boundary state H[i]\mathbf{H}_{[i]}, but the intra-chunk term becomes a single causally masked matrix multiply per chunk so the sequential shrinks from TT individual tokens to TC\frac{T}{C} chunks, and the bulk of the compute runs as large matrix multiplies that the tensor cores and MXUs are built for.

Mamba-2’s (Dao and Gu, 2024) chunkwise-parallel form can be written as:

H(i+1)C=F(i+1)CFiCHiC+K[i]⊤diag(F(i+1)CF[i])V[i]Y[i]=diag(F[i]FiC)Q[i]HiC+[diag(F[i])Q[i]K[i]⊤diag(F[i]−1)⊙M]V[i]\begin{aligned} \mathbf{H}_{(i + 1)C} &= \frac{F_{(i + 1)C}}{F_{iC}} \mathbf{H}_{iC} + \mathbf{K}_{[i]}^\top \mathrm{diag}\left(\frac{F_{(i+1)C}}{F_{[i]}}\right) \mathbf{V}_{[i]} \\ \mathbf{Y}_{[i]} &= \mathrm{diag}\left(\frac{F_{[i]}}{F_{iC}}\right)\mathbf{Q}_{[i]} \mathbf{H}_{iC} + \left[ \mathrm{diag}(F_{[i]}) \mathbf{Q}_{[i]} \mathbf{K}_{[i]}^\top \mathrm{diag}(F_{[i]}^{-1}) \odot \mathbf{M} \right] \mathbf{V}_{[i]} \end{aligned}

where CC is the chunk size and Ft=∏s≤tfsF_t = \prod_{s \le t} f_s denotes the scalar prefix decay. Note that we use the notation [i][i] to denote the ithi^{th} chunk.

Exclusive Self Attention (XSA)

Rigel‘s attention layers also use exclusive self attention (XSA) (Zhai, 2026), a two-line modification to standard causal self-attention. XSA (Zhai, 2026) observes an attention similarity bias in trained Transformers: a token’s attention output zi=∑j≤iai,jvjz_i = \sum_{j \le i} a_{i,j} v_j tends to have high cosine similarity with its own value vector viv_i, i.e. the attention block spends capacity re-deriving a point-wise feature the residual path’s MLP could already supply. XSA removes exactly that component by projecting it out of the attention output:

yi=zi−(zi⊤vi)vi∥vi∥22y_i = z_i - (z_i^\top v_i) \frac{v_i}{\lVert v_i \rVert_2^2}

an orthogonal rejection of ziz_i against viv_i, implementable as a two-line addition after the usual softmax-attention output. XSA (Zhai, 2026) reports gains that grow with both model size and context length, and frames XSA as an implicit attention sink (Xiao et al., 2024). Rigel enables it on every GQA (Ainslie et al., 2023) layer.

We ablated Grouped Query Attention (GQA) (Ainslie et al., 2023), gated attention (Qiu et al., 2025), Exclusive Self Attention (XSA) (Zhai, 2026) and combining Qwen’s gated attention (Qiu et al., 2025) with XSA. We find that XSA overall performs the best with XSA and gated attention being close in language modeling loss. Note that these ablations are done in the same hybrid Mamba-2 (Dao and Gu, 2024) MoE setting with the only changes between the runs being the attention layers. We keep the total parameters across all runs fixed. While both XSA and Gated Attention + XSA track closely, we choose XSA.

Following Kimi (Kimi Team et al., 2026) and our prior experience on training Granite models (Granite Team, IBM, 2024), we use no positional embeddings (NoPE) (Kazemnejad et al., 2023) which have been shown to be effective in hybrid models since the recurrent layers (Mamba-2 (Dao and Gu, 2024) in our case) already carry positional information. We choose to use 40 layers which is unusually high for a 2.3B model. While a wider model would have been more efficient to train since AI accelerators (especially TPUs (Norrie et al., 2020; Jouppi et al., 2023)) prefer larger matrix multiplies, we choose a deeper model purely for better downstream performance. We summarize our architecture in the following table:

PropertyValue
Sequence mixers30×30\times Mamba-2 (Dao and Gu, 2024), 10×10\times GQA-XSA (Ainslie et al., 2023; Zhai, 2026) attention (3:1, repeating every 4th layer)
Hidden size1,024
MLP (every layer)MoE, 128 experts, top-2, no shared expert
Vocab size100,352
Position embeddingNoPE
Initialization / width scalingμ\muP (Yang et al., 2022)
Layers40

Pretraining Data

Rigel is pretrained in 6 different phases using the following data mixtures:

  1. an initial web data and code heavy launch mix
  2. a swap towards STEM-heavy reasoning data
  3. a return to web text with a heavier math weight
  4. Nemotron-CC-v2 (NVIDIA et al., 2025) is introduced due to its higher quality
  5. increasing the ratio of datasets of higher quality with focus on finepdfs (Kydlíček et al., 2025), Nemotron-CC-v2 (NVIDIA et al., 2025) etc.
  6. a long-context phase where long chain-of-thought reasoning QA displaces nearly everything else.

Phase 1

the original base-training mix

Phase 2

web swapped for STEM-heavy data

Phase 3

web source swapped again, math share grows

Phase 4

Nemotron-CC-v2 introduced; math/code balloon to 35% each

Phase 5

Nemotron-CC becomes the dominant data source

Phase 6

long-context stage (294,912 tokens): long-CoT QA takes over

General web
Code
Math
Multilingual
STEM / specialized reasoning
Nemotron-CC-v2
FinePDF
Long-CoT reasoning QA
Other
Rigel’s different pretraining data-mixture phases.

Training Cluster/Infrastructure Kekw laughing-face meme.

Cluster

We trained Rigel on H100s, A100s (both 40GB and 80GB), V100s (32GB), TPU v5p(s), and TPU v6e(s), with most of the time split between TPU v6e and A100s. We used the lm-engine codebase for training Rigel due to its compatibility with both NVIDIA GPUs and TPUs.NoteThe TPU code for Mamba-2 (Dao and Gu, 2024) is being prepared for open-source release, together with a report on the lm-engine codebase.

Kernels

For training our models on H100 GPUs, we use the SonicMoE (Guo et al., 2026) for efficient MoE training and ScatterMoE (Tan et al., 2024) on A100 GPUs. We further develop kernels for use on V100s.

For TPUs (Norrie et al., 2020; Jouppi et al., 2023), we used torch-xla instead of JAX. We chose to use torch-xla just for having a unified model code to avoid any potential incorrectness in the model code across a PyTorch and a JAX codebase. Although, this doesn’t guarantee an op-level precision match across hardware, we found our training curves to be nearly perfectly identical to not care about this at the 2.3B parameter scale. This was also a good exercise for us to write Pallas kernels for TPUs for Mamba-2 (Dao and Gu, 2024) which we plan to open source.

We used mixed-precision training (Micikevicius et al., 2018) in BF16 (Kalamkar et al., 2019) for Rigel on A100, H100 and TPU with any all-reduce and reduce-scatter operations in FP32 precision with the exception of V100 GPUs. V100s don’t support the BF16 precision format and hence we used FP16 with loss scaling to avoid small gradients becoming zeros.

Our models were trained purely with DDP (Li et al., 2020) on H100s and A100s with 80GB memory and HSDP-2 on V100s 32GBs and A100 40GBs. We used purely FSDP (Fully Sharded Data Parallel) (Rajbhandari et al., 2020; Zhao et al., 2023) on TPUs since we were not able to get HSDP (Hybrid Sharded Data Parallel) working with torch-xla on TPUs 🥲.

We used efficient kernels wherever applicable: CODA kernels (Guo et al., 2026)NoteThe CODA project has since been added to the lm-engine org! Thanks to Han Guo. for H100 GPUs, quack kernels, causal-conv1d, mamba-ssm etc. We also relied on the XLA compiler heavily to generate efficient TPU kernels and torch.compile for GPU kernels.

All our long-context training was conducted on H100 GPUs with context parallelism (CP) (Liu et al., 2023; Jacobs et al., 2023) communication within a server node.

Context parallelism for attention

For full causal attention, ring attention (Liu et al., 2023) partitions the sequence across NCPN_\textrm{CP} ranks and exchanges KV blocks around the ring during the forward computation. Causal attention is computationally asymmetric: positions near the end of a sequence attend to more tokens and incur higher compute than positions near the beginning, so naively dividing the sequence into NCPN_\textrm{CP} contiguous chunks yields unbalanced GPU utilization: early ranks sit idle while late ranks are compute-bound. The fix is a load balancer that splits the sequence into 2NCP2N_\textrm{CP} chunks, where rank ii receives the ithi^\textrm{th} chunk and the (NCP−1−i)th(N_\textrm{CP}-1-i)^\textrm{th} chunk: one compute-light (early) segment and one compute-heavy (late) segment per rank, which balances FLOPs perfectly across all CP ranks.

Ring attention load balancing across NCPN_\textrm{CP} ranks. For full causal attention, naive contiguous chunking (left) assigns heavier attention compute to later ranks, while zigzag ring attention (middle) pairs each rank’s head chunk with a tail chunk to balance FLOPs perfectly. For causal sliding window attention (right), naive contiguous chunking is already near-balanced: every rank past the first does the same window-capped amount of work, so no zigzag reshuffling is needed.

However, sliding window attention (SWA) needs none of this reshuffling. Since each query attends to at most WW (window size) tokens regardless of position, per-rank compute is already near-uniform under a naive contiguous split: only the ranks whose queries sit within the first WW tokens of the sequence see truncated windows and do less work. Formally, with a contiguous split every rank owns L=S/NCPL = S / N_\textrm{CP} consecutive tokens, and a query at position tt attends to min⁡(t+1,W)\min(t + 1, W) key-value pairs, so a (0-indexed) rank rr carries the full, saturated load of LWLW once its first token has a complete window behind it, i.e. once

r≥NCPWS,r \ge \frac{N_\textrm{CP} W}{S},

where WW is the sliding window size and SS is the full context length. Every rank from ⌈NCPW/S⌉\lceil N_\textrm{CP} W / S \rceil onward therefore does identical work; only the ranks before it are lighter, and the sole light rank in the right panel of the figure above is the r=0r = 0 case. For Rigel, W=4,096W = 4{,}096, S=294,912S = 294{,}912 and NCP=8N_\textrm{CP} = 8 within a node, so NCPW/S=1/9N_\textrm{CP} W / S = 1/9: rank 0 is the only under-loaded rank, and its deficit is roughly W2/2W^2/2 out of LWLW, i.e. W/2L≈5.6%W / 2L \approx 5.6\% below the others. Hence, we adopt the naive contiguous chunking for sequences since Rigel uses SWA.

Context parallelism for Mamba-2

Mamba-2 needs no ring exchange of per-token keys and values under CP at all. Recall the chunkwise-parallel recurrence from the Architecture section, H(i+1)C=AiHiC+B[i]\mathbf{H}_{(i+1)C} = A_i \mathbf{H}_{iC} + \mathbf{B}_{[i]} with Ai=F(i+1)C/FiCA_i = F_{(i+1)C}/F_{iC} and B[i]=K[i]⊤diag(F(i+1)C/F[i])V[i]\mathbf{B}_{[i]} = \mathbf{K}_{[i]}^\top \mathrm{diag}(F_{(i+1)C}/F_{[i]})\mathbf{V}_{[i]}.

Note that both AiA_i and B[i]\mathbf{B}_{[i]} can be computed independently on each accelerator, since there’s no inter-chunk dependence in their computation. Once every rank has done this local computation, an all-gather of just these small (Ai,B[i])(A_i, \mathbf{B}_{[i]}) pairs suffices to compute the full Mamba-2 recurrence outputs:

Algorithm 1 Mamba-2 context-parallel forward. Accelerator i=0,...,NCP−1i = 0, ..., N_\textrm{CP} - 1 owns chunk ii (tokens [iC,(i+1)C)[iC, (i+1)C)) and executes the same procedure; the coloured quantities are the only ones exchanged across accelerators.
Input: local chunk Q[i],K[i]∈RC×K\mathbf{Q}_{[i]}, \mathbf{K}_{[i]} \in \mathbb{R}^{C \times K}, V[i]∈RC×V\mathbf{V}_{[i]} \in \mathbb{R}^{C \times V}, decays F[i]F_{[i]}, base state H0\mathbf{H}_0 (shared model buffer, requires no communication)
Output: local output chunk Y[i]∈RC×V\mathbf{Y}_{[i]} \in \mathbb{R}^{C \times V}
1Compute AiA_i ←F(i+1)C/FiC\leftarrow F_{(i+1)C} / F_{iC}
2Compute B[i]\mathbf{B}_{[i]} ←K[i]⊤diag(F(i+1)C/F[i])V[i]\leftarrow \mathbf{K}_{[i]}^\top \mathrm{diag}\left(F_{(i+1)C}/F_{[i]}\right) \mathbf{V}_{[i]}
3All-gather {Aj,B[j]}j=0NCP−1\{A_j, \mathbf{B}_{[j]}\}_{j=0}^{N_\textrm{CP}-1} across accelerators
4for j=0j = 0 to i−1i-1 do
5H(j+1)C←\mathbf{H}_{(j+1)C} \leftarrow AjA_j HjC+\, \mathbf{H}_{jC} + B[j]\mathbf{B}_{[j]}
end for
6Compute Y[i]\mathbf{Y}_{[i]} via the chunkwise-parallel form above, using the corrected starting state HiC\mathbf{H}_{iC}
7Return Y[i]\mathbf{Y}_{[i]}

The correction loop in step 4 is just the recurrence H(i+1)C=AiHiC+B[i]\mathbf{H}_{(i+1)C} = A_i \mathbf{H}_{iC} + \mathbf{B}_{[i]} unrolled across ranks: rank ii recovers the true starting state for its own chunk by chaining the all-gathered Aj,B[j]A_j, \mathbf{B}_{[j]} sequentially across every preceding rank j<ij < i, starting from the shared base state H0\mathbf{H}_0. The shared base state is already resident on every rank and needs no communication at all. Because this correction is a recurrence over only NCPN_\textrm{CP} small per-rank states rather than over the sequence itself, its cost is small compared to the local SSD computations. Similarly a simple state passing can be used to implement context parallelism for causal convolution as well. And unlike attention’s triangular compute imbalance, per-token SSD compute is chunk-local and independent of position, so a naive contiguous split already balances (approximately) load across ranks for Rigel‘s 30 Mamba-2 layers.

Model Training

We pretrain Rigel with 4,096 context length with a batch size of 1,152 sequences per training step (an effective global batch size of 4.7M tokens). We keep the token count per batch constant during the entire course of training. During the long-context extension phase, we adjusted the batch size to 16 sequences with 294,912 tokens each.

We used μ\muP (Yang et al., 2022) for hyper-parameter transfer across model widths. While we wanted to use the Hyperball optimizer (Wen et al., 2026) for training our models, we didn’t have the compute budget to carefully ablate it with hybrid Mamba-2 (Dao and Gu, 2024) layers. While we believe the Marin project recipes have ablated this at scale, we didn’t want to risk our training run since we haven’t tested these on our training stack yet.

V100×24×48×56×64
A100 40GB×32×192
A100 80GB×24
H100×16×24
TPU v5p×32
TPU v6e×4×16
Rigel’s logged training loss across the full run, excluding the long-context phase. Colour denotes the cluster (accelerator type and chip count) each stretch of the run was trained on. The discontinuities are where we change the sampling proportions.

We finally went with AdamW (Kingma and Ba, 2017; Loshchilov and Hutter, 2019) optimizer (we weren’t able to ablate Muon (Jordan et al., 2024; Liu et al., 2025) with Mamba-2 (Dao and Gu, 2024) layers due to a lack of compute) with a learning rate schedule that follows power law (Shen et al., 2024) followed by linear decay (Bergsma et al., 2025).

Assuming ss as the current training steps, Sw=5,000S_w = 5{,}000 as the warmup steps and ηmax⁡=0.01\eta_{\max} = 0.01 as the maximum learning rate, SdS_d as the step at which the linear decay begins and StS_t as the final training step, our LR schedule is given as:

η(s)={ηmax⁡⋅sSw,0≤s≤Swmin⁡{ηmax⁡,1.819 s−0.51},Sw≤s≤Sd1.819 Sd−0.51St−Sd(St−s),Sd≤s≤St\eta(s) = \begin{cases} \eta_{\max} \cdot \dfrac{s}{S_w}, & 0 \le s \le S_w \\[6pt] \min\{ \eta_{\max}, 1.819\, s^{-0.51} \}, & S_w \le s \le S_d \\[6pt] \frac{1.819\, S_d^{−0.51}}{S_t - S_d} (S_t - s), & S_d \le s \le S_t \end{cases}
V100×24×48×56×64
A100 40GB×32×192
A100 80GB×24
H100×16×24
TPU v5p×32
TPU v6e×4×16
Rigel’s learning rate across the first 5 different phases of pretraining: linear warmup to the peak LR of 0.01, constant phase (to cap the learning rate), power law decay followed by a linear decay towards the end. The colours denote the cluster’s accelerator type and accelerator count.

Long-context extension

The sixth and final phase extends the context from 4,096 to 294,912 tokens with 4,096-token sliding window attention, entirely on H100s using context parallelism (Liu et al., 2023). It runs for 25,000 steps and about 118B tokens, on a fresh, much smaller learning-rate schedule: a 2,500 step linear warmup to a peak of 5×10−55 \times 10^{-5} (200×200\times below the main run’s peak) followed by a linear decay to zero.

Rigel’s training loss over the long-context extension phase, at 294,912 tokens per sequence.
Learning rate over the long-context extension phase: a linear warmup over the first 2,500 steps to 5×10−55 \times 10^{-5}, then a linear decay to zero.

Evaluation

We compare zero-shot accuracy for Rigel against open source dense models: Granite-4.2-3B, Llama-3.2-1B and Llama-3.2-3B (base and instruct variants), and SmolLM3-3B (base and instruct variants), across a standard suite of common-sense, reading-comprehension, and knowledge benchmarks via lm-evaluation-harness. Rigel lands within a few points of Granite-4.2-3B and Llama-3.2-3B, above Llama-3.2-1B, and behind SmolLM3-3B, which used 35×35\times its pretraining compute. Since Granite-4.2-3B doesn’t release the base model, we directly compare to the variant which has undergone alignment.

All our results are demonstrated in the figure belowNoteClick a model’s icon in the figure’s legend to show or hide it.:

Zero-shot accuracy (%) of Rigel against Granite-4.2-3B, Llama-3.2-1B/3B (base and instruct) and SmolLM3-3B (base and instruct) on common-sense reasoning and knowledge benchmarks.

Rigel reaches these numbers on a fraction of the pretraining compute its baselines used:

Total non-embedding FLOPs used for pretraining: Rigel vs. Granite-4.2-3B, SmolLM3-3B, and Llama-3.2-1B/3B (both Llama-3.2 models are pruned and distilled from Llama-3.1-8B, so they share its training run). Since Llama-3.2 doesn’t report the training FLOPs for distillation/pruning from their 8B model, we only report the 8B FLOPs assuming the distillation is free.

Conclusion

Rigel is a 2.3B-total, 360M-active hybrid built from a fine-grained 128-expert, top-2 MoE with 40-layers of which 30 layers are Mamba-2 (Dao and Gu, 2024) and 10 are Grouped Query Attention(GQA) (Ainslie et al., 2023) with Exclusive Self Attention (XSA) (Zhai, 2026) and sliding window attention (SWA). Rigel keeps the inference state size bounded (constant state size for both Mamba-2 and sliding window attention). We further extend the context length of Rigel from its 4,096 tokens (the base pretraining length) to a 294,912 token context.

The Rigel run moved across H100s, A100s, V100s, and TPU v5p and v6e as we got access to these chips on one codebase, with kernels written or adapted for every one of those chips, FP16 loss scaling where BF16 wasn’t available.

Rigel was also trained with far less compute compared to Granite-4.2-3B, Llama-3.2 models while achieving almost the same performance across commonsense reasoning benchmarks.

What’s next: instruction-tuned variants, further baselines at similar active-parameter budgets, and the open-source release of the TPU kernels for Mamba-2 (Dao and Gu, 2024) and Gated Linear Attention (Yang et al., 2024) and their integration into the lm-engine training stack that made the multi-chip run possible.

Citation

If you’d like to cite this work:

@misc{mishra2026rigel,
  title  = {Rigel Base: Reaching Llama-3.2 Quality with <1% of its Compute},
  author = {Mishra, Mayank and Runwal, Bharat and Stoica, Ion and Dao, Tri and Gonzalez, Joseph E.},
  year   = {2026},
  url    = {https://open-lm-engine.github.io/blog/rigel/}
}

References

  1. Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality.” arXiv preprint arXiv:2405.21060 (2024).

  2. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. “Attention Is All You Need.” arXiv preprint arXiv:1706.03762 (2017).

  3. William Fedus, Barret Zoph, and Noam Shazeer. “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.” arXiv preprint arXiv:2101.03961 (2022).

  4. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.” arXiv preprint arXiv:2305.13245 (2023).

  5. Shuangfei Zhai. “Exclusive Self Attention.” arXiv preprint arXiv:2603.09078 (2026).

  6. Songlin Yang, Jan Kautz, and Ali Hatamizadeh. “Gated Delta Networks: Improving Mamba2 with Delta Rule.” arXiv preprint arXiv:2412.06464 (2025).

  7. Kimi Team, et al. “Kimi Linear: An Expressive, Efficient Attention Architecture.” arXiv preprint arXiv:2510.26692 (2025).

  8. DeepSeek-AI, et al. “DeepSeek-V4.1-Flash: Pushing the Limits of KV Cache Compression.” arXiv preprint arXiv:2609.19969 (2026).

  9. Damai Dai, et al. “DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models.” arXiv preprint arXiv:2401.06066 (2024).

  10. Kimi Team, et al. “Kimi K3: Open Frontier Intelligence.” arXiv preprint arXiv:2607.24653 (2026).

  11. An Yang, et al. “Qwen3 Technical Report.” arXiv preprint arXiv:2505.09388 (2025).

  12. Noam Shazeer. “GLU Variants Improve Transformer.” arXiv preprint arXiv:2002.05202 (2020).

  13. Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, and Tri Dao. “SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations.” arXiv preprint arXiv:2512.14080 (2026).

  14. Xiaofeng Zhang, Yikang Shen, Zeyu Huang, Jie Zhou, Wenge Rong, and Zhang Xiong. “Mixture of Attention Heads: Selecting Attention Heads Per Token.” arXiv preprint arXiv:2210.05144 (2022).

  15. Albert Gu and Tri Dao. “Mamba: Linear-Time Sequence Modeling with Selective State Spaces.” arXiv preprint arXiv:2312.00752 (2024).

  16. Thomas Norrie, Nishant Patil, Doe Hyun Yoon, George Kurian, Sheng Li, James Laudon, Cliff Young, Norman P. Jouppi, and David Patterson. “Google's Training Chips Revealed: TPUv2 and TPUv3.” 2020 IEEE Hot Chips 32 Symposium (HCS) (2020): 1-70.

  17. Norm Jouppi, et al. “TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings.” Proceedings of the 50th Annual International Symposium on Computer Architecture (2023).

  18. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention.” arXiv preprint arXiv:2006.16236 (2020).

  19. Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. “Gated Linear Attention Transformers with Hardware-Efficient Training.” arXiv preprint arXiv:2312.06635 (2024).

  20. Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. “Efficient Streaming Language Models with Attention Sinks.” arXiv preprint arXiv:2309.17453 (2024).

  21. Zihan Qiu, et al. “Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free.” arXiv preprint arXiv:2505.06708 (2025).

  22. Granite Team, IBM. “Granite 3.0 Language Models.” (2024).

  23. Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. “The Impact of Positional Encoding on Length Generalization in Transformers.” arXiv preprint arXiv:2305.19466 (2023).

  24. Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. “Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer.” arXiv preprint arXiv:2203.03466 (2022).

  25. NVIDIA, et al. “NVIDIA Nemotron Nano 2: An Accurate and Efficient Hybrid Mamba-Transformer Reasoning Model.” arXiv preprint arXiv:2508.14444 (2025).

  26. Hynek Kydlíček, Guilherme Penedo, and Leandro von Werra. “FinePDFs.” Hugging Face repository (2025).

  27. Shawn Tan, Yikang Shen, Rameswar Panda, and Aaron Courville. “Scattered Mixture-of-Experts Implementation.” arXiv preprint arXiv:2403.08245 (2024).

  28. Paulius Micikevicius, et al. “Mixed Precision Training.” arXiv preprint arXiv:1710.03740 (2018).

  29. Dhiraj Kalamkar, et al. “A Study of BFLOAT16 for Deep Learning Training.” arXiv preprint arXiv:1905.12322 (2019).

  30. Shen Li, et al. “PyTorch Distributed: Experiences on Accelerating Data Parallel Training.” arXiv preprint arXiv:2006.15704 (2020).

  31. Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. “ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.” arXiv preprint arXiv:1910.02054 (2020).

  32. Yanli Zhao, et al. “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel.” arXiv preprint arXiv:2304.11277 (2023).

  33. Han Guo, Jack Zhang, Arjun Menon, Driss Guessous, Vijay Thakkar, Yoon Kim, and Tri Dao. “CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs.” arXiv preprint arXiv:2605.19269 (2026).

  34. Hao Liu, Matei Zaharia, and Pieter Abbeel. “Ring Attention with Blockwise Transformers for Near-Infinite Context.” arXiv preprint arXiv:2310.01889 (2023).

  35. Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, and Yuxiong He. “DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models.” arXiv preprint arXiv:2309.14509 (2023).

  36. Kaiyue Wen, Xingyu Dang, Kaifeng Lyu, Tengyu Ma, and Percy Liang. “Fantastic Pretraining Optimizers and Where to Find Them II: Hyperball Optimization.” arXiv preprint arXiv:2606.16899 (2026).

  37. Diederik P. Kingma and Jimmy Ba. “Adam: A Method for Stochastic Optimization.” arXiv preprint arXiv:1412.6980 (2017).

  38. Ilya Loshchilov and Frank Hutter. “Decoupled Weight Decay Regularization.” arXiv preprint arXiv:1711.05101 (2019).

  39. Keller Jordan, Yuchen Jin, Vlado Boza, Jiacheng You, Franz Cesista, Laker Newhouse, and Jeremy Bernstein. “Muon: An optimizer for hidden layers in neural networks.” (2024).

  40. Jingyuan Liu, et al. “Muon is Scalable for LLM Training.” arXiv preprint arXiv:2502.16982 (2025).

  41. Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, and Rameswar Panda. “Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler.” arXiv preprint arXiv:2408.13359 (2024).

  42. Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. “Straight to Zero: Why Linearly Decaying the Learning Rate to Zero Works Best for LLMs.” arXiv preprint arXiv:2502.15938 (2025).