参数副本确定性计算Bug修复 & UT补充构建-MindSpeed-AtomGit

已合并

参数副本确定性计算Bug修复 & UT补充构建 #1452

AtomGit-Bot创建于 2024年11月25日从refs/pull/1452/head 合入到 master

参数副本确定性计算Bug修复 & UT补充构建 #1452

已合并

AtomGit-Bot创建于 2024年11月25日

从refs/pull/1452/head合入到master

共 3 个文件变更+249-33

Mmindspeed/optimizer/distrib_optimizer.py+50-32

	@@ -622,11 +622,11 @@ def fp16_tensor_convert_to_fp32_tensor(self):
622	shard_fp32_main_param_view[-remain * 2:].copy_(	622	shard_fp32_main_param_view[-remain * 2:].copy_(
623	workspace_convert_view.view(2, -1).transpose(1, 0).reshape(-1).contiguous())	623	workspace_convert_view.view(2, -1).transpose(1, 0).reshape(-1).contiguous())
624		624
625	if not self.first_sub_flag:	625	if not self.first_sub_flag and not npu_deterministic:
626	if npu_deterministic:	626	shard_main_param_int32_view_bucket[:param_data_dp_numel].sub_(32768)
		A A aijgnem12024年11月27日加个注释吧 `32768`
627	fp16_tensor_convert_to_fp32_tensor_deterministic(self.shard_fp32_from_float16_groups, self.optimizer)	627
628	else:	628	if not self.first_sub_flag and npu_deterministic:
629	shard_main_param_int32_view_bucket[:param_data_dp_numel].sub_(32768)	629	fp16_tensor_convert_to_fp32_tensor_deterministic(self.shard_fp32_from_float16_groups, self.optimizer)
630		630
631		631
632	def fp32_tensor_convert_to_fp16_tensor(self):	632	def fp32_tensor_convert_to_fp16_tensor(self):
	@@ -701,6 +701,51 @@ def fp32_tensor_convert_to_fp16_tensor(self):
701	shard_fp32_main_param_view[param_data_dp_numel:param_data_dp_numel * 2].copy_(shard_fp32_main_param_view[:param_data_dp_numel])	701	shard_fp32_main_param_view[param_data_dp_numel:param_data_dp_numel * 2].copy_(shard_fp32_main_param_view[:param_data_dp_numel])
702		702
703		703
		704	def fp16_tensor_convert_to_fp32_tensor_deterministic(shard_fp32_from_float16_groups, optimizer):
		705	assert hasattr(optimizer, "state")
		706	for shard_fp32_from_float16_group in shard_fp32_from_float16_groups:
		707	for shard_fp32_param in shard_fp32_from_float16_group:
		708	if "exp_avg_sq" not in optimizer.state[shard_fp32_param]:
		709	continue
		710	shard_int32_tensor = shard_fp32_param.view(torch.int32)
		711	assert shard_int32_tensor.numel() == shard_fp32_param.numel()
		712	loops = shard_int32_tensor.numel() // TRANSPOSE_BF16_BLOCK_SIZE
		713	remain = shard_int32_tensor.numel() % TRANSPOSE_BF16_BLOCK_SIZE
		714	exp_avg_sq_flatten = optimizer.state[shard_fp32_param]["exp_avg_sq"].reshape(-1)
		715	for loop in range(loops):
		716	odd_even_tensor = torch.sign(exp_avg_sq_flatten[loop * TRANSPOSE_BF16_BLOCK_SIZE: (loop + 1) * TRANSPOSE_BF16_BLOCK_SIZE] > 0)
		717	shard_int32_tensor[loop * TRANSPOSE_BF16_BLOCK_SIZE: (loop + 1) * TRANSPOSE_BF16_BLOCK_SIZE].add_(odd_even_tensor)
		718	if remain > 0:
		719	odd_even_tensor = torch.sign(exp_avg_sq_flatten[-remain:] > 0)
		720	shard_int32_tensor[-remain:].add_(odd_even_tensor)
		721	shard_int32_tensor.sub_(32768)
		722	optimizer.state[shard_fp32_param]["exp_avg_sq"].abs_()
		723
		724
		725	def fp32_tensor_convert_to_fp16_tensor_deterministic(shard_fp32_from_float16_groups, optimizer):
		726	assert hasattr(optimizer, "state")
		727	for shard_fp32_from_float16_group in shard_fp32_from_float16_groups:
		728	for shard_fp32_param in shard_fp32_from_float16_group:
		729	if "exp_avg_sq" not in optimizer.state[shard_fp32_param]:
		730	continue
		731	shard_int32_tensor = shard_fp32_param.view(torch.int32)
		732	assert shard_int32_tensor.numel() == shard_fp32_param.numel()
		733	loops = shard_int32_tensor.numel() // TRANSPOSE_BF16_BLOCK_SIZE
		734	remain = shard_int32_tensor.numel() % TRANSPOSE_BF16_BLOCK_SIZE
		735	exp_avg_sq_flatten = optimizer.state[shard_fp32_param]["exp_avg_sq"].reshape(-1)
		736	shard_int32_tensor.add_(32768)
		737	for loop in range(loops):
		738	odd_even_tensor = ((shard_int32_tensor[loop * TRANSPOSE_BF16_BLOCK_SIZE: (loop + 1) * TRANSPOSE_BF16_BLOCK_SIZE] & 131071) == 65536).int()
		739	shard_int32_tensor[loop * TRANSPOSE_BF16_BLOCK_SIZE: (loop + 1) * TRANSPOSE_BF16_BLOCK_SIZE].sub_(odd_even_tensor)
		740	sign_tensor = torch.sign(odd_even_tensor - 0.5)
		741	exp_avg_sq_flatten[loop * TRANSPOSE_BF16_BLOCK_SIZE: (loop + 1) * TRANSPOSE_BF16_BLOCK_SIZE].mul_(sign_tensor)
		742	if remain > 0:
		743	odd_even_tensor = ((shard_int32_tensor[-remain:] & 131071) == 65536).int()
		744	shard_int32_tensor[-remain:].sub_(odd_even_tensor)
		745	sign_tensor = torch.sign(odd_even_tensor - 0.5)
		746	exp_avg_sq_flatten[-remain:].mul_(sign_tensor)
		747
		748
704	def get_parameter_state_dp_zero_hccl(self):	749	def get_parameter_state_dp_zero_hccl(self):
705	"""	750	"""
706	Replace the communication method of gather from gloo to hccl.	751	Replace the communication method of gather from gloo to hccl.
	@@ -908,30 +953,3 @@ def load_parameter_state_from_dp_zero_hccl(self, state_dict):
908	tensor_to_copy_into.data.copy_(	953	tensor_to_copy_into.data.copy_(
909	recv_tensor[gbuf_local_start:gbuf_local_end]	954	recv_tensor[gbuf_local_start:gbuf_local_end]
910	)	955	)
911
912
913	def fp16_tensor_convert_to_fp32_tensor_deterministic(shard_fp32_from_float16_groups, optimizer):
914	assert hasattr(optimizer, "state")
915	for shard_fp32_from_float16_group in shard_fp32_from_float16_groups:
916	for shard_fp32_param in shard_fp32_from_float16_group:
917	if "exp_avg_sq" in optimizer.state[shard_fp32_param]:
918	shard_int32_tensor = torch.tensor(shard_fp32_param.untyped_storage(), dtype=torch.int32, device=shard_fp32_param.device)
919	assert shard_int32_tensor.numel() == shard_fp32_param.numel()
920	odd_even_tensor = (torch.sign(optimizer.state[shard_fp32_param]["exp_avg_sq"]) > 0).reshape(-1)
921	optimizer.state[shard_fp32_param]["exp_avg_sq"].abs_()
922	shard_int32_tensor.add_(odd_even_tensor)
923	shard_int32_tensor.sub_(32768)
924
925
926	def fp32_tensor_convert_to_fp16_tensor_deterministic(shard_fp32_from_float16_groups, optimizer):
927	assert hasattr(optimizer, "state")
928	for shard_fp32_from_float16_group in shard_fp32_from_float16_groups:
929	for shard_fp32_param in shard_fp32_from_float16_group:
930	if "exp_avg_sq" in optimizer.state[shard_fp32_param]:
931	shard_int32_tensor = torch.tensor(shard_fp32_param.untyped_storage(), dtype=torch.int32, device=shard_fp32_param.device)
932	assert shard_int32_tensor.numel() == shard_fp32_param.numel()
933	odd_even_tensor = ((shard_int32_tensor & 131071) == 32768).int()
934	shard_int32_tensor.add_(32768)
935	shard_int32_tensor.sub_(odd_even_tensor)
936	sign_tensor = torch.sign(odd_even_tensor - 0.5).reshape(optimizer.state[shard_fp32_param]["exp_avg_sq"].shape)
937	optimizer.state[shard_fp32_param]["exp_avg_sq"].mul_(sign_tensor)

Mtests_extend/unit_tests/features/algorithm/test_reuse_dptr.py+41-1

	@@ -1,15 +1,55 @@
		1	from unittest import mock
		2
1	import pytest	3	import pytest
2	import torch	4	import torch
3	import torch_npu	5	import torch_npu
4		6
		7	from mindspeed import megatron_adaptor
5	from mindspeed.op_builder import AlgorithmOpBuilder	8	from mindspeed.op_builder import AlgorithmOpBuilder
		9	from mindspeed.optimizer.optimizer import (
		10	bf16_tensors_to_fp32_tensors, \
		11	fp32_tensors_to_bf16_tensors, \
		12	bf16_tensors_to_fp32_tensors_deterministic, \
		13	fp32_tensors_to_bf16_tensors_deterministic
		14	)
6		15
7		16
8	class TestReuseDataPtr:	17	class TestReuseDataPtr:
9
10	def test_reuse_dptr(self):	18	def test_reuse_dptr(self):
11	mindspeed_ops = AlgorithmOpBuilder().load()	19	mindspeed_ops = AlgorithmOpBuilder().load()
12	fp32_tensor = torch.randn(10, dtype=torch.float).npu()	20	fp32_tensor = torch.randn(10, dtype=torch.float).npu()
13	int8_tensor = torch.empty(fp32_tensor.numel() * 4, dtype=torch.int8).npu()	21	int8_tensor = torch.empty(fp32_tensor.numel() * 4, dtype=torch.int8).npu()
14	mindspeed_ops.reuse_data_ptr(int8_tensor, fp32_tensor, 0)	22	mindspeed_ops.reuse_data_ptr(int8_tensor, fp32_tensor, 0)
15	assert fp32_tensor.data_ptr() == int8_tensor.data_ptr()	23	assert fp32_tensor.data_ptr() == int8_tensor.data_ptr()
		24
		25
		26	class TestReuseFp32Param:
		27	fp32_tensors = [torch.randn(512, 512, dtype=torch.float32).npu()]
		28	fp32_truth_tensors = [tensor.clone() for tensor in fp32_tensors]
		29	bf16_truth_tensors = [tensor.bfloat16().reshape(-1) for tensor in fp32_tensors]
		30
		31	int32_tensors = [tensor.view(torch.int32).reshape(-1) for tensor in fp32_tensors]
		32	bf16_fp32_tensors = [tensor.view(torch.bfloat16).reshape(-1) for tensor in fp32_tensors]
		33	bf16_tensors = [tensor.view(torch.bfloat16).reshape(-1)[tensor.numel():] for tensor in fp32_tensors]
		34	optimizer = mock.MagicMock()
		35	optimizer.state = {tensor: {"exp_avg_sq": torch.ones_like(tensor)} for tensor in fp32_tensors}
		36
		37	def test_reuse_fp32_param_deterministic(self):
		38	fp32_tensors_to_bf16_tensors_deterministic(self.int32_tensors, self.bf16_fp32_tensors, self.fp32_tensors, self.optimizer)
		39	for reuse, truth in zip(self.bf16_tensors, self.bf16_truth_tensors):
		40	assert torch.allclose(reuse, truth, rtol=0, atol=0)
		41	bf16_tensors_to_fp32_tensors_deterministic(self.int32_tensors, self.bf16_fp32_tensors, self.fp32_tensors, self.optimizer)
		42	for reuse, truth in zip(self.fp32_tensors, self.fp32_truth_tensors):
		43	assert torch.allclose(reuse, truth, rtol=0, atol=0)
		44
		45
		46	def test_reuse_fp32_param(self):
		47	# fp32 -> bf16
		48	fp32_tensors_to_bf16_tensors(int32_tensors=self.int32_tensors, bf16_fp32_tensors=self.bf16_fp32_tensors)
		49	for reuse, truth in zip(self.bf16_tensors, self.bf16_truth_tensors):
		50	assert torch.allclose(reuse, truth, rtol=0.005, atol=0.005)
		51
		52	# bf16 -> fp32
		53	bf16_tensors_to_fp32_tensors(int32_tensors=self.int32_tensors, bf16_fp32_tensors=self.bf16_fp32_tensors)
		54	for reuse, truth in zip(self.fp32_tensors, self.fp32_truth_tensors):
		55	assert torch.allclose(reuse, truth, rtol=0, atol=0)

Atests_extend/unit_tests/features/algorithm/test_reuse_optimizer.py+158-0

@@ -0,0 +1,158 @@
	1	# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
	2	from functools import partial
	3	from time import sleep
	4	import os
	5	import copy
	6	import itertools
	7	from unittest import mock
	8
	9	import pytest
	10	import torch
	11	import torch_npu
	12	from mindspeed import megatron_adaptor
	13
	14	from megatron.training.arguments import parse_args
	15	from megatron.training.global_vars import set_args
	16	from megatron.core.models.gpt import GPTModel
	17	from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
	18	from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer
	19	from megatron.core.timers import DummyTimer
	20	from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed
	21	from megatron.core.transformer import TransformerConfig
	22	from megatron.training.training import get_model
	23	from megatron.training.utils import unwrap_model
	24
	25	from unit_tests.common import DistributedTest
	26	from commons import set_random_seed, initialize_model_parallel
	27
	28
	29	def initialize_gpt_model(pre_process=True, post_process=True, seed=0, **config_kwargs):
	30	torch.manual_seed(seed)
	31	model_parallel_cuda_manual_seed(seed)
	32
	33	default_config_kwargs = dict(num_layers=8, hidden_size=512, num_attention_heads=32, use_cpu_initialization=True)
	34	default_config_kwargs.update(**config_kwargs)
	35	transformer_config = TransformerConfig(**default_config_kwargs)
	36	model = GPTModel(config=transformer_config, transformer_layer_spec=get_gpt_layer_local_spec(), vocab_size=1024, max_sequence_length=64, pre_process=pre_process, post_process=post_process)
	37
	38	model.bfloat16()
	39	with torch.no_grad():
	40	for p in model.parameters():
	41	p.random_()
	42	return model
	43
	44
	45	def init_mock_args(args, use_distributed_optimizer=False, reuse_fp32_param=False):
	46	args.data_parallel_random_init = False
	47	args.virtual_pipeline_model_parallel_size = None
	48	args.bf16 = True
	49	args.accumulate_allreduce_grads_in_fp32 = True
	50	args.use_distributed_optimizer = use_distributed_optimizer
	51	args.ddp_bucket_size = None
	52	args.reuse_fp32_param = reuse_fp32_param
	53	return args
	54
	55
	56	def setup_model_and_optimizer(seed, use_distributed_optimizer=False):
	57	model = get_model(partial(initialize_gpt_model, seed=seed))
	58	set_random_seed(seed)
	59	config = OptimizerConfig(lr=1e-4, bf16=True, params_dtype=torch.bfloat16, use_distributed_optimizer=use_distributed_optimizer)
	60	config.timers = Timers()
	61	optimizer = get_megatron_optimizer(config, model)
	62
	63	for group in optimizer.optimizer.param_groups:
	64	for p in group['params']:
	65	if len(optimizer.optimizer.state[p]) == 0:
	66	optimizer.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data)
	67	optimizer.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data)
	68	optimizer.reload_model_params()
	69	return unwrap_model(model), optimizer
	70
	71
	72	class Timers:
	73	def __init__(self, args, *kwargs):
	74	self._dummy_timer = DummyTimer()
	75
	76	def __call__(self, args, *kwargs):
	77	return self._dummy_timer
	78
	79
	80	class TestDistributedOptimizer(DistributedTest):
	81	world_size = 8
	82
	83	@pytest.mark.parametrize("is_deterministic", [False])
	84	@pytest.mark.parametrize("tp_pp", [(4, 1), (2, 2), (8, 1)])
	85	def test_reuse_float16_params_optimizer(self, tp_pp, is_deterministic):
	86	args = parse_args(None, True)
	87	args.npu_deterministic = is_deterministic
	88	set_args(args)
	89	# truth
	90	init_mock_args(args)
	91	initialize_model_parallel(*tp_pp)
	92	_, optimizer = setup_model_and_optimizer(seed=2)
	93	for _ in range(10):
	94	for float16_group in optimizer.float16_groups:
	95	for p in float16_group:
	96	p.grad = torch.randn_like(p.data, dtype=p.data.dtype)
	97	optimizer.step()
	98	truth_params = copy.deepcopy(list(itertools.chain(*optimizer.float16_groups)))
	99
	100	# reuse
	101	init_mock_args(args, reuse_fp32_param=True)
	102	_, optimizer = setup_model_and_optimizer(seed=2)
	103	for _ in range(10):
	104	for float16_group in optimizer.float16_groups:
	105	for p in float16_group:
	106	p.grad = torch.randn_like(p.data, dtype=p.data.dtype)
	107	optimizer.step()
	108	reuse_params = copy.deepcopy(list(itertools.chain(*optimizer.float16_groups)))
	109
	110	for p, reuse_p in zip(truth_params, reuse_params):
	111	if is_deterministic:
	112	assert torch.allclose(p.data, reuse_p.data, rtol=0, atol=0)
	113	else:
	114	assert torch.allclose(p.data, reuse_p.data, rtol=0.005, atol=0.005)
	115
	116
	117	@pytest.mark.parametrize("is_deterministic", [False])
	118	@pytest.mark.parametrize("overlap_grad_reduce", [True, False])
	119	@pytest.mark.parametrize("overlap_param_gather", [True, False])
	120	@pytest.mark.parametrize("tp_pp", [(4, 1), (2, 2), (8, 1)])
	121	def test_reuse_distributed_optimizer(self, tp_pp, is_deterministic, overlap_grad_reduce, overlap_param_gather):
	122	args = parse_args(None, True)
	123	args.npu_deterministic = is_deterministic
	124	args.overlap_grad_reduce = overlap_grad_reduce
	125	args.overlap_param_gather = overlap_param_gather
	126	set_args(args)
	127
	128	# truth
	129	init_mock_args(args, use_distributed_optimizer=True)
	130	initialize_model_parallel(*tp_pp)
	131	_, optimizer = setup_model_and_optimizer(seed=5, use_distributed_optimizer=True)
	132	for _ in range(10):
	133	for float16_group in optimizer.model_float16_groups:
	134	for p in float16_group:
	135	p.grad = torch.randn_like(p.data, dtype=p.data.dtype)
	136	optimizer.step()
	137	if overlap_param_gather:
	138	torch.cuda.synchronize()
	139	truth_params = copy.deepcopy(list(itertools.chain(*optimizer.model_float16_groups)))
	140
	141	# reuse
	142	init_mock_args(args, use_distributed_optimizer=True, reuse_fp32_param=True)
	143	initialize_model_parallel(*tp_pp)
	144	_, optimizer = setup_model_and_optimizer(seed=5, use_distributed_optimizer=True)
	145	for _ in range(10):
	146	for float16_group in optimizer.model_float16_groups:
	147	for p in float16_group:
	148	p.grad = torch.randn_like(p.data, dtype=p.data.dtype)
	149	optimizer.step()
	150	if overlap_param_gather:
	151	torch.cuda.synchronize()
	152	reuse_params = copy.deepcopy(list(itertools.chain(*optimizer.model_float16_groups)))
	153
	154	for p, reuse_p in zip(truth_params, reuse_params):
	155	if is_deterministic:
	156	assert torch.allclose(p.data, reuse_p.data, rtol=0, atol=0)
	157	else:
	158	assert torch.allclose(p.data, reuse_p.data, rtol=0.005, atol=0.005)