[inductor] Parallelize Max Autotune step 2: Use multiple GPUs

[masnesral]

Stack from ghstack (oldest at bottom):

• -> [inductor] Parallelize Max Autotune step 2: Use multiple GPUs #109127
• [inductor] Parallelize Max Autotune step 1: refactor autotune_process #109126

Test Plan:
python test/inductor/test_max_autotune.py
TORCHINDUCTOR_AUTOTUNE_IN_SUBPROC=1 TORCHINDUCTOR_MAX_AUTOTUNE=1 python benchmarks/dynamo/torchbench.py --device cuda --performance --backend inductor --inference --only hf_Bart
TORCHINDUCTOR_AUTOTUNE_MULTI_DEVICE=1 TORCHINDUCTOR_AUTOTUNE_IN_SUBPROC=1 TORCHINDUCTOR_MAX_AUTOTUNE=1 python benchmarks/dynamo/torchbench.py --device cuda --performance --backend inductor --inference --only hf_Bart

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 cdd5ed3 with merge base 264f1e7 ():
💚 Looks good so far! There are no failures yet. 💚

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

[masnesral]

@shunting314 FYI this is a redo of #107983 which we had to revert since it didn't work well in fbcode.

@aakhundov, FYI I took your suggestions from that PR.

[eellison]

nice!

[eellison]

cool, so this includes the fix @aakhundov mentioned

[aakhundov]

@masnesral Thanks for the fix!

[masnesral]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[aakhundov]

Interesting that this works. I guess, it means that pytorch is able to picks up the device value automatically from the newly update os.environ["CUDA_VISIBLE_DEVICES"]?

[masnesral]

@aakhundov Hmmmmm, good catch. It does not work as I expected. I saw speedups comparable to the first implementation, so blindly assumed it was due to using multiple GPUs. But I guess the use of multiple subprocesses (even if using the same GPU) provides perf benefits? I found nvidia-smi pmon and it's reporting all my parallel subprocesses on the same device. The reason for organizing it this way is the multiprocessing API doesn't seem to allow passing a new environment before spawning the new procees. Lemme dig some more.

[shunting314]

the benchmark result will be very much distorted in this case. We have to make sure different subprocess are running on different devices.

[eellison]

I wonder if anyone has any good idea for testing this

[aakhundov]

I agree it's incorrect to run more than one profiling on the same device in parallel. So in AITemplate we do this to launch a process with a specific CUDA_VISIBLE_DEVICES env var (dev_select_flag is "CUDA_VISIBLE_DEVICES"; device is a zero-based int id of the device). And this seems to work: whatever the process does on a GPU ends up on the #device. Could we do something similar here?

[aakhundov]

@masnesral I believe, the reason you saw speedups was because the actual GPU work during benchmarking is not that intensive. As I understand, most of the time may be spent in launching and processing. But we can't guarantee that the kernels will not overlap (hence distort the measurements), so still must serialize them on each device.

[masnesral]

So in AITemplate we do this

Yes, setting that env var works as expected, generally speaking. The problem in this current implementation is that setting the env var after the sub-process has already started does not have the desired affect. The previous impl that used a Popen (basically the same approach as the AIT impl) worked correctly in this regard. AFAICT, if we stick with the multiprocessing lib (which has other advantages) the only way to manipulate the env is to do it in the parent process before starting the sub-process. I'll put up an RFC soon.

[masnesral]

#109500