fix: fix seed with multiple ranks

Fix deepmodeling#4440.

deepmd/pd/entrypoints/main.py

@@ -123,7 +123,7 @@ def prepare_trainer_input_single(
 
         # validation and training data
         # avoid the same batch sequence among devices
-        rank_seed = (seed + rank) % (2**32) if seed is not None else None
+        rank_seed = [rank, seed % (2**32)] if seed is not None else None
         validation_data_single = (
             DpLoaderSet(
                 validation_systems,

deepmd/pt/entrypoints/main.py

@@ -138,7 +138,7 @@ def prepare_trainer_input_single(
 
         # validation and training data
         # avoid the same batch sequence among devices
-        rank_seed = (seed + rank) % (2**32) if seed is not None else None
+        rank_seed = [rank, seed % (2**32)] if seed is not None else None
         validation_data_single = (
             DpLoaderSet(
                 validation_systems,

deepmd/pt/utils/dataloader.py

@@ -50,8 +50,10 @@
 
 
 def setup_seed(seed) -> None:
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
+    if isinstance(seed, (list, tuple)):
+        mixed_seed = mix_entropy(seed)
+    torch.manual_seed(mixed_seed)
+    torch.cuda.manual_seed_all(mixed_seed)
     torch.backends.cudnn.deterministic = True
     dp_random.seed(seed)
 

deepmd/tf/entrypoints/train.py

@@ -220,9 +220,10 @@ def _do_work(
     seed = jdata["training"].get("seed", None)
     if seed is not None:
         # avoid the same batch sequence among workers
-        seed += run_opt.my_rank
         seed = seed % (2**32)
-    dp_random.seed(seed)
+        dp_random.seed([run_opt.my_rank, seed])
+    else:
+        dp_random.seed(seed)
 
     # setup data modifier
     modifier = get_modifier(jdata["model"].get("modifier", None))

deepmd/utils/random.py

@@ -56,7 +56,7 @@ def random(size=None):
     return _RANDOM_GENERATOR.random_sample(size)
 
 
-def seed(val: Optional[int] = None) -> None:
+def seed(val: Optional[Union[int, list[int]]] = None) -> None:
     """Seed the generator.
 
     Parameters