Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,141 @@
----
-license: mit
----
+---
+license: apache-2.0
+base_model: Qwen/Qwen3-4B
+tags:
+  - SAT
+  - combinatorial-optimization
+  - classification
+  - cube-and-conquer
+  - data-augmentation
+language:
+  - en
+pipeline_tag: text-classification
+---
+
+# Qwen3-4B-SAT-VarSelector-Sym-Aug
+
+A Qwen3-4B model fine-tuned for **SAT branching variable selection** using **symmetry-based data augmentation**.
+
+## Model Description
+
+This model predicts which variable to branch/cube on next, given a SAT CNF formula state. It was trained with **5x augmented data** using CNF symmetry transformations, achieving **21.8% top-1 accuracy** (vs 19% for Qwen3-0.6B).
+
+### Architecture
+
+- **Base**: `Qwen/Qwen3-4B` (causal language model)
+- **Head**: LayerNorm → Linear(hidden_size, 601)
+- **Max Variables**: 600
+- **Pooling**: Last non-pad token hidden state
+- **Masking**: Invalid variables (not in CNF) are masked to -10000 before softmax
+- **Size**: ~8GB (bfloat16)
+
+### Training with Symmetry Augmentation
+
+This model was trained with **5x data augmentation** using semantically-safe CNF transformations:
+
+| Augmentation | Description | Effect |
+|-------------|-------------|--------|
+| **Variable Permutation** | Bijective remapping of variable IDs | Prevents memorizing specific variable numbers |
+| **Clause Shuffling** | Random reordering of clauses | Teaches position-independence |
+| **Literal Reordering** | Shuffle literals within clauses | Token-level variation |
+| **Polarity Flipping** | Flip signs of random variable subset | Teaches structural vs. polarity features |
+
+### Training Details
+
+| Parameter | Value |
+|-----------|-------|
+| Original training samples | 8,110 |
+| Augmented training samples | **40,550** (5x) |
+| Validation samples | 902 (unaugmented) |
+| Epochs | 3 |
+| Hardware | 8×H100 GPUs |
+| Training framework | DeepSpeed ZeRO-3 |
+| Peak learning rate | 5e-6 |
+| Training time | ~4 hours |
+| Best checkpoint | Step 1850 (epoch 2.92) |
+
+### Performance Comparison
+
+| Model | Parameters | Training Data | Top-1 Accuracy |
+|-------|-----------|--------------|----------------|
+| Qwen3-0.6B (baseline) | 600M | 8,110 samples | ~12% |
+| Qwen3-0.6B (augmented) | 600M | 40,550 samples | ~19% |
+| **Qwen3-4B (augmented)** | **4B** | **40,550 samples** | **~22%** |
+
+### Training Curve Highlights
+
+- Peak accuracy: **22.0%** at epoch 2.76
+- Final accuracy: **21.8%** at epoch 2.92
+- Eval loss: 3.35 (vs 3.37 for 0.6B)
+
+## Usage
+
+```python
+import torch
+from transformers import AutoTokenizer
+from sft_qwen_var_classifier import QwenVarClassifier, cnf_valid_mask
+
+# Load model
+model = QwenVarClassifier("Qwen/Qwen3-4B", max_vars=600)
+state_dict = torch.load("pytorch_model.bin", map_location="cpu")
+model.load_state_dict(state_dict, strict=False)
+model = model.to("cuda", dtype=torch.bfloat16)
+model.eval()
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
+
+# Prepare CNF input
+cnf_text = """p cnf 100 250
+1 -2 3 0
+-1 2 -4 0
+...
+"""
+
+# Tokenize
+inputs = tokenizer(cnf_text, return_tensors="pt", truncation=True, max_length=8192)
+inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+# Get valid variable mask
+valid_mask = torch.tensor([cnf_valid_mask(cnf_text, max_vars=600)], dtype=torch.bool, device="cuda")
+
+# Predict
+with torch.no_grad():
+    outputs = model(**inputs)
+    logits = outputs["logits"]
+    logits = logits.masked_fill(~valid_mask, -1e4)
+    predicted_var = logits.argmax(dim=-1).item()
+
+print(f"Predicted branching variable: {predicted_var}")
+```
+
+## Files
+
+- `pytorch_model.bin` - Model weights (~8GB, bfloat16)
+- `sft_qwen_var_classifier.py` - Model class definition (required for loading)
+
+## When to Use This Model
+
+- **Higher accuracy** than 0.6B version (+3pp)
+- **Production use** when accuracy matters more than speed
+- **Cube-and-Conquer** style SAT solving
+
+## Limitations
+
+- Maximum 600 variables
+- Maximum 8192 tokens for CNF input
+- Larger model size (~8GB vs 1.2GB for 0.6B)
+- Slower inference (~6x slower than 0.6B)
+
+## Related Models
+
+- [Qwen3-0.6B-SAT-VarSelector-Sym-Aug](https://huggingface.co/Yale-ROSE/Qwen3-0.6B-SAT-VarSelector-Sym-Aug) - Smaller, faster version
+- [Qwen3-0.6B-SAT-VarSelector](https://huggingface.co/Yale-ROSE/Qwen3-0.6B-SAT-VarSelector) - Non-augmented baseline
+
+## Citation
+
+If you use this model, please cite the Transformer-CnC paper.
+
+## License
+
+Apache 2.0

added_tokens.json

@@ -0,0 +1,28 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

chat_template.jinja

@@ -0,0 +1,89 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

global_step1850/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_1_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_2_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_3_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_4_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_5_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_6_mp_rank_00_optim_states.pt

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

global_step1850/bf16_zero_pp_rank_7_mp_rank_00_optim_states.pt

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

global_step1850/zero_pp_rank_0_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_1_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_2_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_3_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_4_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_5_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_6_mp_rank_00_model_states.pt

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

global_step1850/zero_pp_rank_7_mp_rank_00_model_states.pt

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

latest

@@ -0,0 +1 @@
+global_step1850

pytorch_model.bin

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

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

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

rng_state_4.pth

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

rng_state_5.pth

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

rng_state_6.pth

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

rng_state_7.pth

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

scheduler.pt

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

sft_qwen_var_classifier.py

@@ -0,0 +1,728 @@
+"""
+Qwen Variable Classifier for SAT Cube-and-Conquer
+
+This script trains a transformer-based policy to select the next branching variable
+for SAT (Boolean Satisfiability) solving using the Cube-and-Conquer approach.
+
+== Problem Overview ==
+In Cube-and-Conquer SAT solving, we split a hard SAT problem into subproblems ("cubes")
+by choosing variables to branch on. The quality of variable selection significantly
+affects solving performance. This model learns to predict good branching variables
+from expert demonstrations.
+
+== Architecture ==
+- Backbone: Qwen3-4B (pretrained causal language model)
+- Head: LayerNorm + Linear classifier over variable IDs (1 to max_vars)
+- The model reads a CNF formula as text and outputs logits for each possible variable
+
+== Training Approach ==
+- Supervised Fine-Tuning (SFT) on expert variable choices
+- Masked classification: only variables appearing in the CNF are valid choices
+- Loss: Cross-entropy with invalid variable logits masked to -infinity
+
+== Data Format ==
+JSONL with fields:
+  - "cnf": DIMACS-format CNF text (e.g., "p cnf 100 200\n1 -2 3 0\n...")
+  - "label": integer variable ID to branch on (1 to max_vars)
+"""
+
+import os
+import argparse
+from dataclasses import dataclass
+from typing import Any, Dict, List
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from datasets import load_dataset
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
+    set_seed,
+)
+
+
+# =============================================================================
+# DEBUG FLAG: Set to True to enable verbose debug output, False to disable
+# Can also be controlled via environment variable: DEBUG_TRAINING=1
+# =============================================================================
+DEBUG_TRAINING = os.environ.get("DEBUG_TRAINING", "0") == "1"
+
+
+# =============================================================================
+# CNF PARSING: Extract valid variables from DIMACS CNF text
+# =============================================================================
+
+def cnf_valid_mask(cnf_text: str, max_vars: int) -> List[int]:
+    """
+    Build a binary mask indicating which variable IDs appear in the CNF.
+    
+    This is crucial for masked classification:
+    - A variable that doesn't appear in the (simplified) CNF cannot be branched on
+    - By masking invalid variables, we ensure the model only learns over valid choices
+    
+    Args:
+        cnf_text: DIMACS-format CNF string. Format example:
+            p cnf 100 200       # header: 100 variables, 200 clauses
+            1 -2 3 0            # clause: (x1 OR NOT x2 OR x3)
+            -1 4 0              # clause: (NOT x1 OR x4)
+            ...
+        max_vars: Maximum variable ID supported (typically 500)
+    
+    Returns:
+        List of length (max_vars + 1) where:
+        - mask[0] = 0 (unused, variables are 1-indexed)
+        - mask[v] = 1 if variable v appears in any clause
+        - mask[v] = 0 if variable v does not appear
+    
+    Note: We skip the header line "p cnf ..." to avoid capturing the clause count
+    as a valid variable (which was a bug in the original regex-based approach).
+    """
+    mask = [0] * (max_vars + 1)
+    
+    for line in cnf_text.split('\n'):
+        line = line.strip()
+        
+        # Skip empty lines, comment lines (start with 'c'), and header line (starts with 'p')
+        # The header "p cnf <num_vars> <num_clauses>" would incorrectly add num_clauses as a variable
+        if not line or line.startswith('c') or line.startswith('p'):
+            continue
+        
+        # Parse clause: space-separated integers ending with 0
+        # Each integer is a literal: positive = variable, negative = negated variable
+        # Example: "1 -2 3 0" means (x1 OR NOT x2 OR x3)
+        for tok in line.split():
+            try:
+                lit = int(tok)
+                v = abs(lit)  # Variable ID is absolute value of literal
+                if 1 <= v <= max_vars:
+                    mask[v] = 1
+            except ValueError:
+                continue  # Skip non-integer tokens (shouldn't happen in valid DIMACS)
+
+    # Fallback: if no variables found (e.g., truncated/malformed input), allow all
+    # This prevents the model from having zero valid outputs
+    if sum(mask) == 0:
+        for v in range(1, max_vars + 1):
+            mask[v] = 1
+    
+    return mask
+
+
+# =============================================================================
+# MODEL: Qwen backbone with classification head for variable selection
+# =============================================================================
+
+class QwenVarClassifier(nn.Module):
+    """
+    Transformer-based variable classifier for SAT branching.
+    
+    Architecture:
+        Input (CNF text) 
+            → Tokenize 
+            → Qwen3-4B backbone (frozen initially, fine-tuned with small LR)
+            → Extract last token's hidden state (sequence pooling)
+            → LayerNorm (stabilizes hidden state magnitude)
+            → Linear head (hidden_dim → num_classes)
+            → Logits for each variable ID
+    
+    Why this architecture?
+    1. Pretrained LLM backbone understands text structure and can learn CNF patterns
+    2. Last-token pooling: the final token has attended to the entire input
+    3. LayerNorm: Qwen's hidden states have large magnitudes; normalizing prevents
+       exploding gradients when combined with randomly-initialized head
+    4. Single linear head: simple, interpretable, efficient
+    """
+
+    def __init__(self, base_model_name: str, max_vars: int):
+        """
+        Initialize the classifier.
+        
+        Args:
+            base_model_name: HuggingFace model ID (e.g., "Qwen/Qwen3-4B")
+            max_vars: Maximum variable ID to classify (e.g., 500)
+                      Output dimension will be max_vars + 1 (index 0 unused)
+        """
+        super().__init__()
+        self.max_vars = max_vars
+        
+        # Load Qwen configuration and enable hidden state output
+        cfg = AutoConfig.from_pretrained(base_model_name)
+        cfg.output_hidden_states = True  # We need hidden states, not just logits
+        
+        # Load pretrained Qwen model
+        # Using bfloat16 for memory efficiency on modern GPUs (H100, A100)
+        self.backbone = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            config=cfg,
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        )
+        
+        hidden = self.backbone.config.hidden_size  # e.g., 2560 for Qwen3-4B
+        
+        # LayerNorm to normalize hidden states before classification
+        # This is critical for stable training:
+        # - Qwen's hidden states can have large magnitude (std >> 1)
+        # - Randomly initialized linear head expects normalized inputs
+        # - Without LayerNorm, initial logits can be huge → high loss → exploding gradients
+        self.head_ln = nn.LayerNorm(hidden)
+        
+        # Classification head: maps hidden state to variable logits
+        # Output shape: [batch, max_vars + 1]
+        # Index 0 is unused (variables are 1-indexed in DIMACS)
+        self.head = nn.Linear(hidden, max_vars + 1)
+        
+        # Initialize head with standard small weights
+        # LayerNorm ensures the input has unit variance, so this init is appropriate
+        nn.init.normal_(self.head.weight, std=0.02)
+        nn.init.zeros_(self.head.bias)
+        
+        # Expose backbone config for DeepSpeed compatibility
+        # DeepSpeed checks model.config.hidden_size for auto-configuration
+        self.config = self.backbone.config
+
+    def forward(self, input_ids, attention_mask, **kwargs):
+        """
+        Forward pass: CNF tokens → variable logits.
+        
+        Args:
+            input_ids: [batch, seq_len] token IDs from tokenizer
+            attention_mask: [batch, seq_len] binary mask (1 = real token, 0 = padding)
+            **kwargs: ignored (allows passing 'labels' without error during eval)
+        
+        Returns:
+            dict with "logits": [batch, max_vars + 1] raw classification logits
+        """
+        # Run through Qwen backbone
+        out = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=True,  # Need hidden states, not LM logits
+            use_cache=False,  # Disable KV cache (not needed for training)
+        )
+        
+        # Get hidden states from the last transformer layer
+        # Shape: [batch, seq_len, hidden_dim]
+        h = out.hidden_states[-1]
+
+        # Pool by taking the last non-padding token's hidden state
+        # This is the standard approach for causal LMs (like using [CLS] for BERT)
+        # 
+        # Why last token?
+        # - In causal attention, each token only sees previous tokens
+        # - The last token has attended to the entire input sequence
+        # - It's a natural "summary" of the input
+        #
+        # Compute index of last real token: sum of attention mask minus 1
+        last_idx = attention_mask.sum(dim=1) - 1  # [batch]
+        last_idx = last_idx.clamp(min=0)  # Safety: ensure non-negative
+        
+        # Gather hidden state at the last token position for each batch element
+        b = torch.arange(h.size(0), device=h.device)
+        pooled = h[b, last_idx]  # [batch, hidden_dim]
+        
+        # DEBUG: Check hidden state stats
+        if DEBUG_TRAINING:
+            if not hasattr(self, '_debug_count'):
+                self._debug_count = 0
+            if self._debug_count < 3:
+                print(f"[DEBUG {self._debug_count}] pooled dtype={pooled.dtype}, mean={pooled.float().mean():.2f}, std={pooled.float().std():.2f}")
+                self._debug_count += 1
+        
+        # Normalize hidden states for stable classification
+        pooled = self.head_ln(pooled)
+        
+        # DEBUG: Check after LayerNorm
+        if DEBUG_TRAINING and hasattr(self, '_debug_count') and self._debug_count <= 3:
+            print(f"[DEBUG] after LN: dtype={pooled.dtype}, mean={pooled.float().mean():.4f}, std={pooled.float().std():.4f}")
+        
+        # Project to variable logits
+        logits = self.head(pooled)  # [batch, max_vars + 1]
+        
+        # DEBUG: Check logits
+        if DEBUG_TRAINING and hasattr(self, '_debug_count') and self._debug_count <= 3:
+            print(f"[DEBUG] logits: dtype={logits.dtype}, mean={logits.float().mean():.2f}, std={logits.float().std():.2f}, min={logits.float().min():.2f}, max={logits.float().max():.2f}")
+        
+        return {"logits": logits}
+
+
+# =============================================================================
+# DATA COLLATOR: Batch preparation with padding and mask handling
+# =============================================================================
+
+@dataclass
+class Collator:
+    """
+    Custom data collator for variable classification.
+    
+    Responsibilities:
+    1. Pad variable-length token sequences to the same length within a batch
+    2. Stack labels and valid_mask tensors
+    3. Create proper attention masks for padded sequences
+    
+    Why custom collator?
+    - We have custom fields (valid_mask) that need special handling
+    - Standard HF collators don't know about our mask format
+    """
+    tokenizer: Any  # Tokenizer for padding configuration
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        """
+        Collate a list of examples into a batch.
+        
+        Args:
+            features: List of dicts, each with:
+                - input_ids: List[int] - token IDs
+                - attention_mask: List[int] - attention mask
+                - label: int - target variable ID
+                - valid_mask: List[int] - binary mask of valid variables
+        
+        Returns:
+            Dict with batched tensors:
+                - input_ids: [batch, max_seq_len]
+                - attention_mask: [batch, max_seq_len]
+                - labels: [batch]
+                - valid_mask: [batch, max_vars + 1]
+        """
+        # Convert to tensors
+        input_ids = [torch.tensor(f["input_ids"], dtype=torch.long) for f in features]
+        attention_mask = [torch.tensor(f["attention_mask"], dtype=torch.long) for f in features]
+        labels = torch.tensor([f["label"] for f in features], dtype=torch.long)
+        valid_mask = torch.tensor([f["valid_mask"] for f in features], dtype=torch.bool)
+        
+        # Pad sequences to same length within batch
+        # Using pad_sequence pads shorter sequences with padding_value
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, 
+            batch_first=True, 
+            padding_value=self.tokenizer.pad_token_id
+        )
+        attention_mask = torch.nn.utils.rnn.pad_sequence(
+            attention_mask, 
+            batch_first=True, 
+            padding_value=0  # Padding positions get 0 attention
+        )
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels,
+            "valid_mask": valid_mask,
+        }
+
+
+# =============================================================================
+# TRAINER: Custom loss computation with variable masking
+# =============================================================================
+
+class MaskedVarTrainer(Trainer):
+    """
+    Custom HuggingFace Trainer with masked cross-entropy loss.
+    
+    The key modification: before computing cross-entropy, we mask out logits
+    for invalid variables (those not appearing in the CNF). This ensures:
+    1. The model cannot predict invalid variables
+    2. No gradient flows to invalid variable logits
+    3. Training focuses only on distinguishing valid choices
+    
+    NOTE on displayed metrics:
+    - 'loss' shown by Trainer is summed across GPUs (loss × world_size)
+      We add 'true_loss' which is the actual per-sample loss
+    - 'grad_norm' is the L2 norm across ALL ~4B parameters BEFORE clipping
+      Values of 100-200 are normal for large models; it gets clipped to max_grad_norm
+    """
+
+    def __init__(self, *args, max_vars: int, **kwargs):
+        """
+        Args:
+            max_vars: Maximum variable ID (for sanity checking labels)
+            *args, **kwargs: Passed to parent Trainer
+        """
+        super().__init__(*args, **kwargs)
+        self.max_vars = max_vars
+        self._accumulated_loss = 0.0
+        self._loss_count = 0
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        """
+        Compute masked cross-entropy loss for variable classification.
+        
+        Algorithm:
+        1. Extract labels and valid_mask from inputs
+        2. Forward pass to get logits
+        3. Set logits for invalid variables to -inf (or -1e4 for bf16 stability)
+        4. Compute cross-entropy loss
+        
+        Args:
+            model: The QwenVarClassifier
+            inputs: Dict with input_ids, attention_mask, labels, valid_mask
+            return_outputs: If True, return (loss, outputs) tuple
+            num_items_in_batch: Unused (for API compatibility)
+        
+        Returns:
+            loss: Scalar loss value, or (loss, outputs) tuple if return_outputs=True
+        """
+        # Get labels and mask (don't pop - prediction_loop needs labels for compute_metrics)
+        labels = inputs.get("labels")          # [batch]
+        valid_mask = inputs.get("valid_mask")  # [batch, max_vars + 1] boolean
+        
+        # Remove from inputs for model.forward (which doesn't expect them)
+        model_inputs = {k: v for k, v in inputs.items() if k not in ["labels", "valid_mask"]}
+        
+        # Forward pass
+        outputs = model(**model_inputs)
+        logits = outputs["logits"]  # [batch, max_vars + 1]
+        
+        # DEBUG: Check if label is in valid_mask
+        if DEBUG_TRAINING:
+            if not hasattr(self, '_loss_debug_count'):
+                self._loss_debug_count = 0
+            if self._loss_debug_count < 5:
+                for i, (lbl, vmask) in enumerate(zip(labels, valid_mask)):
+                    label_in_mask = vmask[lbl].item()
+                    valid_count = vmask.sum().item()
+                    logit_at_label = logits[i, lbl].item()
+                    print(f"[LOSS DEBUG {self._loss_debug_count}] label={lbl.item()}, in_mask={label_in_mask}, valid_vars={valid_count}, logit_at_label={logit_at_label:.2f}")
+                self._loss_debug_count += 1
+        
+        # Mask invalid variables by setting their logits to a large negative value
+        # After softmax, these will have probability ≈ 0
+        # 
+        # Why -1e4 instead of -inf or -1e9?
+        # - bfloat16 has limited dynamic range
+        # - -1e9 can cause NaN issues when computing softmax/cross-entropy
+        # - -1e4 is small enough to give ~0 probability while staying numerically stable
+        logits = logits.masked_fill(~valid_mask.to(logits.device), -1e4)
+        
+        # Sanity check: labels must be valid variable IDs (1 to max_vars)
+        # This catches data bugs early
+        if torch.any(labels <= 0) or torch.any(labels > self.max_vars):
+            bad = labels[(labels <= 0) | (labels > self.max_vars)].detach().cpu().tolist()
+            raise ValueError(f"Out-of-range labels detected (showing up to 20): {bad[:20]}")
+        
+        # DEBUG: Check logit at label after masking
+        if DEBUG_TRAINING and hasattr(self, '_loss_debug_count') and self._loss_debug_count <= 5:
+            for i, lbl in enumerate(labels):
+                masked_logit = logits[i, lbl].item()
+                print(f"[LOSS DEBUG] after mask: logit_at_label={masked_logit:.2f}")
+        
+        # Standard cross-entropy loss
+        # PyTorch's cross_entropy expects logits, not probabilities
+        loss = F.cross_entropy(logits, labels.to(logits.device))
+        
+        # Track true loss for accurate logging
+        self._accumulated_loss += loss.item()
+        self._loss_count += 1
+        
+        # DEBUG: Print loss
+        if DEBUG_TRAINING and hasattr(self, '_loss_debug_count') and self._loss_debug_count <= 5:
+            print(f"[LOSS DEBUG] loss={loss.item():.2f}")
+        
+        # Return masked logits in outputs (so compute_metrics gets properly masked predictions)
+        masked_outputs = {"logits": logits}
+        return (loss, masked_outputs) if return_outputs else loss
+
+    def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None):
+        """
+        Override prediction_step to properly return loss and logits for evaluation.
+        
+        The default HF Trainer prediction_step doesn't work well with custom compute_loss,
+        so we implement our own that properly computes masked loss and returns logits.
+        """
+        model.eval()
+        
+        with torch.no_grad():
+            # Get labels and mask
+            labels = inputs.get("labels")
+            valid_mask = inputs.get("valid_mask")
+            
+            # Forward pass
+            model_inputs = {k: v for k, v in inputs.items() if k not in ["labels", "valid_mask"]}
+            outputs = model(**model_inputs)
+            logits = outputs["logits"]
+            
+            # Mask invalid variables
+            logits = logits.masked_fill(~valid_mask.to(logits.device), -1e4)
+            
+            # Compute loss
+            loss = F.cross_entropy(logits, labels.to(logits.device))
+        
+        # Return (loss, logits, labels) - this is what compute_metrics expects
+        return (loss, logits.detach(), labels.detach())
+
+    def log(self, logs: Dict[str, float], start_time: float = None) -> None:
+        """
+        Override log to add true_loss and ensure eval metrics are logged to W&B.
+        
+        The default 'loss' in HF Trainer is summed across GPUs in DDP/DeepSpeed.
+        We track the actual per-sample loss and report it as 'true_loss'.
+        """
+        if self._loss_count > 0:
+            # Calculate true average loss on this device
+            true_loss = self._accumulated_loss / self._loss_count
+            logs["true_loss"] = round(true_loss, 4)
+            
+            # Reset for next logging interval
+            self._accumulated_loss = 0.0
+            self._loss_count = 0
+        
+        # Let HF Trainer handle W&B logging - it manages step ordering correctly
+        super().log(logs, start_time)
+
+
+def compute_metrics(eval_pred):
+    """
+    Compute accuracy for evaluation.
+    
+    Args:
+        eval_pred: (logits, labels) from Trainer's prediction_loop
+            - logits: [num_samples, max_vars + 1] (already masked with -1e4 for invalid vars)
+            - labels: [num_samples]
+    
+    Returns:
+        Dict with "accuracy" (Trainer will prefix with "eval_")
+    
+    Note: eval_loss is computed automatically by Trainer from prediction_step's loss.
+    We don't need to compute it here.
+    
+    Since invalid variables have logits ≈ -1e4, argmax will naturally avoid them.
+    """
+    logits, labels = eval_pred
+    
+    # Accuracy: argmax prediction vs true label
+    preds = np.argmax(logits, axis=-1)
+    accuracy = float((preds == labels).mean())
+    
+    return {"accuracy": accuracy}
+
+
+def get_wandb_report_to():
+    """
+    Determine if this process should log to W&B.
+    
+    Only the main process (rank 0) should log to W&B to avoid creating multiple runs.
+    Other ranks should not log to any external service.
+    
+    Returns:
+        ["wandb"] for rank 0, [] for other ranks
+    """
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    
+    if local_rank == 0:
+        return ["wandb"]
+    else:
+        return []
+
+
+# =============================================================================
+# MAIN: Training pipeline
+# =============================================================================
+    
+def main():
+    """
+    Main training function.
+    
+    Pipeline:
+    1. Parse command line arguments
+    2. Load tokenizer and datasets
+    3. Preprocess: tokenize CNF text, compute valid masks
+    4. Initialize model with pretrained backbone + new classification head
+    5. Configure training (optimizer, scheduler, logging, etc.)
+    6. Train and evaluate
+    """
+    ap = argparse.ArgumentParser(
+        description="Train a Qwen-based variable classifier for SAT branching"
+    )
+    
+    # Model and data arguments
+    ap.add_argument("--model_name", type=str, default="Qwen/Qwen3-4B",
+                    help="HuggingFace model ID for the backbone")
+    ap.add_argument("--train_jsonl", type=str, required=True,
+                    help="Path to training data (JSONL with 'cnf' and 'label' fields)")
+    ap.add_argument("--valid_jsonl", type=str, required=True,
+                    help="Path to validation data (same format)")
+    ap.add_argument("--output_dir", type=str, default="./out_qwen_var_sft",
+                    help="Directory for checkpoints and logs")
+    ap.add_argument("--max_vars", type=int, default=500,
+                    help="Maximum variable ID (determines output dimension)")
+    ap.add_argument("--max_length", type=int, default=8192,
+                    help="Maximum sequence length in tokens (truncates longer CNFs)")
+    ap.add_argument("--seed", type=int, default=0,
+                    help="Random seed for reproducibility")
+    
+    # Training hyperparameters
+    ap.add_argument("--per_device_train_batch_size", type=int, default=1,
+                    help="Batch size per GPU for training")
+    ap.add_argument("--per_device_eval_batch_size", type=int, default=1,
+                    help="Batch size per GPU for evaluation")
+    ap.add_argument("--gradient_accumulation_steps", type=int, default=8,
+                    help="Accumulate gradients over this many steps (effective batch = this * batch_size * num_gpus)")
+    ap.add_argument("--learning_rate", type=float, default=5e-6,
+                    help="Peak learning rate (after warmup). Lower than typical fine-tuning due to classification head")
+    ap.add_argument("--num_train_epochs", type=float, default=3.0,
+                    help="Total training epochs")
+    ap.add_argument("--warmup_ratio", type=float, default=0.03,
+                    help="Fraction of training steps for learning rate warmup")
+    ap.add_argument("--weight_decay", type=float, default=0.0,
+                    help="Weight decay (L2 regularization)")
+    ap.add_argument("--logging_steps", type=int, default=10,
+                    help="Log training metrics every N steps")
+    ap.add_argument("--eval_steps", type=int, default=200,
+                    help="Evaluate every N steps")
+    ap.add_argument("--save_steps", type=int, default=200,
+                    help="Save checkpoint every N steps")
+    ap.add_argument("--report_to", type=str, default="wandb", 
+                    choices=["wandb", "tensorboard", "none"],
+                    help="Logging backend")
+    ap.add_argument("--deepspeed", type=str, default=None,
+                    help="Path to DeepSpeed config JSON for distributed training")
+    ap.add_argument("--resume_from_checkpoint", type=str, default=None,
+                    help="Path to checkpoint directory to resume from. If a directory is given, "
+                         "the latest checkpoint in that directory will be used.")
+    
+    args = ap.parse_args()
+
+    # Set random seeds for reproducibility
+    set_seed(args.seed)
+
+    # Load tokenizer
+    # Qwen uses a byte-level BPE tokenizer
+    tok = AutoTokenizer.from_pretrained(args.model_name, use_fast=True)
+    if tok.pad_token is None:
+        # Qwen doesn't have a dedicated pad token; use eos as pad
+        tok.pad_token = tok.eos_token
+
+    # Load datasets from JSONL files
+    ds = load_dataset(
+        "json",
+        data_files={"train": args.train_jsonl, "validation": args.valid_jsonl},
+    )
+
+    def preprocess(ex):
+        """
+        Preprocess a single example.
+        
+        Steps:
+        1. Tokenize the CNF text
+        2. Compute valid variable mask
+        3. Return features for training
+        
+        Args:
+            ex: Dict with 'cnf' (str) and 'label' (int)
+        
+        Returns:
+            Dict with input_ids, attention_mask, label, valid_mask
+        """
+        cnf = ex["cnf"]
+        label = int(ex["label"])
+        
+        # Tokenize CNF text
+        # No special prompt/instruction - the model learns to interpret raw CNF
+        enc = tok(
+            cnf, 
+            truncation=True, 
+            max_length=args.max_length, 
+            padding=False  # We handle padding in the collator
+        )
+        
+        return {
+            "input_ids": enc["input_ids"],
+            "attention_mask": enc["attention_mask"],
+            "label": label,
+            "valid_mask": cnf_valid_mask(cnf, args.max_vars),
+        }
+    
+    # Apply preprocessing to all examples
+    # remove_columns drops original fields (cnf, label) since we've extracted what we need
+    ds = ds.map(preprocess, remove_columns=ds["train"].column_names)
+    
+    # Initialize model
+    model = QwenVarClassifier(args.model_name, max_vars=args.max_vars)
+    
+    # Enable gradient checkpointing to save memory on long sequences
+    # This trades compute for memory by recomputing activations during backward pass
+    model.backbone.gradient_checkpointing_enable()
+    
+    # Configure W&B logging (only rank 0 logs to avoid duplicate runs)
+    report_to = get_wandb_report_to()
+
+    # Configure training
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        overwrite_output_dir=True,
+        
+        # Precision settings for modern GPUs
+        bf16=True,   # Use bfloat16 for training (good for H100/A100)
+        tf32=True,   # Enable TF32 for faster matmuls on Ampere+
+        
+        # Batch configuration
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        
+        # Optimizer settings
+        learning_rate=args.learning_rate,
+        warmup_ratio=args.warmup_ratio,
+        num_train_epochs=args.num_train_epochs,
+        weight_decay=args.weight_decay,
+        
+        # Gradient clipping for training stability
+        # Clips gradient norm to this value if it exceeds it
+        # This prevents exploding gradients from destabilizing training
+        max_grad_norm=1.0,
+        
+        # Logging and evaluation
+        logging_steps=args.logging_steps,
+        eval_strategy="steps",
+        eval_steps=args.eval_steps,
+        
+        # Checkpointing - keep best checkpoints based on validation accuracy
+        save_strategy="steps",
+        save_steps=args.save_steps,
+        save_total_limit=3,  # Keep best 3 checkpoints
+        load_best_model_at_end=True,  # Load best checkpoint at end of training
+        metric_for_best_model="eval_accuracy",  # Use validation accuracy to determine best
+        greater_is_better=True,  # Higher accuracy is better
+        
+        # Logging backend
+        report_to=report_to,
+        run_name=os.environ.get("WANDB_RUN_NAME", "qwen-var-sft") if args.report_to == "wandb" else None,
+        logging_dir=os.path.join(args.output_dir, "logs"),
+        
+        # Important: don't remove valid_mask column (we need it in compute_loss)
+        remove_unused_columns=False,
+        
+        # DDP settings (for multi-GPU)
+        ddp_find_unused_parameters=False,
+        
+        # DeepSpeed for efficient distributed training
+        deepspeed=args.deepspeed,
+        
+        # Use pickle format for saving (safetensors has issues with some weight tying configs)
+        save_safetensors=False,
+    )
+
+    # Create trainer with custom loss computation
+    trainer = MaskedVarTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=ds["train"],
+        eval_dataset=ds["validation"],
+        tokenizer=tok,
+        data_collator=Collator(tok),
+        compute_metrics=compute_metrics,
+        max_vars=args.max_vars,
+    )
+
+    # Train! (resume from checkpoint if specified)
+    trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
+    
+    # Final evaluation
+    trainer.evaluate()
+
+
+if __name__ == "__main__":
+    main()

special_tokens_map.json

@@ -0,0 +1,31 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

tokenizer_config.json

@@ -0,0 +1,239 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

trainer_state.json

@@ -0,0 +1,1847 @@
+{
+  "best_global_step": 1750,
+  "best_metric": 0.21951219512195122,
+  "best_model_checkpoint": "out_qwen_4b_sft_augmented/checkpoint-1750",
+  "epoch": 2.9185243637798384,
+  "eval_steps": 50,
+  "global_step": 1850,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.015782205563227462,
+      "grad_norm": 324.9078480271927,
+      "learning_rate": 2.3560209424083772e-07,
+      "loss": 36.8443,
+      "step": 10,
+      "true_loss": 4.5233
+    },
+    {
+      "epoch": 0.031564411126454923,
+      "grad_norm": 168.6194251999255,
+      "learning_rate": 4.973821989528796e-07,
+      "loss": 36.0645,
+      "step": 20,
+      "true_loss": 4.3899
+    },
+    {
+      "epoch": 0.04734661668968238,
+      "grad_norm": 192.6442003874084,
+      "learning_rate": 7.591623036649215e-07,
+      "loss": 37.1343,
+      "step": 30,
+      "true_loss": 4.7917
+    },
+    {
+      "epoch": 0.06312882225290985,
+      "grad_norm": 169.69405108448686,
+      "learning_rate": 1.0209424083769635e-06,
+      "loss": 35.8286,
+      "step": 40,
+      "true_loss": 4.5888
+    },
+    {
+      "epoch": 0.0789110278161373,
+      "grad_norm": 177.6752080525516,
+      "learning_rate": 1.2827225130890052e-06,
+      "loss": 35.1516,
+      "step": 50,
+      "true_loss": 4.1833
+    },
+    {
+      "epoch": 0.0789110278161373,
+      "eval_accuracy": 0.02328159645232816,
+      "eval_loss": 4.40457820892334,
+      "eval_runtime": 23.7263,
+      "eval_samples_per_second": 38.017,
+      "eval_steps_per_second": 4.763,
+      "step": 50
+    },
+    {
+      "epoch": 0.09469323337936476,
+      "grad_norm": 175.29116271759068,
+      "learning_rate": 1.5445026178010472e-06,
+      "loss": 34.7314,
+      "step": 60,
+      "true_loss": 4.224
+    },
+    {
+      "epoch": 0.11047543894259222,
+      "grad_norm": 753.0264561952334,
+      "learning_rate": 1.8062827225130891e-06,
+      "loss": 33.9206,
+      "step": 70,
+      "true_loss": 4.0565
+    },
+    {
+      "epoch": 0.1262576445058197,
+      "grad_norm": 200.47588314223879,
+      "learning_rate": 2.068062827225131e-06,
+      "loss": 34.0037,
+      "step": 80,
+      "true_loss": 4.3621
+    },
+    {
+      "epoch": 0.14203985006904715,
+      "grad_norm": 164.18077066091436,
+      "learning_rate": 2.329842931937173e-06,
+      "loss": 33.3705,
+      "step": 90,
+      "true_loss": 4.1432
+    },
+    {
+      "epoch": 0.1578220556322746,
+      "grad_norm": 144.60587833889255,
+      "learning_rate": 2.591623036649215e-06,
+      "loss": 33.3276,
+      "step": 100,
+      "true_loss": 4.1424
+    },
+    {
+      "epoch": 0.1578220556322746,
+      "eval_accuracy": 0.043237250554323724,
+      "eval_loss": 4.096348285675049,
+      "eval_runtime": 23.3422,
+      "eval_samples_per_second": 38.643,
+      "eval_steps_per_second": 4.841,
+      "step": 100
+    },
+    {
+      "epoch": 0.17360426119550207,
+      "grad_norm": 270.635578631285,
+      "learning_rate": 2.853403141361257e-06,
+      "loss": 32.6581,
+      "step": 110,
+      "true_loss": 4.1497
+    },
+    {
+      "epoch": 0.18938646675872953,
+      "grad_norm": 100.33088317059614,
+      "learning_rate": 3.115183246073299e-06,
+      "loss": 32.8629,
+      "step": 120,
+      "true_loss": 4.1273
+    },
+    {
+      "epoch": 0.20516867232195699,
+      "grad_norm": 108.96506382691295,
+      "learning_rate": 3.3769633507853404e-06,
+      "loss": 33.0279,
+      "step": 130,
+      "true_loss": 4.1729
+    },
+    {
+      "epoch": 0.22095087788518444,
+      "grad_norm": 101.57985399798042,
+      "learning_rate": 3.6387434554973826e-06,
+      "loss": 33.3713,
+      "step": 140,
+      "true_loss": 4.2379
+    },
+    {
+      "epoch": 0.2367330834484119,
+      "grad_norm": 109.81214136700734,
+      "learning_rate": 3.900523560209425e-06,
+      "loss": 33.2539,
+      "step": 150,
+      "true_loss": 4.388
+    },
+    {
+      "epoch": 0.2367330834484119,
+      "eval_accuracy": 0.04434589800443459,
+      "eval_loss": 4.049750328063965,
+      "eval_runtime": 23.5037,
+      "eval_samples_per_second": 38.377,
+      "eval_steps_per_second": 4.808,
+      "step": 150
+    },
+    {
+      "epoch": 0.2525152890116394,
+      "grad_norm": 131.31761293379333,
+      "learning_rate": 4.1623036649214665e-06,
+      "loss": 32.8656,
+      "step": 160,
+      "true_loss": 4.1088
+    },
+    {
+      "epoch": 0.26829749457486685,
+      "grad_norm": 93.21133961816744,
+      "learning_rate": 4.424083769633508e-06,
+      "loss": 32.7579,
+      "step": 170,
+      "true_loss": 4.092
+    },
+    {
+      "epoch": 0.2840797001380943,
+      "grad_norm": 99.10049214777489,
+      "learning_rate": 4.68586387434555e-06,
+      "loss": 32.777,
+      "step": 180,
+      "true_loss": 4.1288
+    },
+    {
+      "epoch": 0.29986190570132176,
+      "grad_norm": 103.48370344024212,
+      "learning_rate": 4.947643979057592e-06,
+      "loss": 32.983,
+      "step": 190,
+      "true_loss": 3.9901
+    },
+    {
+      "epoch": 0.3156441112645492,
+      "grad_norm": 101.42410788589872,
+      "learning_rate": 4.976621858562245e-06,
+      "loss": 32.8281,
+      "step": 200,
+      "true_loss": 4.1047
+    },
+    {
+      "epoch": 0.3156441112645492,
+      "eval_accuracy": 0.06873614190687362,
+      "eval_loss": 4.009181022644043,
+      "eval_runtime": 23.3653,
+      "eval_samples_per_second": 38.604,
+      "eval_steps_per_second": 4.836,
+      "step": 200
+    },
+    {
+      "epoch": 0.3314263168277767,
+      "grad_norm": 83.2901468921201,
+      "learning_rate": 4.94739918176505e-06,
+      "loss": 33.1028,
+      "step": 210,
+      "true_loss": 4.0052
+    },
+    {
+      "epoch": 0.34720852239100414,
+      "grad_norm": 170.23336077552057,
+      "learning_rate": 4.9181765049678555e-06,
+      "loss": 32.644,
+      "step": 220,
+      "true_loss": 4.1984
+    },
+    {
+      "epoch": 0.3629907279542316,
+      "grad_norm": 73.10967930914595,
+      "learning_rate": 4.888953828170661e-06,
+      "loss": 32.9998,
+      "step": 230,
+      "true_loss": 4.1035
+    },
+    {
+      "epoch": 0.37877293351745905,
+      "grad_norm": 75.66187004852954,
+      "learning_rate": 4.859731151373466e-06,
+      "loss": 32.6017,
+      "step": 240,
+      "true_loss": 3.9774
+    },
+    {
+      "epoch": 0.3945551390806865,
+      "grad_norm": 101.5303291876841,
+      "learning_rate": 4.830508474576272e-06,
+      "loss": 32.5215,
+      "step": 250,
+      "true_loss": 4.107
+    },
+    {
+      "epoch": 0.3945551390806865,
+      "eval_accuracy": 0.07095343680709534,
+      "eval_loss": 3.933732271194458,
+      "eval_runtime": 23.3822,
+      "eval_samples_per_second": 38.576,
+      "eval_steps_per_second": 4.833,
+      "step": 250
+    },
+    {
+      "epoch": 0.41033734464391397,
+      "grad_norm": 87.68390524050062,
+      "learning_rate": 4.801285797779077e-06,
+      "loss": 32.8537,
+      "step": 260,
+      "true_loss": 4.0852
+    },
+    {
+      "epoch": 0.42611955020714143,
+      "grad_norm": 83.56798040873913,
+      "learning_rate": 4.772063120981883e-06,
+      "loss": 32.452,
+      "step": 270,
+      "true_loss": 4.1138
+    },
+    {
+      "epoch": 0.4419017557703689,
+      "grad_norm": 76.69106361052454,
+      "learning_rate": 4.742840444184687e-06,
+      "loss": 31.8737,
+      "step": 280,
+      "true_loss": 3.8671
+    },
+    {
+      "epoch": 0.45768396133359635,
+      "grad_norm": 79.70560950668154,
+      "learning_rate": 4.713617767387494e-06,
+      "loss": 32.1678,
+      "step": 290,
+      "true_loss": 3.9949
+    },
+    {
+      "epoch": 0.4734661668968238,
+      "grad_norm": 77.14726547694464,
+      "learning_rate": 4.684395090590298e-06,
+      "loss": 32.2324,
+      "step": 300,
+      "true_loss": 4.1541
+    },
+    {
+      "epoch": 0.4734661668968238,
+      "eval_accuracy": 0.07206208425720621,
+      "eval_loss": 3.944307804107666,
+      "eval_runtime": 24.0666,
+      "eval_samples_per_second": 37.479,
+      "eval_steps_per_second": 4.695,
+      "step": 300
+    },
+    {
+      "epoch": 0.4892483724600513,
+      "grad_norm": 74.39357075334843,
+      "learning_rate": 4.655172413793104e-06,
+      "loss": 32.7538,
+      "step": 310,
+      "true_loss": 3.9691
+    },
+    {
+      "epoch": 0.5050305780232788,
+      "grad_norm": 69.2140729496397,
+      "learning_rate": 4.625949736995909e-06,
+      "loss": 32.241,
+      "step": 320,
+      "true_loss": 4.1026
+    },
+    {
+      "epoch": 0.5208127835865062,
+      "grad_norm": 76.34039893901215,
+      "learning_rate": 4.596727060198715e-06,
+      "loss": 32.7047,
+      "step": 330,
+      "true_loss": 4.2664
+    },
+    {
+      "epoch": 0.5365949891497337,
+      "grad_norm": 76.48584051124504,
+      "learning_rate": 4.56750438340152e-06,
+      "loss": 32.2899,
+      "step": 340,
+      "true_loss": 4.1581
+    },
+    {
+      "epoch": 0.5523771947129611,
+      "grad_norm": 75.38836139042635,
+      "learning_rate": 4.5382817066043256e-06,
+      "loss": 32.8474,
+      "step": 350,
+      "true_loss": 4.0243
+    },
+    {
+      "epoch": 0.5523771947129611,
+      "eval_accuracy": 0.07317073170731707,
+      "eval_loss": 3.8937127590179443,
+      "eval_runtime": 23.6141,
+      "eval_samples_per_second": 38.198,
+      "eval_steps_per_second": 4.785,
+      "step": 350
+    },
+    {
+      "epoch": 0.5681594002761886,
+      "grad_norm": 85.00115390444991,
+      "learning_rate": 4.509059029807131e-06,
+      "loss": 31.7532,
+      "step": 360,
+      "true_loss": 3.988
+    },
+    {
+      "epoch": 0.583941605839416,
+      "grad_norm": 76.08205167618928,
+      "learning_rate": 4.479836353009936e-06,
+      "loss": 31.7753,
+      "step": 370,
+      "true_loss": 3.9777
+    },
+    {
+      "epoch": 0.5997238114026435,
+      "grad_norm": 68.74246027512453,
+      "learning_rate": 4.450613676212742e-06,
+      "loss": 31.4578,
+      "step": 380,
+      "true_loss": 3.9916
+    },
+    {
+      "epoch": 0.6155060169658709,
+      "grad_norm": 71.81332443942438,
+      "learning_rate": 4.4213909994155465e-06,
+      "loss": 32.0656,
+      "step": 390,
+      "true_loss": 4.1254
+    },
+    {
+      "epoch": 0.6312882225290984,
+      "grad_norm": 71.77472037885094,
+      "learning_rate": 4.392168322618352e-06,
+      "loss": 32.2929,
+      "step": 400,
+      "true_loss": 3.8869
+    },
+    {
+      "epoch": 0.6312882225290984,
+      "eval_accuracy": 0.1130820399113082,
+      "eval_loss": 3.7887723445892334,
+      "eval_runtime": 23.904,
+      "eval_samples_per_second": 37.734,
+      "eval_steps_per_second": 4.727,
+      "step": 400
+    },
+    {
+      "epoch": 0.647070428092326,
+      "grad_norm": 71.4147738866297,
+      "learning_rate": 4.3629456458211574e-06,
+      "loss": 31.9759,
+      "step": 410,
+      "true_loss": 3.9893
+    },
+    {
+      "epoch": 0.6628526336555534,
+      "grad_norm": 73.99693807240743,
+      "learning_rate": 4.333722969023963e-06,
+      "loss": 31.8449,
+      "step": 420,
+      "true_loss": 3.9648
+    },
+    {
+      "epoch": 0.6786348392187809,
+      "grad_norm": 74.124412349921,
+      "learning_rate": 4.304500292226768e-06,
+      "loss": 31.9984,
+      "step": 430,
+      "true_loss": 3.8608
+    },
+    {
+      "epoch": 0.6944170447820083,
+      "grad_norm": 80.24673254296576,
+      "learning_rate": 4.275277615429574e-06,
+      "loss": 31.8185,
+      "step": 440,
+      "true_loss": 4.033
+    },
+    {
+      "epoch": 0.7101992503452358,
+      "grad_norm": 83.65131551190922,
+      "learning_rate": 4.246054938632379e-06,
+      "loss": 31.7604,
+      "step": 450,
+      "true_loss": 3.9851
+    },
+    {
+      "epoch": 0.7101992503452358,
+      "eval_accuracy": 0.11862527716186252,
+      "eval_loss": 3.731326103210449,
+      "eval_runtime": 23.8861,
+      "eval_samples_per_second": 37.762,
+      "eval_steps_per_second": 4.731,
+      "step": 450
+    },
+    {
+      "epoch": 0.7259814559084632,
+      "grad_norm": 82.91785549569512,
+      "learning_rate": 4.216832261835184e-06,
+      "loss": 31.9212,
+      "step": 460,
+      "true_loss": 4.0262
+    },
+    {
+      "epoch": 0.7417636614716907,
+      "grad_norm": 77.3700837190592,
+      "learning_rate": 4.18760958503799e-06,
+      "loss": 31.558,
+      "step": 470,
+      "true_loss": 3.9037
+    },
+    {
+      "epoch": 0.7575458670349181,
+      "grad_norm": 79.2567711758356,
+      "learning_rate": 4.158386908240795e-06,
+      "loss": 31.9821,
+      "step": 480,
+      "true_loss": 3.9674
+    },
+    {
+      "epoch": 0.7733280725981456,
+      "grad_norm": 79.15209026205977,
+      "learning_rate": 4.1291642314436e-06,
+      "loss": 31.6287,
+      "step": 490,
+      "true_loss": 4.0236
+    },
+    {
+      "epoch": 0.789110278161373,
+      "grad_norm": 82.91621808879624,
+      "learning_rate": 4.0999415546464065e-06,
+      "loss": 31.1936,
+      "step": 500,
+      "true_loss": 3.8798
+    },
+    {
+      "epoch": 0.789110278161373,
+      "eval_accuracy": 0.1419068736141907,
+      "eval_loss": 3.6688694953918457,
+      "eval_runtime": 23.5535,
+      "eval_samples_per_second": 38.296,
+      "eval_steps_per_second": 4.798,
+      "step": 500
+    },
+    {
+      "epoch": 0.8048924837246005,
+      "grad_norm": 81.33416413304444,
+      "learning_rate": 4.070718877849211e-06,
+      "loss": 31.9352,
+      "step": 510,
+      "true_loss": 3.9601
+    },
+    {
+      "epoch": 0.8206746892878279,
+      "grad_norm": 84.38809367146038,
+      "learning_rate": 4.0414962010520166e-06,
+      "loss": 31.4177,
+      "step": 520,
+      "true_loss": 4.0845
+    },
+    {
+      "epoch": 0.8364568948510555,
+      "grad_norm": 71.30053276240011,
+      "learning_rate": 4.012273524254822e-06,
+      "loss": 31.7603,
+      "step": 530,
+      "true_loss": 3.997
+    },
+    {
+      "epoch": 0.8522391004142829,
+      "grad_norm": 80.60566610130194,
+      "learning_rate": 3.9830508474576275e-06,
+      "loss": 31.2764,
+      "step": 540,
+      "true_loss": 3.8168
+    },
+    {
+      "epoch": 0.8680213059775104,
+      "grad_norm": 91.53898913592603,
+      "learning_rate": 3.953828170660433e-06,
+      "loss": 31.4915,
+      "step": 550,
+      "true_loss": 3.9316
+    },
+    {
+      "epoch": 0.8680213059775104,
+      "eval_accuracy": 0.1419068736141907,
+      "eval_loss": 3.6105031967163086,
+      "eval_runtime": 23.4993,
+      "eval_samples_per_second": 38.384,
+      "eval_steps_per_second": 4.809,
+      "step": 550
+    },
+    {
+      "epoch": 0.8838035115407378,
+      "grad_norm": 80.87596019476308,
+      "learning_rate": 3.924605493863238e-06,
+      "loss": 31.1347,
+      "step": 560,
+      "true_loss": 3.7968
+    },
+    {
+      "epoch": 0.8995857171039653,
+      "grad_norm": 80.68070986181071,
+      "learning_rate": 3.895382817066044e-06,
+      "loss": 31.7808,
+      "step": 570,
+      "true_loss": 3.9588
+    },
+    {
+      "epoch": 0.9153679226671927,
+      "grad_norm": 82.14608732630818,
+      "learning_rate": 3.8661601402688484e-06,
+      "loss": 31.0608,
+      "step": 580,
+      "true_loss": 3.9779
+    },
+    {
+      "epoch": 0.9311501282304202,
+      "grad_norm": 79.11284293573352,
+      "learning_rate": 3.836937463471655e-06,
+      "loss": 30.9633,
+      "step": 590,
+      "true_loss": 4.0027
+    },
+    {
+      "epoch": 0.9469323337936476,
+      "grad_norm": 74.83038134855214,
+      "learning_rate": 3.8077147866744598e-06,
+      "loss": 31.2218,
+      "step": 600,
+      "true_loss": 3.8863
+    },
+    {
+      "epoch": 0.9469323337936476,
+      "eval_accuracy": 0.1419068736141907,
+      "eval_loss": 3.6132075786590576,
+      "eval_runtime": 24.0285,
+      "eval_samples_per_second": 37.539,
+      "eval_steps_per_second": 4.703,
+      "step": 600
+    },
+    {
+      "epoch": 0.9627145393568751,
+      "grad_norm": 77.6288128447819,
+      "learning_rate": 3.7784921098772652e-06,
+      "loss": 31.104,
+      "step": 610,
+      "true_loss": 3.8553
+    },
+    {
+      "epoch": 0.9784967449201026,
+      "grad_norm": 78.73674435660126,
+      "learning_rate": 3.7492694330800707e-06,
+      "loss": 31.4592,
+      "step": 620,
+      "true_loss": 3.8146
+    },
+    {
+      "epoch": 0.99427895048333,
+      "grad_norm": 84.19142784542193,
+      "learning_rate": 3.7200467562828757e-06,
+      "loss": 31.4694,
+      "step": 630,
+      "true_loss": 3.8633
+    },
+    {
+      "epoch": 1.0094693233379364,
+      "grad_norm": 83.97120538549059,
+      "learning_rate": 3.6908240794856816e-06,
+      "loss": 29.9569,
+      "step": 640,
+      "true_loss": 3.9506
+    },
+    {
+      "epoch": 1.0252515289011639,
+      "grad_norm": 78.03329090311782,
+      "learning_rate": 3.6616014026884866e-06,
+      "loss": 30.6145,
+      "step": 650,
+      "true_loss": 3.8972
+    },
+    {
+      "epoch": 1.0252515289011639,
+      "eval_accuracy": 0.13303769401330376,
+      "eval_loss": 3.5806682109832764,
+      "eval_runtime": 23.4297,
+      "eval_samples_per_second": 38.498,
+      "eval_steps_per_second": 4.823,
+      "step": 650
+    },
+    {
+      "epoch": 1.0410337344643914,
+      "grad_norm": 90.48371602841107,
+      "learning_rate": 3.6323787258912916e-06,
+      "loss": 30.3281,
+      "step": 660,
+      "true_loss": 4.0951
+    },
+    {
+      "epoch": 1.056815940027619,
+      "grad_norm": 80.37129983569977,
+      "learning_rate": 3.6031560490940975e-06,
+      "loss": 30.4516,
+      "step": 670,
+      "true_loss": 4.0298
+    },
+    {
+      "epoch": 1.0725981455908462,
+      "grad_norm": 89.0421974632656,
+      "learning_rate": 3.5739333722969025e-06,
+      "loss": 30.5053,
+      "step": 680,
+      "true_loss": 3.838
+    },
+    {
+      "epoch": 1.0883803511540737,
+      "grad_norm": 87.6238382263665,
+      "learning_rate": 3.544710695499708e-06,
+      "loss": 30.1458,
+      "step": 690,
+      "true_loss": 3.4672
+    },
+    {
+      "epoch": 1.1041625567173012,
+      "grad_norm": 93.33199336725635,
+      "learning_rate": 3.5154880187025135e-06,
+      "loss": 29.9857,
+      "step": 700,
+      "true_loss": 3.8479
+    },
+    {
+      "epoch": 1.1041625567173012,
+      "eval_accuracy": 0.1574279379157428,
+      "eval_loss": 3.5641186237335205,
+      "eval_runtime": 23.6052,
+      "eval_samples_per_second": 38.212,
+      "eval_steps_per_second": 4.787,
+      "step": 700
+    },
+    {
+      "epoch": 1.1199447622805287,
+      "grad_norm": 92.22724838108498,
+      "learning_rate": 3.486265341905319e-06,
+      "loss": 29.8985,
+      "step": 710,
+      "true_loss": 3.6761
+    },
+    {
+      "epoch": 1.1357269678437563,
+      "grad_norm": 84.13971221580583,
+      "learning_rate": 3.457042665108124e-06,
+      "loss": 30.7157,
+      "step": 720,
+      "true_loss": 3.6518
+    },
+    {
+      "epoch": 1.1515091734069836,
+      "grad_norm": 85.8960190061776,
+      "learning_rate": 3.42781998831093e-06,
+      "loss": 30.2104,
+      "step": 730,
+      "true_loss": 3.7526
+    },
+    {
+      "epoch": 1.167291378970211,
+      "grad_norm": 86.69215126012148,
+      "learning_rate": 3.398597311513735e-06,
+      "loss": 29.405,
+      "step": 740,
+      "true_loss": 4.0226
+    },
+    {
+      "epoch": 1.1830735845334386,
+      "grad_norm": 87.86033437742292,
+      "learning_rate": 3.3693746347165403e-06,
+      "loss": 30.511,
+      "step": 750,
+      "true_loss": 3.7897
+    },
+    {
+      "epoch": 1.1830735845334386,
+      "eval_accuracy": 0.14523281596452328,
+      "eval_loss": 3.493298292160034,
+      "eval_runtime": 23.8274,
+      "eval_samples_per_second": 37.856,
+      "eval_steps_per_second": 4.742,
+      "step": 750
+    },
+    {
+      "epoch": 1.198855790096666,
+      "grad_norm": 111.93866332819323,
+      "learning_rate": 3.3401519579193458e-06,
+      "loss": 29.8573,
+      "step": 760,
+      "true_loss": 3.7307
+    },
+    {
+      "epoch": 1.2146379956598934,
+      "grad_norm": 78.19039410059621,
+      "learning_rate": 3.310929281122151e-06,
+      "loss": 30.2595,
+      "step": 770,
+      "true_loss": 3.7585
+    },
+    {
+      "epoch": 1.230420201223121,
+      "grad_norm": 136.24794475282891,
+      "learning_rate": 3.2817066043249562e-06,
+      "loss": 30.1392,
+      "step": 780,
+      "true_loss": 3.6031
+    },
+    {
+      "epoch": 1.2462024067863484,
+      "grad_norm": 83.26617285476617,
+      "learning_rate": 3.252483927527762e-06,
+      "loss": 29.7774,
+      "step": 790,
+      "true_loss": 3.8831
+    },
+    {
+      "epoch": 1.261984612349576,
+      "grad_norm": 82.12832696286748,
+      "learning_rate": 3.223261250730567e-06,
+      "loss": 30.2729,
+      "step": 800,
+      "true_loss": 3.7239
+    },
+    {
+      "epoch": 1.261984612349576,
+      "eval_accuracy": 0.15964523281596452,
+      "eval_loss": 3.5300114154815674,
+      "eval_runtime": 23.4485,
+      "eval_samples_per_second": 38.467,
+      "eval_steps_per_second": 4.819,
+      "step": 800
+    },
+    {
+      "epoch": 1.2777668179128034,
+      "grad_norm": 90.3512257461839,
+      "learning_rate": 3.194038573933372e-06,
+      "loss": 30.1961,
+      "step": 810,
+      "true_loss": 3.784
+    },
+    {
+      "epoch": 1.2935490234760307,
+      "grad_norm": 84.55475078707286,
+      "learning_rate": 3.164815897136178e-06,
+      "loss": 29.3181,
+      "step": 820,
+      "true_loss": 3.5114
+    },
+    {
+      "epoch": 1.3093312290392582,
+      "grad_norm": 89.4992221052368,
+      "learning_rate": 3.135593220338983e-06,
+      "loss": 30.2988,
+      "step": 830,
+      "true_loss": 3.7831
+    },
+    {
+      "epoch": 1.3251134346024858,
+      "grad_norm": 92.43123125189919,
+      "learning_rate": 3.1063705435417885e-06,
+      "loss": 30.0873,
+      "step": 840,
+      "true_loss": 3.8634
+    },
+    {
+      "epoch": 1.340895640165713,
+      "grad_norm": 89.04832868320074,
+      "learning_rate": 3.0771478667445944e-06,
+      "loss": 30.2665,
+      "step": 850,
+      "true_loss": 3.8983
+    },
+    {
+      "epoch": 1.340895640165713,
+      "eval_accuracy": 0.16851441241685144,
+      "eval_loss": 3.462191343307495,
+      "eval_runtime": 23.8924,
+      "eval_samples_per_second": 37.753,
+      "eval_steps_per_second": 4.73,
+      "step": 850
+    },
+    {
+      "epoch": 1.3566778457289406,
+      "grad_norm": 92.16494941680725,
+      "learning_rate": 3.0479251899473994e-06,
+      "loss": 29.7738,
+      "step": 860,
+      "true_loss": 4.0234
+    },
+    {
+      "epoch": 1.372460051292168,
+      "grad_norm": 85.57074344764075,
+      "learning_rate": 3.0187025131502045e-06,
+      "loss": 30.0651,
+      "step": 870,
+      "true_loss": 3.7162
+    },
+    {
+      "epoch": 1.3882422568553956,
+      "grad_norm": 97.12213474754635,
+      "learning_rate": 2.9894798363530103e-06,
+      "loss": 30.0534,
+      "step": 880,
+      "true_loss": 3.969
+    },
+    {
+      "epoch": 1.404024462418623,
+      "grad_norm": 85.14532114581857,
+      "learning_rate": 2.9602571595558154e-06,
+      "loss": 30.1717,
+      "step": 890,
+      "true_loss": 3.8654
+    },
+    {
+      "epoch": 1.4198066679818504,
+      "grad_norm": 89.01307343902022,
+      "learning_rate": 2.931034482758621e-06,
+      "loss": 30.2877,
+      "step": 900,
+      "true_loss": 3.6224
+    },
+    {
+      "epoch": 1.4198066679818504,
+      "eval_accuracy": 0.16186252771618626,
+      "eval_loss": 3.4574382305145264,
+      "eval_runtime": 23.3454,
+      "eval_samples_per_second": 38.637,
+      "eval_steps_per_second": 4.84,
+      "step": 900
+    },
+    {
+      "epoch": 1.435588873545078,
+      "grad_norm": 87.24388793474637,
+      "learning_rate": 2.9018118059614263e-06,
+      "loss": 30.8903,
+      "step": 910,
+      "true_loss": 3.8262
+    },
+    {
+      "epoch": 1.4513710791083054,
+      "grad_norm": 84.04918766102597,
+      "learning_rate": 2.8725891291642317e-06,
+      "loss": 29.8196,
+      "step": 920,
+      "true_loss": 3.918
+    },
+    {
+      "epoch": 1.4671532846715327,
+      "grad_norm": 83.16951966482195,
+      "learning_rate": 2.8433664523670368e-06,
+      "loss": 29.6097,
+      "step": 930,
+      "true_loss": 3.4664
+    },
+    {
+      "epoch": 1.4829354902347602,
+      "grad_norm": 96.29394268139959,
+      "learning_rate": 2.8141437755698426e-06,
+      "loss": 29.7147,
+      "step": 940,
+      "true_loss": 3.6577
+    },
+    {
+      "epoch": 1.4987176957979877,
+      "grad_norm": 101.35324066409235,
+      "learning_rate": 2.7849210987726477e-06,
+      "loss": 29.4593,
+      "step": 950,
+      "true_loss": 3.6391
+    },
+    {
+      "epoch": 1.4987176957979877,
+      "eval_accuracy": 0.15188470066518847,
+      "eval_loss": 3.448209762573242,
+      "eval_runtime": 23.74,
+      "eval_samples_per_second": 37.995,
+      "eval_steps_per_second": 4.76,
+      "step": 950
+    },
+    {
+      "epoch": 1.5144999013612153,
+      "grad_norm": 91.01958083657513,
+      "learning_rate": 2.7556984219754535e-06,
+      "loss": 29.4563,
+      "step": 960,
+      "true_loss": 3.5778
+    },
+    {
+      "epoch": 1.5302821069244428,
+      "grad_norm": 94.87158519297142,
+      "learning_rate": 2.7264757451782586e-06,
+      "loss": 29.8641,
+      "step": 970,
+      "true_loss": 3.9124
+    },
+    {
+      "epoch": 1.5460643124876703,
+      "grad_norm": 84.01959219280997,
+      "learning_rate": 2.697253068381064e-06,
+      "loss": 29.8265,
+      "step": 980,
+      "true_loss": 3.5047
+    },
+    {
+      "epoch": 1.5618465180508976,
+      "grad_norm": 91.64799353240991,
+      "learning_rate": 2.6680303915838695e-06,
+      "loss": 29.858,
+      "step": 990,
+      "true_loss": 3.701
+    },
+    {
+      "epoch": 1.577628723614125,
+      "grad_norm": 88.05009850532113,
+      "learning_rate": 2.638807714786675e-06,
+      "loss": 29.9069,
+      "step": 1000,
+      "true_loss": 3.6317
+    },
+    {
+      "epoch": 1.577628723614125,
+      "eval_accuracy": 0.14412416851441243,
+      "eval_loss": 3.4386119842529297,
+      "eval_runtime": 24.0921,
+      "eval_samples_per_second": 37.44,
+      "eval_steps_per_second": 4.69,
+      "step": 1000
+    },
+    {
+      "epoch": 1.5934109291773524,
+      "grad_norm": 88.61468815912133,
+      "learning_rate": 2.60958503798948e-06,
+      "loss": 29.7709,
+      "step": 1010,
+      "true_loss": 3.6406
+    },
+    {
+      "epoch": 1.60919313474058,
+      "grad_norm": 85.25549864190619,
+      "learning_rate": 2.580362361192286e-06,
+      "loss": 29.8656,
+      "step": 1020,
+      "true_loss": 4.0205
+    },
+    {
+      "epoch": 1.6249753403038074,
+      "grad_norm": 95.17123739206981,
+      "learning_rate": 2.551139684395091e-06,
+      "loss": 29.3146,
+      "step": 1030,
+      "true_loss": 3.7517
+    },
+    {
+      "epoch": 1.640757545867035,
+      "grad_norm": 94.05266971651754,
+      "learning_rate": 2.521917007597896e-06,
+      "loss": 29.5985,
+      "step": 1040,
+      "true_loss": 3.8767
+    },
+    {
+      "epoch": 1.6565397514302624,
+      "grad_norm": 84.23330948930467,
+      "learning_rate": 2.4926943308007014e-06,
+      "loss": 29.6116,
+      "step": 1050,
+      "true_loss": 3.6806
+    },
+    {
+      "epoch": 1.6565397514302624,
+      "eval_accuracy": 0.1574279379157428,
+      "eval_loss": 3.4331889152526855,
+      "eval_runtime": 24.0418,
+      "eval_samples_per_second": 37.518,
+      "eval_steps_per_second": 4.7,
+      "step": 1050
+    },
+    {
+      "epoch": 1.67232195699349,
+      "grad_norm": 95.80967931580453,
+      "learning_rate": 2.463471654003507e-06,
+      "loss": 29.4546,
+      "step": 1060,
+      "true_loss": 3.5296
+    },
+    {
+      "epoch": 1.6881041625567172,
+      "grad_norm": 89.21975858244956,
+      "learning_rate": 2.4342489772063123e-06,
+      "loss": 29.4728,
+      "step": 1070,
+      "true_loss": 3.8411
+    },
+    {
+      "epoch": 1.7038863681199448,
+      "grad_norm": 93.36233191510571,
+      "learning_rate": 2.4050263004091177e-06,
+      "loss": 29.5254,
+      "step": 1080,
+      "true_loss": 3.7742
+    },
+    {
+      "epoch": 1.7196685736831723,
+      "grad_norm": 107.3636400374849,
+      "learning_rate": 2.375803623611923e-06,
+      "loss": 29.5684,
+      "step": 1090,
+      "true_loss": 3.6389
+    },
+    {
+      "epoch": 1.7354507792463996,
+      "grad_norm": 1304.6259656582167,
+      "learning_rate": 2.3465809468147286e-06,
+      "loss": 29.638,
+      "step": 1100,
+      "true_loss": 3.7105
+    },
+    {
+      "epoch": 1.7354507792463996,
+      "eval_accuracy": 0.16518847006651885,
+      "eval_loss": 3.4311046600341797,
+      "eval_runtime": 23.6704,
+      "eval_samples_per_second": 38.107,
+      "eval_steps_per_second": 4.774,
+      "step": 1100
+    },
+    {
+      "epoch": 1.751232984809627,
+      "grad_norm": 96.2130469861494,
+      "learning_rate": 2.3173582700175337e-06,
+      "loss": 29.3766,
+      "step": 1110,
+      "true_loss": 3.67
+    },
+    {
+      "epoch": 1.7670151903728546,
+      "grad_norm": 95.92673596052384,
+      "learning_rate": 2.288135593220339e-06,
+      "loss": 29.3919,
+      "step": 1120,
+      "true_loss": 3.8554
+    },
+    {
+      "epoch": 1.782797395936082,
+      "grad_norm": 88.05948564467495,
+      "learning_rate": 2.2589129164231446e-06,
+      "loss": 29.2821,
+      "step": 1130,
+      "true_loss": 3.7715
+    },
+    {
+      "epoch": 1.7985796014993096,
+      "grad_norm": 90.55869365935298,
+      "learning_rate": 2.22969023962595e-06,
+      "loss": 29.7986,
+      "step": 1140,
+      "true_loss": 3.6651
+    },
+    {
+      "epoch": 1.8143618070625371,
+      "grad_norm": 94.20405819159555,
+      "learning_rate": 2.2004675628287555e-06,
+      "loss": 29.8496,
+      "step": 1150,
+      "true_loss": 3.7782
+    },
+    {
+      "epoch": 1.8143618070625371,
+      "eval_accuracy": 0.16186252771618626,
+      "eval_loss": 3.4140689373016357,
+      "eval_runtime": 24.0501,
+      "eval_samples_per_second": 37.505,
+      "eval_steps_per_second": 4.699,
+      "step": 1150
+    },
+    {
+      "epoch": 1.8301440126257644,
+      "grad_norm": 87.96745013401917,
+      "learning_rate": 2.171244886031561e-06,
+      "loss": 29.5695,
+      "step": 1160,
+      "true_loss": 3.8398
+    },
+    {
+      "epoch": 1.845926218188992,
+      "grad_norm": 92.10246815773178,
+      "learning_rate": 2.142022209234366e-06,
+      "loss": 29.8099,
+      "step": 1170,
+      "true_loss": 3.9006
+    },
+    {
+      "epoch": 1.8617084237522192,
+      "grad_norm": 83.12415816117908,
+      "learning_rate": 2.1127995324371714e-06,
+      "loss": 30.2325,
+      "step": 1180,
+      "true_loss": 3.623
+    },
+    {
+      "epoch": 1.8774906293154467,
+      "grad_norm": 107.24249122552868,
+      "learning_rate": 2.083576855639977e-06,
+      "loss": 29.8222,
+      "step": 1190,
+      "true_loss": 3.5215
+    },
+    {
+      "epoch": 1.8932728348786743,
+      "grad_norm": 91.4088303372265,
+      "learning_rate": 2.054354178842782e-06,
+      "loss": 29.498,
+      "step": 1200,
+      "true_loss": 3.653
+    },
+    {
+      "epoch": 1.8932728348786743,
+      "eval_accuracy": 0.16851441241685144,
+      "eval_loss": 3.4037177562713623,
+      "eval_runtime": 23.8616,
+      "eval_samples_per_second": 37.801,
+      "eval_steps_per_second": 4.736,
+      "step": 1200
+    },
+    {
+      "epoch": 1.9090550404419018,
+      "grad_norm": 100.69059595068175,
+      "learning_rate": 2.0251315020455873e-06,
+      "loss": 29.1097,
+      "step": 1210,
+      "true_loss": 3.5445
+    },
+    {
+      "epoch": 1.9248372460051293,
+      "grad_norm": 96.05062058766768,
+      "learning_rate": 1.995908825248393e-06,
+      "loss": 29.5919,
+      "step": 1220,
+      "true_loss": 3.6302
+    },
+    {
+      "epoch": 1.9406194515683568,
+      "grad_norm": 89.1551204065743,
+      "learning_rate": 1.9666861484511982e-06,
+      "loss": 29.0917,
+      "step": 1230,
+      "true_loss": 3.5472
+    },
+    {
+      "epoch": 1.956401657131584,
+      "grad_norm": 101.09204426369772,
+      "learning_rate": 1.9374634716540037e-06,
+      "loss": 28.8413,
+      "step": 1240,
+      "true_loss": 3.8083
+    },
+    {
+      "epoch": 1.9721838626948116,
+      "grad_norm": 105.57133293543227,
+      "learning_rate": 1.908240794856809e-06,
+      "loss": 29.533,
+      "step": 1250,
+      "true_loss": 3.7014
+    },
+    {
+      "epoch": 1.9721838626948116,
+      "eval_accuracy": 0.17516629711751663,
+      "eval_loss": 3.385300397872925,
+      "eval_runtime": 23.5478,
+      "eval_samples_per_second": 38.305,
+      "eval_steps_per_second": 4.799,
+      "step": 1250
+    },
+    {
+      "epoch": 1.987966068258039,
+      "grad_norm": 96.01274952042326,
+      "learning_rate": 1.8790181180596146e-06,
+      "loss": 29.2844,
+      "step": 1260,
+      "true_loss": 3.8112
+    },
+    {
+      "epoch": 2.0031564411126457,
+      "grad_norm": 94.04158733000145,
+      "learning_rate": 1.8497954412624196e-06,
+      "loss": 27.161,
+      "step": 1270,
+      "true_loss": 3.399
+    },
+    {
+      "epoch": 2.0189386466758728,
+      "grad_norm": 113.00306213639786,
+      "learning_rate": 1.820572764465225e-06,
+      "loss": 27.5121,
+      "step": 1280,
+      "true_loss": 3.5377
+    },
+    {
+      "epoch": 2.0347208522391003,
+      "grad_norm": 128.85739550143202,
+      "learning_rate": 1.7913500876680305e-06,
+      "loss": 27.3831,
+      "step": 1290,
+      "true_loss": 3.3378
+    },
+    {
+      "epoch": 2.0505030578023278,
+      "grad_norm": 135.357215369742,
+      "learning_rate": 1.7621274108708358e-06,
+      "loss": 28.0219,
+      "step": 1300,
+      "true_loss": 3.6161
+    },
+    {
+      "epoch": 2.0505030578023278,
+      "eval_accuracy": 0.17849223946784923,
+      "eval_loss": 3.400832414627075,
+      "eval_runtime": 24.2954,
+      "eval_samples_per_second": 37.126,
+      "eval_steps_per_second": 4.651,
+      "step": 1300
+    },
+    {
+      "epoch": 2.0662852633655553,
+      "grad_norm": 133.18820676570508,
+      "learning_rate": 1.7329047340736412e-06,
+      "loss": 27.7807,
+      "step": 1310,
+      "true_loss": 3.5333
+    },
+    {
+      "epoch": 2.082067468928783,
+      "grad_norm": 132.12454176042834,
+      "learning_rate": 1.7036820572764467e-06,
+      "loss": 27.6285,
+      "step": 1320,
+      "true_loss": 3.122
+    },
+    {
+      "epoch": 2.0978496744920103,
+      "grad_norm": 121.63901914437486,
+      "learning_rate": 1.674459380479252e-06,
+      "loss": 27.2663,
+      "step": 1330,
+      "true_loss": 3.2256
+    },
+    {
+      "epoch": 2.113631880055238,
+      "grad_norm": 137.06803144579217,
+      "learning_rate": 1.6452367036820574e-06,
+      "loss": 27.6847,
+      "step": 1340,
+      "true_loss": 3.3579
+    },
+    {
+      "epoch": 2.1294140856184653,
+      "grad_norm": 127.41062741726843,
+      "learning_rate": 1.6160140268848628e-06,
+      "loss": 27.7085,
+      "step": 1350,
+      "true_loss": 3.45
+    },
+    {
+      "epoch": 2.1294140856184653,
+      "eval_accuracy": 0.18070953436807094,
+      "eval_loss": 3.376422643661499,
+      "eval_runtime": 23.8791,
+      "eval_samples_per_second": 37.774,
+      "eval_steps_per_second": 4.732,
+      "step": 1350
+    },
+    {
+      "epoch": 2.1451962911816924,
+      "grad_norm": 121.2757103113971,
+      "learning_rate": 1.586791350087668e-06,
+      "loss": 27.6917,
+      "step": 1360,
+      "true_loss": 3.3579
+    },
+    {
+      "epoch": 2.16097849674492,
+      "grad_norm": 128.57906764071424,
+      "learning_rate": 1.5575686732904735e-06,
+      "loss": 26.6196,
+      "step": 1370,
+      "true_loss": 3.2661
+    },
+    {
+      "epoch": 2.1767607023081474,
+      "grad_norm": 144.7381021633525,
+      "learning_rate": 1.528345996493279e-06,
+      "loss": 28.1214,
+      "step": 1380,
+      "true_loss": 3.5539
+    },
+    {
+      "epoch": 2.192542907871375,
+      "grad_norm": 145.5114518133331,
+      "learning_rate": 1.4991233196960842e-06,
+      "loss": 27.5329,
+      "step": 1390,
+      "true_loss": 3.5611
+    },
+    {
+      "epoch": 2.2083251134346025,
+      "grad_norm": 137.75864220522647,
+      "learning_rate": 1.4699006428988897e-06,
+      "loss": 27.6466,
+      "step": 1400,
+      "true_loss": 3.4073
+    },
+    {
+      "epoch": 2.2083251134346025,
+      "eval_accuracy": 0.18514412416851442,
+      "eval_loss": 3.3832011222839355,
+      "eval_runtime": 23.5091,
+      "eval_samples_per_second": 38.368,
+      "eval_steps_per_second": 4.807,
+      "step": 1400
+    },
+    {
+      "epoch": 2.22410731899783,
+      "grad_norm": 132.61505430150723,
+      "learning_rate": 1.4406779661016951e-06,
+      "loss": 27.1518,
+      "step": 1410,
+      "true_loss": 3.4681
+    },
+    {
+      "epoch": 2.2398895245610575,
+      "grad_norm": 140.85574741341603,
+      "learning_rate": 1.4114552893045006e-06,
+      "loss": 27.7041,
+      "step": 1420,
+      "true_loss": 3.6406
+    },
+    {
+      "epoch": 2.255671730124285,
+      "grad_norm": 245.85007563586552,
+      "learning_rate": 1.3822326125073058e-06,
+      "loss": 27.6303,
+      "step": 1430,
+      "true_loss": 3.5352
+    },
+    {
+      "epoch": 2.2714539356875125,
+      "grad_norm": 577.1875021093457,
+      "learning_rate": 1.3530099357101113e-06,
+      "loss": 28.509,
+      "step": 1440,
+      "true_loss": 3.5396
+    },
+    {
+      "epoch": 2.28723614125074,
+      "grad_norm": 147.03022521178113,
+      "learning_rate": 1.3237872589129167e-06,
+      "loss": 27.6055,
+      "step": 1450,
+      "true_loss": 3.4989
+    },
+    {
+      "epoch": 2.28723614125074,
+      "eval_accuracy": 0.18514412416851442,
+      "eval_loss": 3.4081857204437256,
+      "eval_runtime": 24.0487,
+      "eval_samples_per_second": 37.507,
+      "eval_steps_per_second": 4.699,
+      "step": 1450
+    },
+    {
+      "epoch": 2.303018346813967,
+      "grad_norm": 158.25616293531198,
+      "learning_rate": 1.2945645821157218e-06,
+      "loss": 27.6877,
+      "step": 1460,
+      "true_loss": 3.6034
+    },
+    {
+      "epoch": 2.3188005523771946,
+      "grad_norm": 155.95509000199831,
+      "learning_rate": 1.2653419053185272e-06,
+      "loss": 27.8664,
+      "step": 1470,
+      "true_loss": 3.5089
+    },
+    {
+      "epoch": 2.334582757940422,
+      "grad_norm": 146.83540699519045,
+      "learning_rate": 1.2361192285213327e-06,
+      "loss": 27.413,
+      "step": 1480,
+      "true_loss": 3.5116
+    },
+    {
+      "epoch": 2.3503649635036497,
+      "grad_norm": 141.1565994249431,
+      "learning_rate": 1.2068965517241381e-06,
+      "loss": 27.1606,
+      "step": 1490,
+      "true_loss": 3.426
+    },
+    {
+      "epoch": 2.366147169066877,
+      "grad_norm": 159.8393327569506,
+      "learning_rate": 1.1776738749269434e-06,
+      "loss": 27.7131,
+      "step": 1500,
+      "true_loss": 3.3148
+    },
+    {
+      "epoch": 2.366147169066877,
+      "eval_accuracy": 0.19623059866962306,
+      "eval_loss": 3.372471332550049,
+      "eval_runtime": 23.4809,
+      "eval_samples_per_second": 38.414,
+      "eval_steps_per_second": 4.812,
+      "step": 1500
+    },
+    {
+      "epoch": 2.3819293746301047,
+      "grad_norm": 168.70353622158657,
+      "learning_rate": 1.1484511981297488e-06,
+      "loss": 26.9731,
+      "step": 1510,
+      "true_loss": 3.3077
+    },
+    {
+      "epoch": 2.397711580193332,
+      "grad_norm": 163.63685602929394,
+      "learning_rate": 1.1192285213325543e-06,
+      "loss": 27.206,
+      "step": 1520,
+      "true_loss": 3.2603
+    },
+    {
+      "epoch": 2.4134937857565593,
+      "grad_norm": 171.43412737818744,
+      "learning_rate": 1.0900058445353595e-06,
+      "loss": 26.3347,
+      "step": 1530,
+      "true_loss": 3.2123
+    },
+    {
+      "epoch": 2.4292759913197868,
+      "grad_norm": 167.19508362270676,
+      "learning_rate": 1.0607831677381648e-06,
+      "loss": 27.482,
+      "step": 1540,
+      "true_loss": 3.4389
+    },
+    {
+      "epoch": 2.4450581968830143,
+      "grad_norm": 165.6072657994525,
+      "learning_rate": 1.0315604909409702e-06,
+      "loss": 26.6002,
+      "step": 1550,
+      "true_loss": 3.4554
+    },
+    {
+      "epoch": 2.4450581968830143,
+      "eval_accuracy": 0.1984478935698448,
+      "eval_loss": 3.3719024658203125,
+      "eval_runtime": 23.4592,
+      "eval_samples_per_second": 38.45,
+      "eval_steps_per_second": 4.817,
+      "step": 1550
+    },
+    {
+      "epoch": 2.460840402446242,
+      "grad_norm": 173.48714728794272,
+      "learning_rate": 1.0023378141437757e-06,
+      "loss": 27.0777,
+      "step": 1560,
+      "true_loss": 3.4451
+    },
+    {
+      "epoch": 2.4766226080094693,
+      "grad_norm": 169.41551174748028,
+      "learning_rate": 9.731151373465811e-07,
+      "loss": 27.1345,
+      "step": 1570,
+      "true_loss": 3.3625
+    },
+    {
+      "epoch": 2.492404813572697,
+      "grad_norm": 179.8847072240023,
+      "learning_rate": 9.438924605493864e-07,
+      "loss": 27.0216,
+      "step": 1580,
+      "true_loss": 3.2843
+    },
+    {
+      "epoch": 2.5081870191359243,
+      "grad_norm": 166.02108936347767,
+      "learning_rate": 9.146697837521917e-07,
+      "loss": 26.7401,
+      "step": 1590,
+      "true_loss": 3.2657
+    },
+    {
+      "epoch": 2.523969224699152,
+      "grad_norm": 454.3584387218693,
+      "learning_rate": 8.854471069549972e-07,
+      "loss": 26.1234,
+      "step": 1600,
+      "true_loss": 3.3808
+    },
+    {
+      "epoch": 2.523969224699152,
+      "eval_accuracy": 0.20288248337028825,
+      "eval_loss": 3.380685567855835,
+      "eval_runtime": 24.1535,
+      "eval_samples_per_second": 37.344,
+      "eval_steps_per_second": 4.678,
+      "step": 1600
+    },
+    {
+      "epoch": 2.5397514302623794,
+      "grad_norm": 164.3488500545829,
+      "learning_rate": 8.562244301578025e-07,
+      "loss": 26.6329,
+      "step": 1610,
+      "true_loss": 3.4502
+    },
+    {
+      "epoch": 2.555533635825607,
+      "grad_norm": 224.71536684248224,
+      "learning_rate": 8.270017533606079e-07,
+      "loss": 27.4489,
+      "step": 1620,
+      "true_loss": 3.4381
+    },
+    {
+      "epoch": 2.571315841388834,
+      "grad_norm": 189.92117872676408,
+      "learning_rate": 7.977790765634133e-07,
+      "loss": 27.2729,
+      "step": 1630,
+      "true_loss": 3.3647
+    },
+    {
+      "epoch": 2.5870980469520615,
+      "grad_norm": 172.51933992524627,
+      "learning_rate": 7.685563997662187e-07,
+      "loss": 27.1534,
+      "step": 1640,
+      "true_loss": 3.5251
+    },
+    {
+      "epoch": 2.602880252515289,
+      "grad_norm": 166.6926362336406,
+      "learning_rate": 7.393337229690241e-07,
+      "loss": 26.9745,
+      "step": 1650,
+      "true_loss": 3.3921
+    },
+    {
+      "epoch": 2.602880252515289,
+      "eval_accuracy": 0.2006651884700665,
+      "eval_loss": 3.3468966484069824,
+      "eval_runtime": 23.7179,
+      "eval_samples_per_second": 38.03,
+      "eval_steps_per_second": 4.764,
+      "step": 1650
+    },
+    {
+      "epoch": 2.6186624580785165,
+      "grad_norm": 160.91931691807488,
+      "learning_rate": 7.101110461718295e-07,
+      "loss": 26.7737,
+      "step": 1660,
+      "true_loss": 3.4577
+    },
+    {
+      "epoch": 2.634444663641744,
+      "grad_norm": 181.67048483012417,
+      "learning_rate": 6.808883693746347e-07,
+      "loss": 26.7581,
+      "step": 1670,
+      "true_loss": 3.2471
+    },
+    {
+      "epoch": 2.6502268692049715,
+      "grad_norm": 217.44579386829224,
+      "learning_rate": 6.516656925774401e-07,
+      "loss": 26.4481,
+      "step": 1680,
+      "true_loss": 3.0002
+    },
+    {
+      "epoch": 2.6660090747681986,
+      "grad_norm": 195.64586609651414,
+      "learning_rate": 6.224430157802455e-07,
+      "loss": 26.7862,
+      "step": 1690,
+      "true_loss": 3.1025
+    },
+    {
+      "epoch": 2.681791280331426,
+      "grad_norm": 179.8032160582305,
+      "learning_rate": 5.93220338983051e-07,
+      "loss": 27.108,
+      "step": 1700,
+      "true_loss": 3.1252
+    },
+    {
+      "epoch": 2.681791280331426,
+      "eval_accuracy": 0.2017738359201774,
+      "eval_loss": 3.3589749336242676,
+      "eval_runtime": 23.6074,
+      "eval_samples_per_second": 38.208,
+      "eval_steps_per_second": 4.787,
+      "step": 1700
+    },
+    {
+      "epoch": 2.6975734858946536,
+      "grad_norm": 189.69952810255015,
+      "learning_rate": 5.639976621858563e-07,
+      "loss": 26.6985,
+      "step": 1710,
+      "true_loss": 3.3933
+    },
+    {
+      "epoch": 2.713355691457881,
+      "grad_norm": 242.2858845486621,
+      "learning_rate": 5.347749853886616e-07,
+      "loss": 26.7437,
+      "step": 1720,
+      "true_loss": 3.1209
+    },
+    {
+      "epoch": 2.7291378970211086,
+      "grad_norm": 185.37642956971033,
+      "learning_rate": 5.05552308591467e-07,
+      "loss": 26.7456,
+      "step": 1730,
+      "true_loss": 3.1121
+    },
+    {
+      "epoch": 2.744920102584336,
+      "grad_norm": 201.70311652816994,
+      "learning_rate": 4.763296317942724e-07,
+      "loss": 26.8516,
+      "step": 1740,
+      "true_loss": 3.397
+    },
+    {
+      "epoch": 2.7607023081475637,
+      "grad_norm": 181.35331376149924,
+      "learning_rate": 4.4710695499707774e-07,
+      "loss": 27.6208,
+      "step": 1750,
+      "true_loss": 3.6422
+    },
+    {
+      "epoch": 2.7607023081475637,
+      "eval_accuracy": 0.21951219512195122,
+      "eval_loss": 3.3623762130737305,
+      "eval_runtime": 23.6173,
+      "eval_samples_per_second": 38.192,
+      "eval_steps_per_second": 4.785,
+      "step": 1750
+    },
+    {
+      "epoch": 2.776484513710791,
+      "grad_norm": 207.9779142213798,
+      "learning_rate": 4.1788427819988314e-07,
+      "loss": 26.674,
+      "step": 1760,
+      "true_loss": 3.221
+    },
+    {
+      "epoch": 2.7922667192740187,
+      "grad_norm": 203.7549680190773,
+      "learning_rate": 3.8866160140268854e-07,
+      "loss": 26.5723,
+      "step": 1770,
+      "true_loss": 3.2956
+    },
+    {
+      "epoch": 2.808048924837246,
+      "grad_norm": 196.13189611302036,
+      "learning_rate": 3.594389246054939e-07,
+      "loss": 26.4406,
+      "step": 1780,
+      "true_loss": 3.4648
+    },
+    {
+      "epoch": 2.8238311304004737,
+      "grad_norm": 198.30227589105047,
+      "learning_rate": 3.3021624780829924e-07,
+      "loss": 26.1588,
+      "step": 1790,
+      "true_loss": 3.2933
+    },
+    {
+      "epoch": 2.839613335963701,
+      "grad_norm": 199.85226446602726,
+      "learning_rate": 3.0099357101110464e-07,
+      "loss": 26.2669,
+      "step": 1800,
+      "true_loss": 3.2296
+    },
+    {
+      "epoch": 2.839613335963701,
+      "eval_accuracy": 0.2073170731707317,
+      "eval_loss": 3.3512346744537354,
+      "eval_runtime": 23.6919,
+      "eval_samples_per_second": 38.072,
+      "eval_steps_per_second": 4.77,
+      "step": 1800
+    },
+    {
+      "epoch": 2.8553955415269283,
+      "grad_norm": 217.9595424207247,
+      "learning_rate": 2.7177089421391004e-07,
+      "loss": 26.0995,
+      "step": 1810,
+      "true_loss": 3.0267
+    },
+    {
+      "epoch": 2.871177747090156,
+      "grad_norm": 222.64507007990645,
+      "learning_rate": 2.425482174167154e-07,
+      "loss": 25.9481,
+      "step": 1820,
+      "true_loss": 3.4288
+    },
+    {
+      "epoch": 2.8869599526533833,
+      "grad_norm": 201.21510324225665,
+      "learning_rate": 2.1332554061952078e-07,
+      "loss": 25.9752,
+      "step": 1830,
+      "true_loss": 3.1737
+    },
+    {
+      "epoch": 2.902742158216611,
+      "grad_norm": 192.09070144987018,
+      "learning_rate": 1.8410286382232613e-07,
+      "loss": 26.059,
+      "step": 1840,
+      "true_loss": 3.024
+    },
+    {
+      "epoch": 2.9185243637798384,
+      "grad_norm": 940.254161469039,
+      "learning_rate": 1.5488018702513153e-07,
+      "loss": 25.8304,
+      "step": 1850,
+      "true_loss": 2.8998
+    },
+    {
+      "epoch": 2.9185243637798384,
+      "eval_accuracy": 0.21840354767184036,
+      "eval_loss": 3.3478667736053467,
+      "eval_runtime": 23.6868,
+      "eval_samples_per_second": 38.08,
+      "eval_steps_per_second": 4.771,
+      "step": 1850
+    }
+  ],
+  "logging_steps": 10,
+  "max_steps": 1902,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 3,
+  "save_steps": 50,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": false
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 0.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:6ef326857202175de846b7a9b400ab5118100fabc46f9e836eb5cd734f7d1cd7
+size 7249

zero_to_fp32.py

@@ -0,0 +1,760 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example:
+#   python zero_to_fp32.py . output_dir/
+#   or
+#   python zero_to_fp32.py . output_dir/ --safe_serialization
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+import gc
+import json
+import numpy as np
+from tqdm import tqdm
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device, weights_only=False)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+    total_files = len(files)
+    state_dicts = []
+    for f in tqdm(files, desc='Loading checkpoint shards'):
+        state_dict = torch.load(f, map_location=device, mmap=True, weights_only=False)
+        # immediately discard the potentially huge 2 optimizer states as we only care for fp32 master weights
+        # and also handle the case where it was already removed by another helper script
+        state_dict["optimizer_state_dict"].pop("optimizer_state_dict", None)
+        state_dicts.append(state_dict)
+
+    if ZERO_STAGE not in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                                          exclude_frozen_parameters)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _has_callable(obj, fn):
+    attr = getattr(obj, fn, None)
+    return callable(attr)
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel() if _has_callable(shape, 'numel') else math.prod(shape)
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+class GatheredTensor:
+    """
+    A pseudo tensor that collects partitioned weights.
+    It is more memory efficient when there are multiple groups.
+    """
+
+    def __init__(self, flat_groups, flat_groups_offset, offset, partitioned_numel, shape):
+        self.flat_groups = flat_groups
+        self.flat_groups_offset = flat_groups_offset
+        self.offset = offset
+        self.partitioned_numel = partitioned_numel
+        self.shape = shape
+        self.dtype = self.flat_groups[0][0].dtype
+
+    def contiguous(self):
+        """
+        Merge partitioned weights from flat_groups into a single tensor.
+        """
+        end_idx = self.offset + self.partitioned_numel
+        world_size = len(self.flat_groups)
+        pad_flat_param_chunks = []
+
+        for rank_i in range(world_size):
+            # for each rank, we need to collect weights from related group/groups
+            flat_groups_at_rank_i = self.flat_groups[rank_i]
+            start_group_id = None
+            end_group_id = None
+            for group_id in range(len(self.flat_groups_offset)):
+                if self.flat_groups_offset[group_id] <= self.offset < self.flat_groups_offset[group_id + 1]:
+                    start_group_id = group_id
+                if self.flat_groups_offset[group_id] < end_idx <= self.flat_groups_offset[group_id + 1]:
+                    end_group_id = group_id
+                    break
+            # collect weights from related group/groups
+            for group_id in range(start_group_id, end_group_id + 1):
+                flat_tensor = flat_groups_at_rank_i[group_id]
+                start_offset = self.offset - self.flat_groups_offset[group_id]
+                end_offset = min(end_idx, self.flat_groups_offset[group_id + 1]) - self.flat_groups_offset[group_id]
+                pad_flat_param_chunks.append(flat_tensor[start_offset:end_offset])
+
+        # collect weights from all ranks
+        pad_flat_param = torch.cat(pad_flat_param_chunks, dim=0)
+        param = pad_flat_param[:self.shape.numel()].view(self.shape).contiguous()
+        return param
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = sum([flat_group.numel() for flat_group in fp32_flat_groups[0]]) * world_size
+
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    flat_groups_offset = [0] + list(np.cumsum([flat_tensor.numel() for flat_tensor in fp32_flat_groups[0]]))
+    for name, shape in tqdm(param_shapes.items(), desc='Gathering sharded weights'):
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # memory efficient tensor
+        tensor = GatheredTensor(fp32_flat_groups, flat_groups_offset, offset, partitioned_numel, shape)
+        state_dict[name] = tensor
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states,
+                                               exclude_frozen_parameters):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    if not exclude_frozen_parameters:
+        _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def to_torch_tensor(state_dict, return_empty_tensor=False):
+    """
+    Convert state_dict of GatheredTensor to torch tensor
+    """
+    torch_state_dict = {}
+    converted_tensors = {}
+    for name, tensor in state_dict.items():
+        tensor_id = id(tensor)
+        if tensor_id in converted_tensors:  # shared tensors
+            shared_tensor = torch_state_dict[converted_tensors[tensor_id]]
+            torch_state_dict[name] = shared_tensor
+        else:
+            converted_tensors[tensor_id] = name
+            if return_empty_tensor:
+                torch_state_dict[name] = torch.empty(tensor.shape, dtype=tensor.dtype)
+            else:
+                torch_state_dict[name] = tensor.contiguous()
+    return torch_state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                             tag=None,
+                                             exclude_frozen_parameters=False,
+                                             lazy_mode=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+        - ``lazy_mode``: get state_dict in lazy mode. It returns a dict of pesduo tensor instead of torch tensor, which is more memory efficient.
+          Convert the pesduo tensor to torch tensor by ``.contiguous()``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    Note: the above usage may not work if your application doesn't have sufficient free CPU memory.
+    You may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint. Or you can load state_dict in lazy mode ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, lazy_mode=True) # not on cpu
+        for name, lazy_tensor in state_dict.item():
+            tensor = lazy_tensor.contiguous()  # to cpu
+            print(name, tensor)
+            # del tensor to release memory if it no longer in use
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    state_dict = _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir, exclude_frozen_parameters)
+    if lazy_mode:
+        return state_dict
+    else:
+        return to_torch_tensor(state_dict)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir,
+                                               output_dir,
+                                               max_shard_size="5GB",
+                                               safe_serialization=False,
+                                               tag=None,
+                                               exclude_frozen_parameters=False):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_dir``: directory to the pytorch fp32 state_dict output files
+        - ``max_shard_size``: the maximum size for a checkpoint before being sharded, default value is 5GB
+        - ``safe_serialization``:  whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+        - ``exclude_frozen_parameters``: exclude frozen parameters
+    """
+
+    # Dependency pre-check
+    if safe_serialization:
+        try:
+            from safetensors.torch import save_file
+        except ImportError:
+            print('If you want to use `safe_serialization`, please `pip install safetensors`')
+            raise
+    if max_shard_size is not None:
+        try:
+            from huggingface_hub import split_torch_state_dict_into_shards
+        except ImportError:
+            print('If you want to use `max_shard_size`, please `pip install huggingface_hub`')
+            raise
+
+    # Convert zero checkpoint to state_dict
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir,
+                                                          tag,
+                                                          exclude_frozen_parameters,
+                                                          lazy_mode=True)
+
+    # Shard the model if it is too big.
+    weights_name = "model.safetensors" if safe_serialization else "pytorch_model.bin"
+    if max_shard_size is not None:
+        filename_pattern = weights_name.replace(".bin", "{suffix}.bin").replace(".safetensors", "{suffix}.safetensors")
+        # an memory-efficient approach for sharding
+        empty_state_dict = to_torch_tensor(state_dict, return_empty_tensor=True)
+        state_dict_split = split_torch_state_dict_into_shards(empty_state_dict,
+                                                              filename_pattern=filename_pattern,
+                                                              max_shard_size=max_shard_size)
+    else:
+        from collections import namedtuple
+        StateDictSplit = namedtuple("StateDictSplit", ["is_sharded", "filename_to_tensors"])
+        state_dict_split = StateDictSplit(is_sharded=False,
+                                          filename_to_tensors={weights_name: list(state_dict.keys())})
+
+    # Save the model by shard
+    os.makedirs(output_dir, exist_ok=True)
+    filename_to_tensors = state_dict_split.filename_to_tensors.items()
+    for shard_file, tensors in tqdm(filename_to_tensors, desc="Saving checkpoint shards"):
+        shard_state_dict = {tensor_name: state_dict[tensor_name] for tensor_name in tensors}
+        shard_state_dict = to_torch_tensor(shard_state_dict)
+        output_path = os.path.join(output_dir, shard_file)
+        if safe_serialization:
+            save_file(shard_state_dict, output_path, metadata={"format": "pt"})
+        else:
+            torch.save(shard_state_dict, output_path)
+        # release the memory of current shard
+        for tensor_name in list(shard_state_dict.keys()):
+            del state_dict[tensor_name]
+            del shard_state_dict[tensor_name]
+        del shard_state_dict
+        gc.collect()
+
+    # Save index if sharded
+    if state_dict_split.is_sharded:
+        index = {
+            "metadata": state_dict_split.metadata,
+            "weight_map": state_dict_split.tensor_to_filename,
+        }
+        save_index_file = "model.safetensors.index.json" if safe_serialization else "pytorch_model.bin.index.json"
+        save_index_file = os.path.join(output_dir, save_index_file)
+        with open(save_index_file, "w", encoding="utf-8") as f:
+            content = json.dumps(index, indent=2, sort_keys=True) + "\n"
+            f.write(content)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info("Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info("Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument("output_dir",
+                        type=str,
+                        help="directory to the pytorch fp32 state_dict output files"
+                        "(e.g. path/checkpoint-12-output/)")
+    parser.add_argument(
+        "--max_shard_size",
+        type=str,
+        default="5GB",
+        help="The maximum size for a checkpoint before being sharded. Checkpoints shard will then be each of size"
+        "lower than this size. If expressed as a string, needs to be digits followed by a unit (like `5MB`"
+        "We default it to 5GB in order for models to be able to run easily on free-tier google colab instances"
+        "without CPU OOM issues.")
+    parser.add_argument(
+        "--safe_serialization",
+        default=False,
+        action='store_true',
+        help="Whether to save the model using `safetensors` or the traditional PyTorch way (that uses `pickle`).")
+    parser.add_argument("-t",
+                        "--tag",
+                        type=str,
+                        default=None,
+                        help="checkpoint tag used as a unique identifier for checkpoint. e.g., global_step1")
+    parser.add_argument("--exclude_frozen_parameters", action='store_true', help="exclude frozen parameters")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir,
+                                               args.output_dir,
+                                               max_shard_size=args.max_shard_size,
+                                               safe_serialization=args.safe_serialization,
+                                               tag=args.tag,
+                                               exclude_frozen_parameters=args.exclude_frozen_parameters)