Accelerate documentation
Sử dụng nhiều mô hình với DeepSpeed
Sử dụng nhiều mô hình với DeepSpeed
Hướng dẫn này giả định rằng bạn đã đọc và hiểu hướng dẫn sử dụng DeepSpeed.
Việc chạy nhiều mô hình với Accelerate và DeepSpeed hữu ích cho:
- Chưng cất tri thức
- Các kỹ thuật hậu huấn luyện như RLHF (xem thư viện TRL để biết thêm các ví dụ)
- Huấn luyện nhiều mô hình cùng lúc
Hiện tại, Accelerate có một API rất thử nghiệm để giúp bạn sử dụng nhiều mô hình.
Hướng dẫn này sẽ tập trung vào hai trường hợp sử dụng phổ biến:
- Knowledge distillation, trong đó một mô hình student nhỏ hơn được huấn luyện để bắt chước một mô hình teacher lớn hơn, có hiệu suất tốt hơn. Nếu mô hình student vừa trên một GPU duy nhất, chúng ta có thể sử dụng ZeRO-2 để huấn luyện và ZeRO-3 để chia nhỏ (shard) mô hình teacher cho suy luận. Cách này nhanh hơn đáng kể so với việc sử dụng ZeRO-3 cho cả hai mô hình.
- Huấn luyện nhiều mô hình rời rạc cùng lúc.
Chưng cất tri thức
Knowledge distillation là một ví dụ tốt về việc sử dụng nhiều mô hình nhưng chỉ huấn luyện một trong số chúng.
Thông thường, bạn sẽ sử dụng một utils.DeepSpeedPlugin duy nhất cho cả hai mô hình. Tuy nhiên, trong trường hợp này, có hai cấu hình riêng biệt. Accelerate cho phép bạn tạo và sử dụng nhiều plugin chỉ khi chúng nằm trong một dict để bạn có thể tham chiếu và kích hoạt plugin phù hợp khi cần.
from accelerate.utils import DeepSpeedPlugin
zero2_plugin = DeepSpeedPlugin(hf_ds_config="zero2_config.json")
zero3_plugin = DeepSpeedPlugin(hf_ds_config="zero3_config.json")
deepspeed_plugins = {"student": zero2_plugin, "teacher": zero3_plugin}zero2_config.json nên được cấu hình cho huấn luyện đầy đủ (do đó hãy chỉ định scheduler và optimizer nếu bạn không sử dụng của riêng mình), trong khi zero3_config.json chỉ nên được cấu hình cho mô hình suy luận, như được minh họa trong ví dụ bên dưới.
{
"bf16": {
"enabled": "auto"
},
"zero_optimization": {
"stage": 3,
"overlap_comm": true,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": "auto",
"stage3_max_reuse_distance": "auto",
},
"train_micro_batch_size_per_gpu": 1
}Dưới đây là một ví dụ về cấu hình zero2_config.json.
{
"bf16": {
"enabled": "auto"
},
"optimizer": {
"type": "AdamW",
"params": {
"lr": "auto",
"weight_decay": "auto",
"torch_adam": true,
"adam_w_mode": true
}
},
"scheduler": {
"type": "WarmupLR",
"params": {
"warmup_min_lr": "auto",
"warmup_max_lr": "auto",
"warmup_num_steps": "auto"
}
},
"zero_optimization": {
"stage": 2,
"offload_optimizer": {
"device": "cpu",
"pin_memory": true
},
},
"gradient_accumulation_steps": 1,
"gradient_clipping": "auto",
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
}DeepSpeed sẽ gây ra lỗi nếu
train_micro_batch_size_per_gpukhông được chỉ định, ngay cả khi mô hình cụ thể này không được huấn luyện.
Từ đây, tạo một Accelerator duy nhất và truyền vào cả hai cấu hình.
from accelerate import Accelerator
accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)Bây giờ hãy xem cách sử dụng chúng.
Mô hình sinh viên
Mặc định, Accelerate đặt mục đầu tiên trong dict làm plugin mặc định hoặc được bật (plugin "student"). Hãy xác minh điều này bằng cách sử dụng hàm utils.deepspeed.get_active_deepspeed_plugin() để xem plugin nào đang được bật.
active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]AcceleratorState cũng lưu trữ plugin DeepSpeed đang hoạt động trong state.deepspeed_plugin.
assert active_plugin is accelerator.deepspeed_pluginVì student là plugin đang hoạt động, chúng ta sẽ tiến hành chuẩn bị mô hình, optimizer và scheduler.
student_model, optimizer, scheduler = ... student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)
Bây giờ là lúc xử lý mô hình giáo viên.
Mô hình giáo viên
Đầu tiên, bạn cần chỉ định trong Accelerator rằng cấu hình zero3_config.json sẽ được sử dụng.
accelerator.state.select_deepspeed_plugin("teacher")Điều này vô hiệu hóa plugin "student" và thay vào đó kích hoạt plugin "teacher". Cấu hình có trạng thái của DeepSpeed bên trong Transformers sẽ được cập nhật, và nó thay đổi cấu hình plugin nào được gọi khi sử dụng deepspeed.initialize(). Điều này cho phép bạn sử dụng tích hợp context manager deepspeed.zero.Init tự động mà Transformers cung cấp.
teacher_model = AutoModel.from_pretrained(...) teacher_model = accelerator.prepare(teacher_model)
Ngược lại, bạn nên khởi tạo thủ công mô hình bằng deepspeed.zero.Init.
with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
model = MyModel(...)Huấn luyện
Từ đây, vòng lặp huấn luyện của bạn có thể là bất kỳ thứ gì bạn muốn, miễn là teacher_model không bao giờ được huấn luyện.
teacher_model.eval()
student_model.train()
for batch in train_dataloader:
with torch.no_grad():
output_teacher = teacher_model(**batch)
output_student = student_model(**batch)
# Combine the losses or modify it in some way
loss = output_teacher.loss + output_student.loss
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()Huấn luyện nhiều mô hình rời rạc
Việc huấn luyện nhiều mô hình là một kịch bản phức tạp hơn. Ở trạng thái hiện tại, chúng tôi giả định rằng mỗi mô hình hoàn toàn tách biệt với các mô hình khác trong quá trình huấn luyện.
Kịch bản này vẫn yêu cầu tạo hai utils.DeepSpeedPlugin. Tuy nhiên, bạn cũng cần một Accelerator thứ hai, vì các deepspeed engine khác nhau được gọi vào những thời điểm khác nhau. Một Accelerator duy nhất chỉ có thể mang một instance tại một thời điểm.
Tuy nhiên, vì state.AcceleratorState là một đối tượng có trạng thái, nó đã nhận thức được cả hai utils.DeepSpeedPlugin khả dụng. Bạn chỉ cần khởi tạo một Accelerator thứ hai mà không cần thêm bất kỳ đối số nào.
first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins) second_accelerator = Accelerator()
Bạn có thể gọi first_accelerator.state.select_deepspeed_plugin() để bật hoặc tắt một plugin cụ thể, sau đó gọi prepare.
# can be `accelerator_0`, `accelerator_1`, or by calling `AcceleratorState().select_deepspeed_plugin(...)`
first_accelerator.state.select_deepspeed_plugin("first_model")
first_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
first_optimizer, first_scheduler, train_dl, eval_dl = get_training_items(model1)
first_model, first_optimizer, first_scheduler, train_dl, eval_dl = accelerator.prepare(
first_model, first_optimizer, first_scheduler, train_dl, eval_dl
)
second_accelerator.state.select_deepspeed_plugin("second_model")
second_model = AutoModel.from_pretrained(...)
# For this example, `get_training_items` is a nonexistent function that gets the setup we need for training
second_optimizer, second_scheduler, _, _ = get_training_items(model2)
second_model, second_optimizer, second_scheduler = accelerator.prepare(
second_model, second_optimizer, second_scheduler
)Và bây giờ bạn có thể huấn luyện:
for batch in dl:
outputs1 = first_model(**batch)
first_accelerator.backward(outputs1.loss)
first_optimizer.step()
first_scheduler.step()
first_optimizer.zero_grad()
outputs2 = model2(**batch)
second_accelerator.backward(outputs2.loss)
second_optimizer.step()
second_scheduler.step()
second_optimizer.zero_grad()Tài nguyên
Để xem thêm các ví dụ, vui lòng kiểm tra các bài kiểm tra liên quan hiện có trong [Accelerate].
Update on GitHub