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,104 @@
----
-license: mit
----
+---
+license: apache-2.0
+base_model: Qwen/Qwen3-4B
+tags:
+  - SAT
+  - combinatorial-optimization
+  - classification
+  - cube-and-conquer
+language:
+  - en
+pipeline_tag: text-classification
+---
+
+# Qwen3-4B-SAT-VarSelector
+
+A Qwen3-4B model fine-tuned for **SAT branching variable selection** in Cube-and-Conquer (CnC) solvers.
+
+## Model Description
+
+This model predicts which variable to branch/cube on next, given a SAT CNF formula state. Instead of generating text, it outputs a **classification over variable IDs** (1-500).
+
+### Architecture
+
+- **Base**: `Qwen/Qwen3-4B` (causal language model)
+- **Head**: LayerNorm → Linear(hidden_size, 501)
+- **Pooling**: Last non-pad token hidden state
+- **Masking**: Invalid variables (not in CNF) are masked to -10000 before softmax
+
+### Training
+
+- **Dataset**: 3,898 training / 434 validation samples
+- **Task**: Predict expert-selected branching variable
+- **Best validation accuracy**: 16.36% (16x better than random ~1%)
+- **Training**: 8 epochs, 8×H100 GPUs, DeepSpeed ZeRO-3
+
+## 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=500)
+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=500)], 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)
+- `inference_demo.py` - Example inference script
+
+## Metrics
+
+| Metric | Value |
+|--------|-------|
+| Validation Accuracy | 16.36% |
+| Validation Loss | 3.87 |
+| Random Baseline | ~1% |
+| Improvement | 16x |
+
+## Limitations
+
+- Maximum 500 variables
+- Maximum 8192 tokens for CNF input
+- Trained on specific CNF distribution (may not generalize to all SAT instances)
+
+## 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_step440/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:02c35b3103bdd6d1e454a19691d8eaa8fe9d5760adf760caf0c7b111bebc046f
+size 6035639921

global_step440/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:08ecf013ea2d99840f2ac2171762ff5d925629b1a54a72d6f69a6e1976a48ac6
+size 6035639921

global_step440/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:df72b0c2e5a1057f7656aab1e15b633f6ba4fd14e8b1bfe49947308053835ffb
+size 6035639921

global_step440/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:10681e851bf7b379ebc47c322a89f47598eb8b4b88bfbf061ed84349370c960f
+size 6035639921

global_step440/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:900494879db0429dcef1ac17e5c620440428b4f02e657d69ef2d49dc283a7ef7
+size 6035639921

global_step440/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:3f295bc241d107b583819043ddb9ac92ec6ae2cf3181677a8c8ec9c1d951ba03
+size 6035639921

global_step440/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:f01e8b4d005cb73765ada0b2c05ef68385a047935475e525894c6f535fb7bfe8
+size 6035639921

global_step440/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:1317e2cd5882643804b4609614c2d4498e5b2952a88608e30bb6e5a80d2c1c07
+size 6035639921

global_step440/zero_pp_rank_0_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_1_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_2_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_3_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_4_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_5_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_6_mp_rank_00_model_states.pt

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

global_step440/zero_pp_rank_7_mp_rank_00_model_states.pt

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

latest

@@ -0,0 +1 @@
+global_step440

pytorch_model.bin

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

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

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

rng_state_4.pth

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

rng_state_5.pth

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

rng_state_6.pth

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

rng_state_7.pth

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

scheduler.pt

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

sft_qwen_var_classifier.py

@@ -0,0 +1,725 @@
+"""
+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")
+    
+    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!
+    trainer.train()
+    
+    # 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,1134 @@
+{
+  "best_global_step": 440,
+  "best_metric": 0.16359447004608296,
+  "best_model_checkpoint": "out_qwen_var_sft/checkpoint-440",
+  "epoch": 7.213114754098361,
+  "eval_steps": 10,
+  "global_step": 440,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.08196721311475409,
+      "grad_norm": 2001.4421126814825,
+      "learning_rate": 4.0816326530612243e-07,
+      "loss": 42.6172,
+      "step": 5,
+      "true_loss": 5.0281
+    },
+    {
+      "epoch": 0.16393442622950818,
+      "grad_norm": 158.37906608693382,
+      "learning_rate": 9.183673469387756e-07,
+      "loss": 43.1594,
+      "step": 10,
+      "true_loss": 5.1953
+    },
+    {
+      "epoch": 0.16393442622950818,
+      "eval_accuracy": 0.01152073732718894,
+      "eval_loss": 5.286218166351318,
+      "eval_runtime": 15.0885,
+      "eval_samples_per_second": 28.764,
+      "eval_steps_per_second": 3.645,
+      "step": 10
+    },
+    {
+      "epoch": 0.2459016393442623,
+      "grad_norm": 203.23818895534905,
+      "learning_rate": 1.4285714285714286e-06,
+      "loss": 42.6562,
+      "step": 15,
+      "true_loss": 5.2555
+    },
+    {
+      "epoch": 0.32786885245901637,
+      "grad_norm": 176.6457829157802,
+      "learning_rate": 1.938775510204082e-06,
+      "loss": 42.1543,
+      "step": 20,
+      "true_loss": 5.1922
+    },
+    {
+      "epoch": 0.32786885245901637,
+      "eval_accuracy": 0.01152073732718894,
+      "eval_loss": 5.156322002410889,
+      "eval_runtime": 14.9518,
+      "eval_samples_per_second": 29.027,
+      "eval_steps_per_second": 3.678,
+      "step": 20
+    },
+    {
+      "epoch": 0.4098360655737705,
+      "grad_norm": 157.3725914328133,
+      "learning_rate": 2.4489795918367347e-06,
+      "loss": 41.0754,
+      "step": 25,
+      "true_loss": 4.9437
+    },
+    {
+      "epoch": 0.4918032786885246,
+      "grad_norm": 177.5762016006674,
+      "learning_rate": 2.959183673469388e-06,
+      "loss": 38.793,
+      "step": 30,
+      "true_loss": 5.1707
+    },
+    {
+      "epoch": 0.4918032786885246,
+      "eval_accuracy": 0.02304147465437788,
+      "eval_loss": 4.853038787841797,
+      "eval_runtime": 14.9796,
+      "eval_samples_per_second": 28.973,
+      "eval_steps_per_second": 3.672,
+      "step": 30
+    },
+    {
+      "epoch": 0.5737704918032787,
+      "grad_norm": 224.10088573018953,
+      "learning_rate": 3.469387755102041e-06,
+      "loss": 38.6121,
+      "step": 35,
+      "true_loss": 4.7551
+    },
+    {
+      "epoch": 0.6557377049180327,
+      "grad_norm": 302.8526334743389,
+      "learning_rate": 3.979591836734694e-06,
+      "loss": 37.8027,
+      "step": 40,
+      "true_loss": 4.9977
+    },
+    {
+      "epoch": 0.6557377049180327,
+      "eval_accuracy": 0.06221198156682028,
+      "eval_loss": 4.661470413208008,
+      "eval_runtime": 15.006,
+      "eval_samples_per_second": 28.922,
+      "eval_steps_per_second": 3.665,
+      "step": 40
+    },
+    {
+      "epoch": 0.7377049180327869,
+      "grad_norm": 177.15837283975523,
+      "learning_rate": 4.489795918367348e-06,
+      "loss": 36.4664,
+      "step": 45,
+      "true_loss": 4.5766
+    },
+    {
+      "epoch": 0.819672131147541,
+      "grad_norm": 190.75273796127124,
+      "learning_rate": 5e-06,
+      "loss": 36.8304,
+      "step": 50,
+      "true_loss": 4.7605
+    },
+    {
+      "epoch": 0.819672131147541,
+      "eval_accuracy": 0.052995391705069124,
+      "eval_loss": 4.522649765014648,
+      "eval_runtime": 14.8422,
+      "eval_samples_per_second": 29.241,
+      "eval_steps_per_second": 3.706,
+      "step": 50
+    },
+    {
+      "epoch": 0.9016393442622951,
+      "grad_norm": 185.07218548511187,
+      "learning_rate": 4.943052391799545e-06,
+      "loss": 36.585,
+      "step": 55,
+      "true_loss": 4.5461
+    },
+    {
+      "epoch": 0.9836065573770492,
+      "grad_norm": 169.35214138176565,
+      "learning_rate": 4.886104783599089e-06,
+      "loss": 35.3022,
+      "step": 60,
+      "true_loss": 4.3983
+    },
+    {
+      "epoch": 0.9836065573770492,
+      "eval_accuracy": 0.06221198156682028,
+      "eval_loss": 4.3924431800842285,
+      "eval_runtime": 14.9823,
+      "eval_samples_per_second": 28.967,
+      "eval_steps_per_second": 3.671,
+      "step": 60
+    },
+    {
+      "epoch": 1.0655737704918034,
+      "grad_norm": 186.19471703324027,
+      "learning_rate": 4.829157175398634e-06,
+      "loss": 34.9871,
+      "step": 65,
+      "true_loss": 4.534
+    },
+    {
+      "epoch": 1.1475409836065573,
+      "grad_norm": 178.50065430802542,
+      "learning_rate": 4.772209567198178e-06,
+      "loss": 34.8587,
+      "step": 70,
+      "true_loss": 4.0799
+    },
+    {
+      "epoch": 1.1475409836065573,
+      "eval_accuracy": 0.06682027649769585,
+      "eval_loss": 4.316955089569092,
+      "eval_runtime": 14.9903,
+      "eval_samples_per_second": 28.952,
+      "eval_steps_per_second": 3.669,
+      "step": 70
+    },
+    {
+      "epoch": 1.2295081967213115,
+      "grad_norm": 256.1945567078422,
+      "learning_rate": 4.7152619589977225e-06,
+      "loss": 33.4347,
+      "step": 75,
+      "true_loss": 3.7832
+    },
+    {
+      "epoch": 1.3114754098360657,
+      "grad_norm": 178.23899873402075,
+      "learning_rate": 4.658314350797267e-06,
+      "loss": 33.5913,
+      "step": 80,
+      "true_loss": 4.3762
+    },
+    {
+      "epoch": 1.3114754098360657,
+      "eval_accuracy": 0.08064516129032258,
+      "eval_loss": 4.25508975982666,
+      "eval_runtime": 14.9497,
+      "eval_samples_per_second": 29.031,
+      "eval_steps_per_second": 3.679,
+      "step": 80
+    },
+    {
+      "epoch": 1.3934426229508197,
+      "grad_norm": 176.89010340519602,
+      "learning_rate": 4.601366742596811e-06,
+      "loss": 33.1532,
+      "step": 85,
+      "true_loss": 4.1607
+    },
+    {
+      "epoch": 1.4754098360655736,
+      "grad_norm": 146.21731787991052,
+      "learning_rate": 4.544419134396356e-06,
+      "loss": 33.0214,
+      "step": 90,
+      "true_loss": 4.3359
+    },
+    {
+      "epoch": 1.4754098360655736,
+      "eval_accuracy": 0.0967741935483871,
+      "eval_loss": 4.17578125,
+      "eval_runtime": 14.9139,
+      "eval_samples_per_second": 29.1,
+      "eval_steps_per_second": 3.688,
+      "step": 90
+    },
+    {
+      "epoch": 1.5573770491803278,
+      "grad_norm": 163.96782239209782,
+      "learning_rate": 4.4874715261959e-06,
+      "loss": 31.8735,
+      "step": 95,
+      "true_loss": 4.1047
+    },
+    {
+      "epoch": 1.639344262295082,
+      "grad_norm": 172.0759692604596,
+      "learning_rate": 4.4305239179954446e-06,
+      "loss": 33.1337,
+      "step": 100,
+      "true_loss": 4.4055
+    },
+    {
+      "epoch": 1.639344262295082,
+      "eval_accuracy": 0.09907834101382489,
+      "eval_loss": 4.103317737579346,
+      "eval_runtime": 15.1423,
+      "eval_samples_per_second": 28.661,
+      "eval_steps_per_second": 3.632,
+      "step": 100
+    },
+    {
+      "epoch": 1.721311475409836,
+      "grad_norm": 163.50986481723115,
+      "learning_rate": 4.373576309794989e-06,
+      "loss": 30.4056,
+      "step": 105,
+      "true_loss": 3.5892
+    },
+    {
+      "epoch": 1.8032786885245902,
+      "grad_norm": 174.28722424027927,
+      "learning_rate": 4.316628701594533e-06,
+      "loss": 32.4443,
+      "step": 110,
+      "true_loss": 3.8227
+    },
+    {
+      "epoch": 1.8032786885245902,
+      "eval_accuracy": 0.11059907834101383,
+      "eval_loss": 4.104156970977783,
+      "eval_runtime": 14.9928,
+      "eval_samples_per_second": 28.947,
+      "eval_steps_per_second": 3.668,
+      "step": 110
+    },
+    {
+      "epoch": 1.8852459016393444,
+      "grad_norm": 599.8963091372194,
+      "learning_rate": 4.259681093394078e-06,
+      "loss": 32.7365,
+      "step": 115,
+      "true_loss": 4.3132
+    },
+    {
+      "epoch": 1.9672131147540983,
+      "grad_norm": 176.15071401035615,
+      "learning_rate": 4.202733485193622e-06,
+      "loss": 30.8238,
+      "step": 120,
+      "true_loss": 3.9128
+    },
+    {
+      "epoch": 1.9672131147540983,
+      "eval_accuracy": 0.09216589861751152,
+      "eval_loss": 4.044179439544678,
+      "eval_runtime": 14.9286,
+      "eval_samples_per_second": 29.072,
+      "eval_steps_per_second": 3.684,
+      "step": 120
+    },
+    {
+      "epoch": 2.0491803278688523,
+      "grad_norm": 193.228785849168,
+      "learning_rate": 4.145785876993167e-06,
+      "loss": 31.3295,
+      "step": 125,
+      "true_loss": 4.0655
+    },
+    {
+      "epoch": 2.1311475409836067,
+      "grad_norm": 167.07209609143695,
+      "learning_rate": 4.088838268792711e-06,
+      "loss": 29.0552,
+      "step": 130,
+      "true_loss": 3.727
+    },
+    {
+      "epoch": 2.1311475409836067,
+      "eval_accuracy": 0.12211981566820276,
+      "eval_loss": 4.029242992401123,
+      "eval_runtime": 15.3137,
+      "eval_samples_per_second": 28.341,
+      "eval_steps_per_second": 3.592,
+      "step": 130
+    },
+    {
+      "epoch": 2.2131147540983607,
+      "grad_norm": 208.65171046196372,
+      "learning_rate": 4.0318906605922555e-06,
+      "loss": 28.6111,
+      "step": 135,
+      "true_loss": 3.5737
+    },
+    {
+      "epoch": 2.2950819672131146,
+      "grad_norm": 198.8811989422508,
+      "learning_rate": 3.9749430523918e-06,
+      "loss": 28.6129,
+      "step": 140,
+      "true_loss": 3.4636
+    },
+    {
+      "epoch": 2.2950819672131146,
+      "eval_accuracy": 0.1152073732718894,
+      "eval_loss": 4.056672096252441,
+      "eval_runtime": 14.9575,
+      "eval_samples_per_second": 29.016,
+      "eval_steps_per_second": 3.677,
+      "step": 140
+    },
+    {
+      "epoch": 2.3770491803278686,
+      "grad_norm": 199.51861074084005,
+      "learning_rate": 3.917995444191344e-06,
+      "loss": 28.8025,
+      "step": 145,
+      "true_loss": 3.5059
+    },
+    {
+      "epoch": 2.459016393442623,
+      "grad_norm": 193.67355565121787,
+      "learning_rate": 3.861047835990889e-06,
+      "loss": 28.5974,
+      "step": 150,
+      "true_loss": 2.9942
+    },
+    {
+      "epoch": 2.459016393442623,
+      "eval_accuracy": 0.0967741935483871,
+      "eval_loss": 4.023684978485107,
+      "eval_runtime": 15.0321,
+      "eval_samples_per_second": 28.872,
+      "eval_steps_per_second": 3.659,
+      "step": 150
+    },
+    {
+      "epoch": 2.540983606557377,
+      "grad_norm": 205.5841428370514,
+      "learning_rate": 3.804100227790433e-06,
+      "loss": 28.4969,
+      "step": 155,
+      "true_loss": 3.5412
+    },
+    {
+      "epoch": 2.6229508196721314,
+      "grad_norm": 182.43691547474708,
+      "learning_rate": 3.7471526195899776e-06,
+      "loss": 29.1582,
+      "step": 160,
+      "true_loss": 3.5827
+    },
+    {
+      "epoch": 2.6229508196721314,
+      "eval_accuracy": 0.10368663594470046,
+      "eval_loss": 4.006450176239014,
+      "eval_runtime": 15.0347,
+      "eval_samples_per_second": 28.867,
+      "eval_steps_per_second": 3.658,
+      "step": 160
+    },
+    {
+      "epoch": 2.7049180327868854,
+      "grad_norm": 209.23223838323548,
+      "learning_rate": 3.690205011389522e-06,
+      "loss": 28.4554,
+      "step": 165,
+      "true_loss": 3.8515
+    },
+    {
+      "epoch": 2.7868852459016393,
+      "grad_norm": 181.70496603318674,
+      "learning_rate": 3.6332574031890664e-06,
+      "loss": 28.0033,
+      "step": 170,
+      "true_loss": 3.8595
+    },
+    {
+      "epoch": 2.7868852459016393,
+      "eval_accuracy": 0.1336405529953917,
+      "eval_loss": 3.9438364505767822,
+      "eval_runtime": 15.1184,
+      "eval_samples_per_second": 28.707,
+      "eval_steps_per_second": 3.638,
+      "step": 170
+    },
+    {
+      "epoch": 2.8688524590163933,
+      "grad_norm": 189.43537999115608,
+      "learning_rate": 3.5763097949886104e-06,
+      "loss": 27.3218,
+      "step": 175,
+      "true_loss": 3.4404
+    },
+    {
+      "epoch": 2.9508196721311473,
+      "grad_norm": 203.83162067201448,
+      "learning_rate": 3.519362186788155e-06,
+      "loss": 28.3321,
+      "step": 180,
+      "true_loss": 3.1337
+    },
+    {
+      "epoch": 2.9508196721311473,
+      "eval_accuracy": 0.10829493087557604,
+      "eval_loss": 3.9445700645446777,
+      "eval_runtime": 15.1698,
+      "eval_samples_per_second": 28.61,
+      "eval_steps_per_second": 3.626,
+      "step": 180
+    },
+    {
+      "epoch": 3.0327868852459017,
+      "grad_norm": 188.7733943154734,
+      "learning_rate": 3.4624145785876997e-06,
+      "loss": 27.27,
+      "step": 185,
+      "true_loss": 3.334
+    },
+    {
+      "epoch": 3.1147540983606556,
+      "grad_norm": 299.860020893839,
+      "learning_rate": 3.405466970387244e-06,
+      "loss": 23.8397,
+      "step": 190,
+      "true_loss": 2.6139
+    },
+    {
+      "epoch": 3.1147540983606556,
+      "eval_accuracy": 0.11059907834101383,
+      "eval_loss": 3.9798836708068848,
+      "eval_runtime": 15.4586,
+      "eval_samples_per_second": 28.075,
+      "eval_steps_per_second": 3.558,
+      "step": 190
+    },
+    {
+      "epoch": 3.19672131147541,
+      "grad_norm": 261.2265032273993,
+      "learning_rate": 3.3485193621867885e-06,
+      "loss": 24.7364,
+      "step": 195,
+      "true_loss": 3.0539
+    },
+    {
+      "epoch": 3.278688524590164,
+      "grad_norm": 244.13814588236283,
+      "learning_rate": 3.291571753986333e-06,
+      "loss": 25.6726,
+      "step": 200,
+      "true_loss": 3.0328
+    },
+    {
+      "epoch": 3.278688524590164,
+      "eval_accuracy": 0.11981566820276497,
+      "eval_loss": 3.931368350982666,
+      "eval_runtime": 14.9343,
+      "eval_samples_per_second": 29.061,
+      "eval_steps_per_second": 3.683,
+      "step": 200
+    },
+    {
+      "epoch": 3.360655737704918,
+      "grad_norm": 223.86848695043577,
+      "learning_rate": 3.2346241457858773e-06,
+      "loss": 24.6349,
+      "step": 205,
+      "true_loss": 2.9529
+    },
+    {
+      "epoch": 3.442622950819672,
+      "grad_norm": 239.830671087524,
+      "learning_rate": 3.1776765375854217e-06,
+      "loss": 24.5938,
+      "step": 210,
+      "true_loss": 3.3097
+    },
+    {
+      "epoch": 3.442622950819672,
+      "eval_accuracy": 0.1175115207373272,
+      "eval_loss": 3.9209704399108887,
+      "eval_runtime": 14.9774,
+      "eval_samples_per_second": 28.977,
+      "eval_steps_per_second": 3.672,
+      "step": 210
+    },
+    {
+      "epoch": 3.5245901639344264,
+      "grad_norm": 188.8971324623384,
+      "learning_rate": 3.120728929384966e-06,
+      "loss": 24.3125,
+      "step": 215,
+      "true_loss": 2.8892
+    },
+    {
+      "epoch": 3.6065573770491803,
+      "grad_norm": 211.46305164676468,
+      "learning_rate": 3.0637813211845106e-06,
+      "loss": 24.9786,
+      "step": 220,
+      "true_loss": 3.4745
+    },
+    {
+      "epoch": 3.6065573770491803,
+      "eval_accuracy": 0.11290322580645161,
+      "eval_loss": 3.892094373703003,
+      "eval_runtime": 15.0091,
+      "eval_samples_per_second": 28.916,
+      "eval_steps_per_second": 3.664,
+      "step": 220
+    },
+    {
+      "epoch": 3.6885245901639343,
+      "grad_norm": 209.57415125751984,
+      "learning_rate": 3.0068337129840546e-06,
+      "loss": 24.4776,
+      "step": 225,
+      "true_loss": 3.0013
+    },
+    {
+      "epoch": 3.7704918032786887,
+      "grad_norm": 219.2212947656317,
+      "learning_rate": 2.949886104783599e-06,
+      "loss": 25.0714,
+      "step": 230,
+      "true_loss": 3.3653
+    },
+    {
+      "epoch": 3.7704918032786887,
+      "eval_accuracy": 0.12211981566820276,
+      "eval_loss": 3.88911509513855,
+      "eval_runtime": 15.0058,
+      "eval_samples_per_second": 28.922,
+      "eval_steps_per_second": 3.665,
+      "step": 230
+    },
+    {
+      "epoch": 3.8524590163934427,
+      "grad_norm": 215.0136859889706,
+      "learning_rate": 2.892938496583144e-06,
+      "loss": 25.0681,
+      "step": 235,
+      "true_loss": 2.7688
+    },
+    {
+      "epoch": 3.9344262295081966,
+      "grad_norm": 212.6143565856701,
+      "learning_rate": 2.8359908883826882e-06,
+      "loss": 24.524,
+      "step": 240,
+      "true_loss": 3.0097
+    },
+    {
+      "epoch": 3.9344262295081966,
+      "eval_accuracy": 0.10599078341013825,
+      "eval_loss": 3.901205062866211,
+      "eval_runtime": 14.9742,
+      "eval_samples_per_second": 28.983,
+      "eval_steps_per_second": 3.673,
+      "step": 240
+    },
+    {
+      "epoch": 4.016393442622951,
+      "grad_norm": 203.1242362392191,
+      "learning_rate": 2.7790432801822326e-06,
+      "loss": 22.3977,
+      "step": 245,
+      "true_loss": 2.4811
+    },
+    {
+      "epoch": 4.098360655737705,
+      "grad_norm": 219.93955042504686,
+      "learning_rate": 2.722095671981777e-06,
+      "loss": 20.6891,
+      "step": 250,
+      "true_loss": 2.6231
+    },
+    {
+      "epoch": 4.098360655737705,
+      "eval_accuracy": 0.12903225806451613,
+      "eval_loss": 3.918625831604004,
+      "eval_runtime": 15.0247,
+      "eval_samples_per_second": 28.886,
+      "eval_steps_per_second": 3.661,
+      "step": 250
+    },
+    {
+      "epoch": 4.180327868852459,
+      "grad_norm": 269.83943945736723,
+      "learning_rate": 2.6651480637813215e-06,
+      "loss": 19.9576,
+      "step": 255,
+      "true_loss": 2.4069
+    },
+    {
+      "epoch": 4.262295081967213,
+      "grad_norm": 285.3088807670382,
+      "learning_rate": 2.608200455580866e-06,
+      "loss": 19.9079,
+      "step": 260,
+      "true_loss": 2.5739
+    },
+    {
+      "epoch": 4.262295081967213,
+      "eval_accuracy": 0.12672811059907835,
+      "eval_loss": 3.955899477005005,
+      "eval_runtime": 14.9722,
+      "eval_samples_per_second": 28.987,
+      "eval_steps_per_second": 3.673,
+      "step": 260
+    },
+    {
+      "epoch": 4.344262295081967,
+      "grad_norm": 235.65294642592144,
+      "learning_rate": 2.5512528473804103e-06,
+      "loss": 19.9913,
+      "step": 265,
+      "true_loss": 2.6661
+    },
+    {
+      "epoch": 4.426229508196721,
+      "grad_norm": 243.25746669728042,
+      "learning_rate": 2.4943052391799547e-06,
+      "loss": 21.5919,
+      "step": 270,
+      "true_loss": 2.931
+    },
+    {
+      "epoch": 4.426229508196721,
+      "eval_accuracy": 0.12672811059907835,
+      "eval_loss": 3.8867075443267822,
+      "eval_runtime": 14.9423,
+      "eval_samples_per_second": 29.045,
+      "eval_steps_per_second": 3.681,
+      "step": 270
+    },
+    {
+      "epoch": 4.508196721311475,
+      "grad_norm": 268.3143937900749,
+      "learning_rate": 2.437357630979499e-06,
+      "loss": 21.3721,
+      "step": 275,
+      "true_loss": 2.6628
+    },
+    {
+      "epoch": 4.590163934426229,
+      "grad_norm": 233.6702831549346,
+      "learning_rate": 2.3804100227790436e-06,
+      "loss": 21.651,
+      "step": 280,
+      "true_loss": 2.3591
+    },
+    {
+      "epoch": 4.590163934426229,
+      "eval_accuracy": 0.1359447004608295,
+      "eval_loss": 3.855499029159546,
+      "eval_runtime": 15.1421,
+      "eval_samples_per_second": 28.662,
+      "eval_steps_per_second": 3.632,
+      "step": 280
+    },
+    {
+      "epoch": 4.672131147540983,
+      "grad_norm": 219.0257023920148,
+      "learning_rate": 2.323462414578588e-06,
+      "loss": 21.5054,
+      "step": 285,
+      "true_loss": 2.5704
+    },
+    {
+      "epoch": 4.754098360655737,
+      "grad_norm": 287.08072005658994,
+      "learning_rate": 2.2665148063781324e-06,
+      "loss": 20.7173,
+      "step": 290,
+      "true_loss": 3.0756
+    },
+    {
+      "epoch": 4.754098360655737,
+      "eval_accuracy": 0.14055299539170507,
+      "eval_loss": 3.8604798316955566,
+      "eval_runtime": 15.1432,
+      "eval_samples_per_second": 28.66,
+      "eval_steps_per_second": 3.632,
+      "step": 290
+    },
+    {
+      "epoch": 4.836065573770492,
+      "grad_norm": 264.8569129204395,
+      "learning_rate": 2.209567198177677e-06,
+      "loss": 19.9181,
+      "step": 295,
+      "true_loss": 2.2299
+    },
+    {
+      "epoch": 4.918032786885246,
+      "grad_norm": 239.4181479000276,
+      "learning_rate": 2.1526195899772212e-06,
+      "loss": 22.0119,
+      "step": 300,
+      "true_loss": 3.0984
+    },
+    {
+      "epoch": 4.918032786885246,
+      "eval_accuracy": 0.1313364055299539,
+      "eval_loss": 3.848721742630005,
+      "eval_runtime": 15.1331,
+      "eval_samples_per_second": 28.679,
+      "eval_steps_per_second": 3.634,
+      "step": 300
+    },
+    {
+      "epoch": 5.0,
+      "grad_norm": 216.31247846676033,
+      "learning_rate": 2.0956719817767656e-06,
+      "loss": 20.7322,
+      "step": 305,
+      "true_loss": 2.2975
+    },
+    {
+      "epoch": 5.081967213114754,
+      "grad_norm": 266.0474875386491,
+      "learning_rate": 2.03872437357631e-06,
+      "loss": 17.2354,
+      "step": 310,
+      "true_loss": 2.0884
+    },
+    {
+      "epoch": 5.081967213114754,
+      "eval_accuracy": 0.1382488479262673,
+      "eval_loss": 3.835157871246338,
+      "eval_runtime": 15.1115,
+      "eval_samples_per_second": 28.72,
+      "eval_steps_per_second": 3.64,
+      "step": 310
+    },
+    {
+      "epoch": 5.163934426229508,
+      "grad_norm": 258.5339989822872,
+      "learning_rate": 1.9817767653758545e-06,
+      "loss": 17.442,
+      "step": 315,
+      "true_loss": 2.1051
+    },
+    {
+      "epoch": 5.245901639344262,
+      "grad_norm": 308.2469254205616,
+      "learning_rate": 1.924829157175399e-06,
+      "loss": 17.0414,
+      "step": 320,
+      "true_loss": 2.5328
+    },
+    {
+      "epoch": 5.245901639344262,
+      "eval_accuracy": 0.15207373271889402,
+      "eval_loss": 3.899200201034546,
+      "eval_runtime": 15.7217,
+      "eval_samples_per_second": 27.605,
+      "eval_steps_per_second": 3.498,
+      "step": 320
+    },
+    {
+      "epoch": 5.327868852459017,
+      "grad_norm": 297.7046229142417,
+      "learning_rate": 1.8678815489749433e-06,
+      "loss": 16.4162,
+      "step": 325,
+      "true_loss": 2.2139
+    },
+    {
+      "epoch": 5.409836065573771,
+      "grad_norm": 262.17760383527263,
+      "learning_rate": 1.8109339407744877e-06,
+      "loss": 17.9956,
+      "step": 330,
+      "true_loss": 2.3699
+    },
+    {
+      "epoch": 5.409836065573771,
+      "eval_accuracy": 0.15668202764976957,
+      "eval_loss": 3.869380235671997,
+      "eval_runtime": 14.9258,
+      "eval_samples_per_second": 29.077,
+      "eval_steps_per_second": 3.685,
+      "step": 330
+    },
+    {
+      "epoch": 5.491803278688525,
+      "grad_norm": 267.855545172113,
+      "learning_rate": 1.753986332574032e-06,
+      "loss": 18.2262,
+      "step": 335,
+      "true_loss": 2.1083
+    },
+    {
+      "epoch": 5.573770491803279,
+      "grad_norm": 263.032871409478,
+      "learning_rate": 1.6970387243735763e-06,
+      "loss": 17.7507,
+      "step": 340,
+      "true_loss": 2.0572
+    },
+    {
+      "epoch": 5.573770491803279,
+      "eval_accuracy": 0.14516129032258066,
+      "eval_loss": 3.860694646835327,
+      "eval_runtime": 14.9606,
+      "eval_samples_per_second": 29.01,
+      "eval_steps_per_second": 3.676,
+      "step": 340
+    },
+    {
+      "epoch": 5.655737704918033,
+      "grad_norm": 251.13017134627248,
+      "learning_rate": 1.640091116173121e-06,
+      "loss": 18.9093,
+      "step": 345,
+      "true_loss": 2.552
+    },
+    {
+      "epoch": 5.737704918032787,
+      "grad_norm": 283.50724134889464,
+      "learning_rate": 1.5831435079726654e-06,
+      "loss": 17.0755,
+      "step": 350,
+      "true_loss": 2.621
+    },
+    {
+      "epoch": 5.737704918032787,
+      "eval_accuracy": 0.15668202764976957,
+      "eval_loss": 3.821519613265991,
+      "eval_runtime": 15.0984,
+      "eval_samples_per_second": 28.745,
+      "eval_steps_per_second": 3.643,
+      "step": 350
+    },
+    {
+      "epoch": 5.8196721311475414,
+      "grad_norm": 751.7678054126261,
+      "learning_rate": 1.5261958997722096e-06,
+      "loss": 18.1668,
+      "step": 355,
+      "true_loss": 1.9259
+    },
+    {
+      "epoch": 5.901639344262295,
+      "grad_norm": 318.53273858625954,
+      "learning_rate": 1.469248291571754e-06,
+      "loss": 17.0525,
+      "step": 360,
+      "true_loss": 2.2937
+    },
+    {
+      "epoch": 5.901639344262295,
+      "eval_accuracy": 0.1497695852534562,
+      "eval_loss": 3.812807083129883,
+      "eval_runtime": 15.0494,
+      "eval_samples_per_second": 28.838,
+      "eval_steps_per_second": 3.655,
+      "step": 360
+    },
+    {
+      "epoch": 5.983606557377049,
+      "grad_norm": 259.86569786349736,
+      "learning_rate": 1.4123006833712984e-06,
+      "loss": 16.5851,
+      "step": 365,
+      "true_loss": 1.926
+    },
+    {
+      "epoch": 6.065573770491803,
+      "grad_norm": 236.16491616521617,
+      "learning_rate": 1.355353075170843e-06,
+      "loss": 14.0549,
+      "step": 370,
+      "true_loss": 1.6638
+    },
+    {
+      "epoch": 6.065573770491803,
+      "eval_accuracy": 0.14516129032258066,
+      "eval_loss": 3.8400564193725586,
+      "eval_runtime": 15.0501,
+      "eval_samples_per_second": 28.837,
+      "eval_steps_per_second": 3.654,
+      "step": 370
+    },
+    {
+      "epoch": 6.147540983606557,
+      "grad_norm": 267.4679343061016,
+      "learning_rate": 1.2984054669703875e-06,
+      "loss": 13.1733,
+      "step": 375,
+      "true_loss": 1.4947
+    },
+    {
+      "epoch": 6.229508196721311,
+      "grad_norm": 303.4138473047156,
+      "learning_rate": 1.2414578587699317e-06,
+      "loss": 14.2548,
+      "step": 380,
+      "true_loss": 1.8772
+    },
+    {
+      "epoch": 6.229508196721311,
+      "eval_accuracy": 0.14516129032258066,
+      "eval_loss": 3.897352933883667,
+      "eval_runtime": 15.0678,
+      "eval_samples_per_second": 28.803,
+      "eval_steps_per_second": 3.65,
+      "step": 380
+    },
+    {
+      "epoch": 6.311475409836065,
+      "grad_norm": 304.09080880251156,
+      "learning_rate": 1.1845102505694763e-06,
+      "loss": 14.1155,
+      "step": 385,
+      "true_loss": 1.4
+    },
+    {
+      "epoch": 6.39344262295082,
+      "grad_norm": 328.52917098962524,
+      "learning_rate": 1.1275626423690205e-06,
+      "loss": 14.0041,
+      "step": 390,
+      "true_loss": 1.5224
+    },
+    {
+      "epoch": 6.39344262295082,
+      "eval_accuracy": 0.1313364055299539,
+      "eval_loss": 3.8768956661224365,
+      "eval_runtime": 14.9485,
+      "eval_samples_per_second": 29.033,
+      "eval_steps_per_second": 3.679,
+      "step": 390
+    },
+    {
+      "epoch": 6.475409836065574,
+      "grad_norm": 374.87828963261063,
+      "learning_rate": 1.0706150341685651e-06,
+      "loss": 13.4204,
+      "step": 395,
+      "true_loss": 1.7657
+    },
+    {
+      "epoch": 6.557377049180328,
+      "grad_norm": 290.7307487199588,
+      "learning_rate": 1.0136674259681093e-06,
+      "loss": 14.282,
+      "step": 400,
+      "true_loss": 1.8787
+    },
+    {
+      "epoch": 6.557377049180328,
+      "eval_accuracy": 0.14285714285714285,
+      "eval_loss": 3.83672833442688,
+      "eval_runtime": 15.0933,
+      "eval_samples_per_second": 28.755,
+      "eval_steps_per_second": 3.644,
+      "step": 400
+    },
+    {
+      "epoch": 6.639344262295082,
+      "grad_norm": 293.9907784197036,
+      "learning_rate": 9.567198177676538e-07,
+      "loss": 13.0058,
+      "step": 405,
+      "true_loss": 1.9678
+    },
+    {
+      "epoch": 6.721311475409836,
+      "grad_norm": 329.0206706011657,
+      "learning_rate": 8.997722095671982e-07,
+      "loss": 15.0129,
+      "step": 410,
+      "true_loss": 1.8787
+    },
+    {
+      "epoch": 6.721311475409836,
+      "eval_accuracy": 0.15668202764976957,
+      "eval_loss": 3.854370594024658,
+      "eval_runtime": 15.0162,
+      "eval_samples_per_second": 28.902,
+      "eval_steps_per_second": 3.663,
+      "step": 410
+    },
+    {
+      "epoch": 6.80327868852459,
+      "grad_norm": 435.30888223467264,
+      "learning_rate": 8.428246013667427e-07,
+      "loss": 14.3213,
+      "step": 415,
+      "true_loss": 2.0839
+    },
+    {
+      "epoch": 6.885245901639344,
+      "grad_norm": 283.4882127418948,
+      "learning_rate": 7.858769931662871e-07,
+      "loss": 13.6361,
+      "step": 420,
+      "true_loss": 1.5802
+    },
+    {
+      "epoch": 6.885245901639344,
+      "eval_accuracy": 0.1497695852534562,
+      "eval_loss": 3.839586019515991,
+      "eval_runtime": 14.9583,
+      "eval_samples_per_second": 29.014,
+      "eval_steps_per_second": 3.677,
+      "step": 420
+    },
+    {
+      "epoch": 6.967213114754099,
+      "grad_norm": 353.6661268712156,
+      "learning_rate": 7.289293849658314e-07,
+      "loss": 15.2028,
+      "step": 425,
+      "true_loss": 1.7352
+    },
+    {
+      "epoch": 7.049180327868853,
+      "grad_norm": 273.2552710608739,
+      "learning_rate": 6.71981776765376e-07,
+      "loss": 12.2664,
+      "step": 430,
+      "true_loss": 1.4752
+    },
+    {
+      "epoch": 7.049180327868853,
+      "eval_accuracy": 0.14746543778801843,
+      "eval_loss": 3.839120388031006,
+      "eval_runtime": 15.1557,
+      "eval_samples_per_second": 28.636,
+      "eval_steps_per_second": 3.629,
+      "step": 430
+    },
+    {
+      "epoch": 7.131147540983607,
+      "grad_norm": 278.2724305429665,
+      "learning_rate": 6.150341685649204e-07,
+      "loss": 10.7664,
+      "step": 435,
+      "true_loss": 1.3734
+    },
+    {
+      "epoch": 7.213114754098361,
+      "grad_norm": 316.474377641746,
+      "learning_rate": 5.580865603644648e-07,
+      "loss": 10.5371,
+      "step": 440,
+      "true_loss": 1.6166
+    },
+    {
+      "epoch": 7.213114754098361,
+      "eval_accuracy": 0.16359447004608296,
+      "eval_loss": 3.866461753845215,
+      "eval_runtime": 14.9602,
+      "eval_samples_per_second": 29.01,
+      "eval_steps_per_second": 3.676,
+      "step": 440
+    }
+  ],
+  "logging_steps": 5,
+  "max_steps": 488,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 8,
+  "save_steps": 10,
+  "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:fd25a69663dfac6f6868d122593f03b846718e390a0c72360a1c186d8e0bb2af
+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)