Nâng cao hiệu suất trong huấn luyện LLM quy mô lớn bằng cách sử dụng song song tensor không đồng nhất

Việc huấn luyện các LLM ở quy mô lớn mang lại những thách thức cơ sở hạ tầng độc đáo, đặc biệt khi các công việc trải dài hàng nghìn GPU và chạy trong thời gian dài. Công việc càng chạy lâu, khả năng gặp phải các sự cố gián đoạn ngoài kế hoạch hoặc biến động tài nguyên càng lớn. Ngay cả việc thiết bị không khả dụng không thường xuyên cũng có thể gây ra những tác động lớn đến các cụm được kết nối chặt chẽ, dẫn đến sự chậm lại cho một lần chạy huấn luyện nhất định.

Đối với các công việc huấn luyện quy mô lớn, việc điều chỉnh linh hoạt công việc theo số lượng GPU khả dụng là một phương pháp mạnh mẽ để cải thiện Goodput. Trong bối cảnh huấn luyện AI, Goodput đại diện cho thước đo quan trọng về khối lượng công việc hữu ích, thúc đẩy sự hội tụ được hoàn thành, thay vì chỉ là thông lượng phần cứng thô.

Các phương pháp hiệu quả để mở rộng linh hoạt hiện nay bao gồm việc loại bỏ một bản sao dữ liệu, sử dụng tính năng khởi động lại điểm kiểm tra nhanh, hoặc chuyển sang các thiết bị dự phòng nóng. Những phương pháp này cho phép các công việc LLM thích ứng với những thay đổi về tính khả dụng của GPU trong khi vẫn duy trì sự cân bằng trên toàn hệ thống. Tuy nhiên, chúng cũng gây ra một lượng thông lượng bị mất và chi phí cao hơn trong khoảng thời gian công việc huấn luyện chạy trong điều kiện khả dụng giảm.

Một bài báo gần đây về Nonuniform Tensor Parallelism (NTP) giới thiệu một khuôn khổ thử nghiệm hướng tới tương lai, xây dựng trên các phương pháp này theo cách giảm thiểu chi phí thông lượng. Kết hợp với khả năng tăng cường điện năng động tiềm năng để bù đắp bất kỳ sự suy giảm hiệu suất nào, thông lượng vẫn ổn định, biến các sự gián đoạn thành các sự kiện có thể quản lý và phục hồi.

Ad

Đóng góp cốt lõi của NTP là khả năng duy trì Goodput cao bằng cách ngăn chặn các sự cố thiết bị tạm thời làm đình trệ các tác vụ huấn luyện lớn, được kết nối. Bằng cách điều chỉnh động mức độ song song tensor và chồng chéo thông minh việc phân chia lại dữ liệu cần thiết, NTP giảm thiểu thời gian bị mất và nỗ lực tính toán.

Điều này đảm bảo rằng cụm máy chủ dành phần lớn thời gian hoạt động của mình để thực hiện công việc hữu ích, duy trì tính toàn vẹn và hiệu quả tổng thể của quá trình huấn luyện ngay cả khi các điều kiện phần cứng dao động.

Những thách thức với việc huấn luyện quy mô lớn

Việc huấn luyện mô hình AI là một công việc song song, trải rộng trên hàng nghìn GPU. Một kỹ thuật phổ biến để song song hóa các khối lượng công việc này là tensor parallelism (TP), trong đó các lớp của mạng nơ-ron được chia sẻ trên một nhóm GPU được kết nối chặt chẽ. Số lượng GPU trong nhóm này trùng với miền mở rộng quy mô (scale-up domain), vốn được kết nối với các kết nối tốc độ cao như NVIDIA NVLink. Trên các hệ thống NVIDIA Blackwell và NVIDIA Blackwell Ultra, NVLink kết nối tối đa 72 GPU ở tốc độ 1.800 GB/s, hỗ trợ giao tiếp tất cả-đến-tất cả (all-to-all) trong một bước nhảy.

Trong một kịch bản điển hình, một LLM tiên tiến được huấn luyện trên một cụm các giá đỡ, mỗi giá đỡ chứa nhiều máy chủ. Một giá đỡ máy chủ đơn lẻ tạo thành miền mở rộng và tạo thành một nhóm TP duy nhất. Song song dữ liệu (DP) sau đó nhân bản mô hình trên nhiều miền mở rộng như vậy, mỗi miền xử lý một lô dữ liệu khác nhau. Sự thay đổi trạng thái GPU trong một miền mở rộng có thể ảnh hưởng đến hiệu suất của nhóm TP đó. Vì các GPU trong cùng một nhóm TP chia sẻ các phép tính được ghép nối chặt chẽ, một sự cố với một thiết bị có thể làm giảm hiệu suất huấn luyện hoặc yêu cầu cân bằng lại tạm thời để duy trì tiến độ.

Khi kiến trúc trung tâm dữ liệu phát triển để hỗ trợ các miền mở rộng lớn hơn, từ tám đến 72 GPU và hơn thế nữa, việc tối đa hóa thời gian hoạt động hiệu quả của mọi thiết bị khỏe mạnh trở thành chìa khóa để đạt được Goodput cao.

Bài báo lưu ý rằng thay vì để các gián đoạn cục bộ, nhất thời quyết định thông lượng huấn luyện tổng thể, các hệ thống có thể được thiết kế để giữ cho phần lớn các GPU đang hoạt động tiếp tục xử lý. Bằng cách thích ứng với những biến động phần cứng này, các cụm có thể duy trì việc sử dụng tài nguyên hiệu quả cao cần thiết để tối ưu hóa Goodput trong quá trình huấn luyện mô hình lớn ở quy mô lớn.

Huấn luyện thường liên quan đến một quy trình mà mỗi giai đoạn phụ thuộc vào việc hoàn thành kịp thời các bước trước đó. Nếu một GPU gặp sự chậm trễ trong một nhóm TP, nó có thể làm chậm quá trình đồng bộ hóa hoặc xử lý trên nhóm đó, gây ra tình trạng dừng tạm thời hoặc giảm thông lượng cho đến khi hệ thống phục hồi hoặc phân bổ lại khối lượng công việc.

Cách NTP duy trì quá trình huấn luyện 

Ý tưởng cơ bản đằng sau NTP là đảm bảo một bản sao DP vẫn hoạt động ngay cả trong các gián đoạn phần cứng nhất thời. Khi tiếp tục từ điểm kiểm tra mới nhất, mô hình tự động điều chỉnh cấu hình của nó theo phần cứng có sẵn, duy trì chức năng một phần để giữ cho quy trình tiếp tục mà không làm mất hoàn toàn sự đóng góp của bản sao đó.

Phần sau đi sâu hơn vào cách NTP đạt được khả năng phục hồi này.

Thích ứng bậc TP động

Khi một GPU trong miền mở rộng (scale-up domain) gặp sự cố, hệ thống sẽ xác định nhóm bị ảnh hưởng và tự động cấu hình lại song song tensor (tensor parallelism) của nó để chỉ sử dụng các GPU còn hoạt động. Ví dụ, nếu một nhóm TP gồm tám GPU bị một GPU bị loại bỏ, nó có thể chuyển động sang bậc TP là bảy. Phần mô hình đó tiếp tục tính toán, ngăn chặn sự mất mát hoàn toàn đóng góp của nó. Các GPU còn lại trong nhóm sẽ tăng khối lượng công việc cá nhân, cho phép công việc huấn luyện duy trì Goodput và tính khả dụng cao ngay cả khi một tập hợp tài nguyên gặp sự cố.

Tăng cường nguồn điện để bù đắp hiệu suất

Việc giảm bậc TP là không đủ để duy trì thông lượng toàn cục. Một bản sao DP với ít GPU hơn sẽ chạy chậm hơn một cách cố hữu, khiến toàn bộ hệ thống DP bị đình trệ, chờ bản sao chậm nhất. Để khắc phục điều này, nghiên cứu đề xuất một thiết kế rack tích hợp các khả năng điện và nhiệt được cải tiến. Thiết kế này cho phép tăng cường nguồn điện (power-boosting) cho các miền mở rộng với tính khả dụng giảm. Bằng cách tăng động nguồn điện cung cấp cho các GPU đang hoạt động trong một miền bị ảnh hưởng, tần số xung nhịp và thông lượng tính toán có thể được nâng cao tạm thời.

Bản sao DP với bậc TP giảm có thể bắt kịp và theo kịp các bản sao khác đang hoạt động đầy đủ, ngăn ngừa các nút thắt đồng bộ hóa toàn cục và đảm bảo Goodput của cụm được tối ưu hóa cao mặc dù có sự dao động phần cứng cục bộ.

Tái phân chia hiệu quả

Việc điều chỉnh động bậc TP đòi hỏi một cơ chế hiệu quả để phân phối lại các mảnh tensor (tensor shards) của mô hình giữa các GPU còn lại. NTP sử dụng kỹ thuật tái phân chia (resharding) thông minh được chồng lên các giai đoạn tính toán khác.

Bằng cách thực hiện việc tái phân chia này trong các giai đoạn tính toán ngược và đồng bộ hóa tham số, chi phí bổ sung đối với các bản sao khỏe mạnh được giảm thiểu, thường dưới 1%. Việc lập lịch trình cẩn thận này tối đa hóa hiệu quả tính toán, duy trì liền mạch Goodput tối ưu mà bản thân cơ chế thích ứng không trở thành nút thắt cổ chai về hiệu suất.

NTP xây dựng một con đường kiên cường hơn cho việc huấn luyện AI quy mô lớn

Công trình này nhấn mạnh tầm quan trọng sống còn của việc đồng thiết kế phần cứng và phần mềm để giải quyết các thách thức cố hữu trong việc huấn luyện AI quy mô lớn. Việc tích hợp NTP với các thiết kế giá đỡ tiên tiến, cung cấp khả năng dự trữ điện và nhiệt cần thiết cho việc tăng cường công suất động, là một ví dụ điển hình về cách các đổi mới phần cứng chu đáo có thể bổ sung sâu sắc cho các giải pháp phần mềm tinh vi.

Mối quan hệ cộng sinh giữa phần cứng và phần mềm này là điều cần thiết để mở khóa các mức hiệu suất, hiệu quả, Goodput ổn định và khả năng phục hồi cao hơn trong thế hệ hệ thống AI tiếp theo. Là một tính năng thử nghiệm hướng tới tương lai, NTP chứng minh những gì có thể xảy ra khi khả năng phục hồi được tích hợp trực tiếp vào chiến lược song song.

Dựa trên nền tảng này, nghiên cứu đang được tiến hành để mở rộng các khái niệm này sang Song song Chuyên gia Không đồng nhất (NEP), tối ưu hóa khả năng phục hồi cho các mô hình Mixture-of-Experts (MoE) nơi song song tensor tiêu chuẩn kém lý tưởng hơn.

Hãy xem các tính năng chịu lỗi và khả năng phục hồi sẵn sàng cho sản xuất có trong NVIDIA Resiliency Extension (NVRx) hoặc Nonuniform Tensor Parallelism Readme để tìm hiểu thêm về sự bổ sung gần đây của nó vào nhánh nhà phát triển của NVIDIA Megatron Core.

____
Bài viết liên quan