From d605e02400c21ae1ce3746f962e24e298cddf7d8 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Sat, 29 Nov 2025 16:13:30 -0800 Subject: [PATCH 1/2] changed performance mode to deterministic mode -- which is off by default, so we can make more effective use of compute by default and only run deterministically when explicitly requested by the user --- src/gfn/env.py | 2 +- src/gfn/utils/common.py | 30 +++++++++++++++--------------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/gfn/env.py b/src/gfn/env.py index 646094f0..4b119322 100644 --- a/src/gfn/env.py +++ b/src/gfn/env.py @@ -271,7 +271,7 @@ def reset( assert not (random and sink) if random and seed is not None: - set_seed(seed, performance_mode=True) + set_seed(seed, deterministic_mode=False) # TODO: configurable? if isinstance(batch_shape, int): batch_shape = (batch_shape,) diff --git a/src/gfn/utils/common.py b/src/gfn/utils/common.py index 6f4601ee..bd66b8c6 100644 --- a/src/gfn/utils/common.py +++ b/src/gfn/utils/common.py @@ -145,12 +145,12 @@ def filter_kwargs_for_callable( # ----------------------------------------------------------------------------- -def set_seed(seed: int, performance_mode: bool = False) -> None: +def set_seed(seed: int, deterministic_mode: bool = False) -> None: """Used to control randomness for both single and distributed training. Args: seed: The seed to use for all random number generators - performance_mode: If True, disables deterministic behavior for better performance. + deterministic_mode: If True, uses deterministic behavior for better performance. In multi-GPU settings, this only affects cuDNN. In multi-CPU settings, this allows parallel processing in NumPy. """ @@ -186,19 +186,22 @@ def set_seed(seed: int, performance_mode: bool = False) -> None: # Set device-specific environment variables if torch.cuda.is_available(): - # For GPU training, we can use multiple threads for CPU operations - if performance_mode: - os.environ["OMP_NUM_THREADS"] = str(num_cpus) - os.environ["MKL_NUM_THREADS"] = str(num_cpus) - else: + if deterministic_mode: # For reproducibility in GPU training, we still want deterministic # CPU operations os.environ["OMP_NUM_THREADS"] = "1" os.environ["MKL_NUM_THREADS"] = "1" + else: + # For GPU training, we can use multiple threads for CPU operations + os.environ["OMP_NUM_THREADS"] = str(num_cpus) + os.environ["MKL_NUM_THREADS"] = str(num_cpus) else: # For CPU-only training, we need to be more careful with threading - if performance_mode: - + if deterministic_mode: + # For perfect reproducibility in CPU training, disable parallel processing + os.environ["OMP_NUM_THREADS"] = "1" + os.environ["MKL_NUM_THREADS"] = "1" + else: # Allow parallel processing but with controlled number of threads # Different backends might handle threading differently if backend in ["mpi", "ccl"]: @@ -211,10 +214,6 @@ def set_seed(seed: int, performance_mode: bool = False) -> None: os.environ["OMP_NUM_THREADS"] = str(num_threads) os.environ["MKL_NUM_THREADS"] = str(num_threads) - else: - # For perfect reproducibility in CPU training, disable parallel processing - os.environ["OMP_NUM_THREADS"] = "1" - os.environ["MKL_NUM_THREADS"] = "1" else: # Non-distributed training - use the global seed @@ -237,7 +236,7 @@ def set_seed(seed: int, performance_mode: bool = False) -> None: threading.current_thread()._seed = seed # These are only set when we care about reproducibility over performance - if not performance_mode: + if deterministic_mode: # GPU-specific settings if torch.cuda.is_available(): torch.backends.cudnn.deterministic = True @@ -332,7 +331,8 @@ def make_dataloader_seed_fns( def _worker_init_fn(worker_id: int) -> None: # pragma: no cover # Each worker gets a distinct seed in the same pattern used for ranks. - set_seed(base_seed + worker_id, performance_mode=False) + # TODO: Can this be false? + set_seed(base_seed + worker_id, deterministic_mode=True) gen = torch.Generator() gen.manual_seed(base_seed) From 7c22f79c080d05afc004ad8911b26dba2683cd73 Mon Sep 17 00:00:00 2001 From: Joseph Viviano Date: Wed, 10 Dec 2025 12:54:31 -0500 Subject: [PATCH 2/2] set_seed is no longer by default deterministic --- src/gfn/utils/common.py | 9 +++++++-- testing/test_estimators.py | 2 +- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/gfn/utils/common.py b/src/gfn/utils/common.py index bd66b8c6..5c7fd44d 100644 --- a/src/gfn/utils/common.py +++ b/src/gfn/utils/common.py @@ -313,9 +313,15 @@ def temporarily_set_seed(seed): def make_dataloader_seed_fns( base_seed: int, + deterministic_mode: bool = False, ) -> Tuple[Callable[[int], None], torch.Generator]: """Return `(worker_init_fn, generator)` for DataLoader reproducibility. + Args: + base_seed: The base seed to use for the DataLoader. + deterministic_mode: If True, uses deterministic behavior for better + reproducibility at the cost of performance. + Example ------- >>> w_init, g = make_dataloader_seed_fns(process_seed) @@ -331,8 +337,7 @@ def make_dataloader_seed_fns( def _worker_init_fn(worker_id: int) -> None: # pragma: no cover # Each worker gets a distinct seed in the same pattern used for ranks. - # TODO: Can this be false? - set_seed(base_seed + worker_id, deterministic_mode=True) + set_seed(base_seed + worker_id, deterministic_mode=deterministic_mode) gen = torch.Generator() gen.manual_seed(base_seed) diff --git a/testing/test_estimators.py b/testing/test_estimators.py index d0e99fc4..ca28c08a 100644 --- a/testing/test_estimators.py +++ b/testing/test_estimators.py @@ -476,7 +476,7 @@ def test_uniform_log_probs_method(): def test_mix_with_uniform_in_log_space(): """Test the _mix_with_uniform_in_log_space static method.""" batch_size, n_actions = 3, 4 - set_seed(123) + set_seed(123, deterministic_mode=True) # Create log-softmax values logits = torch.randn(batch_size, n_actions)