split-safetensors-into-shards.py

# Author: Daniel de Kok
# Usage: python shard.py --safetensors-path /fsx/danieldk/4bit-gptq-instruct/gptq_model-4bit-128g.safetensors --framework torch --output-path /fsx/danieldk/4bit-gptq-instruct/gptq-sharded

import argparse

import safetensors
import huggingface_hub


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--safetensors-path", type=str)
    parser.add_argument("--framework", type=str)
    parser.add_argument("--output-path", type=str)
    return parser.parse_args()


class FakeStateDict:
    def __init__(self, weights):
        self.weights = weights

    def items(self):
        for key in self.weights.keys():
            yield key, self.weights.get_tensor(key)

    def __getitem__(self, item):
        return self.weights.get_tensor(item)


if __name__ == "__main__":
    args = get_args()

    weights = safetensors.safe_open(args.safetensors_path, framework=args.framework)
    huggingface_hub.save_torch_state_dict(FakeStateDict(weights), args.output_path)  # type: ignore