Benchmark backprop of a te.TransformerLayer.

[wujingyue]

$ mpirun -np 1 pytest tests/python/test_transformer_engine.py --only-mpi

------------------------------------------------------------------------------------------------- benchmark: 2 tests ------------------------------------------------------------------------------------------------
Name (time in us)                           Min                   Max                  Mean              StdDev                Median                 IQR            Outliers       OPS            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
test_transformer_layer[forward]        952.7383 (1.0)      1,200.1079 (1.0)      1,016.5755 (1.0)      103.4595 (1.0)        976.7041 (1.0)       79.6057 (1.0)           1;1  983.6948 (1.0)           5           1
test_transformer_layer[backward]     1,070.4808 (1.12)     1,347.0454 (1.12)     1,188.3112 (1.17)     136.6049 (1.32)     1,125.9178 (1.15)     257.2080 (3.23)          1;0  841.5304 (0.86)          5           1
---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------

[wujingyue]

!build

[wujingyue]

!build

[wujingyue]

!build

[wujingyue]

I'm getting a weird error from jit_python_distributed_tests_17_A100 based on this PR. I'll fix that before merging this PR. I've yet to reproduce that locally unfortunately.

[wujingyue]

I believe the error has something to do with calling init_process_group and destroy_process_group for each test. Due to race conditions, some rank calls init_process_group from the second test before all ranks finished destroying the default process group.

[cowanmeg]

How about we don't add this to CI for now? We can benchmark locally and then work on fixing the CI problems