DeepSpeed Tutorial
#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
    Support Python 3.8
    @author: Lou Xiao(louxiao@i32n.com)
    @maintainer: Lou Xiao(louxiao@i32n.com)
    @copyright: Copyright 2018~2023
    @created time: 2023-09-05 15:37:05 CST
    @updated time: 2023-09-05 15:37:05 CST
"""

import time

import torch
import torch.nn as nn
import torch.nn.functional as tnf
import torch.utils.data as tud
import numpy as np
from loguru import logger as logging


# Your Dataset
class MyDataset(tud.Dataset):

    def __init__(self, image_shape: tuple, num_category: int, sample_count: int = 10000):
        self.image_shape = image_shape
        self.num_category = num_category
        self.sample_count = sample_count

    def __len__(self):
        return self.sample_count

    def __getitem__(self, index: int):
        xx = torch.randn(self.image_shape, dtype=torch.float32)
        yy = torch.randint(low=0, high=self.num_category, size=[1])
        return xx, yy


# Your Neural Network
class ConvBlock(nn.Module):

    def __init__(self, num_channels: int, layer_scale_init: float = 1e-6):
        super().__init__()
        self.residual = nn.Sequential(
            nn.GroupNorm(1, num_channels),  # LayerNorm
            nn.GELU(),
            nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(1, num_channels),  # LayerNorm
            nn.GELU(),
            nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1),
        )
        self.layer_scale = nn.Parameter(torch.tensor(layer_scale_init))

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        h = inputs + self.layer_scale * self.residual(inputs)
        return h


class MyClassifier(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, ch_multi: int = 32):
        super().__init__()

        self.stage1 = nn.Sequential(
            # downscale
            nn.Sequential(
                nn.Conv2d(in_channels, ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, ch_multi),  # LayerNorm
            ),
            ConvBlock(ch_multi),
            ConvBlock(ch_multi),
        )

        self.stage2 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(ch_multi, 2 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 2 * ch_multi),  # LayerNorm
            ),
            ConvBlock(2 * ch_multi),
            ConvBlock(2 * ch_multi),
        )

        self.stage3 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(2 * ch_multi, 4 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 4 * ch_multi),  # LayerNorm
            ),
            ConvBlock(4 * ch_multi),
            ConvBlock(4 * ch_multi),
        )

        self.stage4 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(4 * ch_multi, 8 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 8 * ch_multi),  # LayerNorm
            ),
            ConvBlock(8 * ch_multi),
            ConvBlock(8 * ch_multi),
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(1),
            nn.Linear(8 * ch_multi, out_channels),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        h = self.stage1(inputs)
        h = self.stage2(h)
        h = self.stage3(h)
        h = self.stage4(h)
        h = self.classifier(h)
        return h


def train():
    # hyper arguments

    default_device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    default_type = torch.float32

    # init model
    model = MyClassifier(3, 100, ch_multi=128)
    model.to(device=default_device, dtype=default_type)

    # init optimizer
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    # init dataset
    ds = MyDataset((3, 512, 512), 100, sample_count=int(1e6))
    dl = tud.DataLoader(ds, shuffle=True, batch_size=1, drop_last=True, num_workers=4, pin_memory=True)

    # train
    last_time = time.time()
    loss_list = []
    echo_interval = 10
    model.train()
    for step, (xx, yy) in enumerate(dl):
        step += 1
        xx = xx.to(device=default_device, dtype=default_type)
        yy = yy.to(device=default_device, dtype=torch.long).reshape(-1)

        optimizer.zero_grad()
        outputs = model(xx)
        loss = tnf.cross_entropy(outputs, yy)
        loss.backward()
        optimizer.step()
        loss_list.append(loss.detach().cpu().numpy())

        if step % echo_interval == 0:
            loss_avg = np.mean(loss_list[-echo_interval:])
            used_time = time.time() - last_time
            time_p_step = used_time / echo_interval
            logging.info(
                "[Train Step] Step:{:10d}  Loss:{:8.4f} | Time/Batch: {:6.4f}s",
                step, loss_avg, time_p_step,
            )
            last_time = time.time()


def main():
    train()


if __name__ == '__main__':
    main()

base_script.py
OvJat/DeepSpeedTutorial
#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
    Support Python 3.8
    @author: Lou Xiao(louxiao@i32n.com)
    @maintainer: Lou Xiao(louxiao@i32n.com)
    @copyright: Copyright 2018~2023
    @created time: 2023-09-07 16:51:48 CST
    @updated time: 2023-09-07 16:51:48 CST
"""

import json
import os.path

deepspeed_config = {
    # "train_batch_size": 128,
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 0.001,
            "betas": [
                0.8,
                0.999
            ],
            "eps": 1e-8,
            "weight_decay": 3e-7
        }
    },
    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 0.001,
            "warmup_num_steps": 1000
        }
    },
    "activation_checkpointing": {
        "partition_activations": True,
        "cpu_checkpointing": True,
        "contiguous_memory_optimization": False,
        "number_checkpoints": None,
        "synchronize_checkpoint_boundary": False,
        "profile": True,
    },
    "fp16": {
        "enabled": True,
        "auto_cast": False,
        "loss_scale": 0,
        "initial_scale_power": 16,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "consecutive_hysteresis": False,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            "pin_memory": True,
        },
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": True,
        },
        "contiguous_gradients": True,
        "overlap_comm": True,
    },
}

config_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), "deepspeed_config.json")
with open(config_file, 'w') as f:
    f.write(json.dumps(deepspeed_config, indent=4, ensure_ascii=False))

deepspeed_config.py
OvJat/DeepSpeedTutorial
#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
    Support Python 3.8
    @author: Lou Xiao(louxiao@i32n.com)
    @maintainer: Lou Xiao(louxiao@i32n.com)
    @copyright: Copyright 2018~2023
    @created time: 2023-09-05 15:37:05 CST
    @updated time: 2023-09-05 15:37:05 CST
"""

import time

import torch
import torch.nn as nn
import torch.nn.functional as tnf
import torch.utils.data as tud
import deepspeed
import numpy as np
from loguru import logger as logging


# Your Dataset
class MyDataset(tud.Dataset):

    def __init__(self, image_shape: tuple, num_category: int, sample_count: int = 10000):
        self.image_shape = image_shape
        self.num_category = num_category
        self.sample_count = sample_count

    def __len__(self):
        return self.sample_count

    def __getitem__(self, index: int):
        xx = torch.randn(self.image_shape, dtype=torch.float32)
        yy = torch.randint(low=0, high=self.num_category, size=[1])
        return xx, yy


# Your Neural Network
class ConvBlock(nn.Module):

    def __init__(self, num_channels: int, layer_scale_init: float = 1e-6):
        super().__init__()
        self.residual = nn.Sequential(
            nn.GroupNorm(1, num_channels),  # LayerNorm
            nn.GELU(),
            nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1),
            nn.GroupNorm(1, num_channels),  # LayerNorm
            nn.GELU(),
            nn.Conv2d(num_channels, num_channels, kernel_size=3, stride=1, padding=1),
        )
        self.layer_scale = nn.Parameter(torch.tensor(layer_scale_init))

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        h = inputs + self.layer_scale * self.residual(inputs)
        return h


class MyClassifier(nn.Module):

    def __init__(self, in_channels: int, out_channels: int, ch_multi: int = 32):
        super().__init__()

        self.stage1 = nn.Sequential(
            # downscale
            nn.Sequential(
                nn.Conv2d(in_channels, ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, ch_multi),  # LayerNorm
            ),
            ConvBlock(ch_multi),
            ConvBlock(ch_multi),
        )

        self.stage2 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(ch_multi, 2 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 2 * ch_multi),  # LayerNorm
            ),
            ConvBlock(2 * ch_multi),
            ConvBlock(2 * ch_multi),
        )

        self.stage3 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(2 * ch_multi, 4 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 4 * ch_multi),  # LayerNorm
            ),
            ConvBlock(4 * ch_multi),
            ConvBlock(4 * ch_multi),
        )

        self.stage4 = nn.Sequential(
            nn.Sequential(
                nn.Conv2d(4 * ch_multi, 8 * ch_multi, kernel_size=2, stride=2, padding=0),
                nn.GroupNorm(1, 8 * ch_multi),  # LayerNorm
            ),
            ConvBlock(8 * ch_multi),
            ConvBlock(8 * ch_multi),
        )

        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            nn.Flatten(1),
            nn.Linear(8 * ch_multi, out_channels),
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        h = self.stage1(inputs)
        h = self.stage2(h)
        h = self.stage3(h)
        h = self.stage4(h)
        h = self.classifier(h)
        return h


def parse_arguments():
    import argparse
    parser = argparse.ArgumentParser(description='deepspeed training script.')
    parser.add_argument('--local_rank', type=int, default=-1,
                        help='local rank passed from distributed launcher')
    # Include DeepSpeed configuration arguments
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    return args


def train():
    args = parse_arguments()

    # init distributed
    deepspeed.init_distributed()

    # init model
    model = MyClassifier(3, 100, ch_multi=128)

    # init dataset
    ds = MyDataset((3, 512, 512), 100, sample_count=int(1e6))

    # init engine
    engine, optimizer, training_dataloader, lr_scheduler = deepspeed.initialize(
        args=args,
        model=model,
        model_parameters=model.parameters(),
        training_data=ds,
        # config=deepspeed_config,
    )

    # load checkpoint
    engine.load_checkpoint("./data/checkpoints/MyClassifier/")

    # train
    last_time = time.time()
    loss_list = []
    echo_interval = 10

    engine.train()
    for step, (xx, yy) in enumerate(training_dataloader):
        step += 1
        xx = xx.to(device=engine.device, dtype=torch.float16)
        yy = yy.to(device=engine.device, dtype=torch.long).reshape(-1)

        outputs = engine(xx)
        loss = tnf.cross_entropy(outputs, yy)
        engine.backward(loss)
        engine.step()
        loss_list.append(loss.detach().cpu().numpy())

        if step % echo_interval == 0:
            loss_avg = np.mean(loss_list[-echo_interval:])
            used_time = time.time() - last_time
            time_p_step = used_time / echo_interval
            if args.local_rank == 0:
                logging.info(
                    "[Train Step] Step:{:10d}  Loss:{:8.4f} | Time/Batch: {:6.4f}s",
                    step, loss_avg, time_p_step,
                )
            last_time = time.time()
    # save checkpoint
    engine.save_checkpoint("./data/checkpoints/MyClassifier/")


def main():
    train()


if __name__ == '__main__':
    main()

deepspeed_script.py
OvJat/DeepSpeedTutorial
#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
    Support Python 3.8
    @author: Lou Xiao(louxiao@i32n.com)
    @maintainer: Lou Xiao(louxiao@i32n.com)
    @copyright: Copyright 2018~2023
    @created time: 2023-09-18 14:07:23 CST
    @updated time: 2023-09-18 14:07:23 CST
"""

import os.path
import random
import time

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as tnf
from torchvision import transforms
import deepspeed

from diffusers import StableDiffusionPipeline, DDPMScheduler
from loguru import logger as logging

SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))


def init_dataset(
        dataset_dir: str,
        tokenizer=None,
        resolution: int = 512,
        center_crop: bool = False,
        random_flip: bool = False,
):
    from datasets import load_dataset
    # Downloading and loading a dataset from the hub.
    dataset = load_dataset(dataset_dir)
    column_names = dataset["train"].column_names
    image_column, text_column = column_names
    print(column_names)

    # preprocess
    def text_processor(text_list):
        texts = []
        for text in text_list:
            if isinstance(text, str):
                texts.append(text)
            elif isinstance(text, (list, np.ndarray)):
                # take a random caption if there are multiple
                texts.append(random.choice(text))
            else:
                raise ValueError(
                    f"Caption column `{text}` should contain either strings or lists of strings."
                )
        r = tokenizer(
            texts,
            max_length=tokenizer.model_max_length,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        return r.input_ids

    train_transforms = transforms.Compose(
        [
            transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(resolution) if center_crop else transforms.RandomCrop(resolution),
            transforms.RandomHorizontalFlip() if random_flip else transforms.Lambda(lambda x: x),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )

    def custom_transform(samples):
        images = [image.convert("RGB") for image in samples[image_column]]
        samples["image"] = [train_transforms(image) for image in images]
        samples["text"] = text_processor(samples[text_column])
        return samples

    dataset = dataset['train'].with_transform(custom_transform)
    return dataset


def collate_fn(samples):
    image = torch.stack([sample["image"] for sample in samples])
    image = image.to(memory_format=torch.contiguous_format).float()
    input_ids = torch.stack([example["text"] for example in samples])
    return {"image": image, "text": input_ids}


class Diffusion(nn.Module):

    def __init__(self, model_path: str):
        super().__init__()

        # init components
        pipe = StableDiffusionPipeline.from_pretrained(
            model_path,
            local_files_only=True,
            use_auth_token=False,
            torch_dtype=torch.float16,
            use_safetensors=True,
        )
        self.text_encoder = pipe.text_encoder
        self.vae = pipe.vae
        self.vae.decoder = None
        # we're only training this
        self.unet = pipe.unet
        self.noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
        self.tokenizer = pipe.tokenizer

    def forward(self, images, texts):
        # train step
        with torch.no_grad():
            # Convert images to latent space
            latents = self.vae.encode(images).latent_dist.sample()
            latents = latents * self.vae.config.scaling_factor

            # process text
            text_tokens = self.text_encoder(texts)[0]

            noise = torch.randn_like(latents)
            # Sample a random timestep for each image
            batch_size = latents.shape[0]
            timesteps = torch.randint(0, self.noise_scheduler.config.num_train_timesteps, (batch_size,), device=latents.device)
            timesteps = timesteps.long()
            noisy_latents = self.noise_scheduler.add_noise(latents, noise, timesteps)

            # epsilon
            target = noise

        # Predict the noise residual and compute loss
        predications = self.unet(noisy_latents, timesteps, text_tokens).sample

        loss = tnf.mse_loss(predications, target, reduction="mean")
        return loss


def parse_arguments():
    import argparse
    parser = argparse.ArgumentParser(description='deepspeed training script.')
    parser.add_argument('--local_rank', type=int, default=-1,
                        help='local rank passed from distributed launcher')
    # Include DeepSpeed configuration arguments
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    return args


def main():
    args = parse_arguments()

    # init distributed
    deepspeed.init_distributed()

    # init model
    logging.debug("init model")
    model_dir = os.path.join(SCRIPT_PATH, 'data', 'stable-diffusion-2-1')
    model = Diffusion(model_dir)

    # init dataset
    logging.debug("init dataset")
    dataset_dir = os.path.join(SCRIPT_PATH, 'data', 'pokemon-blip-captions')
    dataset = init_dataset(
        dataset_dir,
        tokenizer=model.tokenizer,
        resolution=512,
        center_crop=False,
        random_flip=False,
    )

    # init engine
    logging.debug("init engine")
    engine, optimizer, training_dataloader, lr_scheduler = deepspeed.initialize(
        args=args,
        model=model,
        model_parameters=model.unet.parameters(),
        training_data=dataset,
        collate_fn=collate_fn,
        # config=deepspeed_config,
    )

    # load checkpoint
    engine.load_checkpoint("./data/SavedModels/FineTunedStableDiffusion/")

    # train
    last_time = time.time()
    loss_list = []
    echo_interval = 1

    engine.train()
    for step, samples in enumerate(training_dataloader):
        step += 1
        images = samples['image'].to(device=engine.device, dtype=torch.float16)
        texts = samples['text'].to(device=engine.device, dtype=torch.long)

        loss = engine(images, texts)
        engine.backward(loss)
        engine.step()
        loss_list.append(loss.detach().cpu().numpy())

        if step % echo_interval == 0:
            loss_avg = np.mean(loss_list[-echo_interval:])
            used_time = time.time() - last_time
            time_p_step = used_time / echo_interval
            if args.local_rank == 0:
                logging.info(
                    "[Train Step] Step:{:10d}  Loss:{:8.4f} | Time/Batch: {:6.4f}s",
                    step, loss_avg, time_p_step,
                )
            last_time = time.time()
    # save checkpoint
    engine.save_checkpoint("./data/SavedModels/FineTunedStableDiffusion/")


if __name__ == '__main__':
    main()

deepspeed_stable_diffusion.py
OvJat/DeepSpeedTutorial