From 94078196511538052dbf0d7c83bbd76ddb786aaf Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Sun, 16 Nov 2025 01:10:43 +0800 Subject: [PATCH] fix train_iters --- swift/megatron/trainers/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swift/megatron/trainers/base.py b/swift/megatron/trainers/base.py index 2b6d938cc4..999745acbf 100644 --- a/swift/megatron/trainers/base.py +++ b/swift/megatron/trainers/base.py @@ -87,7 +87,7 @@ def initialize_megatron(*_args, **kwargs): args = get_args() data_parallel_size = mpu.get_data_parallel_world_size() step_batch_size = args.micro_batch_size * data_parallel_size - num_generations = args.num_generations if hasattr(args, 'num_generations') else 1 + num_generations = args.num_generations if args.rlhf_type == 'grpo' else 1 if args.train_iters is None and args.max_epochs is not None: if hasattr(train_dataset, '__len__'): dataset_sample = len(train_dataset) // step_batch_size * step_batch_size