Skip to content

Commit 45c8463

Browse files
committed
add hash based mixing
1 parent 79ffe50 commit 45c8463

File tree

1 file changed

+11
-2
lines changed
  • src/lightning/fabric/utilities

1 file changed

+11
-2
lines changed

src/lightning/fabric/utilities/seed.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,10 +111,19 @@ def pl_worker_init_function(worker_id: int, rank: Optional[int] = None) -> None:
111111

112112

113113
def _generate_seed_sequence(base_seed: int, worker_id: int, global_rank: int, count: int) -> list[int]:
114-
"""Generates a sequence of seeds from a base seed, worker id and rank using the linear congruential generator (LCG)
115-
algorithm."""
114+
"""Generates a sequence of seeds from a base seed, worker id and rank using hash-based mixing followed by the
115+
linear congruential generator (LCG) algorithm."""
116116
# Combine base seed, worker id and rank into a unique 64-bit number
117117
combined_seed = (base_seed << 32) | (worker_id << 16) | global_rank
118+
119+
# Apply hash-based mixing (MurmurHash3 finalizer) to distribute bits uniformly
120+
# This ensures that small base seeds don't result in zeros in lower bits
121+
combined_seed ^= combined_seed >> 33
122+
combined_seed = (combined_seed * 0xFF51AFD7ED558CCD) & ((1 << 64) - 1)
123+
combined_seed ^= combined_seed >> 33
124+
combined_seed = (combined_seed * 0xC4CEB9FE1A85EC53) & ((1 << 64) - 1)
125+
combined_seed ^= combined_seed >> 33
126+
118127
seeds = []
119128
for _ in range(count):
120129
# x_(n+1) = (a * x_n + c) mod m. With c=1, m=2^64 and a is D. Knuth's constant

0 commit comments

Comments
 (0)