refactor resample

Author: hjh0119

swift/megatron/trainers/gkd_trainer.py

@@ -58,16 +58,11 @@ def __init__(self, args: MegatronArguments, template, **kwargs):
         self.truncation_strategy = args.truncation_strategy
         self.max_completion_length = args.max_completion_length
 
-        # Resample iterator will be initialized lazily
         self.resample_data_iterator = None
-        self._train_dataset = None
 
     def train(self, train_dataset, val_dataset):
-        """Override train to initialize resample iterator for truncation_strategy='delete'."""
-        # Store dataset provider for lazy resample iterator initialization
         if self.truncation_strategy == 'delete':
-            self._train_dataset = train_dataset
-
+            self.resample_data_iterator = self._init_resample_data_iterator(train_dataset)
         super().train(train_dataset, val_dataset)
 
     def prepare_model(self):
@@ -184,13 +179,16 @@ def _determine_data_source(self) -> DataSource:
             # Mode 3: Off-Policy learning, use dataset responses
             return DataSource.DATASET
 
-    def _init_resample_data_iterator(self):
-        """Initialize an independent data iterator for dynamic resampling (lazy initialization).
+    def _init_resample_data_iterator(self, train_dataset):
+        """Initialize an independent data iterator for resampling.
 
         Uses a different seed (args.seed + 1) to avoid overlapping with training samples.
 
+        Args:
+            train_dataset: The training dataset to create the resample iterator from.
+
         Returns:
-            train_data_iterator: Independent data iterator with different random seed
+            The resample data iterator (first element of the iterator tuple).
         """
         args = self.args
         resample_seed = getattr(args, 'seed', 42) + 1
@@ -200,8 +198,7 @@ def _init_resample_data_iterator(self):
                 args.data_parallel_random_init,
                 args.te_rng_tracker,
             )
-            resample_data_iterator = self._prepare_data_iterator(self._train_dataset, use_origin_cyclic=True)
-            self._train_dataset = None
+            resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0]
         finally:
             set_random_seed(
                 args.seed,
@@ -225,10 +222,6 @@ def resample_encode_failed_inputs(self, inputs: List[Dict], max_resample_rounds:
         valid_samples = []
         pending_samples = list(inputs)
 
-        # Lazy initialization of resample_data_iterator
-        if self.resample_data_iterator is None:
-            self.resample_data_iterator = self._init_resample_data_iterator()[0]
-
         for _ in range(max_resample_rounds + 1):
             still_needed = required_count - len(valid_samples)
             if still_needed <= 0:
@@ -283,7 +276,7 @@ def _replace_data_iterator(self, data_iterator):
             raw_batch = next(data_iterator)
 
             # Resample for encoding failed data when truncation_strategy is 'delete'
-            if self.truncation_strategy == 'delete' and self._train_dataset is not None:
+            if self.truncation_strategy == 'delete' and self.resample_data_iterator is not None:
                 raw_batch = self.resample_encode_failed_inputs(raw_batch)
 
             global_batch.extend(raw_batch)

swift/megatron/trainers/grpo_trainer.py

@@ -54,13 +54,11 @@ def __init__(self, args: MegatronRLHFArguments, template: Template, **kwargs):
         self._init_rollout_engine()
         self._prepare_rewards()
         self._prepare_scheduler()
-        self._train_dataset = None
+        self.resample_data_iterator = None
 
     def train(self, train_dataset, val_dataset):
-        # Store dataset provider for lazy resample iterator initialization
-        # Used by both dynamic_sample and truncation_strategy='delete'
         if self.dynamic_sample or self.truncation_strategy == 'delete':
-            self._train_dataset = train_dataset
+            self.resample_data_iterator = self._init_resample_data_iterator(train_dataset)
         super().train(train_dataset, val_dataset)
 
     def _init_grpo_params(self):
@@ -215,31 +213,28 @@ def _prepare_scheduler(self):
                 assert isinstance(args.multi_turn_scheduler, MultiTurnScheduler)
                 self.multi_turn_scheduler: MultiTurnScheduler = args.multi_turn_scheduler
 
-    def _init_resample_data_iterator(self):
-        """Initialize an independent data iterator for dynamic resampling (lazy initialization).
+    def _init_resample_data_iterator(self, train_dataset):
+        """Initialize an independent data iterator for resampling.
 
         Uses a different seed (args.seed + 1) to avoid overlapping with training samples.
 
+        Args:
+            train_dataset: The training dataset to create the resample iterator from.
+
         Returns:
-            train_data_iterator: Independent data iterator with different random seed
+            The resample data iterator (first element of the iterator tuple).
         """
         args = self.args
-        # Use different seed for resample iterator (offset by 1 to avoid overlap)
         resample_seed = getattr(args, 'seed', 42) + 1
         try:
-            # Set new seed for resample iterator creation
             set_random_seed(
                 resample_seed,
                 args.data_parallel_random_init,
                 args.te_rng_tracker,
             )
-
-            # Build data iterators with new seed
             # TODO: VPP (Virtual Pipeline Parallelism)
-            resample_data_iterator = self._prepare_data_iterator(self._train_dataset, use_origin_cyclic=True)
-            self._train_dataset = None
+            resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0]
         finally:
-            # Restore original random states to avoid affecting training
             set_random_seed(
                 args.seed,
                 args.data_parallel_random_init,
@@ -909,9 +904,6 @@ def _dynamic_sampling(self, rollout_batch: DataType,
             if len(valid_samples) >= self.generation_batch_size:
                 break
 
-            # Lazy initialization of resample_data_iterator
-            if not hasattr(self, 'resample_data_iterator') or self.resample_data_iterator is None:
-                self.resample_data_iterator = self._init_resample_data_iterator()[0]
             num_iters_per_step = self.get_num_iters_per_step()
             next_rollout_prompt_batch = []
             for _ in range(num_iters_per_step):
@@ -1562,11 +1554,7 @@ def resample_encode_failed_inputs(self, inputs: DataType, max_resample_rounds: i
         required_count = len(inputs)
         valid_samples = []
 
-        # Buffer for samples waiting to be validated
         pending_samples = list(inputs)
-        # Lazy initialization of resample_data_iterator
-        if not hasattr(self, 'resample_data_iterator') or self.resample_data_iterator is None:
-            self.resample_data_iterator = self._init_resample_data_iterator()[0]
         for _ in range(max_resample_rounds + 1):
             # Calculate how many more samples we need
             still_needed = required_count - len(valid_samples)