[Huggingface] Safetensor sharding

Huggingface Transformers를 사용하여 모델을 저장하면 bin 파일로 모델이 저장된다.

하지만 최근 많은 모델들이 Safetensor를 사용하여 모델을 저장하고 있는데, 특히 모델의 weight를 sharding하여 저장한다.

방법은 간단한데, safetensor 라이브러리에서 지원하는 기능은 따로 없는듯하고, transformers 라이브러리에서 저장할때 몇 가지 옵션만 추가하면 된다.

모델을 저장할때, max_shard_size와 safe_serialization 옵션을 설정하면 된다. max_shard_size에는 str로 얼마만큼의 크기로 모델을 sharding 할지 적어주면 된다. safe_serialization은 True로 설정하도록하자.

model.save_pretrained(
        out_folder, 
        max_shard_size="1GB", 
        safe_serialization=True
    )