diff --git a/swift/megatron/trainers/base.py b/swift/megatron/trainers/base.py index da6a53129d..4a4b9cbdc6 100644 --- a/swift/megatron/trainers/base.py +++ b/swift/megatron/trainers/base.py @@ -626,8 +626,12 @@ def copy_path(src_path: str, tgt_path: str): else: raise ValueError(f'Source path is neither a file nor a directory: {src_path}') - def _prepare_data_iterator(self, train_dataset, val_dataset=None, use_origin_cyclic: bool = False): - train_dataloader, val_dataloader = self._prepare_dataloader(train_dataset, val_dataset) + def _prepare_data_iterator(self, + train_dataset, + val_dataset=None, + use_origin_cyclic: bool = False, + seed: Optional[int] = None): + train_dataloader, val_dataloader = self._prepare_dataloader(train_dataset, val_dataset, seed=seed) train_data_iterator = iter(self.cyclic_iter(train_dataloader, use_origin_cyclic=use_origin_cyclic)) val_data_iterator = None if val_dataset is not None: @@ -1036,11 +1040,15 @@ def _aggregated_metrics(self, metrics, total_metrics): total_metrics[key] = torch.tensor([0.0, 0.0], dtype=torch.float32, device=torch.cuda.current_device()) total_metrics[key] += val - def _prepare_dataloader(self, train_dataset, val_dataset=None): + def _prepare_dataloader(self, train_dataset, val_dataset=None, seed: Optional[int] = None): args = self.args val_dataloader = None + generator = None + if seed is not None: + generator = torch.Generator() + generator.manual_seed(seed) if args.streaming: - train_dataloader = build_streaming_dataloader(args, train_dataset, self.data_collator) + train_dataloader = build_streaming_dataloader(args, train_dataset, self.data_collator, generator=generator) if val_dataset is not None: val_dataloader = build_streaming_dataloader(args, val_dataset, self.data_collator) return train_dataloader, val_dataloader @@ -1054,8 +1062,9 @@ def _prepare_dataloader(self, train_dataset, val_dataset=None): data_sharding=args.data_sharding, shuffle=args.train_dataloader_shuffle, group_by_length=args.group_by_length, + seed=seed or 0, ) - train_dataloader = self._create_dataloader(train_dataset, train_batch_sampler) + train_dataloader = self._create_dataloader(train_dataset, train_batch_sampler, generator=generator) if val_dataset is not None: val_batch_sampler = MegatronPretrainingSampler( total_samples=len(val_dataset), @@ -1067,7 +1076,7 @@ def _prepare_dataloader(self, train_dataset, val_dataset=None): val_dataloader = self._create_dataloader(val_dataset, val_batch_sampler) return train_dataloader, val_dataloader - def _create_dataloader(self, dataset, batch_sampler): + def _create_dataloader(self, dataset, batch_sampler, generator=None): args = self.args dataloader_kwargs = {} @@ -1082,6 +1091,7 @@ def _create_dataloader(self, dataset, batch_sampler): persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False, prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None, collate_fn=self.data_collator, + generator=generator, **dataloader_kwargs, ) return dataloader diff --git a/swift/megatron/trainers/batch_sampler.py b/swift/megatron/trainers/batch_sampler.py index dfbc788bbe..4e616ab970 100644 --- a/swift/megatron/trainers/batch_sampler.py +++ b/swift/megatron/trainers/batch_sampler.py @@ -75,6 +75,7 @@ def __init__( data_sharding, shuffle: bool = True, group_by_length: bool = False, + seed: int = 0, ): # Keep a copy of input params for later use. self.dataset = dataset @@ -93,6 +94,7 @@ def __init__( self.data_sharding = data_sharding self.shuffle = shuffle self.group_by_length = group_by_length + self.seed = seed self.lengths = self.dataset['lengths'] if group_by_length else None if self.lengths is not None: self.lengths = [max(length) if isinstance(length, list) else length for length in self.lengths] @@ -124,14 +126,14 @@ def __iter__(self): start_idx = self.data_parallel_rank * bucket_size g = torch.Generator() - g.manual_seed(self.epoch) + g.manual_seed(self.seed + self.epoch) random_idx = torch.randperm(bucket_size, generator=g).tolist() idx_range = [start_idx + x for x in random_idx[bucket_offset:]] else: full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size full_bucket_offset = current_epoch_samples g = torch.Generator() - g.manual_seed(self.epoch) + g.manual_seed(self.seed + self.epoch) if self.group_by_length: from transformers.trainer_pt_utils import get_length_grouped_indices idx_range_total = get_length_grouped_indices( diff --git a/swift/megatron/trainers/gkd_trainer.py b/swift/megatron/trainers/gkd_trainer.py index 1d558ed76f..3fc3f8de9f 100644 --- a/swift/megatron/trainers/gkd_trainer.py +++ b/swift/megatron/trainers/gkd_trainer.py @@ -5,7 +5,6 @@ import torch.nn.functional as F from contextlib import contextmanager from functools import partial -from mcore_bridge import set_random_seed from megatron.core import mpu from transformers.utils import ContextManagers from typing import Dict, List, Optional @@ -167,20 +166,7 @@ def _init_resample_data_iterator(self, train_dataset): """ args = self.args resample_seed = getattr(args, 'seed', 42) + 1 - try: - set_random_seed( - resample_seed, - args.data_parallel_random_init, - args.te_rng_tracker, - ) - resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0] - finally: - set_random_seed( - args.seed, - args.data_parallel_random_init, - args.te_rng_tracker, - ) - return resample_data_iterator + return self._prepare_data_iterator(train_dataset, use_origin_cyclic=True, seed=resample_seed)[0] def resample_encode_failed_inputs(self, inputs: List[Dict], max_resample_rounds: int = 10) -> List[Dict]: """Attempt to encode each input. If encoding fails, resample until we have enough valid samples. diff --git a/swift/megatron/trainers/grpo_trainer.py b/swift/megatron/trainers/grpo_trainer.py index 6532546261..56993eae59 100644 --- a/swift/megatron/trainers/grpo_trainer.py +++ b/swift/megatron/trainers/grpo_trainer.py @@ -5,7 +5,6 @@ from contextlib import contextmanager from copy import copy, deepcopy from functools import partial -from mcore_bridge import set_random_seed from megatron.core import mpu from typing import Any, Dict, List, Optional, Tuple, Union @@ -209,21 +208,8 @@ def _init_resample_data_iterator(self, train_dataset): """ args = self.args resample_seed = getattr(args, 'seed', 42) + 1 - try: - set_random_seed( - resample_seed, - args.data_parallel_random_init, - args.te_rng_tracker, - ) - # TODO: VPP (Virtual Pipeline Parallelism) - resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0] - finally: - set_random_seed( - args.seed, - args.data_parallel_random_init, - args.te_rng_tracker, - ) - return resample_data_iterator + # TODO: VPP (Virtual Pipeline Parallelism) + return self._prepare_data_iterator(train_dataset, use_origin_cyclic=True, seed=resample_seed)[0] def _build_rollout_buffer(self, data_iterator): num_gen_steps = self.steps_per_generation if self.unwrapped_models[0].training else 1 diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index cd7f2b9bee..1987632268 100644 --- a/swift/megatron/trainers/utils.py +++ b/swift/megatron/trainers/utils.py @@ -334,7 +334,7 @@ def group(self): return mpu.get_data_parallel_group() -def build_streaming_dataloader(args, dataset, collate_fn): +def build_streaming_dataloader(args, dataset, collate_fn, generator=None): dataloader_kwargs = {} mp_context = getattr(args, 'dataloader_multiprocessing_context', None) if mp_context is not None and args.dataloader_num_workers > 0: @@ -347,6 +347,7 @@ def build_streaming_dataloader(args, dataset, collate_fn): batch_size=args.micro_batch_size, prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None, persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False, + generator=generator, **dataloader_kwargs, ) return MegatronDataLoaderDispatcher(base_dataloader) diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py index 8d641bb02c..c7668d63e9 100644 --- a/swift/megatron/utils/megatron_lm_utils.py +++ b/swift/megatron/utils/megatron_lm_utils.py @@ -314,7 +314,7 @@ def save_mcore_checkpoint( _get_rng_state( fsdp_dtensor=fsdp_dtensor, data_parallel_random_init=args.data_parallel_random_init, - ) if models else None) + ) if not args.no_save_rng else None) checkpoint_dir = os.path.join(output_dir, f'iter_{iteration:07d}') sharded_sd_metadata = get_sharded_sd_metadata(args) os.makedirs(checkpoint_dir, exist_ok=True) @@ -333,7 +333,7 @@ def save_mcore_checkpoint( _filter_adapter_state_dict(state_dict, peft_format) kwargs = {'content_metadata': sharded_sd_metadata} async_save = args.async_save - if not models: # save GPU memory + if not models and rng_state is None: # save GPU memory when only common state remains assert 'optimizer' not in state_dict async_save = False common_path = os.path.join(checkpoint_dir, 'common.pt') diff --git a/tests/megatron/test_checkpoint_symlink.py b/tests/megatron/test_checkpoint_symlink.py index 8b70225d30..97465c1278 100644 --- a/tests/megatron/test_checkpoint_symlink.py +++ b/tests/megatron/test_checkpoint_symlink.py @@ -53,7 +53,8 @@ def _save(self, async_save: bool): checkpoint_dir = os.path.join(root, 'checkpoint-2') os.makedirs(checkpoint_dir) request = FakeAsyncRequest() - args = types.SimpleNamespace(output_dir=root, async_save=async_save) + args = types.SimpleNamespace( + output_dir=root, async_save=async_save, no_save_rng=False, data_parallel_random_init=False) dist_checkpointing = mock.MagicMock() dist_checkpointing.save.return_value = request if async_save else None