mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-22 18:09:06 +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:
@@ -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"]
|
||||
|
@@ -1,5 +1,4 @@
|
||||
import argparse
|
||||
from contextlib import nullcontext
|
||||
from typing import Callable, List, Union
|
||||
|
||||
import evaluate
|
||||
@@ -7,7 +6,7 @@ import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn as nn
|
||||
from data import GLUEDataBuilder
|
||||
from torch.optim import Adam, Optimizer
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
@@ -17,7 +16,6 @@ import colossalai
|
||||
from colossalai.booster import Booster
|
||||
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
|
||||
from colossalai.cluster import DistCoordinator
|
||||
from colossalai.lazy import LazyInitContext
|
||||
from colossalai.nn.optimizer import HybridAdam
|
||||
from colossalai.utils import get_current_device
|
||||
|
||||
@@ -104,7 +102,7 @@ def evaluate_model(
|
||||
results = metric.compute()
|
||||
dist.all_reduce(accum_loss.div_(len(dataloader)))
|
||||
if coordinator.is_master() and results is not None:
|
||||
results['loss'] = accum_loss.item() / coordinator.world_size
|
||||
results["loss"] = accum_loss.item() / coordinator.world_size
|
||||
|
||||
return results
|
||||
|
||||
@@ -115,13 +113,20 @@ def evaluate_model(
|
||||
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, _criterion: Callable, lr_scheduler: LRScheduler,
|
||||
train_dataloader: DataLoader, booster: Booster, coordinator: DistCoordinator):
|
||||
|
||||
def train_epoch(
|
||||
epoch: int,
|
||||
model: nn.Module,
|
||||
optimizer: Optimizer,
|
||||
_criterion: Callable,
|
||||
lr_scheduler: LRScheduler,
|
||||
train_dataloader: DataLoader,
|
||||
booster: Booster,
|
||||
coordinator: DistCoordinator,
|
||||
):
|
||||
use_pipeline = isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size > 1
|
||||
is_pp_last_stage = use_pipeline and booster.plugin.stage_manager.is_last_stage()
|
||||
total_step = len(train_dataloader)
|
||||
@@ -129,22 +134,21 @@ def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, _criterion:
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
train_dataloader_iter = iter(train_dataloader)
|
||||
with tqdm(range(total_step),
|
||||
desc=f'Epoch [{epoch + 1}/{NUM_EPOCHS}]',
|
||||
disable=not (coordinator.is_master() or is_pp_last_stage)) as pbar:
|
||||
with tqdm(
|
||||
range(total_step),
|
||||
desc=f"Epoch [{epoch + 1}/{NUM_EPOCHS}]",
|
||||
disable=not (coordinator.is_master() or is_pp_last_stage),
|
||||
) as pbar:
|
||||
# Forward pass
|
||||
for _ in pbar:
|
||||
if use_pipeline:
|
||||
outputs = booster.execute_pipeline(train_dataloader_iter,
|
||||
model,
|
||||
_criterion,
|
||||
optimizer,
|
||||
return_loss=True,
|
||||
return_outputs=True)
|
||||
outputs = booster.execute_pipeline(
|
||||
train_dataloader_iter, model, _criterion, optimizer, return_loss=True, return_outputs=True
|
||||
)
|
||||
# Backward and optimize
|
||||
if is_pp_last_stage:
|
||||
loss = outputs['loss']
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
loss = outputs["loss"]
|
||||
pbar.set_postfix({"loss": loss.item()})
|
||||
else:
|
||||
data = next(train_dataloader_iter)
|
||||
data = move_to_cuda(data)
|
||||
@@ -152,7 +156,7 @@ def train_epoch(epoch: int, model: nn.Module, optimizer: Optimizer, _criterion:
|
||||
loss = _criterion(outputs, None)
|
||||
# Backward
|
||||
booster.backward(loss, optimizer)
|
||||
pbar.set_postfix({'loss': loss.item()})
|
||||
pbar.set_postfix({"loss": loss.item()})
|
||||
|
||||
optimizer.step()
|
||||
optimizer.zero_grad()
|
||||
@@ -164,24 +168,26 @@ 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', 'hybrid_parallel'],
|
||||
help="plugin to use")
|
||||
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", "hybrid_parallel"],
|
||||
help="plugin to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_type",
|
||||
type=str,
|
||||
default="gpt2",
|
||||
help="only gpt2 now",
|
||||
)
|
||||
parser.add_argument('--target_f1', type=float, default=None, help="target f1 score. Raise exception if not reached")
|
||||
parser.add_argument('--use_lazy_init', type=bool, default=False, help="for initiating lazy init context")
|
||||
parser.add_argument("--target_f1", type=float, default=None, help="target f1 score. Raise exception if not reached")
|
||||
parser.add_argument("--use_lazy_init", type=bool, default=False, help="for initiating lazy init context")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.model_type == 'gpt2':
|
||||
if args.model_type == "gpt2":
|
||||
model_name = "gpt2"
|
||||
else:
|
||||
raise RuntimeError
|
||||
@@ -198,36 +204,35 @@ 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':
|
||||
elif args.plugin == "gemini":
|
||||
plugin = GeminiPlugin(initial_scale=2**5)
|
||||
elif args.plugin == 'low_level_zero':
|
||||
elif args.plugin == "low_level_zero":
|
||||
plugin = LowLevelZeroPlugin(initial_scale=2**5)
|
||||
elif args.plugin == 'hybrid_parallel':
|
||||
|
||||
elif args.plugin == "hybrid_parallel":
|
||||
# modify the param accordingly for finetuning test cases
|
||||
plugin = HybridParallelPlugin(tp_size=1,
|
||||
pp_size=2,
|
||||
num_microbatches=None,
|
||||
microbatch_size=1,
|
||||
enable_all_optimization=True,
|
||||
zero_stage=1,
|
||||
precision='fp16',
|
||||
initial_scale=1)
|
||||
plugin = HybridParallelPlugin(
|
||||
tp_size=1,
|
||||
pp_size=2,
|
||||
num_microbatches=None,
|
||||
microbatch_size=1,
|
||||
enable_all_optimization=True,
|
||||
zero_stage=1,
|
||||
precision="fp16",
|
||||
initial_scale=1,
|
||||
)
|
||||
|
||||
booster = Booster(plugin=plugin, **booster_kwargs)
|
||||
|
||||
# ==============================
|
||||
# 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()
|
||||
|
||||
@@ -275,10 +280,9 @@ 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
|
||||
)
|
||||
|
||||
# ==============================
|
||||
# Train model
|
||||
@@ -286,14 +290,22 @@ def main():
|
||||
for epoch in range(NUM_EPOCHS):
|
||||
train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, train_dataloader, booster, coordinator)
|
||||
|
||||
results = evaluate_model(model, _criterion, test_dataloader, data_builder.num_labels, args.task,
|
||||
data_builder.eval_splits, booster, coordinator)
|
||||
results = evaluate_model(
|
||||
model,
|
||||
_criterion,
|
||||
test_dataloader,
|
||||
data_builder.num_labels,
|
||||
args.task,
|
||||
data_builder.eval_splits,
|
||||
booster,
|
||||
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