Merge branch 'te_test_fix' into 'main'

fix TE test. See merge request ADLR/megatron-lm!1026
NVIDIA · Dec 20, 2023 · e4c73d6 · e4c73d6
2 parents d33e359 + d0e3b23
commit e4c73d6
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