Accumulating local gradients to global shared model

jleguina0: <p>Hi everyone!</p>
<p>First post here so apologies if it is unclear. I am currently facing a problem similar to what was mentioned here (<a href="https://discuss.pytorch.org/t/share-gradients-in-multiprocessing/15181" class="inline-onebox">Share gradients in multiprocessing</a>). Given that no answer was provided I will try to further specify the problem.</p>
<p>I have implemented a Hogwild model training procedure. I am trying to modify it such that gradients calculated with local models over N cores can be accumulated/averaged onto the shared model. After aggregating the gradients from N-1 cores, I want the Nth core to update the shared model given such accumulated/averaged gradients.</p>
<p>An overview of what I am considering is as follows:</p>
<pre><code class="lang-auto">import torch.multiprocessing as mp
from model import MyModel

def train(shared_model, shared_optimizer, rank):
    # Construct data_loader, optimizer, etc.
    local_model = MyModel()
    for data, labels in data_loader:
        local_model.load_state_dict(shared_model.state_dict())  # Ensure local model up to date
        loss_fn(local_model(data), labels).backward()
        sync_grads(local_model, shared_model)
        if rank == 0:
            shared_optimizer.step()  # This will update the shared parameters
            shared_optimizer.zero_grad()

def sync_grads(local_model, shared_model):
    for local_param, shared_param in zip(local_model.parameters(), shared_model.parameters()):
        shared_param._grad = shared_param.grad + local_param.grad  # Accumulating in this case
</code></pre>
<pre><code class="lang-auto">if __name__ == '__main__':
    num_processes = 4
    shared_model = MyModel()
    shared_model.share_memory()
    
    # SharedAdam: https://github.com/ikostrikov/pytorch-a3c/blob/master/my_optim.py
    shared_optimizer = SharedAdam(shared_model.parameters())  # Note shared model parameters
    shared_optimizer.share_memory()
    processes = []
    for rank in range(num_processes):
        p = mp.Process(target=train, args=(shared_model, shared_optimizer, rank, ))
        p.start()
        processes.append(p)
    for p in processes:
        p.join()
</code></pre>
<p>My questions are:</p>
<ul>
<li>Will <code>sync_grads()</code> upload gradients to the shared model (and thus grads will be shared among processes)?</li>
<li>Or will this only aggregate gradients to a <em>local</em> version of the <code>shared_model</code>, leading to the update only really occurring with gradients computed for the rank0 process?</li>
<li>In other words, does putting a <code>shared_model</code> on shared memory share parameters <em>and</em> gradients or only parameters?</li>
</ul>
<p>Finally, if my fears are correct, how can I do this? I have seen <a href="https://pytorch.org/tutorials/intermediate/dist_tuto.html#distributed-training" rel="noopener nofollow ugc">this</a>:</p>
<pre><code class="lang-auto">""" Gradient averaging. """
def average_gradients(model):
    size = float(dist.get_world_size())
    for param in model.parameters():
        dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
        param.grad.data /= size
</code></pre>
<p>However, I am skeptical of whether it will work (last update to the tutorial is from 2017). Else, I am also considering moving grads to an <code>mp.Queue</code> for all processes, then query, average &amp; assign to the local version of <code>shared_model</code> in rank0 and gradient step on that process only.</p>
<p>Thanks a lot!</p>

gcramer23: <p><a class="mention" href="/u/jleguina0">@jleguina0</a> Does DDP fit your need <a href="https://pytorch.org/tutorials/intermediate/ddp_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with Distributed Data Parallel — PyTorch Tutorials 1.9.0+cu102 documentation</a>? You will have one model replica per process. Instead of storing the gradients in a shared model DDP will update all the model replicas. You can define a communication hook <a href="https://pytorch.org/docs/stable/ddp_comm_hooks.html" class="inline-onebox" rel="noopener nofollow ugc">DDP Communication Hooks — PyTorch 1.9.0 documentation</a> for gradient communication logic.</p>

jleguina0: <p><a class="mention" href="/u/gcramer23">@gcramer23</a> That is exactly what I was looking for! Thanks.</p>