Model save

README.md

@@ -1,303 +1,74 @@
 ---
-language: en
-license: apache-2.0
+library_name: transformers
 tags:
-- causal-lm
-- research
-- bf16
-- attention
-- normalization
-- neollm
-datasets:
-- HuggingFaceFW/fineweb-edu
+- generated_from_trainer
+model-index:
+- name: NeoLLM
+  results: []
 ---
 
-# NeoLLM
-
-NeoLLM is a **84.57 M parameter** decoder-only language model trained from scratch on
-[FineWeb-Edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) with
-**BF16** compute, completing training in approximately
-**~1h 11m** on **NVIDIA GeForce RTX 5090**.
-It integrates a collection of recently published attention and normalization techniques
-into a single architecture, with the goal of studying how they interact during
-pretraining. The model is actively being developed and the current checkpoint represents
-an intermediate training state.
-
-> **Author / contact:** [@Kyokopom](https://x.com/Kyokopom) on X
-> **Repository:** [KitsuVp/NeoLLM](https://huggingface.co/KitsuVp/NeoLLM)
-
----
-
-## Architecture
-
-NeoLLM is a decoder-only transformer with the following configuration:
-
-| Parameter | Value |
-|---|---|
-| Hidden size | 512 |
-| Layers | 12 |
-| Attention heads | 8 |
-| KV heads (GQA) | 4 |
-| Head dim | 64 |
-| Intermediate size | 1536 |
-| Vocabulary | LiquidAI/LFM2.5-1.2B-Thinking tokenizer (64,402 tokens) |
-| Context length | 512 tokens |
-
-### Parameter breakdown
-
-| Parameter bucket | Count |
-|---|---|
-| **Total parameters** | 84.57M (84,567,928) |
-| **Embedding parameters** (tied) | 32.97M (32,973,824) |
-| **Non-embedding parameters** | 51.59M (51,594,104) |
-| **Effective trainable parameters** | 84.57M (84,567,928) |
-
-> Weight tying is **enabled**: the input embedding matrix and the language-model head
-> share the same parameters, so the effective trainable budget is
-> `total − embed = 51.59M`.
-
-### Integrated techniques
-
-NeoLLM combines architecture modules, optional auxiliary objectives, and
-training-time optimizer/stability components from the following papers.
-
-**Embedding and token representation**
-
-- **Learnable Multipliers** ([arXiv:2601.04890](https://arxiv.org/abs/2601.04890)) — Adds
-  per-row and per-column learnable scalar parameters to selected matrix layers and, when
-  enabled, embeddings.
-- **Leviathan** ([arXiv:2601.22040](https://arxiv.org/abs/2601.22040)) — Optional
-  continuous token embedding generator that can replace the discrete input lookup table.
-- **KHRONOS** ([arXiv:2505.13315](https://arxiv.org/abs/2505.13315)) — Kernel/basis
-  reference used by the Leviathan continuous token generator implementation.
-- **Spelling Bee Embeddings** ([arXiv:2601.18030](https://arxiv.org/abs/2601.18030)) —
-  Augments token embeddings with character-level spelling information.
-- **Token Embedding Manifold analysis** ([arXiv:2504.01002](https://arxiv.org/abs/2504.01002)) —
-  Reference motivation for treating token embeddings as structured objects rather than
-  unconstrained lookup rows.
-
-**Attention, positions, and output projection**
-
-- **FAN** ([arXiv:2502.21309](https://arxiv.org/abs/2502.21309)) — Fourier Analysis Networks.
-  A portion of the projection channels are dedicated to periodic cosine/sine features.
-- **MEA** ([arXiv:2601.19611](https://arxiv.org/abs/2601.19611)) — Explicit Multi-head
-  Attention. Adds small learnable interaction matrices between attention heads for K and V.
-- **LUCID** ([arXiv:2602.10410](https://arxiv.org/abs/2602.10410)) — Applies a learned
-  lower-triangular preconditioner to V before attention, decorrelating value representations
-  across positions.
-- **Affine-Scaled Attention** ([arXiv:2602.23057](https://arxiv.org/abs/2602.23057)) — Adds
-  two learnable per-head scalars (α and β) to the softmax weights:
-  `[α·softmax(QKᵀ) + β]·V`.
-- **XSA** ([arXiv:2603.09078](https://arxiv.org/abs/2603.09078)) — Exclusive Self Attention.
-  After computing attention, removes the component of the output aligned with the token's
-  own value vector.
-- **Directional Routing** ([arXiv:2603.14923](https://arxiv.org/abs/2603.14923)) — Each head
-  learns K=4 directions in the output space; a learned router suppresses the attention output
-  along each direction per input.
-- **Gated Attention** ([arXiv:2505.06708](https://arxiv.org/abs/2505.06708)) — A sigmoid gate
-  is applied to the attention output before the output projection, introducing non-linearity
-  and preventing attention sinks.
-- **Momentum Attention** ([arXiv:2411.03884](https://arxiv.org/abs/2411.03884)) — Modifies Q
-  and K by subtracting a fraction of the previous position's Q and K values (causal
-  first-difference).
-- **Interleaved Head Attention / IHA** ([arXiv:2602.21371](https://arxiv.org/abs/2602.21371)) —
-  Builds pseudo-heads from learned cross-head mixtures to create multiple attention patterns
-  per original head.
-- **REPO** ([arXiv:2512.14391](https://arxiv.org/abs/2512.14391)) — Context re-positioning
-  module that learns contextual position coordinates above a configurable start layer.
-- **GRAPE** ([arXiv:2512.07805](https://arxiv.org/abs/2512.07805)) — Group representational
-  position encoding used by the REPO-GRAPE positional path.
-- **GOAT priors** ([arXiv:2601.15380](https://arxiv.org/abs/2601.15380)) — Optional
-  factorized attention log-prior channels inspired by trainable attention priors.
-- **Hadamard output projection** ([arXiv:2603.08343](https://arxiv.org/abs/2603.08343)) —
-  Replaces dense attention output projection with a structured Hadamard transform plus
-  lightweight scaling.
-
-**Normalization, residual flow, and MLP**
-
-- **SeeDNorm** ([arXiv:2510.22777](https://arxiv.org/abs/2510.22777)) — Applied to Q and K
-  projections. Dynamically rescales normalization from the input's own statistics.
-- **LayerNorm Scaling / LNS** ([arXiv:2502.05795](https://arxiv.org/abs/2502.05795)) — Each
-  layer's output is scaled by 1/√ℓ where ℓ is the layer index.
-- **GPAS** ([arXiv:2506.22049](https://arxiv.org/abs/2506.22049)) — Gradient-Preserving
-  Activation Scaling for residual junctions.
-- **PolyNorm** ([arXiv:2602.04902](https://arxiv.org/abs/2602.04902)) — Replaces the standard
-  MLP activation with normalized linear, quadratic, and cubic branches.
-- **SimpleGPT** ([arXiv:2602.01212](https://arxiv.org/abs/2602.01212)) — Second-order
-  geometry-inspired normalization strategy applied inside MLP projections.
-- **StackMemory / STACKTRANS** ([NeurIPS 2025](https://openreview.net/forum?id=2bbDg587uh)) —
-  Optional differentiable hidden-state stack between decoder layers.
-- **Attention Residuals / AttnRes** ([arXiv:2603.15031](https://arxiv.org/abs/2603.15031)) —
-  Optional learned depth-wise aggregation over previous layer outputs or block summaries.
-- **LAUREL** ([arXiv:2411.07501](https://arxiv.org/abs/2411.07501)) — Optional learned
-  augmented residual layer with residual-weight and low-rank variants.
+<!-- This model card has been generated automatically according to the information the Trainer had access to. You
+should probably proofread and complete it, then remove this comment. -->
 
-**Training objectives and training-time regularizers**
-
-- **Cut Cross Entropy** ([Apple repository](https://github.com/apple/ml-cross-entropy)) —
-  Memory-efficient next-token loss that avoids materializing the full token-by-vocabulary
-  logits tensor. NeoLLM remains compatible with the upstream package when the extensions
-  below are disabled.
-- **MiLe Loss** ([arXiv:2310.19531](https://arxiv.org/abs/2310.19531)) — Optional detached,
-  mean-normalized predictive-entropy weighting of token losses, implemented inside the
-  extended CCE path.
-- **Output Embedding Centering / mu-loss**
-  ([arXiv:2601.02031](https://arxiv.org/abs/2601.02031)) — Optional
-  `lambda * ||mean(output_embeddings)||^2` regularizer for output-logit stability.
-- **MEAP** ([arXiv:2502.07490](https://arxiv.org/abs/2502.07490)) — Optional training-only
-  input corruption that masks a fixed fraction of eligible tokens while preserving clean
-  next-token labels, causal attention, and the inference path.
-- **TWEO** ([arXiv:2511.23225](https://arxiv.org/abs/2511.23225)) — Optional
-  Transformers Without Extreme Outliers activation regularizer for FP8/low-bit-friendly
-  training.
-- **NITP** ([arXiv:2605.24956](https://arxiv.org/abs/2605.24956)) — Optional Next Implicit
-  Token Prediction auxiliary objective using shallow-layer implicit token targets and a
-  cosine loss.
-- **NextLat** ([arXiv:2511.05963](https://arxiv.org/abs/2511.05963)) — Optional next-latent
-  prediction objective using latent dynamics, Smooth L1 supervision, and frozen-head KL.
-
-### Optional extended-CCE configuration
-
-| Feature | Enabled | Value |
-|---|---:|---:|
-| MiLe Loss | True | gamma=1.0 |
-| mu-loss | True | lambda=0.0001 |
-| MEAP | True | ratio=0.15 |
-
-MiLe, mu-loss, and MEAP require the extended
-[`Kitsunp/ml-cross-entropy`](https://github.com/Kitsunp/ml-cross-entropy) package only when
-their corresponding flags are enabled. With all flags disabled, NeoLLM calls upstream CCE
-without extension-specific arguments. When any extension is active, CCE reports three compact
-scalars: unweighted NTP cross entropy, the MiLe reweighting delta, and the mu-loss penalty.
-Their sum reconstructs `ntp_loss` exactly. MEAP reports its eligible and selected counts from
-the masking kernel; the trainer logs the selected count and exact fraction. No diagnostic path
-materializes full-vocabulary logits or a token mask outside the kernels.
-
-**Optimizer and training stability**
-
-- **Adam-family optimizer mode** — Active variant: **AdamW**.
-  **AdEMAMix** ([arXiv:2409.03137](https://arxiv.org/abs/2409.03137)) contributes the optional
-  beta3 slow-memory branch when `ADEMAMIX=True`. **DeltaMomentum**
-  ([arXiv:2608.19491](https://arxiv.org/abs/2608.19491)) optionally replaces only the fast beta1
-  first-moment update on hidden Linear weights when `DELTA=True`; beta2 remains Adam's
-  ordinary squared-gradient EMA in every mode.
-- **BF16 stochastic rounding** ([arXiv:2502.20566](https://arxiv.org/abs/2502.20566)) —
-  Optional master-free BF16 model update with stochastic rather than nearest rounding; the same
-  write-back path is shared by AdamW, DeltaAdamW, AdEMAMix, and DeltaAdEMAMix.
-- **Cautious Weight Decay** ([arXiv:2510.12402](https://arxiv.org/abs/2510.12402)) —
-  Sign-selective weight decay variant used by the custom optimizer logic.
-- **Correction of Decoupled Weight Decay** ([arXiv:2512.08217](https://arxiv.org/abs/2512.08217)) —
-  Adapts decoupled weight decay during learning-rate decay.
-- **AdamHD** ([arXiv:2511.14721](https://arxiv.org/abs/2511.14721)) —
-  Decoupled Huber decay regularization reference used by the optimizer.
-- **GradientStabilizer** ([arXiv:2502.17055](https://arxiv.org/abs/2502.17055)) —
-  Active; repository per-parameter tensor norms; bias correction=True; no stabilizer warmup; EMA state stored in optimizer checkpoints.
----
-
-## Training
-
-| Setting | Value |
-|---|---|
-| Dataset | FineWeb-Edu (sample-10BT) |
-| Tokens seen | ~0.51B (15,625 steps × batch 64 × length 512) |
-| Precision | BF16 model state (bfloat16: 84,565,816 elements, float32: 2,112 elements) + optimizer policy unavailable |
-| Optimizer | AdamW |
-| Gradient transform | Active; repository per-parameter tensor norms; bias correction=True; no stabilizer warmup; EMA state stored in optimizer checkpoints |
-| Learning rate | 1e-03 with 4,000-step linear warmup |
-| Weight decay | coefficients [0.005, 0.01, 0.1]; WD 5/5 groups; CWD 2; Huber 0; Correction 0 |
-| Training time | ~1h 11m |
-| Hardware | NVIDIA GeForce RTX 5090 |
+# NeoLLM
 
-### Training curve
+This model is a fine-tuned version of [](https://huggingface.co/) on an unknown dataset.
+It achieves the following results on the evaluation set:
+- Loss: 4.8333
+- Ntp Loss: 4.2436
+- Ntp Ce Unweighted: 3.5923
+- Mile Reweighting Delta: 0.6512
+- Mu Loss: 0.0001
+- Tweo Loss: 0.0016
+- Nitp Loss: 0.4902
+- Total Model Loss: 4.7338
+- Meap Mask Embedding Rms: 0.0432
+- Meap Mask Embedding Max Abs: 0.5781
+- Gradient Stabilizer Target Global Norm: 18.9319
+- Gradient Stabilizer Mean Tensor Scale: 330.9539
+- Gradient Stabilizer Min Tensor Scale: 1.3281
+- Gradient Stabilizer Max Tensor Scale: 1536.0
+- Gradient Stabilizer Active Parameter Tensors: 366.0
+- Gradient Stabilizer Skipped Parameter Tensors: 0.0
 
-| Step | Train Loss | Val Loss |
-|---|---|---|
-| 5,000 | 5.772 | 5.396 |
-| 10,000 | 5.516 | 5.115 |
-| 15,000 | 5.305 | 4.875 |
-| 15,625 | — | 4.842 |
+## Model description
 
----
+More information needed
 
-## Limitations
+## Intended uses & limitations
 
-- **Token budget** — ~1.5 B tokens seen; below estimated optimum. Knowledge-intensive tasks
-  will improve with more training.
-- **Gradient spike at step 40k** — Reorganized the attention pattern in layer 9 that
-  previously captured long-range token correlations. A checkpoint from ~step 38k is expected
-  to have better aggregate benchmark scores.
-- **PolyNorm exclusivity** — The quadratic branch has become partially redundant with the
-  linear branch. Will be corrected in the next training run.
-- **Base model only** — Not instruction-tuned or aligned; purely a next-token-prediction
-  base model.
+More information needed
 
----
+## Training and evaluation data
 
-## References
+More information needed
 
-All papers whose techniques are integrated into NeoLLM's architecture,
-training objective, or training stack:
+## Training procedure
 
-| Area | Technique | Paper title | Reference |
-|---|---|---|---|
-| Embeddings | Learnable Multipliers | Freeing the Scale of Language Model Matrix Layers | [arXiv:2601.04890](https://arxiv.org/abs/2601.04890) |
-| Embeddings | Leviathan | A Separable Architecture for Continuous Token Representation in Language Models | [arXiv:2601.22040](https://arxiv.org/abs/2601.22040) |
-| Embeddings | KHRONOS | KHRONOS: a Kernel-Based Neural Architecture for Rapid, Resource-Efficient Scientific Computation | [arXiv:2505.13315](https://arxiv.org/abs/2505.13315) |
-| Embeddings | Spelling Bee | Spelling Bee Embeddings for Language Modeling | [arXiv:2601.18030](https://arxiv.org/abs/2601.18030) |
-| Embeddings | Token embedding analysis | Token Embeddings Violate the Manifold Hypothesis | [arXiv:2504.01002](https://arxiv.org/abs/2504.01002) |
-| Attention / positions | FAN | Fourier Analysis Networks | [arXiv:2502.21309](https://arxiv.org/abs/2502.21309) |
-| Attention / positions | MEA | Explicit Multi-head Attention for Inter-head Interaction in Large Language Models | [arXiv:2601.19611](https://arxiv.org/abs/2601.19611) |
-| Attention / positions | LUCID | Attention with Preconditioned Representations | [arXiv:2602.10410](https://arxiv.org/abs/2602.10410) |
-| Attention / positions | Affine-Scaled Attention | Affine-Scaled Attention: Towards Flexible and Stable Transformer Attention | [arXiv:2602.23057](https://arxiv.org/abs/2602.23057) |
-| Attention / positions | XSA | Exclusive Self Attention | [arXiv:2603.09078](https://arxiv.org/abs/2603.09078) |
-| Attention / positions | Directional Routing | Directional Routing in Transformers | [arXiv:2603.14923](https://arxiv.org/abs/2603.14923) |
-| Attention / positions | Gated Attention | Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free | [arXiv:2505.06708](https://arxiv.org/abs/2505.06708) |
-| Attention / positions | Momentum Attention | Momentum Attention | [arXiv:2411.03884](https://arxiv.org/abs/2411.03884) |
-| Attention / positions | IHA | Interleaved Head Attention | [arXiv:2602.21371](https://arxiv.org/abs/2602.21371) |
-| Attention / positions | REPO | Language Models with Context Re-Positioning | [arXiv:2512.14391](https://arxiv.org/abs/2512.14391) |
-| Attention / positions | GRAPE | Group Representational Position Encoding | [arXiv:2512.07805](https://arxiv.org/abs/2512.07805) |
-| Attention / positions | GOAT priors | You Need Better Attention Priors | [arXiv:2601.15380](https://arxiv.org/abs/2601.15380) |
-| Attention / positions | Hadamard o_proj | Rethinking Attention Output Projection: Structured Hadamard Transforms for Efficient Transformers | [arXiv:2603.08343](https://arxiv.org/abs/2603.08343) |
-| Residual / normalization | SeeDNorm | Self-Rescaled Dynamic Normalization | [arXiv:2510.22777](https://arxiv.org/abs/2510.22777) |
-| Residual / normalization | LNS | The Curse of Depth in LLMs | [arXiv:2502.05795](https://arxiv.org/abs/2502.05795) |
-| Residual / normalization | GPAS | Gradient-Preserving Activation Scaling | [arXiv:2506.22049](https://arxiv.org/abs/2506.22049) |
-| Residual / normalization | PolyNorm | PolyNorm / PolyCom | [arXiv:2602.04902](https://arxiv.org/abs/2602.04902) |
-| Residual / normalization | SimpleGPT | SimpleGPT | [arXiv:2602.01212](https://arxiv.org/abs/2602.01212) |
-| Residual / normalization | StackMemory / STACKTRANS | Recursive Transformer: Boosting Reasoning Ability with State Stack | [NeurIPS 2025](https://openreview.net/forum?id=2bbDg587uh) |
-| Residual / normalization | Attention Residuals | Attention Residuals | [arXiv:2603.15031](https://arxiv.org/abs/2603.15031) |
-| Residual / normalization | LAUREL | LAUREL: Learned Augmented Residual Layer | [arXiv:2411.07501](https://arxiv.org/abs/2411.07501) |
-| Objectives | TWEO | Transformers Without Extreme Outliers Enables FP8 Training And Quantization For Dummies | [arXiv:2511.23225](https://arxiv.org/abs/2511.23225) |
-| Objectives | NITP | Next Implicit Token Prediction for LLM Pre-training | [arXiv:2605.24956](https://arxiv.org/abs/2605.24956) |
-| Objectives | NextLat | Next-Latent Prediction Transformers Learn Compact World Models | [arXiv:2511.05963](https://arxiv.org/abs/2511.05963) |
-| Optimizer / training | AdEMAMix | The AdEMAMix Optimizer: Better, Faster, Older | [arXiv:2409.03137](https://arxiv.org/abs/2409.03137) |
-| Optimizer / training | DeltaMomentum | DeltaMomentum: A Key-Value based Anisotropic Momentum Update via Delta Rule | [arXiv:2608.19491](https://arxiv.org/abs/2608.19491) |
-| Optimizer / training | BF16 + SR | Stochastic Rounding for LLM Training: Theory and Practice | [arXiv:2502.20566](https://arxiv.org/abs/2502.20566) |
-| Optimizer / training | CWD | Cautious Weight Decay | [arXiv:2510.12402](https://arxiv.org/abs/2510.12402) |
-| Optimizer / training | WD correction | Correction of Decoupled Weight Decay | [arXiv:2512.08217](https://arxiv.org/abs/2512.08217) |
-| Optimizer / training | AdamHD | AdamHD: Decoupled Huber Decay Regularization for Language Model Pre-Training | [arXiv:2511.14721](https://arxiv.org/abs/2511.14721) |
-| Optimizer / training | GradientStabilizer | GradientStabilizer | [arXiv:2502.17055](https://arxiv.org/abs/2502.17055) |
+### Training hyperparameters
 
----
+The following hyperparameters were used during training:
+- learning_rate: 0.001
+- train_batch_size: 64
+- eval_batch_size: 64
+- seed: 42
+- optimizer: Use OptimizerNames.ADAMW_TORCH_FUSED with betas=(0.9,0.999) and epsilon=1e-08 and optimizer_args=No additional optimizer arguments
+- lr_scheduler_type: linear
+- num_epochs: 1
 
-## Citation
+### Training results
 
-```bibtex
-@misc{neollm2026,
-  title  = {NeoLLM: A Research Language Model Integrating Recent Attention and Normalization Techniques},
-  author = {KitsuVp},
-  year   = {2026},
-  url    = {https://huggingface.co/KitsuVp/NeoLLM}
-}
-```
+| Training Loss | Epoch | Step  | Validation Loss | Loss   | Ce Unweighted | Reweighting Delta | Model Loss | Mask Embedding Rms | Mask Embedding Max Abs | Stabilizer Target Global Norm | Stabilizer Mean Tensor Scale | Stabilizer Min Tensor Scale | Stabilizer Max Tensor Scale | Stabilizer Active Parameter Tensors | Stabilizer Skipped Parameter Tensors |
+|:-------------:|:-----:|:-----:|:---------------:|:------:|:-------------:|:-----------------:|:----------:|:------------------:|:----------------------:|:-----------------------------:|:----------------------------:|:---------------------------:|:---------------------------:|:-----------------------------------:|:------------------------------------:|
+| 5.7803        | 0.32  | 5000  | 5.3890          | 0.5638 | 4.1975        | 0.5496            | 5.3116     | 0.0483             | 0.5820                 | 18.6946                       | 495.0521                     | 2.4380                      | 3248.0                      | 366.0                               | 0.0                                  |
+| 5.5108        | 0.64  | 10000 | 5.1084          | 0.5126 | 3.8947        | 0.6121            | 5.0197     | 0.0525             | 0.7188                 | 18.9658                       | 369.8585                     | 2.5312                      | 1800.0                      | 366.0                               | 0.0                                  |
+| 5.3037        | 0.96  | 15000 | 4.8615          | 0.4918 | 3.6319        | 0.6389            | 4.7627     | 0.0444             | 0.6016                 | 18.7181                       | 330.5682                     | 1.5781                      | 1600.0                      | 366.0                               | 0.0                                  |
+| 5.2719        | 1.0   | 15625 | 4.8333          | 0.4902 | 3.5923        | 0.6512            | 4.7338     | 0.0432             | 0.5781                 | 18.9319                       | 330.9539                     | 1.3281                      | 1536.0                      | 366.0                               | 0.0                                  |
 
----
-
-## Author
-
-[@Kyokopom](https://x.com/Kyokopom) on X
-
----
 
-## License
+### Framework versions
 
-Apache 2.0
+- Transformers 5.16.1
+- Pytorch 2.14.0+cu132
+- Datasets 5.0.1
+- Tokenizers 0.23.2

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e2821a8346cf052244d8a914302540f11ab4dc30e558c2b8f33aa22ddba7d693
+oid sha256:64483bfee890f57b9da0bf3e361619f03f9a75f9db76f7dd87434e251efc08e7
 size 169183584

optimizer_precision.json

@@ -1,18 +1,18 @@
 {
   "ademamix_enabled": false,
-  "bf16_stochastic_round_active": false,
-  "bf16_stochastic_round_elements": 0,
-  "bf16_stochastic_round_groups": 0,
-  "bf16_stochastic_round_tensors": 0,
+  "bf16_stochastic_round_active": true,
+  "bf16_stochastic_round_elements": 84565816,
+  "bf16_stochastic_round_groups": 6,
+  "bf16_stochastic_round_tensors": 350,
   "delta_eligible_linear_weights": 0,
   "delta_enabled": false,
   "delta_eta": 0.37,
-  "fp32_master_active": true,
-  "fp32_master_elements": 84567928,
-  "fp32_master_groups": 6,
-  "fp32_master_tensors": 366,
-  "fp32_passthrough_elements": 0,
-  "mode": "fp32_master",
+  "fp32_master_active": false,
+  "fp32_master_elements": 0,
+  "fp32_master_groups": 0,
+  "fp32_master_tensors": 0,
+  "fp32_passthrough_elements": 2112,
+  "mode": "bf16_stochastic_rounding",
   "mxfp8_active": false,
   "mxfp8_available": false,
   "mxfp8_linear_layers": 0,
@@ -25,6 +25,6 @@
     "float32": 2112
   },
   "parameter_storage_bytes": 169140080,
-  "runtime_precision_mode": "BF16_FP32_MASTER",
+  "runtime_precision_mode": "BF16_BF16_SR_MASTER_FREE",
   "stochastic_round_seed": 0
 }

training_args.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3a4f920c57b0390c72c9efad34abd3c9cf10157d9058955d0a31b31a393ba8b6
+oid sha256:ffb37840e69056ed9d86138ec65b160df0dbc9fe277466cbe65de6223d6d3f1c
 size 5329