Does DistributedOptimizer support zero_grad and lr_scheduling?

Kunlin_Yang: <p>I try to use distributed_rpc to implement parameter server. I would like to use zero_grad and lr_scheduling feature in the trainer. But it seems that the DistributedOptimizer does not support this. Is there any workaround?</p>

mrshenli: <p>Hey <a class="mention" href="/u/kunlin_yang">@Kunlin_Yang</a></p>
<h3><code>zero_grad</code></h3>
<p>The reason <code>DistributedOptimizer</code> does not provide a <code>zero_grad</code> API is because the gradients of each backward pass is stored in its own dedicated context (instead of in <code>param.grad</code>), and the context will be cleared when exiting the <code>with dist_autograd.context() as context_id:</code> scope. So <code>zero_grad</code> is not needed here. We can certainly add it if necessary, and it won’t be too hard to implement it in the application code either. E.g.,</p>
<pre><code class="lang-python">def zero_grad(pr, context_id):
    dist_autograd.get_gradients(context_id)[pr.local_value()].zero_()


with dist_autograd.context() as context_id:
    # omitting forward-backward-optstep here
    futs = []
    for pr in param_rrefs:
        futs.append(rpc.rpc_async(pr.owner(), zero_grad, args=(pr, context_id)))
    [fut.wait() for fut in futs]
</code></pre>
<p>Or is there a different reason you would like to use the <code>zero_grad</code> API?</p>
<h3><code>lr_scheduling</code></h3>
<p>Currently, there is no distributed implementation or lr scheduling yet. I created an issue to track this: <a href="https://github.com/pytorch/pytorch/issues/38548" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/38548</a></p>
<p>For now, you will need to do that using the raw RPC API. You can access the RRefs of the remote optimizers through <code>DistributedOptimizer().remote_optimizers</code>, so it can be sth like:</p>
<pre><code class="lang-python">def create_lr_schheduler(opt_rref):
     # create and return lr_schheduler

def lrs_step(lrs_rref):
    lrs_rref.local_value().step()

opt = DistributedOptimizer(...)
lrs_rrefs = []
for opt_rref in opt.remote_optimizers:
    lrs_rrefs = rpc.remote(opt_rref.owner(), create_lr_schheduler, args=(opt_rref,))

with dist_autograd.context() as context_id:
    # omitting forward-backward-optstep here
    futs = []
    for lrs_rref iin lrs_rrefs:
        futs.append(rpc.rpc_async(lrs_rref.owner(), lrs_step, args=(lrs_rref,)))
    [fut.wait() for fut in futs]
</code></pre>
<p>If you are using master branch, the above code can be simplified with <a href="https://pytorch.org/docs/master/rpc.html#torch.distributed.rpc.RRef.rpc_async" rel="nofollow noopener"><code>RRef.rpc_async()</code></a> API.</p>

user090: <p>I am also using distributed rpc and trying to implement lr_scheduler. I tried your suggestion, however I am not able to update the learning rate. It still prints the initial learning rate. It would be great if you can help.</p>
<p>I am inheriting the Distributed optimizer as follows:</p>
<pre><code class="lang-auto">from torch.distributed.optim import DistributedOptimizer

class MyDistributedOptimizer(DistributedOptimizer):  
    def __init__(self, *args, **kwargs):  
        super().__init__(*args, **kwargs)  
      
    def fetch_learning_rate(self, optimizer_rref):  
        optimizer_local = optimizer_rref.local_value()  
        return optimizer_local.optim.defaults['lr']  
  
    def get_learning_rates(self):  
        learning_rates = {}  
        for optimizer in self.remote_optimizers:  
            if optimizer.owner() == rpc.get_worker_info():  
                optimizer_local = optimizer.local_value()  
                learning_rates_on_owner = optimizer_local.optim.defaults['lr']  
                learning_rates[f'worker_{optimizer.owner().name}'] = learning_rates_on_owner  
            else:  
                learning_rates_on_owner = rpc.rpc_sync(  
                    optimizer.owner(),  
                    self.fetch_learning_rate,  
                    args=(optimizer,)  
                )  
                learning_rates[f'worker_{optimizer.owner().name}'] = learning_rates_on_owner  
        return learning_rates  
</code></pre>
<p>The optimizer is called as follows:</p>
<pre><code class="lang-auto">    optimizer = MyDistributedOptimizer( 
        optim.AdamW,
        model.parameter_rrefs(),
        lr=config.learning_rate,) 
</code></pre>
<p>I have defined the custom scheduler as follows:</p>
<pre><code class="lang-auto">class CustomLRScheduler:
    def __init__(self, distributed_optimizer_rref, num_warmup_steps, num_training_steps, num_cycles=0.5):
        self.distributed_optimizer_rref = distributed_optimizer_rref
        self.num_warmup_steps = num_warmup_steps
        self.num_training_steps = num_training_steps
        self.num_cycles = num_cycles
        self.last_epoch = 0
        self.learning_rate = 0.0

    def step(self):
        self.last_epoch += 1
        self.learning_rate = self._get_lr(self.last_epoch)
        print(f"Updating learning rates to: {self.learning_rate}")
        try:
            rpc.rpc_sync(
                self.distributed_optimizer_rref.owner(),
                CustomLRScheduler.update_learning_rates,
                args=(self.distributed_optimizer_rref, self.learning_rate)
            )
        except Exception as e:
            print(f"Failed to update learning rates: {e}")

    def get_current_lr(self):
        return f"{self.learning_rate:.1e}"

    def _get_lr(self, current_step):
        if current_step &lt; self.num_warmup_steps:
            return float(current_step) / float(max(1, self.num_warmup_steps))
        progress = float(current_step - self.num_warmup_steps) / float(max(1, self.num_training_steps - self.num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(self.num_cycles) * 2.0 * progress)))

    @staticmethod
    def update_learning_rates(optimizer_rref, new_lr):
        optimizer_local = optimizer_rref.local_value()
        print(f"Worker updating learning rate to: {new_lr}")
        
        # Update the defaults dictionary
        optimizer_local.optim.defaults['lr'] = new_lr

def get_lr(lrs_rref):  
    return lrs_rref.rpc_sync().get_current_lr()  
  
def create_lr_scheduler(distributed_optimizer_rref, num_warmup_steps, num_training_steps, num_cycles=0.5):  
    scheduler = CustomLRScheduler(distributed_optimizer_rref, num_warmup_steps, num_training_steps, num_cycles)  
    return scheduler  
  
def lrs_step(lrs_rref):  
    lrs = lrs_rref.local_value()  
    lrs.step()
</code></pre>
<p>In the training loop, I am doing the following:</p>
<pre><code class="lang-auto">    lrs_rrefs = [  
        rpc.remote(  
            opt_rref.owner(),  
            create_lr_scheduler,  
            args=(opt_rref, num_warmup_steps, num_training_steps)  
        ) for opt_rref in optimizer.remote_optimizers  
    ] 

with dist_autograd.context() as context_id:
                #forward/backward steps.........
                # Step the learning rate scheduler  
                futs = []  
                for lrs_rref in lrs_rrefs:  
                    futs.append(rpc.rpc_async(lrs_rref.owner(), lrs_step, args=(lrs_rref,)))  
                [fut.wait() for fut in futs]  

                current_lrs = optimizer.get_learning_rates()
                print(f"Current learning rates: {current_lrs}")  
                logs = {"loss": loss.detach().item(),  "learning rates": current_lrs, "step": global_step}
</code></pre>