diff --git a/mindformers/models/modeling_utils.py b/mindformers/models/modeling_utils.py index 207a3bd20d80cd8b17d906d14bc17232fce58005..849d7b4bcddb84b2ea173a6ba118c08617bac764 100644 --- a/mindformers/models/modeling_utils.py +++ b/mindformers/models/modeling_utils.py @@ -301,6 +301,27 @@ class PreTrainedModel(nn.Cell, ModuleUtilsMixin, GenerationMixin, PushToHubMixin """ return getattr(self, self.base_model_prefix, self) + def check_pipeline_stage(self): + """check pipeline_stage and num_layers""" + config = self.config + parallel_mode = ms.get_auto_parallel_context("parallel_mode") + pp = config.parallel_config.pipeline_stage + if parallel_mode in ["semi_auto_parallel"]: + num_layers = config.num_layers + if num_layers and num_layers < pp: + raise ValueError( + f"num_layers of model should be greater than or equal to pipeline_stage, " + f"but get num_layers ({num_layers}) < pp({pp})" + ) + pipeline_interleave_enabled = ms.get_auto_parallel_context("pipeline_interleave") + pp_interleave_num = getattr(config, 'pp_interleave_num', 0) or 0 + if pipeline_interleave_enabled and pp_interleave_num * pp > num_layers: + raise ValueError( + f"num_layers should be greater than `pp * pp_interleave_num`, " + f"but got num_layers : {num_layers} " + f"and pp * pp_interleave_num = {pp * pp_interleave_num}." + ) + @classmethod def can_generate(cls) -> bool: """ @@ -582,7 +603,6 @@ class PreTrainedModel(nn.Cell, ModuleUtilsMixin, GenerationMixin, PushToHubMixin val.__dict__.pop("type") config.__dict__.update({key: val}) - def prepare_inputs_for_predict_layout(self, input_ids, **kwargs): """ prepare inputs for transform ckpt. diff --git a/mindformers/tools/check_rules.py b/mindformers/tools/check_rules.py index eedd02f6f798706fb815700f84cf8841752cfca9..2a256186ac300135929cd644e45e940d9b9e2fc7 100644 --- a/mindformers/tools/check_rules.py +++ b/mindformers/tools/check_rules.py @@ -125,17 +125,6 @@ def _check_full_batch(): f"but get {parallel_mode}, full_batch has been forced to False") -def _check_pipeline_interleave(config, pp): - """check vpp config""" - pipeline_interleave_enabled = getattr(config.parallel.pipeline_config, 'pipeline_interleave', False) - if not pipeline_interleave_enabled: - return False - # Set pp_interleave_num to 0 if there is no pp_interleave_num in model_config - # or if pp_interleave_num is set to None. - pp_interleave_num = getattr(config.model.model_config, 'pp_interleave_num', 0) or 0 - return pp_interleave_num * pp > config.model.model_config.num_layers - - def _check_context_parallel_algo_valid(config, cp, mp): """check cp config""" n_kv_heads = getattr(config.model.model_config, 'n_kv_heads', None) @@ -172,10 +161,6 @@ def _check_parallel(config): f"be equal to device_num, but get dp*mp*sp*pp = {dp}*{mp}*{cp}*{pp} = {dp * mp * cp * pp} " f"!= device_num({device_num})") - if config.model.model_config.num_layers and config.model.model_config.num_layers < pp: - raise ValueError(f"num_layers of model should be greater than or equal to pipeline_stage, but get " - f"num_layers ({config.model.model_config.num_layers}) < pp({pp})") - if server_num > 1: if server_num % pp != 0: logger.warning(f"server_num % pipeline_stage = {server_num} % {pp} = {server_num % pp} != 0, " @@ -196,10 +181,6 @@ def _check_parallel(config): f"use_flash_attention {config.model.model_config.use_flash_attention}, please " f"set use_flash_attention=True") - if _check_pipeline_interleave(config, pp): - raise ValueError(f"num_layers should be greater than `pp * pp_interleave_num`, " - f"but got num_layers : {config.model.model_config.num_layers} " - f"and pp * pp_interleave_num = {pp * config.model.model_config.pp_interleave_num}.") if cp > 1: _check_context_parallel_algo_valid(config, cp, mp) diff --git a/mindformers/trainer/base_trainer.py b/mindformers/trainer/base_trainer.py index 285ba893b9f8f75a0069e8fbd028fc23b504a333..d6ef29abd4d13b76f4d551485f32de37fb357b7a 100644 --- a/mindformers/trainer/base_trainer.py +++ b/mindformers/trainer/base_trainer.py @@ -452,7 +452,10 @@ class BaseTrainer: def create_network(self, default_args: dict = None): """Create the network for task trainer.""" logger.info(".........Build Network From Config..........") - return build_network(self.config.model, default_args=default_args) + network = build_network(self.config.model, default_args=default_args) + if hasattr(network, "check_pipeline_stage") and callable(network.check_pipeline_stage): + network.check_pipeline_stage() + return network def create_network_without_param_init(self, default_args: dict = None): """Create the network for task trainer without initialize parameters.""" diff --git a/research/deepseek3/deepseek3_model_train.py b/research/deepseek3/deepseek3_model_train.py index 76da73df827b58f321e73b42c6710c75209091a2..2a1b691622ec6ef59934a30c36aaca22358a46f3 100644 --- a/research/deepseek3/deepseek3_model_train.py +++ b/research/deepseek3/deepseek3_model_train.py @@ -13,6 +13,7 @@ # limitations under the License. # ============================================================================ """DeepseekV3 models' APIs.""" +import mindspore as ms from mindspore.common import dtype as mstype from mindspore.ops import operations as P @@ -122,3 +123,24 @@ class TrainingDeepseekV3ForCausalLM(DeepseekV2ForCausalLM): x = self.concat_2d((x, self.cast(pad_zeros, x.dtype))) return x + + def check_pipeline_stage(self): + """check pipeline_stage and num_layers""" + config = self.config + parallel_mode = ms.get_auto_parallel_context("parallel_mode") + pp = config.parallel_config.pipeline_stage + if parallel_mode in ["semi_auto_parallel"]: + num_layers = config.num_layers + if num_layers and num_layers + config.mtp_depth < pp: + raise ValueError( + f"num_layers + mtp_depth of model should be greater than or equal to pipeline_stage, " + f"but get num_layers ({num_layers}) + mtp_depth ({config.mtp_depth}) < pp({pp})" + ) + pipeline_interleave_enabled = ms.get_auto_parallel_context("pipeline_interleave") + pp_interleave_num = getattr(config, 'pp_interleave_num', 0) or 0 + if pipeline_interleave_enabled and pp_interleave_num * pp > num_layers + config.mtp_depth: + raise ValueError( + f"num_layers + mtp_depth should be greater than `pp * pp_interleave_num`, " + f"but got num_layers + mtp_depth : {num_layers} + {config.mtp_depth} " + f"and pp * pp_interleave_num = {pp * pp_interleave_num}." + )