oslo
oslo copied to clipboard
EleutherAI
[Feature] Add GradScaler for ZeroOptim
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)
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)
