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,111 @@
----
-license: mit
----
+---
+license: apache-2.0
+base_model: Qwen/Qwen3-0.6B
+tags:
+  - SAT
+  - combinatorial-optimization
+  - classification
+  - cube-and-conquer
+language:
+  - en
+pipeline_tag: text-classification
+---
+
+# Qwen3-0.6B-SAT-VarSelector
+
+A lightweight Qwen3-0.6B 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-0.6B` (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
+- **Size**: ~1.2GB (bfloat16)
+
+### Training
+
+- **Dataset**: 3,898 training / 434 validation samples
+- **Task**: Predict expert-selected branching variable
+- **Training**: 8 epochs, 8×H100 GPUs, DeepSpeed ZeRO-3
+
+### Comparison with 4B Model
+
+| Model | Size | Top-1 Acc | Top-5 Acc | Inference Speed |
+|-------|------|-----------|-----------|-----------------|
+| Qwen3-4B | 8GB | 24% | 48% | ~150ms/sample |
+| **Qwen3-0.6B** | 1.2GB | ~12% | ~32% | ~45ms/sample |
+
+## Usage
+
+```python
+import torch
+from transformers import AutoTokenizer
+from sft_qwen_var_classifier import QwenVarClassifier, cnf_valid_mask
+
+# Load model
+model = QwenVarClassifier("Qwen/Qwen3-0.6B", 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-0.6B")
+
+# 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 (~1.2GB, bfloat16)
+- `sft_qwen_var_classifier.py` - Model class definition (required for loading)
+
+## When to Use
+
+- **Production/Deployment**: Faster inference, smaller memory footprint
+- **Edge devices**: Can run on smaller GPUs
+- **Rapid prototyping**: Quick experiments
+- **CPU inference**: More practical than 4B model
+
+For maximum accuracy, use the [4B model](https://huggingface.co/Yale-ROSE/Qwen3-4B-SAT-VarSelector).
+
+## Limitations
+
+- Maximum 500 variables
+- Maximum 8192 tokens for CNF input
+- Lower accuracy than 4B model
+- Trained on specific CNF distribution
+
+## 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_step320/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:5ba4aaa5f340e1af930196eecba2e968e0606fb442aa489de142bf4e9e65daf4
+size 894853637

global_step320/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:7270f33db7c1b6d305d7efdb3f52e887898e2850dbcc53395c8b2449e4036e43
+size 894853637

global_step320/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:2334db35f3c0f64730919f0dd513fc3abdd055081b6d7ffd9919ec8d4176ec6d
+size 894853637

global_step320/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:18958fd5923f85803693c5fd7ab0914f9b618e01cd9db9f758a7c6dc36c78215
+size 894853637

global_step320/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:222f91a14301182368b6438fc8ac2c6440681c980160bf484ca283b54a7fb024
+size 894853637

global_step320/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:66587836b6898a2e51b2458ed6958c7ab42719d9aa72bdeff6606adc3d724026
+size 894853637

global_step320/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:883951900ba1059e95f5dcd4dfa3e6653dac450aa6350f1cb239e1c151561da0
+size 894853637

global_step320/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:744199262c8e9b1f18cd284b84e3dac9ed9ea4f15d0c74eb0b57c5aeb565510b
+size 894853637

global_step320/zero_pp_rank_0_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_1_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_2_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_3_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_4_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_5_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_6_mp_rank_00_model_states.pt

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

global_step320/zero_pp_rank_7_mp_rank_00_model_states.pt

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

latest

@@ -0,0 +1 @@
+global_step320

pytorch_model.bin

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

rng_state_0.pth

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

rng_state_1.pth

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

rng_state_2.pth

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

rng_state_3.pth

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

rng_state_4.pth

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

rng_state_5.pth

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

rng_state_6.pth

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

rng_state_7.pth

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

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2e73fcdd32d7fba0b375b46ae068e7b467fe50ce8c7455a317f26909f7d97b5b
+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,834 @@
+{
+  "best_global_step": 320,
+  "best_metric": 0.16129032258064516,
+  "best_model_checkpoint": "out_qwen_0.6b_sft/checkpoint-320",
+  "epoch": 5.245901639344262,
+  "eval_steps": 10,
+  "global_step": 320,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.08196721311475409,
+      "grad_norm": 128.95440298135514,
+      "learning_rate": 4.0816326530612243e-07,
+      "loss": 40.2266,
+      "step": 5,
+      "true_loss": 5.0094
+    },
+    {
+      "epoch": 0.16393442622950818,
+      "grad_norm": 115.65629290597954,
+      "learning_rate": 9.183673469387756e-07,
+      "loss": 39.852,
+      "step": 10,
+      "true_loss": 4.8969
+    },
+    {
+      "epoch": 0.16393442622950818,
+      "eval_accuracy": 0.0069124423963133645,
+      "eval_loss": 5.0297017097473145,
+      "eval_runtime": 7.6937,
+      "eval_samples_per_second": 56.41,
+      "eval_steps_per_second": 7.149,
+      "step": 10
+    },
+    {
+      "epoch": 0.2459016393442623,
+      "grad_norm": 121.33698181204981,
+      "learning_rate": 1.4285714285714286e-06,
+      "loss": 39.7914,
+      "step": 15,
+      "true_loss": 5.0324
+    },
+    {
+      "epoch": 0.32786885245901637,
+      "grad_norm": 123.59741878237826,
+      "learning_rate": 1.938775510204082e-06,
+      "loss": 40.1176,
+      "step": 20,
+      "true_loss": 5.0516
+    },
+    {
+      "epoch": 0.32786885245901637,
+      "eval_accuracy": 0.016129032258064516,
+      "eval_loss": 4.985995292663574,
+      "eval_runtime": 7.5658,
+      "eval_samples_per_second": 57.364,
+      "eval_steps_per_second": 7.27,
+      "step": 20
+    },
+    {
+      "epoch": 0.4098360655737705,
+      "grad_norm": 119.26551170123645,
+      "learning_rate": 2.4489795918367347e-06,
+      "loss": 39.4852,
+      "step": 25,
+      "true_loss": 4.8977
+    },
+    {
+      "epoch": 0.4918032786885246,
+      "grad_norm": 121.17042135033064,
+      "learning_rate": 2.959183673469388e-06,
+      "loss": 38.3961,
+      "step": 30,
+      "true_loss": 4.7777
+    },
+    {
+      "epoch": 0.4918032786885246,
+      "eval_accuracy": 0.02534562211981567,
+      "eval_loss": 4.856062889099121,
+      "eval_runtime": 9.8239,
+      "eval_samples_per_second": 44.178,
+      "eval_steps_per_second": 5.599,
+      "step": 30
+    },
+    {
+      "epoch": 0.5737704918032787,
+      "grad_norm": 148.64472538819717,
+      "learning_rate": 3.469387755102041e-06,
+      "loss": 38.4152,
+      "step": 35,
+      "true_loss": 4.7766
+    },
+    {
+      "epoch": 0.6557377049180327,
+      "grad_norm": 157.8362042375912,
+      "learning_rate": 3.979591836734694e-06,
+      "loss": 38.5402,
+      "step": 40,
+      "true_loss": 4.8457
+    },
+    {
+      "epoch": 0.6557377049180327,
+      "eval_accuracy": 0.04838709677419355,
+      "eval_loss": 4.698588848114014,
+      "eval_runtime": 7.6407,
+      "eval_samples_per_second": 56.801,
+      "eval_steps_per_second": 7.198,
+      "step": 40
+    },
+    {
+      "epoch": 0.7377049180327869,
+      "grad_norm": 146.82607440660624,
+      "learning_rate": 4.489795918367348e-06,
+      "loss": 36.5211,
+      "step": 45,
+      "true_loss": 4.6023
+    },
+    {
+      "epoch": 0.819672131147541,
+      "grad_norm": 158.20262885169598,
+      "learning_rate": 5e-06,
+      "loss": 37.591,
+      "step": 50,
+      "true_loss": 4.9816
+    },
+    {
+      "epoch": 0.819672131147541,
+      "eval_accuracy": 0.07373271889400922,
+      "eval_loss": 4.536938190460205,
+      "eval_runtime": 7.6075,
+      "eval_samples_per_second": 57.049,
+      "eval_steps_per_second": 7.23,
+      "step": 50
+    },
+    {
+      "epoch": 0.9016393442622951,
+      "grad_norm": 147.980428218924,
+      "learning_rate": 4.943052391799545e-06,
+      "loss": 37.0629,
+      "step": 55,
+      "true_loss": 4.6125
+    },
+    {
+      "epoch": 0.9836065573770492,
+      "grad_norm": 126.54677633437188,
+      "learning_rate": 4.886104783599089e-06,
+      "loss": 36.1198,
+      "step": 60,
+      "true_loss": 4.4017
+    },
+    {
+      "epoch": 0.9836065573770492,
+      "eval_accuracy": 0.07603686635944701,
+      "eval_loss": 4.475734233856201,
+      "eval_runtime": 7.662,
+      "eval_samples_per_second": 56.643,
+      "eval_steps_per_second": 7.178,
+      "step": 60
+    },
+    {
+      "epoch": 1.0655737704918034,
+      "grad_norm": 122.80613718263169,
+      "learning_rate": 4.829157175398634e-06,
+      "loss": 36.1342,
+      "step": 65,
+      "true_loss": 4.6152
+    },
+    {
+      "epoch": 1.1475409836065573,
+      "grad_norm": 136.54531412889096,
+      "learning_rate": 4.772209567198178e-06,
+      "loss": 36.0707,
+      "step": 70,
+      "true_loss": 4.284
+    },
+    {
+      "epoch": 1.1475409836065573,
+      "eval_accuracy": 0.07373271889400922,
+      "eval_loss": 4.457517147064209,
+      "eval_runtime": 7.7142,
+      "eval_samples_per_second": 56.26,
+      "eval_steps_per_second": 7.13,
+      "step": 70
+    },
+    {
+      "epoch": 1.2295081967213115,
+      "grad_norm": 139.59274630104497,
+      "learning_rate": 4.7152619589977225e-06,
+      "loss": 34.9596,
+      "step": 75,
+      "true_loss": 4.2389
+    },
+    {
+      "epoch": 1.3114754098360657,
+      "grad_norm": 142.6891577520439,
+      "learning_rate": 4.658314350797267e-06,
+      "loss": 35.7543,
+      "step": 80,
+      "true_loss": 4.5697
+    },
+    {
+      "epoch": 1.3114754098360657,
+      "eval_accuracy": 0.07603686635944701,
+      "eval_loss": 4.3655314445495605,
+      "eval_runtime": 7.6899,
+      "eval_samples_per_second": 56.437,
+      "eval_steps_per_second": 7.152,
+      "step": 80
+    },
+    {
+      "epoch": 1.3934426229508197,
+      "grad_norm": 159.00089279785396,
+      "learning_rate": 4.601366742596811e-06,
+      "loss": 34.967,
+      "step": 85,
+      "true_loss": 4.449
+    },
+    {
+      "epoch": 1.4754098360655736,
+      "grad_norm": 150.6608564699539,
+      "learning_rate": 4.544419134396356e-06,
+      "loss": 34.9076,
+      "step": 90,
+      "true_loss": 4.6
+    },
+    {
+      "epoch": 1.4754098360655736,
+      "eval_accuracy": 0.07373271889400922,
+      "eval_loss": 4.319250583648682,
+      "eval_runtime": 7.6466,
+      "eval_samples_per_second": 56.757,
+      "eval_steps_per_second": 7.193,
+      "step": 90
+    },
+    {
+      "epoch": 1.5573770491803278,
+      "grad_norm": 154.24917515552164,
+      "learning_rate": 4.4874715261959e-06,
+      "loss": 34.3756,
+      "step": 95,
+      "true_loss": 4.4023
+    },
+    {
+      "epoch": 1.639344262295082,
+      "grad_norm": 158.50353361820606,
+      "learning_rate": 4.4305239179954446e-06,
+      "loss": 34.9543,
+      "step": 100,
+      "true_loss": 4.5141
+    },
+    {
+      "epoch": 1.639344262295082,
+      "eval_accuracy": 0.08525345622119816,
+      "eval_loss": 4.283203125,
+      "eval_runtime": 7.5617,
+      "eval_samples_per_second": 57.394,
+      "eval_steps_per_second": 7.273,
+      "step": 100
+    },
+    {
+      "epoch": 1.721311475409836,
+      "grad_norm": 126.79955873407555,
+      "learning_rate": 4.373576309794989e-06,
+      "loss": 32.6608,
+      "step": 105,
+      "true_loss": 3.9732
+    },
+    {
+      "epoch": 1.8032786885245902,
+      "grad_norm": 152.75650023991204,
+      "learning_rate": 4.316628701594533e-06,
+      "loss": 33.611,
+      "step": 110,
+      "true_loss": 4.0678
+    },
+    {
+      "epoch": 1.8032786885245902,
+      "eval_accuracy": 0.0944700460829493,
+      "eval_loss": 4.219686031341553,
+      "eval_runtime": 7.5831,
+      "eval_samples_per_second": 57.233,
+      "eval_steps_per_second": 7.253,
+      "step": 110
+    },
+    {
+      "epoch": 1.8852459016393444,
+      "grad_norm": 154.91519081206442,
+      "learning_rate": 4.259681093394078e-06,
+      "loss": 34.0898,
+      "step": 115,
+      "true_loss": 4.3053
+    },
+    {
+      "epoch": 1.9672131147540983,
+      "grad_norm": 154.26448928608832,
+      "learning_rate": 4.202733485193622e-06,
+      "loss": 33.0864,
+      "step": 120,
+      "true_loss": 4.2828
+    },
+    {
+      "epoch": 1.9672131147540983,
+      "eval_accuracy": 0.10368663594470046,
+      "eval_loss": 4.198147773742676,
+      "eval_runtime": 7.7523,
+      "eval_samples_per_second": 55.984,
+      "eval_steps_per_second": 7.095,
+      "step": 120
+    },
+    {
+      "epoch": 2.0491803278688523,
+      "grad_norm": 152.56212863173033,
+      "learning_rate": 4.145785876993167e-06,
+      "loss": 33.29,
+      "step": 125,
+      "true_loss": 4.1834
+    },
+    {
+      "epoch": 2.1311475409836067,
+      "grad_norm": 162.83913288979136,
+      "learning_rate": 4.088838268792711e-06,
+      "loss": 32.6221,
+      "step": 130,
+      "true_loss": 4.0537
+    },
+    {
+      "epoch": 2.1311475409836067,
+      "eval_accuracy": 0.11059907834101383,
+      "eval_loss": 4.172513961791992,
+      "eval_runtime": 7.5194,
+      "eval_samples_per_second": 57.718,
+      "eval_steps_per_second": 7.314,
+      "step": 130
+    },
+    {
+      "epoch": 2.2131147540983607,
+      "grad_norm": 179.4233086801099,
+      "learning_rate": 4.0318906605922555e-06,
+      "loss": 31.9488,
+      "step": 135,
+      "true_loss": 4.0262
+    },
+    {
+      "epoch": 2.2950819672131146,
+      "grad_norm": 171.46365679764654,
+      "learning_rate": 3.9749430523918e-06,
+      "loss": 32.0664,
+      "step": 140,
+      "true_loss": 3.8796
+    },
+    {
+      "epoch": 2.2950819672131146,
+      "eval_accuracy": 0.10368663594470046,
+      "eval_loss": 4.146646499633789,
+      "eval_runtime": 7.4997,
+      "eval_samples_per_second": 57.869,
+      "eval_steps_per_second": 7.334,
+      "step": 140
+    },
+    {
+      "epoch": 2.3770491803278686,
+      "grad_norm": 178.73613610852706,
+      "learning_rate": 3.917995444191344e-06,
+      "loss": 32.302,
+      "step": 145,
+      "true_loss": 4.0002
+    },
+    {
+      "epoch": 2.459016393442623,
+      "grad_norm": 200.35242389508278,
+      "learning_rate": 3.861047835990889e-06,
+      "loss": 32.0721,
+      "step": 150,
+      "true_loss": 3.6775
+    },
+    {
+      "epoch": 2.459016393442623,
+      "eval_accuracy": 0.11290322580645161,
+      "eval_loss": 4.112561225891113,
+      "eval_runtime": 7.5207,
+      "eval_samples_per_second": 57.708,
+      "eval_steps_per_second": 7.313,
+      "step": 150
+    },
+    {
+      "epoch": 2.540983606557377,
+      "grad_norm": 189.21153835823273,
+      "learning_rate": 3.804100227790433e-06,
+      "loss": 31.7859,
+      "step": 155,
+      "true_loss": 3.8711
+    },
+    {
+      "epoch": 2.6229508196721314,
+      "grad_norm": 179.8665671736523,
+      "learning_rate": 3.7471526195899776e-06,
+      "loss": 32.339,
+      "step": 160,
+      "true_loss": 3.9259
+    },
+    {
+      "epoch": 2.6229508196721314,
+      "eval_accuracy": 0.1175115207373272,
+      "eval_loss": 4.069781303405762,
+      "eval_runtime": 7.5234,
+      "eval_samples_per_second": 57.687,
+      "eval_steps_per_second": 7.311,
+      "step": 160
+    },
+    {
+      "epoch": 2.7049180327868854,
+      "grad_norm": 184.87406326616818,
+      "learning_rate": 3.690205011389522e-06,
+      "loss": 31.6148,
+      "step": 165,
+      "true_loss": 3.9918
+    },
+    {
+      "epoch": 2.7868852459016393,
+      "grad_norm": 194.9656301652979,
+      "learning_rate": 3.6332574031890664e-06,
+      "loss": 30.9461,
+      "step": 170,
+      "true_loss": 4.1431
+    },
+    {
+      "epoch": 2.7868852459016393,
+      "eval_accuracy": 0.11981566820276497,
+      "eval_loss": 4.030214786529541,
+      "eval_runtime": 7.5282,
+      "eval_samples_per_second": 57.65,
+      "eval_steps_per_second": 7.306,
+      "step": 170
+    },
+    {
+      "epoch": 2.8688524590163933,
+      "grad_norm": 168.19816225086294,
+      "learning_rate": 3.5763097949886104e-06,
+      "loss": 31.2117,
+      "step": 175,
+      "true_loss": 4.0352
+    },
+    {
+      "epoch": 2.9508196721311473,
+      "grad_norm": 177.32224733459006,
+      "learning_rate": 3.519362186788155e-06,
+      "loss": 31.3617,
+      "step": 180,
+      "true_loss": 3.6492
+    },
+    {
+      "epoch": 2.9508196721311473,
+      "eval_accuracy": 0.1313364055299539,
+      "eval_loss": 3.996732711791992,
+      "eval_runtime": 7.498,
+      "eval_samples_per_second": 57.882,
+      "eval_steps_per_second": 7.335,
+      "step": 180
+    },
+    {
+      "epoch": 3.0327868852459017,
+      "grad_norm": 192.7334527443661,
+      "learning_rate": 3.4624145785876997e-06,
+      "loss": 31.2043,
+      "step": 185,
+      "true_loss": 3.8145
+    },
+    {
+      "epoch": 3.1147540983606556,
+      "grad_norm": 204.32213875438026,
+      "learning_rate": 3.405466970387244e-06,
+      "loss": 28.996,
+      "step": 190,
+      "true_loss": 3.1655
+    },
+    {
+      "epoch": 3.1147540983606556,
+      "eval_accuracy": 0.11290322580645161,
+      "eval_loss": 4.006399631500244,
+      "eval_runtime": 7.5951,
+      "eval_samples_per_second": 57.142,
+      "eval_steps_per_second": 7.242,
+      "step": 190
+    },
+    {
+      "epoch": 3.19672131147541,
+      "grad_norm": 219.38722327347378,
+      "learning_rate": 3.3485193621867885e-06,
+      "loss": 30.1186,
+      "step": 195,
+      "true_loss": 3.619
+    },
+    {
+      "epoch": 3.278688524590164,
+      "grad_norm": 232.33818268663353,
+      "learning_rate": 3.291571753986333e-06,
+      "loss": 30.4023,
+      "step": 200,
+      "true_loss": 3.7051
+    },
+    {
+      "epoch": 3.278688524590164,
+      "eval_accuracy": 0.12442396313364056,
+      "eval_loss": 3.9848430156707764,
+      "eval_runtime": 7.5257,
+      "eval_samples_per_second": 57.669,
+      "eval_steps_per_second": 7.308,
+      "step": 200
+    },
+    {
+      "epoch": 3.360655737704918,
+      "grad_norm": 229.39700824019172,
+      "learning_rate": 3.2346241457858773e-06,
+      "loss": 29.5804,
+      "step": 205,
+      "true_loss": 3.5024
+    },
+    {
+      "epoch": 3.442622950819672,
+      "grad_norm": 225.9592509617873,
+      "learning_rate": 3.1776765375854217e-06,
+      "loss": 29.0619,
+      "step": 210,
+      "true_loss": 3.5379
+    },
+    {
+      "epoch": 3.442622950819672,
+      "eval_accuracy": 0.15207373271889402,
+      "eval_loss": 3.9506587982177734,
+      "eval_runtime": 7.8653,
+      "eval_samples_per_second": 55.179,
+      "eval_steps_per_second": 6.993,
+      "step": 210
+    },
+    {
+      "epoch": 3.5245901639344264,
+      "grad_norm": 209.06281793043178,
+      "learning_rate": 3.120728929384966e-06,
+      "loss": 29.9575,
+      "step": 215,
+      "true_loss": 3.71
+    },
+    {
+      "epoch": 3.6065573770491803,
+      "grad_norm": 213.36792913688217,
+      "learning_rate": 3.0637813211845106e-06,
+      "loss": 29.4311,
+      "step": 220,
+      "true_loss": 3.8148
+    },
+    {
+      "epoch": 3.6065573770491803,
+      "eval_accuracy": 0.12211981566820276,
+      "eval_loss": 3.9277074337005615,
+      "eval_runtime": 7.5116,
+      "eval_samples_per_second": 57.778,
+      "eval_steps_per_second": 7.322,
+      "step": 220
+    },
+    {
+      "epoch": 3.6885245901639343,
+      "grad_norm": 215.89135789777671,
+      "learning_rate": 3.0068337129840546e-06,
+      "loss": 30.0775,
+      "step": 225,
+      "true_loss": 3.7924
+    },
+    {
+      "epoch": 3.7704918032786887,
+      "grad_norm": 216.66007631870477,
+      "learning_rate": 2.949886104783599e-06,
+      "loss": 29.6202,
+      "step": 230,
+      "true_loss": 3.6933
+    },
+    {
+      "epoch": 3.7704918032786887,
+      "eval_accuracy": 0.12672811059907835,
+      "eval_loss": 3.9045398235321045,
+      "eval_runtime": 7.6392,
+      "eval_samples_per_second": 56.813,
+      "eval_steps_per_second": 7.2,
+      "step": 230
+    },
+    {
+      "epoch": 3.8524590163934427,
+      "grad_norm": 216.8488298790556,
+      "learning_rate": 2.892938496583144e-06,
+      "loss": 29.7172,
+      "step": 235,
+      "true_loss": 3.3319
+    },
+    {
+      "epoch": 3.9344262295081966,
+      "grad_norm": 220.90823999032682,
+      "learning_rate": 2.8359908883826882e-06,
+      "loss": 29.235,
+      "step": 240,
+      "true_loss": 3.5549
+    },
+    {
+      "epoch": 3.9344262295081966,
+      "eval_accuracy": 0.11290322580645161,
+      "eval_loss": 3.904125928878784,
+      "eval_runtime": 7.6954,
+      "eval_samples_per_second": 56.398,
+      "eval_steps_per_second": 7.147,
+      "step": 240
+    },
+    {
+      "epoch": 4.016393442622951,
+      "grad_norm": 242.3023416449528,
+      "learning_rate": 2.7790432801822326e-06,
+      "loss": 27.9219,
+      "step": 245,
+      "true_loss": 3.2678
+    },
+    {
+      "epoch": 4.098360655737705,
+      "grad_norm": 255.00674520341965,
+      "learning_rate": 2.722095671981777e-06,
+      "loss": 27.9275,
+      "step": 250,
+      "true_loss": 3.5343
+    },
+    {
+      "epoch": 4.098360655737705,
+      "eval_accuracy": 0.1313364055299539,
+      "eval_loss": 3.898968458175659,
+      "eval_runtime": 7.517,
+      "eval_samples_per_second": 57.735,
+      "eval_steps_per_second": 7.317,
+      "step": 250
+    },
+    {
+      "epoch": 4.180327868852459,
+      "grad_norm": 273.1905409329504,
+      "learning_rate": 2.6651480637813215e-06,
+      "loss": 27.1873,
+      "step": 255,
+      "true_loss": 3.3906
+    },
+    {
+      "epoch": 4.262295081967213,
+      "grad_norm": 227.69914032682405,
+      "learning_rate": 2.608200455580866e-06,
+      "loss": 27.2035,
+      "step": 260,
+      "true_loss": 3.4531
+    },
+    {
+      "epoch": 4.262295081967213,
+      "eval_accuracy": 0.1359447004608295,
+      "eval_loss": 3.8699867725372314,
+      "eval_runtime": 7.5321,
+      "eval_samples_per_second": 57.62,
+      "eval_steps_per_second": 7.302,
+      "step": 260
+    },
+    {
+      "epoch": 4.344262295081967,
+      "grad_norm": 254.45648837084522,
+      "learning_rate": 2.5512528473804103e-06,
+      "loss": 27.0482,
+      "step": 265,
+      "true_loss": 3.4579
+    },
+    {
+      "epoch": 4.426229508196721,
+      "grad_norm": 266.3782057050062,
+      "learning_rate": 2.4943052391799547e-06,
+      "loss": 27.514,
+      "step": 270,
+      "true_loss": 3.5607
+    },
+    {
+      "epoch": 4.426229508196721,
+      "eval_accuracy": 0.1336405529953917,
+      "eval_loss": 3.8537046909332275,
+      "eval_runtime": 7.546,
+      "eval_samples_per_second": 57.514,
+      "eval_steps_per_second": 7.289,
+      "step": 270
+    },
+    {
+      "epoch": 4.508196721311475,
+      "grad_norm": 270.0666125278816,
+      "learning_rate": 2.437357630979499e-06,
+      "loss": 27.5337,
+      "step": 275,
+      "true_loss": 3.4807
+    },
+    {
+      "epoch": 4.590163934426229,
+      "grad_norm": 274.734392454001,
+      "learning_rate": 2.3804100227790436e-06,
+      "loss": 28.3342,
+      "step": 280,
+      "true_loss": 3.3558
+    },
+    {
+      "epoch": 4.590163934426229,
+      "eval_accuracy": 0.12903225806451613,
+      "eval_loss": 3.859492063522339,
+      "eval_runtime": 7.6721,
+      "eval_samples_per_second": 56.569,
+      "eval_steps_per_second": 7.169,
+      "step": 280
+    },
+    {
+      "epoch": 4.672131147540983,
+      "grad_norm": 272.2658673500175,
+      "learning_rate": 2.323462414578588e-06,
+      "loss": 27.877,
+      "step": 285,
+      "true_loss": 3.4582
+    },
+    {
+      "epoch": 4.754098360655737,
+      "grad_norm": 268.9659407625432,
+      "learning_rate": 2.2665148063781324e-06,
+      "loss": 27.6915,
+      "step": 290,
+      "true_loss": 3.9791
+    },
+    {
+      "epoch": 4.754098360655737,
+      "eval_accuracy": 0.1336405529953917,
+      "eval_loss": 3.843111038208008,
+      "eval_runtime": 7.6659,
+      "eval_samples_per_second": 56.615,
+      "eval_steps_per_second": 7.175,
+      "step": 290
+    },
+    {
+      "epoch": 4.836065573770492,
+      "grad_norm": 263.12731442923484,
+      "learning_rate": 2.209567198177677e-06,
+      "loss": 26.7233,
+      "step": 295,
+      "true_loss": 3.1607
+    },
+    {
+      "epoch": 4.918032786885246,
+      "grad_norm": 295.63911480933103,
+      "learning_rate": 2.1526195899772212e-06,
+      "loss": 28.7988,
+      "step": 300,
+      "true_loss": 3.9041
+    },
+    {
+      "epoch": 4.918032786885246,
+      "eval_accuracy": 0.1313364055299539,
+      "eval_loss": 3.8396456241607666,
+      "eval_runtime": 7.7433,
+      "eval_samples_per_second": 56.049,
+      "eval_steps_per_second": 7.103,
+      "step": 300
+    },
+    {
+      "epoch": 5.0,
+      "grad_norm": 276.8145069400469,
+      "learning_rate": 2.0956719817767656e-06,
+      "loss": 27.155,
+      "step": 305,
+      "true_loss": 3.1979
+    },
+    {
+      "epoch": 5.081967213114754,
+      "grad_norm": 277.6991231127756,
+      "learning_rate": 2.03872437357631e-06,
+      "loss": 25.5137,
+      "step": 310,
+      "true_loss": 3.413
+    },
+    {
+      "epoch": 5.081967213114754,
+      "eval_accuracy": 0.14285714285714285,
+      "eval_loss": 3.8058576583862305,
+      "eval_runtime": 7.6657,
+      "eval_samples_per_second": 56.616,
+      "eval_steps_per_second": 7.175,
+      "step": 310
+    },
+    {
+      "epoch": 5.163934426229508,
+      "grad_norm": 284.7257730508428,
+      "learning_rate": 1.9817767653758545e-06,
+      "loss": 26.0606,
+      "step": 315,
+      "true_loss": 3.2927
+    },
+    {
+      "epoch": 5.245901639344262,
+      "grad_norm": 292.26300972505453,
+      "learning_rate": 1.924829157175399e-06,
+      "loss": 25.2526,
+      "step": 320,
+      "true_loss": 3.2581
+    },
+    {
+      "epoch": 5.245901639344262,
+      "eval_accuracy": 0.16129032258064516,
+      "eval_loss": 3.8149120807647705,
+      "eval_runtime": 7.6305,
+      "eval_samples_per_second": 56.877,
+      "eval_steps_per_second": 7.208,
+      "step": 320
+    }
+  ],
+  "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:0e3b6e1d011240a6e67dc34aa90b77ce8b3b5df9a9a443de1096cce60d60c346
+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)