How to load a checkpoint model with SHARDED_STATE_DICT?
How to load a checkpoint model with SHARDED_STATE_DICT?I have a checkpoint which is place in a folderpytorch_model_0, which contains multipledistcpfiles.
2024-01-08T14:44:09Z
Ihave the same questiondid you find out something ?
2024-01-10T11:05:38Z
I have found the solution for llama or you can edit for other modelsyou can use this :llama-recipes/docs/inference.md at main · facebookresearch/llama-recipes · GitHubi have edited some files inside it to load mistral instead of llama
2024-01-11T08:56:24Z
Yep. I used the same solution with you. I tried to find this code the day you ask me but I can not remember where it is. So glad you find it yourself.
2024-01-11T08:59:04Z
I will post my code here: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(REDACTED_TOKEN: str, fsdp_model_path: str, output_path: str):
    '''
    REDACTED_TOKEN: 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(REDACTED_TOKEN, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained(REDACTED_TOKEN, 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)
2024-01-11T20:59:26Z
This topic was automatically closed 12 hours after the last reply. New replies are no longer allowed.