mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-11 22:10:37 +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:
@@ -2,36 +2,35 @@ from colossalai import get_default_parser
|
||||
|
||||
|
||||
def parse_demo_args():
|
||||
|
||||
parser = get_default_parser()
|
||||
parser.add_argument("--model_name_or_path",
|
||||
type=str,
|
||||
default="facebook/opt-350m",
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.")
|
||||
parser.add_argument("--output_path",
|
||||
type=str,
|
||||
default="./output_model.bin",
|
||||
help="The path of your saved model after finetuning.")
|
||||
parser.add_argument(
|
||||
"--model_name_or_path",
|
||||
type=str,
|
||||
default="facebook/opt-350m",
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output_path", type=str, default="./output_model.bin", help="The path of your saved model after finetuning."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--plugin",
|
||||
type=str,
|
||||
default="gemini",
|
||||
help=
|
||||
"Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero', 'hybrid_parallel'."
|
||||
help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero', 'hybrid_parallel'.",
|
||||
)
|
||||
parser.add_argument("--num_epoch", type=int, default=10, help="Number of epochs.")
|
||||
parser.add_argument("--batch_size",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Batch size (per dp group) for the training dataloader.")
|
||||
parser.add_argument("--learning_rate",
|
||||
type=float,
|
||||
default=5e-5,
|
||||
help="Initial learning rate (after the potential warmup period) to use.")
|
||||
parser.add_argument("--warmup_ratio",
|
||||
type=float,
|
||||
default=0.1,
|
||||
help="Ratio of warmup steps against total training steps.")
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=32, help="Batch size (per dp group) for the training dataloader."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--learning_rate",
|
||||
type=float,
|
||||
default=5e-5,
|
||||
help="Initial learning rate (after the potential warmup period) to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup_ratio", type=float, default=0.1, help="Ratio of warmup steps against total training steps."
|
||||
)
|
||||
parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay to use.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
|
||||
|
||||
@@ -40,25 +39,28 @@ def parse_demo_args():
|
||||
|
||||
|
||||
def parse_benchmark_args():
|
||||
|
||||
parser = get_default_parser()
|
||||
parser.add_argument("--model_name_or_path",
|
||||
type=str,
|
||||
default="facebook/opt-125m",
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.")
|
||||
parser.add_argument(
|
||||
"--model_name_or_path",
|
||||
type=str,
|
||||
default="facebook/opt-125m",
|
||||
help="Path to pretrained model or model identifier from huggingface.co/models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--plugin",
|
||||
type=str,
|
||||
default="gemini",
|
||||
help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'.")
|
||||
parser.add_argument("--batch_size",
|
||||
type=int,
|
||||
default=32,
|
||||
help="Batch size (per dp group) for the training dataloader.")
|
||||
parser.add_argument("--learning_rate",
|
||||
type=float,
|
||||
default=5e-5,
|
||||
help="Initial learning rate (after the potential warmup period) to use.")
|
||||
help="Plugin to use. Valid plugins include 'torch_ddp','torch_ddp_fp16','gemini','low_level_zero'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch_size", type=int, default=32, help="Batch size (per dp group) for the training dataloader."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--learning_rate",
|
||||
type=float,
|
||||
default=5e-5,
|
||||
help="Initial learning rate (after the potential warmup period) to use.",
|
||||
)
|
||||
parser.add_argument("--weight_decay", type=float, default=0.0, help="Weight decay to use.")
|
||||
parser.add_argument("--max_train_steps", type=int, default=20, help="Total number of training steps to perform.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="A seed for reproducible training.")
|
||||
|
@@ -1,37 +1,38 @@
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
|
||||
class NetflixDataset(Dataset):
|
||||
|
||||
def __init__(self, tokenizer):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.input_ids = []
|
||||
self.attn_masks = []
|
||||
self.labels = []
|
||||
self.txt_list = netflix_descriptions = load_dataset("hugginglearners/netflix-shows", split="train")['description']
|
||||
self.txt_list = netflix_descriptions = load_dataset("hugginglearners/netflix-shows", split="train")[
|
||||
"description"
|
||||
]
|
||||
self.max_length = max([len(self.tokenizer.encode(description)) for description in netflix_descriptions])
|
||||
|
||||
for txt in self.txt_list:
|
||||
encodings_dict = self.tokenizer('</s>' + txt + '</s>',
|
||||
truncation=True,
|
||||
max_length=self.max_length,
|
||||
padding="max_length")
|
||||
self.input_ids.append(torch.tensor(encodings_dict['input_ids']))
|
||||
self.attn_masks.append(torch.tensor(encodings_dict['attention_mask']))
|
||||
encodings_dict = self.tokenizer(
|
||||
"</s>" + txt + "</s>", truncation=True, max_length=self.max_length, padding="max_length"
|
||||
)
|
||||
self.input_ids.append(torch.tensor(encodings_dict["input_ids"]))
|
||||
self.attn_masks.append(torch.tensor(encodings_dict["attention_mask"]))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.input_ids)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.input_ids[idx], self.attn_masks[idx]
|
||||
|
||||
|
||||
|
||||
def netflix_collator(data):
|
||||
return {'input_ids': torch.stack([x[0] for x in data]),
|
||||
'attention_mask': torch.stack([x[1] for x in data]),
|
||||
'labels': torch.stack([x[0] for x in data])}
|
||||
return {
|
||||
"input_ids": torch.stack([x[0] for x in data]),
|
||||
"attention_mask": torch.stack([x[1] for x in data]),
|
||||
"labels": torch.stack([x[0] for x in data]),
|
||||
}
|
||||
|
@@ -35,6 +35,7 @@ def get_data(batch_size, seq_len, vocab_size):
|
||||
|
||||
def colo_memory_cap(size_in_GB):
|
||||
from colossalai.utils import colo_device_memory_capacity, colo_set_process_memory_fraction, get_current_device
|
||||
|
||||
cuda_capacity = colo_device_memory_capacity(get_current_device())
|
||||
if size_in_GB * (1024**3) < cuda_capacity:
|
||||
colo_set_process_memory_fraction(size_in_GB * (1024**3) / cuda_capacity)
|
||||
@@ -42,7 +43,6 @@ def colo_memory_cap(size_in_GB):
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
args = parse_benchmark_args()
|
||||
|
||||
# Launch ColossalAI
|
||||
@@ -72,13 +72,13 @@ def main():
|
||||
|
||||
# Set plugin
|
||||
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(offload_optim_frac=1.0, pin_memory=True, initial_scale=2**5)
|
||||
elif args.plugin == 'low_level_zero':
|
||||
elif args.plugin == "low_level_zero":
|
||||
plugin = LowLevelZeroPlugin(initial_scale=2**5)
|
||||
logger.info(f"Set plugin as {args.plugin}", ranks=[0])
|
||||
|
||||
@@ -101,11 +101,10 @@ def main():
|
||||
start_time = time.time()
|
||||
|
||||
for _ in range(args.max_train_steps):
|
||||
|
||||
input_ids, attn_mask = get_data(args.batch_size, SEQ_LEN, VOCAB_SIZE)
|
||||
optimizer.zero_grad()
|
||||
outputs = model(input_ids=input_ids, attention_mask=attn_mask, labels=input_ids, use_cache=False)
|
||||
loss = outputs['loss']
|
||||
loss = outputs["loss"]
|
||||
booster.backward(loss, optimizer)
|
||||
optimizer.step()
|
||||
|
||||
@@ -123,7 +122,8 @@ def main():
|
||||
f"plugin: {args.plugin}, "
|
||||
f"throughput: {throughput}, "
|
||||
f"maximum memory usage per gpu: {max_mem}.",
|
||||
ranks=[0])
|
||||
ranks=[0],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
@@ -1,5 +1,3 @@
|
||||
import time
|
||||
|
||||
import datasets
|
||||
import torch
|
||||
import transformers
|
||||
@@ -12,7 +10,6 @@ from transformers.utils.versions import require_version
|
||||
import colossalai
|
||||
from colossalai.booster import Booster
|
||||
from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin
|
||||
from colossalai.booster.plugin.hybrid_parallel_plugin import HybridParallelModule
|
||||
from colossalai.cluster import DistCoordinator
|
||||
from colossalai.logging import disable_existing_loggers, get_dist_logger
|
||||
from colossalai.nn.optimizer import HybridAdam
|
||||
@@ -29,7 +26,6 @@ def move_to_cuda(batch, device):
|
||||
|
||||
|
||||
def train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, dataloader, booster, coordinator):
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
use_pipeline = isinstance(booster.plugin, HybridParallelPlugin) and booster.plugin.pp_size > 1
|
||||
@@ -39,22 +35,19 @@ def train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, dataloader, b
|
||||
model.train()
|
||||
optimizer.zero_grad()
|
||||
dataloader = iter(dataloader)
|
||||
with tqdm(range(total_step), desc=f'Epoch [{epoch + 1}]',
|
||||
disable=not (coordinator.is_master() or is_pp_last_stage)) as pbar:
|
||||
|
||||
with tqdm(
|
||||
range(total_step), desc=f"Epoch [{epoch + 1}]", disable=not (coordinator.is_master() or is_pp_last_stage)
|
||||
) as pbar:
|
||||
# Forward pass
|
||||
for _ in pbar:
|
||||
if use_pipeline:
|
||||
outputs = booster.execute_pipeline(dataloader,
|
||||
model,
|
||||
_criterion,
|
||||
optimizer,
|
||||
return_loss=True,
|
||||
return_outputs=True)
|
||||
outputs = booster.execute_pipeline(
|
||||
dataloader, 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(dataloader)
|
||||
data = move_to_cuda(data)
|
||||
@@ -62,7 +55,7 @@ def train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, dataloader, b
|
||||
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()
|
||||
@@ -70,7 +63,6 @@ def train_epoch(epoch, model, optimizer, _criterion, lr_scheduler, dataloader, b
|
||||
|
||||
|
||||
def main():
|
||||
|
||||
args = parse_demo_args()
|
||||
|
||||
# Launch ColossalAI
|
||||
@@ -98,34 +90,34 @@ def main():
|
||||
|
||||
# Set plugin
|
||||
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(offload_optim_frac=1.0, pin_memory=True, 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=2,
|
||||
pp_size=2,
|
||||
num_microbatches=2,
|
||||
enable_all_optimization=True,
|
||||
zero_stage=0,
|
||||
precision='fp16',
|
||||
initial_scale=1)
|
||||
plugin = HybridParallelPlugin(
|
||||
tp_size=2,
|
||||
pp_size=2,
|
||||
num_microbatches=2,
|
||||
enable_all_optimization=True,
|
||||
zero_stage=0,
|
||||
precision="fp16",
|
||||
initial_scale=1,
|
||||
)
|
||||
|
||||
logger.info(f"Set plugin as {args.plugin}", ranks=[0])
|
||||
|
||||
# Prepare tokenizer and dataloader
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
|
||||
dataset = NetflixDataset(tokenizer)
|
||||
dataloader = plugin.prepare_dataloader(dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
drop_last=True,
|
||||
collate_fn=netflix_collator)
|
||||
dataloader = plugin.prepare_dataloader(
|
||||
dataset, batch_size=args.batch_size, shuffle=True, drop_last=True, collate_fn=netflix_collator
|
||||
)
|
||||
|
||||
# Set optimizer
|
||||
optimizer = HybridAdam(model.parameters(), lr=(args.learning_rate * world_size), weight_decay=args.weight_decay)
|
||||
@@ -133,9 +125,9 @@ def main():
|
||||
# Set lr scheduler
|
||||
total_steps = len(dataloader) * args.num_epoch
|
||||
num_warmup_steps = int(args.warmup_ratio * total_steps)
|
||||
lr_scheduler = get_linear_schedule_with_warmup(optimizer,
|
||||
num_warmup_steps=num_warmup_steps,
|
||||
num_training_steps=len(dataloader) * args.num_epoch)
|
||||
lr_scheduler = get_linear_schedule_with_warmup(
|
||||
optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=len(dataloader) * args.num_epoch
|
||||
)
|
||||
|
||||
# Define criterion
|
||||
def _criterion(outputs, inputs):
|
||||
@@ -145,11 +137,9 @@ def main():
|
||||
|
||||
# Set booster
|
||||
booster = Booster(plugin=plugin, **booster_kwargs)
|
||||
model, optimizer, _criterion, dataloader, lr_scheduler = booster.boost(model=model,
|
||||
optimizer=optimizer,
|
||||
dataloader=dataloader,
|
||||
criterion=_criterion,
|
||||
lr_scheduler=lr_scheduler)
|
||||
model, optimizer, _criterion, dataloader, lr_scheduler = booster.boost(
|
||||
model=model, optimizer=optimizer, dataloader=dataloader, criterion=_criterion, lr_scheduler=lr_scheduler
|
||||
)
|
||||
|
||||
# Start finetuning
|
||||
logger.info(f"Start finetuning", ranks=[0])
|
||||
|
@@ -24,7 +24,7 @@ torchrun \
|
||||
--mem_cap ${MEMCAP} \
|
||||
--plugin ${PLUGIN} \
|
||||
--batch_size ${BS}
|
||||
|
||||
|
||||
done
|
||||
done
|
||||
done
|
||||
|
Reference in New Issue
Block a user