Tối ưu Huấn luyện Tiền xử lý Mô hình Kiểu Llama 3 trên Amazon SageMaker với TorchTitan

Tối ưu Huấn luyện Tiền xử lý Mô hình Kiểu Llama 3 trên Amazon SageMaker với TorchTitan

Mục lục:

1. Giới thiệu
2. TorchTitan: Kiến trúc tham chiếu cho huấn luyện LLM quy mô lớn 2.1. Chuyển đổi từ FSDP1 sang FSDP2 2.2. Hỗ trợ torch.compile 2.3. Hỗ trợ phép toán tuyến tính FP8 2.4. Hỗ trợ FP8 all-gather
3. Huấn luyện Tiền xử lý Meta Llama 3 8B với TorchTitan trên Amazon SageMaker 3.1. Tổng quan về giải pháp 3.2. Điều kiện tiên quyết 3.3. Xây dựng ảnh tùy chỉnh TorchTitan 3.4. Chuẩn bị bộ dữ liệu (tùy chọn) 3.5. Khởi chạy quá trình huấn luyện với TorchTitan
4. Kết quả hiệu suất
5. Kết luận
6. Về tác giả

1. Giới thiệu

Huấn luyện tiền xử lý (pre-training) là bước đầu tiên và vô cùng quan trọng trong quá trình phát triển các mô hình ngôn ngữ lớn (LLM) mạnh mẽ. Quá trình này giúp mô hình học hỏi các mẫu ngôn ngữ chung, kiến thức thế giới và khả năng lập luận thông qua việc tiếp xúc với một lượng dữ liệu khổng lồ và đa dạng. Bài viết này, hợp tác cùng đội ngũ PyTorch tại Meta, sẽ trình bày cách thư viện torchtitan giúp tăng tốc và đơn giản hóa quá trình huấn luyện tiền xử lý cho các kiến trúc mô hình kiểu Meta Llama 3 trên Amazon SageMaker. Chúng ta sẽ chứng minh hiệu quả của torchtitan thông qua việc huấn luyện tiền xử lý mô hình Meta Llama 3 8B trên các instance p5.48xlarge của Amazon SageMaker, mỗi instance được trang bị 8 GPU Nvidia H100.

2. TorchTitan: Kiến trúc tham chiếu cho huấn luyện LLM quy mô lớn

Torchtitan là một kiến trúc tham chiếu cho huấn luyện LLM quy mô lớn sử dụng PyTorch gốc. Mục tiêu của nó là trình diễn các tính năng huấn luyện phân tán mới nhất của PyTorch trong một cơ sở mã sạch sẽ và tối giản. Thư viện này được thiết kế đơn giản, dễ hiểu, sử dụng và mở rộng cho các mục đích huấn luyện khác nhau, với những thay đổi tối thiểu cần thiết cho mã mô hình khi áp dụng các kỹ thuật xử lý song song.

2.1. Chuyển đổi từ FSDP1 sang FSDP2

FSDP1 và FSDP2 là hai phương pháp huấn luyện song song dữ liệu được chia sẻ đầy đủ. FSDP1 sử dụng phương pháp chia sẻ tham số phẳng, trong khi FSDP2 xử lý từng tham số riêng lẻ, cho phép thao tác linh hoạt hơn và tối ưu hóa hiệu quả hơn.

2.2. Hỗ trợ torch.compile

torch.compile là một tính năng quan trọng trong PyTorch giúp tăng đáng kể hiệu năng mô hình với những thay đổi mã tối thiểu. Torchtitan hỗ trợ torch.compile, mang lại tốc độ đáng kể, đặc biệt là đối với các mô hình lớn và kiến trúc phức tạp.

2.3. Hỗ trợ phép toán tuyến tính FP8

Torchtitan hỗ trợ tính toán FP8 (điểm nổi 8 bit), giúp giảm đáng kể dung lượng bộ nhớ và tăng hiệu năng trong huấn luyện LLM. FP8 có hai định dạng, E4M3 và E5M2, mỗi định dạng được tối ưu cho các khía cạnh khác nhau của huấn luyện. Việc sử dụng FP8 không ảnh hưởng đến độ chính xác của mô hình.

2.4. Hỗ trợ FP8 all-gather

Tính năng này cho phép truyền tensor FP8 hiệu quả giữa nhiều GPU, giảm đáng kể băng thông mạng so với các hoạt động all-gather bfloat16.

3. Huấn luyện Tiền xử lý Meta Llama 3 8B với TorchTitan trên Amazon SageMaker

Amazon SageMaker cung cấp môi trường được quản lý đầy đủ, đơn giản hóa huấn luyện phân tán quy mô lớn trên nhiều instance. SageMaker hỗ trợ container tùy chỉnh, tích hợp liền mạch với thư viện torchtitan và các phụ thuộc của nó.

3.1. Tổng quan về giải pháp

Bài viết hướng dẫn cách chuẩn bị một ảnh tùy chỉnh với thư viện torchtitan, sau đó cấu hình hàm ước lượng công việc huấn luyện để khởi chạy quá trình huấn luyện tiền xử lý mô hình Meta Llama 3 8B với bộ dữ liệu c4 (Colossal Clean Crawled Corpus) trên SageMaker.

3.2. Điều kiện tiên quyết

• Tài khoản AWS
• Miền SageMaker và Amazon SageMaker Studio
• Token truy cập Hugging Face
• Yêu cầu tăng hạn ngạch ít nhất 1 instance ml.p5.48xlarge

3.3. Xây dựng ảnh tùy chỉnh TorchTitan

SageMaker BYOC (Bring Your Own Container) cho phép sử dụng container Docker tùy chỉnh để huấn luyện và triển khai mô hình ML. Bài viết hướng dẫn cách sử dụng Amazon SageMaker Studio Image Build để xây dựng ảnh container tùy chỉnh.

3.4. Chuẩn bị bộ dữ liệu (tùy chọn)

Bài viết cung cấp notebook mẫu để tải xuống bộ dữ liệu allenai/c4 “en” từ Hugging Face về bucket S3.

3.5. Khởi chạy quá trình huấn luyện với TorchTitan

Bài viết hướng dẫn từng bước cách khởi chạy huấn luyện, bao gồm: nhập các mô-đun SageMaker cần thiết, chuẩn bị môi trường huấn luyện, tải xuống bộ phân từ Meta Llama 3, cấu hình tệp TOML cho torchtitan (bật TensorBoard, torch.compile, FP8), thiết lập kênh dữ liệu, tạo trình ước lượng SageMaker PyTorch và bắt đầu huấn luyện.

4. Kết quả hiệu suất

Các kết quả cho thấy sự cải thiện hiệu suất đáng kể khi kết hợp các tối ưu hóa: torch.compile, phép toán tuyến tính FP8 và FP8 all-gather. Hiệu năng tăng lên đến 38.23% so với dòng cơ sở mà không ảnh hưởng đến chất lượng huấn luyện (độ mất mát vẫn ổn định).

5. Kết luận

Bài viết đã chứng minh cách huấn luyện tiền xử lý hiệu quả các mô hình Meta Llama 3 bằng thư viện torchtitan trên SageMaker. Sự kết hợp giữa các tính năng huấn luyện phân tán mới nhất của PyTorch trong torchtitan và khả năng mở rộng của SageMaker cho phép các tổ chức tạo ra các mô hình AI mạnh mẽ và hiệu quả cao.

6. Về tác giả

Bài viết được đồng tác giả bởi Roy Allela, Trevor Harvey, Kanwaljit Khurmi, Less Wright và Wei Feng, những chuyên gia hàng đầu trong lĩnh vực AI/ML và PyTorch. Thông tin chi tiết về tác giả được cung cấp trong bài viết gốc.