Accelerate documentation
DeepSpeed를 사용하여 여러 모델 사용
DeepSpeed를 사용하여 여러 모델 사용
이 가이드는 DeepSpeed 사용 가이드를 읽고 이해하셨다고 가정합니다.
Accelerate와 DeepSpeed를 사용하여 여러 모델을 실행하는 것은 다음 경우에 유용합니다:
- 지식 증류
- RLHF와 같은 사후 학습 기법(TRL 라이브러리를 참고하여 더 많은 예시를 확인하세요)
- 여러 모델을 동시에 학습
현재 Accelerate에는 여러 모델을 사용하는 데 도움이 되는 매우 실험적인 API가 있습니다.
이 튜토리얼은 두 가지 일반적인 사용 사례에 초점을 맞춥니다:
- 지식 증류(Knowledge distillation)는 더 크고 성능이 좋은 teacher 모델의 동작을 모방하도록 더 작은 student 모델을 학습시키는 기법입니다. student 모델이 단일 GPU에 들어맞는 경우, 학습에는 ZeRO-2를 사용하고 teacher의 추론을 위한 샤딩(sharding)에는 ZeRO-3를 사용할 수 있습니다. 이는 두 모델 모두에 ZeRO-3를 사용하는 것보다 훨씬 빠릅니다.
- 여러 개의 서로 다른 모델을 동시에 학습합니다.
지식 증류
Knowledge distillation은 여러 모델을 사용하되 그중 하나만 학습하는 좋은 예시입니다.
일반적으로 두 모델 모두에 대해 단일 utils.DeepSpeedPlugin을 사용합니다. 그러나 이 경우 두 개의 독립적인 구성이 있습니다. Accelerate는 필요한 시점에 올바른 플러그인을 참조하고 활성화할 수 있도록, dict 안에 여러 플러그인을 생성하고 사용할 수 있도록 허용합니다. 단, 반드시 dict 형태여야 합니다.
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은 전체 학습을 위해 구성해야 하며(자체적으로 사용하지 않는 경우 scheduler와 optimizer를 지정해야 함), zero3_config.json은 아래 예시와 같이 추론 모델에만 구성해야 합니다.
{
"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
}아래에 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",
}
train_micro_batch_size_per_gpu가 지정되지 않으면 해당 모델이 학습되지 않는 경우에도 DeepSpeed는 오류를 발생시킵니다.
여기서 Accelerator를 하나 생성하고 두 구성을 모두 전달합니다.
from accelerate import Accelerator
accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins)이제 어떻게 사용하는지 살펴보겠습니다.
Student model
기본적으로 Accelerate는 dict의 첫 번째 항목을 기본 또는 활성화된 플러그인("student" 플러그인)으로 설정합니다. utils.deepspeed.get_active_deepspeed_plugin() 함수를 사용하여 활성화된 플러그인이 무엇인지 확인해 보세요.
active_plugin = get_active_deepspeed_plugin(accelerator.state)
assert active_plugin is deepspeed_plugins["student"]AcceleratorState는 활성 DeepSpeed 플러그인을 state.deepspeed_plugin에 저장합니다.
assert active_plugin is accelerator.deepspeed_pluginstudent가 현재 활성화된 플러그인이므로, 이제 모델, 옵티마이저, 스케줄러를 준비해 보겠습니다.
student_model, optimizer, scheduler = ... student_model, optimizer, scheduler, train_dataloader = accelerator.prepare(student_model, optimizer, scheduler, train_dataloader)
이제 teacher model을 다뤄야 할 때입니다.
교사 모델
먼저 Accelerator에서 zero3_config.json 구성을 사용하도록 지정해야 합니다.
accelerator.state.select_deepspeed_plugin("teacher")이 설정은 "student" 플러그인을 비활성화하고 대신 "teacher" 플러그인을 활성화합니다. Transformers 내부의 DeepSpeed stateful config가 업데이트되며, deepspeed.initialize()를 사용할 때 호출되는 플러그인 구성이 변경됩니다. 이를 통해 Transformers가 제공하는 자동 deepspeed.zero.Init 컨텍스트 매니저 통합을 사용할 수 있습니다.
teacher_model = AutoModel.from_pretrained(...) teacher_model = accelerator.prepare(teacher_model)
그렇지 않은 경우, deepspeed.zero.Init로 모델을 수동으로 초기화해야 합니다.
with deepspeed.zero.Init(accelerator.deepspeed_plugin.config):
model = MyModel(...)학습
여기서부터는 teacher_model이 학습되지 않는 한, 학습 루프는 원하는 대로 구성할 수 있습니다.
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()여러 개의 서로 다른 모델 학습
여러 모델을 학습하는 것은 더 복잡한 시나리오입니다. 현재 상태에서는 학습 중 각 모델이 서로 완전히 독립적이라고 가정합니다.
이 시나리오에서는 여전히 두 개의 utils.DeepSpeedPlugin을 생성해야 합니다. 또한, 서로 다른 시점에 서로 다른 deepspeed 엔진이 호출되므로 두 번째 Accelerator도 필요합니다. 하나의 Accelerator는 한 번에 하나의 인스턴스만 처리할 수 있습니다.
다만 state.AcceleratorState는 상태가 있는 객체이므로 이미 사용 가능한 utils.DeepSpeedPlugin 두 개를 모두 인지하고 있습니다. 따라서 추가 인자 없이 두 번째 Accelerator를 인스턴스화하면 됩니다.
first_accelerator = Accelerator(deepspeed_plugins=deepspeed_plugins) second_accelerator = Accelerator()
first_accelerator.state.select_deepspeed_plugin()을 호출하여 특정 플러그인을 활성화하거나 비활성화한 다음, 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
)이제 학습을 시작할 수 있습니다:
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()리소스
더 많은 예시를 보려면 [Accelerate]에 현재 포함된 관련 테스트를 확인해 보세요.
Update on GitHub