Fix FP8 checkpoint resumption with onnx export flag (mosaicml#2907)

composer/core/precision.py

@@ -78,7 +78,11 @@ def get_precision_context(precision: Union[str, Precision],
                 }
             fp8_recipe = DelayedScaling(**precision_config)
             with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
-                yield
+                # The te.onnx_export flag ensures that we save all fp8 buffers
+                # as tensors instead of bytes. This is necessary for proper
+                # saving and resumption of checkpoints.
+                with te.onnx_export(enabled=True):
+                    yield
         else:
             if te_installed:
                 raise RuntimeError('AMP_FP8 precision is used but current device does not support it.')