🚀[FEA]: Upstream distributed profiling context utilities from GLOBE AirFRANS example to DistributedManager.

[peterdsharpe]

Is this a new feature, an improvement, or a change to existing functionality?

New Feature

How would you describe the priority of this feature request

Low (would be nice)

Please provide a clear description of problem you would like to solve.

The GLOBE AirFRANS example as part of #1401 adds some utilities for doing profiling during distributed training. @laserkelvin suggested we upstream these, which I think is a great idea. Creating this issue until we can land the PR for this.

Related thread:
#1401 (comment)

Describe any alternatives you have considered

No response

[coreyjadams]

Hey @peterdsharpe and @laserkelvin I'm sorry; I'm having trouble finding the context in the GLOBE PR for this and I also reviewed it 😇 . Can you highlight the relevant sections?

Note that we've got some profiling utilities baked into physicsnemo already (torch profiling, line_profiling) and they also support distributed profiling (I use it all the time) so there's totally space to add more functionality. I would be happy to help. Do you want to spin it off into a separate PR so we don't hold up GLOBE?

[laserkelvin]

IIRC it's in relation to bits like this: https://github.com/NVIDIA/physicsnemo/pull/1401/changes#diff-71df5750e40509c44c7c8067f1a3e123dd8bc2c017da19325055d1a833fac33eR556

In my head I think it was a combination of context managers, doing stuff on rank 0, that kind of thing

[coreyjadams]

We can get into the details on the right method, but I actually think it's better to be profiling on every rank. There are tricky nuances - clocks are not necessarily in sync - but you can actually, in principle, synthesize traces across processes and see collectives slow down because some ranks are slower, etc.

Anyways, the physicsnemo profiling tools support arbitrary profiling extensions natively with distributed profiling. It controls output mapping to avoid collisions:

https://github.com/NVIDIA/physicsnemo/blob/main/physicsnemo/utils/profiling/core.py#L185-L208

And things like the torch profiler are built in:

https://github.com/NVIDIA/physicsnemo/blob/main/physicsnemo/utils/profiling/torch.py

the tools I made, if you want to look at them, are basically a registry pattern: profilers are very much stateful so you don't want to be constantly making new ones, instead you have a thread-safe static getter and it will enable -> initialize -> finalize any profilers the user wants smoothly. I think it works really nicely but I'm biased. You can do things like decorate functions with @profile from physicsnemo, and any tools that are decorate based you can turn on and off in one spot. You can do similar with contexts.

I guess this is not widely used or known about!

[github-actions]

This issue is stale because it has been open for 60 days with no activity. Remove the stale label or comment to keep it open, otherwise this will be closed in 14 days.