attempt to address lucidrains/audiolm-pytorch#279 again

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.18.6"
+version = "1.18.7"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

vector_quantize_pytorch/residual_fsq.py

@@ -1,6 +1,6 @@
 import random
 from math import log2
-from functools import partial, cache
+from functools import partial
 
 from typing import List
 
@@ -33,10 +33,17 @@ def round_up_multiple(num, mult):
 
 # distributed helpers
 
-@cache
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
+def get_maybe_sync_seed(max_size = 10_000):
+    rand_int = torch.randint(0, max_size, ())
+
+    if is_distributed():
+        dist.all_reduce(rand_int)
+
+    return rand_int.item()
+
 # main class
 
 class ResidualFSQ(Module):
@@ -175,18 +182,12 @@ def forward(
 
         if should_quantize_dropout:
 
-            if exists(rand_quantize_dropout_fixed_seed):
-                # seed is manually passed in
-                rand = random.Random(rand_quantize_dropout_fixed_seed)
+            # check if seed is manually passed in
 
-            elif is_distributed():
-                # in distributed environment, synchronize a random seed value if not given
-                t = torch.tensor(random.randrange(10_000), device = device)
-                dropout_seed = dist.all_reduce(t).item()
-                rand = random.Random(dropout_seed)
+            if not exists(rand_quantize_dropout_fixed_seed):
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
 
-            else:
-                rand = random
+            rand = random.Random(rand_quantize_dropout_fixed_seed)
 
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 
@@ -304,7 +305,7 @@ def forward(
 
         forward_kwargs = dict(
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = random.randint(0, int(1e7))
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
         )
 
         # invoke residual vq on each group

vector_quantize_pytorch/residual_lfq.py

@@ -28,10 +28,17 @@ def round_up_multiple(num, mult):
 
 # distributed helpers
 
-@cache
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
+def get_maybe_sync_seed(max_size = 10_000):
+    rand_int = torch.randint(0, max_size, ())
+
+    if is_distributed():
+        dist.all_reduce(rand_int)
+
+    return rand_int.item()
+
 # main class
 
 class ResidualLFQ(Module):
@@ -152,18 +159,12 @@ def forward(
 
         if should_quantize_dropout:
 
-            if exists(rand_quantize_dropout_fixed_seed):
-                # seed is manually passed in
-                rand = random.Random(rand_quantize_dropout_fixed_seed)
+            # check if seed is manually passed in
 
-            elif is_distributed():
-                # in distributed environment, synchronize a random seed value if not given
-                t = torch.tensor(random.randrange(10_000), device = device)
-                dropout_seed = dist.all_reduce(t).item()
-                rand = random.Random(dropout_seed)
+            if not exists(rand_quantize_dropout_fixed_seed):
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
 
-            else:
-                rand = random
+            rand = random.Random(rand_quantize_dropout_fixed_seed)
 
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 
@@ -271,7 +272,7 @@ def forward(
         forward_kwargs = dict(
             mask = mask,
             return_all_codes = return_all_codes,
-            rand_quantize_dropout_fixed_seed = random.randint(0, int(1e7))
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
         )
 
         # invoke residual vq on each group

vector_quantize_pytorch/residual_vq.py

@@ -33,10 +33,17 @@ def round_up_multiple(num, mult):
 
 # distributed helpers
 
-@cache
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
+def get_maybe_sync_seed(max_size = 10_000):
+    rand_int = torch.randint(0, max_size, ())
+
+    if is_distributed():
+        dist.all_reduce(rand_int)
+
+    return rand_int.item()
+
 # the mlp for generating the neural implicit codebook
 # from Huijben et al. https://arxiv.org/abs/2401.14732
 
@@ -286,18 +293,13 @@ def forward(
 
         if should_quantize_dropout:
 
-            if exists(rand_quantize_dropout_fixed_seed):
-                # seed is manually passed in
-                rand = random.Random(rand_quantize_dropout_fixed_seed)
+            # check if seed is manually passed in
+
+            if not exists(rand_quantize_dropout_fixed_seed):
+                rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
 
-            elif is_distributed():
-                # in distributed environment, synchronize a random seed value if not given
-                t = torch.tensor(random.randrange(10_000), device = device)
-                dropout_seed = dist.all_reduce(t).item()
-                rand = random.Random(dropout_seed)
+            rand = random.Random(rand_quantize_dropout_fixed_seed)
 
-            else:
-                rand = random
 
             rand_quantize_dropout_index = rand.randrange(self.quantize_dropout_cutoff_index, num_quant)
 
@@ -466,7 +468,7 @@ def forward(
             sample_codebook_temp = sample_codebook_temp,
             mask = mask,
             freeze_codebook = freeze_codebook,
-            rand_quantize_dropout_fixed_seed = random.randint(0, int(1e7))
+            rand_quantize_dropout_fixed_seed = get_maybe_sync_seed()
         )
 
         # invoke residual vq on each group