rndubs

Add docs/DATA_GUIDE.md

d07165b verified 3 months ago

preview code

Raw

History Blame Contribute Delete

9.22 kB

Data Guide: Datasets, Mixing, and Internal Data Curation

Pre-Training Data Sources

Tier 1: Core Datasets (must-have)

Dataset	What	Size	HF Link	License
FineWeb-Edu	High-quality educational web text	~4.3TB	HuggingFaceFW/fineweb-edu	ODC-By
The Stack v2 (dedup)	Deduplicated source code (600+ languages)	~67TB raw	bigcode/the-stack-v2-dedup	Various (per-file)
Your Internal Code	Internal Python libraries + apps	varies	Local	Internal

Tier 2: Domain-Specific (recommended)

Dataset	What	Size	HF Link
OpenR1-Math-220k	Math problems with CoT reasoning	~2GB	open-r1/OpenR1-Math-220k
Proof-Pile-2	Math/science papers + textbooks	~53GB	EleutherAI/proof-pile-2
Wikipedia	Structured knowledge	~20GB	wikimedia/wikipedia
Slurm/Bash data (curated)	HPC scripts + documentation	5-20GB	Self-curated

Tier 3: SFT-Quality (for annealing phase)

Dataset	What	Size	HF Link
Magicoder-Evol-Instruct-110K	Code instruction pairs	110K examples	ise-uiuc/Magicoder-Evol-Instruct-110K
OpenR1-Math-220k (messages)	Math instruction w/ reasoning	220K examples	Same as above
Synthetic internal data	Your library Q&A, docs, examples	50-200K examples	Self-generated

Data Mixing Strategy

Phase 1a: General Pre-Training (70% of compute)

data_mix = {
    "fineweb-edu":          0.30,   # General knowledge
    "the-stack-v2-python":  0.20,   # Python code (primary language)
    "the-stack-v2-other":   0.15,   # Other code (bash, JS, C, etc.)
    "openr1-math":          0.05,   # Math reasoning
    "proof-pile":           0.05,   # Math/science
    "wikipedia":            0.05,   # Structured knowledge
    "github-docs":          0.05,   # GitHub issues, READMEs, docs
    "slurm-bash":           0.05,   # HPC/shell data
    "internal-libraries":   0.10,   # YOUR internal code + docs
}

Phase 1b: Code-Heavy Annealing (20% of compute)

data_mix_annealing = {
    "the-stack-v2-python":  0.35,   # Increase Python
    "internal-libraries":   0.20,   # Increase internal code
    "magicoder-evol":       0.10,   # High-quality code instructions
    "openr1-math":          0.10,   # Math reasoning
    "slurm-bash":           0.10,   # Increase HPC data
    "fineweb-edu":          0.15,   # Reduced general text
}

Phase 1c: Long-Context Extension (10% of compute)

# Same mix as annealing but with long-context examples
# Filter for documents > 8K tokens
# Progressively increase from 4K → 32K → 131K context

Why These Ratios?

Based on published training recipes:

Model	Code %	Math %	Web %	Other %
DeepSeek-V3	Not published but heavy code emphasis
StarCoder2	85% code + 10% math + 5% web
Llama 3	50% general + 25% code + 25% other
OLMoE	~40% web + 30% code + 20% math/sci + 10% other
Ours (recommended)	~45% code	~10% math	~30% web	~15% other

The high code percentage (45%) reflects your primary use case. The 10% internal data ensures the model deeply understands your libraries.

Internal Data Curation Pipeline

Step 1: Crawl Your Repositories

python data_curation/crawl_internal_repos.py \
    --repos \
        /path/to/your/library-core \
        /path/to/your/library-utils \
        /path/to/your/app-1 \
        /path/to/your/app-2 \
        /path/to/your/slurm-scripts \
    --output data/raw/internal_raw.jsonl \
    --include-tests \
    --include-docs \
    --repo-level

This produces:

File-level examples: Individual Python files with metadata headers
Repo-level examples: Combined package views showing cross-file dependencies
Test files: Unit tests that teach the model expected behavior
Documentation: READMEs, docstrings, markdown docs

Step 2: Generate Synthetic Data with an API Model

export OPENAI_API_KEY="sk-..."
python data_curation/generate_synthetic_data.py \
    --input data/raw/internal_raw.jsonl \
    --output data/raw/internal_synthetic.jsonl \
    --api openai \
    --model gpt-4o \
    --tasks docs qa completion bugfix slurm reasoning \
    --num-per-file 5 \
    --max-parallel 10

What this generates for each source file:

Documentation — Full API docs, usage guides, tutorials
Q&A Pairs — "How do I use X?" → answer with code examples
Code Completion — Function stubs → full implementations
Bug-Fix Pairs — Buggy code → explanation → fix
Slurm Scripts — Job scripts using your library
Reasoning — Step-by-step problem solving with <think> tags

Cost estimate: At ~$0.03/generation (GPT-4o), 500 files × 6 tasks × 5 examples = 15K generations = ~$450

Step 3: Quality Filter and Deduplicate

python data_curation/filter_and_dedup.py \
    --input data/raw/internal_raw.jsonl data/raw/internal_synthetic.jsonl \
    --output data/filtered/internal_filtered.jsonl \
    --min-length 100 \
    --max-length 100000 \
    --min-quality 0.3 \
    --dedup-threshold 0.7

Step 4: Tokenize for Megatron-LM

python data_curation/tokenize_for_megatron.py \
    --input data/filtered/internal_filtered.jsonl \
    --output-prefix data/tokenized/internal-libraries \
    --tokenizer openai/gpt-oss-120b \
    --workers 32

This produces internal-libraries_text_document.bin and .idx files that Megatron-LM reads directly during training.

Step 5: Prepare SFT Data

For SFT (Phase 2), you need data in ChatML format:

python data_curation/generate_synthetic_data.py \
    --input data/raw/internal_raw.jsonl \
    --output data/sft/internal_library_qa.jsonl \
    --tasks qa reasoning \
    --num-per-file 10

# Convert to ChatML messages format:
python -c "
import json

with open('data/sft/internal_library_qa.jsonl') as f, \
     open('data/sft/internal_sft_messages.jsonl', 'w') as out:
    for line in f:
        item = json.loads(line)
        text = item['text']

        # Parse Q&A into messages
        if 'Question:' in text and 'Answer:' in text:
            parts = text.split('Answer:', 1)
            question = parts[0].replace('Question:', '').strip()
            answer = parts[1].strip()
            messages = [
                {'role': 'system', 'content': 'You are a helpful coding assistant with deep knowledge of our internal libraries and HPC systems.'},
                {'role': 'user', 'content': question},
                {'role': 'assistant', 'content': answer}
            ]
            out.write(json.dumps({'messages': messages}) + '\n')
"

Slurm/Bash Data Curation

Since there's no standard Slurm dataset:

1. Your own cluster scripts

python data_curation/prepare_slurm_data.py \
    --output data/filtered/slurm_bash.jsonl \
    --search-dirs ~/slurm-scripts /shared/job-scripts \
    --include-manpages \
    --include-synthetic

2. GitHub mining (optional but valuable)

Search GitHub for Slurm scripts:

# Use GitHub API to find .sbatch files
# Example search: "SBATCH language:shell filename:*.sbatch"
# Collect, filter, and add to training data

3. Documentation scraping

Slurm docs: https://slurm.schedmd.com/documentation.html
Convert HTML → Markdown for training

4. Synthetic generation (most effective)

Use the API model to generate diverse Slurm scripts:

python data_curation/generate_synthetic_data.py \
    --input data/raw/internal_raw.jsonl \
    --output data/raw/slurm_synthetic.jsonl \
    --tasks slurm \
    --num-per-file 10

Data Quality Tips

Upweight high-quality data in later epochs. During annealing, increase the mix of curated, instruction-quality data and decrease noisy web data.
Repeat internal data. Your internal code is ~10% of the mix but might be <1% of total tokens. The model sees it many more times than public data. This is intentional and beneficial — it's similar to "curriculum learning" where important data is repeated.
Include test files. Unit tests are incredibly valuable training data — they show the model the expected input/output behavior of your libraries. The model learns "when I call X with Y, it should return Z."
Include git history selectively. Commit messages and diffs teach the model about code evolution and debugging patterns. But raw git logs are noisy — filter for meaningful commits.
Decontaminate evaluation data. Remove any overlap between your training data and evaluation benchmarks (HumanEval, MBPP, etc.) to get honest benchmark scores.