1 पॉइंट द्वारा GN⁺ 2023-12-17 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • ML मॉडल जितने बड़े होते हैं, उनका execution efficiency सिर्फ hardware पर नहीं बल्कि compiler optimization पर भी काफी निर्भर करता है, और Google Research व Google DeepMind ने इस क्षेत्र को फिर से ML के जरिए बेहतर बनाने का तरीका पेश किया है
  • TensorFlow, JAX, PyTorch जैसे frameworks के नीचे लिए जाने वाले compiler decisions एक ही मॉडल के लिए भी execution time और resource usage को काफी बदल सकते हैं
  • TpuGraphs TPU के लिए ML programs के computational graphs, compiler settings और execution time को समेटने वाला dataset है, जिसे learning-based cost model research में इस्तेमाल किया जा सकता है
  • बड़े computational graph को एक साथ train करने की कठिनाई कम करने के लिए Graph Segment Training graph को segments में बांटता है, memory usage घटाता है और training time को 3 गुना कम करता है
  • Kaggle competition में 66 देशों की 616 teams और 792 लोगों ने हिस्सा लिया, और graph compression, padding value adjustment, node features जोड़ना, configurations के बीच attention जैसी वास्तविक cost prediction model improvement techniques को validate किया

ML compiler execution performance को क्यों प्रभावित करता है

  • आधुनिक ML मॉडल natural language understanding, conversation, image generation, video generation जैसे tasks करते हैं, और TensorFlow, JAX, PyTorch जैसे ML programming frameworks से लिखे और train किए जाते हैं
  • Frameworks matrix multiplication, convolution जैसे linear algebra operations और 2D convolution layers, transformer layers जैसी neural network layers प्रदान करते हैं
  • Users को hardware पर model को efficiently चलाने के detail तरीके सीधे संभालने की जरूरत नहीं होती; framework के नीचे मौजूद compiler model को अपने-आप optimize करता है
  • हालांकि compiler अक्सर complex optimization problems को heuristics से solve करता है, इसलिए हमेशा optimal performance देना मुश्किल हो सकता है

Computational graph और दो-स्तरीय optimization

  • ML compiler user द्वारा लिखे mathematical commands को ऐसे commands में बदलता है जिन्हें वास्तविक hardware पर चलाया जा सके
  • ML program को computational graph के रूप में व्यक्त किया जा सकता है
    • Nodes matrix multiplication जैसे tensor operations को दर्शाते हैं
    • Edges एक node से दूसरे node तक flow होने वाले tensors को दर्शाते हैं
  • Compiler optimization को broadly दो प्रकारों में बांटा जाता है
    • Graph-level optimization: पूरे graph के context को देखते हुए decisions लेता है और पूरे graph को transform करता है
    • Kernel-level optimization: fused subgraph यानी एक kernel को अन्य kernels से स्वतंत्र रूप से transform करता है

Memory layout के performance trade-offs

  • Matrix जैसे 2D tensors memory में [A B C a b c] या [A a B b C c] रूप में store किए जा सकते हैं, जो क्रमशः row-major और column-major layouts के अनुरूप हैं
  • ML compiler की एक महत्वपूर्ण optimization प्रक्रिया program के सभी intermediate tensors को memory layout assign करना है
  • कोई खास layout किसी individual operation के लिए सबसे efficient हो सकता है, लेकिन अगर add और convolution के बीच layout match नहीं करता, तो compiler को अतिरिक्त copy operation insert करना पड़ता है
  • इसके उलट, भले ही individual operation performance थोड़ी कम हो, layout conversion की जरूरत न पड़ने वाली configuration overall execution में बेहतर हो सकती है
  • XLA benchmark suite में default compiler settings के बजाय optimal layout configuration चुनने पर अधिकतम 32% speedup देखा गया

TpuGraphs dataset

  • TpuGraphs Google के custom TPU पर चलने वाले programs के लिए learning-based cost model dataset है
  • लक्ष्य ऐसा cost model train करना है जो input program और compiler settings लेकर program का execution time predict करे
  • Dataset दो XLA compiler settings को target करता है
    • layout: matrices के row-major·column-major concepts को higher-dimensional tensors तक generalize करने वाली setting
    • tiling: tile size setting
  • हर example में ML workload का computational graph, compiler settings, और उन settings से compile करने पर execution time शामिल होता है
  • Graphs open-source ML programs से collect किए गए हैं और इनमें ResNet, EfficientNet, Mask R-CNN, Transformer जैसे model architectures शामिल हैं
  • Download method और starter code TpuGraphs GitHub पर उपलब्ध हैं
  • TpuGraphs में समान graph size वाले मौजूदा सबसे बड़े graph property prediction dataset की तुलना में graph की संख्या 25 गुना ज्यादा है, और मौजूदा ML program performance prediction dataset की तुलना में average graph size 770 गुना बड़ा है

Baseline cost model और GNN structure

  • TpuGraphs baseline learning-based cost model भी प्रदान करता है, और input program graph के रूप में represented होने के कारण GNN का उपयोग करता है
  • Node features दो हिस्सों से बने होते हैं
    • opcode id: tensor operation type को दर्शाने वाली सबसे महत्वपूर्ण node information
    • अन्य node features
  • Baseline model opcode id को embedding lookup table के जरिए opcode embedding में बदलता है
  • Opcode embedding और बाकी node features को combine करके GNN input के रूप में इस्तेमाल किया जाता है
  • GNN द्वारा बनाए गए node embeddings को sum और mean जैसे simple graph pooling reduction से fixed-size graph embedding में combine किया जाता है
  • Final graph embedding feedforward layer से होकर एक scalar output में बदलता है

Graph Segment Training से बड़े graphs की training

  • Graph Segment Training सीमित memory capacity वाले devices पर बड़े graphs संभालने के लिए GNN training scaling technique है
  • यह तरीका उन graph-level prediction situations के लिए है जहां prediction target node या edge नहीं, बल्कि पूरा graph होता है
  • Computational graph में लाखों तक nodes हो सकते हैं, इसलिए पूरे graph को एक साथ इस्तेमाल करने वाला Full Graph Training computationally infeasible हो सकता है
  • GST बड़े graph को छोटे segments में बांटता है, और model update करने के लिए segments का केवल random subset चुनता है
  • बाकी segments intermediate activation store किए बिना embedding generate करते हैं, जिससे memory usage घटता है
  • सभी segment embeddings को combine करके मूल बड़े graph का embedding बनाया जाता है, और इसे prediction में इस्तेमाल किया जाता है
  • Historical embedding की staleness कम करने के लिए historical embedding table और segment dropout भी introduce किए गए हैं
  • पूरा method end-to-end training time को 3 गुना कम करता है

Kaggle competition में validate हुई improvement techniques

  • Fast or Slow? Predict AI Model Runtime Kaggle competition TpuGraphs dataset पर आधारित था, और इसमें 66 देशों की 616 teams व 792 लोगों ने हिस्सा लिया
  • कुल submissions 10,507 थीं; 153 लोगों के लिए यह पहला Kaggle competition था, और इनमें से 47 लोग top 100 में आए
  • Participating teams ने कई techniques experiment कीं
    • Graph pruning·compression: GST के बजाय बड़े graphs को compress करने के तरीके test किए गए, और configurable nodes व उनके immediate neighbours को शामिल करने वाले subgraph को ही रखने का तरीका इस्तेमाल हुआ
    • Padding value change: default padding value 0 valid feature values से conflict करती है, इसलिए -1 इस्तेमाल करने से model accuracy काफी बेहतर हो सकती थी
    • Node features जोड़ना·encoding बदलना: dot general के contracting dimensions जैसे additional node features महत्वपूर्ण हैं, और node feature encoding method भी results को प्रभावित कर सकता है
    • Cross-configuration attention: winning team ने एक simple layer design की जिससे model configurations की explicitly comparison कर सके, और इसने हर configuration को अलग-अलग infer कराने वाले तरीके से काफी बेहतर result दिया
  • Competition results और winning solutions 16 दिसंबर 2023 को NeurIPS के ML for Systems workshop competition session में कवर किए जाने वाले हैं

NeurIPS Expo related session

  • Structured data और artificial intelligence research में रुचि रखने वाले readers के लिए NeurIPS Expo panel Graph Learning Meets Artificial Intelligence 9 दिसंबर 2023 को आयोजित होगा
  • यह panel learning-based cost models की प्रगति आदि को cover करेगा

1 टिप्पणियां

 
GN⁺ 2023-12-17
Hacker News की राय
  • ML compiler को बढ़ा-चढ़ाकर पेश किया जा रहा है। यह पारंपरिक compiler जैसा ही trade-off है: किसी expert performance programmer को रखने की तुलना में आपको throughput बहुत ज्यादा मिलता है, लेकिन बाद वाला आम तौर पर काफी तेज होता है और कुछ मामलों में कई orders of magnitude आगे हो सकता है
    कई स्तरों पर कमी है। algorithm के स्तर पर, यह network को तेज बनाने की तरकीबें इंसान को वापस feedback नहीं करता, बस बहुत basic signals देता है। intent भी खो जाता है। ML network designer Python में structure specify करता है, लेकिन कई stages की lowering के बाद नतीजा पूरी तरह अलग हो सकता है। हाल में मैंने देखा कि एक compiler slice update करते समय array के सभी संभावित index ranges बना रहा था, फिर update करने वाले indices पाने के लिए उसे slice कर रहा था, और फिर scatter कर रहा था; मैंने उसे एक single memcpy call से replace कर दिया। kernels भी inefficient हैं। जब भी ऐसे compiler output का सामना किसी skilled assembly programmer से होता है, compiler हारता है, आम तौर पर 30% से ज्यादा अंतर से। यह आसानी से हल होने जैसा लगता है, लेकिन अगर पिछले 50 साल में किसी ने इसे ठीक से solve नहीं किया, तो जाहिर है कि यह सुनने में जितना सरल लगता है उतना नहीं है

    • chess engine Stockfish को देखें: board evaluation में इसने वर्षों से इंसानों द्वारा लिखे गए heuristics छोड़कर एक छोटे neural network से replace कर दिया, और वह बेहतर निकला
      compilers में भी inlining, loop unrolling, vectorization जैसे कई heuristics होते हैं, इसलिए neural networks मदद कर सकते हैं, और इंसानों द्वारा लिखे ढेर सारे heuristics की तुलना में maintain करना भी आसान हो सकता है
    • आपने कहा कि यह पारंपरिक compiler जैसा trade-off है, पर मजेदार बात यह है कि वे पारंपरिक compilers बेहद उपयोगी रहे हैं
    • यह बहुत ज्यादा definitive और बंद सोच वाला नजर आता है
    • सही। जब hand-written assembly लिखने के लिए किसी को hire किया जा सकता है, तो कोई gcc/clang क्यों इस्तेमाल करेगा?
    • वही throughput असली बात है। हर ML task के लिए performance expert नहीं लगाया जा सकता
      ऐसी optimization बिल्कुल न होने से यह अब भी कहीं बेहतर है
  • क्या इसे थोड़ा practical terms में समझा सकते हैं? मौजूदा ML compilers की वास्तविक स्थिति क्या है, और निकट भविष्य में हम क्या उम्मीद कर सकते हैं?

    • सबसे आसान approaches में से एक torch.compile है। यह PyTorch compiler का latest iteration है; पहले TorchScript और FX Tracing जैसे तरीके थे
      बस model = torch.compile(model) लिखना होता है। “इन 163 open source models में, torch.compile 93% cases में चला, और NVIDIA A100 GPU पर training 43% तेज हुई। Float32 precision में average 21%, और AMP precision में average 51% तेजी आई।”[1] लगता है Google इस तरह की R&D में ज्यादा लोगों को शामिल कराना चाहता है
      [1] https://pytorch.org/get-started/pytorch-2.0/
    • निकट भविष्य की उम्मीद यह है कि जिस framework में model develop हुआ है, उसके लिए explicit vendor support न होने पर भी AMD, CUDA, TPU, CPU आदि इस्तेमाल किए जा सकें
      वास्तविकता जटिल है, इसलिए काफी सरल करके कहें तो तरीका यह है कि computation graph को किसी intermediate representation में compile किया जाए और उसके लिए उपयुक्त backend implement किया जाए। संबंधित projects के लिए stableHLO, IREE, openXLA देखें। Jax का jit compiler भी ऐसे compiler का एक रूप माना जा सकता है। यह traced operations को XLA में lower करता है, और XLA फिर backend पर चलाने के लिए तरह-तरह का जादू करता है। आखिरकार नीचे जाते-जाते यह लगातार transformation और abstraction ही है
    • torch.compile देखें
  • सार यह है कि computation graph के execution-time performance prediction को graph neural network (GNN) से बेहतर बनाने का काम है। इसमें हर node के opcode के लिए embedding dictionary और shape, bits, window size जैसे दूसरे node features साथ इस्तेमाल होते हैं([1])
    उन्होंने अलग-अलग XLA compile settings और TPU पर मिली performance वाला बड़ा graph dataset [2] में release किया, और graph को partition करने के तरीके (METIS graph partition मैंने पहली बार देखा) तथा कई learning techniques से [3] में पहले से बड़े graphs के prediction को बेहतर किया। यह किसी दिए गए graph की performance predict करने की बात है, equivalent नया graph सुधारने, suggest करने या modify करने की नहीं। FunSearch की तरह, ठीक-ठाक predictive power वाला model evolutionary search के साथ इस्तेमाल किया जा सकता है
    [1] https://github.com/google-research-datasets/tpu_graphs#featu...
    [2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
    [3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322

  • क्या कोई समझा सकता है कि उस graph में convolution कैसे काम करता है? shape [2,4,16] वाले tensor पर shape [4,16,8] के kernel से convolution किया गया और [2,8] tensor निकला—यह कैसे संभव है?

    • पता नहीं मदद करेगा या नहीं, लेकिन input tensor [2,4,16] में 2 को batch size, 4 को input feature dimension, और 16 को input channel dimension माना जा सकता है
      kernel [4,16,8] में 4 filter window का size है, 16 input channel dimension से match करता है, और 8 output channel dimension है। output [2,8] में 2 batch size के रूप में बना रहता है, और 8 kernel के output channel dimension से match करता है। ऊपर से dimensions mismatch लगते हैं, लेकिन graph में convolution neighborhood structure का उपयोग करता है। kernel graph पर move करता है, current node और किसी खास radius के अंदर पड़ोसी features पर weights apply करता है, और उन weighted sums को जोड़कर हर output channel के लिए नए features बनाता है। graph structure और edge weights, padding और stride जैसी implementation details भी output shape को प्रभावित कर सकती हैं
  • Gemini किस स्थिति में है?

    • यह दिलचस्प है कि GPT-4 लगातार आगे है: https://twitter.com/lmsysorg/status/1735729398672716114
      अभी याद आने भर से ही Llama, Claude, Gemini, Falcon, Mistral जैसे कम से कम पांच foundation models हैं और वे एक-दूसरे से आगे-पीछे होते रहते हैं, लेकिन GPT अब भी एक level ऊपर है और एक साल से ऐसा ही है। पता चला है कि Transformer-based large language models इतने सरल हैं कि अगर कोई GPU time पर करीब एक million dollars खर्च कर सके तो लगभग कोई भी बना सकता है, लेकिन OpenAI को पूरी तरह catch up नहीं कर पा रहा। उनका खास secret क्या है?
  • Transformer खुद कैसा है? क्या किसी अर्थ में इसके optimal होने का कोई संकेत है?

  • लगता है पहले paragraph में core point दब गया है, लेकिन बाकी content शानदार है

  • अभी ML की progress speed हैरान करने वाली है। मैं singularity पर विश्वास नहीं करता, लेकिन यह software और समाज को ऐसे तरीकों से बदल रहा है जिनका कोई अनुमान नहीं लगा सकता

    • इसे और FunSearch को देखकर लगता है singularity करीब है
      https://deepmind.google/discover/blog/funsearch-making-new-d...
    • मुझे यह dot-com, mobile, cloud, VR के बाद आई एक और gold rush लगती है
    • 5 साल बाद लोग शायद आज की तरह programming नहीं करेंगे
    • मैं पहले ऐसी बीमारियों के इलाज देखना चाहूंगा जिनका इलाज मुश्किल है। singularity अपने-आप में तब तक मायने नहीं रखती जब तक उससे इंसानों को फायदा न हो, और वह फायदा मुख्यतः स्वास्थ्य सुधार और पीड़ा कम करने में होना चाहिए