ამ ბლოგ-პოსტში განვიხილავთ, თუ როგორ უნდა დავხვეწოთ Llama 2 70B PyTorch FSDP-ის გამოყენებით და მასთან დაკავშირებულ საუკეთესო პრაქტიკებს. ამისთვის გამოვიყენებთ Hugging Face Transformers-ს, Accelerate-ს და TRL-ს. ასევე ვისწავლით, თუ როგორ გამოვიყენოთ Accelerate SLURM-თან ერთად. მონაცემთა სრულად შარდირებული პარალელიზმი (FSDP) არის პარადიგმა, რომლის დროსაც ოპტიმიზატორის მდგომარეობები, გრადიენტები და პარამეტრები ნაწილდება (შარდირდება) მოწყობილობებს შორის. „პირდაპირი გადაცემის“ (forward pass) დროს, FSDP-ის ყოველი ერთეული ასრულებს „ყველას შეკრების“ (all-gather) ოპერაციას სრული წონების მისაღებად, რის შემდეგაც ხორციელდება გამოთვლა და სხვა მოწყობილობებიდან მიღებული ნაწილები (shards) უგულებელყოფილია. „პირდაპირი გადაცემის“ შემდეგ გამოითვლება დანაკარგი (loss), რასაც მოსდევს „უკუ გადაცემა“ (backward pass). „უკუ გადაცემის“ დროს, FSDP-ის ყოველი ერთეული კვლავ ასრულებს „ყველას შეკრების“ ოპერაციას სრული წონების მისაღებად, გამოთვლა კი ხორციელდება ლოკალური გრადიენტების მისაღებად. ეს ლოკალური გრადიენტები საშუალოვდება და ნაწილდება მოწყობილობებს შორის „შემცირება-გაფანტვის“ (reduce-scatter) ოპერაციის მეშვეობით, რათა ყოველმა მოწყობილობამ შეძლოს თავისი ნაწილის პარამეტრების განახლება. მეტი ინფორმაციისთვის PyTorch FSDP-ის შესახებ, იხილეთ ბლოგ-პოსტი: „მსხვილი მოდელების ტრეინინგის დაჩქარება PyTorch Fully Sharded Data Parallel-ის გამოყენებით“. (წყარო: link) კვანძების რაოდენობა: 2. მინიმუმ 1. GPU-ების რაოდენობა კვანძზე: 8 GPU ტიპი: A100 GPU მეხსიერება: 80GB კვანძშიდა კავშირი: NVLink RAM კვანძზე: 1TB CPU ბირთვები კვანძზე: 96 კვანძთაშორისი კავშირი: Elastic Fabric Adapter LLaMa 70B-ის FSDP-ით დახვეწის მცდელობისას შეგვხვდა სამი ძირითადი გამოწვევა: FSDP მოდელს ახვევს (wraps) წინასწარ გაწვრთნილი მოდელის ჩატვირთვის შემდეგ. თუ კვანძში არსებული ყოველი პროცესი/რანგი Llama-70B მოდელს ჩატვირთავს, ამას დასჭირდება დაახლოებით 70*4*8 GB ~ 2TB CPU RAM, სადაც 4 არის ბაიტების რაოდენობა ყოველ პარამეტრზე, ხოლო 8 არის GPU-ების რაოდენობა თითოეულ კვანძზე. ეს გამოიწვევდა CPU RAM-ის ამოწურვას და პროცესების შეწყვეტას. შუალედური საკონტროლო წერტილების სრულად შენახვას FULL_STATE_DICT-ის გამოყენებით, CPU offloading-ით rank 0-ზე, დიდი დრო სჭირდება და ხშირად იწვევს NCCL Timeout შეცდომებს, რაც გამოწვეულია მაუწყებლობის დროს გაურკვეველი შეფერხებებით. თუმცა, ტრეინინგის დასრულებისას ჩვენ გვინდა მოდელის სრული მდგომარეობის ლექსიკონი (state dict) და არა შარდირებული, რომელიც მხოლოდ FSDP-თან არის თავსებადი. საჭიროა გავაუმჯობესოთ სიჩქარე და შევამციროთ VRAM-ის გამოყენება, რათა დავაჩქაროთ ტრეინინგი და დავზოგოთ გამოთვლითი ხარჯები. მოდით ვნახოთ, როგორ გადავჭრათ ზემოთ ხსენებული გამოწვევები და დავხვეწოთ 70B მოდელი! სანამ დავიწყებთ, წარმოგიდგენთ ყველა საჭირო რესურსს ჩვენი შედეგების რეპროდუცირებისთვის: Codebase: https://github.com/pacman100/DHS-LLM-Workshop/tree/main/chat_assistant/sft/training Flash Attention V2-ით FSDP კონფიგურაცია: https://github.com/pacman100/DHS-LLM-Workshop/blob/main/chat_assistant/training/configs/fsdp_config.yaml SLURM სკრიპტი launch.slurm: https://gist.github.com/pacman100/1cb1f17b2f1b3139a63b764263e70b25 მოდელი: meta-llama/Llama-2-70b-chat-hf მონაცემთა ნაკრები: smangrul/code-chat-assistant-v1 (LIMA+GUANACO-ს ნარევი სწორი ფორმატირებით, ტრეინინგისთვის მზა ფორმატში) პირველ რიგში, მიჰყევით ამ ნაბიჯებს Flash Attention V2-ის ინსტალაციისთვის: Dao-AILab/flash-attention: Fast and memory-efficient exact attention (github.com). დააინსტალირეთ PyTorch-ის უახლესი ღამის ვერსიები CUDA ≥11.8-ით. დააინსტალირეთ დარჩენილი მოთხოვნები DHS-LLM-Workshop/code_assistant/training/requirements.txt-ის შესაბამისად. აქ ჩვენ დავაინსტალირებთ 🤗 Accelerate-ს და 🤗 Transformers-ს მთავარი ბრენჩიდან. PR-ები huggingface/transformers#25107 და huggingface/accelerate#1777 პირველ გამოწვევას წყვეტენ და მომხმარებლის მხრიდან კოდის ცვლილებებს არ საჭიროებენ. ეს შემდეგს აკეთებს: ქვემოთ მოცემულია 7B მოდელის გამოსავლის ნაწილი 2 GPU-ზე, რომელიც ზომავს მოხმარებულ მეხსიერებას და მოდელის პარამეტრებს სხვადასხვა ეტაპზე. ვხედავთ, რომ წინასწარ გაწვრთნილი მოდელის ჩატვირთვისას rank 0-ს და rank 1-ს ჰქონდათ CPU-ს ჯამური პიკური მეხსიერება 32744 MB და 1506 MB, შესაბამისად. მაშასადამე, მხოლოდ rank 0 ტვირთავს წინასწარ გაწვრთნილ მოდელს, რაც იწვევს CPU RAM-ის ეფექტურ გამოყენებას. სრული ლოგები შეგიძლიათ იხილოთ აქ. ეს პრობლემა გადაწყდება SHARDED_STATE_DICT მდგომარეობის ლექსიკონის ტიპის არჩევით FSDP კონფიგურაციის შექმნისას. SHARDED_STATE_DICT ცალკე ინახავს ყოველ ნაწილს (shard) GPU-ზე, რაც აჩქარებს შენახვას ან ტრეინინგის განახლებას შუალედური საკონტროლო წერტილიდან. როდესაც გამოიყენება FULL_STATE_DICT, პირველი პროცესი (rank 0) აგროვებს მთელ მოდელს CPU-ზე და შემდეგ ინახავს მას სტანდარტულ ფორმატში. მოდით შევქმნათ Accelerate კონფიგურაცია ქვემოთ მოცემული ბრძანების გამოყენებით: მიღებული კონფიგურაცია ხელმისაწვდომია აქ: fsdp_config.yaml. აქ შარდირების სტრატეგია არის FULL_SHARD. ავტომატური შეხვევის (auto wrap) პოლიტიკისთვის ვიყენებთ TRANSFORMER_BASED_WRAP-ს, რომელიც _no_split_module-ს იყენებს Transformer ბლოკის სახელის მოსაძებნად ბუდობრივი FSDP ავტომატური შეხვევისთვის. შუალედური საკონტროლო წერტილებისა და ოპტიმიზატორის მდგომარეობების შესანახად ვიყენებთ SHARDED_STATE_DICT-ს, PyTorch გუნდის მიერ რეკომენდებული ფორმატით. დარწმუნდით, რომ ჩართეთ მოდულის პარამეტრების მაუწყებლობა rank 0-დან თავიდანვე, როგორც აღინიშნა ზემოთ მოცემულ პარაგრაფში გამოწვევა 1-ის მოგვარების შესახებ. ჩვენ ვრთავთ bf16 შერეული სიზუსტის ტრეინინგს. საბოლოო საკონტროლო წერტილის მთელი მოდელის მდგომარეობის ლექსიკონად მისაღებად, გამოიყენება ქვემოთ მოცემული კოდის ნაწილი: Flash Attention-ის და გრადიენტების შენახვის (gradient checkpointing) ჩართვა აუცილებელია ტრეინინგის დაჩქარებისთვის და VRAM-ის გამოყენების შემცირებისთვის, რათა შესაძლებელი გახდეს დახვეწა და გამოთვლითი ხარჯების დაზოგვა. ამჟამად, codebase იყენებს monkey patching-ს.