gpt2-gsm8k-teacher-only

GPT-2 GSM8K teacher-only checkpoint

This repository contains the final teacher-only checkpoint used by the CoDi research code. It was trained as the visible-chain-of-thought teacher for the state-autoencoder and trajectory- supervised student experiments.

Checkpoint details

• Base model: gpt2
• Training data: GSM8K-Aug (data_name=icot)
• Training mode: teacher_only=True
• LoRA: rank 128, alpha 32, targets c_attn, c_proj, and c_fc
• Context length: 512
• Training: 40 epochs, 30,000 optimizer steps, learning rate 0.003
• Seed: 11
• Final weight file: pytorch_model.bin
• SHA-256: 2ff73124230aef272da686e78c0239c8d9228c425391a16e34302130d6565e3c

The weight file is the original, unmodified final CODI.state_dict. It contains the GPT-2 base weights and the unmerged LoRA weights. This is not an adapter-only checkpoint and is not a standalone Transformers checkpoint, so do not load this repository with AutoModelForCausalLM.from_pretrained("hanseungwook/gpt2-gsm8k-teacher-only"). Reconstruct the CoDi wrapper around gpt2 and then load this state dict, as the project code does.

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-gsm8k-teacher-only",
    local_dir="checkpoints/gpt2-gsm8k-teacher-only",
)
PY

export TEACHER_CKPT="$PWD/checkpoints/gpt2-gsm8k-teacher-only"

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

Use for state-autoencoder training

The canonical launcher reconstructs the teacher as GPT-2 plus rank-128 LoRA, loads the checkpoint, freezes it, and extracts hidden states for the autoencoder:

TEACHER_CKPT="$TEACHER_CKPT" \
  bash scripts/train_gpt2_state_ae_gsm8k-aug.sh

For a new dataset loader, call the underlying entry point and keep the teacher arguments unchanged:

python train_state_ae.py \
  --model_name_or_path gpt2 \
  --teacher_ckpt "$TEACHER_CKPT" \
  --teacher_use_lora True \
  --teacher_lora_r 128 \
  --teacher_lora_alpha 32 \
  --teacher_layer -1 \
  --data_name YOUR_DATASET_NAME \
  --output_dir outputs/state_ae_YOUR_DATASET \
  --num_slots 6 \
  --bf16 True

The relevant loader is src.teacher_states.FrozenTeacher. It accepts either a checkpoint directory or a weight-file path and reports the number of missing and unexpected keys after loading. Both counts should be zero.

Use for student training

For the autoencoder-supervised student, first train the state autoencoder as above. Then provide its state_ae.pt to the trajectory-supervised student:

AE_CKPT=outputs/state_ae_YOUR_DATASET/state_ae.pt \
  ARM=ae \
  bash scripts/train_gpt2_loop_trajectory_gsm8k-aug.sh

The stage-B cache rebuilds this frozen teacher from the path stored in the autoencoder checkpoint. Keep the downloaded teacher directory available at that path, or update the autoencoder/pipeline configuration when moving the run between machines.

For raw fixed-teacher anchor supervision instead of autoencoder targets, pass this directory as --traj_teacher_ckpt (the ARM=steps and ARM=random launchers do this through TEACHER_CKPT).

Low-level loading contract

If integrating a new pipeline directly, mirror src.teacher_states.FrozenTeacher:

1. Instantiate the CoDi wrapper with base model gpt2, teacher_only=True, use_lora=True, use_prj=False, and num_latent=0.
2. Use a PEFT LoraConfig with rank 128, alpha 32, and target modules c_attn, c_proj, and c_fc.
3. Load pytorch_model.bin into the wrapper and require zero missing and zero unexpected keys.
4. Freeze the teacher and put it in evaluation mode before extracting states.

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