Safety Mid-Training: Internalizing LLM Safety as a Foundational Capability

This repository contains the code for Safety Mid-Training: Internalizing LLM Safety as a Foundational Capability.

🤗 Hugging Face Artifacts

Type	Artifact
🤗 Dataset	Ethan271/midtrain-document-data
🤗 Dataset	Ethan271/midtrain-reasoning-data
🤗 Model	Ethan271/qwen3-1.7b-safety-midtrain
🤗 Model	Ethan271/qwen3-4b-safety-midtrain
🤗 Model	Ethan271/qwen3-8b-safety-midtrain

The default base model for the released code path is Qwen/Qwen3-4B.

Setup

conda create -n safety-midtrain python=3.10 -y
conda activate safety-midtrain
pip install --upgrade pip
pip install -r requirements.txt

For GRPO's online Bedrock-compatible reward judge:

export AWS_BEARER_TOKEN_BEDROCK=...

For eval scoring with Azure OpenAI:

export AZURE_OPENAI_API_KEY=...
export AZURE_OPENAI_ENDPOINT=...
export AZURE_OPENAI_DEPLOYMENT=...

End-to-End Script

The example script covers released data download, optional midtrain data generation and filtering, midtraining, SFT data construction/self-sampling, SFT training, GRPO, and eval:

bash scripts/example_full_pipeline.sh

By default, the script downloads the released Hugging Face datasets and uses their user_prompt fields as the prompt source for SFT self-sampling and GRPO. To regenerate midtrain data before training, place the source corpora at:

data/corpora/document_seed.jsonl
data/corpora/reasoning_seed.jsonl

The generation code reads the training prompts from each dataset row's user_prompt field; no separate hard-coded prompt list is used.

Then run:

RUN_DATA_GENERATION=1 bash scripts/example_full_pipeline.sh

Useful switches:

SFT_MODE=direct bash scripts/example_full_pipeline.sh
RUN_GRPO=0 RUN_EVAL=0 bash scripts/example_full_pipeline.sh
EVAL_EXTRA_ARGS="--generate-only --limit 50" bash scripts/example_full_pipeline.sh

Midtrain Data

Download the released midtrain data:

python - <<'PY'
from pathlib import Path
import shutil
from huggingface_hub import hf_hub_download

pairs = [
    ("Ethan271/midtrain-document-data", "document_midtrain_deduped.jsonl", "data/document_midtrain.jsonl"),
    ("Ethan271/midtrain-reasoning-data", "reasoning_midtrain_deduped.jsonl", "data/reasoning_midtrain.jsonl"),
]
for repo_id, filename, target in pairs:
    src = hf_hub_download(repo_id=repo_id, filename=filename, repo_type="dataset")
    Path(target).parent.mkdir(parents=True, exist_ok=True)
    shutil.copyfile(src, target)
PY

Regenerate document-format midtrain data:

python src/data/documents.py generate \
  --model Qwen/Qwen3-4B \
  --n-candidates 2 \
  --out data/documents_round1.jsonl

python src/data/documents.py judge \
  --input data/documents_round1.jsonl \
  --out data/documents_round1_judged.jsonl

python src/data/documents.py failed \
  --input data/documents_round1_judged.jsonl \
  --list-out data/document_resample_ids.txt \
  --min-quality 4

python src/data/documents.py generate \
  --model Qwen/Qwen3-4B \
  --source-list-file data/document_resample_ids.txt \
  --n-candidates 4 \
  --out data/documents_resample.jsonl

python src/data/documents.py judge \
  --input data/documents_resample.jsonl \
  --out data/documents_resample_judged.jsonl

python src/data/documents.py select \
  --input data/documents_round1_judged.jsonl data/documents_resample_judged.jsonl \
  --out data/document_midtrain.jsonl \
  --min-quality 4 \
  --require-pass

Regenerate reasoning-format midtrain data:

python src/data/reasoning.py generate \
  --model Qwen/Qwen3-4B \
  --out-candidates data/reasoning_round1.jsonl

python src/data/reasoning.py judge \
  --candidates data/reasoning_round1.jsonl \
  --out data/reasoning_round1_judged.jsonl

python src/data/reasoning.py failed \
  --judged data/reasoning_round1_judged.jsonl \
  --list-out data/reasoning_resample_ids.txt

python src/data/reasoning.py generate \
  --model Qwen/Qwen3-4B \
  --source-list-file data/reasoning_resample_ids.txt \
  --out-candidates data/reasoning_resample.jsonl

python src/data/reasoning.py judge \
  --candidates data/reasoning_resample.jsonl \
  --out data/reasoning_resample_judged.jsonl

python src/data/reasoning.py select \
  --judged data/reasoning_round1_judged.jsonl data/reasoning_resample_judged.jsonl \
  --out data/reasoning_selected.jsonl \
  --require-pass

python src/data/reasoning.py render \
  --selected data/reasoning_selected.jsonl \
  --out data/reasoning_midtrain.jsonl

The document and reasoning judges default to Qwen/Qwen3-32B.

Training

Midtrain on document + reasoning data:

python src/train_midtrain.py --config configs/midtrain.yaml \
  --model_name_or_path Qwen/Qwen3-4B \
  --data_file data/document_midtrain.jsonl data/reasoning_midtrain.jsonl \
  --output_dir output/midtrain

configs/midtrain.yaml defaults to data_variant: reasoning_plus_document. The supported variants are reasoning_only, document_only, and reasoning_plus_document.

Build SFT data directly from reasoning midtrain data:

python src/sft/prepare_sft_data.py \
  --input data/reasoning_midtrain.jsonl \
  --output-dir data/sft/reasoning_midtraining_direct

Self-sample SFT data from dataset prompts:

python - <<'PY'
import json
from pathlib import Path

src = Path("data/reasoning_midtrain.jsonl")
out = Path("data/sft_seed/prompts.jsonl")
out.parent.mkdir(parents=True, exist_ok=True)
with src.open() as f, out.open("w") as g:
    for i, line in enumerate(f):
        row = json.loads(line)
        if row.get("user_prompt"):
            g.write(json.dumps({
                "sample_id": row.get("sample_id", f"row_{i}"),
                "prompt": row["user_prompt"],
                "reference_answer": row.get("answer") or row.get("assistant_content", ""),
                "dataset_group": "safety" if row.get("source_label") == "unsafe" else "general",
                "source_dataset": row.get("source_dataset", ""),
            }, ensure_ascii=False) + "\n")
PY

python src/sft/generate_sft_data.py \
  --model output/midtrain \
  --input data/sft_seed/prompts.jsonl \
  --output-dir data/sft/self_sampled \
  --no-few-shot

The SFT generator stops on <|im_end|>, <|im_start|>, and textual user/assistant/system turn markers so a completed assistant response cannot continue into another dialogue turn.

Train SFT:

python src/sft/train_sft.py --config configs/sft.yaml \
  --model_name_or_path output/midtrain \
  --sft_data self_sampled \
  --output_dir output/sft

Train GRPO using prompts from the reasoning dataset:

python src/grpo/train_grpo.py \
  --model-path output/midtrain \
  --data-file data/reasoning_selected.jsonl \
  --output-dir output/grpo

Eval

Run inference and scoring for base, midtrain, SFT, and GRPO models:

python src/eval/run_eval.py --config configs/eval.json --model-key base --gpu 0
python src/eval/run_eval.py --config configs/eval.json --model-key midtrain --gpu 0
python src/eval/run_eval.py --config configs/eval.json --model-key sft --gpu 0
python src/eval/run_eval.py --config configs/eval.json --model-key grpo --gpu 0

Use --generate-only to run inference and save predictions.jsonl without judge scoring. Benchmark JSONL files are included under eval/benchmarks/<benchmark>/<benchmark>.jsonl.