import transformer_engine initializes CUDA

[szmigacz]

>>> import torch
>>> torch.cuda.is_initialized()
False
>>> import transformer_engine
>>> torch.cuda.is_initialized()
True

Import alone shouldn't initialize CUDA. Custom subprocess launchers are going to fail with RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method if CUDA is initialized.

[timmoon10]

The root cause is because torch.compile initializes CUDA in https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/jit.py. If I run with NVTE_TORCH_COMPILE=0 in the environment (to use nvFuser instead of torch.compile), importing TE doesn't initialize CUDA.

Possible next steps:

• If you just need a quick workaround, run with NVTE_TORCH_COMPILE=0.
• We could modify https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/jit.py so that it performs fusions lazily. This would be messier, but would make initialization faster since we only fuse the functions we use.
• File a bug with PyTorch.