diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py index 0ae59a21b4..017ab1e565 100644 --- a/llmfoundry/callbacks/hf_checkpointer.py +++ b/llmfoundry/callbacks/hf_checkpointer.py @@ -17,7 +17,7 @@ import numpy as np import torch import torch.nn as nn -from composer.core import Callback, Event, State, Time, TimeUnit +from composer.core import Callback, Event, Precision, State, Time, TimeUnit from composer.core.state import fsdp_state_dict_type_context from composer.loggers import Logger, MLFlowLogger from composer.models import HuggingFaceModel @@ -496,7 +496,8 @@ def dtensor_to_tensor_hook( # Needed for proper hf ckpt saving. context_manager = te.onnx_export( True, - ) if is_te_imported else contextlib.nullcontext() + ) if is_te_imported and state.precision == Precision.AMP_FP8 else contextlib.nullcontext( + ) with context_manager: new_model_instance.save_pretrained(temp_save_dir) if original_tokenizer is not None: