Unverified 提交 8289338c 作者: Tyler Luan 提交者: GitHub

Multi-gpu lora load bugfix (#312)

上级 31655339
...@@ -11,7 +11,7 @@ DEVICE_ID = "0" if torch.cuda.is_available() else None ...@@ -11,7 +11,7 @@ DEVICE_ID = "0" if torch.cuda.is_available() else None
DEVICE = f"{DEVICE_}:{DEVICE_ID}" if DEVICE_ID else DEVICE_ DEVICE = f"{DEVICE_}:{DEVICE_ID}" if DEVICE_ID else DEVICE_
def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: def auto_configure_device_map(num_gpus: int, use_lora: bool) -> Dict[str, int]:
# transformer.word_embeddings 占用1层 # transformer.word_embeddings 占用1层
# transformer.final_layernorm 和 lm_head 占用1层 # transformer.final_layernorm 和 lm_head 占用1层
# transformer.layers 占用 28 层 # transformer.layers 占用 28 层
...@@ -19,14 +19,21 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: ...@@ -19,14 +19,21 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
num_trans_layers = 28 num_trans_layers = 28
per_gpu_layers = 30 / num_gpus per_gpu_layers = 30 / num_gpus
# bugfix: PEFT加载lora模型出现的层命名不同
if LLM_LORA_PATH and use_lora:
layer_prefix = 'base_model.model.transformer'
else:
layer_prefix = 'transformer'
# bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError # bugfix: 在linux中调用torch.embedding传入的weight,input不在同一device上,导致RuntimeError
# windows下 model.device 会被设置成 transformer.word_embeddings.device # windows下 model.device 会被设置成 transformer.word_embeddings.device
# linux下 model.device 会被设置成 lm_head.device # linux下 model.device 会被设置成 lm_head.device
# 在调用chat或者stream_chat时,input_ids会被放到model.device上 # 在调用chat或者stream_chat时,input_ids会被放到model.device上
# 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError # 如果transformer.word_embeddings.device和model.device不同,则会导致RuntimeError
# 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上 # 因此这里将transformer.word_embeddings,transformer.final_layernorm,lm_head都放到第一张卡上
device_map = {'transformer.word_embeddings': 0, device_map = {f'{layer_prefix}.word_embeddings': 0,
'transformer.final_layernorm': 0, 'lm_head': 0} f'{layer_prefix}.final_layernorm': 0, 'lm_head': 0,
f'base_model.model.lm_head': 0, }
used = 2 used = 2
gpu_target = 0 gpu_target = 0
...@@ -35,7 +42,7 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]: ...@@ -35,7 +42,7 @@ def auto_configure_device_map(num_gpus: int) -> Dict[str, int]:
gpu_target += 1 gpu_target += 1
used = 0 used = 0
assert gpu_target < num_gpus assert gpu_target < num_gpus
device_map[f'transformer.layers.{i}'] = gpu_target device_map[f'{layer_prefix}.layers.{i}'] = gpu_target
used += 1 used += 1
return device_map return device_map
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论