add TransformerEngine in rocm6.2.x.dockerfile

dockerfile/rocm6.2.x.dockerfile

@@ -180,6 +180,13 @@ RUN cp -r /opt/superbench/third_party/hipBLASLt/build/release/hipblaslt-install/
 RUN cd third_party/Megatron/Megatron-DeepSpeed && \
     git apply ../megatron_deepspeed_rocm6.patch
 
+# Install transformer_engine
+git clone --recursive https://github.com/ROCm/TransformerEngine.git
+cd TransformerEngine
+export NVTE_FRAMEWORK=pytorch #optionally set framework, currently only support pytorch and jax
+export PYTORCH_ROCM_ARCH=gfx942 # CK fused attn only support MI200 and MI300 and fp8 features are only supported on MI300
+pip install .
+
 ADD . .
 ENV USE_HIP_DATATYPE=1
 ENV USE_HIPBLAS_COMPUTETYPE=1