# Avoid saving deepspeed optimizer and model states at checkpoints

**URL:** <https://discuss.huggingface.co/t/avoid-saving-deepspeed-optimizer-and-model-states-at-checkpoints/141178>\
**Category:** Beginners\
**Created:** [February 17, 2025, 12:35am UTC](https://discuss.huggingface.co/t/avoid-saving-deepspeed-optimizer-and-model-states-at-checkpoints/141178 "2025-02-17T00:35:27Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![esualp](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/esualp/32/39451_2.png) [@esualp](https://discuss.huggingface.co/u/esualp)\
**Post date:** [February 17, 2025, 12:35am UTC](https://discuss.huggingface.co/t/avoid-saving-deepspeed-optimizer-and-model-states-at-checkpoints/141178/1 "2025-02-17T00:35:27Z")

</div>

Hello all,

I am using Deepspeed ZeRO-3 optimizer to finetune 7B llama-based model for 3 epochs. After 1st epoch, the model starts to overfit since evaluation loss becomes much higher than previous steps. I would like to check the inference performance of models in each checkpoint (save\_strategy = ‘epochs’ already).

However, “_\_optim\_states.pt" and "_\_model\_states.pt” are taking huge place in terms of storage. I’ll never continue to finetune these checkpoints, do i still need to keep them for inference?

I think I won’t need optim\_states since it is optimizer-related, what about **model\_states**? I guess hf trainer uses pytorch\_model-00001-of-00002.bin and pytorch\_model-00002-of-00002.bin during inference, but want to be sure about it

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](https://discuss.huggingface.co/u/John6666)\
**Post date:** [February 17, 2025, 6:30am UTC](https://discuss.huggingface.co/t/avoid-saving-deepspeed-optimizer-and-model-states-at-checkpoints/141178/2 "2025-02-17T06:30:46Z")

</div>

It looks really hard…

> [@Trainer option to disable saving DeepSpeed checkpoints](https://discuss.huggingface.co/t/trainer-option-to-disable-saving-deepspeed-checkpoints/13262):
>
> I’d like to ask for opinions about adding a Trainer configuration option to disable saving of DeepSpeed checkpoints (potentially only keeping the model weights). Context: I’m finetuning gpt-j-6b for basic translation phrases on consumer hardware (128GB System RAM and Nvidia GPU with 24GB RAM). I use the DeepSpeed Zero optimizer, stages 2 and 3 so 99% of my system memory is fully allocated (I also have a huge swap file on an NVMe). Arguments for adding this flag: The DeepSpeed checkpoints are…

> <https://github.com/huggingface/transformers/issues/22822>
>
> \### System Info
> 
> \- \`transformers\` version: 4.28.0.dev0
> \- Platform: Linux-4.18….0-372.32.1.el8\_6.x86\_64-x86\_64-with-glibc2.17
> \- Python version: 3.8.16
> \- Huggingface\_hub version: 0.13.3
> \- Safetensors version: not installed
> \- PyTorch version (GPU?): 1.12.1+cu116 (True)
> \- Tensorflow version (GPU?): not installed (NA)
> \- Flax version (CPU?/GPU?/TPU?): not installed (NA)
> \- Jax version: not installed
> \- JaxLib version: not installed
> \- Using GPU in script?: yes
> \- Using distributed or parallel set-up in script?: yes
> 
> 
> \### Who can help?
> 
> @stas00 @sgugger
> 
> 
> 
> \### Information
> 
> \- \[X\] The official example scripts
> \- \[X\] My own modified scripts
> 
> \### Tasks
> 
> \- \[x\] An officially supported task in the \`examples\` folder (such as GLUE/SQuAD, ...)
> \- \[\] My own task or dataset (give details below)
> 
> \### Reproduction
> 
> I'm using Trainer with deepspeed integration to fine-tune a Llama model.
> 
> This is the stage2 config im using:
> \`\`\`json
> {
> "optimizer": {
> "type": "AdamW",
> "params": {
> "lr": "auto",
> "betas": "auto",
> "eps": "auto",
> "weight\_decay": "auto"
> }
> },
> 
> 
> "zero\_optimization": {
> "stage": 2,
> "offload\_optimizer": {
> "device": "cpu",
> "pin\_memory": true
> },
> "allgather\_partitions": true,
> "allgather\_bucket\_size": 2e8,
> "overlap\_comm": true,
> "reduce\_scatter": true,
> "reduce\_bucket\_size": 2e8,
> "contiguous\_gradients": true
> },
> 
> "gradient\_accumulation\_steps": "auto",
> "gradient\_clipping": "auto",
> "train\_batch\_size": "auto",
> "train\_micro\_batch\_size\_per\_gpu": "auto"
> }
> \`\`\` 
> 
> So I'm using zero2 with optimizer offload. I found the size of the model checkpoints after \`trainer.train()\` become much larger than what they should be. 
> 
> Using official \[run\_clm.py\](https://github.com/huggingface/transformers/blob/main/examples/pytorch/language-modeling/run\_clm.py) script as an example :
> 
> \`\`\`bash
> deepspeed --num\_gpus=1 run\_clm.py \\
> --num\_train\_epochs 0.01 \\
> --model\_name\_or\_path decapoda-research/llama-7b-hf \\
> --dataset\_name wikitext \\
> --dataset\_config\_name wikitext-2-raw-v1 \\
> --per\_device\_train\_batch\_size 2 \\
> --do\_train \\
> --output\_dir /tmp/test-plm \\
> --deepspeed ds\_config.json
> \`\`\`
> 
> I add these two save\_model lines around \`trainer.train()\` for testing:
> \`\`\`python
> trainer.save\_model("test1")
> train\_result = trainer.train(resume\_from\_checkpoint=checkpoint)
> trainer.save\_model("test2")
> \`\`\`
> 
> Now check the size:
> \`\`\`bash
> du -sh test1
> 26G test1
> 
> du -sh test2
> 76G test2
> \`\`\`
> Note, I have deleted \`global\_step\*\` folder in \`test2\` before calculating the size.
> 
> I believe 26G is the correct size for an fp32 llama 7b. So, after training with trainer, the model size is wrong? Interestingly, seems the wrong size model still works with \`.from\_pretrain\`. 
> 
> 
> I have located the issue raised after this \[line\](https://github.com/huggingface/transformers/blob/dacd34568d1a27b91f84610eab526640ed8f94e0/src/transformers/deepspeed.py#L378), which changed the model assignment in trainer \`\_inner\_training\_loop\` \[here\](https://github.com/huggingface/transformers/blob/dacd34568d1a27b91f84610eab526640ed8f94e0/src/transformers/trainer.py#L1733) afterward. After this the model saved by \`trainer.\_save()\` will have the wrong size.
> 
> Does deepspeed engine add some extra things to pytorch\_model.bin? is this expected?
> 
> My current solution to this is always using \`self.deepspeed.save\_16bit\_model()\` in \[trainer.save\_model()\](https://github.com/huggingface/transformers/blob/dacd34568d1a27b91f84610eab526640ed8f94e0/src/transformers/trainer.py#L2771) for zerostage2:
> \`\`\`python
> elif self.deepspeed:
> # this takes care of everything as long as we aren't under zero3
> if self.args.should\_save:
> self.\_save(output\_dir)
> if is\_deepspeed\_zero3\_enabled():
> # It's too complicated to try to override different places where the weights dump gets
> # saved, so since under zero3 the file is bogus, simply delete it. The user should
> # either user deepspeed checkpoint to resume or to recover full weights use
> # zero\_to\_fp32.py stored in the checkpoint.
> if self.args.should\_save:
> file = os.path.join(output\_dir, WEIGHTS\_NAME)
> if os.path.isfile(file):
> # logger.info(f"deepspeed zero3: removing {file}, see zero\_to\_fp32.py to recover weights")
> os.remove(file)
> 
> # now save the real model if stage3\_gather\_16bit\_weights\_on\_model\_save=True
> # if false it will not be saved.
> # This must be called on all ranks
> if not self.deepspeed.save\_16bit\_model(output\_dir, WEIGHTS\_NAME):
> logger.warning(
> "deepspeed.save\_16bit\_model didn't save the model, since"
> " stage3\_gather\_16bit\_weights\_on\_model\_save=false. Saving the full checkpoint instead, use"
> " zero\_to\_fp32.py to recover weights"
> )
> self.deepspeed.save\_checkpoint(output\_dir)
> else:
> if self.args.should\_save:
> for filename in os.listdir(output\_dir):
> full\_filename = os.path.join(output\_dir, filename)
> # If we have a shard file that is not going to be replaced, we delete it, but only from the main process
> # in distributed settings to avoid race conditions.
> weights\_no\_suffix = WEIGHTS\_NAME.replace(".bin", "").replace(".safetensors", "")
> 
> # delete everything start with weights\_no\_suffix, usually are "pytorch\_model".
> if (
> filename.startswith(weights\_no\_suffix)
> and os.path.isfile(full\_filename)
> ):
> os.remove(full\_filename)
> self.deepspeed.save\_16bit\_model(output\_dir, WEIGHTS\_NAME)
> \`\`\`
> 
> 
> 
> 
> \### Expected behavior
> 
> Model checkpoint size should be unchanged after \`trainer.train()\`

---

<div class="post-metadata">

**Author:** ![esualp](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/esualp/32/39451_2.png) [@esualp](https://discuss.huggingface.co/u/esualp)\
**Post date:** [February 19, 2025, 9:36pm UTC](https://discuss.huggingface.co/t/avoid-saving-deepspeed-optimizer-and-model-states-at-checkpoints/141178/3 "2025-02-19T21:36:02Z")

</div>

Thanks for the reply! I would like to understand what “\_model\_states.pt” for? I am removing optim\_states (they’re really huge) with this callback, if anyone interested:

```auto
class CleanupCallback(TrainerCallback):
    def on_save(self, args, state, control, **kwargs):
        """
        Deletes optimizer checkpoint files from DeepSpeed checkpoints 
        inside `checkpoint-*/global_step*` folders.
        """
        # Get all DeepSpeed checkpoint folders
        checkpoint_dirs = sorted(glob.glob(os.path.join(args.output_dir, "checkpoint-*/global_step*")), 
                                 key=os.path.getctime, reverse=True)

        if not checkpoint_dirs:
            print("No DeepSpeed checkpoints found for cleanup.")
            return control

        latest_checkpoint = checkpoint_dirs[0] # Get the latest saved checkpoint

        # Build file paths safely
        optimizer_files = [
            os.path.join(latest_checkpoint, "optimizer.pt"),
            *glob.glob(os.path.join(latest_checkpoint, "*_optim_states.pt")),
        ]

        deleted_count = 0
        for file in optimizer_files:
            try:
                os.remove(file)
                print(f"Deleted: {file}")
                deleted_count += 1
            except Exception as e:
                print(f"Error deleting {file}: {e}")

        # Print summary
        if deleted_count > 0:
            print(f"Deleted {deleted_count} optimizer-related files from: {latest_checkpoint}")

        return control

```

I am wondering whether I need to keep model\_states.pt for inference