commit change

mosaicml · Jun 8, 2024 · 1c80786 · 1c80786
1 parent 4dd6218
commit 1c80786
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -486,10 +486,12 @@ def dtensor_to_tensor_hook(
             )
 
             log.debug('Saving Hugging Face checkpoint to disk')
-            new_model_instance.save_pretrained(temp_save_dir)
-            if original_tokenizer is not None:
-                assert isinstance(original_tokenizer, PreTrainedTokenizerBase)
-                original_tokenizer.save_pretrained(temp_save_dir)
+            import transformer_engine.pytorch as te
+            with te.onnx_export(True):
+                new_model_instance.save_pretrained(temp_save_dir)
+                if original_tokenizer is not None:
+                    assert isinstance(original_tokenizer, PreTrainedTokenizerBase)
+                    original_tokenizer.save_pretrained(temp_save_dir)
 
             # Only need to edit files for MPT because it has custom code
             if original_model.config.model_type == 'mpt':