File size: 876 Bytes
3738348 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | """Convert the latest checkpoint to bf16 to halve file size for Colab upload."""
import torch
import json
import os
def main():
ckpt = torch.load("checkpoints/latest.pt", map_location="cpu", weights_only=False)
print(f"Original checkpoint: step={ckpt['step']}, loss={ckpt['loss']}")
print(f" Embedding shape: {ckpt['model_state_dict']['token_embedding.weight'].shape}")
# Convert all float32 tensors to bf16
state = ckpt["model_state_dict"]
for key in state:
if state[key].dtype == torch.float32:
state[key] = state[key].to(torch.bfloat16)
ckpt["model_state_dict"] = state
out_path = "checkpoints/latest_bf16.pt"
torch.save(ckpt, out_path)
size_mb = os.path.getsize(out_path) / 1e6
print(f"bf16 checkpoint saved: {out_path} ({size_mb:.1f} MB)")
if __name__ == "__main__":
main()
|