Завантажені графічні процесори: метод вибірки та конвеєрної обробки прискорює глибоке навчання на великих графіках

Графи, потенційно розгалужена мережа вузлів, з’єднаних ребрами, можна використовувати для вираження та опитування зв’язків між даними, такими як соціальні зв’язки, фінансові операції, трафік, енергетичні мережі та молекулярні взаємодії. Оскільки дослідники збирають більше даних і будують ці графічні зображення, дослідникам знадобляться швидші та ефективніші методи, а також більша обчислювальна потужність, щоб проводити глибоке навчання на них за допомогою графових нейронних мереж (GNN).  

Тепер новий метод під назвою SALIENT (Sampling, SLIcing, and data movemeNT), розроблений дослідниками з Массачусетського технологічного інституту та IBM Research, покращує ефективність навчання та висновків, усуваючи три ключові вузькі місця в обчисленнях. Це значно скорочує час виконання GNN на великих наборах даних, які, наприклад, містять у масштабі 100 мільйонів вузлів і 1 мільярд ребер. Крім того, команда виявила, що ця техніка добре масштабується, коли додається обчислювальна потужність від одного до 16 графічних процесорів (GPU). Робота була представлена ​​на П'ятій конференції з машинного навчання та систем.

«Ми почали розглядати проблеми, з якими стикаються поточні системи, коли масштабують найсучасніші методи машинного навчання для графіків до дійсно великих наборів даних. Виявилося, що потрібно зробити багато роботи, тому що багато існуючих систем досягали високої продуктивності в основному на менших наборах даних, які вміщуються в пам’ять графічного процесора», – каже Тім Калер, провідний автор і постдокторант MIT Computer Science. та Лабораторія штучного інтелекту (CSAIL).

Під величезними наборами даних експерти мають на увазі такі масштаби, як уся мережа біткойн, де певні закономірності та зв’язки даних можуть вказувати на тенденції чи нечесну гру. «Існує майже мільярд транзакцій біткойн у блокчейні, і якщо ми хочемо виявити незаконну діяльність у такій об’єднаній мережі, то перед нами буде графік такого масштабу», — каже співавтор Джі Чен, старший науковий співробітник і менеджер. IBM Research і MIT-IBM Watson AI Lab. «Ми хочемо побудувати систему, яка здатна обробляти такий тип графіків і дозволити обробці бути максимально ефективним, тому що кожного дня ми хочемо не відставати від нових даних, які генеруються».

Співавтори Калера та Чена включають Ніколаса Статаса Менга '21 з Jump Trading, який розробив SALIENT як частину своєї дипломної роботи; колишня стажерка MIT-IBM Watson AI Lab та аспірантка MIT Енн Уянг; MIT CSAIL postdoc Александрос-Ставрос Іліопулос; Дослідник MIT CSAIL Тао Б. Шардл; і Чарльз Е. Лейзерсон, професор електротехніки Едвіна Сіблі Вебстера в Массачусетському технологічному інституті та дослідник лабораторії MIT-IBM Watson AI Lab.     

Для вирішення цієї проблеми команда застосувала системно-орієнтований підхід у розробці свого методу: SALIENT, каже Калер. Для цього дослідники запровадили важливі, на їхню думку, базові оптимізації компонентів, які вписуються в існуючі структури машинного навчання, такі як PyTorch Geometric і бібліотека глибоких графів (DGL), які є інтерфейсами для побудови моделі машинного навчання. Статас каже, що цей процес схожий на заміну двигунів на створення швидшого автомобіля. Їхній метод був розроблений таким чином, щоб відповідати існуючій архітектурі GNN, щоб експерти з домену могли легко застосувати цю роботу до своїх визначених полів, щоб пришвидшити навчання моделям і швидше отримати інформацію під час логічного висновку. Команда визначила, що хитрість полягала в тому, щоб усе обладнання (ЦП, канали передачі даних і графічні процесори) було зайнято в будь-який час: поки ЦП пробує графік і готує міні-пакети даних, які потім будуть передані через канал передачі даних , тим критичніший GPU працює над навчанням моделі машинного навчання або проведенням висновків. 

Дослідники почали з аналізу продуктивності широко використовуваної бібліотеки машинного навчання для GNN (PyTorch Geometric), яка показала вражаюче низьке використання доступних ресурсів GPU. Застосувавши просту оптимізацію, дослідники покращили використання графічного процесора з 10 до 30 відсотків, що призвело до підвищення продуктивності в 1.4-50.4 рази порівняно з загальнодоступними тестовими кодами. Цей швидкий базовий код може виконати один повний прохід через великий навчальний набір даних через алгоритм (епоху) за XNUMX секунди.                          

Прагнучи подальшого покращення продуктивності, дослідники вирішили вивчити вузькі місця, які виникають на початку конвеєра даних: алгоритми для вибірки графів і підготовки міні-пакетів. На відміну від інших нейронних мереж, GNN виконує операцію агрегування сусідства, яка обчислює інформацію про вузол, використовуючи інформацію, наявну в інших сусідніх вузлах на графі — наприклад, у графі соціальної мережі, інформацію від друзів друзів користувача. Зі збільшенням кількості рівнів у GNN кількість вузлів, до яких мережа має охопити інформацію, може вибухнути, перевищивши межі комп’ютера. Алгоритми вибірки сусідства допомагають шляхом вибору меншої випадкової підмножини вузлів для збору; однак дослідники виявили, що поточні реалізації цього надто повільні, щоб не відставати від швидкості обробки сучасних графічних процесорів. У відповідь вони визначили поєднання структур даних, алгоритмічної оптимізації тощо, що покращило швидкість вибірки, зрештою покращивши саму операцію вибірки приблизно втричі, піднявши час виконання за епоху з 50.4 до 34.6 секунди. Вони також виявили, що вибірка з відповідною частотою може бути зроблена під час висновку, покращуючи загальну енергоефективність і продуктивність, що було пропущено в літературі, зазначає команда.      

У попередніх системах цей етап вибірки був багатопроцесним підходом, що створював додаткові дані та непотрібне переміщення даних між процесами. Дослідники зробили свій метод SALIENT більш гнучким, створивши єдиний процес із легкими потоками, які зберігали дані на ЦП у спільній пам’яті. Крім того, SALIENT використовує переваги кеш-пам’яті сучасних процесорів, каже Статас, розпаралелюючи нарізку функцій, яка витягує відповідну інформацію з вузлів, що цікавлять, і їх оточуючих сусідів і країв у спільній пам’яті кешу ядра ЦП. Це знову зменшило загальний час роботи за епоху з 34.6 до 27.8 секунд.

Останнім вузьким місцем, на яке звернулися дослідники, була конвеєрна передача міні-пакетів даних між процесором і графічним процесором за допомогою етапу попередньої вибірки, який готував би дані безпосередньо перед тим, як вони знадобляться. Команда підрахувала, що це максимізує використання пропускної здатності в каналі передачі даних і доведе метод до ідеального використання; однак вони бачили лише близько 90 відсотків. Вони виявили та виправили помилку продуктивності в популярній бібліотеці PyTorch, яка спричиняла непотрібний обмін даними між процесором і графічним процесором. Після виправлення цієї помилки команда досягла 16.5 секунди на епоху роботи з SALIENT.

«Я вважаю, що наша робота показала, що диявол криється в деталях», — каже Калер. «Якщо ви приділяєте пильну увагу деталям, які впливають на продуктивність під час навчання графової нейронної мережі, ви можете вирішити величезну кількість проблем продуктивності. З нашими рішеннями ми виявилися повністю обмеженими обчисленнями GPU, що є ідеальною метою такої системи».

Швидкість SALIENT оцінювалася за трьома стандартними наборами даних ogbn-arxiv, ogbn-products і ogbn-papers100M, а також у налаштуваннях на кількох машинах з різними рівнями fanout (кількість даних, які центральний процесор підготує для графічного процесора) і на кількох архітектурах, включаючи найновішу, GraphSAGE-RI. У кожному налаштуванні SALIENT перевершував PyTorch Geometric, особливо на великому наборі даних ogbn-papers100M, який містив 100 мільйонів вузлів і понад мільярд граней. Тут він був утричі швидшим, працюючи на одному графічному процесорі, ніж оптимізована базова лінія, яка спочатку була створена для ця робота; з 16 графічним процесором SALIENT був у вісім разів швидшим. 

У той час як інші системи мали дещо інше апаратне забезпечення та експериментальні налаштування, тому це не завжди було прямим порівнянням, SALIENT все одно перевершував їх. Серед систем, які досягли подібної точності, репрезентативні показники продуктивності включають 99 секунд при використанні одного GPU і 32 ЦП і 13 секунд при використанні 1,536 ЦП. Навпаки, час роботи SALIENT з використанням одного графічного процесора та 20 процесорів становив 16.5 секунди та лише дві секунди з використанням 16 графічних процесорів та 320 процесорів. «Якщо ви подивитеся на підсумкові цифри, які звітує попередня робота, наш час роботи з 16 графічним процесором (дві секунди) на порядок швидший, ніж інші цифри, які повідомлялися раніше на цьому наборі даних», — каже Калер. Дослідники пояснюють підвищення продуктивності частково своїм підходом до оптимізації свого коду для однієї машини перед переходом до розподіленого налаштування. Статас каже, що урок полягає в тому, що за ваші гроші «доцільніше використовувати апаратне забезпечення, яке у вас є, максимально ефективно, перш ніж почати масштабування до кількох комп’ютерів», що може забезпечити значну економію витрат і викидів вуглецю. які можна отримати разом із навчанням моделі.

Ця нова здатність тепер дозволить дослідникам глибше вивчати все більші графіки. Наприклад, згадана раніше мережа Bitcoin містила 100,000 1,000 вузлів; система SALIENT здатна обробляти графік у XNUMX разів (або на три порядки) більший.

«У майбутньому ми плануємо не просто запустити систему навчання графової нейронної мережі на існуючих алгоритмах, які ми реалізували для класифікації або прогнозування властивостей кожного вузла, але ми також хочемо виконувати більш глибокі завдання, такі як виявлення загальних шаблонів у графі (шаблони підграфів), [які] можуть бути насправді цікавими для вказівки на фінансові злочини», — каже Чен. «Ми також хочемо ідентифікувати вузли на графіку, які схожі в тому сенсі, що вони, можливо, відповідатимуть тому самому поганому учаснику фінансового злочину. Ці завдання вимагатимуть розробки додаткових алгоритмів і, можливо, також архітектури нейронних мереж».

Це дослідження було підтримано MIT-IBM Watson AI Lab і частково Дослідницькою лабораторією ВПС США та Прискорювачем штучного інтелекту ВПС США.