[Feature] Add GradScaler for ZeroOptim

[nijkah]

Add GradScaler for ZeroOptim

Numerical Precision is not tested yet.

Description

Test Script

import os

import torch, time, gc
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

from oslo.torch.utils import get_free_port, set_seed
from oslo.torch.distributed.parallel_context import ParallelContext
from oslo.torch.nn.parallel.data_parallel.zero import ZeroRedundancyOptimizer as OsloZeroRedundancyOptimizer
from oslo.torch.nn.parallel.data_parallel.grad_scaler import DynamicGradScaler
from torch.distributed.optim import ZeroRedundancyOptimizer

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_WORLD_SIZE"] = str(world_size)

def cleanup():
    dist.destroy_process_group()

def main_print(args):
    if dist.get_rank() != 0:
        return
    print(args)

# Timing utilities
start_time = None

def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.time()


def end_timer_and_print(local_msg):
    torch.cuda.synchronize()
    end_time = time.time()
    if dist.get_rank() != 0:
        return
    print("\n" + local_msg)
    print("Total execution time = {:.3f} sec".format(end_time - start_time))
    print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))

def make_model(in_size, out_size, num_layers):
    layers = []
    for _ in range(num_layers - 1):
        layers.append(torch.nn.Linear(in_size, in_size))
        layers.append(torch.nn.ReLU())
    layers.append(torch.nn.Linear(in_size, out_size))
    return torch.nn.Sequential(*tuple(layers))

def train(rank, world_size):
    print(f"Running oslo DDP example on rank {rank}.")
    setup(rank, world_size)
    parallel_context = ParallelContext.from_torch(data_parallel_size=world_size)

    batch_size = 512 # Try, for example, 128, 256, 513.
    in_size = 4096
    out_size = 4096
    num_layers = 3
    num_batches = 1
    epochs = 1
    use_zero = True
    use_oslo = True
    if use_oslo:
        use_zero = False

    local_rank = torch.distributed.get_rank()

    # Creates data in default precision.
    # The same data is used for both default and mixed precision trials below.
    # You don't need to manually change inputs' ``dtype`` when enabling mixed precision.
    set_seed(2021 + local_rank)
    data = [torch.randn(batch_size, in_size, device="cuda") for _ in range(num_batches)]
    targets = [torch.randn(batch_size, out_size, device="cuda") for _ in range(num_batches)]


    net = make_model(in_size, out_size, num_layers).to(rank)
    model = DDP(net, device_ids=[rank])
    loss_fn = torch.nn.MSELoss()
    if use_zero:
        opt = ZeroRedundancyOptimizer(
            model.parameters(),
            optimizer_class=torch.optim.Adam,
            lr=0.01
        )
    elif use_oslo:
        opt = OsloZeroRedundancyOptimizer(
            torch.optim.Adam(model.parameters(), lr=0.01),
            parallel_context=parallel_context,
            overlap_communication=True,
        )
    else:
        opt = torch.optim.Adam(model.parameters(), lr=0.01)

    if use_oslo:
        scaler = DynamicGradScaler(growth_interval=2000)
    else:
        scaler = torch.cuda.amp.GradScaler()
    start_timer()
    for epoch in range(epochs):
        for input, target in zip(data, targets):
            with torch.autocast(device_type='cuda', dtype=torch.float16):
                output = net(input)
                assert output.dtype is torch.float16
                loss = loss_fn(output, target)
                assert loss.dtype is torch.float32

            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
            opt.zero_grad() # set_to_none=True here can modestly improve performance
    print('rank:', rank, 'param:', net[0].weight)
    end_timer_and_print("Default precision:")


def main(world_size):
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main(2)

# with torch.cuda.amp.GradScaler
tensor([[-0.0017,  0.0104, -0.0024,  ..., -0.0023,  0.0195,  0.0028],
        [-0.0129,  0.0044,  0.0003,  ..., -0.0158, -0.0250,  0.0110],
        [-0.0020, -0.0109, -0.0240,  ..., -0.0088, -0.0186, -0.0056],
        ...,
        [-0.0038, -0.0212, -0.0238,  ...,  0.0218,  0.0037, -0.0159],
        [-0.0031, -0.0091, -0.0030,  ..., -0.0242,  0.0247, -0.0192],
        [-0.0030,  0.0172,  0.0144,  ..., -0.0131, -0.0018, -0.0049]],

# with oslo.torch.nn.parallel.data_parallel.grad_scaler.DynamicGradScaler
tensor([[-0.0015,  0.0104, -0.0024,  ..., -0.0023,  0.0197,  0.0028],
        [-0.0130,  0.0044,  0.0002,  ..., -0.0158, -0.0051, -0.0090],
        [-0.0020,  0.0086, -0.0041,  ..., -0.0088, -0.0187, -0.0056],
        ...,
        [-0.0037, -0.0014, -0.0238,  ...,  0.0218,  0.0035, -0.0160],
        [-0.0231, -0.0091, -0.0029,  ..., -0.0044,  0.0049,  0.0007],
        [-0.0029,  0.0173,  0.0146,  ...,  0.0067,  0.0181, -0.0049]],
       device='cuda:0', requires_grad=True)

[CLAassistant]

All committers have signed the CLA.

[nijkah]

Valid Test script

import os

import torch, time, gc
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader, DistributedSampler
import torch.nn as nn

from oslo.torch.utils import get_free_port, set_seed
from oslo.torch.distributed.parallel_context import ParallelContext
from oslo.torch.nn.parallel.data_parallel.zero import ZeroRedundancyOptimizer as OsloZeroRedundancyOptimizer
from oslo.torch.nn.parallel.data_parallel.grad_scaler import DynamicGradScaler
from torch.distributed.optim import ZeroRedundancyOptimizer

def setup(rank, world_size):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12345"
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["LOCAL_WORLD_SIZE"] = str(world_size)

def cleanup():
    dist.destroy_process_group()

def main_print(args):
    if dist.get_rank() != 0:
        return
    print(args)

# Timing utilities
start_time = None

def start_timer():
    global start_time
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.synchronize()
    start_time = time.time()


def end_timer_and_print(local_msg):
    torch.cuda.synchronize()
    end_time = time.time()
    if dist.get_rank() != 0:
        return
    print("\n" + local_msg)
    print("Total execution time = {:.3f} sec".format(end_time - start_time))
    print("Max memory used by tensors = {} bytes".format(torch.cuda.max_memory_allocated()))

def make_model(in_size, out_size, num_layers):
    layers = []
    for _ in range(num_layers - 1):
        layers.append(torch.nn.Linear(in_size, in_size))
        layers.append(torch.nn.ReLU())
    layers.append(torch.nn.Linear(in_size, out_size))
    return torch.nn.Sequential(*tuple(layers))

def train(rank, world_size):
    print(f"Running oslo DDP example on rank {rank}.")
    setup(rank, world_size)
    parallel_context = ParallelContext.from_torch(data_parallel_size=world_size)


    epochs = 10
    use_zero = True
    use_oslo = True
    if use_oslo:
        use_zero = False

    local_rank = torch.distributed.get_rank()

    # Define the data transformation
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    set_seed(42)
    # Load the CIFAR10 dataset
    train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    val_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=rank)
    train_loader = DataLoader(train_dataset, num_workers=4, batch_size=16, sampler=train_sampler)
    val_loader = DataLoader(val_dataset, num_workers=4, batch_size=64)

    model = models.resnet50(weights=models.resnet.ResNet50_Weights.IMAGENET1K_V2)
    model.fc = nn.Linear(2048, 10) # CIFAR10 has 10 classes


    net = model.to(rank)
    model = DDP(net, device_ids=[rank])
    loss_fn = nn.CrossEntropyLoss()
    if use_zero:
        opt = ZeroRedundancyOptimizer(
            model.parameters(),
            optimizer_class=torch.optim.AdamW,
            lr=1e-4
        )
    elif use_oslo:
        opt = OsloZeroRedundancyOptimizer(
            torch.optim.AdamW(model.parameters(), lr=1e-4),
            parallel_context=parallel_context,
            overlap_communication=True,
        )
    else:
        opt = torch.optim.AdamW(model.parameters(), lr=1e-4)

    if use_oslo:
        scaler = DynamicGradScaler(growth_interval=2000)
    else:
        scaler = torch.cuda.amp.GradScaler()
    start_timer()
    for epoch in range(epochs):
        for idx, (input, target) in enumerate(train_loader):
            opt.zero_grad() # set_to_none=True here can modestly improve performance
            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=True):
                output = net(input.half().to(rank))
                assert output.dtype is torch.float16
                loss = loss_fn(output, target.to(rank))
                if (idx % 20) == 0 and rank == 0:
                    print('epoch', epoch, 'idx:', idx, 'loss:', loss)
                assert loss.dtype is torch.float32
            scaler.scale(loss).backward()
            scaler.step(opt)
            scaler.update()
    # print('rank:', rank, 'param:', net[0].weight)
    end_timer_and_print("Default precision:")


def main(world_size):
    mp.spawn(train, args=(world_size,), nprocs=world_size, join=True)


if __name__ == "__main__":
    main(2)

[hyunwoongko]

Please modify this to non-draft pr when you are done :)