```import math
import torch
from torch.optim.optimizer import Optimizer
import torch.distributed as dist

"""Implements AdamW algorithm with distributed settings.

It has been proposed in `Fixing Weight Decay Regularization in Adam`.

Arguments:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay (not L2 penalty) (default: 0)
"""

def __init__(self, params, local_rank, world_size, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay=0):
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay)
self.local_rank = local_rank
self.world_size = world_size

[docs]    def step(self, closure=None):
"""Performs a single optimization step.

Arguments:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
loss = closure()

for group in self.param_groups:
for p in group['params']:
continue

for group in self.param_groups:
for p in group['params']:
continue
state = self.state[p]

# State initialization
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
# Exponential moving average of squared gradient values

exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
beta1, beta2 = group['betas']

state['step'] += 1

# Decay the first and second moment running average coefficient

bias_correction1 = 1 - beta1 ** state['step']
bias_correction2 = 1 - beta2 ** state['step']
step_size = group['lr'] * \
math.sqrt(bias_correction2) / bias_correction1

if group['weight_decay'] != 0: