feat(decoder): do not split checkpoint files

huggingface · Apr 15, 2024 · bfa8fdb · bfa8fdb
1 parent f936089
commit bfa8fdb
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
@@ -42,7 +42,6 @@
 
 if is_transformers_neuronx_available():
     from transformers_neuronx.config import ContinuousBatchingConfig, NeuronConfig
-    from transformers_neuronx.module import save_split
 
 
 if TYPE_CHECKING:
@@ -249,9 +248,7 @@ def _create_checkpoint(
 
         # Save the model checkpoint in a temporary directory
         checkpoint_dir = TemporaryDirectory()
-        model.save_pretrained(
-            checkpoint_dir.name, save_function=save_split, safe_serialization=False, max_shard_size="10000GB"
-        )
+        model.save_pretrained(checkpoint_dir.name)
         return checkpoint_dir
 
     @classmethod