Fork safety

[Flamefire]

It looks like deepspeed is not fork-safe:

• During import it checks builder.is_compatible() for all ops
• Some ops use torch.cuda.get_device_properties
• This initializes the CUDA context

See

DeepSpeed/deepspeed/git_version_info.py

Line 30 in 38bd11a

compatible_ops[op_name] = op_compatible

When the process then forks to run in parallel any access to CUDA through PyTorch will fail:

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

I'm not sure why this doesn't seem to be fully consistent but I can reproduce it with:
pytest --forked tests/unit/ops/deepspeed4science/test_DS4Sci_EvoformerAttention.py -k 'test_DS4Sci_EvoformerAttention[tensor_shape1-dtype1]' -s

Where it will fail when invoking skip_on_arch(8 if dtype == torch.bfloat16 else 7) which calls torch.cuda.get_device_properties, now in the forked subprocess.

This is an issue in general: Fork-multiprocessing can not be used after importing deepspeed

But it also contradicts the documentation:

Note that pytest-forked and the --forked flag are required to test CUDA functionality in distributed tests.

It seems the opposite is true: The flag must not be used.

Or am I missing anything?

[vibeyclaw]

Ran into similar CUDA fork issues before. The root cause is clear from your trace: any torch.cuda.* call during import locks the CUDA context to the parent process, making fork() unusable.

One workaround that helped us: lazy-wrap the compatibility checks so they only fire when the op is actually used, not at import time. Something like caching is_compatible() behind a property that defers the CUDA call.

For the test suite specifically, switching pytest to --forked with mp.set_start_method("spawn") in conftest avoids it, though that is more of a band-aid.

Might be worth checking if CUDA_VISIBLE_DEVICES="" during import + deferred re-check could work as a less invasive fix.