Động lực

Trong lĩnh vực học máy và trí tuệ nhân tạo, việc sử dụng GPU trong quá trình huấn luyện và kiểm thử mô hình đã trở nên phổ biến. Việc này giúp giảm thiểu thời gian đào tạo đáng kể so với sử dụng CPU. Tuy nhiên, khi mô hình hoặc tập dữ liệu trở nên lớn hơn, việc sử dụng một GPU thường không đủ. Đặc biệt là đối với các mô hình ngôn ngữ lớn hiện nay, cần phải sử dụng nhiều GPU mạnh mẽ và thời gian đào tạo có thể kéo dài từ vài ngày đến nhiều tháng. Chính vì vậy, cách thức tối ưu để đào tạo mô hình trên nhiều GPU đã trở thành chủ đề quan tâm lớn. Pytorch cung cấp hai phương pháp chủ yếu để phân phối mô hình và dữ liệu trên nhiều GPU: nn.DataParallel và nn.DistributedDataParallel. Trong bài viết này, chúng ta sẽ cùng khám phá từng phương pháp, cách ứng dụng cũng như những lưu ý quan trọng.

Giới thiệu về nn.DataParallel và nn.DistributedDataParallel

Khi quyết định sử dụng nn.DataParallel hay nn.DistributedDataParallel, ta cần tìm hiểu sự khác nhau giữa hai phương pháp này.

DataParallel

• Chạy trên một máy: DataParallel hoạt động trong môi trường một tiến trình với nhiều luồng trên cùng một máy.
• Sao chép mô hình: Mỗi GPU sẽ sao chép mô hình và xử lý phần dữ liệu khác nhau.
• Vấn đề GIL: Việc sử dụng nhiều luồng có thể gặp phải vấn đề GIL (Global Interpreter Lock), dẫn đến hiệu suất giảm.
• Tiêu tốn bộ nhớ: Mỗi GPU sẽ cần tối đa bộ nhớ để lưu trữ bản sao mô hình, điều này có thể gây khó khăn khi mô hình lớn.

DistributedDataParallel

• Chạy trên nhiều máy: DistributedDataParallel hoạt động trong môi trường nhiều tiến trình, có thể tồn tại trên một hoặc nhiều máy.
• Giao tiếp giữa các tiến trình: Sử dụng phương pháp giao tiếp hiệu quả giữa các tiến trình để phân phối quá trình đào tạo.
• Giảm thiểu chi phí GIL: Giải pháp không yêu cầu sao chép nhiều bản sao của mô hình, giúp tiết kiệm bộ nhớ.

Cách sử dụng DataParallel

Để sử dụng DataParallel, ta có thể theo các bước dưới đây:

1. Cài đặt môi trường

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader

# Các tham số cần thiết
input_size = 5
output_size = 2
batch_size = 30
data_size = 100

2. Khởi tạo thiết bị

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

3. Tạo dataset

class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len

rand_loader = DataLoader(dataset=RandomDataset(input_size, data_size),
                         batch_size=batch_size, shuffle=True)

4. Định nghĩa mô hình

class Model(nn.Module):
    def __init__(self, input_size, output_size):
        super(Model, self).__init__()
        self.fc = nn.Linear(input_size, output_size)

    def forward(self, input):
        output = self.fc(input)
        print("\tIn Model: input size", input.size(), "output size", output.size())
        return output

5. Huấn luyện mô hình với nhiều GPU

model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
  print("Let's use", torch.cuda.device_count(), "GPUs!")
  model = nn.DataParallel(model)

model.to(device)

for data in rand_loader:
    input = data.to(device)
    output = model(input)
    print("Outside: input size", input.size(), "output_size", output.size())

Mô hình Parallel trên một máy

Nếu mô hình quá lớn so với dung lượng của một GPU, thì phương pháp DataParallel sẽ không còn hiệu quả. Ta có thể sử dụng phương pháp gọi là model parallel, trong đó mô hình được chia thành nhiều phần và mỗi phần thực hiện trên một GPU riêng. Điều này tối ưu hóa khả năng xử lý và sử dụng tài nguyên hiệu quả hơn.

Ví dụ mô hình là hai lớp linear

class RandomModel(nn.Module):
    def __init__(self):
        super(RandomModel, self).__init__()
        self.net1 = nn.Linear(10, 5).to('cuda:0')
        self.relu = nn.LeakyReLU()
        self.net2 = nn.Linear(5, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

Sử dụng DistributedDataParallel

Khởi tạo môi trường phân tán

Để bắt đầu với DistributedDataParallel, cần cài đặt nhóm tiến trình và sử dụng DistributedSampler để đảm bảo dữ liệu được chia đều cho các GPU.

import os
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader, DistributedSampler

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

Huấn luyện với DistributedDataParallel

class RandomModel(nn.Module):
    def __init__(self):
        super(RandomModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))

Ta có thể sử dụng hàm mp.spawn để chạy nhiều tiến trình:

def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size, ),
             nprocs=world_size,
             join=True)

Kết luận

Việc thành thạo các phương pháp đào tạo mô hình với nhiều GPU trở nên vô cùng quan trọng khi mà kích thước của mô hình ngày càng lớn. Bài viết đã trình bày những kiến thức cơ bản về DataParallel và DistributedDataParallel, từ đó giúp người đọc có cái nhìn rõ hơn về cách tối ưu tốc độ đào tạo mô hình.

Tham khảo

[1] https://pytorch.org/tutorials/intermediate/ddp_tutorial.html
[2] https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html
[3] https://pytorch.org/docs/stable/data.html?highlight=distributedsampler#torch.utils.data.distributed.DistributedSampler
source: viblo