Kỹ thuật đào tạo mạng lưới nơ-ron lớn

Các mạng nơ-ron lớn là cốt lõi của nhiều tiến bộ gần đây trong AI, nhưng việc đào tạo chúng là một thách thức khó khăn về kỹ thuật và nghiên cứu, đòi hỏi phải sắp xếp một cụm GPU để thực hiện một phép tính đồng bộ duy nhất. Khi kích thước cụm và mô hình tăng lên, những người thực hành học máy đã phát triển nhiều kỹ thuật đa dạng hơn để song song hóa quá trình đào tạo mô hình trên nhiều GPU. Thoạt nhìn, việc hiểu các kỹ thuật song song này có vẻ khó khăn, nhưng chỉ cần một vài giả định về cấu trúc tính toán, các kỹ thuật này trở nên rõ ràng hơn nhiều—vào thời điểm đó, bạn chỉ đang di chuyển các bit mờ đục từ A đến B giống như một công tắc mạng di chuyển các gói tin.

Không có sự song song

Đào tạo mạng nơ-ron là một quá trình lặp đi lặp lại. Trong mỗi lần lặp lại, chúng tôi thực hiện chuyển tiếp qua các lớp của mô hình (mở trong cửa sổ mới) để tính toán đầu ra cho mỗi ví dụ đào tạo trong một lô dữ liệu. Sau đó, một lần chạy khác tiến hành ngược lại (mở trong cửa sổ mới) qua các lớp, truyền bá mức độ ảnh hưởng của từng tham số đến đầu ra cuối cùng bằng cách tính toán một gradient (mở trong cửa sổ mới) đối với từng tham số. Độ dốc trung bình cho lô, các tham số và một số trạng thái tối ưu hóa cho mỗi tham số được chuyển đến một thuật toán tối ưu hóa, chẳng hạn như Adam (mở trong cửa sổ mới), tính toán các tham số của lần lặp tiếp theo (có thể có hiệu suất tốt hơn một chút trên dữ liệu của bạn) và trạng thái tối ưu hóa mới cho từng tham số. Khi quá trình đào tạo lặp lại qua các đợt dữ liệu, mô hình sẽ phát triển để tạo ra các đầu ra ngày càng chính xác.

Nhiều kỹ thuật song song chia quá trình đào tạo này thành nhiều chiều khác nhau, bao gồm:

+ Song song dữ liệu—chạy các tập hợp con khác nhau của lô trên các GPU khác nhau;

+ Song song hóa đường ống—chạy các lớp khác nhau của mô hình trên các GPU khác nhau;

+ Song song tenxơ—chia nhỏ phép toán cho một phép toán đơn lẻ như phép nhân ma trận để chia trên các GPU;

+ Hỗn hợp chuyên gia—xử lý từng ví dụ chỉ bằng một phần nhỏ của mỗi lớp.

(Trong bài đăng này, chúng tôi sẽ giả định rằng bạn đang sử dụng GPU để đào tạo mạng nơ-ron của mình, nhưng những ý tưởng tương tự cũng áp dụng cho những người sử dụng bất kỳ trình tăng tốc mạng nơ-ron nào khác (mở trong cửa sổ mới).)

Dữ liệu song song

Đào tạo song song dữ liệu có nghĩa là sao chép cùng một tham số vào nhiều GPU (thường được gọi là "người làm việc") và chỉ định các ví dụ khác nhau cho từng GPU để xử lý đồng thời. Chỉ riêng tính song song dữ liệu vẫn yêu cầu mô hình của bạn phải phù hợp với bộ nhớ của một GPU duy nhất, nhưng cho phép bạn sử dụng tính toán của nhiều GPU với chi phí lưu trữ nhiều bản sao trùng lặp của các tham số của bạn. Điều đó nói rằng, có những chiến lược để tăng RAM hiệu quả có sẵn cho GPU của bạn, chẳng hạn như tạm thời chuyển các tham số sang bộ nhớ CPU giữa các lần sử dụng.

Khi mỗi công nhân song song dữ liệu cập nhật bản sao các tham số của mình, họ cần phối hợp để đảm bảo rằng mỗi công nhân tiếp tục có các tham số tương tự. Cách tiếp cận đơn giản nhất là giới thiệu giao tiếp chặn giữa các công nhân: độc lập tính toán độ dốc trên mỗi công nhân; trung bình các độ dốc giữa các công nhân (mở trong cửa sổ mới); và tính toán độc lập các tham số mới giống nhau trên mỗi công nhân. Bước (2) là một giá trị trung bình chặn đòi hỏi phải truyền khá nhiều dữ liệu (tỷ lệ với số lượng công nhân nhân với kích thước của các tham số của bạn), điều này có thể làm giảm thông lượng đào tạo của bạn. Có nhiều lược đồ đồng bộ hóa không đồng bộ khác nhau (mở trong cửa sổ mới) để loại bỏ chi phí này, nhưng lại làm giảm hiệu quả học tập; trên thực tế, mọi người thường áp dụng phương pháp đồng bộ.

Đường ống song song

Với đào tạo Pipeline Parallel, chúng tôi phân vùng các khối tuần tự của mô hình trên các GPU. Mỗi GPU chỉ giữ một phần nhỏ các tham số và do đó cùng một mô hình tiêu thụ ít bộ nhớ hơn theo tỷ lệ trên mỗi GPU.

Thật dễ dàng để chia một mô hình lớn thành các khối lớp liên tiếp. Tuy nhiên, có một sự phụ thuộc tuần tự giữa đầu vào và đầu ra của các lớp, vì vậy một triển khai ngây thơ có thể dẫn đến một lượng lớn thời gian nhàn rỗi trong khi một công nhân chờ đầu ra từ máy trước được sử dụng làm đầu vào của nó. Các khối thời gian chờ này được gọi là "bong bóng", lãng phí việc tính toán có thể được thực hiện bởi các máy nhàn rỗi.

Chúng ta có thể tái sử dụng các ý tưởng từ tính song song dữ liệu để giảm chi phí của bong bóng bằng cách để mỗi công nhân chỉ xử lý một tập hợp con các phần tử dữ liệu tại một thời điểm, cho phép chúng ta khéo léo chồng chéo tính toán mới với thời gian chờ. Ý tưởng cốt lõi là chia một lô thành nhiều lô nhỏ; mỗi lô nhỏ sẽ được xử lý nhanh hơn theo tỷ lệ và mỗi công nhân bắt đầu làm việc trên lô nhỏ tiếp theo ngay khi có sẵn, do đó đẩy nhanh quá trình thực hiện đường ống. Với đủ số lượng lô nhỏ, công nhân có thể được sử dụng hầu hết thời gian với một bong bóng tối thiểu ở đầu và cuối bước. Các gradient được tính trung bình trên các lô nhỏ và các bản cập nhật cho các tham số chỉ diễn ra sau khi tất cả các lô nhỏ đã hoàn tất.

Số lượng công nhân mà mô hình được chia ra thường được gọi là độ sâu đường ống .

Trong quá trình chuyển tiếp, các công nhân chỉ cần gửi đầu ra (gọi là kích hoạt) của khối lớp của mình cho công nhân tiếp theo; trong quá trình chuyển tiếp ngược, nó chỉ gửi các gradient trên các kích hoạt đó cho công nhân trước đó. Có một không gian thiết kế lớn về cách lên lịch cho các lần chuyển tiếp này và cách tổng hợp các gradient trên các microbatch. GPipe (mở trong cửa sổ mới) có mỗi quy trình công nhân thực hiện liên tục các bước tiến và lùi, sau đó tổng hợp các gradient từ nhiều mẻ nhỏ một cách đồng bộ vào cuối. PipeDream (mở trong cửa sổ mới) thay vào đó, hãy lên lịch cho từng công nhân để lần lượt xử lý các lượt đi tới và đi lui.

Song song tenxơ

Song song đường ống chia tách một mô hình theo "chiều dọc" theo từng lớp. Cũng có thể chia tách "theo chiều ngang" một số hoạt động trong một lớp, thường được gọi là đào tạo song song Tensor . Đối với nhiều mô hình hiện đại (như Transformer(mở trong cửa sổ mới)), nút thắt tính toán là nhân ma trận lô kích hoạt với ma trận trọng số lớn. Phép nhân ma trận (mở trong cửa sổ mới) có thể được coi là tích vô hướng giữa các cặp hàng và cột; có thể tính tích vô hướng độc lập trên các GPU khác nhau hoặc tính các phần của từng tích vô hướng trên các GPU khác nhau và tính tổng kết quả. Với bất kỳ chiến lược nào, chúng ta có thể cắt ma trận trọng số thành các "mảnh" có kích thước bằng nhau, lưu trữ từng mảnh trên một GPU khác nhau và sử dụng mảnh đó để tính phần có liên quan của tích ma trận tổng thể trước khi giao tiếp sau đó để kết hợp các kết quả.

Một ví dụ là Megatron-LM (mở trong cửa sổ mới), song song hóa phép nhân ma trận trong các lớp tự chú ý và MLP của Transformer. PTD-P (mở trong cửa sổ mới) sử dụng tenxơ, dữ liệu và song song đường ống; lịch trình đường ống của nó gán nhiều lớp không liên tiếp cho mỗi thiết bị, giảm chi phí bong bóng với cái giá phải trả là nhiều giao tiếp mạng hơn.

Đôi khi đầu vào của mạng có thể được song song hóa trên một chiều với mức độ tính toán song song cao so với giao tiếp chéo. Song song trình tự (mở trong cửa sổ mới) là một ý tưởng như vậy, trong đó chuỗi đầu vào được chia theo thời gian thành nhiều ví dụ phụ, giảm tỷ lệ tiêu thụ bộ nhớ tối đa bằng cách cho phép tính toán tiến hành với các ví dụ có kích thước chi tiết hơn.

Hỗn hợp chuyên gia (MoE)

Với sự kết hợp của các chuyên gia (MoE) (mở trong cửa sổ mới) cách tiếp cận, chỉ một phần nhỏ của mạng được sử dụng để tính toán đầu ra cho bất kỳ một đầu vào nào. Một ví dụ về cách tiếp cận là có nhiều bộ trọng số và mạng có thể chọn bộ nào để sử dụng thông qua cơ chế gating tại thời điểm suy luận. Điều này cho phép nhiều tham số hơn mà không làm tăng chi phí tính toán. Mỗi bộ trọng số được gọi là "chuyên gia", với hy vọng rằng mạng sẽ học cách chỉ định các kỹ năng và tính toán chuyên biệt cho từng chuyên gia. Các chuyên gia khác nhau có thể được lưu trữ trên các GPU khác nhau, cung cấp một cách rõ ràng để mở rộng quy mô số lượng GPU được sử dụng cho một mô hình.

GS cứng (mở trong cửa sổ mới) mở rộng một máy biến áp MoE lên tới 600 tỷ tham số với một sơ đồ trong đó chỉ có các lớp MoE được chia thành nhiều thiết bị TPU và các lớp khác được sao chép hoàn toàn. Máy biến áp chuyển mạch(mở trong cửa sổ mới) mở rộng quy mô mô hình lên hàng nghìn tỷ tham số với độ thưa thớt thậm chí còn cao hơn bằng cách định tuyến một đầu vào cho một chuyên gia duy nhất.

Các thiết kế tiết kiệm bộ nhớ khác

Có nhiều chiến lược tính toán khác để giúp việc đào tạo mạng nơ-ron ngày càng lớn trở nên dễ dàng hơn. Ví dụ:

+ Để tính toán gradient, bạn cần phải lưu các kích hoạt ban đầu, điều này có thể tiêu tốn rất nhiều RAM của thiết bị. Điểm kiểm tra(mở trong cửa sổ mới) (còn được gọi là tính toán lại kích hoạt) lưu trữ bất kỳ tập hợp con nào của các kích hoạt và tính toán lại các kích hoạt trung gian ngay trong quá trình truyền ngược. Điều này tiết kiệm rất nhiều bộ nhớ với chi phí tính toán tối đa là một lần truyền ngược đầy đủ bổ sung. Người ta cũng có thể liên tục cân bằng giữa chi phí tính toán và bộ nhớ bằng cách tính toán lại kích hoạt có chọn lọc(mở trong cửa sổ mới), là các tập hợp con kiểm tra các hoạt động kích hoạt có chi phí lưu trữ tương đối cao nhưng lại rẻ hơn khi tính toán.

+ Đào tạo chính xác hỗn hợp(mở trong cửa sổ mới) là đào tạo các mô hình sử dụng các số có độ chính xác thấp hơn (thường là FP16) (mở trong cửa sổ mới)). Các máy gia tốc hiện đại có thể đạt được số lượng FLOP cao hơn nhiều với số có độ chính xác thấp hơn và bạn cũng tiết kiệm được RAM của thiết bị. Với sự chăm sóc thích hợp, mô hình kết quả có thể gần như không mất đi độ chính xác.

+ Offloading là tạm thời offload dữ liệu chưa sử dụng đến CPU hoặc giữa các thiết bị khác nhau và sau đó đọc lại khi cần. Các triển khai ngây thơ sẽ làm chậm quá trình đào tạo rất nhiều, nhưng các triển khai tinh vi sẽ tải trước dữ liệu để thiết bị không bao giờ phải chờ đợi dữ liệu đó. Một triển khai của ý tưởng này là ZeRO (mở trong cửa sổ mới) phân tách các tham số, độ dốc và trạng thái tối ưu hóa trên tất cả phần cứng có sẵn và hiện thực hóa chúng khi cần thiết.

+ Các trình tối ưu hóa hiệu quả bộ nhớ đã được đề xuất để giảm dấu chân bộ nhớ của trạng thái đang chạy được duy trì bởi trình tối ưu hóa, chẳng hạn như Adafactor (mở trong cửa sổ mới).

+ Nén cũng có thể được sử dụng để lưu trữ các kết quả trung gian trong mạng. Ví dụ, Gist (mở trong cửa sổ mới) nén các kích hoạt được lưu cho lần truyền ngược; DALL·E nén các gradient trước khi đồng bộ hóa chúng.

Tại OpenAI, chúng tôi đang đào tạo và cải thiện các mô hình lớn từ cơ sở hạ tầng cơ bản cho đến khi triển khai chúng cho các vấn đề thực tế. Nếu bạn muốn đưa các ý tưởng từ bài đăng này vào thực tế—đặc biệt liên quan đến nhóm Nghiên cứu ứng dụng và Mở rộng quy mô của chúng tôi—chúng tôi đang tuyển dụng !

Xem thêm: mua tài khoản ChatGPT Plus chính hãng giá rẻ với nhiều ưu đãi đặc biệt!