Skip to content

allow saving ckpt w/o ckpt json + bloom copy fix #2237

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jeffra merged 4 commits into from
Aug 19, 2022
Merged

allow saving ckpt w/o ckpt json + bloom copy fix #2237

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cf0acbb
allow saving ckpt w/o ckpt json + bloom copy fix
jeffra Aug 18, 2022
4bb6558
save name as bloom if model is a bloom model
jeffra Aug 18, 2022
6294eab
formatting
jeffra Aug 19, 2022
abae054
guard hf import
jeffra Aug 19, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 23 additions & 5 deletions deepspeed/module_inject/replace_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -478,8 +478,16 @@ def _transpose(x):
attn_block.attn_ow = mp_replace.copy(attn_block.attn_ow, dense_w)
attn_block.attn_ob = mp_replace.copy(attn_block.attn_ob, dense_b)
else:
attn_block.attn_qkvw = mp_replace.qkv_copy(attn_block.attn_qkvw, qkvw)
attn_block.attn_qkvb = mp_replace.qkv_copy(attn_block.attn_qkvb, qkvb)
if bigscience_bloom:
attn_block.attn_qkvw = mp_replace.copy(attn_block.attn_qkvw, qkvw)
attn_block.attn_qkvb = mp_replace.copy(attn_block.attn_qkvb, qkvb)
else:
attn_block.attn_qkvw = mp_replace.qkv_copy(
attn_block.attn_qkvw,
qkvw)
attn_block.attn_qkvb = mp_replace.qkv_copy(
attn_block.attn_qkvb,
qkvb)

attn_block.attn_ow = mp_replace.copy(attn_block.attn_ow, dense_w)
attn_block.attn_ob = mp_replace.copy(attn_block.attn_ob, dense_b)
Expand Down Expand Up @@ -773,10 +781,10 @@ def replace_fn(child, _policy, layer_id=0):
replace_fn=replace_fn,
_replace_policy=policy)

world_size = dist.get_world_size() if dist.is_initialized() else 1
rank = dist.get_rank() if dist.is_initialized() else 0
if checkpoint_dict is not None:
start_time = time.time()
rank = dist.get_rank() if dist.is_initialized() else 0
world_size = dist.get_world_size() if dist.is_initialized() else 1
checkpoint = checkpoint_dict['checkpoints']
ckpt_type = checkpoint_dict.get('parallelization', 'pp')
ckpt_mp_size = checkpoint_dict.get('mp_size', mp_size)
Expand Down Expand Up @@ -818,12 +826,22 @@ def replace_fn(child, _policy, layer_id=0):
from collections import OrderedDict
import json

ckpt_name = checkpoint_dict['type']
if checkpoint_dict is None:
ckpt_name = "ds_model"
try:
from transformers.models.bloom.modeling_bloom import BloomForCausalLM
if isinstance(model, BloomForCausalLM):
ckpt_name = "bloom"
except ImportError:
ckpt_name = "ds_model"
else:
ckpt_name = checkpoint_dict['type']
if dist.is_initialized():
dist.barrier()
transformer_name = get_transformer_name(replaced_module)
non_tp_ckpt_name = f'{ckpt_name}-non-tp.pt'
ckpt_files = [non_tp_ckpt_name] * world_size
os.makedirs(save_mp_checkpoint_path, exist_ok=True)
if not dist.is_initialized() or dist.get_rank() == 0:
print("Saving tp-sharded checkpoints")
torch.save(
Expand Down