sft-pipeline-backup

Reasoning Distillation Pipeline for Llama-3.2-3B

This repository contains the complete pipeline for distilling reasoning capabilities from DeepSeek-R1 into a Llama-3.2-3B model. It covers every stage of the process: dataset curation, supervised fine-tuning (SFT), and benchmark evaluation with LLM-as-Judge analysis.

The resulting model, Llama-3.2-3B-thinking-v3-13.9K, learns to produce structured chain-of-thought reasoning wrapped in <think>...</think> tags before generating a final answer. The trained weights and all checkpoints are publicly available on HuggingFace.

---

Table of Contents

1. Repository Structure
2. Prerequisites
3. Trained Model
4. Stage 1: Dataset Preparation
5. Stage 2: Training
6. Stage 3: Benchmark Evaluation
7. Results Summary

---

Repository Structure

the-entire-pipeline/
├── dataset/                          # Stage 1: Data preparation
│   ├── extract_sources.py            # Extract and count source labels from raw data
│   ├── stratified_sampler.py         # Stratified sampling across task categories
│   ├── convert_to_training_format.py # Convert to clean chat-style training format
│   ├── filter_by_length.py           # Remove token-length outliers
│   ├── analyze_token_lengths.py      # Token distribution analysis
│   ├── stratified_sample_14k.jsonl   # Intermediate: raw stratified samples
│   ├── training_dataset_14k.jsonl    # Intermediate: converted training format
│   └── training_dataset_filtered_16k.jsonl  # Final training data (13,772 samples)
│
├── training-script/                  # Stage 2: Model training
│   ├── train_sft.py                  # Full-weight SFT training script
│   ├── requirements.txt              # Python dependencies
│   └── training-logs/                # Training run logs
│       ├── train_full.log            # Initial run (crashed at step 1000)
│       └── train_resume.log          # Resumed run (completed)
│
└── benchmark-results/                # Stage 3: Evaluation
    ├── bbh/                          # BIG-Bench Hard raw outputs (4 configurations)
    ├── gsm8k/                        # GSM8K results (0-shot and 8-shot)
    ├── llm_judge_data/               # Prepared data for LLM judge (finetuned model)
    ├── llm_judge_data_base/          # Prepared data for LLM judge (base model)
    ├── llm_judge_results/            # Judge verdicts for finetuned model
    ├── llm_judge_results_base/       # Judge verdicts for base model
    ├── analysis_output/              # Generated figures and reports (finetuned)
    ├── analysis_output_base/         # Generated figures and reports (base)
    ├── comprehensive_llm_judge_analyzer.py  # Main analysis script
    └── model_comparison_report.json  # Base vs. finetuned comparison

---

Prerequisites

• Python 3.11+
• A single GPU with at least 80 GB VRAM (A100 80GB or equivalent recommended)
• The AM-DeepSeek-R1-Distilled-1.4M dataset from HuggingFace
• lm-eval harness for benchmarking

Install training dependencies:

cd training-script
pip install -r requirements.txt

---

Trained Model

The final trained model and all intermediate checkpoints are hosted on HuggingFace:

To load the model for inference:

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "x1nx3r/Llama-3.2-3B-thinking-v3-13.9K"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")

---

Stage 1: Dataset Preparation

The goal of this stage is to produce a curated, balanced training set from the 1.4M-sample AM-DeepSeek-R1-Distilled dataset. The raw dataset contains reasoning traces distilled from DeepSeek-R1-671B across math, code, NLP, and general reasoning domains.

Pre-built intermediate and final dataset files from each step below are available on HuggingFace at x1nx3r/dataset-checkpoint. If you prefer to skip regeneration, download the files directly from there.

1.1 Acquiring the Dataset

Clone the AM-DeepSeek-R1-Distilled-1.4M dataset from HuggingFace. The full dataset is split into two JSONL files; this pipeline uses only the am_0.9M.jsonl split, which contains approximately 900,000 samples distilled by the AM team from DeepSeek-R1-671B.

# Requires git-lfs
git lfs install
git clone https://huggingface.co/datasets/a-m-team/AM-DeepSeek-R1-Distilled-1.4M

After cloning, the relevant file is located at:

AM-DeepSeek-R1-Distilled-1.4M/am_0.9M.jsonl

The other split (open_0.5M.jsonl) contains open-source reasoning data and is not used in this pipeline.

1.2 Source Analysis

Inspect the distribution of instruction and response sources in am_0.9M.jsonl:

cd dataset
python3 extract_sources.py \
    --file /path/to/AM-DeepSeek-R1-Distilled-1.4M/am_0.9M.jsonl \
    --action counts \
    --role-split \
    --out source_counts

This produces per-role source counts, revealing the composition of the dataset. The largest instruction sources are natural_reasoning (319K), InfinityInstruct (307K), and KodCode (211K).

1.3 Stratified Sampling

To ensure representation across task types, the sampler draws from five predefined categories with target counts:

Run the sampler:

python3 stratified_sampler.py \
    --input /path/to/AM-DeepSeek-R1-Distilled-1.4M/am_0.9M.jsonl \
    --output stratified_sample_14k.jsonl

The script uses reservoir sampling weighted by per-source counts, producing approximately 14,000 samples.

1.4 Format Conversion

The raw dataset uses a nested format with metadata fields (info, source, reference_answer, etc.). The conversion script strips metadata and produces a clean chat format with:

• A system prompt instructing the model to reason step-by-step inside <think> tags
• User messages containing the raw problem
• Assistant messages preserving the original <think>...</think> and <answer>...</answer> structure

python3 convert_to_training_format.py \
    --input stratified_sample_14k.jsonl \
    --output training_dataset_14k.jsonl

Content is cleaned without destroying LaTeX formatting, markdown structure, or code indentation.

1.5 Length Filtering

Extremely long conversations (above the 99th percentile in estimated token count) are removed to prevent memory issues during training. The filter uses a character-based token approximation (characters / 4):

python3 filter_by_length.py \
    --input training_dataset_14k.jsonl \
    --output training_dataset_filtered_16k.jsonl \
    --max-tokens 16384

Final dataset: 13,772 conversations (~143 MB), suitable for training with a 16,384-token context window.

---

Stage 2: Training

The training stage performs full-weight supervised fine-tuning on Llama-3.2-3B-Instruct (3.21 billion parameters). All parameters are updated -- no adapters or LoRA are used.

2.1 Training Configuration

2.2 Label Masking

The ChatDataset class in train_sft.py constructs training examples by applying the chat template and creating label masks. Loss is computed only on assistant tokens -- system and user tokens are masked with -100. This ensures the model learns to generate responses without memorizing prompts.

The masking logic locates assistant turn boundaries using special tokens (<|start_header_id|>assistant<|end_header_id|> through <|eot_id|>) and sets all non-assistant positions to the ignore index.

2.3 Running Training

cd training-script

python3 train_sft.py \
    --base /path/to/Llama-3.2-3B-Instruct \
    --data /path/to/training_dataset_filtered_16k.jsonl \
    --output /path/to/output_sft \
    --max_length 16384 \
    --epochs 2 \
    --per_device_batch 1 \
    --grad_accum 8 \
    --lr 1e-5 \
    --save_steps 500 \
    --logging_steps 20 \
    --save_total_limit 3

2.4 Training Notes

The initial training run encountered a RuntimeError at step 1000/3442 caused by a disk write failure during optimizer state serialization. Training was resumed from the last saved checkpoint (step 500) and completed without further issues. Both log files are preserved in training-logs/ for reference.

Average throughput was approximately 4--5 seconds per step on the training hardware.

---

Stage 3: Benchmark Evaluation

The finetuned model was evaluated on two benchmarks: BIG-Bench Hard (BBH) and GSM8K. An LLM-as-Judge pipeline using Gemini Pro 3 Preview was employed for semantic evaluation of BBH responses.

3.1 BBH Evaluation

BBH consists of 27 challenging reasoning tasks spanning logical deduction, arithmetic, language understanding, and more. The model was evaluated under four configurations:

Evaluation was run using the lm-eval harness:

lm-eval \
    --model hf \
    --model_args pretrained=x1nx3r/Llama-3.2-3B-thinking-v3-13.9K,dtype=float16,trust_remote_code=True \
    --tasks bbh \
    --apply_chat_template \
    --system_instruction "You are a helpful reasoning assistant. Think step-by-step inside <think></think> before giving your final answer." \
    --num_fewshot 0 \
    --batch_size auto \
    --gen_kwargs '{"temperature": 0.0, "repetition_penalty": 1.0, "do_sample": false, "max_new_tokens": 2048}' \
    --write_out \
    --log_samples \
    --output_path benchmark-results/bbh/0shot-nosysprompt

3.2 GSM8K Evaluation

GSM8K (grade-school math) was evaluated at 0-shot and 8-shot:

lm-eval \
    --model hf \
    --model_args pretrained=x1nx3r/Llama-3.2-3B-thinking-v3-13.9K,dtype=float16,trust_remote_code=True \
    --tasks gsm8k \
    --apply_chat_template \
    --system_instruction "You are a helpful reasoning assistant. Think step-by-step inside <think></think> before giving your final answer." \
    --num_fewshot 0 \
    --batch_size auto \
    --gen_kwargs '{"temperature": 0.0, "repetition_penalty": 1.0, "do_sample": false, "max_new_tokens": 2048, "until": ["<|eot_id|>", "<|end_of_text|>"]}' \
    --write_out \
    --log_samples \
    --output_path benchmark-results/gsm8k/testrun-complete-0-shot

3.3 LLM-as-Judge Pipeline

Because BBH tasks often involve free-form answers that cannot be evaluated with exact string matching, an LLM-as-Judge framework was used. The pipeline proceeds in two steps:

Step 1 -- Prepare judge input. Convert raw BBH outputs into a structured JSONL format containing the question, model response, and ground-truth answer:

cd benchmark-results
python3 convert_for_judge.py

This produces 27 JSONL files in llm_judge_data/, one per task.

Step 2 -- Run the judge. Each sample is evaluated by Gemini Pro 3 Preview, which assigns:

• is_correct (boolean) -- whether the model's answer matches the ground truth under semantic equivalence rules
• reasoning_score (1--5) -- quality of the chain-of-thought reasoning
• error_type -- classification of the failure mode (none, wrong_logic, hallucination, no_answer, format_error, calculation_error)
• comment -- brief explanation of the verdict

Semantic equivalence rules handle format variations (e.g., (A) matches A, Option A, [A], \boxed{A}; True matches yes, valid, correct). The full judge prompt is documented in llm_judge_prompt.md.

Step 3 -- Analyze results. Generate figures and a markdown report:

python3 comprehensive_llm_judge_analyzer.py

Output is written to analysis_output/, including accuracy bar charts, error distribution plots, reasoning score histograms, and a per-task error heatmap.

---

Results Summary

BBH Performance (27 Tasks, 6,510 Samples)

Overall accuracy: 60.1%

Error Analysis

Key Observations

1. wrong_logic dominates failure modes (38.8%). The model generally follows correct reasoning structures but makes logical or computational errors within them. This suggests the distillation successfully transferred the reasoning format but not all of the logical precision.

2. Strong correlation between accuracy and reasoning scores. Tasks with accuracy above 70% consistently show average reasoning scores of 4.0 or higher, indicating that the judge's reasoning assessment aligns well with correctness.

3. Task complexity degrades performance predictably. The tracking shuffled objects family shows clear degradation with increasing object count: 57.2% (3 objects), 54.4% (5 objects), 53.6% (7 objects).

4. Symbolic reasoning remains the weakest domain. Dyck Languages (bracket matching) achieved only 16.1% accuracy, suggesting that small distilled models struggle with purely syntactic pattern matching.

5. Arithmetic reasoning shows an interesting gap. Multistep Arithmetic achieved 61.6% accuracy despite a high average reasoning score of 4.42, indicating the model demonstrates sound methodology but commits calculation slips.

Base Model Comparison (Partial)

The finetuned model improves reasoning quality (higher average scores) on tasks where the base model underperforms, though some regression is observed on tasks where the base model was already strong.

License

This repository contains benchmark evaluation results and training scripts for research purposes. The training dataset is released under CC-BY-NC-4.0. Commercial use is not permitted.