[BUG] - Multiple 5090s failing on deepspeed.initialize()

[Oruli]

Describe the bug

The developer of the training code Diffusion-pipe helped me debug this, the issue on that repository has all the relevant information that I have now. His summary:

So plain PyTorch GPU communication ops work. But deepspeed.initialize() is always failing when it does its version of cross-GPU communication. Myself and other users have this working, but it fails specifically with multiple 5090s, and you are probably the only person who has tried that setup.

I would raise an issue with Deepspeed. I don't think I've done anything wrong in the application code, and it is likely an internal Deepspeed problem. Without being able to reproduce the error myself, there's not much more I can do.

Full issue: tdrussell/diffusion-pipe#235 (comment)

To Reproduce
Steps to reproduce the behavior:

Run deepspeed.initialize() with 2 x 5090 GPUs

System info (please complete the following information):

• OS: ubuntu 24.04
• GPU count and types: 1 machine with 2 x 5090s
• Python version: 3.12
• Any other relevant info about your setup: All latest Nvidia drivers, pytorch nightly etc.

[AutoAwesome]

I have the same issue. Just with an 5090 + 4090. I can run training on each card individually. But using both together throws illegal memory access errors.

[loadams]

Hi @Oruli - we don't have any 5090s so we cannot test this but I do not see this on a machine with 2 A6000s. Could you perhaps share your DeepSpeed and cuda version as well here? As well as just the simple script you are running?

[Oruli]

Hi @Oruli - we don't have any 5090s so we cannot test this but I do not see this on a machine with 2 A6000s. Could you perhaps share your DeepSpeed and cuda version as well here? As well as just the simple script you are running?

I would love to help test this so we can get it fixed for me and other users.

I am not on that box right now but I believe I am using the latest of everything, pytorch nightly, CUDA 12.8, as for deepspeed I am not specifying version so again assume it's latest.

It's installed via https://github.com/tdrussell/diffusion-pipe/blob/main/requirements.txt

I am not using my own script, I am running Diffusion Pipe, I posted the developers reply in my OP.

Let me know if you have a script I can run to test and provide logs for etc.

[Oruli]

@loadams Still no reply even after my offer to debug and get this fixed? I'm not the only person with 2 x 5090s.

[loadams]

@Oruli - I re-read the thread. Are you still seeing this with the latest DeepSpeed version? Just so we can narrow our search to the current commits.

And can you share a minimal repro for the issue, since I don't see one on the linked issue. That way we can look at the code to be sure and see the launcher you're using as well as if deepspeed.init_distributed() is required here.

[Oruli]

@Oruli - I re-read the thread. Are you still seeing this with the latest DeepSpeed version? Just so we can narrow our search to the current commits.

And can you share a minimal repro for the issue, since I don't see one on the linked issue. That way we can look at the code to be sure and see the launcher you're using as well as if deepspeed.init_distributed() is required here.

Yes it's the same, and I already answered the question about code, I don't write and I'm not using my own code.

I am using the diffusion-pipe repo, that I've linked and even provided the developers response that it relates to your code, I'm not sure what else you want me to provide at this point.

[loadams]

@Oruli - can you please run a deepspeed example that uses multiple GPUs to see if it is the DeepSpeed integration in the other library or related to DeepSpeed and 5090s specifically?

Any training script should do, lets test with the DeepSpeed launcher:

deepspeed --num_gpus=2 your_script.py --deepspeed_config ds_config.json

Dummy ds_config:

{
  "train_batch_size": 4,
  "gradient_accumulation_steps": 1,
  "fp16": {
    "enabled": false
  },
  "zero_optimization": {
    "stage": 0
  }
}

Dummy script:

# test_deepspeed_two_gpus.py

import torch
import deepspeed
import argparse

def parse_args():
    parser = argparse.ArgumentParser()
    parser = deepspeed.add_config_arguments(parser)
    parser.add_argument('--deepspeed_config', type=str, default='ds_config.json')
    return parser.parse_args()

def main():
    args = parse_args()

    # Dummy model for testing
    model = torch.nn.Linear(10, 10)

    # Initialize DeepSpeed
    model_engine, _, _, _ = deepspeed.initialize(
        args=args,
        model=model,
        model_parameters=model.parameters()
    )

    # Print which GPU the model is on
    print(f"Model is on device: {model_engine.device}")

    # Dummy input
    input_data = torch.randn(4, 10).to(model_engine.device)
    output = model_engine(input_data)
    print(f"Output shape: {output.shape}")

if __name__ == "__main__":
    main()

[Oruli]

@loadams Thanks for the help.

So to be clear my setup is using Conda, I'm installing as follows to be able to run your script:

conda create -n deepspeed python=3.12
conda activate diffusion-pipe
conda install pip

pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
conda install nvidia/label/cuda-12.8.1::cuda-nvcc

pip install --upgrade pip wheel setuptools
pip install deepspeed

result: (currently getting error running the command you gave)

deepspeed --num_gpus=2 your_script.py --deepspeed_config ds_config.json

argparse.ArgumentError: argument --deepspeed_config: conflicting option string: --deepspeed_config

[tjruwase]

@Oruli, I noticed in your OP that the failure occurs during a send/recv operation. Can you also try the p2p tests in the communication benchmark suite? https://github.com/deepspeedai/DeepSpeedExamples/tree/master/benchmarks/communication