Torch paylaşımlı tensörler

TL;DR

Çoğu durumda sizin için işe yarayacak belirli işlevleri kullanın. Bu, yan etkisiz değildir.

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"))

Paylaşılan tensörler nedir?

Pytorch, bazı hesaplamalar için paylaşılan tensörler kullanır. Bu, genel olarak bellek kullanımını azaltmak açısından son derece ilginçtir.

Çok klasik bir kullanım senaryosu, transformers'ta embeddings'in lm_head ile paylaşılmasıdır. Aynı matrisi kullanarak model daha az parametre kullanır ve gradyanlar embeddings'e (modelin başlangıcı olduğu için oraya kolayca akamazlar, oysa lm_head modelin sonunda olduğu için gradyanlar orada son derece iyidir, aynı tensörler olduğu için her ikisi de bundan faydalanır) çok daha iyi akar.

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

Neden paylaşılan tensörler safetensors'ta kaydedilmez?

Bunun birden fazla nedeni var:

• Çerçevelerin tümü bunları desteklemez, örneğin tensorflow desteklemez. Dolayısıyla biri torch'ta paylaşılan tensörleri kaydederse, bunları benzer bir şekilde yüklemenin bir yolu yoktur ve bu nedenle aynı Dict[str, Tensor] API'sini koruyamazdık.

• Bu, tembel yüklemeyi çok hızlı hale getirir. Tembel yükleme, belirli bir dosya için yalnızca bazı tensörleri veya tensörlerin bir kısmını yükleme yeteneğidir. Tensör paylaşımı olmadan bunu yapmak çok basittir ancak tensör paylaşımı ile

  Kopyalandı

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

  Şu an, verilen bu kodla sonradan tamponları “yeniden paylaşmak” imkânsızdır. a tensörünü verdiğimizde, siz b istediğinizde aynı belleği geri verme yolumuz yoktur. (Bu özel örnekte verilen tamponları takip edebilirdik, ancak bu genel olarak geçerli değildir; çünkü a üzerinde b istenmeden önce başka bir cihaza gönderme gibi keyfi işlemler yapabilirsiniz)

• Gereksiz yere çok daha büyük dosyalara yol açabilir. Yalnızca daha büyük bir tensörün bir kısmı olan paylaşılan bir tensörü kaydediyorsanız, pytorch ile kaydetme işlemi, yalnızca gerekli olanı kaydetmek yerine tüm tamponu kaydeder.

  Kopyalandı

  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.

Tüm bu gerekçeler belirtilmiş olsa da, bu konuda kesin bir karar yoktur. Paylaşılan tensörler güvensizliğe veya hizmetin reddi potansiyeline yol açmadığından, mevcut geçici çözümler tatmin edici değilse bu karar yeniden değerlendirilebilir.

Nasıl çalışır?

Tasarım oldukça basittir. Tüm paylaşılan tensörleri, ardından tamponun tamamını kapsayan tüm tensörleri (böyle birden fazla tensör olabilir) arıyoruz. Bu, kaydedilebilecek birden fazla ad verir; biz de basitçe ilkini seçeriz

load_model sırasında, load_state_dict gibi yükleriz; ancak paylaşılan tamponları kontrol etmek için modelin kendisine bakar ve aslında tampon paylaşımı sayesinde kapsanmış olan “eksik anahtarları” yoksayarız (altta yatan bir tamponun yüklendiği için bunlar zaten doğru şekilde yüklenmiştir). Diğer tüm hatalar olduğu gibi fırlatılır.

Not: Bu, dosya içindeki bazı anahtarların atılacağı anlamına gelir. Yani diske kaydedilen anahtarları kontrol ediyorsanız bazı “eksik tensörler” görebilirsiniz veya load_state_dict kullanıyorsanız aynı durumla karşılaşabilirsiniz. Format içinde paylaşılan tensörleri doğrudan desteklemeye başlamadığımız sürece bunun gerçek bir çözümü yoktur.