Merge branch 'main' into shanmugamr-main-patch-23323

NVIDIA · Dec 20, 2023 · d44375d · d44375d
2 parents b7f8a07 + e4c73d6
commit d44375d
Showing 1 changed file with 1 addition and 0 deletions.
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
@@ -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