Đã có rất nhiều cuộc thảo luận sau bài trình bày về nhân Jalapeño MLA tại HotChips và các bài bình luận tiếp theo của SemiAnalysis. Với tư cách là đội ngũ phần cứng của OpenAI, chúng tôi chỉ chạm nhẹ vào một điểm vàng nhỏ này: sự thật rằng AI đang viết các nhân (kernel) của chúng tôi, và khi nó làm vậy, chúng tôi không thực sự cần hiểu từng dòng mã nhân đó làm gì. Chúng tôi đã bỏ qua một cách rõ ràng: làm thế nào mà điều như vậy lại có thể xảy ra? Đâu là cách suy nghĩ đúng đắn về vấn đề này, so với phương pháp tạo mã truyền thống hơn? Liệu nhân đã được tối ưu hóa có đáng tin cậy như nhân chưa được tối ưu hóa không?
Về nền tảng của tôi, tôi đã làm việc với các trình biên dịch cho các bộ tăng tốc (accelerator) trong hơn một thập kỷ. Tôi đã khởi xướng XLA, một cơ sở hạ tầng trình biên dịch xuất sắc với một đội ngũ và nỗ lực liên công ty tuyệt vời. Trong hơn 2 năm qua tại OpenAI, tôi đã cố gắng tái định hình cách các trình biên dịch nên hoạt động trong kỷ nguyên AI. Các công thức trình biên dịch mới sẽ dựa trên những thế mạnh hiện có, nhưng không thể phủ nhận rằng có một công cụ mới mạnh mẽ để tận dụng trong bộ công cụ.
Đây sẽ là một hành trình dài, nhưng tôi hy vọng sẽ làm sáng tỏ cách AI đang được sử dụng để tự động hóa việc cải tiến chương trình máy tính; tức là biên dịch tối ưu hóa. Tôi thực sự tin rằng, nhờ AI, chúng ta có thể trải nghiệm điều mà chúng ta gọi là "trình biên dịch 2.0". AI ít bị giới hạn về mặt cơ bản hơn trong những gì nó có thể đề xuất, và những gì nó đề xuất là kết quả của quá trình huấn luyện và ngữ cảnh của mô hình, điều này khiến tôi phân loại nó như một "bộ tối ưu hóa ngẫu nhiên" (stochastic optimizer) – điều này có thể đặt ra những thách thức nhưng, như chúng ta sẽ thấy, cũng là nguồn gốc của những thế mạnh to lớn...
Rất nhiều nghiên cứu học thuật và ứng dụng trong ngành đã và đang hướng tới hướng đi này, và nhanh chóng khám phá ra tiềm năng cho sự tham gia của AI trong lĩnh vực trình biên dịch tối ưu hóa, nhưng chúng ta đang ở một thời điểm cần một lời giải thích tổng quan.
Bối cảnh
Trình biên dịch nhận đầu vào là các chương trình và tạo ra các phiên bản đã được dịch hoặc cải tiến của các chương trình đó.
Các chương trình, cả ở phía đầu vào và đầu ra, đều có ngữ nghĩa (semantics) cho chúng ta biết chương trình đó có nghĩa là gì, chúng có thể làm gì, và cách suy luận về những điều chúng có thể làm.
Những người trong chúng ta làm việc với trình biên dịch thường nghĩ về chúng như các hàm thuần túy (pure functions) – chúng nhận một cấu trúc dữ liệu và tạo ra một cấu trúc dữ liệu có ngữ nghĩa tương ứng.
Đôi khi trình biên dịch của chúng ta tập trung vào việc "hạ thấp mức độ" (lowering) hoặc "dịch thuật" (translating). Ví dụ, chúng có thể nhận mã C và tạo ra mã hợp ngữ x86-64, mà chúng ta thường coi là "cấp thấp hơn". Nhưng thông thường, chúng làm nhiều hơn là chỉ dịch thuật như một phần nhỏ của quá trình đó...
Trong thực tế, trình biên dịch của chúng ta tập trung vào việc "tối ưu hóa". Chúng có thể nhận một cấu trúc dữ liệu đại diện cho chương trình – theo thuật ngữ của chúng tôi là "Biểu diễn trung gian" (Intermediate Representation - IR) – và cố gắng tạo ra một phiên bản tốt hơn của chương trình đó. Đôi khi "tốt hơn" có nghĩa là chạy ít chu kỳ hơn, đôi khi có nghĩa là có ít mã không cần thiết hơn, đôi khi có nghĩa là chuyên biệt hóa cho những thứ chúng ta có thể chứng minh là "phải đúng" về chương trình (đánh giá một phần - partial evaluation).
Bây giờ, hãy xem xét một cách ngắn gọn rằng LLM ban đầu được tạo ra để dịch văn bản của con người từ ngôn ngữ này sang ngôn ngữ khác. Rõ ràng việc dịch thuật nằm trong thế mạnh của chúng. Và chúng ta có thể thấy qua việc sử dụng LLM trong các tác vụ hàng ngày rằng chúng cũng có thể viết các giải pháp mới và cải thiện các giải pháp hiện có. Nhiều lập trình viên trong chúng ta cũng có kinh nghiệm yêu cầu LLM "tối ưu hóa đoạn mã này" và chúng có thể làm được một cách đáng kinh ngạc. (Tuy nhiên, chúng ta cần biết rằng chúng đã tối ưu hóa mã một cách chính xác, chúng ta sẽ đề cập đến vấn đề này sau!) Điều này chỉ đơn giản là để nhấn mạnh rằng LLM có những khả năng mà chúng ta tìm kiếm ở một trình biên dịch tối ưu hóa.
Tối ưu hóa và Tính tối ưu
Các trình biên dịch tối ưu hóa, không có gì ngạc nhiên, đang cố gắng tăng tính tối ưu của chương trình mà chúng đang xử lý, dựa trên một mục tiêu nào đó (thường là thời gian thực thi). Điều đó rất khó thực hiện trong trường hợp tổng quát, đối với một chương trình bất kỳ, đến nỗi có một định lý gọi là định lý toàn dụng lao động cho các kỹ sư trình biên dịch. (Tôi chỉ phát hiện ra điều này sau khi chọn làm kỹ sư trình biên dịch, nhưng nó vẫn mang lại cho tôi sự an ủi!)
"Siêu tối ưu hóa" (Superoptimizers) là một tiểu lĩnh vực tuyệt vời của các trình biên dịch tối ưu hóa. Hãy tưởng tượng có một chương trình nhất định, và chúng ta có thể nói nó làm gì thông qua ngữ nghĩa. Đâu là chương trình tối ưu nhất có cùng ngữ nghĩa đó? Đó là những gì siêu tối ưu hóa cố gắng giải quyết, và về bản chất, đó là một bài toán tìm kiếm...
Hãy tưởng tượng tôi đang cố gắng tìm chương trình ngắn nhất có cùng ngữ nghĩa đó, và tôi có một cách để hỏi liệu một chương trình ứng viên có cùng ngữ nghĩa hay không. Về mặt giả thuyết, tôi có thể liệt kê mọi chương trình theo thứ tự mục tiêu, và chọn chương trình nhỏ nhất có cùng ngữ nghĩa.
Tuy nhiên, việc liệt kê mọi chương trình theo thứ tự mục tiêu nghe có vẻ khó khả thi. Một trong những bài báo học thuật yêu thích của tôi, được thực hiện vào năm 2013 có tựa đề "STOKE" (Siêu tối ưu hóa ngẫu nhiên), đã đặt câu hỏi: "nếu chúng ta chỉ điều chỉnh các chương trình một cách ngẫu nhiên nhiều lần, liệu cuối cùng chúng ta có quan sát được chương trình tốt nhất không?" Họ đề xuất rằng thông qua một bước đi ngẫu nhiên (random walk) (và với người bạn học máy OG của chúng ta là Markov Chain Monte Carlo / Metropolis-Hastings), cuối cùng bạn sẽ thấy chương trình tối ưu đó.
Việc điều chỉnh Monte Carlo thường là ngẫu nhiên (bạn chọn ngẫu nhiên một sự điều chỉnh), nhưng cũng nhanh. LLM thì rất thông minh (nhiều token suy luận), nhưng tương đối chậm.
Điều gì sẽ xảy ra nếu, thay vì điều chỉnh Monte Carlo ngẫu nhiên/nhanh, chúng ta để LLM tìm ra các hướng để đưa chương trình đi? Chúng ta sẽ có một bộ tối ưu hóa ngẫu nhiên rất thông minh, dẫn dắt chương trình của chúng ta đi qua không gian chương trình đã được tối ưu hóa.
Trực giác về Tối ưu hóa
Hãy lùi lại một bước. Hãy nghĩ về người mà bạn biết, người thể hiện rõ nhất việc "tối ưu hóa triệt để các đoạn mã". Nói ngắn gọn, hãy gọi họ là "Ollie tối ưu". Ollie có lẽ có một trực giác bẩm sinh về những loại điều chỉnh mã nào có thể mang lại kết quả. Ollie có thể thử một vài thứ để xem chúng có hiệu quả không, và nếu không, sẽ quay lại và thử thứ khác. Nhưng họ có một số trực giác về những loại điều gì là khả thi, và làm thế nào họ có thể đánh bại trình biên dịch.
Những trực giác mà Ollie có thường vượt xa những gì trình biên dịch làm. Mặc dù các trình biên dịch tối ưu hóa hiện đại khá ấn tượng về kết quả của chúng, nhưng chúng dựa trên các quy tắc và heuristic khá đơn giản. Theo thuật ngữ kỹ thuật, chúng dựa trên ý tưởng về một phép biến đổi luồng dữ liệu cục bộ được chạy đến điểm cố định. Chúng ta cũng sắp xếp thứ tự các giai đoạn; tức là chúng ta xây dựng các pipeline trình biên dịch để xem xét A và sau đó là B, nhưng không phải là bài toán tổng hợp AB. Các bộ lập lịch (scheduler) và bộ cấp phát thanh ghi (register allocator) là một ví dụ khét tiếng về điều này, nhiều bằng tiến sĩ đã được thực hiện về bộ lập lịch-cấp phát thanh ghi tổng hợp (để có được lợi ích của việc loại bỏ thứ tự giai đoạn), nhưng chúng rất khó để hoạt động trong thực tế.
Đây là lý do tại sao chuyên môn của Ollie rất có giá trị. Ollie thường biết cách cân bằng một số bài toán NP-khó với các heuristic phù hợp với từng tình huống. Vì vậy, có sự nhạy bén và nhận thức về ngữ cảnh phù hợp hơn. Ollie cũng có thể sử dụng các kỹ thuật mà trình biên dịch tối ưu hóa có thể không áp dụng một cách có lợi, đặc biệt là kết hợp, như các kỹ thuật tách nhánh (outlining) hoặc tạo các ABI tùy chỉnh hoặc các phép biến đổi để kích hoạt vector hóa, hoặc vô số thứ khác khiến chúng ta càu nhàu "ước gì trình biên dịch có cách để làm điều này..."
Bây giờ hãy xem xét rằng AI, thông qua các cơ chế suy luận của nó, có thể hoạt động như một Ollie thu nhỏ. Nó có thể không có trực giác phù hợp về những gì sẽ thành công, nhưng nó có một ý niệm về những gì có thể mang lại lợi ích, và nó có thể thực hiện rất, rất nhiều lần thử.
Với cách tiếp cận này, không giống như trong bài báo STOKE, chúng ta không thể đảm bảo rằng khi thời gian tiến đến vô cùng, chúng ta có thể thấy chương trình tối ưu, nhưng vì AI có các cơ chế suy luận "giống con người hơn", nó thực sự có thể đạt được lực kéo đáng kể giống con người trên mỗi đơn vị thời gian.
Kết nối lại: Nhân MLA
Hãy để tôi bắt đầu bằng cách nói: Tôi không biết mã cấp thấp mà AI đã tạo ra cho nhân Jalapeño MLA của chúng tôi, nhưng tôi biết cách gõ numpy cho MLA.
Trong trình biên dịch XLA mà tôi đã từng làm việc, chúng tôi sẽ hợp nhất các phép toán numpy đó thành các khối, và sau đó sử dụng một siêu chương trình gọi là "emitter" để hạ thấp nó xuống các vòng lặp, lệnh và các nguyên thủy cấp thấp hơn.
Khi trình biên dịch / chương trình emitter XLA làm điều đó, tôi không cần quan tâm mã hợp ngữ nào được tạo ra ở đầu ra. Đối với bộ tối ưu hóa ngẫu nhiên của chúng tôi, AI về mặt khái niệm thay thế vai trò của siêu chương trình emitter – nó vừa hạ thấp mức độ vừa tối ưu hóa, và chúng tôi có thể yêu cầu nó tối ưu hóa ngày càng gần hơn với roofline.
Tôi hy vọng điều này làm rõ vị trí của AI và cách nó tương tự như một thành phần trong một hệ thống trình biên dịch tối ưu hóa hiện có. Cũng hữu ích khi nghĩ: lớp mà chúng ta coi là "mã hợp ngữ" giờ đây đang được nâng lên. Khi bạn gõ mã C++ thông thường và biên dịch nó ở mức -O3 (mức tối ưu hóa điển hình cao nhất), bạn không mong đợi hiểu mã hợp ngữ được tạo ra, ngay cả khi bạn hiểu mã C++ bạn đã gõ. Chúng tôi đang làm điều tương tự ở đây, nhưng với một đặc tả đầu vào toán học và cấp cao hơn.
Bây giờ, một câu hỏi quan trọng là làm thế nào chúng tôi kiểm tra rằng chương trình nhận được từ AI thực sự tương đương với mô tả cấp cao / numpy. Cơ chế kiểm tra đó thiết lập tính đúng đắn (soundness) của quá trình tối ưu hóa ngẫu nhiên bằng AI. Tôi hy vọng một bài đăng blog trong tương lai sẽ đi vào chi tiết hơn về vấn đề này, nhưng hiện tại, chỉ cần nói rằng việc kiểm tra tính tương đương về ngữ nghĩa là khả thi và chúng tôi làm điều đó. Các chương trình tăng tốc (accelerator programs) đặc biệt phù hợp với các hợp đồng mạnh mẽ, hoàn chỉnh mà chúng tôi có thể xác minh "chính xác những gì chương trình được AI tối ưu hóa làm", vì chúng khá toán học và định hướng luồng dữ liệu trong bối cảnh rộng của chúng.
Lưu ý rằng nhiều kỹ thuật liên quan trong lĩnh vực này đã được tiên phong bởi các nỗ lực trong tiểu lĩnh vực tổng hợp chương trình (program synthesis). Trong khi các trình biên dịch tối ưu hóa nói, "đây là một chương trình với ngữ nghĩa, hãy làm cho nó tốt hơn nhưng với ngữ nghĩa tương đương!", thì tổng hợp chương trình nói, "tôi tin rằng có một chương trình với những ngữ nghĩa này, hãy cố gắng tìm ra chương trình tốt nhất có thể". Tổng hợp chương trình là một bài toán khó hơn biên dịch tối ưu hóa, nhưng nó cũng ít bị giới hạn về mặt cơ bản hơn. Về bản chất, đó là những gì con người như Ollie làm khi họ làm tốt hơn trình biên dịch tối ưu hóa, và đó là điều mà AI giờ đây có thể giúp chúng ta tự động hóa. AI có thể lấy "cảm hứng" từ chương trình gốc, nhưng nó không nhất thiết chỉ thực hiện các phép biến đổi cục bộ nhỏ trên nó. Các trình biên dịch tối ưu hóa cổ điển sẽ không thấy "ồ, bạn đã viết một bubble sort" và, bằng cách hiểu hợp đồng, chuyển nó thành quick-sort, nhưng cả Ollie và AI đều có thể làm điều đó. Đây là điều đưa chúng ta vào chế độ tổng hợp chương trình với tối ưu hóa ngẫu nhiên hơn là chế độ trình biên dịch tối ưu hóa cổ điển.
Tất cả điều này kết hợp lại với nhau trong thực tế là bạn có thể bắt đầu với một thứ "không quá xa so với numpy", đợi 48 giờ, và có một nhân đã được tối ưu hóa với cùng ngữ nghĩa, như chúng tôi đã trình bày trong bài nói chuyện tại HotChips:
Như slide cũng lưu ý, trên máy của chúng tôi, chúng tôi thường có thể quan sát thấy AI vượt qua các chuyên gia con người của chúng tôi ngay cả trên các nhân mà chúng tôi cảm thấy đã được tinh chỉnh khá tốt. Thường vẫn còn một tỷ lệ phần trăm khả thi đáng kể chỉ do có nhiều loại kết hợp / hoán vị có thể cần được khám phá. Những điều này thường rất tẻ nhạt và khó khăn đối với một kỹ sư hiệu năng con người.
Tóm tắt & Kết luận
Một trình biên dịch, xét cho cùng, chỉ là một hàm. Chúng ta đưa chương trình của mình cho hàm đó, và chúng ta nhận lại một phiên bản tốt hơn của chương trình. Chúng ta mong đợi chương trình nhận được và chương trình đưa vào có cùng ngữ nghĩa.
Các trình biên dịch tối ưu hóa truyền thống làm cho chương trình tốt hơn thông qua các quy tắc luồng dữ liệu và heuristic. Chúng hoàn toàn có thể hiểu được về nguồn gốc của mình, nhưng cũng có thể bị hạn chế hơn trong những thay đổi chúng có thể thực hiện.
Ngược lại, AI, với tư cách là một bộ tối ưu hóa ngẫu nhiên, chỉ cần "suy nghĩ kỹ" và đưa ra một kết quả. Các thay đổi của nó không bị giới hạn về mặt cơ bản, khiến nó trở nên tương tự hơn với chuyên gia tối ưu hóa con người của chúng ta. Chúng ta cần có những cách để kiểm tra rằng các chương trình nó đưa ra là đúng đắn và thực hiện cùng một ngữ nghĩa mà chúng ta đưa vào, và chúng ta đã có những cách đó. Và loại tối ưu hóa AI này đặc biệt phù hợp với các phép toán có các hợp đồng rất mạnh mẽ. Các hợp đồng này loại bỏ nhu cầu phải hiểu từng dòng mã nhân làm gì.
Đó là cách chúng tôi có được nhân MLA do AI tạo ra!
