bởi Robert Van Dusen, Gautam Kumar, Rahul Huilgol, và Luis Quintela | vào ngày 22 THÁNG 12 năm 2023 |
Huấn luyện mô hình ngôn ngữ lớn (LLM) đã tăng mạnh về phổ biến trong năm qua với việc phát hành một số mô hình phổ biến như Llama 2, Falcon và Mistral. Khách hàng hiện đang tiền xử lý và điều chỉnh LLM từ 1 tỷ đến hơn 175 tỷ tham số để tối ưu hóa hiệu suất mô hình cho các ứng dụng trải rộng từ lĩnh vực chăm sóc sức khỏe đến tài chính và tiếp thị.
Việc huấn luyện mô hình hiệu suất ở quy mô này có thể là một thách thức. LLM chính xác cao có thể yêu cầu tới terabytes dữ liệu huấn luyện và hàng ngàn hoặc thậm chí hàng triệu giờ tính toán trên thiết bị tăng tốc để đạt được độ chính xác mục tiêu. Để hoàn thành quá trình đào tạo và ra mắt sản phẩm đúng thời hạn, khách hàng phải phụ thuộc vào các kỹ thuật song song để phân phối công việc khổng lồ này trên đến hàng nghìn thiết bị tăng tốc. Tuy nhiên, những kỹ thuật song song này có thể khó sử dụng: các kỹ thuật và thư viện khác nhau chỉ tương thích với một số công việc cụ thể hoặc bị hạn chế đối với một số kiến trúc mô hình, hiệu suất đào tạo có thể rất nhạy cảm với cấu hình mơ hồ, và tình trạng nghệ thuật đang nhanh chóng phát triển. Do đó, các chuyên gia máy học phải dành nhiều tuần để mở rộng công việc LLM của họ lên các cụm GPU lớn.
Trong bài viết này, chúng tôi tôn chỉ các tính năng mới của thư viện chia sẻ mô hình Amazon SageMaker giúp đơn giản hóa quá trình huấn luyện mô hình lớn và giúp bạn huấn luyện LLM nhanh hơn. Cụ thể, chúng tôi đề cập đến trải nghiệm người dùng mới đơn giản hóa của thư viện chia sẻ mô hình, xây dựng trên API Fully Sharded Data Parallel (FSDP) nguồn mở của PyTorch, tính năng chia sẻ tensor mở rộng giúp huấn luyện mô hình với hàng trăm tỷ tham số, và tối ưu hiệu suất giảm thời gian và chi phí huấn luyện mô hình lên đến 20%.
Để biết thêm thông tin về thư viện chia sẻ mô hình SageMaker, tham khảo tài liệu về SageMaker model parallelism library v2. Bạn cũng có thể tham khảo các sổ ghi chú ví dụ của chúng tôi để bắt đầu.
Các tính năng mới giúp đơn giản hóa và tăng tốc quá trình huấn luyện mô hình lớn
Bài viết này thảo luận về các tính năng mới được bao gồm trong phiên bản v2.0 của thư viện chia sẻ mô hình SageMaker. Các tính năng này cải thiện khả năng sử dụng của thư viện, mở rộng chức năng và tăng tốc quá trình huấn luyện. Trong các phần tiếp theo, chúng tôi tóm tắt các tính năng mới và thảo luận về cách bạn có thể sử dụng thư viện để tăng tốc quá trình huấn luyện mô hình lớn của mình.
Đồng bộ SMP với PyTorch mã nguồn mở
Kể từ khi ra mắt vào năm 2020, SMP đã cho phép huấn luyện quy mô lớn, hiệu suất cao trên các máy chủ tính toán SageMaker. Với phiên bản SMP mới nhất này, thư viện đơn giản hóa trải nghiệm người dùng bằng cách đồng bộ hóa các API của nó với PyTorch mã nguồn mở.
PyTorch cung cấp Phân bổ Dữ liệu Đa Phân (FSDP) làm phương pháp chính để hỗ trợ công việc huấn luyện lớn trên nhiều thiết bị tính toán. Như được minh họa trong đoạn mã nguồn dưới đây, các API đã được cập nhật của SMP cho các kỹ thuật như phân phối dữ liệu phân mảnh phản ánh những của PyTorch. Bạn chỉ cần chạy import torch.sagemaker và sử dụng nó thay cho torch.
## training_script.py
import torch.sagemaker as tsm
tsm.init()
# Set up a PyTorch model
model = …
# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
model,
…
)
optimizer = …
…
Với những cập nhật này đối với các API của SMP, bạn có thể hiện thực được những lợi ích về hiệu suất của SageMaker và thư viện SMP mà không cần thay đổi hoàn toàn các kịch bản huấn luyện PyTorch FSDP hiện tại của bạn. Mô hình này cũng cho phép bạn sử dụng cùng mã nguồn khi huấn luyện trên cơ sở và trên SageMaker, giúp đơn giản hóa trải nghiệm người dùng đối với khách hàng thực hiện huấn luyện trong nhiều môi trường.
Để biết thêm thông tin về cách bật SMP với các kịch bản huấn luyện PyTorch FSDP hiện tại của bạn, tham khảo Get started with SMP.
Với những cập nhật cho API của SMP, bạn hiện có thể thực hiện được những lợi ích về hiệu suất của SageMaker và thư viện SMP mà không cần phải làm đổi kịch bản huấn luyện PyTorch FSDP hiện tại của bạn. Mô hình này cũng cho phép bạn sử dụng cùng mã nguồn khi huấn luyện trên cơ sở và trên SageMaker, giúp đơn giản hóa trải nghiệm người dùng đối với khách hàng thực hiện huấn luyện trong nhiều môi trường.
Để biết thêm thông tin về cách bật SMP với các kịch bản huấn luyện PyTorch FSDP hiện tại của bạn, tham khảo Get started with SMP.
Kết hợp phân bổ dữ liệu phân mảnh để tăng cường khả năng huấn luyện trên các cụm lớn
Bản phát hành SMP này cũng mở rộng khả năng của PyTorch FSDP để bao gồm các kỹ thuật phân bổ dữ liệu phân mảnh. Một vấn đề khi chỉ sử dụng phân bổ dữ liệu phân mảnh một mình là bạn có thể gặp vấn đề hội tụ khi tăng kích thước cụm máy chủ của mình. Điều này xảy ra vì việc phân mảnh các tham số, gradient và trạng thái tối ưu hóa qua các hạng phân mảnh dữ liệu cũng làm tăng kích thước tổng của bạn; trên các cụm lớn, kích thước tổng này có thể vượt quá ngưỡng dưới đó mô hình sẽ hội tụ. Bạn cần tích hợp một kỹ thuật phân bổ song song bổ sung mà không đòi hỏi tăng kích thước tổng khi mở rộng cụm máy chủ của mình.
Để giảm thiểu vấn đề này, SMP v2.0 giới thiệu khả năng kết hợp phân bổ dữ liệu phân mảnh với phân bổ dữ liệu tensor. Tensor parallelism cho phép tăng kích thước cụm mà không làm thay đổi kích thước tổng toàn cục hoặc ảnh hưởng đến sự hội tụ của mô hình. Với tính năng này, bạn có thể tăng hiệu suất huấn luyện an toàn bằng cách cung cấp các cụm với 256 nút trở lên.
Hiện nay, phân bổ dữ liệu tensor với PyTorch FSDP chỉ có sẵn với SMP v2. SMP v2 cho phép bạn bật tính năng này với vài dòng mã thay đổi và mở khóa việc huấn luyện ổn định ngay cả trên các cụm lớn. SMP v2 tích hợp với Transformer Engine cho việc triển khai phân bổ dữ liệu tensor và làm cho nó tương thích với các API PyTorch FSDP. Bạn có thể kích hoạt PyTorch FSDP và phân bổ dữ liệu tensor SMP đồng thời mà không cần thay đổi bất kỳ điều gì trong mô hình PyTorch hoặc cấu hình PyTorch FSDP. Đoạn mã sau đây mô tả cách thiết lập từ điển cấu hình SMP dưới dạng JSON và thêm mô-đun khởi tạo SMP torch.sagemaker.init(), chấp nhận từ điển cấu hình khi bạn bắt đầu công việc huấn luyện, vào kịch bản huấn luyện của bạn.
Cấu hình SMP như sau:
{
“tensor_parallel_degree”: 8,
“tensor_parallel_seed”: 0
}
In your training script, use the following code:
import torch.sagemaker as tsm
tsm.init()
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_config(..)
model = tsm.transform(model)
Để tìm hiểu thêm về việc sử dụng phân bổ dữ liệu tensor trong SMP, tham khảo phần phân bổ dữ liệu tensor trong tài liệu hướng dẫn của chúng tôi.
Sử dụng các tính năng tiên tiến để tăng tốc quá trình huấn luyện mô hình lên đến 20%
Ngoài việc cho phép huấn luyện phân tán trên các cụm với hàng trăm instance, SMP còn cung cấp các kỹ thuật tối ưu hóa có thể gia tăng quá trình huấn luyện mô hình lên đến 20%. Ở phần này, chúng tôi sẽ nổi bật một số trong số những tối ưu hóa này. Để biết thêm thông tin, tham khảo phần các tính năng cốt lõi trong tài liệu hướng dẫn của chúng tôi.
Phân mảnh kết hợp
Phân mảnh dữ liệu phân mảnh là một kỹ thuật huấn luyện phân tán tiết kiệm bộ nhớ, nó chia state của một mô hình (tham số mô hình, gradient và trạng thái tối ưu hóa) qua các thiết bị. Dung lượng bộ nhớ nhỏ hơn này cho phép bạn đưa một mô hình lớn hơn vào cụm của mình hoặc tăng kích thước batch. Tuy nhiên, phân mảnh dữ liệu phân mảnh cũng làm tăng yêu cầu về giao tiếp của công việc huấn luyện của bạn vì các tệp mô hình phân mảnh thường xuyên được thu thập từ các thiết bị khác nhau trong quá trình huấn luyện. Điều này làm cho mức độ phân mảnh trở thành một cấu hình quan trọng, đặt ra sự đánh đổi giữa việc tiêu thụ bộ nhớ và chi phí giao tiếp.
Mặc định, PyTorch FSDP phân mảnh các tệp mô hình trên tất cả các thiết bị gia tốc trong cụm của bạn. Tùy thuộc vào công việc huấn luyện của bạn, phương pháp phân mảnh này có thể làm tăng chi phí giao tiếp và tạo ra một điểm bottleneck. Để hỗ trợ trong việc này, thư viện SMP cung cấp phân mảnh dữ liệu phân mảnh kết hợp có thể cấu hình trên cơ sở PyTorch FSDP. Tính năng này cho phép bạn đặt mức độ phân mảnh tối ưu cho công việc huấn luyện của bạn. Chỉ cần chỉ định mức độ phân mảnh trong một đối tượng cấu hình JSON và bao gồm nó trong kịch bản huấn luyện SMP của bạn.
Cấu hình SMP như sau:
{ “hybrid_shard_degree”: 16 }
Để tìm hiểu thêm về ưu điểm của phân mảnh dữ liệu phân mảnh kết hợp, tham khảo bài viết Near-linear scaling of gigantic-model training on AWS. Đối với thông tin chi tiết về cách triển khai phân mảnh kết hợp với kịch bản huấn luyện FSDP hiện tại của bạn, xem phân mảnh dữ liệu phân mảnh kết hợp trong tài liệu hướng dẫn của chúng tôi.
Để tăng tốc quá trình huấn luyện phân tán của bạn, bạn có thể sử dụng thư viện SMP cùng với thư viện SageMaker distributed data parallelism (SMDDP). SMDDP bao gồm một toán tử giao tiếp tổng hợp AllGather được tối ưu hóa để đạt được hiệu suất tốt nhất trên các instance được tăng tốc p4d và p4de của SageMaker. Trong huấn luyện phân tán, các toán tử giao tiếp tổng hợp được sử dụng để đồng bộ thông tin giữa các GPU worker. AllGather là một trong những toán tử giao tiếp tổng hợp cốt lõi thường được sử dụng trong phân mảnh dữ liệu phân mảnh để thể hiện các tham số của lớp trước các bước tính toán tiến và lùi. Đối với các công việc huấn luyện mà gặp bottleneck từ giao tiếp, các hoạt động giao tiếp tổng hợp nhanh hơn có thể giảm thời gian và chi phí huấn luyện mà không ảnh hưởng đến sự hội tụ.
Để sử dụng thư viện SMDDP, bạn chỉ cần thêm hai dòng mã vào kịch bản huấn luyện của mình:
import torch.distributed as dist
# Initialize with SMDDP
import smdistributed.dataparallel.torch.torch_smddp
dist.init_process_group(backend=”smddp”) # Replacing “nccl”
# Initialize with SMP
import torch.sagemaker as tsm
tsm.init()
Ngoài SMP, SMDDP hỗ trợ PyTorch FSDP mã nguồn mở và DeepSpeed. Để biết thêm về thư viện SMDDP, hãy xem Chạy huấn luyện phân tán với thư viện phân tán dữ liệu của SageMaker.
Típicamente, qua trình đi qua (forward pass) của việc huấn luyện mô hình tính các kích hoạt tại mỗi lớp và giữ chúng trong bộ nhớ GPU cho đến khi qua trình đi lại (backward pass) cho lớp tương ứng kết thúc. Những kích hoạt đã được lưu này có thể tiêu thụ bộ nhớ GPU đáng kể trong suốt quá trình huấn luyện. Offloading kích hoạt là một kỹ thuật thay thế chuyển những tensor này vào bộ nhớ CPU sau qua trình đi qua và sau đó đưa chúng trở lại GPU khi chúng cần thiết. Phương pháp này có thể giảm đáng kể việc sử dụng bộ nhớ GPU trong suốt quá trình huấn luyện.
Mặc dù PyTorch hỗ trợ offloading kích hoạt, việc triển khai của nó không hiệu quả và có thể làm cho GPU bận rộn trong khi kích hoạt được đưa trở lại từ CPU trong qua trình đi lại. Điều này có thể dẫn đến sự suy giảm đáng kể về hiệu suất khi sử dụng offloading kích hoạt.
SMP v2 cung cấp một thuật toán offloading kích hoạt được tối ưu hóa có thể cải thiện hiệu suất huấn luyện. Việc triển khai của SMP tiên đoán trước kích hoạt trước khi chúng cần thiết trên GPU, giảm thời gian không hoạt động.
Vì SMP được xây dựng trên các API của PyTorch, việc kích hoạt offloading được tối ưu hóa chỉ cần một số dòng mã thay đổi. Chỉ cần thêm các cấu hình liên quan (các tham số sm_activation_offloading và activation_loading_horizon) và bao gồm chúng trong kịch bản huấn luyện của bạn.
Cấu hình SMP như sau:
{
“activation_loading_horizon”: 2,
“sm_activation_offloading”: True
}
Trong kịch bản huấn luyện, sử dụng đoạn mã sau:
import torch.sagemaker as tsm
tsm.init()
# Native PyTorch module for activation offloading
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
apply_activation_checkpointing,
offload_wrapper,
)
model = FSDP(…)
# Activation offloading requires activation checkpointing.
apply_activation_checkpointing(
model,
check_fn=checkpoint_tformer_layers_policy,
)
model = offload_wrapper(model)
Để tìm hiểu thêm về các công cụ kiểm tra cho offloading kích hoạt PyTorch mã nguồn mở, hãy xem kịch bản checkpoint_wrapper.py trong kho GitHub PyTorch và bài đăng blog Activation Checkpointing trong PyTorch về Scaling Multimodal Foundation Models in TorchMultimodal with Pytorch Distributed. Để biết thêm về cách SMP tối ưu hóa triển khai offloading kích hoạt, hãy xem phần offloading kích hoạt trong tài liệu của chúng tôi.
Ngoài hybrid sharding, SMDDP và offloading kích hoạt, SMP còn cung cấp các tối ưu hóa khác có thể tăng tốc công việc huấn luyện mô hình lớn của bạn. Điều này bao gồm kiểm tra kích hoạt được tối ưu hóa, khởi tạo tham số trễ, và những tính năng khác. Để tìm hiểu thêm, hãy tham khảo phần các tính năng cốt lõi của tài liệu của chúng tôi.
Kết luận
Khi bộ dữ liệu, kích thước mô hình và các cụm huấn luyện tiếp tục phát triển, huấn luyện phân tán hiệu quả ngày càng trở nên quan trọng để giao hàng mô hình và sản phẩm đúng hẹn và giá cả phải chăng. Bản phát hành mới nhất của thư viện SageMaker model parallel giúp bạn đạt được điều này bằng cách giảm thiểu sự thay đổi mã và làm cho nó phù hợp với API PyTorch FSDP, cho phép huấn luyện trên các cụm lớn thông qua tensor parallelism và các tối ưu hóa có thể giảm thời gian huấn luyện lên đến 20%.
Để bắt đầu với SMP v2, tham khảo tài liệu của chúng tôi và các kịch bản mẫu của chúng tôi.
About the Authors
Robert Van Dusen là một Quản lý Sản phẩm Cao cấp với Amazon SageMaker. Ông dẫn đầu về frameworks, trình biên dịch và các kỹ thuật tối ưu hóa cho huấn luyện học sâu.
AWS Study Group 