From d0e3b238ac42d74cb6c634e8fa70d1b23cbc8ddd Mon Sep 17 00:00:00 2001 From: Jared Casper Date: Tue, 19 Dec 2023 14:56:29 -0800 Subject: [PATCH] fix TE test. --- .../test_scripts/gpt3/pretrain_gpt3_distributed_test.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/functional_tests/test_scripts/gpt3/pretrain_gpt3_distributed_test.sh b/tests/functional_tests/test_scripts/gpt3/pretrain_gpt3_distributed_test.sh index 3cad97cc60..e3f9626707 100755 --- a/tests/functional_tests/test_scripts/gpt3/pretrain_gpt3_distributed_test.sh +++ b/tests/functional_tests/test_scripts/gpt3/pretrain_gpt3_distributed_test.sh @@ -42,6 +42,7 @@ if [[ $USE_TE -eq 1 ]]; then echo "Running with TransformerEngine ..." TRANSFORMER_IMPL=transformer_engine TRAINING_DTYPE=bf16 + ADDITIONAL_PARAMS+=" --attention-softmax-in-fp32" else echo "Running with local transformer implementation ..." fi