diff --git a/modelopt/torch/utils/random.py b/modelopt/torch/utils/random.py index b7ab30f2392..0c8249ff980 100644 --- a/modelopt/torch/utils/random.py +++ b/modelopt/torch/utils/random.py @@ -44,9 +44,10 @@ def _get_generator(seed: int | None = None) -> _random.Random: delattr(_get_generator, "generator") if not hasattr(_get_generator, "generator"): # synchronizing random seed and initialize generator - seed = dist.broadcast(seed or _random.getrandbits(64)) + is_manual = seed is not None + seed = dist.broadcast(seed if seed is not None else _random.getrandbits(64)) _get_generator.generator = _random.Random(seed) # type: ignore[attr-defined] - _get_generator.is_manual = seed is not None # type: ignore[attr-defined] + _get_generator.is_manual = is_manual # type: ignore[attr-defined] _get_generator.is_synced = dist.size() > 1 # type: ignore[attr-defined] return _get_generator.generator # type: ignore[attr-defined] @@ -157,7 +158,12 @@ def _deterministic_seed(): Resets the random state to prior upon exit. """ old_random_generator = _get_generator() + old_is_manual = getattr(_get_generator, "is_manual", False) + old_is_synced = getattr(_get_generator, "is_synced", False) _set_deterministic_seed(1024) - yield - delattr(_get_generator, "generator") - setattr(_get_generator, "generator", old_random_generator) + try: + yield + finally: + _get_generator.generator = old_random_generator # type: ignore[attr-defined] + _get_generator.is_manual = old_is_manual # type: ignore[attr-defined] + _get_generator.is_synced = old_is_synced # type: ignore[attr-defined] diff --git a/tests/unit/torch/utils/test_random_seed.py b/tests/unit/torch/utils/test_random_seed.py new file mode 100644 index 00000000000..240efd70158 --- /dev/null +++ b/tests/unit/torch/utils/test_random_seed.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Regression tests for seed handling in ``modelopt.torch.utils.random``.""" + +import random as stdlib_random + +import pytest + +from modelopt.torch.utils import random as mtu_random +from modelopt.torch.utils.random import _get_generator + + +def _clear_generator_cache(): + for attr in ("generator", "is_manual", "is_synced"): + if hasattr(_get_generator, attr): + delattr(_get_generator, attr) + + +@pytest.fixture(autouse=True) +def _fresh_generator_cache(): + """Clear the cached generator state on ``_get_generator`` around each test.""" + _clear_generator_cache() + yield + _clear_generator_cache() + + +def test_explicit_seed_zero_is_honored(): + """An explicit seed of 0 must not be silently replaced by random bits.""" + _get_generator(seed=0) + stream_0a = [mtu_random.random() for _ in range(5)] + + _get_generator(seed=0) + stream_0b = [mtu_random.random() for _ in range(5)] + + _get_generator(seed=1) + stream_1 = [mtu_random.random() for _ in range(5)] + + reference = stdlib_random.Random(0) + assert stream_0a == stream_0b + assert stream_0a == [reference.random() for _ in range(5)] + assert stream_0a != stream_1 + + +def test_is_manual_flag_reflects_original_argument(): + """``is_manual`` must be False for auto-seeding and True for a manual seed.""" + _get_generator() + assert _get_generator.is_manual is False + + _get_generator(seed=42) + assert _get_generator.is_manual is True + + +def test_deterministic_seed_restores_generator_on_exception(): + """``_deterministic_seed`` must restore the outer generator even if the body raises.""" + _get_generator(seed=7) + reference = stdlib_random.Random(7) + assert [mtu_random.random() for _ in range(3)] == [reference.random() for _ in range(3)] + + with pytest.raises(RuntimeError, match="boom"), mtu_random._deterministic_seed(): + mtu_random.random() + raise RuntimeError("boom") + + # The outer stream must continue exactly where it left off. + assert [mtu_random.random() for _ in range(3)] == [reference.random() for _ in range(3)]