Override cast_to_fp8 in te.module.linear (#120)

**Description**
Fix a bug in TE integration. 
Currently we only override te.cpp_extensions.cast_to_fp8 with our own
cat_to_fp8 in msamp.te.extension.

We also need to override te.module.linear.cast_to_fp8. Otherwise, it
will use the original function which does not support ScalingTensor and
will raise an exception in Megatron-LM

**Major Revision**
- Override cast_to_fp8 in te.module.linear

msamp/te/extension.py

@@ -124,6 +124,7 @@ def override():
         """Override transformer engine extension functions."""
         tex.fused_cast_transpose = TeExtensionOverrider.fused_cast_transpose
         te.cpp_extensions.cast_to_fp8 = TeExtensionOverrider.cast_to_fp8
+        te.module.linear.cast_to_fp8 = TeExtensionOverrider.cast_to_fp8
         te.cpp_extensions.fp8_cast_transpose_fused = TeExtensionOverrider.fp8_cast_transpose_fused