Sử dụng Local SGD với Accelerate

Local SGD là một kỹ thuật cho huấn luyện phân tán, trong đó các gradient không được đồng bộ hóa ở mỗi bước. Do đó, mỗi tiến trình cập nhật phiên bản trọng số mô hình của riêng mình và sau một số bước nhất định, các trọng số này được đồng bộ hóa bằng cách lấy trung bình trên tất cả các tiến trình. Điều này cải thiện hiệu quả truyền thông và có thể dẫn đến tăng tốc huấn luyện đáng kể, đặc biệt khi máy tính thiếu kết nối nhanh hơn như NVLink. Khác với tích lũy gradient (nơi việc cải thiện hiệu quả truyền thông đòi hỏi phải tăng kích thước lô hiệu dụng), Local SGD không yêu cầu thay đổi kích thước lô hoặc tốc độ học / lịch trình. Tuy nhiên, nếu cần, Local SGD cũng có thể được kết hợp với tích lũy gradient.

Trong hướng dẫn này, bạn sẽ thấy cách thiết lập nhanh Local SGD Accelerate. So với cấu hình Accelerate tiêu chuẩn, điều này chỉ cần thêm hai dòng mã.

Ví dụ này sẽ sử dụng một vòng lặp huấn luyện PyTorch rất đơn giản, thực hiện tích lũy gradient mỗi hai lô:

device = "cuda"
model.to(device)

gradient_accumulation_steps = 2

for index, batch in enumerate(training_dataloader):
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss = loss / gradient_accumulation_steps
    loss.backward()
    if (index + 1) % gradient_accumulation_steps == 0:
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()

Chuyển đổi sang Accelerate

Trước tiên, đoạn mã được hiển thị ở trên sẽ được chuyển đổi để sử dụng Accelerate mà không có LocalSGD hay trình trợ giúp tích lũy gradient:

+ from accelerate import Accelerator
+ accelerator = Accelerator()

+ model, optimizer, training_dataloader, scheduler = accelerator.prepare(
+     model, optimizer, training_dataloader, scheduler
+ )

  for index, batch in enumerate(training_dataloader):
      inputs, targets = batch
-     inputs = inputs.to(device)
-     targets = targets.to(device)
      outputs = model(inputs)
      loss = loss_function(outputs, targets)
      loss = loss / gradient_accumulation_steps
+     accelerator.backward(loss)
      if (index+1) % gradient_accumulation_steps == 0:
          optimizer.step()
          scheduler.step()

Để Accelerate xử lý đồng bộ hóa mô hình

Những gì còn lại bây giờ là để Accelerate xử lý việc đồng bộ hóa tham số mô hình và tích lũy gradient cho chúng ta. Để đơn giản, hãy giả sử chúng ta cần đồng bộ hóa mỗi 8 bước. Điều này được thực hiện bằng cách thêm một câu lệnh with LocalSGD và một lời gọi local_sgd.step() sau mỗi bước tối ưu hóa:

+local_sgd_steps=8

+with LocalSGD(accelerator=accelerator, model=model, local_sgd_steps=8, enabled=True) as local_sgd:
    for batch in training_dataloader:
        with accelerator.accumulate(model):
            inputs, targets = batch
            outputs = model(inputs)
            loss = loss_function(outputs, targets)
            accelerator.backward(loss)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
+           local_sgd.step()

Bên trong, mã Local SGD tắt đồng bộ hóa gradient tự động (nhưng tích lũy vẫn hoạt động như mong đợi!). Thay vào đó, nó trung bình hóa các tham số mô hình mỗi local_sgd_steps bước (cũng như ở cuối vòng lặp huấn luyện).

Hạn chế

Cài đặt hiện tại chỉ hoạt động với huấn luyện đa GPU (hoặc đa CPU) cơ bản, không có, ví dụ, DeepSpeed.

Tham khảo

Mặc dù chúng tôi không rõ nguồn gốc thực sự của cách tiếp cận đơn giản này, nhưng ý tưởng về SGD cục bộ đã khá lâu đời và có thể truy ngược về ít nhất:

Zhang, J., De Sa, C., Mitliagkas, I., & Ré, C. (2016). SGD song song: Khi nào việc trung bình hóa có ích?. Bản thảo arXiv arXiv:1606.07365.

Chúng tôi ghi nhận thuật ngữ Local SGD cho bài báo sau (nhưng có thể có các tài liệu tham khảo sớm hơn mà chúng tôi không biết đến).

Stich, Sebastian Urban. “Local SGD Converges Fast and Communicates Little.” ICLR 2019-International Conference on Learning Representations. No. CONF. 2019.