[BUG] Engine returned by deepspeed.initialize() on the wrong device

[skpig]

Describe the bug
According to #662, the --include arguments can set the CUDA_VISIBLE_DEVICES properly. But the engine returned by 1deepspeed.initialize()1 is on the wrong device.
To Reproduce
launch the code below with deepspeed --include="localhost:3" train.py --deepspeed --deepspeed_config config.json

def debug():
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", default=0, type=int,
                        help="local_rank for distributed training on gpus")
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    os.environ['MASTER_ADDR'] = 'localhost'  #
    os.environ['MASTER_PORT'] = '12345'  #
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
    args.local_rank = int(os.environ['LOCAL_RANK'])

    model = torch.nn.Linear(10,10)
    engine, _, _, _ = deepspeed.initialize(args=args, model=model, model_parameters=model.parameters())
    print("Engine is on device: ", next(engine.parameters()).device)

Expected behavior
I'm willing to run the model on device 3, but the engine is on device 0.

[skpig]

I've checked the source code of class DeepSpeedEngine. The device is set by args.local_rank or os.environ['LOCAL_RANK']
https://github.com/microsoft/DeepSpeed/blob/41ab660b5df3567966935fe8ac3672497bc8690a/deepspeed/runtime/engine.py#L502-L526

So I run the other code snippet as below

def debug():
    parser = argparse.ArgumentParser()
    parser.add_argument("--local_rank", default=0, type=int,
                        help="local_rank for distributed training on gpus")
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    os.environ['MASTER_ADDR'] = 'localhost'  #
    os.environ['MASTER_PORT'] = '12345'  #
    os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
    os.environ['LOCAL_RANK'] = '3'
    args.local_rank = int(os.environ['LOCAL_RANK'])

    model = torch.nn.Linear(1000,1000)
    engine, _, _, _ = deepspeed.initialize(args=args, model=model, model_parameters=model.parameters())
    print("Engine is on device: ", next(engine.parameters()).device)

Now the print() shows the engine is on device 3 properly. But nvidia-smi shows some memory leak on device 0.

[loadams]

Hi @skpig - apologies for not replying sooner, is this still a bug that you believe you're hitting? And if so, could you confirm it happens with the latest DeepSpeed too?

[loadams]

Hi @skpig - closing this issue for now given the older version of DeepSpeed. If you are hitting this still, please reply and we can re-open and will gladly work on debugging it now. Apologies for the long delay in replying the first time.

[motata]

I'm facing the same problem. Even with --include local_host there's still memory being located on GPU:0.