mixed precision training क्या है और उसमें क्या गड़बड़ हो सकती है?

mixed precision ज्यादातर operations 16 bit (fp16 या bf16) में चलाती है, जबकि weights की एक master copy और कुछ reductions fp32 में रहती हैं। इससे मेमोरी मोटे तौर पर आधी हो जाती है और tensor cores इस्तेमाल होते हैं, इसलिए वह काफी तेज है। fp16 का खतरा underflow है: छोटे gradients घटकर शून्य हो जाते हैं, और इसीलिए loss scaling मौजूद है। bf16 fp32 वाली exponent range रखता है, इसलिए वह इससे काफी हद तक बच जाता है।

weights की एक master copy आप fp32 में रखते हैं, लेकिन forward और backward pass 16 bit में चलाते हैं, इसलिए activations और gradients आधे आकार के होते हैं और matmuls tensor cores पर जाते हैं। व्यवहार में यह मोटे तौर पर आधी मेमोरी और अक्सर करीब दोगुना throughput है, जो आपके hardware पर run समाने और न समाने का फर्क बन सकता है। क्लासिक नाकामी fp16 underflow है। fp16 में सिर्फ पांच exponent bits हैं, इसलिए बहुत छोटे gradients शून्य हो जाते हैं, और छोटे gradients ठीक वही होते हैं जो ट्रेनिंग के आखिर में मायने रखते हैं। dynamic loss scaling यही ठीक करती है: backward से पहले loss को एक बड़े factor से गुणा कीजिए ताकि gradients दिखाने लायक range में उतरें, optimizer step से पहले unscale कीजिए, और infinities दिखते ही factor घटा दीजिए। bf16 आठ exponent bits रखता है, fp32 जितनी ही range, और इसकी कीमत वह कम mantissa bits से चुकाता है, इसलिए वह कहीं कम underflow करता है और आम तौर पर उसे loss scaling चाहिए ही नहीं। मेरा अब का डिफॉल्ट Ampere या उससे नए किसी भी hardware पर bf16 है। और मैं softmax, layer norm और खुद loss जैसे reductions fp32 में रखता हूं, क्योंकि 16 bit में हजारों terms जोड़ने पर accuracy तेजी से गिरती है।