Use a Generator object to specify the seed for the random embeddings.

[oscarkey]

Currently we set the global pytorch random seed. When multi-device inference is enabled, this can result in a race condition because the global seed is shared between threads. Use a Generator instead to avoid this.

Testing:

• Save the sampled embeddings from the old and the new versions and check they are equal.
• Existing consistency checks, plus the additional two I have added to verify multi-device inference.

[gemini-code-assist]

Code Review

This pull request effectively addresses a potential race condition in multi-device inference by replacing the global PyTorch random seed with a thread-safe torch.Generator object. The changes are well-implemented and the accompanying tests for multi-device consistency are a great addition. I've found one minor issue: a leftover debugging statement that should be removed.

[oscarkey]

hey @LeoGrin , I updated the implementation a bit to disable the seed during onnx export. This fixes the tests, and adds a warning that the seed won't be set correctly. What do you think?

[LeoGrin]

LGTM! Just one question to check I'm understanding things correctly: this should give use the exact same random output as before the change right? (I think yes because the consistency tests are passing?)

[oscarkey]

yes, and I saved + compared the actual embedding tensors before and after the change as well. But I'll do that one more time before merging just in case haha

[oscarkey]

I did the following on macos mps, cpu, cuda:

• checked that the value of the embeddings here was equal before and after the change, when performing classification on the iris dataset
• used the script below

from contextlib import contextmanager
from typing import Generator
import torch

device = torch.device("cuda")


@contextmanager
def isolate_torch_rng(seed: int, device: torch.device) -> Generator[None, None, None]:
    torch_rng_state = torch.get_rng_state()
    if torch.cuda.is_available():
        torch_cuda_rng_state = torch.cuda.get_rng_state(device=device)
    torch.manual_seed(seed)
    try:
        yield
    finally:
        torch.set_rng_state(torch_rng_state)
        if torch.cuda.is_available():
            torch.cuda.set_rng_state(torch_cuda_rng_state, device=device)


def global_rng():
    with isolate_torch_rng(seed=0, device=device):
        a = torch.randn(10_000, 10, device=device)
    with isolate_torch_rng(seed=0, device=device):
        b = torch.randn(10_000, 10, device=device)
    with isolate_torch_rng(seed=1, device=device):
        c = torch.randn(10_000, 10, device=device)
    return a, b, c


def generator_rng():
    a = torch.randn(
        10_000,
        10,
        generator=torch.Generator(device=device).manual_seed(0),
        device=device,
    )
    b = torch.randn(
        10_000,
        10,
        generator=torch.Generator(device=device).manual_seed(0),
        device=device,
    )
    c = torch.randn(
        10_000,
        10,
        generator=torch.Generator(device=device).manual_seed(1),
        device=device,
    )
    return a, b, c


global_a, global_b, global_c = global_rng()
generator_a, generator_b, generator_c = generator_rng()

print(torch.allclose(global_a, generator_a))
print(torch.allclose(global_b, generator_b))
print(torch.allclose(global_c, generator_c))

It all passed, so seems safe to merge.