Currently this PR creates 2 workers, each with 2 GPUs

pytest test --verbose -s --dist load --tx popen//env:CUDA_VISIBLE_DEVICES=0,1 --tx popen//env:CUDA_VISIBLE_DEVICES=2,3

I think only FSDP tests require 2 GPUs

Thus, CI runtime can be further improved by

1. Mark these tests as "require 2 GPUs"
2. Run the remaining tests with 4 workers, each with 1 GPU
3. Run the 2-GPU tests in a separate step. Can either be 2 workers with 2 GPUs each, or just normal, non-parallel pytest since there are not a lot of 2-GPU tests right now.

UPDATE: parallel tests with 4 1-GPU workers + sequential multi-GPU tests is implemented now. Some problems I encountered along the way (for future reference and improvement). pytest-xdist provides multiple methods to distribute work across workers. Refer to here for the full list. I have tried the following

• --dist load (send each test to next available worker):
  • TestSaveLoadMeta use the same filename test.pth which is problematic when multiple tests write/read the same file. I made a unique name for each test to solve it for now, though the better way is to use pytest tmp_path fixture (not compatible with unittest.TestCase).
  • test_quant_llm_quantize_ebits_2_mbits_2 gets a strange torch.compile warning in PyTorch 2.3 only torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py:436] [0/48] failed to eagerly compile backwards for dynamic, suppressing in case backwards not needed. Looks harmless though.
• --dist loadscope (send a TestClass to next available worker):
  • TestSaveLoadMeta gets a strange error - RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent forward run. Stack trace: [Could not find stack trace]. To prevent overwriting, clone the tensor outside of torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.. No traceback. Since tests within a class are run by the same worker, there shouldn't be filename conflict test.py. Perhaps some strange interaction between how pytest-xdist sends test to the subprocess worker and CUDA graphs + torch.compile?
• --dist loadfile (send each file to next available worker): this should be the least error-prone. However, there is little speedup (24min vs 35min) since test_integration.py file is huge -> load imbalance (probably a good idea to split this into smaller files anyway).

The current approach of this PR is using --dist load. To summarize

I also notice that it takes 3min to pull the Docker image and another 3-4min to download and install packages. So it seems like two 2-GPU workers have pretty good scaling (7 + 14 vs 7 + 28), while four 1-GPU workers is not much of an improvement (7 + 12 vs 7 + 28). I guess the 2-GPU FSDP tests can be run in parallel too for the latest configuration.