[DTensor] Decide / Document RNG semantics

[wconstab]

Following the work in #159933, DTensor will utilize a 'generator' passed as kwarg to random ops instead of using its own built-in generator.

This exposed a question of semantics, which I want to hear opinions on before implementing and documenting the behavior.

Question
If a user passes their own RNG to DTensor random ops, should the state-offset advancement be visible to the user on that RNG object?

My intuition is, yes: it is very surprising if the passed-in RNG object does not advance after being used.

However, implementing it this way introduces another discrepancy. DTensor ALREADY does this deviation, on the default RNG. Currently, DTensor keeps an alternate copy of the original default RNG, and only advances that copy.

Example: test_init_with_user_generator

Example: First, the test uses the passed-in RNG to sample twice in a row in a loop, and ensures that both loop iterations match the values from the default RNG which wsa seeded with the same value. This part already works fine after #159933 lands.

Then, it runs into trouble with trying to establish a sane UX if someone re-seeds the RNG, or if someone inspects the state of the passed-in RNG.

    def test_init_with_user_generator(self):
        device_mesh = self.build_device_mesh()
        torch.manual_seed(42)
        rng = torch.Generator(device="cuda").manual_seed(42)
        t1 = torch.distributed.tensor.empty(
            (2, 3), device_mesh=device_mesh, placements=[Shard(0)]
        )
        t2 = torch.distributed.tensor.empty(
            (2, 3), device_mesh=device_mesh, placements=[Shard(0)]
        )
        for i in range(2):
            print(f"{i=}")
            # run a second time, to make sure that `rng`'s offset-state is advancing on the second usage
            torch.nn.init.uniform_(t1, 0.0, 1.0)
            torch.nn.init.uniform_(t2, 0.0, 1.0, rng)
            self.assertEqual(t1.full_tensor(), t2.full_tensor(), f"Failed at {i=}")

        # Problem area:
        # (1)
        # torch.distributed.tensor._random._rng_tracker._manual_seed(55)
        # torch.manual_seed(55)
        # (2)
        # rng.manual_seed(55)
        torch.nn.init.uniform_(t1, 0.0, 1.0)
        torch.nn.init.uniform_(t2, 0.0, 1.0, rng)
        self.assertEqual(t1.full_tensor(), t2.full_tensor())

Problems, in this example:

1. If someone wants to affect the state of DTensor's RNG after init, they have to know to call this hidden manual_seed API torch.distributed.tensor._random._rng_tracker._manual_seed instead of calling torch.manual_seed which worked in the first place.

2. if someone wanted to re-seed the passed-in RNG object, I made that ineffective in [DTensor] Support user-supplied Generator for random ops #159933 so that it is consistent with the behavior of the default RNG. I also made it so rng.get_state() is consistent with torch.cuda.get_rng_state() (e.g. both will show the initial state, not reflecting any activity from DTensor)

Proposal:
We could make DTensor update the global RNG state so that (1) does not require a hidden seed API, and DTensor always uses the latest RNG state from torch. We can then easily also make the passed-in RNG state get updated whenever its used.

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @d4l3k @pragupta @svekars @sekyondaMeta @AlannaBurke @tianyu-l @XilunWu

[XilunWu]

Do I understand that the proposal means:

1. DTensor RNG no longer maintains a copy of RNG states. Instead it always uses torch RNG state.
2. When entering the random ops context, DTensor reads the current torch RNG state then modify the philox offset. (DTensor assume the users guarantee that the torch RNG state is identical across ranks)
3. When exiting the random ops context, DTensor advances the RNG state in a way that after executing the op, the RNG state remains identical across ranks.

In summary, users now are responsible for making sure the RNG passed to DTensor random ops is consistent across ranks, and if so DTensor will guarantee that they remain consistent after executing random ops. And this applies to the default CUDA RNG and user created RNG.

[lucasb-eyer]

The way I understand the proposal, from a user perspective, it would match my intuition/expectation and SGTM.

[tianyu-l]

When entering the random ops context, DTensor reads the current torch RNG state then modify the philox offset. (DTensor assume the users guarantee that the torch RNG state is identical across ranks)

I'm assuming this "modify the philox offset" only happens within op. E.g. if users perform random ops around

1. Replicate dtensor (requires same RNG state) ->
2. Sharded dtensor (requires different RNG state) ->
3. Replicate dtensor (requires same RNG state)

The 3rd Replicate dtensor still gets the same RNG state.

[wconstab]

I'm assuming this "modify the philox offset" only happens within op. E.g. if users perform random ops around

correct, after the op is complete, then the RNG would be synchronized again on all SPMD ranks, assuming it was synchronized to begin with.

And the thing I would document is that users are responsible for ensuring the RNG is synchronized across all SPMD ranks. (e.g. if they are doing pipeline parallelism, its fine to have different seeds on stage0 ranks and stage1 ranks. But all the ranks that are collectively using DTensor should have the same RNG state.)

[ngimel]

I'm a bit worried about "users now are responsible for making sure the RNG passed to DTensor random ops is consistent across ranks", it's easy for ranks to go out of sync. It's a reasonable requirement though, so as a compromise I'd propose having some debug mode that would check it.

[wconstab]

@ngimel - 2 comments

1. yes we could definitely do a debug mode
2. would you prefer actually that we keep the current behavior for the default rng, which is that dtensor makes an initial copy and is isolated from whatever happens on the default rng thereafter? Even if we do something different with user-passed RNG? I think it could be argued that user-passed rng is a special case.

[ngimel]

I'd rather the behavior of user-passed rng and default rng was the same.

[wanchaol]

My feeling is yes, we could switch to use the current default pytorch RNG state. The reason why we forked one was because dtensor broadcasts the rng to all ranks first to make sure everyone start the same, so using the default rng means all ranks would need to overwrite their current RNG to the broadcasted RNG, it felt risky in the beginning.

But given that user can pass their own RNG, and we would need to alter the user passed RNG with DTensor anyways (if we don't alter it then it does not seem to make sense), it make sense to make the behavior consistent to the global RNG state.

correct, after the op is complete, then the RNG would be synchronized again on all SPMD ranks, assuming it was synchronized to begin with.

To make sure we are aligned, the passed in RNG from user must be the same, and dtensor would advance the RNG in the same way, so after user provide an correct RNG, we would not need to synchornize inside/outside the op, because dtensor would make sure they advanced to the same state.

[tianyu-l]

To make sure we are aligned, the passed in RNG from user must be the same, and dtensor would advance the RNG in the same way, so after user provide an correct RNG, we would not need to synchornize inside/outside the op, because dtensor would make sure they advanced to the same state.

@wconstab
One more clarifying question regarding granularity of RNG state management:

Imagine there are two runs on same number of devices, with seed fixed in the beginning (e.g. both with seed=42) on all ranks

• first run is with plain tensor
• second run is with DTensor

Two runs call the exact same sequence of random ops. In your proposal, would the end states of RNG between two runs always the same?

[wconstab]

Well, I'll have to think about that.

For sure the ranks will all have the same rng value at the end of the dtensor run. Whether that value matched the non-dtensor run, I think depends if dtensor used up more offsets. Keep in mind dtensor already does not provide matching unit values for a sharded tensor vs a singe gpu run, because of how the rng state advances in fixed size chunks irrespective of how many values were used by a kernel.

[XilunWu]

I think depends if dtensor used up more offsets

Yes, DTensor advances the offset more than required. And as Will said, it's not relevant since we do not provide matching results.

[wconstab]

@ngimel I started working on a util for debugging. If the basic util looks good then I thought about making a either a contextmanager or a torchdispatch mode that runs it automatically after random ops. (thoughts on what would be most useful?)
#160283