codecheck整改-2.0.0-MindSpeed-LLM-AtomGit

Mexamples/legacy/llama2/pretrain_llama2_7b_ptd.sh+0-1

	@@ -89,7 +89,6 @@ torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \
89	$DATA_ARGS \	89	$DATA_ARGS \
90	$OUTPUT_ARGS \	90	$OUTPUT_ARGS \
91	--distributed-backend nccl \	91	--distributed-backend nccl \
92	- --jit-compile \
93	--load $CKPT_LOAD_DIR \	92	--load $CKPT_LOAD_DIR \
94	--save $CKPT_SAVE_DIR \	93	--save $CKPT_SAVE_DIR \
95	\| tee logs/train_llama2_7b.log	94	\| tee logs/train_llama2_7b.log

Mmindspeed_llm/core/datasets/gpt_dataset.py+8-2

	@@ -35,8 +35,8 @@ def gpt_dataset_getitem_wrapper(fn):
35		35
36	if mpu.get_context_parallel_rank() == 0 and mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == 0:	36	if mpu.get_context_parallel_rank() == 0 and mpu.get_tensor_model_parallel_rank() == 0 and mpu.get_pipeline_model_parallel_rank() == 0:
37	batch_idx = numpy.array([idx], dtype=numpy.int64)	37	batch_idx = numpy.array([idx], dtype=numpy.int64)
38	- document_ids = numpy.pad(document_ids, (0, len(text) - len(document_ids)),'constant',constant_values=(-100, -100))	38	+ document_ids = numpy.pad(document_ids, (0, len(text) - len(document_ids)), 'constant', constant_values=(-100, -100))
39	- batch_idx = numpy.pad(batch_idx, (0, len(text) - len(batch_idx)),'constant',constant_values=(-100, -100))	39	+ batch_idx = numpy.pad(batch_idx, (0, len(text) - len(batch_idx)), 'constant', constant_values=(-100, -100))
40	batch["document_ids"] = document_ids	40	batch["document_ids"] = document_ids
41	batch["idx"] = batch_idx	41	batch["idx"] = batch_idx
42		42
	@@ -75,7 +75,13 @@ def _build_document_sample_shuffle_indices(
75	# here we change from (class)GPTDataset._build_document_sample_shuffle_indices	75	# here we change from (class)GPTDataset._build_document_sample_shuffle_indices
76	# end of megatron_adaptation	76	# end of megatron_adaptation
77		77
		78	+ def get_path_temp_func(suffix):
		79	+ return os.path.join(
		80	+ path_to_cache, f"{self.unique_description_hash}-{type(self).__name__}-{suffix}"
		81	+ )
		82	+
78	if path_to_cache:	83	if path_to_cache:
		84	+ get_path_to = get_path_temp_func
79	get_path_to = lambda suffix: os.path.join(	85	get_path_to = lambda suffix: os.path.join(
80	path_to_cache, f"{self.unique_description_hash}-{type(self).__name__}-{suffix}"	86	path_to_cache, f"{self.unique_description_hash}-{type(self).__name__}-{suffix}"
81	)	87	)

Mmindspeed_llm/core/models/common/embeddings/rotary_pos_embedding.py+4-4

	@@ -258,10 +258,10 @@ def yarn_get_mscale(scale=1, mscale=1):
258	return 0.1 * mscale * math.log(scale) + 1.0	258	return 0.1 * mscale * math.log(scale) + 1.0
259		259
260		260
261	-def yarn_linear_ramp_mask(min, max, dim):	261	+def yarn_linear_ramp_mask(min_, max_, dim):
262	- if min == max:	262	+ if min_ == max_:
263	- max += 0.001 # Prevent singularity	263	+ max_ += 0.001 # Prevent singularity
264		264
265	- linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)	265	+ linear_func = (torch.arange(dim, dtype=torch.float32) - min_) / (max_ - min_)
266	ramp_func = torch.clamp(linear_func, 0, 1)	266	ramp_func = torch.clamp(linear_func, 0, 1)
267	return ramp_func	267	return ramp_func

Mmindspeed_llm/core/models/gpt/gpt_model.py+2-1

	@@ -314,7 +314,8 @@ def gpt_model_forward(self, input_ids: Tensor,
314	# [s b h] => [b s h]	314	# [s b h] => [b s h]
315	return logits.transpose(0, 1).contiguous()	315	return logits.transpose(0, 1).contiguous()
316	if args.is_instruction_dataset:	316	if args.is_instruction_dataset:
317	- for i in range(len(labels)):	317	+ label_length = len(labels)
		318	+ for i in range(label_length):
318	labels[i] = labels[i][:, 1:].contiguous()	319	labels[i] = labels[i][:, 1:].contiguous()
319	logits = logits[:-1, :, :].contiguous()	320	logits = logits[:-1, :, :].contiguous()
320	loss += self.compute_language_model_loss(labels[0], logits)	321	loss += self.compute_language_model_loss(labels[0], logits)

Mmindspeed_llm/core/tensor_parallel/layers.py+3-4

	@@ -184,10 +184,9 @@ class SegmentedColumnParallelLinear(ColumnParallelLinear):
184	)	184	)
185		185
186	if self.config._cpu_offloading_context is not None:	186	if self.config._cpu_offloading_context is not None:
187	- if self.config._cpu_offloading_context.inside_context == True:	187	+ if self.config._cpu_offloading_context.inside_context:
188	- assert (	188	+ if self.config.cpu_offloading:
189	- self.config.cpu_offloading == False	189	+ raise ValueError("CPU Offloading cannot be enabled while using non-TE modules")
190	- ), "CPU Offloading cannot be enabled while using non-TE modules"
191		190
192	bias = self.bias if not self.skip_bias_add else None	191	bias = self.bias if not self.skip_bias_add else None
193		192

Mmindspeed_llm/core/transformer/mlp.py+1-1

	@@ -77,7 +77,7 @@ def core_mlp_init(self, config, submodules, is_expert=False, input_size=None, sh
77		77
78	self.config: TransformerConfig = config	78	self.config: TransformerConfig = config
79		79
80	- self.input_size = input_size if input_size != None else self.config.hidden_size	80	+ self.input_size = input_size if input_size else self.config.hidden_size
81	# geglu activation function	81	# geglu activation function
82	_args = get_args()	82	_args = get_args()
83	if _args.geglu:	83	if _args.geglu:

Mmindspeed_llm/core/transformer/transformer_block.py+9-9

Mmindspeed_llm/inference/text_generation/generation.py+2-1

	@@ -248,7 +248,8 @@ def beam_search_and_return_on_first_stage(
248	tokenizer = get_tokenizer()	248	tokenizer = get_tokenizer()
249		249
250	batch_size = tokens.size(0)	250	batch_size = tokens.size(0)
251	- assert(batch_size == 1)	251	+ if batch_size != 1:
		252	+ raise ValueError(f"batch_size must be 1, but current value is {batch_size}")
252	prompt_length = lengths.item()	253	prompt_length = lengths.item()
253	final_sequence_length = tokens.size(1)	254	final_sequence_length = tokens.size(1)
254	final_sequence_length = min(final_sequence_length, args.max_position_embeddings)	255	final_sequence_length = min(final_sequence_length, args.max_position_embeddings)

Mmindspeed_llm/legacy/model/gpt_model.py+4-4

	@@ -34,15 +34,15 @@ def post_language_model_processing(lm_output, labels, logit_weights,
34		34
35	if labels is None:	35	if labels is None:
36	# [s b h] => [b s h]	36	# [s b h] => [b s h]
37	- return output.transpose(0,1).contiguous()	37	+ return output.transpose(0, 1).contiguous()
38	else:	38	else:
39	# [b s] => [s b]	39	# [b s] => [s b]
40	- labels = labels.transpose(0,1).contiguous()	40	+ labels = labels.transpose(0, 1).contiguous()
41		41
42	args = get_args()	42	args = get_args()
43	if args.is_instruction_dataset:	43	if args.is_instruction_dataset:
44	labels = labels[1:, ...].contiguous()	44	labels = labels[1:, ...].contiguous()
45	- output = output[:-1, : , :].contiguous()	45	+ output = output[:-1, :, :].contiguous()
46		46
47	if fp16_lm_cross_entropy:	47	if fp16_lm_cross_entropy:
48	assert output.dtype == torch.half	48	assert output.dtype == torch.half
	@@ -50,7 +50,7 @@ def post_language_model_processing(lm_output, labels, logit_weights,
50	else:	50	else:
51	loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels)	51	loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels)
52	# [s b] => [b, s]	52	# [s b] => [b, s]
53	- loss = loss.transpose(0,1).contiguous()	53	+ loss = loss.transpose(0, 1).contiguous()
54	return loss	54	return loss
55		55
56		56

Mmindspeed_llm/legacy/model/transformer.py+2-2

	@@ -779,11 +779,11 @@ def ParallelAttentionForward(self, hidden_states, attention_mask,
779	if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:	779	if self.num_attention_heads_per_partition // self.num_query_groups_per_partition > 1:
780	key_layer = key_layer.repeat_interleave(	780	key_layer = key_layer.repeat_interleave(
781	self.num_attention_heads_per_partition // self.num_query_groups_per_partition,	781	self.num_attention_heads_per_partition // self.num_query_groups_per_partition,
782	- dim = 2	782	+ dim=2
783	)	783	)
784	value_layer = value_layer.repeat_interleave(	784	value_layer = value_layer.repeat_interleave(
785	self.num_attention_heads_per_partition // self.num_query_groups_per_partition,	785	self.num_attention_heads_per_partition // self.num_query_groups_per_partition,
786	- dim = 2	786	+ dim=2
787	)	787	)
788		788
789	# apply relative positional encoding (rotary embedding)	789	# apply relative positional encoding (rotary embedding)

Mmindspeed_llm/tasks/checkpoint/convert_param.py+12-4

	@@ -329,7 +329,7 @@ class ConvertBase:
329		329
330	class ConvertHf2Mg(ConvertBase):	330	class ConvertHf2Mg(ConvertBase):
331	def __init__(self, args_cmd):	331	def __init__(self, args_cmd):
332	- ConvertBase.__init__(self, args_cmd)	332	+ super().__init__(args_cmd)
333		333
334	def _set_dense_mg_model(self, hf_model, tp_rank, pp_rank):	334	def _set_dense_mg_model(self, hf_model, tp_rank, pp_rank):
335	"""	335	"""
	@@ -568,7 +568,7 @@ class ConvertHf2Mg(ConvertBase):
568		568
569	class ConvertMg2Hf(ConvertBase):	569	class ConvertMg2Hf(ConvertBase):
570	def __init__(self, args_cmd):	570	def __init__(self, args_cmd):
571	- ConvertBase.__init__(self, args_cmd)	571	+ super().__init__(args_cmd)
572		572
573	# --- setup dense HuggingFace model	573	# --- setup dense HuggingFace model
574	def _set_dense_hf_model(self, pp_rank):	574	def _set_dense_hf_model(self, pp_rank):
	@@ -682,7 +682,9 @@ class ConvertMg2Hf(ConvertBase):
682	hf_model[ParamKey.get_hf_attn_value_weight_key(self.model_name, hf_layer_id)] = vw	682	hf_model[ParamKey.get_hf_attn_value_weight_key(self.model_name, hf_layer_id)] = vw
683		683
684	dense_tp_weights = [	684	dense_tp_weights = [
685	- m[f'decoder.layers.{mg_layer_id}.self_attention.linear_proj.weight'] for m in vp_mg_tp_models]	685	+ m[f'decoder.layers.{mg_layer_id}.self_attention.linear_proj.weight']
		686	+ for m in vp_mg_tp_models
		687	+ ]
686	dense_w = self.get_tp2d_merge_matrix_weight(dense_tp_weights,	688	dense_w = self.get_tp2d_merge_matrix_weight(dense_tp_weights,
687	tp_x=self.args_cmd.tp_x,	689	tp_x=self.args_cmd.tp_x,
688	tp_y=self.args_cmd.tp_y,	690	tp_y=self.args_cmd.tp_y,
	@@ -694,7 +696,13 @@ class ConvertMg2Hf(ConvertBase):
694		696
695	def _set_hf_model_layer_mlp(self, vp_mg_tp_models, mg_layer_id, hf_layer_id, hf_model):	697	def _set_hf_model_layer_mlp(self, vp_mg_tp_models, mg_layer_id, hf_layer_id, hf_model):
696	gate_up_tp_weights = [	698	gate_up_tp_weights = [
697	- torch.chunk(m[f'decoder.layers.{mg_layer_id}.mlp.linear_fc1.weight'], 2, dim=0) for m in vp_mg_tp_models]	699	+ torch.chunk(
		700	+ m[f'decoder.layers.{mg_layer_id}.mlp.linear_fc1.weight'],
		701	+ 2,
		702	+ dim=0
		703	+ )
		704	+ for m in vp_mg_tp_models
		705	+ ]
698	gate_tp_weights = [m[0] for m in gate_up_tp_weights]	706	gate_tp_weights = [m[0] for m in gate_up_tp_weights]
699	up_tp_weights = [m[1] for m in gate_up_tp_weights]	707	up_tp_weights = [m[1] for m in gate_up_tp_weights]
700		708

Mmindspeed_llm/tasks/checkpoint/optim.py+0-2

	@@ -652,7 +652,6 @@ class OptimTargetProcessor(OptimBaseProcessor):
652	if ep_rank == 0:	652	if ep_rank == 0:
653	for idx in range(len(ckpt[0].keys())):	653	for idx in range(len(ckpt[0].keys())):
654	if key == "param":	654	if key == "param":
655	- # TODO multi-bucket process
656	non_expert_ckpt[idx][(torch.bfloat16, torch.float32)]['numel_unpadded'] = ckpt[0][idx][0].numel()	655	non_expert_ckpt[idx][(torch.bfloat16, torch.float32)]['numel_unpadded'] = ckpt[0][idx][0].numel()
657	non_expert_ckpt[idx][(torch.bfloat16, torch.float32)][key] = ckpt[0][idx][0]	656	non_expert_ckpt[idx][(torch.bfloat16, torch.float32)][key] = ckpt[0][idx][0]
658	else:	657	else:
	@@ -661,7 +660,6 @@ class OptimTargetProcessor(OptimBaseProcessor):
661	if self.ep_size > 1:	660	if self.ep_size > 1:
662	for idx in range(len(ckpt[1].keys())):	661	for idx in range(len(ckpt[1].keys())):
663	if key == "param":	662	if key == "param":
664	- # TODO multi-bucket process
665	expert_ckpt[idx][(torch.bfloat16, torch.float32)]['numel_unpadded'] = ckpt[1][idx][0].numel()	663	expert_ckpt[idx][(torch.bfloat16, torch.float32)]['numel_unpadded'] = ckpt[1][idx][0].numel()
666	expert_ckpt[idx][(torch.bfloat16, torch.float32)][key] = ckpt[1][idx][0]	664	expert_ckpt[idx][(torch.bfloat16, torch.float32)][key] = ckpt[1][idx][0]
667		665

Mmindspeed_llm/tasks/checkpoint/optim_converter.py+4-5

	@@ -253,7 +253,6 @@ class OptimConverter(abc.ABC):
253	ep_rank_list = self.src_optim.ep_ranks	253	ep_rank_list = self.src_optim.ep_ranks
254		254
255	for i in layer_nums:	255	for i in layer_nums:
256	- # TODO:Noop_layer特性适配: d = GLOBAL_LAYER[i]
257	d = i	256	d = i
258	state_dict = OrderedDict()	257	state_dict = OrderedDict()
259	if "embedding.word_embeddings.weight" in src_data[(i, tp_rank_list[0], ep_rank_list[0])]:	258	if "embedding.word_embeddings.weight" in src_data[(i, tp_rank_list[0], ep_rank_list[0])]:
	@@ -468,7 +467,6 @@ class OptimConverter(abc.ABC):
468	shared_experts_linear_fc1_weight = mlp_moe.pop("mlp shared experts linear fc1 weight")	467	shared_experts_linear_fc1_weight = mlp_moe.pop("mlp shared experts linear fc1 weight")
469	shared_experts_linear_fc2_weight = mlp_moe.pop("mlp shared experts linear fc2 weight")	468	shared_experts_linear_fc2_weight = mlp_moe.pop("mlp shared experts linear fc2 weight")
470	if self.target_optim.moe_grouped_gemm:	469	if self.target_optim.moe_grouped_gemm:
471	- # TODO: check TP
472	weight1 = torch.chunk(mlp_moe.pop("mlp experts weight1 module").view(self.target_optim.hidden_size, -1),	470	weight1 = torch.chunk(mlp_moe.pop("mlp experts weight1 module").view(self.target_optim.hidden_size, -1),
473	self.target_optim.ep_size, dim=0)	471	self.target_optim.ep_size, dim=0)
474	weight2 = torch.chunk(mlp_moe.pop("mlp experts weight2 module").view(-1, self.target_optim.hidden_size),	472	weight2 = torch.chunk(mlp_moe.pop("mlp experts weight2 module").view(-1, self.target_optim.hidden_size),
	@@ -486,7 +484,6 @@ class OptimConverter(abc.ABC):
486	dst_data[(tp_rank, ep_rank)][layer_num][	484	dst_data[(tp_rank, ep_rank)][layer_num][
487	module_layer + "mlp.shared_experts.linear_fc2.weight"] = shared_experts_linear_fc2_weight	485	module_layer + "mlp.shared_experts.linear_fc2.weight"] = shared_experts_linear_fc2_weight
488	if self.target_optim.moe_grouped_gemm:	486	if self.target_optim.moe_grouped_gemm:
489	- # TODO: check TP
490	dst_data[(tp_rank_list[0], ep_rank)][layer_num][	487	dst_data[(tp_rank_list[0], ep_rank)][layer_num][
491	module_layer + "mlp.experts.weight1"] = weight1[ep_rank].view(self.target_optim.hidden_size, -1)	488	module_layer + "mlp.experts.weight1"] = weight1[ep_rank].view(self.target_optim.hidden_size, -1)
492	dst_data[(tp_rank_list[0], ep_rank)][layer_num][	489	dst_data[(tp_rank_list[0], ep_rank)][layer_num][
	@@ -533,8 +530,10 @@ class OptimConverter(abc.ABC):
533	layer_nums = sorted(list(src_data.keys()))	530	layer_nums = sorted(list(src_data.keys()))
534	tp_rank_list = self.target_optim.tp_ranks	531	tp_rank_list = self.target_optim.tp_ranks
535	ep_rank_list = self.target_optim.ep_ranks	532	ep_rank_list = self.target_optim.ep_ranks
536	- dst_dict = {(tp_rank, ep_rank): defaultdict(OrderedDict) for tp_rank, ep_rank in	533	+ dst_dict = {
537	- product(tp_rank_list, ep_rank_list)}	534	+ (tp_rank, ep_rank): defaultdict(OrderedDict)
		535	+ for tp_rank, ep_rank in product(tp_rank_list, ep_rank_list)
		536	+ }
538	out_word_embed_list = []	537	out_word_embed_list = []
539	for i in layer_nums:	538	for i in layer_nums:
540		539

Mmindspeed_llm/tasks/evaluation/eval_impl/bbh_eval.py+21-22

	@@ -24,11 +24,10 @@ import pandas as pd
24	from torch import distributed as dist	24	from torch import distributed as dist
25	from megatron.training import get_args	25	from megatron.training import get_args
26	from mindspeed_llm.tasks.preprocess.templates import Role	26	from mindspeed_llm.tasks.preprocess.templates import Role
27	-from .template import BBH_TEMPLATE_DIR, get_eval_template
28	from mindspeed_llm.tasks.evaluation.eval_api.dataset_eval import DatasetEval	27	from mindspeed_llm.tasks.evaluation.eval_api.dataset_eval import DatasetEval
29	from mindspeed_llm.tasks.evaluation.eval_api.chat import Chat	28	from mindspeed_llm.tasks.evaluation.eval_api.chat import Chat
30	from mindspeed_llm.tasks.utils.error_utils import check_divisible_by_zero	29	from mindspeed_llm.tasks.utils.error_utils import check_divisible_by_zero
31	-from .template import BBH_TEMPLATE_DIR, get_eval_template	30	+from mindspeed_llm.tasks.evaluation.eval_impl.template import BBH_TEMPLATE_DIR, get_eval_template
32		31
33		32
34	logger = logging.getLogger(__name__)	33	logger = logging.getLogger(__name__)
	@@ -196,30 +195,30 @@ class BBHEval(DatasetEval):
196	pass	195	pass
197		196
198	def format_instructions(self, instruction_set, instruction_list) -> (list, int):	197	def format_instructions(self, instruction_set, instruction_list) -> (list, int):
199	- for idx in range(1, len(instruction_set) - 1):	198	+ for idx in range(1, len(instruction_set) - 1):
200	- if idx == 1:	199	+ if idx == 1:
201	- prompt = instruction_set[0] + " Question: " + instruction_set[1]	200	+ prompt = instruction_set[0] + " Question: " + instruction_set[1]
202	- else:	201	+ else:
203	- prompt = instruction_set[idx]	202	+ prompt = instruction_set[idx]
204		203
205	- if prompt[-3:] != '(A)':	204	+ if prompt[-3:] != '(A)':
206	- answer_index = prompt.rfind('A')	205	+ answer_index = prompt.rfind('A')
207	- else:	206	+ else:
208	- answer_index = prompt[:-2].rfind('A')	207	+ answer_index = prompt[:-2].rfind('A')
209		208
210	- instruction_list[-1].extend([	209	+ instruction_list[-1].extend([
211	- {'role': Role.USER.value, 'content': prompt[:answer_index] + 'Answer: '},	210	+ {'role': Role.USER.value, 'content': prompt[:answer_index] + 'Answer: '},
212	- {'role': Role.ASSISTANT.value, 'content': prompt[answer_index + 3:].strip()}	211	+ {'role': Role.ASSISTANT.value, 'content': prompt[answer_index + 3:].strip()}
213	- ])	212	+ ])
214		213
215	- final_answer_index = instruction_set[-1].rfind('A')	214	+ final_answer_index = instruction_set[-1].rfind('A')
216	- instruction_list[-1].append({	215	+ instruction_list[-1].append({
217	- 'role': Role.USER.value,	216	+ 'role': Role.USER.value,
218	- 'content': instruction_set[-1][:final_answer_index] + 'Answer: '	217	+ 'content': instruction_set[-1][:final_answer_index] + 'Answer: '
219	- })	218	+ })
220		219
221	- options = re.findall(r'$([A-Z])$', instruction_set[-1][:final_answer_index])	220	+ options = re.findall(r'$([A-Z])$', instruction_set[-1][:final_answer_index])
222	- return options, final_answer_index	221	+ return options, final_answer_index
223		222
224	def get_best_choice(self, idx, model, instruction_set, options) -> str:	223	def get_best_choice(self, idx, model, instruction_set, options) -> str:
225	loss_records = []	224	loss_records = []

Mmindspeed_llm/tasks/evaluation/eval_impl/ceval_exam.py+2-2

	@@ -97,8 +97,8 @@ class CEvalExam(DatasetEval):
97	instruction = None	97	instruction = None
98	# 5-shot	98	# 5-shot
99	if self.prompt_type is not None:	99	if self.prompt_type is not None:
100	- train_dir = os.path.dirname(self.test_dir) + "/dev/"	100	+ train_dir = os.path.join(os.path.dirname(self.test_dir), "dev")
101	- train_file_path = os.path.join(train_dir, subject_name + "_dev.csv")	101	+ train_file_path = os.path.join(train_dir, f"{subject_name}_dev.csv")
102		102
103	if not os.path.exists(train_file_path):	103	if not os.path.exists(train_file_path):
104	raise FileExistsError("The file ({}) does not exist !".format(train_file_path))	104	raise FileExistsError("The file ({}) does not exist !".format(train_file_path))

Mmindspeed_llm/tasks/evaluation/eval_impl/needlebench_eval.py+1-1

	@@ -111,7 +111,7 @@ class NeedleBenchEval(DatasetEval):
111	queries = batch["prompt"]	111	queries = batch["prompt"]
112		112
113	chat_results, rank = chat.chat(instruction=queries, history=[])	113	chat_results, rank = chat.chat(instruction=queries, history=[])
114	- for idx, ans in enumerate(batch['answer']):	114	+ for _, ans in enumerate(batch['answer']):
115	if rank == 0:	115	if rank == 0:
116	acc = score(chat_results, ans).get('score')	116	acc = score(chat_results, ans).get('score')
117	logger.info("#################acc: {0}, chat_results: {1}, ans: {2}#################".format(acc,	117	logger.info("#################acc: {0}, chat_results: {1}, ans: {2}#################".format(acc,

Mmindspeed_llm/tasks/evaluation/eval_utils/cmmlu_utils.py+1-1

	@@ -25,7 +25,7 @@ def cmmlu_format_example(
25	prefex = f'There is a single choice question about {subject_name.replace("_", " ")}. Answer the question by replying A, B, C or D. \nQuestion: '	25	prefex = f'There is a single choice question about {subject_name.replace("_", " ")}. Answer the question by replying A, B, C or D. \nQuestion: '
26	else:	26	else:
27	prefex = f'以下是关于{subject_name}的单项选择题，请直接给出正确答案的选项。\n题目：'	27	prefex = f'以下是关于{subject_name}的单项选择题，请直接给出正确答案的选项。\n题目：'
28	- for idx, row in support_set.iterrows():	28	+ for _, row in support_set.iterrows():
29	prompt, response = _parse_example(row)	29	prompt, response = _parse_example(row)
30	messages += prefex + prompt + response + "\n"	30	messages += prefex + prompt + response + "\n"
31		31

Mmindspeed_llm/tasks/evaluation/eval_utils/mmlu_utils.py+1-1

	@@ -72,7 +72,7 @@ def _format_example(
72	"""	72	"""
73	messages = ""	73	messages = ""
74	prefex = f'There is a single choice question about {subject_name.replace("_", " ")}. Answer the question by replying A, B, C or D. \nQuestion: '	74	prefex = f'There is a single choice question about {subject_name.replace("_", " ")}. Answer the question by replying A, B, C or D. \nQuestion: '
75	- for idx, row in support_set.iterrows():	75	+ for _, row in support_set.iterrows():
76	prompt, response = _parse_example(row)	76	prompt, response = _parse_example(row)
77	messages += prefex + prompt + "\n" + response + "\n\n"	77	messages += prefex + prompt + "\n" + response + "\n\n"
78		78

Mmindspeed_llm/tasks/models/common/alibi.py+1-1

	@@ -59,7 +59,7 @@ class Alibi:
59	self.alibi_tensor = alibi_tensor	59	self.alibi_tensor = alibi_tensor
60		60
61	if args.square_alibi_mask and args.fill_neg_inf:	61	if args.square_alibi_mask and args.fill_neg_inf:
62	- if(len(attention_mask.size()) == 4):	62	+ if (len(attention_mask.size()) == 4):
63	if attention_mask.shape[0] < batch_size:	63	if attention_mask.shape[0] < batch_size:
64	attention_mask = attention_mask.repeat(batch_size, 1, 1, 1)[:batch_size, :, :, :]	64	attention_mask = attention_mask.repeat(batch_size, 1, 1, 1)[:batch_size, :, :, :]
65	else:	65	else:

Mmindspeed_llm/tasks/models/spec/deepseek_spec.py+5-3

	@@ -1,4 +1,9 @@
1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.	1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
		2	+
		3	+"""
		4	+MultiHeadLatent Layer Specification, which is mainly for Deepseek.
		5	+"""
		6	+
2	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add	7	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
3	from megatron.training import get_args	8	from megatron.training import get_args
4	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear	9	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
	@@ -15,9 +20,6 @@ from mindspeed_llm.tasks.models.transformer.multi_head_latent_attention import (
15	from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention	20	from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention
16	from mindspeed_llm.core import PTNorm	21	from mindspeed_llm.core import PTNorm
17		22
18	-"""
19	-MultiHeadLatent Layer Specification, which is mainly for Deepseek.
20	-"""
21		23
22	args = get_args()	24	args = get_args()
23	num_experts, moe_grouped_gemm, qk_layernorm, mla_mm_split = (	25	num_experts, moe_grouped_gemm, qk_layernorm, mla_mm_split = (

Mmindspeed_llm/tasks/models/spec/hunyuan_spec.py+10-7

	@@ -1,3 +1,9 @@
		1	+# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
		2	+
		3	+"""
		4	+Layer Specification for hunyuan-large.
		5	+"""
		6	+
1	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add	7	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
2	from megatron.training import get_args	8	from megatron.training import get_args
3	from megatron.core.transformer.enums import AttnMaskType	9	from megatron.core.transformer.enums import AttnMaskType
	@@ -9,15 +15,12 @@ from megatron.core.models.gpt.gpt_layer_specs import _get_mlp_module_spec
9	from mindspeed_llm.core import PTNorm	15	from mindspeed_llm.core import PTNorm
10	from mindspeed_llm.tasks.models.transformer.transformer_layer_hunyuan import HunyuanLargeTransformerLayer	16	from mindspeed_llm.tasks.models.transformer.transformer_layer_hunyuan import HunyuanLargeTransformerLayer
11	from mindspeed_llm.tasks.models.transformer.hunyuan_large_attention import HunyuanLargeAttention, HunyuanLargeAttentionSubmodules	17	from mindspeed_llm.tasks.models.transformer.hunyuan_large_attention import HunyuanLargeAttention, HunyuanLargeAttentionSubmodules
12	-	18	+
13	-"""	19	+
14	-Layer Specification for hunyuan-large.
15	-"""
16	-
17	args = get_args()	20	args = get_args()
18	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm	21	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm
19	-	22	+
20	-	23	+
21	layer_spec = ModuleSpec(	24	layer_spec = ModuleSpec(
22	module=HunyuanLargeTransformerLayer,	25	module=HunyuanLargeTransformerLayer,
23	submodules=TransformerLayerSubmodules(	26	submodules=TransformerLayerSubmodules(

Mmindspeed_llm/tasks/models/spec/minicpm_spec.py+5-3

	@@ -1,4 +1,9 @@
1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.	1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
		2	+
		3	+"""
		4	+Layer Specification for MiniCPM.
		5	+"""
		6	+
2	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add	7	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
3	from megatron.training import get_args	8	from megatron.training import get_args
4	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear	9	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
	@@ -10,9 +15,6 @@ from mindspeed_llm.tasks.models.transformer.multi_head_latent_attention import M
10	from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention	15	from mindspeed_llm.tasks.models.transformer.mla_dot_product_attention import MlaDotProductAttention
11	from mindspeed_llm.core import PTNorm	16	from mindspeed_llm.core import PTNorm
12		17
13	-"""
14	-Layer Specification for MiniCPM.
15	-"""
16		18
17	args = get_args()	19	args = get_args()
18	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm	20	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm

Mmindspeed_llm/tasks/models/spec/mtp_spec.py+5-3

	@@ -1,13 +1,15 @@
1	# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.	1	# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
		2	+
		3	+"""
		4	+Multi Token Predication Layer Specification.
		5	+"""
		6	+
2	from megatron.core.tensor_parallel import ColumnParallelLinear	7	from megatron.core.tensor_parallel import ColumnParallelLinear
3	from megatron.core.transformer import ModuleSpec	8	from megatron.core.transformer import ModuleSpec
4	from mindspeed_llm.core.transformer.custom_layers.transformer_engine import PTNorm	9	from mindspeed_llm.core.transformer.custom_layers.transformer_engine import PTNorm
5	from mindspeed_llm.tasks.models.transformer.multi_token_predication import MultiTokenPredicationSubmodules, \	10	from mindspeed_llm.tasks.models.transformer.multi_token_predication import MultiTokenPredicationSubmodules, \
6	MultiTokenPredication	11	MultiTokenPredication
7		12
8	-"""
9	-Multi Token Predication Layer Specification.
10	-"""
11		13
12	# Use this spec for multi token predication	14	# Use this spec for multi token predication
13	mtp_sepc = ModuleSpec(	15	mtp_sepc = ModuleSpec(

Mmindspeed_llm/tasks/models/spec/phi35_moe_spec.py+5-3

	@@ -1,4 +1,9 @@
1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.	1	# Copyright (c) 2024, HUAWEI CORPORATION. All rights reserved.
		2	+
		3	+"""
		4	+Layer Specification for Phi3.5-MoE
		5	+"""
		6	+
2	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add	7	from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
3	from megatron.training import get_args	8	from megatron.training import get_args
4	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear	9	from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear
	@@ -11,9 +16,6 @@ from megatron.core.transformer.dot_product_attention import DotProductAttention
11	from mindspeed_llm.core import PTNorm	16	from mindspeed_llm.core import PTNorm
12	from mindspeed_llm.tasks.models.transformer.attention import SelfAttentionWithDenseBias	17	from mindspeed_llm.tasks.models.transformer.attention import SelfAttentionWithDenseBias
13		18
14	-"""
15	-Layer Specification for Phi3.5-MoE
16	-"""
17		19
18	args = get_args()	20	args = get_args()
19	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm	21	num_experts, moe_grouped_gemm, qk_layernorm = args.num_experts, args.moe_grouped_gemm, args.qk_layernorm

Mmindspeed_llm/tasks/models/transformer/hunyuan_large_attention.py+6-4

	@@ -134,7 +134,8 @@ class HunyuanLargeAttention(SelfAttention):
134	)	134	)
135		135
136	def run_realtime_tests(self):	136	def run_realtime_tests(self):
137	- """Performs a consistency check.	137	+ """
		138	+ Performs a consistency check.
138		139
139	This function makes sure that tensors across devices are the same during an experiment.	140	This function makes sure that tensors across devices are the same during an experiment.
140	This is often not guaranteed to be so because of silent hardware failures (eg, memory	141	This is often not guaranteed to be so because of silent hardware failures (eg, memory
	@@ -142,7 +143,8 @@ class HunyuanLargeAttention(SelfAttention):
142		143
143	(TODO) In the future, more tensors should be checked across the training run and	144	(TODO) In the future, more tensors should be checked across the training run and
144	checked every X iterations. This is left for future work. Equality of tensors is probably not	145	checked every X iterations. This is left for future work. Equality of tensors is probably not
145	- required; transmitting hashes is sufficient."""	146	+ required; transmitting hashes is sufficient.
		147	+ """
146		148
147	if not self.config.qk_layernorm:	149	if not self.config.qk_layernorm:
148	return	150	return
	@@ -169,7 +171,7 @@ class HunyuanLargeAttention(SelfAttention):
169	src == tgt	171	src == tgt
170	), f"Discrepancy between {name} in {parallelism} ranks {i} and {rank}. Diff: {torch.norm(src - tgt)}"	172	), f"Discrepancy between {name} in {parallelism} ranks {i} and {rank}. Diff: {torch.norm(src - tgt)}"
171		173
172	- for i, dp in enumerate(dp_list):	174	+ for _, dp in enumerate(dp_list):
173	q_w, q_b, k_w, k_b = torch.unbind(dp)	175	q_w, q_b, k_w, k_b = torch.unbind(dp)
174	_compare(	176	_compare(
175	[q_w, q_b, k_w, k_b],	177	[q_w, q_b, k_w, k_b],
	@@ -188,7 +190,7 @@ class HunyuanLargeAttention(SelfAttention):
188	tp_list[rank] = inputs	190	tp_list[rank] = inputs
189	torch.distributed.all_gather(tp_list, inputs, group=get_tensor_model_parallel_group())	191	torch.distributed.all_gather(tp_list, inputs, group=get_tensor_model_parallel_group())
190		192
191	- for i, tp in enumerate(tp_list):	193	+ for _, tp in enumerate(tp_list):
192	q_w, q_b, k_w, k_b = torch.unbind(tp)	194	q_w, q_b, k_w, k_b = torch.unbind(tp)
193	_compare(	195	_compare(
194	[q_w, q_b, k_w, k_b],	196	[q_w, q_b, k_w, k_b],

Mmindspeed_llm/tasks/posttrain/lora/qlora.py+6-1

	@@ -80,7 +80,12 @@ def parallel_linear_save_to_state_dict_wrapper(fn):
80	def parallel_linear_load_from_state_dict_wrapper(fn):	80	def parallel_linear_load_from_state_dict_wrapper(fn):
81	def wrapper(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):	81	def wrapper(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
82	if any(['bitsandbytes' in i for i in state_dict.keys()]): # is quantized linear	82	if any(['bitsandbytes' in i for i in state_dict.keys()]): # is quantized linear
83	- qs_dict = {key: v for k, v in state_dict.items() if (key := k.replace(prefix, "")) != '_extra_state'}	83	+ qs_dict = {}
		84	+ for k, v in state_dict.items():
		85	+ key = k.replace(prefix, "")
		86	+ if key != '_extra_state':
		87	+ qs_dict[key] = v
		88	+
84	self.weight = bnb.nn.Params4bit.from_prequantized(	89	self.weight = bnb.nn.Params4bit.from_prequantized(
85	data=qs_dict.get('weight'),	90	data=qs_dict.get('weight'),
86	quantized_stats={key.replace('weight.', ''): qs_dict[key] for key in qs_dict if key != 'weight' and key != 'bias'},	91	quantized_stats={key.replace('weight.', ''): qs_dict[key] for key in qs_dict if key != 'weight' and key != 'bias'},

Mmindspeed_llm/tasks/posttrain/rejection_sampling/rejection_sampling.py+5-1

	@@ -20,6 +20,10 @@ def clean_up():
20	destroy_distributed_environment()	20	destroy_distributed_environment()
21	gc.collect()	21	gc.collect()
22	torch.npu.empty_cache()	22	torch.npu.empty_cache()
		23	+
		24	+
		25	+def dummy_is_rank_0():
		26	+ return True
23		27
24		28
25	def batch_generate_vllm(args):	29	def batch_generate_vllm(args):
	@@ -28,7 +32,7 @@ def batch_generate_vllm(args):
28		32
29	dummy_strategy = Empty()	33	dummy_strategy = Empty()
30	dummy_strategy.print = print	34	dummy_strategy.print = print
31	- dummy_strategy.is_rank_0 = lambda: True	35	+ dummy_strategy.is_rank_0 = dummy_is_rank_0
32	dummy_strategy.args = args	36	dummy_strategy.args = args
33		37
34	# configure tokenizer	38	# configure tokenizer

Mmindspeed_llm/tasks/posttrain/rejection_sampling/utils.py+1-1

	@@ -20,7 +20,7 @@ def blending_datasets(
20	raise ValueError(f"Length of probabilities ({len(probabilities)}) must match the length of datasets ({len(datasets)})")	20	raise ValueError(f"Length of probabilities ({len(probabilities)}) must match the length of datasets ({len(datasets)})")
21		21
22	train_data_list = []	22	train_data_list = []
23	- for i, dataset in enumerate(datasets):	23	+ for _, dataset in enumerate(datasets):
24	dataset = dataset.strip()	24	dataset = dataset.strip()
25	strategy.print(f"dataset: {dataset}")	25	strategy.print(f"dataset: {dataset}")
26		26

Mmindspeed_llm/tasks/posttrain/rlxf/single_controller/base/decorator.py+15-15

	@@ -142,15 +142,15 @@ def _concat_data_proto_or_future(output: List):
142	import ray	142	import ray
143		143
144	# make sure all the elements in output has the same type	144	# make sure all the elements in output has the same type
145	- for o in output:	145	+ for single_output in output:
146	- if type(o) != type(output[0]):	146	+ if not isinstance(single_output, type(output[0])):
147	- raise TypeError(f"All elements in output must have the same type. Found {type(o)} and {type(output[0])}")	147	+ raise TypeError(f"All elements in output must have the same type. Found {type(single_output)} and {type(output[0])}")
148		148
149	- o = output[0]	149	+ output_prime = output[0]
150		150
151	- if isinstance(o, DataProto):	151	+ if isinstance(output_prime, DataProto):
152	return DataProto.concat(output)	152	return DataProto.concat(output)
153	- elif isinstance(o, ray.ObjectRef):	153	+ elif isinstance(output_prime, ray.ObjectRef):
154	return DataProtoFuture.concat(output)	154	return DataProtoFuture.concat(output)
155	else:	155	else:
156	raise NotImplementedError	156	raise NotImplementedError
	@@ -164,9 +164,9 @@ def collect_megatron_compute_data_proto(worker_group, output):
164	import ray	164	import ray
165		165
166	output = collect_megatron_compute(worker_group, output)	166	output = collect_megatron_compute(worker_group, output)
167	- for o in output:	167	+ for single_output in output:
168	- if not isinstance(o, (DataProto, ray.ObjectRef)):	168	+ if not isinstance(single_output, (DataProto, ray.ObjectRef)):
169	- raise TypeError(f"Expecting {o} to be DataProto or ray.ObjectRef, but got {type(o)}")	169	+ raise TypeError(f"Expecting {single_output} to be DataProto or ray.ObjectRef, but got {type(single_output)}")
170		170
171	return _concat_data_proto_or_future(output)	171	return _concat_data_proto_or_future(output)
172		172
	@@ -278,7 +278,7 @@ def dispatch_dp_compute(worker_group, args, *kwargs):
278	for arg in args:	278	for arg in args:
279	if not isinstance(arg, (Tuple, List)) or len(arg) != worker_group.world_size:	279	if not isinstance(arg, (Tuple, List)) or len(arg) != worker_group.world_size:
280	raise ValueError(f'Each argument in args must be a Tuple or List of length {worker_group.world_size}')	280	raise ValueError(f'Each argument in args must be a Tuple or List of length {worker_group.world_size}')
281	- for k, v in kwargs.items():	281	+ for _, v in kwargs.items():
282	if not isinstance(v, (Tuple, List)) or len(v) != worker_group.world_size:	282	if not isinstance(v, (Tuple, List)) or len(v) != worker_group.world_size:
283	raise ValueError(f'Each argument in kwargs must be a Tuple or List of length {worker_group.world_size}')	283	raise ValueError(f'Each argument in kwargs must be a Tuple or List of length {worker_group.world_size}')
284	return args, kwargs	284	return args, kwargs
	@@ -318,9 +318,9 @@ def dispatch_dp_compute_data_proto_with_func(worker_group, args, *kwargs):
318	def collect_dp_compute_data_proto(worker_group, output):	318	def collect_dp_compute_data_proto(worker_group, output):
319	import ray	319	import ray
320	from mindspeed_llm.tasks.posttrain.rlxf.utils.protocol import DataProto	320	from mindspeed_llm.tasks.posttrain.rlxf.utils.protocol import DataProto
321	- for o in output:	321	+ for single_output in output:
322	- if not isinstance(o, (DataProto, ray.ObjectRef)):	322	+ if not isinstance(single_output, (DataProto, ray.ObjectRef)):
323	- raise TypeError(f"Expecting {o} to be DataProto or ray.ObjectRef, but got {type(o)}")	323	+ raise TypeError(f"Expecting {single_output} to be DataProto or ray.ObjectRef, but got {type(single_output)}")
324		324
325	output = collect_dp_compute(worker_group, output)	325	output = collect_dp_compute(worker_group, output)
326	return _concat_data_proto_or_future(output)	326	return _concat_data_proto_or_future(output)
	@@ -409,7 +409,7 @@ def get_predefined_dispatch_fn(dispatch_mode):
409	'collect_fn': collect_dp_infer,	409	'collect_fn': collect_dp_infer,
410	},	410	},
411	}	411	}
412	- return predefined_dispatch_mode_fn[dispatch_mode]	412	+ return predefined_dispatch_mode_fn.get(dispatch_mode)
413		413
414		414
415	def get_predefined_execute_fn(execute_mode):	415	def get_predefined_execute_fn(execute_mode):
	@@ -431,7 +431,7 @@ def get_predefined_execute_fn(execute_mode):
431	'execute_fn_name': 'execute_train'	431	'execute_fn_name': 'execute_train'
432	}	432	}
433	}	433	}
434	- return predefined_execute_mode_fn[execute_mode]	434	+ return predefined_execute_mode_fn.get(execute_mode)
435		435
436		436
437	def _check_dispatch_mode(dispatch_mode):	437	def _check_dispatch_mode(dispatch_mode):

Mmindspeed_llm/tasks/posttrain/rlxf/single_controller/base/worker_group.py+11-4

	@@ -47,9 +47,16 @@ class ResourcePool:
47	return self._store	47	return self._store
48		48
49	def local_world_size_list(self) -> List[int]:	49	def local_world_size_list(self) -> List[int]:
50	- nested_local_world_size_list = [	50	+ nested_local_world_size_list = []
51	- [local_world_size for _ in range(local_world_size)] for local_world_size in self._store	51	+
52	- ]	52	+ for local_world_size in self._store:
		53	+ inner_list = []
		54	+
		55	+ for _ in range(local_world_size):
		56	+ inner_list.append(local_world_size)
		57	+
		58	+ nested_local_world_size_list.append(inner_list)
		59	+
53	return [item for row in nested_local_world_size_list for item in row]	60	return [item for row in nested_local_world_size_list for item in row]
54		61
55	def local_rank_list(self) -> List[int]:	62	def local_rank_list(self) -> List[int]:
	@@ -212,4 +219,4 @@ class WorkerGroup:
212	try:	219	try:
213	setattr(self, method_name, func)	220	setattr(self, method_name, func)
214	except Exception as e:	221	except Exception as e:
215	- raise ValueError(f'Fail to set method_name {method_name}')	222	+ raise ValueError(f'Fail to set method_name {method_name}') from e

Mmindspeed_llm/tasks/posttrain/rlxf/single_controller/ray/base.py+8-10

	@@ -11,22 +11,23 @@
11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.	11	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12	# See the License for the specific language governing permissions and	12	# See the License for the specific language governing permissions and
13	# limitations under the License.	13	# limitations under the License.
		14	+
		15	+__all__ = ['Worker']
		16	+
14	import os	17	import os
15	import time	18	import time
16	from typing import Dict, List, Any	19	from typing import Dict, List, Any
		20	+from unittest.mock import patch
17		21
18	import ray	22	import ray
19	from ray.util import list_named_actors	23	from ray.util import list_named_actors
20	from ray.util.placement_group import placement_group	24	from ray.util.placement_group import placement_group
21	from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy, NodeAffinitySchedulingStrategy	25	from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy, NodeAffinitySchedulingStrategy
22	from ray.experimental.state.api import get_actor	26	from ray.experimental.state.api import get_actor
23	-from unittest.mock import patch
24		27
25	from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base import WorkerGroup, ResourcePool, ClassWithInitArgs, Worker	28	from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base import WorkerGroup, ResourcePool, ClassWithInitArgs, Worker
26	from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base.decorator import MAGIC_ATTR	29	from mindspeed_llm.tasks.posttrain.rlxf.single_controller.base.decorator import MAGIC_ATTR
27		30
28	-__all__ = ['Worker']
29	-
30	ACTOR_INFER_WORLD_SIZE = None	31	ACTOR_INFER_WORLD_SIZE = None
31	ACTOR_TRAIN_WORLD_SIZE = None	32	ACTOR_TRAIN_WORLD_SIZE = None
32		33
	@@ -387,14 +388,11 @@ class RayWorkerGroup(WorkerGroup):
387	return self._world_size	388	return self._world_size
388		389
389		390
390	-"""
391	-Utilities that enables creating workers inside the same ray.Actor,
392	-with code written in separate ray.Actors.
393	-"""
394	-
395	-
396	def _bind_workers_method_to_parent(cls, key, user_defined_cls):	391	def _bind_workers_method_to_parent(cls, key, user_defined_cls):
397	"""	392	"""
		393	+ Utilities that enables creating workers inside the same ray.Actor,
		394	+ with code written in separate ray.Actors.
		395	+
398	Binds the methods of each worker to the WorkerDict.	396	Binds the methods of each worker to the WorkerDict.
399	Note that we only bind public methods that are decorated by register	397	Note that we only bind public methods that are decorated by register
400	"""	398	"""
	@@ -426,7 +424,7 @@ def _bind_workers_method_to_parent(cls, key, user_defined_cls):
426	method_name_with_prefix = key + '_' + method_name	424	method_name_with_prefix = key + '_' + method_name
427	setattr(cls, method_name_with_prefix, func)	425	setattr(cls, method_name_with_prefix, func)
428	except Exception as e:	426	except Exception as e:
429	- raise ValueError(f'Fail to set method_name {method_name}')	427	+ raise ValueError(f'Fail to set method_name {method_name}') from e
430		428
431		429
432	def _unwrap_ray_remote(cls):	430	def _unwrap_ray_remote(cls):

Mmindspeed_llm/tasks/posttrain/rlxf/training/core_algos.py+3-4

	@@ -17,9 +17,10 @@ The function implemented in this file should be used by trainer with different d
17	implement PPO	17	implement PPO
18	"""	18	"""
19		19
		20	+from copy import deepcopy
		21	+
20	import numpy as np	22	import numpy as np
21	import torch	23	import torch
22	-from copy import deepcopy
23	from transformers import AutoTokenizer	24	from transformers import AutoTokenizer
24		25
25	import mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional as F	26	import mindspeed_llm.tasks.posttrain.rlxf.utils.torch_functional as F
	@@ -306,7 +307,6 @@ def compute_advantage(data: DataProto, config):
306	response_mask = attention_mask[:, -response_length:]	307	response_mask = attention_mask[:, -response_length:]
307	token_level_rewards = data.batch['token_level_rewards']	308	token_level_rewards = data.batch['token_level_rewards']
308		309
309	- # TODO: add other ways to estimate advantages
310	if config.algorithm.adv_estimator == 'gae':	310	if config.algorithm.adv_estimator == 'gae':
311	values = data.batch['values']	311	values = data.batch['values']
312	advantages, returns = compute_gae_advantage_return(token_level_rewards=token_level_rewards,	312	advantages, returns = compute_gae_advantage_return(token_level_rewards=token_level_rewards,
	@@ -398,7 +398,7 @@ def compute_verifier_score(batch, metrics, config, valid_mask):
398		398
399	logger.logger.info("=" * 50)	399	logger.logger.info("=" * 50)
400		400
401	- labels = [label for label, mask in zip(extra_data["labels"], valid_mask) if mask]	401	+ labels = [label for label, mask in zip(extra_data.get("labels"), valid_mask) if mask]
402	scores = verifier(str_responses, labels, config, metrics, infos=None)	402	scores = verifier(str_responses, labels, config, metrics, infos=None)
403		403
404	scores = torch.tensor(	404	scores = torch.tensor(
	@@ -485,7 +485,6 @@ def reduce_metrics(metrics: dict):
485		485
486		486
487	def compute_data_metrics(batch):	487	def compute_data_metrics(batch):
488	- # TODO: add response length
489	sequence_score = batch.batch['rm_scores'].sum(-1)	488	sequence_score = batch.batch['rm_scores'].sum(-1)
490	sequence_reward = batch.batch['token_level_rewards'].sum(-1)	489	sequence_reward = batch.batch['token_level_rewards'].sum(-1)
491		490

Mmindspeed_llm/tasks/posttrain/rlxf/training/parallel_state.py+3-3

	@@ -51,11 +51,11 @@ def initialize_model_parallel_2megatron(
51	if nccl_communicator_config_path is not None:	51	if nccl_communicator_config_path is not None:
52	try:	52	try:
53	import yaml	53	import yaml
54	- except ImportError:	54	+ except ImportError as e:
55	- raise RuntimeError(	55	+ raise ImportError(
56	"Cannot import `yaml`. Setting custom nccl communicator configs "	56	"Cannot import `yaml`. Setting custom nccl communicator configs "
57	"requires the yaml package."	57	"requires the yaml package."
58	- )	58	+ ) from e
59		59
60	with open(nccl_communicator_config_path, "r") as stream:	60	with open(nccl_communicator_config_path, "r") as stream:
61	nccl_comm_cfgs = yaml.safe_load(stream)	61	nccl_comm_cfgs = yaml.safe_load(stream)

Mmindspeed_llm/tasks/posttrain/rlxf/training/parameter_mapping.py+1-1

	@@ -156,7 +156,7 @@ def recv_model_from_train_model(moudle: torch.nn.Module):
156	idx = 0	156	idx = 0
157	if args.md5_validate:	157	if args.md5_validate:
158	hash_value = hashlib.md5()	158	hash_value = hashlib.md5()
159	- for name, param in moudle.named_parameters():	159	+ for _, param in moudle.named_parameters():
160	if flag:	160	if flag:
161	cur_num = int(recv_param_nums[idx])	161	cur_num = int(recv_param_nums[idx])
162	cur_group = model_receive_groups[idx]	162	cur_group = model_receive_groups[idx]

Mmindspeed_llm/tasks/posttrain/rlxf/utils/protocol.py+5-6

	@@ -17,6 +17,8 @@ Implement base data transfer protocol between any two functions, modules.
17	We can subclass Protocol to define more detailed batch info with specific keys	17	We can subclass Protocol to define more detailed batch info with specific keys
18	"""	18	"""
19		19
		20	+__all__ = ['DataProto', 'union_tensor_dict']
		21	+
20	import copy	22	import copy
21	from dataclasses import dataclass, field	23	from dataclasses import dataclass, field
22	from typing import Callable, Dict, List, Union	24	from typing import Callable, Dict, List, Union
	@@ -28,8 +30,6 @@ import tensordict
28	from tensordict import TensorDict	30	from tensordict import TensorDict
29	from torch.utils.data import DataLoader	31	from torch.utils.data import DataLoader
30		32
31	-__all__ = ['DataProto', 'union_tensor_dict']
32	-
33	try:	33	try:
34	tensordict.set_lazy_legacy(False).set()	34	tensordict.set_lazy_legacy(False).set()
35	except Exception as e:	35	except Exception as e:
	@@ -168,7 +168,6 @@ class DataProto:
168	raise ValueError('only support num_batch_dims=1')	168	raise ValueError('only support num_batch_dims=1')
169		169
170	if len(self.non_tensor_batch) != 0:	170	if len(self.non_tensor_batch) != 0:
171	- # TODO: we can actually lift this restriction if needed
172	if len(self.batch.batch_size) != 1:	171	if len(self.batch.batch_size) != 1:
173	raise ValueError('only support num_batch_dims=1 when non_tensor_batch is not empty.')	172	raise ValueError('only support num_batch_dims=1 when non_tensor_batch is not empty.')
174		173
	@@ -515,9 +514,9 @@ class DataProtoFuture:
515		514
516	def get(self):	515	def get(self):
517	output = ray.get(self.futures) # dp_size.	516	output = ray.get(self.futures) # dp_size.
518	- for o in output:	517	+ for single_output in output:
519	- if not isinstance(o, DataProto):	518	+ if not isinstance(single_output, DataProto):
520	- raise TypeError(f"Expected instance of DataProto, but got {type(o)}.")	519	+ raise TypeError(f"Expected instance of DataProto, but got {type(single_output)}.")
521	output = self.collect_fn(output) # select dp, concat	520	output = self.collect_fn(output) # select dp, concat
522	if self.dispatch_fn is not None:	521	if self.dispatch_fn is not None:
523	output = self.dispatch_fn(output) # split in batch dim, select using dp	522	output = self.dispatch_fn(output) # split in batch dim, select using dp

Mmindspeed_llm/tasks/posttrain/rlxf/workers/actor_train_infer.py+5-6

	@@ -111,7 +111,6 @@ class PPOActorWorker(MegatronWorker):
111		111
112	metrics = self.node.actor.update_policy(dataloader=dataloader)	112	metrics = self.node.actor.update_policy(dataloader=dataloader)
113		113
114	- # TODO: here, we should return all metrics
115	output = DataProto(meta_info={'metrics': metrics})	114	output = DataProto(meta_info={'metrics': metrics})
116	output = output.to('cpu')	115	output = output.to('cpu')
117	torch.cuda.empty_cache()	116	torch.cuda.empty_cache()
	@@ -189,7 +188,8 @@ def pad_to_tensor_dict(data, padding_side="right", pad_multi_of=16):
189	pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id else tokenizer.eos_token_id	188	pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id else tokenizer.eos_token_id
190	context_lengths = [len(val) for val in data]	189	context_lengths = [len(val) for val in data]
191		190
192	- for i in range(len(data)):	191	+ data_length = len(data)
		192	+ for i in range(data_length):
193	if context_lengths[i] < max_length:	193	if context_lengths[i] < max_length:
194	if padding_side == "right":	194	if padding_side == "right":
195	data[i].extend([pad_id] * (max_length - context_lengths[i]))	195	data[i].extend([pad_id] * (max_length - context_lengths[i]))
	@@ -365,7 +365,7 @@ class PPOActorInferWorker(BaseTrainer):
365	additional_val = batch.get(additional_key).view(-1).cpu().numpy().tolist()	365	additional_val = batch.get(additional_key).view(-1).cpu().numpy().tolist()
366		366
367	for _ in range(args.n_samples_per_prompt):	367	for _ in range(args.n_samples_per_prompt):
368	- additional_dict_per_step[additional_key].append(copy.deepcopy(additional_val))	368	+ additional_dict_per_step.get(additional_key).append(copy.deepcopy(additional_val))
369		369
370	for _ in range(args.n_samples_per_prompt):	370	for _ in range(args.n_samples_per_prompt):
371	idx_list_per_step.append(copy.deepcopy(tokens_list))	371	idx_list_per_step.append(copy.deepcopy(tokens_list))
	@@ -406,7 +406,7 @@ class PPOActorInferWorker(BaseTrainer):
406	)	406	)
407		407
408	for additional_key in self.args.dataset_additional_keys:	408	for additional_key in self.args.dataset_additional_keys:
409	- tmp_val = additional_dict[additional_key]	409	+ tmp_val = additional_dict.get(additional_key)
410	pad_to_tensor_dict(	410	pad_to_tensor_dict(
411	tmp_val,	411	tmp_val,
412	pad_multi_of=args.pad_to_multiple_of	412	pad_multi_of=args.pad_to_multiple_of
	@@ -590,7 +590,7 @@ class MegatronPPOActor():
590	output = self.forward_backward_batch(data, forward_only=True, post_process_fn=compute_logprobs_fn)	590	output = self.forward_backward_batch(data, forward_only=True, post_process_fn=compute_logprobs_fn)
591	if mpu.is_pipeline_last_stage(ignore_virtual=True):	591	if mpu.is_pipeline_last_stage(ignore_virtual=True):
592	# only on last rank. It should be on every tp rank	592	# only on last rank. It should be on every tp rank
593	- log_probs = torch.cat([o['log_probs'] for o in output], dim=0) # (bs, seq_size)	593	+ log_probs = torch.cat([single_output['log_probs'] for single_output in output], dim=0) # (bs, seq_size)
594	log_probs = log_probs.to(torch.float32)	594	log_probs = log_probs.to(torch.float32)
595	else:	595	else:
596	log_probs = None	596	log_probs = None
	@@ -641,7 +641,6 @@ class MegatronPPOActor():
641	- The communication shape is (total_nnz_pad_to_sp // tp_size, 1, hidden_size) if sequence parallel is enabled	641	- The communication shape is (total_nnz_pad_to_sp // tp_size, 1, hidden_size) if sequence parallel is enabled
642	"""	642	"""
643	# broadcast from last pp rank to all other pp ranks	643	# broadcast from last pp rank to all other pp ranks
644	- # TODO: actually, we just need to control the sampling order.
645		644
646	data.batch['attention_mask'] = data.batch['attention_mask'].to(bool)	645	data.batch['attention_mask'] = data.batch['attention_mask'].to(bool)
647		646

Mmindspeed_llm/tasks/posttrain/trl_ppo/utils.py+2-1

	@@ -110,7 +110,8 @@ def pad_to_tensor_dict(data, padding_side="right", pad_multi_of=16):
110	else:	110	else:
111	ori_context_lengths.append(torch.nonzero(torch.tensor(val) == pad_id).min().item() + 1)	111	ori_context_lengths.append(torch.nonzero(torch.tensor(val) == pad_id).min().item() + 1)
112		112
113	- for i in range(len(data)):	113	+ data_length = len(data)
		114	+ for i in range(data_length):
114	if context_lengths[i] < max_length:	115	if context_lengths[i] < max_length:
115	if padding_side == "right":	116	if padding_side == "right":
116	data[i].extend([pad_id] * (max_length - context_lengths[i]))	117	data[i].extend([pad_id] * (max_length - context_lengths[i]))

Mmindspeed_llm/tasks/posttrain/verifier/math_eval_toolkit/grader.py+21-18

	@@ -1,4 +1,5 @@
1	import re	1	import re
		2	+import warnings
2	import multiprocessing	3	import multiprocessing
3	from math import isclose	4	from math import isclose
4	from typing import Union	5	from typing import Union
	@@ -53,7 +54,8 @@ def str_to_pmatrix(input_str):
53		54
54	for m in matrix_str:	55	for m in matrix_str:
55	m = m.strip("{}")	56	m = m.strip("{}")
56	- pmatrix = r"\begin{pmatrix}" + m.replace(",", "\\") + r"\end{pmatrix}"	57	+ m_processed = m.replace(",", "\\")
		58	+ pmatrix = fr"\begin{{{m_processed}}}\end{{pmatrix}}"
57	pmatrix_list.append(pmatrix)	59	pmatrix_list.append(pmatrix)
58		60
59	return ", ".join(pmatrix_list)	61	return ", ".join(pmatrix_list)
	@@ -98,7 +100,8 @@ def math_equal(
98	else:	100	else:
99	if item == prediction:	101	if item == prediction:
100	return True	102	return True
101	- except Exception:	103	+ except Exception as e:
		104	+ warnings.warn(f"An exception occurred during comparison: {e}")
102	continue	105	continue
103	return False	106	return False
104	except (ValueError, TypeError, AttributeError) as e:	107	except (ValueError, TypeError, AttributeError) as e:
	@@ -164,20 +167,20 @@ def math_equal(
164	reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}")	167	reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}")
165	)	168	)
166	):	169	):
167	- pred_lines = [	170	+ pred_lines = []
168	- line.strip()	171	+ prediction_processed = prediction[len("\\begin{pmatrix}"): -len("\\end{pmatrix}")]
169	- for line in prediction[	172	+ for line in prediction_processed.split("\\\\"):
170	- len("\\begin{pmatrix}") : -len("\\end{pmatrix}")	173	+ stripped_line = line.strip()
171	- ].split("\\\\")	174	+ if stripped_line:
172	- if line.strip()	175	+ pred_lines.append(stripped_line)
173	- ]	176	+
174	- ref_lines = [	177	+ ref_lines = []
175	- line.strip()	178	+ reference_processed = reference[len("\\begin{pmatrix}"): -len("\\end{pmatrix}")]
176	- for line in reference[	179	+ for line in reference_processed.split("\\\\"):
177	- len("\\begin{pmatrix}") : -len("\\end{pmatrix}")	180	+ stripped_line = line.strip()
178	- ].split("\\\\")	181	+ if stripped_line:
179	- if line.strip()	182	+ ref_lines.append(stripped_line)
180	- ]	183	+
181	matched = True	184	matched = True
182	if len(pred_lines) == len(ref_lines):	185	if len(pred_lines) == len(ref_lines):
183	for pred_line, ref_line in zip(pred_lines, ref_lines):	186	for pred_line, ref_line in zip(pred_lines, ref_lines):
	@@ -255,10 +258,10 @@ def symbolic_equal(a, b):
255	for f in [parse_latex, parse_expr, latex2sympy]:	258	for f in [parse_latex, parse_expr, latex2sympy]:
256	try:	259	try:
257	return f(s.replace("\\\\", "\\"))	260	return f(s.replace("\\\\", "\\"))
258	- except:	261	+ except (SyntaxError, ValueError) as e1:
259	try:	262	try:
260	return f(s)	263	return f(s)
261	- except:	264	+ except (SyntaxError, ValueError) as e2:
262	pass	265	pass
263	return s	266	return s
264		267

Mmindspeed_llm/tasks/posttrain/verifier/math_eval_toolkit/parser.py+10-2

	@@ -7,7 +7,15 @@ import sympy
7	from latex2sympy2 import latex2sympy	7	from latex2sympy2 import latex2sympy
8	from word2number import w2n	8	from word2number import w2n
9		9
10	-from mindspeed_llm.tasks.posttrain.verifier.math_eval_toolkit.utils import *	10	+from mindspeed_llm.tasks.posttrain.verifier.math_eval_toolkit.utils import (
		11	+ set_seed,
		12	+ load_jsonl,
		13	+ save_jsonl,
		14	+ lower_keys,
		15	+ PROMPT_TEMPLATES,
		16	+ key_map,
		17	+ show_sample
		18	+)
11		19
12		20
13	def _fix_fracs(string):	21	def _fix_fracs(string):
	@@ -52,7 +60,7 @@ def _fix_a_slash_b(string):
52	b = int(b)	60	b = int(b)
53	if string != "{}/{}".format(a, b):	61	if string != "{}/{}".format(a, b):
54	raise ValueError(f"String does not match the expected format: {string}")	62	raise ValueError(f"String does not match the expected format: {string}")
55	- new_string = "\\frac{" + str(a) + "}{" + str(b) + "}"	63	+ new_string = f"\\frac{{{a}}}{{{b}}}"
56	return new_string	64	return new_string
57	except ValueError:	65	except ValueError:
58	return string	66	return string

Mmindspeed_llm/tasks/posttrain/verifier/rule_verifier.py+23-9

	@@ -157,15 +157,29 @@ def validate_response_structure(processed_str: str) -> bool:
157	if count != expected_count:	157	if count != expected_count:
158	validation_passed = False	158	validation_passed = False
159		159
160	- if (positions['think_start'] > positions['think_end'] or	160	+ think_start = positions.get('think_start')
161	- positions['think_end'] > positions['answer_start'] or	161	+ think_end = positions.get('think_end')
162	- positions['answer_start'] > positions['boxed_start'] or	162	+ answer_start = positions.get('answer_start')
163	- positions['boxed_start'] > positions['answer_end'] or	163	+ boxed_start = positions.get('boxed_start')
164	- not processed_str.startswith('<think>') or	164	+ answer_end = positions.get('answer_end')
165	- not processed_str.endswith('</answer>')	165	+
166	- ):	166	+ is_think_start_valid = think_start > think_end
167	- validation_passed = False	167	+ is_think_end_valid = think_end > answer_start
168	- else:	168	+ is_answer_start_valid = answer_start > boxed_start
		169	+ is_boxed_start_valid = boxed_start > answer_end
		170	+ is_start_with_think = not processed_str.startswith('<think>')
		171	+ is_end_with_answer = not processed_str.endswith('</answer>')
		172	+
		173	+ validation_passed = not (
		174	+ is_think_start_valid or
		175	+ is_think_end_valid or
		176	+ is_answer_start_valid or
		177	+ is_boxed_start_valid or
		178	+ is_start_with_think or
		179	+ is_end_with_answer
		180	+ )
		181	+
		182	+ if not validation_passed:
169	pass	183	pass
170		184
171	return validation_passed	185	return validation_passed

Mmindspeed_llm/tasks/preprocess/data_handler.py+3-3

	@@ -13,6 +13,8 @@
13	# See the License for the specific language governing permissions and	13	# See the License for the specific language governing permissions and
14	# limitations under the License.	14	# limitations under the License.
15		15
		16	+__all__ = ["get_dataset_handler", "build_dataset"]
		17	+
16	import os	18	import os
17	import sys	19	import sys
18	import time	20	import time
	@@ -42,8 +44,6 @@ from .utils import (
42	logging.basicConfig(level=logging.INFO)	44	logging.basicConfig(level=logging.INFO)
43	logger = logging.getLogger(__name__)	45	logger = logging.getLogger(__name__)
44		46
45	-__all__ = ["get_dataset_handler", "build_dataset"]
46	-
47		47
48	class BaseDatasetHandler(object):	48	class BaseDatasetHandler(object):
49	"""	49	"""
	@@ -623,7 +623,7 @@ class PPOAlpacaStyleInstructionHandler(BaseDatasetHandler):
623	else:	623	else:
624	messages = example["prompt"] + example["response"]	624	messages = example["prompt"] + example["response"]
625		625
626	- for source_ids, target_ids in self.llama_factory_template.encode_multiturn(	626	+ for source_ids, _ in self.llama_factory_template.encode_multiturn(
627	tokenizer, messages, example["system"][0], example["tools"][0]	627	tokenizer, messages, example["system"][0], example["tools"][0]
628	):	628	):
629	input_ids += source_ids	629	input_ids += source_ids

Mmindspeed_llm/tasks/preprocess/formatter.py+1-1

	@@ -186,7 +186,7 @@ class ToolFormatter(Formatter):
186	content = kwargs.pop("content")	186	content = kwargs.pop("content")
187	try:	187	try:
188	tools = json.loads(content)	188	tools = json.loads(content)
189	- if not len(tools):	189	+ if not tools:
190	return [""]	190	return [""]
191		191
192	if self.tool_format == "default":	192	if self.tool_format == "default":

Mmindspeed_llm/tasks/preprocess/utils.py+12-3

	@@ -110,7 +110,7 @@ def get_dataset_list(data_args) -> List["InstructionDatasetAttr"]:
110	if len(dataset_names) != 0:	110	if len(dataset_names) != 0:
111	raise ValueError(	111	raise ValueError(
112	"Cannot open {} due to {}.".format(os.path.join(data_args.dataset_dir, DATA_CONFIG), str(err))	112	"Cannot open {} due to {}.".format(os.path.join(data_args.dataset_dir, DATA_CONFIG), str(err))
113	- )	113	+ ) from err
114	dataset_info = None	114	dataset_info = None
115		115
116	if dataset_info is not None:	116	if dataset_info is not None:
	@@ -222,10 +222,19 @@ def convert_alpaca_to_intermediate(sample: Dict[str, List[Any]], dataset_attr: "
222	else:	222	else:
223	if dataset_attr.response and isinstance(sample[dataset_attr.response], list):	223	if dataset_attr.response and isinstance(sample[dataset_attr.response], list):
224	response = [	224	response = [
225	- {"role": Role.ASSISTANT.value, "content": content} for content in sample[dataset_attr.response]	225	+ {
		226	+ "role": Role.ASSISTANT.value,
		227	+ "content": content
		228	+ }
		229	+ for content in sample[dataset_attr.response]
226	]	230	]
227	elif dataset_attr.response and isinstance(sample[dataset_attr.response], str):	231	elif dataset_attr.response and isinstance(sample[dataset_attr.response], str):
228	- response = [{"role": Role.ASSISTANT.value, "content": sample[dataset_attr.response]}]	232	+ response = [
		233	+ {
		234	+ "role": Role.ASSISTANT.value,
		235	+ "content": sample[dataset_attr.response]
		236	+ }
		237	+ ]
229	else:	238	else:
230	response = []	239	response = []
231		240

Mmindspeed_llm/training/arguments.py+1-2

	@@ -1081,7 +1081,7 @@ def _validate_inference_args(args):
1081	def _validate_evaluation_args(args):	1081	def _validate_evaluation_args(args):
1082	# five shot only supported on mmlu and ceval now	1082	# five shot only supported on mmlu and ceval now
1083	if args.prompt_type is not None and hasattr(args, "task") and (args.task == "mmlu" or args.task == "ceval"):	1083	if args.prompt_type is not None and hasattr(args, "task") and (args.task == "mmlu" or args.task == "ceval"):
1084	- train_dir = os.path.dirname(args.task_data_path) + "/dev/"	1084	+ train_dir = os.path.join(os.path.dirname(args.task_data_path), "dev")
1085	if not os.path.isdir(train_dir) or not os.path.isdir(args.task_data_path):	1085	if not os.path.isdir(train_dir) or not os.path.isdir(args.task_data_path):
1086	raise ValueError(f"Test and dev directory must exists when specify prompt_type in evaluation")	1086	raise ValueError(f"Test and dev directory must exists when specify prompt_type in evaluation")
1087		1087
	@@ -1350,7 +1350,6 @@ def _add_dummy_args(args):
1350	args.npu_deterministic = False	1350	args.npu_deterministic = False
1351		1351
1352		1352
1353	-
1354	def _validate_noop_layer(args):	1353	def _validate_noop_layer(args):
1355	if isinstance(args.noop_layers, str):	1354	if isinstance(args.noop_layers, str):
1356	noop_layers = set()	1355	noop_layers = set()

Mmindspeed_llm/training/checkpointing.py+7-4

	@@ -25,7 +25,7 @@ from megatron.core.dist_checkpointing.strategies.fully_parallel import \
25	FullyParallelSaveStrategyWrapper	25	FullyParallelSaveStrategyWrapper
26	from megatron.training.utils import print_rank_0, unwrap_model, append_to_progress_log, is_last_rank	26	from megatron.training.utils import print_rank_0, unwrap_model, append_to_progress_log, is_last_rank
27	from megatron.training.async_utils import schedule_async_save	27	from megatron.training.async_utils import schedule_async_save
28	-from megatron.training.checkpointing import (_load_base_checkpoint,get_rng_state, get_checkpoint_name,	28	+from megatron.training.checkpointing import (_load_base_checkpoint, get_rng_state, get_checkpoint_name,
29	get_distributed_optimizer_checkpoint_name,	29	get_distributed_optimizer_checkpoint_name,
30	ensure_directory_exists, generate_state_dict, get_checkpoint_tracker_filename)	30	ensure_directory_exists, generate_state_dict, get_checkpoint_tracker_filename)
31	from megatron.training.one_logger_utils import on_save_checkpoint_start, on_save_checkpoint_success	31	from megatron.training.one_logger_utils import on_save_checkpoint_start, on_save_checkpoint_success
	@@ -239,7 +239,8 @@ def save_checkpoint_wrapper(fn):
239	torch.save(state_dict, checkpoint_name)	239	torch.save(state_dict, checkpoint_name)
240	start_misc = time()	240	start_misc = time()
241	if not args.async_save:	241	if not args.async_save:
242	- assert async_save_request is None	242	+ if async_save_request is not None:
		243	+ raise ValueError("async_save_request should be None")
243	# Wait so everyone is done (necessary)	244	# Wait so everyone is done (necessary)
244	if torch.distributed.is_initialized():	245	if torch.distributed.is_initialized():
245	torch.distributed.barrier()	246	torch.distributed.barrier()
	@@ -259,7 +260,8 @@ def save_checkpoint_wrapper(fn):
259	barrier=False)	260	barrier=False)
260		261
261	if args.async_save:	262	if args.async_save:
262	- assert async_save_request is not None	263	+ if async_save_request is not None:
		264	+ raise ValueError("async_save_request should be None")
263	async_save_request.add_finalize_fn(iter_finalize_fn)	265	async_save_request.add_finalize_fn(iter_finalize_fn)
264	else:	266	else:
265	iter_finalize_fn()	267	iter_finalize_fn()
	@@ -271,7 +273,8 @@ def save_checkpoint_wrapper(fn):
271	on_save_checkpoint_success(productive_metrics, args.async_save)	273	on_save_checkpoint_success(productive_metrics, args.async_save)
272		274
273	if args.async_save:	275	if args.async_save:
274	- assert async_save_request is not None	276	+ if async_save_request is not None:
		277	+ raise ValueError("async_save_request should be None")
275	async_save_request.add_finalize_fn(onelogger_finalize_fn)	278	async_save_request.add_finalize_fn(onelogger_finalize_fn)
276	else:	279	else:
277	onelogger_finalize_fn()	280	onelogger_finalize_fn()

	@@ -433,32 +433,32 @@ def share_kvstates_checkpointed_forward_func(
433	# Uniformly divide the total number of Transformer layers and checkpoint	433	# Uniformly divide the total number of Transformer layers and checkpoint
434	# the input activation of each divided chunk.	434	# the input activation of each divided chunk.
435	# A method to further reduce memory usage reducing checkpoints.	435	# A method to further reduce memory usage reducing checkpoints.
436	- l = 0	436	+ layer = 0
437	- while l < self.num_layers_per_pipeline_rank:	437	+ while layer < self.num_layers_per_pipeline_rank:
438	hidden_states, context, key_value_states = checkpoint_handler(	438	hidden_states, context, key_value_states = checkpoint_handler(
439	- custom(l, l + self.config.recompute_num_layers)	439	+ custom(layer, layer + self.config.recompute_num_layers)
440	)	440	)
441		441
442	- l += self.config.recompute_num_layers	442	+ layer += self.config.recompute_num_layers
443		443
444	elif self.config.recompute_method == 'block':	444	elif self.config.recompute_method == 'block':
445	# Checkpoint the input activation of only a set number of individual	445	# Checkpoint the input activation of only a set number of individual
446	# Transformer layers and skip the rest.	446	# Transformer layers and skip the rest.
447	# A method fully use the device memory removing redundant re-computation.	447	# A method fully use the device memory removing redundant re-computation.
448	recompute_skip_num_layers = 0	448	recompute_skip_num_layers = 0
449	- for l in range(self.num_layers_per_pipeline_rank):	449	+ for layer in range(self.num_layers_per_pipeline_rank):
450	# Skip recomputation when input grad computation is not needed.	450	# Skip recomputation when input grad computation is not needed.
451	# Need to have at least one input tensor with gradient computation	451	# Need to have at least one input tensor with gradient computation
452	# for re-enterant autograd engine.	452	# for re-enterant autograd engine.
453	if self.config.fp8 and not hidden_states.requires_grad:	453	if self.config.fp8 and not hidden_states.requires_grad:
454	recompute_skip_num_layers += 1	454	recompute_skip_num_layers += 1
455	if (	455	if (
456	- l >= recompute_skip_num_layers	456	+ layer >= recompute_skip_num_layers
457	- and l < self.config.recompute_num_layers + recompute_skip_num_layers	457	+ and layer < self.config.recompute_num_layers + recompute_skip_num_layers
458	):	458	):
459	- hidden_states, context, key_value_states = checkpoint_handler(custom(l, l + 1))	459	+ hidden_states, context, key_value_states = checkpoint_handler(custom(layer, layer + 1))
460	else:	460	else:
461	- hidden_states, context, key_value_states = custom(l, l + 1)(	461	+ hidden_states, context, key_value_states = custom(layer, layer + 1)(
462	hidden_states,	462	hidden_states,
463	attention_mask,	463	attention_mask,
464	context,	464	context,