Google Cloud-ის TPUs არის სპეციალურად შექმნილი AI ამაჩქარებლები, რომლებიც ოპტიმიზებულია დიდი ხელოვნური ინტელექტის მოდელების, მათ შორის უახლესი LLM-ებისა და გენერაციული AI მოდელების, როგორიცაა SDXL, ტრეინინგისა და ინფერენსისთვის. ახალი Cloud TPU v5e შექმნილია იმ მიზნით, რომ უზრუნველყოს დიდი მასშტაბის AI ტრეინინგისა და ინფერენსისთვის საჭირო ხარჯთეფექტურობა და წარმადობა. TPU v4-ის ღირებულების ნახევარზე ნაკლები ფასით, TPU v5e მეტ ორგანიზაციას აძლევს საშუალებას მოახდინოს AI მოდელების ტრეინინგი და დეპლოირება. Diffusers-ის JAX ინტეგრაცია გთავაზობთ მოსახერხებელ გზას SDXL-ის TPU-ზე XLA-ს საშუალებით გასაშვებად, და ჩვენ შევქმენით დემო ვერსია ამის საჩვენებლად. შეგიძლიათ სცადოთ ეს Space-ში ან ქვემოთ ჩაშენებულ playground-ში: შიგნიდან, ეს დემო მუშაობს რამდენიმე TPU v5e-4 ინსტანციაზე (თითოეულ ინსტანციას აქვს 4 TPU ჩიპი) და იყენებს პარალელიზაციას ოთხი დიდი 1024×1024 ზომის გამოსახულების დაახლოებით 4 წამში მოსამსახურებლად. ეს დრო მოიცავს ფორმატის კონვერტაციებს, კომუნიკაციის დროს და ფრონტენდ დამუშავებას; რეალური გენერაციის დრო დაახლოებით 2.3 წამია, როგორც ქვემოთ ვნახავთ! ამ ბლოგპოსტში, SDXL-ის JAX-ით სერვისინგი Cloud TPU v5e-ზე მაღალი წარმადობითა და ხარჯთეფექტურობით შესაძლებელია სპეციალურად შექმნილი TPU აპარატურისა და წარმადობისთვის ოპტიმიზებული პროგრამული უზრუნველყოფის სტეკის კომბინაციის წყალობით. ქვემოთ ხაზს ვუსვამთ ორ ძირითად ფაქტორს: JAX-ის just-in-time (jit) კომპილაციას და XLA კომპილატორზე დაფუძნებულ პარალელიზაციას JAX pmap-თან ერთად. JAX-ის გამორჩეული მახასიათებელია მისი just-in-time (jit) კომპილაცია. JIT კომპილატორი კოდს აკონტროლებს პირველი გაშვებისას და ქმნის მაღალოპტიმიზებულ TPU ბინარებს, რომლებიც ხელახლა გამოიყენება შემდგომი გამოძახებისას. ამ პროცესის თავისებურება ის არის, რომ ის მოითხოვს ყველა შემავალი, შუალედური და გამომავალი ფორმის სტატიკურობას, რაც ნიშნავს, რომ ისინი წინასწარ უნდა იყოს ცნობილი. ყოველ ჯერზე, როცა ფორმებს შევცვლით, ახალი და ძვირადღირებული კომპილაციის პროცესი ხელახლა დაიწყება. JIT კომპილაცია იდეალურია სერვისებისთვის, რომლებიც შეიძლება დაპროექტდეს სტატიკური ფორმების გარშემო: კომპილაცია ერთხელ სრულდება, შემდეგ კი სუპერ-სწრაფი ინფერენსის დროებით ვსარგებლობთ. გამოსახულების გენერაცია კარგად შეესაბამება JIT კომპილაციას. თუ ყოველთვის ერთსა და იმავე რაოდენობის გამოსახულებებს ვაგენერირებთ და მათ ერთნაირი ზომა აქვთ, მაშინ გამომავალი ფორმები მუდმივია და წინასწარ ცნობილი. ტექსტური შეტანებიც მუდმივია: დიზაინით, Stable Diffusion და SDXL იყენებენ ფიქსირებული ფორმის ჩაშენებულ ვექტორებს (padding-ით) მომხმარებლის მიერ აკრეფილი მოთხოვნების წარმოსადგენად. აქედან გამომდინარე, შეგვიძლია დავწეროთ JAX კოდი, რომელიც ეყრდნობა ფიქსირებულ ფორმებს და რომლის ოპტიმიზაცია დიდად არის შესაძლებელი! სამუშაო დატვირთვების სკალირება შესაძლებელია მრავალ მოწყობილობაზე JAX-ის pmap-ის გამოყენებით, რომელიც გამოხატავს ერთი პროგრამის მრავალ მონაცემთა (SPMD) პროგრამებს. pmap-ის ფუნქციაზე გამოყენება მოახდენს ფუნქციის კომპილაციას XLA-ით, შემდეგ კი მას პარალელურად შეასრულებს სხვადასხვა XLA მოწყობილობაზე. ტექსტიდან გამოსახულების გენერაციის სამუშაო დატვირთვებისთვის ეს ნიშნავს, რომ ერთდროულად რენდერირებული გამოსახულებების რაოდენობის გაზრდა მარტივად ხორციელდება და არ აზიანებს წარმადობას. მაგალითად, SDXL-ის 8-ჩიპიანი TPU-ზე გაშვება მოახდენს 8 გამოსახულების გენერაციას იმავე დროში, რაც 1 ჩიპს სჭირდება ერთი გამოსახულების შესაქმნელად. TPU v5e ინსტანციები წარმოდგენილია მრავალი ფორმით, მათ შორის 1, 4 და 8-ჩიპიანი ფორმებით, 256 ჩიპამდე (სრული TPU v5e pod), ჩიპებს შორის ულტრა-სწრაფი ICI კავშირებით. ეს საშუალებას გაძლევთ აირჩიოთ TPU-ის ის ფორმა, რომელიც საუკეთესოდ შეესაბამება თქვენს გამოყენების შემთხვევას და მარტივად გამოიყენოთ JAX-ისა და TPUs-ის მიერ მოწოდებული პარალელიზაცია. ეტაპობრივად განვიხილავთ კოდს, რომელიც გჭირდებათ JAX-ის გამოყენებით სუპერ-სწრაფი ინფერენსის გასაშვებად! პირველ რიგში, შემოვიტანოთ დამოკიდებულებები. ახლა ჩავტვირთავთ ძირითად SDXL მოდელს და ინფერენსისთვის საჭირო დანარჩენ კომპონენტებს. Diffusers-ის pipeline ზრუნავს ყველაფრის ჩამოტვირთვასა და ქეშირებაზე ჩვენთვის. JAX-ის ფუნქციონალური მიდგომის გათვალისწინებით, მოდელის პარამეტრები ცალკე ბრუნდება და ინფერენსის დროს უნდა გადაეცეს pipeline-ს: მოდელის პარამეტრები ნაგულისხმევად ჩამოიტვირთება 32-ბიტიანი სიზუსტით. მეხსიერების შესანარჩუნებლად და გამოთვლების დასაჩქარებლად მათ გადავიყვანთ bfloat16-ში, ეფექტურ 16-ბიტიან წარმოდგენაში. თუმცა, არის ერთი გაფრთხილება: საუკეთესო შედეგებისთვის, scheduler-ის მდგომარეობა უნდა შევინარჩუნოთ float32-ში, წინააღმდეგ შემთხვევაში სიზუსტის შეცდომები დაგროვდება და დაბალი ხარისხის ან თუნდაც შავ გამოსახულებებს გამოიწვევს. ახლა მზად ვართ, მოვაწყოთ ჩვენი prompt და pipeline-ის დანარჩენი შეტანები. prompt-ები უნდა მიეწოდოს pipeline-ს ტენზორების სახით, და მათ ყოველთვის უნდა ჰქონდეთ ერთი და იგივე განზომილებები გამოძახებებს შორის. ეს საშუალებას აძლევს ინფერენსის გამოძახებას, რომ მოხდეს მისი კომპილაცია. pipeline-ის prepare_inputs მეთოდი ასრულებს ჩვენთვის ყველა საჭირო ნაბიჯს, ამიტომ შევქმნით დამხმარე ფუნქციას, რათა მოვამზადოთ როგორც ჩვენი prompt, ასევე ნეგატიური prompt ტენზორების სახით. მოგვიანებით მას ჩვენი generate ფუნქციიდან გამოვიყენებთ: პარალელიზაციის უპირატესობის გამოსაყენებლად, შეტანებს მოწყობილობებზე გავამრავლებთ. Cloud TPU v5e-4-ს აქვს 4 ჩიპი, ამიტომ შეტანების გამრავლებით თითოეული ჩიპი პარალელურად ქმნის განსხვავებულ გამოსახულებას. ყურადღება უნდა მივაქციოთ, რომ თითოეულ ჩიპს მივაწოდოთ განსხვავებული შემთხვევითი "seed", რათა 4 გამოსახულება განსხვავებული იყოს: ახლა მზად ვართ, ყველაფერი ერთად მოვაქციოთ generate ფუნქციაში: jit=True მიუთითებს, რომ გვინდა pipeline-ის გამოძახების კომპილაცია. ეს მოხდება generate-ის პირველი გამოძახებისას, და ის იქნება ძალიან ნელი