Đồng bộ hóa gradient

Module phân tán của PyTorch hoạt động bằng cách liên tục trao đổi thông tin giữa tất cả các GPU trong hệ thống của bạn. Quá trình liên lạc này mất thời gian, và việc đảm bảo tất cả các tiến trình đều biết được trạng thái của nhau xảy ra tại các điểm kích hoạt cụ thể khi sử dụng module ddp.

Các triggerpoint này được thêm vào mô hình PyTorch, cụ thể là các phương thức forward() và backward() của chúng. Điều này xảy ra khi mô hình được bọc bởi DistributedDataParallel:

import torch.nn as nn
from torch.nn.parallel import DistributedDataParallel

model = nn.Linear(10, 10)
ddp_model = DistributedDataParallel(model)

Trong Accelerate, quá trình chuyển đổi này diễn ra tự động khi gọi prepare() và truyền vào model của bạn.

+ from accelerate import Accelerator
+ accelerator = Accelerator()
  import torch.nn as nn
- from torch.nn.parallel import DistributedDataParallel

  model = nn.Linear(10,10)
+ model = accelerator.prepare(model)

Sự chậm lại trong tích lũy gradient

Bạn đã hiểu rằng PyTorch thêm các hook vào phương thức forward và backward của mô hình PyTorch khi huấn luyện trong cấu hình phân tán. Nhưng điều này có thể làm chậm mã của bạn như thế nào?

Trong DDP (phân tán song song dữ liệu), thứ tự cụ thể mà các tiến trình được thực hiện và chạy được kỳ vọng tại các thời điểm nhất định và những tiến trình này cũng phải xảy ra vào khoảng thời gian tương đương trước khi tiếp tục.

Ví dụ trực tiếp nhất là khi bạn cập nhật tham số mô hình thông qua optimizer.step(). Nếu không có gradient accumulation, tất cả các instance của mô hình cần phải tính toán, gom nhóm và cập nhật gradient của chúng trước khi chuyển sang batch dữ liệu tiếp theo. Khi thực hiện gradient accumulation, bạn tích lũy n gradient của loss và bỏ qua optimizer.step() cho đến khi đạt n batch. Vì tất cả các tiến trình huấn luyện chỉ cần đồng bộ hóa khi optimizer.step() được gọi, mà không có bất kỳ thay đổi nào đối với bước huấn luyện của bạn, việc giao tiếp giữa các tiến trình không cần thiết này có thể gây ra sự chậm trễ đáng kể.

Bạn có thể tránh chi phí này như thế nào?

Giải quyết vấn đề chậm

Vì bạn đang bỏ qua việc cập nhật tham số mô hình khi huấn luyện trên các batch này, gradient của chúng không cần được đồng bộ hóa cho đến khi optimizer.step() thực sự được gọi. PyTorch không thể tự động xác định thời điểm bạn cần thực hiện điều này, nhưng họ cung cấp một công cụ hỗ trợ thông qua context manager no_sync được thêm vào mô hình của bạn sau khi chuyển đổi nó sang DDP.

Trong trình quản lý ngữ cảnh này, PyTorch sẽ bỏ qua việc đồng bộ hóa gradient khi .backward() được gọi, và lần gọi đầu tiên của .backward() bên ngoài trình quản lý ngữ cảnh này sẽ kích hoạt việc đồng bộ hóa. Xem ví dụ bên dưới:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for index, batch in enumerate(dataloader):
    inputs, targets = batch
    # Trigger gradient synchronization on the last batch
    if index != (len(dataloader) - 1):
        with ddp_model.no_sync():
            # Gradients only accumulate
            outputs = ddp_model(inputs)
            loss = loss_func(outputs)
            accelerator.backward(loss)
    else:
        # Gradients finally sync
        outputs = ddp_model(inputs)
        loss = loss_func(outputs)
        accelerator.backward(loss)
        optimizer.step()

Trong Accelerate, để biến điều này thành một API có thể được gọi bất kể thiết bị huấn luyện nào (mặc dù nó có thể không làm gì nếu bạn không đang ở trong một hệ thống phân tán!), ddp_model.no_sync được thay thế bằng no_sync() và hoạt động theo cách tương tự:

  ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

  for index, batch in enumerate(dataloader):
      inputs, targets = batch
      # Trigger gradient synchronization on the last batch
      if index != (len(dataloader)-1):
-         with ddp_model.no_sync():
+         with accelerator.no_sync(model):
              # Gradients only accumulate
              outputs = ddp_model(inputs)
              loss = loss_func(outputs, targets)
              accelerator.backward(loss)
      else:
          # Gradients finally sync
          outputs = ddp_model(inputs)
          loss = loss_func(outputs)
          accelerator.backward(loss)
          optimizer.step()
          optimizer.zero_grad()

Như bạn có thể dự đoán, hàm accumulate() bao quanh việc kiểm tra điều kiện này bằng cách theo dõi số lô hiện tại, để lại cho bạn API tích lũy gradient cuối cùng:

ddp_model, dataloader, optimizer = accelerator.prepare(model, dataloader, optimizer)

for batch in dataloader:
    with accelerator.accumulate(model):
        optimizer.zero_grad()
        inputs, targets = batch
        outputs = model(inputs)
        loss = loss_function(outputs, targets)
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

Do đó, bạn nên sử dụng accelerator.accumulate hoặc accelerator.no_sync khi lựa chọn API.

Mức độ chậm đi cụ thể là bao nhiêu, và những sai lầm dễ mắc phải

Để thiết lập một ví dụ thực tế, hãy xem xét cấu hình sau:

• Hai nút T4 đơn GPU và một nút có hai GPU
• Mỗi GPU là một T4 và được lưu trữ trên GCP
• Script được sử dụng là một bản sửa đổi của script Ví dụ NLP
• Kích thước lô trên mỗi GPU là 16, và gradient được tích lũy mỗi 4 bước

Tất cả các script đều có sẵn trong kho này.

Nếu không chú ý đến việc đồng bộ hóa gradient và giao tiếp GPU, một lượng lớn thời gian có thể bị lãng phí khi các GPU này giao tiếp với nhau trong những khoảng thời gian không cần thiết.

Bao nhiêu?

Tham khảo:

• Baseline: không sử dụng các phương pháp đồng bộ hóa được thảo luận ở đây
• no_sync không đúng cách: no_sync chỉ bao quanh lời gọi backward, không phải forward
• no_sync: sử dụng đúng mô hình no_sync
• accumulate: sử dụng accumulate() đúng cách

Dưới đây là số giây trung bình cho mỗi batch khi lặp qua 29 batch dữ liệu cho mỗi cấu hình trên cả một nút đơn và cấu hình hai nút:

Như bạn có thể thấy, nếu không cẩn thận trong cách thiết lập đồng bộ gradient, bạn có thể gặp phải tình trạng chậm hơn gấp đôi hoặc hơn trong quá trình huấn luyện!

Nếu bạn lo lắng về việc đảm bảo mọi thứ được thực hiện đúng cách, chúng tôi rất khuyến khích bạn sử dụng hàm accumulate() và truyền gradient_accumulation_steps hoặc gradient_accumulation_plugin vào đối tượng Accelerator để Accelerate có thể xử lý việc này cho bạn.

no_sync yêu cầu thêm bộ nhớ GPU khi sử dụng FSDP

Lưu ý rằng việc không đồng bộ gradient có thể gây ra những tác động tiêu cực trong quá trình huấn luyện FSDP. Như đã được cảnh báo trong torch, trình quản lý ngữ cảnh no_sync cho FSDP sẽ yêu cầu thêm bộ nhớ.

Do đó, trong các tình huống sử dụng nhiều bộ nhớ khi dùng FSDP, chúng tôi khuyến nghị đặt sync_each_batch thành True trong GradientAccumulationPlugin để tắt no_sync.

Xem ví dụ bên dưới, nơi chúng ta tinh chỉnh Mixtral (47B tham số) trên 8 GPU A100-80GB. Chúng ta thấy rằng ngay cả với gradient_accumulation_steps=2 khiêm tốn, bộ nhớ sẽ nhanh chóng bị tràn (OOM) nếu no_sync được bật. Một lần nữa, điều này là do chi phí bộ nhớ bổ sung do no_sync của FSDP gây ra. Tuy nhiên, nếu no_sync bị tắt thông qua sync_each_batch=True, thì mức tiêu thụ bộ nhớ cho gradient_accumulation_steps=16 sẽ trở về mức của gradient_accumulation_steps=1.

Việc tắt no_sync đồng nghĩa với việc sẽ chậm hơn do các thao tác đồng bộ hóa dữ liệu bổ sung, như đã được giải thích trong các phần trước của hướng dẫn này.