Upload checkpoint-15858 (SFT, final step)

config.json

@@ -0,0 +1,63 @@
+{
+  "architectures": [
+    "PragyaForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_pragya.PragyaConfig",
+    "AutoModelForCausalLM": "modeling_pragya.PragyaForCausalLM"
+  },
+  "bias": false,
+  "block_size": 1024,
+  "bos_token_id": 12287,
+  "d_rope": null,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "engram_compressed_vocab": 9807,
+  "engram_embed_per_ngram": 128,
+  "engram_kernel_size": 4,
+  "engram_layer_ids": [
+    1,
+    8
+  ],
+  "engram_max_ngram": 3,
+  "engram_n_heads": 4,
+  "engram_seed": 0,
+  "engram_vocab_size": [
+    4096,
+    4096
+  ],
+  "eos_token_id": 12287,
+  "ffn_mult": 2,
+  "gen_headroom": 128,
+  "gpt_top_k": 64,
+  "gradient_checkpointing": false,
+  "int8_group_size": 128,
+  "max_gen_tokens": 1024,
+  "mod_budget_coeff": 0.05,
+  "mod_capacity": 0.25,
+  "mod_gate_entropy_coeff": 0.0,
+  "mod_hard_eval": false,
+  "model_type": "pragya_e2",
+  "moe_fixed_capacity": null,
+  "moe_top_k": 2,
+  "mtp_depth": 1,
+  "mtp_lambda": 0.1,
+  "n_embd": 384,
+  "n_experts": 8,
+  "n_head": 12,
+  "n_layer": 16,
+  "n_shared": 1,
+  "num_kv_heads": 6,
+  "pad_token_id": 12287,
+  "pattern": "dense",
+  "sliding_window_size": 512,
+  "tp_size": 1,
+  "transformers_version": "5.22.1",
+  "use_cache": false,
+  "use_engram": true,
+  "use_int8": false,
+  "use_moe": false,
+  "use_mtp": false,
+  "use_nvfp4": false,
+  "vocab_size": 12288
+}

config.py

@@ -0,0 +1,121 @@
+
+use_liger=False
+xla_on_gpu=False
+# ---------------------------------------------------------------------------
+# Core shape -- chosen for GPU efficiency, not just "a number that fits"
+# ---------------------------------------------------------------------------
+CONTEXT = 1024
+max_gen_tokens = 1024
+gen_headroom = 128
+vocab_size = 8192+4096
+# ── Engram (static n-gram memory) ──────────────────────────────────
+use_engram              = True
+engram_layer_ids        = (1, 8)         # two layers, one early, one mid
+engram_max_ngram        = 3              # 2-gram + 3-gram tables
+engram_vocab_size       = (4096, 4096)   # halved from default
+engram_embed_per_ngram  = 128
+engram_n_heads          = 4
+engram_kernel_size      = 4
+engram_seed             = 0
+engram_compressed_vocab = 0              # 0 → set from tokenizer lookup
+bias = False
+d_rope = None
+top_k = 64
+moe_top_k = 2
+use_triton_fusions = False
+
+
+gradient_checkpointing = False
+num_kv_heads = 4
+pattern = "dense"
+sliding_window_size = 512
+
+tensor_parallel_size = 1
+
+# ─────────────────────────────────────────────────────────────────────
+# Curriculum learning (token-level mixing across sources)
+# ─────────────────────────────────────────────────────────────────────
+# Master switch. False = old behavior (uniform random over the whole
+# flat token array). True = per-source anchor schedule below.
+CURRICULUM_ENABLED = False
+
+# Number of training steps the anchor schedule spans. Set equal to
+# MAX_STEPS in train.py so the ramp completes at run end.
+CURRICULUM_TOTAL_STEPS = 20000
+
+# Sources, in curriculum order (easiest first). Each entry:
+#   (name, [glob_patterns], w@0, w@T/3, w@2T/3, w@T)
+# Between anchors, weights interpolate linearly. Columns should sum
+# to ~1.0 (the code renormalizes if not).
+# Files matching none of the globs are EXCLUDED from training.
+CURRICULUM_SOURCES = [
+    ("tinystories", ["tinystories*.txt", "tiny_*.txt"],
+     0.70, 0.10, 0.03, 0.02),
+    ("simplewiki",  ["simplewiki*.txt",  "wiki_*.txt"],
+     0.10, 0.70, 0.10, 0.03),
+    ("cosmopedia",  ["cosmopedia*.txt",  "cosmo_*.txt"],
+     0.10, 0.10, 0.70, 0.10),
+    ("finephrase",  ["finephrase*.txt",  "fine_*.txt"],
+     0.10, 0.10, 0.17, 0.85),
+]
+
+# Hard floor: no source ever drops below this fraction at any step.
+# 0.0 disables. 0.02 = 2% minimum, prevents catastrophic forgetting.
+CURRICULUM_FLOOR = 0.02
+
+mod_capacity = 0.25
+use_mtp    = False
+mtp_depth  = 1      
+mtp_lambda = 0.1 
+use_nvfp4 = False
+# config.py
+use_fp8 = False
+fp8_recipe = "rowwise_with_gw_hp"        # default
+fp8_bf16_tail_frac = 0.15                # default
+use_int8 = False
+int8_group_size = 128
+mod_budget_coeff = 0.05
+mod_gate_entropy_coeff = 0.0
+mod_hard_eval = False
+use_flexattention = True
+numberoflayers = 16     
+numberofheads = 12
+D_MODEL = 384        
+ALPHA_COMPILE_MODE="reduce-overhead"  
+num_kv_heads = 6 
+sliding_window_size = 512  
+use_sdpa=False
+optimizer_8bit = True
+use_wsd = True
+wsd_decay_steps = 3500 
+ffn_mult = 8 // 3           
+use_moe = False            
+n_experts = 8
+n_shared = 1
+# T4 path
+
+PRECISION = 'bf16'         
+GLOBAL_DTYPE = PRECISION
+USE_ASYNC_LAYER_PREFETCH = True
+intra_op_threads = 2
+num_workers = 2
+GPU_SIDE_BATCHING = True
+DISABLE_MMAP = False
+
+batch_size = 320
+weight_decay = 0.01        
+use_static_lr = False
+static_lr = 1e-3        
+learning_rate = 1e-3
+
+
+attn_gate_window = 16
+max_steps = 19000
+data_parallel_only = True
+use_fsdp = False
+dropout = 0.0                   
+use_muon = True
+DATA_PATH = './training_data'
+data_path = DATA_PATH
+TOKENIZER_MODEL_PATH = './tokenizer'
+use_streaming_batcher = False

configuration_pragya.py

@@ -0,0 +1,106 @@
+from transformers import PretrainedConfig
+
+
+class PragyaConfig(PretrainedConfig):
+    """Config for Pragya E2. Field names match model.GPTConfig, except gpt_top_k
+    (GPTConfig.top_k), renamed because Transformers reserves top_k for generation.
+
+    Defaults are generated from the trained run's config.py, so PragyaConfig()
+    is the real model. engram_compressed_vocab is 0 by design: the model wrapper
+    reads it from the checkpoint's Engram lookup buffer.
+    """
+
+    model_type = "pragya_e2"
+
+    def __init__(
+        self,
+        block_size=1024,
+        vocab_size=12288,
+        n_layer=16,
+        n_head=12,
+        n_embd=384,
+        dropout=0.0,
+        bias=False,
+        ffn_mult=2,
+        d_rope=None,
+        gpt_top_k=64,
+        n_experts=8,
+        n_shared=1,
+        use_moe=False,
+        moe_top_k=2,
+        gradient_checkpointing=False,
+        num_kv_heads=6,
+        pattern='dense',
+        sliding_window_size=512,
+        tp_size=1,
+        gen_headroom=128,
+        max_gen_tokens=1024,
+        mod_capacity=0.25,
+        use_mtp=False,
+        mtp_depth=1,
+        mtp_lambda=0.1,
+        use_nvfp4=False,
+        use_int8=False,
+        int8_group_size=128,
+        mod_gate_entropy_coeff=0.0,
+        mod_budget_coeff=0.05,
+        mod_hard_eval=False,
+        moe_fixed_capacity=None,
+        use_engram=True,
+        engram_layer_ids=(1, 8),
+        engram_max_ngram=3,
+        engram_vocab_size=(4096, 4096),
+        engram_embed_per_ngram=128,
+        engram_n_heads=4,
+        engram_kernel_size=4,
+        engram_seed=0,
+        engram_compressed_vocab=0,
+        **kwargs,
+    ):
+        # Documents are [BOS ... EOS] and this tokenizer uses one id (12287) for all three.
+        # Without eos_token_id, generate() has no stop token and fills the whole context.
+        kwargs.setdefault("bos_token_id", 12287)
+        kwargs.setdefault("eos_token_id", 12287)
+        kwargs.setdefault("pad_token_id", 12287)
+        super().__init__(**kwargs)
+        self.block_size = block_size
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.dropout = dropout
+        self.bias = bias
+        self.ffn_mult = ffn_mult
+        self.d_rope = d_rope
+        self.gpt_top_k = gpt_top_k
+        self.n_experts = n_experts
+        self.n_shared = n_shared
+        self.use_moe = use_moe
+        self.moe_top_k = moe_top_k
+        self.gradient_checkpointing = gradient_checkpointing
+        self.num_kv_heads = num_kv_heads
+        self.pattern = pattern
+        self.sliding_window_size = sliding_window_size
+        self.tp_size = tp_size
+        self.gen_headroom = gen_headroom
+        self.max_gen_tokens = max_gen_tokens
+        self.mod_capacity = mod_capacity
+        self.use_mtp = use_mtp
+        self.mtp_depth = mtp_depth
+        self.mtp_lambda = mtp_lambda
+        self.use_nvfp4 = use_nvfp4
+        self.use_int8 = use_int8
+        self.int8_group_size = int8_group_size
+        self.mod_gate_entropy_coeff = mod_gate_entropy_coeff
+        self.mod_budget_coeff = mod_budget_coeff
+        self.mod_hard_eval = mod_hard_eval
+        self.moe_fixed_capacity = moe_fixed_capacity
+        self.use_engram = use_engram
+        self.engram_layer_ids = list(engram_layer_ids)
+        self.engram_max_ngram = engram_max_ngram
+        self.engram_vocab_size = list(engram_vocab_size)
+        self.engram_embed_per_ngram = engram_embed_per_ngram
+        self.engram_n_heads = engram_n_heads
+        self.engram_kernel_size = engram_kernel_size
+        self.engram_seed = engram_seed
+        self.engram_compressed_vocab = engram_compressed_vocab

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 12287,
+  "eos_token_id": [
+    12287
+  ],
+  "pad_token_id": 12287,
+  "transformers_version": "5.22.1",
+  "use_mtp": false
+}

kernel.py

@@ -0,0 +1,163 @@
+"""Attention kernel with two backends, chosen once at import.
+
+* FlexAttention (default): CUDA with compute capability >= 8.0 (sm80+), torch >= 2.5.
+  Masks are BlockMasks; the mask logic lives in mask_mod.
+* xFormers memory-efficient attention: CUDA GPUs below sm80 (no Triton/Flex support).
+  Masks are dense boolean (T, T) or (B, T, T) tensors converted to additive biases.
+
+Every attention layer calls `fused_attention(q, k, v, mask)`; the mask object comes from
+`causal_masks` or `doc_masks`, so the backend is invisible to model.py.
+
+Attention-weight dropout is not supported on either backend (DenseMHA refuses dropout > 0).
+"""
+import functools
+import importlib
+import warnings
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import create_block_mask, flex_attention
+
+_IS_CUDA = torch.cuda.is_available() and torch.version.cuda is not None
+_USE_XFORMERS = _IS_CUDA and torch.cuda.get_device_capability() < (8, 0)
+
+_XFORMERS_OK = _USE_XFORMERS  # cleared on the first failure; then SDPA is used
+
+if _USE_XFORMERS:
+    # Loaded dynamically: Transformers scans source files for static imports and would
+    # require xformers on every machine, including ones that never use it.
+    xops = importlib.import_module("xformers.ops")  # pip install xformers (match your torch)
+    _flex = None
+else:
+    xops = None
+    _flex = torch.compile(flex_attention, dynamic=False) if _IS_CUDA else flex_attention
+
+
+class _XMask:
+    """Dense allowed-positions mask for the xFormers path. allowed: bool (T, T) or (B, T, T)."""
+    __slots__ = ("allowed",)
+
+    def __init__(self, allowed: torch.Tensor):
+        self.allowed = allowed
+
+
+def fused_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask) -> torch.Tensor:
+    """q: (B, Hq, T, D); k, v: (B, Hkv, T, D). GQA is handled natively (Hq a multiple of Hkv)."""
+    if isinstance(mask, _XMask):
+        return _xformers_attention(q, k, v, mask.allowed)
+    return _flex(q, k, v, block_mask=mask, enable_gqa=q.shape[1] != k.shape[1])
+
+
+def _xformers_attention(q, k, v, allowed):
+    B, Hq, T, D = q.shape
+    Hkv = k.shape[1]
+    if Hq != Hkv:
+        # query head h uses kv head h // group, the same grouping FlexAttention's enable_gqa uses
+        group = Hq // Hkv
+        k = k.repeat_interleave(group, dim=1)
+        v = v.repeat_interleave(group, dim=1)
+    bias = torch.zeros(allowed.shape, dtype=q.dtype, device=q.device).masked_fill(~allowed, float("-inf"))
+    bias = bias[None, None] if bias.dim() == 2 else bias[:, None]
+    bias = bias.expand(B, Hq, T, T)
+    global _XFORMERS_OK
+    if _XFORMERS_OK:
+        try:
+            # xFormers expects (B, T, H, D)
+            out = xops.memory_efficient_attention(
+                q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), attn_bias=bias)
+            _record_backend("xformers")
+            return out.transpose(1, 2)
+        except Exception as e:
+            # Typical cause: xformers built for a different torch/CUDA than the runtime.
+            _XFORMERS_OK = False
+            warnings.warn(f"xformers memory_efficient_attention unavailable ({type(e).__name__}: {e}); "
+                          "using torch scaled_dot_product_attention with the same mask.")
+    # Same math with PyTorch's built-in attention (works on any CUDA GPU, no extra package).
+    out = F.scaled_dot_product_attention(q, k, v, attn_mask=bias)
+    _record_backend("sdpa_masked")
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Masks. One rule for every layer: window = window_size or seq_len; since q - kv <= seq_len - 1,
+# "no window" is the same expression as "window = seq_len".
+# ---------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=None)
+def _causal_mask(seq_len: int, window_size, device_key: str):
+    window = window_size or seq_len
+
+    def mask_mod(b, h, q_idx, kv_idx):
+        return (kv_idx <= q_idx) & (q_idx - kv_idx < window)
+
+    return create_block_mask(mask_mod, B=None, H=None, Q_LEN=seq_len, KV_LEN=seq_len, device=device_key)
+
+
+@functools.lru_cache(maxsize=None)
+def _xcausal_mask(seq_len: int, window_size, device_key: str):
+    idx = torch.arange(seq_len, device=device_key)
+    q_idx = idx[:, None]
+    kv_idx = idx[None, :]
+    window = window_size or seq_len
+    return _XMask((kv_idx <= q_idx) & (q_idx - kv_idx < window))
+
+
+@torch.compiler.disable
+def causal_masks(seq_len: int, windows, device) -> dict:
+    """{window_size: mask} for plain causal attention; cached per (seq_len, window, device)."""
+    if _USE_XFORMERS:
+        return {w: _xcausal_mask(seq_len, w, str(device)) for w in windows}
+    return {w: _causal_mask(seq_len, w, str(device)) for w in windows}
+
+
+@torch.compiler.disable
+def doc_masks(doc_ids: torch.Tensor, windows) -> dict:
+    """{window_size: mask} for causal + same-document attention. doc_ids: (B, T) int, per batch."""
+    b_, t_ = doc_ids.shape
+
+    if _USE_XFORMERS:
+        idx = torch.arange(t_, device=doc_ids.device)
+        q_idx = idx[:, None]
+        kv_idx = idx[None, :]
+        same_doc = doc_ids[:, :, None] == doc_ids[:, None, :]          # (B, T, T)
+
+        def build_x(window_size):
+            window = window_size or t_
+            return _XMask(same_doc & (kv_idx <= q_idx) & (q_idx - kv_idx < window))
+
+        return {w: build_x(w) for w in windows}
+
+    def build(window_size):
+        window = window_size or t_
+
+        def mask_mod(b, h, q_idx, kv_idx):
+            return (doc_ids[b, q_idx] == doc_ids[b, kv_idx]) & (kv_idx <= q_idx) & (q_idx - kv_idx < window)
+
+        return create_block_mask(mask_mod, B=b_, H=None, Q_LEN=t_, KV_LEN=t_,
+                                 device=doc_ids.device, _compile=True)
+
+    return {w: build(w) for w in windows}
+
+
+# ---------------------------------------------------------------------------
+# Telemetry counters (model.py records its Liger/Triton success/fallback events here).
+# ---------------------------------------------------------------------------
+
+_backend_stats: dict[str, int] = {}
+
+
+def _record_backend(name: str) -> None:
+    # Inside a compiled region this body runs only at trace time, so counting here under-reports.
+    # Record nothing while compiling rather than report wrong numbers.
+    if torch.compiler.is_compiling():
+        return
+    _backend_stats[name] = _backend_stats.get(name, 0) + 1
+
+
+def get_backend_stats() -> dict[str, int]:
+    return dict(_backend_stats)
+
+
+def reset_backend_stats() -> None:
+    _backend_stats.clear()
+    

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7869aa4bc65771b579bd83bcbc2b621268fb0995302a429ce4c6f17f8c4dd9df
+size 114127392

modeling_pragya.py

@@ -0,0 +1,175 @@
+import dataclasses
+
+import torch
+import torch.nn.functional as F
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_pragya import PragyaConfig
+from .model import GPT, GPTConfig, KVCache
+# The Hub loader only copies sibling files named in "from .name import" form. model.py
+# falls back to "from . import config/kernel/ngram", so import them this way too.
+from .config import *  # noqa: F401,F403
+from .kernel import *  # noqa: F401,F403
+from .ngram import *  # noqa: F401,F403
+
+# Config field names that differ between PragyaConfig and GPTConfig.
+_RENAMED = {"top_k": "gpt_top_k"}
+# Fields GPTConfig stores as tuples, which PragyaConfig stores as lists.
+_TUPLE_FIELDS = ("engram_layer_ids", "engram_vocab_size")
+
+
+def to_gpt_config(cfg: PragyaConfig) -> GPTConfig:
+    """Build a GPTConfig from every field, so no value comes from config.py defaults."""
+    kwargs = {}
+    for f in dataclasses.fields(GPTConfig):
+        value = getattr(cfg, _RENAMED.get(f.name, f.name))
+        if f.name in _TUPLE_FIELDS:
+            value = tuple(value)
+        kwargs[f.name] = value
+    return GPTConfig(**kwargs)
+
+
+def _non_persistent_buffers(module):
+    """{name: tensor} for every non-persistent buffer in a torch module tree."""
+    out = {}
+    for prefix, sub in module.named_modules():
+        for buf_name in sub._non_persistent_buffers_set:
+            full = f"{prefix}.{buf_name}" if prefix else buf_name
+            out[full] = getattr(sub, buf_name)
+    return out
+
+
+class PragyaForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = PragyaConfig
+    base_model_prefix = "model"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = False
+    # wte.weight and lm_head.weight are one tensor; save only lm_head.weight.
+    _tied_weights_keys = {"model.transformer.wte.weight": "model.lm_head.weight"}
+
+    def __init__(self, config: PragyaConfig):
+        super().__init__(config)
+        if config.use_engram and config.engram_compressed_vocab <= 0:
+            raise ValueError(
+                "engram_compressed_vocab must be set before loading. It is computed from the "
+                "tokenizer when the model is first built, so the export step must record it."
+            )
+        self.model = GPT(to_gpt_config(config))
+        self.post_init()
+
+    def _move_missing_keys_from_meta_to_device(self, missing_keys, device_map, device_mesh, hf_quantizer):
+        # Transformers replaces every non-persistent buffer (RoPE tables, XSA masks, Engram
+        # `mods`) with uninitialised memory when loading. They are derived from the config,
+        # so rebuild them from a freshly constructed GPT.
+        super()._move_missing_keys_from_meta_to_device(missing_keys, device_map, device_mesh, hf_quantizer)
+        fresh = GPT(to_gpt_config(self.config))
+        own = _non_persistent_buffers(self.model)
+        with torch.no_grad():
+            for name, buf in _non_persistent_buffers(fresh).items():
+                own[name].copy_(buf)
+        del fresh
+        # wte and lm_head are one tensor. Checkpoints store one of the two, so tie from the
+        # one that was actually loaded, not from whichever happens to be first.
+        if "model.lm_head.weight" in missing_keys:
+            self.model.lm_head.weight = self.model.transformer.wte.weight
+        else:
+            self.model.transformer.wte.weight = self.model.lm_head.weight
+
+    def _init_weights(self, module):
+        # Weights come from GPT.__init__ (Engram tables, zero-init attention gate, tied
+        # embeddings) or from the checkpoint. Transformers' default would overwrite them.
+        pass
+
+    def get_input_embeddings(self):
+        return self.model.transformer.wte
+
+    def get_output_embeddings(self):
+        return self.model.lm_head
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None,
+                logits_to_keep=0, past_key_values=None, use_cache=None, **kwargs):
+        if attention_mask is not None and not bool(torch.all(attention_mask == 1)):
+            if labels is not None:
+                raise NotImplementedError(
+                    "attention_mask with padding is not supported for training. Use unpadded batches."
+                )
+            # Inference with padding: the model has no mask support, so each row is run on its own
+            # real tokens (no pads attended to, positions start at 0) and the logits are put back
+            # at the row's original positions. Padded positions get zeros.
+            return self._forward_unpadded_rows(input_ids, attention_mask, logits_to_keep)
+        if labels is not None:
+            # Training path. GPT compares logits[t] with targets[t] without shifting, so shift
+            # here: position t predicts token t+1. Trainer passes labels == input_ids.
+            logits, loss = self.model(input_ids[:, :-1], targets=labels[:, 1:])
+            return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+        cache_wanted = bool(use_cache) or isinstance(past_key_values, KVCache)
+        if cache_wanted and self.model.supports_kv_cache:
+            return self._cached_forward(input_ids, past_key_values)
+
+        # Inference path. logits_to_keep=0 means logits for every position (B, T, V),
+        # which lm-evaluation-harness and loglikelihood scoring need. logits_to_keep=k
+        # returns the last k positions. Hidden states are captured after ln_f, and the
+        # projection is the same F.linear(x, lm_head.weight) GPT uses internally.
+        captured = {}
+        handle = self.model.transformer.ln_f.register_forward_hook(
+            lambda module, inp, out: captured.__setitem__("x", out))
+        try:
+            self.model(input_ids)
+        finally:
+            handle.remove()
+        x = captured["x"]
+        if logits_to_keep:
+            x = x[:, -logits_to_keep:, :]
+        logits = F.linear(x, self.model.lm_head.weight)
+        return CausalLMOutputWithPast(logits=logits)
+
+    def _forward_unpadded_rows(self, input_ids, attention_mask, logits_to_keep):
+        rows = []
+        for b in range(input_ids.size(0)):
+            keep = attention_mask[b].bool()
+            row = input_ids[b:b + 1, keep]
+            lg = self.forward(row).logits                     # (1, n, V), full logits
+            full = lg.new_zeros(1, input_ids.size(1), lg.size(-1))
+            full[:, keep] = lg
+            rows.append(full)
+        logits = torch.cat(rows, dim=0)
+        if logits_to_keep:
+            logits = logits[:, -logits_to_keep:]
+        return CausalLMOutputWithPast(logits=logits)
+
+    def _cached_forward(self, input_ids, past):
+        """Decoding with the model's own KVCache (prefill on the first call, one token per call).
+
+        Returns last-position logits, which is all generate() needs. Past the context window the
+        model's sliding-window behaviour applies, so the cache is dropped and the last block_size
+        tokens are re-encoded, matching GPT._generate_loop.
+        """
+        block = self.config.block_size
+        cache = past if isinstance(past, KVCache) else None
+        if cache is None:
+            cache = self.model.new_kv_cache()
+        if cache.pos + input_ids.size(1) > block:
+            window = input_ids[:, -block:]
+            logits, _ = self.model(window)
+            return CausalLMOutputWithPast(logits=logits, past_key_values=None)
+        logits, _ = self.model(input_ids, kv_cache=cache)
+        return CausalLMOutputWithPast(logits=logits, past_key_values=cache)
+
+    def _supports_default_dynamic_cache(self):
+        # Transformers would build its own DynamicCache, which needs num_hidden_layers and other
+        # standard attributes. Decoding uses the model's KVCache, so opt out.
+        return False
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        block = self.config.block_size
+        cache = past_key_values if isinstance(past_key_values, KVCache) else None
+        if cache is None or not self.model.supports_kv_cache or input_ids.size(1) > block:
+            # No usable cache, or the sequence has passed the window: re-encode the last
+            # block_size tokens, as GPT._generate_loop does.
+            return {"input_ids": input_ids[:, -block:]}
+        # Only tokens the cache has not seen yet: the whole prompt on the first call, then one token.
+        return {"input_ids": input_ids[:, cache.pos:],
+                "past_key_values": cache,
+                "use_cache": True}

ngram.py

@@ -0,0 +1,459 @@
+"""
+Engram -- static n-gram memory (arXiv:2601.07372, DeepSeek engram_demo_v1.py),
+adapted for Veylon/Arya's standard transformer (hc_mult = 1, tiktoken).
+
+Fixes vs. the drafted port (all verified in test_engram.py):
+
+  1. Hashing ran on CPU/numpy INSIDE forward(): a GPU->CPU sync + torch.compile
+     graph break on every Engram layer, every step (and fatal on XLA). Now the
+     hash is pure int64 torch ops on the model device, compile-friendly.
+  2. Every Engram layer built its own NgramHashMapping (re-decoding the whole
+     vocab) and each call hashed ALL layers, then kept one. Now ONE NgramHasher
+     lives on GPT, hashes all layers once per forward.
+  3. pad_id=None crashed np.pad; pad_id=-1 silently indexed the LAST lookup
+     entry. Left-context padding is now a dedicated reserved id (never
+     collides with a real token class).
+  4. nn.RMSNorm (torch>=2.4 only, eps=finfo(fp16)=1e-3 default) replaced by a local
+     _RMSNorm (eps=1e-5, weight cast to the input dtype) matching model.RMSNorm's eager
+     path. It does NOT upcast to fp32.
+  5. Gate dot-product computed in fp32 (sum over D in fp16 can overflow/lose
+     precision), cast back after the sigmoid.
+  6. The compressed-vocab lookup is now a persistent buffer, so a checkpoint
+     carries it and inference needs no tokenizer object.
+  7. sympy dependency dropped (trial-division primality is plenty for table sizes).
+  8. Special-id discovery no longer assumes TokenizerWrapper._special_ids is a
+     set of ints (dict / missing attr both handled).
+  9. embed_per_ngram % n_heads != 0 silently produced a wrong-width table ->
+     hard assert. Same for len(engram_vocab_size) < max_ngram-1.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+import unicodedata
+import warnings
+from typing import List, Sequence
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Primes
+# ─────────────────────────────────────────────────────────────────────────────
+
+def _is_prime(n: int) -> bool:
+    if n < 2:
+        return False
+    if n < 4:
+        return True
+    if n % 2 == 0 or n % 3 == 0:
+        return False
+    i = 5
+    while i * i <= n:
+        if n % i == 0 or n % (i + 2) == 0:
+            return False
+        i += 6
+    return True
+
+
+def find_next_prime(start: int, seen_primes: set) -> int:
+    candidate = start + 1
+    while True:
+        if _is_prime(candidate) and candidate not in seen_primes:
+            return candidate
+        candidate += 1
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# CompressedTokenizer -- build-time only (numpy). Output is a lookup table.
+# ─────────────────────────────────────────────────────────────────────────────
+
+class CompressedTokenizer:
+    """Maps raw token IDs -> compressed IDs (NFKC / strip-accents / lowercase /
+    whitespace-collapse of each token's surface string). Build-time only: the
+    result (`lookup_table`) is handed to the model as a persistent buffer."""
+
+    _SENTINEL = "\uE000"
+    _WS_RE = re.compile(r"[ \t\r\n]+")
+
+    def __init__(self, tokenizer_wrapper):
+        self.tokenizer = tokenizer_wrapper
+        self._special_ids = self._collect_special_ids(tokenizer_wrapper)
+        self.lookup_table, self.num_new_token = self._build_lookup_table()
+
+    def __len__(self):
+        return self.num_new_token
+
+    @staticmethod
+    def _collect_special_ids(tw) -> set:
+        ids = set()
+        raw = getattr(tw, "_special_ids", None)
+        if raw is not None:
+            if isinstance(raw, dict):
+                for k, v in raw.items():
+                    for cand in (k, v):
+                        if isinstance(cand, (int, np.integer)):
+                            ids.add(int(cand))
+            else:
+                for v in raw:
+                    if isinstance(v, (int, np.integer)):
+                        ids.add(int(v))
+        vs = int(tw.vocab_size)
+        for name in ("bos_id", "eos_id", "pad_id"):
+            v = getattr(tw, name, None)
+            if isinstance(v, (int, np.integer)) and 0 <= int(v) < vs:
+                ids.add(int(v))
+        return ids
+
+    @classmethod
+    def _normalize(cls, text: str) -> str:
+        text = unicodedata.normalize("NFKC", text)
+        text = unicodedata.normalize("NFD", text)
+        text = "".join(c for c in text if unicodedata.category(c) != "Mn")
+        text = text.lower()
+        text = cls._WS_RE.sub(" ", text)
+        if text == " ":                      # lone space survives strip()
+            text = cls._SENTINEL
+        text = text.strip()
+        return text.replace(cls._SENTINEL, " ")
+
+    def _decode_id(self, tid: int) -> str:
+        tw = self.tokenizer
+        try:
+            enc = getattr(tw, "enc", None)
+            if enc is not None:
+                return enc.decode([tid])
+            return tw.decode([tid])
+        except Exception:
+            return ""
+
+    def _build_lookup_table(self):
+        key2new, new_tokens = {}, []
+        vocab_size = int(self.tokenizer.vocab_size)
+        lookup = np.empty(vocab_size, dtype=np.int64)
+        for tid in range(vocab_size):
+            if tid in self._special_ids:
+                key = f"__special_{tid}"
+            else:
+                text = self._decode_id(tid)
+                if not text or "\ufffd" in text:
+                    key = f"__raw_{tid}"
+                else:
+                    norm = self._normalize(text)
+                    key = norm if norm else f"__raw_{tid}"
+            nid = key2new.get(key)
+            if nid is None:
+                nid = len(new_tokens)
+                key2new[key] = nid
+                new_tokens.append(key)
+            lookup[tid] = nid
+        return lookup, len(new_tokens)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# NgramHasher -- static, deterministic, layer-specific, pure torch int64
+# ─────────────────────────────────────────────────────────────────────────────
+
+class NgramHasher(nn.Module):
+    """(B, T) raw ids -> per-layer embedding indices of shape (B, T, n_cols).
+
+    No parameters, no gradients. `lookup` and the hash multipliers `mult` are
+    persistent (state_dict); the prime moduli `mods` are re-derived from config.
+    `mult` comes from NumPy's Generator, which is not guaranteed stable across
+    NumPy releases, so it is saved with the checkpoint instead of re-derived.
+    """
+
+    def __init__(self, layer_ids: Sequence[int], max_ngram: int,
+                 vocab_size_per_ngram: Sequence[int], n_heads: int,
+                 compressed_vocab: int, raw_vocab_size: int, seed: int,
+                 lookup: torch.Tensor | None = None):
+        super().__init__()
+        assert max_ngram >= 2, "engram_max_ngram must be >= 2"
+        assert len(vocab_size_per_ngram) >= max_ngram - 1, (
+            f"engram_vocab_size needs {max_ngram - 1} entries (ngram 2..{max_ngram}), "
+            f"got {len(vocab_size_per_ngram)}"
+        )
+        assert compressed_vocab > 0, "engram_compressed_vocab must be set (>0)"
+
+        self.layer_ids = tuple(int(l) for l in layer_ids)
+        self.max_ngram = int(max_ngram)
+        self.n_heads = int(n_heads)
+        self.n_cols = (self.max_ngram - 1) * self.n_heads
+        self.compressed_vocab = int(compressed_vocab)
+        # Reserved id for left-context padding: distinct from every token class.
+        self.pad_cid = self.compressed_vocab
+
+        if lookup is None:
+            # Every token would hash to class 0 and Engram would silently ignore its input. This is only
+            # valid if a checkpoint (which carries `lookup`) is loaded right after construction.
+            warnings.warn(
+                "NgramHasher built without a lookup table: all tokens map to class 0 until a checkpoint "
+                "containing `lookup` is loaded. Pass engram_lookup from CompressedTokenizer for training.",
+                stacklevel=2)
+            lookup = torch.zeros(int(raw_vocab_size), dtype=torch.int64)
+        else:
+            lookup = lookup.detach().to(torch.int64).clone()
+            assert lookup.numel() == int(raw_vocab_size), (
+                f"lookup has {lookup.numel()} entries, vocab_size={raw_vocab_size}")
+            assert int(lookup.max()) < self.compressed_vocab
+        self.register_buffer("lookup", lookup, persistent=True)
+
+        # Multipliers: r*2+1, bounded so tok*mult can never overflow int64.
+        # (+1 vocab slot for the reserved pad id.)
+        max_long = int(np.iinfo(np.int64).max)
+        M_max = max_long // (self.compressed_vocab + 1)
+        half_bound = max(1, M_max // 2)
+        PRIME_1 = 10007
+        mult = np.empty((len(self.layer_ids), self.max_ngram), dtype=np.int64)
+        for li, layer_id in enumerate(self.layer_ids):
+            g = np.random.default_rng(int(seed + PRIME_1 * int(layer_id)))
+            r = g.integers(low=0, high=half_bound, size=(self.max_ngram,), dtype=np.int64)
+            mult[li] = r * 2 + 1
+        self.register_buffer("mult", torch.from_numpy(mult), persistent=True)
+
+        # Distinct prime moduli per (layer, ngram order, head); unique globally.
+        seen: set = set()
+        self.head_sizes: List[List[int]] = []      # per layer: flat list, len n_cols
+        for _ in self.layer_ids:
+            flat = []
+            for ngram in range(2, self.max_ngram + 1):
+                search_start = int(vocab_size_per_ngram[ngram - 2]) - 1
+                for _h in range(self.n_heads):
+                    found = find_next_prime(search_start, seen)
+                    seen.add(found)
+                    flat.append(found)
+                    search_start = found
+            self.head_sizes.append(flat)
+        self.register_buffer(
+            "mods", torch.tensor(self.head_sizes, dtype=torch.int64), persistent=False)
+        head_offsets = []
+        for sizes in self.head_sizes:
+            offsets = []
+            offset = 0
+            for size in sizes:
+                offsets.append(offset)
+                offset += size
+            head_offsets.append(offsets)
+        self.register_buffer(
+            "head_offsets", torch.tensor(head_offsets, dtype=torch.int64), persistent=False)
+
+    def compress(self, idx: torch.Tensor) -> torch.Tensor:
+        return self.lookup[idx.long()]            # idx is never negative (the -1 ignore index lives in targets)
+
+    def forward(self, idx: torch.Tensor, doc_ids: torch.Tensor | None = None,
+                cache: dict | None = None):
+        """Hash tokens, optionally using cached history for incremental decoding.
+
+        doc_ids (B, T), when supplied, replaces cross-document history with the
+        left-padding id so packed documents hash as if processed separately.
+        """
+        ids = self.compress(idx)                                 # (B, T_new)
+        B, T_new = ids.shape
+        K = self.max_ngram - 1
+
+        if cache is not None:
+            hist_ids = cache.get("hash_hist_ids")
+            hist_doc = cache.get("hash_hist_doc")
+            if hist_ids is None:
+                hist_ids = ids.new_full((B, K), self.pad_cid)
+                hist_doc = (
+                    torch.full((B, K), -1, dtype=doc_ids.dtype, device=doc_ids.device)
+                    if doc_ids is not None else None
+                )
+            else:
+                hist_ids = hist_ids.to(device=ids.device)
+                if hist_doc is not None:
+                    hist_doc = hist_doc.to(device=ids.device)
+            full_ids = torch.cat([hist_ids, ids], dim=1)          # (B, K + T_new)
+            if doc_ids is not None and hist_doc is not None:
+                full_doc = torch.cat([hist_doc, doc_ids], dim=1)
+            else:
+                full_doc = doc_ids
+            cache["hash_hist_ids"] = full_ids[:, -K:].detach()
+            if full_doc is not None:
+                cache["hash_hist_doc"] = full_doc[:, -K:].detach()
+        else:
+            full_ids, full_doc = ids, doc_ids
+
+        T_full = full_ids.shape[1]
+        # Keep only one running hash per layer and one shifted-id tensor per order;
+        # broadcasting all layers and orders at once creates a large int64 product.
+        mixes = [full_ids * self.mult[li, 0] for li in range(len(self.layer_ids))]
+        out = {
+            layer_id: torch.empty(
+                (B, T_full, self.n_cols), dtype=torch.int64, device=full_ids.device)
+            for layer_id in self.layer_ids
+        }
+        for n in range(2, self.max_ngram + 1):
+            shifted = F.pad(full_ids, (n - 1, 0), value=self.pad_cid)[:, :T_full]
+            if full_doc is not None:
+                same = F.pad(full_doc, (n - 1, 0), value=-1)[:, :T_full] == full_doc
+                shifted = torch.where(same, shifted, torch.full_like(shifted, self.pad_cid))
+            lo = (n - 2) * self.n_heads
+            hi = lo + self.n_heads
+            for li, layer_id in enumerate(self.layer_ids):
+                mix = torch.bitwise_xor(mixes[li], shifted * self.mult[li, n - 1])
+                mixes[li] = mix
+                col = mix.unsqueeze(-1) % self.mods[li, lo:hi]
+                col += self.head_offsets[li, lo:hi]
+                out[layer_id][..., lo:hi] = col
+        if T_full > T_new:
+            out = {layer_id: values[:, -T_new:].contiguous() for layer_id, values in out.items()}
+        return out
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Modules
+# ─────────────────────────────────────────────────────────────────────────────
+
+class _RMSNorm(nn.Module):
+    """RMSNorm with fused accumulation for the input dtype."""
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps)
+
+
+class ShortConv(nn.Module):
+    """(B, T, D) -> (B, T, D). Depthwise causal conv, dilation = max_ngram."""
+
+    def __init__(self, hidden_size: int, kernel_size: int = 4,
+                 dilation: int = 1, activation: bool = True):
+        super().__init__()
+        self.activation = activation
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.conv = nn.Conv1d(
+            hidden_size, hidden_size, kernel_size=kernel_size,
+            groups=hidden_size, bias=False,
+            padding=(kernel_size - 1) * dilation, dilation=dilation,
+        )
+        self.norm = _RMSNorm(hidden_size)
+        self.act_fn = nn.SiLU() if activation else None
+
+    def forward(self, x: torch.Tensor, doc_ids: torch.Tensor | None = None,
+                cache: dict | None = None) -> torch.Tensor:
+        T = x.shape[1]
+        K, d = self.kernel_size, self.dilation
+        rf = (K - 1) * d
+        B, _, D = x.shape
+        xn = self.norm(x)                                        # (B, T, D)
+
+        if cache is not None:
+            hist_xn = cache.get("sc_hist_x")
+            hist_doc = cache.get("sc_hist_doc")
+            if hist_xn is None:
+                hist_xn = xn.new_zeros(B, rf, D)
+                hist_doc = (
+                    torch.full((B, rf), -1, dtype=doc_ids.dtype, device=doc_ids.device)
+                    if doc_ids is not None else None
+                )
+            else:
+                hist_xn = hist_xn.to(dtype=xn.dtype, device=xn.device)
+                if hist_doc is not None:
+                    hist_doc = hist_doc.to(device=xn.device)
+            full_xn = torch.cat([hist_xn, xn], dim=1) if rf > 0 else xn
+            if doc_ids is not None and hist_doc is not None:
+                full_doc = torch.cat([hist_doc, doc_ids], dim=1) if rf > 0 else doc_ids
+            else:
+                full_doc = doc_ids
+            if rf > 0:
+                cache["sc_hist_x"] = full_xn[:, -rf:].detach()
+                if full_doc is not None:
+                    cache["sc_hist_doc"] = full_doc[:, -rf:].detach()
+        else:
+            full_xn, full_doc = xn, doc_ids
+
+        T_full = full_xn.shape[1]
+        w = self.conv.weight[:, 0, :].to(full_xn.dtype)          # (D, K)
+        xp = F.pad(full_xn, (0, 0, rf, 0))
+
+        tap = xp[:, :T_full]
+        if full_doc is not None and K > 1:
+            s = (K - 1) * d
+            same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc
+            tap = tap * same.unsqueeze(-1)
+        y = tap * w[:, 0]
+        for j in range(1, K):
+            tap = xp[:, j * d:j * d + T_full]
+            if full_doc is not None:
+                s = (K - 1 - j) * d
+                same = F.pad(full_doc, (s, 0), value=-1)[:, :T_full] == full_doc
+                tap = tap * same.unsqueeze(-1)
+            y.addcmul_(tap, w[:, j])
+        if T_full > T:
+            y = y[:, -T:].contiguous()
+        if self.activation:
+            y = self.act_fn(y)
+        return y
+
+
+class MultiHeadEmbedding(nn.Module):
+    def __init__(self, list_of_N: List[int], D: int):
+        super().__init__()
+        self.embedding = nn.Embedding(sum(list_of_N), D)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.normal_(self.embedding.weight, std=0.01)
+
+    def forward(self, hash_ids: torch.Tensor) -> torch.Tensor:
+        return self.embedding(hash_ids)
+
+
+class Engram(nn.Module):
+    """Static n-gram memory with conditional gated injection (hc_mult = 1).
+
+    forward() returns the DELTA to add to the residual stream:
+        x = x + engram(x, hash_ids)
+    """
+
+    def __init__(self, hidden_size: int, head_sizes: List[int], max_ngram: int,
+                 embed_per_ngram: int, n_heads: int, kernel_size: int):
+        super().__init__()
+        assert embed_per_ngram % n_heads == 0, (
+            f"engram_embed_per_ngram={embed_per_ngram} must be divisible by "
+            f"engram_n_heads={n_heads}")
+        self.hidden_size = hidden_size
+        d_head = embed_per_ngram // n_heads
+        engram_hidden = (max_ngram - 1) * embed_per_ngram
+
+        self.multi_head_embedding = MultiHeadEmbedding(list(head_sizes), d_head)
+        self.value_proj = nn.Linear(engram_hidden, hidden_size, bias=False)
+        self.key_proj = nn.Linear(engram_hidden, hidden_size, bias=False)
+        self.norm1 = _RMSNorm(hidden_size)
+        self.norm2 = _RMSNorm(hidden_size)
+        self.short_conv = ShortConv(hidden_size, kernel_size=kernel_size,
+                                    dilation=max_ngram)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.multi_head_embedding.reset_parameters()
+        nn.init.normal_(self.value_proj.weight, std=0.02)
+        nn.init.normal_(self.key_proj.weight, std=0.02)
+        # Zero-init the conv: ShortConv RMS-normalises its input, so without this the gated value is
+        # blown back up to ~1 rms and the initial delta is ~0.13 rms (larger than the residual stream at
+        # D<=256). With zeros the branch starts at ~0 and the conv weight gets gradient on step 1.
+        nn.init.zeros_(self.short_conv.conv.weight)
+
+    def forward(self, hidden_states: torch.Tensor, hash_ids: torch.Tensor,
+                doc_ids: torch.Tensor | None = None, cache: dict | None = None):
+        emb = self.multi_head_embedding(hash_ids).flatten(start_dim=-2)  # (B,T,engram_hidden)
+        key = self.key_proj(emb)
+        nk = self.norm1(key)
+        nq = self.norm2(hidden_states)
+        # Product AND sum in fp32: with bf16 inputs, rounding nk*nq to bf16 first leaves gate errors up to
+        # ~0.02 (sqrt(|g|) amplifies small pre-activations); the casts fuse under torch.compile.
+        gate = (nk.float() * nq.float()).sum(dim=-1) / math.sqrt(self.hidden_size)
+        gate = gate.abs().clamp_min(1e-6).sqrt() * gate.sign()
+        gate = gate.sigmoid().unsqueeze(-1).to(hidden_states.dtype)      # (B,T,1)
+        value = gate * self.value_proj(emb)
+        return value + self.short_conv(value, doc_ids, cache=cache)

optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f0ba3c44137a3fb75324ac0a6a1cd877aeadc13c324b7210af17721ba18d133
+size 228177995

rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61c19bab1174704a4a4441475683bf1270277af15d2e2c95e964789128e482c4
+size 14645

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:10a3ff0aa1e5c1862bafecb07e5f9585646ede86fefc8837cd3a4073cdde1205
+size 1465

tokenization_pragya.py

@@ -0,0 +1,182 @@
+import json
+import os
+
+import tiktoken
+from transformers import PreTrainedTokenizer
+
+
+def bytes_to_unicode():
+    """Reversible byte -> printable-char map (GPT-2 style). Every byte maps to
+    a unique character, so token strings never lose bytes."""
+    bs = (list(range(ord("!"), ord("~") + 1))
+          + list(range(ord("¡"), ord("¬") + 1))
+          + list(range(ord("®"), ord("ÿ") + 1)))
+    cs = bs[:]
+    n = 0
+    for b in range(256):
+        if b not in bs:
+            bs.append(b)
+            cs.append(256 + n)
+            n += 1
+    return {b: chr(c) for b, c in zip(bs, cs)}
+
+
+_BYTE_ENCODER = bytes_to_unicode()
+_BYTE_DECODER = {v: k for k, v in _BYTE_ENCODER.items()}
+
+
+class PragyaTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"tokenizer_file": "pragya.tokenizer.json"}
+
+    def __init__(self, tokenizer_file=None, bos_token="<|bos|>",
+                 eos_token="<|bos|>", pad_token="<|bos|>",
+                 unk_token="<|bos|>", **kwargs):
+        if tokenizer_file is None:
+            raise ValueError(
+                "PragyaTokenizer needs tokenizer_file=<path to pragya.tokenizer.json>, "
+                "or load it with from_pretrained()."
+            )
+        with open(tokenizer_file, "r", encoding="utf-8") as f:
+            self._raw = json.load(f)
+
+        mergeable_ranks = {bytes.fromhex(k): int(v)
+                           for k, v in self._raw["mergeable_ranks"].items()}
+        special_tokens = {k: int(v) for k, v in self._raw["special_tokens"].items()}
+        self._enc = tiktoken.Encoding(
+            name=self._raw["name"],
+            pat_str=self._raw["pat_str"],
+            mergeable_ranks=mergeable_ranks,
+            special_tokens=special_tokens,
+        )
+
+        # id <-> token string. Ordinary tokens use the reversible byte map;
+        # special tokens keep their literal string.
+        self._special_set = set(special_tokens)
+        self._id_to_tok = {}
+        for i in range(self._enc.n_vocab):
+            if i in set(special_tokens.values()):
+                continue
+            raw = self._enc.decode_single_token_bytes(i)
+            self._id_to_tok[i] = "".join(_BYTE_ENCODER[b] for b in raw)
+        for name, i in special_tokens.items():
+            self._id_to_tok[i] = name
+        self._tok_to_id = {t: i for i, t in self._id_to_tok.items()}
+
+        kwargs.setdefault("clean_up_tokenization_spaces", False)
+        kwargs.setdefault("padding_side", "left")  # generation needs left padding
+        kwargs.setdefault("model_max_length", 1024)  # CONTEXT in config__1_.py
+        super().__init__(bos_token=bos_token, eos_token=eos_token,
+                         pad_token=pad_token, unk_token=unk_token, **kwargs)
+
+    @classmethod
+    def _from_pretrained(cls, resolved_vocab_files, pretrained_model_name_or_path,
+                         *args, **kwargs):
+        # Transformers 5 can resolve our file to a default name (tokenizer.json) that
+        # we never ship. Find pragya.tokenizer.json in the local folder, or download it
+        # from the Hub repo, so loading works for both paths and repo ids.
+        if resolved_vocab_files.get("tokenizer_file") is None:
+            name = "pragya.tokenizer.json"
+            local_dir = str(pretrained_model_name_or_path)
+            if os.path.isdir(local_dir) and os.path.isfile(os.path.join(local_dir, name)):
+                resolved_vocab_files["tokenizer_file"] = os.path.join(local_dir, name)
+            else:
+                from huggingface_hub import hf_hub_download
+                from huggingface_hub.errors import EntryNotFoundError
+                try:
+                    resolved_vocab_files["tokenizer_file"] = hf_hub_download(
+                        repo_id=local_dir, filename=name,
+                        revision=kwargs.get("revision"), token=kwargs.get("token"),
+                    )
+                except EntryNotFoundError as e:
+                    raise FileNotFoundError(
+                        f"{name} is missing from {local_dir} on the Hub. Upload it to the "
+                        f"repository root alongside the other tokenizer files."
+                    ) from e
+        return super()._from_pretrained(resolved_vocab_files,
+                                        pretrained_model_name_or_path,
+                                        *args, **kwargs)
+
+    def _add_tokens(self, new_tokens, special_tokens=False):
+        # Transformers assigns ids past the end of the vocab to unknown tokens,
+        # which the model embedding cannot look up. Refuse them instead.
+        unknown = [t for t in new_tokens
+                   if str(t) not in self._tok_to_id and str(t) not in self.added_tokens_encoder]
+        if unknown:
+            raise ValueError(
+                f"Cannot add tokens {unknown!r}: this tokenizer has a fixed vocab "
+                f"of {self._enc.n_vocab} tokens and the model embedding cannot grow "
+                f"without retraining."
+            )
+        return super()._add_tokens(new_tokens, special_tokens=special_tokens)
+
+    def tokenize(self, text, **kwargs):
+        # Bypass HF's special-token splitting: the training pipeline encoded
+        # "<|bos|>" in text as ordinary characters, so match that exactly.
+        return self._tokenize(text, **kwargs)
+
+    @property
+    def vocab_size(self):
+        return self._enc.n_vocab
+
+    def get_vocab(self):
+        return {**self._tok_to_id, **self.added_tokens_encoder}
+
+    def _tokenize(self, text, **kwargs):
+        return [self._id_to_tok[i] for i in self._enc.encode_ordinary(text)]
+
+    def _convert_token_to_id(self, token):
+        if token in self._tok_to_id:
+            return self._tok_to_id[token]
+        if token in self.added_tokens_encoder:
+            return self.added_tokens_encoder[token]
+        raise ValueError(
+            f"Unknown token {token!r}. This tokenizer ships with a fixed vocab "
+            f"of {self._enc.n_vocab} tokens; new special tokens cannot be added "
+            f"without retraining the model embedding."
+        )
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        # Training wrapper prepends <|bos|> by default (add_bos=True).
+        ids = [self.bos_token_id] + list(token_ids_0)
+        if token_ids_1 is not None:
+            ids += list(token_ids_1)
+        return ids
+
+    def _convert_id_to_token(self, index):
+        index = int(index)
+        if index not in self._id_to_tok:
+            raise ValueError(
+                f"Token id {index} is outside the vocab of {self.vocab_size} tokens."
+            )
+        return self._id_to_tok[index]
+
+    def _decode(self, token_ids, skip_special_tokens=False, **kwargs):
+        # Accept a single id, a list, or a tensor/array (callers pass all three).
+        if hasattr(token_ids, "tolist"):
+            token_ids = token_ids.tolist()
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        valid = [int(i) for i in token_ids if 0 <= int(i) < self.vocab_size]
+        return super()._decode(valid, skip_special_tokens=skip_special_tokens, **kwargs)
+
+    def convert_tokens_to_string(self, tokens):
+        special = self._special_set | set(self.added_tokens_encoder)
+        out = bytearray()
+        parts = []
+        for t in tokens:
+            if t in special:
+                parts.append(out.decode("utf-8", errors="replace"))
+                out = bytearray()
+                parts.append(t)
+            else:
+                out.extend(_BYTE_DECODER[c] for c in t)
+        parts.append(out.decode("utf-8", errors="replace"))
+        return "".join(parts)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        name = (filename_prefix + "-" if filename_prefix else "") + "pragya.tokenizer.json"
+        path = os.path.join(save_directory, name)
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self._raw, f, indent=2, ensure_ascii=False)
+        return (path,)
+        

tokenizer_config.json

@@ -0,0 +1,29 @@
+{
+  "added_tokens_decoder": {
+    "12287": {
+      "content": "<|bos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_pragya.PragyaTokenizer",
+      null
+    ]
+  },
+  "backend": "custom",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|bos|>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<|bos|>",
+  "padding_side": "left",
+  "tokenizer_class": "PragyaTokenizer",
+  "unk_token": "<|bos|>"
+}

trainer_state.json

@@ -0,0 +1,2253 @@
+{
+  "best_global_step": null,
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 500,
+  "global_step": 15858,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.003152982721654685,
+      "grad_norm": 4.853567123413086,
+      "learning_rate": 9.9999880707815e-05,
+      "loss": 2.872132568359375,
+      "step": 50
+    },
+    {
+      "epoch": 0.00630596544330937,
+      "grad_norm": 8.673802375793457,
+      "learning_rate": 9.999633156224002e-05,
+      "loss": 2.0706358337402344,
+      "step": 100
+    },
+    {
+      "epoch": 0.009458948164964056,
+      "grad_norm": 3.1262404918670654,
+      "learning_rate": 9.998785338773455e-05,
+      "loss": 2.0303359985351563,
+      "step": 150
+    },
+    {
+      "epoch": 0.01261193088661874,
+      "grad_norm": 5.10759162902832,
+      "learning_rate": 9.997444702014325e-05,
+      "loss": 1.960927734375,
+      "step": 200
+    },
+    {
+      "epoch": 0.015764913608273427,
+      "grad_norm": 5.117474555969238,
+      "learning_rate": 9.995611378117058e-05,
+      "loss": 1.7625619506835937,
+      "step": 250
+    },
+    {
+      "epoch": 0.018917896329928113,
+      "grad_norm": 3.3129849433898926,
+      "learning_rate": 9.993285547825046e-05,
+      "loss": 1.963189697265625,
+      "step": 300
+    },
+    {
+      "epoch": 0.0220708790515828,
+      "grad_norm": 4.3984375,
+      "learning_rate": 9.990467440436808e-05,
+      "loss": 1.8761654663085938,
+      "step": 350
+    },
+    {
+      "epoch": 0.02522386177323748,
+      "grad_norm": 3.8485946655273438,
+      "learning_rate": 9.987157333783385e-05,
+      "loss": 1.8703598022460937,
+      "step": 400
+    },
+    {
+      "epoch": 0.028376844494892167,
+      "grad_norm": 3.3610949516296387,
+      "learning_rate": 9.98335555420095e-05,
+      "loss": 1.8849630737304688,
+      "step": 450
+    },
+    {
+      "epoch": 0.03152982721654685,
+      "grad_norm": 9.949187278747559,
+      "learning_rate": 9.979062476498629e-05,
+      "loss": 1.897828369140625,
+      "step": 500
+    },
+    {
+      "epoch": 0.03468280993820154,
+      "grad_norm": 3.609001398086548,
+      "learning_rate": 9.974278523921565e-05,
+      "loss": 1.9846488952636718,
+      "step": 550
+    },
+    {
+      "epoch": 0.037835792659856225,
+      "grad_norm": 2.9588377475738525,
+      "learning_rate": 9.96900416810917e-05,
+      "loss": 1.7952581787109374,
+      "step": 600
+    },
+    {
+      "epoch": 0.04098877538151091,
+      "grad_norm": 4.958306789398193,
+      "learning_rate": 9.963239929048644e-05,
+      "loss": 1.9661424255371094,
+      "step": 650
+    },
+    {
+      "epoch": 0.0441417581031656,
+      "grad_norm": 3.1150712966918945,
+      "learning_rate": 9.956986375023703e-05,
+      "loss": 1.8595201110839843,
+      "step": 700
+    },
+    {
+      "epoch": 0.04729474082482028,
+      "grad_norm": 4.197102069854736,
+      "learning_rate": 9.950244122558557e-05,
+      "loss": 1.8694996643066406,
+      "step": 750
+    },
+    {
+      "epoch": 0.05044772354647496,
+      "grad_norm": 3.935776472091675,
+      "learning_rate": 9.94301383635712e-05,
+      "loss": 1.8022630310058594,
+      "step": 800
+    },
+    {
+      "epoch": 0.05360070626812965,
+      "grad_norm": 2.621243953704834,
+      "learning_rate": 9.93529622923749e-05,
+      "loss": 1.8342015075683593,
+      "step": 850
+    },
+    {
+      "epoch": 0.056753688989784334,
+      "grad_norm": 5.246140956878662,
+      "learning_rate": 9.927092062061666e-05,
+      "loss": 1.8499333190917968,
+      "step": 900
+    },
+    {
+      "epoch": 0.05990667171143902,
+      "grad_norm": 2.5088658332824707,
+      "learning_rate": 9.918402143660537e-05,
+      "loss": 1.9792045593261718,
+      "step": 950
+    },
+    {
+      "epoch": 0.0630596544330937,
+      "grad_norm": 11.082776069641113,
+      "learning_rate": 9.909227330754144e-05,
+      "loss": 1.897413787841797,
+      "step": 1000
+    },
+    {
+      "epoch": 0.06621263715474839,
+      "grad_norm": 5.241158962249756,
+      "learning_rate": 9.899568527867216e-05,
+      "loss": 1.7708767700195311,
+      "step": 1050
+    },
+    {
+      "epoch": 0.06936561987640308,
+      "grad_norm": 4.000182151794434,
+      "learning_rate": 9.889426687239994e-05,
+      "loss": 1.9385163879394531,
+      "step": 1100
+    },
+    {
+      "epoch": 0.07251860259805776,
+      "grad_norm": 7.7123003005981445,
+      "learning_rate": 9.878802808734356e-05,
+      "loss": 1.6802894592285156,
+      "step": 1150
+    },
+    {
+      "epoch": 0.07567158531971245,
+      "grad_norm": 4.490638256072998,
+      "learning_rate": 9.867697939735235e-05,
+      "loss": 1.8175428771972657,
+      "step": 1200
+    },
+    {
+      "epoch": 0.07882456804136713,
+      "grad_norm": 5.000014781951904,
+      "learning_rate": 9.856113175047368e-05,
+      "loss": 1.7184861755371095,
+      "step": 1250
+    },
+    {
+      "epoch": 0.08197755076302182,
+      "grad_norm": 3.2123770713806152,
+      "learning_rate": 9.844049656787357e-05,
+      "loss": 1.8914837646484375,
+      "step": 1300
+    },
+    {
+      "epoch": 0.0851305334846765,
+      "grad_norm": 6.576638698577881,
+      "learning_rate": 9.831508574271071e-05,
+      "loss": 1.7806048583984375,
+      "step": 1350
+    },
+    {
+      "epoch": 0.0882835162063312,
+      "grad_norm": 7.240126132965088,
+      "learning_rate": 9.818491163896392e-05,
+      "loss": 1.707454833984375,
+      "step": 1400
+    },
+    {
+      "epoch": 0.09143649892798587,
+      "grad_norm": 2.1819403171539307,
+      "learning_rate": 9.804998709021327e-05,
+      "loss": 1.8345822143554686,
+      "step": 1450
+    },
+    {
+      "epoch": 0.09458948164964057,
+      "grad_norm": 2.9357962608337402,
+      "learning_rate": 9.791032539837478e-05,
+      "loss": 1.6845474243164062,
+      "step": 1500
+    },
+    {
+      "epoch": 0.09774246437129525,
+      "grad_norm": 3.118610382080078,
+      "learning_rate": 9.776594033238906e-05,
+      "loss": 1.7468299865722656,
+      "step": 1550
+    },
+    {
+      "epoch": 0.10089544709294992,
+      "grad_norm": 2.261528491973877,
+      "learning_rate": 9.761684612686381e-05,
+      "loss": 1.8356428527832032,
+      "step": 1600
+    },
+    {
+      "epoch": 0.10404842981460462,
+      "grad_norm": 2.1853625774383545,
+      "learning_rate": 9.746305748067054e-05,
+      "loss": 1.624322509765625,
+      "step": 1650
+    },
+    {
+      "epoch": 0.1072014125362593,
+      "grad_norm": 2.855556011199951,
+      "learning_rate": 9.730458955549537e-05,
+      "loss": 1.7296275329589843,
+      "step": 1700
+    },
+    {
+      "epoch": 0.11035439525791399,
+      "grad_norm": 3.9297707080841064,
+      "learning_rate": 9.714145797434428e-05,
+      "loss": 1.8347354125976563,
+      "step": 1750
+    },
+    {
+      "epoch": 0.11350737797956867,
+      "grad_norm": 2.9527971744537354,
+      "learning_rate": 9.69736788200029e-05,
+      "loss": 1.7922471618652345,
+      "step": 1800
+    },
+    {
+      "epoch": 0.11666036070122336,
+      "grad_norm": 4.253082275390625,
+      "learning_rate": 9.680126863345096e-05,
+      "loss": 1.723883056640625,
+      "step": 1850
+    },
+    {
+      "epoch": 0.11981334342287804,
+      "grad_norm": 2.4197330474853516,
+      "learning_rate": 9.662424441223148e-05,
+      "loss": 1.6388121032714844,
+      "step": 1900
+    },
+    {
+      "epoch": 0.12296632614453273,
+      "grad_norm": 6.219167709350586,
+      "learning_rate": 9.644262360877507e-05,
+      "loss": 1.878433074951172,
+      "step": 1950
+    },
+    {
+      "epoch": 0.1261193088661874,
+      "grad_norm": 5.6117706298828125,
+      "learning_rate": 9.625642412867937e-05,
+      "loss": 1.6546157836914062,
+      "step": 2000
+    },
+    {
+      "epoch": 0.1292722915878421,
+      "grad_norm": 4.747069835662842,
+      "learning_rate": 9.606566432894364e-05,
+      "loss": 1.7719760131835938,
+      "step": 2050
+    },
+    {
+      "epoch": 0.13242527430949677,
+      "grad_norm": 8.663575172424316,
+      "learning_rate": 9.587036301615917e-05,
+      "loss": 1.6918789672851562,
+      "step": 2100
+    },
+    {
+      "epoch": 0.13557825703115148,
+      "grad_norm": 3.6302502155303955,
+      "learning_rate": 9.567053944465504e-05,
+      "loss": 1.8181585693359374,
+      "step": 2150
+    },
+    {
+      "epoch": 0.13873123975280616,
+      "grad_norm": 14.437663078308105,
+      "learning_rate": 9.546621331459991e-05,
+      "loss": 1.7362814331054688,
+      "step": 2200
+    },
+    {
+      "epoch": 0.14188422247446084,
+      "grad_norm": 4.509804725646973,
+      "learning_rate": 9.525740477005986e-05,
+      "loss": 1.7941587829589845,
+      "step": 2250
+    },
+    {
+      "epoch": 0.14503720519611552,
+      "grad_norm": 5.562815189361572,
+      "learning_rate": 9.50441343970124e-05,
+      "loss": 1.6657884216308594,
+      "step": 2300
+    },
+    {
+      "epoch": 0.14819018791777022,
+      "grad_norm": 3.0274181365966797,
+      "learning_rate": 9.482642322131695e-05,
+      "loss": 1.8234332275390626,
+      "step": 2350
+    },
+    {
+      "epoch": 0.1513431706394249,
+      "grad_norm": 3.96207857131958,
+      "learning_rate": 9.460429270664193e-05,
+      "loss": 1.7132150268554687,
+      "step": 2400
+    },
+    {
+      "epoch": 0.15449615336107958,
+      "grad_norm": 3.421797752380371,
+      "learning_rate": 9.437776475234874e-05,
+      "loss": 1.7809776306152343,
+      "step": 2450
+    },
+    {
+      "epoch": 0.15764913608273426,
+      "grad_norm": 3.259612798690796,
+      "learning_rate": 9.414686169133271e-05,
+      "loss": 1.6762347412109375,
+      "step": 2500
+    },
+    {
+      "epoch": 0.16080211880438894,
+      "grad_norm": 3.375741720199585,
+      "learning_rate": 9.39116062878214e-05,
+      "loss": 1.6602528381347657,
+      "step": 2550
+    },
+    {
+      "epoch": 0.16395510152604365,
+      "grad_norm": 6.738144874572754,
+      "learning_rate": 9.367202173513025e-05,
+      "loss": 1.741137237548828,
+      "step": 2600
+    },
+    {
+      "epoch": 0.16710808424769832,
+      "grad_norm": 4.628149032592773,
+      "learning_rate": 9.34281316533761e-05,
+      "loss": 1.7104803466796874,
+      "step": 2650
+    },
+    {
+      "epoch": 0.170261066969353,
+      "grad_norm": 3.664982318878174,
+      "learning_rate": 9.317996008714844e-05,
+      "loss": 1.674332275390625,
+      "step": 2700
+    },
+    {
+      "epoch": 0.17341404969100768,
+      "grad_norm": 2.6816370487213135,
+      "learning_rate": 9.292753150313892e-05,
+      "loss": 1.6296810913085937,
+      "step": 2750
+    },
+    {
+      "epoch": 0.1765670324126624,
+      "grad_norm": 4.442898750305176,
+      "learning_rate": 9.267087078772934e-05,
+      "loss": 1.7200907897949218,
+      "step": 2800
+    },
+    {
+      "epoch": 0.17972001513431707,
+      "grad_norm": 3.3474419116973877,
+      "learning_rate": 9.241000324453798e-05,
+      "loss": 1.781136474609375,
+      "step": 2850
+    },
+    {
+      "epoch": 0.18287299785597175,
+      "grad_norm": 4.042556285858154,
+      "learning_rate": 9.214495459192516e-05,
+      "loss": 1.8000424194335938,
+      "step": 2900
+    },
+    {
+      "epoch": 0.18602598057762643,
+      "grad_norm": 3.282440662384033,
+      "learning_rate": 9.187575096045755e-05,
+      "loss": 1.7477403259277344,
+      "step": 2950
+    },
+    {
+      "epoch": 0.18917896329928113,
+      "grad_norm": 3.451063632965088,
+      "learning_rate": 9.160241889033212e-05,
+      "loss": 1.6878971862792969,
+      "step": 3000
+    },
+    {
+      "epoch": 0.1923319460209358,
+      "grad_norm": 3.1991782188415527,
+      "learning_rate": 9.132498532875962e-05,
+      "loss": 1.721937255859375,
+      "step": 3050
+    },
+    {
+      "epoch": 0.1954849287425905,
+      "grad_norm": 2.625821113586426,
+      "learning_rate": 9.104347762730778e-05,
+      "loss": 1.7239382934570313,
+      "step": 3100
+    },
+    {
+      "epoch": 0.19863791146424517,
+      "grad_norm": 9.162576675415039,
+      "learning_rate": 9.075792353920494e-05,
+      "loss": 1.806209259033203,
+      "step": 3150
+    },
+    {
+      "epoch": 0.20179089418589985,
+      "grad_norm": 2.865137815475464,
+      "learning_rate": 9.04683512166038e-05,
+      "loss": 1.7318582153320312,
+      "step": 3200
+    },
+    {
+      "epoch": 0.20494387690755456,
+      "grad_norm": 2.1771979331970215,
+      "learning_rate": 9.017478920780605e-05,
+      "loss": 1.7237417602539062,
+      "step": 3250
+    },
+    {
+      "epoch": 0.20809685962920924,
+      "grad_norm": 5.151927471160889,
+      "learning_rate": 8.987726645444773e-05,
+      "loss": 1.7477003479003905,
+      "step": 3300
+    },
+    {
+      "epoch": 0.21124984235086391,
+      "grad_norm": 2.5384881496429443,
+      "learning_rate": 8.957581228864616e-05,
+      "loss": 1.7482113647460937,
+      "step": 3350
+    },
+    {
+      "epoch": 0.2144028250725186,
+      "grad_norm": 6.593316555023193,
+      "learning_rate": 8.927045643010785e-05,
+      "loss": 1.6856721496582032,
+      "step": 3400
+    },
+    {
+      "epoch": 0.2175558077941733,
+      "grad_norm": 2.9909262657165527,
+      "learning_rate": 8.896122898319878e-05,
+      "loss": 1.61490478515625,
+      "step": 3450
+    },
+    {
+      "epoch": 0.22070879051582798,
+      "grad_norm": 9.641027450561523,
+      "learning_rate": 8.864816043397632e-05,
+      "loss": 1.7636630249023437,
+      "step": 3500
+    },
+    {
+      "epoch": 0.22386177323748266,
+      "grad_norm": 2.434013605117798,
+      "learning_rate": 8.833128164718368e-05,
+      "loss": 1.86688232421875,
+      "step": 3550
+    },
+    {
+      "epoch": 0.22701475595913734,
+      "grad_norm": 5.4786696434021,
+      "learning_rate": 8.801062386320707e-05,
+      "loss": 1.761875457763672,
+      "step": 3600
+    },
+    {
+      "epoch": 0.23016773868079202,
+      "grad_norm": 2.9664249420166016,
+      "learning_rate": 8.76862186949958e-05,
+      "loss": 1.688774871826172,
+      "step": 3650
+    },
+    {
+      "epoch": 0.23332072140244672,
+      "grad_norm": 6.82322883605957,
+      "learning_rate": 8.735809812494556e-05,
+      "loss": 1.5707682800292968,
+      "step": 3700
+    },
+    {
+      "epoch": 0.2364737041241014,
+      "grad_norm": 3.971310615539551,
+      "learning_rate": 8.702629450174538e-05,
+      "loss": 1.8241322326660157,
+      "step": 3750
+    },
+    {
+      "epoch": 0.23962668684575608,
+      "grad_norm": 2.892256259918213,
+      "learning_rate": 8.669084053718843e-05,
+      "loss": 1.6581959533691406,
+      "step": 3800
+    },
+    {
+      "epoch": 0.24277966956741076,
+      "grad_norm": 4.926596164703369,
+      "learning_rate": 8.635176930294706e-05,
+      "loss": 1.64531982421875,
+      "step": 3850
+    },
+    {
+      "epoch": 0.24593265228906547,
+      "grad_norm": 6.136980056762695,
+      "learning_rate": 8.600911422731237e-05,
+      "loss": 1.7906488037109376,
+      "step": 3900
+    },
+    {
+      "epoch": 0.24908563501072015,
+      "grad_norm": 2.399834156036377,
+      "learning_rate": 8.566290909189842e-05,
+      "loss": 1.5979080200195312,
+      "step": 3950
+    },
+    {
+      "epoch": 0.2522386177323748,
+      "grad_norm": 3.5449185371398926,
+      "learning_rate": 8.531318802831202e-05,
+      "loss": 1.8152336120605468,
+      "step": 4000
+    },
+    {
+      "epoch": 0.25539160045402953,
+      "grad_norm": 4.556742191314697,
+      "learning_rate": 8.495998551478759e-05,
+      "loss": 1.7356678771972656,
+      "step": 4050
+    },
+    {
+      "epoch": 0.2585445831756842,
+      "grad_norm": 2.8251750469207764,
+      "learning_rate": 8.460333637278808e-05,
+      "loss": 1.5453990173339844,
+      "step": 4100
+    },
+    {
+      "epoch": 0.2616975658973389,
+      "grad_norm": 2.4924354553222656,
+      "learning_rate": 8.424327576357201e-05,
+      "loss": 1.7409744262695312,
+      "step": 4150
+    },
+    {
+      "epoch": 0.26485054861899354,
+      "grad_norm": 2.6736836433410645,
+      "learning_rate": 8.387983918472697e-05,
+      "loss": 1.75020263671875,
+      "step": 4200
+    },
+    {
+      "epoch": 0.26800353134064825,
+      "grad_norm": 10.364773750305176,
+      "learning_rate": 8.351306246667006e-05,
+      "loss": 1.7393624877929688,
+      "step": 4250
+    },
+    {
+      "epoch": 0.27115651406230296,
+      "grad_norm": 5.713336944580078,
+      "learning_rate": 8.314298176911533e-05,
+      "loss": 1.538885040283203,
+      "step": 4300
+    },
+    {
+      "epoch": 0.2743094967839576,
+      "grad_norm": 4.43858003616333,
+      "learning_rate": 8.276963357750896e-05,
+      "loss": 1.821629180908203,
+      "step": 4350
+    },
+    {
+      "epoch": 0.2774624795056123,
+      "grad_norm": 23.337303161621094,
+      "learning_rate": 8.239305469943226e-05,
+      "loss": 1.6454342651367186,
+      "step": 4400
+    },
+    {
+      "epoch": 0.280615462227267,
+      "grad_norm": 3.1600468158721924,
+      "learning_rate": 8.201328226097279e-05,
+      "loss": 1.6436555480957031,
+      "step": 4450
+    },
+    {
+      "epoch": 0.28376844494892167,
+      "grad_norm": 2.385958433151245,
+      "learning_rate": 8.163035370306429e-05,
+      "loss": 1.8098031616210937,
+      "step": 4500
+    },
+    {
+      "epoch": 0.2869214276705764,
+      "grad_norm": 3.7839107513427734,
+      "learning_rate": 8.124430677779537e-05,
+      "loss": 1.7188877868652344,
+      "step": 4550
+    },
+    {
+      "epoch": 0.29007441039223103,
+      "grad_norm": 3.354954719543457,
+      "learning_rate": 8.085517954468768e-05,
+      "loss": 1.6609513854980469,
+      "step": 4600
+    },
+    {
+      "epoch": 0.29322739311388574,
+      "grad_norm": 2.1363513469696045,
+      "learning_rate": 8.046301036694368e-05,
+      "loss": 1.6629969787597656,
+      "step": 4650
+    },
+    {
+      "epoch": 0.29638037583554044,
+      "grad_norm": 2.8907933235168457,
+      "learning_rate": 8.006783790766448e-05,
+      "loss": 1.5718318176269532,
+      "step": 4700
+    },
+    {
+      "epoch": 0.2995333585571951,
+      "grad_norm": 2.4533779621124268,
+      "learning_rate": 7.966970112603815e-05,
+      "loss": 1.6219189453125,
+      "step": 4750
+    },
+    {
+      "epoch": 0.3026863412788498,
+      "grad_norm": 5.0952558517456055,
+      "learning_rate": 7.926863927349883e-05,
+      "loss": 1.8162387084960938,
+      "step": 4800
+    },
+    {
+      "epoch": 0.30583932400050445,
+      "grad_norm": 3.4920997619628906,
+      "learning_rate": 7.8864691889857e-05,
+      "loss": 1.8339678955078125,
+      "step": 4850
+    },
+    {
+      "epoch": 0.30899230672215916,
+      "grad_norm": 6.650540828704834,
+      "learning_rate": 7.845789879940129e-05,
+      "loss": 1.536930389404297,
+      "step": 4900
+    },
+    {
+      "epoch": 0.31214528944381387,
+      "grad_norm": 2.608124256134033,
+      "learning_rate": 7.804830010697239e-05,
+      "loss": 1.6511822509765626,
+      "step": 4950
+    },
+    {
+      "epoch": 0.3152982721654685,
+      "grad_norm": 3.004046678543091,
+      "learning_rate": 7.763593619400913e-05,
+      "loss": 1.7977149963378907,
+      "step": 5000
+    },
+    {
+      "epoch": 0.3184512548871232,
+      "grad_norm": 2.8951213359832764,
+      "learning_rate": 7.722084771456742e-05,
+      "loss": 1.69650390625,
+      "step": 5050
+    },
+    {
+      "epoch": 0.3216042376087779,
+      "grad_norm": 2.372918128967285,
+      "learning_rate": 7.680307559131212e-05,
+      "loss": 1.7052632141113282,
+      "step": 5100
+    },
+    {
+      "epoch": 0.3247572203304326,
+      "grad_norm": 3.3484256267547607,
+      "learning_rate": 7.638266101148277e-05,
+      "loss": 1.7605667114257812,
+      "step": 5150
+    },
+    {
+      "epoch": 0.3279102030520873,
+      "grad_norm": 2.811953544616699,
+      "learning_rate": 7.595964542283286e-05,
+      "loss": 1.6529649353027345,
+      "step": 5200
+    },
+    {
+      "epoch": 0.33106318577374194,
+      "grad_norm": 3.4263999462127686,
+      "learning_rate": 7.553407052954367e-05,
+      "loss": 1.4979841613769531,
+      "step": 5250
+    },
+    {
+      "epoch": 0.33421616849539665,
+      "grad_norm": 5.087120532989502,
+      "learning_rate": 7.51059782881127e-05,
+      "loss": 1.600594482421875,
+      "step": 5300
+    },
+    {
+      "epoch": 0.33736915121705136,
+      "grad_norm": 3.6007165908813477,
+      "learning_rate": 7.467541090321735e-05,
+      "loss": 1.7187171936035157,
+      "step": 5350
+    },
+    {
+      "epoch": 0.340522133938706,
+      "grad_norm": 3.5818512439727783,
+      "learning_rate": 7.424241082355393e-05,
+      "loss": 1.6174636840820313,
+      "step": 5400
+    },
+    {
+      "epoch": 0.3436751166603607,
+      "grad_norm": 4.568120956420898,
+      "learning_rate": 7.380702073765284e-05,
+      "loss": 1.98239501953125,
+      "step": 5450
+    },
+    {
+      "epoch": 0.34682809938201536,
+      "grad_norm": 3.639190673828125,
+      "learning_rate": 7.336928356966996e-05,
+      "loss": 1.6210305786132813,
+      "step": 5500
+    },
+    {
+      "epoch": 0.34998108210367007,
+      "grad_norm": 5.252638816833496,
+      "learning_rate": 7.29292424751549e-05,
+      "loss": 1.669630889892578,
+      "step": 5550
+    },
+    {
+      "epoch": 0.3531340648253248,
+      "grad_norm": 2.4306838512420654,
+      "learning_rate": 7.248694083679623e-05,
+      "loss": 1.5243975830078125,
+      "step": 5600
+    },
+    {
+      "epoch": 0.35628704754697943,
+      "grad_norm": 3.1585371494293213,
+      "learning_rate": 7.204242226014471e-05,
+      "loss": 1.6251898193359375,
+      "step": 5650
+    },
+    {
+      "epoch": 0.35944003026863414,
+      "grad_norm": 2.6724190711975098,
+      "learning_rate": 7.159573056931415e-05,
+      "loss": 1.5321110534667968,
+      "step": 5700
+    },
+    {
+      "epoch": 0.3625930129902888,
+      "grad_norm": 2.7691164016723633,
+      "learning_rate": 7.114690980266092e-05,
+      "loss": 1.5984962463378907,
+      "step": 5750
+    },
+    {
+      "epoch": 0.3657459957119435,
+      "grad_norm": 4.595975875854492,
+      "learning_rate": 7.069600420844234e-05,
+      "loss": 1.469564208984375,
+      "step": 5800
+    },
+    {
+      "epoch": 0.3688989784335982,
+      "grad_norm": 2.72231388092041,
+      "learning_rate": 7.024305824045429e-05,
+      "loss": 1.40793212890625,
+      "step": 5850
+    },
+    {
+      "epoch": 0.37205196115525285,
+      "grad_norm": 7.415063858032227,
+      "learning_rate": 6.978811655364864e-05,
+      "loss": 1.7532875061035156,
+      "step": 5900
+    },
+    {
+      "epoch": 0.37520494387690756,
+      "grad_norm": 2.393813133239746,
+      "learning_rate": 6.933122399973086e-05,
+      "loss": 1.5745999145507812,
+      "step": 5950
+    },
+    {
+      "epoch": 0.37835792659856227,
+      "grad_norm": 6.05046272277832,
+      "learning_rate": 6.88724256227381e-05,
+      "loss": 1.6636524963378907,
+      "step": 6000
+    },
+    {
+      "epoch": 0.3815109093202169,
+      "grad_norm": 2.008434772491455,
+      "learning_rate": 6.84117666545984e-05,
+      "loss": 1.6576199340820312,
+      "step": 6050
+    },
+    {
+      "epoch": 0.3846638920418716,
+      "grad_norm": 2.7076210975646973,
+      "learning_rate": 6.794929251067147e-05,
+      "loss": 1.6187727355957031,
+      "step": 6100
+    },
+    {
+      "epoch": 0.3878168747635263,
+      "grad_norm": 9.455795288085938,
+      "learning_rate": 6.748504878527123e-05,
+      "loss": 1.5639329528808594,
+      "step": 6150
+    },
+    {
+      "epoch": 0.390969857485181,
+      "grad_norm": 2.4191970825195312,
+      "learning_rate": 6.701908124717071e-05,
+      "loss": 1.71536376953125,
+      "step": 6200
+    },
+    {
+      "epoch": 0.3941228402068357,
+      "grad_norm": 9.594432830810547,
+      "learning_rate": 6.655143583508992e-05,
+      "loss": 1.5790968322753907,
+      "step": 6250
+    },
+    {
+      "epoch": 0.39727582292849034,
+      "grad_norm": 2.7539901733398438,
+      "learning_rate": 6.608215865316676e-05,
+      "loss": 1.6756141662597657,
+      "step": 6300
+    },
+    {
+      "epoch": 0.40042880565014505,
+      "grad_norm": 2.871234178543091,
+      "learning_rate": 6.561129596641177e-05,
+      "loss": 1.537653045654297,
+      "step": 6350
+    },
+    {
+      "epoch": 0.4035817883717997,
+      "grad_norm": 7.255851745605469,
+      "learning_rate": 6.513889419614691e-05,
+      "loss": 1.5986100769042968,
+      "step": 6400
+    },
+    {
+      "epoch": 0.4067347710934544,
+      "grad_norm": 2.803478956222534,
+      "learning_rate": 6.466499991542905e-05,
+      "loss": 1.5736700439453124,
+      "step": 6450
+    },
+    {
+      "epoch": 0.4098877538151091,
+      "grad_norm": 4.441070079803467,
+      "learning_rate": 6.418965984445839e-05,
+      "loss": 1.5858750915527344,
+      "step": 6500
+    },
+    {
+      "epoch": 0.41304073653676376,
+      "grad_norm": 3.3662540912628174,
+      "learning_rate": 6.371292084597241e-05,
+      "loss": 1.7305607604980469,
+      "step": 6550
+    },
+    {
+      "epoch": 0.41619371925841847,
+      "grad_norm": 4.166026592254639,
+      "learning_rate": 6.323482992062579e-05,
+      "loss": 1.5626411437988281,
+      "step": 6600
+    },
+    {
+      "epoch": 0.4193467019800731,
+      "grad_norm": 7.287445068359375,
+      "learning_rate": 6.275543420235679e-05,
+      "loss": 1.5980224609375,
+      "step": 6650
+    },
+    {
+      "epoch": 0.42249968470172783,
+      "grad_norm": 2.271918773651123,
+      "learning_rate": 6.227478095374027e-05,
+      "loss": 1.6121871948242188,
+      "step": 6700
+    },
+    {
+      "epoch": 0.42565266742338254,
+      "grad_norm": 2.7370519638061523,
+      "learning_rate": 6.179291756132826e-05,
+      "loss": 1.7519253540039061,
+      "step": 6750
+    },
+    {
+      "epoch": 0.4288056501450372,
+      "grad_norm": 3.291109800338745,
+      "learning_rate": 6.130989153097825e-05,
+      "loss": 1.593614044189453,
+      "step": 6800
+    },
+    {
+      "epoch": 0.4319586328666919,
+      "grad_norm": 2.3776819705963135,
+      "learning_rate": 6.08257504831697e-05,
+      "loss": 1.6327113342285156,
+      "step": 6850
+    },
+    {
+      "epoch": 0.4351116155883466,
+      "grad_norm": 3.4966859817504883,
+      "learning_rate": 6.034054214830909e-05,
+      "loss": 1.524584503173828,
+      "step": 6900
+    },
+    {
+      "epoch": 0.43826459831000125,
+      "grad_norm": 6.587697982788086,
+      "learning_rate": 5.985431436202447e-05,
+      "loss": 1.7271713256835937,
+      "step": 6950
+    },
+    {
+      "epoch": 0.44141758103165596,
+      "grad_norm": 5.027621746063232,
+      "learning_rate": 5.9367115060449364e-05,
+      "loss": 1.4599151611328125,
+      "step": 7000
+    },
+    {
+      "epoch": 0.4445705637533106,
+      "grad_norm": 2.3149125576019287,
+      "learning_rate": 5.887899227549689e-05,
+      "loss": 1.6622496032714844,
+      "step": 7050
+    },
+    {
+      "epoch": 0.4477235464749653,
+      "grad_norm": 2.9743404388427734,
+      "learning_rate": 5.838999413012432e-05,
+      "loss": 1.5690019226074219,
+      "step": 7100
+    },
+    {
+      "epoch": 0.45087652919662,
+      "grad_norm": 3.092113494873047,
+      "learning_rate": 5.79001688335888e-05,
+      "loss": 1.6651068115234375,
+      "step": 7150
+    },
+    {
+      "epoch": 0.4540295119182747,
+      "grad_norm": 2.21628737449646,
+      "learning_rate": 5.7409564676694535e-05,
+      "loss": 1.6519837951660157,
+      "step": 7200
+    },
+    {
+      "epoch": 0.4571824946399294,
+      "grad_norm": 2.658684730529785,
+      "learning_rate": 5.691823002703187e-05,
+      "loss": 1.5685855102539064,
+      "step": 7250
+    },
+    {
+      "epoch": 0.46033547736158403,
+      "grad_norm": 6.6410017013549805,
+      "learning_rate": 5.642621332420879e-05,
+      "loss": 1.6819776916503906,
+      "step": 7300
+    },
+    {
+      "epoch": 0.46348846008323874,
+      "grad_norm": 5.142892837524414,
+      "learning_rate": 5.5933563075075425e-05,
+      "loss": 1.5651889038085938,
+      "step": 7350
+    },
+    {
+      "epoch": 0.46664144280489345,
+      "grad_norm": 2.901061534881592,
+      "learning_rate": 5.544032784894187e-05,
+      "loss": 1.6200245666503905,
+      "step": 7400
+    },
+    {
+      "epoch": 0.4697944255265481,
+      "grad_norm": 2.4046876430511475,
+      "learning_rate": 5.4946556272789806e-05,
+      "loss": 1.6270751953125,
+      "step": 7450
+    },
+    {
+      "epoch": 0.4729474082482028,
+      "grad_norm": 2.8165345191955566,
+      "learning_rate": 5.4452297026478486e-05,
+      "loss": 1.5757356262207032,
+      "step": 7500
+    },
+    {
+      "epoch": 0.4761003909698575,
+      "grad_norm": 2.545131206512451,
+      "learning_rate": 5.3957598837945525e-05,
+      "loss": 1.6749734497070312,
+      "step": 7550
+    },
+    {
+      "epoch": 0.47925337369151216,
+      "grad_norm": 4.947996139526367,
+      "learning_rate": 5.3462510478402825e-05,
+      "loss": 1.5667103576660155,
+      "step": 7600
+    },
+    {
+      "epoch": 0.48240635641316687,
+      "grad_norm": 10.719531059265137,
+      "learning_rate": 5.2967080757528495e-05,
+      "loss": 1.6775833129882813,
+      "step": 7650
+    },
+    {
+      "epoch": 0.4855593391348215,
+      "grad_norm": 3.945706844329834,
+      "learning_rate": 5.247135851865461e-05,
+      "loss": 1.541700439453125,
+      "step": 7700
+    },
+    {
+      "epoch": 0.48871232185647623,
+      "grad_norm": 4.664523124694824,
+      "learning_rate": 5.197539263395201e-05,
+      "loss": 1.6849325561523438,
+      "step": 7750
+    },
+    {
+      "epoch": 0.49186530457813094,
+      "grad_norm": 2.7066149711608887,
+      "learning_rate": 5.1479231999612046e-05,
+      "loss": 1.6272938537597657,
+      "step": 7800
+    },
+    {
+      "epoch": 0.4950182872997856,
+      "grad_norm": 3.3770575523376465,
+      "learning_rate": 5.098292553102598e-05,
+      "loss": 1.781251220703125,
+      "step": 7850
+    },
+    {
+      "epoch": 0.4981712700214403,
+      "grad_norm": 2.614497661590576,
+      "learning_rate": 5.048652215796257e-05,
+      "loss": 1.6480955505371093,
+      "step": 7900
+    },
+    {
+      "epoch": 0.501324252743095,
+      "grad_norm": 3.470582962036133,
+      "learning_rate": 4.999007081974418e-05,
+      "loss": 1.6021221923828124,
+      "step": 7950
+    },
+    {
+      "epoch": 0.5044772354647497,
+      "grad_norm": 1.8953889608383179,
+      "learning_rate": 4.949362046042195e-05,
+      "loss": 1.5847406005859375,
+      "step": 8000
+    },
+    {
+      "epoch": 0.5076302181864043,
+      "grad_norm": 6.359344482421875,
+      "learning_rate": 4.899722002395052e-05,
+      "loss": 1.6197129821777343,
+      "step": 8050
+    },
+    {
+      "epoch": 0.5107832009080591,
+      "grad_norm": 3.1751062870025635,
+      "learning_rate": 4.8500918449362686e-05,
+      "loss": 1.6575222778320313,
+      "step": 8100
+    },
+    {
+      "epoch": 0.5139361836297137,
+      "grad_norm": 4.892060279846191,
+      "learning_rate": 4.800476466594477e-05,
+      "loss": 1.8009135437011718,
+      "step": 8150
+    },
+    {
+      "epoch": 0.5170891663513684,
+      "grad_norm": 3.5821950435638428,
+      "learning_rate": 4.7508807588412607e-05,
+      "loss": 1.503452911376953,
+      "step": 8200
+    },
+    {
+      "epoch": 0.5202421490730231,
+      "grad_norm": 6.790828704833984,
+      "learning_rate": 4.7013096112089274e-05,
+      "loss": 1.71809814453125,
+      "step": 8250
+    },
+    {
+      "epoch": 0.5233951317946778,
+      "grad_norm": 2.2684848308563232,
+      "learning_rate": 4.651767910808453e-05,
+      "loss": 1.7311285400390626,
+      "step": 8300
+    },
+    {
+      "epoch": 0.5265481145163324,
+      "grad_norm": 3.4118940830230713,
+      "learning_rate": 4.602260541847674e-05,
+      "loss": 1.52655029296875,
+      "step": 8350
+    },
+    {
+      "epoch": 0.5297010972379871,
+      "grad_norm": 3.000192403793335,
+      "learning_rate": 4.552792385149774e-05,
+      "loss": 1.7028633117675782,
+      "step": 8400
+    },
+    {
+      "epoch": 0.5328540799596418,
+      "grad_norm": 39.9377555847168,
+      "learning_rate": 4.503368317672073e-05,
+      "loss": 1.5430813598632813,
+      "step": 8450
+    },
+    {
+      "epoch": 0.5360070626812965,
+      "grad_norm": 2.711533546447754,
+      "learning_rate": 4.453993212025242e-05,
+      "loss": 1.5495613098144532,
+      "step": 8500
+    },
+    {
+      "epoch": 0.5391600454029511,
+      "grad_norm": 3.9591002464294434,
+      "learning_rate": 4.4046719359929065e-05,
+      "loss": 1.6500083923339843,
+      "step": 8550
+    },
+    {
+      "epoch": 0.5423130281246059,
+      "grad_norm": 3.133467435836792,
+      "learning_rate": 4.355409352051747e-05,
+      "loss": 1.6975883483886718,
+      "step": 8600
+    },
+    {
+      "epoch": 0.5454660108462606,
+      "grad_norm": 3.454503059387207,
+      "learning_rate": 4.3062103168921264e-05,
+      "loss": 1.5262863159179687,
+      "step": 8650
+    },
+    {
+      "epoch": 0.5486189935679152,
+      "grad_norm": 3.6723296642303467,
+      "learning_rate": 4.2570796809392644e-05,
+      "loss": 1.7020167541503906,
+      "step": 8700
+    },
+    {
+      "epoch": 0.55177197628957,
+      "grad_norm": 5.7662529945373535,
+      "learning_rate": 4.208022287875054e-05,
+      "loss": 1.5918016052246093,
+      "step": 8750
+    },
+    {
+      "epoch": 0.5549249590112246,
+      "grad_norm": 2.1820693016052246,
+      "learning_rate": 4.159042974160537e-05,
+      "loss": 1.5883863830566407,
+      "step": 8800
+    },
+    {
+      "epoch": 0.5580779417328793,
+      "grad_norm": 3.121429920196533,
+      "learning_rate": 4.110146568559077e-05,
+      "loss": 1.6320164489746094,
+      "step": 8850
+    },
+    {
+      "epoch": 0.561230924454534,
+      "grad_norm": 6.376467704772949,
+      "learning_rate": 4.061337891660314e-05,
+      "loss": 1.635543212890625,
+      "step": 8900
+    },
+    {
+      "epoch": 0.5643839071761887,
+      "grad_norm": 7.583218574523926,
+      "learning_rate": 4.012621755404901e-05,
+      "loss": 1.5718000793457032,
+      "step": 8950
+    },
+    {
+      "epoch": 0.5675368898978433,
+      "grad_norm": 3.7945668697357178,
+      "learning_rate": 3.9640029626101146e-05,
+      "loss": 1.6808647155761718,
+      "step": 9000
+    },
+    {
+      "epoch": 0.570689872619498,
+      "grad_norm": 2.8824946880340576,
+      "learning_rate": 3.915486306496356e-05,
+      "loss": 1.685601806640625,
+      "step": 9050
+    },
+    {
+      "epoch": 0.5738428553411528,
+      "grad_norm": 6.4140472412109375,
+      "learning_rate": 3.867076570214584e-05,
+      "loss": 1.5931394958496095,
+      "step": 9100
+    },
+    {
+      "epoch": 0.5769958380628074,
+      "grad_norm": 2.3989524841308594,
+      "learning_rate": 3.8187785263747804e-05,
+      "loss": 1.6603244018554688,
+      "step": 9150
+    },
+    {
+      "epoch": 0.5801488207844621,
+      "grad_norm": 3.4079809188842773,
+      "learning_rate": 3.7705969365754e-05,
+      "loss": 1.6268783569335938,
+      "step": 9200
+    },
+    {
+      "epoch": 0.5833018035061168,
+      "grad_norm": 4.664431095123291,
+      "learning_rate": 3.722536550933953e-05,
+      "loss": 1.4689956665039063,
+      "step": 9250
+    },
+    {
+      "epoch": 0.5864547862277715,
+      "grad_norm": 4.000642776489258,
+      "learning_rate": 3.674602107618701e-05,
+      "loss": 1.5993995666503906,
+      "step": 9300
+    },
+    {
+      "epoch": 0.5896077689494261,
+      "grad_norm": 3.3489816188812256,
+      "learning_rate": 3.6267983323815226e-05,
+      "loss": 1.6080624389648437,
+      "step": 9350
+    },
+    {
+      "epoch": 0.5927607516710809,
+      "grad_norm": 2.956134796142578,
+      "learning_rate": 3.5791299380920164e-05,
+      "loss": 1.562136688232422,
+      "step": 9400
+    },
+    {
+      "epoch": 0.5959137343927355,
+      "grad_norm": 2.9019055366516113,
+      "learning_rate": 3.531601624272873e-05,
+      "loss": 1.596844482421875,
+      "step": 9450
+    },
+    {
+      "epoch": 0.5990667171143902,
+      "grad_norm": 3.2570741176605225,
+      "learning_rate": 3.484218076636553e-05,
+      "loss": 1.5718740844726562,
+      "step": 9500
+    },
+    {
+      "epoch": 0.602219699836045,
+      "grad_norm": 3.1301848888397217,
+      "learning_rate": 3.436983966623337e-05,
+      "loss": 1.5843309020996095,
+      "step": 9550
+    },
+    {
+      "epoch": 0.6053726825576996,
+      "grad_norm": 2.8983476161956787,
+      "learning_rate": 3.389903950940777e-05,
+      "loss": 1.501625518798828,
+      "step": 9600
+    },
+    {
+      "epoch": 0.6085256652793543,
+      "grad_norm": 2.442763566970825,
+      "learning_rate": 3.3429826711045996e-05,
+      "loss": 1.5783216857910156,
+      "step": 9650
+    },
+    {
+      "epoch": 0.6116786480010089,
+      "grad_norm": 3.317591905593872,
+      "learning_rate": 3.29622475298112e-05,
+      "loss": 1.6270458984375,
+      "step": 9700
+    },
+    {
+      "epoch": 0.6148316307226637,
+      "grad_norm": 5.122899532318115,
+      "learning_rate": 3.249634806331168e-05,
+      "loss": 1.551365509033203,
+      "step": 9750
+    },
+    {
+      "epoch": 0.6179846134443183,
+      "grad_norm": 3.8709020614624023,
+      "learning_rate": 3.203217424355643e-05,
+      "loss": 1.5474314880371094,
+      "step": 9800
+    },
+    {
+      "epoch": 0.621137596165973,
+      "grad_norm": 4.906006813049316,
+      "learning_rate": 3.156977183242666e-05,
+      "loss": 1.5910145568847656,
+      "step": 9850
+    },
+    {
+      "epoch": 0.6242905788876277,
+      "grad_norm": 2.5910794734954834,
+      "learning_rate": 3.110918641716428e-05,
+      "loss": 1.4721426391601562,
+      "step": 9900
+    },
+    {
+      "epoch": 0.6274435616092824,
+      "grad_norm": 11.209244728088379,
+      "learning_rate": 3.0650463405877574e-05,
+      "loss": 1.5800926208496093,
+      "step": 9950
+    },
+    {
+      "epoch": 0.630596544330937,
+      "grad_norm": 2.9779834747314453,
+      "learning_rate": 3.019364802306447e-05,
+      "loss": 1.521623077392578,
+      "step": 10000
+    },
+    {
+      "epoch": 0.6337495270525918,
+      "grad_norm": 4.262292385101318,
+      "learning_rate": 2.973878530515397e-05,
+      "loss": 1.553704071044922,
+      "step": 10050
+    },
+    {
+      "epoch": 0.6369025097742465,
+      "grad_norm": 3.476924419403076,
+      "learning_rate": 2.928592009606619e-05,
+      "loss": 1.6724470520019532,
+      "step": 10100
+    },
+    {
+      "epoch": 0.6400554924959011,
+      "grad_norm": 2.8566722869873047,
+      "learning_rate": 2.883509704279114e-05,
+      "loss": 1.580164031982422,
+      "step": 10150
+    },
+    {
+      "epoch": 0.6432084752175558,
+      "grad_norm": 5.948812007904053,
+      "learning_rate": 2.8386360590987166e-05,
+      "loss": 1.5428895568847656,
+      "step": 10200
+    },
+    {
+      "epoch": 0.6463614579392105,
+      "grad_norm": 3.7444076538085938,
+      "learning_rate": 2.79397549805992e-05,
+      "loss": 1.5748348999023438,
+      "step": 10250
+    },
+    {
+      "epoch": 0.6495144406608652,
+      "grad_norm": 2.652961015701294,
+      "learning_rate": 2.749532424149709e-05,
+      "loss": 1.590621337890625,
+      "step": 10300
+    },
+    {
+      "epoch": 0.6526674233825198,
+      "grad_norm": 3.2727580070495605,
+      "learning_rate": 2.705311218913499e-05,
+      "loss": 1.5267787170410156,
+      "step": 10350
+    },
+    {
+      "epoch": 0.6558204061041746,
+      "grad_norm": 10.58151912689209,
+      "learning_rate": 2.6613162420231484e-05,
+      "loss": 1.4954409790039063,
+      "step": 10400
+    },
+    {
+      "epoch": 0.6589733888258292,
+      "grad_norm": 3.0570061206817627,
+      "learning_rate": 2.6175518308471642e-05,
+      "loss": 1.4097785949707031,
+      "step": 10450
+    },
+    {
+      "epoch": 0.6621263715474839,
+      "grad_norm": 2.9748072624206543,
+      "learning_rate": 2.574022300023084e-05,
+      "loss": 1.4549171447753906,
+      "step": 10500
+    },
+    {
+      "epoch": 0.6652793542691386,
+      "grad_norm": 4.236860275268555,
+      "learning_rate": 2.5307319410320994e-05,
+      "loss": 1.7020782470703124,
+      "step": 10550
+    },
+    {
+      "epoch": 0.6684323369907933,
+      "grad_norm": 2.981145143508911,
+      "learning_rate": 2.4876850217759813e-05,
+      "loss": 1.6508351135253907,
+      "step": 10600
+    },
+    {
+      "epoch": 0.671585319712448,
+      "grad_norm": 3.248321771621704,
+      "learning_rate": 2.444885786156308e-05,
+      "loss": 1.588943328857422,
+      "step": 10650
+    },
+    {
+      "epoch": 0.6747383024341027,
+      "grad_norm": 4.321737766265869,
+      "learning_rate": 2.4023384536560638e-05,
+      "loss": 1.5426924133300781,
+      "step": 10700
+    },
+    {
+      "epoch": 0.6778912851557574,
+      "grad_norm": 2.416674852371216,
+      "learning_rate": 2.3600472189236693e-05,
+      "loss": 1.4259295654296875,
+      "step": 10750
+    },
+    {
+      "epoch": 0.681044267877412,
+      "grad_norm": 2.6175670623779297,
+      "learning_rate": 2.3180162513594123e-05,
+      "loss": 1.4130958557128905,
+      "step": 10800
+    },
+    {
+      "epoch": 0.6841972505990667,
+      "grad_norm": 2.6447396278381348,
+      "learning_rate": 2.276249694704417e-05,
+      "loss": 1.459014892578125,
+      "step": 10850
+    },
+    {
+      "epoch": 0.6873502333207214,
+      "grad_norm": 4.287785530090332,
+      "learning_rate": 2.2347516666321168e-05,
+      "loss": 1.600723876953125,
+      "step": 10900
+    },
+    {
+      "epoch": 0.6905032160423761,
+      "grad_norm": 4.545198440551758,
+      "learning_rate": 2.1935262583422917e-05,
+      "loss": 1.58876708984375,
+      "step": 10950
+    },
+    {
+      "epoch": 0.6936561987640307,
+      "grad_norm": 7.753457546234131,
+      "learning_rate": 2.1525775341577405e-05,
+      "loss": 1.5217156982421876,
+      "step": 11000
+    },
+    {
+      "epoch": 0.6968091814856855,
+      "grad_norm": 5.363049030303955,
+      "learning_rate": 2.1119095311235797e-05,
+      "loss": 1.5814126586914063,
+      "step": 11050
+    },
+    {
+      "epoch": 0.6999621642073401,
+      "grad_norm": 5.7626848220825195,
+      "learning_rate": 2.0715262586092428e-05,
+      "loss": 1.5985421752929687,
+      "step": 11100
+    },
+    {
+      "epoch": 0.7031151469289948,
+      "grad_norm": 4.6263017654418945,
+      "learning_rate": 2.0314316979132035e-05,
+      "loss": 1.6171270751953124,
+      "step": 11150
+    },
+    {
+      "epoch": 0.7062681296506496,
+      "grad_norm": 2.6141841411590576,
+      "learning_rate": 1.9916298018704665e-05,
+      "loss": 1.6464085388183594,
+      "step": 11200
+    },
+    {
+      "epoch": 0.7094211123723042,
+      "grad_norm": 2.904296636581421,
+      "learning_rate": 1.952124494462873e-05,
+      "loss": 1.5681950378417968,
+      "step": 11250
+    },
+    {
+      "epoch": 0.7125740950939589,
+      "grad_norm": 17.208356857299805,
+      "learning_rate": 1.9129196704322382e-05,
+      "loss": 1.4759814453125,
+      "step": 11300
+    },
+    {
+      "epoch": 0.7157270778156136,
+      "grad_norm": 4.990921497344971,
+      "learning_rate": 1.874019194896376e-05,
+      "loss": 1.4149720764160156,
+      "step": 11350
+    },
+    {
+      "epoch": 0.7188800605372683,
+      "grad_norm": 3.007599353790283,
+      "learning_rate": 1.8354269029680516e-05,
+      "loss": 1.3602633666992188,
+      "step": 11400
+    },
+    {
+      "epoch": 0.7220330432589229,
+      "grad_norm": 2.493307590484619,
+      "learning_rate": 1.7971465993768815e-05,
+      "loss": 1.7200181579589844,
+      "step": 11450
+    },
+    {
+      "epoch": 0.7251860259805776,
+      "grad_norm": 2.6334586143493652,
+      "learning_rate": 1.759182058094235e-05,
+      "loss": 1.4771337890625,
+      "step": 11500
+    },
+    {
+      "epoch": 0.7283390087022323,
+      "grad_norm": 4.686981201171875,
+      "learning_rate": 1.721537021961167e-05,
+      "loss": 1.5164639282226562,
+      "step": 11550
+    },
+    {
+      "epoch": 0.731491991423887,
+      "grad_norm": 7.735423564910889,
+      "learning_rate": 1.6842152023194145e-05,
+      "loss": 1.4624899291992188,
+      "step": 11600
+    },
+    {
+      "epoch": 0.7346449741455416,
+      "grad_norm": 3.569797992706299,
+      "learning_rate": 1.6472202786455136e-05,
+      "loss": 1.6836318969726562,
+      "step": 11650
+    },
+    {
+      "epoch": 0.7377979568671964,
+      "grad_norm": 3.5440757274627686,
+      "learning_rate": 1.610555898188039e-05,
+      "loss": 1.5107637023925782,
+      "step": 11700
+    },
+    {
+      "epoch": 0.740950939588851,
+      "grad_norm": 2.3150668144226074,
+      "learning_rate": 1.5742256756080358e-05,
+      "loss": 1.5355830383300781,
+      "step": 11750
+    },
+    {
+      "epoch": 0.7441039223105057,
+      "grad_norm": 3.0666487216949463,
+      "learning_rate": 1.53823319262265e-05,
+      "loss": 1.6443894958496095,
+      "step": 11800
+    },
+    {
+      "epoch": 0.7472569050321605,
+      "grad_norm": 3.2240090370178223,
+      "learning_rate": 1.5025819976520244e-05,
+      "loss": 1.5667915344238281,
+      "step": 11850
+    },
+    {
+      "epoch": 0.7504098877538151,
+      "grad_norm": 4.273111343383789,
+      "learning_rate": 1.4672756054694636e-05,
+      "loss": 1.4949490356445312,
+      "step": 11900
+    },
+    {
+      "epoch": 0.7535628704754698,
+      "grad_norm": 2.6397128105163574,
+      "learning_rate": 1.4323174968549147e-05,
+      "loss": 1.580462188720703,
+      "step": 11950
+    },
+    {
+      "epoch": 0.7567158531971245,
+      "grad_norm": 4.470552444458008,
+      "learning_rate": 1.3977111182518148e-05,
+      "loss": 1.6157127380371095,
+      "step": 12000
+    },
+    {
+      "epoch": 0.7598688359187792,
+      "grad_norm": 2.7804784774780273,
+      "learning_rate": 1.3634598814273075e-05,
+      "loss": 1.4012615966796875,
+      "step": 12050
+    },
+    {
+      "epoch": 0.7630218186404338,
+      "grad_norm": 2.2358453273773193,
+      "learning_rate": 1.3295671631358853e-05,
+      "loss": 1.4852983093261718,
+      "step": 12100
+    },
+    {
+      "epoch": 0.7661748013620885,
+      "grad_norm": 4.683980941772461,
+      "learning_rate": 1.2960363047864849e-05,
+      "loss": 1.6236349487304687,
+      "step": 12150
+    },
+    {
+      "epoch": 0.7693277840837432,
+      "grad_norm": 2.738664150238037,
+      "learning_rate": 1.2628706121130573e-05,
+      "loss": 1.4560935974121094,
+      "step": 12200
+    },
+    {
+      "epoch": 0.7724807668053979,
+      "grad_norm": 4.65731143951416,
+      "learning_rate": 1.2300733548486748e-05,
+      "loss": 1.6234567260742188,
+      "step": 12250
+    },
+    {
+      "epoch": 0.7756337495270526,
+      "grad_norm": 5.1576104164123535,
+      "learning_rate": 1.1976477664031682e-05,
+      "loss": 1.596574249267578,
+      "step": 12300
+    },
+    {
+      "epoch": 0.7787867322487073,
+      "grad_norm": 3.1281590461730957,
+      "learning_rate": 1.1655970435443458e-05,
+      "loss": 1.4229803466796875,
+      "step": 12350
+    },
+    {
+      "epoch": 0.781939714970362,
+      "grad_norm": 2.9372220039367676,
+      "learning_rate": 1.1339243460828503e-05,
+      "loss": 1.4996551513671874,
+      "step": 12400
+    },
+    {
+      "epoch": 0.7850926976920166,
+      "grad_norm": 3.2248833179473877,
+      "learning_rate": 1.1026327965606147e-05,
+      "loss": 1.4965032958984374,
+      "step": 12450
+    },
+    {
+      "epoch": 0.7882456804136714,
+      "grad_norm": 2.159266471862793,
+      "learning_rate": 1.0717254799430398e-05,
+      "loss": 1.5509339904785155,
+      "step": 12500
+    },
+    {
+      "epoch": 0.791398663135326,
+      "grad_norm": 2.4573779106140137,
+      "learning_rate": 1.0412054433148421e-05,
+      "loss": 1.6131498718261719,
+      "step": 12550
+    },
+    {
+      "epoch": 0.7945516458569807,
+      "grad_norm": 4.4481658935546875,
+      "learning_rate": 1.011075695579648e-05,
+      "loss": 1.4983843994140624,
+      "step": 12600
+    },
+    {
+      "epoch": 0.7977046285786354,
+      "grad_norm": 2.8198156356811523,
+      "learning_rate": 9.813392071633615e-06,
+      "loss": 1.6985472106933595,
+      "step": 12650
+    },
+    {
+      "epoch": 0.8008576113002901,
+      "grad_norm": 9.23430347442627,
+      "learning_rate": 9.519989097213094e-06,
+      "loss": 1.5060746765136719,
+      "step": 12700
+    },
+    {
+      "epoch": 0.8040105940219447,
+      "grad_norm": 2.075164794921875,
+      "learning_rate": 9.230576958492121e-06,
+      "loss": 1.5431832885742187,
+      "step": 12750
+    },
+    {
+      "epoch": 0.8071635767435994,
+      "grad_norm": 27.49349594116211,
+      "learning_rate": 8.945184187980233e-06,
+      "loss": 1.5060545349121093,
+      "step": 12800
+    },
+    {
+      "epoch": 0.8103165594652542,
+      "grad_norm": 9.251137733459473,
+      "learning_rate": 8.663838921926164e-06,
+      "loss": 1.4995553588867188,
+      "step": 12850
+    },
+    {
+      "epoch": 0.8134695421869088,
+      "grad_norm": 2.920060396194458,
+      "learning_rate": 8.38656889754409e-06,
+      "loss": 1.7685847473144531,
+      "step": 12900
+    },
+    {
+      "epoch": 0.8166225249085635,
+      "grad_norm": 2.4019455909729004,
+      "learning_rate": 8.113401450279019e-06,
+      "loss": 1.5964656066894531,
+      "step": 12950
+    },
+    {
+      "epoch": 0.8197755076302182,
+      "grad_norm": 8.358434677124023,
+      "learning_rate": 7.844363511111835e-06,
+      "loss": 1.489500732421875,
+      "step": 13000
+    },
+    {
+      "epoch": 0.8229284903518729,
+      "grad_norm": 3.5551395416259766,
+      "learning_rate": 7.579481603904282e-06,
+      "loss": 1.587592315673828,
+      "step": 13050
+    },
+    {
+      "epoch": 0.8260814730735275,
+      "grad_norm": 12.914484024047852,
+      "learning_rate": 7.318781842784e-06,
+      "loss": 1.5424423217773438,
+      "step": 13100
+    },
+    {
+      "epoch": 0.8292344557951823,
+      "grad_norm": 2.483072519302368,
+      "learning_rate": 7.0622899295699764e-06,
+      "loss": 1.6078681945800781,
+      "step": 13150
+    },
+    {
+      "epoch": 0.8323874385168369,
+      "grad_norm": 7.478156566619873,
+      "learning_rate": 6.810031151238733e-06,
+      "loss": 1.678245849609375,
+      "step": 13200
+    },
+    {
+      "epoch": 0.8355404212384916,
+      "grad_norm": 3.201509475708008,
+      "learning_rate": 6.5620303774312274e-06,
+      "loss": 1.5973834228515624,
+      "step": 13250
+    },
+    {
+      "epoch": 0.8386934039601462,
+      "grad_norm": 2.6549758911132812,
+      "learning_rate": 6.3183120580010745e-06,
+      "loss": 1.5435140991210938,
+      "step": 13300
+    },
+    {
+      "epoch": 0.841846386681801,
+      "grad_norm": 3.7113656997680664,
+      "learning_rate": 6.0789002206041444e-06,
+      "loss": 1.4877342224121093,
+      "step": 13350
+    },
+    {
+      "epoch": 0.8449993694034557,
+      "grad_norm": 7.131357669830322,
+      "learning_rate": 5.843818468329593e-06,
+      "loss": 1.6756939697265625,
+      "step": 13400
+    },
+    {
+      "epoch": 0.8481523521251103,
+      "grad_norm": 3.209019660949707,
+      "learning_rate": 5.613089977372993e-06,
+      "loss": 1.4503678894042968,
+      "step": 13450
+    },
+    {
+      "epoch": 0.8513053348467651,
+      "grad_norm": 3.1845173835754395,
+      "learning_rate": 5.38673749475142e-06,
+      "loss": 1.6333909606933594,
+      "step": 13500
+    },
+    {
+      "epoch": 0.8544583175684197,
+      "grad_norm": 4.371488571166992,
+      "learning_rate": 5.164783336060819e-06,
+      "loss": 1.4515260314941407,
+      "step": 13550
+    },
+    {
+      "epoch": 0.8576113002900744,
+      "grad_norm": 2.6283791065216064,
+      "learning_rate": 4.947249383276065e-06,
+      "loss": 1.6989479064941406,
+      "step": 13600
+    },
+    {
+      "epoch": 0.8607642830117291,
+      "grad_norm": 3.4831080436706543,
+      "learning_rate": 4.7341570825935485e-06,
+      "loss": 1.4910855102539062,
+      "step": 13650
+    },
+    {
+      "epoch": 0.8639172657333838,
+      "grad_norm": 3.8321712017059326,
+      "learning_rate": 4.525527442316929e-06,
+      "loss": 1.49341552734375,
+      "step": 13700
+    },
+    {
+      "epoch": 0.8670702484550384,
+      "grad_norm": 5.276515960693359,
+      "learning_rate": 4.321381030785942e-06,
+      "loss": 1.5483000183105469,
+      "step": 13750
+    },
+    {
+      "epoch": 0.8702232311766932,
+      "grad_norm": 3.14579176902771,
+      "learning_rate": 4.1217379743485885e-06,
+      "loss": 1.3947171020507811,
+      "step": 13800
+    },
+    {
+      "epoch": 0.8733762138983479,
+      "grad_norm": 2.654825448989868,
+      "learning_rate": 3.92661795537696e-06,
+      "loss": 1.4749700927734375,
+      "step": 13850
+    },
+    {
+      "epoch": 0.8765291966200025,
+      "grad_norm": 3.828669786453247,
+      "learning_rate": 3.736040210326791e-06,
+      "loss": 1.4993699645996095,
+      "step": 13900
+    },
+    {
+      "epoch": 0.8796821793416572,
+      "grad_norm": 2.1799724102020264,
+      "learning_rate": 3.55002352784094e-06,
+      "loss": 1.5674758911132813,
+      "step": 13950
+    },
+    {
+      "epoch": 0.8828351620633119,
+      "grad_norm": 4.14724588394165,
+      "learning_rate": 3.368586246897126e-06,
+      "loss": 1.534229736328125,
+      "step": 14000
+    },
+    {
+      "epoch": 0.8859881447849666,
+      "grad_norm": 5.016180992126465,
+      "learning_rate": 3.1917462549998556e-06,
+      "loss": 1.6074850463867187,
+      "step": 14050
+    },
+    {
+      "epoch": 0.8891411275066212,
+      "grad_norm": 1.8124960660934448,
+      "learning_rate": 3.0195209864169847e-06,
+      "loss": 1.395457763671875,
+      "step": 14100
+    },
+    {
+      "epoch": 0.892294110228276,
+      "grad_norm": 6.43141508102417,
+      "learning_rate": 2.8519274204609038e-06,
+      "loss": 1.4506192016601562,
+      "step": 14150
+    },
+    {
+      "epoch": 0.8954470929499306,
+      "grad_norm": 2.8766117095947266,
+      "learning_rate": 2.688982079814545e-06,
+      "loss": 1.5080227661132812,
+      "step": 14200
+    },
+    {
+      "epoch": 0.8986000756715853,
+      "grad_norm": 4.454990386962891,
+      "learning_rate": 2.5307010289025103e-06,
+      "loss": 1.5115219116210938,
+      "step": 14250
+    },
+    {
+      "epoch": 0.90175305839324,
+      "grad_norm": 2.8794169425964355,
+      "learning_rate": 2.3770998723072646e-06,
+      "loss": 1.5923085021972656,
+      "step": 14300
+    },
+    {
+      "epoch": 0.9049060411148947,
+      "grad_norm": 3.060925006866455,
+      "learning_rate": 2.2281937532307506e-06,
+      "loss": 1.5018788146972657,
+      "step": 14350
+    },
+    {
+      "epoch": 0.9080590238365494,
+      "grad_norm": 6.880637168884277,
+      "learning_rate": 2.08399735200141e-06,
+      "loss": 1.5022019958496093,
+      "step": 14400
+    },
+    {
+      "epoch": 0.9112120065582041,
+      "grad_norm": 4.857387065887451,
+      "learning_rate": 1.9445248846269194e-06,
+      "loss": 1.3896197509765624,
+      "step": 14450
+    },
+    {
+      "epoch": 0.9143649892798588,
+      "grad_norm": 7.964158058166504,
+      "learning_rate": 1.809790101392661e-06,
+      "loss": 1.6113026428222657,
+      "step": 14500
+    },
+    {
+      "epoch": 0.9175179720015134,
+      "grad_norm": 2.4885189533233643,
+      "learning_rate": 1.6798062855061024e-06,
+      "loss": 1.4576806640625,
+      "step": 14550
+    },
+    {
+      "epoch": 0.9206709547231681,
+      "grad_norm": 5.294617652893066,
+      "learning_rate": 1.5545862517872156e-06,
+      "loss": 1.538426513671875,
+      "step": 14600
+    },
+    {
+      "epoch": 0.9238239374448228,
+      "grad_norm": 3.588104248046875,
+      "learning_rate": 1.4341423454051373e-06,
+      "loss": 1.6244303894042968,
+      "step": 14650
+    },
+    {
+      "epoch": 0.9269769201664775,
+      "grad_norm": 3.241389274597168,
+      "learning_rate": 1.318486440661043e-06,
+      "loss": 1.5407508850097655,
+      "step": 14700
+    },
+    {
+      "epoch": 0.9301299028881321,
+      "grad_norm": 2.7573251724243164,
+      "learning_rate": 1.207629939817534e-06,
+      "loss": 1.5604298400878907,
+      "step": 14750
+    },
+    {
+      "epoch": 0.9332828856097869,
+      "grad_norm": 2.5889720916748047,
+      "learning_rate": 1.1015837719744348e-06,
+      "loss": 1.5485455322265624,
+      "step": 14800
+    },
+    {
+      "epoch": 0.9364358683314415,
+      "grad_norm": 3.1077330112457275,
+      "learning_rate": 1.0003583919914016e-06,
+      "loss": 1.617424774169922,
+      "step": 14850
+    },
+    {
+      "epoch": 0.9395888510530962,
+      "grad_norm": 4.1840314865112305,
+      "learning_rate": 9.039637794571509e-07,
+      "loss": 1.6456863403320312,
+      "step": 14900
+    },
+    {
+      "epoch": 0.942741833774751,
+      "grad_norm": 3.3467092514038086,
+      "learning_rate": 8.124094377056135e-07,
+      "loss": 1.4491183471679687,
+      "step": 14950
+    },
+    {
+      "epoch": 0.9458948164964056,
+      "grad_norm": 2.538663387298584,
+      "learning_rate": 7.25704392879012e-07,
+      "loss": 1.4048564147949218,
+      "step": 15000
+    },
+    {
+      "epoch": 0.9490477992180603,
+      "grad_norm": 4.097651481628418,
+      "learning_rate": 6.438571930379833e-07,
+      "loss": 1.527132568359375,
+      "step": 15050
+    },
+    {
+      "epoch": 0.952200781939715,
+      "grad_norm": 3.757591724395752,
+      "learning_rate": 5.668759073188756e-07,
+      "loss": 1.515183868408203,
+      "step": 15100
+    },
+    {
+      "epoch": 0.9553537646613697,
+      "grad_norm": 7.4409074783325195,
+      "learning_rate": 4.947681251382063e-07,
+      "loss": 1.5174710083007812,
+      "step": 15150
+    },
+    {
+      "epoch": 0.9585067473830243,
+      "grad_norm": 7.107412338256836,
+      "learning_rate": 4.2754095544442765e-07,
+      "loss": 1.3609051513671875,
+      "step": 15200
+    },
+    {
+      "epoch": 0.961659730104679,
+      "grad_norm": 4.804090976715088,
+      "learning_rate": 3.652010260171146e-07,
+      "loss": 1.5015805053710938,
+      "step": 15250
+    },
+    {
+      "epoch": 0.9648127128263337,
+      "grad_norm": 4.099542140960693,
+      "learning_rate": 3.0775448281348816e-07,
+      "loss": 1.6412095642089843,
+      "step": 15300
+    },
+    {
+      "epoch": 0.9679656955479884,
+      "grad_norm": 4.165719985961914,
+      "learning_rate": 2.5520698936254975e-07,
+      "loss": 1.6391523742675782,
+      "step": 15350
+    },
+    {
+      "epoch": 0.971118678269643,
+      "grad_norm": 2.3679826259613037,
+      "learning_rate": 2.0756372620669428e-07,
+      "loss": 1.549954376220703,
+      "step": 15400
+    },
+    {
+      "epoch": 0.9742716609912978,
+      "grad_norm": 3.1526825428009033,
+      "learning_rate": 1.6482939039099143e-07,
+      "loss": 1.702782440185547,
+      "step": 15450
+    },
+    {
+      "epoch": 0.9774246437129525,
+      "grad_norm": 4.235635280609131,
+      "learning_rate": 1.2700819500009454e-07,
+      "loss": 1.4687045288085938,
+      "step": 15500
+    },
+    {
+      "epoch": 0.9805776264346071,
+      "grad_norm": 5.9085822105407715,
+      "learning_rate": 9.410386874291188e-08,
+      "loss": 1.4969650268554688,
+      "step": 15550
+    },
+    {
+      "epoch": 0.9837306091562619,
+      "grad_norm": 3.3030614852905273,
+      "learning_rate": 6.611965558495637e-08,
+      "loss": 1.5178204345703126,
+      "step": 15600
+    },
+    {
+      "epoch": 0.9868835918779165,
+      "grad_norm": 2.6117329597473145,
+      "learning_rate": 4.305831442857344e-08,
+      "loss": 1.456710205078125,
+      "step": 15650
+    },
+    {
+      "epoch": 0.9900365745995712,
+      "grad_norm": 4.085411548614502,
+      "learning_rate": 2.4922118840936538e-08,
+      "loss": 1.658289031982422,
+      "step": 15700
+    },
+    {
+      "epoch": 0.9931895573212258,
+      "grad_norm": 3.8749728202819824,
+      "learning_rate": 1.1712856829870777e-08,
+      "loss": 1.5386245727539063,
+      "step": 15750
+    },
+    {
+      "epoch": 0.9963425400428806,
+      "grad_norm": 2.440901279449463,
+      "learning_rate": 3.4318306676217602e-09,
+      "loss": 1.6089360046386718,
+      "step": 15800
+    },
+    {
+      "epoch": 0.9994955227645352,
+      "grad_norm": 3.1938250064849854,
+      "learning_rate": 7.985676244159912e-11,
+      "loss": 1.4393968200683593,
+      "step": 15850
+    }
+  ],
+  "logging_steps": 50,
+  "max_steps": 15858,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 500,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 656938445174016.0,
+  "train_batch_size": 1,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

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