GPU bận rộn: Phương pháp lấy mẫu và đường ống tăng tốc độ học sâu trên các biểu đồ lớn

Đồ thị, một mạng lưới các nút có khả năng mở rộng được kết nối bởi các cạnh, có thể được sử dụng để biểu thị và thẩm vấn các mối quan hệ giữa dữ liệu, như kết nối xã hội, giao dịch tài chính, lưu lượng truy cập, lưới năng lượng và tương tác phân tử. Khi các nhà nghiên cứu thu thập thêm dữ liệu và xây dựng những bức tranh đồ họa này, các nhà nghiên cứu sẽ cần các phương pháp nhanh hơn và hiệu quả hơn, cũng như sức mạnh tính toán cao hơn, để tiến hành học sâu trên chúng, theo cách của mạng thần kinh đồ thị (GNN).  

Giờ đây, một phương pháp mới, được gọi là SALIENT (SAmpling, sLIcing, and data movemeNT), được phát triển bởi các nhà nghiên cứu tại MIT và IBM Research, cải thiện hiệu suất đào tạo và suy luận bằng cách giải quyết ba tắc nghẽn chính trong tính toán. Điều này cắt giảm đáng kể thời gian chạy của GNN trên các bộ dữ liệu lớn, ví dụ, chứa trên quy mô 100 triệu nút và 1 tỷ cạnh. Hơn nữa, nhóm nhận thấy rằng kỹ thuật này có quy mô tốt khi sức mạnh tính toán được thêm từ một đến 16 đơn vị xử lý đồ họa (GPU). Công trình đã được trình bày tại Hội nghị lần thứ năm về Máy học và Hệ thống.

“Chúng tôi bắt đầu xem xét những thách thức mà các hệ thống hiện tại gặp phải khi nhân rộng các kỹ thuật máy học tiên tiến nhất dành cho đồ thị thành các tập dữ liệu thực sự lớn. Hóa ra còn rất nhiều việc phải làm, bởi vì rất nhiều hệ thống hiện có đang đạt được hiệu suất tốt chủ yếu dựa trên các bộ dữ liệu nhỏ hơn phù hợp với bộ nhớ GPU,” Tim Kaler, tác giả chính và một postdoc tại MIT Computer Science cho biết. và Phòng thí nghiệm trí tuệ nhân tạo (CSAIL).

Theo các bộ dữ liệu khổng lồ, các chuyên gia có nghĩa là các quy mô giống như toàn bộ mạng Bitcoin, nơi các mẫu và mối quan hệ dữ liệu nhất định có thể đánh vần các xu hướng hoặc hành vi chơi xấu. Đồng tác giả Jie Chen, nhà khoa học nghiên cứu cấp cao và quản lý cho biết: “Có gần một tỷ giao dịch Bitcoin trên chuỗi khối và nếu chúng tôi muốn xác định các hoạt động bất hợp pháp bên trong một mạng chung như vậy, thì chúng tôi đang phải đối mặt với một biểu đồ có quy mô như vậy”. của IBM Research và MIT-IBM Watson AI Lab. “Chúng tôi muốn xây dựng một hệ thống có thể xử lý loại biểu đồ đó và cho phép xử lý hiệu quả nhất có thể, bởi vì mỗi ngày chúng tôi muốn bắt kịp tốc độ của dữ liệu mới được tạo ra.”

Các đồng tác giả của Kaler và Chen bao gồm Nickolas Stathas MEng '21 của Jump Trading, người đã phát triển SALIENT như một phần của công việc tốt nghiệp của mình; cựu thực tập sinh MIT-IBM Watson AI Lab và sinh viên tốt nghiệp MIT Anne Ouyang; Tiến sĩ hậu tiến sĩ MIT CSAIL Alexandros-Stavros Iliopoulos; Nhà khoa học nghiên cứu MIT CSAIL Tao B. Schardl; và Charles E. Leiserson, Giáo sư Kỹ thuật Điện Edwin Sibley Webster tại MIT và là nhà nghiên cứu của Phòng thí nghiệm trí tuệ nhân tạo MIT-IBM Watson.     

Đối với vấn đề này, nhóm đã áp dụng cách tiếp cận hướng hệ thống trong việc phát triển phương pháp của họ: SALIENT, Kaler nói. Để làm được điều này, các nhà nghiên cứu đã triển khai những gì họ cho là tối ưu hóa cơ bản, quan trọng đối với các thành phần phù hợp với các khung máy học hiện có, chẳng hạn như PyTorch Geometric và thư viện đồ thị sâu (DGL), là các giao diện để xây dựng mô hình máy học. Stathas cho biết quá trình này giống như hoán đổi động cơ để chế tạo một chiếc ô tô nhanh hơn. Phương pháp của họ được thiết kế để phù hợp với kiến ​​trúc GNN hiện có, để các chuyên gia miền có thể dễ dàng áp dụng công việc này cho các trường được chỉ định của họ để đẩy nhanh quá trình đào tạo mô hình và đưa ra những hiểu biết sâu sắc trong quá trình suy luận nhanh hơn. Nhóm đã xác định được mẹo là giữ cho tất cả phần cứng (CPU, liên kết dữ liệu và GPU) luôn bận rộn: trong khi CPU lấy mẫu biểu đồ và chuẩn bị các lô dữ liệu nhỏ sau đó sẽ được truyền qua liên kết dữ liệu , GPU quan trọng hơn đang hoạt động để đào tạo mô hình máy học hoặc tiến hành suy luận. 

Các nhà nghiên cứu đã bắt đầu bằng cách phân tích hiệu suất của một thư viện máy học thường được sử dụng cho GNN (PyTorch Geometric), cho thấy mức sử dụng tài nguyên GPU có sẵn thấp một cách đáng kinh ngạc. Áp dụng các tối ưu hóa đơn giản, các nhà nghiên cứu đã cải thiện mức sử dụng GPU từ 10 lên 30 phần trăm, dẫn đến cải thiện hiệu suất từ ​​1.4 đến hai lần so với các mã điểm chuẩn công khai. Mã cơ sở nhanh này có thể thực hiện một lượt hoàn chỉnh trên tập dữ liệu huấn luyện lớn thông qua thuật toán (một kỷ nguyên) trong 50.4 giây.                          

Để tìm kiếm những cải tiến hiệu suất hơn nữa, các nhà nghiên cứu đã bắt đầu kiểm tra các nút thắt cổ chai xảy ra ở phần đầu của đường dẫn dữ liệu: các thuật toán để lấy mẫu biểu đồ và chuẩn bị lô nhỏ. Không giống như các mạng thần kinh khác, GNN thực hiện thao tác tổng hợp vùng lân cận, tính toán thông tin về một nút bằng cách sử dụng thông tin có trong các nút lân cận khác trong biểu đồ — ví dụ: trong biểu đồ mạng xã hội, thông tin từ bạn của bạn của người dùng. Khi số lớp trong GNN tăng lên, số lượng nút mà mạng phải tiếp cận để lấy thông tin có thể bùng nổ, vượt quá giới hạn của máy tính. Các thuật toán lấy mẫu vùng lân cận trợ giúp bằng cách chọn một tập hợp con ngẫu nhiên nhỏ hơn của các nút để thu thập; tuy nhiên, các nhà nghiên cứu nhận thấy rằng việc triển khai tính năng này hiện tại quá chậm để theo kịp tốc độ xử lý của các GPU hiện đại. Để đáp lại, họ đã xác định được sự kết hợp của các cấu trúc dữ liệu, tối ưu hóa thuật toán, v.v. để cải thiện tốc độ lấy mẫu, cuối cùng là cải thiện riêng hoạt động lấy mẫu khoảng ba lần, kéo dài thời gian chạy mỗi kỷ nguyên từ 50.4 lên 34.6 giây. Họ cũng phát hiện ra rằng việc lấy mẫu, với tốc độ phù hợp, có thể được thực hiện trong quá trình suy luận, cải thiện hiệu suất và hiệu suất năng lượng tổng thể, một điểm đã bị bỏ qua trong tài liệu, nhóm nghiên cứu lưu ý.      

Trong các hệ thống trước đây, bước lấy mẫu này là một cách tiếp cận đa quy trình, tạo ra dữ liệu bổ sung và chuyển động dữ liệu không cần thiết giữa các quy trình. Các nhà nghiên cứu đã làm cho phương pháp SALIENT của họ linh hoạt hơn bằng cách tạo ra một quy trình duy nhất với các luồng nhẹ giữ dữ liệu trên CPU trong bộ nhớ dùng chung. Hơn nữa, SALIENT tận dụng bộ đệm của các bộ xử lý hiện đại, Stathas cho biết, song song hóa tính năng cắt lát, giúp trích xuất thông tin liên quan từ các nút quan tâm cũng như các nút và cạnh xung quanh của chúng, trong bộ nhớ dùng chung của bộ đệm lõi CPU. Điều này một lần nữa làm giảm tổng thời gian chạy mỗi kỷ nguyên từ 34.6 xuống 27.8 giây.

Nút cổ chai cuối cùng mà các nhà nghiên cứu đã giải quyết là dẫn truyền dữ liệu theo lô nhỏ giữa CPU và GPU bằng cách sử dụng bước tìm nạp trước, bước này sẽ chuẩn bị dữ liệu ngay trước khi cần. Nhóm đã tính toán rằng điều này sẽ tối đa hóa việc sử dụng băng thông trong liên kết dữ liệu và đưa phương pháp này đến mức sử dụng hoàn hảo; tuy nhiên, họ chỉ nhìn thấy khoảng 90 phần trăm. Họ đã xác định và sửa một lỗi hiệu suất trong thư viện PyTorch phổ biến gây ra giao tiếp khứ hồi không cần thiết giữa CPU và GPU. Sau khi sửa lỗi này, nhóm đã đạt được thời gian chạy mỗi kỷ nguyên là 16.5 giây với SALIENT.

Kaler nói: “Tôi nghĩ, công việc của chúng tôi cho thấy rằng ma quỷ nằm trong các chi tiết. “Khi bạn chú ý kỹ đến các chi tiết ảnh hưởng đến hiệu suất khi đào tạo mạng lưới thần kinh đồ thị, bạn có thể giải quyết một số lượng lớn các vấn đề về hiệu suất. Với các giải pháp của chúng tôi, cuối cùng chúng tôi hoàn toàn bị tắc nghẽn do tính toán GPU, đó là mục tiêu lý tưởng của một hệ thống như vậy.”

Tốc độ của SALIENT được đánh giá trên ba bộ dữ liệu tiêu chuẩn ogbn-arxiv, ogbn-products và ogbn-papers100M, cũng như trong cài đặt nhiều máy, với các mức phân tán khác nhau (lượng dữ liệu mà CPU sẽ chuẩn bị cho GPU) và trên một số kiến ​​trúc, bao gồm kiến ​​trúc hiện đại nhất gần đây, GraphSAGE-RI. Trong mỗi cài đặt, SALIENT vượt trội so với PyTorch Geometric, đáng chú ý nhất là trên bộ dữ liệu ogbn-papers100M lớn, chứa 100 triệu nút và hơn một tỷ cạnh. Tại đây, nó chạy trên một GPU nhanh hơn ba lần so với đường cơ sở được tối ưu hóa ban đầu được tạo cho công việc này; với 16 GPU, SALIENT nhanh hơn gấp tám lần. 

Mặc dù các hệ thống khác có thiết lập thử nghiệm và phần cứng hơi khác nhau nên không phải lúc nào cũng là so sánh trực tiếp, SALIENT vẫn vượt trội so với các hệ thống đó. Trong số các hệ thống đạt được độ chính xác tương tự, các con số hiệu suất đại diện bao gồm 99 giây khi sử dụng một GPU và 32 CPU và 13 giây khi sử dụng 1,536 CPU. Ngược lại, thời gian chạy của SALIENT khi sử dụng một GPU và 20 CPU là 16.5 giây và chỉ là 16 giây với 320 GPU và 16 CPU. Kaler cho biết: “Nếu bạn nhìn vào những con số cuối cùng mà công việc trước đó báo cáo, thì thời gian chạy XNUMX GPU của chúng tôi (hai giây) nhanh hơn rất nhiều so với những con số khác đã được báo cáo trước đây trên tập dữ liệu này”. Các nhà nghiên cứu cho rằng những cải tiến về hiệu suất của họ một phần là nhờ cách tiếp cận tối ưu hóa mã của họ cho một máy duy nhất trước khi chuyển sang cài đặt phân tán. Stathas nói rằng bài học ở đây là đối với tiền của bạn, “sẽ hợp lý hơn nếu sử dụng phần cứng bạn có một cách hiệu quả và đến mức tối đa, trước khi bạn bắt đầu mở rộng quy mô cho nhiều máy tính”, điều này có thể giúp tiết kiệm đáng kể chi phí và lượng khí thải carbon có thể đi kèm với đào tạo người mẫu.

Khả năng mới này giờ đây sẽ cho phép các nhà nghiên cứu giải quyết và tìm hiểu sâu hơn về các biểu đồ ngày càng lớn hơn. Ví dụ: mạng Bitcoin đã được đề cập trước đó có 100,000 nút; hệ thống SALIENT có thể xử lý một biểu đồ lớn hơn 1,000 lần (hoặc ba bậc độ lớn).

“Trong tương lai, chúng tôi sẽ xem xét không chỉ chạy hệ thống đào tạo mạng thần kinh đồ thị này trên các thuật toán hiện có mà chúng tôi đã triển khai để phân loại hoặc dự đoán các thuộc tính của từng nút, mà chúng tôi còn muốn thực hiện các tác vụ chuyên sâu hơn, chẳng hạn như xác định các mẫu phổ biến trong một biểu đồ (các mẫu biểu đồ con), [mà] có thể thực sự thú vị để chỉ ra các tội phạm tài chính,” Chen nói. “Chúng tôi cũng muốn xác định các nút trong biểu đồ giống nhau theo nghĩa là chúng có thể tương ứng với cùng một kẻ xấu trong tội phạm tài chính. Những nhiệm vụ này sẽ yêu cầu phát triển các thuật toán bổ sung và có thể cả kiến ​​trúc mạng thần kinh.”

Nghiên cứu này được hỗ trợ bởi Phòng thí nghiệm trí tuệ nhân tạo MIT-IBM Watson và một phần bởi Phòng thí nghiệm nghiên cứu của Lực lượng Không quân Hoa Kỳ và Máy gia tốc Trí tuệ Nhân tạo của Lực lượng Không quân Hoa Kỳ.