PyTorch-TPU პროექტი დაიწყო როგორც Facebook PyTorch-ისა და Google TPU-ის გუნდების თანამშრომლობის შედეგი და ოფიციალურად გაეშვა 2019 წლის PyTorch დეველოპერების კონფერენციაზე. მას შემდეგ, ჩვენ ვითანამშრომლეთ Hugging Face-ის გუნდთან, რათა PyTorch / XLA-ს გამოყენებით Cloud TPU-ებზე ვარჯიშს უმაღლესი დონის მხარდაჭერა მივცეთ. ეს ახალი ინტეგრაცია საშუალებას აძლევს PyTorch-ის მომხმარებლებს, გაუშვან და გააფართოვონ თავიანთი მოდელები Cloud TPU-ებზე, Hugging Face-ის ტრენერების ზუსტად იგივე ინტერფეისის შენარჩუნებით. ეს ბლოგ-პოსტი გვაწვდის მიმოხილვას Hugging Face-ის ბიბლიოთეკაში შეტანილი ცვლილებების შესახებ, თუ რას აკეთებს PyTorch / XLA ბიბლიოთეკა, მოჰყავს მაგალითი იმისა, თუ როგორ უნდა დაიწყოთ თქვენი საყვარელი ტრანსფორმერების ვარჯიში Cloud TPU-ებზე, და წარმოადგენს წარმადობის რამდენიმე ბენჩმარკს. თუ ვერ ითმენთ TPU-ებთან მუშაობის დაწყებას, გთხოვთ, პირდაპირ გადახვიდეთ სექციაზე „ავარჯიშეთ თქვენი ტრანსფორმერი Cloud TPU-ებზე“ – ჩვენ PyTorch / XLA-ის მთელ მექანიკას თქვენთვის Trainer მოდულის ფარგლებში ვმართავთ! PyTorch / XLA PyTorch-ს უმატებს ახალ xla მოწყობილობის ტიპს. ეს მოწყობილობის ტიპი მუშაობს ისევე, როგორც PyTorch-ის სხვა მოწყობილობის ტიპები. მაგალითად, აი, როგორ შევქმნათ და დავბეჭდოთ XLA ტენზორი: ეს კოდი ნაცნობი უნდა იყოს. PyTorch / XLA იყენებს იმავე ინტერფეისს, რასაც ჩვეულებრივი PyTorch, მცირე დამატებებით. torch_xla-ს იმპორტი ახდენს PyTorch / XLA-ის ინიციალიზაციას, ხოლო xm.xla_device() აბრუნებს მიმდინარე XLA მოწყობილობას. ეს შეიძლება იყოს CPU, GPU ან TPU თქვენი გარემოდან გამომდინარე, მაგრამ ამ ბლოგ-პოსტში ძირითადად TPU-ზე გავამახვილებთ ყურადღებას. Trainer მოდული იყენებს TrainingArguments მონაცემთა კლასს ვარჯიშის სპეციფიკის განსაზღვრისთვის. ის ამუშავებს მრავალ არგუმენტს, დაწყებული batch sizes-დან, learning rate-დან, gradient accumulation-დან და სხვა პარამეტრებიდან, დამთავრებული გამოყენებული მოწყობილობებით. ზემოაღნიშნულის საფუძველზე, TrainingArguments._setup_devices()-ში XLA:TPU მოწყობილობების გამოყენებისას, ჩვენ უბრალოდ ვაბრუნებთ TPU მოწყობილობას, რომელიც Trainer-ის მიერ იქნება გამოყენებული: ტიპურ XLA:TPU ვარჯიშის სცენარში, ჩვენ პარალელურად ვავარჯიშებთ მრავალ TPU ბირთვზე (ერთი Cloud TPU მოწყობილობა მოიცავს 8 TPU ბირთვს). ამიტომ უნდა დავრწმუნდეთ, რომ ყველა გრადიენტი გაიცვლება მონაცემთა პარალელურ რეპლიკებს შორის გრადიენტების კონსოლიდაციით და ოპტიმიზატორის ნაბიჯის გადადგმით. ამისათვის ჩვენ გთავაზობთ xm.optimizer_step(optimizer)-ს, რომელიც ახორციელებს გრადიენტების კონსოლიდაციას და ნაბიჯების გადადგმას. Hugging Face-ის ტრენერში, ჩვენ შესაბამისად ვაახლებთ ვარჯიშის ნაბიჯს PyTorch / XLA API-ების გამოსაყენებლად: PyTorch / XLA მოდელის გასაშვებად ორი ძირითადი ნაწილია: (1) თქვენი მოდელის გრაფის ზარმაცად (lazily) მიკვლევა და შესრულება (უფრო დეტალური ახსნისთვის იხილეთ ქვემოთ მოცემული „PyTorch / XLA ბიბლიოთეკა“ სექცია) და (2) თქვენი მოდელის შეყვანა (feeding). ყოველგვარი ოპტიმიზაციის გარეშე, თქვენი მოდელის მიკვლევა/შესრულება და შეყვანის მიწოდება სერიულად შესრულდებოდა, რაც დროს ტოვებდა, რა დროსაც თქვენი ჰოსტი CPU და თქვენი TPU ამაჩქარებლები უმოქმედო იქნებოდნენ, შესაბამისად. ამის თავიდან ასაცილებლად, ჩვენ გთავაზობთ API-ს, რომელიც ამ ორს ერთმანეთთან აკავშირებს და ამგვარად შეუძლია n+1 ნაბიჯის მიკვლევის გადაფარვა, სანამ n ნაბიჯი ჯერ კიდევ სრულდება. როდესაც ტენზორი XLA მოწყობილობიდან კონტროლდება (checkpointed) და შემდეგ უკან იტვირთება კონტროლიდან, ის უკანვე ჩაიტვირთება თავდაპირველ მოწყობილობაში. მოდელში ტენზორების კონტროლამდე, უნდა დარწმუნდეთ, რომ თქვენი ყველა ტენზორი არის CPU მოწყობილობებზე და არა XLA მოწყობილობებზე. ამ გზით, როდესაც ტენზორებს უკან ჩატვირთავთ, მათ ჩატვირთავთ CPU მოწყობილობების მეშვეობით და შემდეგ გექნებათ შესაძლებლობა, განათავსოთ ისინი ნებისმიერ XLA მოწყობილობაზე, რომელიც გსურთ. ამისათვის ჩვენ გთავაზობთ xm.save() API-ს, რომელიც უკვე ზრუნავს იმაზე, რომ მხოლოდ ერთი პროცესით ჩაიწეროს მონაცემები შენახვის ადგილას თითოეულ ჰოსტზე (ან ერთი გლობალურად, თუ გამოყენებულია საერთო ფაილური სისტემა ჰოსტებს შორის). PyTorch / XLA არის Python პაკეტი, რომელიც იყენებს XLA წრფივი ალგებრის კომპილერს PyTorch ღრმა სწავლის ფრეიმვორკის XLA მოწყობილობებთან დასაკავშირებლად, რაც მოიცავს CPU, GPU და Cloud TPU-ებს. ქვემოთ მოცემული შინაარსის ნაწილი ასევე ხელმისაწვდომია ჩვენს API_GUIDE.md-ში. XLA ტენზორებისა და მოწყობილობების გამოყენება მხოლოდ რამდენიმე სტრიქონის კოდის შეცვლას მოითხოვს. თუმცა, მიუხედავად იმისა, რომ XLA ტენზორები ძალიან ჰგავს CPU და CUDA ტენზორებს, მათი შიდა სტრუქტურა განსხვავებულია. CPU და CUDA ტენზორები ოპერაციებს დაუყოვნებლივ ან ენთუზიაზმით (eagerly) იწყებენ. XLA ტენზორები, მეორე მხრივ, ზარმაცია (lazy). ისინი ოპერაციებს გრაფაში აღრიცხავენ, სანამ შედეგები არ იქნება საჭირო. შესრულების ასეთი გადადება XLA-ს საშუალებას აძლევს, ოპტიმიზაცია მოახდინოს. მრავალი ცალკეული ოპერაციის გრაფი შეიძლება გაერთიანდეს ერთ ოპტიმიზებულ ოპერაციაში. ზარმაცი შესრულება ზოგადად უხილავია გამომძახებლისთვის. PyTorch / XLA ავტომატურად აგებს გრაფებს, აგზავნის მათ XLA მოწყობილობებზე და ახდენს სინქრონიზაციას XLA მოწყობილობასა და CPU-ს შორის მონაცემების კოპირებისას. ბარიერის ჩასმა ოპტიმიზატორის ნაბიჯის გადადგმისას აშკარად ახდენს CPU-სა და XLA მოწყობილობის სინქრონიზაციას. ეს ნიშნავს, რომ როდესაც თქვენ გამოიძახებთ model(input) forward pass-ს, გამოთვლით თქვენს დანაკარგს loss.backward()-ით და გადადგამთ ოპტიმიზაციის ნაბიჯს xm.optimizer_step(optimizer)-ით, ყველა ოპერაციის გრაფი ფონზე შენდება. მხოლოდ მაშინ, როდესაც თქვენ ან აშკარად შეაფასებთ ტენზორს (მაგ. ტენზორის დაბეჭდვა ან მისი CPU მოწყობილობაზე გადატანა) ან მონიშნავთ ნაბიჯს (რასაც MpDeviceLoader გააკეთებს ყოველ ჯერზე, როდესაც მასზე გაივლით), სრულდება მთელი ნაბიჯი. მომხმარებლის თვალსაზრისით, ტიპური ვარჯიშის რეჟიმი