Chạy PyTorch Lightning và PyTorch DDP gốc trên Amazon SageMaker Training, có Amazon Search

Quá nhiều dữ liệu, quá ít thời gian. Các chuyên gia học máy (ML), nhà khoa học dữ liệu, kỹ sư và những người đam mê đã gặp phải vấn đề này trên toàn thế giới. Từ xử lý ngôn ngữ tự nhiên đến thị giác máy tính, bảng biểu đến chuỗi thời gian và mọi thứ ở giữa, vấn đề lâu đời là tối ưu hóa tốc độ khi chạy dữ liệu với càng nhiều GPU càng tốt đã truyền cảm hứng cho vô số giải pháp. Hôm nay, chúng tôi vui mừng thông báo các tính năng dành cho các nhà phát triển PyTorch sử dụng các khuôn khổ nguồn mở gốc, như Tia chớp PyTorch và PyTorch DDP, điều đó sẽ hợp lý hóa đường dẫn của họ đến đám mây.

Amazon SageMaker là một dịch vụ được quản lý hoàn toàn cho ML và đào tạo mô hình SageMaker là một môi trường máy tính được tối ưu hóa để đào tạo hiệu suất cao trên quy mô lớn. Đào tạo mô hình SageMaker cung cấp trải nghiệm đào tạo từ xa với một mặt phẳng điều khiển liền mạch để dễ dàng đào tạo và tái tạo mô hình ML với hiệu suất cao và chi phí thấp. Chúng tôi rất vui mừng được thông báo về các tính năng mới trong danh mục đào tạo SageMaker giúp việc chạy PyTorch trên quy mô lớn thậm chí còn dễ dàng và dễ tiếp cận hơn:

1. PyTorch Lightning hiện có thể được tích hợp để Thư viện song song dữ liệu phân tán của SageMaker chỉ với một dòng mã thay đổi.
2. Đào tạo mô hình SageMaker hiện đã hỗ trợ cho dữ liệu phân tán PyTorch gốc Song song với phụ trợ NCCL, cho phép các nhà phát triển chuyển sang SageMaker dễ dàng hơn bao giờ hết.

Trong bài đăng này, chúng tôi thảo luận về các tính năng mới này và cũng tìm hiểu cách Amazon Search đã chạy PyTorch Lightning với chương trình phụ trợ đào tạo phân tán được tối ưu hóa trong SageMaker để tăng tốc thời gian đào tạo mô hình.

Trước khi đi sâu vào nghiên cứu điển hình về Tìm kiếm của Amazon, đối với những người không quen thuộc, chúng tôi muốn cung cấp một số thông tin cơ bản về Thư viện song song dữ liệu phân tán của SageMaker. Vào năm 2020, chúng tôi đã phát triển và ra mắt cấu hình cụm tùy chỉnh cho phân tán độ dốc xuống theo quy mô nhằm tăng hiệu quả tổng thể của cụm, được giới thiệu trên Amazon Science như Cá trích. Sử dụng tốt nhất của cả máy chủ tham số và cấu trúc liên kết dựa trên vòng, SageMaker Phân tán Dữ liệu Song song (SMDDP) được tối ưu hóa cho Đám mây điện toán đàn hồi Amazon Cấu trúc liên kết mạng (Amazon EC2), bao gồm EFA. Đối với các kích thước cụm lớn hơn, SMDDP có thể cải thiện thông lượng 20–40% so với Horovod (TensorFlow) và PyTorch Phân phối dữ liệu song song. Đối với các kích thước cụm nhỏ hơn và các mô hình được hỗ trợ, chúng tôi khuyên bạn nên Trình biên dịch đào tạo SageMaker, có thể giảm thời gian công việc tổng thể lên đến 50%.

Tiêu điểm của khách hàng: PyTorch Lightning trên chương trình phụ trợ được tối ưu hóa của SageMaker với Amazon Search

Amazon Search chịu trách nhiệm về trải nghiệm tìm kiếm và khám phá trên Amazon.com. Nó cung cấp trải nghiệm tìm kiếm cho khách hàng đang tìm kiếm sản phẩm để mua trên Amazon. Ở cấp độ cao, Amazon Search xây dựng chỉ mục cho tất cả các sản phẩm được bán trên Amazon.com. Khi khách hàng nhập một truy vấn, Amazon Search sau đó sẽ sử dụng nhiều kỹ thuật ML khác nhau, bao gồm cả các mô hình học sâu, để đối sánh các sản phẩm có liên quan và thú vị với truy vấn của khách hàng. Sau đó, nó xếp hạng các sản phẩm trước khi hiển thị kết quả cho khách hàng.

Các nhà khoa học của Amazon Search đã sử dụng PyTorch Lightning làm một trong những khuôn khổ chính để đào tạo các mô hình học sâu hỗ trợ xếp hạng Tìm kiếm do các tính năng hữu dụng được bổ sung trên PyTorch. SMDDP không được hỗ trợ cho các mô hình học sâu được viết bằng PyTorch Lightning trước khi ra mắt SageMaker mới này. Điều này đã ngăn cản các nhà khoa học của Amazon Search, những người thích sử dụng PyTorch Lightning mở rộng quy mô đào tạo mô hình của họ bằng cách sử dụng các kỹ thuật song song dữ liệu, làm chậm đáng kể thời gian đào tạo của họ và ngăn họ thử nghiệm các thử nghiệm mới yêu cầu đào tạo có khả năng mở rộng hơn.

Kết quả đo điểm chuẩn ban đầu của nhóm cho thấy thời gian đào tạo nhanh hơn 7.3 lần đối với mô hình mẫu khi được đào tạo trên tám nút so với đường cơ sở đào tạo một nút. Mô hình cơ sở được sử dụng trong các phép đo điểm chuẩn này là một mạng nơron perceptron nhiều lớp với bảy lớp được kết nối đầy đủ dày đặc và hơn 200 tham số. Bảng sau đây tóm tắt kết quả đo điểm chuẩn trên các trường hợp đào tạo SageMaker trong ml.p3.16xlarge.

Tiếp theo, chúng tôi đi sâu vào chi tiết về các phiên bản mới ra mắt. Nếu bạn thích, bạn có thể bước qua sổ tay ví dụ.

Chạy PyTorch Lightning với thư viện đào tạo được phân phối của SageMaker

Chúng tôi vui mừng thông báo rằng SageMaker Data Parallel hiện đã tích hợp liền mạch với PyTorch Lightning trong quá trình đào tạo SageMaker.

PyTorch Lightning là một khung công tác mã nguồn mở cung cấp sự đơn giản hóa cho việc viết các mô hình tùy chỉnh trong PyTorch. Theo một số cách tương tự như những gì Keras đã làm cho TensorFlow, hoặc thậm chí được cho là Ôm khuôn mặt, PyTorch Lightning cung cấp một API cấp cao với các nội dung trừu tượng cho phần lớn chức năng cấp thấp hơn của chính PyTorch. Điều này bao gồm xác định mô hình, lập hồ sơ, đánh giá, cắt tỉa, song song mô hình, cấu hình siêu tham số, học chuyển giao, v.v.

Trước đây, các nhà phát triển PyTorch Lightning không chắc chắn về cách di chuyển liền mạch mã đào tạo của họ sang các cụm GPU SageMaker hiệu suất cao. Ngoài ra, không có cách nào để họ tận dụng lợi thế về hiệu quả được giới thiệu bởi SageMaker Data Parallel.

Nói chung, đối với PyTorch Lightning, sẽ có những thay đổi mã từ ít đến nhiều để chỉ cần chạy các API này trên SageMaker Training. Trong các sổ ghi chép ví dụ, chúng tôi sử dụng chiến lược DDP và DDPPplugin phương pháp.

Có ba bước để sử dụng PyTorch Lightning với SageMaker Data Parallel làm chương trình phụ trợ được tối ưu hóa:

1. Sử dụng một được hỗ trợ Vùng chứa học tập sâu AWS (DLC) làm hình ảnh cơ sở của bạn hoặc tùy chọn tạo vùng chứa của riêng bạn và tự cài đặt chương trình phụ trợ song song dữ liệu SageMaker. Đảm bảo rằng bạn có PyTorch Lightning được bao gồm trong các gói cần thiết của bạn, chẳng hạn như requirements.txt tập tin.
2. Thực hiện một vài thay đổi mã nhỏ đối với tập lệnh đào tạo của bạn để kích hoạt chương trình phụ trợ được tối ưu hóa. Bao gồm các:
   1. Nhập thư viện SM DDP:

      import smdistributed.dataparallel.torch.torch_smddp
      

   2. Thiết lập môi trường PyTorch Lightning cho SageMaker:

      from pytorch_lightning.plugins.environments.lightning_environment 
        import LightningEnvironment
      
      env = LightningEnvironment()
      env.world_size = lambda: int(os.environ["WORLD_SIZE"])
      env.global_rank = lambda: int(os.environ["RANK"])

   3. Nếu bạn đang sử dụng phiên bản PyTorch Lightning cũ hơn 1.5.10, bạn sẽ cần thêm một vài bước nữa.
      1. Đầu tiên, hãy thêm biến môi trường:

         os.environ["PL_TORCH_DISTRIBUTED_BACKEND"] = "smddp"

      2. Thứ hai, đảm bảo bạn sử dụng DDPPlugin, Chứ không phải là DDPStrategy. Nếu bạn đang sử dụng phiên bản mới hơn, bạn có thể dễ dàng đặt phiên bản này bằng cách đặt requirements.txt trong source_dir cho công việc của bạn, thì điều này không cần thiết. Xem đoạn mã sau:

         ddp = DDPPlugin(parallel_devices=[torch.device("cuda", d) for d in range(num_gpus)], cluster_environment=env)

   4. Theo tùy chọn, xác định phụ trợ nhóm quy trình của bạn là "smddp" trong DDPSTrategy sự vật. Tuy nhiên, nếu bạn đang sử dụng PyTorch Lightning với phần phụ trợ PyTorch DDP, cũng được hỗ trợ, chỉ cần xóa phần này 'process_group_backendtham số `. Xem đoạn mã sau:

      ddp = DDPStrategy(
        cluster_environment=env, 
        process_group_backend="smddp", 
        accelerator="gpu")

3. Đảm bảo rằng bạn đã lưu ý phương pháp phân phối trong công cụ ước tính, chẳng hạn như distribution={"smdistributed":{"dataparallel":{"enabled":True} nếu bạn đang sử dụng phần phụ trợ Herring, hoặc distribution={"pytorchddp":{"enabled":True}.

• Để có danh sách đầy đủ các thông số phù hợp trong distribution tham số, xem tài liệu của chúng tôi tại đây.

Bây giờ bạn có thể bắt đầu công việc đào tạo SageMaker của mình! Bạn có thể khởi chạy công việc đào tạo của mình thông qua Python SDK, Boto3, bảng điều khiển SageMaker, Giao diện dòng lệnh AWS (AWS CLI), và vô số phương pháp khác. Từ góc độ AWS, đây là một lệnh API duy nhất: create-training-job. Cho dù bạn khởi chạy lệnh này từ thiết bị đầu cuối cục bộ của mình, AWS Lambda chức năng, một Xưởng sản xuất Amazon SageMaker sổ ghi chép, đường dẫn KubeFlow hoặc bất kỳ môi trường máy tính nào khác là hoàn toàn tùy thuộc vào bạn.

Xin lưu ý rằng tích hợp giữa PyTorch Lightning và SageMaker Data Parallel hiện chỉ được hỗ trợ cho các phiên bản mới hơn của PyTorch, bắt đầu từ 1.11. Ngoài ra, bản phát hành này chỉ có sẵn trong AWS DLC cho SageMaker bắt đầu từ PyTorch 1.12. Hãy chắc chắn rằng bạn chỉ vào hình ảnh này làm cơ sở của bạn. Trong us-east-1, địa chỉ này như sau:

ecr_image = '763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:1.12.0-gpu-py38-cu113-ubuntu20.04-sagemaker'

Sau đó, bạn có thể mở rộng vùng chứa Docker của mình bằng cách sử dụng nó làm hình ảnh cơ sở của bạn hoặc bạn có thể chuyển nó dưới dạng một biến vào image_uri đối số của Công cụ ước tính đào tạo SageMaker.

Do đó, bạn sẽ có thể chạy mã PyTorch Lightning của mình trên các GPU được tối ưu hóa của SageMaker Training, với hiệu suất tốt nhất hiện có trên AWS.

Chạy song song dữ liệu phân tán PyTorch trên SageMaker

Vấn đề lớn nhất Dữ liệu phân tán PyTorch song song (DDP) giải quyết được đánh lừa là đơn giản: tốc độ. Một khung đào tạo phân tán tốt phải cung cấp sự ổn định, độ tin cậy và quan trọng nhất là hiệu suất tuyệt vời trên quy mô lớn. PyTorch DDP thực hiện điều này thông qua việc cung cấp cho các nhà phát triển ngọn đuốc các API để sao chép các mô hình của họ trên nhiều thiết bị GPU, trong cả cài đặt một nút và nhiều nút. Sau đó, khung công tác quản lý việc phân loại các đối tượng khác nhau từ tập dữ liệu đào tạo thành từng bản sao mô hình, lấy trung bình các độ dốc cho mỗi bản sao mô hình để đồng bộ hóa chúng ở mỗi bước. Điều này tạo ra một mô hình khi hoàn thành toàn bộ quá trình đào tạo. Sơ đồ sau đây minh họa quá trình này.

PyTorch DDP thường gặp trong các dự án sử dụng bộ dữ liệu lớn. Kích thước chính xác của mỗi tập dữ liệu sẽ rất khác nhau, nhưng nguyên tắc chung là chia tỷ lệ tập dữ liệu, tính toán kích thước và kích thước mô hình theo các tỷ lệ tương tự. Còn được gọi là luật nhân rộng, sự kết hợp tối ưu của ba điều này còn rất nhiều tranh luận và sẽ khác nhau tùy theo ứng dụng. Tại AWS, dựa trên việc làm việc với nhiều khách hàng, chúng tôi có thể thấy rõ lợi ích từ chiến lược song song dữ liệu khi kích thước tập dữ liệu tổng thể tối thiểu là vài chục GB. Khi các tập dữ liệu thậm chí còn lớn hơn, việc triển khai một số loại chiến lược song song dữ liệu là một kỹ thuật quan trọng để tăng tốc thử nghiệm tổng thể và cải thiện thời gian của bạn để nâng cao giá trị.

Trước đây, những khách hàng đang sử dụng PyTorch DDP để đào tạo phân tán tại cơ sở hoặc trong các môi trường máy tính khác thiếu khuôn khổ để dễ dàng di chuyển dự án của họ sang Đào tạo SageMaker để tận dụng GPU hiệu suất cao với mặt phẳng điều khiển liền mạch. Cụ thể, họ cần phải di chuyển khung dữ liệu song song của mình sang SMDDP hoặc phát triển và kiểm tra các khả năng của PyTorch DDP trên SageMaker Training theo cách thủ công. Hôm nay, SageMaker Training rất vui được cung cấp trải nghiệm liền mạch cho khách hàng khi giới thiệu mã PyTorch DDP của họ.

Để sử dụng điều này một cách hiệu quả, bạn không cần thực hiện bất kỳ thay đổi nào đối với các tập lệnh đào tạo của mình.

Bạn có thể thấy thông số mới này trong đoạn mã sau. bên trong distribution tham số, chỉ cần thêm pytorchddp và đặt kích hoạt là true.

estimator = PyTorch(
    base_job_name="pytorch-dataparallel-mnist",
    source_dir="code",
    entry_point = "my_model.py",
    ... 
    # Training using SMDataParallel Distributed Training Framework
    distribution = {"pytorchddp": {"enabled": "true"}}
)

Cấu hình mới này bắt đầu từ phiên bản SageMaker Python SDK 2.102.0 và PyTorch DLC's 1.11.

Đối với các nhà phát triển PyTorch DDP, những người đã quen thuộc với ngọn đuốc , thật hữu ích khi biết rằng điều này không cần thiết trong môi trường đào tạo SageMaker, môi trường đã cung cấp khả năng chịu lỗi mạnh mẽ. Tuy nhiên, để giảm thiểu việc viết lại mã, bạn có thể sử dụng một tập lệnh trình khởi chạy khác chạy lệnh này làm điểm nhập của mình.

Giờ đây, các nhà phát triển PyTorch có thể dễ dàng di chuyển các tập lệnh của họ vào SageMaker, đảm bảo các tập lệnh và vùng chứa của họ có thể chạy liền mạch trên nhiều môi trường máy tính.

Điều này chuẩn bị cho họ, trong tương lai, tận dụng lợi thế của Thư viện đào tạo phân tán của SageMaker cung cấp cấu trúc liên kết đào tạo được tối ưu hóa AWS để cung cấp các cải tiến tăng tốc lên đến 40%. Đối với các nhà phát triển PyTorch, đây là một dòng mã duy nhất! Đối với mã PyTorch DDP, bạn có thể chỉ cần đặt phần phụ trợ thành smddp trong quá trình khởi tạo (xem Sửa đổi tập lệnh đào tạo PyTorch), như được hiển thị trong đoạn mã sau:

import smdistributed.dataparallel.torch.torch_smddp
import torch.distributed as dist
dist.init_process_group(backend='smddp')

Như chúng ta đã thấy ở trên, bạn cũng có thể đặt phần phụ trợ của DDPStrategy đến smddp khi sử dụng Lightning. Điều này có thể dẫn đến Tăng tốc tổng thể 40% cho các cụm lớn! Để tìm hiểu thêm về đào tạo phân tán trên SageMaker, hãy xem hội thảo trên web theo yêu cầu của chúng tôi, hỗ trợ máy tính xách tay, liên quan, thích hợp tài liệu hướng dẫnvà giấy tờ.

Kết luận

Trong bài đăng này, chúng tôi đã giới thiệu hai tính năng mới trong gia đình Đào tạo SageMaker. Những điều này giúp các nhà phát triển PyTorch sử dụng mã hiện có của họ trên SageMaker, cả PyTorch DDP và PyTorch Lightning dễ dàng hơn nhiều.

Chúng tôi cũng cho thấy cách Amazon Search sử dụng SageMaker Training để đào tạo mô hình học sâu của họ, và cụ thể là PyTorch Lightning với thư viện tập thể được tối ưu hóa SageMaker Data Parallel làm chương trình phụ trợ. Chuyển sang đào tạo phân tán tổng thể đã giúp Amazon Search đạt được thời gian đào tạo nhanh hơn gấp 7.3 lần.

Giới thiệu về tác giả

Emily Webber đã tham gia AWS ngay sau khi SageMaker ra mắt và kể từ đó đã cố gắng nói cho cả thế giới biết về điều đó! Ngoài việc xây dựng trải nghiệm ML mới cho khách hàng, Emily thích thiền định và nghiên cứu Phật giáo Tây Tạng.

Karan Diman là Kỹ sư phát triển phần mềm tại AWS, có trụ sở tại Toronto, Canada. Anh ấy rất đam mê không gian Học máy và xây dựng các giải pháp để tăng tốc khối lượng công việc tính toán phân tán.

Vishwa Karia là Kỹ sư phát triển phần mềm tại AWS Deep Engine. Sở thích của cô nằm ở điểm giao nhau giữa Học máy và Hệ thống phân tán, đồng thời cô cũng đam mê trao quyền cho phụ nữ trong lĩnh vực công nghệ và AI.

Eiman Elnahrawy là Kỹ sư phần mềm chính tại Amazon Search, dẫn đầu các nỗ lực về tăng tốc, mở rộng quy mô và tự động hóa Học máy. Chuyên môn của cô trải dài trên nhiều lĩnh vực, bao gồm Học máy, Hệ thống phân tán và Cá nhân hóa.