Star208
313
代码
代码
Issues59
Pull Requests91
项目讨论
Wiki
分析
Star208
313
已合并
delete old code of ckpt #1837
温一盏创建于 2024年11月1日
delete old code of ckpt #1837
已合并
温一盏创建于 2024年11月1日
从refs/pull/1837/head合入到master
共 10 个文件变更+33-1656
Mexamples/legacy/llama2/ckpt_convert_llama2_hf2legacy.sh+4-3
@@ -4,10 +4,11 @@ source /usr/local/Ascend/ascend-toolkit/set_env.sh
4# 权重格式转换，设置需要的并行配置，--num-layers-per-virtual-pipeline-stage 5，--params-dtype bf16 结合需要使用4# 权重格式转换，设置需要的并行配置，--num-layers-per-virtual-pipeline-stage 5，--params-dtype bf16 结合需要使用
5python convert_ckpt.py \5python convert_ckpt.py \
6 --model-type GPT \6 --model-type GPT \
7- --loader llama2_hf \7+ --load-model-type hf \
8- --saver megatron \8+ --save-model-type mg \
9 --target-tensor-parallel-size 1 \9 --target-tensor-parallel-size 1 \
10 --target-pipeline-parallel-size 2 \10 --target-pipeline-parallel-size 2 \
11 --load-dir ./model_from_hf/llama-2-hf/ \11 --load-dir ./model_from_hf/llama-2-hf/ \
12 --save-dir ./model_weights/llama-2-legacy/ \12 --save-dir ./model_weights/llama-2-legacy/ \
13- --tokenizer-model ./model_from_hf/llama-2-hf/tokenizer.model13+ --tokenizer-model ./model_from_hf/llama-2-hf/tokenizer.model \
14+ --model-type-hf llama2
Mexamples/legacy/llama2/ckpt_convert_llama2_legacy2hf.sh+2-3
@@ -4,9 +4,8 @@ source /usr/local/Ascend/ascend-toolkit/set_env.sh
4# --num-layers-per-virtual-pipeline-stage 5 \ 结合需要使用4# --num-layers-per-virtual-pipeline-stage 5 \ 结合需要使用
5python convert_ckpt.py \5python convert_ckpt.py \
6 --model-type GPT \6 --model-type GPT \
7- --loader megatron \7+ --load-model-type mg \
8- --saver megatron \8+ --save-model-type hf \
9- --save-model-type save_huggingface_llama \
10 --load-dir ./model_weights/llama2-legacy/ \9 --load-dir ./model_weights/llama2-legacy/ \
11 --target-tensor-parallel-size 1 \10 --target-tensor-parallel-size 1 \
12 --target-pipeline-parallel-size 1 \11 --target-pipeline-parallel-size 1 \
Mexamples/legacy/llama2/ckpt_convert_llama2_legacy2hf_lora.sh+2-3
@@ -3,9 +3,8 @@ source /usr/local/Ascend/ascend-toolkit/set_env.sh
3 3 
4python convert_ckpt.py \4python convert_ckpt.py \
5 --model-type GPT \5 --model-type GPT \
6- --loader megatron \6+ --load-model-type mg \
7- --saver megatron \7+ --save-model-type hf \
8- --save-model-type save_huggingface_llama \
9 --load-dir ./model_weights/llama-2-7b-hf-v0.1-tp8-pp1/ \8 --load-dir ./model_weights/llama-2-7b-hf-v0.1-tp8-pp1/ \
10 --lora-load ./ckpt/llama-2-7b-lora \9 --lora-load ./ckpt/llama-2-7b-lora \
11 --lora-r 8 \10 --lora-r 8 \
Mexamples/legacy/llama2/ckpt_convert_llama2_legacy2legacy_lora.sh+2-2
@@ -3,8 +3,8 @@ source /usr/local/Ascend/ascend-toolkit/set_env.sh
3 3 
4python convert_ckpt.py \4python convert_ckpt.py \
5 --model-type GPT \5 --model-type GPT \
6- --loader megatron \6+ --load-model-type mg \
7- --saver megatron \7+ --save-model-type mg \
8 --load-dir ./model_weights/llama-2-7b-hf-v0.1-tp8-pp1/ \8 --load-dir ./model_weights/llama-2-7b-hf-v0.1-tp8-pp1/ \
9 --lora-load ./ckpt/llama-2-7b-lora \9 --lora-load ./ckpt/llama-2-7b-lora \
10 --lora-r 8 \10 --lora-r 8 \
Dmodellink/tasks/checkpoint/loader_llama2_hf.py+0-440
@@ -1,440 +0,0 @@
1-# coding=utf-8
2-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
3-#
4-# Licensed under the Apache License, Version 2.0 (the "License");
5-# you may not use this file except in compliance with the License.
6-# You may obtain a copy of the License at
7-#
8-# http://www.apache.org/licenses/LICENSE-2.0
9-#
10-# Unless required by applicable law or agreed to in writing, software
11-# distributed under the License is distributed on an "AS IS" BASIS,
12-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-# See the License for the specific language governing permissions and
14-# limitations under the License.
15-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
16- 
17-import json
18-import os
19-import sys
20-import types
21-import torch
22-import transformers
23-from tqdm import tqdm
24- 
25- 
26-def add_arguments(parser):
27- group = parser.add_argument_group(title='Llama-2 HF loader.')
28- 
29- group.add_argument('--true-vocab-size', type=int, default=None,
30- help='original size of vocab, if specified will trim padding from embedding table.')
31- group.add_argument('--vocab-file', type=str, default=None,
32- help='Path to the vocab file. If specified will use this to get vocab size and '
33- 'trim padding from the embedding table.')
34- group.add_argument('--tokenizer-model', required=True,
35- help='Sentencepiece tokenizer model.')
36- group.add_argument('--megatron-path', type=str, default=None,
37- help='Base directory of deepspeed repository')
38- group.add_argument("--w-pack", type=bool,
39- help='True is w_pack weight for llm',
40- default=False)
41- parser.add_argument('--add-qkv-bias', action='store_true',
42- help='Add bias for attention qkv',
43- default=False)
44- parser.add_argument('--add-dense-bias', action='store_true',
45- help='Add bias for attention dense',
46- default=False)
47- parser.add_argument('--params-dtype', type=str,
48- help='Set weight dtype',
49- default='fp16')
50- group.add_argument('--make-vocab-size-divisible-by', type=int, default=1,
51- help='Pad the vocab size to be divisible by this value.'
52- 'This is added for computational efficieny reasons.')
53- 
54- 
55-def verify_transformers_version():
56- major, minor, patch = map(int, transformers.__version__.split('.'))
57- if major < 4 or minor < 31:
58- raise ValueError("the version transformers should greater or equal 4.31")
59- 
60- 
61-def load_args_from_checkpoint(args):
62- # Read Llama args.
63- llama_args_path = os.path.join(args.load, "config.json")
64- with open(llama_args_path) as f:
65- llama_args = json.load(f)
66- 
67- # Update Megatron args.
68- args.seq_length = 4096
69- args.max_position_embeddings = 4096
70- args.hidden_size = llama_args["hidden_size"]
71- args.num_attention_heads = llama_args["num_attention_heads"]
72- args.num_layers = llama_args["num_hidden_layers"]
73- args.global_batch_size = 1024
74- args.norm_epsilon = llama_args["rms_norm_eps"]
75- args.iteration = 1 # '0', 'release' don't work
76- args.add_position_embedding = True
77- args.use_rotary_position_embeddings = True
78- args.swiglu = True
79- args.tokenizer_type = "Llama2Tokenizer"
80- args.normalization = "RMSNorm"
81- args.add_bias_linear = False
82- args.untie_embeddings_and_output_weights = not llama_args.get("tie_word_embeddings", False)
83- args.vocab_size = llama_args["vocab_size"]
84- args.padded_vocab_size = llama_args["vocab_size"]
85- args.llama = llama_args
86- args.ffn_hidden_size = llama_args["intermediate_size"]
87- args.gradient_accumulation_fusion = False
88- if args.add_dense_bias:
89- args.skip_bias_add = False
90- 
91- if "num_key_value_heads" in llama_args \
92- and llama_args["num_attention_heads"] != llama_args["num_key_value_heads"] \
93- and llama_args["num_key_value_heads"] != 1:
94- args.group_query_attention = True
95- args.num_query_groups = llama_args["num_key_value_heads"]
96- 
97- 
98-def set_preprocess_state(args, model, hf_model):
99- '''Set embedding params.'''
100- model.language_model.embedding.word_embeddings.weight.data.copy_(
101- hf_model.model.embed_tokens.weight)
102- 
103- 
104-def set_postprocess_state(args, model, hf_model):
105- '''Set output layer & norm params.'''
106- model.language_model.encoder.final_norm.weight.data.copy_(hf_model.model.norm.weight)
107- if args.untie_embeddings_and_output_weights:
108- model.language_model.output_layer.weight.data.copy_(hf_model.lm_head.weight)
109- 
110- 
111-def set_attn_state(args, layer, hf_layer):
112- '''Set self-attention params.'''
113- 
114- # Get attention layer & state.
115- attn = layer.self_attention
116- hf_attn = hf_layer.self_attn
117- 
118- # Reshape loaded weights.
119- nh = args.num_attention_heads
120- ng = (args.num_query_groups if args.group_query_attention \
121- else args.num_attention_heads)
122- dim = args.kv_channels
123- if not nh % ng == 0:
124- raise ValueError("nh % ng should equal 0")
125- 
126- if args.w_pack:
127- w_pack = hf_attn.W_pack.weight
128- wq, wk, wv = w_pack.chunk(3, dim=0)
129- attn.query_key_value.weight.data.copy_(torch.cat([
130- wq.reshape((ng, dim * nh // ng, -1)),
131- wk.reshape((ng, dim, -1)),
132- wv.reshape((ng, dim, -1)),
133- ], dim=1).reshape((-1, args.hidden_size)))
134- 
135- else:
136- attn.query_key_value.weight.data.copy_(torch.cat([
137- hf_attn.q_proj.weight.reshape((ng, dim * nh // ng, -1)),
138- hf_attn.k_proj.weight.reshape((ng, dim, -1)),
139- hf_attn.v_proj.weight.reshape((ng, dim, -1)),
140- ], dim=1).reshape((-1, args.hidden_size)))
141- 
142- if args.add_qkv_bias:
143- attn.query_key_value.bias.data.copy_(torch.cat([
144- hf_attn.q_proj.bias.reshape((ng, dim * nh // ng)),
145- hf_attn.k_proj.bias.reshape((ng, dim)),
146- hf_attn.v_proj.bias.reshape((ng, dim)),
147- ], dim=1).reshape((-1)))
148- 
149- if args.add_dense_bias:
150- attn.dense.bias.data.copy_(hf_attn.o_proj.bias)
151- 
152- attn.dense.weight.data.copy_(hf_attn.o_proj.weight)
153- 
154- 
155-def set_mlp_state(args, layer, hf_layer):
156- '''Set MLP params.'''
157- 
158- mlp = layer.mlp
159- hf_mlp = hf_layer.mlp
160- 
161- mlp.dense_h_to_4h.weight.data.copy_(torch.cat([
162- hf_mlp.gate_proj.weight,
163- hf_mlp.up_proj.weight,
164- ], dim=0))
165- mlp.dense_4h_to_h.weight.data.copy_(hf_mlp.down_proj.weight)
166- 
167- 
168-def set_layer_state(args, model, hf_model, layer_idx):
169- '''Set transformer layer params.'''
170- 
171- layer = model.language_model.encoder.layers[layer_idx]
172- hf_layer = hf_model.model.layers[layer_idx]
173- 
174- set_attn_state(args, layer, hf_layer)
175- set_mlp_state(args, layer, hf_layer)
176- layer.input_norm.weight.data.copy_(hf_layer.input_layernorm.weight)
177- layer.post_attention_norm.weight.data.copy_(hf_layer.post_attention_layernorm.weight)
178- 
179- 
180-def load_checkpoint_to_model(model_provider, args):
181- '''Set model params.'''
182- 
183- from transformers import AutoModelForCausalLM
184- 
185- # Load Huggingface model.
186- hf_model = AutoModelForCausalLM.from_pretrained(args.load, device_map="cpu", trust_remote_code=True)
187- 
188- # Init Megatron model.
189- model = model_provider(True, True).to(args.params_dtype)
190- 
191- # Set model state.
192- set_preprocess_state(args, model, hf_model)
193- set_postprocess_state(args, model, hf_model)
194- for layer_idx in tqdm(range(args.num_layers), "set layer states"):
195- set_layer_state(args, model, hf_model, layer_idx)
196- 
197- return model
198- 
199- 
200-def _load_checkpoint(model_provider, queue, args):
201- # Llama-2 requires HF transformers >=4.31.0.
202- verify_transformers_version()
203- 
204- # Search in directory above this.
205- sys.path.append(os.path.abspath(
206- os.path.join(os.path.dirname(__file__),
207- os.path.pardir,
208- os.path.pardir)))
209- if args.megatron_path is not None:
210- sys.path.insert(0, args.megatron_path)
211- 
212- from megatron.training.arguments import validate_args
213- from modellink.training.utils import parse_args
214- from megatron.training.global_vars import set_global_variables
215- from megatron.legacy.model import module
216- from megatron.core import mpu
217- from megatron.core.enums import ModelType
218- 
219- # We want all arguments to come from us.
220- sys.argv = ['script.py',
221- '--no-masked-softmax-fusion',
222- '--no-bias-gelu-fusion',
223- '--no-bias-dropout-fusion',
224- '--no-async-tensor-model-parallel-allreduce',
225- '--use-cpu-initialization',
226- '--micro-batch-size', '1',
227- '--no-load-optim',
228- '--no-load-rng',
229- '--no-save-optim',
230- '--no-save-rng',
231- '--no-initialization',
232- '--load', args.load_dir
233- ]
234- 
235- margs = parse_args()
236- margs.w_pack = args.w_pack
237- margs.add_qkv_bias = args.add_qkv_bias
238- margs.add_dense_bias = args.add_dense_bias
239- margs.tokenizer_model = args.tokenizer_model
240- margs.make_vocab_size_divisible_by = args.make_vocab_size_divisible_by
241- if args.params_dtype == 'bf16':
242- margs.bf16 = True
243- elif args.params_dtype == 'fp16':
244- margs.fp16 = True
245- load_args_from_checkpoint(margs)
246- 
247- # Arguments do sanity checks on the world size, but we don't care,
248- # so trick it into thinking we are plenty of processes.
249- margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size
250- 
251- margs = validate_args(margs)
252- 
253- def check_for_arg(arg_name, default=None):
254- if getattr(margs, arg_name, None) is None:
255- if default is not None:
256- setattr(margs, arg_name, default)
257- else:
258- print(f"Checkpoint does not specify the argument {arg_name}. Exiting.")
259- print(f"Arguments: {margs}")
260- queue.put("exit")
261- exit(1)
262- 
263- check_for_arg('tensor_model_parallel_size')
264- check_for_arg('pipeline_model_parallel_size')
265- check_for_arg('num_layers')
266- check_for_arg('hidden_size')
267- check_for_arg('seq_length')
268- check_for_arg('num_attention_heads')
269- check_for_arg('max_position_embeddings')
270- check_for_arg('position_embedding_type')
271- check_for_arg('tokenizer_type')
272- check_for_arg('iteration')
273- check_for_arg('bert_binary_head')
274- check_for_arg('disable_bias_linear', False)
275- check_for_arg('params_dtype')
276- check_for_arg('swiglu', False)
277- 
278- # Determine how to make our models.
279- if not args.model_type == 'GPT':
280- raise ValueError("Llama-2 is a GPT model.")
281- margs.model_type = ModelType.encoder_or_decoder
282- 
283- # Suppress warning about torch.distributed not being initialized.
284- module.MegatronModule.embedding_warning_printed = True
285- 
286- set_global_variables(margs, build_tokenizer=False)
287- mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
288- mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
289- mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
290- 
291- # Short aliases.
292- tp_size = margs.tensor_model_parallel_size
293- pp_size = margs.pipeline_model_parallel_size
294- vp_size = margs.virtual_pipeline_model_parallel_size
295- if vp_size is None:
296- vp_size = 1
297- 
298- # Metadata.
299- md = types.SimpleNamespace()
300- md.model_type = args.model_type
301- md.spec = args.spec
302- md.num_layers = margs.num_layers
303- md.hidden_size = margs.hidden_size
304- md.seq_length = margs.seq_length
305- md.num_attention_heads = margs.num_attention_heads
306- md.max_position_embeddings = margs.max_position_embeddings
307- md.tokenizer_type = margs.tokenizer_type
308- md.iteration = margs.iteration
309- md.params_dtype = margs.params_dtype
310- md.bert_binary_head = margs.bert_binary_head
311- md.output_layer = margs.untie_embeddings_and_output_weights
312- md.position_embedding_type = margs.position_embedding_type
313- md.linear_bias = margs.add_bias_linear
314- md.norm_has_bias = False
315- if args.loader in ['loader_bloom_hf', 'bloom_hf']:
316- md.norm_has_bias = True
317- md.swiglu = margs.swiglu
318- md.previous_tensor_parallel_size = margs.tensor_model_parallel_size
319- md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size
320- md.true_vocab_size = margs.vocab_size # skips padding in saver
321- md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by
322- md.checkpoint_args = margs
323- md.consumed_train_samples = 0
324- md.consumed_valid_samples = 0
325- md.embed_layernorm = margs.embed_layernorm
326- 
327- # Get first pipe stage.
328- mpu.set_tensor_model_parallel_rank(0)
329- mpu.set_pipeline_model_parallel_rank(0)
330- model = load_checkpoint_to_model(model_provider, margs)
331- 
332- queue.put(md)
333- 
334- def queue_put(name, msg):
335- print(f"sending {name}")
336- msg["name"] = name
337- queue.put(msg)
338- 
339- # Send embeddings.
340- message = {
341- "word embeddings": model.language_model.embedding.word_embeddings.weight.data
342- }
343- 
344- # bloom
345- if hasattr(model.language_model.embedding.word_embeddings, 'norm'):
346- message["word embeddings norm_w"] = model.language_model.embedding.word_embeddings.norm.weight.data
347- message["word embeddings norm_b"] = model.language_model.embedding.word_embeddings.norm.bias.data
348- 
349- if md.position_embedding_type == 'learned_absolute':
350- message["position embeddings"] = model.language_model.embedding.position_embeddings.weight.data
351- else:
352- if hasattr(model.language_model.embedding, 'position_embeddings'):
353- raise ValueError("model should have position_embeddings")
354- 
355- queue_put("embeddings", message)
356- 
357- for layer_num in range(margs.num_layers):
358- message = {}
359- 
360- # Get non-parallel tensors from tp_rank 0.
361- layer = model.language_model.encoder.layers[layer_num]
362- message["input norm weight"] = layer.input_norm.weight.data
363- message["post norm weight"] = layer.post_attention_norm.weight.data
364- if md.linear_bias:
365- message["dense bias"] = layer.self_attention.dense.bias.data
366- message["mlp l1 bias"] = layer.mlp.dense_4h_to_h.bias.data
367- 
368- if md.norm_has_bias:
369- message["input norm bias"] = layer.input_norm.bias.data
370- message["post norm bias"] = layer.post_attention_norm.bias.data
371- 
372- # Grab all parallel tensors for this layer.
373- qkv_weight = []
374- qkv_bias = []
375- dense_weight = []
376- mlp_l0_weight = []
377- mlp_l0_bias = []
378- mlp_l1_weight = []
379- layer = model.language_model.encoder.layers[layer_num]
380- qkv_weight.append(layer.self_attention.query_key_value.weight.data)
381- dense_weight.append(layer.self_attention.dense.weight.data)
382- mlp_l0_weight.append(layer.mlp.dense_h_to_4h.weight.data)
383- mlp_l1_weight.append(layer.mlp.dense_4h_to_h.weight.data)
384- if md.linear_bias:
385- qkv_bias.append(layer.self_attention.query_key_value.bias.data)
386- mlp_l0_bias.append(layer.mlp.dense_h_to_4h.bias.data)
387- if args.add_qkv_bias:
388- message["qkv bias"] = layer.self_attention.query_key_value.bias.data
389- if args.add_dense_bias:
390- message["dense bias"] = layer.self_attention.dense.bias.data
391- # Handle gated linear units.
392- if md.swiglu:
393- # Concat all the first halves ('W's) and all the second halves ('V's).
394- for tp_rank in range(tp_size):
395- mlp_l0_weight[tp_rank] = torch.chunk(mlp_l0_weight[tp_rank], 2, dim=0)
396- message["mlp l0 weight W"] = torch.cat([w[0] for w in mlp_l0_weight], dim=0)
397- message["mlp l0 weight V"] = torch.cat([w[1] for w in mlp_l0_weight], dim=0)
398- else:
399- message["mlp l0 weight"] = torch.cat(mlp_l0_weight, dim=0)
400- 
401- # Simple concat of the rest.
402- message["qkv weight"] = torch.cat(qkv_weight, dim=0)
403- message["dense weight"] = torch.cat(dense_weight, dim=1)
404- message["mlp l1 weight"] = torch.cat(mlp_l1_weight, dim=1)
405- if md.linear_bias:
406- message["qkv bias"] = torch.cat(qkv_bias, dim=0)
407- if md.swiglu:
408- for tp_rank in range(tp_size):
409- mlp_l0_bias[tp_rank] = torch.chunk(mlp_l0_bias[tp_rank], 2, dim=0)
410- message["mlp l0 bias W"] = torch.cat([b[0] for b in mlp_l0_bias], dim=0)
411- message["mlp l0 bias V"] = torch.cat([b[1] for b in mlp_l0_bias], dim=0)
412- else:
413- message["mlp l0 bias"] = torch.cat(mlp_l0_bias, dim=0)
414- 
415- queue_put(f"transformer layer {layer_num}", message)
416- 
417- # Send final norm from tp_rank 0.
418- message = {
419- "weight": model.language_model.encoder.final_norm.weight.data,
420- }
421- if md.norm_has_bias:
422- message["bias"] = model.language_model.encoder.final_norm.bias.data
423- queue_put("final norm", message)
424- 
425- if md.output_layer:
426- message = {
427- "weight": model.language_model.output_layer.weight.data
428- }
429- queue_put("output layer", message)
430- 
431- queue.put("done")
432- 
433- 
434-def load_checkpoint(model_provider, queue, args):
435- try:
436- _load_checkpoint(model_provider, queue, args)
437- except:
438- queue.put("exit")
439- raise
440- 
Dmodellink/tasks/checkpoint/loader_megatron.py+0-409
@@ -1,409 +0,0 @@
1-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
2- 
3-import json
4-import os
5-import sys
6-import types
7- 
8-import torch
9- 
10- 
11-def add_arguments(parser):
12- group = parser.add_argument_group(title='Megatron loader')
13- 
14- group.add_argument('--true-vocab-size', type=int, default=None,
15- help='original size of vocab, if specified will trim padding from embedding table.')
16- group.add_argument('--vocab-file', type=str, default=None,
17- help='Path to the vocab file. If specified will use this to get vocab size and '
18- 'trim padding from the embedding table.')
19- group.add_argument('--megatron-path', type=str, default=None,
20- help='Base directory of deepspeed repository')
21- parser.add_argument('--add-qkv-bias', action='store_true',
22- help='Add bias for attention qkv', default=False,
23- )
24- parser.add_argument('--add-dense-bias', action='store_true',
25- help='Add bias for attention dense', default=False,
26- )
27- parser.add_argument('--embed-layernorm', action='store_true',
28- help='Add embed layernorm for word embedding', default=False,
29- )
30- parser.add_argument('--params-dtype', type=str,
31- help='Set weight dtype', default='fp16',
32- )
33- parser.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,
34- help='Number of layers per virtual pipeline stage')
35- group.add_argument('--lora-target-modules', nargs='+', type=str, default=[],
36- help='Lora target modules.')
37- group.add_argument('--lora-load', type=str, default=None,
38- help='Directory containing a lora model checkpoint.')
39- group.add_argument('--lora-r', type=int, default=16,
40- help='Lora r.')
41- group.add_argument('--lora-alpha', type=int, default=32,
42- help='Lora alpha.')
43- 
44- 
45-def _load_checkpoint(model_provider, queue, args):
46- 
47- # Search in directory above this
48- sys.path.append(os.path.abspath(
49- os.path.join(os.path.dirname(__file__),
50- os.path.pardir)))
51- if args.megatron_path is not None:
52- sys.path.insert(0, args.megatron_path)
53- 
54- try:
55- from modellink.training.utils import parse_args
56- from megatron.training.arguments import validate_args
57- from megatron.training.global_vars import set_args, set_global_variables
58- from megatron.training.checkpointing import load_args_from_checkpoint
59- from megatron.training.checkpointing import load_checkpoint as load_checkpoint_mg
60- from megatron.legacy.model import module
61- from megatron.core import mpu
62- from megatron.core.enums import ModelType
63- from modellink.training.checkpointing import load_checkpoint_wrapper
64- load_checkpoint_mg = load_checkpoint_wrapper(load_checkpoint_mg)
65- except ModuleNotFoundError:
66- print("Unable to import Megatron, please specify the path to Megatron using --megatron-path. Exiting.")
67- queue.put("exit")
68- 
69- # We want all arguments to come from us
70- sys.argv = ['script.py',
71- '--no-masked-softmax-fusion',
72- '--no-bias-gelu-fusion',
73- '--no-bias-dropout-fusion',
74- '--no-async-tensor-model-parallel-allreduce',
75- '--use-cpu-initialization',
76- '--micro-batch-size', '1',
77- '--no-load-optim',
78- '--no-load-rng',
79- '--no-save-optim',
80- '--no-save-rng',
81- '--no-initialization',
82- '--load', args.load_dir,
83- '--finetune'
84- ]
85- 
86- margs = parse_args()
87- margs.embed_layernorm = args.embed_layernorm
88- set_args(margs)
89- margs, checkpoint_args = load_args_from_checkpoint(margs)
90- margs.add_qkv_bias = args.add_qkv_bias
91- margs.add_dense_bias = args.add_dense_bias
92- margs.num_layers_per_virtual_pipeline_stage = args.num_layers_per_virtual_pipeline_stage
93- if args.add_dense_bias:
94- margs.skip_bias_add = False
95- if args.params_dtype == 'bf16':
96- margs.bf16 = True
97- elif args.params_dtype == 'fp16':
98- margs.fp16 = True
99- 
100- margs.lora_target_modules = args.lora_target_modules
101- margs.lora_load = args.lora_load
102- margs.lora_r = args.lora_r
103- margs.lora_alpha = args.lora_alpha
104- 
105- # Arguments do sanity checks on the world size, but we don't care,
106- # so trick it into thinking we are plenty of processes
107- margs.world_size = margs.tensor_model_parallel_size * margs.pipeline_model_parallel_size
108- 
109- margs = validate_args(margs)
110- 
111- def check_for_arg(arg_name, default=None):
112- if getattr(margs, arg_name, None) is None:
113- if default is not None:
114- setattr(margs, arg_name, default)
115- else:
116- print(f"Checkpoint does not specify the argument {arg_name}. Exiting.")
117- print(f"Arguments: {margs}")
118- queue.put("exit")
119- 
120- check_for_arg('tensor_model_parallel_size')
121- check_for_arg('pipeline_model_parallel_size')
122- check_for_arg('num_layers')
123- check_for_arg('hidden_size')
124- check_for_arg('seq_length')
125- check_for_arg('num_attention_heads')
126- check_for_arg('max_position_embeddings')
127- check_for_arg('position_embedding_type')
128- check_for_arg('tokenizer_type')
129- check_for_arg('iteration')
130- check_for_arg('bert_binary_head')
131- check_for_arg('disable_bias_linear', False)
132- check_for_arg('params_dtype')
133- check_for_arg('swiglu', False)
134- 
135- # Determine how to make our models
136- if args.model_type == 'GPT':
137- margs.model_type = ModelType.encoder_or_decoder
138- elif args.model_type == 'BERT':
139- margs.model_type = ModelType.encoder_or_decoder
140- else:
141- raise Exception(f'unrecognized model type: {args.model_type}')
142- 
143- from modellink.training import model_provider_func_wrapper
144- model_provider = model_provider_func_wrapper(model_provider)
145- 
146- # supress warning about torch.distributed not being initialized
147- module.MegatronModule.embedding_warning_printed = True
148- 
149- consumed_train_samples = None
150- consumed_valid_samples = None
151- 
152- def get_models(count, dtype):
153- nonlocal consumed_train_samples
154- nonlocal consumed_valid_samples
155- model_array_len = margs.virtual_pipeline_model_parallel_size
156- if model_array_len is None:
157- model_array_len = 1
158- models = [[] for _ in range(model_array_len)]
159- pre_process = mpu.is_pipeline_first_stage()
160- post_process = mpu.is_pipeline_last_stage()
161- for rank in range(count):
162- mpu.set_tensor_model_parallel_rank(rank)
163- if margs.virtual_pipeline_model_parallel_size is not None:
164- model_ = []
165- for i in range(margs.virtual_pipeline_model_parallel_size):
166- mpu.set_virtual_pipeline_model_parallel_rank(i)
167- # Set pre_process and post_process only after virtual rank is set.
168- pre_process = mpu.is_pipeline_first_stage()
169- post_process = mpu.is_pipeline_last_stage()
170- this_model = model_provider(
171- pre_process=pre_process,
172- post_process=post_process
173- ).to(dtype)
174- model_.append(this_model)
175- else:
176- pre_process = mpu.is_pipeline_first_stage()
177- post_process = mpu.is_pipeline_last_stage()
178- model_rank = 0
179- model_ = [model_provider(pre_process, post_process).to(dtype)]
180- margs.consumed_train_samples = 0
181- margs.consumed_valid_samples = 0
182- load_checkpoint_mg(model_, None, None)
183- 
184- if margs.lora_target_modules:
185- for model_item in model_:
186- model_item.merge_and_unload()
187- if consumed_train_samples is not None:
188- if margs.consumed_train_samples != consumed_train_samples:
189- return None
190- else:
191- consumed_train_samples = margs.consumed_train_samples
192- if consumed_valid_samples is not None:
193- if margs.consumed_valid_samples != consumed_valid_samples:
194- return None
195- else:
196- consumed_valid_samples = margs.consumed_valid_samples
197- for vp_rank in range(model_array_len):
198- models[vp_rank].append(model_[vp_rank])
199- return models
200- 
201- set_args(margs)
202- mpu.set_tensor_model_parallel_world_size(margs.tensor_model_parallel_size)
203- mpu.set_pipeline_model_parallel_world_size(margs.pipeline_model_parallel_size)
204- mpu.set_virtual_pipeline_model_parallel_world_size(margs.virtual_pipeline_model_parallel_size)
205- 
206- # Get true (non-padded) vocab size
207- if args.true_vocab_size is not None:
208- true_vocab_size = args.true_vocab_size
209- elif args.vocab_file is not None:
210- vb_file = open(args.vocab_file)
211- vocab = json.load(vb_file)
212- true_vocab_size = len(vocab)
213- if args.true_vocab_size is not None and true_vocab_size != args.true_vocab_size:
214- print("Both --true-vocab-size and --vocab-file specified and the vocab size does not match, aborting.")
215- queue.put("exit")
216- vb_file.close()
217- else:
218- true_vocab_size = None
219- 
220- # short aliases
221- tp_size = margs.tensor_model_parallel_size
222- pp_size = margs.pipeline_model_parallel_size
223- vp_size = margs.virtual_pipeline_model_parallel_size
224- if vp_size is None:
225- vp_size = 1
226- 
227- # Layernorm has bias; RMSNorm does not.
228- if hasattr(checkpoint_args, 'normalization'):
229- norm_has_bias = checkpoint_args.normalization == "LayerNorm"
230- else:
231- # older models only supported LayerNorm
232- norm_has_bias = True
233- 
234- # metadata
235- md = types.SimpleNamespace()
236- md.model_type = args.model_type
237- md.num_layers = margs.num_layers
238- md.hidden_size = margs.hidden_size
239- md.seq_length = margs.seq_length
240- md.num_attention_heads = margs.num_attention_heads
241- md.max_position_embeddings = margs.max_position_embeddings
242- md.tokenizer_type = margs.tokenizer_type
243- md.iteration = margs.iteration
244- md.params_dtype = margs.params_dtype
245- md.bert_binary_head = margs.bert_binary_head
246- md.output_layer = margs.untie_embeddings_and_output_weights
247- md.position_embedding_type = margs.position_embedding_type
248- md.linear_bias = margs.add_bias_linear
249- md.norm_has_bias = norm_has_bias
250- md.swiglu = margs.swiglu
251- md.previous_tensor_parallel_size = margs.tensor_model_parallel_size
252- md.previous_pipeline_parallel_size = margs.pipeline_model_parallel_size
253- md.true_vocab_size = true_vocab_size
254- md.make_vocab_size_divisible_by = margs.make_vocab_size_divisible_by
255- md.checkpoint_args = checkpoint_args
256- md.embed_layernorm = margs.embed_layernorm
257- 
258- # Get first pipe stage
259- mpu.set_pipeline_model_parallel_rank(0)
260- all_models = [get_models(tp_size, md.params_dtype)]
261- models = all_models[0][0]
262- 
263- md.consumed_train_samples = consumed_train_samples
264- md.consumed_valid_samples = consumed_valid_samples
265- queue.put(md)
266- 
267- def queue_put(name, msg):
268- print(f"sending {name}")
269- msg["name"] = name
270- queue.put(msg)
271- 
272- # Send embeddings
273- message = {
274- "word embeddings": torch.cat(
275- [models[tp_rank].language_model.embedding.word_embeddings.weight.data for tp_rank in range(tp_size)],
276- dim=0)
277- }
278- if md.position_embedding_type == 'learned_absolute':
279- message["position embeddings"] = models[0].language_model.embedding.position_embeddings.weight.data
280- if md.embed_layernorm:
281- message["word embeddings norm_w"] = models[0].language_model.embedding.word_embeddings.norm.weight.data
282- message["word embeddings norm_b"] = models[0].language_model.embedding.word_embeddings.norm.bias.data
283- queue_put("embeddings", message)
284- 
285- total_layer_num = 0
286- for vp_rank in range(vp_size):
287- mpu.set_virtual_pipeline_model_parallel_rank(vp_rank)
288- for pp_rank in range(pp_size):
289- if pp_rank > 0:
290- mpu.set_pipeline_model_parallel_rank(pp_rank)
291- if vp_rank == 0:
292- all_models.append(get_models(tp_size, md.params_dtype))
293- models = all_models[pp_rank][vp_rank]
294- for layer_num, _ in enumerate(models[0].language_model.encoder.layers):
295- message = {}
296- 
297- # Get non-parallel tensors from tp_rank 0
298- layer = models[0].language_model.encoder.layers[layer_num]
299- message["input norm weight"] = layer.input_norm.weight.data
300- if norm_has_bias:
301- message["input norm bias"] = layer.input_norm.bias.data
302- message["post norm weight"] = layer.post_attention_norm.weight.data
303- if norm_has_bias:
304- message["post norm bias"] = layer.post_attention_norm.bias.data
305- if md.linear_bias:
306- message["dense bias"] = layer.self_attention.dense.bias.data
307- message["mlp l1 bias"] = layer.mlp.dense_4h_to_h.bias.data
308- if args.add_dense_bias:
309- message["dense bias"] = layer.self_attention.dense.bias.data
310- 
311- # Grab all parallel tensors for this layer
312- qkv_weight = []
313- qkv_bias = []
314- dense_weight = []
315- mlp_l0_weight = []
316- mlp_l1_weight = []
317- mlp_l0_bias = []
318- 
319- for tp_rank, model in enumerate(models):
320- layer = model.language_model.encoder.layers[layer_num]
321- qkv_weight.append(layer.self_attention.query_key_value.weight.data)
322- dense_weight.append(layer.self_attention.dense.weight.data)
323- mlp_l0_weight.append(layer.mlp.dense_h_to_4h.weight.data)
324- mlp_l1_weight.append(layer.mlp.dense_4h_to_h.weight.data)
325- 
326- if md.linear_bias:
327- qkv_bias.append(layer.self_attention.query_key_value.bias.data)
328- mlp_l0_bias.append(layer.mlp.dense_h_to_4h.bias.data)
329- if args.add_qkv_bias:
330- qkv_bias.append(layer.self_attention.query_key_value.bias.data)
331-
332- # Handle gated linear units
333- if md.swiglu:
334- # concat all the first halves ('W's) and all the second halves ('V's)
335- for tp_rank in range(tp_size):
336- mlp_l0_weight[tp_rank] = torch.chunk(mlp_l0_weight[tp_rank], 2, dim=0)
337- message["mlp l0 weight W"] = torch.cat([w[0] for w in mlp_l0_weight], dim=0)
338- message["mlp l0 weight V"] = torch.cat([w[1] for w in mlp_l0_weight], dim=0)
339- else:
340- message["mlp l0 weight"] = torch.cat(mlp_l0_weight, dim=0)
341- 
342- # simple concat of the rest
343- message["qkv weight"] = torch.cat(qkv_weight, dim=0)
344- message["dense weight"] = torch.cat(dense_weight, dim=1)
345- message["mlp l1 weight"] = torch.cat(mlp_l1_weight, dim=1)
346- if md.linear_bias:
347- message["qkv bias"] = torch.cat(qkv_bias, dim=0)
348- if md.swiglu:
349- for tp_rank in range(tp_size):
350- mlp_l0_bias[tp_rank] = torch.chunk(mlp_l0_bias[tp_rank], 2, dim=0)
351- message["mlp l0 bias W"] = torch.cat([b[0] for b in mlp_l0_bias], dim=0)
352- message["mlp l0 bias V"] = torch.cat([b[1] for b in mlp_l0_bias], dim=0)
353- else:
354- message["mlp l0 bias"] = torch.cat(mlp_l0_bias, dim=0)
355- if args.add_qkv_bias:
356- message["qkv bias"] = torch.cat(qkv_bias, dim=0)
357- 
358- queue_put(f"transformer layer {total_layer_num}", message)
359- 
360- total_layer_num = total_layer_num + 1
361- 
362- # Send final norm from tp_rank 0
363- message = {
364- "weight": models[0].language_model.encoder.final_norm.weight.data,
365- }
366- if norm_has_bias:
367- message["bias"] = models[0].language_model.encoder.final_norm.bias.data
368- queue_put("final norm", message)
369- 
370- if md.output_layer:
371- message = {
372- "weight": torch.cat(
373- [models[tp_rank].language_model.output_layer.weight.data for tp_rank in range(tp_size)],
374- dim=0)
375- }
376- queue_put("output layer", message)
377- 
378- # Send BERT lm head and binary head if it exists
379- if md.model_type == 'BERT':
380- message = {
381- "weight": models[0].language_model.pooler.dense.weight.data,
382- "bias": models[0].language_model.pooler.dense.bias.data,
383- }
384- queue_put("pooler", message)
385- 
386- message = {
387- "dense weight": models[0].lm_head.dense.weight.data,
388- "dense bias": models[0].lm_head.dense.bias.data,
389- "norm weight": models[0].lm_head.norm.weight.data,
390- }
391- if norm_has_bias:
392- message["norm bias"] = models[0].lm_head.norm.bias.data
393- queue_put("lm head", message)
394- 
395- if md.bert_binary_head:
396- message = {
397- "weight": models[0].binary_head.weight.data,
398- "bias": models[0].binary_head.bias.data,
399- }
400- queue_put("binary head", message)
401- queue.put("done")
402- 
403- 
404-def load_checkpoint(model_provider, queue, args):
405- try:
406- _load_checkpoint(model_provider, queue, args)
407- except:
408- queue.put("exit")
409- raise
Dmodellink/tasks/checkpoint/saver_megatron.py+0-774
@@ -1,774 +0,0 @@
1-# coding=utf-8
2-# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
3-#
4-# Licensed under the Apache License, Version 2.0 (the "License");
5-# you may not use this file except in compliance with the License.
6-# You may obtain a copy of the License at
7-#
8-# http://www.apache.org/licenses/LICENSE-2.0
9-#
10-# Unless required by applicable law or agreed to in writing, software
11-# distributed under the License is distributed on an "AS IS" BASIS,
12-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13-# See the License for the specific language governing permissions and
14-# limitations under the License.
15-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
16- 
17-import os
18-import sys
19-import copy
20-import torch
21- 
22- 
23-def add_arguments(parser):
24- group = parser.add_argument_group(title='Megatron saver')
25- 
26- group.add_argument('--megatron-path', type=str, default=None,
27- help='Base directory of Megatron repository')
28- 
29- group.add_argument('--target-tensor-parallel-size', type=int,
30- help='Target tensor model parallel size, defaults to the tensor parallel size '
31- 'in the input checkpoint if provided by the loader, otherwise to 1')
32- group.add_argument('--target-pipeline-parallel-size', type=int,
33- help='Target tensor model parallel size, default to the pipeline parall size '
34- 'in the input checkpoint if provided by the loader, otherwise to 1')
35- group.add_argument('--save-model-type', type=str, default='megatron',
36- help='Save model type')
37- group.add_argument("--w-pack", type=bool,
38- help='True is w_pack weight for llm',
39- default=False)
40- group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,
41- help='Number of layers per virtual pipeline stage')
42- group.add_argument('--num-layer-list',
43- type=str, help='a list of number of layers, seperated by comma; e.g., 4,4,4,4')
44- 
45- 
46-def save_huggingface(args, model):
47- '''Set model params.'''
48- from transformers import BloomForCausalLM
49- 
50- # Load Huggingface model.
51- hf_model = BloomForCausalLM.from_pretrained(args.save_dir, device_map="cpu", torch_dtype="auto")
52- 
53- for name_param_h, name_param_m in zip(hf_model.named_parameters(), model.named_parameters()):
54- name_param_h[1].data.copy_(name_param_m[1])
55- 
56- save_dir = os.path.join(args.save_dir, 'mg2hf')
57- print(f'save weight to {save_dir}')
58- hf_model.save_pretrained(save_dir)
59- 
60- 
61-def save_huggingface_llama(args, model, model_args):
62- '''Set model params.'''
63- from transformers import AutoModelForCausalLM
64- 
65- # Load Huggingface model.
66- hf_model = AutoModelForCausalLM.from_pretrained(args.save_dir, device_map="cpu", trust_remote_code=True, torch_dtype="auto")
67- hf2mg_map = {}
68- for name_param_m in model.named_parameters():
69- layer_num = name_param_m[0].split(".")[3] if len(name_param_m[0].split(".")) > 3 else name_param_m[0].split(".")[1]
70- nh = model_args.num_attention_heads
71- ng = (
72- model_args.checkpoint_args.num_query_groups
73- if model_args.checkpoint_args.group_query_attention
74- else model_args.num_attention_heads
75- )
76- repeats = nh // ng
77- if name_param_m[0] == "language_model.embedding.word_embeddings.weight":
78- hf2mg_map["model.embed_tokens.weight"] = name_param_m[1]
79- continue
80- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.post_attention_norm.weight":
81- hf2mg_map[f"model.layers.{layer_num}.post_attention_layernorm.weight"] = name_param_m[1]
82- continue
83- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.input_norm.weight":
84- hf2mg_map[f"model.layers.{layer_num}.input_layernorm.weight"] = name_param_m[1]
85- continue
86- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.post_attention_norm.weight":
87- hf2mg_map[f"model.layers.{layer_num}.post_attention_layernorm.weight"] = name_param_m[1]
88- continue
89- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.weight":
90- qkv_weight = name_param_m[1].reshape(
91- ng,
92- repeats + 2,
93- name_param_m[1].shape[0] // ng // (repeats + 2),
94- name_param_m[1].shape[1],
95- )
96- w = qkv_weight.shape[-1]
97- qw = qkv_weight[:, :repeats, ...].reshape(-1, w)
98- kw = qkv_weight[:, repeats : repeats + 1, ...].reshape(-1, w)
99- vw = qkv_weight[:, repeats + 1 :, ...].reshape(-1, w)
100- if args.w_pack:
101- qkv = torch.cat((qw, kw, vw), dim=0)
102- hf2mg_map[f"model.layers.{layer_num}.self_attn.W_pack.weight"] = qkv
103- else:
104- hf2mg_map[f"model.layers.{layer_num}.self_attn.q_proj.weight"] = qw
105- hf2mg_map[f"model.layers.{layer_num}.self_attn.k_proj.weight"] = kw
106- hf2mg_map[f"model.layers.{layer_num}.self_attn.v_proj.weight"] = vw
107- continue
108- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.bias":
109- bias_weight = name_param_m[1].reshape(
110- ng, repeats + 2, name_param_m[1].shape[0] // ng // (repeats + 2)
111- )
112- w = bias_weight.shape[-1]
113- qw = bias_weight[:, :repeats, ...].reshape(-1)
114- kw = bias_weight[:, repeats : repeats + 1, ...].reshape(-1)
115- vw = bias_weight[:, repeats + 1 :, ...].reshape(-1)
116- hf2mg_map[f"model.layers.{layer_num}.self_attn.q_proj.bias"] = qw
117- hf2mg_map[f"model.layers.{layer_num}.self_attn.k_proj.bias"] = kw
118- hf2mg_map[f"model.layers.{layer_num}.self_attn.v_proj.bias"] = vw
119- continue
120- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.dense.bias":
121- hf2mg_map[f"model.layers.{layer_num}.self_attn.dense.bias"] = name_param_m[1]
122- continue
123- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.dense.weight":
124- hf2mg_map[f"model.layers.{layer_num}.self_attn.o_proj.weight"] = name_param_m[1]
125- continue
126- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_h_to_4h.weight":
127- proj_read_h_half = name_param_m[1].shape[0] // 2
128- hf2mg_map[f"model.layers.{layer_num}.mlp.gate_proj.weight"] = name_param_m[1][:proj_read_h_half, ...]
129- hf2mg_map[f"model.layers.{layer_num}.mlp.up_proj.weight"] = name_param_m[1][proj_read_h_half:, ...]
130- continue
131- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_4h_to_h.weight":
132- hf2mg_map[f"model.layers.{layer_num}.mlp.down_proj.weight"] = name_param_m[1]
133- continue
134- if name_param_m[0] == "language_model.encoder.final_norm.weight":
135- hf2mg_map[f"model.norm.weight"] = name_param_m[1]
136- continue
137- if name_param_m[0] == "language_model.output_layer.weight":
138- hf2mg_map[f"lm_head.weight"] = name_param_m[1]
139- continue
140- for name_param_h in hf_model.named_parameters():
141- if name_param_h[0] in hf2mg_map.keys():
142- name_param_h[1].data.copy_(hf2mg_map[name_param_h[0]])
143- 
144- save_dir = os.path.join(args.save_dir, 'mg2hf')
145- print(f'save weight to {save_dir}')
146- hf_model.save_pretrained(save_dir)
147- 
148- 
149-def save_huggingface_qwen(args, model, model_args):
150- """Set model params."""
151- from transformers import AutoModelForCausalLM
152- from accelerate import init_empty_weights
153- 
154- # Load Huggingface model.
155- with init_empty_weights():
156- hf_model = AutoModelForCausalLM.from_pretrained(
157- args.save_dir, device_map="cpu", trust_remote_code=True, torch_dtype="auto"
158- )
159- hf2mg_map = {}
160- for name_param_m in model.named_parameters():
161- layer_num = (
162- name_param_m[0].split(".")[3]
163- if len(name_param_m[0].split(".")) > 3
164- else name_param_m[0].split(".")[1]
165- )
166- nh = model_args.num_attention_heads
167- ng = (
168- model_args.checkpoint_args.num_query_groups
169- if model_args.checkpoint_args.group_query_attention
170- else model_args.num_attention_heads
171- )
172- repeats = nh // ng
173- if name_param_m[0] == "language_model.embedding.word_embeddings.weight":
174- hf2mg_map["transformer.wte.weight"] = name_param_m[1]
175- continue
176- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.post_attention_norm.weight":
177- hf2mg_map[f"transformer.h.{layer_num}.ln_2.weight"] = name_param_m[1]
178- continue
179- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.input_norm.weight":
180- hf2mg_map[f"transformer.h.{layer_num}.ln_1.weight"] = name_param_m[1]
181- continue
182- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.post_attention_norm.weight":
183- hf2mg_map[f"transformer.h.{layer_num}.ln_2.weight"] = name_param_m[1]
184- continue
185- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.weight":
186- qkv_weight = name_param_m[1].reshape(
187- ng,
188- repeats + 2,
189- name_param_m[1].shape[0] // ng // (repeats + 2),
190- name_param_m[1].shape[1],
191- )
192- w = qkv_weight.shape[-1]
193- qw = qkv_weight[:, :repeats, ...].reshape(-1, w)
194- kw = qkv_weight[:, repeats : repeats + 1, ...].reshape(-1, w)
195- vw = qkv_weight[:, repeats + 1 :, ...].reshape(-1, w)
196- qkv = torch.cat((qw, kw, vw), dim=0)
197- hf2mg_map[f"transformer.h.{layer_num}.attn.c_attn.weight"] = qkv
198- continue
199- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.bias":
200- bias_weight = name_param_m[1].reshape(
201- ng, repeats + 2, name_param_m[1].shape[0] // ng // (repeats + 2)
202- )
203- w = bias_weight.shape[-1]
204- qw = bias_weight[:, :repeats, ...].reshape(-1)
205- kw = bias_weight[:, repeats : repeats + 1, ...].reshape(-1)
206- vw = bias_weight[:, repeats + 1 :, ...].reshape(-1)
207- hf2mg_map[f"transformer.h.{layer_num}.attn.c_attn.bias"] = torch.cat(
208- [qw, kw, vw], dim=0
209- )
210- continue
211- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.dense.weight":
212- hf2mg_map[f"transformer.h.{layer_num}.attn.c_proj.weight"] = name_param_m[1]
213- continue
214- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_h_to_4h.weight":
215- proj_read_h_half = name_param_m[1].shape[0] // 2
216- hf2mg_map[f"transformer.h.{layer_num}.mlp.w2.weight"] = name_param_m[1][
217- :proj_read_h_half, ...
218- ]
219- hf2mg_map[f"transformer.h.{layer_num}.mlp.w1.weight"] = name_param_m[1][
220- proj_read_h_half:, ...
221- ]
222- continue
223- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_4h_to_h.weight":
224- hf2mg_map[f"transformer.h.{layer_num}.mlp.c_proj.weight"] = name_param_m[1]
225- continue
226- if name_param_m[0] == "language_model.encoder.final_norm.weight":
227- hf2mg_map[f"transformer.ln_f.weight"] = name_param_m[1]
228- continue
229- if name_param_m[0] == "language_model.output_layer.weight":
230- hf2mg_map[f"lm_head.weight"] = name_param_m[1]
231- continue
232- for name_param_h in hf_model.named_parameters():
233- if name_param_h[0] in hf2mg_map.keys():
234- name_param_h[1].data.copy_(hf2mg_map[name_param_h[0]])
235- 
236- save_dir = os.path.join(args.save_dir, "mg2hf")
237- print(f"save weight to {save_dir}")
238- hf_model.save_pretrained(save_dir)
239- 
240- 
241- 
242-def save_huggingface_chatglm3(args, model, model_args):
243- '''Set model params.'''
244- from transformers import AutoModelForCausalLM
245- 
246- # Load Huggingface model.
247- hf_model = AutoModelForCausalLM.from_pretrained(args.save_dir, device_map="cpu", trust_remote_code=True, torch_dtype="auto")
248- hf2mg_map = {}
249- for name_param_m in model.named_parameters():
250- layer_num = name_param_m[0].split(".")[3] if len(name_param_m[0].split(".")) > 3 else name_param_m[0].split(".")[1]
251- nh = model_args.num_attention_heads
252- ng = (
253- model_args.checkpoint_args.num_query_groups
254- if model_args.checkpoint_args.group_query_attention
255- else model_args.num_attention_heads
256- )
257- repeats = nh // ng
258- if name_param_m[0] == "language_model.embedding.word_embeddings.weight":
259- hf2mg_map["transformer.embedding.word_embeddings.weight"] = name_param_m[1]
260- continue
261- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.input_norm.weight":
262- hf2mg_map[f"model.layers.{layer_num}.input_layernorm.weight"] = name_param_m[1]
263- continue
264- # query_key_value weight
265- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.weight":
266- qkv_weight = name_param_m[1].reshape(
267- ng,
268- repeats + 2,
269- name_param_m[1].shape[0] // ng // (repeats + 2),
270- name_param_m[1].shape[1],
271- )
272- w = qkv_weight.shape[-1]
273- qw = qkv_weight[:, :repeats, ...].reshape(-1, w)
274- kw = qkv_weight[:, repeats : repeats + 1, ...].reshape(-1, w)
275- vw = qkv_weight[:, repeats + 1 :, ...].reshape(-1, w)
276- qkv = torch.cat((qw, kw, vw), dim=0)
277- hf2mg_map[f"transformer.encoder.layers.{layer_num}.self_attention.query_key_value.weight"] = qkv
278- continue
279- # query_key_value bias
280- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.query_key_value.bias":
281- bias_weight = name_param_m[1].reshape(
282- ng, repeats + 2, name_param_m[1].shape[0] // ng // (repeats + 2)
283- )
284- w = bias_weight.shape[-1]
285- qw = bias_weight[:, :repeats, ...].reshape(-1)
286- kw = bias_weight[:, repeats : repeats + 1, ...].reshape(-1)
287- vw = bias_weight[:, repeats + 1 :, ...].reshape(-1)
288- hf2mg_map[f"transformer.encoder.layers.{layer_num}.self_attention.query_key_value.bias"] = torch.cat(
289- [qw, kw, vw], dim=0
290- )
291- continue
292- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.self_attention.dense.weight":
293- hf2mg_map[f"transformer.encoder.layers.{layer_num}.self_attention.dense.weight"] = name_param_m[1]
294- continue
295- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.post_attention_norm.weight":
296- hf2mg_map[f"transformer.encoder.layers.{layer_num}.post_attention_layernorm.weight"] = name_param_m[1]
297- continue
298- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_h_to_4h.weight":
299- hf2mg_map[f"transformer.encoder.layers.{layer_num}.mlp.dense_h_to_4h.weight"] = name_param_m[1]
300- continue
301- if name_param_m[0] == f"language_model.encoder.layers.{layer_num}.mlp.dense_4h_to_h.weight":
302- hf2mg_map[f"transformer.encoder.layers.{layer_num}.mlp.dense_4h_to_h.weight"] = name_param_m[1]
303- continue
304- if name_param_m[0] == "language_model.encoder.final_norm.weight":
305- hf2mg_map[f"transformer.encoder.final_layernorm.weight"] = name_param_m[1]
306- continue
307- if name_param_m[0] == "language_model.output_layer.weight":
308- hf2mg_map[f"transformer.output_layer.weight"] = name_param_m[1]
309- continue
310- for name_param_h in hf_model.named_parameters():
311- if name_param_h[0] in hf2mg_map.keys():
312- name_param_h[1].data.copy_(hf2mg_map[name_param_h[0]])
313- 
314- save_dir = os.path.join(args.save_dir, 'mg2hf')
315- print(f'save weight to {save_dir}')
316- hf_model.save_pretrained(save_dir)
317- 
318- 
319-def vocab_padding(md, margs, orig_tensor, _vocab_size_with_padding):
320- # figure out what our padded vocab size is
321- orig_vocab_size = orig_tensor.shape[0]
322- margs.padded_vocab_size = _vocab_size_with_padding(md.true_vocab_size, margs)
323- 
324- # Cut out extra padding we don't need
325- if orig_vocab_size > margs.padded_vocab_size:
326- full_word_embed = orig_tensor[0:margs.padded_vocab_size, :]
327- 
328- # Expanding embedding to larger size by replicating final entry
329- elif orig_vocab_size < margs.padded_vocab_size:
330- padding_size = margs.padded_vocab_size - orig_vocab_size
331- 
332- full_word_embed = torch.cat((
333- orig_tensor,
334- orig_tensor[-1].unsqueeze(0).expand(padding_size, -1)))
335- 
336- # Same size!
337- else:
338- full_word_embed = orig_tensor
339- 
340- return full_word_embed
341- 
342- 
343-def save_model_checkpoint(model_provider, queue, args):
344- 
345- # Search in directory above this
346- sys.path.append(os.path.abspath(
347- os.path.join(os.path.dirname(__file__),
348- os.path.pardir,
349- os.path.pardir)))
350- if args.megatron_path is not None:
351- sys.path.insert(0, args.megatron_path)
352- 
353- from megatron.training.arguments import validate_args
354- from modellink.training.utils import parse_args
355- from megatron.training.checkpointing import save_checkpoint
356- from megatron.training.global_vars import set_global_variables, get_args
357- from megatron.core.enums import ModelType
358- from megatron.training.tokenizer.tokenizer import _vocab_size_with_padding
359- from megatron.core import mpu
360- 
361- def queue_get(name=None):
362- val = queue.get()
363- if val == "exit":
364- print("Loader exited, exiting saver")
365- exit(1)
366- if name is not None and args.checking and val["name"] != name:
367- val_name = val["name"]
368- print(f'Unexpected message. Expecting "{name}" but got "{val_name}". Exiting saver.')
369- exit(1)
370- if name is not None:
371- print(f"received {name}")
372- return val
373- 
374- def check_message(msg):
375- if not args.checking:
376- return
377- msg_name = msg.pop("name")
378- if len(msg.keys()) > 0:
379- print(f"Unexpected values in {msg_name}:")
380- for key in msg.keys():
381- print(f" {key}")
382- print(f"Exiting. If you want to ignore this, use the argument --no-checking.")
383- exit(1)
384- 
385- 
386- md = queue_get()
387- 
388- if args.target_tensor_parallel_size is None:
389- if hasattr(md, 'previous_tensor_parallel_size'):
390- args.target_tensor_parallel_size = md.previous_tensor_parallel_size
391- else:
392- print("loader did not provide a tensor parallel size and --target-tensor-parallel-size not provided on command line. "
393- "Default to 1.")
394- args.target_tensor_parallel_size = 1
395- 
396- if args.target_pipeline_parallel_size is None:
397- if hasattr(md, 'previous_pipeline_parallel_size'):
398- args.target_pipeline_parallel_size = md.previous_pipeline_parallel_size
399- else:
400- print("loader did not provide a pipeline parallel size and --target-pipeline-parallel-size not provided on command line. "
401- "Default to 1.")
402- args.target_pipeline_parallel_size = 1
403- 
404- 
405- # Arguments do sanity checks on the world size, but we don't care,
406- # so trick it into thinking we are plenty of processes
407- if args.target_tensor_parallel_size is not None and args.target_pipeline_parallel_size is not None:
408- os.environ["WORLD_SIZE"] = f'{args.target_tensor_parallel_size * args.target_pipeline_parallel_size}'
409- 
410- # We want all arguments to come from us
411- sys.argv = ['script.py',
412- '--num-layers', str(md.num_layers),
413- '--hidden-size', str(md.hidden_size),
414- '--seq-length', str(md.seq_length),
415- '--num-attention-heads', str(md.num_attention_heads),
416- '--max-position-embeddings', str(md.max_position_embeddings),
417- '--position-embedding-type', str(md.position_embedding_type),
418- '--tokenizer-type', str(md.tokenizer_type),
419- '--tensor-model-parallel-size', str(args.target_tensor_parallel_size),
420- '--pipeline-model-parallel-size', str(args.target_pipeline_parallel_size),
421- '--no-masked-softmax-fusion',
422- '--no-bias-gelu-fusion',
423- '--no-bias-dropout-fusion',
424- '--no-async-tensor-model-parallel-allreduce',
425- '--use-cpu-initialization',
426- '--micro-batch-size', '1',
427- '--no-load-optim',
428- '--no-load-rng',
429- '--no-save-optim',
430- '--no-save-rng',
431- '--no-initialization',
432- '--save-interval', '1',
433- '--save', args.save_dir
434- ]
435- if md.make_vocab_size_divisible_by is not None:
436- sys.argv.extend(['--make-vocab-size-divisible-by', str(md.make_vocab_size_divisible_by)])
437- if md.params_dtype == torch.float16:
438- sys.argv.append('--fp16')
439- elif md.params_dtype == torch.bfloat16:
440- sys.argv.append('--bf16')
441- 
442- if md.output_layer:
443- sys.argv.append('--untie-embeddings-and-output-weights')
444- if not md.linear_bias:
445- sys.argv.append('--disable-bias-linear')
446- 
447- if md.model_type == 'BERT' and not md.bert_binary_head:
448- sys.argv.append('--bert-no-binary-head')
449- 
450- margs = parse_args()
451- margs.make_vocab_size_divisible_by = 1
452- margs.w_pack = args.w_pack
453- margs.num_layer_list = args.num_layer_list
454- 
455- if hasattr(md, 'checkpoint_args'):
456- # These are arguments that we are either changing, or cause problems for validation if they are set
457- # Note that some of these deal with T5 so will need to be changed if we support T5.
458- args_to_keep = ['tensor_model_parallel_size', 'pipeline_model_parallel_size', 'world_size', 'params_dtype',
459- 'num_layers_per_virtual_pipeline_stage', 'virtual_pipeline_model_parallel_size',
460- 'masked_softmax_fusion', 'bias_gelu_fusion', 'bias_dropout_fusion',
461- 'sequence_parallel', 'async_tensor_model_parallel_allreduce',
462- 'no_load_optim', 'no_load_rng', 'no_save_optim', 'no_save_rng',
463- 'vocab_file', 'tokenizer_model',
464- 'save_interval', 'save',
465- 'perform_initialization', 'use_cpu_initialization',
466- 'recompute_granularity', 'recompute_num_layers', 'recompute_method',
467- 'encoder_num_layers', 'encoder_seq_length',
468- 'distribute_saved_activations',
469- 'train_iters', 'lr_decay_iters', 'lr_warmup_iters', 'lr_warmup_fraction',
470- 'start_weight_decay', 'end_weight_decay', 'make_vocab_size_divisible_by',
471- 'num_layer_list', 'lora_target_modules']
472- 
473- 
474- for arg, value in vars(md.checkpoint_args).items():
475- if arg in args_to_keep:
476- continue
477- if not hasattr(margs, arg):
478- print(f"Checkpoint had argument {arg} but new arguments does not have this.")
479- continue
480- if getattr(margs, arg) != value:
481- print(f"Overwriting default {arg} value {getattr(margs, arg)} with value from checkpoint {value}.")
482- setattr(margs, arg, value)
483- 
484- validate_args(margs)
485- 
486- set_global_variables(margs, build_tokenizer=False)
487- 
488- # margs = megatron args
489- margs = get_args()
490- 
491- if hasattr(md, 'consumed_train_samples'):
492- margs.consumed_train_samples = md.consumed_train_samples
493- margs.consumed_valid_samples = md.consumed_valid_samples
494- print(f"Setting consumed_train_samples to {margs.consumed_train_samples}"
495- f" and consumed_valid_samples to {margs.consumed_valid_samples}")
496- else:
497- print("consumed_train_samples not provided.")
498- 
499- # Determine how to make our models
500- if md.model_type == 'GPT':
501- margs.model_type = ModelType.encoder_or_decoder
502- elif md.model_type == 'BERT':
503- margs.model_type = ModelType.encoder_or_decoder
504- else:
505- raise Exception(f'unrecognized model type: {args.model_type}')
506- 
507- def get_models(count, dtype, pre_process, post_process):
508- models = [model_provider(pre_process, post_process).to(dtype) for _ in range(count)]
509- return models
510- 
511- # fake initializing distributed
512- mpu.set_tensor_model_parallel_world_size(args.target_tensor_parallel_size)
513- mpu.set_pipeline_model_parallel_world_size(args.target_pipeline_parallel_size)
514- mpu.set_tensor_model_parallel_rank(0)
515- mpu.set_pipeline_model_parallel_rank(0)
516- 
517- # Embeddings
518- #-----------
519- embeddings_msg = queue_get("embeddings")
520- 
521- pos_embed = None
522- if md.position_embedding_type == 'learned_absolute':
523- pos_embed = embeddings_msg.pop("position embeddings")
524- orig_word_embed = embeddings_msg.pop("word embeddings")
525- orig_word_embed_n_w, orig_word_embed_n_b = None, None
526- if "word embeddings norm_w" in embeddings_msg and "word embeddings norm_b" in embeddings_msg:
527- orig_word_embed_n_w = embeddings_msg.pop("word embeddings norm_w")
528- orig_word_embed_n_b = embeddings_msg.pop("word embeddings norm_b")
529- check_message(embeddings_msg)
530- 
531- # Deal with padding
532- if md.true_vocab_size is not None:
533- full_word_embed = vocab_padding(md, margs, orig_word_embed, _vocab_size_with_padding)
534- else:
535- print("Original vocab size not specified, leaving embedding table as-is. "
536- "If you've changed the tensor parallel size this could cause problems.")
537- margs.padded_vocab_size = orig_word_embed.shape[0]
538- full_word_embed = orig_word_embed
539- 
540- # Split into new tensor model parallel sizes
541- out_word_embed = torch.chunk(full_word_embed, args.target_tensor_parallel_size, dim=0)
542- 
543- # Make models for first pipeline stage and fill in embeddings
544- mpu.set_pipeline_model_parallel_rank(0)
545- post_process = args.target_pipeline_parallel_size == 1
546- models = get_models(args.target_tensor_parallel_size, md.params_dtype, True, post_process)
547- for tp_rank, model in enumerate(models):
548- model.language_model.embedding.word_embeddings.weight.data.copy_(out_word_embed[tp_rank])
549- if orig_word_embed_n_w is not None:
550- model.language_model.embedding.word_embeddings.norm.weight.data.copy_(orig_word_embed_n_w)
551- model.language_model.embedding.word_embeddings.norm.bias.data.copy_(orig_word_embed_n_b)
552- if pos_embed is not None:
553- model.language_model.embedding.position_embeddings.weight.data.copy_(pos_embed)
554- else:
555- if hasattr(model.language_model.embedding, 'position_embeddings'):
556- raise ValueError("model should have position_embeddings")
557- 
558- # Transformer layers
559- #-------------------
560- total_layer_num = 0
561- lst = []
562- if args.num_layers_per_virtual_pipeline_stage and args.save_model_type == 'megatron':
563- while queue.qsize() > 3:
564- lst.append(queue.get())
565- for pp_rank in range(args.target_pipeline_parallel_size):
566- # For later pipeline parallel ranks, make the new models
567- if pp_rank > 0:
568- mpu.set_pipeline_model_parallel_rank(pp_rank)
569- post_process = pp_rank == args.target_pipeline_parallel_size - 1
570- models = get_models(args.target_tensor_parallel_size, md.params_dtype, False, post_process)
571- 
572- if args.num_layers_per_virtual_pipeline_stage and args.save_model_type == 'megatron':
573- vp_size = margs.num_layers // args.target_pipeline_parallel_size // args.num_layers_per_virtual_pipeline_stage
574- else:
575- vp_size = 1
576- for vpp_rank in range(vp_size):
577- for layer in range(len(models[0].language_model.encoder.layers) // vp_size):
578- if args.num_layers_per_virtual_pipeline_stage and args.save_model_type == 'megatron':
579- # vpp model的layer间执行顺序与pp model不同，这里需要计算索引按实际执行顺序排列layer
580- total_layer_num = args.target_pipeline_parallel_size * vpp_rank * args.num_layers_per_virtual_pipeline_stage + pp_rank * args.num_layers_per_virtual_pipeline_stage + layer
581- msg = lst[total_layer_num]
582- else:
583- msg = queue_get(f"transformer layer {total_layer_num}")
584- 
585- # duplicated tensors
586- input_norm_weight = msg.pop("input norm weight")
587- if md.norm_has_bias:
588- input_norm_bias = msg.pop("input norm bias")
589- post_norm_weight = msg.pop("post norm weight")
590- if md.norm_has_bias:
591- post_norm_bias = msg.pop("post norm bias")
592- if md.linear_bias:
593- dense_bias = msg.pop("dense bias")
594- mlp_l1_bias = msg.pop("mlp l1 bias")
595- 
596- if args.add_qkv_bias:
597- qkv_bias = torch.chunk(msg.pop("qkv bias"), args.target_tensor_parallel_size, dim=0)
598- if args.add_dense_bias:
599- dense_bias = msg.pop("dense bias")
600- 
601- qkv_org = msg.pop("qkv weight")
602- qkv_weight = torch.chunk(qkv_org, args.target_tensor_parallel_size, dim=0)
603- 
604- # Split up the parallel tensors
605- dense_weight = torch.chunk(msg.pop("dense weight"), args.target_tensor_parallel_size, dim=1)
606- mlp_l1_weight = torch.chunk(msg.pop("mlp l1 weight"), args.target_tensor_parallel_size, dim=1)
607- 
608- # Special handling for swiglu
609- if md.swiglu:
610- mlp_l0_weight_W = torch.chunk(msg.pop("mlp l0 weight W"), args.target_tensor_parallel_size, dim=0)
611- mlp_l0_weight_V = torch.chunk(msg.pop("mlp l0 weight V"), args.target_tensor_parallel_size, dim=0)
612- mlp_l0_weight = [torch.cat(weights, dim=0) for weights in zip(mlp_l0_weight_W, mlp_l0_weight_V)]
613- else:
614- mlp_l0_weight = torch.chunk(msg.pop("mlp l0 weight"), args.target_tensor_parallel_size, dim=0)
615- 
616- if md.linear_bias:
617- qkv_bias = torch.chunk(msg.pop("qkv bias"), args.target_tensor_parallel_size, dim=0)
618- if md.swiglu:
619- mlp_l0_bias_W = torch.chunk(msg.pop("mlp l0 bias W"), args.target_tensor_parallel_size, dim=0)
620- mlp_l0_bias_V = torch.chunk(msg.pop("mlp l0 bias V"), args.target_tensor_parallel_size, dim=0)
621- mlp_l0_bias = [torch.cat(bias, dim=0) for bias in zip(mlp_l0_bias_W, mlp_l0_bias_V)]
622- else:
623- mlp_l0_bias = torch.chunk(msg.pop("mlp l0 bias"), args.target_tensor_parallel_size, dim=0)
624- 
625- # Save them to the model
626- for tp_rank in range(args.target_tensor_parallel_size):
627- if args.num_layers_per_virtual_pipeline_stage and args.save_model_type == 'megatron':
628- l = models[tp_rank].language_model.encoder.layers[
629- vpp_rank * args.num_layers_per_virtual_pipeline_stage + layer]
630- else:
631- l = models[tp_rank].language_model.encoder.layers[layer]
632- l.input_norm.weight.data.copy_(input_norm_weight)
633- if md.norm_has_bias:
634- l.input_norm.bias.data.copy_(input_norm_bias)
635- l.self_attention.query_key_value.weight.data.copy_(qkv_weight[tp_rank])
636- l.self_attention.dense.weight.data.copy_(dense_weight[tp_rank])
637- l.post_attention_norm.weight.data.copy_(post_norm_weight)
638- if md.norm_has_bias:
639- l.post_attention_norm.bias.data.copy_(post_norm_bias)
640- l.mlp.dense_h_to_4h.weight.data.copy_(mlp_l0_weight[tp_rank])
641- l.mlp.dense_4h_to_h.weight.data.copy_(mlp_l1_weight[tp_rank])
642- if md.linear_bias:
643- l.self_attention.query_key_value.bias.data.copy_(qkv_bias[tp_rank])
644- l.self_attention.dense.bias.data.copy_(dense_bias)
645- l.mlp.dense_h_to_4h.bias.data.copy_(mlp_l0_bias[tp_rank])
646- l.mlp.dense_4h_to_h.bias.data.copy_(mlp_l1_bias)
647- if args.add_qkv_bias:
648- l.self_attention.query_key_value.bias.data.copy_(qkv_bias[tp_rank])
649- if args.add_dense_bias:
650- l.self_attention.dense.bias.data.copy_(dense_bias)
651- 
652- total_layer_num = total_layer_num + 1
653- check_message(msg)
654- 
655- if post_process:
656- msg = queue_get("final norm")
657- final_norm_weight = msg.pop("weight")
658- if md.norm_has_bias:
659- final_norm_bias = msg.pop("bias")
660- for tp_rank in range(args.target_tensor_parallel_size):
661- models[tp_rank].language_model.encoder.final_norm.weight.data.copy_(final_norm_weight)
662- if md.norm_has_bias:
663- models[tp_rank].language_model.encoder.final_norm.bias.data.copy_(final_norm_bias)
664- if pp_rank != 0 and not md.output_layer:
665- # Copy word embeddings to final pipeline rank
666- models[tp_rank].word_embeddings.weight.data.copy_(out_word_embed[tp_rank])
667- del final_norm_weight
668- if md.norm_has_bias:
669- del final_norm_bias
670- check_message(msg)
671- 
672- if md.output_layer:
673- msg = queue_get("output layer")
674- if not hasattr(models[0].language_model, 'output_layer'):
675- print("ERROR: got an output layer, but model does not have one")
676- exit(1)
677-
678- if md.true_vocab_size is not None:
679- weight = vocab_padding(md, margs, msg.pop("weight"), _vocab_size_with_padding)
680- else:
681- weight = msg.pop("weight")
682- output_layer_weight = torch.chunk(weight, args.target_tensor_parallel_size, dim=0)
683- 
684- for tp_rank in range(args.target_tensor_parallel_size):
685- models[tp_rank].language_model.output_layer.weight.data.copy_(output_layer_weight[tp_rank])
686- del output_layer_weight
687- check_message(msg)
688- msg = queue_get()
689- if msg != "done" and msg["name"] == "pooler":
690- if not hasattr(models[0].language_model, 'pooler'):
691- print("ERROR: got a pooler, but model does not have one")
692- exit(1)
693- print("received pooler")
694- pooler_weight = msg.pop("weight")
695- pooler_bias = msg.pop("bias")
696- for tp_rank in range(args.target_tensor_parallel_size):
697- models[tp_rank].language_model.pooler.dense.weight.data.copy_(pooler_weight)
698- models[tp_rank].language_model.pooler.dense.bias.data.copy_(pooler_bias)
699- del pooler_weight
700- del pooler_bias
701- check_message(msg)
702- msg = queue_get()
703- 
704- if msg != "done" and msg["name"] == "lm head":
705- if not hasattr(models[0], 'lm_head'):
706- print("ERROR: got an lm head, but model does not have one")
707- exit(1)
708- print("received lm head")
709- lm_head_dense_weight = msg.pop("dense weight")
710- lm_head_dense_bias = msg.pop("dense bias")
711- lm_head_norm_weight = msg.pop("norm weight")
712- if md.norm_has_bias:
713- lm_head_norm_bias = msg.pop("norm bias")
714- for tp_rank in range(args.target_tensor_parallel_size):
715- models[tp_rank].lm_head.dense.weight.data.copy_(lm_head_dense_weight)
716- models[tp_rank].lm_head.dense.bias.data.copy_(lm_head_dense_bias)
717- models[tp_rank].lm_head.norm.weight.data.copy_(lm_head_norm_weight)
718- if md.norm_has_bias:
719- models[tp_rank].lm_head.norm.bias.data.copy_(lm_head_norm_bias)
720- check_message(msg)
721- msg = queue_get()
722- 
723- if msg != "done" and msg["name"] == "binary head":
724- if not hasattr(models[0], 'binary_head'):
725- print("ERROR: got a binary head, but model does not have one")
726- exit(1)
727- print("received binary head")
728- binary_head_weight = msg.pop("weight")
729- binary_head_bias = msg.pop("bias")
730- for tp_rank in range(args.target_tensor_parallel_size):
731- models[tp_rank].binary_head.weight.data.copy_(binary_head_weight)
732- models[tp_rank].binary_head.bias.data.copy_(binary_head_bias)
733- check_message(msg)
734- msg = queue_get()
735- 
736- if msg != "done":
737- print("ERROR: got some more data but was expecting to be done")
738- 
739- for tp_rank in range(args.target_tensor_parallel_size):
740- mpu.set_tensor_model_parallel_rank(tp_rank)
741- if args.save_model_type == 'megatron':
742- # 将pp拆分成多个vpp,通过复制和删除选取每个vpp对应的层数
743- if args.num_layers_per_virtual_pipeline_stage:
744- vp_models = []
745- layers = margs.num_layers // args.target_pipeline_parallel_size
746- for vp_rank in range(vp_size):
747- model = copy.deepcopy(models[tp_rank])
748- left = vp_rank * args.num_layers_per_virtual_pipeline_stage
749- right = (vp_rank + 1) * args.num_layers_per_virtual_pipeline_stage
750- for i in range(layers - 1, -1, -1):
751- if i >= right or i < left:
752- del model.language_model.encoder.layers[i]
753- if right < layers and pp_rank == args.target_pipeline_parallel_size - 1:
754- del model.language_model.encoder.final_norm
755- if getattr(model.language_model, "output_layer", None):
756- model.language_model.post_process = False
757- del model.language_model.output_layer
758- if pp_rank == 0 and vp_rank > 0:
759- model.language_model.pre_process = False
760- del model.language_model.embedding
761- vp_models.append(model)
762- save_checkpoint(md.iteration, vp_models, None, None, 0)
763- else:
764- save_checkpoint(md.iteration, [models[tp_rank]], None, None, 0)
765- elif args.save_model_type == 'huggingface_bloom':
766- save_huggingface(args, models[tp_rank])
767- elif args.save_model_type == "save_huggingface_llama":
768- save_huggingface_llama(args, models[tp_rank], md)
769- elif args.save_model_type == "save_huggingface_qwen":
770- save_huggingface_qwen(args, models[tp_rank], md)
771- elif args.save_model_type == "save_huggingface_chatglm3":
772- save_huggingface_chatglm3(args, models[tp_rank], md)
773- 
774- print("Done!")
Mtests/ut/checkpoint/test_convert_ckpt_from_huggingface.py+15-12
@@ -12,8 +12,8 @@ from tests.test_tools.utils import weight_compare
12 12 
13class CovertCkptFromHuggingfaceArgs:13class CovertCkptFromHuggingfaceArgs:
14 model_type = "GPT"14 model_type = "GPT"
15- loader = "llama2_hf"15+ load_model_type = "hf"
16- saver = "megatron"16+ save_model_type = "mg"
17 target_tensor_parallel_size = "8"17 target_tensor_parallel_size = "8"
18 load_dir = "/data/llama-2-7b-hf"18 load_dir = "/data/llama-2-7b-hf"
19 save_dir = "/data/llama2-7B-tp8-pp1"19 save_dir = "/data/llama2-7B-tp8-pp1"
@@ -22,8 +22,8 @@ class CovertCkptFromHuggingfaceArgs:
22 22 
23class CovertVppCkptFromHuggingfaceArgs:23class CovertVppCkptFromHuggingfaceArgs:
24 model_type = "GPT"24 model_type = "GPT"
25- loader = "llama2_hf"25+ load_model_type = "hf"
26- saver = "megatron"26+ save_model_type = "mg"
27 target_tensor_parallel_size = "2"27 target_tensor_parallel_size = "2"
28 target_pipeline_parallel_size = "4"28 target_pipeline_parallel_size = "4"
29 load_dir = "/data/llama-2-7b-hf"29 load_dir = "/data/llama-2-7b-hf"
@@ -34,8 +34,8 @@ class CovertVppCkptFromHuggingfaceArgs:
34 34 
35class CovertDynamicCkptFromHuggingfaceArgs:35class CovertDynamicCkptFromHuggingfaceArgs:
36 model_type = "GPT"36 model_type = "GPT"
37- loader = "llama2_hf"37+ load_model_type = "hf"
38- saver = "megatron"38+ save_model_type = "mg"
39 target_tensor_parallel_size = "2"39 target_tensor_parallel_size = "2"
40 target_pipeline_parallel_size = "4"40 target_pipeline_parallel_size = "4"
41 load_dir = "/data/llama-2-7b-hf/"41 load_dir = "/data/llama-2-7b-hf/"
@@ -222,13 +222,14 @@ class TestConvertCkptFromHuggingface:
222 file_path = os.path.join(base_dir, "convert_ckpt.py")222 file_path = os.path.join(base_dir, "convert_ckpt.py")
223 arguments = [223 arguments = [
224 "--model-type", args.model_type,224 "--model-type", args.model_type,
225- "--loader", args.loader,225+ "--load-model-type", args.load_model_type,
226+ "--save-model-type", args.save_model_type,
226 "--num-layer-list", args.num_layer_list,227 "--num-layer-list", args.num_layer_list,
227- "--saver", args.saver,
228 "--target-tensor-parallel-size", args.target_tensor_parallel_size,228 "--target-tensor-parallel-size", args.target_tensor_parallel_size,
229 "--target-pipeline-parallel-size", args.target_pipeline_parallel_size,229 "--target-pipeline-parallel-size", args.target_pipeline_parallel_size,
230 "--load-dir", args.load_dir,230 "--load-dir", args.load_dir,
231 "--save-dir", args.save_dir,231 "--save-dir", args.save_dir,
232+ "--model-type-hf", "llama2",
232 "--tokenizer-model", args.tokenizer_model233 "--tokenizer-model", args.tokenizer_model
233 ]234 ]
234 exit_code = subprocess.run(["python", file_path] + arguments).returncode235 exit_code = subprocess.run(["python", file_path] + arguments).returncode
@@ -294,11 +295,12 @@ class TestConvertCkptFromHuggingface:
294 file_path = os.path.join(base_dir, "convert_ckpt.py")295 file_path = os.path.join(base_dir, "convert_ckpt.py")
295 arguments = [296 arguments = [
296 "--model-type", args.model_type,297 "--model-type", args.model_type,
297- "--loader", args.loader,298+ "--load-model-type", args.load_model_type,
298- "--saver", args.saver,299+ "--save-model-type", args.save_model_type,
299 "--target-tensor-parallel-size", args.target_tensor_parallel_size,300 "--target-tensor-parallel-size", args.target_tensor_parallel_size,
300 "--load-dir", args.load_dir,301 "--load-dir", args.load_dir,
301 "--save-dir", args.save_dir,302 "--save-dir", args.save_dir,
303+ "--model-type-hf", "llama2",
302 "--tokenizer-model", args.tokenizer_model304 "--tokenizer-model", args.tokenizer_model
303 ]305 ]
304 subprocess.run(["python", file_path] + arguments)306 subprocess.run(["python", file_path] + arguments)
@@ -333,13 +335,14 @@ class TestConvertCkptFromHuggingface:
333 file_path = os.path.join(base_dir, "convert_ckpt.py")335 file_path = os.path.join(base_dir, "convert_ckpt.py")
334 arguments = [336 arguments = [
335 "--model-type", args.model_type,337 "--model-type", args.model_type,
336- "--loader", args.loader,338+ "--load-model-type", args.load_model_type,
337- "--saver", args.saver,339+ "--save-model-type", args.save_model_type,
338 "--target-tensor-parallel-size", args.target_tensor_parallel_size,340 "--target-tensor-parallel-size", args.target_tensor_parallel_size,
339 "--target-pipeline-parallel-size", args.target_pipeline_parallel_size,341 "--target-pipeline-parallel-size", args.target_pipeline_parallel_size,
340 "--load-dir", args.load_dir,342 "--load-dir", args.load_dir,
341 "--save-dir", args.save_dir,343 "--save-dir", args.save_dir,
342 "--tokenizer-model", args.tokenizer_model,344 "--tokenizer-model", args.tokenizer_model,
345+ "--model-type-hf", "llama2",
343 "--num-layers-per-virtual-pipeline-stage", args.num_layers_per_virtual_pipeline_stage346 "--num-layers-per-virtual-pipeline-stage", args.num_layers_per_virtual_pipeline_stage
344 ]347 ]
345 subprocess.run(["python", file_path] + arguments)348 subprocess.run(["python", file_path] + arguments)
Mtests/ut/checkpoint/test_convert_ckpt_from_megatron.py+4-5
@@ -11,8 +11,8 @@ from tests.test_tools.utils import judge_expression
11 11 
12class CovertCkptFromMegatronArgs:12class CovertCkptFromMegatronArgs:
13 model_type = "GPT"13 model_type = "GPT"
14- loader = "megatron"14+ load_model_type = "mg"
15- saver = "megatron"15+ save_model_type = "mg"
16 target_tensor_parallel_size = "1"16 target_tensor_parallel_size = "1"
17 save_dir = "/data/llama2-7B-tp1-pp1"17 save_dir = "/data/llama2-7B-tp1-pp1"
18 load_dir = "/data/llama2-7B-tp8-pp1"18 load_dir = "/data/llama2-7B-tp8-pp1"
@@ -27,9 +27,8 @@ class TestConvertCkptFromMegatron:
27 file_path = os.path.join(base_dir, "convert_ckpt.py")27 file_path = os.path.join(base_dir, "convert_ckpt.py")
28 arguments = [28 arguments = [
29 "--model-type", args.model_type,29 "--model-type", args.model_type,
30- "--loader", args.loader,30+ "--load-model-type", args.load_model_type,
31- "--saver", args.saver,31+ "--save-model-type", args.save_model_type,
32- "--save-model-type", "megatron",
33 "--load-dir", args.load_dir,32 "--load-dir", args.load_dir,
34 "--target-tensor-parallel-size", "1",33 "--target-tensor-parallel-size", "1",
35 "--target-pipeline-parallel-size", "1",34 "--target-pipeline-parallel-size", "1",
Mtests/ut/checkpoint/test_convert_ckpt_to_huggingface.py+4-5
@@ -13,8 +13,8 @@ from tests.test_tools.utils import judge_expression
13 13 
14class CovertCkptToHuggingfaceArgs:14class CovertCkptToHuggingfaceArgs:
15 model_type = "GPT"15 model_type = "GPT"
16- loader = "megatron"16+ load_model_type = "mg"
17- saver = "megatron"17+ save_model_type = "mg"
18 save_dir = "/data/llama-2-7b-hf"18 save_dir = "/data/llama-2-7b-hf"
19 lora_dir = "/data/llama2-7B-lora-ckpt"19 lora_dir = "/data/llama2-7B-lora-ckpt"
20 load_dir = "/data/llama2-7B-tp8-pp1"20 load_dir = "/data/llama2-7B-tp8-pp1"
@@ -38,9 +38,8 @@ class TestConvertCkptFromHuggingface:
38 file_path = os.path.join(base_dir, "convert_ckpt.py")38 file_path = os.path.join(base_dir, "convert_ckpt.py")
39 arguments = [39 arguments = [
40 "--model-type", args.model_type,40 "--model-type", args.model_type,
41- "--loader", args.loader,41+ "--load-model-type", args.load_model_type,
42- "--saver", args.saver,42+ "--save-model-type", args.save_model_type,
43- "--save-model-type", "save_huggingface_llama",
44 "--load-dir", args.load_dir,43 "--load-dir", args.load_dir,
45 "--lora-load", args.lora_dir,44 "--lora-load", args.lora_dir,
46 "--target-tensor-parallel-size", "1",45 "--target-tensor-parallel-size", "1",