mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2026-01-25 14:55:10 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
Hongxin Liu
GitHub
@@ -1,7 +1,6 @@
|
||||
import argparse
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
import torchvision.transforms as transforms
|
||||
|
||||
@@ -9,15 +8,15 @@ import torchvision.transforms as transforms
|
||||
# Parse Arguments
|
||||
# ==============================
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-e', '--epoch', type=int, default=80, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory")
|
||||
parser.add_argument("-e", "--epoch", type=int, default=80, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument("-c", "--checkpoint", type=str, default="./checkpoint", help="checkpoint directory")
|
||||
args = parser.parse_args()
|
||||
|
||||
# ==============================
|
||||
# Prepare Test Dataset
|
||||
# ==============================
|
||||
# CIFAR-10 dataset
|
||||
test_dataset = torchvision.datasets.CIFAR10(root='./data/', train=False, transform=transforms.ToTensor())
|
||||
test_dataset = torchvision.datasets.CIFAR10(root="./data/", train=False, transform=transforms.ToTensor())
|
||||
|
||||
# Data loader
|
||||
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128, shuffle=False)
|
||||
@@ -26,7 +25,7 @@ test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=128,
|
||||
# Load Model
|
||||
# ==============================
|
||||
model = torchvision.models.resnet18(num_classes=10).cuda()
|
||||
state_dict = torch.load(f'{args.checkpoint}/model_{args.epoch}.pth')
|
||||
state_dict = torch.load(f"{args.checkpoint}/model_{args.epoch}.pth")
|
||||
model.load_state_dict(state_dict)
|
||||
|
||||
# ==============================
|
||||
@@ -45,4 +44,4 @@ with torch.no_grad():
|
||||
total += labels.size(0)
|
||||
correct += (predicted == labels).sum().item()
|
||||
|
||||
print('Accuracy of the model on the test images: {} %'.format(100 * correct / total))
|
||||
print("Accuracy of the model on the test images: {} %".format(100 * correct / total))
|
||||
|
||||
@@ -30,23 +30,19 @@ LEARNING_RATE = 1e-3
|
||||
def build_dataloader(batch_size: int, coordinator: DistCoordinator, plugin: DPPluginBase):
|
||||
# transform
|
||||
transform_train = transforms.Compose(
|
||||
[transforms.Pad(4),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.RandomCrop(32),
|
||||
transforms.ToTensor()])
|
||||
[transforms.Pad(4), transforms.RandomHorizontalFlip(), transforms.RandomCrop(32), transforms.ToTensor()]
|
||||
)
|
||||
transform_test = transforms.ToTensor()
|
||||
|
||||
# CIFAR-10 dataset
|
||||
data_path = os.environ.get('DATA', './data')
|
||||
data_path = os.environ.get("DATA", "./data")
|
||||
with coordinator.priority_execution():
|
||||
train_dataset = torchvision.datasets.CIFAR10(root=data_path,
|
||||
train=True,
|
||||
transform=transform_train,
|
||||
download=True)
|
||||
test_dataset = torchvision.datasets.CIFAR10(root=data_path,
|
||||
train=False,
|
||||
transform=transform_test,
|
||||
download=True)
|
||||
train_dataset = torchvision.datasets.CIFAR10(
|
||||
root=data_path, train=True, transform=transform_train, download=True
|
||||
)
|
||||
test_dataset = torchvision.datasets.CIFAR10(
|
||||
root=data_path, train=False, transform=transform_test, download=True
|
||||
)
|
||||
|
||||
# Data loader
|
||||
train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
|
||||
@@ -70,14 +66,21 @@ def evaluate(model: nn.Module, test_dataloader: DataLoader, coordinator: DistCoo
|
||||
dist.all_reduce(total)
|
||||
accuracy = correct.item() / total.item()
|
||||
if coordinator.is_master():
|
||||
print(f'Accuracy of the model on the test images: {accuracy * 100:.2f} %')
|
||||
print(f"Accuracy of the model on the test images: {accuracy * 100:.2f} %")
|
||||
return accuracy
|
||||
|
||||
|
||||
def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: nn.Module, train_dataloader: DataLoader,
|
||||
booster: Booster, coordinator: DistCoordinator):
|
||||
def train_epoch(
|
||||
epoch: int,
|
||||
model: nn.Module,
|
||||
optimizer: Optimizer,
|
||||
criterion: nn.Module,
|
||||
train_dataloader: DataLoader,
|
||||
booster: Booster,
|
||||
coordinator: DistCoordinator,
|
||||
):
|
||||
model.train()
|
||||
with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar:
|
||||
with tqdm(train_dataloader, desc=f"Epoch [{epoch + 1}/{NUM_EPOCHS}]", disable=not coordinator.is_master()) as pbar:
|
||||
for images, labels in pbar:
|
||||
images = images.cuda()
|
||||
labels = labels.cuda()
|
||||
@@ -91,7 +94,7 @@ def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: n
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Print log info
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
pbar.set_postfix({"loss": loss.item()})
|
||||
|
||||
|
||||
def main():
|
||||
@@ -100,19 +103,20 @@ def main():
|
||||
# ==============================
|
||||
parser = argparse.ArgumentParser()
|
||||
# FIXME(ver217): gemini is not supported resnet now
|
||||
parser.add_argument('-p',
|
||||
'--plugin',
|
||||
type=str,
|
||||
default='torch_ddp',
|
||||
choices=['torch_ddp', 'torch_ddp_fp16', 'low_level_zero'],
|
||||
help="plugin to use")
|
||||
parser.add_argument('-r', '--resume', type=int, default=-1, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory")
|
||||
parser.add_argument('-i', '--interval', type=int, default=5, help="interval of saving checkpoint")
|
||||
parser.add_argument('--target_acc',
|
||||
type=float,
|
||||
default=None,
|
||||
help="target accuracy. Raise exception if not reached")
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--plugin",
|
||||
type=str,
|
||||
default="torch_ddp",
|
||||
choices=["torch_ddp", "torch_ddp_fp16", "low_level_zero"],
|
||||
help="plugin to use",
|
||||
)
|
||||
parser.add_argument("-r", "--resume", type=int, default=-1, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument("-c", "--checkpoint", type=str, default="./checkpoint", help="checkpoint directory")
|
||||
parser.add_argument("-i", "--interval", type=int, default=5, help="interval of saving checkpoint")
|
||||
parser.add_argument(
|
||||
"--target_acc", type=float, default=None, help="target accuracy. Raise exception if not reached"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# ==============================
|
||||
@@ -136,13 +140,13 @@ def main():
|
||||
# Instantiate Plugin and Booster
|
||||
# ==============================
|
||||
booster_kwargs = {}
|
||||
if args.plugin == 'torch_ddp_fp16':
|
||||
booster_kwargs['mixed_precision'] = 'fp16'
|
||||
if args.plugin.startswith('torch_ddp'):
|
||||
if args.plugin == "torch_ddp_fp16":
|
||||
booster_kwargs["mixed_precision"] = "fp16"
|
||||
if args.plugin.startswith("torch_ddp"):
|
||||
plugin = TorchDDPPlugin()
|
||||
elif args.plugin == 'gemini':
|
||||
plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == 'low_level_zero':
|
||||
elif args.plugin == "gemini":
|
||||
plugin = GeminiPlugin(placement_policy="cuda", strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == "low_level_zero":
|
||||
plugin = LowLevelZeroPlugin(initial_scale=2**5)
|
||||
|
||||
booster = Booster(plugin=plugin, **booster_kwargs)
|
||||
@@ -168,18 +172,17 @@ def main():
|
||||
# ==============================
|
||||
# Boost with ColossalAI
|
||||
# ==============================
|
||||
model, optimizer, criterion, _, lr_scheduler = booster.boost(model,
|
||||
optimizer,
|
||||
criterion=criterion,
|
||||
lr_scheduler=lr_scheduler)
|
||||
model, optimizer, criterion, _, lr_scheduler = booster.boost(
|
||||
model, optimizer, criterion=criterion, lr_scheduler=lr_scheduler
|
||||
)
|
||||
|
||||
# ==============================
|
||||
# Resume from checkpoint
|
||||
# ==============================
|
||||
if args.resume >= 0:
|
||||
booster.load_model(model, f'{args.checkpoint}/model_{args.resume}.pth')
|
||||
booster.load_optimizer(optimizer, f'{args.checkpoint}/optimizer_{args.resume}.pth')
|
||||
booster.load_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{args.resume}.pth')
|
||||
booster.load_model(model, f"{args.checkpoint}/model_{args.resume}.pth")
|
||||
booster.load_optimizer(optimizer, f"{args.checkpoint}/optimizer_{args.resume}.pth")
|
||||
booster.load_lr_scheduler(lr_scheduler, f"{args.checkpoint}/lr_scheduler_{args.resume}.pth")
|
||||
|
||||
# ==============================
|
||||
# Train model
|
||||
@@ -191,14 +194,14 @@ def main():
|
||||
|
||||
# save checkpoint
|
||||
if args.interval > 0 and (epoch + 1) % args.interval == 0:
|
||||
booster.save_model(model, f'{args.checkpoint}/model_{epoch + 1}.pth')
|
||||
booster.save_optimizer(optimizer, f'{args.checkpoint}/optimizer_{epoch + 1}.pth')
|
||||
booster.save_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{epoch + 1}.pth')
|
||||
booster.save_model(model, f"{args.checkpoint}/model_{epoch + 1}.pth")
|
||||
booster.save_optimizer(optimizer, f"{args.checkpoint}/optimizer_{epoch + 1}.pth")
|
||||
booster.save_lr_scheduler(lr_scheduler, f"{args.checkpoint}/lr_scheduler_{epoch + 1}.pth")
|
||||
|
||||
accuracy = evaluate(model, test_dataloader, coordinator)
|
||||
if args.target_acc is not None:
|
||||
assert accuracy >= args.target_acc, f'Accuracy {accuracy} is lower than target accuracy {args.target_acc}'
|
||||
assert accuracy >= args.target_acc, f"Accuracy {accuracy} is lower than target accuracy {args.target_acc}"
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -32,35 +32,37 @@ LEARNING_RATE = 1e-3
|
||||
def vit_cifar(**kwargs):
|
||||
pretrained_cfg = _cfg(num_classes=10, input_size=(3, 32, 32), crop_pct=1.0)
|
||||
model_kwargs = dict(patch_size=4, embed_dim=512, depth=6, num_heads=8, drop_rate=0.1, mlp_ratio=1.0, **kwargs)
|
||||
model = _create_vision_transformer('vit_cifar', pretrained_cfg=pretrained_cfg, **model_kwargs)
|
||||
model = _create_vision_transformer("vit_cifar", pretrained_cfg=pretrained_cfg, **model_kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def build_dataloader(batch_size: int, coordinator: DistCoordinator, plugin: DPPluginBase):
|
||||
# transform
|
||||
transform_train = transforms.Compose([
|
||||
transforms.RandomCrop(32, padding=4),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
|
||||
])
|
||||
transform_test = transforms.Compose([
|
||||
transforms.Resize(32),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
|
||||
])
|
||||
transform_train = transforms.Compose(
|
||||
[
|
||||
transforms.RandomCrop(32, padding=4),
|
||||
transforms.RandomHorizontalFlip(),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
|
||||
]
|
||||
)
|
||||
transform_test = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(32),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize((0.49139968, 0.48215827, 0.44653124), (0.24703233, 0.24348505, 0.26158768)),
|
||||
]
|
||||
)
|
||||
|
||||
# CIFAR-10 dataset
|
||||
data_path = os.environ.get('DATA', './data')
|
||||
data_path = os.environ.get("DATA", "./data")
|
||||
with coordinator.priority_execution():
|
||||
train_dataset = torchvision.datasets.CIFAR10(root=data_path,
|
||||
train=True,
|
||||
transform=transform_train,
|
||||
download=True)
|
||||
test_dataset = torchvision.datasets.CIFAR10(root=data_path,
|
||||
train=False,
|
||||
transform=transform_test,
|
||||
download=True)
|
||||
train_dataset = torchvision.datasets.CIFAR10(
|
||||
root=data_path, train=True, transform=transform_train, download=True
|
||||
)
|
||||
test_dataset = torchvision.datasets.CIFAR10(
|
||||
root=data_path, train=False, transform=transform_test, download=True
|
||||
)
|
||||
|
||||
# Data loader
|
||||
train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
|
||||
@@ -84,14 +86,21 @@ def evaluate(model: nn.Module, test_dataloader: DataLoader, coordinator: DistCoo
|
||||
dist.all_reduce(total)
|
||||
accuracy = correct.item() / total.item()
|
||||
if coordinator.is_master():
|
||||
print(f'Accuracy of the model on the test images: {accuracy * 100:.2f} %')
|
||||
print(f"Accuracy of the model on the test images: {accuracy * 100:.2f} %")
|
||||
return accuracy
|
||||
|
||||
|
||||
def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: nn.Module, train_dataloader: DataLoader,
|
||||
booster: Booster, coordinator: DistCoordinator):
|
||||
def train_epoch(
|
||||
epoch: int,
|
||||
model: nn.Module,
|
||||
optimizer: Optimizer,
|
||||
criterion: nn.Module,
|
||||
train_dataloader: DataLoader,
|
||||
booster: Booster,
|
||||
coordinator: DistCoordinator,
|
||||
):
|
||||
model.train()
|
||||
with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar:
|
||||
with tqdm(train_dataloader, desc=f"Epoch [{epoch + 1}/{NUM_EPOCHS}]", disable=not coordinator.is_master()) as pbar:
|
||||
for images, labels in pbar:
|
||||
images = images.cuda()
|
||||
labels = labels.cuda()
|
||||
@@ -105,7 +114,7 @@ def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, criterion: n
|
||||
optimizer.zero_grad()
|
||||
|
||||
# Print log info
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
pbar.set_postfix({"loss": loss.item()})
|
||||
|
||||
|
||||
def main():
|
||||
@@ -114,19 +123,20 @@ def main():
|
||||
# ==============================
|
||||
parser = argparse.ArgumentParser()
|
||||
# FIXME(ver217): gemini is not supported resnet now
|
||||
parser.add_argument('-p',
|
||||
'--plugin',
|
||||
type=str,
|
||||
default='torch_ddp',
|
||||
choices=['torch_ddp', 'torch_ddp_fp16', 'low_level_zero'],
|
||||
help="plugin to use")
|
||||
parser.add_argument('-r', '--resume', type=int, default=-1, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument('-c', '--checkpoint', type=str, default='./checkpoint', help="checkpoint directory")
|
||||
parser.add_argument('-i', '--interval', type=int, default=5, help="interval of saving checkpoint")
|
||||
parser.add_argument('--target_acc',
|
||||
type=float,
|
||||
default=None,
|
||||
help="target accuracy. Raise exception if not reached")
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--plugin",
|
||||
type=str,
|
||||
default="torch_ddp",
|
||||
choices=["torch_ddp", "torch_ddp_fp16", "low_level_zero"],
|
||||
help="plugin to use",
|
||||
)
|
||||
parser.add_argument("-r", "--resume", type=int, default=-1, help="resume from the epoch's checkpoint")
|
||||
parser.add_argument("-c", "--checkpoint", type=str, default="./checkpoint", help="checkpoint directory")
|
||||
parser.add_argument("-i", "--interval", type=int, default=5, help="interval of saving checkpoint")
|
||||
parser.add_argument(
|
||||
"--target_acc", type=float, default=None, help="target accuracy. Raise exception if not reached"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# ==============================
|
||||
@@ -150,13 +160,13 @@ def main():
|
||||
# Instantiate Plugin and Booster
|
||||
# ==============================
|
||||
booster_kwargs = {}
|
||||
if args.plugin == 'torch_ddp_fp16':
|
||||
booster_kwargs['mixed_precision'] = 'fp16'
|
||||
if args.plugin.startswith('torch_ddp'):
|
||||
if args.plugin == "torch_ddp_fp16":
|
||||
booster_kwargs["mixed_precision"] = "fp16"
|
||||
if args.plugin.startswith("torch_ddp"):
|
||||
plugin = TorchDDPPlugin()
|
||||
elif args.plugin == 'gemini':
|
||||
plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == 'low_level_zero':
|
||||
elif args.plugin == "gemini":
|
||||
plugin = GeminiPlugin(placement_policy="cuda", strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == "low_level_zero":
|
||||
plugin = LowLevelZeroPlugin(initial_scale=2**5)
|
||||
|
||||
booster = Booster(plugin=plugin, **booster_kwargs)
|
||||
@@ -182,19 +192,17 @@ def main():
|
||||
# ==============================
|
||||
# Boost with ColossalAI
|
||||
# ==============================
|
||||
model, optimizer, criterion, train_dataloader, lr_scheduler = booster.boost(model,
|
||||
optimizer,
|
||||
criterion=criterion,
|
||||
dataloader=train_dataloader,
|
||||
lr_scheduler=lr_scheduler)
|
||||
model, optimizer, criterion, train_dataloader, lr_scheduler = booster.boost(
|
||||
model, optimizer, criterion=criterion, dataloader=train_dataloader, lr_scheduler=lr_scheduler
|
||||
)
|
||||
|
||||
# ==============================
|
||||
# Resume from checkpoint
|
||||
# ==============================
|
||||
if args.resume >= 0:
|
||||
booster.load_model(model, f'{args.checkpoint}/model_{args.resume}.pth')
|
||||
booster.load_optimizer(optimizer, f'{args.checkpoint}/optimizer_{args.resume}.pth')
|
||||
booster.load_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{args.resume}.pth')
|
||||
booster.load_model(model, f"{args.checkpoint}/model_{args.resume}.pth")
|
||||
booster.load_optimizer(optimizer, f"{args.checkpoint}/optimizer_{args.resume}.pth")
|
||||
booster.load_lr_scheduler(lr_scheduler, f"{args.checkpoint}/lr_scheduler_{args.resume}.pth")
|
||||
|
||||
# ==============================
|
||||
# Train model
|
||||
@@ -206,14 +214,14 @@ def main():
|
||||
|
||||
# save checkpoint
|
||||
if args.interval > 0 and (epoch + 1) % args.interval == 0:
|
||||
booster.save_model(model, f'{args.checkpoint}/model_{epoch + 1}.pth')
|
||||
booster.save_optimizer(optimizer, f'{args.checkpoint}/optimizer_{epoch + 1}.pth')
|
||||
booster.save_lr_scheduler(lr_scheduler, f'{args.checkpoint}/lr_scheduler_{epoch + 1}.pth')
|
||||
booster.save_model(model, f"{args.checkpoint}/model_{epoch + 1}.pth")
|
||||
booster.save_optimizer(optimizer, f"{args.checkpoint}/optimizer_{epoch + 1}.pth")
|
||||
booster.save_lr_scheduler(lr_scheduler, f"{args.checkpoint}/lr_scheduler_{epoch + 1}.pth")
|
||||
|
||||
accuracy = evaluate(model, test_dataloader, coordinator)
|
||||
if args.target_acc is not None:
|
||||
assert accuracy >= args.target_acc, f'Accuracy {accuracy} is lower than target accuracy {args.target_acc}'
|
||||
assert accuracy >= args.target_acc, f"Accuracy {accuracy} is lower than target accuracy {args.target_acc}"
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
@@ -5,7 +5,6 @@ from colossalai.booster.plugin.dp_plugin_base import DPPluginBase
|
||||
|
||||
|
||||
class GLUEDataBuilder:
|
||||
|
||||
task_text_field_map = {
|
||||
"cola": ["sentence"],
|
||||
"sst2": ["sentence"],
|
||||
@@ -84,10 +83,9 @@ class GLUEDataBuilder:
|
||||
AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)
|
||||
|
||||
def train_dataloader(self):
|
||||
return self.plugin.prepare_dataloader(self.dataset["train"],
|
||||
batch_size=self.train_batch_size,
|
||||
shuffle=True,
|
||||
drop_last=True)
|
||||
return self.plugin.prepare_dataloader(
|
||||
self.dataset["train"], batch_size=self.train_batch_size, shuffle=True, drop_last=True
|
||||
)
|
||||
|
||||
def val_dataloader(self):
|
||||
if len(self.eval_splits) == 1:
|
||||
@@ -108,7 +106,6 @@ class GLUEDataBuilder:
|
||||
]
|
||||
|
||||
def convert_to_features(self, example_batch):
|
||||
|
||||
# Either encode single sentence or sentence pairs
|
||||
if len(self.text_fields) > 1:
|
||||
texts_or_text_pairs = list(zip(example_batch[self.text_fields[0]], example_batch[self.text_fields[1]]))
|
||||
@@ -116,10 +113,9 @@ class GLUEDataBuilder:
|
||||
texts_or_text_pairs = example_batch[self.text_fields[0]]
|
||||
|
||||
# Tokenize the text/text pairs
|
||||
features = self.tokenizer.batch_encode_plus(texts_or_text_pairs,
|
||||
max_length=self.max_seq_length,
|
||||
padding='max_length',
|
||||
truncation=True)
|
||||
features = self.tokenizer.batch_encode_plus(
|
||||
texts_or_text_pairs, max_length=self.max_seq_length, padding="max_length", truncation=True
|
||||
)
|
||||
|
||||
# Rename label to labels to make it easier to pass to model forward
|
||||
features["labels"] = example_batch["label"]
|
||||
|
||||
@@ -33,8 +33,14 @@ def move_to_cuda(batch):
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoader]], num_labels: int, task_name: str,
|
||||
eval_splits: List[str], coordinator: DistCoordinator):
|
||||
def evaluate(
|
||||
model: nn.Module,
|
||||
test_dataloader: Union[DataLoader, List[DataLoader]],
|
||||
num_labels: int,
|
||||
task_name: str,
|
||||
eval_splits: List[str],
|
||||
coordinator: DistCoordinator,
|
||||
):
|
||||
metric = datasets.load_metric("glue", task_name, process_id=coordinator.rank, num_process=coordinator.world_size)
|
||||
model.eval()
|
||||
|
||||
@@ -58,7 +64,7 @@ def evaluate(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoade
|
||||
results = metric.compute()
|
||||
dist.all_reduce(accum_loss.div_(len(dataloader)))
|
||||
if coordinator.is_master():
|
||||
results['loss'] = accum_loss.item() / coordinator.world_size
|
||||
results["loss"] = accum_loss.item() / coordinator.world_size
|
||||
return results
|
||||
|
||||
if isinstance(test_dataloader, DataLoader):
|
||||
@@ -68,14 +74,21 @@ def evaluate(model: nn.Module, test_dataloader: Union[DataLoader, List[DataLoade
|
||||
final_results = {}
|
||||
for split, sub_loader in zip(eval_splits, test_dataloader):
|
||||
results = evaluate_subset(sub_loader)
|
||||
final_results.update({f'{k}_{split}': v for k, v in results.items()})
|
||||
final_results.update({f"{k}_{split}": v for k, v in results.items()})
|
||||
return final_results
|
||||
|
||||
|
||||
def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, lr_scheduler, train_dataloader: DataLoader,
|
||||
booster: Booster, coordinator: DistCoordinator):
|
||||
def train_epoch(
|
||||
epoch: int,
|
||||
model: nn.Module,
|
||||
optimizer: Optimizer,
|
||||
lr_scheduler,
|
||||
train_dataloader: DataLoader,
|
||||
booster: Booster,
|
||||
coordinator: DistCoordinator,
|
||||
):
|
||||
model.train()
|
||||
with tqdm(train_dataloader, desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]', disable=not coordinator.is_master()) as pbar:
|
||||
with tqdm(train_dataloader, desc=f"Epoch [{epoch + 1}/{NUM_EPOCHS}]", disable=not coordinator.is_master()) as pbar:
|
||||
for batch in pbar:
|
||||
# Forward pass
|
||||
batch = move_to_cuda(batch)
|
||||
@@ -89,7 +102,7 @@ def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, lr_scheduler
|
||||
lr_scheduler.step()
|
||||
|
||||
# Print log info
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
pbar.set_postfix({"loss": loss.item()})
|
||||
|
||||
|
||||
def main():
|
||||
@@ -97,14 +110,16 @@ def main():
|
||||
# Parse Arguments
|
||||
# ==============================
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('-t', '--task', default='mrpc', help="GLUE task to run")
|
||||
parser.add_argument('-p',
|
||||
'--plugin',
|
||||
type=str,
|
||||
default='torch_ddp',
|
||||
choices=['torch_ddp', 'torch_ddp_fp16', 'gemini', 'low_level_zero'],
|
||||
help="plugin to use")
|
||||
parser.add_argument('--target_f1', type=float, default=None, help="target f1 score. Raise exception if not reached")
|
||||
parser.add_argument("-t", "--task", default="mrpc", help="GLUE task to run")
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--plugin",
|
||||
type=str,
|
||||
default="torch_ddp",
|
||||
choices=["torch_ddp", "torch_ddp_fp16", "gemini", "low_level_zero"],
|
||||
help="plugin to use",
|
||||
)
|
||||
parser.add_argument("--target_f1", type=float, default=None, help="target f1 score. Raise exception if not reached")
|
||||
args = parser.parse_args()
|
||||
|
||||
# ==============================
|
||||
@@ -115,19 +130,19 @@ def main():
|
||||
|
||||
# local_batch_size = BATCH_SIZE // coordinator.world_size
|
||||
lr = LEARNING_RATE * coordinator.world_size
|
||||
model_name = 'bert-base-uncased'
|
||||
model_name = "bert-base-uncased"
|
||||
|
||||
# ==============================
|
||||
# Instantiate Plugin and Booster
|
||||
# ==============================
|
||||
booster_kwargs = {}
|
||||
if args.plugin == 'torch_ddp_fp16':
|
||||
booster_kwargs['mixed_precision'] = 'fp16'
|
||||
if args.plugin.startswith('torch_ddp'):
|
||||
if args.plugin == "torch_ddp_fp16":
|
||||
booster_kwargs["mixed_precision"] = "fp16"
|
||||
if args.plugin.startswith("torch_ddp"):
|
||||
plugin = TorchDDPPlugin()
|
||||
elif args.plugin == 'gemini':
|
||||
plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == 'low_level_zero':
|
||||
elif args.plugin == "gemini":
|
||||
plugin = GeminiPlugin(placement_policy="cuda", strict_ddp_mode=True, initial_scale=2**5)
|
||||
elif args.plugin == "low_level_zero":
|
||||
plugin = LowLevelZeroPlugin(initial_scale=2**5)
|
||||
|
||||
booster = Booster(plugin=plugin, **booster_kwargs)
|
||||
@@ -135,11 +150,9 @@ def main():
|
||||
# ==============================
|
||||
# Prepare Dataloader
|
||||
# ==============================
|
||||
data_builder = GLUEDataBuilder(model_name,
|
||||
plugin,
|
||||
args.task,
|
||||
train_batch_size=BATCH_SIZE,
|
||||
eval_batch_size=BATCH_SIZE)
|
||||
data_builder = GLUEDataBuilder(
|
||||
model_name, plugin, args.task, train_batch_size=BATCH_SIZE, eval_batch_size=BATCH_SIZE
|
||||
)
|
||||
train_dataloader = data_builder.train_dataloader()
|
||||
test_dataloader = data_builder.test_dataloader()
|
||||
|
||||
@@ -185,14 +198,15 @@ def main():
|
||||
for epoch in range(NUM_EPOCHS):
|
||||
train_epoch(epoch, model, optimizer, lr_scheduler, train_dataloader, booster, coordinator)
|
||||
|
||||
results = evaluate(model, test_dataloader, data_builder.num_labels, args.task, data_builder.eval_splits,
|
||||
coordinator)
|
||||
results = evaluate(
|
||||
model, test_dataloader, data_builder.num_labels, args.task, data_builder.eval_splits, coordinator
|
||||
)
|
||||
|
||||
if coordinator.is_master():
|
||||
print(results)
|
||||
if args.target_f1 is not None and 'f1' in results:
|
||||
assert results['f1'] >= args.target_f1, f'f1 score {results["f1"]} is lower than target {args.target_f1}'
|
||||
if args.target_f1 is not None and "f1" in results:
|
||||
assert results["f1"] >= args.target_f1, f'f1 score {results["f1"]} is lower than target {args.target_f1}'
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user