Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit e831abd

Browse files
zhangting2020zhangting2020
authored
disable kernel promotion for amp training (PaddlePaddle#5922)
* disable kernel promotion for amp training * improve amp training performance * fix eval
1 parent 5540c0c commit e831abd

3 files changed

Lines changed: 87 additions & 32 deletions

File tree

‎examples/language_model/moe/dygraph/run_moe_pretrain.py‎

Lines changed: 56 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -80,16 +80,31 @@ def run_evaluate(args, data_loader, model, criterion, iter_steps, log_writer, gl
8080
local_time = time.time()
8181
for eval_step, batch in enumerate(data_loader):
8282
tokens, loss_mask, labels = batch
83-
with paddle.amp.auto_cast(
84-
args.use_pure_fp16,
85-
custom_black_list=[
86-
"reduce_sum",
87-
"c_softmax_with_cross_entropy",
88-
"elementwise_div",
89-
],
90-
level="O2",
91-
):
92-
preds = model(tokens)
83+
# paddle version >= 2.5.0 or develop
84+
paddle_version = float(paddle.__version__[:3])
85+
if (paddle_version == 0.0) or (paddle_version >= 2.5):
86+
with paddle.amp.auto_cast(
87+
args.use_pure_fp16,
88+
custom_black_list=[
89+
"reduce_sum",
90+
"c_softmax_with_cross_entropy",
91+
"elementwise_div",
92+
],
93+
level="O2",
94+
use_promote=False,
95+
):
96+
preds = model(tokens)
97+
else:
98+
with paddle.amp.auto_cast(
99+
args.use_pure_fp16,
100+
custom_black_list=[
101+
"reduce_sum",
102+
"c_softmax_with_cross_entropy",
103+
"elementwise_div",
104+
],
105+
level="O2",
106+
):
107+
preds = model(tokens)
93108
preds = paddle.cast(preds, dtype="float32")
94109
loss = criterion(preds, labels, loss_mask)
95110

@@ -517,7 +532,6 @@ def do_train(args):
517532
# many times. and start a new random dataloader.
518533
valid_data_loader = valid_data_loader()
519534
test_data_loader = test_data_loader()
520-
521535
for step, batch in enumerate(train_data_loader()):
522536
# to remove the train data that has been studyed.
523537
if step < global_step - pass_num:
@@ -534,19 +548,37 @@ def do_train(args):
534548
start_index = i * args.micro_batch_size
535549
end_index = start_index + args.micro_batch_size
536550
timers("forward-compute").start()
537-
with paddle.amp.auto_cast(
538-
args.use_pure_fp16,
539-
custom_black_list=[
540-
"reduce_sum",
541-
"c_softmax_with_cross_entropy",
542-
"elementwise_div",
543-
],
544-
level="O2",
545-
):
546-
preds = model(tokens[start_index:end_index, :])
547-
loss_mbs = criterion(
548-
preds, labels[start_index:end_index, :], loss_mask[start_index:end_index, :]
549-
)
551+
# paddle version >= 2.5.0 or develop
552+
paddle_version = float(paddle.__version__[:3])
553+
if (paddle_version == 0.0) or (paddle_version >= 2.5):
554+
with paddle.amp.auto_cast(
555+
args.use_pure_fp16,
556+
custom_black_list=[
557+
"reduce_sum",
558+
"c_softmax_with_cross_entropy",
559+
"elementwise_div",
560+
],
561+
level="O2",
562+
use_promote=False,
563+
):
564+
preds = model(tokens[start_index:end_index, :])
565+
loss_mbs = criterion(
566+
preds, labels[start_index:end_index, :], loss_mask[start_index:end_index, :]
567+
)
568+
else:
569+
with paddle.amp.auto_cast(
570+
args.use_pure_fp16,
571+
custom_black_list=[
572+
"reduce_sum",
573+
"c_softmax_with_cross_entropy",
574+
"elementwise_div",
575+
],
576+
level="O2",
577+
):
578+
preds = model(tokens[start_index:end_index, :])
579+
loss_mbs = criterion(
580+
preds, labels[start_index:end_index, :], loss_mask[start_index:end_index, :]
581+
)
550582
timers("forward-compute").stop()
551583

552584
if args.gate != "naive" and args.balance_loss_weight:

‎tests/test_tipc/benchmark/options.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ def get_parser():
132132
parser.add_argument("--use_amp", type=str2bool, nargs="?", const=False, help="Enable AMP. ")
133133
parser.add_argument("--scale_loss", type=float, default=128, help="Loss scale. ")
134134
parser.add_argument("--amp_level", type=str, default="O2", help="AMP LEVEL. O1 or O2. ")
135+
parser.add_argument("--amp_use_promote", action="store_true", help="Enable kernel promotion for AMP training. ")
135136
parser.add_argument("--custom_black_list", type=str, nargs="+", default=None, help="Custom black list for AMP. ")
136137

137138
parser.add_argument("--to_static", action="store_true", help="Enable to static. ")

‎tests/test_tipc/train.py‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -103,10 +103,21 @@ def do_generated_inputs(args):
103103
cloned_inputs = clone_inputs(example_inputs)
104104

105105
if args.use_amp:
106-
with paddle.amp.auto_cast(
107-
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {}, level=args.amp_level
108-
):
109-
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
106+
# paddle version >= 2.5.0 or develop
107+
paddle_version = float(paddle.__version__[:3])
108+
if (paddle_version == 0.0) or (paddle_version >= 2.5):
109+
with paddle.amp.auto_cast(
110+
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
111+
level=args.amp_level,
112+
use_promote=args.amp_use_promote,
113+
):
114+
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
115+
else:
116+
with paddle.amp.auto_cast(
117+
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
118+
level=args.amp_level,
119+
):
120+
loss, sample_per_cards = benchmark_model.forward(model, args, cloned_inputs)
110121

111122
scaled = scaler.scale(loss)
112123
scaled.backward()
@@ -247,10 +258,21 @@ def do_train(args):
247258
train_reader_cost = time.time() - batch_start
248259

249260
if args.use_amp:
250-
with paddle.amp.auto_cast(
251-
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {}, level=args.amp_level
252-
):
253-
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
261+
# paddle version >= 2.5.0 or develop
262+
paddle_version = float(paddle.__version__[:3])
263+
if (paddle_version == 0.0) or (paddle_version >= 2.5):
264+
with paddle.amp.auto_cast(
265+
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
266+
level=args.amp_level,
267+
use_promote=args.amp_use_promote,
268+
):
269+
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
270+
else:
271+
with paddle.amp.auto_cast(
272+
custom_black_list=args.custom_black_list if args.amp_level == "O2" else {},
273+
level=args.amp_level,
274+
):
275+
loss, sample_per_cards = benchmark_model.forward(model, args, input_data)
254276

255277
scaled = scaler.scale(loss)
256278
scaled.backward()

0 commit comments

Comments
 (0)