Tensor bersama Torch

TL;DR

Menggunakan fungsi-fungsi tertentu, yang seharusnya berfungsi dalam kebanyakan kasus Anda. Hal ini tidak tanpa efek samping.

from safetensors.torch import load_model, save_model

save_model(model, "model.safetensors")
# Instead of save_file(model.state_dict(), "model.safetensors")

load_model(model, "model.safetensors")
# Instead of model.load_state_dict(load_file("model.safetensors"))

Apa itu tensor bersama?

Pytorch menggunakan tensor bersama untuk beberapa komputasi. Hal ini sangat menarik untuk mengurangi penggunaan memori secara umum.

Salah satu kasus penggunaan yang sangat klasik adalah pada transformers, di mana embeddings dibagikan dengan lm_head. Dengan menggunakan matriks yang sama, model menggunakan lebih sedikit parameter, dan gradien mengalir jauh lebih baik ke embeddings (yang merupakan bagian awal model, sehingga gradien tidak mudah mengalir ke sana, sedangkan lm_head berada di bagian akhir model, sehingga gradien sangat baik di sana; karena keduanya adalah tensor yang sama, keduanya mendapat manfaat)

from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(100, 100)
        self.b = self.a

    def forward(self, x):
        return self.b(self.a(x))


model = Model()
print(model.state_dict())
# odict_keys(['a.weight', 'a.bias', 'b.weight', 'b.bias'])
torch.save(model.state_dict(), "model.bin")
# This file is now 41k instead of ~80k, because A and B are the same weight hence only 1 is saved on disk with both `a` and `b` pointing to the same buffer

Mengapa tensor bersama tidak disimpan dalam safetensors?

Ada beberapa alasan untuk itu:

• Bukan semua framework mendukungnya, misalnya tensorflow tidak mendukung. Jadi, jika seseorang menyimpan tensor bersama di torch, tidak ada cara untuk memuatnya dengan cara yang serupa, sehingga kami tidak dapat mempertahankan API Dict[str, Tensor] yang sama.

• Hal ini membuat lazy loading menjadi sangat cepat. Lazy loading adalah kemampuan untuk memuat hanya beberapa tensor, atau bagian dari tensor, untuk file tertentu. Hal ini mudah dilakukan tanpa berbagi tensor, tetapi dengan berbagi tensor

  Disalin

  with safe_open("model.safetensors", framework="pt") as f:
      a = f.get_tensor("a")
      b = f.get_tensor("b")

  Sekarang, dengan kode yang diberikan ini, tidak mungkin untuk “reshare” buffer setelahnya. Setelah kita memberikan tensor a, kita tidak punya cara untuk mengembalikan memori yang sama ketika Anda meminta b. (Pada contoh khusus ini, kita bisa melacak buffer yang diberikan, tetapi ini tidak berlaku secara umum, karena Anda bisa melakukan pekerjaan arbitrer dengan a seperti mengirimkannya ke perangkat lain sebelum meminta b)

• Dapat menghasilkan file yang jauh lebih besar dari yang diperlukan. Jika Anda menyimpan tensor bersama yang hanya merupakan sebagian kecil dari tensor yang lebih besar, maka menyimpannya dengan pytorch akan menyimpan seluruh buffer alih-alih hanya menyimpan bagian yang diperlukan.

  Disalin

  a = torch.zeros((100, 100))
  b = a[:1, :]
  torch.save({"b": b}, "model.bin")
  # File is 41k instead of the expected 400 bytes
  # In practice it could happen that you save several 10GB instead of 1GB.

Dengan semua alasan tersebut telah disebutkan, tidak ada yang bersifat tetap di sana. Tensor bersama tidak menyebabkan ketidakamanan, atau potensi penolakan layanan, sehingga keputusan ini dapat ditinjau kembali jika solusi sementara saat ini tidak memuaskan.

Bagaimana cara kerjanya?

Desainnya cukup sederhana. Kita akan mencari semua tensor bersama, kemudian mencari semua tensor yang mencakup seluruh buffer (dapat ada beberapa tensor seperti itu). Hal ini memberikan beberapa nama yang dapat disimpan, kita cukup memilih yang pertama

Selama load_model, kami memuat data dengan cara yang mirip seperti yang dilakukan load_state_dict, kecuali kami memeriksa model itu sendiri untuk mengecek buffer yang dibagikan, dan mengabaikan “missed keys” yang sebenarnya sudah tercakup karena berbagi buffer (mereka dimuat dengan benar karena ada buffer yang dimuat di balik layar). Setiap error lainnya dilempar apa adanya

Perhatian: Ini berarti kami menghapus beberapa kunci di dalam file. Artinya, jika Anda memeriksa kunci yang disimpan di disk, Anda akan melihat beberapa “tensor hilang” atau jika Anda menggunakan load_state_dict. Kecuali kami mulai mendukung tensor bersama secara langsung dalam format tersebut, tidak ada cara nyata untuk menghindarinya.