ray.torch.train.prepare_data_loader

[MMorente]

What happened + What you expected to happen

I have a local ray cluster and I am using pytorch dataloaders. As described in the docs, I am using ray.train.torch.prepare_data_loader to prepare the dataloader for ray, this works well when the workers only use one GPU, however if in the ScalingConfig I specify more than one GPU per worker, prepare_data_loader fails with the message:

device = get_device() 
self._auto_transfer = auto_transfer if device.type == "cuda" else False
AttributeError: 'list' object has no attribute 'type'

My cluster has two nodes and each one has 4 GPUs.

Versions / Dependencies

Ray: 2.6.1.
Python: 3.9.5
OS: Unix

Reproduction script

from torch.utils.data import DataLoader, Dataset
import numpy as np

import ray
from ray.train.torch import TorchTrainer
from ray.air.config import ScalingConfig

class TestDataset(Dataset):
    def __init__(self):
        super().__init__()
        self.data = np.random.random_sample((10,1))

    def __getitem__(self, index):
        return self.data[index]

def train_func(config: dict) -> None:
   dataset = TestDataset()
   dataloader = DataLoader(dataset)
   dataloader = ray.train.torch.prepare_data_loader(dataloader)
   print("done")

if __name__ == "__main__":
    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        scaling_config=ScalingConfig(num_workers=1, use_gpu=True, resources_per_worker={"GPU":2})
    )
    trainer.fit()

Issue Severity

High: It blocks me from completing my task.

[xwjiang2010]

I verified that the bug existed on 2.6.1. And probably on master as well. We didn't do a good job testing multiple GPU per worker case.

The following needs to happen:

• crystalize get_device semantics. Across multiple versions of this file, I see the semantics change - whether it's a list v.s. it's an int.
• update get_device doc string to reflect that clearly.
• update the logic in WrappedDataLoader to accommodate the fact that it can be a list.

cc @matthewdeng for triaging visibility. Marking this as P1 for now.

@MMorente would you like to contribute a PR to fix this? I think the fix can be as simple as

device = get_device()
if isinstance(device, list):
    device = device[0]

[matthewdeng]

@xwjiang2010 can you update the docs to reflect the behavior?

[MMorente]

I verified that the bug existed on 2.6.1. And probably on master as well. We didn't do a good job testing multiple GPU per worker case.

The following needs to happen:

• crystalize get_device semantics. Across multiple versions of this file, I see the semantics change - whether it's a list v.s. it's an int.
• update get_device doc string to reflect that clearly.
• update the logic in WrappedDataLoader to accommodate the fact that it can be a list.

cc @matthewdeng for triaging visibility. Marking this as P1 for now.

@MMorente would you like to contribute a PR to fix this? I think the fix can be as simple as

device = get_device()
if isinstance(device, list):
    device = device[0]

@xwjiang2010 I have just created a pull request #38127 with the change.

[anyscalesam]

@matthewdeng can you re-review and advise on next steps?

[matthewdeng]

This will be fixed by #42314, cc @woshiyyya to double check.

[woshiyyya]

Yes this will be fixed by the this PR.