[BUG]: GeminiStrategy save_pretrained CAN NOT be loaded with map_location

[cwher]

🐛 Describe the bug

These scripts can be used to reproduce this error. Just run run.sh.

  File "load.py", line 6, in <module>
    state_dict = torch.load("bloom/pytorch_model.bin", map_location=torch.device('cpu'))
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_tensor.py", line 60, in _rebuild_from_type_v2
    ret = func(*args)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_utils.py", line 153, in _rebuild_tensor_v2
    tensor = _rebuild_tensor(storage, storage_offset, size, stride)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_utils.py", line 147, in _rebuild_tensor
    return t.set_(storage.untyped(), storage_offset, size, stride)
RuntimeError: Trying to resize storage that is not resizable

Environment

No response

[cwher]

@ver217 @Fridge003

[Issues-translate-bot]

Bot detected the issue body's language is not English, translate it automatically. 👯👭🏻🧑‍🤝‍🧑👫🧑🏿‍🤝‍🧑🏻👩🏾‍🤝‍👨🏿👬🏿

---

@ver217 @Fridge003

[ver217]

Can you make the scipt simpler and only remain save/load logic?

[cwher]

The following are much simpler.

save.py

import argparse

from coati.models import convert_to_lora_module
from coati.trainer.strategies import GeminiStrategy
from transformers import BloomForCausalLM, BloomTokenizerFast


def train(args):
    # configure strategy
    strategy = GeminiStrategy(placement_policy='cuda')

    # configure model
    with strategy.model_init_context():
        model = convert_to_lora_module(
            BloomForCausalLM.from_pretrained(args.pretrain),
            0
        ).half().cuda()
    tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')

    model = strategy.prepare(model)

    # save model checkpoint after fitting on only rank0
    strategy.save_pretrained(model, path=args.save_path, only_rank0=True, tokenizer=tokenizer)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--pretrain', type=str, default=None)
    parser.add_argument('--save_path', type=str, default='output')
    args = parser.parse_args()
    train(args)

load.py

import torch

if __name__ == "__main__":
    state_dict = torch.load("bloom/pytorch_model.bin")
    print("[INFO]: torch.load w/o map_location")
    state_dict = torch.load("bloom/pytorch_model.bin", map_location=torch.device('cpu'))
    print("[INFO]: torch.load w map_location")
    # RuntimeError: Trying to resize storage that is not resizable

run.sh

set -xue

torchrun --standalone --nproc_per_node=1 save.py \
    --pretrain bigscience/bloom-560m --save_path bloom
python load.py

run.sh yields on my environment,

Traceback (most recent call last):
  File "load.py", line 6, in <module>
    state_dict = torch.load("bloom/pytorch_model.bin", map_location=torch.device('cpu'))
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/serialization.py", line 789, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/serialization.py", line 1131, in _load
    result = unpickler.load()
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_tensor.py", line 60, in _rebuild_from_type_v2
    ret = func(*args)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_utils.py", line 153, in _rebuild_tensor_v2
    tensor = _rebuild_tensor(storage, storage_offset, size, stride)
  File "/home/lccwh/.conda/envs/cls-ai/lib/python3.8/site-packages/torch/_utils.py", line 147, in _rebuild_tensor
    return t.set_(storage.untyped(), storage_offset, size, stride)
RuntimeError: Trying to resize storage that is not resizable

[Fridge003]

Please use the checkpointing methods we customized for GeminiPlugin. Original huggingface from_pretrained method isn't compatible with GeminiStrategy.