Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Nodes] Add a ToDevice node, or combine with pin memory #1407

Open
andrewkho opened this issue Dec 13, 2024 · 2 comments
Open

[Nodes] Add a ToDevice node, or combine with pin memory #1407

andrewkho opened this issue Dec 13, 2024 · 2 comments
Labels

Comments

@andrewkho
Copy link
Contributor

🚀 The feature

We should add a node that will send batches to device (probably one at a time). We could either separate this, add it on to pre-fetcher (ie always call .to(device) on the head of the queue, or maybe part of pin-memory

Motivation, pitch

Sending data to device can be slow, and often users want this done in a background thread. DataLoader should do this in the backgroudn as it consolidates state management

Alternatives

No response

Additional context

No response

@divyanshk
Copy link
Contributor

I wonder, how different can this be from doing the transfer within a Mapper, similar to a collate_fn doing tensor.to(device)

@divyanshk
Copy link
Contributor

For cases where we have multiple threads reading from data, we might be able to create multiple thread local CUDA streams to transfer data onto the GPU. WDYT @andrewkho ?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

No branches or pull requests

2 participants