🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/139445

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 16503c4 with merge base failed to retrieve merge base, please contact dev infra:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Does "parallel" mode need help from torch.pipelining?

Does "parallel" mode need help from torch.pipelining?

not strictly. It only needs to know 'pp_rank' to set different seeds per pp-rank, and then do init as usual. However, once adding an API to pipelining, I think the only use case for 'parallel' would be to get back some performance if serial was too slow, and having it be a flag-flip seems nicer than having to edit your code to do init differently.

I don't have problem landing the proposed as experimental API. The description in the RFC indicates that a pipeline schedule can be of some help for serialized init and hence why attaching a util here. (Although, it could be still debatable whether pipeline schedule -- a "run" thing -- should get caught in model init, but anyway). The reason for suggesting making it experimental is so that we can be off the hook of maintenance for now.

That said, just thinking out loud:
Re "parallel" mode:
Per @wconstab's comment above, users can probably do it easily / cleanly in their code. So this does not motivate the proposed API originally.

Re "serialized" mode:
The proposed API has an assumption that user provide a single init function for all the stage models. How easy / general is that (e.g. would that need some amount of if-else in the init function? Anything else?)

Re customer base:
If there are only a small set of users that need serialized init, what's wrong with recommending checkpoint based init to them? It seems similar to a training recovery step or an initial load step in modeling tuning. And would it be more performant than serialized init since all ranks can then load from the checkpoint concurrently?

My point being, for the majority of users, should we not give them a feeling that they should/could change they are doing today? i.e. set different seeds, init the modules, hook it with pipeline schedule, run, etc.

I think we should run experiments at larger scale to see if serialized init is noticeably slow. If not, I would be happy to recommend it as the default way, and just consider 'parallel' mode as an escape hatch.

My point being, for the majority of users, should we not give them a feeling that they should/could change they are doing today? i.e. set different seeds, init the modules, hook it with pipeline schedule, run, etc.

Users can freely ignore this helper API and do their own seeding if they want.

The proposed API has an assumption that user provide a single init function for all the stage models.

Agreed that this is a bit weak in general. However, if users follow the practice of writing 'init_weights' per submodule then it is quite easy to have a generalized top-level initializer that delegates. That is what we have in titan.

In the titan code,

I was thinking maybe we can get rid of the 'for m in model_parts' loop here, define a helper 'parallelize_and_initialize(module)' that can be reused in PP and non PP branch, then just have the overall code look like this:

    def parallelize_and_initialize(m):
            # apply SPMD-style PT-D techniques
            models_parallelize_fns[model_name](m, world_mesh, parallel_dims, job_config)
            m.to_empty(device=init_device)
            m.init_weights(buffer_device=buffer_device)
            m.train()

    # apply parallelisms and initialization
    if parallel_dims.pp_enabled:
        # apply PT-D Pipeline Parallel
        pp_schedule, model_parts = models_pipelining_fns[model_name](
            model, pp_mesh, parallel_dims, job_config, device, model_config, loss_fn
        )

       schedule.seeded_module_init(parallelize_and_initialize)

    else:
        parallelize_and_initialize(model)

        model_parts = [model]

i'm not totally sure its good to bundle the parallelize and initialize in this way, but it might be ok?

Anyway.. I think there are 2 points
(a) should we offer serial init?
(b) if we are going to have serial init, what should the API be?

After going through this PR, I believe DTensor RNG won't have impact on this because DTensor RNG is only used on DTensor random ops (e.g. torch.distributed.tensor.randn) while PP RNG is directly calling torch.manual_seed. There's no interaction between them, unless we're using DTensor ops to initialize tensors.

code pointer: https://github.com/pytorch/pytorch/blob/main/torch/distributed/tensor/_random.py#L91-L99
You can see that DTensor RNG only sets up the seed for DTensor Random ops' python context.

Abandoning this as not widely useful. It works fine for PP only (1D parallel) but will not work as intended when combined with SPMD/DTensor parallelism. That's because DTensor/SPMD parallelism does not have a way to gaurantee same initialization behavior under different sharding configurations.