[ad_1]
Làm thế nào nó hoạt động?
Như tôi đã nói, khi đào tạo trên một số GPU, mỗi quy trình có các bản sao chính xác của cùng một dữ liệu khi đào tạo với DDP. Chúng tôi có thể tối ưu hóa nó bằng cách triển khai một số cải tiến:
Trạng thái tối ưu hóa phân đoạn (ZeRO 1)
Khi đào tạo với DDP, mỗi quy trình sẽ giữ một bản sao hoàn chỉnh của các trạng thái tối ưu hóa. Với ZeRO1, chúng tôi phân chia các trạng thái tối ưu hóa này trên tất cả các cấp sao cho mỗi cấp chỉ giữ một phần trạng thái tối ưu hóa. Trong quá trình chuyển ngược, mỗi cấp bậc chỉ cần thu thập các trạng thái tối ưu hóa có liên quan đến các tham số của nó để thực hiện bước tối ưu hóa. Việc giảm dư thừa này giúp bảo tồn bộ nhớ.
💡 Trong trường hợp Adam chứa các tham số có kích thước mô hình gần gấp đôi, việc phân chia trạng thái tối ưu hóa thành 8 cấp có nghĩa là mỗi cấp chỉ lưu trữ một phần tư (2/8) tổng kích thước trạng thái.
Độ dốc phân đoạn (ZeRO 2)
Chúng tôi phân chia các trạng thái tối ưu hóa. Bây giờ, chúng tôi cũng sẽ sửa đổi bước tối ưu hóa để phân chia độ dốc. Nếu một thứ hạng có trạng thái tối ưu hóa cho một phần tham số thì chúng tôi sẽ:
- tổng hợp tất cả các độ dốc có liên quan đến các trạng thái mà thứ hạng nắm giữ
- tính toán bước tối ưu hóa
- gửi bước tối ưu hóa cho một phần tham số tới tất cả các cấp bậc khác
Như bạn đã nhận thấy, giờ đây mỗi cấp bậc không cần phải chứa bản sao đầy đủ của độ dốc. Chúng tôi có thể gửi gradient đến thứ hạng phù hợp ngay khi chúng có sẵn. Vì vậy, chúng ta có thể giảm mức tiêu thụ bộ nhớ cao nhất hơn nữa.
Thông số mô hình phân đoạn (ZeRO 3)
Điều này sắp trở nên hoành tráng.
Tại sao chúng ta cần lưu trữ bản sao đầy đủ của mô hình trên mỗi cấp bậc? Hãy phân chia các tham số mô hình giữa tất cả các cấp bậc. Sau đó, chúng ta sẽ tìm nạp các tham số cần thiết kịp thời trong quá trình tiến và lùi.
💡 Đối với các mô hình lớn, những tối ưu hóa này có thể giảm đáng kể mức tiêu thụ bộ nhớ
Làm thế nào để sử dụng FSDP?
Thực sự khá đơn giản. Tất cả những gì chúng ta cần là bọc mô hình bằng FSDP:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributed.fsdp import FullyShardedDataParallel as FSDPmannequin = FSDP(mannequin)
# it's vital to get parameters from the wrapped mannequin
# as solely a portion of them returned (sharded half)
optimizer = optim.Adam(mannequin.parameters())
# consuct coaching as ordinary
practice(mannequin, optimizer)
Bạn cũng có thể chỉ định chiến lược sharding của FSDP. Ví dụ: chúng ta có thể chọn SHARD_GRAD_OP chiến lược để đạt được hành vi tương tự như ZeRO2. Bạn có thể tìm hiểu về các chiến lược khác tại đây:
Ngoài ra, bạn có thể gói bằng các mô-đun con FSDP. Trong ví dụ trên, chỉ có một mô-đun FSDP được sử dụng, điều này sẽ làm giảm hiệu quả tính toán và hiệu quả bộ nhớ. Cách thức hoạt động là giả sử mô hình của bạn chứa 100 lớp Tuyến tính. Nếu bạn thực hiện FSDP(mannequin), sẽ chỉ có một đơn vị FSDP bao bọc toàn bộ mô hình. Trong trường hợp đó, allgather sẽ thu thập đầy đủ tham số cho tất cả 100 lớp tuyến tính và do đó sẽ không lưu bộ nhớ CUDA để phân chia tham số.
Bạn có thể gói các mô-đun con một cách rõ ràng hoặc xác định chính sách tự động gói. Để tìm hiểu thêm về FSDP, hãy đọc hướng dẫn PyTorch:
[ad_2]
Source link