已开启
adding topology search and auto-parallel for Qwen3b model #900
adding topology search and auto-parallel for Qwen3b model #900
已开启
aminmalekisadr创建于 1月29日
3 个文件变更+2060-36
@@ -1,4 +1,5 @@
1# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.1# Copyright (c) 2025, HUAWEI CORPORATION. All rights reserved.
2+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
2"""3"""
3Note that we don't combine the main with trainer as trainer is used by other main.4Note that we don't combine the main with trainer as trainer is used by other main.
4"""5"""
@@ -18,7 +19,7 @@ from mindspeed_rl.config_cls.validate_config import validate_rl_args
18from mindspeed_rl.utils import get_tokenizer19from mindspeed_rl.utils import get_tokenizer
19from mindspeed_rl.datasets.build_dataset import build_train_valid_test_datasets20from mindspeed_rl.datasets.build_dataset import build_train_valid_test_datasets
20from mindspeed_rl.utils import seed_all21from mindspeed_rl.utils import seed_all
21-from mindspeed_rl.utils.utils import MsProbe, get_node_nums22+from mindspeed_rl.utils.utils import MsProbe, get_total_node_nums, get_required_node_nums
22from mindspeed_rl.utils.loggers import Loggers23from mindspeed_rl.utils.loggers import Loggers
23from mindspeed_rl.utils.utils import parse_args_from_config, init_torch_compile24from mindspeed_rl.utils.utils import parse_args_from_config, init_torch_compile
24from mindspeed_rl.config_cls.megatron_config import MegatronConfig25from mindspeed_rl.config_cls.megatron_config import MegatronConfig
@@ -124,8 +125,7 @@ def train(config):
124 125 
125 reward_list.append(reward_worker)126 reward_list.append(reward_worker)
126 127 
127- actor_config.max_prompt_length = rl_config.max_prompt_length128+ rule_reward_num_process = min(get_total_node_nums(), get_required_node_nums(rl_config))
128- rule_reward_num_process = get_node_nums()
129 if rl_config.rule_reward:129 if rl_config.rule_reward:
130 pg = placement_group(130 pg = placement_group(
131 [{"CPU": rl_config.num_cpus_for_local_task} for _ in range(rule_reward_num_process)],131 [{"CPU": rl_config.num_cpus_for_local_task} for _ in range(rule_reward_num_process)],
@@ -217,6 +217,88 @@ def parse_training_config(config: Dict):
217 reward_config = MegatronConfig({**config.get("megatron_training"), **config.get("reward_config")},217 reward_config = MegatronConfig({**config.get("megatron_training"), **config.get("reward_config")},
218 config.get('model'))218 config.get('model'))
219 generate_config = GenerateConfig(config.get("generate_config"))219 generate_config = GenerateConfig(config.get("generate_config"))
220+
221+ # Auto-parallelization for training (if enabled)
222+ # Wrap in try-except to prevent any side effects from breaking dataset loading
223+ try:
224+ if getattr(actor_config, "auto_parallel", False):
225+ from mindspeed_rl.workers.scheduler.launcher import auto_parallelize_training, get_npu_deployment
226+ from mindspeed_rl.workers.actor_hybrid_worker import ActorHybridWorker
227+ num_npus = get_npu_deployment(rl_config, ActorHybridWorker)
228+ if num_npus > 0:
229+ logger.info(f"Auto-parallelization (training): Enabled. num_npus={num_npus}")
230+ # check_memory=None will use megatron_config.auto_parallel_check_memory
231+ tp, pp, cp, ep = auto_parallelize_training(
232+ megatron_config=actor_config,
233+ num_npus=num_npus,
234+ prefer_tp_within_node=True,
235+ check_memory=None, # Use config value
236+ )
237+ logger.info(
238+ f"Auto-parallelization (training): Updated actor_config - "
239+ f"TP={tp}, PP={pp}, CP={cp}, EP={ep}, SP={actor_config.sequence_parallel}"
240+ )
241+ # Also update ref_config and reward_config if they share the same config
242+ if rl_config.use_integrated_worker:
243+ # ref_config and reward_config point to actor_config, so they're already updated
244+ pass
245+ else:
246+ # Update ref_config and reward_config separately if auto_parallel is enabled there too
247+ if getattr(ref_config, "auto_parallel", False):
248+ # check_memory=None will use megatron_config.auto_parallel_check_memory
249+ auto_parallelize_training(
250+ megatron_config=ref_config,
251+ num_npus=num_npus,
252+ prefer_tp_within_node=True,
253+ check_memory=None, # Use config value
254+ )
255+ if getattr(reward_config, "auto_parallel", False):
256+ # check_memory=None will use megatron_config.auto_parallel_check_memory
257+ auto_parallelize_training(
258+ megatron_config=reward_config,
259+ num_npus=num_npus,
260+ prefer_tp_within_node=True,
261+ check_memory=None, # Use config value
262+ )
263+ else:
264+ logger.warning("Auto-parallelization (training): num_npus is 0 or could not be determined. Skipping auto-parallelization.")
265+ except Exception as e:
266+ logger.warning(f"Auto-parallelization (training): Failed with error: {e}. Continuing without auto-parallelization.")
267+ logger.exception(e)
268+ 
269+ # Auto-parallelization for inference (if enabled)
270+ # Wrap in try-except to prevent any side effects from breaking dataset loading
271+ try:
272+ if getattr(generate_config, "auto_parallel", False):
273+ from mindspeed_rl.workers.scheduler.launcher import auto_parallelize_inference, get_npu_deployment
274+ from mindspeed_rl.workers.actor_hybrid_worker import ActorHybridWorker
275+
276+ # Get num_npus for inference (same as training)
277+ num_npus = get_npu_deployment(rl_config, ActorHybridWorker)
278+ if num_npus > 0:
279+ logger.info(f"Auto-parallelization (inference): Enabled. num_npus={num_npus}")
280+ check_memory = getattr(generate_config, "auto_parallel_check_memory", True)
281+
282+ # Get tokenizer path for memory checking
283+ tokenizer_name_or_path = actor_config.tokenizer_name_or_path
284+
285+ tp, cp, ep = auto_parallelize_inference(
286+ generate_config=generate_config,
287+ num_npus=num_npus,
288+ megatron_config=actor_config,
289+ check_memory=check_memory if check_memory else False, # Only enable if explicitly True
290+ tokenizer_name_or_path=tokenizer_name_or_path if check_memory else None,
291+ prefer_tp_within_node=True,
292+ )
293+ logger.info(
294+ f"Auto-parallelization (inference): Updated generate_config - "
295+ f"TP={tp}, CP={cp}, EP={ep}, DP={num_npus // tp}"
296+ )
297+ else:
298+ logger.warning("Auto-parallelization (inference): num_npus is 0 or could not be determined. Skipping auto-parallelization.")
299+ except Exception as e:
300+ logger.warning(f"Auto-parallelization (inference): Failed with error: {e}. Continuing without auto-parallelization.")
301+ logger.exception(e)
220 302 
221 validate_rl_args(actor_config, ref_config, reward_config, rl_config, generate_config)303 validate_rl_args(actor_config, ref_config, reward_config, rl_config, generate_config)
222 304 
@@ -227,6 +309,7 @@ def parse_training_config(config: Dict):
227 role="integrated"309 role="integrated"
228 ),310 ),
229 })311 })
312+ actor_config.max_prompt_length = rl_config.max_prompt_length
230 313 
231 msprobe_config = MsprobeConfig(314 msprobe_config = MsprobeConfig(
232 config.get("msprobe_config", {}),315 config.get("msprobe_config", {}),
@@ -253,15 +336,14 @@ def get_megatron_module():
253 from megatron.training import get_args336 from megatron.training import get_args
254 from megatron.core.pipeline_parallel import get_forward_backward_func337 from megatron.core.pipeline_parallel import get_forward_backward_func
255 from megatron.core import DistributedDataParallel as LocalDDP338 from megatron.core import DistributedDataParallel as LocalDDP
256- from megatron.core.transformer.module import Float16Module339+ from megatron.legacy.model import Float16Module
257 from megatron.training.training import get_model, unwrap_model340 from megatron.training.training import get_model, unwrap_model
258 from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig341 from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
259 from megatron.core.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy342 from megatron.core.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy
260 from megatron.training.training import setup_model_and_optimizer343 from megatron.training.training import setup_model_and_optimizer
261 from megatron.core.enums import ModelType344 from megatron.core.enums import ModelType
262 from megatron.core.distributed import finalize_model_grads345 from megatron.core.distributed import finalize_model_grads
263- from mindspeed.utils import set_position_ids346+ from mindspeed.utils import set_actual_seq_len, set_position_ids, get_actual_seq_len
264- from mindspeed.core.context_parallel.get_batch_utils import set_actual_seq_len, get_actual_seq_len
265 from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer347 from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
266 from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params348 from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params
267 349 
@@ -356,7 +438,7 @@ def rm_model_provider(pre_process, post_process):
356 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec438 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_local_spec
357 from megatron.core.transformer.spec_utils import import_module439 from megatron.core.transformer.spec_utils import import_module
358 from megatron.training.arguments import core_transformer_config_from_args440 from megatron.training.arguments import core_transformer_config_from_args
359- from mindspeed_rl.models.orm.orm_model import GPTRewardModel441+ from mindspeed_llm.tasks.posttrain.orm.orm_model import GPTRewardModel
360 args = get_args()442 args = get_args()
361 logger.info('building RM GPT model ...')443 logger.info('building RM GPT model ...')
362 # Experimental loading arguments from configs444 # Experimental loading arguments from configs
@@ -397,8 +479,6 @@ def initialize_megatron(
397 ignore_unknown_args=False,479 ignore_unknown_args=False,
398 allow_no_cuda=False,480 allow_no_cuda=False,
399 skip_mpu_initialization=False,481 skip_mpu_initialization=False,
400- get_embedding_ranks=None,
401- get_position_embedding_ranks=None,
402 config=None,482 config=None,
403):483):
404 """Set global variables, initialize distributed, and484 """Set global variables, initialize distributed, and
@@ -419,7 +499,8 @@ def initialize_megatron(
419 from mindspeed_llm.training.arguments import parse_args_decorator499 from mindspeed_llm.training.arguments import parse_args_decorator
420 import megatron500 import megatron
421 501 
422- args = megatron.training.arguments.parse_args()502+ parse_args = parse_args_decorator(megatron.training.arguments.parse_args)
503+ args = parse_args(extra_args_provider, ignore_unknown_args)
423 sys.argv = origin_sys_argv504 sys.argv = origin_sys_argv
424 505 
425 if not allow_no_cuda:506 if not allow_no_cuda:
@@ -447,7 +528,7 @@ def initialize_megatron(
447 from mindspeed.core.tensor_parallel.lcal_coc.user_config import initialize_coc_from_cfg528 from mindspeed.core.tensor_parallel.lcal_coc.user_config import initialize_coc_from_cfg
448 initialize_coc_from_cfg(args)529 initialize_coc_from_cfg(args)
449 530
450- if args.npu_deterministic:531+ if args.use_deter_comp:
451 seed_all(args.seed)532 seed_all(args.seed)
452 logger.info("deterministic computing is applied for npu.")533 logger.info("deterministic computing is applied for npu.")
453 534 
@@ -455,15 +536,12 @@ def initialize_megatron(
455 def finish_mpu_init():536 def finish_mpu_init():
456 args = get_args()537 args = get_args()
457 # Pytorch distributed.538 # Pytorch distributed.
458- _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks)539+ _initialize_distributed()
459 540 
460 # Random seeds for reproducibility.541 # Random seeds for reproducibility.
461 if args.rank == 0:542 if args.rank == 0:
462 logger.info("> setting random seeds to {} ...".format(args.seed))543 logger.info("> setting random seeds to {} ...".format(args.seed))
463 _set_random_seed(args.seed, args.data_parallel_random_init)544 _set_random_seed(args.seed, args.data_parallel_random_init)
464- if args.use_ascend_mc2:
465- from mindspeed.core.tensor_parallel.ascend_turbo.initialize import initialize_cfg_from_args
466- initialize_cfg_from_args(args)
467 545 
468 if skip_mpu_initialization:546 if skip_mpu_initialization:
469 return None547 return None
@@ -495,7 +573,7 @@ def initialize_megatron(
495 return None573 return None
496 574 
497 575 
498-def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks):576+def _initialize_distributed():
499 """Initialize torch.distributed and core model parallel."""577 """Initialize torch.distributed and core model parallel."""
500 from megatron.core import parallel_state578 from megatron.core import parallel_state
501 from megatron.training import get_args579 from megatron.training import get_args
@@ -542,20 +620,11 @@ def _initialize_distributed(get_embedding_ranks, get_position_embedding_ranks):
542 args.pipeline_model_parallel_size,620 args.pipeline_model_parallel_size,
543 args.virtual_pipeline_model_parallel_size,621 args.virtual_pipeline_model_parallel_size,
544 args.pipeline_model_parallel_split_rank,622 args.pipeline_model_parallel_split_rank,
545- pipeline_model_parallel_comm_backend=args.pipeline_model_parallel_comm_backend,
546 context_parallel_size=args.context_parallel_size,623 context_parallel_size=args.context_parallel_size,
547- hierarchical_context_parallel_sizes=args.hierarchical_context_parallel_sizes,
548 expert_model_parallel_size=args.expert_model_parallel_size,624 expert_model_parallel_size=args.expert_model_parallel_size,
549- num_distributed_optimizer_instances=args.num_distributed_optimizer_instances,
550- expert_tensor_parallel_size=args.expert_tensor_parallel_size,
551 distributed_timeout_minutes=args.distributed_timeout_minutes,625 distributed_timeout_minutes=args.distributed_timeout_minutes,
552 nccl_communicator_config_path=args.nccl_communicator_config_path,626 nccl_communicator_config_path=args.nccl_communicator_config_path,
553- order='tp-cp-ep-dp-pp' if not args.use_tp_pp_dp_mapping else 'tp-cp-ep-pp-dp',627+ order='tp-cp-ep-dp-pp' if not args.use_tp_pp_dp_mapping else 'tp-pp-dp',
554- encoder_tensor_model_parallel_size=args.encoder_tensor_model_parallel_size,
555- encoder_pipeline_model_parallel_size=args.encoder_pipeline_model_parallel_size,
556- get_embedding_ranks=get_embedding_ranks,
557- get_position_embedding_ranks=get_position_embedding_ranks,
558- create_gloo_process_groups=args.enable_gloo_process_groups,
559 )628 )
560 if args.rank == 0:629 if args.rank == 0:
561 logger.info(630 logger.info(
@@ -576,12 +645,7 @@ def main(config):
576 rl_config = RLConfig(config.get("rl_config"))645 rl_config = RLConfig(config.get("rl_config"))
577 with open(os.path.join(cur_file_dir, rl_config.runtime_env_path)) as file:646 with open(os.path.join(cur_file_dir, rl_config.runtime_env_path)) as file:
578 runtime_env = yaml.safe_load(file)647 runtime_env = yaml.safe_load(file)
579- generate_config = GenerateConfig(config.get("generate_config"))
580- enable_expert_parallel = getattr(generate_config, "enable_expert_parallel", None)
581- if enable_expert_parallel:
582- runtime_env["env_vars"]["VLLM_DP_SIZE"] = str(generate_config.infer_expert_parallel_size)
583 runtime_env["env_vars"]["IS_MULTIMODAL"] = str(rl_config.is_multimodal)648 runtime_env["env_vars"]["IS_MULTIMODAL"] = str(rl_config.is_multimodal)
584- runtime_env["env_vars"]["HCCL_BUFFSIZE"] = str(rl_config.hccl_buffersize)
585 logger.info(f"ray init with runtime_env: {runtime_env}")649 logger.info(f"ray init with runtime_env: {runtime_env}")
586 ray.init(runtime_env=runtime_env)650 ray.init(runtime_env=runtime_env)
587 651 
@@ -590,3 +654,4 @@ def main(config):
590 654 
591if __name__ == '__main__':655if __name__ == '__main__':
592 main()656 main()
657+ 
@@ -0,0 +1,135 @@
1+defaults:
2+ - model:
3+ - qwen25_3b
4+ 
5+megatron_training:
6+ model: qwen25_3b
7+ use_fused_rmsnorm: true
8+ use_mcore_models: true
9+ sequence_parallel: true
10+ use_flash_attn: true
11+ no_masked_softmax_fusion: true
12+ attention_softmax_in_fp32: true
13+ no_gradient_accumulation_fusion: true
14+ use_fused_swiglu: true
15+ use_fused_rotary_pos_emb: true
16+ bf16: true
17+ use_distributed_optimizer: true
18+ tokenizer_type: PretrainedFromHF
19+ tokenizer_name_or_path: ./model_from_hf/qwen25-3b-instruct
20+ global_batch_size: 128
21+ seq_length: 2048
22+ save_interval: 200
23+ train_iters: 20000
24+ stage: ray_ppo
25+ attention_dropout: 0.0
26+ init_method_std: 0.01
27+ hidden_dropout: 0.0
28+ distributed_backend: nccl
29+ no_shared_storage: true
30+ variable_seq_lengths: true
31+ dataset_additional_keys: ['labels']
32+ data_path: ./dataset/dapo-math-17k/data
33+ split: 100,0,0
34+ no_shuffle: false
35+ full_shuffle_instruction_dataset: false
36+ seed: 1
37+ 
38+actor_config:
39+ model: qwen25_3b
40+ micro_batch_size: 4
41+ tensor_model_parallel_size: 2
42+ pipeline_model_parallel_size: 4
43+ lr: 1e-6
44+ lr_decay_style: constant
45+ min_lr: 0
46+ weight_decay: 0.01
47+ lr_warmup_fraction: 0.0
48+ clip_grad: 1.0
49+ adam_beta1: 0.9
50+ adam_beta2: 0.95
51+ finetune: true
52+ load: ./model_weights/qwen25-3b-mcore/
53+ save: ./ckpt
54+ no_load_optim: true
55+ no_load_rng: true
56+ 
57+critic_config:
58+ model: qwen25_3b
59+ tensor_model_parallel_size: 2
60+ pipeline_model_parallel_size: 4
61+ micro_batch_size: 4
62+ lr: 5e-6
63+ lr_decay_style: constant
64+ min_lr: 0
65+ weight_decay: 0.01
66+ lr_warmup_fraction: 0.0
67+ clip_grad: 1.0
68+ adam_beta1: 0.9
69+ adam_beta2: 0.95
70+ finetune: true
71+ no_load_optim: True
72+ no_load_rng: True
73+ load: ./model_weights/qwen25-3b-reward-mcore/
74+ save: ./ckpt
75+ 
76+rl_config:
77+ guarantee_order: true
78+ use_integrated_worker: true
79+ blocking: true
80+ actor_forward_micro_batch_size: 4
81+ ref_forward_micro_batch_size: 4
82+ gamma: 1.0
83+ lam: 1.0
84+ adv_estimator: gae
85+ kl_penalty: kl
86+ kl_ctrl_type: fixed
87+ init_kl_coef: 0.0
88+ use_kl_in_reward: False
89+ mini_batch_size: 128
90+ max_prompt_length: 2048
91+ epochs: 1
92+ clip_ratio: 0.2
93+ cliprange_value: 0.5
94+ entropy_coeff: 0
95+ shuffle_mini_batch: false
96+ n_samples_per_prompt: 8
97+ rule_reward: true
98+ verifier_function: ["acc_for_ppo"]
99+ verifier_weight: [1.0]
100+ num_cpus_for_local_task: 1.0
101+ use_tensorboard: false
102+ actor_resource:
103+ num_npus: 8
104+ critic_resource:
105+ num_npus: 8
106+ 
107+generate_config:
108+ enforce_eager: True
109+ trust_remote_code: true
110+ offload_train_optimizer: true
111+ offload_train_grad: true
112+ offload_train_param: true
113+ 
114+ # 推理时的并行配置
115+ infer_tensor_parallel_size: 2
116+ infer_pipeline_parallel_size: 1
117+ infer_expert_parallel_size: 1
118+ 
119+ # vllm 模型相关设置
120+ max_num_seqs: 1024
121+ max_model_len: 4096
122+ max_num_batched_tokens: 8192
123+ dtype: "bfloat16"
124+ gpu_memory_utilization: 0.6
125+ 
126+ # 采样配置
127+ sampling_config:
128+ seed: 1
129+ logprobs: 1
130+ max_tokens: 2048
131+ top_p: 1
132+ top_k: -1
133+ min_p: 0.0
134+ temperature: 1.0
135+ detokenize: false