Google-მა და Meta-მ წარადგინეს TorchTPU, PyTorch-ის ნატიური ბექენდი Google-ის TPU-ებისთვის

TorchTPU Google-ისა და Meta-ის ღია კოდის ერთობლივი პროექტია, რომელიც Google-ის TPU-ებისთვის PyTorch-ის ნატიურ ბექენდს ქმნის. ის საშუალებას იძლევა მოდელები ყოველგვარი დამატებითი მორგების გარეშე გაეშვას და TPU-ზე გადავიდეს, ამავდროულად მიზნად ისახავს ინდუსტრიაში კონკურენტუნარიან ინფერენსისა და წვრთნის წარმადობას eager execution-ის ოპტიმიზაციების, torch.compile-ის ინტეგრაციისა და აპარატურის ნატიური აჩქარების მეშვეობით.

Google-ის TPU-ები უზარმაზარი მასშტაბით გამოიყენებოდა დიდი მანქანური სწავლების მოდელების წვრთნისა და გაშვებისთვის. მართალია, TPU-ები ისტორიულად TensorFlow-სა და JAX-ზე იყვნენ მიბმული, ეს პროექტი Google-ში 2025 წლის შუაში დაიწყო Meta-ის PyTorch გუნდთან ერთად. ადრინდელი PyTorch/XLA ინტეგრაცია, რომელიც 2020 წელს დანერგეს და Lazy Tensors-ს ეფუძნებოდა, საწყის კოდში ღრმა ჩარევას მოითხოვდა, კომპილაციის დროს სერიოზული შეყოვნება ჰქონდა, eager execution-ის მხარდაჭერა არ ჰქონდა და განაწილებულ პრიმიტივებში ჩამორჩებოდა.

დიზაინის პრინციპები და გამოყენებადობა

TorchTPU შეიქმნა როგორც პრინციპული TPU ბექენდი PyTorch-ისთვის და არა JAX-ის ადაპტაცია. მის პრინციპებს შორისაა PyTorch-ის სრულფასოვანი მოქალაქეობა (runtime-ში JAX-ზე დამოკიდებულების არარსებობა, საიმედო eager mode, Torch Dynamo-ზე დაფუძნებული torch.compile და DDP-ის, FSDP-ისა და DTensor-ის მხარდაჭერა), ეკოსისტემასთან უწყვეტი ინტეგრაცია და საზოგადოების ჩართულობა. მომხმარებლებზე ორიენტირება პორტირების სიმარტივეს გულისხმობს: model.to("cuda")-ის model.to("tpu")-ით ჩანაცვლებით მოდელი ყოველგვარი დამატებითი მორგების გარეშე უნდა გაეშვას, თუმცა პიკური წარმადობისთვის შესაძლოა დამატებითი ტუნინგი დაგჭირდეს. პროექტი ასევე მისდევს შემოსაზღვრულ დინამიურობას, რათა თავიდან აიცილოს ფორმის ცვლილებების გამო XLA-ს ხელახალი კომპილაცია, და იკვლევს დინამიკურ TPU kernel-ებს, რომელთა წინასწარ კომპილირებაც შესაძლებელია. TPU მანქანაზე, სადაც TorchTPU-ის wheel პაკეტი არის დაყენებული, PyTorch-ის იმპორტი მხოლოდ უმნიშვნელო ცვლილებებს მოითხოვს PyTorch-ის autoload მექანიზმის მეშვეობით.

Eager execution და კომპილაცია

TorchTPU PyTorch-ის დისპეტჩერს აფართოებს C++ ბექენდით, რომელიც PrivateUse1-ად არის რეგისტრირებული, და TPU ტენზორის შენახვას განსაზღვრავს სტანდარტული at::Tensor მექანიზმით. ის გვთავაზობს eager რეჟიმის სამ ვარიანტს: Strict (ნაგულისხმევი) ყოველ ჯერზე ერთ Aten op-ს ასინქრონულად აგზავნის; Debug (TORCH_TPU_LAUNCH_BLOCKING=1) სინქრონულია და დამატებითი საზღვრების შემოწმებასაც მოიცავს; ხოლო DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) მრავალ op-ს აყოვნებს და აერთიანებს, რაზეც ნათქვამია, რომ Strict eager-ზე 50%-ით ან მეტით სწრაფია და პიკურ HBM მეხსიერებასაც ამცირებს. მატერიალიზაციის მომენტში, მაგალითად .to("cpu") გამოძახებისას, TorchTPU ატარებს liveness ანალიზს, op-ების გრაფისთვის ქეშის გასაღებს ქმნის და ან დაქეშირებულ შესრულებადს აგზავნის, ან ასინქრონულად კომპილირებს. Repeated Ops Heuristic ამოიცნობს განმეორებით მიმდევრობებს, მაგალითად ტრანსფორმერის ბლოკების დასტებს, და მატერიალიზაციას ადრე იწყებს. სამდონიანი კომპილაციის ქეში მოიცავს პროცესის ლოკალურ მეხსიერებას, მანქანის ლოკალურ საერთო მეხსიერებას და მუდმივ დისტანციურ ფაილურ სისტემის ქეშს, რათა ხელახლა გამოყენებადი იყოს ჰოსტებსა და გადატვირთვებს შორის.

მხარდაჭერილი დატვირთვები და სიზუსტე

TorchTPU მხარს უჭერს დაახლოებით 450 Aten op-ს, მათ შორის ძირითადი არქიტექტურებისთვის საჭირო Core ATen op-ების forward და backward გავლებს, ასევე მორგებულ kernel-ებს JAX/Pallas-ის ან Helion DSL-ის მეშვეობით. ის ყოველგვარი დამატებითი მორგების გარეშე შემოწმდა 5 000 ყველაზე ხშირად ჩამოტვირთულ Hugging Face მოდელზე, რომლებიც მოიცავს LLM-ებს, აუდიოსა და მეტყველებას, ხედვასა და ტექსტურ ენკოდერებს. ასევე შემოწმდა DTensor-ზე დაფუძნებული sharding, მრავალჰოსტიანი inference vLLM-ითა და SGLang-ით, fine-tuning და წვრთნა TorchTitan-ით და ინტეგრაცია Ray-სა და TorchRec-თან. TPU-ის მატრიცული გამრავლების ერთეულები იყენებენ bfloat16 შეყვანას FP32-ის დაგროვებით, ხოლო non-contracting op-ები ვექტორულ ერთეულზე FP32-ში მუშაობენ. TorchTPU PyTorch-ის სიზუსტის API-ებს ამ აპარატურაზე ასახავს: medium, რომელიც ნაგულისხმევია, იყენებს BF16 გამრავლებას FP32-ის დაგროვებით; high იყენებს 3-გავლიან ალგორითმს დაახლოებით 3-ჯერ მეტი გამოთვლითი ხარჯის ფასად; highest იყენებს 6 გავლას დაახლოებით 6-ჯერ მეტი დანახარჯის ფასად. high-ის ან highest-ის არჩევა წარმადობის გაფრთხილებას იწვევს.

წყაროები: PyTorch · Google Cloud Blog · Hugging Face Blog