आपका मॉडल अब एक GPU पर नहीं समाता। training बांटने के आपके पास क्या विकल्प हैं?

अगर मॉडल समा जाता है और आपको सिर्फ throughput चाहिए, तो data parallelism इस्तेमाल कीजिए: मॉडल की नकल कीजिए, batch बांटिए, gradients का all reduce कीजिए। अगर वह समाता ही नहीं, तो उसे shard कीजिए। tensor parallelism अलग अलग layers को devices में बांटता है, pipeline parallelism अलग layers को अलग stages पर रखता है, और fully sharded data parallel parameters, gradients और optimizer state को ranks में बांट देता है। बड़े runs आम तौर पर इन्हें मिलाकर चलाते हैं।

पहला सवाल यह है कि मॉडल समाता नहीं है या बस धीमा है, क्योंकि ये अलग समस्याएं हैं। अगर समा जाता है तो data parallel आसान जीत है: हर GPU पर पूरी नकल रहती है, हर एक batch का एक हिस्सा लेता है, और step से पहले आप gradients का all reduce कर लेते हैं। यह तब तक अच्छा scale करता है जब तक all reduce आपके interconnect को भर न दे। अगर वह सचमुच समाता ही नहीं, तो मैं पहले मेमोरी की बनावट देखता हूं, क्योंकि Adam के साथ fp32 में optimizer state मोटे तौर पर parameter count का दोगुना होता है, यानी अक्सर कुल जगह में weights अल्पमत में होते हैं। यह सीधे sharded data parallel की तरफ इशारा करता है, FSDP या ZeRO stage three, जो parameters, gradients और optimizer state को ranks में बांटकर ऐन वक्त पर इकट्ठा करता है। आम तौर पर इतना काफी होता है और programming model सरल बना रहता है। उससे आगे आप tensor parallel पर जाते हैं, यानी एक layer के अंदर के matmuls को devices में बांटना, जिसके लिए तेज interconnect चाहिए क्योंकि वह हर layer पर communicate करता है, इसलिए मैं उसे एक node के अंदर ही रखता हूं। pipeline parallel अलग layers को अलग devices पर रखता है और micro batching से आप bubble overhead से लड़ते हैं। इन सबसे पहले मैं gradient checkpointing और mixed precision आजमाता हूं, क्योंकि कभी कभी अकेले उतने से ही आप सीमा के अंदर लौट आते हैं।