Puro-2B: Poor Lab's Qwen2-1.5B Trained on RTX 5090 within $5090 — part 3/7

Cleaned-text code points [31304, 46994). Bibliography omitted here; source.md includes it.

FP8 kernels complete their arithmetic more quickly while introducing additional quantization operations, so they become memory-bound more readily than BF16 kernels. This effect is amplified by the relatively small hidden dimension of our model. Increasing the micro-batch size (MBS) improves the arithmetic intensity and Tensor Core utilization, but an excessive MBS exhausts device memory: DP does not partition activations; and 1F1B pipeline schedule Narayanan et al. [2021] helps control the number of simultaneously resident microbatches, but it does not eliminate the per-microbatch activation-memory growth caused by increasing MBS. We therefore extend Megatron-LM’s analytical FLOP estimator to report the invocation shapes of GEMM and FlashAttention kernels. For each shape induced by our model configuration, we benchmark the corresponding Transformer Engine kernel and select the smallest MBS beyond which achieved FLOP/s no longer increases materially, analogous to locating the knee of a roofline curve.

##### Overcoming load imbalance.

Although our model is relatively small, its vocabulary is large; consequently, the embedding and LM head account for a non-negligible fraction of both computation and memory. FP8 GEMMs accelerate the internal Transformer layers, further increasing their performance gap with the LM head. As a result, the LM head alone entails computation comparable to several Transformer layers, creating substantial pipeline imbalance. Guided by our extended FLOP estimator, which attributes theoretical work by layer, submodule, operator, and precision, we assign fewer Transformer layers to the stage containing the LM head. This balances pipeline computation at the cost of greater memory pressure on the first stage. Tensors assigned to Muon are preferably kept intact rather than sharded. We replace the zigzag placement policy of the layer-wise distributed optimizer in this version of Megatron Core, which over-concentrates the embedding and LM-head tensors, with a memory-aware allocation DeepSeek-AI et al. [2026] : Muon tensors are greedily assigned by memory footprint like bin packing, after which flattened Adam tensors fill the remaining per-device capacity as evenly as possible. This hybrid placement balances optimizer memory across GPUs and is important for fitting the training state into the 32 GB memory available on each RTX 5090.

These principles yield the same fastest configuration as an exhaustive enumeration of the candidate parallel strategies. Together with the hardware optimizations described above and standard Megatron Core options such as the [T,H,D] layout for packed-sequence attention, the Phase 1 production run sustains a median 238 TFLOP/s per GPU with global batch size 1536 on 24 GPUs across three nodes. The best configuration is MBS = 2 {}=2 , PP = 2 {}=2 with layout (18|10) 7 7 7 The PP layout lists the numbers of Transformer layers assigned to successive stages. In (18|10) , the first stage contains the embedding and the first 18 Transformer layers, while the second contains the remaining 10 layers and the LM head. , and DP = 12 {}=12 . The measured throughput corresponds to an equivalent MFU 8 8 8 Model FLOPs utilization (MFU) is conventionally defined as achieved model FLOP/s divided by the accelerator’s peak FLOP/s Korthikanti et al. [2023] . For mixed-precision training, we weight the peak Tensor Core performance by the fraction of theoretical work at each precision. In our recipe, FP8 accounts for 72% of the theoretical Tensor Core operations and BF16 for the remaining 28%, giving an effective peak of P eff = ( 0.72 / 419 + 0.28 / 209.5 ) − 1 ≈ 327 P_{\mathrm{eff}}=(0.72/419+0.28/209.5)^{-1}\approx 327 TFLOP/s per GPU. of approximately 73%. For larger-scale Phase 2 training, strong scaling is limited by Amdahl’s law, and gradient synchronization accounts for a larger fraction of the step time. We consequently use PP = 4 {}=4 with layout (9|9|9|1) and DP = 24 {}=24 on 96 GPUs, while still achieving median throughput 192 TFLOP/s per GPU. The run-level setup and final losses are reported in Table 7 .

### 3.2 FP8 Mixed-Precision Training

We use FP8 mixed precision from random initialization onward, without a BF16 warm-up stage or a later precision switch, using our training stack NVIDIA Corporation [2026a] and its blockwise FP8 support NVIDIA Corporation [2026b] . The training state and numerically sensitive operations retain the standard BF16/FP32 mixed-precision path, while Transformer linear layers use E4M3 operands. FP8 is therefore an online compute and activation-storage format rather than the persistent dtype of the whole model.

E4M3 Micikevicius et al. [2022] provides more mantissa bits than E5M2 but has a narrower dynamic range. With one scale for an entire tensor, a few outliers can enlarge the quantization interval and reduce the effective precision of ordinary values DeepSeek-AI et al. [2024] . We instead compute scales online for local groups: activations and activation gradients use one-dimensional groups of 128 consecutive values along the GEMM reduction dimension, while weights use two-dimensional 128 × 128 128\times 128 blocks Yan et al. [2026] . This fine-grained numerical design follows the recipe introduced and validated at scale by DeepSeek-V3 DeepSeek-AI et al. [2024] .

On our RTX 5090 (SM 120) Corporation [2025a] , each logical block scale is constrained to a power of two and therefore has the exponent-only expressiveness of E8M0. Transformer Engine NVIDIA Corporation [2026b] implements these blocks through Blackwell’s native MXFP8 path. Thus, our logical block sizes follow DeepSeek-V3 DeepSeek-AI et al. [2024] , while the scaling-factor representation and execution path follow MXFP8; the fact that DeepSeek-V3 was trained on Hopper GPUs may account for this implementation difference.

The complete precision flow is shown in Figure 14 . Here, Transformer linear layers include the QKV projections, attention output projection, and MLP projections; their Fprop, Dgrad, and Wgrad GEMMs all use the FP8 path DeepSeek-AI et al. [2024] . Core attention refers only to the FlashAttention/SDPA computation between these projections and remains in BF16 Yan et al. [2026] . A linear layer quantizes its BF16 inputs and weights immediately before GEMM and materializes a BF16 output for the surrounding model. Its quantized input is retained for Wgrad, reducing the dominant saved-activation footprint Yan et al. [2026] . FP8-accelerated operations account for 72% of the overall computation; the cost impact and net benefit of FP8 are analyzed in Figure 2 .

The model-weight portion of each training checkpoint consequently remains BF16; the desired release or deployment representation is produced by a separate post-training conversion. In the matched 20-token-per-parameter scaling ladder, blockwise FP8 increases validation loss by 0.0031–0.0039 relative to BF16 across the five tested model sizes. The shared-shape fit maps this difference to 98.0% BF16-equivalent compute retention. At the 1.7B scale, however, blockwise FP8 improves median training throughput by 1.36 × 1.36\times , yielding a quality-adjusted net gain of 1.34 × 1.34\times . The complete setup, fitting procedure, and size-wise results are reported in Section 4.3.1 .

### 3.3 Hyperball Optimization

In our recipe, we adopt Muon with the Hyperball wrapper (MuonH), introduced by Wen et al. [2026] , as the training optimizer for selected weight matrices. We apply MuonH to approximately scale-invariant weights, including attention and MLP matrices, while embeddings, normalization layers, the language-model head, and the remaining parameters use AdamW. Following Wen et al. [2026] , MuonH normalizes the Muon update and constrains each wrapped matrix to a fixed-radius sphere. More specifically, for each wrapped matrix W t W_{t} , MuonH fixes the radius at R = ∥ W 0 ∥ F R=\lVert W_{0}\rVert_{F} and normalizes the Muon update u t u_{t} as u ^ t = u t / ∥ u t ∥ F \widehat{u}_{t}=u_{t}/\lVert u_{t}\rVert_{F} . It then applies W ~ t + 1 = W t − η t ​ R ​ u ^ t , W t + 1 = R ​ Normalize ⁡ ( W ~ t + 1 ) , \widetilde{W}_{t+1}=W_{t}-\eta_{t}R\widehat{u}_{t},\qquad W_{t+1}=R\,\operatorname{Normalize}(\widetilde{W}_{t+1}), (1) where Normalize ⁡ ( A ) = A / ∥ A ∥ F \operatorname{Normalize}(A)=A/\lVert A\rVert_{F} . The first operation normalizes the update scale, while the second projects the updated weight matrix back to its initial Frobenius radius. The illustration of MuonH optimizer is shown in Figure 4 .

We vary the learning rate η t \eta_{t} according to a two-phase learning rate schedule. The design of this schedule is motivated by an analysis of how the loss gap between MuonH and Muon arises from the perspective of effective learning rate, as we elaborate below.

#### 3.3.1 Effective Learning Rate

First, we briefly revisit the rationale behind Hyperball optimization and explain how it makes the effective LR schedule explicit [ Wen et al., 2026 ] . We then examine the loss gap between ordinary Muon and MuonH runs and find that the effective LR schedule strongly affects the loss convergence rate. This observation leads us to focus on learning rate schedule design for the MuonH optimizer.

##### Scale invariance motivates the effective learning rate.

Hyperball is motivated by approximately scale-invariant parameter matrices. For a selected weight matrix W W , we call the parameter approximately scale invariant when positively rescaling it has little effect on the model function or loss, ℒ ⁡ ( c ​ W ) ≈ ℒ ⁡ ( W ) , c > 0 , \mathcal{L}(cW)\approx\mathcal{L}(W),\qquad c>0, (2) with the remaining parameters held fixed [ Wen et al., 2026 ] . Normalization in modern Transformers makes this approximation relevant for several attention and MLP weight matrices. In this regime, the update magnitude is naturally interpreted relative to the current weight scale. We therefore consider the matrix-wise effective LR, also called the effective learning rate (ELR). For ordinary Muon, with Muon update u t u_{t} , scalar learning rate η t \eta_{t} , and decoupled weight-decay coefficient λ \lambda , W t + 1 = ( 1 − λ ​ η t ) ​ W t − η t ​ u t , W_{t+1}=\bigl(1-\lambda\eta_{t}\bigr)W_{t}-\eta_{t}u_{t}, (3) and, ignoring the radial shrinkage from weight decay, the effective LR is ρ t ​ ( W ) = η t ​ ∥ u t ∥ F ∥ W t ∥ F . \rho_{t}(W)=\frac{\eta_{t}\lVert u_{t}\rVert_{F}}{\lVert W_{t}\rVert_{F}}. (4) Thus, ordinary Muon’s effective LR is not determined by the scalar learning rate alone, because both the weight norm and the Muon-update norm evolve during training.

##### MuonH makes the effective LR schedule explicit.

Under the Hyperball update in Equation 1 , the pre-projection displacement η t ​ R ​ u ^ t \eta_{t}R\widehat{u}_{t} has norm η t ​ R \eta_{t}R , while the wrapped matrix W t W_{t} has norm R R . Its effective LR is therefore ρ t ​ ( W ) = η t ​ ∥ u t ∥ F / ∥ W t ∥ F = η t ​ R / R = η t \rho_{t}(W)=\eta_{t}\lVert u_{t}\rVert_{F}/\lVert W_{t}\rVert_{F}=\eta_{t}R/R=\eta_{t} , numerically equal to the Hyperball weight learning rate, and the subsequent projection restores the fixed radius. Thus, ordinary Muon induces its effective LR indirectly through the scalar learning rate and evolving norms, whereas MuonH directly controls the effective LR schedule through its learning rate schedule. This explicit control is the property of Hyperball that is most relevant to our subsequent schedule design.

In our experiments, we use different learning rates for the MuonH-wrapped and AdamW-optimized parameter groups. We define a base learning rate schedule η t base \eta_{t}^{\mathrm{base}} for the AdamW parameter groups and set the Hyperball weight learning rate to a fixed multiple, η t H = m ​ η t base . \eta_{t}^{H}=m\,\eta_{t}^{\mathrm{base}}. For MuonH-wrapped matrices, η t H \eta_{t}^{H} is numerically equal to the prescribed effective LR. In the production run, we use m = 10 m=10 , while the remaining parameter groups follow η t base \eta_{t}^{\mathrm{base}} directly with AdamW.

##### Matching the effective LR schedule also aligns validation loss.

The MuonH optimizer can make training more efficient than the ordinary Muon optimizer, and we find that the efficiency gap can be attributed to the difference in their effective LR schedule. We compare three 170M-parameter BF16 runs with the same architecture, batch size, sequence length, seed, and evaluation protocol. The MuonH run uses a predefined effective LR schedule ρ t H \rho_{t}^{H} . The ordinary-LR run uses the base LR schedule of the MuonH run as the LR schedule and does not compensate for the evolving weight and update norms. Then, as shown in Figure 5 , the ordinary-LR Muon run produces an effective LR trace that decays rapidly early in training and finishes close to zero. The MuonH run instead follows the predefined linear-decay effective-LR schedule. Although the ordinary-LR Muon run achieves lower validation loss early in training, MuonH overtakes it near the end. The MuonH run finishes at a validation loss of 3.029 3.029 , compared with 3.073 3.073 for ordinary-LR Muon.

For comparison, we conduct an effective-LR aligned run, which uses the ordinary Muon optimizer but adjusts its scalar learning rate online to follow the same effective LR curve, η t = ρ t H ​ ∥ W t ∥ F ∥ u t ∥ F , \eta_{t}=\rho_{t}^{H}\frac{\lVert W_{t}\rVert_{F}}{\lVert u_{t}\rVert_{F}}, (5) so that the effective LR ρ t ​ ( W ) = ρ t H \rho_{t}(W)=\rho_{t}^{H} without projecting W t + 1 W_{t+1} back to a fixed-radius sphere. As shown in Figure 5 , effective-LR aligned Muon follows the prescribed linear post-warmup LR decay and reaches a final validation loss of 3.030 3.030 , close to MuonH at 3.029 3.029 . Thus, in this selected diagnostic, ordinary Muon largely recovers the MuonH loss curve when its effective LR schedule is aligned . This comparison supports the effective-LR schedule as an important factor in MuonH’s training behavior. It also suggests that better control of the effective-LR schedule can improve loss convergence. Recent work also discusses related behavior from the perspective of angular update size [ Zhou et al., 2026b , Xiao et al., 2026 ] . We provide an additional ordinary Muon control in Section F.3 that can even use a counterintuitive hill-like scalar learning rate schedule to achieve loss convergence while maintaining the prescribed effective LR trajectory.

##### Effective LR provides a more informative descriptor of the loss trend.

The experiments above show a close relationship between effective LR and loss dynamics. We further test whether effective LR provides a more informative descriptor than the raw optimizer learning rate. We apply the Multi-Power Law (MPL), a LR-schedule scaling law that predicts loss-curve shape from the LR schedule, to both the effective-LR schedule and the raw learning rate schedule [ Luo et al., 2025b ] . For the two ordinary-Muon runs above, we fit the post-warmup validation curves using either the scalar learning rate or the induced effective LR, while holding out the final 20 % 20\% of validation targets. As shown in Figure 16 , using the effective LR reduces the mean held-out RMSE across the two runs from 0.0265 0.0265 to 0.0210 0.0210 . The improvement is driven mainly by the base-Muon run; for effective-LR aligned Muon, the two LR signals are comparably predictive. This small diagnostic therefore supports effective LR as a more informative descriptor for schedule-dependent training dynamics. The MPL formulation and fitting results are provided in Section F.1 .

##### Effective LR MPL helps interpret the late-stage crossover.