github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity
Browse Source

[example] remove redundant texts & update roberta (#3493) 

* update roberta example

* update roberta example

* modify conflict & update roberta
pull/3497/head
mandoxzhang 2 years ago committed by GitHub
parent
ab5fd127e3
commit
8f2c55f9c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 0 additions and 93 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 93
      examples/language/roberta/pretraining/run_pretraining.py

93
examples/language/roberta/pretraining/run_pretraining.py
Unescape View File

@ -4,7 +4,6 @@ import time
from functools import partial
import torch
<<<<<<< HEAD
from tqdm import tqdm
import os
import time
@ -20,15 +19,9 @@ from colossalai.utils.model.colo_init_context import ColoInitContext
from colossalai.zero import ZeroOptimizer
from colossalai.tensor import ColoParameter, ComputePattern, ComputeSpec, ProcessGroup, ReplicaSpec, ShardSpec
=======
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
from arguments import parse_args
from evaluation import evaluate
from loss import LossForPretraining
<<<<<<< HEAD
from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider
=======
from nvidia_bert_dataset_provider import NvidiaBertDatasetProvider
from pretrain_utils import get_lr_scheduler, get_model, get_optimizer, save_ckpt
from tqdm import tqdm
@ -37,20 +30,6 @@ from utils.exp_util import get_mem_info, get_tflops, log_args, throughput_calcul
from utils.global_vars import get_tensorboard_writer, get_timers, set_global_variables
from utils.logger import Logger
import colossalai
import colossalai.nn as col_nn
from colossalai.context import ParallelMode
from colossalai.core import global_context as gpc
from colossalai.nn.optimizer import HybridAdam
from colossalai.nn.parallel import ZeroDDP
from colossalai.tensor import ProcessGroup
from colossalai.utils import get_current_device
from colossalai.zero import ZeroOptimizer
from colossalai.zero.gemini import ChunkManager, ColoInitContext, GeminiManager
from colossalai.zero.legacy import ShardedModelV2, ShardedOptimizerV2, ZeroInitContext
from colossalai.zero.legacy.shard_utils import TensorShardStrategy
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
def main():
@ -59,13 +38,8 @@ def main():
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
<<<<<<< HEAD
# os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
=======
os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
logger = Logger(os.path.join(args.log_path, launch_time), cuda=torch.cuda.is_available(), debug=args.vscode_debug)
if args.vscode_debug:
@ -78,11 +52,7 @@ def main():
args.local_rank = -1
args.log_interval = 1
else:
<<<<<<< HEAD
colossalai.launch_from_torch(config={}) #args.colossal_config
=======
colossalai.launch_from_torch(args.colossal_config) # args.colossal_config
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
args.local_rank = int(os.environ["LOCAL_RANK"])
logger.info(
f'launch_from_torch, world size: {torch.distributed.get_world_size()} | ' +
@ -93,17 +63,11 @@ def main():
args.tokenizer = tokenizer
args.logger = logger
set_global_variables(launch_time, args.tensorboard_path)
<<<<<<< HEAD
=======
use_zero = hasattr(gpc.config, 'zero')
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
world_size = torch.distributed.get_world_size()
init_dev = get_current_device()
# build model, optimizer and criterion
<<<<<<< HEAD
if args.distplan.startswith("CAI"):
# all param must use the same process group.
world_size = torch.distributed.get_world_size()
@ -118,13 +82,6 @@ def main():
dtype=torch.half,
default_dist_spec=default_dist_spec,
default_pg=shard_pg):
=======
if use_zero:
shard_strategy = TensorShardStrategy()
with ZeroInitContext(target_device=torch.cuda.current_device(), shard_strategy=shard_strategy,
shard_param=True):
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
config, model, numel = get_model(args, logger)
# asign running configurations
@ -170,14 +127,9 @@ def main():
logger.info(f'Model numel: {numel}')
get_tflops_func = partial(get_tflops, numel, args.train_micro_batch_size_per_gpu, args.max_seq_length)
<<<<<<< HEAD
# 144003367 is is the length of the entire dataset
steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size #len(dataloader)
=======
# len(dataloader)
steps_per_epoch = 144003367 // world_size // args.train_micro_batch_size_per_gpu // args.gradient_accumulation_steps // args.refresh_bucket_size
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
total_steps = steps_per_epoch * args.epoch
lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=-1)
@ -204,32 +156,14 @@ def main():
start_epoch = o_l_state_dict['epoch']
start_shard = o_l_state_dict['shard'] + 1
# global_step = o_l_state_dict['global_step'] + 1
<<<<<<< HEAD
logger.info(f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}')
=======
logger.info(
f'resume from epoch {start_epoch} shard {start_shard} step {lr_scheduler.last_epoch} lr {lr_scheduler.get_last_lr()[0]}'
)
else:
optimizer = get_optimizer(model, lr=args.lr)
lr_scheduler = get_lr_scheduler(optimizer, total_steps=total_steps, last_epoch=-1)
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
criterion = LossForPretraining(config.vocab_size)
# build dataloader
pretrain_dataset_provider = NvidiaBertDatasetProvider(args)
<<<<<<< HEAD
=======
# initialize with colossalai
engine, _, _, lr_scheduelr = colossalai.initialize(model=model,
optimizer=optimizer,
criterion=criterion,
lr_scheduler=lr_scheduler)
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
logger.info(get_mem_info(prefix='After init model, '))
best_loss = None
@ -254,15 +188,9 @@ def main():
else:
iterator_data = enumerate(dataset_iterator)
<<<<<<< HEAD
model.train()
for step, batch_data in iterator_data:
=======
engine.train()
for step, batch_data in iterator_data:
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
# batch_data = pretrain_dataset_provider.get_batch(batch_index)
input_ids = batch_data[0].cuda(f"cuda:{torch.cuda.current_device()}")
@ -271,31 +199,18 @@ def main():
mlm_label = batch_data[3].cuda(f"cuda:{torch.cuda.current_device()}")
# nsp_label = batch_data[5].cuda()
<<<<<<< HEAD
output = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
loss = criterion(output.logits, mlm_label)
=======
output = engine(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
loss = engine.criterion(output.logits, mlm_label)
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
pretrain_dataset_provider.prefetch_batch()
optimizer.backward(loss)
train_loss += loss.float().item()
# if (step + 1) % args.accumulation_step == 0:
<<<<<<< HEAD
optimizer.step()
lr_scheduler.step()
optimizer.zero_grad()
=======
engine.step()
lr_scheduelr.step()
engine.zero_grad()
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
global_step += 1
if global_step % args.log_interval == 0 and global_step != 0 \
@ -326,18 +241,10 @@ def main():
logger.info(f'epoch {epoch} shard {shard} has cost {timers("shard_time").elapsed() / 60 :.3f} mins')
logger.info('*' * 100)
<<<<<<< HEAD
eval_loss += evaluate(model, args, logger, global_step, criterion)
save_ckpt(model, optimizer, lr_scheduler, os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch, shard, global_step)
=======
eval_loss += evaluate(engine, args, logger, global_step)
save_ckpt(engine.model, optimizer, lr_scheduelr,
os.path.join(args.ckpt_path, launch_time, f'epoch-{epoch}_shard-{shard}_' + launch_time), epoch,
shard, global_step)
>>>>>>> 52a933e17509c71811e919b165de38cb3d5d6d41
eval_loss /= len(os.listdir(args.data_path_prefix))
logger.info(
f'epoch {epoch} | shard_length {len(os.listdir(args.data_path_prefix))} | elapsed_time: {timers("epoch_time").elapsed() / 60 :.3f} mins'

Write Preview
Loading…
Cancel
Save