Rref can only be used once

username: 111282
cooked: <p>hi,<br>
I implement a RPC distributed model like <a href="https://pytorch.org/tutorials/intermediate/rpc_tutorial.html#distributed-rnn-using-distributed-autograd-and-distributed-optimizer" rel="noopener nofollow ugc">the tutorial example(rnn part)</a></p>
<p>I put an embedding module locally and a pretrained bert model remotely. The whole model is as follows:</p>
<pre><code class="lang-auto">class DistCPM(nn.Module):
    def __init__(self,
                 save_prompt_embedding_path,
                 device="cuda:0",
                 ps="ps"):
        super(DistCPM, self).__init__()
        # setup prompt locally(I want to add some prompt tokens to improve the accuracy)
        self.prompt_embedding_policy = PromptEmbeddingPolicy(
            prompt_embeds_path=save_prompt_embedding_path,
            device=device
        )
        # setup bert remotely
        # In addition to create a LanguageModelService model(e.g bert/bart), I also load its checkpoint. This will take a while
        # And the requires_grad is set to false to make the remote model has no grad.
        self.language_model_service_rref = rpc.remote(ps, LanguageModelService, args=())

    def forward(self,
                word_embeds,
                pad_emnbeds,
                batch_target_sequence_token):
        # make input embeddings and so on
        model_batch, _ = self.prompt_embedding_policy(
            word_embeds, pad_emnbeds, batch_target_sequence_token
        )
        logits = _remote_method(LanguageModelService.forward,
                                self.language_model_service_rref,
                                model_batch)
        return logits
    # return all params(local+remote) just as the example
    def parameter_rrefs(self):
        remote_params = []
        # create RRefs for local parameters
        remote_params.extend(_parameter_rrefs(self.prompt_embedding_policy))
        # get RRefs of bert(or bart and so on.Given I set bert.requires_grad=False, the following line may not be necessary)
        remote_params.extend(_remote_method(_parameter_rrefs, self.language_model_service_rref))
        return remote_params

# trainer, also like that in the example
def _run_trainer():
    r"""
    The trainer creates a distributed RNNModel and a DistributedOptimizer. Then,
    it performs training using random input data.
    """
    train_dataset_path = 'data/diagnose/train_1000.jsonl'
    test_dataset_path = 'data/diagnose/test.jsonl'
    save_prompt_embedding_path = 'data/checkpoint/embedding.pth'  # if exists, will be load

    device0 = "cuda:0"

    model = DistCPM(save_prompt_embedding_path, device=device0, ps="ps") 

    criterion = torch.nn.CrossEntropyLoss()

    opt = DistributedOptimizer(
        optim.SGD,
        model.parameter_rrefs(),
        lr=0.05,
    )

    # load bert word embeddings
    # I will concate the prompt tokens embedding and bert-tokens embedding
    word_embedding_service = WordEmbeddingService(device0)

    batch_size = 4
    batch_train_dataset = load_dataset(train_dataset_path, batch_size)

    for k in range(100):
        for bix, batch in enumerate(batch_train_dataset):
            print(f"epoch-{k} batch-{bix}")
            with dist_autograd.context() as context_id:
                batch_information = {'special_vocab_string_list': ['&lt;pad&gt;', '&lt;mask&gt;'],
                                     'batch_source_sequence_string': batch['source'],
                                     'batch_target_sequence_string': batch['target']}
                
                # The next few lines create input embeddings.
                batch_information = word_embedding_service.get_word_embedding(batch_information)
                word_embeds = batch_information['batch_sequence_embedding']
                pad_emnbeds = batch_information['special_vocab_embedding_list'][0]
                batch_target_sequence_token = batch_information['batch_target_sequence_token']
                
                # logits
                logits = model(word_embeds, pad_emnbeds, batch_target_sequence_token)
                logits = logits.float()

                target = batch_target_sequence_token

                loss = 0.0
                target_token_len = []
                # calculate loss trivially
                for index in range(len(target)):
                    if logits[index].shape[0] != target[index].shape[0]:
                        pass
                    else:
                        loss += criterion(logits[index], target[index])
                    target_token_len.append(target[index].shape[0])
                loss = loss / sum(target_token_len)
                print('batch_train_loss:', loss)

                # run distributed backward pass
                dist_autograd.backward(context_id, [loss])
                # run distributed optimizer
                opt.step(context_id)

</code></pre>
<p>the error is: <strong>RuntimeError: RRef creation via rpc.remote() timed out, and it is possible that the RRef on the owner node does not exist.</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/6/76916145b9a298d8244d1e16d8ed3230080a7e67.png" data-download-href="https://discuss.pytorch.org/uploads/default/76916145b9a298d8244d1e16d8ed3230080a7e67" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/76916145b9a298d8244d1e16d8ed3230080a7e67_2_690x289.png" alt="image" data-base62-sha1="gUTODfyyR24IaSXRJNa5R6B798X" width="690" height="289" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/76916145b9a298d8244d1e16d8ed3230080a7e67_2_690x289.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/6/76916145b9a298d8244d1e16d8ed3230080a7e67_2_1035x433.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/6/76916145b9a298d8244d1e16d8ed3230080a7e67.png 2x" data-dominant-color="EEF2F5"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1230×516 82.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I find the RREF object(namely self.language_model_service_rref) only can be used once.<br>
The first usage occurs in optimizer:</p>
<pre><code class="lang-auto">opt = DistributedOptimizer(
        optim.SGD,
        # [RRef(prompt_embedding_policy.prompt_embeds)],
        model.parameter_rrefs(),  # get the remote params will use self.language_model_service_rref)
        lr=0.05,
    )
</code></pre>
<p>The second usage is in model(*args) where the error occurs:</p>
<pre><code class="lang-auto">logits = model(word_embeds, pad_emnbeds, batch_target_sequence_token)

# model.forward()
def forward(self,
                word_embeds,
                pad_emnbeds,
                batch_target_sequence_token):
        # make input embeddings and so on
        model_batch, _ = self.prompt_embedding_policy(
            word_embeds, pad_emnbeds, batch_target_sequence_token
        )
        logits = _remote_method(LanguageModelService.forward,  # the second usage
                                self.language_model_service_rref,
                                model_batch)  
</code></pre>
<p>Any comment would be appreciated.</p>

username: wayi
cooked: <p>Does this mean that this error won’t occur if you don’t use <code>DistributedOptimizer</code>?</p>
<p>cc: <a class="mention" href="/u/mrshenli">@mrshenli</a></p>

username: 111282
cooked: <p>In the tutorial, the distributedOptimizer is necessary.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/9/39819ca942c0f2732fb1614afa51e45fbe36268b.png" data-download-href="https://discuss.pytorch.org/uploads/default/39819ca942c0f2732fb1614afa51e45fbe36268b" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39819ca942c0f2732fb1614afa51e45fbe36268b_2_690x264.png" alt="image" data-base62-sha1="8cIUiWB0t9Jywjt2LYb5TmsoFqr" width="690" height="264" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39819ca942c0f2732fb1614afa51e45fbe36268b_2_690x264.png, https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39819ca942c0f2732fb1614afa51e45fbe36268b_2_1035x396.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/3/9/39819ca942c0f2732fb1614afa51e45fbe36268b_2_1380x528.png 2x" data-dominant-color="E5E3E3"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1620×620 150 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>On the other hand, even the rref object(self.language_model_service_rref) is only used in the model(used in model.forward() exactly), the same error occurs when the second batch begins to be handled.</p>
<pre><code class="lang-auto"># remove the rref object from optimizer
opt = DistributedOptimizer(
        optim.SGD,

        # Just get the local params and wrap it with RRef,
        # because I need no update to the remote bert params
        [RRef(prompt_embedding_policy.params)],  

        # model.parameter_rrefs(),  # get the remote params, will use self.language_model_service_rref)
        lr=0.05,
    )
...
# when the model begins to deal with the second batch, the same error occurs(the rref object does not exist).
for batch in data_loader:
    logits = model(batch)  # note just an example here
</code></pre>

username: pbelevich
cooked: <p>hi <a class="mention" href="/u/111282">@111282</a>, when you create <code>self.language_model_service_rref</code> it takes sometime as you mentioned, but the result rref is returned to you immediately, although the underlying remote object is not created yet.<br>
This creation takes more time than is allowed by the default timeout and thus it fails silently. But when you try to call <code>_remote_method(LanguageModelService.forward, self.language_model_service_rref,...)</code> this error is shown to you. This behavior is described in <a href="https://pytorch.org/docs/stable/rpc.html#torch.distributed.rpc.remote" rel="noopener nofollow ugc">remote method timeout argument</a>:</p>
<blockquote>
<p>timeout in seconds for this remote call. If the creation of this <a href="https://pytorch.org/docs/stable/rpc.html#torch.distributed.rpc.RRef" rel="noopener nofollow ugc"><code>RRef</code></a> on worker <code>to</code> is not successfully processed on this worker within this timeout, then the next time there is an attempt to use the RRef (such as <code>to_here()</code> ), a timeout will be raised indicating this failure.</p>
</blockquote>
<p>What you need to do is to add some timeout greater than default one(60) here:</p>
<pre><code class="lang-auto">self.language_model_service_rref = rpc.remote(ps, LanguageModelService, args=(), timeout=...)
</code></pre>

username: 111282
cooked: <p><a class="mention" href="/u/pbelevich">@pbelevich</a> You are right. But I am still wondering why the error occurs on the second use?</p>
<p>The above error occurs when the second usage(namely model(*args)) happens. The first usage is getting the remote params when creating DistributedOptimizer.</p>
<p>Also even I remove the first usage(code as follows), the same error occurs when the model begins to deal with the second batch(The first batch is ok). So I thought the rref object can only be used once.</p>
<pre><code class="lang-auto"># remove the rref object from optimizer
opt = DistributedOptimizer(
        optim.SGD,

        # Just get the local params and wrap it with RRef,
        # because I need no update to the remote bert params
        # and bert.requires_grad = False
        [RRef(prompt_embedding_policy.params)],  

        # model.parameter_rrefs(),  # get the remote params, will use self.language_model_service_rref)
        lr=0.05,
    )
...
# Now the same error occurs when the second batch begins to be dealt with.
for batch in data_loader:
    logits = model(batch)  # note just an example here
</code></pre>