From 2cec34b80bcfb499d1316028ddf8adacf46cad23 Mon Sep 17 00:00:00 2001 From: Wang Huan Date: Wed, 8 May 2024 02:59:57 +0000 Subject: [PATCH] pre copy pinned data to gpu --- model_zoo/bert/run_pretrain.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/model_zoo/bert/run_pretrain.py b/model_zoo/bert/run_pretrain.py index edc7d952d042..36e6f51ed03d 100644 --- a/model_zoo/bert/run_pretrain.py +++ b/model_zoo/bert/run_pretrain.py @@ -398,6 +398,14 @@ def do_train(args): next_sentence_labels, masked_lm_scale, ) = batch + input_ids = input_ids.cuda(blocking=False) + segment_ids = segment_ids.cuda(blocking=False) + input_mask = input_mask.cuda(blocking=False) + masked_lm_positions = masked_lm_positions.cuda(blocking=False) + masked_lm_labels = masked_lm_labels.cuda(blocking=False) + next_sentence_labels = next_sentence_labels.cuda(blocking=False) + masked_lm_scale = masked_lm_scale.cuda(blocking=False) + with paddle.amp.auto_cast( args.use_amp, custom_white_list=["layer_norm", "softmax", "gelu", "fused_attention", "fused_feedforward"],