Đào tạo các mô hình khổng lồ với tỷ lệ gần tuyến tính bằng cách sử dụng song song dữ liệu phân đoạn trên Amazon SageMaker

Để theo đuổi độ chính xác vượt trội, các mô hình học sâu trong các lĩnh vực như xử lý ngôn ngữ tự nhiên và thị giác máy tính đã phát triển đáng kể về quy mô trong vài năm qua, thường được tính bằng hàng chục đến hàng trăm tỷ tham số. Việc đào tạo các mô hình khổng lồ này là một thách thức và đòi hỏi các chiến lược phân phối phức tạp. Các nhà khoa học dữ liệu và kỹ sư học máy liên tục tìm kiếm cách tốt nhất để tối ưu hóa máy tính đào tạo của họ, nhưng họ đang vật lộn với chi phí giao tiếp có thể tăng cùng với kích thước cụm tổng thể.

Đây là lý do tại sao gần đây chúng tôi đã khởi chạy ssong song dữ liệu cứng on Amazon SageMaker, một kỹ thuật đào tạo phân tán tiết kiệm bộ nhớ mới trong Thư viện mô hình song song (SMP) của SageMaker. Tính song song dữ liệu phân mảnh được xây dựng có mục đích cho các mô hình quy mô cực lớn và sử dụng Amazon nội bộ MiCS công nghệ ẩn, một nỗ lực khoa học nhằm giảm thiểu quy mô truyền thông bằng cách giảm chi phí truyền thông đắt đỏ bắt nguồn từ việc thu thập tham số và đồng bộ hóa gradient. Với mẫu GPT-30 tham số 2B có độ dài chuỗi 2048, tính năng mới này đạt được 141 TFLOP, tăng 39.7% tốc độ so với DeepSpeed ​​ZeRO-3. Đối với mô hình 10B GPT-2 với độ dài chuỗi 512, tính năng mới này cũng đạt được 564 mẫu mỗi giây, tăng 13.9% so với PyTorch's Full Sharded Data Parallel (FSDP). Hãy nhớ rằng trong đào tạo mô hình khổng lồ, mỗi tỷ lệ phần trăm tăng tốc chuyển thành số tiền tiết kiệm được và năng suất đạt được trong nhóm của bạn.

Trong bài đăng trên blog này, trước tiên chúng ta sẽ xem xét kỹ hơn các điểm khác biệt chính của tính song song dữ liệu phân mảnh và khi nào sử dụng nó. Sau đó, bạn sẽ học cách huấn luyện mô hình GPT-30 tham số 2B trên SageMaker một cách dễ dàng với tính năng mới này. Cuối cùng, chúng tôi sẽ so sánh hiệu suất với các tùy chọn mã nguồn mở khác, đáng chú ý là vượt trội hơn DeepSpeed ​​ZeRO tới 39.7% trên 256 GPU.

Cách thức hoạt động của tính năng song song dữ liệu phân đoạn và khi nào sử dụng nó

Trước khi chúng tôi giới thiệu song song dữ liệu phân mảnh, chúng ta hãy xem xét họ kỹ thuật rộng hơn của nó. Các cách tiếp cận đào tạo phân tán gần đây cho các mô hình lớn đã chuyển sang một mô hình trong đó các thông số mô hình, độ dốc và trạng thái trình tối ưu hóa được chia sẻ qua các nút song song dữ liệu. Không giống như Pipeline Parallelism có sự phức tạp bẩm sinh của việc chọn các lớp để phân vùng trên các thiết bị, đặc biệt là khi khung của bạn không hỗ trợ tách mô hình tự động, mô hình này bảo tồn một cách trang nhã sự đơn giản của tính song song dữ liệu, đồng thời loại bỏ ràng buộc của tính song song dữ liệu trong đó một mô hình phải vừa với một GPU duy nhất.

Trong các khuôn khổ hiện có thuộc mô hình này, đặc biệt là DeepSpeed ​​ZeRO-3 và FSDP của PyTorch được ngược dòng từ FairScale, các trạng thái mô hình được chia nhỏ theo tất cả các GPU, một chiến lược làm giảm mức tiêu thụ bộ nhớ trên mỗi GPU với chi phí phát sinh chi phí truyền thông lớn, tăng theo kích thước cụm và do đó làm cho khả năng mở rộng giảm đáng kể trên quy mô. Ngược lại, dữ liệu phân đoạn song song trong mô hình phân vùng thư viện SMP trạng thái nhận biết quy mô bằng cách phân vùng từng bản sao của các trạng thái mô hình chỉ trong một tập hợp con của GPU.

Hãy xem xét kỹ hơn phân vùng mô hình nhận biết quy mô trong MiCS, công nghệ cốt lõi đằng sau dữ liệu phân mảnh song song. Trực giác đằng sau thiết kế này là các trạng thái huấn luyện phân vùng trên toàn bộ nhóm song song dữ liệu có thể không cần thiết để huấn luyện một mô hình với hàng chục tỷ tham số. Ví dụ: 8 GPU V100 (mỗi GPU 32GB) đủ để giữ trạng thái mô hình bản sao của mô hình thông số 10B, cần khoảng 200 GB bộ nhớ khi đào tạo với trình tối ưu hóa Adam sử dụng độ chính xác hỗn hợp. Bằng cách giới hạn một bản sao hoàn chỉnh của các trạng thái mô hình trong nhỏ nhất tập hợp con GPU, chúng tôi có thể giảm quy mô tổng chi phí giao tiếp một cách hiệu quả so với DeepSpeed ​​và PyTorch FSDP. Dữ liệu được chia nhỏ song song cũng thúc đẩy các kỹ thuật khác trong MiCS như Giao tiếp phân cấp và Đồng bộ hóa Gradient 2 bước. Để biết thêm thông tin, hãy kiểm tra Mở rộng quy mô gần tuyến tính của đào tạo mô hình khổng lồ trên AWS or MiCS: Tỷ lệ gần tuyến tính để đào tạo mô hình khổng lồ trên đám mây công cộng.

Bây giờ, làm thế nào để bạn biết khi nào nên chọn dữ liệu phân mảnh song song với các kỹ thuật đào tạo phân tán khác? Nguyên tắc chung là nếu mô hình của bạn có ít hơn 1 tỷ thông số và có thể vừa với bộ nhớ GPU, Thư viện song song dữ liệu SageMaker or Trình biên dịch đào tạo SageMaker có thể đủ cho bạn. Nếu bạn có các mô hình thị giác máy tính hoặc ngôn ngữ lớn hơn, đề xuất của chúng tôi là đào tạo nó bằng kỹ thuật phân đoạn dữ liệu song song kết hợp với kiểm tra kích hoạt và kích hoạt giảm tải trong thư viện song song mô hình SageMaker trước, trước khi các kỹ thuật khác như tính song song tensor hoặc song song đường ống.

Sử dụng song song dữ liệu phân đoạn để đào tạo GPT-2 trên Amazon SageMaker

Bây giờ chúng ta hãy tìm hiểu cách đào tạo mô hình GPT-2 với dữ liệu phân đoạn song song, với SMP sẽ gói gọn sự phức tạp cho bạn. Đây hoàn thành sổ tay hướng dẫn hướng dẫn bạn toàn bộ quy trình, từ xử lý dữ liệu, xác định và gửi các công việc đào tạo đến theo dõi nhật ký đào tạo. Sau đây là tổng quan ngắn gọn nêu bật các bước chính để sử dụng tính năng này.

1. Bắt đầu

Tính song song dữ liệu được chia nhỏ có sẵn trong PyTorch v1.12.0 + và hoạt động với cả FP16 và BF16. Cách dễ nhất để sử dụng thư viện SMP là thông qua Bộ chứa AWS Deep Learning được xây dựng sẵn cho PyTorch. Tuy nhiên, nếu bạn muốn mang theo hộp đựng Docker của riêng mình, bạn có thể tham khảo Tạo vùng chứa Docker của riêng bạn với Thư viện song song mô hình phân tán SageMaker. Để bắt đầu, hãy làm theo Sửa đổi tập lệnh đào tạo PyTorch để điều chỉnh các API của SMP trong tập lệnh đào tạo của bạn. Trong phần này, chúng tôi chỉ nêu ra một số bước chính với các đoạn mã từ tập lệnh đào tạo sẵn sàng sử dụng train_gpt_simple.py. Bạn có thể làm theo các nhận xét trong kịch bản và Tài liệu API để tìm hiểu thêm về nơi sử dụng các API SMP.

Đầu tiên, nhập và khởi tạo thư viện bằng cách gọi smdistributed.modelparallel.torch.init() ở đầu tập lệnh đào tạo:

import smdistributed.modelparallel.torch as smp

smp.init(smp_config)

Thứ hai, bọc mô hình được phân vùng với smdistribution.modelparallel.torch.DistributionModel và sử dụng trả lại DistributedModel đối tượng đi về phía trước:

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_config(model_config)
model = smp.DistributedModel(model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation)

Kết hợp trình tối ưu hóa với smdistributed.modelparallel.torch.DistributedOptimizer để lưu và tải các trạng thái của trình tối ưu hóa.

from torch import optim

optimizer = optim.Adam(
    param_groups, betas=(args.beta1, args.beta2), lr=args.lr, weight_decay=args.weight_decay
)

optimizer = smp.DistributedOptimizer(
        optimizer, 
        static_loss_scale=None, 
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
        )

Đặt logic tiến và lùi trong một hàm bước và trang trí nó bằng smdistributed.modelparallel.torch.step.  Bất kỳ tính toán nào được xác định bên trong smp.step-decorated chức năng được thực thi theo cách phân tán.

@smp.step
def train_step(model, optimizer, input_ids, attention_mask, args):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    model.backward(loss)

    return loss

@smp.step
def test_step(model, input_ids, attention_mask):
    loss = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)["loss"]
    
    return loss

2. Chuẩn bị tập dữ liệu

Chúng tôi sử dụng văn bản mở là tập dữ liệu mà chúng tôi sử dụng trong ví dụ này. Sổ tay sử dụng tập lệnh data_prep_512.py để tải xuống và xử lý trước tập dữ liệu. Bạn cũng có thể đào tạo với các bộ dữ liệu khác bằng cách sửa đổi data_pipeline.py. Khi xử lý tập dữ liệu và mô hình lớn, bạn có thể tăng tốc công việc đào tạo bằng cách sử dụng dữ liệu được lưu trữ trong Amazon FSx cho ánh, cung cấp hệ thống tệp hiệu suất cao được tích hợp nguyên bản với Dịch vụ lưu trữ đơn giản của Amazon (S3). Vui lòng xem hướng dẫn từ Định cấu hình kênh nhập dữ liệu để sử dụng Amazon FSx cho Lustre để được hướng dẫn về cách đặt hệ thống tệp FSx Luster làm kênh đầu vào dữ liệu.

3. Bắt đầu các công việc đào tạo

Bước này giả sử bạn đã đã sửa đổi kịch bản đào tạo của bạn và chuẩn bị tập dữ liệu như đã đề cập trong các phần trước. Đến cho phép phân đoạn dữ liệu song song, chỉ cần đặt sharded_data_parallel_degree trong Công cụ ước tính PyTorch. Trong hướng dẫn này, chúng tôi đặt sharded_data_parallel_degree=128 và instace_count=32 đối với các nút p4d.24xlarge, chỉ ra rằng trạng thái mô hình sẽ được chia nhỏ trên 128 GPU trong tổng số 256 GPU. Dựa trên giá trị đã chọn này, SMP sau đó sẽ tự động đặt mức độ song song của dữ liệu thành 2 (vì 256/128 = 2), nghĩa là chúng ta sẽ có hai bản sao cho dữ liệu song song. Quy tắc chung để chọn một giá trị lý tưởng cho sharded_data_parallel_degree là thêm một nút nữa vào nhóm chia sẻ trên mỗi 3B tham số mô hình. Trong hướng dẫn này, kích thước mô hình của chúng tôi là 30B, vì vậy chúng tôi nên sử dụng ít nhất 10 nút cho sharding. Và bởi vì 16 nút (128 GPU) là công suất nhỏ nhất của 2 trên ngưỡng, chúng tôi đặt sharded_data_parallel_degree=128.

Để kiểm tra, chúng tôi cũng cung cấp một tập hợp các tiện ích kiểm tra trong sharded_data_parallel_checkpoint.py , bao gồm một tiện ích để tái tạo lại toàn bộ state_dict cho các trường hợp sử dụng nâng cao. Cuối cùng, chúng tôi có thể khởi chạy công việc đào tạo phân tán bằng cách gọi fit () trên Công cụ ước tính.

smp_estimator = PyTorch(
    entry_point="train_gpt_simple.py",
    instance_type="ml.p4d.24xlarge",
    source_dir=os.getcwd(),
    volume_size=500,
    instance_count=32,
    distribution={
        "mpi": {
            "enabled": True,
            "processes_per_host": processes_per_host,
            "custom_mpi_options": mpioptions,
        },
        "smdistributed": {
            "modelparallel": {
                "enabled": True,
                "parameters": {
                    "ddp": True,
                    "skip_tracing": True,
                    "delayed_parameter_initialization": True,
                    "offload_activations": True,
                    "activation_loading_horizon": 4,
                    # To enable sharded data parallelism.
                    # Here we shard model states across 128 GPUs. 
                    "sharded_data_parallel_degree": 128, 
                    "fp16": False,
                    "bf16": True,
                    # This is to disable pipeline parallelism.
                    "partitions": 1,
                },
            }
        },
    },
    framework_version="1.12",
    py_version="py38",
    hyperparameters=hyperparameters,
    checkpoint_s3_uri=checkpoint_s3_uri if not use_fsx else None,
    checkpoint_local_path=hyperparameters["checkpoint-dir"] if use_fsx else None,
    ...
)

smp_estimator.fit(inputs=data_channels)

4. Giám sát các công việc đào tạo

Bạn có thể truy cập nhật ký đào tạo và theo dõi việc sử dụng GPU và bộ nhớ trên amazoncloudwatch. Đảm bảo xem nhật ký của “algo-1” bởi vì đó là nút chính có luồng đầu ra có nhật ký công việc đào tạo từ tất cả các trường hợp.

Hiệu suất đo điểm chuẩn

Chúng tôi đã đánh giá chuẩn độ song song dữ liệu phân đoạn trong thư viện SMP trên cả hai nút 16 và 32 p4d.24xlarge cho độ dài chuỗi lần lượt là 512 và 2048. Mô hình GPT30 tham số 2B được định cấu hình để sử dụng chiều rộng ẩn là 7168, 48 lớp và 64 đầu. Bạn có thể áp dụng cùng một cấu hình trong đó độ dài chuỗi là 2048 bằng cách cài đặt model_config = "gpt2-30b" trong sổ tay hướng dẫn. Với thiết lập này, SMP đạt được 73.52 mẫu mỗi giây, tăng 39.7% so với DeepSpeed ​​ZeRO-3. Nếu kích thước mã thông báo của bạn là 500 tỷ, tốc độ tăng này đồng nghĩa với việc tiết kiệm gần 367 giờ trên các nút p4d.24xlarge, tương đương với hơn 12,000 đô la ngân sách tiết kiệm cho mỗi lần đào tạo! Bảng sau đây tóm tắt kết quả điểm chuẩn của chúng tôi.

1 / Đối với mỗi cấu hình mô hình, chúng tôi đã thử nghiệm các tính năng, giai đoạn và cấu hình khác nhau trong DeepSpeed ​​ZeRO và chọn một cấu hình cung cấp thông lượng tốt nhất làm đường cơ sở DeepSpeed. Điểm chuẩn đã được chạy trên Đám mây điện toán đàn hồi Amazon (Amazon EC2). 2 / Những kết quả này dựa trên các tập thể giao tiếp được cải tiến được tối ưu hóa cho AWS sẽ sớm được cung cấp. 3 / Thời gian đào tạo được dự đoán từ tốc độ dựa trên số lượng token được xử lý.

Tóm lại, chúng tôi đã quan sát thấy thông lượng cao hơn một cách nhất quán với tính năng song song dữ liệu phân đoạn trong SMP khi so sánh với DeepSpeed ​​trên một loạt các mô hình và cấu hình. Tính năng mới này cũng chứng tỏ hiệu quả bộ nhớ tốt hơn so với DeepSpeed, cho phép SMP phù hợp với kích thước lô lớn hơn và giảm mức độ tích lũy gradient cần thiết để phù hợp với kích thước lô toàn cầu cụ thể.

Kết luận

Trong bài đăng này, chúng tôi đã giới thiệu một kỹ thuật đào tạo phân tán mới - phân đoạn dữ liệu song song - và cách nó tăng tốc đào tạo mô hình khổng lồ với tỷ lệ gần tuyến tính trên Amazon SageMaker. Chúng tôi cũng đã hướng dẫn cách đào tạo mô hình GPT-2 với kỹ thuật mới sau đây hoàn thành ví dụ. Bạn có thể làm theo Amazon SageMaker Ví dụ về repo GitHub để theo dõi tất cả các ví dụ song song của mô hình SageMaker hoặc tham dự phân phối hội thảo đào tạo. Để tìm hiểu thêm về tính song song dữ liệu phân đoạn, vui lòng xem tài liệu hướng dẫ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.

Có thể Karakus là Nhà khoa học ứng dụng cao cấp tại AWS, tối ưu hóa quá trình học sâu phân tán trên quy mô lớn trên AWS. Các mối quan tâm nghiên cứu của anh ấy bao gồm học sâu, tối ưu hóa phân tán, hệ thống phân tán và lý thuyết thông tin. Ngoài công việc, anh thích đạp xe, du lịch, đọc sách và học hỏi.

Rahul Huilgol là Kỹ sư phần mềm cao cấp tại AWS. Anh ấy làm việc trên các hệ thống học sâu phân tán, hướng tới việc đào tạo các mô hình học sâu lớn trên đám mây một cách dễ dàng và hiệu quả. Khi rảnh rỗi, anh thích chụp ảnh, đạp xe và làm vườn.

Suhit Kodgule là Kỹ sư phát triển phần mềm của nhóm Trí tuệ nhân tạo AWS làm việc trên các khuôn khổ học sâu. Khi rảnh rỗi, anh ấy thích đi bộ đường dài, đi du lịch và nấu ăn.

Erin Hồ là Giám đốc sản phẩm của AWS Deep Learning. Cô ấy nghiên cứu các sản phẩm giúp khách hàng đào tạo mô hình học sâu trên AWS dễ dàng hơn. Đối với công việc bên ngoài, cô thích đi bộ đường dài và trượt tuyết.