Thêm Accelerate vào mã của bạn

Mỗi framework huấn luyện phân tán đều có cách làm việc riêng, có thể đòi hỏi viết nhiều mã tùy chỉnh để thích ứng với mã huấn luyện PyTorch và môi trường huấn luyện của bạn. Accelerate cung cấp một cách tiếp cận thân thiện để tương tác với các framework huấn luyện phân tán này mà không cần phải học chi tiết cụ thể của từng framework. Accelerate lo liệu các chi tiết đó cho bạn, giúp bạn tập trung vào mã huấn luyện và mở rộng nó sang bất kỳ môi trường huấn luyện phân tán nào.

Trong hướng dẫn này, bạn sẽ học cách điều chỉnh mã PyTorch hiện có của mình với Accelerate và bắt đầu hành trình huấn luyện trên các hệ thống phân tán một cách dễ dàng! Bạn sẽ bắt đầu với một vòng lặp huấn luyện PyTorch cơ bản (giả định rằng tất cả các đối tượng huấn luyện như model và optimizer đã được thiết lập sẵn) và dần tích hợp Accelerate vào đó.

device = "cuda"
model.to(device)

for batch in training_dataloader:
    optimizer.zero_grad()
    inputs, targets = batch
    inputs = inputs.to(device)
    targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
    loss.backward()
    optimizer.step()
    scheduler.step()

Bộ tăng tốc

Accelerator là lớp chính để điều chỉnh mã của bạn nhằm hoạt động với Accelerate. Lớp này nắm rõ cấu hình phân tán mà bạn đang sử dụng, chẳng hạn như số lượng tiến trình khác nhau và loại phần cứng. Lớp này cũng cung cấp quyền truy cập vào nhiều phương thức cần thiết để cho phép mã PyTorch của bạn hoạt động trong bất kỳ môi trường huấn luyện phân tán nào và để quản lý cũng như thực thi các tiến trình trên các thiết bị.

Đó là lý do bạn luôn nên bắt đầu bằng cách import và tạo một instance Accelerator trong script của mình.

from accelerate import Accelerator

accelerator = Accelerator()

Accelerator cũng biết thiết bị nào để di chuyển các đối tượng PyTorch của bạn, vì vậy khuyến nghị để Accelerate xử lý việc này cho bạn.

- device = "cuda"
+ device = accelerator.device
  model.to(device)

Chuẩn bị các đối tượng PyTorch

Tiếp theo, bạn cần chuẩn bị các đối tượng PyTorch của mình (model, optimizer, scheduler, v.v.) cho huấn luyện phân tán. Phương thức prepare() sẽ lo việc đặt model của bạn vào container phù hợp (như single GPU hoặc multi-GPU) cho cấu hình huấn luyện, điều chỉnh optimizer và scheduler để sử dụng AcceleratedOptimizer và AcceleratedScheduler của Accelerate, và tạo một dataloader mới có thể được chia nhỏ (sharded) giữa các tiến trình.

Accelerate chỉ chuẩn bị các đối tượng kế thừa từ các lớp PyTorch tương ứng của chúng, chẳng hạn như torch.optim.Optimizer.

Các đối tượng PyTorch được trả về theo đúng thứ tự chúng được gửi đi.

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

Vòng lặp huấn luyện

Cuối cùng, hãy xóa các lệnh gọi to(device) cho inputs và targets trong vòng lặp huấn luyện vì các lớp DataLoader của Accelerate sẽ tự động đặt chúng lên đúng thiết bị. Bạn cũng nên thay thế bước backward() thông thường bằng phương thức backward() của Accelerate, phương thức này sẽ tự động điều chỉnh tỷ lệ cho các gradient và sử dụng phương thức backward() phù hợp tùy theo cấu hình phân tán của bạn (ví dụ: DeepSpeed hoặc Megatron).

-   inputs = inputs.to(device)
-   targets = targets.to(device)
    outputs = model(inputs)
    loss = loss_function(outputs, targets)
-   loss.backward()
+   accelerator.backward(loss)

Ghép tất cả lại và vòng lặp huấn luyện Accelerate mới của bạn sẽ trông như thế này!

from accelerate import Accelerator
accelerator = Accelerator()

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

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

Tính năng huấn luyện

Accelerate cung cấp các tính năng bổ sung - như tích lũy gradient, cắt gradient, huấn luyện độ chính xác hỗn hợp và nhiều hơn nữa - bạn có thể thêm vào script của mình để cải thiện quá trình huấn luyện. Hãy cùng khám phá ba tính năng này.

Gradient accumulation

Gradient accumulation cho phép bạn huấn luyện với kích thước batch lớn hơn bằng cách tích lũy gradient qua nhiều batch trước khi cập nhật trọng số. Tính năng này có thể hữu ích để vượt qua các giới hạn về bộ nhớ. Để bật tính năng này trong Accelerate, hãy chỉ định tham số gradient_accumulation_steps trong lớp Accelerator và thêm context manager accumulate() vào script của bạn.

+ accelerator = Accelerator(gradient_accumulation_steps=2)
  model, optimizer, training_dataloader = accelerator.prepare(model, optimizer, training_dataloader)

  for input, label in training_dataloader:
+     with accelerator.accumulate(model):
          predictions = model(input)
          loss = loss_function(predictions, label)
          accelerator.backward(loss)
          optimizer.step()
          scheduler.step()
          optimizer.zero_grad()

Cắt gradient

Cắt gradient là một kỹ thuật để ngăn chặn hiện tượng “gradient phát nổ”, và Accelerate cung cấp:

• clipgrad_value() để cắt gradient về một giá trị tối thiểu và tối đa
• clipgrad_norm() để chuẩn hóa gradient về một giá trị nhất định

Độ chính xác hỗn hợp

Mixed precision tăng tốc độ huấn luyện bằng cách sử dụng kiểu dữ liệu có độ chính xác thấp hơn như fp16 (half-precision) để tính toán các gradient. Để đạt hiệu suất tốt nhất với Accelerate, loss nên được tính toán bên trong model của bạn (giống như trong các model của Transformers) vì các phép tính bên ngoài model được thực hiện với độ chính xác đầy đủ.

Đặt kiểu độ chính xác hỗn hợp cần sử dụng trong Accelerator, sau đó dùng trình quản lý ngữ cảnh autocast() để tự động chuyển đổi các giá trị sang kiểu dữ liệu đã chỉ định.

Accelerate cho phép tự động hóa độ chính xác hỗn hợp, do đó autocast() chỉ cần thiết khi có các thao tác độ chính xác hỗn hợp khác ngoài những thao tác được thực hiện trên loss bởi backward(), vốn đã xử lý việc mở rộng (scaling).

+ accelerator = Accelerator(mixed_precision="fp16")
+ with accelerator.autocast():
      loss = complex_loss_function(outputs, target)

Lưu và tải

Accelerate cũng có thể lưu và tải mô hình sau khi huấn luyện hoàn tất, hoặc bạn cũng có thể lưu trạng thái của mô hình và bộ tối ưu hóa state, điều này có thể hữu ích để tiếp tục huấn luyện.

Mô hình

Khi tất cả các tiến trình hoàn tất, hãy giải bọc mô hình bằng phương thức unwrap_model() trước khi lưu, vì phương thức prepare() đã bọc mô hình của bạn vào giao diện phù hợp cho huấn luyện phân tán. Nếu không giải bọc mô hình, việc lưu từ điển trạng thái mô hình cũng sẽ lưu lại bất kỳ lớp bổ sung nào từ mô hình lớn hơn, và bạn sẽ không thể nạp lại trọng số vào mô hình cơ sở của mình.

Bạn nên sử dụng phương thức save_model() để tách và lưu trữ từ điển trạng thái của mô hình. Phương thức này cũng có thể lưu mô hình vào các checkpoint phân mảnh hoặc vào định dạng safetensors.

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory)

unwrapped_model = accelerator.unwrap_model(model)
path_to_checkpoint = os.path.join(save_directory,"pytorch_model.bin")
unwrapped_model.load_state_dict(torch.load(path_to_checkpoint))

accelerator.wait_for_everyone()
accelerator.save_model(model, save_directory, max_shard_size="1GB", safe_serialization=True)

load_checkpoint_in_model(unwrapped_model, save_directory, device_map={"":device})

Trạng thái

Trong quá trình huấn luyện, bạn có thể muốn lưu trạng thái hiện tại của mô hình, optimizer, các bộ sinh số ngẫu nhiên và có thể cả learning rate scheduler để khôi phục chúng trong cùng một script. Bạn nên thêm các phương thức save_state() và load_state() vào script của mình để lưu và nạp trạng thái.

Để tùy chỉnh thêm vị trí và cách lưu trạng thái thông qua save_state(), hãy sử dụng lớp ProjectConfiguration. Ví dụ, nếu automatic_checkpoint_naming được bật, mỗi checkpoint được lưu sẽ được đặt tại Accelerator.project_dir/checkpoints/checkpoint_{checkpoint_number}.

Các mục có trạng thái khác cần lưu trữ phải được đăng ký với phương thức register_for_checkpointing() để có thể được lưu và tải. Mọi đối tượng được truyền vào phương thức này để lưu trữ phải có hàm load_state_dict và state_dict.

Nếu bạn đã cài đặt torchdata>=0.8.0, bạn có thể truyền thêm use_stateful_dataloader=True vào DataLoaderConfiguration của mình. Điều này mở rộng các lớp DataLoader của Accelerate bằng cách thêm hàm load_state_dict và state_dict, đồng thời cho phép Accelerator.save_state và Accelerator.load_state theo dõi mức độ đã đọc trong bộ dữ liệu huấn luyện khi lưu mô hình.