Skip to content

Support convert checkpoint from any directory #148

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shuheng-liu merged 5 commits into from
Mar 26, 2026
Merged

Support convert checkpoint from any directory #148

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f2c5090
fix: support ckpt conversion from a different location than repo root
shuheng-liu Mar 18, 2026
944eab5
feat: skip conversion if model.safetensors is already present
shuheng-liu Mar 18, 2026
cfd806a
Update src/opentau/scripts/convert_checkpoint.sh
shuheng-liu Mar 23, 2026
6f5e4fd
chore: cursor review
shuheng-liu Mar 25, 2026
ded6b4e
chore: merge upstream changes with cursor
shuheng-liu Mar 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions src/opentau/scripts/convert_checkpoint.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,15 +36,23 @@ if [ ! -d "$CHECKPOINT_DIR" ]; then
exit 1
fi

if [ -f "$CHECKPOINT_DIR/model.safetensors" ]; then
echo "model.safetensors already exists in '$CHECKPOINT_DIR'. Skipping conversion."
exit 0
fi

echo "Converting checkpoint in directory: $CHECKPOINT_DIR"

# Step 1: Convert sharded checkpoint to full state dict
echo "Step 1: Converting sharded checkpoint to full state dict..."
python src/opentau/scripts/zero_to_fp32.py "$CHECKPOINT_DIR" "$CHECKPOINT_DIR/full_state_dict" --max_shard_size 1000GB
python -m opentau.scripts.zero_to_fp32 "$CHECKPOINT_DIR" "$CHECKPOINT_DIR/full_state_dict" --max_shard_size 1000GB

# Step 2: Convert pytorch_model.bin to model.safetensors
echo "Step 2: Converting pytorch_model.bin to model.safetensors..."
python src/opentau/scripts/bin_to_safetensors.py "$CHECKPOINT_DIR/full_state_dict/pytorch_model.bin" --output_file "$CHECKPOINT_DIR/model.safetensors"
python -m opentau.scripts.bin_to_safetensors "$CHECKPOINT_DIR/full_state_dict/pytorch_model.bin" --output_file "$CHECKPOINT_DIR/model.safetensors"

echo "Step 3: Cleaning up intermediate files..."
rm -rf "$CHECKPOINT_DIR/full_state_dict"

echo "Conversion completed successfully!"
echo "Model saved as: $CHECKPOINT_DIR/model.safetensors"
Loading