merge.py的坑

[Ravansper]

分享一个有关文件“xtuner/xtuner/tools/model_converters/merge.py”的坑，具体来说就是原merge.py文件不能支持多gpu，让你想合并一共参数较大的模型权重时会报错，例如qwen2.5-72b,原因在于
parser.add_argument(
'--device',
default='cuda',
choices=('cuda', 'cpu', 'auto'),
help='Indicate the device')
这代代码中，默认值为CUDA，这就导致，你想用两张或多张卡（CUDA_VISIBLE_DEVICES=0,1,...）不起作用，还是会只使用单卡，需要把默认值改为auto，修改后的完整代码如下：
import argparse
import os
import torch
from peft import PeftModel
from transformers import (AutoModelForCausalLM, AutoTokenizer,
CLIPImageProcessor, CLIPVisionModel)

from xtuner.model.utils import LoadWoInit

def parse_args():
parser = argparse.ArgumentParser(
description='Merge a HuggingFace adapter to base model')
parser.add_argument('model_name_or_path', help='model name or path')
parser.add_argument('adapter_name_or_path', help='adapter name or path')
parser.add_argument(
'save_dir', help='the directory to save the merged model')
parser.add_argument(
'--max-shard-size',
type=str,
default='2GB',
help='Only applicable for LLM. The maximum size for '
'each sharded checkpoint.')
parser.add_argument(
'--is-clip',
action='store_true',
help='Indicate if the model is a clip model')
parser.add_argument(
'--safe-serialization',
action='store_true',
help='Indicate if using safe_serialization')
parser.add_argument(
'--device',
default='auto', # 设置默认设备为 'auto' 以支持多 GPU
choices=('cuda', 'cpu', 'auto'),
help='Indicate the device')

args, unknown = parser.parse_known_args()  # 使用 parse_known_args 以忽略未知参数
return args

def main():
args = parse_args()

if args.is_clip:
    with LoadWoInit():
        model = CLIPVisionModel.from_pretrained(
            args.model_name_or_path, device_map=args.device)
    processor = CLIPImageProcessor.from_pretrained(args.model_name_or_path)
else:
    with LoadWoInit():
        model = AutoModelForCausalLM.from_pretrained(
            args.model_name_or_path,
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            device_map=args.device,  # 'auto' 会自动分配到多个 GPU
            trust_remote_code=True)
    processor = AutoTokenizer.from_pretrained(
        args.model_name_or_path, trust_remote_code=True)

model_unmerged = PeftModel.from_pretrained(
    model,
    args.adapter_name_or_path,
    device_map=args.device,  # 确保 PEFT 模型也使用 'auto' 分配
    is_trainable=False,
    trust_remote_code=True)

model_merged = model_unmerged.merge_and_unload()
print(f'Saving to {args.save_dir}...')

model_merged.save_pretrained(
    args.save_dir,
    safe_serialization=args.safe_serialization,
    max_shard_size=args.max_shard_size)

processor.save_pretrained(args.save_dir)
print('All done!')

if name == 'main':
main()