Fix FSDP RNG restore across DP resharding

Author: addsubmuldiv

swift/megatron/utils/megatron_fsdp_checkpoint.py

@@ -32,18 +32,42 @@ def build_rng_state(rng_state, data_parallel_random_init: bool = False, rng_key:
     return {rng_key: rng_state_list}
 
 
-def get_rng_load_key(checkpoint_dir: str) -> str:
+def get_rng_load_key(checkpoint_dir: str, data_parallel_random_init: bool = False) -> Optional[str]:
     metadata_keys = FileSystemReader(checkpoint_dir).read_metadata().state_dict_metadata
-    rank_key = f'global_rank_{torch.distributed.get_rank()}'
-    if any(key.startswith(f'rng_state.{rank_key}.') for key in metadata_keys):
-        return rank_key
+    global_rank_prefix = 'rng_state.global_rank_'
+    saved_global_ranks = set()
+    for key in metadata_keys:
+        if key.startswith(global_rank_prefix):
+            rank = key[len(global_rank_prefix):].split('.', 1)[0]
+            if rank.isdigit():
+                saved_global_ranks.add(int(rank))
+
+    if saved_global_ranks:
+        # Per-rank RNG is exact only when the saved and current world-rank sets match.
+        # A topology-changing reshard should still load model/optimizer state without RNG.
+        current_world_size = torch.distributed.get_world_size()
+        if saved_global_ranks != set(range(current_world_size)):
+            return None
+        current_rank = torch.distributed.get_rank()
+        rank_key = f'global_rank_{current_rank}'
+        return rank_key if current_rank in saved_global_ranks else None
 
     pp_rank = mpu.get_pipeline_model_parallel_rank()
     tp_rank = mpu.get_tensor_model_parallel_rank()
     legacy_key = f'({pp_rank}, {tp_rank})'
-    if any(key.startswith(f'rng_state.{legacy_key}.') for key in metadata_keys):
-        return legacy_key
-    raise RuntimeError(f'RNG state for global rank {torch.distributed.get_rank()} was not found in `{checkpoint_dir}`.')
+    legacy_prefix = f'rng_state.{legacy_key}.'
+    legacy_metadata_keys = [key for key in metadata_keys if key.startswith(legacy_prefix)]
+    if not legacy_metadata_keys:
+        return None
+    if data_parallel_random_init:
+        saved_dp_ranks = set()
+        for key in legacy_metadata_keys:
+            dp_rank = key[len(legacy_prefix):].split('.', 1)[0]
+            if dp_rank.isdigit():
+                saved_dp_ranks.add(int(dp_rank))
+        if saved_dp_ranks != set(range(mpu.get_data_parallel_world_size())):
+            return None
+    return legacy_key
 
 
 def select_rng_state(rng_state, rng_key: str, data_parallel_random_init: bool):

swift/megatron/utils/megatron_lm_utils.py

@@ -468,12 +468,25 @@ def load_mcore_checkpoint(args,
     if (ckpt_tp_pp == run_tp_pp and not finetune and not no_load_rng
             and not getattr(state_dict['args'], 'no_save_rng', False)):
         if fsdp_dtensor:
-            fsdp_rng_key = fsdp_checkpoint.get_rng_load_key(checkpoint_dir)
-        gen_sd_rng_state = _get_rng_state(
-            fsdp_dtensor=fsdp_dtensor,
-            data_parallel_random_init=args.data_parallel_random_init,
-            fsdp_rng_key=fsdp_rng_key,
-        )  # we can load the rng state
+            fsdp_rng_key = fsdp_checkpoint.get_rng_load_key(checkpoint_dir, args.data_parallel_random_init)
+            if fsdp_rng_key is None:
+                gen_sd_rng_state = None
+                logger.warning(
+                    f'Megatron-FSDP RNG state in `{checkpoint_dir}` is incompatible with the current distributed '
+                    f'topology/world size. Model, optimizer, and scheduler checkpoint loading will continue, but '
+                    f'exact RNG restore is skipped; resumed training is not guaranteed to be bitwise or stepwise '
+                    f'deterministic.')
+            else:
+                gen_sd_rng_state = _get_rng_state(
+                    fsdp_dtensor=True,
+                    data_parallel_random_init=args.data_parallel_random_init,
+                    fsdp_rng_key=fsdp_rng_key,
+                )
+        else:
+            gen_sd_rng_state = _get_rng_state(
+                fsdp_dtensor=False,
+                data_parallel_random_init=args.data_parallel_random_init,
+            )  # we can load the rng state
     else:
         gen_sd_rng_state = None
         if ckpt_tp_pp != run_tp_pp:
@@ -543,7 +556,7 @@ def load_mcore_checkpoint(args,
     elif (args.fp16 or args.bf16) and optimizer is not None:
         optimizer.reload_model_params()
 
-    if not finetune and not no_load_rng:
+    if not finetune and not no_load_rng and (not fsdp_dtensor or fsdp_rng_key is not None):
         if 'rng_state' in state_dict:
             if fsdp_dtensor:
                 rng_state = fsdp_checkpoint.select_rng_state(