safetensors-layer-grabber

Extracts all weights corresponding to a specific layer or block from large models saved in safetensors format, and exports them as a PyTorch-compatible state dict.

Example

safetensors_layer_grabber \
  --model-dir ~/models/my_model \
  --layer model.layers.12.mlp \
  --output layer12_mlp.ckpt

Features

• Targeted Extraction: Extract weights corresponding to any layer or nested component by name prefix (e.g. model.layers.8.self_attn)
• Flexible Output: Choose your output filename, or default to <layer>.ckpt
• Efficient: Works directly with multipart safetensors files; no need to load the entire model

Installation

pip install safetensors-layer-grabber

After installation, run it as safetensors_layer_grabber.

Arguments

Understanding Parameter Names

In safetensors files

Safetensors model shards store parameters as key–tensor pairs, with keys like:

model.layers.8.self_attn.q_proj.weight
model.layers.8.self_attn.k_proj.weight
model.layers.8.self_attn.v_proj.bias

In the extracted .ckpt file

When you use a prefix with --layer, that prefix is stripped from the keys in the saved state dict.

• Example:
  • Input safetensors key: model.layers.8.self_attn.q_proj.weight
  • Layer argument: model.layers.8.self_attn
  • Output key in .ckpt: q_proj.weight
• Another example:
  • Input safetensors key: model.layers.10.mlp.fc1.weight
  • Layer argument: model.layers.10.mlp
  • Output key in .ckpt: fc1.weight

Key Mapping Table

This makes the exported layer weights easy to load with torch.load or directly into submodules.

FAQ & Notes

• If the specified --layer prefix doesn't match any weights exactly, the script will raise an exception and not write a file.
• If no .safetensors files are found in the directory, an error will be raised.

Contributing

Contributions are welcome! Please submit pull requests or open issues on the GitHub repository.

License

This project is licensed under the MIT License.