Οι ατομιστικές προσομοιώσεις αποτελούν κρίσιμο εργαλείο για την κατανόηση της συμπεριφοράς μορίων και υλικών, από χημικές αντιδράσεις μέχρι μεταβάσεις φάσης και μηχανική αστοχία. Το βασικό πρόβλημα σε αυτό το πεδίο είναι η γνωστή ισορροπία ανάμεσα στην ακρίβεια και το υπολογιστικό κόστος. Οι κβαντομηχανικές μέθοδοι, όπως η density functional theory, προσφέρουν υψηλή πιστότητα, αλλά συχνά είναι υπερβολικά δαπανηρές για μεγάλα συστήματα ή για προσομοιώσεις μεγάλης διάρκειας. Στην άλλη πλευρά, τα κλασικά interatomic potentials εκτελούνται πολύ γρηγορότερα, αλλά δυσκολεύονται να γενικεύσουν σωστά σε διαφορετικά ατομικά περιβάλλοντα.
Σε αυτό το κενό τοποθετούνται τα machine learning interatomic potentials, γνωστά ως MLIPs. Αυτά τα μοντέλα χρησιμοποιούν νευρωνικά δίκτυα για να προβλέπουν ενέργειες και διατομικές δυνάμεις με πολύ χαμηλότερο κόστος από τις αμιγώς κβαντομηχανικές προσεγγίσεις. Έτσι, επιτρέπουν μεγαλύτερα συστήματα και περισσότερα χρονικά βήματα, διατηρώντας σημαντικό μέρος της απαιτούμενης ακρίβειας. Το ερώτημα, όμως, είναι αν μπορούν να αξιοποιήσουν καλύτερα το σύγχρονο hardware χωρίς να θυσιάσουν τη φυσική αξιοπιστία των αποτελεσμάτων.
Γιατί η ακρίβεια παραμένει κρίσιμη στα scientific workloads
Στις προσομοιώσεις μοριακής δυναμικής, οι μικρές αριθμητικές αποκλίσεις δεν μένουν τοπικές. Οι δυνάμεις υπολογίζονται ως παράγωγοι της προβλεπόμενης ενέργειας και χρησιμοποιούνται επαναληπτικά για την ενημέρωση ταχυτήτων και θέσεων των ατόμων. Αυτό σημαίνει ότι ακόμη και μικρά σφάλματα μπορούν να συσσωρευτούν σε πολλά χρονικά βήματα και να αλλοιώσουν τη συνολική τροχιά του συστήματος. Γι’ αυτόν τον λόγο, πολλά MLIPs συνεχίζουν να εκτελούν inference σε FP32.
Η επιλογή αυτή, όμως, έχει κόστος σε απόδοση. Οι σύγχρονες Nvidia GPUs προσφέρουν σημαντικά υψηλότερο throughput για πράξεις χαμηλότερης ακρίβειας μέσω Tensor Cores, σε σύγκριση με τους κλασικούς FP32 πολλαπλασιασμούς πινάκων. Η απόσταση ανάμεσα σε αυτό που απαιτεί το scientific computing και σε αυτό που αποδίδει καλύτερα το hardware είναι ακριβώς το σημείο που επιχειρούν να καλύψουν οι μέθοδοι Tensor Core FP32. Η βασική ιδέα είναι να διατηρηθεί συμπεριφορά κοντά στο FP32, αξιοποιώντας παράλληλα την ταχύτητα των Tensor Cores.
Η προσέγγιση της Nvidia για FP32 μέσω Tensor Cores
Αντί να αντικαταστήσει πλήρως το FP32 με χαμηλότερη ακρίβεια, η Nvidia περιγράφει μια ενδιάμεση προσέγγιση. Οι μέθοδοι Tensor Core FP32 διασπούν κάθε FP32 τελεστέο σε μικρότερα στοιχεία χαμηλότερης ακρίβειας, υπολογίζουν επιλεγμένα γινόμενα στους Tensor Cores και επανασυνθέτουν το αποτέλεσμα με συσσώρευση σε FP32. Με αυτόν τον τρόπο επιδιώκουν να περιορίσουν το αριθμητικό σφάλμα, χωρίς αλλαγές στα βάρη ή στην αρχιτεκτονική του μοντέλου. Σύμφωνα με την ανάλυση, αυτό μπορεί να αυξήσει αισθητά το throughput των πολλαπλασιασμών πινάκων.
Η αξιολόγηση έγινε πάνω στο Point Edge Transformer, ή PET, ένα transformer-based MLIP που αναπτύχθηκε από το Lab COSMO στο EPFL. Η επιλογή μόνο τυχαία δεν είναι, καθώς το PET αφιερώνει μεγάλο μέρος του χρόνου του σε ήδη βελτιστοποιημένα GEMMs. Με άλλα λόγια, πρόκειται για απαιτητικό τεστ, επειδή υπάρχουν περιορισμένα περιθώρια για πρόσθετη επιτάχυνση. Αν οι μέθοδοι Tensor Core FP32 προσφέρουν κέρδος εδώ, αυτό έχει ιδιαίτερη σημασία για αντίστοιχα ευαίσθητα scientific workloads.
BF16x9 σε Blackwell, επιτάχυνση χωρίς αλλαγές στον κώδικα
Στις Blackwell GPUs, η Nvidia αξιολόγησε τη μέθοδο BF16x9 μέσω του cuBLAS. Σε αυτή την προσέγγιση, κάθε FP32 τελεστέος διασπάται σε τρία BF16 components και προκύπτουν εννέα επιμέρους γινόμενα, τα οποία επανασυνδυάζονται με FP32 accumulation. Το πρακτικό πλεονέκτημα είναι ότι αυτή η διαδρομή μπορεί να ενεργοποιηθεί μέσω μεταβλητής περιβάλλοντος ή των αντίστοιχων ρυθμίσεων math mode και compute type. Αυτό σημαίνει ότι τα συμβατά PyTorch GEMMs μπορούν να τη χρησιμοποιήσουν χωρίς αλλαγές στην αρχιτεκτονική του μοντέλου ή προσαρμοσμένες υλοποιήσεις GEMM.
Στα αποτελέσματα, τα μεγάλα GEMMs του PET επιταχύνθηκαν κατά 2,2 έως 2,5 φορές για workload 17.280 ατόμων και κατά 1,2 έως 1,4 φορές για workload 192 ατόμων. Σε επίπεδο MD step, η performant στρατηγική έδωσε επιτάχυνση 1,2x στο μικρό σύστημα και 1,7x στο μεγάλο, ενώ η eager στρατηγική έφτασε το 1,1x και το 1,4x αντίστοιχα. Ιδιαίτερο βάρος έχει το συμπέρασμα ότι το σφάλμα που εισάγει η μέθοδος BF16x9 βρίσκεται κάτω από το ίδιο το επίπεδο μη ντετερμινισμού του μοντέλου. Αυτό δείχνει ότι υπάρχει πραγματικό περιθώριο αξιοποίησης των Tensor Cores ακόμη και σε ένα ήδη βελτιστοποιημένο μοντέλο.
TF32x3 σε Ampere, προσαρμοσμένη υλοποίηση με CUTLASS
Για τις Ampere GPUs, η προσέγγιση ήταν διαφορετική και πιο παρεμβατική. Η Nvidia υλοποίησε το TF32x3 αντικαθιστώντας τα τυπικά nn.Linear layers του PET με ένα προσαρμοσμένο TF32Linear layer που χρησιμοποιεί CUTLASS C++ kernel. Το TF32x3 αναπαριστά κάθε FP32 τελεστέο με ένα κύριο TF32 component και ένα υπολειμματικό component. Στη συνέχεια υπολογίζει τρία γινόμενα στους Tensor Cores και παραλείπει τον μικρότερο residual-by-residual όρο, ενώ οι είσοδοι και οι έξοδοι παραμένουν σε FP32.
Στον Nvidia A100, τα βασικά σχήματα linear layer του PET εμφάνισαν επιτάχυνση περίπου 1,7 έως 1,9 φορές στα GEMMs. Για το μικρό checkpoint των 26 εκατομμυρίων παραμέτρων, το TF32x3 διατήρησε το σφάλμα στις δυνάμεις περίπου τρεις φορές πάνω από τη διαφορά ανάμεσα σε FP32 runs μετά από recompilation και έδωσε επιτάχυνση 1,5x στο μεγάλο σύστημα. Για το μεγαλύτερο checkpoint των 255 εκατομμυρίων παραμέτρων, στο σύστημα 192 ατόμων, πέτυχε speedup 1,4x, έναντι 1,2x στο μικρότερο μοντέλο. Το one-pass TF32 ήταν ταχύτερο, αλλά με αισθητά υψηλότερο κόστος σε ακρίβεια και σταθερότητα τροχιάς.
Έλεγχος φυσικής συνέπειας πέρα από τα απλά benchmarks
Η Nvidia δεν περιορίστηκε σε μετρήσεις τύπου throughput ή single-point error. Σε τέτοιες εφαρμογές, η πιστότητα κρίνεται τελικά από το αν διατηρείται ομαλή η επιφάνεια δυναμικής ενέργειας και αν οι προσομοιώσεις παραμένουν σταθερές στον χρόνο. Για να ελέγξει αυτά τα σημεία, η ομάδα εξέτασε την ομαλότητα της δυναμικής ενέργειας και τη διατήρηση της ενέργειας σε ξεχωριστά τεστ. Η λογική είναι σαφής, επειδή οι μικρές αριθμητικές ανωμαλίες μπορεί να μη φανούν αμέσως σε μια μεμονωμένη πρόβλεψη, αλλά να εκδηλωθούν σε μια μακρά προσομοίωση.
Στο τεστ ομαλότητας, ένα μόριο CO μετακινήθηκε άκαμπτα μακριά από μια επιφάνεια Pt(111), με αξιολόγηση της ολικής ενέργειας σε 60 θέσεις. Το TF32x3 ακολούθησε την ομαλή καμπύλη του FP32, ενώ το one-pass TF32 εμφάνισε ορατές τοπικές διακυμάνσεις. Σε ξεχωριστή προσομοίωση NVE διάρκειας 30 ps για κουτί νερού 192 ατόμων, το FP32 και το TF32x3 παρουσίασαν συγκρίσιμες, περιορισμένες διακυμάνσεις ενέργειας. Αντίθετα, το one-pass TF32 έδειξε μεγαλύτερη συστηματική μετατόπιση, κάτι που το καθιστά λιγότερο κατάλληλο για τέτοιες ευαίσθητες χρήσεις.
Τι σημαίνουν τα αποτελέσματα για το scientific machine learning
Το βασικό συμπέρασμα είναι ότι οι μέθοδοι Tensor Core FP32 μπορούν να μειώσουν το χάσμα ανάμεσα στις ανάγκες ακρίβειας των επιστημονικών μοντέλων και στις δυνατότητες απόδοσης του σύγχρονου hardware. Στην περίπτωση του PET, αυτό επιτεύχθηκε χωρίς retraining, χωρίς αλλαγή των βαρών και, στην περίπτωση του BF16x9, χωρίς τροποποιήσεις στον κώδικα του μοντέλου. Για τους ερευνητές που εκτελούν MLIPs, αυτή είναι μια πρακτική προοπτική βελτίωσης της απόδοσης σε workloads όπου η πλήρης μετάβαση σε χαμηλότερη ακρίβεια θεωρείται ριψοκίνδυνη. Παράλληλα, η μελέτη δείχνει ότι η αξιολόγηση δεν μπορεί να περιορίζεται σε απλές μετρήσεις speedup.
Εξίσου σημαντικό είναι ότι η Nvidia τονίζει την ανάγκη για end-to-end δοκιμές σε κάθε workload. Ο συνδυασμός αυτών των μεθόδων με βελτιστοποιημένη graph construction, fused operations, custom kernels ή CUDA graphs δεν μετρήθηκε συνολικά στην παρούσα αξιολόγηση. Άρα, τα πραγματικά κέρδη θα πρέπει να επιβεβαιώνονται ξεχωριστά ανά μοντέλο και εφαρμογή. Για τον χώρο του scientific machine learning, αυτή η εργασία λειτουργεί λιγότερο ως γενική υπόσχεση και περισσότερο ως σαφές πλαίσιο αξιολόγησης για workloads όπου η αριθμητική ακρίβεια παραμένει αδιαπραγμάτευτη.
