From d15828e2541666ea3eac0896d346e051e6e930e1 Mon Sep 17 00:00:00 2001 From: arham766 Date: Sun, 5 Jul 2026 18:47:47 -0700 Subject: [PATCH] fix: honor explicit seed=0, compute is_manual before reassignment, exception-safe _deterministic_seed Three related defects in the shared RNG helper: - is_manual was computed AFTER seed was reassigned by dist.broadcast(seed or getrandbits(64)), so it was always True and the resync branch for auto-seeded generators created before process-group init was dead code - distributed ranks could sample divergently (e.g. AutoNAS RandomSearcher subnets). - The same falsy check silently replaced an explicit seed=0 with random bits. - _deterministic_seed yielded without try/finally, leaking the seed-1024 generator (and the is_manual/is_synced flags) if the body raised, e.g. a failing sample callback in nas algorithms; the restore now also covers both flags, matching the docstring's 'resets the random state to prior upon exit'. Regression tests live in a fix-scoped file; the fuller suite for this module is proposed separately in #1916. Part of the findings in #1902 (wave 2). Signed-off-by: arham766 --- modelopt/torch/utils/random.py | 16 +++-- tests/unit/torch/utils/test_random_seed.py | 77 ++++++++++++++++++++++ 2 files changed, 88 insertions(+), 5 deletions(-) create mode 100644 tests/unit/torch/utils/test_random_seed.py 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)]