Checking we use fused kernels to compute scaled masked softmax on prefix lm

[thomasw21]

• Fixes: [PrefixLM] Improve test to test out custom cuda kernel #206

There seems to be no issues with cuda kernels, as tests pass locally.

[stas00]

@thomasw21, this PR appears to have broken the test suite:

>               self.assertIn("Using fused softmax", cs.out)
E               AssertionError: 'Using fused softmax' not found in 'using world size: 1, data-parallel-size: 1, tensor-model-parallel size: 1, pipeline-model-parallel size: 1 \nusing torch.float16 for parameters ...\n------------------------ arguments ------------------------\n  accumulate_allreduce_grads_in_fp32 .......

=========================== short test summary info ============================
FAILED tests/test_model.py::MyTestCase::test_prefix_lm_wo_reset_attention_mask
FAILED tests/test_training.py::MegDSTestTraining::test_training_prefix_lm_all_0
FAILED tests/test_training.py::MegDSTestTraining::test_training_prefix_lm_all_1
FAILED tests/test_training.py::MegDSTestTraining::test_training_prefix_lm_all_2
FAILED tests/test_training.py::MegDSTestTraining::test_training_prefix_lm_all_3

Full logs:

https://github.com/bigscience-workshop/Megatron-DeepSpeed/runs/4339199658?check_suite_focus=true

In general please try to run the test suite locally if AWS doesn't give resources to run the CI (which unfortunately sucks :( )

[thomasw21]

Hmm reverted in c9afebc Though this passed locally. I'll have second look at it on monday.

[stas00]

Thank you, Thomas! That's helpful to other PRs.