# How do I load a trained checkpoint model?

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-do-i-load-a-trained-checkpoint-model%2F155932%3E%5C
**Category:** 🤗Transformers\
**Created:** [May 19, 2025, 5:58pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-do-i-load-a-trained-checkpoint-model%2F155932 "2025-05-19T17:58:38Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![Lamsheeper](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lamsheeper/32/47867_2.png) [@Lamsheeper](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FLamsheeper)%5C
**Post date:** [May 19, 2025, 5:58pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-do-i-load-a-trained-checkpoint-model%2F155932%2F1 "2025-05-19T17:58:38Z")

</div>

I’ve trained an OLMo model from a revision with the state dict sharded across multiple .distcp files and the model weights in rank.pt files. I’m trying to run inference with this checkpoint. How can I load this model into a .from\_pretrained format?

I’m trying to use this code but getting a lot of key mismatch/missing key errors:

```auto
import fire

import torch.distributed.checkpoint as dist_cp

from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoModel

def load_sharded_model_single_gpu(model, model_path):
    
    state_dict = {
        "model": model.state_dict()
    }
    
    dist_cp.load_state_dict(
                state_dict=state_dict,
                storage_reader=dist_cp.FileSystemReader(model_path),
                no_dist=True,
            )
    
    result = model.load_state_dict(state_dict["model"])
    
    print(f"Sharded state checkpoint loaded from {model_path}")
    print(result)
    return model

def convert_checkpoint(hf_model: str, fsdp_model_path: str, output_path: str):
    '''
    hf_model: transformers path.
    fsdp_model_path: path to the fsdp checkpoint, for example `/x/checkpoint-xxx/pytorch_model_x`
    output_path: output path to save the converted checkpoint
    '''
    config = AutoConfig.from_pretrained(hf_model, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(hf_model, trust_remote_code=True)
    model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
    model = load_sharded_model_single_gpu(model, fsdp_model_path)
    model.save_pretrained(output_path, max_shard_size="10GB")
    tokenizer.save_pretrained(output_path)

if __name__ == " __main__":
    fire.Fire(convert_checkpoint)

```

---

<div class="post-metadata">

**Author:** ![John6666](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/john6666/32/27664_2.png) [@John6666](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FJohn6666)%5C
**Post date:** [May 20, 2025, 5:55am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-do-i-load-a-trained-checkpoint-model%2F155932%2F2 "2025-05-20T05:55:10Z")

</div>

I think that method is correct, but there seem to be reports of tensor mismatch issues when training with FSDP.

> [@Key errors when trying to load an accelerate-FSDP model checkpoint](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fkey-errors-when-trying-to-load-an-accelerate-fsdp-model-checkpoint%2F85829)%3A
>
> Hi, I have been training a pythia-2.8b model, using the DPOTrainer in Huggingface TRL. I am training a LORA adaptor on top of the reference SFT model. In order to speed it up and address memory bottlenecks, I used fsdp and launched my code using accelerate launch --config\_file=fsdp.yaml --num\_processes=4 run\_dpo.py After training, I observed my checkpoints to be of the following form optimizer\_0/ pytorch\_model\_fsdp\_0/ rng\_state\_0.pth rng\_state\_1.pth scheduler.pt trainer\_…

> [@How to load a checkpoint model with SHARDED\_STATE\_DICT?](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fhow-to-load-a-checkpoint-model-with-sharded-state-dict%2F62448)%3A
>
> How to load a checkpoint model with SHARDED\_STATE\_DICT? I have a checkpoint which is place in a folder pytorch\_model\_0, which contains multiple distcp files.

> **[Serialization](https://huggingface.co/docs/huggingface_hub/main/en/package_reference/serialization)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.