gpt2-proofwriter-cwa-depth5-teacher-only

GPT-2 ProofWriter CWA depth-5 teacher-only checkpoint

This repository contains the final ProofWriter CWA depth-5 teacher-only checkpoint used by the CoDi research code. It is intended as a frozen teacher for state-autoencoder, trajectory-target, and student training.

Checkpoint details

• Base model: gpt2
• Training data: ProofWriter V2020.12.3, CWA, depth-5, train split
• CoDi data mode: proofwriter
• Training mode: teacher_only=True
• LoRA: rank 128, alpha 32, targets c_attn, c_proj, and c_fc
• Context length: 1024
• Training: 10 epochs, 1,370 optimizer steps, learning rate 0.003
• Seed: 11
• Final weight file: pytorch_model.bin
• SHA-256: 5c234e68c36be837ad837fe73ef9dd6843967eec39c9175ca89696905b7ba853

The weight file is the original, unmodified final CODI.state_dict. It contains the GPT-2 base weights and unmerged LoRA weights. This is not an adapter-only checkpoint or a standalone Transformers checkpoint. Do not load this repository with AutoModelForCausalLM.from_pretrained(); reconstruct the CoDi wrapper around gpt2 and load the state dict as described below.

Download

git clone https://github.com/hanseungwook/codi.git
cd codi
pip install -r requirements.txt

python - <<'PY'
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id="hanseungwook/gpt2-proofwriter-cwa-depth5-teacher-only",
    local_dir="checkpoints/gpt2-proofwriter-cwa-depth5-teacher-only",
)
PY

export TEACHER_CKPT="$PWD/checkpoints/gpt2-proofwriter-cwa-depth5-teacher-only"

TEACHER_CKPT may point to the downloaded directory or directly to its pytorch_model.bin file.

Load as a frozen teacher

import torch
from src.teacher_states import FrozenTeacher

teacher = FrozenTeacher(
    model_name_or_path="gpt2",
    teacher_ckpt="checkpoints/gpt2-proofwriter-cwa-depth5-teacher-only",
    layer=-1,
    torch_dtype=torch.bfloat16,
    use_lora=True,
    lora_r=128,
    lora_alpha=32,
)

FrozenTeacher reconstructs GPT-2 plus the LoRA modules, loads this state dict, freezes the model, and exposes hidden-state extraction. It reports missing and unexpected keys after loading; both counts should be zero.

Autoencoder and student integration

The current src/ae_data.py implementation is GSM-specific. For ProofWriter, add a dataset adapter that preserves the project's formatted source, proof trace/CoT, and answer fields and produces the tokenized inputs expected by the state-AE pipeline. The teacher reconstruction above does not need to change.

Invoke train_state_ae.py with this directory as --teacher_ckpt, keeping:

--model_name_or_path gpt2
--teacher_use_lora True
--teacher_lora_r 128
--teacher_lora_alpha 32
--teacher_layer -1

For an AE-supervised student, pass the resulting state_ae.pt as --ae_ckpt. The stage-B cache rebuilds the frozen teacher from the path stored in the AE checkpoint, so keep this download available or update the pipeline configuration when moving runs between machines.

For raw fixed-teacher trajectory anchors, pass this directory directly as --traj_teacher_ckpt.

Tokenizer files, training_args.bin, and trainer_state.json are included unchanged alongside the final weights for provenance.