शीर्षक: बग केवल एक axis का था

PyTorch के chunked-scan implementation में एक छिपा हुआ बग है जो कुछ hybrid मॉडल्स को भविष्य के tokens को देखने की अनुमति देता है जब fast fused kernels अनुपस्थित होते हैं, जिससे Zamba2-1.2B और Nemotron-H-8B जैसे checkpoints प्रभावित होते हैं। यह खामी गलत axis पर किए गए एक single reduction के कारण उत्पन्न होती है, और यह उस execution path पर सामने आती है जिस पर अधिकांश CI pipelines और CPU runs निर्भर करते हैं।

यह बग क्यों महत्वपूर्ण है

Hybrid और state-space मॉडल्स अब केवल attention पर निर्भर नहीं हैं; वे linear recurrences, scans, और convolutions का भी उपयोग करते हैं। Attention matrix में भविष्य के tokens को रोकने वाला masking operation उन अतिरिक्त operations के लिए causality की गारंटी स्वचालित रूप से नहीं देता है। जब वे fast fused kernels अनुपस्थित होते हैं जो सामान्य रूप से scan को सही ढंग से संभालते हैं, तो PyTorch एक pure-Python chunked-scan path पर वापस चला जाता है। उस fallback में axis-mixup होता है, जिससे inference के दौरान बाद की स्थितियों (later positions) से जानकारी पीछे की ओर प्रवाहित हो सकती है।

यह लीक शांत है। यह मॉडल को क्रैश नहीं करता या कोई स्पष्ट त्रुटि (error) नहीं दिखाता। इसके बजाय, यह loss और perplexity को कृत्रिम रूप से कम दिखाता है क्योंकि मॉडल प्रभावी रूप से धोखाधड़ी कर रहा है—उसी tokens को देख रहा है जिनका उसे अनुमान लगाना है। इसलिए, कोई भी downstream evaluation जो इन metrics पर भरोसा करता है, वह एक कमजोर आधार पर बना होता है।

समस्या का पता कैसे चला

शोधकर्ताओं ने एक ही मॉडल के माध्यम से दो forward passes की तुलना की:

  1. एक रैंडम token sequence.
  2. एक ही sequence जिसमें एक single token बदला गया हो।

उन्होंने layer-by-layer hidden states में अंतर को मापा। Mask inspection में कुछ भी संदिग्ध नहीं मिला, लेकिन faults इंजेक्ट करने वाले per-layer audit ने इंजेक्ट किए गए 192 में से 192 errors का पता लगाया, जिससे लीक की पुष्टि हुई।

transformers library के त्वरित सर्वेक्षण से पता चला:

  • Zamba2-1.2B तब लीक होता है जब chunk size 256 पर सेट होता है।
  • Nemotron-H-8B 128 के chunk size पर लीक होता है।
  • अधिकांश अन्य checkpoints ने समान परिस्थितियों में कोई लीकेज नहीं दिखाया।

यह दोष chunked-scan code path में है जो तब चलता है जब भी optional fused kernels अनुपस्थित होते हैं। इसमें शामिल हैं:

  • सभी CPU execution.
  • वे GPU environments जिनमें विशिष्ट fused-kernel packages नहीं हैं।
  • Standard PyTorch installations जिनमें अतिरिक्त dependencies नहीं हैं।

क्योंकि यह बग केवल तभी दिखाई देता है जब वे kernels अनुपस्थित होते हैं, यह CI environments और CPUs पर सामने आ सकता है।

जोखिम में कौन है और इसकी क्या कीमत है

कोई भी टीम जो fused kernels के बिना hybrid models को train, fine-tune, या evaluate करती है, उसे बढ़ा-चढ़ाकर दिखाए गए performance numbers प्रकाशित करने का जोखिम है। Loss या perplexity में दिखने वाला सुधार भ्रमपूर्ण है; मॉडल ने प्रभावी रूप से "आगे देख लिया" है। रिसर्च ग्रुप्स के लिए, यह state-of-the-art परिणामों के बारे में भ्रामक दावे कर सकता है। कमर्शियल deployments के लिए, यह generation tasks में ऐसी downstream errors का कारण बन सकता है जो वास्तव में कभी सीखी ही नहीं गई थीं।

प्रति-तर्क

कुछ डेवलपर्स का तर्क है कि भविष्य के token leakage को रोकने के लिए सही ढंग से लागू किया गया causal mask पर्याप्त है। यह बग उस धारणा को गलत साबित करता है: scans, convolutions, और कुछ normalization layers मास्क को पूरी तरह से बायपास कर सकते हैं। Chunked scan में axis-mixup यह दर्शाता है कि causality को हर उस जगह लागू किया जाना चाहिए जहाँ डेटा प्रवाहित होता है, न कि केवल attention matrix में।

आगे क्या ध्यान रखें

  • Dependency hygiene: सभी training और inference nodes पर fast fused-kernel packages इंस्टॉल करें, विशेष रूप से CI pipelines में।
  • Audit scripts: खोजकर्ताओं द्वारा अनुशंसित दो-चरणीय ऑडिट का पालन करें:
    • एक positive control के रूप में उस checkpoint में एक ज्ञात fault इंजेक्ट करें जिसका आप परीक्षण कर रहे हैं।
    • मॉडल के chunk या window size से लंबे sequences चलाएं; छोटे sequences कभी भी लीक को उजागर नहीं करेंगे।

Chunk size से अधिक sequence length वाले किसी भी hybrid या state-space checkpoint पर ऑडिट चलाने से पता चल जाएगा कि मॉडल अभी भी vulnerable है या नहीं।

निष्कर्ष (Takeaway): यहाँ तक कि एक मॉडल जो हर मानक परीक्षण (standard test) को पास कर लेता है, वह भी तब चुपचाप धोखाधड़ी कर सकता है जब execution path एक buggy implementation पर वापस चला जाता है। किसी भी hybrid model के metrics पर भरोसा करने से पहले यह सत्यापित करना कि fast fused kernels मौजूद हैं—और attention masks के परे लीकेज के लिए स्पष्ट रूप से परीक्षण करना—अब आवश्यक कदम हैं।