# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: PySpark
#     language: python
#     name: pysparkkernel
# ---

# ## MNIST training and DeepSpeed ZeRO
# Maggy enables you to train with Microsoft's DeepSpeed ZeRO optimizer. Since DeepSpeed does not follow the common PyTorch programming model, Maggy is unable to provide full distribution transparency to the user. This means that if you want to use DeepSpeed for your training, you will have to make small changes to your code. In this notebook, we will show you what exactly you have to change in order to make DeepSpeed run with Maggy.

from hops import hdfs
import torch
import torch.nn.functional as F


# ### Define the model
# First off, we have to define our model. Since DeepSpeed's ZeRO is meant to reduce the memory consumption of our model, we will use an unreasonably large CNN for this example.

class CNN(torch.nn.Module):
    
    def __init__(self):
        super().__init__()
        self.l1 = torch.nn.Conv2d(1,1000,3)
        self.l2 = torch.nn.Conv2d(1000,3000,5)
        self.l3 = torch.nn.Conv2d(3000,3000,5)
        self.l4 = torch.nn.Linear(3000*18*18,10)
        
    def forward(self, x):
        x = F.relu(self.l1(x))
        x = F.relu(self.l2(x))
        x = F.relu(self.l3(x))
        x = F.softmax(self.l4(x.flatten(start_dim=1)), dim=0)
        return x


# ### Adapting the training function
# There are a few minor changes that have to be done in order to train with DeepSpeed:
# - There is no need for an optimizer anymore. You can configure your optimizer later in the DeepSpeed config.
# - DeepSpeed's ZeRO _requires_ you to use FP16 training. Therefore, convert your data to half precision!
# - The backward call is not executed on the loss, but on the model (`model.backward(loss)` instead of `loss.backward()`).
# - The step call is not executed on the optimizer, but also on the model (`model.step()` instead of `optimizer.step()`).
# - As we have no optimizer anymore, there is also no need to call `optimizer.zero_grad()`.
# You do not have to worry about the implementation of these calls, Maggy configures your model at runtime to act as a DeepSpeed engine.

def train_fn(module, hparams, train_set, test_set):
    
    import time
    import torch
        
    from maggy.core.patching import MaggyPetastormDataLoader
    
    model = module(**hparams)
    
    batch_size = 4
    lr_base = 0.1 * batch_size/256
    
    # Parameters as in https://arxiv.org/pdf/1706.02677.pdf
    loss_criterion = torch.nn.CrossEntropyLoss()

    train_loader = MaggyPetastormDataLoader(train_set, batch_size=batch_size)
                            
    model.train()
    for idx, data in enumerate(train_loader):
        img, label = data["image"].half(), data["label"].half()
        prediction = model(img)
        loss = loss_criterion(prediction, label.long())
        model.backward(loss)

        m1 = torch.cuda.max_memory_allocated(0)
        model.step()
        m2 = torch.cuda.max_memory_allocated(0)
        print("Optimizer pre: {}MB\n Optimizer post: {}MB".format(m1//1e6,m2//1e6))
        print(f"Finished batch {idx}")
    return float(1)


train_ds = hdfs.project_path() + "/DataSets/MNIST/PetastormMNIST/train_set"
test_ds = hdfs.project_path() + "/DataSets/MNIST/PetastormMNIST/test_set"
print(hdfs.exists(train_ds), hdfs.exists(test_ds))

# ### Configuring DeepSpeed
# In order to use DeepSpeed's ZeRO, the `deepspeed` backend has to be chosen. This backend also requires its own config. You can read a full specification of the possible settings [here](https://www.deepspeed.ai/docs/config-json/#zero-optimizations-for-fp16-training).

# +
from maggy import experiment
from maggy.experiment_config import TorchDistributedConfig

ds_config = {"train_micro_batch_size_per_gpu": 1,
 "gradient_accumulation_steps": 1,
 "optimizer": {"type": "Adam", "params": {"lr": 0.1}},
 "fp16": {"enabled": True},
 "zero_optimization": {"stage": 2},
}

config = TorchDistributedConfig(name='DS_ZeRO', module=CNN, train_set=train_ds, test_set=test_ds, backend="deepspeed", deepspeed_config=ds_config)
# -

# ### Starting the training
# You can now launch training with DS ZeRO. Note that the overhead of DeepSpeed is considerably larger than PyTorch's build in sharding, albeit also more efficient for a larger number of GPUs. DS will also jit compile components on the first run. If you want to compare memory efficiency with the default training, you can rewrite this notebook to work with standard PyTorch training.

result = experiment.lagom(train_fn, config)



StarcoderdataJupyterScriptsDedupFiltered