Implementation:LaurentMazare Tch rs Convert Checkpoint
| Knowledge Sources | |
|---|---|
| Domains | NLP, Model_Serialization |
| Last Updated | 2026-02-08 14:00 GMT |
Overview
External Python tool for converting original LLaMA checkpoints to safetensors format compatible with the tch-rs Rust implementation.
Description
convert_checkpoint.py is a Python script that loads LLaMA weights from a .pth file, renames keys to match the Rust model's VarStore path hierarchy, concatenates separate Q/K/V weights, converts to float16 precision, and saves as a safetensors file. This is a prerequisite step before running LLaMA inference in Rust.
Usage
Run this Python script once per checkpoint before Rust inference. Requires PyTorch and safetensors Python packages.
Code Reference
Source Location
- Repository: tch-rs
- File: examples/llama/convert_checkpoint.py
- Lines: 51-58
Signature
def convert_weights(llama_ckpt, *, output_st=Path("llama.safetensors"), dtype="float16"):
"""Convert LLaMA checkpoint to safetensors format."""
Import
import torch
from safetensors.torch import save_file
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| llama_ckpt | Path | Yes | Path to original LLaMA .pth checkpoint (e.g., consolidated.00.pth) |
| output_st | Path | No | Output safetensors path (default: llama.safetensors) |
| dtype | str | No | Target precision (default: float16) |
Outputs
| Name | Type | Description |
|---|---|---|
| safetensors file | File | Renamed and converted weights in safetensors format |
Usage Examples
# Command line usage
python examples/llama/convert_checkpoint.py /path/to/consolidated.00.pth
# Programmatic usage
from convert_checkpoint import convert_weights
convert_weights("/path/to/consolidated.00.pth", output_st="llama7b.safetensors")
Related Pages
Implements Principle
Requires Environment
Page Connections
Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment