From b45d6923be1e08959597c14af8ff2b5709dfc165 Mon Sep 17 00:00:00 2001 From: lkm2835 Date: Sun, 27 Oct 2024 18:46:14 +0900 Subject: [PATCH] fix load_model_on_cpu on qwen/convert_checkpoint.py --- examples/qwen/convert_checkpoint.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/examples/qwen/convert_checkpoint.py b/examples/qwen/convert_checkpoint.py index 2bdb8d42a..bd3f04fd1 100644 --- a/examples/qwen/convert_checkpoint.py +++ b/examples/qwen/convert_checkpoint.py @@ -196,7 +196,8 @@ def args_to_build_options(args): 'embedding_sharding_dim': args.embedding_sharding_dim, 'share_embedding_table': args.use_embedding_sharing, 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin + args.disable_weight_only_quant_plugin, + 'load_model_on_cpu': args.load_model_on_cpu, } @@ -227,6 +228,7 @@ def convert_and_save_hf(args): dtype=args.dtype, mapping=mapping, quant_config=quant_config, + device='cpu' if args.load_model_on_cpu else 'cuda', calib_dataset=args.calib_dataset, **override_fields) else: