Using 2 GPUs for Different Parts of the Model

[DuaneNielsen]

One approach…

Start 2 python programs, in separate interpreters to avoid the dreaded GIL lock.

Processor 1

1. Put tensor on cuda:0, get the output.
2. Serialize and push the output to shared redis database

Processor 2

1. Consumer picks up from database, pushes to cuda:1
2. Consumer runs the next step of calculation.

If you need to send gradients for backprop you can store and reload them also.

Bit of fun with gradients

Thought I’d share this code. Learned a few things about autograd doing this. import torch # unbroken gradient, backward goes all the way to x x = torch.ones(2, 2, requires_grad=True) y = 2 * x + 2 z = y * y * 3 out = z.mean() out.backward() print(x.grad) baseline_x = x.grad # broken gradient, ends at _y x = torch.ones(2, 2, requires_grad=True) y = 2 * x + 2 _y = torch.tensor(y.detach(), requires_grad=True) z = _y * _y * 3 out = z.mean() out.backward() print(x.grad) print(_y.grad) # we can …

That’s one way… Not easy though. I spent easy a month just trying to distribute calculations over multiple processors.

If you can pull it off… then it’s an awesome skill.

Also, there is the Ray project GitHub - ray-project/ray: Ray is a unified framework for scaling AI and Python applications. Ray consists of a core distributed runtime and a set of AI Libraries for accelerating ML workloads.

I tried using it. It had great promise, but ended up being a bit too new at the time. It might be a bit more mature now.