Upload 12 files

README.md

@@ -0,0 +1,67 @@
+---
+language:
+- en
+library_name: transformers
+pipeline_tag: text-generation
+datasets:
+- epfml/FineWeb-HQ
+- HuggingFaceTB/smollm-corpus
+tags:
+- causal-lm
+- base-model
+- muon
+- xsa-refresh
+- custom-code
+- trust-remote-code
+---
+
+# BananaMind 2 Pico Test - 5%
+
+This is the **5% checkpoint** of a 900,002-parameter
+base causal language model. It is not instruction tuned.
+
+## Architecture
+
+| Field | Value |
+|---|---:|
+| Parameters | 900,002 |
+| Layers / hidden size | 6 / 96 |
+| SwiGLU intermediate size | 380 |
+| Query / KV heads | 6 / 2 |
+| Head dimension | 16 |
+| Context | 4,096 |
+| Vocabulary | 384, tied |
+| Refresh layers | 4 and 6 |
+| Refresh kernel | Causal depthwise, width 9 |
+
+The selective XSA refresh gate reads detached attention output as its signal,
+reinjects the original input embedding as its value, and carries convolution
+history alongside the K/V cache. Its learned residual scalar starts at zero.
+
+## Training
+
+| Field | Value |
+|---|---:|
+| Progress | 5% |
+| Tokens seen | 10,000,269,312 |
+| Target tokens | 200,000,000,000 |
+| Hardware | 4 x NVIDIA H200 |
+| Matrix optimizer | Stock `torch.optim.Muon` |
+| Muon peak LR | 0.07 |
+| Embedding/control optimizer | AdamW, LR 0.004 |
+| Precision | bfloat16 autocast |
+
+| Token range | FineWeb-HQ | Cosmopedia v2 |
+|---|---:|---:|
+| 0.00B-100.00B | 80% | 20% |
+| 100.00B-200.00B | 60% | 40% |
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "Banaxi-Tech/pico-test"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
+```

checkpoint_metadata.json

@@ -0,0 +1,72 @@
+{
+  "parameters": 900002,
+  "architecture": {
+    "vocab_size": 384,
+    "hidden_size": 96,
+    "num_hidden_layers": 6,
+    "num_attention_heads": 6,
+    "num_key_value_heads": 2,
+    "head_dim": 16,
+    "intermediate_size": 380,
+    "refresh_layers": [
+      3,
+      5
+    ],
+    "refresh_kernel_size": 9,
+    "refresh_alpha_init": 0.0,
+    "max_position_embeddings": 4096,
+    "rope_theta": 100000.0,
+    "rms_norm_eps": 1e-06,
+    "initializer_range": 0.02,
+    "use_cache": true,
+    "transformers_version": "5.15.0",
+    "architectures": null,
+    "output_hidden_states": false,
+    "return_dict": true,
+    "dtype": null,
+    "chunk_size_feed_forward": 0,
+    "is_encoder_decoder": false,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "problem_type": null,
+    "_name_or_path": "",
+    "tie_word_embeddings": true,
+    "bos_token_id": 1,
+    "eos_token_id": 2,
+    "pad_token_id": 0,
+    "unk_token_id": 3,
+    "model_type": "bananamind2_pico",
+    "output_attentions": false
+  },
+  "training_percent": 5,
+  "step": 1122,
+  "total_steps": 22439,
+  "tokens_seen": 10000269312,
+  "target_tokens": 200000000000,
+  "trained_token_target": 199996473344,
+  "tokens_per_step": 8912896,
+  "local_batch": 544,
+  "effective_global_batch": 2176,
+  "world_size": 4,
+  "gpu_name": "NVIDIA H200",
+  "optimizer": "Muon + AdamW",
+  "muon_peak_lr": 0.07,
+  "adamw_peak_lr": 0.004,
+  "architecture_revision": "ee5739850035c878ce4d9fc29fc38b075987e03b",
+  "tokenizer_revision": "ee5739850035c878ce4d9fc29fc38b075987e03b",
+  "dataset_revisions": {
+    "fineweb_hq": "e58199cdd52438d94405df1a4d8630cc5f13bf84",
+    "cosmopedia_v2": "3ba9d605774198c5868892d7a8deda78031a781f"
+  },
+  "source_tokens": {
+    "fineweb_hq": 8003780608,
+    "cosmopedia_v2": 1996488704
+  },
+  "elapsed_seconds": 1523.7226781845093
+}

config.json

@@ -0,0 +1,51 @@
+{
+  "vocab_size": 384,
+  "hidden_size": 96,
+  "num_hidden_layers": 6,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 2,
+  "head_dim": 16,
+  "intermediate_size": 380,
+  "refresh_layers": [
+    3,
+    5
+  ],
+  "refresh_kernel_size": 9,
+  "refresh_alpha_init": 0.0,
+  "max_position_embeddings": 4096,
+  "rope_theta": 100000.0,
+  "rms_norm_eps": 1e-06,
+  "initializer_range": 0.02,
+  "use_cache": true,
+  "transformers_version": "5.15.0",
+  "architectures": [
+    "BananaMind2PicoForCausalLM"
+  ],
+  "output_hidden_states": false,
+  "return_dict": true,
+  "dtype": null,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "problem_type": null,
+  "_name_or_path": "Banaxi-Tech/pico-test",
+  "tie_word_embeddings": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "unk_token_id": 3,
+  "model_type": "bananamind2_pico",
+  "output_attentions": false,
+  "auto_map": {
+    "AutoConfig": "configuration_bananamind2pico.BananaMind2PicoConfig",
+    "AutoModelForCausalLM": "modeling_bananamind2pico.BananaMind2PicoForCausalLM"
+  },
+  "torch_dtype": "float32"
+}

configuration_bananamind2pico.py

@@ -0,0 +1,111 @@
+"""Configuration for BananaMind 2 Pico."""
+
+from __future__ import annotations
+
+from transformers import PretrainedConfig
+
+
+class BananaMind2PicoConfig(PretrainedConfig):
+    model_type = "bananamind2_pico"
+
+    def __init__(
+        self,
+        vocab_size: int = 384,
+        hidden_size: int = 96,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 6,
+        num_key_value_heads: int = 2,
+        head_dim: int = 16,
+        intermediate_size: int = 380,
+        refresh_layers: tuple[int, ...] | list[int] = (3, 5),
+        refresh_kernel_size: int = 9,
+        refresh_alpha_init: float = 0.0,
+        max_position_embeddings: int = 4096,
+        rope_theta: float = 100000.0,
+        rms_norm_eps: float = 1e-6,
+        initializer_range: float = 0.02,
+        tie_word_embeddings: bool = True,
+        use_cache: bool = True,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        pad_token_id: int = 0,
+        unk_token_id: int = 3,
+        **kwargs,
+    ):
+        if hidden_size != num_attention_heads * head_dim:
+            raise ValueError("hidden_size must equal num_attention_heads * head_dim")
+        if num_attention_heads % num_key_value_heads:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if intermediate_size <= 0:
+            raise ValueError("intermediate_size must be positive")
+        if refresh_kernel_size <= 0 or refresh_kernel_size % 2 == 0:
+            raise ValueError("refresh_kernel_size must be a positive odd integer")
+        refresh_layers = tuple(int(index) for index in refresh_layers)
+        if len(refresh_layers) != len(set(refresh_layers)):
+            raise ValueError("refresh_layers must not contain duplicates")
+        if any(index < 0 or index >= num_hidden_layers for index in refresh_layers):
+            raise ValueError("refresh_layers contains an out-of-range zero-based index")
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.refresh_layers = list(refresh_layers)
+        self.refresh_kernel_size = refresh_kernel_size
+        self.refresh_alpha_init = refresh_alpha_init
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            unk_token_id=unk_token_id,
+            **kwargs,
+        )
+
+    def parameter_counts(self) -> dict[str, int]:
+        """Return an analytical parameter breakdown for drift checks."""
+        hidden = self.hidden_size
+        query_width = self.num_attention_heads * self.head_dim
+        kv_width = self.num_key_value_heads * self.head_dim
+        embedding = self.vocab_size * hidden
+        attention = hidden * query_width + 2 * hidden * kv_width + query_width * hidden
+        qk_norms = 2 * self.head_dim
+        mlp = 3 * hidden * self.intermediate_size
+        block_norms = 2 * hidden
+        refresh = (
+            3 * hidden * hidden
+            + hidden * self.refresh_kernel_size
+            + 3 * hidden
+            + 1
+        )
+        base_layer = attention + qk_norms + mlp + block_norms
+        total = (
+            embedding
+            + self.num_hidden_layers * base_layer
+            + len(self.refresh_layers) * refresh
+            + hidden
+        )
+        if not self.tie_word_embeddings:
+            total += embedding
+        return {
+            "total": total,
+            "embedding": embedding,
+            "base_layer": base_layer,
+            "attention_per_layer": attention + qk_norms,
+            "mlp_per_layer": mlp,
+            "refresh_per_selected_layer": refresh,
+            "refresh_layer_count": len(self.refresh_layers),
+        }
+
+
+BananaMind2PicoConfig.register_for_auto_class("AutoConfig")
+
+__all__ = ["BananaMind2PicoConfig"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "transformers_version": "5"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:27e72a094ac7acdbced367943b7cee235132784f210df46fe7580b646827d290
+size 3756344

modeling_bananamind2pico.py

@@ -0,0 +1,456 @@
+"""BananaMind 2 Pico causal LM with selective XSA refresh gates."""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_bananamind2pico import BananaMind2PicoConfig
+except ImportError:  # Standalone training jobs import the mounted/downloaded files.
+    from configuration_bananamind2pico import BananaMind2PicoConfig
+
+
+class BananaMind2PicoRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        states = hidden_states.float()
+        states = states * torch.rsqrt(states.square().mean(-1, keepdim=True) + self.eps)
+        return (states * self.weight.float()).to(hidden_states.dtype)
+
+
+def _rope_cos_sin(
+    head_dim: int,
+    positions: torch.Tensor,
+    theta: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        theta
+        ** (
+            torch.arange(0, head_dim, 2, dtype=torch.float32, device=positions.device)
+            / head_dim
+        )
+    )
+    frequencies = torch.outer(positions.float(), inv_freq)
+    return frequencies.cos(), frequencies.sin()
+
+
+def _apply_rope(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cosine: torch.Tensor,
+    sine: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    query_dtype, key_dtype = query.dtype, key.dtype
+    cosine = cosine[None, None, :, :]
+    sine = sine[None, None, :, :]
+    query_pairs = query.float().reshape(*query.shape[:-1], -1, 2)
+    key_pairs = key.float().reshape(*key.shape[:-1], -1, 2)
+    query_even, query_odd = query_pairs.unbind(-1)
+    key_even, key_odd = key_pairs.unbind(-1)
+    query = torch.stack(
+        (query_even * cosine - query_odd * sine, query_even * sine + query_odd * cosine),
+        dim=-1,
+    ).flatten(-2)
+    key = torch.stack(
+        (key_even * cosine - key_odd * sine, key_even * sine + key_odd * cosine),
+        dim=-1,
+    ).flatten(-2)
+    return query.to(query_dtype), key.to(key_dtype)
+
+
+class BananaMind2PicoCache(DynamicCache):
+    """K/V cache plus causal-convolution history for selected refresh layers."""
+
+    def __init__(self, config: BananaMind2PicoConfig):
+        try:
+            super().__init__(config=config)
+        except TypeError:
+            super().__init__()
+        self.refresh_states: list[torch.Tensor | None] = [
+            None for _ in range(config.num_hidden_layers)
+        ]
+
+    def refresh_input(
+        self,
+        layer_idx: int,
+        current: torch.Tensor,
+        history_size: int,
+    ) -> torch.Tensor:
+        history = self.refresh_states[layer_idx]
+        if history is None:
+            history = current.new_zeros(current.size(0), current.size(1), history_size)
+        elif history.size(-1) < history_size:
+            history = F.pad(history, (history_size - history.size(-1), 0))
+        convolution_input = torch.cat((history, current), dim=-1)
+        self.refresh_states[layer_idx] = convolution_input[..., -history_size:]
+        return convolution_input
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        super().reorder_cache(beam_idx)
+        self.refresh_states = [
+            None if state is None else state.index_select(0, beam_idx.to(state.device))
+            for state in self.refresh_states
+        ]
+
+    def batch_repeat_interleave(self, repeats: int):
+        super().batch_repeat_interleave(repeats)
+        self.refresh_states = [
+            None if state is None else state.repeat_interleave(repeats, dim=0)
+            for state in self.refresh_states
+        ]
+
+    def batch_select_indices(self, indices: torch.Tensor):
+        super().batch_select_indices(indices)
+        self.refresh_states = [
+            None if state is None else state.index_select(0, indices.to(state.device))
+            for state in self.refresh_states
+        ]
+
+    def crop(self, max_length: int):
+        current_length = self.get_seq_length()
+        target_length = current_length + max_length if max_length < 0 else max_length
+        if target_length < current_length:
+            raise NotImplementedError(
+                "BananaMind2PicoCache cannot roll back causal refresh state"
+            )
+        super().crop(max_length)
+
+
+class BananaMind2PicoAttention(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.o_proj.BANANAMIND_SCALE_INIT = True
+        self.q_norm = BananaMind2PicoRMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = BananaMind2PicoRMSNorm(self.head_dim, config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cosine: torch.Tensor,
+        sine: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+    ) -> torch.Tensor:
+        batch_size, query_length, _ = hidden_states.shape
+        query = self.q_proj(hidden_states).view(
+            batch_size, query_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key = self.k_proj(hidden_states).view(
+            batch_size, query_length, self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        value = self.v_proj(hidden_states).view(
+            batch_size, query_length, self.num_kv_heads, self.head_dim
+        ).transpose(1, 2)
+        query = self.q_norm(query)
+        key = self.k_norm(key)
+        query, key = _apply_rope(query, key, cosine, sine)
+
+        past_length = 0
+        if past_key_values is not None:
+            past_length = past_key_values.get_seq_length(self.layer_idx)
+            key, value = past_key_values.update(key, value, self.layer_idx)
+
+        key_length = key.size(-2)
+        key = key.repeat_interleave(self.num_kv_groups, dim=1)
+        value = value.repeat_interleave(self.num_kv_groups, dim=1)
+        is_causal = query_length > 1 and past_length == 0 and attention_mask is None
+        sdpa_mask = None
+        if not is_causal and query_length > 1:
+            query_positions = past_length + torch.arange(query_length, device=query.device)
+            key_positions = torch.arange(key_length, device=query.device)
+            sdpa_mask = (key_positions[None, :] <= query_positions[:, None])[None, None]
+        if attention_mask is not None:
+            key_padding = attention_mask.to(torch.bool)
+            if key_padding.size(-1) < key_length:
+                key_padding = F.pad(
+                    key_padding,
+                    (key_length - key_padding.size(-1), 0),
+                    value=True,
+                )
+            else:
+                key_padding = key_padding[:, -key_length:]
+            key_padding = key_padding[:, None, None, :]
+            sdpa_mask = key_padding if sdpa_mask is None else sdpa_mask & key_padding
+            is_causal = False
+
+        output = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=sdpa_mask,
+            is_causal=is_causal,
+        )
+        output = output.transpose(1, 2).contiguous().view(
+            batch_size, query_length, self.num_heads * self.head_dim
+        )
+        return self.o_proj(output)
+
+
+class BananaMind2PicoRefreshGate(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        hidden = config.hidden_size
+        self.layer_idx = layer_idx
+        self.kernel_size = config.refresh_kernel_size
+        self.attention_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.embedding_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.output_norm = BananaMind2PicoRMSNorm(hidden, config.rms_norm_eps)
+        self.gate_proj = nn.Linear(hidden, hidden, bias=False)
+        self.value_proj = nn.Linear(hidden, hidden, bias=False)
+        self.out_proj = nn.Linear(hidden, hidden, bias=False)
+        self.out_proj.BANANAMIND_SCALE_INIT = True
+        # A real 2D parameter keeps the depthwise kernel eligible for stock Muon.
+        self.depthwise_kernel = nn.Parameter(torch.empty(hidden, self.kernel_size))
+        nn.init.normal_(self.depthwise_kernel, mean=0.0, std=config.initializer_range)
+        self.alpha = nn.Parameter(torch.tensor(float(config.refresh_alpha_init)))
+
+    def _causal_depthwise_conv(
+        self,
+        attention_signal: torch.Tensor,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        signal = attention_signal.transpose(1, 2)
+        history_size = self.kernel_size - 1
+        if past_key_values is None:
+            convolution_input = F.pad(signal, (history_size, 0))
+        else:
+            convolution_input = past_key_values.refresh_input(
+                self.layer_idx,
+                signal,
+                history_size,
+            )
+        convolved = F.conv1d(
+            convolution_input,
+            self.depthwise_kernel.unsqueeze(1),
+            groups=signal.size(1),
+        )
+        return convolved.transpose(1, 2)
+
+    def forward(
+        self,
+        attention_output: torch.Tensor,
+        original_embedding: torch.Tensor,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        attention_signal = self.attention_norm(attention_output.detach())
+        embedding_value = self.embedding_norm(original_embedding)
+        gate = self.gate_proj(attention_signal) + self._causal_depthwise_conv(
+            attention_signal,
+            past_key_values,
+        )
+        value = self.value_proj(embedding_value)
+        refreshed = self.output_norm(self.out_proj(F.silu(gate) * value))
+        return self.alpha * refreshed
+
+
+class BananaMind2PicoMLP(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+        self.down_proj.BANANAMIND_SCALE_INIT = True
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(
+            F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
+        )
+
+
+class BananaMind2PicoBlock(nn.Module):
+    def __init__(self, config: BananaMind2PicoConfig, layer_idx: int):
+        super().__init__()
+        self.input_norm = BananaMind2PicoRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.attention = BananaMind2PicoAttention(config, layer_idx)
+        self.refresh = (
+            BananaMind2PicoRefreshGate(config, layer_idx)
+            if layer_idx in config.refresh_layers
+            else None
+        )
+        self.post_attention_norm = BananaMind2PicoRMSNorm(
+            config.hidden_size,
+            config.rms_norm_eps,
+        )
+        self.mlp = BananaMind2PicoMLP(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        original_embedding: torch.Tensor,
+        cosine: torch.Tensor,
+        sine: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        past_key_values: BananaMind2PicoCache | None,
+    ) -> torch.Tensor:
+        attention_output = self.attention(
+            self.input_norm(hidden_states),
+            cosine,
+            sine,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+        )
+        hidden_states = hidden_states + attention_output
+        if self.refresh is not None:
+            hidden_states = hidden_states + self.refresh(
+                attention_output,
+                original_embedding,
+                past_key_values,
+            )
+        return hidden_states + self.mlp(self.post_attention_norm(hidden_states))
+
+
+class BananaMind2PicoPreTrainedModel(PreTrainedModel):
+    config_class = BananaMind2PicoConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["BananaMind2PicoBlock"]
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module: nn.Module):
+        std = self.config.initializer_range
+        if hasattr(module, "BANANAMIND_SCALE_INIT"):
+            std *= (2 * self.config.num_hidden_layers) ** -0.5
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=std)
+
+
+class BananaMind2PicoForCausalLM(BananaMind2PicoPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        return False
+
+    def __init__(self, config: BananaMind2PicoConfig):
+        super().__init__(config)
+        self.transformer = nn.ModuleDict(
+            {
+                "wte": nn.Embedding(config.vocab_size, config.hidden_size),
+                "h": nn.ModuleList(
+                    [
+                        BananaMind2PicoBlock(config, index)
+                        for index in range(config.num_hidden_layers)
+                    ]
+                ),
+                "ln_f": BananaMind2PicoRMSNorm(
+                    config.hidden_size,
+                    config.rms_norm_eps,
+                ),
+            }
+        )
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.embedding_scale = math.sqrt(config.hidden_size)
+        self.post_init()
+        if config.tie_word_embeddings:
+            self.tie_weights()
+
+    def get_input_embeddings(self):
+        return self.transformer["wte"]
+
+    def set_input_embeddings(self, value):
+        self.transformer["wte"] = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        if use_cache is None:
+            use_cache = self.config.use_cache and labels is None
+        if use_cache and past_key_values is None:
+            past_key_values = BananaMind2PicoCache(self.config)
+        if use_cache and not isinstance(past_key_values, BananaMind2PicoCache):
+            raise TypeError("BananaMind 2 Pico requires its cache for refresh state")
+        if not use_cache:
+            past_key_values = None
+
+        past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        sequence_length = input_ids.size(1)
+        total_length = past_length + sequence_length
+        if total_length > self.config.max_position_embeddings:
+            raise ValueError(
+                f"Sequence length {total_length} exceeds "
+                f"{self.config.max_position_embeddings}"
+            )
+
+        original_embedding = self.transformer["wte"](input_ids) * self.embedding_scale
+        hidden_states = original_embedding
+        positions = torch.arange(
+            past_length,
+            total_length,
+            dtype=torch.float32,
+            device=input_ids.device,
+        )
+        cosine, sine = _rope_cos_sin(
+            self.config.head_dim,
+            positions,
+            self.config.rope_theta,
+        )
+        for block in self.transformer["h"]:
+            hidden_states = block(
+                hidden_states,
+                original_embedding,
+                cosine,
+                sine,
+                attention_mask,
+                past_key_values,
+            )
+        hidden_states = self.transformer["ln_f"](hidden_states)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits[..., :-1, :].float().reshape(-1, logits.size(-1)),
+                labels[..., 1:].reshape(-1),
+                ignore_index=-100,
+            )
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values,
+        )
+
+
+BananaMind2PicoForCausalLM.register_for_auto_class("AutoModelForCausalLM")
+
+__all__ = [
+    "BananaMind2PicoCache",
+    "BananaMind2PicoForCausalLM",
+    "BananaMind2PicoPreTrainedModel",
+]

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "pad_token": "<|pad|>",
+  "bos_token": "<|bos|>",
+  "eos_token": "<|eos|>",
+  "unk_token": "<|unk|>"
+}

tokenizer.json

@@ -0,0 +1,991 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "<|pad|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "<|bos|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "<|eos|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "<|unk|>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": {
+    "type": "NFKC"
+  },
+  "pre_tokenizer": {
+    "type": "Sequence",
+    "pretokenizers": [
+      {
+        "type": "Split",
+        "pattern": {
+          "Regex": "\\d"
+        },
+        "behavior": "Isolated",
+        "invert": false
+      },
+      {
+        "type": "ByteLevel",
+        "add_prefix_space": false,
+        "trim_offsets": true,
+        "use_regex": true
+      }
+    ]
+  },
+  "post_processor": {
+    "type": "TemplateProcessing",
+    "single": [
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      }
+    ],
+    "pair": [
+      {
+        "Sequence": {
+          "id": "A",
+          "type_id": 0
+        }
+      },
+      {
+        "Sequence": {
+          "id": "B",
+          "type_id": 1
+        }
+      }
+    ],
+    "special_tokens": {}
+  },
+  "decoder": {
+    "type": "ByteLevel",
+    "add_prefix_space": true,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "<|unk|>",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "ignore_merges": false,
+    "vocab": {
+      "<|pad|>": 0,
+      "<|bos|>": 1,
+      "<|eos|>": 2,
+      "<|unk|>": 3,
+      "!": 4,
+      "\"": 5,
+      "#": 6,
+      "$": 7,
+      "%": 8,
+      "&": 9,
+      "'": 10,
+      "(": 11,
+      ")": 12,
+      "*": 13,
+      "+": 14,
+      ",": 15,
+      "-": 16,
+      ".": 17,
+      "/": 18,
+      "0": 19,
+      "1": 20,
+      "2": 21,
+      "3": 22,
+      "4": 23,
+      "5": 24,
+      "6": 25,
+      "7": 26,
+      "8": 27,
+      "9": 28,
+      ":": 29,
+      ";": 30,
+      "<": 31,
+      "=": 32,
+      ">": 33,
+      "?": 34,
+      "@": 35,
+      "A": 36,
+      "B": 37,
+      "C": 38,
+      "D": 39,
+      "E": 40,
+      "F": 41,
+      "G": 42,
+      "H": 43,
+      "I": 44,
+      "J": 45,
+      "K": 46,
+      "L": 47,
+      "M": 48,
+      "N": 49,
+      "O": 50,
+      "P": 51,
+      "Q": 52,
+      "R": 53,
+      "S": 54,
+      "T": 55,
+      "U": 56,
+      "V": 57,
+      "W": 58,
+      "X": 59,
+      "Y": 60,
+      "Z": 61,
+      "[": 62,
+      "\\": 63,
+      "]": 64,
+      "^": 65,
+      "_": 66,
+      "`": 67,
+      "a": 68,
+      "b": 69,
+      "c": 70,
+      "d": 71,
+      "e": 72,
+      "f": 73,
+      "g": 74,
+      "h": 75,
+      "i": 76,
+      "j": 77,
+      "k": 78,
+      "l": 79,
+      "m": 80,
+      "n": 81,
+      "o": 82,
+      "p": 83,
+      "q": 84,
+      "r": 85,
+      "s": 86,
+      "t": 87,
+      "u": 88,
+      "v": 89,
+      "w": 90,
+      "x": 91,
+      "y": 92,
+      "z": 93,
+      "{": 94,
+      "|": 95,
+      "}": 96,
+      "~": 97,
+      "¡": 98,
+      "¢": 99,
+      "£": 100,
+      "¤": 101,
+      "¥": 102,
+      "¦": 103,
+      "§": 104,
+      "¨": 105,
+      "©": 106,
+      "ª": 107,
+      "«": 108,
+      "¬": 109,
+      "®": 110,
+      "¯": 111,
+      "°": 112,
+      "±": 113,
+      "²": 114,
+      "³": 115,
+      "´": 116,
+      "µ": 117,
+      "¶": 118,
+      "·": 119,
+      "¸": 120,
+      "¹": 121,
+      "º": 122,
+      "»": 123,
+      "¼": 124,
+      "½": 125,
+      "¾": 126,
+      "¿": 127,
+      "À": 128,
+      "Á": 129,
+      "Â": 130,
+      "Ã": 131,
+      "Ä": 132,
+      "Å": 133,
+      "Æ": 134,
+      "Ç": 135,
+      "È": 136,
+      "É": 137,
+      "Ê": 138,
+      "Ë": 139,
+      "Ì": 140,
+      "Í": 141,
+      "Î": 142,
+      "Ï": 143,
+      "Ð": 144,
+      "Ñ": 145,
+      "Ò": 146,
+      "Ó": 147,
+      "Ô": 148,
+      "Õ": 149,
+      "Ö": 150,
+      "×": 151,
+      "Ø": 152,
+      "Ù": 153,
+      "Ú": 154,
+      "Û": 155,
+      "Ü": 156,
+      "Ý": 157,
+      "Þ": 158,
+      "ß": 159,
+      "à": 160,
+      "á": 161,
+      "â": 162,
+      "ã": 163,
+      "ä": 164,
+      "å": 165,
+      "æ": 166,
+      "ç": 167,
+      "è": 168,
+      "é": 169,
+      "ê": 170,
+      "ë": 171,
+      "ì": 172,
+      "í": 173,
+      "î": 174,
+      "ï": 175,
+      "ð": 176,
+      "ñ": 177,
+      "ò": 178,
+      "ó": 179,
+      "ô": 180,
+      "õ": 181,
+      "ö": 182,
+      "÷": 183,
+      "ø": 184,
+      "ù": 185,
+      "ú": 186,
+      "û": 187,
+      "ü": 188,
+      "ý": 189,
+      "þ": 190,
+      "ÿ": 191,
+      "Ā": 192,
+      "ā": 193,
+      "Ă": 194,
+      "ă": 195,
+      "Ą": 196,
+      "ą": 197,
+      "Ć": 198,
+      "ć": 199,
+      "Ĉ": 200,
+      "ĉ": 201,
+      "Ċ": 202,
+      "ċ": 203,
+      "Č": 204,
+      "č": 205,
+      "Ď": 206,
+      "ď": 207,
+      "Đ": 208,
+      "đ": 209,
+      "Ē": 210,
+      "ē": 211,
+      "Ĕ": 212,
+      "ĕ": 213,
+      "Ė": 214,
+      "ė": 215,
+      "Ę": 216,
+      "ę": 217,
+      "Ě": 218,
+      "ě": 219,
+      "Ĝ": 220,
+      "ĝ": 221,
+      "Ğ": 222,
+      "ğ": 223,
+      "Ġ": 224,
+      "ġ": 225,
+      "Ģ": 226,
+      "ģ": 227,
+      "Ĥ": 228,
+      "ĥ": 229,
+      "Ħ": 230,
+      "ħ": 231,
+      "Ĩ": 232,
+      "ĩ": 233,
+      "Ī": 234,
+      "ī": 235,
+      "Ĭ": 236,
+      "ĭ": 237,
+      "Į": 238,
+      "į": 239,
+      "İ": 240,
+      "ı": 241,
+      "IJ": 242,
+      "ij": 243,
+      "Ĵ": 244,
+      "ĵ": 245,
+      "Ķ": 246,
+      "ķ": 247,
+      "ĸ": 248,
+      "Ĺ": 249,
+      "ĺ": 250,
+      "Ļ": 251,
+      "ļ": 252,
+      "Ľ": 253,
+      "ľ": 254,
+      "Ŀ": 255,
+      "ŀ": 256,
+      "Ł": 257,
+      "ł": 258,
+      "Ń": 259,
+      "Ġt": 260,
+      "Ġa": 261,
+      "in": 262,
+      "he": 263,
+      "re": 264,
+      "on": 265,
+      "Ġthe": 266,
+      "er": 267,
+      "Ġs": 268,
+      "at": 269,
+      "Ġo": 270,
+      "en": 271,
+      "Ġc": 272,
+      "Ġw": 273,
+      "es": 274,
+      "nd": 275,
+      "it": 276,
+      "ing": 277,
+      "is": 278,
+      "or": 279,
+      "ou": 280,
+      "al": 281,
+      "Ġp": 282,
+      "ed": 283,
+      "Ġf": 284,
+      "an": 285,
+      "Ġb": 286,
+      "Ġin": 287,
+      "ar": 288,
+      "Ġof": 289,
+      "Ġm": 290,
+      "Ġand": 291,
+      "Ġto": 292,
+      "Ġd": 293,
+      "ion": 294,
+      "ic": 295,
+      "le": 296,
+      "om": 297,
+      "Ġh": 298,
+      "as": 299,
+      "Ġth": 300,
+      "ent": 301,
+      "Ġe": 302,
+      "Ġre": 303,
+      "ct": 304,
+      "il": 305,
+      "ro": 306,
+      "Ġl": 307,
+      "st": 308,
+      "ve": 309,
+      "Ġn": 310,
+      "ly": 311,
+      "et": 312,
+      "Ġy": 313,
+      "Ġg": 314,
+      "id": 315,
+      "im": 316,
+      "ĠT": 317,
+      "se": 318,
+      "ation": 319,
+      "Ġbe": 320,
+      "ol": 321,
+      "ce": 322,
+      "ur": 323,
+      "ut": 324,
+      "ow": 325,
+      "ra": 326,
+      "ot": 327,
+      "ig": 328,
+      "Ġyou": 329,
+      "Ġis": 330,
+      "Ġu": 331,
+      "ĠS": 332,
+      "ĠA": 333,
+      "ĠI": 334,
+      "Ġfor": 335,
+      "ch": 336,
+      "ver": 337,
+      "Ġon": 338,
+      "Ġst": 339,
+      "am": 340,
+      "ay": 341,
+      "ĠC": 342,
+      "Ġthat": 343,
+      "ir": 344,
+      "ith": 345,
+      "el": 346,
+      "Ġhe": 347,
+      "ul": 348,
+      "ad": 349,
+      "âĢ": 350,
+      "Ġcon": 351,
+      "Ġwith": 352,
+      "Ġas": 353,
+      "ter": 354,
+      "if": 355,
+      "Ġwh": 356,
+      "Ġan": 357,
+      "Ġit": 358,
+      "ĠM": 359,
+      "Ġor": 360,
+      "Ġpro": 361,
+      "Ġal": 362,
+      "her": 363,
+      "ers": 364,
+      "em": 365,
+      "od": 366,
+      "ess": 367,
+      "Ġex": 368,
+      "ate": 369,
+      "ĠB": 370,
+      "ag": 371,
+      "ist": 372,
+      "Ġcom": 373,
+      "ri": 374,
+      "ĠP": 375,
+      "us": 376,
+      "ill": 377,
+      "res": 378,
+      "pp": 379,
+      "ab": 380,
+      "ke": 381,
+      "th": 382,
+      "ĠH": 383
+    },
+    "merges": [
+      [
+        "Ġ",
+        "t"
+      ],
+      [
+        "Ġ",
+        "a"
+      ],
+      [
+        "i",
+        "n"
+      ],
+      [
+        "h",
+        "e"
+      ],
+      [
+        "r",
+        "e"
+      ],
+      [
+        "o",
+        "n"
+      ],
+      [
+        "Ġt",
+        "he"
+      ],
+      [
+        "e",
+        "r"
+      ],
+      [
+        "Ġ",
+        "s"
+      ],
+      [
+        "a",
+        "t"
+      ],
+      [
+        "Ġ",
+        "o"
+      ],
+      [
+        "e",
+        "n"
+      ],
+      [
+        "Ġ",
+        "c"
+      ],
+      [
+        "Ġ",
+        "w"
+      ],
+      [
+        "e",
+        "s"
+      ],
+      [
+        "n",
+        "d"
+      ],
+      [
+        "i",
+        "t"
+      ],
+      [
+        "in",
+        "g"
+      ],
+      [
+        "i",
+        "s"
+      ],
+      [
+        "o",
+        "r"
+      ],
+      [
+        "o",
+        "u"
+      ],
+      [
+        "a",
+        "l"
+      ],
+      [
+        "Ġ",
+        "p"
+      ],
+      [
+        "e",
+        "d"
+      ],
+      [
+        "Ġ",
+        "f"
+      ],
+      [
+        "a",
+        "n"
+      ],
+      [
+        "Ġ",
+        "b"
+      ],
+      [
+        "Ġ",
+        "in"
+      ],
+      [
+        "a",
+        "r"
+      ],
+      [
+        "Ġo",
+        "f"
+      ],
+      [
+        "Ġ",
+        "m"
+      ],
+      [
+        "Ġa",
+        "nd"
+      ],
+      [
+        "Ġt",
+        "o"
+      ],
+      [
+        "Ġ",
+        "d"
+      ],
+      [
+        "i",
+        "on"
+      ],
+      [
+        "i",
+        "c"
+      ],
+      [
+        "l",
+        "e"
+      ],
+      [
+        "o",
+        "m"
+      ],
+      [
+        "Ġ",
+        "h"
+      ],
+      [
+        "a",
+        "s"
+      ],
+      [
+        "Ġt",
+        "h"
+      ],
+      [
+        "en",
+        "t"
+      ],
+      [
+        "Ġ",
+        "e"
+      ],
+      [
+        "Ġ",
+        "re"
+      ],
+      [
+        "c",
+        "t"
+      ],
+      [
+        "i",
+        "l"
+      ],
+      [
+        "r",
+        "o"
+      ],
+      [
+        "Ġ",
+        "l"
+      ],
+      [
+        "s",
+        "t"
+      ],
+      [
+        "v",
+        "e"
+      ],
+      [
+        "Ġ",
+        "n"
+      ],
+      [
+        "l",
+        "y"
+      ],
+      [
+        "e",
+        "t"
+      ],
+      [
+        "Ġ",
+        "y"
+      ],
+      [
+        "Ġ",
+        "g"
+      ],
+      [
+        "i",
+        "d"
+      ],
+      [
+        "i",
+        "m"
+      ],
+      [
+        "Ġ",
+        "T"
+      ],
+      [
+        "s",
+        "e"
+      ],
+      [
+        "at",
+        "ion"
+      ],
+      [
+        "Ġb",
+        "e"
+      ],
+      [
+        "o",
+        "l"
+      ],
+      [
+        "c",
+        "e"
+      ],
+      [
+        "u",
+        "r"
+      ],
+      [
+        "u",
+        "t"
+      ],
+      [
+        "o",
+        "w"
+      ],
+      [
+        "r",
+        "a"
+      ],
+      [
+        "o",
+        "t"
+      ],
+      [
+        "i",
+        "g"
+      ],
+      [
+        "Ġy",
+        "ou"
+      ],
+      [
+        "Ġ",
+        "is"
+      ],
+      [
+        "Ġ",
+        "u"
+      ],
+      [
+        "Ġ",
+        "S"
+      ],
+      [
+        "Ġ",
+        "A"
+      ],
+      [
+        "Ġ",
+        "I"
+      ],
+      [
+        "Ġf",
+        "or"
+      ],
+      [
+        "c",
+        "h"
+      ],
+      [
+        "v",
+        "er"
+      ],
+      [
+        "Ġ",
+        "on"
+      ],
+      [
+        "Ġs",
+        "t"
+      ],
+      [
+        "a",
+        "m"
+      ],
+      [
+        "a",
+        "y"
+      ],
+      [
+        "Ġ",
+        "C"
+      ],
+      [
+        "Ġth",
+        "at"
+      ],
+      [
+        "i",
+        "r"
+      ],
+      [
+        "it",
+        "h"
+      ],
+      [
+        "e",
+        "l"
+      ],
+      [
+        "Ġ",
+        "he"
+      ],
+      [
+        "u",
+        "l"
+      ],
+      [
+        "a",
+        "d"
+      ],
+      [
+        "â",
+        "Ģ"
+      ],
+      [
+        "Ġc",
+        "on"
+      ],
+      [
+        "Ġw",
+        "ith"
+      ],
+      [
+        "Ġa",
+        "s"
+      ],
+      [
+        "t",
+        "er"
+      ],
+      [
+        "i",
+        "f"
+      ],
+      [
+        "Ġw",
+        "h"
+      ],
+      [
+        "Ġa",
+        "n"
+      ],
+      [
+        "Ġ",
+        "it"
+      ],
+      [
+        "Ġ",
+        "M"
+      ],
+      [
+        "Ġo",
+        "r"
+      ],
+      [
+        "Ġp",
+        "ro"
+      ],
+      [
+        "Ġa",
+        "l"
+      ],
+      [
+        "he",
+        "r"
+      ],
+      [
+        "er",
+        "s"
+      ],
+      [
+        "e",
+        "m"
+      ],
+      [
+        "o",
+        "d"
+      ],
+      [
+        "es",
+        "s"
+      ],
+      [
+        "Ġe",
+        "x"
+      ],
+      [
+        "at",
+        "e"
+      ],
+      [
+        "Ġ",
+        "B"
+      ],
+      [
+        "a",
+        "g"
+      ],
+      [
+        "is",
+        "t"
+      ],
+      [
+        "Ġc",
+        "om"
+      ],
+      [
+        "r",
+        "i"
+      ],
+      [
+        "Ġ",
+        "P"
+      ],
+      [
+        "u",
+        "s"
+      ],
+      [
+        "il",
+        "l"
+      ],
+      [
+        "re",
+        "s"
+      ],
+      [
+        "p",
+        "p"
+      ],
+      [
+        "a",
+        "b"
+      ],
+      [
+        "k",
+        "e"
+      ],
+      [
+        "t",
+        "h"
+      ],
+      [
+        "Ġ",
+        "H"
+      ]
+    ]
+  }
+}

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "model_max_length": 4096,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}

tokenizer_training_manifest.json

@@ -0,0 +1,29 @@
+{
+  "format": "bananamind2_pico_tokenizer_v1",
+  "stage": "tokenizer_complete",
+  "vocab_size": 384,
+  "model_max_length": 8192,
+  "target_gib": 1.0,
+  "sources": [
+    {
+      "name": "fineweb_hq",
+      "dataset_id": "epfml/FineWeb-HQ",
+      "config": null,
+      "text_field": "text",
+      "byte_share": 0.75,
+      "revision": "e58199cdd52438d94405df1a4d8630cc5f13bf84",
+      "target_bytes": 805306368,
+      "consumed_bytes": 805307468
+    },
+    {
+      "name": "cosmopedia_v2",
+      "dataset_id": "HuggingFaceTB/smollm-corpus",
+      "config": "cosmopedia-v2",
+      "text_field": "text",
+      "byte_share": 0.25,
+      "revision": "3ba9d605774198c5868892d7a8deda78031a781f",
+      "target_bytes": 268435456,
+      "consumed_bytes": 268438475
+    }
+  ]
+}

training_metrics.jsonl

@@ -0,0 +1,113 @@
+{"step": 1, "total_steps": 22439, "tokens": 8912896, "source": "fineweb_hq", "phase": 1, "loss": 5.967588186264038, "ce_loss": 5.96400511264801, "logit_z_loss": 35.83174705505371, "muon_lr": 7.000000000000001e-05, "adamw_lr": 4e-06, "grad_norm": 0.2956637740135193, "tokens_per_second": 189609.30176880484, "eta_seconds": 1054734.9659662247, "data_wait_seconds": 14.140621960163116}
+{"step": 10, "total_steps": 22439, "tokens": 89128960, "source": "fineweb_hq", "phase": 1, "loss": 5.960334777832031, "ce_loss": 5.956752300262451, "logit_z_loss": 35.8258056640625, "muon_lr": 0.0007000000000000001, "adamw_lr": 4e-05, "grad_norm": 0.31783637404441833, "tokens_per_second": 5420500.847961092, "eta_seconds": 36879.86590006616, "data_wait_seconds": 1.4603137969970703e-05}
+{"step": 20, "total_steps": 22439, "tokens": 178257920, "source": "fineweb_hq", "phase": 1, "loss": 5.931681752204895, "ce_loss": 5.928101181983948, "logit_z_loss": 35.80507564544678, "muon_lr": 0.0014000000000000002, "adamw_lr": 8e-05, "grad_norm": 0.4683179557323456, "tokens_per_second": 8017463.216447492, "eta_seconds": 24922.872737860678, "data_wait_seconds": 2.5212764739990234e-05}
+{"step": 30, "total_steps": 22439, "tokens": 267386880, "source": "fineweb_hq", "phase": 1, "loss": 5.882287383079529, "ce_loss": 5.878711581230164, "logit_z_loss": 35.75691223144531, "muon_lr": 0.0021000000000000003, "adamw_lr": 0.00011999999999999999, "grad_norm": 0.8383325934410095, "tokens_per_second": 9117651.124056717, "eta_seconds": 21905.761006474495, "data_wait_seconds": 1.71661376953125e-05}
+{"step": 40, "total_steps": 22439, "tokens": 356515840, "source": "fineweb_hq", "phase": 1, "loss": 5.780414581298828, "ce_loss": 5.776850700378418, "logit_z_loss": 35.63748264312744, "muon_lr": 0.0028000000000000004, "adamw_lr": 0.00016, "grad_norm": 1.4044978618621826, "tokens_per_second": 6684341.910831861, "eta_seconds": 29866.808156609535, "data_wait_seconds": 0.3183640241622925}
+{"step": 50, "total_steps": 22439, "tokens": 445644800, "source": "fineweb_hq", "phase": 1, "loss": 5.59889280796051, "ce_loss": 5.595354795455933, "logit_z_loss": 35.38116645812988, "muon_lr": 0.0035000000000000005, "adamw_lr": 0.0002, "grad_norm": 1.7537622451782227, "tokens_per_second": 5099766.386212357, "eta_seconds": 39129.40582602024, "data_wait_seconds": 2.2590160369873047e-05}
+{"step": 60, "total_steps": 22439, "tokens": 534773760, "source": "fineweb_hq", "phase": 1, "loss": 5.355036377906799, "ce_loss": 5.351527690887451, "logit_z_loss": 35.08798122406006, "muon_lr": 0.004200000000000001, "adamw_lr": 0.00023999999999999998, "grad_norm": 1.1905395984649658, "tokens_per_second": 2275357.7469289554, "eta_seconds": 87661.68742177486, "data_wait_seconds": 0.6237474083900452}
+{"step": 70, "total_steps": 22439, "tokens": 623902720, "source": "fineweb_hq", "phase": 1, "loss": 5.146511554718018, "ce_loss": 5.143023490905762, "logit_z_loss": 34.881529808044434, "muon_lr": 0.004900000000000001, "adamw_lr": 0.00028000000000000003, "grad_norm": 0.7038048505783081, "tokens_per_second": 8753024.844715588, "eta_seconds": 22777.562518215178, "data_wait_seconds": 1.5974044799804688e-05}
+{"step": 80, "total_steps": 22439, "tokens": 713031680, "source": "fineweb_hq", "phase": 1, "loss": 4.965500593185425, "ce_loss": 4.962033152580261, "logit_z_loss": 34.67446708679199, "muon_lr": 0.005600000000000001, "adamw_lr": 0.00032, "grad_norm": 0.5216459631919861, "tokens_per_second": 8269934.859357731, "eta_seconds": 24097.341158437728, "data_wait_seconds": 2.2351741790771484e-05}
+{"step": 90, "total_steps": 22439, "tokens": 802160640, "source": "fineweb_hq", "phase": 1, "loss": 4.804619193077087, "ce_loss": 4.801155209541321, "logit_z_loss": 34.64008712768555, "muon_lr": 0.006300000000000001, "adamw_lr": 0.00035999999999999997, "grad_norm": 0.42248982191085815, "tokens_per_second": 5011402.028131011, "eta_seconds": 39748.22047519684, "data_wait_seconds": 1.7404556274414062e-05}
+{"step": 100, "total_steps": 22439, "tokens": 891289600, "source": "fineweb_hq", "phase": 1, "loss": 4.666997671127319, "ce_loss": 4.663511872291565, "logit_z_loss": 34.85922431945801, "muon_lr": 0.007000000000000001, "adamw_lr": 0.0004, "grad_norm": 0.38025227189064026, "tokens_per_second": 964689.3882402931, "eta_seconds": 206393.0485512972, "data_wait_seconds": 8.485594153404236}
+{"step": 110, "total_steps": 22439, "tokens": 980418560, "source": "fineweb_hq", "phase": 1, "loss": 4.543195128440857, "ce_loss": 4.539679646492004, "logit_z_loss": 35.1539421081543, "muon_lr": 0.007700000000000001, "adamw_lr": 0.00044, "grad_norm": 0.3504329323768616, "tokens_per_second": 2999739.504100792, "eta_seconds": 66344.44574668408, "data_wait_seconds": 1.913309097290039e-05}
+{"step": 120, "total_steps": 22439, "tokens": 1069547520, "source": "fineweb_hq", "phase": 1, "loss": 4.42356538772583, "ce_loss": 4.4200040102005005, "logit_z_loss": 35.61503314971924, "muon_lr": 0.008400000000000001, "adamw_lr": 0.00047999999999999996, "grad_norm": 0.3508388102054596, "tokens_per_second": 2185220.094304317, "eta_seconds": 91032.90160222055, "data_wait_seconds": 2.489545226097107}
+{"step": 130, "total_steps": 22439, "tokens": 1158676480, "source": "fineweb_hq", "phase": 1, "loss": 4.309021949768066, "ce_loss": 4.305400371551514, "logit_z_loss": 36.215521812438965, "muon_lr": 0.009100000000000002, "adamw_lr": 0.0005200000000000001, "grad_norm": 0.35900041460990906, "tokens_per_second": 2407092.5167736993, "eta_seconds": 82604.96656377315, "data_wait_seconds": 1.722574234008789e-05}
+{"step": 140, "total_steps": 22439, "tokens": 1247805440, "source": "fineweb_hq", "phase": 1, "loss": 4.20254111289978, "ce_loss": 4.198848009109497, "logit_z_loss": 36.930569648742676, "muon_lr": 0.009800000000000001, "adamw_lr": 0.0005600000000000001, "grad_norm": 0.5181800723075867, "tokens_per_second": 2475828.9904053914, "eta_seconds": 80275.60412056446, "data_wait_seconds": 1.6987323760986328e-05}
+{"step": 150, "total_steps": 22439, "tokens": 1336934400, "source": "fineweb_hq", "phase": 1, "loss": 4.0957911014556885, "ce_loss": 4.092004299163818, "logit_z_loss": 37.86782932281494, "muon_lr": 0.010500000000000002, "adamw_lr": 0.0006, "grad_norm": 0.560067355632782, "tokens_per_second": 7663503.481971913, "eta_seconds": 25922.8092492342, "data_wait_seconds": 0.0004285573959350586}
+{"step": 160, "total_steps": 22439, "tokens": 1426063360, "source": "fineweb_hq", "phase": 1, "loss": 3.9586216807365417, "ce_loss": 3.954736649990082, "logit_z_loss": 38.85059070587158, "muon_lr": 0.011200000000000002, "adamw_lr": 0.00064, "grad_norm": 0.6855910420417786, "tokens_per_second": 7677246.520032939, "eta_seconds": 25864.795335912706, "data_wait_seconds": 0.29149770736694336}
+{"step": 170, "total_steps": 22439, "tokens": 1515192320, "source": "fineweb_hq", "phase": 1, "loss": 3.8495513200759888, "ce_loss": 3.845552682876587, "logit_z_loss": 39.98680782318115, "muon_lr": 0.0119, "adamw_lr": 0.00068, "grad_norm": 0.7350673079490662, "tokens_per_second": 9416576.269628365, "eta_seconds": 21077.860502672196, "data_wait_seconds": 0.2970969080924988}
+{"step": 180, "total_steps": 22439, "tokens": 1604321280, "source": "fineweb_hq", "phase": 1, "loss": 3.7499501705169678, "ce_loss": 3.7458695769309998, "logit_z_loss": 40.8062801361084, "muon_lr": 0.012600000000000002, "adamw_lr": 0.0007199999999999999, "grad_norm": 0.8747629523277283, "tokens_per_second": 9286722.133615695, "eta_seconds": 21362.98999900818, "data_wait_seconds": 1.4424324035644531e-05}
+{"step": 190, "total_steps": 22439, "tokens": 1693450240, "source": "fineweb_hq", "phase": 1, "loss": 3.6615224480628967, "ce_loss": 3.657317876815796, "logit_z_loss": 42.04585933685303, "muon_lr": 0.013300000000000001, "adamw_lr": 0.00076, "grad_norm": 0.9143570065498352, "tokens_per_second": 9076120.159146408, "eta_seconds": 21848.875910282135, "data_wait_seconds": 1.6987323760986328e-05}
+{"step": 200, "total_steps": 22439, "tokens": 1782579200, "source": "fineweb_hq", "phase": 1, "loss": 3.5665494799613953, "ce_loss": 3.562268853187561, "logit_z_loss": 42.80591297149658, "muon_lr": 0.014000000000000002, "adamw_lr": 0.0008, "grad_norm": 0.852935254573822, "tokens_per_second": 10821530.650832502, "eta_seconds": 18316.622716283797, "data_wait_seconds": 2.4199485778808594e-05}
+{"step": 210, "total_steps": 22439, "tokens": 1871708160, "source": "fineweb_hq", "phase": 1, "loss": 3.4870752096176147, "ce_loss": 3.482684314250946, "logit_z_loss": 43.90891075134277, "muon_lr": 0.014700000000000001, "adamw_lr": 0.0008399999999999999, "grad_norm": 1.3818645477294922, "tokens_per_second": 8978432.522564368, "eta_seconds": 22066.743241214754, "data_wait_seconds": 1.4841556549072266e-05}
+{"step": 220, "total_steps": 22439, "tokens": 1960837120, "source": "fineweb_hq", "phase": 1, "loss": 3.3964741230010986, "ce_loss": 3.392018675804138, "logit_z_loss": 44.55475902557373, "muon_lr": 0.015400000000000002, "adamw_lr": 0.00088, "grad_norm": 1.1072789430618286, "tokens_per_second": 11189584.677343598, "eta_seconds": 17698.211500644684, "data_wait_seconds": 1.4662742614746094e-05}
+{"step": 230, "total_steps": 22439, "tokens": 2049966080, "source": "fineweb_hq", "phase": 1, "loss": 3.4376906752586365, "ce_loss": 3.4331058263778687, "logit_z_loss": 45.848862648010254, "muon_lr": 0.0161, "adamw_lr": 0.00092, "grad_norm": 1.1426341533660889, "tokens_per_second": 7749173.003144349, "eta_seconds": 25544.210612368584, "data_wait_seconds": 0.3090166449546814}
+{"step": 240, "total_steps": 22439, "tokens": 2139095040, "source": "fineweb_hq", "phase": 1, "loss": 3.430124819278717, "ce_loss": 3.425500750541687, "logit_z_loss": 46.24059581756592, "muon_lr": 0.016800000000000002, "adamw_lr": 0.0009599999999999999, "grad_norm": 0.7869250178337097, "tokens_per_second": 9606595.123747688, "eta_seconds": 20595.99428884983, "data_wait_seconds": 1.519918441772461e-05}
+{"step": 250, "total_steps": 22439, "tokens": 2228224000, "source": "fineweb_hq", "phase": 1, "loss": 3.3795212507247925, "ce_loss": 3.3748202323913574, "logit_z_loss": 47.00990009307861, "muon_lr": 0.0175, "adamw_lr": 0.001, "grad_norm": 0.7784124612808228, "tokens_per_second": 8801138.248823438, "eta_seconds": 22470.75818521976, "data_wait_seconds": 1.615285873413086e-05}
+{"step": 260, "total_steps": 22439, "tokens": 2317352960, "source": "fineweb_hq", "phase": 1, "loss": 3.32125061750412, "ce_loss": 3.316445231437683, "logit_z_loss": 48.05374717712402, "muon_lr": 0.018200000000000004, "adamw_lr": 0.0010400000000000001, "grad_norm": 0.7271197438240051, "tokens_per_second": 9157215.01433851, "eta_seconds": 21587.253337883947, "data_wait_seconds": 1.817941665649414e-05}
+{"step": 270, "total_steps": 22439, "tokens": 2406481920, "source": "fineweb_hq", "phase": 1, "loss": 3.29115754365921, "ce_loss": 3.286272943019867, "logit_z_loss": 48.84625053405762, "muon_lr": 0.018900000000000004, "adamw_lr": 0.00108, "grad_norm": 0.786897599697113, "tokens_per_second": 10203044.899523327, "eta_seconds": 19365.786720514297, "data_wait_seconds": 1.4066696166992188e-05}
+{"step": 280, "total_steps": 22439, "tokens": 2495610880, "source": "fineweb_hq", "phase": 1, "loss": 3.2658641934394836, "ce_loss": 3.260937750339508, "logit_z_loss": 49.26448059082031, "muon_lr": 0.019600000000000003, "adamw_lr": 0.0011200000000000001, "grad_norm": 0.6971092820167542, "tokens_per_second": 8094953.658517224, "eta_seconds": 24398.02261946201, "data_wait_seconds": 0.2794172763824463}
+{"step": 290, "total_steps": 22439, "tokens": 2584739840, "source": "fineweb_hq", "phase": 1, "loss": 3.262732744216919, "ce_loss": 3.257690668106079, "logit_z_loss": 50.42053508758545, "muon_lr": 0.020300000000000002, "adamw_lr": 0.00116, "grad_norm": 0.6038771867752075, "tokens_per_second": 10462282.390195543, "eta_seconds": 18868.897449088097, "data_wait_seconds": 1.4126300811767578e-05}
+{"step": 300, "total_steps": 22439, "tokens": 2673868800, "source": "fineweb_hq", "phase": 1, "loss": 3.2190935611724854, "ce_loss": 3.213988423347473, "logit_z_loss": 51.05208206176758, "muon_lr": 0.021000000000000005, "adamw_lr": 0.0012, "grad_norm": 0.5001488327980042, "tokens_per_second": 8691702.404163677, "eta_seconds": 22702.411491847037, "data_wait_seconds": 0.20256388187408447}
+{"step": 310, "total_steps": 22439, "tokens": 2762997760, "source": "fineweb_hq", "phase": 1, "loss": 3.212398886680603, "ce_loss": 3.2072192430496216, "logit_z_loss": 51.79640483856201, "muon_lr": 0.021700000000000004, "adamw_lr": 0.00124, "grad_norm": 0.5490106344223022, "tokens_per_second": 6087938.281466302, "eta_seconds": 32397.417067193983, "data_wait_seconds": 1.6748905181884766e-05}
+{"step": 320, "total_steps": 22439, "tokens": 2852126720, "source": "fineweb_hq", "phase": 1, "loss": 3.16773384809494, "ce_loss": 3.1624574065208435, "logit_z_loss": 52.76471710205078, "muon_lr": 0.022400000000000003, "adamw_lr": 0.00128, "grad_norm": 0.4295356869697571, "tokens_per_second": 8704049.034871731, "eta_seconds": 22649.728400444983, "data_wait_seconds": 1.627206802368164e-05}
+{"step": 330, "total_steps": 22439, "tokens": 2941255680, "source": "fineweb_hq", "phase": 1, "loss": 3.1672710180282593, "ce_loss": 3.161931037902832, "logit_z_loss": 53.40028476715088, "muon_lr": 0.023100000000000002, "adamw_lr": 0.00132, "grad_norm": 0.48047566413879395, "tokens_per_second": 3734990.4432077846, "eta_seconds": 52759.22941713333, "data_wait_seconds": 0.5583145022392273}
+{"step": 340, "total_steps": 22439, "tokens": 3030384640, "source": "fineweb_hq", "phase": 1, "loss": 3.1743375062942505, "ce_loss": 3.1689356565475464, "logit_z_loss": 54.01854991912842, "muon_lr": 0.0238, "adamw_lr": 0.00136, "grad_norm": 0.45766308903694153, "tokens_per_second": 7305521.797553259, "eta_seconds": 26961.262201690675, "data_wait_seconds": 0.2721228003501892}
+{"step": 350, "total_steps": 22439, "tokens": 3119513600, "source": "fineweb_hq", "phase": 1, "loss": 3.1624141335487366, "ce_loss": 3.1569427251815796, "logit_z_loss": 54.713884353637695, "muon_lr": 0.024500000000000004, "adamw_lr": 0.0014000000000000002, "grad_norm": 0.3848016858100891, "tokens_per_second": 7100108.619969086, "eta_seconds": 27728.72504940033, "data_wait_seconds": 1.245737075805664e-05}
+{"step": 360, "total_steps": 22439, "tokens": 3208642560, "source": "fineweb_hq", "phase": 1, "loss": 3.1361418962478638, "ce_loss": 3.1305649280548096, "logit_z_loss": 55.769742012023926, "muon_lr": 0.025200000000000004, "adamw_lr": 0.0014399999999999999, "grad_norm": 0.5501356720924377, "tokens_per_second": 5588996.4216903215, "eta_seconds": 35209.86880941391, "data_wait_seconds": 2.110004425048828e-05}
+{"step": 370, "total_steps": 22439, "tokens": 3297771520, "source": "fineweb_hq", "phase": 1, "loss": 3.0984156727790833, "ce_loss": 3.092725694179535, "logit_z_loss": 56.90030860900879, "muon_lr": 0.025900000000000003, "adamw_lr": 0.00148, "grad_norm": 0.5306575894355774, "tokens_per_second": 5312108.901753597, "eta_seconds": 37028.36396276951, "data_wait_seconds": 0.31781405210494995}
+{"step": 380, "total_steps": 22439, "tokens": 3386900480, "source": "fineweb_hq", "phase": 1, "loss": 3.0493186712265015, "ce_loss": 3.043471932411194, "logit_z_loss": 58.46726131439209, "muon_lr": 0.026600000000000002, "adamw_lr": 0.00152, "grad_norm": 0.3818303048610687, "tokens_per_second": 6136950.29388512, "eta_seconds": 32037.015691637996, "data_wait_seconds": 0.43073809146881104}
+{"step": 390, "total_steps": 22439, "tokens": 3476029440, "source": "fineweb_hq", "phase": 1, "loss": 2.9980337619781494, "ce_loss": 2.9921072721481323, "logit_z_loss": 59.264946937561035, "muon_lr": 0.027300000000000005, "adamw_lr": 0.00156, "grad_norm": 0.4464535713195801, "tokens_per_second": 6760350.809793864, "eta_seconds": 29069.563020205496, "data_wait_seconds": 0.42292672395706177}
+{"step": 400, "total_steps": 22439, "tokens": 3565158400, "source": "fineweb_hq", "phase": 1, "loss": 2.925785481929779, "ce_loss": 2.9197596311569214, "logit_z_loss": 60.25872230529785, "muon_lr": 0.028000000000000004, "adamw_lr": 0.0016, "grad_norm": 0.42003414034843445, "tokens_per_second": 4391267.891676532, "eta_seconds": 44732.254963612555, "data_wait_seconds": 0.6731599569320679}
+{"step": 410, "total_steps": 22439, "tokens": 3654287360, "source": "fineweb_hq", "phase": 1, "loss": 2.86074960231781, "ce_loss": 2.8545563220977783, "logit_z_loss": 61.93314456939697, "muon_lr": 0.028700000000000003, "adamw_lr": 0.0016400000000000002, "grad_norm": 0.410231351852417, "tokens_per_second": 4237398.700246017, "eta_seconds": 46335.54684684277, "data_wait_seconds": 1.6510486602783203e-05}
+{"step": 420, "total_steps": 22439, "tokens": 3743416320, "source": "fineweb_hq", "phase": 1, "loss": 2.8084540367126465, "ce_loss": 2.8021368384361267, "logit_z_loss": 63.17254829406738, "muon_lr": 0.029400000000000003, "adamw_lr": 0.0016799999999999999, "grad_norm": 0.3880893290042877, "tokens_per_second": 6775218.231631465, "eta_seconds": 28966.307846403124, "data_wait_seconds": 0.3164432644844055}
+{"step": 430, "total_steps": 22439, "tokens": 3832545280, "source": "fineweb_hq", "phase": 1, "loss": 2.7710108160972595, "ce_loss": 2.7645490169525146, "logit_z_loss": 64.61753463745117, "muon_lr": 0.030100000000000002, "adamw_lr": 0.00172, "grad_norm": 0.3648725748062134, "tokens_per_second": 9626711.265961591, "eta_seconds": 20377.044937205315, "data_wait_seconds": 0.3005521297454834}
+{"step": 440, "total_steps": 22439, "tokens": 3921674240, "source": "fineweb_hq", "phase": 1, "loss": 2.729283392429352, "ce_loss": 2.7226916551589966, "logit_z_loss": 65.91742897033691, "muon_lr": 0.030800000000000004, "adamw_lr": 0.00176, "grad_norm": 0.34309566020965576, "tokens_per_second": 9205855.61983161, "eta_seconds": 21298.92181684971, "data_wait_seconds": 1.5854835510253906e-05}
+{"step": 450, "total_steps": 22439, "tokens": 4010803200, "source": "fineweb_hq", "phase": 1, "loss": 2.684748411178589, "ce_loss": 2.6779785752296448, "logit_z_loss": 67.69777297973633, "muon_lr": 0.0315, "adamw_lr": 0.0018, "grad_norm": 0.3034548759460449, "tokens_per_second": 8553461.507916095, "eta_seconds": 22913.02415549755, "data_wait_seconds": 0.7202497720718384}
+{"step": 460, "total_steps": 22439, "tokens": 4099932160, "source": "fineweb_hq", "phase": 1, "loss": 2.5577516555786133, "ce_loss": 2.550634264945984, "logit_z_loss": 71.17383003234863, "muon_lr": 0.0322, "adamw_lr": 0.00184, "grad_norm": 0.28681713342666626, "tokens_per_second": 6204814.153961051, "eta_seconds": 31571.701637339593, "data_wait_seconds": 1.8358230590820312e-05}
+{"step": 470, "total_steps": 22439, "tokens": 4189061120, "source": "fineweb_hq", "phase": 1, "loss": 2.453768730163574, "ce_loss": 2.446514129638672, "logit_z_loss": 72.54606437683105, "muon_lr": 0.032900000000000006, "adamw_lr": 0.0018800000000000002, "grad_norm": 0.354420006275177, "tokens_per_second": 5931623.626615326, "eta_seconds": 33010.761395144465, "data_wait_seconds": 1.7344951629638672e-05}
+{"step": 480, "total_steps": 22439, "tokens": 4278190080, "source": "fineweb_hq", "phase": 1, "loss": 2.4016764760017395, "ce_loss": 2.3941532969474792, "logit_z_loss": 75.23196601867676, "muon_lr": 0.033600000000000005, "adamw_lr": 0.0019199999999999998, "grad_norm": 0.2954200804233551, "tokens_per_second": 6421854.307163578, "eta_seconds": 30476.911169672014, "data_wait_seconds": 0.452781617641449}
+{"step": 490, "total_steps": 22439, "tokens": 4367319040, "source": "fineweb_hq", "phase": 1, "loss": 2.387361526489258, "ce_loss": 2.379659116268158, "logit_z_loss": 77.02351951599121, "muon_lr": 0.034300000000000004, "adamw_lr": 0.00196, "grad_norm": 0.26878464221954346, "tokens_per_second": 8874608.668837259, "eta_seconds": 22043.693598675727, "data_wait_seconds": 3.3855438232421875e-05}
+{"step": 500, "total_steps": 22439, "tokens": 4456448000, "source": "fineweb_hq", "phase": 1, "loss": 2.351953089237213, "ce_loss": 2.344254434108734, "logit_z_loss": 76.98688316345215, "muon_lr": 0.035, "adamw_lr": 0.002, "grad_norm": 0.32271578907966614, "tokens_per_second": 9141830.602539133, "eta_seconds": 21389.59184932709, "data_wait_seconds": 0.00019431114196777344}
+{"step": 510, "total_steps": 22439, "tokens": 4545576960, "source": "fineweb_hq", "phase": 1, "loss": 2.370932877063751, "ce_loss": 2.3630858063697815, "logit_z_loss": 78.47021675109863, "muon_lr": 0.0357, "adamw_lr": 0.00204, "grad_norm": 0.33357399702072144, "tokens_per_second": 8800208.620950352, "eta_seconds": 22209.802608394624, "data_wait_seconds": 1.3589859008789062e-05}
+{"step": 520, "total_steps": 22439, "tokens": 4634705920, "source": "fineweb_hq", "phase": 1, "loss": 2.459707796573639, "ce_loss": 2.4521210193634033, "logit_z_loss": 75.8684310913086, "muon_lr": 0.03640000000000001, "adamw_lr": 0.0020800000000000003, "grad_norm": 0.4615185558795929, "tokens_per_second": 7705173.841703852, "eta_seconds": 25354.62163963318, "data_wait_seconds": 1.436471939086914e-05}
+{"step": 530, "total_steps": 22439, "tokens": 4723834880, "source": "fineweb_hq", "phase": 1, "loss": 2.5353170037269592, "ce_loss": 2.52760773897171, "logit_z_loss": 77.09314155578613, "muon_lr": 0.0371, "adamw_lr": 0.00212, "grad_norm": 0.2756792902946472, "tokens_per_second": 9547387.503115516, "eta_seconds": 20452.991815853118, "data_wait_seconds": 0.23102569580078125}
+{"step": 540, "total_steps": 22439, "tokens": 4812963840, "source": "fineweb_hq", "phase": 1, "loss": 2.521784484386444, "ce_loss": 2.5141079425811768, "logit_z_loss": 76.76511573791504, "muon_lr": 0.03780000000000001, "adamw_lr": 0.00216, "grad_norm": 0.18886592984199524, "tokens_per_second": 7271273.446975077, "eta_seconds": 26843.09852013588, "data_wait_seconds": 0.3465467691421509}
+{"step": 550, "total_steps": 22439, "tokens": 4902092800, "source": "fineweb_hq", "phase": 1, "loss": 2.5104334354400635, "ce_loss": 2.502769708633423, "logit_z_loss": 76.63683700561523, "muon_lr": 0.038500000000000006, "adamw_lr": 0.0022, "grad_norm": 0.2813076376914978, "tokens_per_second": 7425982.480560611, "eta_seconds": 26271.861138200762, "data_wait_seconds": 0.1284111738204956}
+{"step": 560, "total_steps": 22439, "tokens": 4991221760, "source": "fineweb_hq", "phase": 1, "loss": 2.485075056552887, "ce_loss": 2.4772602319717407, "logit_z_loss": 78.14760780334473, "muon_lr": 0.039200000000000006, "adamw_lr": 0.0022400000000000002, "grad_norm": 0.22887028753757477, "tokens_per_second": 6909855.981844515, "eta_seconds": 28221.319242596626, "data_wait_seconds": 0.46238821744918823}
+{"step": 570, "total_steps": 22439, "tokens": 5080350720, "source": "fineweb_hq", "phase": 1, "loss": 2.4911267161369324, "ce_loss": 2.483340859413147, "logit_z_loss": 77.85906410217285, "muon_lr": 0.039900000000000005, "adamw_lr": 0.0022800000000000003, "grad_norm": 0.227179154753685, "tokens_per_second": 7787497.403233405, "eta_seconds": 25029.365986442564, "data_wait_seconds": 1.3172626495361328e-05}
+{"step": 580, "total_steps": 22439, "tokens": 5169479680, "source": "fineweb_hq", "phase": 1, "loss": 2.466982662677765, "ce_loss": 2.4591829776763916, "logit_z_loss": 77.99713134765625, "muon_lr": 0.040600000000000004, "adamw_lr": 0.00232, "grad_norm": 0.2317526489496231, "tokens_per_second": 7487358.288002979, "eta_seconds": 26020.79213654995, "data_wait_seconds": 1.3768672943115234e-05}
+{"step": 590, "total_steps": 22439, "tokens": 5258608640, "source": "fineweb_hq", "phase": 1, "loss": 2.471463978290558, "ce_loss": 2.4635881781578064, "logit_z_loss": 78.75795936584473, "muon_lr": 0.0413, "adamw_lr": 0.0023599999999999997, "grad_norm": 0.2254444658756256, "tokens_per_second": 8296458.800079717, "eta_seconds": 23472.407854557037, "data_wait_seconds": 1.5795230865478516e-05}
+{"step": 600, "total_steps": 22439, "tokens": 5347737600, "source": "fineweb_hq", "phase": 1, "loss": 2.4517582058906555, "ce_loss": 2.443916440010071, "logit_z_loss": 78.41762161254883, "muon_lr": 0.04200000000000001, "adamw_lr": 0.0024, "grad_norm": 0.19819648563861847, "tokens_per_second": 6818300.184785781, "eta_seconds": 28547.985636997222, "data_wait_seconds": 0.26518261432647705}
+{"step": 610, "total_steps": 22439, "tokens": 5436866560, "source": "fineweb_hq", "phase": 1, "loss": 2.4580911993980408, "ce_loss": 2.4501355290412903, "logit_z_loss": 79.55681991577148, "muon_lr": 0.0427, "adamw_lr": 0.00244, "grad_norm": 0.23898088932037354, "tokens_per_second": 7027915.462642427, "eta_seconds": 27683.82855744362, "data_wait_seconds": 0.28087741136550903}
+{"step": 620, "total_steps": 22439, "tokens": 5525995520, "source": "fineweb_hq", "phase": 1, "loss": 2.4332305788993835, "ce_loss": 2.4252651929855347, "logit_z_loss": 79.6538200378418, "muon_lr": 0.04340000000000001, "adamw_lr": 0.00248, "grad_norm": 0.2116815149784088, "tokens_per_second": 6712336.335789734, "eta_seconds": 28972.09974224567, "data_wait_seconds": 0.34881824254989624}
+{"step": 630, "total_steps": 22439, "tokens": 5615124480, "source": "fineweb_hq", "phase": 1, "loss": 2.3960554599761963, "ce_loss": 2.3879873156547546, "logit_z_loss": 80.68179702758789, "muon_lr": 0.0441, "adamw_lr": 0.00252, "grad_norm": 0.1983681470155716, "tokens_per_second": 9203690.473776149, "eta_seconds": 21119.935466957093, "data_wait_seconds": 1.3589859008789062e-05}
+{"step": 640, "total_steps": 22439, "tokens": 5704253440, "source": "fineweb_hq", "phase": 1, "loss": 2.3983024954795837, "ce_loss": 2.390074670314789, "logit_z_loss": 82.27853965759277, "muon_lr": 0.044800000000000006, "adamw_lr": 0.00256, "grad_norm": 0.1874842494726181, "tokens_per_second": 7166536.72879459, "eta_seconds": 27111.03385870457, "data_wait_seconds": 0.7625116109848022}
+{"step": 650, "total_steps": 22439, "tokens": 5793382400, "source": "fineweb_hq", "phase": 1, "loss": 2.396180212497711, "ce_loss": 2.3880707025527954, "logit_z_loss": 81.0950698852539, "muon_lr": 0.045500000000000006, "adamw_lr": 0.0026, "grad_norm": 0.20300427079200745, "tokens_per_second": 9557785.073881257, "eta_seconds": 20318.838459205628, "data_wait_seconds": 1.2983060479164124}
+{"step": 660, "total_steps": 22439, "tokens": 5882511360, "source": "fineweb_hq", "phase": 1, "loss": 2.3795944452285767, "ce_loss": 2.3716355562210083, "logit_z_loss": 79.58875274658203, "muon_lr": 0.046200000000000005, "adamw_lr": 0.00264, "grad_norm": 0.20361796021461487, "tokens_per_second": 7950054.236303745, "eta_seconds": 24416.683989095687, "data_wait_seconds": 1.5437602996826172e-05}
+{"step": 670, "total_steps": 22439, "tokens": 5971640320, "source": "fineweb_hq", "phase": 1, "loss": 2.3759934306144714, "ce_loss": 2.367850661277771, "logit_z_loss": 81.42828941345215, "muon_lr": 0.046900000000000004, "adamw_lr": 0.00268, "grad_norm": 0.18410010635852814, "tokens_per_second": 8104376.304865796, "eta_seconds": 23940.748272943496, "data_wait_seconds": 1.6987323760986328e-05}
+{"step": 680, "total_steps": 22439, "tokens": 6060769280, "source": "fineweb_hq", "phase": 1, "loss": 2.370010197162628, "ce_loss": 2.3618654012680054, "logit_z_loss": 81.44764518737793, "muon_lr": 0.0476, "adamw_lr": 0.00272, "grad_norm": 0.23240381479263306, "tokens_per_second": 6333958.580314357, "eta_seconds": 30618.404210400582, "data_wait_seconds": 0.3833252787590027}
+{"step": 690, "total_steps": 22439, "tokens": 6149898240, "source": "fineweb_hq", "phase": 1, "loss": 2.393221378326416, "ce_loss": 2.385291337966919, "logit_z_loss": 79.30001449584961, "muon_lr": 0.0483, "adamw_lr": 0.0027600000000000003, "grad_norm": 0.2510574460029602, "tokens_per_second": 7164714.090987016, "eta_seconds": 27055.73071615696, "data_wait_seconds": 1.8715858459472656e-05}
+{"step": 700, "total_steps": 22439, "tokens": 6239027200, "source": "fineweb_hq", "phase": 1, "loss": 2.328338325023651, "ce_loss": 2.320348381996155, "logit_z_loss": 79.8990478515625, "muon_lr": 0.04900000000000001, "adamw_lr": 0.0028000000000000004, "grad_norm": 0.1862829178571701, "tokens_per_second": 7604268.968707431, "eta_seconds": 25480.08847942352, "data_wait_seconds": 1.3887882232666016e-05}
+{"step": 710, "total_steps": 22439, "tokens": 6328156160, "source": "fineweb_hq", "phase": 1, "loss": 2.3147504925727844, "ce_loss": 2.306523859500885, "logit_z_loss": 82.26664924621582, "muon_lr": 0.0497, "adamw_lr": 0.0028399999999999996, "grad_norm": 0.2060781866312027, "tokens_per_second": 8560163.437189277, "eta_seconds": 22624.371439290047, "data_wait_seconds": 0.3879215717315674}
+{"step": 720, "total_steps": 22439, "tokens": 6417285120, "source": "fineweb_hq", "phase": 1, "loss": 2.332974135875702, "ce_loss": 2.3249324560165405, "logit_z_loss": 80.41722297668457, "muon_lr": 0.05040000000000001, "adamw_lr": 0.0028799999999999997, "grad_norm": 0.23572801053524017, "tokens_per_second": 10912151.829293167, "eta_seconds": 17739.781415462494, "data_wait_seconds": 1.5556812286376953e-05}
+{"step": 730, "total_steps": 22439, "tokens": 6506414080, "source": "fineweb_hq", "phase": 1, "loss": 2.330069839954376, "ce_loss": 2.321860134601593, "logit_z_loss": 82.09696006774902, "muon_lr": 0.0511, "adamw_lr": 0.00292, "grad_norm": 0.26685255765914917, "tokens_per_second": 9020825.559421727, "eta_seconds": 21449.262929368022, "data_wait_seconds": 1.6987323760986328e-05}
+{"step": 740, "total_steps": 22439, "tokens": 6595543040, "source": "fineweb_hq", "phase": 1, "loss": 2.299080550670624, "ce_loss": 2.290955424308777, "logit_z_loss": 81.25156021118164, "muon_lr": 0.051800000000000006, "adamw_lr": 0.00296, "grad_norm": 0.20033298432826996, "tokens_per_second": 9869975.372023484, "eta_seconds": 19594.87465918064, "data_wait_seconds": 1.3828277587890625e-05}
+{"step": 750, "total_steps": 22439, "tokens": 6684672000, "source": "fineweb_hq", "phase": 1, "loss": 2.305118501186371, "ce_loss": 2.297184705734253, "logit_z_loss": 79.33795356750488, "muon_lr": 0.052500000000000005, "adamw_lr": 0.003, "grad_norm": 0.23620666563510895, "tokens_per_second": 7993148.869520959, "eta_seconds": 24184.686723542214, "data_wait_seconds": 0.653160035610199}
+{"step": 760, "total_steps": 22439, "tokens": 6773800960, "source": "fineweb_hq", "phase": 1, "loss": 2.3622615337371826, "ce_loss": 2.354522705078125, "logit_z_loss": 77.3880786895752, "muon_lr": 0.053200000000000004, "adamw_lr": 0.00304, "grad_norm": 0.19219592213630676, "tokens_per_second": 8174759.375875979, "eta_seconds": 23636.496623277664, "data_wait_seconds": 0.33434557914733887}
+{"step": 770, "total_steps": 22439, "tokens": 6862929920, "source": "fineweb_hq", "phase": 1, "loss": 2.383073151111603, "ce_loss": 2.3752957582473755, "logit_z_loss": 77.77359580993652, "muon_lr": 0.0539, "adamw_lr": 0.0030800000000000003, "grad_norm": 0.22470791637897491, "tokens_per_second": 7493282.985046385, "eta_seconds": 25774.22256832123, "data_wait_seconds": 1.4543533325195312e-05}
+{"step": 780, "total_steps": 22439, "tokens": 6952058880, "source": "fineweb_hq", "phase": 1, "loss": 2.3486892580986023, "ce_loss": 2.340909242630005, "logit_z_loss": 77.80061531066895, "muon_lr": 0.05460000000000001, "adamw_lr": 0.00312, "grad_norm": 0.19442570209503174, "tokens_per_second": 7974212.17532066, "eta_seconds": 24208.587659788132, "data_wait_seconds": 1.6629695892333984e-05}
+{"step": 790, "total_steps": 22439, "tokens": 7041187840, "source": "fineweb_hq", "phase": 1, "loss": 2.3511235117912292, "ce_loss": 2.3434258103370667, "logit_z_loss": 76.97773361206055, "muon_lr": 0.0553, "adamw_lr": 0.00316, "grad_norm": 0.18331126868724823, "tokens_per_second": 8197510.245438312, "eta_seconds": 23538.279273438453, "data_wait_seconds": 1.627206802368164e-05}
+{"step": 800, "total_steps": 22439, "tokens": 7130316800, "source": "fineweb_hq", "phase": 1, "loss": 2.343844771385193, "ce_loss": 2.3361929059028625, "logit_z_loss": 76.51930809020996, "muon_lr": 0.05600000000000001, "adamw_lr": 0.0032, "grad_norm": 0.1974499672651291, "tokens_per_second": 7005134.072868889, "eta_seconds": 27532.114951372147, "data_wait_seconds": 0.3746710419654846}
+{"step": 810, "total_steps": 22439, "tokens": 7219445760, "source": "fineweb_hq", "phase": 1, "loss": 2.2092705368995667, "ce_loss": 2.201336443424225, "logit_z_loss": 79.34107208251953, "muon_lr": 0.0567, "adamw_lr": 0.0032400000000000003, "grad_norm": 0.1958804875612259, "tokens_per_second": 7538323.863033231, "eta_seconds": 25572.92988290787, "data_wait_seconds": 0.8902894258499146}
+{"step": 820, "total_steps": 22439, "tokens": 7308574720, "source": "fineweb_hq", "phase": 1, "loss": 2.1729026436805725, "ce_loss": 2.16474711894989, "logit_z_loss": 81.55514907836914, "muon_lr": 0.05740000000000001, "adamw_lr": 0.0032800000000000004, "grad_norm": 0.23285837471485138, "tokens_per_second": 9451195.882156769, "eta_seconds": 20387.673795628547, "data_wait_seconds": 0.31566721200942993}
+{"step": 830, "total_steps": 22439, "tokens": 7397703680, "source": "fineweb_hq", "phase": 1, "loss": 2.153770089149475, "ce_loss": 2.145902693271637, "logit_z_loss": 78.67469024658203, "muon_lr": 0.058100000000000006, "adamw_lr": 0.0033200000000000005, "grad_norm": 0.1934303492307663, "tokens_per_second": 9319240.540504752, "eta_seconds": 20666.788117218017, "data_wait_seconds": 1.6450881958007812e-05}
+{"step": 840, "total_steps": 22439, "tokens": 7486832640, "source": "fineweb_hq", "phase": 1, "loss": 2.1568243503570557, "ce_loss": 2.148854672908783, "logit_z_loss": 79.69684410095215, "muon_lr": 0.058800000000000005, "adamw_lr": 0.0033599999999999997, "grad_norm": 0.19662460684776306, "tokens_per_second": 9581271.945798812, "eta_seconds": 20092.284384894374, "data_wait_seconds": 1.3232231140136719e-05}
+{"step": 850, "total_steps": 22439, "tokens": 7575961600, "source": "fineweb_hq", "phase": 1, "loss": 2.152070641517639, "ce_loss": 2.1442551612854004, "logit_z_loss": 78.15550422668457, "muon_lr": 0.059500000000000004, "adamw_lr": 0.0034, "grad_norm": 0.18207232654094696, "tokens_per_second": 9762642.39925466, "eta_seconds": 19709.8801609993, "data_wait_seconds": 1.4126300811767578e-05}
+{"step": 860, "total_steps": 22439, "tokens": 7665090560, "source": "fineweb_hq", "phase": 1, "loss": 2.145963132381439, "ce_loss": 2.137876272201538, "logit_z_loss": 80.86890602111816, "muon_lr": 0.060200000000000004, "adamw_lr": 0.00344, "grad_norm": 0.1876753568649292, "tokens_per_second": 9893073.30349838, "eta_seconds": 19441.014625453947, "data_wait_seconds": 1.537799835205078e-05}
+{"step": 870, "total_steps": 22439, "tokens": 7754219520, "source": "fineweb_hq", "phase": 1, "loss": 2.145168960094452, "ce_loss": 2.137241542339325, "logit_z_loss": 79.27460289001465, "muon_lr": 0.0609, "adamw_lr": 0.00348, "grad_norm": 0.221941277384758, "tokens_per_second": 8483347.529066838, "eta_seconds": 22661.131489109994, "data_wait_seconds": 1.5079975128173828e-05}
+{"step": 880, "total_steps": 22439, "tokens": 7843348480, "source": "fineweb_hq", "phase": 1, "loss": 2.1310052275657654, "ce_loss": 2.123206377029419, "logit_z_loss": 77.9881420135498, "muon_lr": 0.06160000000000001, "adamw_lr": 0.00352, "grad_norm": 0.17409531772136688, "tokens_per_second": 8918937.64257767, "eta_seconds": 21544.396044063567, "data_wait_seconds": 1.329183578491211e-05}
+{"step": 890, "total_steps": 22439, "tokens": 7932477440, "source": "fineweb_hq", "phase": 1, "loss": 2.1497772932052612, "ce_loss": 2.1419413089752197, "logit_z_loss": 78.36023139953613, "muon_lr": 0.0623, "adamw_lr": 0.0035600000000000002, "grad_norm": 0.22990204393863678, "tokens_per_second": 9922804.604194226, "eta_seconds": 19355.817590403556, "data_wait_seconds": 1.5676021575927734e-05}
+{"step": 900, "total_steps": 22439, "tokens": 8021606400, "source": "fineweb_hq", "phase": 1, "loss": 2.128612518310547, "ce_loss": 2.120979964733124, "logit_z_loss": 76.3257942199707, "muon_lr": 0.063, "adamw_lr": 0.0036, "grad_norm": 0.18078908324241638, "tokens_per_second": 9650671.570589593, "eta_seconds": 19892.384228372575, "data_wait_seconds": 0.051418960094451904}
+{"step": 910, "total_steps": 22439, "tokens": 8110735360, "source": "fineweb_hq", "phase": 1, "loss": 2.3088300824165344, "ce_loss": 2.301823318004608, "logit_z_loss": 70.06817817687988, "muon_lr": 0.0637, "adamw_lr": 0.00364, "grad_norm": 0.2104783058166504, "tokens_per_second": 7219887.872922974, "eta_seconds": 26577.384768486023, "data_wait_seconds": 0.29506224393844604}
+{"step": 920, "total_steps": 22439, "tokens": 8199864320, "source": "fineweb_hq", "phase": 1, "loss": 2.3767300248146057, "ce_loss": 2.3701910972595215, "logit_z_loss": 65.3893928527832, "muon_lr": 0.0644, "adamw_lr": 0.00368, "grad_norm": 0.23128600418567657, "tokens_per_second": 9192341.784089683, "eta_seconds": 20864.82569174767, "data_wait_seconds": 1.3172626495361328e-05}
+{"step": 930, "total_steps": 22439, "tokens": 8288993280, "source": "fineweb_hq", "phase": 1, "loss": 2.3693774342536926, "ce_loss": 2.362643301486969, "logit_z_loss": 67.34131240844727, "muon_lr": 0.0651, "adamw_lr": 0.00372, "grad_norm": 0.2287091165781021, "tokens_per_second": 10073121.40196061, "eta_seconds": 19031.586378645898, "data_wait_seconds": 1.4781951904296875e-05}
+{"step": 940, "total_steps": 22439, "tokens": 8378122240, "source": "fineweb_hq", "phase": 1, "loss": 2.3465049266815186, "ce_loss": 2.340116858482361, "logit_z_loss": 63.88089084625244, "muon_lr": 0.06580000000000001, "adamw_lr": 0.0037600000000000003, "grad_norm": 0.22197671234607697, "tokens_per_second": 8754935.55516695, "eta_seconds": 21886.894528985023, "data_wait_seconds": 1.341104507446289e-05}
+{"step": 950, "total_steps": 22439, "tokens": 8467251200, "source": "fineweb_hq", "phase": 1, "loss": 2.33204847574234, "ce_loss": 2.3254379630088806, "logit_z_loss": 66.10550880432129, "muon_lr": 0.0665, "adamw_lr": 0.0038000000000000004, "grad_norm": 0.18201303482055664, "tokens_per_second": 8971161.761053331, "eta_seconds": 21349.433578991888, "data_wait_seconds": 0.00030362606048583984}
+{"step": 960, "total_steps": 22439, "tokens": 8556380160, "source": "fineweb_hq", "phase": 1, "loss": 2.350585162639618, "ce_loss": 2.3441357016563416, "logit_z_loss": 64.49478530883789, "muon_lr": 0.06720000000000001, "adamw_lr": 0.0038399999999999997, "grad_norm": 0.2307051122188568, "tokens_per_second": 8768643.45441727, "eta_seconds": 21832.350029873847, "data_wait_seconds": 1.8477439880371094e-05}
+{"step": 970, "total_steps": 22439, "tokens": 8645509120, "source": "fineweb_hq", "phase": 1, "loss": 2.3355281949043274, "ce_loss": 2.328825294971466, "logit_z_loss": 67.02908897399902, "muon_lr": 0.0679, "adamw_lr": 0.0038799999999999998, "grad_norm": 0.21734148263931274, "tokens_per_second": 7220713.442206703, "eta_seconds": 26500.285014152527, "data_wait_seconds": 0.25003528594970703}
+{"step": 980, "total_steps": 22439, "tokens": 8734638080, "source": "fineweb_hq", "phase": 1, "loss": 2.3283920288085938, "ce_loss": 2.3219318985939026, "logit_z_loss": 64.60088920593262, "muon_lr": 0.06860000000000001, "adamw_lr": 0.00392, "grad_norm": 0.20072227716445923, "tokens_per_second": 9270987.236583356, "eta_seconds": 20630.14761893749, "data_wait_seconds": 1.245737075805664e-05}
+{"step": 990, "total_steps": 22439, "tokens": 8823767040, "source": "fineweb_hq", "phase": 1, "loss": 2.323706567287445, "ce_loss": 2.317197024822235, "logit_z_loss": 65.09517860412598, "muon_lr": 0.06930000000000001, "adamw_lr": 0.00396, "grad_norm": 0.18527741730213165, "tokens_per_second": 9336268.191085013, "eta_seconds": 20476.351192069054, "data_wait_seconds": 1.3947486877441406e-05}
+{"step": 1000, "total_steps": 22439, "tokens": 8912896000, "source": "fineweb_hq", "phase": 1, "loss": 2.3358243703842163, "ce_loss": 2.329273998737335, "logit_z_loss": 65.50436592102051, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.23283950984477997, "tokens_per_second": 9433630.894132858, "eta_seconds": 20255.570679879187, "data_wait_seconds": 1.5676021575927734e-05}
+{"step": 1010, "total_steps": 22439, "tokens": 9002024960, "source": "fineweb_hq", "phase": 1, "loss": 2.315945625305176, "ce_loss": 2.309797704219818, "logit_z_loss": 61.47883129119873, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.20814353227615356, "tokens_per_second": 9852860.27603847, "eta_seconds": 19384.67034273148, "data_wait_seconds": 0.22886061668395996}
+{"step": 1020, "total_steps": 22439, "tokens": 9091153920, "source": "fineweb_hq", "phase": 1, "loss": 2.3088624477386475, "ce_loss": 2.30258309841156, "logit_z_loss": 62.792927742004395, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.21581503748893738, "tokens_per_second": 8651330.612411598, "eta_seconds": 22066.5846650362, "data_wait_seconds": 1.806020736694336e-05}
+{"step": 1030, "total_steps": 22439, "tokens": 9180282880, "source": "fineweb_hq", "phase": 1, "loss": 2.2938343286514282, "ce_loss": 2.2875046133995056, "logit_z_loss": 63.2970085144043, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.2105294018983841, "tokens_per_second": 9340256.124235211, "eta_seconds": 20429.438756918906, "data_wait_seconds": 0.26373594999313354}
+{"step": 1040, "total_steps": 22439, "tokens": 9269411840, "source": "fineweb_hq", "phase": 1, "loss": 2.291613817214966, "ce_loss": 2.2854092717170715, "logit_z_loss": 62.04541778564453, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.21177281439304352, "tokens_per_second": 9291618.76773903, "eta_seconds": 20526.78508143425, "data_wait_seconds": 1.3887882232666016e-05}
+{"step": 1050, "total_steps": 22439, "tokens": 9358540800, "source": "fineweb_hq", "phase": 1, "loss": 2.2960025668144226, "ce_loss": 2.289924383163452, "logit_z_loss": 60.78238868713379, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.1940109133720398, "tokens_per_second": 9229805.947100338, "eta_seconds": 20654.59811794758, "data_wait_seconds": 1.6927719116210938e-05}
+{"step": 1060, "total_steps": 22439, "tokens": 9447669760, "source": "fineweb_hq", "phase": 1, "loss": 2.2542683482170105, "ce_loss": 2.2478742003440857, "logit_z_loss": 63.94139099121094, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.21360020339488983, "tokens_per_second": 7268330.627052976, "eta_seconds": 26216.309268426896, "data_wait_seconds": 1.4007091522216797e-05}
+{"step": 1070, "total_steps": 22439, "tokens": 9536798720, "source": "fineweb_hq", "phase": 1, "loss": 2.226654887199402, "ce_loss": 2.220464825630188, "logit_z_loss": 61.90073108673096, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.18968330323696136, "tokens_per_second": 9982306.738514233, "eta_seconds": 19079.72572002411, "data_wait_seconds": 1.424551010131836e-05}
+{"step": 1080, "total_steps": 22439, "tokens": 9625927680, "source": "fineweb_hq", "phase": 1, "loss": 2.2515273690223694, "ce_loss": 2.2451473474502563, "logit_z_loss": 63.80070877075195, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.2117748260498047, "tokens_per_second": 9098439.298023513, "eta_seconds": 20923.428670382502, "data_wait_seconds": 0.24623018503189087}
+{"step": 1090, "total_steps": 22439, "tokens": 9715056640, "source": "fineweb_hq", "phase": 1, "loss": 2.213638424873352, "ce_loss": 2.207147002220154, "logit_z_loss": 64.91473197937012, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.20987659692764282, "tokens_per_second": 9320515.442111455, "eta_seconds": 20415.331950879096, "data_wait_seconds": 0.277202308177948}
+{"step": 1100, "total_steps": 22439, "tokens": 9804185600, "source": "fineweb_hq", "phase": 1, "loss": 2.1991316080093384, "ce_loss": 2.1925421953201294, "logit_z_loss": 65.89375495910645, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.19404654204845428, "tokens_per_second": 9263205.526926557, "eta_seconds": 20532.01639444828, "data_wait_seconds": 1.4960765838623047e-05}
+{"step": 1110, "total_steps": 22439, "tokens": 9893314560, "source": "fineweb_hq", "phase": 1, "loss": 2.1957609057426453, "ce_loss": 2.1894816160202026, "logit_z_loss": 62.79327964782715, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.18220849335193634, "tokens_per_second": 8321119.599639584, "eta_seconds": 22845.863048553467, "data_wait_seconds": 1.722574234008789e-05}
+{"step": 1120, "total_steps": 22439, "tokens": 9982443520, "source": "fineweb_hq", "phase": 1, "loss": 2.1929144859313965, "ce_loss": 2.1866838335990906, "logit_z_loss": 62.306596755981445, "muon_lr": 0.07, "adamw_lr": 0.004, "grad_norm": 0.16241706907749176, "tokens_per_second": 10415627.523396535, "eta_seconds": 18243.16676044464, "data_wait_seconds": 1.4126300811767578e-05}