258 lines
9.8 KiB
Python
258 lines
9.8 KiB
Python
"""
|
|
German Structured Output SFT Training
|
|
======================================
|
|
EU AI Act Article 53 compliant fine-tuning of OLMo 2 1B Instruct
|
|
on philipp-zettl/german-structured-output dataset.
|
|
|
|
Base model: allenai/OLMo-2-0425-1B-Instruct (Apache 2.0)
|
|
- Pretraining data: allenai/olmo-mix-1124 (ODC-BY, fully documented)
|
|
- Every data source has explicit license: DCLM (CC-BY-4.0), ArXiv (ODC-BY),
|
|
PeS2o (ODC-BY), StarCoder (ODC-BY), Wikipedia (ODC-BY), etc.
|
|
- Full training logs, code, and intermediate checkpoints published by AI2
|
|
|
|
Dataset: philipp-zettl/german-structured-output (CC BY-SA 4.0)
|
|
- All sources documented with licenses per example
|
|
- Zero PII, GDPR compliant
|
|
|
|
Method: Full SFT with gradient checkpointing
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import json
|
|
import torch
|
|
import statistics
|
|
|
|
# Force unbuffered output so logs stream in real-time
|
|
os.environ["PYTHONUNBUFFERED"] = "1"
|
|
|
|
print("=" * 60, flush=True)
|
|
print("Starting German Structured Output SFT Training", flush=True)
|
|
print("=" * 60, flush=True)
|
|
|
|
# Verify HF_TOKEN is available (needed for private dataset + push_to_hub)
|
|
token = os.environ.get("HF_TOKEN")
|
|
if not token:
|
|
print("ERROR: HF_TOKEN not set! Cannot access private dataset or push model.", flush=True)
|
|
print("Set via: export HF_TOKEN=hf_... or pass --env HF_TOKEN=hf_...", flush=True)
|
|
sys.exit(1)
|
|
else:
|
|
print(f"HF_TOKEN: present ({token[:8]}...)", flush=True)
|
|
|
|
# Check GPU
|
|
print(f"CUDA available: {torch.cuda.is_available()}", flush=True)
|
|
if torch.cuda.is_available():
|
|
print(f"GPU: {torch.cuda.get_device_name()}", flush=True)
|
|
print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB", flush=True)
|
|
else:
|
|
print("WARNING: No GPU detected! Training will be very slow.", flush=True)
|
|
|
|
from datasets import load_dataset
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
|
from trl import SFTConfig, SFTTrainer
|
|
|
|
print("All imports OK", flush=True)
|
|
|
|
# ============================================================================
|
|
# Configuration
|
|
# ============================================================================
|
|
|
|
MODEL_ID = "allenai/OLMo-2-0425-1B-Instruct"
|
|
DATASET_ID = "philipp-zettl/german-structured-output"
|
|
OUTPUT_MODEL_ID = "philipp-zettl/german-structured-output-olmo2-1b"
|
|
OUTPUT_DIR = "./output"
|
|
|
|
# ============================================================================
|
|
# Trackio monitoring — writes to HF Bucket for persistent storage
|
|
# ============================================================================
|
|
print("Initializing Trackio...", flush=True)
|
|
import trackio
|
|
trackio.init(
|
|
project="german-structured-output-sft",
|
|
name="olmo2-1b-full-sft",
|
|
space_id="philipp-zettl/german-structured-output-training",
|
|
bucket_id="philipp-zettl/german-structured-output-training-bucket",
|
|
)
|
|
print("Trackio initialized with bucket storage", flush=True)
|
|
|
|
# ============================================================================
|
|
# Load dataset
|
|
# ============================================================================
|
|
print("Loading dataset...", flush=True)
|
|
dataset = load_dataset(DATASET_ID)
|
|
print(f"Train: {len(dataset['train'])}, Val: {len(dataset['validation'])}, Test: {len(dataset['test'])}", flush=True)
|
|
|
|
# Quick data audit
|
|
print("\nSample messages structure:", flush=True)
|
|
sample = dataset["train"][0]
|
|
for msg in sample["messages"]:
|
|
print(f" [{msg['role']}]: {msg['content'][:100]}...", flush=True)
|
|
print(f" task_type: {sample['task_type']}", flush=True)
|
|
print(f" quality_score: {sample['quality_score']}", flush=True)
|
|
|
|
# ============================================================================
|
|
# Load model and tokenizer
|
|
# ============================================================================
|
|
print(f"\nLoading model: {MODEL_ID}", flush=True)
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
MODEL_ID,
|
|
torch_dtype=torch.bfloat16,
|
|
attn_implementation="sdpa",
|
|
)
|
|
|
|
# Verify chat template exists
|
|
print(f"Chat template available: {tokenizer.chat_template is not None}", flush=True)
|
|
print(f"Model parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.0f}M", flush=True)
|
|
|
|
# Check sequence lengths in dataset to set max_length appropriately
|
|
print("\nAnalyzing sequence lengths...", flush=True)
|
|
lengths = []
|
|
for i, example in enumerate(dataset["train"]):
|
|
text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
|
|
tokens = tokenizer(text, return_length=True)
|
|
lengths.append(tokens["length"][0])
|
|
if i >= 200: # Sample first 200
|
|
break
|
|
|
|
print(f" Sampled {len(lengths)} examples", flush=True)
|
|
print(f" Min: {min(lengths)}, Max: {max(lengths)}, Mean: {statistics.mean(lengths):.0f}, Median: {statistics.median(lengths):.0f}", flush=True)
|
|
print(f" P95: {sorted(lengths)[int(len(lengths)*0.95)]}", flush=True)
|
|
print(f" P99: {sorted(lengths)[int(len(lengths)*0.99)]}", flush=True)
|
|
|
|
# Set max_length to cover P99 + margin, cap at 2048
|
|
MAX_LENGTH = min(sorted(lengths)[int(len(lengths)*0.99)] + 128, 2048)
|
|
print(f" Using max_length: {MAX_LENGTH}", flush=True)
|
|
|
|
# ============================================================================
|
|
# Training configuration
|
|
# ============================================================================
|
|
training_args = SFTConfig(
|
|
output_dir=OUTPUT_DIR,
|
|
|
|
# Precision
|
|
bf16=True,
|
|
|
|
# Sequence
|
|
max_length=MAX_LENGTH,
|
|
packing=False, # No packing — variable length structured outputs
|
|
|
|
# Batch size: effective = 4 * 4 = 16
|
|
per_device_train_batch_size=4,
|
|
gradient_accumulation_steps=4,
|
|
|
|
# Optimizer
|
|
learning_rate=2e-5,
|
|
lr_scheduler_type="cosine",
|
|
warmup_ratio=0.05,
|
|
weight_decay=0.01,
|
|
max_grad_norm=1.0,
|
|
|
|
# Epochs — small dataset (4K), more epochs to converge
|
|
num_train_epochs=5,
|
|
|
|
# Evaluation
|
|
eval_strategy="epoch",
|
|
eval_on_start=True,
|
|
|
|
# Saving
|
|
save_strategy="epoch",
|
|
save_total_limit=2,
|
|
load_best_model_at_end=True,
|
|
metric_for_best_model="eval_loss",
|
|
greater_is_better=False,
|
|
|
|
# Logging — plain text, no tqdm
|
|
logging_strategy="steps",
|
|
logging_steps=10,
|
|
logging_first_step=True,
|
|
disable_tqdm=True,
|
|
report_to="trackio",
|
|
run_name="olmo2-1b-german-structured-output",
|
|
|
|
# Memory
|
|
gradient_checkpointing=True,
|
|
|
|
# Hub — CRITICAL: push_to_hub to save model (ephemeral job storage)
|
|
push_to_hub=True,
|
|
hub_model_id=OUTPUT_MODEL_ID,
|
|
hub_strategy="end",
|
|
|
|
# NEFTune for small dataset regularization (paper: arxiv 2310.05914)
|
|
neftune_noise_alpha=5.0,
|
|
)
|
|
|
|
# ============================================================================
|
|
# Trainer
|
|
# ============================================================================
|
|
print("\nInitializing SFTTrainer...", flush=True)
|
|
trainer = SFTTrainer(
|
|
model=model,
|
|
processing_class=tokenizer,
|
|
args=training_args,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset["validation"],
|
|
)
|
|
|
|
# Print training summary
|
|
total_steps = len(dataset["train"]) // (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps) * int(training_args.num_train_epochs)
|
|
print(f"\n{'='*60}", flush=True)
|
|
print(f"Training Summary", flush=True)
|
|
print(f"{'='*60}", flush=True)
|
|
print(f" Base model: {MODEL_ID}", flush=True)
|
|
print(f" Base model license: Apache 2.0", flush=True)
|
|
print(f" Base model training data: allenai/olmo-mix-1124 (ODC-BY)", flush=True)
|
|
print(f" Dataset: {DATASET_ID}", flush=True)
|
|
print(f" Dataset license: CC BY-SA 4.0", flush=True)
|
|
print(f" Output model: {OUTPUT_MODEL_ID}", flush=True)
|
|
print(f" Train examples: {len(dataset['train'])}", flush=True)
|
|
print(f" Val examples: {len(dataset['validation'])}", flush=True)
|
|
print(f" Max length: {MAX_LENGTH}", flush=True)
|
|
print(f" Effective batch size: {training_args.per_device_train_batch_size} x {training_args.gradient_accumulation_steps} = {training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps}", flush=True)
|
|
print(f" Learning rate: {training_args.learning_rate}", flush=True)
|
|
print(f" Epochs: {training_args.num_train_epochs}", flush=True)
|
|
print(f" Estimated steps: ~{total_steps}", flush=True)
|
|
print(f" NEFTune alpha: {training_args.neftune_noise_alpha}", flush=True)
|
|
print(f" EU AI Act: Article 53 compliant (full data provenance chain)", flush=True)
|
|
print(f"{'='*60}", flush=True)
|
|
|
|
# ============================================================================
|
|
# Train!
|
|
# ============================================================================
|
|
print("\nStarting training...", flush=True)
|
|
train_result = trainer.train()
|
|
|
|
# ============================================================================
|
|
# Save and push
|
|
# ============================================================================
|
|
print("\nSaving model...", flush=True)
|
|
trainer.save_model()
|
|
|
|
# Log final metrics
|
|
metrics = train_result.metrics
|
|
print(f"\nTraining complete!", flush=True)
|
|
print(f" Train loss: {metrics.get('train_loss', 'N/A')}", flush=True)
|
|
print(f" Train runtime: {metrics.get('train_runtime', 'N/A'):.0f}s", flush=True)
|
|
print(f" Train samples/sec: {metrics.get('train_samples_per_second', 'N/A'):.1f}", flush=True)
|
|
|
|
# Evaluate
|
|
print("\nRunning final evaluation...", flush=True)
|
|
eval_metrics = trainer.evaluate()
|
|
print(f" Eval loss: {eval_metrics.get('eval_loss', 'N/A')}", flush=True)
|
|
|
|
# Push to hub
|
|
print(f"\nPushing to Hub: {OUTPUT_MODEL_ID}", flush=True)
|
|
trainer.push_to_hub(
|
|
commit_message="EU AI Act compliant SFT: OLMo 2 1B on german-structured-output",
|
|
tags=[
|
|
"german", "structured-output", "json", "function-calling",
|
|
"ner", "relation-extraction", "gdpr-anonymization",
|
|
"eu-ai-act-compliant", "gdpr-compliant", "sft",
|
|
"olmo2", "article-53"
|
|
],
|
|
)
|
|
|
|
print("\n✅ Training complete and model pushed to Hub!", flush=True)
|
|
print(f" Model: https://huggingface.co/{OUTPUT_MODEL_ID}", flush=True)
|
|
print(f" Trackio: https://huggingface.co/spaces/philipp-zettl/german-structured-output-training", flush=True)
|