search

huggingface / transformers-bloom-inference

Fast Inference Solutions for BLOOM
Apache License 2.0
561 stars 114 forks source link

root_dir in TemporaryCheckpointsJSON is redundant #82

Open dc3671 opened 1 year ago

dc3671 commented 1 year ago

In TemporaryCheckpointsJSON(https://github.com/huggingface/transformers-bloom-inference/blob/main/inference_server/models/ds_inference.py#L80) , image

When use glob.glob(f"{self.model_path}/*.bin"), files path in the list will all contain model_path prefix (eg: modelname is bigscience/bloom ). 

{"type": "BLOOM", "checkpoints": ["bigscience/bloom/pytorch_model.bin"], "version": 1.0}

While set it as root_dir (glob.glob("*.bin", root_dir=self.model_path)) will not:

{"type": "BLOOM", "checkpoints": ["pytorch_model.bin"], "version": 1.0}

And it will align to DeepSpeed's loading way (replace_module.py). Because when loading, it will add root_dir again:

sd = [
  torch.load(os.path.join(base_dir1, checkpoint[i]), map_location='cpu')
]

So, current dump way will duplicate model_path.

And I raised a PR: https://github.com/huggingface/transformers-bloom-inference/pull/71 .