<jupyter_start><jupyter_text>Applying BlueFog on Deep Learning problem(High Level API Introduction)All previous sections we focused on the low-level API in the BlueFog, which is great for flexible algorithm design and research. Give a quick summary here:- Basic static topology propertis and how to manipulate it- Collective communication such as broadcast and allreduce- Topology based Neighborhood communication such as neighbor_allreduce.- Blocking versus non-blocking operation- Dynamic topology and its corresponding- Asynchronous operation through window object- Several decentralized algorithms and their performance under different scenarios- etc.However, it can be boilerplate if you want to apply one certain algorithm on different tasks. Further, it is also tricky to write a efficiency code combining above mentioned concepts. This becomes even worse in the deep learning problem. Backpropagation property of neural network makes that the gradient can be efficiently calculated. Backpropagation also implied that the gradient is calculated (approximately) in layer-by-layer style, in contrast to one global stochastic (sub-)gradient we encountered in the optimization. Further, this layer-wise computation provides a great opportunity to overlap the communication and comptuation for minimizing the trainning time, which is a crucial. Clearly, writing the code to address these concern correctly and efficient is not easy. Hence, BlueFog further provides the high level APIs, which can be directly applied on the `torch.Optimizer` directly. In this section, we will focus in applying high level APIs of BlueFog on Deep Learning problem, mainly the decentralized trainning task.Before we demystify how we implement the High-Level API in BlueFog, let's see the example of using them to write distributed trainning of ResNet-18 model over CIFAR-10 easily. (*Note, although this example is relative small, it still can be time-consuming and drain tons of computation resources if you want to train them on CPUs.*)<jupyter_code>import os
import time

import ipyparallel as ipp
import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
import torch
from torchvision import datasets, transforms, models

%matplotlib inline
rc = ipp.Client(profile="bluefog")
rc.block = True
rc.ids<jupyter_output><empty_output><jupyter_text>Prepared the data<jupyter_code># Down the CIFAR10 Dataset if not available
# Since the dataset is smaller enough, we just load it in-memory.
train_dataset = datasets.CIFAR10(
    os.path.join(os.getcwd(), "..", "data"),
    train=True,
    download=True,
    transform=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    ),
)

val_dataset = datasets.CIFAR10(
    os.path.join(os.getcwd(), "..", "data"),
    train=False,
    transform=transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    ),
)
# Distribute the data into each worker.
# Note we push the full dataset into each worker is just for simplicity.
# Each worker only read the partial of dataset later.
_ = rc[:].push({"train_dataset": train_dataset, "val_dataset": val_dataset})
%%px
import math
import torch
import bluefog.torch as bf
from bluefog.common import topology_util

seed = 2021
bf.init()
torch.manual_seed(seed)
run_on_cuda = torch.cuda.is_available()
if run_on_cuda:
    print("using cuda.")
    # Bluefog: pin GPU to local rank.
    device_id = (
        bf.local_rank()
        if bf.nccl_built()
        else bf.local_rank() % torch.cuda.device_count()
    )
    torch.cuda.set_device(device_id)
    torch.cuda.manual_seed(seed)
else:
    print("using cpu")
%%px
print(torch.__file__)
%%px
# Prepare the distributed loader for dataset.
batch_size = 32
val_batch_size = 1024
kwargs = {"num_workers": 4, "pin_memory": True} if run_on_cuda else {}

train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset, num_replicas=bf.size(), rank=bf.rank()
)
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, sampler=train_sampler, **kwargs
)

val_sampler = torch.utils.data.distributed.DistributedSampler(
    val_dataset, num_replicas=bf.size(), rank=bf.rank()
)
val_loader = torch.utils.data.DataLoader(
    val_dataset, batch_size=val_batch_size, sampler=val_sampler, **kwargs
)<jupyter_output><empty_output><jupyter_text>Set up standard torch optimizer<jupyter_code>%%px
from torchvision import models

model = models.resnet18(num_classes=10)
if run_on_cuda:
    model.cuda()

# Scale learning rate by the number of GPUs.
base_lr = 0.0125
momentum = 0.9
weight_decay = 0.00005

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=(base_lr * bf.size()),
    momentum=momentum,
    weight_decay=weight_decay,
)<jupyter_output><empty_output><jupyter_text>Wrap the torch standard optimizer into BlueFog distributed one<jupyter_code>%%px
atc_style = False
base_dist_optimizer = (
    bf.DistributedAdaptThenCombineOptimizer
    if atc_style
    else bf.DistributedAdaptWithCombineOptimizer
)
optimizer = base_dist_optimizer(
    optimizer, model=model, communication_type=bf.CommunicationType.neighbor_allreduce
)

# Bluefog: broadcast parameters & optimizer state.
bf.broadcast_parameters(model.state_dict(), root_rank=0)
bf.broadcast_optimizer_state(optimizer, root_rank=0)<jupyter_output><empty_output><jupyter_text>Define several helper functions for training and validation<jupyter_code>%%px
def accuracy(output, target):
    # get the index of the max log-probability
    pred = output.max(1, keepdim=True)[1]
    return pred.eq(target.view_as(pred)).cpu().float().mean()


class Metric(object):
    def __init__(self, name):
        self.name = name
        self.sum = torch.tensor(0.0)  # pylint: disable=not-callable
        self.n = torch.tensor(0.0)  # pylint: disable=not-callable

    def update(self, val):
        self.sum += bf.allreduce(val.detach().cpu(), name=self.name)
        self.n += 1

    @property
    def avg(self):
        return self.sum / self.n


dynamic_neighbor_allreduce_gen = topology_util.GetDynamicOnePeerSendRecvRanks(
    bf.load_topology(), bf.rank()
)


def dynamic_topology_update(epoch, batch_idx):
    send_neighbors, recv_neighbors = next(dynamic_neighbor_allreduce_gen)
    optimizer.send_neighbors = send_neighbors
    optimizer.neighbor_weights = {
        r: 1 / (len(recv_neighbors) + 1) for r in recv_neighbors
    }
    optimizer.self_weight = 1 / (len(recv_neighbors) + 1)


def adjust_learning_rate(epoch, batch_idx):
    if epoch < 5:  # warmup_epochs
        epoch += float(batch_idx + 1) / len(train_loader)
        lr_adj = 1.0 / bf.size() * (epoch * (bf.size() - 1) / 5 + 1)
    elif epoch < 30:
        lr_adj = 1.0
    elif epoch < 60:
        lr_adj = 1e-1
    elif epoch < 80:
        lr_adj = 1e-2
    else:
        lr_adj = 1e-3
    for param_group in optimizer.param_groups:
        param_group["lr"] = base_lr * bf.size() * lr_adj
%%px
# publish_data is useful for monitoring the status of workers
from ipyparallel.datapub import publish_data


def train(epoch):
    model.train()
    train_sampler.set_epoch(epoch)
    train_loss = Metric("train_loss")
    train_accuracy = Metric("train_accuracy")

    for batch_idx, (data, target) in enumerate(train_loader):
        adjust_learning_rate(epoch, batch_idx)
        dynamic_topology_update(epoch, batch_idx)

        if run_on_cuda:
            data, target = data.cuda(), target.cuda()
        optimizer.zero_grad()
        output = model(data)
        train_accuracy.update(accuracy(output, target))
        loss = torch.nn.functional.cross_entropy(output, target)
        train_loss.update(loss)
        # Average gradients among sub-batches
        loss.div_(math.ceil(float(len(data)) / batch_size))
        loss.backward()
        # Gradient is applied across all ranks
        optimizer.step()
        publish_data(
            dict(
                batch_idx=batch_idx,
                loss=train_loss.avg.item(),
                accuracy=100.0 * train_accuracy.avg.item(),
            )
        )


def validate(epoch):
    model.eval()
    val_loss = Metric("val_loss")
    val_accuracy = Metric("val_accuracy")

    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(val_loader):
            if run_on_cuda:
                data, target = data.cuda(), target.cuda()
            output = model(data)

            val_loss.update(torch.nn.functional.cross_entropy(output, target))
            val_accuracy.update(accuracy(output, target))
            publish_data(
                dict(
                    batch_idx=batch_idx,
                    loss=val_loss.avg.item(),
                    accuracy=100.0 * val_accuracy.avg.item(),
                )
            )<jupyter_output><empty_output><jupyter_text>A fake example of monitoring the status of workersThe trainning of deep learning model can be time consuming. In order to monitor the status of workers, you can use `ipyparallel.datapub.publish_data` to push the dataof the workers to the notebook. Following is a simple fake example to help you understandhow to use publish_data and [tqdm](https://tqdm.github.io/) library.<jupyter_code>%%px
# publish_data is useful for monitoring the status of workers
import time
from ipyparallel.datapub import publish_data
import numpy as np


def fake_training():
    for i in range(n_iter):
        loss, accuracy = np.random.rand(2)
        publish_data(dict(i=i, loss=loss, accuracy=accuracy))
        time.sleep(0.5)
from tqdm import tqdm

dview = rc[:]
dview.block = True
n_iter = 50
dview["n_iter"] = n_iter
ar = dview.apply_async(lambda: fake_training())
epoch = 1

with tqdm(total=n_iter - 1, desc="Train Epoch     #{}".format(epoch + 1)) as t:
    while not ar.ready():
        data = ar.data[0]
        t.set_postfix(
            {
                "loss": data.get("loss", 0),
                "accuracy": f"{data.get('accuracy', 0)*100:.2f}%",
            }
        )
        t.update(data.get("i", 0) - t.n)
        time.sleep(0.5)<jupyter_output>Train Epoch     #2: 100%|██████████| 49/49 [00:25<00:00,  1.95it/s, loss=0.846, accuracy=84.98%]<jupyter_text>Start decentralized trainning<jupyter_code>import time

# from tqdm.notebook import tqdm
from tqdm import tqdm

epochs = 2
for epoch in range(epochs):
    rc[:].push({"epoch": epochs}, block=True)

    # Training
    with tqdm(
        total=len(rc[0]["train_loader"]) - 1,
        desc="Train Epoch     #{}".format(epoch + 1),
    ) as t:
        ar_train = rc[:].apply_async(lambda: train(epoch))
        while not ar_train.ready() or t.n < t.total:
            data = ar_train.data[0]
            t.set_postfix(
                {
                    "loss": data.get("loss", 0),
                    "accuracy": f"{data.get('accuracy', 0):.2f}%",
                }
            )
            t.update(data.get("batch_idx", 0) - t.n)
            time.sleep(0.5)
    # Validation
    with tqdm(
        total=len(rc[0]["val_loader"]) - 1, desc="Validate Epoch  #{}".format(epoch + 1)
    ) as t:
        ar_val = rc[:].apply_async(lambda: validate(epoch))
        while not ar_val.ready() or t.n < t.total:
            data = ar_val.data[0]
            t.set_postfix(
                {
                    "loss": data.get("loss", 0),
                    "accuracy": f"{data.get('accuracy', 0):.2f}%",
                }
            )
            t.update(data.get("batch_idx", 0) - t.n)
            time.sleep(0.2)<jupyter_output>Train Epoch     #1: 100%|██████████| 390/390 [00:32<00:00, 12.12it/s, loss=1.72, accuracy=41.46%]
Validate Epoch  #1: 100%|██████████| 2/2 [00:01<00:00,  1.96it/s, loss=1.51, accuracy=52.60%]
Train Epoch     #2: 100%|██████████| 390/390 [00:32<00:00, 12.12it/s, loss=1.23, accuracy=57.31%]
Validate Epoch  #2: 100%|██████████| 2/2 [00:00<00:00,  2.40it/s, loss=1.75, accuracy=58.47%]<jupyter_text>Demystify the BlueFog High Level APIsHere we show more details about two common BlueFog high level APIs -- `DistributedAdaptWithCombineOptimizer` and `DistributedAdaptThenCombineOptimizer`. They are corresponding to ATC and AWC algorithm we introduced before. After understanding how these works, you should be able to quickly grasp main idea of other high level APIs or even build your own high level APIs with different decentralized algorithm.  From the algorithm part, there is nothing new happened under the these two distributed optimizer. But we update our notation slightly since in neural network it is common to denote the gradient layer by layer:\begin{align}    grad_{l}(w_{k,i})      =& \frac{\partial f_k(w_{k,i})}{\partial w_{k,i}[l]} \\     =& \frac{\partial f_k(w_{k,i})}{\partial w_{k,i}[l+1]}\frac{\partial w_{k,i}[l+1]}{\partial w_{k,i}[l]} \\     =& grad_{l+1}(w_{k,i}) \frac{\partial w_{k,i}[l+1]}{\partial w_{k,i}[l]}\end{align}where we use $l$ for $l-$ layer and $[l]$ means the parameter related to layer $l$. In this notation, it is closer to reflect how torch works. Except that, the AWC and ATC are applied in the same way. However, in the engineering part,it is not the same as we did before. The main difference is that we no longer wait for the gradient computation of whole model finished to start the communication. Instead, we try to maximize the communication and computation through layer-wise style. First, let's review the typical deep learning training under Torch framework:```pythonloss = metric(model(data), target)loss.backward()  <---  it starts to compute the gradientoptimizer.step() <---  it applies the computed gradient on the weights```The `backwards()` function will inversely traverse the neural networks and compute the gradients of each components in the model. The earliest time of executing communication is either at the *forward* or the *backward* step depending on the communication vector required gradient information or not. The latest time of finishing communication is at the `step()` step.A typical communication and computation timeline for Adapt-With-Combine optimizer and Adapt-Then-Combine optimizer is shown in the following figure:  To achieve that the key mechanism is `hook` pattern or so-called [Template method pattern](https://en.wikipedia.org/wiki/Template_method_pattern). For those who are not familar with this design pattern, we provide a simple example:<jupyter_code>class TriviaHookPattern:
    def __init__(self):
        pass

    def register_pre_hook(self, f):
        self._pre_hook = f

    def register_post_hook(self, f):
        self._post_hook = f

    def execute(self):
        self._pre_hook()
        print("Execute some function")
        self._post_hook()

    def _pre_hook(self):
        raise NotImplementedError

    def _post_hook(self):
        raise NotImplementedError
h = TriviaHookPattern()
h.register_pre_hook(lambda: print("I am pre hook"))
h.register_post_hook(lambda: print("I am post hook"))
h.execute()<jupyter_output>I am pre hook
Execute some function
I am post hook<jupyter_text>There are three types of hooks provided by torch APIs.- `register_forward_hook` for one `torch.nn` module- `register_backward_hook` for one `torch.nn` module- `register_hook()` in gradient computation of one `torch.Tensor`We implement distributed optimizers with some proper combination of these hooks. We would not list the full source code due to too many engineering details. But we demonstrate a one-layer neural network example to illustrate how to combine it with bluefog non-blocking communication operations.<jupyter_code>%%px
import torch.nn as nn


class LinearNet(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(LinearNet, self).__init__()
        self.fc = nn.Linear(input_dim, output_dim)
        self.num_parameters = input_dim * output_dim

    def forward(self, x):
        return self.fc(x)
%%px
# register for communication hook for model
data, target = torch.rand(5), torch.rand(1)
model = LinearNet(5, 1)

handles = []


def make_hook(model):
    def hook(model, *unused):
        for name, parameter in model.named_parameters():
            print("Make neighbor_allreduce for ", name)
            handle = bf.neighbor_allreduce_nonblocking(parameter.data)
            handles.append(handle)

    return hook


hook = make_hook(model)
_ = model.register_forward_hook(hook)
%%px
# Start nonblocking communication at first time.
# Wait until it finished at last step.
y = model(data)
loss = nn.MSELoss()(y, target)
loss.backward()
print(handles)
tensors = [bf.wait(handle) for handle in handles]
handles.clear()
print(tensors)<jupyter_output>[stdout:0] 
Make neighbor_allreduce for  fc.weight
Make neighbor_allreduce for  fc.bias
[50250, 50251]
[tensor([[ 0.0330, -0.3186,  0.2665,  0.2265, -0.3765]]), tensor([0.3806])]
[stdout:1] 
Make neighbor_allreduce for  fc.weight
Make neighbor_allreduce for  fc.bias
[50250, 50251]
[tensor([[ 0.0330, -0.3186,  0.2665,  0.2265, -0.3765]]), tensor([0.3806])]
[stdout:2] 
Make neighbor_allreduce for  fc.weight
Make neighbor_allreduce for  fc.bias
[50250, 50251]
[tensor([[ 0.0330, -0.3186,  0.2665,  0.2265, -0.3765]]), tensor([0.3806])]
[stdout:3] 
Make neighbor_allreduce for  fc.weight
Make neighbor_allreduce for  fc.bias
[50250, 50251]
[tensor([[ 0.0330, -0.3186,  0.2665,  0.2265, -0.3765]]), tensor([0.3806])]
{'fraction_non_alphanumeric': 0.10884392732274889, 'fraction_numerical': 0.02705336160944575, 'mean_word_length': 3.5202072538860105, 'pattern_counts': {'":': 11, '<': 46, '<?xml version=': 0, '>': 33, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '555991', 'license': 'Apache-2.0', 'n_tokens_mistral': 5860, 'n_tokens_neox': 5445, 'n_words': 1628, 'path': 'Section 7/Applying BlueFog on Deep Learning problem(High Level API Introduction).ipynb', 'repo_name': 'Bluefog-Lib/bluefog-tutorial', 'size': '17447'}