Profiling only a select group of ranks

[githubsgi]

Is it possible to profile only a select group of ranks. Becomes hard to handle the large number of files when there are many ranks. I understand that there could be imbalances when only a few ranks are profiled. Do not know if there are ways to profile , but not dump the profile output file.

[fegin]

Currently, TorchTitan only support rank0 profiling or all ranks profiling. But this request is reasonable, do you want to submit a PR for this?

[githubsgi]

Sure, I can do a PR.

[tianyu-l]

Currently, TorchTitan only support rank0 profiling or all ranks profiling.

btw, this is for tensorboard / wandb logs. For profiler, all ranks will be profiled.

I think it's OK to not generate profiler files on all ranks.