Những tiến bộ của kỹ thuật machine learning dành cho machine learning
(blog.research.google)- Khi mô hình ML càng lớn, hiệu quả thực thi không chỉ phụ thuộc vào phần cứng mà còn bị chi phối mạnh bởi tối ưu hóa trình biên dịch; Google Research và Google DeepMind đã công bố cách tiếp cận dùng chính ML để cải thiện lĩnh vực này
- Các quyết định của trình biên dịch bên dưới những framework như TensorFlow, JAX, PyTorch có thể làm thay đổi đáng kể thời gian chạy và mức sử dụng tài nguyên ngay cả với cùng một mô hình
- TpuGraphs là một bộ dữ liệu chứa đồ thị tính toán, cấu hình biên dịch và thời gian thực thi của chương trình ML cho TPU, có thể dùng để nghiên cứu mô hình chi phí dựa trên học máy
- Để giảm bớt hạn chế của việc khó huấn luyện các đồ thị tính toán lớn cùng lúc, Graph Segment Training chia đồ thị thành các phân đoạn để giảm sử dụng bộ nhớ và rút ngắn thời gian huấn luyện xuống còn 1/3
- Trong cuộc thi Kaggle, 792 người thuộc 616 đội từ 66 quốc gia đã tham gia để kiểm chứng các kỹ thuật cải thiện mô hình dự đoán chi phí thực tế như nén đồ thị, điều chỉnh giá trị padding, bổ sung đặc trưng nút và attention giữa các cấu hình
Vì sao trình biên dịch ML quyết định hiệu năng thực thi
- Các mô hình ML hiện đại thực hiện những tác vụ như hiểu ngôn ngữ tự nhiên, hội thoại, tạo ảnh và tạo video, đồng thời được viết và huấn luyện bằng các framework lập trình ML như TensorFlow, JAX, PyTorch
- Framework cung cấp các phép toán đại số tuyến tính như nhân ma trận, convolution, cũng như các lớp mạng thần kinh như lớp convolution 2D và lớp transformer
- Người dùng không cần trực tiếp xử lý các chi tiết về cách chạy mô hình hiệu quả trên phần cứng, vì trình biên dịch bên dưới framework sẽ tự động tối ưu hóa mô hình
- Tuy vậy, trình biên dịch thường giải các bài toán tối ưu hóa phức tạp bằng heuristic nên không phải lúc nào cũng đạt hiệu năng tối ưu
Đồ thị tính toán và tối ưu hóa hai giai đoạn
- Trình biên dịch ML chuyển các lệnh toán học do người dùng viết thành các lệnh có thể chạy trên phần cứng thực tế
- Chương trình ML có thể được biểu diễn bằng đồ thị tính toán
- Nút biểu diễn các phép toán tensor như matrix multiplication
- Cạnh biểu diễn tensor chảy từ nút này sang nút khác
- Tối ưu hóa trình biên dịch được chia thành hai loại lớn
- Tối ưu hóa cấp độ đồ thị: đưa ra quyết định dựa trên ngữ cảnh của toàn bộ đồ thị và biến đổi cả đồ thị
- Tối ưu hóa cấp độ kernel: biến đổi một kernel là fused subgraph một cách độc lập với các kernel khác
Trade-off hiệu năng của bố cục bộ nhớ
- Tensor 2D như ma trận có thể được lưu trong bộ nhớ ở dạng
[A B C a b c]hoặc[A a B b C c], tương ứng với bố cục row-major và column-major - Một trong những tối ưu hóa quan trọng của trình biên dịch ML là gán bố cục bộ nhớ cho mọi tensor trung gian trong chương trình
- Một bố cục cụ thể có thể là hiệu quả nhất cho từng phép toán riêng lẻ, nhưng nếu bố cục giữa add và convolution không khớp, trình biên dịch sẽ phải chèn thêm phép toán copy
- Ngược lại, dù hiệu năng của từng phép toán riêng có thể thấp hơn đôi chút, cấu hình không cần chuyển đổi bố cục có thể tốt hơn cho toàn bộ quá trình thực thi
- Trong bộ benchmark XLA, khi chọn cấu hình bố cục tối ưu thay vì thiết lập mặc định của trình biên dịch, đã quan sát thấy tăng tốc tới 32%
Bộ dữ liệu TpuGraphs
- TpuGraphs là bộ dữ liệu mô hình chi phí dựa trên học máy cho các chương trình chạy trên TPU tùy biến của Google
- Mục tiêu là huấn luyện một mô hình chi phí nhận đầu vào là chương trình và cấu hình trình biên dịch để dự đoán thời gian thực thi của chương trình
- Bộ dữ liệu nhắm tới hai cấu hình của trình biên dịch XLA
- layout: cấu hình khái quát hóa khái niệm row-major và column-major của ma trận sang tensor nhiều chiều
- tiling: cấu hình kích thước tile
- Mỗi ví dụ bao gồm đồ thị tính toán của workload ML, cấu hình biên dịch và thời gian thực thi khi biên dịch với cấu hình đó
- Các đồ thị được thu thập từ chương trình ML mã nguồn mở và bao gồm các kiến trúc mô hình như ResNet, EfficientNet, Mask R-CNN và Transformer
- Cách tải về và mã khởi đầu được cung cấp tại TpuGraphs GitHub
- TpuGraphs có số lượng đồ thị nhiều gấp 25 lần bộ dữ liệu dự đoán thuộc tính đồ thị lớn nhất trước đây có kích thước đồ thị tương tự, và kích thước đồ thị trung bình lớn gấp 770 lần so với các bộ dữ liệu dự đoán hiệu năng chương trình ML trước đó
Mô hình chi phí cơ sở và kiến trúc GNN
- TpuGraphs đi kèm một mô hình chi phí dựa trên học máy cơ sở, và vì chương trình đầu vào được biểu diễn dưới dạng đồ thị nên nó sử dụng GNN
- Đặc trưng nút gồm hai phần
- opcode id: thông tin nút quan trọng nhất, biểu thị loại phép toán tensor
- Các đặc trưng nút còn lại
- Mô hình cơ sở chuyển opcode id thành opcode embedding thông qua embedding lookup table
- Opcode embedding được kết hợp với các đặc trưng nút còn lại để dùng làm đầu vào cho GNN
- Node embedding do GNN tạo ra được kết hợp thành graph embedding có kích thước cố định bằng các phép graph pooling reduction đơn giản như sum và mean
- Graph embedding cuối cùng được đưa qua feedforward layer để tạo thành một đầu ra scalar
Huấn luyện đồ thị lớn bằng Graph Segment Training
- Graph Segment Training là một kỹ thuật mở rộng huấn luyện GNN để xử lý đồ thị lớn trên các thiết bị có dung lượng bộ nhớ hạn chế
- Phương pháp này nhắm tới bài toán graph-level prediction, nơi đối tượng dự đoán không phải là nút hay cạnh mà là toàn bộ đồ thị
- Đồ thị tính toán có thể chứa hàng trăm nghìn nút, nên Full Graph Training sử dụng toàn bộ đồ thị cùng lúc có thể là điều không khả thi về mặt tính toán
- GST chia đồ thị lớn thành các phân đoạn nhỏ và chỉ chọn một tập con ngẫu nhiên các phân đoạn để cập nhật mô hình
- Các phân đoạn còn lại tạo embedding mà không lưu activation trung gian, từ đó giảm sử dụng bộ nhớ
- Mọi segment embedding được kết hợp để tạo graph embedding của đồ thị lớn ban đầu, rồi dùng cho dự đoán
- Historical embedding table và segment dropout được đưa vào cùng nhau để giảm bớt tình trạng historical embedding bị stale
- Toàn bộ phương pháp giúp rút ngắn thời gian huấn luyện end-to-end xuống 3 lần
Các kỹ thuật cải thiện được kiểm chứng trong cuộc thi Kaggle
- Cuộc thi Kaggle Fast or Slow? Predict AI Model Runtime được tổ chức dựa trên bộ dữ liệu TpuGraphs, với 792 người thuộc 616 đội từ 66 quốc gia tham gia
- Có 10.507 bài nộp, trong đó 153 người lần đầu tham gia Kaggle và 47 người trong số đó lọt vào top 100
- Các đội tham gia đã thử nghiệm nhiều kỹ thuật
- Pruning·compression đồ thị: thay vì GST, họ thử nghiệm cách nén các đồ thị lớn, trong đó có phương pháp chỉ giữ lại subgraph bao gồm các nút có thể cấu hình và các nút láng giềng trực tiếp của chúng
- Thay đổi giá trị padding: vì giá trị padding mặc định 0 xung đột với giá trị feature hợp lệ, việc dùng -1 có thể cải thiện đáng kể độ chính xác của mô hình
- Bổ sung đặc trưng nút·thay đổi mã hóa: các đặc trưng nút bổ sung như contracting dimensions của dot general là quan trọng, và cách mã hóa đặc trưng nút cũng có thể ảnh hưởng đến kết quả
- Cross-configuration attention: đội chiến thắng đã thiết kế một layer đơn giản để mô hình có thể so sánh tường minh giữa các cấu hình, và cho kết quả tốt hơn nhiều so với cách để mô hình suy luận từng cấu hình một cách riêng lẻ
- Kết quả cuộc thi và lời giải chiến thắng dự kiến sẽ được trình bày trong phiên competition của ML for Systems workshop tại NeurIPS vào ngày 16 tháng 12 năm 2023
Phiên liên quan tại NeurIPS Expo
- Dành cho độc giả quan tâm đến nghiên cứu dữ liệu có cấu trúc và trí tuệ nhân tạo, phiên panel của NeurIPS Expo Graph Learning Meets Artificial Intelligence sẽ diễn ra vào ngày 9 tháng 12 năm 2023
- Phiên này sẽ đề cập đến những tiến bộ như mô hình chi phí dựa trên học máy
1 bình luận
Các ý kiến trên Hacker News
Trình biên dịch ML đang bị thổi phồng. Đây là kiểu đánh đổi giống các trình biên dịch truyền thống: bạn có được thông lượng lớn hơn rất nhiều so với việc thuê lập trình viên hiệu năng chuyên nghiệp, nhưng phương án sau thường nhanh hơn nhiều và trong một số trường hợp có thể vượt trước vài bậc độ lớn
Nó thiếu sót ở nhiều tầng. Ở cấp thuật toán, nó không phản hồi lại cho con người những mẹo giúp mạng chạy nhanh hơn, mà chỉ đưa ra vài tín hiệu rất cơ bản. Ý định cũng bị mất đi. Nhà thiết kế mạng ML chỉ định cấu trúc bằng Python, nhưng qua nhiều bước hạ tầng biểu diễn, kết quả có thể trở nên hoàn toàn khác. Gần đây tôi thấy một trình biên dịch khi thực hiện slice update đã tạo mọi khoảng chỉ số có thể của mảng, rồi cắt lấy các chỉ số cần cập nhật và scatter; tôi đã thay bằng một lệnh gọi
memcpyduy nhất. Kernel cũng kém hiệu quả. Mỗi khi đầu ra của các trình biên dịch kiểu này đối đầu với một lập trình viên assembly lành nghề, trình biên dịch đều thua, thường chênh trên 30%. Trông có vẻ dễ giải quyết, nhưng nếu suốt 50 năm qua chưa ai giải được cho tử tế, thì rõ ràng nó không đơn giản như lời nóiTrình biên dịch cũng có nhiều heuristic như inline hóa, unroll vòng lặp, vector hóa, nên mạng nơ-ron có thể hữu ích, và cũng có thể dễ bảo trì hơn rất nhiều heuristic do con người viết
Tối ưu kiểu này vẫn tốt hơn rất nhiều so với không có gì
Có ai giải thích điều này thực tế hơn một chút được không? Tôi tò mò tình trạng thực tế hiện nay của trình biên dịch ML ra sao và trong tương lai gần có thể kỳ vọng điều gì
Chỉ cần viết
model = torch.compile(model). “Trên 163 mô hình mã nguồn mở này, torch.compile hoạt động trong 93% trường hợp và giúp huấn luyện nhanh hơn 43% trên GPU NVIDIA A100. Với độ chính xác Float32, tốc độ trung bình nhanh hơn 21%; với độ chính xác AMP, trung bình nhanh hơn 51%.”[1] Có vẻ Google muốn có thêm nhiều người tham gia R&D về các phương pháp như vậy[1] https://pytorch.org/get-started/pytorch-2.0/
Thực tế thì phức tạp, nên nói đơn giản hóa khá nhiều là: biên dịch đồ thị tính toán sang một dạng biểu diễn trung gian nào đó rồi triển khai backend tương ứng. Các dự án liên quan có thể xem stableHLO, IREE, openXLA. Trình biên dịch
jitcủa Jax cũng có thể xem là một dạng trình biên dịch như vậy. Nó hạ các phép toán được trace xuống XLA, rồi XLA lại làm đủ trò “ma thuật” để chạy trên backend. Rốt cuộc càng đi xuống dưới thì toàn là chuyển đổi và trừu tượng hóatorch.compileTóm lại, đây là công việc cải thiện dự đoán hiệu năng thời gian chạy của đồ thị tính toán bằng mạng nơ-ron đồ thị (GNN). Nó dùng từ điển embedding cho opcode của từng nút cùng với các đặc trưng nút khác như shape, bits, window size ([1])
Họ đã công bố ở [2] một tập dữ liệu đồ thị lớn chứa các thiết lập biên dịch XLA khác nhau và hiệu năng kết quả trên TPU; và ở [3] cải thiện việc dự đoán trên các đồ thị lớn hơn trước bằng cách phân đoạn đồ thị (lần đầu tôi thấy METIS graph partition) cùng nhiều kỹ thuật học khác nhau. Đây là chuyện dự đoán hiệu năng của một đồ thị đã cho, chứ không phải chuyện cải thiện, đề xuất hay chỉnh sửa một đồ thị mới tương đương. Như FunSearch, một mô hình có khả năng dự đoán khá tốt có thể được dùng cùng với tìm kiếm tiến hóa
[1] https://github.com/google-research-datasets/tpu_graphs#featu...
[2] TpuGraphs: A Performance Prediction Dataset on Large Tensor Computational Graphs https://arxiv.org/abs/2308.13490
[3] Learning Large Graph Property Prediction via Graph Segment Training https://arxiv.org/abs/2305.12322
Có ai giải thích được tích chập trong đồ thị đó hoạt động thế nào không? Họ tích chập một kernel có shape
[4,16,8]lên tensor shape[2,4,16]mà lại ra tensor[2,8]; sao có thể vậy?[2,4,16], có thể xem2là kích thước batch,4là chiều đặc trưng đầu vào, và16là chiều kênh đầu vàoTrong kernel
[4,16,8],4là kích thước cửa sổ lọc,16khớp với chiều kênh đầu vào, và8là chiều kênh đầu ra. Trong đầu ra[2,8],2được giữ lại là kích thước batch, còn8khớp với chiều kênh đầu ra của kernel. Nhìn bề ngoài có vẻ không khớp số chiều, nhưng tích chập trên đồ thị tận dụng cấu trúc láng giềng. Kernel di chuyển trên đồ thị, áp trọng số lên nút hiện tại và các đặc trưng của láng giềng trong một bán kính nhất định, rồi gom tổng có trọng số đó để tạo đặc trưng mới cho từng kênh đầu ra. Cấu trúc đồ thị, trọng số cạnh, cũng như các chi tiết triển khai như padding và stride cũng có thể ảnh hưởng đến shape đầu raGemini đang ở tình trạng nào?
Chỉ tính những cái nghĩ ra ngay đã có ít nhất năm mô hình nền tảng như Llama, Claude, Gemini, Falcon, Mistral, chúng liên tục vượt qua rồi bị vượt lại, nhưng GPT vẫn ở một bậc cao hơn và đã như vậy suốt 1 năm. Hóa ra các mô hình ngôn ngữ lớn dựa trên Transformer đủ đơn giản để bất kỳ ai có thể chi khoảng một triệu đô la tiền thời gian GPU đều làm được, nhưng vẫn chưa hoàn toàn bắt kịp OpenAI. Bí quyết đặc biệt của họ là gì?
Bản thân Transformer thì sao? Có manh mối nào cho thấy nó tối ưu theo nghĩa nào đó không?
Cảm giác đoạn đầu tiên đã chôn mất ý chính, nhưng phần còn lại thì rất hay
Tốc độ phát triển ML hiện nay thật đáng kinh ngạc. Tôi không tin vào điểm kỳ dị, nhưng nó đang thay đổi phần mềm và xã hội theo những cách không ai dự đoán được
https://deepmind.google/discover/blog/funsearch-making-new-d...