Huggingface Pretrained 中 device_map = "auto" 的替代方案

Question

我有一个模型，我正在使用以下代码从 Huggingface 读取该模型：

from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True)

现在我阅读了模型，并对内部层进行了一些修改并添加了更多层。当我开始训练/微调时，我发现并非所有内容都在同一模型上。

现在经过更多调查，我发现我的自定义层并未像原始模型那样分布在多个 GPU 上。所以我需要类似

device_map="auto"

的东西，但在阅读模型之后。

就像这样

tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True)

model.device_map = "auto"

Answer 1

我发现

accelerate

中实际上有几种方法可以实现这一点。第一个用于分析您的模型并计算模型将占用的可用内存总量：

https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.infer_auto_device_map

第二个用于将您的型号与设备进行匹配：

https://huggingface.co/docs/accelerate/en/package_reference/big_modeling#accelerate.dispatch_model

所以基本上，根据您的情况，您可以使用以下代码：

from accelerate import dispatch_model, infer_auto_device_map

model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=True)

***
...
new_model = CustomModel(model)
...
***

device_map_dict = infer_auto_device_map(new_model)
dispatch_model(new_model, device_map_dict)

附注但是，请注意，如果我理解正确，您可能仍然会遇到不同设备的问题，因为

device_map="auto"

主要用于推断大型模型，因为它们变得非常大，仅在一台设备上进行计算。这部分还需要测试。