Multiprocessing CUDA issues when importing transformers

[maxjeblick]

Environment info

• transformers version: 4.2.0dev0
• Platform: Linux-4.15.0-128-generic-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.6.9
• PyTorch version (GPU?): 1.6.0 (True)
• Tensorflow version (GPU?): not installed (NA)
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: Partly

Who can help

@stas00

Information

When using multiprocessing, importing the transformers package causes
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

The problem arises when using:

• import transformers is used in the main process.
  This is due to the following line in modeling_fsmt.py, removing torch.cuda.is_available() call resolves the issue.

To reproduce

import multiprocessing
import transformers  # NOQA
# You can also call torch.cuda instead of import transformers to get the same error
import torch.nn as nn


class Net(nn.Module):

    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 10)

    def forward(self, x):
        return self.linear(x)


def to_cuda(i):
    net = Net().cuda()
    print(f'Called {i} process')


try:
    cpus = multiprocessing.cpu_count()
except NotImplementedError:
    cpus = 2  # arbitrary default

pool = multiprocessing.Pool(processes=cpus)
pool.map(to_cuda, range(10))

Expected behavior

The code snippet above runs without issues.

[patrickvonplaten]

Hey @maxjeblick,

We do need to be able to call torch.cuda in our model code. We cannot "forbid" calls to torch.cuda in order to allow quite specific use cases where .cuda() is mixed with multiprocessing. IMO, using multiprocessing in combination with .cuda() is quite an edge case. Are you trying to run your model on multiple GPUs?

I'd suggest to fork the repo and delete this line if you really need this feature.

[maxjeblick]

Hey @patrickvonplaten thanks for the fast response!

I agree that using multiprocessing with .cuda calls is not that common, one usecase would be multi-gpu inference without DDP. The torch.cuda.is_available() call in modeling_fsmt.py is currently the only place which causes the RuntimeError; all other .cuda calls are either fine (e.g. from torch.cuda.amp import autocast in trainer.py) or not executed during the import transformers statement.

[stas00]

It affects bart as well.

I see ProphetNetLayerNorm solved it by using a runtime wrapper.

transformers/src/transformers/models/prophetnet/modeling_prophetnet.py

Lines 513 to 521 in 912f688

	def ProphetNetLayerNorm(normalized_shape, eps=1e-5, elementwise_affine=True):
	if torch.cuda.is_available():
	try:
	from apex.normalization import FusedProphetNetLayerNorm
	return FusedProphetNetLayerNorm(normalized_shape, eps, elementwise_affine)
	except ImportError:
	pass
	return torch.nn.LayerNorm(normalized_shape, eps, elementwise_affine)

So it's modeling_bart and modeling_fsmt are the only 2 that have this check at import time

I wonder if we can just remove this check altogether. Won't from apex.normalization import FusedLayerNorm fail w/o cuda? And then we actually don't need that check. We need to verify that.

[stas00]

heh, it looks like FusedProphetNetLayerNorm importing is silently failing, since it doesn't exist in apex ;) looks like untested code. Fix proposed at #9349

[stas00]

Hmm, it works just fine without cuda:

CUDA_VISIBLE_DEVICES="" python -c "from apex.normalization import FusedLayerNorm; print(FusedLayerNorm(10))"
FusedLayerNorm(torch.Size([10]), eps=1e-05, elementwise_affine=True)

Do you know why did we need that check in first place?

The doc page https://nvidia.github.io/apex/layernorm.html doesn't say anything about needing cuda.

[stas00]

@t-vi figured it out - apex.normalization.FusedLayerNorm falls back on to non-cuda gracefully: https://github.com/NVIDIA/apex/blob/master/apex/normalization/fused_layer_norm.py#L154 so the if torch.cuda.is_available() check is not needed in first place.

[patrickvonplaten]

needed

Nice, so we could actually remove this import statement then from all FusedLayerNorm classes no?

[stas00]

yes, working on this. - just 3 classes

Done: #9350

[stas00]

so we could actually remove this import statement then from all FusedLayerNorm classes no?

but, wait, what import statement are you referring to?

and I guess modeling_fsmt is waiting for the refactoring, right? So since Bart I see has this import-time check removed already, so it'd follow suit anyway.

[stas00]

maxjeblick, while we are sorting out the nuances you can just remove that check so that you could move forward. No matter the outcome that particular call that was getting in your way won't be there once the dust settles.

[jethrokuan]

Related: #9227

I found that the changes proposed above weren't enough for cuda multiprocessing.