wehub-resource-sync
391 lines
14 KiB
Python
391 lines
14 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import inspect
|
|
import os
|
|
import random
|
|
import time
|
|
from pprint import pprint
|
|
|
|
import numpy as np
|
|
import paddle
|
|
import paddle.distributed as dist
|
|
from benchmark import options
|
|
from benchmark.modules.benchmark_utils import clone_inputs
|
|
from benchmark.options import LR_SCHEDULER_REGISTRY, MODEL_REGISTRY, OPTIMIZER_REGISTRY
|
|
from benchmark.utils.record import AverageStatistical
|
|
|
|
from paddlenlp.utils import profiler
|
|
from paddlenlp.utils.log import logger
|
|
|
|
|
|
def set_seed(seed):
|
|
paddle.seed(seed)
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
|
|
def do_generated_inputs(args):
|
|
if args.device == "cpu":
|
|
rank = 0
|
|
trainer_count = 1
|
|
paddle.set_device("cpu")
|
|
else:
|
|
rank = dist.get_rank()
|
|
trainer_count = dist.get_world_size()
|
|
|
|
if trainer_count > 1:
|
|
dist.init_parallel_env()
|
|
|
|
# Set seed for CE
|
|
if args.seed is not None:
|
|
set_seed(args.seed)
|
|
|
|
benchmark_model = MODEL_REGISTRY[args.model]()
|
|
benchmark_optimizer = OPTIMIZER_REGISTRY[args.optimizer]()
|
|
|
|
if args.max_steps is None or (args.max_steps is not None and args.max_steps < 0):
|
|
args.max_steps = 10000
|
|
|
|
# Define model
|
|
model = benchmark_model.build_model(args)
|
|
|
|
if args.to_static:
|
|
input_spec = benchmark_model.create_input_specs()
|
|
model = paddle.jit.to_static(model, input_spec=input_spec)
|
|
logger.info("Successfully to apply @to_static with specs: {}".format(input_spec))
|
|
|
|
# Define data loader
|
|
example_inputs = benchmark_model.generate_inputs_for_model(args, model)
|
|
|
|
if args.lr_scheduler is not None:
|
|
benchmark_lr_scheduler = LR_SCHEDULER_REGISTRY[args.lr_scheduler]()
|
|
lr = benchmark_lr_scheduler.build_scheculer(args)
|
|
else:
|
|
lr = args.learning_rate
|
|
|
|
optimizer = benchmark_optimizer.build_optimizer(args, lr, model)
|
|
|
|
# for amp training
|
|
if args.use_amp:
|
|
scaler = paddle.amp.GradScaler(enable=True, init_loss_scaling=args.scale_loss)
|
|
model = paddle.amp.decorate(models=model, level=args.amp_level, save_dtype="float32")
|
|
|
|
# for distributed training
|
|
if trainer_count > 1:
|
|
model = paddle.DataParallel(model)
|
|
|
|
step_id = 1
|
|
|
|
# For benchmark
|
|
reader_cost_avg = AverageStatistical()
|
|
batch_cost_avg = AverageStatistical()
|
|
batch_ips_avg = AverageStatistical()
|
|
|
|
# Train loop
|
|
for pass_id in range(args.epoch):
|
|
epoch_start = time.time()
|
|
|
|
batch_start = time.time()
|
|
for batch_id in range(args.max_steps):
|
|
train_reader_cost = time.time() - batch_start
|
|
cloned_inputs = clone_inputs(example_inputs)
|
|
|
|
if args.use_amp:
|
|
# paddle version >= 2.5.0 or develop
|
|
paddle_version = float(paddle.__version__[:3])
|
|
if (paddle_version == 0.0) or (paddle_version >= 2.5):
|
|
with paddle.amp.auto_cast(
|
|
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
|
|
level=args.amp_level,
|
|
use_promote=args.amp_use_promote,
|
|
):
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
|
|
else:
|
|
with paddle.amp.auto_cast(
|
|
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
|
|
level=args.amp_level,
|
|
):
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
|
|
|
|
scaled = scaler.scale(loss)
|
|
scaled.backward()
|
|
|
|
scaler.minimize(optimizer, scaled)
|
|
if "set_to_zero" in inspect.getfullargspec(optimizer.clear_grad).args:
|
|
optimizer.clear_grad(set_to_zero=False)
|
|
else:
|
|
optimizer.clear_grad()
|
|
else:
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
|
|
|
|
loss.backward()
|
|
|
|
optimizer.step()
|
|
optimizer.clear_grad()
|
|
|
|
if args.profiler_options is not None:
|
|
profiler.add_profiler_step(args.profiler_options)
|
|
|
|
if args.max_steps and step_id == args.max_steps:
|
|
if args.save_model and rank == 0:
|
|
model_dir = args.save_model
|
|
if not os.path.exists(model_dir):
|
|
os.makedirs(model_dir)
|
|
paddle.save(model.state_dict(), os.path.join(model_dir, "model.pdparams"))
|
|
paddle.save(optimizer.state_dict(), os.path.join(model_dir, "model.pdopt"))
|
|
return
|
|
|
|
if args.lr_scheduler is not None and not args.scheduler_update_by_epoch:
|
|
lr.step()
|
|
|
|
if step_id % args.logging_steps == 0:
|
|
total_avg_loss = loss.numpy()
|
|
|
|
train_batch_cost = time.time() - batch_start
|
|
reader_cost_avg.record(train_reader_cost)
|
|
batch_cost_avg.record(train_batch_cost)
|
|
batch_ips_avg.record(train_batch_cost, sample_per_cards)
|
|
|
|
benchmark_model.logger(
|
|
args,
|
|
step_id=step_id,
|
|
pass_id=pass_id,
|
|
batch_id=batch_id,
|
|
loss=total_avg_loss,
|
|
batch_cost=batch_cost_avg.get_average(),
|
|
reader_cost=reader_cost_avg.get_average(),
|
|
num_samples=sample_per_cards,
|
|
ips=batch_ips_avg.get_average_per_sec(),
|
|
)
|
|
|
|
reader_cost_avg.reset()
|
|
batch_cost_avg.reset()
|
|
batch_ips_avg.reset()
|
|
else:
|
|
train_batch_cost = time.time() - batch_start
|
|
reader_cost_avg.record(train_reader_cost)
|
|
batch_cost_avg.record(train_batch_cost)
|
|
batch_ips_avg.record(train_batch_cost, sample_per_cards)
|
|
|
|
batch_start = time.time()
|
|
|
|
batch_id += 1
|
|
step_id += 1
|
|
|
|
if args.lr_scheduler is not None and args.scheduler_update_by_epoch:
|
|
lr.step()
|
|
|
|
train_epoch_cost = time.time() - epoch_start
|
|
logger.info("train epoch: %d, epoch_cost: %.5f s" % (pass_id, train_epoch_cost))
|
|
|
|
|
|
def do_train(args):
|
|
if args.device == "gpu":
|
|
rank = dist.get_rank()
|
|
trainer_count = dist.get_world_size()
|
|
else:
|
|
rank = 0
|
|
trainer_count = 1
|
|
paddle.set_device("cpu")
|
|
|
|
if trainer_count > 1:
|
|
dist.init_parallel_env()
|
|
|
|
# Set seed for CE
|
|
if args.seed is not None:
|
|
set_seed(args.seed)
|
|
|
|
benchmark_model = MODEL_REGISTRY[args.model]()
|
|
benchmark_optimizer = OPTIMIZER_REGISTRY[args.optimizer]()
|
|
|
|
# Define data loader
|
|
train_loader, eval_loader = benchmark_model.create_data_loader(args)
|
|
|
|
if args.max_steps is None or (args.max_steps is not None and args.max_steps < 0):
|
|
args.max_steps = len(train_loader) * args.epoch
|
|
|
|
# Define model
|
|
model = benchmark_model.build_model(args)
|
|
|
|
if args.to_static:
|
|
input_spec = benchmark_model.create_input_specs()
|
|
model = paddle.jit.to_static(model, input_spec=input_spec)
|
|
logger.info("Successfully to apply @to_static with specs: {}".format(input_spec))
|
|
|
|
if args.lr_scheduler is not None:
|
|
benchmark_lr_scheduler = LR_SCHEDULER_REGISTRY[args.lr_scheduler]()
|
|
lr = benchmark_lr_scheduler.build_scheculer(args)
|
|
else:
|
|
lr = args.learning_rate
|
|
|
|
optimizer = benchmark_optimizer.build_optimizer(args, lr, model)
|
|
|
|
# for amp training
|
|
if args.use_amp:
|
|
scaler = paddle.amp.GradScaler(enable=True, init_loss_scaling=args.scale_loss)
|
|
model = paddle.amp.decorate(models=model, level=args.amp_level, save_dtype="float32")
|
|
|
|
# for distributed training
|
|
if trainer_count > 1:
|
|
model = paddle.DataParallel(model)
|
|
|
|
step_id = 1
|
|
|
|
# For benchmark
|
|
reader_cost_avg = AverageStatistical()
|
|
batch_cost_avg = AverageStatistical()
|
|
batch_ips_avg = AverageStatistical()
|
|
|
|
# Train loop
|
|
for pass_id in range(args.epoch):
|
|
epoch_start = time.time()
|
|
|
|
batch_id = 0
|
|
batch_start = time.time()
|
|
for input_data in train_loader:
|
|
train_reader_cost = time.time() - batch_start
|
|
|
|
if args.use_amp:
|
|
# paddle version >= 2.5.0 or develop
|
|
paddle_version = float(paddle.__version__[:3])
|
|
if (paddle_version == 0.0) or (paddle_version >= 2.5):
|
|
with paddle.amp.auto_cast(
|
|
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
|
|
level=args.amp_level,
|
|
use_promote=args.amp_use_promote,
|
|
):
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
|
|
else:
|
|
with paddle.amp.auto_cast(
|
|
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
|
|
level=args.amp_level,
|
|
):
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
|
|
|
|
scaled = scaler.scale(loss)
|
|
scaled.backward()
|
|
|
|
scaler.minimize(optimizer, scaled)
|
|
if "set_to_zero" in inspect.getfullargspec(optimizer.clear_grad).args:
|
|
optimizer.clear_grad(set_to_zero=False)
|
|
else:
|
|
optimizer.clear_grad()
|
|
else:
|
|
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
|
|
|
|
loss.backward()
|
|
|
|
optimizer.step()
|
|
optimizer.clear_grad()
|
|
|
|
if args.profiler_options is not None:
|
|
profiler.add_profiler_step(args.profiler_options)
|
|
|
|
if args.max_steps and step_id == args.max_steps:
|
|
if args.save_model and rank == 0:
|
|
model_dir = args.save_model
|
|
if not os.path.exists(model_dir):
|
|
os.makedirs(model_dir)
|
|
paddle.save(model.state_dict(), os.path.join(model_dir, "model.pdparams"))
|
|
paddle.save(optimizer.state_dict(), os.path.join(model_dir, "model.pdopt"))
|
|
return
|
|
|
|
if args.lr_scheduler is not None and not args.scheduler_update_by_epoch:
|
|
lr.step()
|
|
|
|
if step_id % args.logging_steps == 0:
|
|
total_avg_loss = loss.numpy()
|
|
|
|
train_batch_cost = time.time() - batch_start
|
|
reader_cost_avg.record(train_reader_cost)
|
|
batch_cost_avg.record(train_batch_cost)
|
|
batch_ips_avg.record(train_batch_cost, sample_per_cards)
|
|
|
|
benchmark_model.logger(
|
|
args,
|
|
step_id=step_id,
|
|
pass_id=pass_id,
|
|
batch_id=batch_id,
|
|
loss=total_avg_loss,
|
|
batch_cost=batch_cost_avg.get_average(),
|
|
reader_cost=reader_cost_avg.get_average(),
|
|
num_samples=sample_per_cards,
|
|
ips=batch_ips_avg.get_average_per_sec(),
|
|
)
|
|
|
|
reader_cost_avg.reset()
|
|
batch_cost_avg.reset()
|
|
batch_ips_avg.reset()
|
|
else:
|
|
train_batch_cost = time.time() - batch_start
|
|
reader_cost_avg.record(train_reader_cost)
|
|
batch_cost_avg.record(train_batch_cost)
|
|
batch_ips_avg.record(train_batch_cost, sample_per_cards)
|
|
|
|
batch_start = time.time()
|
|
|
|
batch_id += 1
|
|
step_id += 1
|
|
|
|
if args.lr_scheduler is not None and args.scheduler_update_by_epoch:
|
|
lr.step()
|
|
|
|
train_epoch_cost = time.time() - epoch_start
|
|
logger.info("train epoch: %d, epoch_cost: %.5f s" % (pass_id, train_epoch_cost))
|
|
|
|
|
|
def do_hapi(args):
|
|
paddle.set_device(args.device)
|
|
|
|
# Set seed for CE
|
|
if args.seed is not None:
|
|
set_seed(args.seed)
|
|
|
|
benchmark_model = MODEL_REGISTRY[args.model]()
|
|
benchmark_optimizer = OPTIMIZER_REGISTRY[args.optimizer]()
|
|
|
|
# Define data loader
|
|
train_loader, eval_loader = benchmark_model.create_data_loader(args)
|
|
|
|
if args.lr_scheduler is not None:
|
|
benchmark_lr_scheduler = LR_SCHEDULER_REGISTRY[args.lr_scheduler]()
|
|
lr = benchmark_lr_scheduler.build_scheculer(args)
|
|
else:
|
|
lr = args.learning_rate
|
|
|
|
model = benchmark_model.build_model(args)
|
|
|
|
if args.to_static:
|
|
input_spec = benchmark_model.create_input_specs()
|
|
model = paddle.jit.to_static(model, input_spec=input_spec)
|
|
logger.info("Successfully to apply @to_static with specs: {}".format(input_spec))
|
|
|
|
optimizer = benchmark_optimizer.build_optimizer(args, lr, model)
|
|
|
|
benchmark_model.forward(model, args, optimizer=optimizer, train_loader=train_loader, eval_loader=eval_loader)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = options.get_training_parser()
|
|
args = options.parse_args_and_model(parser)
|
|
pprint(args)
|
|
|
|
if args.generated_inputs:
|
|
do_generated_inputs(args)
|
|
elif getattr(args, "use_hapi", False):
|
|
do_hapi(args)
|
|
else:
|
|
do_train(args)
|