Model sharing across torch Processes

[JanRocketMan]

❓ Questions and Help

Hi, I have a model that needs to be shared across different process - e.g. one trains it while the other collects new training data for it (like in RL). In traditional pytorch this can be easily achieved with smth like:

# Create shared model
shared_model = nn.Linear(512, 32).to(device)
shared_model.share_memory()

# Pass it during creation
some_process = Process(
    target=some_fn,
    args=(shared_model)
)
some_process.start()

# do smth else with shared_model

# Subprocess can write/get from shared_model

which is unavailable for XLA tensors since they don't have storage.

Now I may try pushing this state back & forth through Queue, or even cast all weights to numpy arrays and share them, but I afraid this may be way too messy & inefficient. Are there any other approaches? Like is there any way of recasting xla tensors to native tensors only for exchanging, and then casting back for efficient computation?

[dlibenzi]

At the moment it is not possible as TPU cores cannot share HBM.

Running totally separate models on different cores it is also not possible, though we may add support for it.

[JanRocketMan]

What about models that are stored on cpu-s? Do they need conversion of weights to non-tensor types?

[dlibenzi]

Once you model.to(device) the data is moved to the destination device.
For models on CPU it might be easier to share data, though, I am not sure how that would work in pytorch for write operations, as I saw no explicit synchronization when write-accessing storages.
Maybe it is up to the user to do that, but can't say as I never used that.

[JanRocketMan]

Yeah, it seems to be working fine once I cpu-stored weights, thanks.

[JanRocketMan]

Sorry for reopening this, but since running separate models is impossible on a single TPU, is there a way to simultaneously optimize and infer predictions of the same model with different batches on torch-xla? As I get it there is an example of doing this in tensorflow (tensorflow/tpu#48), I'm wondering is there any workarounds for doing it here.

[dlibenzi]

That example (which is gone now), does (did) one epoch training followed by eval.
Which is what one normally do in pytorch manually.

Why don't you describe in short pseudo-code what you are trying to do?

[JanRocketMan]

Yeah, probably. I'm currently doing smth like this:

In process 1:

data = ... # create some data

# Synchronize model weights
model.to('cpu')
model.load_state_dict(shared_model.state_dict())
model.to(xla_device)

outputs = model(data)
# send outputs to process 2

In process 2:

data = ... # receive outputs from process 1

# Synchronize model weights
model.to('cpu')
shared_model.load_state_dict(model.state_dict())
model.to(xla_device)

dataloader = ... # Pack data to loader 
for data, _ in dataloader:
    # predict, zero-grad, backprop
    xm.optimizer_step(optimizer, barrier=True)

The code fails in process 1 when I try to copy the model to cpu with:

  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 450, in to
    return self._apply(convert)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 210, in _apply
    module._apply(fn)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 210, in _apply
    module._apply(fn)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 232, in _apply
    param_applied = fn(param)
  File "/anaconda3/envs/torch-xla-nightly/lib/python3.6/site-packages/torch/nn/modules/module.py", line 448, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
RuntimeError: Not found: From /job:tpu_worker/replica:0/task:0:
XRT memory handle not found: 4021333374585961
         [[{{node XRTExecute}}]]

[JanRocketMan]

Actually, I want to test smth pretty similar to what's happening in SEED RL:

https://ai.googleblog.com/2020/03/massively-scaling-reinforcement.html

but without any super-efficient synchronization layer. They somehow manage to simultaneously make predictions at new training points while optimizing model on old ones.

[dlibenzi]

Why would you want to send data to another device?

Meaning, can't you do things on the same device?

[JanRocketMan]

In the above code both processes use the same device. It's just a way for using some parts of device to infer things, while the other parts perform backprop. Otherwise e.g. backprop will have to wait the inference to finish and it'll work two times slower.

[dlibenzi]

The thing is, when TPU connected by the fast interconnect do collective operations (all-reduce & co.), all the cores needs to perform that.

So if in your case you have heterogeneous tasks, that is not possible.
At that point you are better off exchanging data via normal networking.

[JanRocketMan]

This sounds strange - why can't I specify which of the cores perform these computations & synchronize together, while the others perform different computations? What is point in having multiple cores otherwise? I checked once more Seed-Rl implementation, they seem to be distributing inference/training computations btwn different cores.